Star1603
3278
代码
代码
Issues451
Pull Requests921
Actions
项目讨论
Wiki
分析
Star1603
3278
已合并
增加 MsaIndexScore 算子 950 实现 #10672
oscillated创建于 8月26日
增加 MsaIndexScore 算子 950 实现 #10672
已合并
oscillated创建于 8月26日
从oscillated/ops-transformer-ar-950msa-index-score-a5合入到cann/ops-transformermaster
共 34 个文件变更+3839-288
Mattention/msa_index_score/README.md+173-30
@@ -4,9 +4,9 @@
4 4 
5| 产品 | 是否支持 |5| 产品 | 是否支持 |
6| --------------------------------------------------------- | :------: |6| --------------------------------------------------------- | :------: |
7-| <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品</term> | √ |7+| <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term> | √ |
8| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |8| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |
9-| <term>Ascend 950PR/Ascend 950DT</term> | × |9+| <term>Ascend 950PR/Ascend 950DT</term> | √ |
10 10 
11## 功能说明11## 功能说明
12 12 
@@ -45,53 +45,190 @@
45> - D（Head Dim）表示单个注意力头维度45> - D（Head Dim）表示单个注意力头维度
46> - PageAttention 场景下 $block\_num$ 为物理 block 总数、$block\_size$ 为每个 block 的 token 数，$maxBlockNumPerSeq$ 为每个 batch 最大逻辑 block 数（通常 $\ge\lceil S2/block\_size\rceil$），$M_b=\lceil S2/block\_size\rceil$ 为逻辑 block 总数46> - PageAttention 场景下 $block\_num$ 为物理 block 总数、$block\_size$ 为每个 block 的 token 数，$maxBlockNumPerSeq$ 为每个 batch 最大逻辑 block 数（通常 $\ge\lceil S2/block\_size\rceil$），$M_b=\lceil S2/block\_size\rceil$ 为逻辑 block 总数
47 47 
48-| 参数名 | 输入/输出/属性 | 描述 | 数据类型 | 数据格式 |48+<table style="undefined;table-layout: fixed; width: 1280px"><colgroup>
49-| ------------------------ | -------------- | ------------------------------------------------------------ | ---------------------------- | -------- |49+<col style="width: 160px">
50-| query | 输入 | 公式中的 $Q_{idx}$。当前仅支持 TND，shape 为 $[T1, N1, D]$ | BFLOAT16, FLOAT16 | ND |50+<col style="width: 120px">
51-| key | 输入 | 公式中的 $K_{idx}$。支持 TND（$[T2, N2, D]$）、BNBD（$[block\_num, N2, block\_size, D]$）、BBND（$[block\_num, block\_size, N2, D]$） | BFLOAT16, FLOAT16, INT8 | ND, NZ |51+<col style="width: 560px">
52-| block_table | 可选输入 | PageAttention 的逻辑 block → 物理 page 映射表。PA 场景必须传入，二维，第二维长度不能小于 $maxBlockNumPerSeq$；shape 为 $[B, S2/block\_size]$ | INT32 | ND |52+<col style="width: 320px">
53-| scale | 可选输入 | 公式中的 $scale$，反量化系数。非量化必须为空；量化场景必选。PA 为 $[block\_num, N2, block\_size]$ 或 $[block\_num, block\_size, N2]$；TND 为 $[T2, N2]$ | FLOAT | ND, NZ |53+<col style="width: 120px">
54-| atten_mask | 可选输入 | 控制因果可见的 mask。仅在 `sparse_mode=3` 时使用；取值为 1 表示该位不参与计算，为 0 表示参与计算；shape 为 $[2048, 2048]$ | INT8 | ND |54+</colgroup>
55-| actual_seq_qlen | 可选输入 | 每个 Batch 中 Query 的有效 token 数。query 为 TND 时必须传入，单调不减（前缀和），shape 为 $[B+1]$ | INT32 | ND |55+<thead>
56-| actual_seq_klen | 可选输入 | 每个 Batch 中 Key 的有效 token 数。key 为 TND 时必须传入（前缀和）；PageAttention 场景下为各请求可见 $S2$，shape 为 $[B]$ | INT32 | ND |56+ <tr>
57-| start_loc | 输入 | 当前 query 所在逻辑 block 索引（非 token 前缀），用于生成 $local\_mask$；shape 为 $[B]$ | INT32 | ND |57+ <th>参数名</th>
58-| layout_key | 属性 | key 布局。`"TND"` / `"BBND"` / `"BNBD"`。aclnn 参数名为 `layoutKeyOptional`，不传时默认 `"BBND"` | STRING | - |58+ <th>输入/输出/属性</th>
59-| sparse_mode | 属性 | sparse 模式。0：defaultMask（`atten_mask` 传空）；3：rightDownCausal（须传入 $[2048, 2048]$ 的 `atten_mask`） | INT64 | - |59+ <th>描述</th>
60-| init_blocks | 属性 | $local\_mask$ 强制选中的头部 block 数。对逻辑 block $[0, init\_blocks)$ 写入高分 $1\mathrm{e}30$。可选，默认 $0$ | INT64 | - |60+ <th>数据类型</th>
61-| local_blocks | 属性 | $local\_mask$ 强制选中的局部窗口长度。窗口为 $[max(0, start\_loc+1-local\_blocks), start\_loc]$，写入高分 $1\mathrm{e}29$（覆盖同位置的 `init_blocks`）。可选，默认 $1$（对齐 MiniMax HF）；与 Triton raw score 对齐时置 $0$ | INT64 | - |61+ <th>数据格式</th>
62-| score | 输出 | 公式中的 $score$，逐 block 重要性分数；shape 为 $[N1, T1, RoundUp(maxBlockNumPerSeq, 16)]$ | FLOAT | ND |62+ </tr>
63+</thead>
64+<tbody>
65+ <tr>
66+ <td>query</td>
67+ <td>输入</td>
68+ <td>公式中的 $Q_{idx}$。当前仅支持 TND，shape 为 $[T1, N1, D]$。</td>
69+ <td>BFLOAT16、FLOAT16、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN</td>
70+ <td>ND</td>
71+ </tr>
72+ <tr>
73+ <td>key</td>
74+ <td>输入</td>
75+ <td>公式中的 $K_{idx}$。支持 TND（$[T2, N2, D]$）、BNBD（$[block\_num, N2, block\_size, D]$）、BBND（$[block\_num, block\_size, N2, D]$）。Ascend 950 上 PA（BBND/BNBD）允许首轴（物理 page）按 stride 非连续存放，page 内其余轴须连续。</td>
76+ <td>BFLOAT16、FLOAT16、INT8、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN</td>
77+ <td>ND</td>
78+ </tr>
79+ <tr>
80+ <td>block_table</td>
81+ <td>可选输入</td>
82+ <td>PageAttention 的逻辑 block → 物理 page 映射表。PA 场景必须传入，二维，第二维长度不能小于 $maxBlockNumPerSeq$；shape 为 $[B, S2/block\_size]$。</td>
83+ <td>INT32</td>
84+ <td>ND</td>
85+ </tr>
86+ <tr>
87+ <td>scale</td>
88+ <td>可选输入</td>
89+ <td>公式中的 $scale$，反量化系数。非量化必须为空；量化场景必选。PA 为 $[block\_num, N2, block\_size]$ 或 $[block\_num, block\_size, N2]$；TND 为 $[T2, N2]$。</td>
90+ <td>FLOAT</td>
91+ <td>ND</td>
92+ </tr>
93+ <tr>
94+ <td>atten_mask</td>
95+ <td>可选输入</td>
96+ <td>控制因果可见的 mask。仅在 <code>sparse_mode=3</code> 时使用；取值为 1 表示该位不参与计算，为 0 表示参与计算；shape 为 $[2048, 2048]$。</td>
97+ <td>INT8</td>
98+ <td>ND</td>
99+ </tr>
100+ <tr>
101+ <td>actual_seq_qlen</td>
102+ <td>可选输入</td>
103+ <td>每个 Batch 中 Query 的有效 token 数。query 为 TND 时必须传入，单调不减（前缀和），shape 为 $[B+1]$。</td>
104+ <td>INT32</td>
105+ <td>ND</td>
106+ </tr>
107+ <tr>
108+ <td>actual_seq_klen</td>
109+ <td>可选输入</td>
110+ <td>每个 Batch 中 Key 的有效 token 数。key 为 TND 时必须传入（前缀和）；PageAttention 场景下为各请求可见 $S2$，shape 为 $[B]$。</td>
111+ <td>INT32</td>
112+ <td>ND</td>
113+ </tr>
114+ <tr>
115+ <td>start_loc</td>
116+ <td>输入</td>
117+ <td>当前 query 所在逻辑 block 索引（非 token 前缀），用于生成 $local\_mask$；shape 为 $[B]$。</td>
118+ <td>INT32</td>
119+ <td>ND</td>
120+ </tr>
121+ <tr>
122+ <td>layout_key</td>
123+ <td>属性</td>
124+ <td>key 布局。<code>"TND"</code> / <code>"BBND"</code> / <code>"BNBD"</code>。aclnn 参数名为 <code>layoutKeyOptional</code>，不传时默认 <code>"BBND"</code>。</td>
125+ <td>STRING</td>
126+ <td>-</td>
127+ </tr>
128+ <tr>
129+ <td>sparse_mode</td>
130+ <td>属性</td>
131+ <td>sparse 模式。0：defaultMask（<code>atten_mask</code> 传空）；3：rightDownCausal（须传入 $[2048, 2048]$ 的 <code>atten_mask</code>）。</td>
132+ <td>INT64</td>
133+ <td>-</td>
134+ </tr>
135+ <tr>
136+ <td>init_blocks</td>
137+ <td>属性</td>
138+ <td>$local\_mask$ 强制选中的头部 block 数。对逻辑 block $[0, init\_blocks)$ 写入高分 $1\mathrm{e}30$。可选，默认 $0$。</td>
139+ <td>INT64</td>
140+ <td>-</td>
141+ </tr>
142+ <tr>
143+ <td>local_blocks</td>
144+ <td>属性</td>
145+ <td>$local\_mask$ 强制选中的局部窗口长度。窗口为 $[max(0, start\_loc+1-local\_blocks), start\_loc]$，写入高分 $1\mathrm{e}29$（覆盖同位置的 <code>init_blocks</code>）。可选，默认 $1$（对齐 MiniMax HF）；与 Triton raw score 对齐时置 $0$。</td>
146+ <td>INT64</td>
147+ <td>-</td>
148+ </tr>
149+ <tr>
150+ <td>score</td>
151+ <td>输出</td>
152+ <td>公式中的 $score$，逐 block 重要性分数；shape 为 $[N1, T1, RoundUp(maxBlockNumPerSeq, 16)]$</td>
153+ <td>FLOAT</td>
154+ <td>ND</td>
155+ </tr>
156+</tbody>
157+</table>
63 158 
64## 约束说明159## 约束说明
65 160 
66- 当前 $block\_size$ 仅支持 128。161- 当前 $block\_size$ 仅支持 128。
67- `layout_key` 必须显式指定：`"BBND"` / `"BNBD"` / `"TND"`，与 `key` 实际 shape 一致。162- `layout_key` 必须显式指定：`"BBND"` / `"BNBD"` / `"TND"`，与 `key` 实际 shape 一致。
68- PageAttention（`layout_key` 为 `"BBND"` / `"BNBD"`）场景下，`block_table` 必须传入；TND key 场景不得传入 `block_table`，`actual_seq_klen` 为 `[B+1]` 前缀和。163- PageAttention（`layout_key` 为 `"BBND"` / `"BNBD"`）场景下，`block_table` 必须传入；TND key 场景不得传入 `block_table`，`actual_seq_klen` 为 `[B+1]` 前缀和。
69-- 非量化场景下，`key` dtype 与 `query` 相同（当前为 BFLOAT16 / FLOAT16），`scale` 必须为空；量化场景下仅支持 INT8，`scale` 必选：PA 为 $[block\_num, N2, block\_size]$ 或 $[block\_num, block\_size, N2]$，TND 为 $[T2, N2]$，dtype 为 FLOAT。当前不支持 FP8 与 <term>Ascend 950PR/Ascend 950DT</term>。164+- 非量化场景下，`key` dtype 与 `query` 相同（BFLOAT16 / FLOAT16；Ascend 950 另支持 HIFLOAT8 / FLOAT8_E5M2 / FLOAT8_E4M3FN），`scale` 必须为空。量化场景仅支持 INT8（fp16 query），`scale` 必选：PA 为 $[block\_num, N2, block\_size]$ 或 $[block\_num, block\_size, N2]$，TND 为 $[T2, N2]$，dtype 为 FLOAT。三种 FP8 仅 950，且 query 与 key 必须同型、不得传 `scale`。
70- `sparse_mode` 当前仅支持 0、3：165- `sparse_mode` 当前仅支持 0、3：
71 - 为 0 时，代表 defaultMask 模式，`atten_mask` 传入空；166 - 为 0 时，代表 defaultMask 模式，`atten_mask` 传入空；
72 - 为 3 时，代表 rightDownCausal 模式，`atten_mask` 必须传入，shape 为 $[2048, 2048]$，取值为 1 代表该位不参与计算，为 0 代表该位参与计算。167 - 为 3 时，代表 rightDownCausal 模式，`atten_mask` 必须传入，shape 为 $[2048, 2048]$，取值为 1 代表该位不参与计算，为 0 代表该位参与计算。
73- `init_blocks`、`local_blocks` 必须 $\ge 0$ 且不超过逻辑 block 数（PA 为 `block_table` 第二维；TND 为 score 末维对齐宽度）。两者均为 0 时跳过 $local\_mask$。168- `init_blocks`、`local_blocks` 必须 $\ge 0$ 且不超过逻辑 block 数（PA 为 `block_table` 第二维；TND 为 score 末维对齐宽度）。两者均为 0 时跳过 $local\_mask$。
169+- PageAttention `block_table` 第二维可以大于实际 KV 逻辑 block 数；score 末维为 $\mathrm{RoundUp}(width, 16)$。Ascend 950 C2UB 对超过 256 列按 256 列滑窗 flush。
170+- A2/A3 与 Ascend 950：`q_len` / `kv_len` 允许为 0（含整 batch）。对应请求跳过 QK；空 KV 的 score 填 `-inf`；整 batch `T1=0` 时 `SetBlockDim(1)`。
171+- A2/A3 与 Ascend 950：按估计 M-task 数启动 MIX（`B=1` 时为 $\mathrm{CeilDiv}(T1\cdot N1, 128)$；多 batch 用 packed+$B$ 上界），短 decode 不打满 AIC。
74- 本算子输出止于 block score，**不包含** TopK。172- 本算子输出止于 block score，**不包含** TopK。
75 173 
76## 调用示例174## 调用示例
77 175 
78-| 调用方式 | 样例代码 | 说明 |176+<table style="undefined;table-layout: fixed; width: 1100px"><colgroup>
79-|----------|----------|------|177+<col style="width: 180px">
80-| aclnn 单算子调用 | [test_aclnn_msa_index_score.cpp](./examples/test_aclnn_msa_index_score.cpp) | 内置 CPU golden 的端到端精度自验证 |178+<col style="width: 400px">
81-| 接口文档 | [aclnnMsaIndexScore.md](./docs/aclnnMsaIndexScore.md) | 两段式接口说明 |179+<col style="width: 520px">
82-| 测试说明 | [tests/README.md](./tests/README.md) | 用例矩阵与运行方式 |180+</colgroup>
83-| PyTorch API | - | 通过[msa_index_score](../../attention/msa_index_score/docs/torchapi_msa_index_score.md)接口调用msa_index_score算子。 |181+<thead>
182+ <tr>
183+ <th>调用方式</th>
184+ <th>样例代码</th>
185+ <th>说明</th>
186+ </tr>
187+</thead>
188+<tbody>
189+ <tr>
190+ <td>aclnn 单算子调用</td>
191+ <td><a href="./examples/test_aclnn_msa_index_score.cpp">test_aclnn_msa_index_score.cpp</a></td>
192+ <td>内置 CPU golden 的端到端精度自验证</td>
193+ </tr>
194+ <tr>
195+ <td>接口文档</td>
196+ <td><a href="./docs/aclnnMsaIndexScore.md">aclnnMsaIndexScore.md</a></td>
197+ <td>两段式接口说明</td>
198+ </tr>
199+ <tr>
200+ <td>测试说明</td>
201+ <td><a href="./tests/README.md">tests/README.md</a></td>
202+ <td>用例矩阵与运行方式</td>
203+ </tr>
204+ <tr>
205+ <td>PyTorch API</td>
206+ <td>-</td>
207+ <td>通过<a href="./docs/torchapi_msa_index_score.md">msa_index_score</a>接口调用msa_index_score算子。</td>
208+ </tr>
209+</tbody>
210+</table>
84 211 
85-编译与运行：212+编译与运行（Ascend 950，x86 包名为 `linux-x86_64.run`；A2/A3 将 `--soc` 换成对应平台）：
86 213 
87```bash214```bash
88-bash build.sh --pkg --soc=ascend910b --ops=msa_index_score -j32215+cd /path/to/ops-transformer-ar-950
89-./build_out/cann-ops-transformer-custom_linux-aarch64.run --quiet --install-path=/tmp/msa_opp216+bash build.sh --pkg --soc=ascend950 --ops=msa_index_score -j32
90-export ASCEND_CUSTOM_OPP_PATH=/tmp/msa_opp/vendors/custom_transformer217+bash ./build_out/cann-ops-transformer-custom_linux-x86_64.run --quiet --install-path=/path/to/msa_opp/
91-bash build.sh --run_example msa_index_score eager cust --vendor_name=custom218+source /path/to/msa_opp/vendors/custom_transformer/bin/set_env.bash
219+export ASCEND_CUSTOM_OPP_PATH=/path/to/msa_opp/vendors/custom_transformer
220+ 
221+# aclnn example 必须带 --soc=ascend950（否则默认 910b）
222+bash build.sh --run_example msa_index_score eager cust --vendor_name=custom --soc=ascend950
223+# 通过：末行 [PASS]: 47/47 cases passed
224+# 矩阵：39 条 fp16/bf16/int8（BBND/BNBD/TND，含 pad / 空序列 / key dim0 stride / 宽 block_table / 短 decode）+ 8 条 FP8（D=128）
225+# 容差：fp16/bf16/int8 1e-3，FP8 2e-2
226+# A2/A3：同上矩阵去掉 FP8，期望 [PASS]: 39/39（skipped 8 FP8）
92```227```
93 228 
94-> **实现备注（A2/A3）**229+torch_extension 见 [msa_index_score.md](./docs/torchapi_msa_index_score.md)。`cann/set_env.sh` 会把 cann 自带 `site-packages` 插到 `PYTHONPATH` 前面，必须把本仓 `torch_extension` 再插回最前。
230+ 
231+> **实现备注（A2/A3 / Ascend 950）**
95>232>
96> - key 布局由属性 `layout_key`（aclnn：`layoutKeyOptional`）指定，支持 PageAttention **BBND** / **BNBD**，以及 packed **TND**（无 `block_table`，`actual_seq_klen` 为 `[B+1]` 前缀和）。默认 `"BBND"`。233> - key 布局由属性 `layout_key`（aclnn：`layoutKeyOptional`）指定，支持 PageAttention **BBND** / **BNBD**，以及 packed **TND**（无 `block_table`，`actual_seq_klen` 为 `[B+1]` 前缀和）。默认 `"BBND"`。
97> - `sparse_mode=3` 的 `atten_mask[2048,2048]` 在 host 校验必选；device 侧按 rightDownCausal234> - `sparse_mode=3` 的 `atten_mask[2048,2048]` 在 host 校验必选；device 侧按 rightDownCausal
@@ -99,3 +236,9 @@ bash build.sh --run_example msa_index_score eager cust --vendor_name=custom
99> - `start_loc` 为逻辑 block 索引，与属性 `init_blocks`（默认 0）、`local_blocks`（默认 1）236> - `start_loc` 为逻辑 block 索引，与属性 `init_blocks`（默认 0）、`local_blocks`（默认 1）
100> 一起在 Maxpool 之后施加 `local_mask`。237> 一起在 Maxpool 之后施加 `local_mask`。
101> - 完整公式：`score = Maxpool[(scale·)Q@Kᵀ + atten_mask] + local_mask`。238> - 完整公式：`score = Maxpool[(scale·)Q@Kᵀ + atten_mask] + local_mask`。
239+> - **950 当前交付**（`op_kernel/arch35/`）：计算骨架与 A2 相同（Q 驻留 × K pingpong × **8-page S GM** + MODE 0x2 握手）。Cube **原生** FP8（TilingKey 4/5/6，`hifloat8_t` / `fp8_e5m2_t` / `fp8_e4m3fn_t`，无 scale，禁止 Cast→fp16）。L0C→AIV UB（C_to_UB）尚未作为主路径。`arch22/` 冻结。
240+> - **PA key dim0 stride**（A2/A3 与 950）：key 为 `IgnoreContiguous`；tiling 用 `GetInputStride` / `GetRequiredInputStride` 读首轴元素 stride，写入已有 `strideKvBlock`。TND 不允许非连续。scale 仍按逻辑 page 紧凑布局。
241+> - **950 宽 `block_table`**：C2UB score 暂存单窗 256 列；末维 `RoundUp(width,16)>256` 时按窗 flush 并补写后续 `-inf`，不能关掉 stage。
242+> - **短 decode**（A2/A3 与 950）：host 按估计 M-task 数 `SetBlockDim`（单请求即 $\mathrm{CeilDiv}(T1\cdot N1,128)$），不再默认打满 AIC；多 M-tile / 大 batch 仍可占满 AIC。
243+> - torch 的 `torch_npu.hifloat8.npu()` 当前会打出非法 device id，脚本跳过；kernel / aclnn 已挂 HIFLOAT8。
244+> - 测性能须选 Health=OK 且空闲的卡。
Mattention/msa_index_score/docs/aclnnMsaIndexScore.md+274-42
@@ -4,11 +4,15 @@
4 4 
5## 产品支持情况5## 产品支持情况
6 6 
7-| 产品 | 是否支持 |7+<!-- npu="910b" id1 -->
8-| --------------------------------------------------------- | :------: |8+- <term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>：支持
9-| <term>Atlas A2 训练系列产品/Atlas 800I A2 推理产品</term> | √ |9+<!-- end id1 -->
10-| <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term> | √ |10+<!-- npu="A3" id2 -->
11-| <term>Ascend 950PR/Ascend 950DT</term> | × |11+- <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：支持
12+<!-- end id2 -->
13+<!-- npu="950" id3 -->
14+- <term>Ascend 950PR/Ascend 950DT</term>：支持
15+<!-- end id3 -->
12 16 
13## 功能说明17## 功能说明
14 18 
@@ -75,62 +79,290 @@ aclnnStatus aclnnMsaIndexScore(
75 79 
76## aclnnMsaIndexScoreGetWorkspaceSize80## aclnnMsaIndexScoreGetWorkspaceSize
77 81 
78-### 参数说明82+- **参数说明：**
79 83 
80-| 参数名 | 输入/输出 | 描述 | 使用说明 | 数据类型 | 数据格式 | 维度 |84+ <table style="undefined;table-layout: fixed; width: 1600px"><colgroup>
81-| ----------------------------- | --------- | ------------------------------------------------- | ------------------------------------------------------------ | ---------------------------- | -------- | ------------------------------------------------------------ |85+ <col style="width: 200px">
82-| query | 输入 | 公式中的$Q_{idx}$ | 支持的shape为：TND | BFLOAT16, FLOAT16 | ND | 3（$[T1, N1, D]$） |86+ <col style="width: 100px">
83-| key | 输入 | 公式中的$K_{idx}$ | 支持的shape为：TND、BNBD、BBND | BFLOAT16, FLOAT16, INT8 | ND, NZ | 3（$[T2, N2, D]$）或 4（$[block\_num, N2, block\_size, D]$、$[block\_num, block\_size, N2, D]$） |87+ <col style="width: 280px">
84-| blockTableOptional | 输入 | 表示当前传入的使用PageAttention存储的block映射表 | PageAttention场景下，blockTableOptional需为2维；第二维长度不能小于 $maxBlockNumPerSeq$ | INT32 | ND | 2（$[B, S2/block\_size]$） |88+ <col style="width: 420px">
85-| scaleOptional | 输入 | 公式中的$scale$，反量化系数 | 非量化场景传入`nullptr`；量化场景时必选。PA 为 BNB/BBN；TND 为 $[T2, N2]$（N2=1 时可 $[T2]$） | FLOAT | ND, NZ | 3（$[block\_num, N2, block\_size]$、$[block\_num, block\_size, N2]$）或 2（$[T2, N2]$） |89+ <col style="width: 280px">
86-| attenMaskOptional | 输入 | 控制因果可见的mask掩码。 | 仅在sparseMode=3时使用，作为base mask控制causal可见；取值为1代表该位不参与计算，为0代表该位参与计算 | INT8 | ND | 2（$[2048, 2048]$） |90+ <col style="width: 100px">
87-| actualSeqQlenOptional | 输入 | 每个Batch中，Query的有效token数 | 当传入TND时，该入参必须传入，单调不减（前缀和） | INT32 | ND | 1（$[B+1]$） |91+ <col style="width: 220px">
88-| actualSeqKlenOptional | 输入 | 每个Batch中，Key的有效token数 | key 为 TND 时必须传入，单调不减（前缀和，$[B+1]$）；PageAttention 场景下为各请求可见 $S2$（$[B]$） | INT32 | ND | 1（$[B]$ 或 $[B+1]$） |92+ </colgroup>
89-| startLoc | 输入 | 当前 query 所在逻辑 block 索引（非 token 前缀） | 与 `initBlocks` / `localBlocks` 一起生成 $local\_mask$ | INT32 | ND | 1（$[B]$） |93+ <thead>
90-| layoutKeyOptional | 输入 | key 的数据排布 | 取值 `"TND"` / `"BBND"` / `"BNBD"`。不传或空串时默认 `"BBND"`。必须与 `key` 实际 shape 一致，不可仅凭维度推断 | CHAR* | - | - |94+ <tr>
91-| sparseMode | 输入 | 表示sparse的模式 | 为0时，代表defaultMask模式；为3时，代表rightDownCausal模式的mask，对应以右顶点为划分的下三角场景。 | INT64 | - | - |95+ <th>参数名</th>
92-| initBlocks | 输入 | $local\_mask$ 强制选中的头部 block 数 | 对逻辑 block $[0, initBlocks)$ 写入高分 $1\mathrm{e}30$。可选，默认 $0$；须 $\ge 0$ 且 $\le maxBlockNumPerSeq$ | INT64 | - | - |96+ <th>输入/输出</th>
93-| localBlocks | 输入 | $local\_mask$ 强制选中的局部窗口长度 | 窗口为 $[max(0, startLoc+1-localBlocks), startLoc]$，写入高分 $1\mathrm{e}29$（覆盖同位置 init）。可选，默认 $1$（对齐 MiniMax HF）；与 Triton raw score 对齐时置 $0$ | INT64 | - | - |97+ <th>描述</th>
94-| score | 输出 | 公式中的$score$ | 逐block的重要性分数；末维为对齐后的逻辑 block 数 | FLOAT | ND | 3（$[N1, T1, RoundUp(maxBlockNumPerSeq, 16)]$） |98+ <th>使用说明</th>
95-| workspaceSize | 输出 | 所需 workspace 字节数 | - | uint64_t | - | - |99+ <th>数据类型</th>
96-| executor | 输出 | 算子执行器 | - | aclOpExecutor** | - | - |100+ <th>数据格式</th>
101+ <th>维度(shape)</th>
102+ </tr>
103+ </thead>
104+ <tbody>
105+ <tr>
106+ <td>query</td>
107+ <td>输入</td>
108+ <td>公式中的$Q_{idx}$。</td>
109+ <td>支持的shape为：TND。</td>
110+ <td>BFLOAT16、FLOAT16、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN</td>
111+ <td>ND</td>
112+ <td>3（$[T1, N1, D]$）</td>
113+ </tr>
114+ <tr>
115+ <td>key</td>
116+ <td>输入</td>
117+ <td>公式中的$K_{idx}$。</td>
118+ <td>支持的shape为：TND、BNBD、BBND。A2/A3 与 Ascend 950 上 PA key 允许 dim0（物理 page）非连续，page 内其余轴须连续。TND 不允许非连续。</td>
119+ <td>BFLOAT16、FLOAT16、INT8、HIFLOAT8、FLOAT8_E5M2、FLOAT8_E4M3FN</td>
120+ <td>ND</td>
121+ <td>3（$[T2, N2, D]$）或 4（$[block\_num, N2, block\_size, D]$、$[block\_num, block\_size, N2, D]$）</td>
122+ </tr>
123+ <tr>
124+ <td>blockTableOptional</td>
125+ <td>输入</td>
126+ <td>表示当前传入的使用PageAttention存储的block映射表。</td>
127+ <td>PageAttention场景下，blockTableOptional需为2维；第二维长度不能小于 $maxBlockNumPerSeq$。</td>
128+ <td>INT32</td>
129+ <td>ND</td>
130+ <td>2（$[B, S2/block\_size]$）</td>
131+ </tr>
132+ <tr>
133+ <td>scaleOptional</td>
134+ <td>输入</td>
135+ <td>公式中的$scale$，反量化系数。</td>
136+ <td>非量化场景传入 <code>nullptr</code>；量化场景时必选。PA 为 BNB/BBN；TND 为 $[T2, N2]$（N2=1 时可 $[T2]$）。</td>
137+ <td>FLOAT</td>
138+ <td>ND</td>
139+ <td>3（$[block\_num, N2, block\_size]$、$[block\_num, block\_size, N2]$）或 2（$[T2, N2]$）</td>
140+ </tr>
141+ <tr>
142+ <td>attenMaskOptional</td>
143+ <td>输入</td>
144+ <td>控制因果可见的mask掩码。</td>
145+ <td>仅在sparseMode=3时使用，作为base mask控制causal可见；取值为1代表该位不参与计算，为0代表该位参与计算。</td>
146+ <td>INT8</td>
147+ <td>ND</td>
148+ <td>2（$[2048, 2048]$）</td>
149+ </tr>
150+ <tr>
151+ <td>actualSeqQlenOptional</td>
152+ <td>输入</td>
153+ <td>每个Batch中，Query的有效token数。</td>
154+ <td>当传入TND时，该入参必须传入，单调不减（前缀和）。</td>
155+ <td>INT32</td>
156+ <td>ND</td>
157+ <td>1（$[B+1]$）</td>
158+ </tr>
159+ <tr>
160+ <td>actualSeqKlenOptional</td>
161+ <td>输入</td>
162+ <td>每个Batch中，Key的有效token数。</td>
163+ <td>key 为 TND 时必须传入，单调不减（前缀和，$[B+1]$）；PageAttention 场景下为各请求可见 $S2$（$[B]$）。</td>
164+ <td>INT32</td>
165+ <td>ND</td>
166+ <td>1（$[B]$ 或 $[B+1]$）</td>
167+ </tr>
168+ <tr>
169+ <td>startLoc</td>
170+ <td>输入</td>
171+ <td>当前 query 所在逻辑 block 索引（非 token 前缀）。</td>
172+ <td>与 <code>initBlocks</code> / <code>localBlocks</code> 一起生成 $local\_mask$。</td>
173+ <td>INT32</td>
174+ <td>ND</td>
175+ <td>1（$[B]$）</td>
176+ </tr>
177+ <tr>
178+ <td>layoutKeyOptional</td>
179+ <td>输入</td>
180+ <td>key 的数据排布。</td>
181+ <td>取值 <code>"TND"</code> / <code>"BBND"</code> / <code>"BNBD"</code>。不传或空串时默认 <code>"BBND"</code>。必须与 <code>key</code> 实际 shape 一致，不可仅凭维度推断。</td>
182+ <td>CHAR*</td>
183+ <td>-</td>
184+ <td>-</td>
185+ </tr>
186+ <tr>
187+ <td>sparseMode</td>
188+ <td>输入</td>
189+ <td>表示sparse的模式。</td>
190+ <td>为0时，代表defaultMask模式；为3时，代表rightDownCausal模式的mask，对应以右顶点为划分的下三角场景。</td>
191+ <td>INT64</td>
192+ <td>-</td>
193+ <td>-</td>
194+ </tr>
195+ <tr>
196+ <td>initBlocks</td>
197+ <td>输入</td>
198+ <td>$local\_mask$ 强制选中的头部 block 数。</td>
199+ <td>对逻辑 block $[0, initBlocks)$ 写入高分 $1\mathrm{e}30$。可选，默认 $0$；须 $\ge 0$ 且 $\le maxBlockNumPerSeq$。</td>
200+ <td>INT64</td>
201+ <td>-</td>
202+ <td>-</td>
203+ </tr>
204+ <tr>
205+ <td>localBlocks</td>
206+ <td>输入</td>
207+ <td>$local\_mask$ 强制选中的局部窗口长度。</td>
208+ <td>窗口为 $[max(0, startLoc+1-localBlocks), startLoc]$，写入高分 $1\mathrm{e}29$（覆盖同位置 init）。可选，默认 $1$（对齐 MiniMax HF）；与 Triton raw score 对齐时置 $0$。</td>
209+ <td>INT64</td>
210+ <td>-</td>
211+ <td>-</td>
212+ </tr>
213+ <tr>
214+ <td>score</td>
215+ <td>输出</td>
216+ <td>公式中的$score$。</td>
217+ <td>逐block的重要性分数；末维为对齐后的逻辑 block 数。</td>
218+ <td>FLOAT</td>
219+ <td>ND</td>
220+ <td>3（$[N1, T1, RoundUp(maxBlockNumPerSeq, 16)]$）</td>
221+ </tr>
222+ <tr>
223+ <td>workspaceSize</td>
224+ <td>输出</td>
225+ <td>返回需要在Device侧申请的workspace大小。</td>
226+ <td>-</td>
227+ <td>-</td>
228+ <td>-</td>
229+ <td>-</td>
230+ </tr>
231+ <tr>
232+ <td>executor</td>
233+ <td>输出</td>
234+ <td>返回op执行器，包含了算子计算流程。</td>
235+ <td>-</td>
236+ <td>-</td>
237+ <td>-</td>
238+ <td>-</td>
239+ </tr>
240+ </tbody>
241+ </table>
97 242 
98-### 返回值243+- **返回值：**
99 244 
100-| 返回码 | 错误码 | 说明 |245+ aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn_return_code.md)。
101-| ----------------------- | ------ | ---------------------------------------- |246+ 
102-| ACLNN_SUCCESS | 0 | 执行成功 |247+ 第一段接口会完成入参校验，出现以下场景时报错：
103-| ACLNN_ERR_PARAM_NULLPTR | 161001 | 必选入参或出参为空指针 |248+ 
104-| ACLNN_ERR_PARAM_INVALID | 161002 | 数据类型、数据格式、维度或取值不满足约束 |249+ <table style="undefined;table-layout: fixed;width: 1155px"><colgroup>
250+ <col style="width: 319px">
251+ <col style="width: 144px">
252+ <col style="width: 671px">
253+ </colgroup>
254+ <thead>
255+ <tr>
256+ <th>返回码</th>
257+ <th>错误码</th>
258+ <th>描述</th>
259+ </tr>
260+ </thead>
261+ <tbody>
262+ <tr>
263+ <td>ACLNN_SUCCESS</td>
264+ <td>0</td>
265+ <td>执行成功。</td>
266+ </tr>
267+ <tr>
268+ <td>ACLNN_ERR_PARAM_NULLPTR</td>
269+ <td>161001</td>
270+ <td>必选入参或出参为空指针。</td>
271+ </tr>
272+ <tr>
273+ <td>ACLNN_ERR_PARAM_INVALID</td>
274+ <td>161002</td>
275+ <td>数据类型、数据格式、维度或取值不满足约束。</td>
276+ </tr>
277+ </tbody>
278+ </table>
105 279 
106## aclnnMsaIndexScore280## aclnnMsaIndexScore
107 281 
108-### 参数说明282+- **参数说明：**
109 283 
110-| 参数名 | 输入/输出 | 描述 |284+ <table style="undefined;table-layout: fixed; width: 953px"><colgroup>
111-| ------------- | --------- | ---------------------------------- |285+ <col style="width: 173px">
112-| workspace | 输入 | device 侧 workspace 地址 |286+ <col style="width: 112px">
113-| workspaceSize | 输入 | workspace 字节数，由第一段接口返回 |287+ <col style="width: 668px">
114-| executor | 输入 | 算子执行器，由第一段接口返回 |288+ </colgroup>
115-| stream | 输入 | acl stream |289+ <thead>
290+ <tr>
291+ <th>参数名</th>
292+ <th>输入/输出</th>
293+ <th>描述</th>
294+ </tr>
295+ </thead>
296+ <tbody>
297+ <tr>
298+ <td>workspace</td>
299+ <td>输入</td>
300+ <td>在Device侧申请的workspace内存地址。</td>
301+ </tr>
302+ <tr>
303+ <td>workspaceSize</td>
304+ <td>输入</td>
305+ <td>在Device侧申请的workspace大小，由第一段接口aclnnMsaIndexScoreGetWorkspaceSize获取。</td>
306+ </tr>
307+ <tr>
308+ <td>executor</td>
309+ <td>输入</td>
310+ <td>op执行器，包含了算子计算流程。</td>
311+ </tr>
312+ <tr>
313+ <td>stream</td>
314+ <td>输入</td>
315+ <td>指定执行任务的Stream。</td>
316+ </tr>
317+ </tbody>
318+ </table>
116 319 
117-### 返回值320+- **返回值：**
118 321 
119-| 返回码 | 错误码 | 说明 |322+ aclnnStatus：返回状态码，具体参见[aclnn返回码](../../../docs/zh/context/aclnn_return_code.md)。
120-| ----------------------- | ------ | ---------- |323+ 
121-| ACLNN_SUCCESS | 0 | 执行成功 |324+ <table style="undefined;table-layout: fixed;width: 1155px"><colgroup>
122-| ACLNN_ERR_PARAM_INVALID | 161002 | 参数不合法 |325+ <col style="width: 319px">
326+ <col style="width: 144px">
327+ <col style="width: 671px">
328+ </colgroup>
329+ <thead>
330+ <tr>
331+ <th>返回码</th>
332+ <th>错误码</th>
333+ <th>描述</th>
334+ </tr>
335+ </thead>
336+ <tbody>
337+ <tr>
338+ <td>ACLNN_SUCCESS</td>
339+ <td>0</td>
340+ <td>执行成功。</td>
341+ </tr>
342+ <tr>
343+ <td>ACLNN_ERR_PARAM_INVALID</td>
344+ <td>161002</td>
345+ <td>参数不合法。</td>
346+ </tr>
347+ </tbody>
348+ </table>
123 349 
124## 约束说明350## 约束说明
125 351 
126- 当前 $block\_size$ 值使用128。352- 当前 $block\_size$ 值使用128。
127- `layoutKeyOptional` 必须显式指定 key 布局：`"BBND"`（$[block\_num, block\_size, N2, D]$）、`"BNBD"`（$[block\_num, N2, block\_size, D]$）或 `"TND"`（$[T2, N2, D]$）。不传时默认 `"BBND"`。353- `layoutKeyOptional` 必须显式指定 key 布局：`"BBND"`（$[block\_num, block\_size, N2, D]$）、`"BNBD"`（$[block\_num, N2, block\_size, D]$）或 `"TND"`（$[T2, N2, D]$）。不传时默认 `"BBND"`。
128- PageAttention（`layoutKey` 为 `"BBND"` / `"BNBD"`）场景下，`blockTableOptional` 必须传入；TND key 场景不得传入 `blockTableOptional`，`actualSeqKlenOptional` 为 $[B+1]$ 前缀和。354- PageAttention（`layoutKey` 为 `"BBND"` / `"BNBD"`）场景下，`blockTableOptional` 必须传入；TND key 场景不得传入 `blockTableOptional`，`actualSeqKlenOptional` 为 $[B+1]$ 前缀和。
129-- 非量化场景下，`key` dtype 与 `query` 相同（当前为 BFLOAT16 / FLOAT16），`scaleOptional` 必须为 `nullptr`；量化场景下仅支持 INT8，`scaleOptional` 必选，dtype 为 FLOAT：PA 为 $[block\_num, N2, block\_size]$ 或 $[block\_num, block\_size, N2]$，TND 为 $[T2, N2]$。当前不支持 FP8 与 <term>Ascend 950PR/Ascend 950DT</term>。355+- 非量化场景下，`key` dtype 与 `query` 相同（BFLOAT16 / FLOAT16；Ascend 950 另支持 HIFLOAT8 / FLOAT8_E5M2 / FLOAT8_E4M3FN），`scaleOptional` 必须为 `nullptr`。量化场景仅支持 INT8（fp16 query），`scaleOptional` 必选，dtype 为 FLOAT：PA 为 $[block\_num, N2, block\_size]$ 或 $[block\_num, block\_size, N2]$，TND 为 $[T2, N2]$。三种 FP8 仅 950，query 与 key 必须同型。
130- `sparseMode` 当前仅支持 0、3：356- `sparseMode` 当前仅支持 0、3：
131 - 为 0 时，代表 defaultMask 模式，`attenMaskOptional` 传入 `nullptr`；357 - 为 0 时，代表 defaultMask 模式，`attenMaskOptional` 传入 `nullptr`；
132 - 为 3 时，代表 rightDownCausal 模式，`attenMaskOptional` 必须传入，shape 为 $[2048, 2048]$，取值为 1 代表该位不参与计算，为 0 代表该位参与计算。358 - 为 3 时，代表 rightDownCausal 模式，`attenMaskOptional` 必须传入，shape 为 $[2048, 2048]$，取值为 1 代表该位不参与计算，为 0 代表该位参与计算。
133- `initBlocks`、`localBlocks` 必须 $\ge 0$ 且不超过逻辑 block 数（PA 为 `blockTableOptional` 第二维；TND 为 score 末维对齐宽度）。两者均为 0 时跳过 $local\_mask$。359- `initBlocks`、`localBlocks` 必须 $\ge 0$ 且不超过逻辑 block 数（PA 为 `blockTableOptional` 第二维；TND 为 score 末维对齐宽度）。两者均为 0 时跳过 $local\_mask$。
360+- `q_len` / `kv_len` 允许为 0（含整 batch）。对应请求跳过 QK；空 KV 的 score 填 $-inf$；整 batch $T1=0$ 时 `SetBlockDim(1)`（对齐 FIA PR 9246）。
361+- A2/A3 与 Ascend 950：按估计 M-task 数启动 MIX（单 batch 为 $\mathrm{CeilDiv}(T1\cdot N1,128)$；多 batch 取 packed+$B$ 上界并截到 AIC 数），短 decode 不打满空核。
362+- PageAttention `key`（BBND/BNBD）允许首轴非连续（`key | gap | key | ...`），tiling 通过 `GetInputStride` 读取 dim0 元素 stride 写入 `strideKvBlock`；非首轴必须连续。TND `key` 不允许非连续。`scale` 仍按逻辑 page 紧凑布局。
363+- PageAttention `blockTable` 第二维可以大于实际 KV 逻辑 block 数；score 末维为 $\mathrm{RoundUp}(width, 16)$。Ascend 950 C2UB 路径对超过 256 列的末维按 256 列滑窗 flush，并补写后续 $-inf$。
364+- 精度自验证矩阵：39 条 fp16/bf16/int8（含 PA key dim0 stride、宽 `blockTable`、短 decode）+ 8 条 FP8。950 通过标准末行 `[PASS]: 47/47 cases passed`。A2/A3 跳过 FP8，期望 39/39（skipped 8）。容差 fp16/bf16/int8 为 $1\mathrm{e}{-3}$，FP8 为 $2\mathrm{e}{-2}$。
365+- Ascend 950 核实现位于 `op_kernel/arch35/`：当前与 A2 共用 8-page S workspace（非量化 fp16 / int8 fp32）；Cube 原生三种 FP8。`--run_example` 默认 soc 为 910b，950 必须显式 `--soc=ascend950`。
134 366 
135 367 
136## 调用示例368## 调用示例
Mattention/msa_index_score/docs/torchapi_msa_index_score.md+112-13
@@ -9,7 +9,7 @@
9- <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：支持9- <term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>：支持
10<!-- end id2 -->10<!-- end id2 -->
11<!-- npu="950" id3 -->11<!-- npu="950" id3 -->
12-- <term>Ascend 950PR/Ascend 950DT</term>：不支持12+- <term>Ascend 950PR/Ascend 950DT</term>：支持
13<!-- end id3 -->13<!-- end id3 -->
14 14 
15## 功能说明15## 功能说明
@@ -23,6 +23,12 @@ $$
23 23 
24- `sparse_mode=0`：无因果；`sparse_mode=3`：rightDownCausal（须传 `[2048,2048]` `atten_mask`）。24- `sparse_mode=0`：无因果；`sparse_mode=3`：rightDownCausal（须传 `[2048,2048]` `atten_mask`）。
25- `start_loc[B]`：当前 query 所在逻辑 block 索引，用于 `local_mask`。25- `start_loc[B]`：当前 query 所在逻辑 block 索引，用于 `local_mask`。
26+- `init_blocks` / `local_blocks`：Maxpool 之后的强制高分块；与 Triton raw score 对齐时置 0。默认 `0` / `1`。
27+- Ascend 950 额外支持 query/key 同型 `torch.float8_e4m3fn` / `torch.float8_e5m2` / `torch_npu.hifloat8`（无 `scale`）。`hifloat8.npu()` 当前不安全，测试脚本跳过；kernel 已注册。
28+- `q_len` / `kv_len` 允许为 0（含整 batch）。对应请求跳过 QK；空 KV 的 score 为 `-inf`。
29+- A2/A3 与 Ascend 950 短 decode 按估计 M-task 启动 MIX，不打满空核。
30+- PageAttention `key` 允许 dim0（物理 page）非连续，须保持为 view（不要 `.contiguous()`）；TND 不允许。
31+- PageAttention `block_table` 第二维可以大于实际 KV 逻辑 block 数（例如 vLLM 预分配宽表）。950 C2UB 对 score 末维超过 256 列按窗 flush。
26 32 
27## 函数原型33## 函数原型
28 34 
@@ -39,23 +45,116 @@ cann_ops_transformer.msa_index_score(
39 actual_seq_klen=None,45 actual_seq_klen=None,
40 layout_key="BBND",46 layout_key="BBND",
41 sparse_mode=3,47 sparse_mode=3,
48+ init_blocks=0,
49+ local_blocks=1,
42) -> Tensor50) -> Tensor
43```51```
44 52 
45## 参数说明53## 参数说明
46 54 
47-| 参数名 | 可选/必选 | 描述 | dtype | shape |55+<table style="undefined;table-layout: fixed; width: 1180px"><colgroup>
48-|--------|-----------|------|-------|-------|56+<col style="width: 160px">
49-| query | 必选 | $Q_{idx}$ TND | fp16/bf16 | `[T1,N1,D]` |57+<col style="width: 120px">
50-| key | 必选 | $K_{idx}$ PA BBND/BNBD 或 TND | 同 query 或 int8 | `[NP,P,N2,D]` / `[NP,N2,P,D]` / `[T2,N2,D]` |58+<col style="width: 360px">
51-| start_loc | 必选 | query 所在逻辑 block | int32 | `[B]` |59+<col style="width: 280px">
52-| block_table | PA 必选 | 逻辑→物理 page；TND 不传 | int32 | `[B,MB]` |60+<col style="width: 260px">
53-| scale | 量化必选 | 反量化 | float32 | PA `[NP,N2,P]`；TND `[T2,N2]` |61+</colgroup>
54-| atten_mask | mode=3 必选 | 压缩下三角模板 | int8 | `[2048,2048]` |62+<thead>
55-| actual_seq_qlen | TND query 必选 | query 前缀和 | int32 | `[B+1]` |63+ <tr>
56-| actual_seq_klen | 必选 | PA 为各请求 S2；TND 为 key 前缀和 | int32 | `[B]` / `[B+1]` |64+ <th>参数名</th>
57-| layout_key | 可选 | key 布局：`TND` / `BBND` / `BNBD`，默认 `BBND` | str | - |65+ <th>可选/必选</th>
58-| sparse_mode | 可选 | 0 / 3，默认 3 | int | - |66+ <th>描述</th>
67+ <th>dtype</th>
68+ <th>shape</th>
69+ </tr>
70+</thead>
71+<tbody>
72+ <tr>
73+ <td>query</td>
74+ <td>必选</td>
75+ <td>$Q_{idx}$ TND</td>
76+ <td>fp16/bf16；950 另支持 fp8_e4m3fn / fp8_e5m2 / hifloat8</td>
77+ <td><code>[T1,N1,D]</code></td>
78+ </tr>
79+ <tr>
80+ <td>key</td>
81+ <td>必选</td>
82+ <td>$K_{idx}$ PA BBND/BNBD 或 TND。A2/A3 与 950 上 PA 允许 dim0 非连续。</td>
83+ <td>同 query 或 int8（int8 配 fp16 query）</td>
84+ <td><code>[NP,P,N2,D]</code> / <code>[NP,N2,P,D]</code> / <code>[T2,N2,D]</code></td>
85+ </tr>
86+ <tr>
87+ <td>start_loc</td>
88+ <td>必选</td>
89+ <td>query 所在逻辑 block</td>
90+ <td>int32</td>
91+ <td><code>[B]</code></td>
92+ </tr>
93+ <tr>
94+ <td>block_table</td>
95+ <td>PA 必选</td>
96+ <td>逻辑→物理 page；TND 不传</td>
97+ <td>int32</td>
98+ <td><code>[B,MB]</code></td>
99+ </tr>
100+ <tr>
101+ <td>scale</td>
102+ <td>量化必选</td>
103+ <td>反量化</td>
104+ <td>float32</td>
105+ <td>PA <code>[NP,N2,P]</code>；TND <code>[T2,N2]</code></td>
106+ </tr>
107+ <tr>
108+ <td>atten_mask</td>
109+ <td>mode=3 必选</td>
110+ <td>压缩下三角模板</td>
111+ <td>int8</td>
112+ <td><code>[2048,2048]</code></td>
113+ </tr>
114+ <tr>
115+ <td>actual_seq_qlen</td>
116+ <td>TND query 必选</td>
117+ <td>query 前缀和</td>
118+ <td>int32</td>
119+ <td><code>[B+1]</code></td>
120+ </tr>
121+ <tr>
122+ <td>actual_seq_klen</td>
123+ <td>必选</td>
124+ <td>PA 为各请求 S2；TND 为 key 前缀和</td>
125+ <td>int32</td>
126+ <td><code>[B]</code> / <code>[B+1]</code></td>
127+ </tr>
128+ <tr>
129+ <td>layout_key</td>
130+ <td>可选</td>
131+ <td>key 布局：<code>TND</code> / <code>BBND</code> / <code>BNBD</code>，默认 <code>BBND</code></td>
132+ <td>str</td>
133+ <td>-</td>
134+ </tr>
135+ <tr>
136+ <td>sparse_mode</td>
137+ <td>可选</td>
138+ <td>0 / 3，默认 3</td>
139+ <td>int</td>
140+ <td>-</td>
141+ </tr>
142+ <tr>
143+ <td>init_blocks</td>
144+ <td>可选</td>
145+ <td><code>local_mask</code> 头部强制块数，默认 0</td>
146+ <td>int</td>
147+ <td>-</td>
148+ </tr>
149+ <tr>
150+ <td>local_blocks</td>
151+ <td>可选</td>
152+ <td><code>local_mask</code> 局部窗口，默认 1</td>
153+ <td>int</td>
154+ <td>-</td>
155+ </tr>
156+</tbody>
157+</table>
59 158 
60## 输出159## 输出
61 160 
Mattention/msa_index_score/examples/test_aclnn_msa_index_score.cpp+386-28
@@ -13,12 +13,16 @@
13 * \brief aclnnMsaIndexScore 调用示例，内置 CPU golden 做端到端精度自验证。13 * \brief aclnnMsaIndexScore 调用示例，内置 CPU golden 做端到端精度自验证。
14 *14 *
15 * 用例矩阵覆盖：Prefill 多 M-tile、prefix 非 128 对齐的边界 block、varlen 多 batch、15 * 用例矩阵覆盖：Prefill 多 M-tile、prefix 非 128 对齐的边界 block、varlen 多 batch、
16- * Decode(q_len=1)、投机解码(q_len>1)、长序列多 S-tile 轮转、block_table 乱序、16+ * Decode(q_len=1)、投机解码(q_len>1)、短 decode（Hq=4 / 宽 block_table）、
17- * 无效尾填充、bf16 / fp16 双 dtype、int8 key 前融合反量化、PA BNBD、TND packed key。17+ * 长序列多 S-tile 轮转、block_table 乱序、
18+ * 无效尾填充、q_len/kv_len=0 的 mixed-batch pad、bf16 / fp16 双 dtype、
19+ * int8 key 前融合反量化、PA BNBD、TND packed key、A2/A3 与 950 PA key dim0 非连续。
18 */20 */
19 21 
20#include <algorithm>22#include <algorithm>
21#include <cmath>23#include <cmath>
24+#include <cfenv>
25+#include <fenv.h>
22#include <cstdint>26#include <cstdint>
23#include <cstdio>27#include <cstdio>
24#include <cstdlib>28#include <cstdlib>
Mattention/msa_index_score/op_host/msa_index_score_def.cpp+59-7
@@ -13,11 +13,8 @@
13 * \brief MsaIndexScore 算子原型注册（对齐 docs/aclnnMsaIndexScore.md）。13 * \brief MsaIndexScore 算子原型注册（对齐 docs/aclnnMsaIndexScore.md）。
14 *14 *
15 * dtype 组合（按列表下标对齐）：15 * dtype 组合（按列表下标对齐）：
16- * 0: query=bf16, key=bf16 （非量化）16+ * A2/A3（3 组，全 ND）：0 bf16/bf16、1 fp16/fp16、2 fp16/int8
17- * 1: query=fp16, key=fp16 （非量化）17+ * Ascend 950（6 组）：0–2 同上（ND），3 hifloat8、4 fp8_e5m2、5 fp8_e4m3fn
18- * 2: query=fp16, key=int8 （前融合 int8，scale 为反量化系数）
19- *
20- * A2/A3：PageAttention BBND/BNBD key 与 TND packed key；sparse_mode∈{0,3}。不支持 FP8 / Ascend 950。
21 */18 */
22 19 
23#include <cstdint>20#include <cstdint>
@@ -36,11 +33,12 @@ public:
36 .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})33 .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
37 .AutoContiguous();34 .AutoContiguous();
38 // K_idx：BBND [block_num, block_size, N2, D]、BNBD [block_num, N2, block_size, D]、TND [T2, N2, D]35 // K_idx：BBND [block_num, block_size, N2, D]、BNBD [block_num, N2, block_size, D]、TND [T2, N2, D]
36+ // PA key 允许 dim0（物理 page）非连续；TND 仍须连续（tiling 拒绝）。
39 this->Input("key")37 this->Input("key")
40 .ParamType(REQUIRED)38 .ParamType(REQUIRED)
41 .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_INT8})39 .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_INT8})
42 .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})40 .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
43- .AutoContiguous();41+ .IgnoreContiguous();
44 // PA 逻辑 block → 物理 page；PA 场景必须传入42 // PA 逻辑 block → 物理 page；PA 场景必须传入
45 this->Input("block_table")43 this->Input("block_table")
46 .ParamType(OPTIONAL)44 .ParamType(OPTIONAL)
@@ -98,9 +96,63 @@ public:
98 .DynamicShapeSupportFlag(true)96 .DynamicShapeSupportFlag(true)
99 .NeedCheckSupportFlag(false)97 .NeedCheckSupportFlag(false)
100 .PrecisionReduceFlag(true);98 .PrecisionReduceFlag(true);
101- // Atlas A2（ascend910b）/ A3（ascend910_93）；不注册 Ascend 950。
102 this->AICore().AddConfig("ascend910b", aicoreConfig);99 this->AICore().AddConfig("ascend910b", aicoreConfig);
103 this->AICore().AddConfig("ascend910_93", aicoreConfig);100 this->AICore().AddConfig("ascend910_93", aicoreConfig);
101+ 
102+ OpAICoreConfig aicoreConfig950;
103+ aicoreConfig950.Input("query")
104+ .ParamType(REQUIRED)
105+ .DataType({ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_HIFLOAT8, ge::DT_FLOAT8_E5M2,
106+ ge::DT_FLOAT8_E4M3FN})
107+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
108+ .AutoContiguous();
109+ aicoreConfig950.Input("key")
110+ .ParamType(REQUIRED)
111+ .DataType(
112+ {ge::DT_BF16, ge::DT_FLOAT16, ge::DT_INT8, ge::DT_HIFLOAT8, ge::DT_FLOAT8_E5M2, ge::DT_FLOAT8_E4M3FN})
113+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
114+ .IgnoreContiguous();
115+ aicoreConfig950.Input("block_table")
116+ .ParamType(OPTIONAL)
117+ .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
118+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
119+ .AutoContiguous();
120+ aicoreConfig950.Input("scale")
121+ .ParamType(OPTIONAL)
122+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
123+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
124+ .AutoContiguous();
125+ aicoreConfig950.Input("atten_mask")
126+ .ParamType(OPTIONAL)
127+ .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8})
128+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
129+ .AutoContiguous();
130+ aicoreConfig950.Input("actual_seq_qlen")
131+ .ParamType(OPTIONAL)
132+ .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
133+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
134+ .AutoContiguous();
135+ aicoreConfig950.Input("actual_seq_klen")
136+ .ParamType(OPTIONAL)
137+ .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
138+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
139+ .AutoContiguous();
140+ aicoreConfig950.Input("start_loc")
141+ .ParamType(REQUIRED)
142+ .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
143+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
144+ .AutoContiguous();
145+ aicoreConfig950.Output("score")
146+ .ParamType(REQUIRED)
147+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
148+ .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
149+ aicoreConfig950.DynamicCompileStaticFlag(true)
150+ .DynamicFormatFlag(true)
151+ .DynamicRankSupportFlag(true)
152+ .DynamicShapeSupportFlag(true)
153+ .NeedCheckSupportFlag(false)
154+ .PrecisionReduceFlag(true);
155+ this->AICore().AddConfig("ascend950", aicoreConfig950);
104 }156 }
105};157};
106OP_ADD(MsaIndexScore);158OP_ADD(MsaIndexScore);
Mattention/msa_index_score/op_host/msa_index_score_infershape.cpp+9-2
@@ -36,7 +36,10 @@ constexpr uint32_t MSA_ATTR_LAYOUT_KEY = 0;
36constexpr int64_t MSA_SCORE_STRIDE_ALIGN = 16;36constexpr int64_t MSA_SCORE_STRIDE_ALIGN = 16;
37constexpr int64_t MSA_BLOCK_SIZE = 128;37constexpr int64_t MSA_BLOCK_SIZE = 128;
38 38 
39-inline int64_t RoundUpTo(int64_t value, int64_t align) { return (value + align - 1) / align * align; }39+inline int64_t RoundUpTo(int64_t value, int64_t align)
40+{
41+ return (value + align - 1) / align * align;
42+}
40} // namespace43} // namespace
41 44 
42static ge::graphStatus InferShapeMsaIndexScore(gert::InferShapeContext *context)45static ge::graphStatus InferShapeMsaIndexScore(gert::InferShapeContext *context)
@@ -95,7 +98,11 @@ static ge::graphStatus InferShapeMsaIndexScore(gert::InferShapeContext *context)
95 maxBlocks = (totalK + MSA_BLOCK_SIZE - 1) / MSA_BLOCK_SIZE;98 maxBlocks = (totalK + MSA_BLOCK_SIZE - 1) / MSA_BLOCK_SIZE;
96 }99 }
97 }100 }
98- OP_CHECK_IF(maxBlocks <= 0, OP_LOGE(context, "maxBlocksPerSeq must be positive."), return ge::GRAPH_FAILED);101+ // 不拦截全 kv_len=0：按 1 个 dummy block 再 16 对齐，score 末维为 16；
102+ // 核内跳过 QK，把该请求的 score 填 -inf（T1=0 时输出为空张量）。
103+ if (maxBlocks <= 0) {
104+ maxBlocks = 1;
105+ }
99 106 
100 // maxpool 已经把一个 block 内的 blockSize 个 token 归约成 1 个分数，107 // maxpool 已经把一个 block 内的 blockSize 个 token 归约成 1 个分数，
101 // 因此输出末维是 block 数（16 对齐）而不是 kv token 数。108 // 因此输出末维是 block 数（16 对齐）而不是 kv token 数。
Mattention/msa_index_score/op_host/msa_index_score_tiling.cpp+231-27
@@ -10,13 +10,16 @@
10 10 
11/*!11/*!
12 * \file msa_index_score_tiling.cpp12 * \file msa_index_score_tiling.cpp
13- * \brief MsaIndexScore Tiling 实现（A2/A3：PA BBND/BNBD + TND packed key；sparse_mode 0/3）。13+ * \brief MsaIndexScore Tiling 实现（A2/A3 / 950：PA BBND/BNBD + TND packed key；sparse_mode 0/3；
14+ * PA key 允许 dim0 非连续，strideKvBlock 来自 GetInputStride）。
14 */15 */
15 16 
16#include "msa_index_score_tiling.h"17#include "msa_index_score_tiling.h"
17#include "../op_kernel/msa_index_score_common.h"18#include "../op_kernel/msa_index_score_common.h"
18 19 
20+#include <limits>
19#include <string>21#include <string>
22+#include "graph/utils/type_utils.h"
20 23 
21using namespace ge;24using namespace ge;
22using namespace MsaIndexScoreNs;25using namespace MsaIndexScoreNs;
@@ -37,7 +40,151 @@ constexpr uint32_t MSA_DIM_3 = 3;
37constexpr uint32_t MSA_RANK_1D = 1;40constexpr uint32_t MSA_RANK_1D = 1;
38constexpr uint32_t MSA_QLEN_MIN_DIM0 = 2; // prefix-sum 长度至少为 B+141constexpr uint32_t MSA_QLEN_MIN_DIM0 = 2; // prefix-sum 长度至少为 B+1
39 42 
40-inline uint32_t RoundUpU32(uint32_t value, uint32_t align) { return (value + align - 1) / align * align; }43+inline bool IsFp8ComputeDtype(ge::DataType dt)
44+{
45+ return dt == ge::DT_HIFLOAT8 || dt == ge::DT_FLOAT8_E5M2 || dt == ge::DT_FLOAT8_E4M3FN;
46+}
47+ 
48+inline bool IsNonQuantQueryDtype(ge::DataType dt, bool isAscend950)
49+{
50+ if (dt == ge::DT_BF16 || dt == ge::DT_FLOAT16) {
51+ return true;
52+ }
53+ return isAscend950 && IsFp8ComputeDtype(dt);
54+}
55+ 
56+inline uint32_t RoundUpU32(uint32_t value, uint32_t align)
57+{
58+ return (value + align - 1) / align * align;
59+}
60+ 
61+// MIX：按估计 M-task 数启动 AIC，避免短 decode 打满空核（scheduler.Init GetValue / MODE4 VC 空转）。
62+// Host 读不到 per-request q_len。kernel TotalTasks = Σ CeilDiv(qLen_b * Hq, MSA_ROW_TILE_M)。
63+// Σ ceil(x_i) <= ceil(Σ x_i) + B；B=1 时 packed == actual。B>1 用 packed+B 上界，避免少估把
64+// 多 tile 请求串到已启动核上。多 M-tile / 大 batch 仍截到全量 AIC。
65+inline uint32_t EstLaunchAic(const MsaIndexScoreInfo &info, uint32_t aicNum)
66+{
67+ if (info.totalQ == 0U || aicNum == 0U) {
68+ return 1U;
69+ }
70+ const uint64_t packedRows = static_cast<uint64_t>(info.totalQ) * static_cast<uint64_t>(info.numQHeads);
71+ const uint32_t packedM =
72+ static_cast<uint32_t>((packedRows + static_cast<uint64_t>(MSA_ROW_TILE_M) - 1U) / MSA_ROW_TILE_M);
73+ uint32_t est = packedM;
74+ if (info.batch > 1U) {
75+ const uint32_t add = packedM + info.batch;
76+ est = (add < packedM) ? aicNum : add;
77+ }
78+ if (est < 1U) {
79+ est = 1U;
80+ }
81+ if (est > aicNum) {
82+ est = aicNum;
83+ }
84+ return est;
85+}
86+ 
87+inline uint64_t GetDefaultStride0(const gert::Shape &shape)
88+{
89+ uint64_t stride = 1;
90+ for (size_t dim = 1; dim < shape.GetDimNum(); ++dim) {
91+ stride *= static_cast<uint64_t>(shape.GetDim(dim));
92+ }
93+ return stride;
94+}
95+ 
96+// 逻辑 view：优先 OriginShape（aclCreateTensor 的 viewDims / torch sizes）。
97+// StorageShape 在 dim0 插空时可能是 [NP*gap, ...]，不能当 numPages。
98+inline const gert::Shape &KeyLogicalShape(const gert::StorageShape *keyShape, uint32_t keyLayout)
99+{
100+ const size_t expectRank = (keyLayout == MSA_KEY_LAYOUT_TND) ? MSA_KEY_TND_DIM_NUM : MSA_KEY_PA_DIM_NUM;
101+ const gert::Shape &origin = keyShape->GetOriginShape();
102+ if (origin.GetDimNum() == expectRank) {
103+ return origin;
104+ }
105+ return keyShape->GetStorageShape();
106+}
107+ 
108+// aclnn TensorV2 优先 GetRequiredInputStride；
109+// 文档接口 GetInputStride 作为兼容回退。
110+inline const gert::Stride *GetKeyStrideDesc(gert::TilingContext *context)
111+{
112+ const gert::Stride *stride = context->GetRequiredInputStride(MSA_IDX_KEY);
113+ if (stride != nullptr) {
114+ return stride;
115+ }
116+ return context->GetInputStride(MSA_IDX_KEY);
117+}
118+ 
119+inline bool ValidateKeyInnerAxesContiguous(gert::TilingContext *context, const gert::Stride *stride,
120+ const gert::Shape &shape)
121+{
122+ if (stride == nullptr || stride->GetDimNum() != shape.GetDimNum()) {
123+ return true;
124+ }
125+ uint64_t expectedStride = 1;
126+ for (int64_t i = static_cast<int64_t>(shape.GetDimNum()) - 1; i >= 1; --i) {
127+ const int64_t dimSize = shape.GetDim(static_cast<size_t>(i));
128+ if (dimSize <= 1) {
129+ // size-1 轴不参与寻址；PyTorch 对其 stride 不做连续约束（BNBD 且 N2=1 时常见）。
130+ continue;
131+ }
132+ const uint64_t actualStride = static_cast<uint64_t>(stride->GetStride(static_cast<size_t>(i)));
133+ if (actualStride != expectedStride) {
134+ OP_LOGE(context,
135+ "key dim%ld must be contiguous, actual stride=%lu, expected=%lu. "
136+ "Only dim0 (PA page axis) may be non-contiguous.",
137+ i, actualStride, expectedStride);
138+ return false;
139+ }
140+ expectedStride *= static_cast<uint64_t>(dimSize);
141+ }
142+ return true;
143+}
144+ 
145+// PA：strideKvBlock 取 key dim0 元素 stride。连续输入写入值与 shape 推算相同。
146+inline ge::graphStatus ResolveStrideKvBlock(gert::TilingContext *context, const MsaIndexScoreInfo &info,
147+ uint32_t defaultStrideKvBlock, uint32_t &strideKvBlock)
148+{
149+ strideKvBlock = defaultStrideKvBlock;
150+ 
151+ const gert::StorageShape *keyShape = context->GetInputShape(MSA_IDX_KEY);
152+ OP_CHECK_NULL_WITH_CONTEXT(context, keyShape);
153+ const gert::Shape &kc = KeyLogicalShape(keyShape, info.keyLayout);
154+ if (kc.GetShapeSize() == 0) {
155+ return ge::GRAPH_SUCCESS;
156+ }
157+ 
158+ const uint64_t defaultStride0 = GetDefaultStride0(kc);
159+ const gert::Stride *stride = GetKeyStrideDesc(context);
160+ if (stride == nullptr || stride->GetDimNum() != kc.GetDimNum()) {
161+ OP_LOGD(context->GetNodeName(), "key has no stride desc, treat as contiguous, strideKvBlock=%u.",
162+ strideKvBlock);
163+ return ge::GRAPH_SUCCESS;
164+ }
165+ OP_CHECK_IF(!ValidateKeyInnerAxesContiguous(context, stride, kc),
166+ OP_LOGE(context, "key inner axes must stay contiguous."), return ge::GRAPH_FAILED);
167+ 
168+ const int64_t actualStride0 = stride->GetStride(MSA_DIM_0);
169+ if (info.keyLayout == MSA_KEY_LAYOUT_TND) {
170+ OP_CHECK_IF(actualStride0 < 0 || static_cast<uint64_t>(actualStride0) != defaultStride0,
171+ OP_LOGE(context, "TND key dim0 must be contiguous, actual=%ld expected=%lu.", actualStride0,
172+ defaultStride0),
173+ return ge::GRAPH_FAILED);
174+ strideKvBlock = 0;
175+ return ge::GRAPH_SUCCESS;
176+ }
177+ 
178+ OP_CHECK_IF(actualStride0 < 0 || static_cast<uint64_t>(actualStride0) < defaultStride0,
179+ OP_LOGE(context, "PA key dim0 stride must be >= %lu, got %ld.", defaultStride0, actualStride0),
180+ return ge::GRAPH_FAILED);
181+ OP_CHECK_IF(static_cast<uint64_t>(actualStride0) > static_cast<uint64_t>(std::numeric_limits<uint32_t>::max()),
182+ OP_LOGE(context, "PA key dim0 stride %ld exceeds uint32 strideKvBlock.", actualStride0),
183+ return ge::GRAPH_FAILED);
184+ strideKvBlock = static_cast<uint32_t>(actualStride0);
185+ OP_LOGI(context->GetNodeName(), "PA key strideKvBlock=%u (contiguous default=%lu).", strideKvBlock, defaultStride0);
186+ return ge::GRAPH_SUCCESS;
187+}
41 188 
42ge::graphStatus ParseLayoutKeyAttr(gert::TilingContext *context, const std::string &layoutKey, uint32_t &keyLayout)189ge::graphStatus ParseLayoutKeyAttr(gert::TilingContext *context, const std::string &layoutKey, uint32_t &keyLayout)
43{190{
@@ -59,6 +206,8 @@ ge::graphStatus ParseAndCheck(gert::TilingContext *context, MsaIndexScoreInfo &i
59{206{
60 info.platformInfo = context->GetPlatformInfo();207 info.platformInfo = context->GetPlatformInfo();
61 OP_CHECK_IF(info.platformInfo == nullptr, OP_LOGE(context, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED);208 OP_CHECK_IF(info.platformInfo == nullptr, OP_LOGE(context, "GetPlatformInfo is nullptr."), return ge::GRAPH_FAILED);
209+ auto ascendcPlatformEarly = platform_ascendc::PlatformAscendC(info.platformInfo);
210+ info.isAscend950 = (ascendcPlatformEarly.GetSocVersion() == platform_ascendc::SocVersion::ASCEND950);
62 211 
63 const gert::StorageShape *queryShape = context->GetInputShape(MSA_IDX_QUERY);212 const gert::StorageShape *queryShape = context->GetInputShape(MSA_IDX_QUERY);
64 OP_CHECK_NULL_WITH_CONTEXT(context, queryShape);213 OP_CHECK_NULL_WITH_CONTEXT(context, queryShape);
@@ -75,10 +224,15 @@ ge::graphStatus ParseAndCheck(gert::TilingContext *context, MsaIndexScoreInfo &i
75 OP_CHECK_NULL_WITH_CONTEXT(context, startLocShape);224 OP_CHECK_NULL_WITH_CONTEXT(context, startLocShape);
76 225 
77 const gert::Shape &q = queryShape->GetStorageShape();226 const gert::Shape &q = queryShape->GetStorageShape();
78- const gert::Shape &kc = keyShape->GetStorageShape();
79 const gert::Shape &qlen = actualSeqQlenShape->GetStorageShape();227 const gert::Shape &qlen = actualSeqQlenShape->GetStorageShape();
80 const gert::Shape &klen = actualSeqKlenShape->GetStorageShape();228 const gert::Shape &klen = actualSeqKlenShape->GetStorageShape();
81 229 
230+ const auto *keyDesc = context->GetInputDesc(MSA_IDX_KEY);
231+ OP_CHECK_NULL_WITH_CONTEXT(context, keyDesc);
232+ info.keyDtype = keyDesc->GetDataType();
233+ OP_CHECK_IF(static_cast<ge::Format>(ge::GetPrimaryFormat(keyDesc->GetStorageFormat())) == ge::FORMAT_FRACTAL_NZ,
234+ OP_LOGE(context, "FRACTAL_NZ key is not supported."), return ge::GRAPH_FAILED);
235+ 
82 OP_CHECK_IF(q.GetDimNum() != MSA_QUERY_DIM_NUM,236 OP_CHECK_IF(q.GetDimNum() != MSA_QUERY_DIM_NUM,
83 OP_LOGE(context, "query must be TND(3 dims), got %zu.", q.GetDimNum()), return ge::GRAPH_FAILED);237 OP_LOGE(context, "query must be TND(3 dims), got %zu.", q.GetDimNum()), return ge::GRAPH_FAILED);
84 OP_CHECK_IF(qlen.GetDimNum() != MSA_RANK_1D || qlen.GetDim(MSA_DIM_0) < MSA_QLEN_MIN_DIM0,238 OP_CHECK_IF(qlen.GetDimNum() != MSA_RANK_1D || qlen.GetDim(MSA_DIM_0) < MSA_QLEN_MIN_DIM0,
@@ -96,6 +250,7 @@ ge::graphStatus ParseAndCheck(gert::TilingContext *context, MsaIndexScoreInfo &i
96 if (ParseLayoutKeyAttr(context, layoutKey, info.keyLayout) != ge::GRAPH_SUCCESS) {250 if (ParseLayoutKeyAttr(context, layoutKey, info.keyLayout) != ge::GRAPH_SUCCESS) {
97 return ge::GRAPH_FAILED;251 return ge::GRAPH_FAILED;
98 }252 }
253+ const gert::Shape &kc = KeyLogicalShape(keyShape, info.keyLayout);
99 254 
100 if (info.keyLayout == MSA_KEY_LAYOUT_TND) {255 if (info.keyLayout == MSA_KEY_LAYOUT_TND) {
101 info.totalK = static_cast<uint32_t>(kc.GetDim(MSA_DIM_0));256 info.totalK = static_cast<uint32_t>(kc.GetDim(MSA_DIM_0));
@@ -114,10 +269,6 @@ ge::graphStatus ParseAndCheck(gert::TilingContext *context, MsaIndexScoreInfo &i
114 OP_CHECK_IF(klen.GetDimNum() != MSA_RANK_1D || static_cast<uint32_t>(klen.GetDim(MSA_DIM_0)) != info.batch + 1,269 OP_CHECK_IF(klen.GetDimNum() != MSA_RANK_1D || static_cast<uint32_t>(klen.GetDim(MSA_DIM_0)) != info.batch + 1,
115 OP_LOGE(context, "TND actual_seq_klen must be [B+1] prefix-sum."), return ge::GRAPH_FAILED);270 OP_LOGE(context, "TND actual_seq_klen must be [B+1] prefix-sum."), return ge::GRAPH_FAILED);
116 } else {271 } else {
117- OP_CHECK_IF(kc.GetDimNum() != MSA_KEY_PA_DIM_NUM,
118- OP_LOGE(context, "layout_key=%s requires key rank 4, got %zu.",
119- info.keyLayout == MSA_KEY_LAYOUT_BNBD ? "BNBD" : "BBND", kc.GetDimNum()),
120- return ge::GRAPH_FAILED);
121 OP_CHECK_IF(blockTableShape == nullptr, OP_LOGE(context, "PageAttention (BBND/BNBD) requires block_table."),272 OP_CHECK_IF(blockTableShape == nullptr, OP_LOGE(context, "PageAttention (BBND/BNBD) requires block_table."),
122 return ge::GRAPH_FAILED);273 return ge::GRAPH_FAILED);
123 const gert::Shape &bt = blockTableShape->GetStorageShape();274 const gert::Shape &bt = blockTableShape->GetStorageShape();
@@ -129,6 +280,10 @@ ge::graphStatus ParseAndCheck(gert::TilingContext *context, MsaIndexScoreInfo &i
129 bt.GetDim(MSA_DIM_0), info.batch),280 bt.GetDim(MSA_DIM_0), info.batch),
130 return ge::GRAPH_FAILED);281 return ge::GRAPH_FAILED);
131 info.maxBlocksPerBatch = static_cast<uint32_t>(bt.GetDim(MSA_DIM_1));282 info.maxBlocksPerBatch = static_cast<uint32_t>(bt.GetDim(MSA_DIM_1));
283+ OP_CHECK_IF(kc.GetDimNum() != MSA_KEY_PA_DIM_NUM,
284+ OP_LOGE(context, "layout_key=%s requires key rank 4, got %zu.",
285+ info.keyLayout == MSA_KEY_LAYOUT_BNBD ? "BNBD" : "BBND", kc.GetDimNum()),
286+ return ge::GRAPH_FAILED);
132 info.numPages = static_cast<uint32_t>(kc.GetDim(MSA_DIM_0));287 info.numPages = static_cast<uint32_t>(kc.GetDim(MSA_DIM_0));
133 if (info.keyLayout == MSA_KEY_LAYOUT_BBND) {288 if (info.keyLayout == MSA_KEY_LAYOUT_BBND) {
134 info.blockSize = static_cast<uint32_t>(kc.GetDim(MSA_DIM_1));289 info.blockSize = static_cast<uint32_t>(kc.GetDim(MSA_DIM_1));
@@ -176,8 +331,11 @@ ge::graphStatus ParseAndCheck(gert::TilingContext *context, MsaIndexScoreInfo &i
176 OP_CHECK_IF(info.numKvHeads == 0 || (info.numQHeads % info.numKvHeads) != 0,331 OP_CHECK_IF(info.numKvHeads == 0 || (info.numQHeads % info.numKvHeads) != 0,
177 OP_LOGE(context, "numQHeads(%u) must be divisible by numKvHeads(%u).", info.numQHeads, info.numKvHeads),332 OP_LOGE(context, "numQHeads(%u) must be divisible by numKvHeads(%u).", info.numQHeads, info.numKvHeads),
178 return ge::GRAPH_FAILED);333 return ge::GRAPH_FAILED);
179- OP_CHECK_IF(info.batch == 0 || info.totalQ == 0 || info.maxBlocksPerBatch == 0,334+ OP_CHECK_IF(info.batch == 0, OP_LOGE(context, "batch must be positive."), return ge::GRAPH_FAILED);
180- OP_LOGE(context, "batch/totalQ/maxBlocksPerBatch must be positive."), return ge::GRAPH_FAILED);335+ // q_len/kv_len 全 0 跳过计算。InferShape 已把空 KV 的 score 末维钳到 16。
336+ OP_CHECK_IF(info.maxBlocksPerBatch == 0,
337+ OP_LOGE(context, "maxBlocksPerBatch must be positive (empty kv uses aligned score stride)."),
338+ return ge::GRAPH_FAILED);
181 OP_CHECK_IF(info.keyLayout != MSA_KEY_LAYOUT_TND && info.numPages == 0,339 OP_CHECK_IF(info.keyLayout != MSA_KEY_LAYOUT_TND && info.numPages == 0,
182 OP_LOGE(context, "PA numPages must be positive."), return ge::GRAPH_FAILED);340 OP_LOGE(context, "PA numPages must be positive."), return ge::GRAPH_FAILED);
183 OP_CHECK_IF(static_cast<uint32_t>(startLocShape->GetStorageShape().GetDim(MSA_DIM_0)) != info.batch,341 OP_CHECK_IF(static_cast<uint32_t>(startLocShape->GetStorageShape().GetDim(MSA_DIM_0)) != info.batch,
@@ -186,19 +344,20 @@ ge::graphStatus ParseAndCheck(gert::TilingContext *context, MsaIndexScoreInfo &i
186 const auto *queryDesc = context->GetInputDesc(MSA_IDX_QUERY);344 const auto *queryDesc = context->GetInputDesc(MSA_IDX_QUERY);
187 OP_CHECK_NULL_WITH_CONTEXT(context, queryDesc);345 OP_CHECK_NULL_WITH_CONTEXT(context, queryDesc);
188 info.queryDtype = queryDesc->GetDataType();346 info.queryDtype = queryDesc->GetDataType();
189- OP_CHECK_IF(info.queryDtype != ge::DT_BF16 && info.queryDtype != ge::DT_FLOAT16,347+ OP_CHECK_IF(!IsNonQuantQueryDtype(info.queryDtype, info.isAscend950),
190- OP_LOGE(context, "query dtype must be bf16 or fp16 on A2/A3."), return ge::GRAPH_FAILED);348+ OP_LOGE(context, "query dtype must be bf16 or fp16 on A2/A3; Ascend 950 also allows hifloat8 / "
349+ "float8_e5m2 / float8_e4m3fn."),
350+ return ge::GRAPH_FAILED);
191 351 
192- const auto *keyDesc = context->GetInputDesc(MSA_IDX_KEY);
193- OP_CHECK_NULL_WITH_CONTEXT(context, keyDesc);
194- info.keyDtype = keyDesc->GetDataType();
195 info.isQuant = (info.keyDtype == ge::DT_INT8);352 info.isQuant = (info.keyDtype == ge::DT_INT8);
196 if (info.isQuant) {353 if (info.isQuant) {
197- OP_CHECK_IF(info.queryDtype != ge::DT_FLOAT16, OP_LOGE(context, "int8 key currently requires fp16 query."),354+ OP_CHECK_IF(info.queryDtype != ge::DT_FLOAT16, OP_LOGE(context, "ND int8 key currently requires fp16 query."),
198 return ge::GRAPH_FAILED);355 return ge::GRAPH_FAILED);
199 } else {356 } else {
200 OP_CHECK_IF(info.keyDtype != info.queryDtype, OP_LOGE(context, "non-quant key dtype must match query dtype."),357 OP_CHECK_IF(info.keyDtype != info.queryDtype, OP_LOGE(context, "non-quant key dtype must match query dtype."),
201 return ge::GRAPH_FAILED);358 return ge::GRAPH_FAILED);
359+ OP_CHECK_IF(IsFp8ComputeDtype(info.queryDtype) && !info.isAscend950,
360+ OP_LOGE(context, "FP8 query/key is only supported on Ascend 950."), return ge::GRAPH_FAILED);
202 }361 }
203 362 
204 const gert::StorageShape *scaleShape = context->GetOptionalInputShape(MSA_IDX_SCALE);363 const gert::StorageShape *scaleShape = context->GetOptionalInputShape(MSA_IDX_SCALE);
@@ -282,14 +441,30 @@ ge::graphStatus DoTiling(gert::TilingContext *context, const MsaIndexScoreInfo &
282 OP_CHECK_IF(aicNum == 0, OP_LOGE(context, "aic core num is 0."), return ge::GRAPH_FAILED);441 OP_CHECK_IF(aicNum == 0, OP_LOGE(context, "aic core num is 0."), return ge::GRAPH_FAILED);
283 442 
284 const uint32_t scoreBlockStride = RoundUpU32(info.maxBlocksPerBatch, MSA_SCORE_STRIDE_ALIGN);443 const uint32_t scoreBlockStride = RoundUpU32(info.maxBlocksPerBatch, MSA_SCORE_STRIDE_ALIGN);
285- // S workspace：非量化路径元素为 fp16（AIC fixpipe F322F16 直接写出），int8 路径 fp32。444+ // 950 C_to_UB：S 不进 GM；int8/TND 尾页仍要一页 K scratch。A2 与 950 回退（MSA_A5_USE_C2UB=0）走 8-stage S。
286- const uint32_t sWsBytes =445+ const uint32_t sElemBytes = info.isQuant ? sizeof(float) : sizeof(uint16_t);
287- aicNum * MSA_WORKSPACE_STAGES * MSA_STILE_ELEM_NUM * (info.isQuant ? sizeof(float) : sizeof(uint16_t));446+ const bool c2ub = info.isAscend950 && (MSA_A5_USE_C2UB != 0);
447+ const uint32_t sWsBytes = c2ub ? 0U : (aicNum * MSA_WORKSPACE_STAGES * MSA_STILE_ELEM_NUM * sElemBytes);
288 const bool useKScratch = info.isQuant || (info.keyLayout == MSA_KEY_LAYOUT_TND);448 const bool useKScratch = info.isQuant || (info.keyLayout == MSA_KEY_LAYOUT_TND);
289- const uint32_t kScratchBytes = useKScratch ? (aicNum * MSA_K_SCRATCH_ELEM_NUM * sizeof(uint16_t)) : 0U;449+ uint32_t scratchElemBytes = sizeof(uint16_t);
290- // K scratch 基址以 4B 为单位表达（kernel 内按 float* 折算），S 变 fp16 后偏移减半。450+ if (info.isAscend950 && !info.isQuant && IsFp8ComputeDtype(info.queryDtype)) {
451+ scratchElemBytes = 1U;
452+ }
453+ uint32_t kScratchElems = c2ub ? MSA_A5_K_SCRATCH_ELEM_NUM : MSA_K_SCRATCH_ELEM_NUM;
454+ // A2/A3 int8 四槽 K scratch；950 C2UB / 非量化 TND 尾页保持单槽。
455+ if (!info.isAscend950 && info.isQuant) {
456+ kScratchElems *= MSA_K_SCRATCH_STAGES_A2;
457+ }
458+ const uint32_t kScratchBytes = useKScratch ? (aicNum * kScratchElems * scratchElemBytes) : 0U;
291 const uint32_t kScratchOffsetElems = sWsBytes / sizeof(float);459 const uint32_t kScratchOffsetElems = sWsBytes / sizeof(float);
292 460 
461+ uint32_t launchAic = aicNum;
462+ if (info.totalQ == 0U) {
463+ launchAic = 1U;
464+ } else {
465+ launchAic = EstLaunchAic(info, aicNum);
466+ }
467+ 
293 MsaIndexScoreTilingData tilingData;468 MsaIndexScoreTilingData tilingData;
294 tilingData.set_batch(info.batch);469 tilingData.set_batch(info.batch);
295 tilingData.set_totalQ(info.totalQ);470 tilingData.set_totalQ(info.totalQ);
@@ -300,7 +475,7 @@ ge::graphStatus DoTiling(gert::TilingContext *context, const MsaIndexScoreInfo &
300 tilingData.set_blockSize(info.blockSize);475 tilingData.set_blockSize(info.blockSize);
301 tilingData.set_maxBlocksPerBatch(info.maxBlocksPerBatch);476 tilingData.set_maxBlocksPerBatch(info.maxBlocksPerBatch);
302 tilingData.set_scoreBlockStride(scoreBlockStride);477 tilingData.set_scoreBlockStride(scoreBlockStride);
303- tilingData.set_usedCoreNum(aicNum);478+ tilingData.set_usedCoreNum(launchAic);
304 tilingData.set_isQuant(info.isQuant ? 1U : 0U);479 tilingData.set_isQuant(info.isQuant ? 1U : 0U);
305 tilingData.set_sparseMode(info.sparseMode);480 tilingData.set_sparseMode(info.sparseMode);
306 tilingData.set_initBlocks(info.initBlocks);481 tilingData.set_initBlocks(info.initBlocks);
@@ -311,22 +486,28 @@ ge::graphStatus DoTiling(gert::TilingContext *context, const MsaIndexScoreInfo &
311 486 
312 tilingData.set_strideQt(info.numQHeads * info.headDim);487 tilingData.set_strideQt(info.numQHeads * info.headDim);
313 tilingData.set_strideQn(info.headDim);488 tilingData.set_strideQn(info.headDim);
489+ uint32_t defaultStrideKvBlock = 0;
314 if (info.keyLayout == MSA_KEY_LAYOUT_TND) {490 if (info.keyLayout == MSA_KEY_LAYOUT_TND) {
315- tilingData.set_strideKvBlock(0);491+ defaultStrideKvBlock = 0;
316 tilingData.set_strideKvToken(info.numKvHeads * info.headDim);492 tilingData.set_strideKvToken(info.numKvHeads * info.headDim);
317 tilingData.set_strideScalePage(info.numKvHeads);493 tilingData.set_strideScalePage(info.numKvHeads);
318 tilingData.set_strideScaleHead(1);494 tilingData.set_strideScaleHead(1);
319 } else if (info.keyLayout == MSA_KEY_LAYOUT_BNBD) {495 } else if (info.keyLayout == MSA_KEY_LAYOUT_BNBD) {
320- tilingData.set_strideKvBlock(info.numKvHeads * info.blockSize * info.headDim);496+ defaultStrideKvBlock = info.numKvHeads * info.blockSize * info.headDim;
321 tilingData.set_strideKvToken(info.headDim);497 tilingData.set_strideKvToken(info.headDim);
322 tilingData.set_strideScalePage(info.numKvHeads * info.blockSize);498 tilingData.set_strideScalePage(info.numKvHeads * info.blockSize);
323 tilingData.set_strideScaleHead(info.blockSize);499 tilingData.set_strideScaleHead(info.blockSize);
324 } else {500 } else {
325- tilingData.set_strideKvBlock(info.blockSize * info.numKvHeads * info.headDim);501+ defaultStrideKvBlock = info.blockSize * info.numKvHeads * info.headDim;
326 tilingData.set_strideKvToken(info.numKvHeads * info.headDim);502 tilingData.set_strideKvToken(info.numKvHeads * info.headDim);
327 tilingData.set_strideScalePage(info.numKvHeads * info.blockSize);503 tilingData.set_strideScalePage(info.numKvHeads * info.blockSize);
328 tilingData.set_strideScaleHead(info.blockSize);504 tilingData.set_strideScaleHead(info.blockSize);
329 }505 }
506+ uint32_t strideKvBlock = defaultStrideKvBlock;
507+ if (ResolveStrideKvBlock(context, info, defaultStrideKvBlock, strideKvBlock) != ge::GRAPH_SUCCESS) {
508+ return ge::GRAPH_FAILED;
509+ }
510+ tilingData.set_strideKvBlock(strideKvBlock);
330 tilingData.set_strideOutHead(info.totalQ * scoreBlockStride);511 tilingData.set_strideOutHead(info.totalQ * scoreBlockStride);
331 tilingData.set_strideOutToken(scoreBlockStride);512 tilingData.set_strideOutToken(scoreBlockStride);
332 tilingData.set_kScratchOffsetElems(kScratchOffsetElems);513 tilingData.set_kScratchOffsetElems(kScratchOffsetElems);
@@ -335,22 +516,45 @@ ge::graphStatus DoTiling(gert::TilingContext *context, const MsaIndexScoreInfo &
335 context->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());516 context->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());
336 517 
337 // MIX 1AIC:2AIV：CalcTschBlockDim 的 sliceNum 按 AIV 计数，内部再 / (aiv/aic)。518 // MIX 1AIC:2AIV：CalcTschBlockDim 的 sliceNum 按 AIV 计数，内部再 / (aiv/aic)。
338- // 传入 aicNum 会再除一次得到 blockDim=aic/2（910B3: 20→10），只能打一半 Cube。519+ // 传入 aicNum 会再除一次得到 blockDim=aic/2，只能打一半 Cube。
339- // 与 LightningIndexer 等 MIX 算子一致：sliceNum = aivNum → 910B3 blockDim=20。520+ // 整 batch q_len=0：queryS==0 → BlockDim=1，避免 totalTaskNum=0。
340- context->SetBlockDim(ascendcPlatform.CalcTschBlockDim(aivNum, aicNum, aivNum));521+ // sliceNum = launchAic * 2：短 decode 只起实际 M-task 对应的 MIX；
522+ // 多 M-tile / 大 batch 仍打满 AIC。
523+ if (info.totalQ == 0U) {
524+ context->SetBlockDim(1);
525+ } else {
526+ uint32_t sliceAiv = launchAic * MSA_AIV_PER_AIC;
527+ if (sliceAiv > aivNum) {
528+ sliceAiv = aivNum;
529+ }
530+ context->SetBlockDim(ascendcPlatform.CalcTschBlockDim(sliceAiv, aicNum, aivNum));
531+ }
341 532 
342 size_t *workspaces = context->GetWorkspaceSizes(1);533 size_t *workspaces = context->GetWorkspaceSizes(1);
343 OP_CHECK_NULL_WITH_CONTEXT(context, workspaces);534 OP_CHECK_NULL_WITH_CONTEXT(context, workspaces);
344 workspaces[0] =535 workspaces[0] =
345 ascendcPlatform.GetLibApiWorkSpaceSize() + static_cast<size_t>(sWsBytes) + static_cast<size_t>(kScratchBytes);536 ascendcPlatform.GetLibApiWorkSpaceSize() + static_cast<size_t>(sWsBytes) + static_cast<size_t>(kScratchBytes);
537+ // 非量化 PA 无 S GM / K scratch 时 user workspace 为 0，GetUserWorkspace 可能返回 nullptr。
538+ if (info.isAscend950 && sWsBytes == 0U && kScratchBytes == 0U) {
539+ workspaces[0] += 64U;
540+ }
346 541 
347 uint64_t tilingKey = MSA_TILING_KEY_FP16;542 uint64_t tilingKey = MSA_TILING_KEY_FP16;
348 if (info.isQuant) {543 if (info.isQuant) {
349 tilingKey = MSA_TILING_KEY_FP16_INT8;544 tilingKey = MSA_TILING_KEY_FP16_INT8;
545+ } else if (info.queryDtype == ge::DT_HIFLOAT8) {
546+ tilingKey = MSA_TILING_KEY_HIFLOAT8;
547+ } else if (info.queryDtype == ge::DT_FLOAT8_E5M2) {
548+ tilingKey = MSA_TILING_KEY_FP8_E5M2;
549+ } else if (info.queryDtype == ge::DT_FLOAT8_E4M3FN) {
550+ tilingKey = MSA_TILING_KEY_FP8_E4M3FN;
350 } else {551 } else {
351 tilingKey = (info.queryDtype == ge::DT_BF16) ? MSA_TILING_KEY_BF16 : MSA_TILING_KEY_FP16;552 tilingKey = (info.queryDtype == ge::DT_BF16) ? MSA_TILING_KEY_BF16 : MSA_TILING_KEY_FP16;
352 }553 }
353 context->SetTilingKey(tilingKey);554 context->SetTilingKey(tilingKey);
555+ OP_LOGI(context->GetNodeName(),
556+ "MsaIndexScore tilingKey=%lu headDim=%u layout=%u strideKvBlock=%u launchAic=%u aicNum=%u", tilingKey,
557+ info.headDim, info.keyLayout, strideKvBlock, launchAic, aicNum);
354 return ge::GRAPH_SUCCESS;558 return ge::GRAPH_SUCCESS;
355}559}
356} // namespace560} // namespace
Mattention/msa_index_score/op_host/msa_index_score_tiling.h+2-1
@@ -65,7 +65,7 @@ TILING_DATA_FIELD_DEF(uint32_t, strideOutHead)
65TILING_DATA_FIELD_DEF(uint32_t, strideOutToken)65TILING_DATA_FIELD_DEF(uint32_t, strideOutToken)
66TILING_DATA_FIELD_DEF(uint32_t, kScratchOffsetElems)66TILING_DATA_FIELD_DEF(uint32_t, kScratchOffsetElems)
67TILING_DATA_FIELD_DEF(uint32_t, keyLayout) // MSA_KEY_LAYOUT_BBND / BNBD / TND67TILING_DATA_FIELD_DEF(uint32_t, keyLayout) // MSA_KEY_LAYOUT_BBND / BNBD / TND
68-TILING_DATA_FIELD_DEF(uint32_t, totalK) // TND：key 第 0 维 T2；PA：068+TILING_DATA_FIELD_DEF(uint32_t, totalK) // TND：key 第 0 维 T2； PA：0
69END_TILING_DATA_DEF69END_TILING_DATA_DEF
70REGISTER_TILING_DATA_CLASS(MsaIndexScore, MsaIndexScoreTilingData)70REGISTER_TILING_DATA_CLASS(MsaIndexScore, MsaIndexScoreTilingData)
71 71 
@@ -87,6 +87,7 @@ struct MsaIndexScoreInfo {
87 uint32_t keyLayout = 0;87 uint32_t keyLayout = 0;
88 uint32_t totalK = 0;88 uint32_t totalK = 0;
89 bool isQuant = false;89 bool isQuant = false;
90+ bool isAscend950 = false;
90 ge::DataType queryDtype = ge::DT_BF16;91 ge::DataType queryDtype = ge::DT_BF16;
91 ge::DataType keyDtype = ge::DT_BF16;92 ge::DataType keyDtype = ge::DT_BF16;
92};93};
Mattention/msa_index_score/op_kernel/arch22/msa_block_mmad.h+2-3
@@ -64,9 +64,8 @@ public:
64 using L1BAlignHelper = Catlass::Gemm::helper::L1AlignHelper<ElementB, LayoutB>;64 using L1BAlignHelper = Catlass::Gemm::helper::L1AlignHelper<ElementB, LayoutB>;
65 65 
66 // K (B) 的 L1 流水级数。非量化 3 级：fixpipe 减半（fp16 S）后 MTE2 的逐页延迟66 // K (B) 的 L1 流水级数。非量化 3 级：fixpipe 减半（fp16 S）后 MTE2 的逐页延迟
67- // （~390ns，L2 命中）会成为新的关键路径，加深一级让 MTE2 提前两页预取。67+ // 会成为新的关键路径，加深一级让 MTE2 提前两页预取。
68- // int8 保持 2 级：K 源是每 stile 复用的 per-core scratch，AIV 下一 stile 的 cast68+ // A2/A3 int8 双槽 scratch 后同样 3 级：AIC 消费 slot[i] 时 AIV 只写 slot[i^1]。
69- // 会重写该区，更深的预取与 cast 重写产生竞态（实测 STAGES=3 下 int8 崩溃）。
70 // L1 占用：A 32KB + 3×32KB B = 128KB / 512KB。69 // L1 占用：A 32KB + 3×32KB B = 128KB / 512KB。
71 static constexpr uint32_t STAGES = STAGES_IN;70 static constexpr uint32_t STAGES = STAGES_IN;
72 // 非量化路径开 unit flag：mmad 与 fixpipe 的依赖交给硬件互锁，省掉每页71 // 非量化路径开 unit flag：mmad 与 fixpipe 的依赖交给硬件互锁，省掉每页
Mattention/msa_index_score/op_kernel/arch22/msa_index_score_kernel.h+218-78
@@ -11,11 +11,11 @@
11/*!11/*!
12 * \file msa_index_score_kernel.h12 * \file msa_index_score_kernel.h
13 * \brief MsaIndexScore Atlas A2/A3 核函数：AIC 做 paged QKᵀ，AIV 做（可选反量化）+ mask + 分段 RowMax。13 * \brief MsaIndexScore Atlas A2/A3 核函数：AIC 做 paged QKᵀ，AIV 做（可选反量化）+ mask + 分段 RowMax。
14- * 不支持 Ascend 950 / FP8。
15 *14 *
16 * 非量化：query/key 同为 bf16|fp16，score 侧不乘 scale。15 * 非量化：query/key 同为 bf16|fp16，score 侧不乘 scale。
17- * int8 量化：key 为 int8；AIV 将 page DataCopy+Cast 到 fp 暂存后通知 AIC 做 Mmad；16+ * int8 量化：key 为 int8；AIV 将 page DataCopy+Cast 到四槽 fp 暂存后通知 AIC 做 Mmad；
18- * AIV 在 mask 前按 scale[NP,N_kv,P] 对 S 做列乘完成 per-token 反量化。17+ * gather(st+4) 叠在 cube(st+1..st+3) 上；两 AIV 对分 8 page（每页仍串行 MTE2→V→MTE3）。
18+ * 同 AIV 上 MTE2∥MTE3 / 多 page 并发 MTE2 会写坏多 page int8。S 侧仍为 fp32 反量化。
19 */19 */
20 20 
21#ifndef MSA_INDEX_SCORE_KERNEL_H21#ifndef MSA_INDEX_SCORE_KERNEL_H
@@ -36,7 +36,7 @@
36#include "../msa_index_score_common.h"36#include "../msa_index_score_common.h"
37#include "msa_block_mmad.h"37#include "msa_block_mmad.h"
38#include "msa_index_score_task.h"38#include "msa_index_score_task.h"
39-#include "msa_index_score_epilogue.h"39+#include "msa_seg_row_max_epilogue.h"
40 40 
41namespace MsaIndexScoreNs {41namespace MsaIndexScoreNs {
42 42 
@@ -57,20 +57,23 @@ public:
57 using ElementS = std::conditional_t<IS_QUANT, float, half>;57 using ElementS = std::conditional_t<IS_QUANT, float, half>;
58 using SType = Catlass::Gemm::GemmType<ElementS, LayoutS>;58 using SType = Catlass::Gemm::GemmType<ElementS, LayoutS>;
59 59 
60- // int8 的 K 源是每 stile 复用的 per-core scratch，保持 2 级 L1B 流水（3 级实测竞态崩溃）；60+ // int8：三槽 K scratch 后与非量化同样 3 级 L1B + unit flag（单槽时 3 级会与下一 stile
61- // 同样出于该 scratch 的时序保守性，unit flag 只在非量化路径开启。61+ // 的 cast 重写竞态，950 arch35 仍保持 2 级 / 关 unit flag）。
62- using BlockMmad =62+ using BlockMmad = MsaBlockMmad<QType, KType, SType, MSA_L1B_STAGES_FP16, true>;
63- MsaBlockMmad<QType, KType, SType, IS_QUANT ? MSA_L1B_STAGES_INT8 : MSA_L1B_STAGES_FP16, !IS_QUANT>;
64 63 
65 static constexpr uint32_t REVERSE_DEPTH = MSA_WORKSPACE_STAGES - 1;64 static constexpr uint32_t REVERSE_DEPTH = MSA_WORKSPACE_STAGES - 1;
66 static constexpr uint32_t STILE_WIDTH = MSA_BLOCKS_PER_STILE * MSA_BLOCK_SIZE;65 static constexpr uint32_t STILE_WIDTH = MSA_BLOCKS_PER_STILE * MSA_BLOCK_SIZE;
67 static constexpr bool IS_QUANT_V = IS_QUANT;66 static constexpr bool IS_QUANT_V = IS_QUANT;
68- // 单 page cast 暂存：int8 + half 共用 epilogue 前半段 UB（S 尚未载入）。67+ static_assert(MSA_K_SCRATCH_STAGES_A2 == 4, "A2 int8 K scratch handshake assumes 4 slots");
69- static constexpr uint32_t UB_OFF_CAST_I8 = 0;68+ // int8：cast 放 epilogue 尾部。EVENT_ID2 每页成对 MTE2_V / V_MTE3 / MTE3_MTE2。
69+ // 非量化 TND 仍占用 epilogue S 区（gather 在 WaitS 之前，S 尚未载入）。
70+ static constexpr uint32_t UB_OFF_CAST_I8 = IS_QUANT ? MsaSegRowMaxEpilogue<true>::UB_OFF_TAIL : 0;
70 static constexpr uint32_t UB_SIZE_CAST_I8 = MSA_BLOCK_SIZE * MSA_K_TILE * sizeof(int8_t);71 static constexpr uint32_t UB_SIZE_CAST_I8 = MSA_BLOCK_SIZE * MSA_K_TILE * sizeof(int8_t);
71 static constexpr uint32_t UB_OFF_CAST_FP = UB_OFF_CAST_I8 + UB_SIZE_CAST_I8;72 static constexpr uint32_t UB_OFF_CAST_FP = UB_OFF_CAST_I8 + UB_SIZE_CAST_I8;
72 static constexpr uint32_t UB_SIZE_CAST_FP = MSA_BLOCK_SIZE * MSA_K_TILE * sizeof(half);73 static constexpr uint32_t UB_SIZE_CAST_FP = MSA_BLOCK_SIZE * MSA_K_TILE * sizeof(half);
73 static_assert(UB_OFF_CAST_FP + UB_SIZE_CAST_FP <= ArchTag::UB_SIZE, "cast UB out of bounds");74 static_assert(UB_OFF_CAST_FP + UB_SIZE_CAST_FP <= ArchTag::UB_SIZE, "cast UB out of bounds");
75+ static_assert(!IS_QUANT || UB_OFF_CAST_I8 >= MsaSegRowMaxEpilogue<true>::UB_TOTAL,
76+ "int8 cast UB overlaps epilogue");
74 77 
75 __aicore__ inline MsaIndexScoreKernel() {}78 __aicore__ inline MsaIndexScoreKernel() {}
76 79 
@@ -120,12 +123,16 @@ private:
120 Catlass::Arch::Resource<ArchTag> resource;123 Catlass::Arch::Resource<ArchTag> resource;
121 BlockMmad blockMmad(resource);124 BlockMmad blockMmad(resource);
122 Catlass::Arch::CrossCoreFlagWithReverse<REVERSE_DEPTH> flagSReady{MSA_FLAG_S_READY, MSA_FLAG_S_READY_REVERSE};125 Catlass::Arch::CrossCoreFlagWithReverse<REVERSE_DEPTH> flagSReady{MSA_FLAG_S_READY, MSA_FLAG_S_READY_REVERSE};
123- Catlass::Arch::CrossCoreFlag flagKReady{MSA_FLAG_K_READY};126+ Catlass::Arch::CrossCoreFlag flagK0{MSA_FLAG_K_READY};
127+ Catlass::Arch::CrossCoreFlag flagK1{MSA_FLAG_K_READY_REVERSE};
128+ Catlass::Arch::CrossCoreFlag flagK2{MSA_FLAG_K_READY_2};
129+ Catlass::Arch::CrossCoreFlag flagK3{MSA_FLAG_K_READY_3};
124 130 
125 const uint32_t coreIdx = AscendC::GetBlockIdx();131 const uint32_t coreIdx = AscendC::GetBlockIdx();
126 const uint32_t coreNum = AscendC::GetBlockNum();132 const uint32_t coreNum = AscendC::GetBlockNum();
127 const uint64_t coreWsBase = static_cast<uint64_t>(coreIdx) * MSA_WORKSPACE_STAGES * MSA_STILE_ELEM_NUM;133 const uint64_t coreWsBase = static_cast<uint64_t>(coreIdx) * MSA_WORKSPACE_STAGES * MSA_STILE_ELEM_NUM;
128- const uint64_t coreKScratch = static_cast<uint64_t>(coreIdx) * MSA_K_SCRATCH_ELEM_NUM;134+ const uint32_t kStages = IS_QUANT_V ? MSA_K_SCRATCH_STAGES_A2 : 1U;
135+ const uint64_t coreKScratch = static_cast<uint64_t>(coreIdx) * kStages * MSA_K_SCRATCH_ELEM_NUM;
129 const uint32_t totalTasks = scheduler_.TotalTasks();136 const uint32_t totalTasks = scheduler_.TotalTasks();
130 137 
131 uint32_t tileSeq = 0;138 uint32_t tileSeq = 0;
@@ -138,12 +145,13 @@ private:
138 const bool needKScratch = StileNeedsKScratch(task, st * MSA_BLOCKS_PER_STILE);145 const bool needKScratch = StileNeedsKScratch(task, st * MSA_BLOCKS_PER_STILE);
139 if (needKScratch) {146 if (needKScratch) {
140 // MIX 1AIC:2AIV：AIV→AIC 的 0x2 flag 需两个 AIV 都 Set 后才放行。147 // MIX 1AIC:2AIV：AIV→AIC 的 0x2 flag 需两个 AIV 都 Set 后才放行。
141- Catlass::Arch::CrossCoreWaitFlag(flagKReady);148+ // 绑 PIPE_MTE2，避免 PIPE_ALL 把上一 stile 的 cube/fixpipe 冲掉。
142- AscendC::PipeBarrier<PIPE_ALL>();149+ WaitKScratchReady(flagK0, flagK1, flagK2, flagK3, st);
143 }150 }
144 const uint64_t sBase =151 const uint64_t sBase =
145 coreWsBase + static_cast<uint64_t>(tileSeq % MSA_WORKSPACE_STAGES) * MSA_STILE_ELEM_NUM;152 coreWsBase + static_cast<uint64_t>(tileSeq % MSA_WORKSPACE_STAGES) * MSA_STILE_ELEM_NUM;
146- ComputeSTile(blockMmad, task, st * MSA_BLOCKS_PER_STILE, sBase, coreKScratch, needLoadQ);153+ ComputeSTile(blockMmad, task, st * MSA_BLOCKS_PER_STILE, sBase, KScratchSlot(coreKScratch, st),
154+ needLoadQ);
147 needLoadQ = false;155 needLoadQ = false;
148 Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(flagSReady);156 Catlass::Arch::CrossCoreSetFlagWithReverse<0x2, PIPE_FIX>(flagSReady);
149 ++tileSeq;157 ++tileSeq;
@@ -152,6 +160,37 @@ private:
152 AscendC::PipeBarrier<PIPE_ALL>();160 AscendC::PipeBarrier<PIPE_ALL>();
153 }161 }
154 162 
163+ __aicore__ inline uint64_t KScratchSlot(uint64_t coreKScratch, uint32_t st) const
164+ {
165+ if constexpr (IS_QUANT_V) {
166+ return coreKScratch + static_cast<uint64_t>(st % MSA_K_SCRATCH_STAGES_A2) * MSA_K_SCRATCH_ELEM_NUM;
167+ }
168+ return coreKScratch;
169+ }
170+ 
171+ __aicore__ inline void WaitKScratchReady(Catlass::Arch::CrossCoreFlag &flagK0, Catlass::Arch::CrossCoreFlag &flagK1,
172+ Catlass::Arch::CrossCoreFlag &flagK2, Catlass::Arch::CrossCoreFlag &flagK3,
173+ uint32_t st)
174+ {
175+ if constexpr (IS_QUANT_V) {
176+ const uint32_t slot = st % MSA_K_SCRATCH_STAGES_A2;
177+ if (slot == 0U) {
178+ Catlass::Arch::CrossCoreWaitFlag<0x2, PIPE_MTE2>(flagK0);
179+ } else if (slot == 1U) {
180+ Catlass::Arch::CrossCoreWaitFlag<0x2, PIPE_MTE2>(flagK1);
181+ } else if (slot == 2U) {
182+ Catlass::Arch::CrossCoreWaitFlag<0x2, PIPE_MTE2>(flagK2);
183+ } else {
184+ Catlass::Arch::CrossCoreWaitFlag<0x2, PIPE_MTE2>(flagK3);
185+ }
186+ } else {
187+ (void)flagK1;
188+ (void)flagK2;
189+ (void)flagK3;
190+ Catlass::Arch::CrossCoreWaitFlag<0x2, PIPE_MTE2>(flagK0);
191+ }
192+ }
193+ 
155 __aicore__ inline bool KeyBlockNeedsScratch(const MsaTask &task, uint32_t blk) const194 __aicore__ inline bool KeyBlockNeedsScratch(const MsaTask &task, uint32_t blk) const
156 {195 {
157 if constexpr (IS_QUANT_V) {196 if constexpr (IS_QUANT_V) {
@@ -229,20 +268,25 @@ private:
229 AscendC::DataCopyPad(ub, gm, params, pad);268 AscendC::DataCopyPad(ub, gm, params, pad);
230 }269 }
231 270 
232- /// AIV：把本 S-tile 可见 block 的 K 收进 per-core scratch（int8 cast 或 TND fp copy+尾填充）。271+ /// int8 gather：单 i8 + 单 fp16，每页 MTE2→V→MTE3 串行。subIdx/subBlockNum 对分 page。
233- __aicore__ inline void GatherKeySTileToScratch(Catlass::Arch::Resource<ArchTag> &resource, const MsaTask &task,272+ __aicore__ inline void GatherKeySTileQuantPiped(Catlass::Arch::Resource<ArchTag> &resource, const MsaTask &task,
234- uint32_t blkBase, uint64_t kScratchBase)273+ uint32_t blkBase, uint64_t kScratchBase, uint32_t subIdx,
274+ uint32_t subBlockNum)
235 {275 {
236 AscendC::LocalTensor<int8_t> ubI8 = resource.ubBuf.template GetBufferByByte<int8_t>(UB_OFF_CAST_I8);276 AscendC::LocalTensor<int8_t> ubI8 = resource.ubBuf.template GetBufferByByte<int8_t>(UB_OFF_CAST_I8);
237 AscendC::LocalTensor<half> ubHalf = resource.ubBuf.template GetBufferByByte<half>(UB_OFF_CAST_FP);277 AscendC::LocalTensor<half> ubHalf = resource.ubBuf.template GetBufferByByte<half>(UB_OFF_CAST_FP);
238 const uint32_t headDim = tiling_->headDim;278 const uint32_t headDim = tiling_->headDim;
239 const uint32_t nElem = MSA_BLOCK_SIZE * headDim;279 const uint32_t nElem = MSA_BLOCK_SIZE * headDim;
240 const uint32_t pageStride = MSA_BLOCK_SIZE * MSA_K_TILE;280 const uint32_t pageStride = MSA_BLOCK_SIZE * MSA_K_TILE;
281+ const uint32_t split = (subBlockNum == 0U) ? 1U : subBlockNum;
241 282 
242 for (uint32_t j = 0; j < MSA_BLOCKS_PER_STILE; ++j) {283 for (uint32_t j = 0; j < MSA_BLOCKS_PER_STILE; ++j) {
284+ if ((j % split) != subIdx) {
285+ continue;
286+ }
243 const uint32_t blk = blkBase + j;287 const uint32_t blk = blkBase + j;
244 if (blk >= task.visibleEndBlk) {288 if (blk >= task.visibleEndBlk) {
245- break;289+ continue;
246 }290 }
247 if (!KeyBlockNeedsScratch(task, blk)) {291 if (!KeyBlockNeedsScratch(task, blk)) {
248 continue;292 continue;
@@ -252,53 +296,84 @@ private:
252 const uint32_t nValidTok = KeyBlockValidTokens(task, blk);296 const uint32_t nValidTok = KeyBlockValidTokens(task, blk);
253 const uint32_t nValid = nValidTok * headDim;297 const uint32_t nValid = nValidTok * headDim;
254 298 
255- if constexpr (IS_QUANT_V) {299+ uint32_t copyN = nElem;
256- // 与已通过的 PA int8 路径对齐：整页 DataCopy+Cast（nElem），EVENT_ID0。300+ if (tiling_->keyLayout == MSA_KEY_LAYOUT_TND) {
257- // TND 仅在尾页会越过 T2 时改为部分拷贝；中间整页禁止走 CopyGmToUbPartial。301+ const uint32_t tokenStart = task.cuKStart + blk * tiling_->blockSize;
258- uint32_t copyN = nElem;302+ const uint32_t packedEnd = tokenStart + MSA_BLOCK_SIZE;
259- if (tiling_->keyLayout == MSA_KEY_LAYOUT_TND) {303+ if (tiling_->totalK > 0 && packedEnd > tiling_->totalK) {
260- const uint32_t tokenStart = task.cuKStart + blk * tiling_->blockSize;304+ copyN = nValid;
261- const uint32_t packedEnd = tokenStart + MSA_BLOCK_SIZE;
262- if (tiling_->totalK > 0 && packedEnd > tiling_->totalK) {
263- copyN = nValid;
264- }
265 }305 }
266- if (copyN < nElem) {
267- AscendC::Duplicate(ubHalf, static_cast<half>(0), nElem);
268- AscendC::PipeBarrier<PIPE_V>();
269- AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
270- AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
271- }
272- if (copyN > 0) {
273- if (copyN == nElem) {
274- AscendC::DataCopy(ubI8, gKeyInt8_[kOffset], nElem);
275- } else {
276- CopyGmToUbPartial(ubI8, gKeyInt8_[kOffset], copyN);
277- }
278- AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
279- AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
280- AscendC::Cast(ubHalf, ubI8, AscendC::RoundMode::CAST_NONE, copyN);
281- }
282- AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
283- AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
284- AscendC::DataCopy(gKeyScratch_[scratchOff], ubHalf, nElem);
285- } else {
286- AscendC::LocalTensor<ElementQ> ubK = resource.ubBuf.template GetBufferByByte<ElementQ>(UB_OFF_CAST_FP);
287- if (nValid < nElem) {
288- AscendC::Duplicate(ubK, static_cast<ElementQ>(0), nElem);
289- AscendC::PipeBarrier<PIPE_V>();
290- AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
291- AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
292- }
293- if (nValid > 0) {
294- CopyGmToUbPartial(ubK, gKeyFp_[kOffset], nValid);
295- AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
296- AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
297- }
298- AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);
299- AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);
300- AscendC::DataCopy(gKeyScratch_[scratchOff], ubK, nElem);
301 }306 }
307+ if (copyN < nElem) {
308+ AscendC::Duplicate(ubHalf, static_cast<half>(0), nElem);
309+ AscendC::PipeBarrier<PIPE_V>();
310+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
311+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
312+ }
313+ if (copyN > 0) {
314+ if (copyN == nElem) {
315+ AscendC::DataCopy(ubI8, gKeyInt8_[kOffset], nElem);
316+ } else {
317+ CopyGmToUbPartial(ubI8, gKeyInt8_[kOffset], copyN);
318+ }
319+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
320+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
321+ AscendC::Cast(ubHalf, ubI8, AscendC::RoundMode::CAST_NONE, copyN);
322+ }
323+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);
324+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);
325+ AscendC::DataCopy(gKeyScratch_[scratchOff], ubHalf, nElem);
326+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID2);
327+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID2);
328+ }
329+ }
330+ 
331+ /// AIV：把本 S-tile 可见 block 的 K 收进 per-core scratch（int8 cast 或 TND fp copy+尾填充）。
332+ /// int8：两 AIV 按 j%subBlockNum 对分 page；非量化仍由调用方只让 AIV0 进来。
333+ __aicore__ inline void GatherKeySTileToScratch(Catlass::Arch::Resource<ArchTag> &resource, const MsaTask &task,
334+ uint32_t blkBase, uint64_t kScratchBase, uint32_t subIdx,
335+ uint32_t subBlockNum)
336+ {
337+ if constexpr (IS_QUANT_V) {
338+ GatherKeySTileQuantPiped(resource, task, blkBase, kScratchBase, subIdx, subBlockNum);
339+ return;
340+ }
341+ AscendC::LocalTensor<ElementQ> ubK = resource.ubBuf.template GetBufferByByte<ElementQ>(UB_OFF_CAST_FP);
342+ const uint32_t headDim = tiling_->headDim;
343+ const uint32_t nElem = MSA_BLOCK_SIZE * headDim;
344+ const uint32_t pageStride = MSA_BLOCK_SIZE * MSA_K_TILE;
345+ const uint32_t split = (subBlockNum == 0U) ? 1U : subBlockNum;
346+ 
347+ for (uint32_t j = 0; j < MSA_BLOCKS_PER_STILE; ++j) {
348+ if ((j % split) != subIdx) {
349+ continue;
350+ }
351+ const uint32_t blk = blkBase + j;
352+ if (blk >= task.visibleEndBlk) {
353+ continue;
354+ }
355+ if (!KeyBlockNeedsScratch(task, blk)) {
356+ continue;
357+ }
358+ const uint64_t kOffset = KeyBlockGmOffset(task, blk);
359+ const uint64_t scratchOff = kScratchBase + static_cast<uint64_t>(j) * pageStride;
360+ const uint32_t nValidTok = KeyBlockValidTokens(task, blk);
361+ const uint32_t nValid = nValidTok * headDim;
362+ 
363+ if (nValid < nElem) {
364+ AscendC::Duplicate(ubK, static_cast<ElementQ>(0), nElem);
365+ AscendC::PipeBarrier<PIPE_V>();
366+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
367+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
368+ }
369+ if (nValid > 0) {
370+ CopyGmToUbPartial(ubK, gKeyFp_[kOffset], nValid);
371+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
372+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
373+ }
374+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);
375+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID2);
376+ AscendC::DataCopy(gKeyScratch_[scratchOff], ubK, nElem);
302 AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);377 AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
303 AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);378 AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
304 AscendC::PipeBarrier<PIPE_ALL>();379 AscendC::PipeBarrier<PIPE_ALL>();
@@ -306,6 +381,46 @@ private:
306 AscendC::PipeBarrier<PIPE_ALL>();381 AscendC::PipeBarrier<PIPE_ALL>();
307 }382 }
308 383 
384+ __aicore__ inline void NotifyKScratchReady(Catlass::Arch::CrossCoreFlag &flagK0,
385+ Catlass::Arch::CrossCoreFlag &flagK1,
386+ Catlass::Arch::CrossCoreFlag &flagK2,
387+ Catlass::Arch::CrossCoreFlag &flagK3, uint32_t st)
388+ {
389+ if constexpr (IS_QUANT_V) {
390+ const uint32_t slot = st % MSA_K_SCRATCH_STAGES_A2;
391+ if (slot == 0U) {
392+ Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flagK0);
393+ } else if (slot == 1U) {
394+ Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flagK1);
395+ } else if (slot == 2U) {
396+ Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flagK2);
397+ } else {
398+ Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flagK3);
399+ }
400+ } else {
401+ (void)flagK1;
402+ (void)flagK2;
403+ (void)flagK3;
404+ Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flagK0);
405+ }
406+ }
407+ 
408+ __aicore__ inline void IssueKGatherAndNotify(Catlass::Arch::Resource<ArchTag> &resource, const MsaTask &task,
409+ uint32_t st, uint64_t coreKScratch, uint32_t subIdx,
410+ uint32_t subBlockNum, Catlass::Arch::CrossCoreFlag &flagK0,
411+ Catlass::Arch::CrossCoreFlag &flagK1,
412+ Catlass::Arch::CrossCoreFlag &flagK2,
413+ Catlass::Arch::CrossCoreFlag &flagK3)
414+ {
415+ if constexpr (IS_QUANT_V) {
416+ GatherKeySTileToScratch(resource, task, st * MSA_BLOCKS_PER_STILE, KScratchSlot(coreKScratch, st), subIdx,
417+ subBlockNum);
418+ } else if (subIdx == 0U) {
419+ GatherKeySTileToScratch(resource, task, st * MSA_BLOCKS_PER_STILE, KScratchSlot(coreKScratch, st), 0U, 1U);
420+ }
421+ NotifyKScratchReady(flagK0, flagK1, flagK2, flagK3, st);
422+ }
423+ 
309 __aicore__ inline void ComputeSTile(BlockMmad &blockMmad, const MsaTask &task, uint32_t blkBase, uint64_t sBase,424 __aicore__ inline void ComputeSTile(BlockMmad &blockMmad, const MsaTask &task, uint32_t blkBase, uint64_t sBase,
310 uint64_t kScratchBase, bool needLoadL1)425 uint64_t kScratchBase, bool needLoadL1)
311 {426 {
@@ -345,14 +460,18 @@ private:
345 tiling_->maxBlocksPerBatch, tiling_->strideScalePage, tiling_->strideScaleHead, IS_QUANT_V,460 tiling_->maxBlocksPerBatch, tiling_->strideScalePage, tiling_->strideScaleHead, IS_QUANT_V,
346 tiling_->keyLayout, tiling_->totalK, gScore_, gScale_, gBlockTable_);461 tiling_->keyLayout, tiling_->totalK, gScore_, gScale_, gBlockTable_);
347 Catlass::Arch::CrossCoreFlagWithReverse<REVERSE_DEPTH> flagSReady{MSA_FLAG_S_READY, MSA_FLAG_S_READY_REVERSE};462 Catlass::Arch::CrossCoreFlagWithReverse<REVERSE_DEPTH> flagSReady{MSA_FLAG_S_READY, MSA_FLAG_S_READY_REVERSE};
348- Catlass::Arch::CrossCoreFlag flagKReady{MSA_FLAG_K_READY};463+ Catlass::Arch::CrossCoreFlag flagK0{MSA_FLAG_K_READY};
464+ Catlass::Arch::CrossCoreFlag flagK1{MSA_FLAG_K_READY_REVERSE};
465+ Catlass::Arch::CrossCoreFlag flagK2{MSA_FLAG_K_READY_2};
466+ Catlass::Arch::CrossCoreFlag flagK3{MSA_FLAG_K_READY_3};
349 467 
350 const uint32_t subBlockNum = AscendC::GetSubBlockNum();468 const uint32_t subBlockNum = AscendC::GetSubBlockNum();
351 const uint32_t subIdx = AscendC::GetSubBlockIdx();469 const uint32_t subIdx = AscendC::GetSubBlockIdx();
352 const uint32_t coreIdx = AscendC::GetBlockIdx() / subBlockNum;470 const uint32_t coreIdx = AscendC::GetBlockIdx() / subBlockNum;
353 const uint32_t coreNum = AscendC::GetBlockNum();471 const uint32_t coreNum = AscendC::GetBlockNum();
354 const uint64_t coreWsBase = static_cast<uint64_t>(coreIdx) * MSA_WORKSPACE_STAGES * MSA_STILE_ELEM_NUM;472 const uint64_t coreWsBase = static_cast<uint64_t>(coreIdx) * MSA_WORKSPACE_STAGES * MSA_STILE_ELEM_NUM;
355- const uint64_t coreKScratch = static_cast<uint64_t>(coreIdx) * MSA_K_SCRATCH_ELEM_NUM;473+ const uint32_t kStages = IS_QUANT_V ? MSA_K_SCRATCH_STAGES_A2 : 1U;
474+ const uint64_t coreKScratch = static_cast<uint64_t>(coreIdx) * kStages * MSA_K_SCRATCH_ELEM_NUM;
356 const uint32_t totalTasks = scheduler_.TotalTasks();475 const uint32_t totalTasks = scheduler_.TotalTasks();
357 476 
358 uint32_t tileSeq = 0;477 uint32_t tileSeq = 0;
@@ -360,20 +479,41 @@ private:
360 for (uint32_t taskIdx = coreIdx; taskIdx < totalTasks; taskIdx += coreNum) {479 for (uint32_t taskIdx = coreIdx; taskIdx < totalTasks; taskIdx += coreNum) {
361 scheduler_.Decode(taskIdx, task);480 scheduler_.Decode(taskIdx, task);
362 epilogue.BeginTask(task, subIdx, subBlockNum);481 epilogue.BeginTask(task, subIdx, subBlockNum);
363- for (uint32_t st = 0; st < task.numComputeSTiles; ++st) {482+ if constexpr (IS_QUANT_V) {
364- if (StileNeedsKScratch(task, st * MSA_BLOCKS_PER_STILE)) {483+ // 四槽：先填 0/1/2/3。WaitS(st) 后 slot[st%4] 空闲，gather(st+4) 叠在
365- // AIV0 写 scratch；AIV0/AIV1 barrier 后双方都 Set，满足 1:2 MIX 握手。484+ // cube(st+1..st+3) 上。两 AIV 对分 page，每页串行 MTE2→V→MTE3。
366- if (subIdx == 0) {485+ const uint32_t nSt = task.numComputeSTiles;
367- GatherKeySTileToScratch(resource, task, st * MSA_BLOCKS_PER_STILE, coreKScratch);486+ const uint32_t nPref = (nSt < MSA_K_SCRATCH_STAGES_A2) ? nSt : MSA_K_SCRATCH_STAGES_A2;
487+ for (uint32_t p = 0; p < nPref; ++p) {
488+ if (StileNeedsKScratch(task, p * MSA_BLOCKS_PER_STILE)) {
489+ IssueKGatherAndNotify(resource, task, p, coreKScratch, subIdx, subBlockNum, flagK0, flagK1,
490+ flagK2, flagK3);
368 }491 }
369- Catlass::Arch::CrossCoreBarrier<0x1, PIPE_MTE3>();
370- Catlass::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(flagKReady);
371 }492 }
372- Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE3>(flagSReady);493+ for (uint32_t st = 0; st < nSt; ++st) {
373- const uint64_t sBase =494+ Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE3>(flagSReady);
374- coreWsBase + static_cast<uint64_t>(tileSeq % MSA_WORKSPACE_STAGES) * MSA_STILE_ELEM_NUM;495+ const uint32_t stPref = st + MSA_K_SCRATCH_STAGES_A2;
375- epilogue.ProcessSTile(gWorkspace_[sBase], task, st * MSA_BLOCKS_PER_STILE);496+ const uint64_t sBase =
376- ++tileSeq;497+ coreWsBase + static_cast<uint64_t>(tileSeq % MSA_WORKSPACE_STAGES) * MSA_STILE_ELEM_NUM;
498+ if (stPref < nSt && StileNeedsKScratch(task, stPref * MSA_BLOCKS_PER_STILE)) {
499+ IssueKGatherAndNotify(resource, task, stPref, coreKScratch, subIdx, subBlockNum, flagK0, flagK1,
500+ flagK2, flagK3);
501+ }
502+ epilogue.ProcessSTile(gWorkspace_[sBase], task, st * MSA_BLOCKS_PER_STILE);
503+ ++tileSeq;
504+ }
505+ } else {
506+ for (uint32_t st = 0; st < task.numComputeSTiles; ++st) {
507+ if (StileNeedsKScratch(task, st * MSA_BLOCKS_PER_STILE)) {
508+ IssueKGatherAndNotify(resource, task, st, coreKScratch, subIdx, subBlockNum, flagK0, flagK1,
509+ flagK2, flagK3);
510+ }
511+ Catlass::Arch::CrossCoreWaitFlagWithReverse<0x2, PIPE_MTE3>(flagSReady);
512+ const uint64_t sBase =
513+ coreWsBase + static_cast<uint64_t>(tileSeq % MSA_WORKSPACE_STAGES) * MSA_STILE_ELEM_NUM;
514+ epilogue.ProcessSTile(gWorkspace_[sBase], task, st * MSA_BLOCKS_PER_STILE);
515+ ++tileSeq;
516+ }
377 }517 }
378 // 因果不可见尾：不握手，直接写 -inf（与 AIC 跳过这些 tile 对齐）。518 // 因果不可见尾：不握手，直接写 -inf（与 AIC 跳过这些 tile 对齐）。
379 for (uint32_t st = task.numComputeSTiles; st < task.numSTiles; ++st) {519 for (uint32_t st = task.numComputeSTiles; st < task.numSTiles; ++st) {
Mattention/msa_index_score/op_kernel/arch22/msa_index_score_task.h+12-1
@@ -33,6 +33,17 @@ __aicore__ inline uint32_t MsaCeilDiv(uint32_t a, uint32_t b)
33 return (a + b - 1) / b;33 return (a + b - 1) / b;
34}34}
35 35 
36+// C2UB 非量化 F322F16 单目标：两次 Fixpipe 的 mSize 切分（每段偶数），AIV 必须相同。
37+__aicore__ inline void MsaC2UbHalfSplit(uint32_t mActual, uint32_t &m0, uint32_t &m1)
38+{
39+ const uint32_t mEven = (mActual + 1U) & ~1U;
40+ m0 = mEven / 2U;
41+ if ((m0 & 1U) != 0U) {
42+ m0 += 1U;
43+ }
44+ m1 = mEven - m0;
45+}
46+ 
36__aicore__ inline uint32_t MsaMinU32(uint32_t a, uint32_t b)47__aicore__ inline uint32_t MsaMinU32(uint32_t a, uint32_t b)
37{48{
38 return a < b ? a : b;49 return a < b ? a : b;
@@ -144,7 +155,7 @@ public:
144 task.mStart = (taskIdx - taskBase_) * MSA_ROW_TILE_M;155 task.mStart = (taskIdx - taskBase_) * MSA_ROW_TILE_M;
145 task.mActual = MsaMinU32(MSA_ROW_TILE_M, curRows_ - task.mStart);156 task.mActual = MsaMinU32(MSA_ROW_TILE_M, curRows_ - task.mStart);
146 // Cube L0A fractal 为 MSA_M_ALIGN 行。整请求不足对齐宽度时保持原样（L0 已覆盖）；157 // Cube L0A fractal 为 MSA_M_ALIGN 行。整请求不足对齐宽度时保持原样（L0 已覆盖）；
147- // 否则把末尾短 tile 向前重叠到 MSA_M_ALIGN 行，避免 mActual∈(0,16) 的 int8 路径出错。158+ // 否则把末尾短 tile 向前重叠到 MSA_M_ALIGN 行，避免 mActual ∈ (0,16) 的 int8 路径出错。
148 if (task.mActual > 0U && task.mActual < MSA_M_ALIGN && task.mStart >= MSA_M_ALIGN) {159 if (task.mActual > 0U && task.mActual < MSA_M_ALIGN && task.mStart >= MSA_M_ALIGN) {
149 task.mStart -= (MSA_M_ALIGN - task.mActual);160 task.mStart -= (MSA_M_ALIGN - task.mActual);
150 task.mActual = MSA_M_ALIGN;161 task.mActual = MSA_M_ALIGN;
Rattention/msa_index_score/op_kernel/arch22/msa_index_score_epilogue.h→attention/msa_index_score/op_kernel/arch22/msa_seg_row_max_epilogue.h+48-14
@@ -9,7 +9,7 @@
9 */9 */
10 10 
11/*!11/*!
12- * \file msa_index_score_epilogue.h12+ * \file msa_seg_row_max_epilogue.h
13 * \brief AIV 侧 Epilogue：可选反量化列乘 -> atten_mask -> 分段 RowMax -> local_mask -> 写回。13 * \brief AIV 侧 Epilogue：可选反量化列乘 -> atten_mask -> 分段 RowMax -> local_mask -> 写回。
14 *14 *
15 * 数值路径：15 * 数值路径：
@@ -18,8 +18,8 @@
18 * local_mask 由 start_loc（query 所在逻辑 block）+ init/local_blocks 生成强制 +∞。18 * local_mask 由 start_loc（query 所在逻辑 block）+ init/local_blocks 生成强制 +∞。
19 */19 */
20 20 
21-#ifndef MSA_INDEX_SCORE_EPILOGUE_H21+#ifndef MSA_SEG_ROW_MAX_EPILOGUE_H
22-#define MSA_INDEX_SCORE_EPILOGUE_H22+#define MSA_SEG_ROW_MAX_EPILOGUE_H
23 23 
24#include "kernel_operator.h"24#include "kernel_operator.h"
25 25 
@@ -40,8 +40,9 @@ public:
40 40 
41 // score 暂存容量：1 AIC : 2 AIV 下每个 subcore 最多拿到半个 M-tile 的行；41 // score 暂存容量：1 AIC : 2 AIV 下每个 subcore 最多拿到半个 M-tile 的行；
42 // 一次 flush 覆盖 MSA_STAGE_BLOCKS 个 block（= 每行 1KB 连续写回）。42 // 一次 flush 覆盖 MSA_STAGE_BLOCKS 个 block（= 每行 1KB 连续写回）。
43+ // int8：缩到 160，给 UB 尾部腾出整页 cast（48KB）；非量化保持 256。
43 static constexpr uint32_t MSA_STAGE_ROWS = MSA_ROW_TILE_M / MSA_AIV_PER_AIC;44 static constexpr uint32_t MSA_STAGE_ROWS = MSA_ROW_TILE_M / MSA_AIV_PER_AIC;
44- static constexpr uint32_t MSA_STAGE_BLOCKS = 256;45+ static constexpr uint32_t MSA_STAGE_BLOCKS = IS_QUANT ? 160 : 256;
45 static_assert(MSA_STAGE_BLOCKS % MSA_BLOCKS_PER_STILE == 0, "stage window must hold whole S-tiles");46 static_assert(MSA_STAGE_BLOCKS % MSA_BLOCKS_PER_STILE == 0, "stage window must hold whole S-tiles");
46 47 
47 // UB 手工布局（字节偏移），catlass 不提供 UB allocator。48 // UB 手工布局（字节偏移），catlass 不提供 UB allocator。
@@ -59,18 +60,19 @@ public:
59 // 非量化：WholeReduceMax 的 fp16 输出（每 pass MSA_REDUCE_ROWS 个 half）。60 // 非量化：WholeReduceMax 的 fp16 输出（每 pass MSA_REDUCE_ROWS 个 half）。
60 static constexpr uint32_t UB_OFF_RED16 = UB_OFF_DEQ + UB_SIZE_DEQ;61 static constexpr uint32_t UB_OFF_RED16 = UB_OFF_DEQ + UB_SIZE_DEQ;
61 static constexpr uint32_t UB_SIZE_RED16 = MSA_REDUCE_ROWS * sizeof(half);62 static constexpr uint32_t UB_SIZE_RED16 = MSA_REDUCE_ROWS * sizeof(half);
62- // score 暂存：本 subcore 的行 × 一次 flush 覆盖的 block 数。逐 pass 只写 32B 到 GM63+ // score 暂存：本 subcore 的行 × 一次 flush 覆盖的 block 数。
63- // 会把 MTE3 压到 ~9GB/s；改为在 UB 内按行累积、整行（≥1KB）一次写回。64+ // 在 UB 内按行累积、整行（≥1KB）一次写回。
64 static constexpr uint32_t UB_OFF_STAGE =65 static constexpr uint32_t UB_OFF_STAGE =
65 ((UB_OFF_RED16 + UB_SIZE_RED16 + MSA_UB_ALIGN_BYTES - 1U) / MSA_UB_ALIGN_BYTES) * MSA_UB_ALIGN_BYTES;66 ((UB_OFF_RED16 + UB_SIZE_RED16 + MSA_UB_ALIGN_BYTES - 1U) / MSA_UB_ALIGN_BYTES) * MSA_UB_ALIGN_BYTES;
66 static constexpr uint32_t UB_SIZE_STAGE = MSA_STAGE_ROWS * MSA_STAGE_BLOCKS * sizeof(float);67 static constexpr uint32_t UB_SIZE_STAGE = MSA_STAGE_ROWS * MSA_STAGE_BLOCKS * sizeof(float);
67 // 非量化：S16（fp16 载入缓冲）复用量化路径的 fp32 S 区（64KB = 两级 32KB 乒乓）。68 // 非量化：S16（fp16 载入缓冲）复用量化路径的 fp32 S 区（64KB = 两级 32KB 乒乓）。
68- // v0.7 PipeUtilization：aiv_mte2_wait_ratio≈0.84、aiv_mte2_ratio≈0.33、aiv_vec_ratio≈0.57，
69 // 说明 GM→UB 与归约串行；两级缓冲让下一 pass 的 MTE2 叠在当前 pass 的 V 上。69 // 说明 GM→UB 与归约串行；两级缓冲让下一 pass 的 MTE2 叠在当前 pass 的 V 上。
70 static constexpr uint32_t UB_OFF_S16 = UB_OFF_S;70 static constexpr uint32_t UB_OFF_S16 = UB_OFF_S;
71 static constexpr uint32_t UB_SIZE_S16 = MSA_ROWS_PER_PASS * MSA_BLOCKS_PER_STILE * MSA_BLOCK_SIZE * sizeof(half);71 static constexpr uint32_t UB_SIZE_S16 = MSA_ROWS_PER_PASS * MSA_BLOCKS_PER_STILE * MSA_BLOCK_SIZE * sizeof(half);
72 static constexpr uint32_t S16_STAGES = 2;72 static constexpr uint32_t S16_STAGES = 2;
73 static constexpr uint32_t UB_TOTAL = UB_OFF_STAGE + UB_SIZE_STAGE;73 static constexpr uint32_t UB_TOTAL = UB_OFF_STAGE + UB_SIZE_STAGE;
74+ static constexpr uint32_t UB_OFF_TAIL =
75+ ((UB_TOTAL + MSA_UB_ALIGN_BYTES - 1U) / MSA_UB_ALIGN_BYTES) * MSA_UB_ALIGN_BYTES;
74 static_assert(UB_TOTAL <= ArchTag::UB_SIZE, "MsaSegRowMaxEpilogue UB out of bounds");76 static_assert(UB_TOTAL <= ArchTag::UB_SIZE, "MsaSegRowMaxEpilogue UB out of bounds");
75 static_assert(UB_OFF_S16 + S16_STAGES * UB_SIZE_S16 <= UB_OFF_T64, "S16 ping-pong overlaps T64");77 static_assert(UB_OFF_S16 + S16_STAGES * UB_SIZE_S16 <= UB_OFF_T64, "S16 ping-pong overlaps T64");
76 78 
@@ -149,6 +151,42 @@ public:
149 }151 }
150 }152 }
151 153 
154+ __aicore__ inline uint32_t SubRows() const
155+ {
156+ return mSub_;
157+ }
158+ 
159+ __aicore__ inline uint32_t SubOff() const
160+ {
161+ return mOff_;
162+ }
163+ 
164+ __aicore__ inline void BeginSTile(uint32_t blkBase)
165+ {
166+ if (mSub_ == 0U) {
167+ return;
168+ }
169+ if (stageOn_ && (blkBase >= stageBlkBase_ + MSA_STAGE_BLOCKS)) {
170+ FlushStage();
171+ stageBlkBase_ = blkBase;
172+ stageBlkEnd_ = blkBase;
173+ }
174+ }
175+ 
176+ __aicore__ inline void FinishSTile(uint32_t blkBase)
177+ {
178+ if (mSub_ == 0U) {
179+ return;
180+ }
181+ stageBlkEnd_ = blkBase + MSA_BLOCKS_PER_STILE;
182+ }
183+ 
184+ __aicore__ inline void ProcessOnePass(const AscendC::GlobalTensor<ElementS> &gS, const MsaTask &task,
185+ uint32_t blkBase, uint32_t rowOff, uint32_t rows)
186+ {
187+ ProcessPass(gS, task, blkBase, rowOff, rows);
188+ }
189+ 
152 /// 处理一个 S tile（MSA_BLOCKS_PER_STILE 个 sparse block）中属于本 subcore 的行。190 /// 处理一个 S tile（MSA_BLOCKS_PER_STILE 个 sparse block）中属于本 subcore 的行。
153 __aicore__ inline void ProcessSTile(const AscendC::GlobalTensor<ElementS> &gS, const MsaTask &task,191 __aicore__ inline void ProcessSTile(const AscendC::GlobalTensor<ElementS> &gS, const MsaTask &task,
154 uint32_t blkBase)192 uint32_t blkBase)
@@ -158,11 +196,7 @@ public:
158 }196 }
159 const uint32_t mOff = mOff_;197 const uint32_t mOff = mOff_;
160 const uint32_t mSub = mSub_;198 const uint32_t mSub = mSub_;
161- if (stageOn_ && (blkBase >= stageBlkBase_ + MSA_STAGE_BLOCKS)) {199+ BeginSTile(blkBase);
162- FlushStage();
163- stageBlkBase_ = blkBase;
164- stageBlkEnd_ = blkBase;
165- }
166 200 
167 if constexpr (IS_QUANT) {201 if constexpr (IS_QUANT) {
168 for (uint32_t p0 = 0; p0 < mSub; p0 += MSA_ROWS_PER_PASS) {202 for (uint32_t p0 = 0; p0 < mSub; p0 += MSA_ROWS_PER_PASS) {
@@ -202,7 +236,7 @@ public:
202 ReleaseS16(prevStage);236 ReleaseS16(prevStage);
203 }237 }
204 }238 }
205- stageBlkEnd_ = blkBase + MSA_BLOCKS_PER_STILE;239+ FinishSTile(blkBase);
206 }240 }
207 241 
208private:242private:
@@ -668,4 +702,4 @@ private:
668 702 
669} // namespace MsaIndexScoreNs703} // namespace MsaIndexScoreNs
670 704 
671-#endif // MSA_INDEX_SCORE_EPILOGUE_H705+#endif // MSA_SEG_ROW_MAX_EPILOGUE_H
Aattention/msa_index_score/op_kernel/arch35/msa_block_mmad.h+380-0
@@ -0,0 +1,380 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file msa_block_mmad.h
13+ * \brief Atlas A5 BlockMmad：Q 驻留 + K pingpong；C_to_UB 时 L0C Fixpipe dual-dst → AIV UB。
14+ *
15+ * 基于 Catlass FullLoadA：同一 M-tile 只搬一次 Q；相邻 K page 的 FIXPIPE 与 Cube 重叠。
16+ * Atlas A5：计算路径照搬 A2；L0C=256KB，两块 128x128 fp32 C 各 64KB。
17+ * C_to_UB：int8 dual-dst 写 fp32。非量化 dual-dst 禁止 quantPre!=NoQuant（CANN 断言）；
18+ * 改为两次 dualDstCtl=0 + subBlockId 的 F322F16，写出 packed [M/2,128] fp16（对齐 FA）。
19+ */
20+ 
21+#ifndef MSA_BLOCK_MMAD_ARCH35_H
22+#define MSA_BLOCK_MMAD_ARCH35_H
23+ 
24+#include <type_traits>
25+ 
26+#include "catlass/msa_catlass.hpp"
27+#include "catlass/arch/msa_arch.hpp"
28+#include "catlass/arch/msa_resource.hpp"
29+#include "catlass/msa_coord.hpp"
30+#include "catlass/msa_matrix_coord.hpp"
31+#include "catlass/msa_gemm_coord.hpp"
32+#include "catlass/layout/msa_layout.hpp"
33+#include "catlass/gemm/msa_gemm_helper.hpp"
34+#include "catlass/gemm/tile/msa_tile_copy.hpp"
35+#include "catlass/gemm/tile/msa_tile_mmad.hpp"
36+ 
37+#include "../msa_index_score_common.h"
38+#include "../arch22/msa_index_score_task.h"
39+ 
40+namespace MsaIndexScoreNs {
41+ 
42+template <class AType_, class BType_, class CType_, uint32_t STAGES_IN = MSA_L1B_STAGES_FP16,
43+ bool USE_UNIT_FLAG_ = false>
44+class MsaBlockMmad {
45+public:
46+ using ArchTag = Catlass::Arch::AtlasA5;
47+ using ElementA = typename AType_::Element;
48+ using LayoutA = typename AType_::Layout;
49+ using ElementB = typename BType_::Element;
50+ using LayoutB = typename BType_::Layout;
51+ using ElementC = typename CType_::Element;
52+ using LayoutC = typename CType_::Layout;
53+ using BiasType = void;
54+ using TileCopy = Catlass::Gemm::Tile::TileCopy<ArchTag, AType_, BType_, CType_, BiasType>;
55+ using TileMmad = Catlass::Gemm::Tile::TileMmad<ArchTag, AType_, BType_, BiasType>;
56+ using CopyGmToL1A = typename TileCopy::CopyGmToL1A;
57+ using CopyGmToL1B = typename TileCopy::CopyGmToL1B;
58+ using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A;
59+ using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B;
60+ using CopyL0CToGm = typename TileCopy::CopyL0CToGm;
61+ using ElementAccumulator =
62+ typename Catlass::Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
63+ using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc;
64+ using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc;
65+ using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst;
66+ using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst;
67+ using LayoutCInL0 = Catlass::layout::zN;
68+ using L1AAlignHelper = Catlass::Gemm::helper::L1AlignHelper<ElementA, LayoutA>;
69+ using L1BAlignHelper = Catlass::Gemm::helper::L1AlignHelper<ElementB, LayoutB>;
70+ 
71+ // K (B) 的 L1 流水级数。非量化 3 级：fixpipe 减半（fp16 S）后 MTE2 的逐页延迟
72+ // （~390ns，L2 命中）会成为新的关键路径，加深一级让 MTE2 提前两页预取。
73+ // int8 保持 2 级：K 源是每 stile 复用的 per-core scratch，AIV 下一 stile 的 cast
74+ // 会重写该区，更深的预取与 cast 重写产生竞态（实测 STAGES=3 下 int8 崩溃）。
75+ // L1 占用：A 32KB + 3×32KB B = 128KB / 512KB。
76+ static constexpr uint32_t STAGES = STAGES_IN;
77+ // 非量化路径开 unit flag：mmad 与 fixpipe 的依赖交给硬件互锁，省掉每页
78+ // M_FIX/FIX_M 四次标量 set/wait。标量因此能在 cube 完成前就发出下一页的
79+ // MTE2/MTE1，把 GM→L1 的搬运藏到 cube/fixpipe 后面。
80+ static constexpr bool USE_UNIT_FLAG = USE_UNIT_FLAG_;
81+ static constexpr uint32_t L0C_STAGES = 2;
82+ static constexpr uint32_t L1_M = MSA_ROW_TILE_M;
83+ static constexpr uint32_t L1_N = MSA_BLOCK_SIZE;
84+ static constexpr uint32_t L1_K = MSA_K_TILE;
85+ static constexpr uint32_t L1A_SIZE = L1_M * L1_K * sizeof(ElementA);
86+ static constexpr uint32_t L1B_SIZE = L1_N * L1_K * sizeof(ElementB);
87+ // L0A/L0B 各仅 64KB，B tile 32KB → 乒乓固定 2 级（32KB/级）；与 L1B STAGES 解耦。
88+ static constexpr uint32_t L0A_STAGES = 2;
89+ static constexpr uint32_t L0A_RESIDENT_ID = 0;
90+ static constexpr uint32_t L0A_FALLBACK_ID = 1;
91+ static constexpr uint32_t L0B_STAGES = 2;
92+ static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / L0A_STAGES;
93+ static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / L0B_STAGES;
94+ static constexpr uint32_t L0C_TILE_SIZE = L1_M * L1_N * sizeof(ElementAccumulator);
95+ 
96+ static_assert(std::is_same_v<LayoutC, Catlass::layout::RowMajor>, "LayoutC must be RowMajor (fixpipe nz2nd)");
97+ // A2 的 L0A/L0B 各仅 64KB：A/B tile 32KB，乒乓最多 2 级（32KB/级）。
98+ static_assert(L0A_PINGPONG_BUF_SIZE >= L1A_SIZE, "L0A pingpong buf smaller than A tile");
99+ static_assert(L0B_PINGPONG_BUF_SIZE >= L1B_SIZE, "L0B pingpong buf smaller than B tile");
100+ static_assert(L0C_TILE_SIZE * L0C_STAGES <= ArchTag::L0C_SIZE, "L0C pingpong exceeds L0C");
101+ static_assert(L1A_SIZE + L1B_SIZE * STAGES <= ArchTag::L1_SIZE, "L1 A+B exceeds L1");
102+ 
103+ __aicore__ inline MsaBlockMmad(Catlass::Arch::Resource<ArchTag> &resource)
104+ {
105+ l1ATensor_ = resource.l1Buf.template GetBufferByByte<ElementA>(0);
106+ const uint32_t l1BOffset = L1A_SIZE;
107+ for (uint32_t i = 0; i < STAGES; ++i) {
108+ l1BTensorList_[i] = resource.l1Buf.template GetBufferByByte<ElementB>(l1BOffset + L1B_SIZE * i);
109+ l1BEventList_[i] = static_cast<int32_t>(i + STAGES);
110+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[i]);
111+ }
112+ // L0A 槽 0 专供常驻 Q：不预置 M_MTE1，改用 LoadA 内「Set+Wait 成对」自平衡，
113+ // 避免计数型 flag 让 MTE1 用构造时的余额提前放行、覆盖仍被 mmad 读取的 Q。
114+ // 槽 1 供 kTileCount>1 的回退路径（每 k-tile 重搬 A），沿用预置 + 析构等待。
115+ for (uint32_t i = 0; i < L0A_STAGES; ++i) {
116+ l0ATensorList_[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_PINGPONG_BUF_SIZE * i);
117+ l0AEventList_[i] = static_cast<int32_t>(i);
118+ }
119+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[L0A_FALLBACK_ID]);
120+ for (uint32_t i = 0; i < L0B_STAGES; ++i) {
121+ l0BTensorList_[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_PINGPONG_BUF_SIZE * i);
122+ l0BEventList_[i] = static_cast<int32_t>(i + STAGES + L0A_STAGES);
123+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList_[i]);
124+ }
125+ for (uint32_t i = 0; i < L0C_STAGES; ++i) {
126+ l0CTensorList_[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(L0C_TILE_SIZE * i);
127+ l0CEventList_[i] = static_cast<int32_t>(i);
128+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[i]);
129+ }
130+ }
131+ 
132+ __aicore__ inline ~MsaBlockMmad()
133+ {
134+ for (uint32_t i = 0; i < STAGES; ++i) {
135+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[i]);
136+ }
137+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[L0A_FALLBACK_ID]);
138+ for (uint32_t i = 0; i < L0B_STAGES; ++i) {
139+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList_[i]);
140+ }
141+ for (uint32_t i = 0; i < L0C_STAGES; ++i) {
142+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[i]);
143+ }
144+ }
145+ 
146+ __aicore__ inline void operator()(const AscendC::GlobalTensor<ElementA> &gmA, const LayoutA &layoutA,
147+ const AscendC::GlobalTensor<ElementB> &gmB, const LayoutB &layoutB,
148+ const AscendC::GlobalTensor<ElementC> &gmC, const LayoutC &layoutC,
149+ const Catlass::GemmCoord &actualShape, bool needLoadL1)
150+ {
151+ const uint32_t mRound = RoundUp<L1AAlignHelper::M_ALIGNED>(actualShape.m());
152+ const uint32_t nRound = RoundUp<L1BAlignHelper::N_ALIGNED>(actualShape.n());
153+ auto layoutAInL1 = LayoutAInL1::template MakeLayout<ElementA>(L1_M, actualShape.k());
154+ auto layoutBInL1 = LayoutBInL1::template MakeLayout<ElementB>(L1_K, L1_N);
155+ auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(Catlass::MakeCoord(mRound, nRound));
156+ uint32_t kActual = Min(actualShape.k(), L1_K);
157+ 
158+ const uint32_t kTileCount = CeilDiv<L1_K>(actualShape.k());
159+ const uint32_t mPartLoop = CeilDiv<L1_M>(mRound);
160+ const uint32_t nPartLoop = CeilDiv<L1_N>(nRound);
161+ // headDim <= L1_K（实际场景恒成立）时整块 Q 一次进 L0A，可跨 page 复用；
162+ // 否则退回逐 page 重搬 A 的通用路径。
163+ const bool qResident = (kTileCount == 1U) && (mPartLoop == 1U) && (CeilDiv<L1_K>(kActual) == 1U);
164+ 
165+ if (needLoadL1) {
166+ auto layoutTileA = layoutA.GetTileLayout(Catlass::MakeCoord(actualShape.m(), actualShape.k()));
167+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEvent_);
168+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEvent_);
169+ copyGmToL1A_(l1ATensor_, gmA, layoutAInL1, layoutTileA);
170+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEvent_);
171+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEvent_);
172+ if (qResident) {
173+ // 本 M-tile 的 Q 一次搬进 L0A 并驻留：省掉每页 32KB 的 L1→L0A
174+ // 搬运与两次标量同步。先让 MTE1 等上一 M-tile 的 mmad 读完。
175+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[L0A_RESIDENT_ID]);
176+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[L0A_RESIDENT_ID]);
177+ LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout<ElementA>(mRound, kActual);
178+ copyL1ToL0A_(l0ATensorList_[L0A_RESIDENT_ID], l1ATensor_, layoutAInL0, layoutAInL1);
179+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0AEventList_[L0A_RESIDENT_ID]);
180+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0AEventList_[L0A_RESIDENT_ID]);
181+ }
182+ }
183+ 
184+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[l1ListId_]);
185+ auto layoutTileB = layoutB.GetTileLayout(Catlass::MakeCoord(kActual, actualShape.n()));
186+ copyGmToL1B_(l1BTensorList_[l1ListId_], gmB, layoutBInL1, layoutTileB);
187+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList_[l1ListId_]);
188+ 
189+ // 等本槽位上一轮 FIXPIPE 完成后再写 L0C，从而与另一槽位的 FIXPIPE 重叠。
190+ // unit flag 下该依赖由硬件维护，省掉这次标量等待。
191+ if constexpr (!USE_UNIT_FLAG) {
192+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[l0CListId_]);
193+ }
194+ 
195+ for (uint32_t kLoopIdx = 0; kLoopIdx < kTileCount; ++kLoopIdx) {
196+ const uint32_t l1ListIdNext = (l1ListId_ + 1 < STAGES) ? (l1ListId_ + 1) : 0;
197+ uint32_t kActualNext = 0;
198+ if (kLoopIdx < kTileCount - 1) {
199+ const uint32_t kLoopIdxNext = kLoopIdx + 1;
200+ kActualNext = (kLoopIdxNext < kTileCount - 1) ? L1_K : (actualShape.k() - kLoopIdxNext * L1_K);
201+ Catlass::MatrixCoord gmTileBOffset{kLoopIdxNext * L1_K, 0};
202+ auto gmTileB = gmB[layoutB.GetOffset(gmTileBOffset)];
203+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[l1ListIdNext]);
204+ layoutTileB = layoutB.GetTileLayout(Catlass::MakeCoord(kActualNext, actualShape.n()));
205+ copyGmToL1B_(l1BTensorList_[l1ListIdNext], gmTileB, layoutBInL1, layoutTileB);
206+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList_[l1ListIdNext]);
207+ }
208+ 
209+ auto l1ATensor = l1ATensor_;
210+ auto l1BTensor = l1BTensorList_[l1ListId_];
211+ const uint32_t kPartLoop = CeilDiv<L1_K>(kActual);
212+ 
213+ for (uint32_t mPartIdx = 0; mPartIdx < mPartLoop; ++mPartIdx) {
214+ const uint32_t mPartActual = (mPartIdx < mPartLoop - 1) ? L1_M : (mRound - mPartIdx * L1_M);
215+ for (uint32_t kPartIdx = 0; kPartIdx < kPartLoop; ++kPartIdx) {
216+ const uint32_t kPartActual = (kPartIdx < kPartLoop - 1) ? L1_K : (kActual - kPartIdx * L1_K);
217+ auto l0ATile = l0ATensorList_[qResident ? L0A_RESIDENT_ID : L0A_FALLBACK_ID];
218+ if (!qResident) {
219+ LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout<ElementA>(mPartActual, kPartActual);
220+ Catlass::MatrixCoord l1AOffset{mPartIdx * L1_M, kPartIdx * L1_K + kLoopIdx * L1_K};
221+ auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1AOffset)];
222+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[L0A_FALLBACK_ID]);
223+ copyL1ToL0A_(l0ATile, l1ATile, layoutAInL0, layoutAInL1);
224+ }
225+ 
226+ for (uint32_t nPartIdx = 0; nPartIdx < nPartLoop; ++nPartIdx) {
227+ const uint32_t nPartActual = (nPartIdx < nPartLoop - 1) ? L1_N : (nRound - nPartIdx * L1_N);
228+ auto l0BTile = l0BTensorList_[l0BListId_];
229+ LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout<ElementB>(kPartActual, nPartActual);
230+ Catlass::MatrixCoord l1BOffset{kPartIdx * L1_K, nPartIdx * L1_N};
231+ auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BOffset)];
232+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventList_[l0BListId_]);
233+ if ((kPartIdx == 0) && (nPartIdx == 0)) {
234+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventList_[l1ListId_]);
235+ }
236+ copyL1ToL0B_(l0BTile, l1BTile, layoutBInL0, layoutBInL1);
237+ if ((kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1)) {
238+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventList_[l1ListId_]);
239+ }
240+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
241+ Catlass::MatrixCoord l0COffset{mPartIdx * L1_M, nPartIdx * L1_N};
242+ auto l0CTile = l0CTensorList_[l0CListId_][layoutInL0C.GetOffset(l0COffset)];
243+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
244+ const bool initC = ((kLoopIdx == 0) && (kPartIdx == 0));
245+ uint8_t unitFlag = 0b00;
246+ if constexpr (USE_UNIT_FLAG) {
247+ const bool last = (kLoopIdx == kTileCount - 1) && (mPartIdx == mPartLoop - 1) &&
248+ (kPartIdx == kPartLoop - 1) && (nPartIdx == nPartLoop - 1);
249+ unitFlag = last ? 0b11 : 0b10;
250+ }
251+ tileMmad_(l0CTile, l0ATile, l0BTile, mPartActual, nPartActual, kPartActual, initC, unitFlag);
252+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventList_[l0BListId_]);
253+ l0BListId_ = (l0BListId_ + 1 < L0B_STAGES) ? (l0BListId_ + 1) : 0;
254+ }
255+ if (!qResident) {
256+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventList_[L0A_FALLBACK_ID]);
257+ }
258+ }
259+ }
260+ l1ListId_ = l1ListIdNext;
261+ kActual = kActualNext;
262+ }
263+ 
264+ LayoutC layoutBlock = layoutC.GetTileLayout(actualShape.GetCoordMN());
265+ if (storeToUb_) {
266+ StoreToUb(actualShape);
267+ } else if constexpr (USE_UNIT_FLAG) {
268+ copyL0CToGm_(gmC, l0CTensorList_[l0CListId_], layoutBlock, layoutInL0C, 0b11);
269+ } else {
270+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CEventList_[l0CListId_]);
271+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CEventList_[l0CListId_]);
272+ copyL0CToGm_(gmC, l0CTensorList_[l0CListId_], layoutBlock, layoutInL0C);
273+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[l0CListId_]);
274+ }
275+ l0CListId_ = (l0CListId_ + 1 < L0C_STAGES) ? (l0CListId_ + 1) : 0;
276+ }
277+ 
278+ /// L0C → 配对 AIV UB（dualDstCtl=1，按 M 对半）。ubC 已是 pingpong slot 起点。
279+ __aicore__ inline void ComputeToUb(const AscendC::GlobalTensor<ElementA> &gmA, const LayoutA &layoutA,
280+ const AscendC::GlobalTensor<ElementB> &gmB, const LayoutB &layoutB,
281+ const AscendC::LocalTensor<float> &ubC, const Catlass::GemmCoord &actualShape,
282+ bool needLoadL1)
283+ {
284+ storeToUb_ = true;
285+ ubC_ = ubC;
286+ AscendC::GlobalTensor<ElementC> dummyGm;
287+ const LayoutC dummyLayout(actualShape.m(), actualShape.n(), actualShape.n());
288+ operator()(gmA, layoutA, gmB, layoutB, dummyGm, dummyLayout, actualShape, needLoadL1);
289+ storeToUb_ = false;
290+ }
291+ 
292+private:
293+ static constexpr AscendC::FixpipeConfig MSA_CFG_ROW_MAJOR_UB = {AscendC::CO2Layout::ROW_MAJOR, true};
294+ 
295+ __aicore__ inline void StoreToUb(const Catlass::GemmCoord &actualShape)
296+ {
297+ const uint32_t mEven = (actualShape.m() + 1U) & ~1U;
298+ const uint32_t nAlign = (actualShape.n() + 7U) & ~7U;
299+ const uint32_t mRound = RoundUp<L1AAlignHelper::M_ALIGNED>(actualShape.m());
300+ if constexpr (!USE_UNIT_FLAG) {
301+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CEventList_[l0CListId_]);
302+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CEventList_[l0CListId_]);
303+ }
304+ if constexpr (sizeof(ElementC) == 2) {
305+ uint32_t m0 = 0;
306+ uint32_t m1 = 0;
307+ MsaC2UbHalfSplit(actualShape.m(), m0, m1);
308+ auto ubH = ubC_.template ReinterpretCast<half>();
309+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> fp;
310+ fp.srcStride = mRound;
311+ fp.dualDstCtl = 0;
312+ fp.quantPre = QuantMode_t::F322F16;
313+ fp.nSize = nAlign;
314+ fp.dstStride = nAlign;
315+ fp.params.ndNum = 1;
316+ fp.params.srcNdStride = 0;
317+ fp.params.dstNdStride = 0;
318+ if (m0 > 0U) {
319+ fp.mSize = m0;
320+ fp.subBlockId = false;
321+ AscendC::Fixpipe<half, float, MSA_CFG_ROW_MAJOR_UB>(ubH, l0CTensorList_[l0CListId_], fp);
322+ }
323+ if (m1 > 0U) {
324+ fp.mSize = m1;
325+ fp.subBlockId = true;
326+ AscendC::Fixpipe<half, float, MSA_CFG_ROW_MAJOR_UB>(ubH, l0CTensorList_[l0CListId_][m0 * 16U], fp);
327+ }
328+ } else {
329+ // 950 L0C→UB：nSize 上限 64。ndNum=2 用 panel 布局（QLI）：每 AIV 先 [M/2,64]
330+ // 再紧接第二块 [M/2,64]，总 32KB ping。
331+ constexpr uint32_t N_PER_ND = 64;
332+ const uint32_t rowsPerAiv = mEven / MSA_AIV_PER_AIC;
333+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> fp;
334+ fp.mSize = mEven;
335+ fp.srcStride = mRound;
336+ fp.dstStride = N_PER_ND;
337+ fp.dualDstCtl = 1;
338+ fp.quantPre = QuantMode_t::NoQuant;
339+ fp.nSize = N_PER_ND;
340+ fp.params.ndNum = (nAlign > N_PER_ND) ? (nAlign / N_PER_ND) : 1;
341+ if (fp.params.ndNum > 1) {
342+ fp.params.srcNdStride = ((mEven + 15U) / 16U) * N_PER_ND;
343+ fp.params.dstNdStride = rowsPerAiv * N_PER_ND;
344+ } else {
345+ fp.nSize = nAlign;
346+ fp.params.srcNdStride = 0;
347+ fp.params.dstNdStride = 0;
348+ }
349+ AscendC::Fixpipe<float, float, MSA_CFG_ROW_MAJOR_UB>(ubC_, l0CTensorList_[l0CListId_], fp);
350+ }
351+ if constexpr (!USE_UNIT_FLAG) {
352+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventList_[l0CListId_]);
353+ }
354+ }
355+ AscendC::LocalTensor<ElementA> l1ATensor_;
356+ AscendC::LocalTensor<ElementB> l1BTensorList_[STAGES];
357+ AscendC::LocalTensor<ElementA> l0ATensorList_[STAGES];
358+ AscendC::LocalTensor<ElementB> l0BTensorList_[STAGES];
359+ AscendC::LocalTensor<ElementAccumulator> l0CTensorList_[L0C_STAGES];
360+ int32_t l1AEvent_{0};
361+ int32_t l1BEventList_[STAGES]{};
362+ int32_t l0AEventList_[STAGES]{};
363+ int32_t l0BEventList_[STAGES]{};
364+ int32_t l0CEventList_[L0C_STAGES]{};
365+ uint32_t l1ListId_{0};
366+ uint32_t l0BListId_{0};
367+ uint32_t l0CListId_{0};
368+ TileMmad tileMmad_;
369+ CopyGmToL1A copyGmToL1A_;
370+ CopyGmToL1B copyGmToL1B_;
371+ CopyL1ToL0A copyL1ToL0A_;
372+ CopyL1ToL0B copyL1ToL0B_;
373+ CopyL0CToGm copyL0CToGm_;
374+ bool storeToUb_{false};
375+ AscendC::LocalTensor<float> ubC_;
376+};
377+ 
378+} // namespace MsaIndexScoreNs
379+ 
380+#endif // MSA_BLOCK_MMAD_ARCH35_H
Aattention/msa_index_score/op_kernel/arch35/msa_index_score_kernel.h+530-0
@@ -0,0 +1,530 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+/*!
12+ * \file msa_index_score_kernel.h
13+ * \brief MsaIndexScore Ascend 950：计算路径照搬 A2 核函数：AIC 做 paged QKᵀ，AIV 做（可选反量化）+ mask + 分段 RowMax。
14+ * 第一阶段：计算/握手照搬 A2，ArchTag=AtlasA5。
15+ *
16+ * 非量化：query/key 同为 bf16|fp16|fp8，score 侧不乘 scale。
17+ * int8 量化：key 为 int8；AIV 将 page DataCopy+Cast 到 fp 暂存后通知 AIC 做 Mmad；
18+ * AIV 在 mask 前按 scale[NP,N_kv,P] 对 S 做列乘完成 per-token 反量化。
19+ *
20+ * MIX 编译：`if ASCEND_IS_AIC` 拦不住模板实例化。AIV 再实例化 FP8 Cube
21+ *（Mmad / LoadData2DParamsV2）会让 bisheng 栈砸（exit 134）。Cube 路径仅在
22+ * `__DAV_C310_CUBE__` / `__DAV_310R6_CUBE__` 下编译。
23+ */
24+ 
25+#ifndef MSA_INDEX_SCORE_KERNEL_ARCH35_H
26+#define MSA_INDEX_SCORE_KERNEL_H
27+ 
28+#include "kernel_operator.h"
29+ 
30+#include "catlass/msa_catlass.hpp"
31+#include "catlass/arch/msa_arch.hpp"
32+#include "catlass/arch/msa_resource.hpp"
33+#include "catlass/arch/msa_cross_core_sync.hpp"
34+#include "catlass/msa_coord.hpp"
35+#include "catlass/msa_gemm_coord.hpp"
36+#include "catlass/msa_matrix_coord.hpp"
37+#include "catlass/layout/msa_layout.hpp"
38+#include "catlass/gemm/msa_gemm_type.hpp"
39+ 
40+#include "../msa_index_score_common.h"
41+#if defined(__DAV_C310_CUBE__) || defined(__DAV_310R6_CUBE__)
42+#include "msa_block_mmad.h"
43+#endif
44+#include "../arch22/msa_index_score_task.h"
45+#include "msa_seg_row_max_epilogue.h"
46+ 
47+namespace MsaIndexScoreNs {
48+ 
49+template <class ElementQ_, bool IS_QUANT>
50+class MsaIndexScoreKernel {
51+public:
52+ using ElementQ = ElementQ_;
53+ using ArchTag = Catlass::Arch::AtlasA5;
54+ 
55+ using LayoutQ = Catlass::layout::RowMajor;
56+ using LayoutK = Catlass::layout::ColumnMajor;
57+ using LayoutS = Catlass::layout::RowMajor;
58+ 
59+ using QType = Catlass::Gemm::GemmType<ElementQ, LayoutQ>;
60+ using KType = Catlass::Gemm::GemmType<ElementQ, LayoutK>;
61+ // S workspace 元素：非量化 fp16（fixpipe F322F16 直接写，写读流量减半），int8 fp32
62+ // （S 幅值大、1e-3 容差下 fp16 精度余量不足 15%，保持 fp32 反量化）。
63+ using ElementS = std::conditional_t<IS_QUANT, float, half>;
64+ using SType = Catlass::Gemm::GemmType<ElementS, LayoutS>;
65+ 
66+#if defined(__DAV_C310_CUBE__) || defined(__DAV_310R6_CUBE__)
67+ // int8 的 K 源是每 stile 复用的 per-core scratch，保持 2 级 L1B 流水（3 级实测竞态崩溃）；
68+ // 同样出于该 scratch 的时序保守性，unit flag 只在非量化路径开启。
69+ using BlockMmad = MsaBlockMmad<QType, KType, SType, IS_QUANT ? MSA_L1B_STAGES_INT8 : MSA_L1B_STAGES_FP16, false>;
70+#endif
71+ 
72+ static constexpr uint32_t REVERSE_DEPTH = MSA_WORKSPACE_STAGES - 1;
73+ static constexpr uint32_t STILE_WIDTH = MSA_BLOCKS_PER_STILE * MSA_BLOCK_SIZE;
74+ static constexpr bool IS_QUANT_V = IS_QUANT;
75+ static constexpr bool USE_C2UB = (MSA_A5_USE_C2UB != 0);
76+ static constexpr uint32_t K_SCRATCH_ELEMS = USE_C2UB ? MSA_A5_K_SCRATCH_ELEM_NUM : MSA_K_SCRATCH_ELEM_NUM;
77+ // C_to_UB：cast 必须在 epilogue stage 之后。放在 ping 64KB 处会盖住 T64/deq/stage，
78+ // 下一页 gather 会把尚未 Flush 的 score 清掉（TND/int8 多 page：1e29→0、fill 变 0）。
79+ static constexpr uint32_t UB_OFF_CAST_I8 =
80+ USE_C2UB ? (((MsaSegRowMaxEpilogue<IS_QUANT_V>::UB_TOTAL + MSA_UB_ALIGN_BYTES - 1U) / MSA_UB_ALIGN_BYTES) *
81+ MSA_UB_ALIGN_BYTES) :
82+ 0;
83+ static constexpr uint32_t UB_SIZE_CAST_I8 = MSA_BLOCK_SIZE * MSA_K_TILE * sizeof(int8_t);
84+ static constexpr uint32_t UB_OFF_CAST_FP = UB_OFF_CAST_I8 + UB_SIZE_CAST_I8;
85+ static constexpr uint32_t UB_SIZE_CAST_FP = MSA_BLOCK_SIZE * MSA_K_TILE * sizeof(half);
86+ static_assert(UB_OFF_CAST_FP + UB_SIZE_CAST_FP <= ArchTag::UB_SIZE, "int8 cast UB out of bounds");
87+ 
88+ __aicore__ inline MsaIndexScoreKernel() {}
89+ 
90+ __aicore__ inline void Init(GM_ADDR query, GM_ADDR key, GM_ADDR blockTable, GM_ADDR scale, GM_ADDR actualSeqQlen,
91+ GM_ADDR actualSeqKlen, GM_ADDR startLoc, GM_ADDR score, GM_ADDR workspace,
92+ const MsaIndexScoreTilingData *__restrict tiling)
93+ {
94+ tiling_ = tiling;
95+ gQuery_.SetGlobalBuffer(reinterpret_cast<__gm__ ElementQ *>(query));
96+ if constexpr (IS_QUANT_V) {
97+ gKeyInt8_.SetGlobalBuffer(reinterpret_cast<__gm__ int8_t *>(key));
98+ } else {
99+ gKeyFp_.SetGlobalBuffer(reinterpret_cast<__gm__ ElementQ *>(key));
100+ }
Mattention/msa_index_score/op_kernel/catlass/arch/msa_arch.hpp+11-0
@@ -25,6 +25,17 @@ struct AtlasA2 {
25 static constexpr uint32_t L0C_SIZE = 128 * 1024;25 static constexpr uint32_t L0C_SIZE = 128 * 1024;
26};26};
27 27 
28+/// Ascend 950PR/DT（dav-3510）。容量对齐 950 仓 AtlasA5 / BSA arch。
29+struct AtlasA5 {
30+ static constexpr uint32_t BIAS_SIZE = 4 * 1024;
31+ static constexpr uint32_t FIXBUF_SIZE = 16 * 1024;
32+ static constexpr uint32_t UB_SIZE = 256 * 1024;
33+ static constexpr uint32_t L1_SIZE = 512 * 1024;
34+ static constexpr uint32_t L0A_SIZE = 64 * 1024;
35+ static constexpr uint32_t L0B_SIZE = 64 * 1024;
36+ static constexpr uint32_t L0C_SIZE = 256 * 1024;
37+};
38+ 
28template <AscendC::TPosition POS>39template <AscendC::TPosition POS>
29using PositionType = std::integral_constant<AscendC::TPosition, POS>;40using PositionType = std::integral_constant<AscendC::TPosition, POS>;
30 41 
Mattention/msa_index_score/op_kernel/catlass/arch/msa_cross_core_sync.hpp+6-0
@@ -101,6 +101,12 @@ void CrossCoreWaitFlag(CrossCoreFlag &flag)
101 AscendC::CrossCoreWaitFlag(flag.id);101 AscendC::CrossCoreWaitFlag(flag.id);
102}102}
103 103 
104+template <uint8_t MODE, pipe_t PIPE>
105+CATLASS_DEVICE void CrossCoreWaitFlag(CrossCoreFlag &flag)
106+{
107+ AscendC::CrossCoreWaitFlag<MODE, PIPE>(flag.id);
108+}
109+ 
104template <uint8_t MODE, pipe_t PIPE, uint32_t REVERSE_DEPTH>110template <uint8_t MODE, pipe_t PIPE, uint32_t REVERSE_DEPTH>
105CATLASS_DEVICE void CrossCoreSetFlagWithReverse(CrossCoreFlagWithReverse<REVERSE_DEPTH> &flag)111CATLASS_DEVICE void CrossCoreSetFlagWithReverse(CrossCoreFlagWithReverse<REVERSE_DEPTH> &flag)
106{112{
Mattention/msa_index_score/op_kernel/catlass/arch/msa_local_tensor_buffer.hpp+84-1
@@ -45,6 +45,9 @@ public:
45 CATLASS_DEVICE45 CATLASS_DEVICE
46 LocalTensorBuffer()46 LocalTensorBuffer()
47 {47 {
48+ if ASCEND_IS_AIV {
49+ return;
50+ }
48 AscendC::TBuf<AscendC::TPosition::A1> tbufA1;51 AscendC::TBuf<AscendC::TPosition::A1> tbufA1;
49 GetTPipePtr()->InitBuffer(tbufA1, ArchTag::L1_SIZE);52 GetTPipePtr()->InitBuffer(tbufA1, ArchTag::L1_SIZE);
50 tensor = tbufA1.Get<uint8_t>();53 tensor = tbufA1.Get<uint8_t>();
@@ -62,6 +65,9 @@ public:
62 CATLASS_DEVICE65 CATLASS_DEVICE
63 LocalTensorBuffer()66 LocalTensorBuffer()
64 {67 {
68+ if ASCEND_IS_AIV {
69+ return;
70+ }
65 AscendC::TBuf<AscendC::TPosition::A2> tbufA2;71 AscendC::TBuf<AscendC::TPosition::A2> tbufA2;
66 GetTPipePtr()->InitBuffer(tbufA2, ArchTag::L0A_SIZE);72 GetTPipePtr()->InitBuffer(tbufA2, ArchTag::L0A_SIZE);
67 tensor = tbufA2.Get<uint8_t>();73 tensor = tbufA2.Get<uint8_t>();
@@ -79,6 +85,9 @@ public:
79 CATLASS_DEVICE85 CATLASS_DEVICE
80 LocalTensorBuffer()86 LocalTensorBuffer()
81 {87 {
88+ if ASCEND_IS_AIV {
89+ return;
90+ }
82 AscendC::TBuf<AscendC::TPosition::B1> tbufB1;91 AscendC::TBuf<AscendC::TPosition::B1> tbufB1;
83 GetTPipePtr()->InitBuffer(tbufB1, ArchTag::L1_SIZE);92 GetTPipePtr()->InitBuffer(tbufB1, ArchTag::L1_SIZE);
84 tensor = tbufB1.Get<uint8_t>();93 tensor = tbufB1.Get<uint8_t>();
@@ -96,6 +105,9 @@ public:
96 CATLASS_DEVICE105 CATLASS_DEVICE
97 LocalTensorBuffer()106 LocalTensorBuffer()
98 {107 {
108+ if ASCEND_IS_AIV {
109+ return;
110+ }
99 AscendC::TBuf<AscendC::TPosition::B2> tbufB2;111 AscendC::TBuf<AscendC::TPosition::B2> tbufB2;
100 GetTPipePtr()->InitBuffer(tbufB2, ArchTag::L0B_SIZE);112 GetTPipePtr()->InitBuffer(tbufB2, ArchTag::L0B_SIZE);
101 tensor = tbufB2.Get<uint8_t>();113 tensor = tbufB2.Get<uint8_t>();
@@ -114,6 +126,9 @@ public:
114 CATLASS_DEVICE126 CATLASS_DEVICE
115 LocalTensorBuffer()127 LocalTensorBuffer()
116 {128 {
129+ if ASCEND_IS_AIV {
130+ return;
131+ }
117 AscendC::TBuf<AscendC::TPosition::C1> tbufC1;132 AscendC::TBuf<AscendC::TPosition::C1> tbufC1;
118 GetTPipePtr()->InitBuffer(tbufC1, ArchTag::L1_SIZE);133 GetTPipePtr()->InitBuffer(tbufC1, ArchTag::L1_SIZE);
119 tensor = tbufC1.Get<uint8_t>();134 tensor = tbufC1.Get<uint8_t>();
@@ -132,6 +147,9 @@ public:
132 CATLASS_DEVICE147 CATLASS_DEVICE
133 LocalTensorBuffer()148 LocalTensorBuffer()
134 {149 {
150+ if ASCEND_IS_AIV {
151+ return;
152+ }
135 AscendC::TBuf<AscendC::TPosition::C2> tbufC2;153 AscendC::TBuf<AscendC::TPosition::C2> tbufC2;
136 GetTPipePtr()->InitBuffer(tbufC2, ArchTag::BIAS_SIZE);154 GetTPipePtr()->InitBuffer(tbufC2, ArchTag::BIAS_SIZE);
137 tensor = tbufC2.Get<uint8_t>();155 tensor = tbufC2.Get<uint8_t>();
@@ -149,6 +167,9 @@ public:
149 CATLASS_DEVICE167 CATLASS_DEVICE
150 LocalTensorBuffer()168 LocalTensorBuffer()
151 {169 {
170+ if ASCEND_IS_AIV {
171+ return;
172+ }
152 AscendC::TBuf<AscendC::TPosition::CO1> tbufCO1;173 AscendC::TBuf<AscendC::TPosition::CO1> tbufCO1;
153 GetTPipePtr()->InitBuffer(tbufCO1, ArchTag::L0C_SIZE);174 GetTPipePtr()->InitBuffer(tbufCO1, ArchTag::L0C_SIZE);
154 tensor = tbufCO1.Get<uint8_t>();175 tensor = tbufCO1.Get<uint8_t>();
@@ -167,6 +188,63 @@ public:
167 CATLASS_DEVICE188 CATLASS_DEVICE
168 LocalTensorBuffer()189 LocalTensorBuffer()
169 {190 {
191+ if ASCEND_IS_AIV {
192+ return;
193+ }
194+ AscendC::TBuf<AscendC::TPosition::C2PIPE2GM> tbufC2PIPE2GM;
195+ GetTPipePtr()->InitBuffer(tbufC2PIPE2GM, ArchTag::FIXBUF_SIZE);
196+ tensor = tbufC2PIPE2GM.Get<uint8_t>();
197+ }
198+};
199+ 
200+template <>
201+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::C1> : LocalTensorBufferBase {
202+public:
203+ using ArchTag = Arch::AtlasA5;
204+ static constexpr AscendC::TPosition Position = AscendC::TPosition::C1;
205+ 
206+ CATLASS_DEVICE
207+ LocalTensorBuffer()
208+ {
209+ if ASCEND_IS_AIV {
210+ return;
211+ }
212+ AscendC::TBuf<AscendC::TPosition::C1> tbufC1;
213+ GetTPipePtr()->InitBuffer(tbufC1, ArchTag::L1_SIZE);
214+ tensor = tbufC1.Get<uint8_t>();
215+ }
216+};
217+ 
218+template <>
219+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::C2> : LocalTensorBufferBase {
220+public:
221+ using ArchTag = Arch::AtlasA5;
222+ static constexpr AscendC::TPosition Position = AscendC::TPosition::C2;
223+ 
224+ CATLASS_DEVICE
225+ LocalTensorBuffer()
226+ {
227+ if ASCEND_IS_AIV {
228+ return;
229+ }
230+ AscendC::TBuf<AscendC::TPosition::C2> tbufC2;
231+ GetTPipePtr()->InitBuffer(tbufC2, ArchTag::BIAS_SIZE);
232+ tensor = tbufC2.Get<uint8_t>();
233+ }
234+};
235+ 
236+template <>
237+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::C2PIPE2GM> : LocalTensorBufferBase {
238+public:
239+ using ArchTag = Arch::AtlasA5;
240+ static constexpr AscendC::TPosition Position = AscendC::TPosition::C2PIPE2GM;
241+ 
242+ CATLASS_DEVICE
243+ LocalTensorBuffer()
244+ {
245+ if ASCEND_IS_AIV {
246+ return;
247+ }
170 AscendC::TBuf<AscendC::TPosition::C2PIPE2GM> tbufC2PIPE2GM;248 AscendC::TBuf<AscendC::TPosition::C2PIPE2GM> tbufC2PIPE2GM;
171 GetTPipePtr()->InitBuffer(tbufC2PIPE2GM, ArchTag::FIXBUF_SIZE);249 GetTPipePtr()->InitBuffer(tbufC2PIPE2GM, ArchTag::FIXBUF_SIZE);
172 tensor = tbufC2PIPE2GM.Get<uint8_t>();250 tensor = tbufC2PIPE2GM.Get<uint8_t>();
@@ -219,7 +297,12 @@ public:
219 LocalTensorBuffer()297 LocalTensorBuffer()
220 {298 {
221 AscendC::TBuf<AscendC::TPosition::VECCALC> tbufVECCALC;299 AscendC::TBuf<AscendC::TPosition::VECCALC> tbufVECCALC;
222- GetTPipePtr()->InitBuffer(tbufVECCALC, ArchTag::UB_SIZE);300+ uint32_t ubBytes = ArchTag::UB_SIZE;
301+ if ASCEND_IS_AIC {
302+ // Cube 只需 dual-dst pingpong（2*64*128*4=64KB）；AIC 上 InitBuffer 256KB VECCALC 可能卡住 Fixpipe。
303+ ubBytes = 64 * 1024;
304+ }
305+ GetTPipePtr()->InitBuffer(tbufVECCALC, ubBytes);
223 tensor = tbufVECCALC.Get<uint8_t>();306 tensor = tbufVECCALC.Get<uint8_t>();
224 }307 }
225};308};
Mattention/msa_index_score/op_kernel/catlass/gemm/msa_gemm_helper.hpp+18-0
@@ -109,6 +109,24 @@ struct ElementAccumulatorSelector<bfloat16_t, bfloat16_t> {
109 using ElementAccumulator = float;109 using ElementAccumulator = float;
110};110};
111 111 
112+// FP8 / hifloat8 仅 AIC 310（Ascend 950）。
113+#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 310)
114+template <>
115+struct ElementAccumulatorSelector<fp8_e4m3fn_t, fp8_e4m3fn_t> {
116+ using ElementAccumulator = float;
117+};
118+ 
119+template <>
120+struct ElementAccumulatorSelector<fp8_e5m2_t, fp8_e5m2_t> {
121+ using ElementAccumulator = float;
122+};
123+ 
124+template <>
125+struct ElementAccumulatorSelector<hifloat8, hifloat8> {
126+ using ElementAccumulator = float;
127+};
128+#endif
129+ 
112template <>130template <>
113struct ElementAccumulatorSelector<AscendC::int4b_t, AscendC::int4b_t> {131struct ElementAccumulatorSelector<AscendC::int4b_t, AscendC::int4b_t> {
114 using ElementAccumulator = int32_t;132 using ElementAccumulator = int32_t;
Mattention/msa_index_score/op_kernel/catlass/gemm/tile/msa_copy_gm_to_l1.hpp+8-0
@@ -1141,6 +1141,10 @@ struct CopyGmToL1<Arch::AtlasA2, Gemm::GemmType<Element, layout::RowMajor>> {
1141};1141};
1142/////////////////////////////////1142/////////////////////////////////
1143 1143 
1144+template <class Element>
1145+struct CopyGmToL1<Arch::AtlasA5, Gemm::GemmType<Element, layout::RowMajor>>
1146+ : CopyGmToL1<Arch::AtlasA2, Gemm::GemmType<Element, layout::RowMajor>> {};
1147+ 
1144/// Partial specialization for AtlasA2, ColumnMajor in and nZ out.1148/// Partial specialization for AtlasA2, ColumnMajor in and nZ out.
1145template <class Element>1149template <class Element>
1146struct CopyGmToL1<Arch::AtlasA2, Gemm::GemmType<Element, layout::ColumnMajor>> {1150struct CopyGmToL1<Arch::AtlasA2, Gemm::GemmType<Element, layout::ColumnMajor>> {
@@ -1190,6 +1194,10 @@ struct CopyGmToL1<Arch::AtlasA2, Gemm::GemmType<Element, layout::ColumnMajor>> {
1190 }1194 }
1191};1195};
1192 1196 
1197+template <class Element>
1198+struct CopyGmToL1<Arch::AtlasA5, Gemm::GemmType<Element, layout::ColumnMajor>>
1199+ : CopyGmToL1<Arch::AtlasA2, Gemm::GemmType<Element, layout::ColumnMajor>> {};
1200+ 
1193/// Partial specialization for zN in and zN out.1201/// Partial specialization for zN in and zN out.
1194template <class ArchTag, class Element>1202template <class ArchTag, class Element>
1195struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::zN>> {1203struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::zN>> {
Mattention/msa_index_score/op_kernel/catlass/gemm/tile/msa_copy_l0c_to_gm.hpp+54-0
@@ -38,6 +38,16 @@ struct CopyL0CToGmQuantMode<Catlass::Arch::AtlasA2, float, float, ScaleGranulari
38 static constexpr auto VALUE = QuantMode_t::NoQuant;38 static constexpr auto VALUE = QuantMode_t::NoQuant;
39};39};
40 40 
41+template <>
42+struct CopyL0CToGmQuantMode<Catlass::Arch::AtlasA5, float, float, ScaleGranularity::NO_QUANT> {
43+ static constexpr auto VALUE = QuantMode_t::NoQuant;
44+};
45+ 
46+template <>
47+struct CopyL0CToGmQuantMode<Catlass::Arch::AtlasA5, float, half, ScaleGranularity::NO_QUANT> {
48+ static constexpr auto VALUE = QuantMode_t::F322F16;
49+};
50+ 
41// CopyL0CToGm cast fp32 to fp1651// CopyL0CToGm cast fp32 to fp16
42template <>52template <>
43struct CopyL0CToGmQuantMode<Catlass::Arch::AtlasA2, float, half, ScaleGranularity::NO_QUANT> {53struct CopyL0CToGmQuantMode<Catlass::Arch::AtlasA2, float, half, ScaleGranularity::NO_QUANT> {
@@ -169,6 +179,50 @@ struct CopyL0CToGm<Catlass::Arch::AtlasA2, ElementAccumulator_, Gemm::GemmType<E
169 }179 }
170};180};
171 181 
182+/// AtlasA5 / C310：Fixpipe L0C→GM（950 主路径是 L0C→UB；本特化仅满足 TileCopy 实例化 / 编译期回退）。
183+template <class ElementAccumulator_, class ElementDst_, bool ReluEnable_>
184+struct CopyL0CToGm<Catlass::Arch::AtlasA5, ElementAccumulator_, Gemm::GemmType<ElementDst_, layout::RowMajor>,
185+ ScaleGranularity::NO_QUANT, ReluEnable_> {
186+ using ArchTag = Catlass::Arch::AtlasA5;
187+ using ElementDst = ElementDst_;
188+ using ElementSrc = ElementAccumulator_;
189+ using LayoutSrc = Catlass::layout::zN;
190+ using LayoutDst = Catlass::layout::RowMajor;
191+ static constexpr auto quantPre =
192+ CopyL0CToGmQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::NO_QUANT>::VALUE;
193+ static constexpr auto reluEn = ReluEnable_;
194+ 
195+ struct Params {};
196+ Params params;
197+ 
198+ CATLASS_DEVICE
199+ CopyL0CToGm() = default;
200+ 
201+ CATLASS_DEVICE
202+ CopyL0CToGm(Params const &params_)
203+ : params(params_) {};
204+ 
205+ CATLASS_DEVICE
206+ void operator()(AscendC::GlobalTensor<ElementDst> const &dst, AscendC::LocalTensor<ElementSrc> const &src,
207+ LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0)
208+ {
209+ // 950 L0C→GM：与仓内 KDA 相同的 DataCopyCO12Dst + SetFixpipeNz2ndFlag。
210+ // mSize 用实际 M（dstLayout），srcStride 用 L0C 16 对齐步长。
211+ // 写 mRound 行会把 D=16（M=4）打成全 -inf。
212+ AscendC::DataCopyCO12DstParams copyParams;
213+ copyParams.nSize = dstLayout.shape(1);
214+ copyParams.mSize = dstLayout.shape(0);
215+ copyParams.srcStride = srcLayout.stride(LayoutSrc::RANK - 1) / srcLayout.stride(0);
216+ copyParams.dstStride = dstLayout.stride(0);
217+ copyParams.quantPre = quantPre;
218+ copyParams.nz2ndEn = true;
219+ copyParams.reluPre = reluEn;
220+ copyParams.unitFlag = unitFlag;
221+ AscendC::SetFixpipeNz2ndFlag(1, 1, 1);
222+ AscendC::DataCopy(dst, src, copyParams);
223+ }
224+};
225+ 
172template <class ElementAccumulator_, class ElementDst_, bool ReluEnable_>226template <class ElementAccumulator_, class ElementDst_, bool ReluEnable_>
173struct CopyL0CToGm<Catlass::Arch::AtlasA2, ElementAccumulator_, Gemm::GemmType<ElementDst_, layout::RowMajor>,227struct CopyL0CToGm<Catlass::Arch::AtlasA2, ElementAccumulator_, Gemm::GemmType<ElementDst_, layout::RowMajor>,
174 ScaleGranularity::PER_TENSOR, ReluEnable_> {228 ScaleGranularity::PER_TENSOR, ReluEnable_> {
Mattention/msa_index_score/op_kernel/catlass/gemm/tile/msa_copy_l1_to_l0a.hpp+31-0
@@ -379,6 +379,37 @@ struct CopyL1ToL0A<ArchTag, Gemm::GemmType<float, layout::nZ, AscendC::TPosition
379 }379 }
380};380};
381 381 
382+#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 310)
383+/// AtlasA5：L0A 硬件格式为 zN（A2 为 zZ）。K=16 时两者重合，故 D=16 能过、D=128 失败。
384+/// 参数对齐 chunk_kda_fwd 仓内已验证的 Ascend950 CopyL1ToL0A。
385+template <class Element>
386+struct CopyL1ToL0A<Arch::AtlasA5, Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>> {
387+ using LayoutDst = layout::zN;
388+ using LayoutSrc = layout::zN;
389+ 
390+ static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits<Element>::value;
391+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits<Element>::value;
392+ 
393+ CATLASS_DEVICE
394+ CopyL1ToL0A() {};
395+ 
396+ CATLASS_DEVICE
397+ void operator()(AscendC::LocalTensor<Element> const &dstTensor, AscendC::LocalTensor<Element> const &srcTensor,
398+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
399+ {
400+ AscendC::LoadData2DParamsV2 loadDataParams;
401+ loadDataParams.mStartPosition = 0;
402+ loadDataParams.kStartPosition = 0;
403+ loadDataParams.mStep = layoutDst.shape(1);
404+ loadDataParams.kStep = layoutDst.shape(LayoutDst::RANK - 1);
405+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(layoutSrc.stride(LayoutSrc::RANK - 1));
406+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(layoutDst.stride(LayoutDst::RANK - 1));
407+ loadDataParams.ifTranspose = false;
408+ AscendC::LoadData(dstTensor, srcTensor, loadDataParams);
409+ }
410+};
411+#endif
412+ 
382} // namespace Catlass::Gemm::Tile413} // namespace Catlass::Gemm::Tile
383 414 
384#endif // CATLASS_MSA_GEMM_TILE_COPY_L1_TO_L0A_HPP415#endif // CATLASS_MSA_GEMM_TILE_COPY_L1_TO_L0A_HPP
Mattention/msa_index_score/op_kernel/catlass/gemm/tile/msa_copy_l1_to_l0b.hpp+31-0
@@ -475,6 +475,37 @@ struct CopyL1ToL0B<ArchTag, Gemm::GemmType<Element, layout::nZ, AscendC::TPositi
475 }475 }
476};476};
477 477 
478+#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 310)
479+/// AtlasA5：nZ→nZ 用 LoadData2DParamsV2 一次搬完（A2 LoadData2D 在 950 上只对单 C0 正确）。
480+/// 参数对齐 chunk_kda_fwd 仓内已验证的 Ascend950 CopyL1ToL0B；ifTranspose 必须为 false。
481+template <class Element>
482+struct CopyL1ToL0B<Arch::AtlasA5, Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::A1>> {
483+ using LayoutDst = layout::nZ;
484+ using LayoutSrc = layout::nZ;
485+ 
486+ static constexpr uint32_t ELE_NUM_PER_C0 = BytesToBits(BYTE_PER_C0) / SizeOfBits<Element>::value;
487+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BytesToBits(BYTE_PER_FRACTAL) / SizeOfBits<Element>::value;
488+ 
489+ CATLASS_DEVICE
490+ CopyL1ToL0B() {};
491+ 
492+ CATLASS_DEVICE
493+ void operator()(AscendC::LocalTensor<Element> const &dstTensor, AscendC::LocalTensor<Element> const &srcTensor,
494+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
495+ {
496+ AscendC::LoadData2DParamsV2 loadDataParams;
497+ loadDataParams.mStartPosition = 0;
498+ loadDataParams.kStartPosition = 0;
499+ loadDataParams.mStep = layoutDst.shape(LayoutDst::RANK - 1);
500+ loadDataParams.kStep = layoutDst.shape(1);
501+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(layoutSrc.stride(1));
502+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(layoutDst.stride(1));
503+ loadDataParams.ifTranspose = false;
504+ AscendC::LoadData(dstTensor, srcTensor, loadDataParams);
505+ }
506+};
507+#endif
508+ 
478} // namespace Catlass::Gemm::Tile509} // namespace Catlass::Gemm::Tile
479 510 
480#endif // CATLASS_MSA_GEMM_TILE_COPY_L1_TO_L0B_HPP511#endif // CATLASS_MSA_GEMM_TILE_COPY_L1_TO_L0B_HPP
Mattention/msa_index_score/op_kernel/catlass/gemm/tile/msa_tile_mmad.hpp+12-0
@@ -12,6 +12,7 @@
12#define CATLASS_MSA_GEMM_TILE_TILE_MMAD_HPP12#define CATLASS_MSA_GEMM_TILE_TILE_MMAD_HPP
13 13 
14#include "catlass/msa_catlass.hpp"14#include "catlass/msa_catlass.hpp"
15+#include "catlass/arch/msa_arch.hpp"
15#include "catlass/gemm/msa_gemm_helper.hpp"16#include "catlass/gemm/msa_gemm_helper.hpp"
16namespace Catlass::Gemm::Tile {17namespace Catlass::Gemm::Tile {
17 18 
@@ -48,6 +49,12 @@ struct TileMmad {
48 mmadParams.k = k;49 mmadParams.k = k;
49 mmadParams.unitFlag = unitFlag;50 mmadParams.unitFlag = unitFlag;
50 mmadParams.cmatrixInitVal = initC;51 mmadParams.cmatrixInitVal = initC;
52+ // disableGemv 仅 950 MmadParams 有；if constexpr 挡不住非依赖成员检查。
53+#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 310)
54+ if constexpr (std::is_same_v<ArchTag_, Catlass::Arch::AtlasA5>) {
55+ mmadParams.disableGemv = true;
56+ }
57+#endif
51 if constexpr (std::is_same_v<ElementA, float> &&58 if constexpr (std::is_same_v<ElementA, float> &&
52 (std::is_same_v<typename AType_::Layout, layout::ColumnMajor> ||59 (std::is_same_v<typename AType_::Layout, layout::ColumnMajor> ||
53 std::is_same_v<typename AType_::Layout, layout::nZ>)) {60 std::is_same_v<typename AType_::Layout, layout::nZ>)) {
@@ -74,6 +81,11 @@ struct TileMmad {
74 mmadParams.k = k;81 mmadParams.k = k;
75 mmadParams.unitFlag = unitFlag;82 mmadParams.unitFlag = unitFlag;
76 mmadParams.cmatrixInitVal = false;83 mmadParams.cmatrixInitVal = false;
84+#if defined(__CCE_AICORE__) && (__CCE_AICORE__ == 310)
85+ if constexpr (std::is_same_v<ArchTag_, Catlass::Arch::AtlasA5>) {
86+ mmadParams.disableGemv = true;
87+ }
88+#endif
77 if constexpr (std::is_same_v<ElementA, float> &&89 if constexpr (std::is_same_v<ElementA, float> &&
78 (std::is_same_v<typename AType_::Layout, layout::ColumnMajor> ||90 (std::is_same_v<typename AType_::Layout, layout::ColumnMajor> ||
79 std::is_same_v<typename AType_::Layout, layout::nZ>)) {91 std::is_same_v<typename AType_::Layout, layout::nZ>)) {
Mattention/msa_index_score/op_kernel/msa_index_score.cpp+15-2
@@ -17,10 +17,14 @@
17 * atten_mask 仅做 host 校验；device 侧按 sparse_mode 解析因果，不消费该 GM。17 * atten_mask 仅做 host 校验；device 侧按 sparse_mode 解析因果，不消费该 GM。
18 */18 */
19 19 
20-#include "kernel_operator.h"20+#include "kernel_operator.h" // force-rebuild-arch22-v65-int8-4slot-dual-aiv
21#include "lib/matmul_intf.h"21#include "lib/matmul_intf.h"
22#include "msa_index_score_common.h"22#include "msa_index_score_common.h"
23+#if (__CCE_AICORE__ == 310)
24+#include "arch35/msa_index_score_kernel.h"
25+#else
23#include "arch22/msa_index_score_kernel.h"26#include "arch22/msa_index_score_kernel.h"
27+#endif
24 28 
25using namespace MsaIndexScoreNs;29using namespace MsaIndexScoreNs;
26 30 
@@ -38,9 +42,10 @@ extern "C" __global__ __aicore__ void msa_index_score(GM_ADDR query, GM_ADDR key
38{42{
39 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);43 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
40 (void)attenMask; // host 校验；device 按 sparse_mode 解析 rightDownCausal44 (void)attenMask; // host 校验；device 按 sparse_mode 解析 rightDownCausal
45+ // 950 非量化 PA：S 在 UB、无 K scratch，user workspace 可能为 0。不能因此跳过整个 kernel 流程。
41 GM_ADDR userWs = AscendC::GetUserWorkspace(workspace);46 GM_ADDR userWs = AscendC::GetUserWorkspace(workspace);
42 if (userWs == nullptr) {47 if (userWs == nullptr) {
43- return;48+ userWs = workspace;
44 }49 }
45 GET_TILING_DATA_WITH_STRUCT(MsaIndexScoreTilingData, tilingDataIn, tiling);50 GET_TILING_DATA_WITH_STRUCT(MsaIndexScoreTilingData, tilingDataIn, tiling);
46 const MsaIndexScoreTilingData *__restrict tilingData = &tilingDataIn;51 const MsaIndexScoreTilingData *__restrict tilingData = &tilingDataIn;
@@ -51,5 +56,13 @@ extern "C" __global__ __aicore__ void msa_index_score(GM_ADDR query, GM_ADDR key
51 MSA_INVOKE_KERNEL(half, false);56 MSA_INVOKE_KERNEL(half, false);
52 } else if (TILING_KEY_IS(MSA_TILING_KEY_FP16_INT8)) {57 } else if (TILING_KEY_IS(MSA_TILING_KEY_FP16_INT8)) {
53 MSA_INVOKE_KERNEL(half, true);58 MSA_INVOKE_KERNEL(half, true);
59+#if (__CCE_AICORE__ == 310)
60+ } else if (TILING_KEY_IS(MSA_TILING_KEY_FP8_E4M3FN)) {
61+ MSA_INVOKE_KERNEL(fp8_e4m3fn_t, false);
62+ } else if (TILING_KEY_IS(MSA_TILING_KEY_FP8_E5M2)) {
63+ MSA_INVOKE_KERNEL(fp8_e5m2_t, false);
64+ } else if (TILING_KEY_IS(MSA_TILING_KEY_HIFLOAT8)) {
65+ MSA_INVOKE_KERNEL(hifloat8_t, false);
66+#endif
54 }67 }
55}68}
Mattention/msa_index_score/op_kernel/msa_index_score_common.h+32-3
@@ -41,7 +41,7 @@ constexpr uint32_t MSA_K_TILE = 128;
41constexpr uint32_t MSA_BLOCKS_PER_STILE = 8;41constexpr uint32_t MSA_BLOCKS_PER_STILE = 8;
42 42 
43// S workspace 轮转级数。AIC 领先 AIV 的未消费 tile 数 <= MSA_WORKSPACE_STAGES - 1。43// S workspace 轮转级数。AIC 领先 AIV 的未消费 tile 数 <= MSA_WORKSPACE_STAGES - 1。
44-// 5 -> 8：AIV 乒乓化后服务速度提升，加深缓冲让 AIC 几乎不因 AIV 停顿（wait_id2→0）。44+// 5 -> 8：AIV 乒乓化后速度提升，加深缓冲让 AIC 几乎不因 AIV 停顿（wait_id2→0）。
45// 非量化路径 S 为 fp16，8 级 × 256KB × 20 核 = 40MB，仍可整体驻留 L2。45// 非量化路径 S 为 fp16，8 级 × 256KB × 20 核 = 40MB，仍可整体驻留 L2。
46constexpr uint32_t MSA_WORKSPACE_STAGES = 8;46constexpr uint32_t MSA_WORKSPACE_STAGES = 8;
47 47 
@@ -57,6 +57,18 @@ constexpr uint32_t MSA_STILE_BYTES_FP32 = MSA_STILE_ELEM_NUM * sizeof(float);
57// （AIV 先 cast，AIC 再 Mmad，避免逐 page 细粒度同步）。57// （AIV 先 cast，AIC 再 Mmad，避免逐 page 细粒度同步）。
58// 单页布局 [blockSize, headDim] 行主序，按 j*P*D 排列。58// 单页布局 [blockSize, headDim] 行主序，按 j*P*D 排列。
59constexpr uint32_t MSA_K_SCRATCH_ELEM_NUM = MSA_BLOCKS_PER_STILE * MSA_BLOCK_SIZE * MSA_K_TILE;59constexpr uint32_t MSA_K_SCRATCH_ELEM_NUM = MSA_BLOCKS_PER_STILE * MSA_BLOCK_SIZE * MSA_K_TILE;
60+// A2/A3 int8：四槽 scratch。WaitS(st) 后 slot[st%4] 已空，gather(st+4) 叠在
61+// AIC 的 cube(st+1)+cube(st+2)+cube(st+3) 上；epilogue(st) 再叠在 cube(st+4) 上。950 不用。
62+constexpr uint32_t MSA_K_SCRATCH_STAGES_A2 = 4;
63+// 950 一页握手：每 AIC 只需一页 K scratch。
64+constexpr uint32_t MSA_A5_K_SCRATCH_ELEM_NUM = MSA_BLOCK_SIZE * MSA_K_TILE;
65+constexpr uint32_t MSA_A5_S_STAGES = 2;
66+constexpr uint32_t MSA_A5_FP8_C0 = 32;
67+constexpr uint32_t MSA_A5_S_PAGE_ELEM_NUM = MSA_ROW_TILE_M * MSA_BLOCK_SIZE;
68+ 
69+#ifndef MSA_A5_USE_C2UB
70+#define MSA_A5_USE_C2UB 1
71+#endif
60 72 
61// AIV 单次归约处理的行数。受 level-0 向量 API repeatTimes(uint8_t) 上限约束：73// AIV 单次归约处理的行数。受 level-0 向量 API repeatTimes(uint8_t) 上限约束：
62// MSA_ROWS_PER_PASS * MSA_BLOCKS_PER_STILE <= 255，且乘积需为 8 的倍数。74// MSA_ROWS_PER_PASS * MSA_BLOCKS_PER_STILE <= 255，且乘积需为 8 的倍数。
@@ -74,10 +86,15 @@ constexpr uint32_t MSA_DATABLOCK_BYTES = 32;
74constexpr uint32_t MSA_VEC_MASK_WORDS = 2;86constexpr uint32_t MSA_VEC_MASK_WORDS = 2;
75// MIX 核比：1 AIC : 2 AIV。87// MIX 核比：1 AIC : 2 AIV。
76constexpr uint32_t MSA_AIV_PER_AIC = 2;88constexpr uint32_t MSA_AIV_PER_AIC = 2;
89+// 950 C_to_UB：dual-dst 后每 AIV 半个 M-tile × N=128 fp32 pingpong。
90+constexpr uint32_t MSA_A5_S_PING_ROWS = MSA_ROW_TILE_M / MSA_AIV_PER_AIC;
91+constexpr uint32_t MSA_A5_S_PING_ELEM_NUM = MSA_A5_S_PING_ROWS * MSA_BLOCK_SIZE;
92+constexpr uint32_t MSA_A5_S_PING_BYTES = MSA_A5_S_PING_ELEM_NUM * static_cast<uint32_t>(sizeof(float));
77// UB 手工布局 512B 对齐。93// UB 手工布局 512B 对齐。
78constexpr uint32_t MSA_UB_ALIGN_BYTES = 512;94constexpr uint32_t MSA_UB_ALIGN_BYTES = 512;
79 95 
80-// BlockMmad L1B pingpong 级数：int8 用 scratch 时序保守取 2；非量化取 3。96+// BlockMmad L1B pingpong 级数：950 int8 仍为 2（单页 scratch 竞态）；A2/A3 int8 在
97+// arch22 kernel 里改用 MSA_L1B_STAGES_FP16（双槽后安全）。
81constexpr uint32_t MSA_L1B_STAGES_INT8 = 2;98constexpr uint32_t MSA_L1B_STAGES_INT8 = 2;
82constexpr uint32_t MSA_L1B_STAGES_FP16 = 3;99constexpr uint32_t MSA_L1B_STAGES_FP16 = 3;
83 100 
@@ -110,9 +127,18 @@ constexpr uint32_t MSA_ATTEN_MASK_SIZE = 2048;
110// FFTS cross-core flag id，合法范围 0~7（8/9/10 为保留 barrier）。127// FFTS cross-core flag id，合法范围 0~7（8/9/10 为保留 barrier）。
111constexpr uint16_t MSA_FLAG_S_READY = 1;128constexpr uint16_t MSA_FLAG_S_READY = 1;
112constexpr uint16_t MSA_FLAG_S_READY_REVERSE = 2;129constexpr uint16_t MSA_FLAG_S_READY_REVERSE = 2;
113-// int8 路径：AIV cast K→scratch 完成后通知 AIC。130+// int8 路径：AIV cast K→scratch 完成后通知 AIC。A2 四槽用 3/4/5/6。
114constexpr uint16_t MSA_FLAG_K_READY = 3;131constexpr uint16_t MSA_FLAG_K_READY = 3;
115constexpr uint16_t MSA_FLAG_K_READY_REVERSE = 4;132constexpr uint16_t MSA_FLAG_K_READY_REVERSE = 4;
133+constexpr uint16_t MSA_FLAG_K_READY_2 = 5;
134+constexpr uint16_t MSA_FLAG_K_READY_3 = 6;
135+ 
136+// Ascend 950 MIX：MODE4 下 AIC 分别给两个 AIV 置位，flagId+16。
137+constexpr uint32_t MSA_A5_SYNC_MODE4 = 4;
138+constexpr uint16_t MSA_A5_AIV_FLAG_OFFSET = 16;
139+constexpr uint16_t MSA_A5_FLAG_CV = 0; // Cube→Vector：S 已在 GM/UB
140+constexpr uint16_t MSA_A5_FLAG_VC = 2; // Vector→Cube：UB slot 可覆盖
141+constexpr uint16_t MSA_A5_FLAG_K = 4; // Vector→Cube：K scratch 就绪
116 142 
117} // namespace MsaIndexScoreNs143} // namespace MsaIndexScoreNs
118 144 
@@ -122,5 +148,8 @@ constexpr uint16_t MSA_FLAG_K_READY_REVERSE = 4;
122#define MSA_TILING_KEY_FP16 1148#define MSA_TILING_KEY_FP16 1
123#define MSA_TILING_KEY_BF16_INT8 2149#define MSA_TILING_KEY_BF16_INT8 2
124#define MSA_TILING_KEY_FP16_INT8 3150#define MSA_TILING_KEY_FP16_INT8 3
151+#define MSA_TILING_KEY_HIFLOAT8 4
152+#define MSA_TILING_KEY_FP8_E5M2 5
153+#define MSA_TILING_KEY_FP8_E4M3FN 6
125 154 
126#endif // MSA_INDEX_SCORE_COMMON_H155#endif // MSA_INDEX_SCORE_COMMON_H
Mdocs/zh/op_list.md+10-0