Star1056
2553
代码
代码
Issues177
Pull Requests413
Actions
项目讨论
Wiki
分析
Star1056
2553
已合并
【feature】【BSA】【950】BSA正向/推理算子量化/非量化kernel支持LSE输出 #6565
monologue815创建于 6月8日
【feature】【BSA】【950】BSA正向/推理算子量化/非量化kernel支持LSE输出 #6565
已合并
monologue815创建于 6月8日
从chenyizhou/ops-transformer-bsa-optbsa_lse合入到cann/ops-transformermaster
共 11 个文件变更+402-105
Mattention/block_sparse_attention/docs/aclnnBlockSparseAttention.md+15-11
@@ -30,8 +30,8 @@
30 - D：表示隐藏层最小的单元尺寸，需满足D=H/N（Head-Dim）30 - D：表示隐藏层最小的单元尺寸，需满足D=H/N（Head-Dim）
31 31 
32 当前支持的布局：32 当前支持的布局：
33- - qInputLayout: "TND" "BNSD"33+ - qInputLayout: "TND" "BNSD" "BSND"
34- - kvInputLayout: "TND" "BNSD"34+ - kvInputLayout: "TND" "BNSD" "BSND"
35 35 
36## 函数原型36## 函数原型
37 37 
@@ -106,7 +106,8 @@ aclnnStatus aclnnBlockSparseAttention(
106 <td>公式中的query。</td>106 <td>公式中的query。</td>
107 <td>支持的shape为：107 <td>支持的shape为：
108 <ul><li>TND: [totalQTokens, headNum, headDim]。</li>108 <ul><li>TND: [totalQTokens, headNum, headDim]。</li>
109- <li>BNSD: [batch, headNum, maxQSeqLength, headDim]。</li></ul>109+ <li>BNSD: [batch, headNum, maxQSeqLength, headDim]。</li>
110+ <li>BSND: [batch, maxQSeqLength, headNum, headDim]。</li></ul>
110 </td>111 </td>
111 <td>FLOAT16、BFLOAT16</td>112 <td>FLOAT16、BFLOAT16</td>
112 <td>ND</td>113 <td>ND</td>
@@ -121,6 +122,7 @@ aclnnStatus aclnnBlockSparseAttention(
121 <ul>122 <ul>
122 <li>TND: [totalKTokens, numKeyValueHeads, headDim]。</li>123 <li>TND: [totalKTokens, numKeyValueHeads, headDim]。</li>
123 <li>BNSD: [batch, numKeyValueHeads, maxKvSeqLength, headDim]。</li>124 <li>BNSD: [batch, numKeyValueHeads, maxKvSeqLength, headDim]。</li>
125+ <li>BSND: [batch, maxKvSeqLength, numKeyValueHeads, headDim]。</li>
124 </ul>126 </ul>
125 </td>127 </td>
126 <td>FLOAT16、BFLOAT16</td>128 <td>FLOAT16、BFLOAT16</td>
@@ -137,6 +139,7 @@ aclnnStatus aclnnBlockSparseAttention(
137 <ul>139 <ul>
138 <li>TND: [totalVTokens, numKeyValueHeads, headDim]。</li>140 <li>TND: [totalVTokens, numKeyValueHeads, headDim]。</li>
139 <li>BNSD: [batch, numKeyValueHeads, maxKvSeqLength, headDim]。</li>141 <li>BNSD: [batch, numKeyValueHeads, maxKvSeqLength, headDim]。</li>
142+ <li>BSND: [batch, maxKvSeqLength, numKeyValueHeads, headDim]。</li>
140 </ul>143 </ul>
141 </td>144 </td>
142 <td>FLOAT16、BFLOAT16</td>145 <td>FLOAT16、BFLOAT16</td>
@@ -200,7 +203,7 @@ aclnnStatus aclnnBlockSparseAttention(
200 可选输入，用于变长序列场景：203 可选输入，用于变长序列场景：
201 <ul>204 <ul>
202 <li>当qInputLayout为"TND"时：该项输入必须配置。</li>205 <li>当qInputLayout为"TND"时：该项输入必须配置。</li>
203- <li>当qInputLayout为"BNSD"时：如配置该项输入，算子内会按该输入指定的实际序列长度进行处理；如不配置该项输入（传入nullptr），算子内会按照query的shape中的S进行处理。</li>206+ <li>当qInputLayout为"BNSD"或"BSND"时：如配置该项输入，算子内会按该输入指定的实际序列长度进行处理；如不配置该项输入（传入nullptr），算子内会按照query的shape中的S进行处理。</li>
204 </ul>207 </ul>
205 </td>208 </td>
206 <td>INT64</td>209 <td>INT64</td>
@@ -216,7 +219,7 @@ aclnnStatus aclnnBlockSparseAttention(
216 可选输入，用于变长序列场景：219 可选输入，用于变长序列场景：
217 <ul>220 <ul>
218 <li>当kvInputLayout为"TND"时：该项输入必须配置。</li>221 <li>当kvInputLayout为"TND"时：该项输入必须配置。</li>
219- <li>当kvInputLayout为"BNSD"时：如配置该项输入，算子内会按该输入指定的实际序列长度进行处理；如不配置该项输入（传入nullptr），算子内会按照key/value的shape中的S进行处理。</li>222+ <li>当kvInputLayout为"BNSD"或"BSND"时：如配置该项输入，算子内会按该输入指定的实际序列长度进行处理；如不配置该项输入（传入nullptr），算子内会按照key/value的shape中的S进行处理。</li>
220 </ul>223 </ul>
221 </td>224 </td>
222 <td>INT64</td>225 <td>INT64</td>
@@ -238,7 +241,7 @@ aclnnStatus aclnnBlockSparseAttention(
238 <td>qInputLayout</td>241 <td>qInputLayout</td>
239 <td>输入</td>242 <td>输入</td>
240 <td>代表输入query的数据排布格式。</td>243 <td>代表输入query的数据排布格式。</td>
241- <td>当前仅支持"TND"和"BNSD"，qInputLayout与kvInputLayout需要保持一致。</td>244+ <td>当前仅支持"TND"和"BNSD"和"BSND"，qInputLayout与kvInputLayout需要保持一致。</td>
242 <td>String</td>245 <td>String</td>
243 <td>-</td>246 <td>-</td>
244 <td>-</td>247 <td>-</td>
@@ -248,7 +251,7 @@ aclnnStatus aclnnBlockSparseAttention(
248 <td>kvInputLayout</td>251 <td>kvInputLayout</td>
249 <td>输入</td>252 <td>输入</td>
250 <td>代表输入key、value的数据排布格式。</td>253 <td>代表输入key、value的数据排布格式。</td>
251- <td>当前仅支持"TND"和"BNSD"，qInputLayout与kvInputLayout需要保持一致。</td>254+ <td>当前仅支持"TND"和"BNSD"和"BSND"，qInputLayout与kvInputLayout需要保持一致。</td>
252 <td>String</td>255 <td>String</td>
253 <td>-</td>256 <td>-</td>
254 <td>-</td>257 <td>-</td>
@@ -341,7 +344,7 @@ aclnnStatus aclnnBlockSparseAttention(
341 <td>输入</td>344 <td>输入</td>
342 <td>是否使能softmaxLse输出的标志位。</td>345 <td>是否使能softmaxLse输出的标志位。</td>
343 <td>346 <td>
344- 当前只支持传0或1。其中，<term>Ascend 950PR/Ascend 950DT</term>仅支持配置为0，<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>支持配置为0或1347+ 只支持传0或1。
345 <ul>348 <ul>
346 <li>0：表示不输出softmaxLse。</li>349 <li>0：表示不输出softmaxLse。</li>
347 <li>1：表示输出softmaxLse，相比不输出softmaxLse可能存在性能损失。</li>350 <li>1：表示输出softmaxLse，相比不输出softmaxLse可能存在性能损失。</li>
@@ -371,6 +374,7 @@ aclnnStatus aclnnBlockSparseAttention(
371 <ul>374 <ul>
372 <li>query为"TND": [totalQTokens, headNum, 1]。</li>375 <li>query为"TND": [totalQTokens, headNum, 1]。</li>
373 <li>query为"BNSD": [batch, headNum, maxQSeqLength, 1]。</li>376 <li>query为"BNSD": [batch, headNum, maxQSeqLength, 1]。</li>
377+ <li>query为"BSND": [batch, maxQSeqLength, headNum, 1]。</li>
374 </ul>378 </ul>
375 </td>379 </td>
376 <td>FLOAT</td>380 <td>FLOAT</td>
@@ -490,8 +494,8 @@ aclnnStatus aclnnBlockSparseAttention(
490- 确定性计算：494- 确定性计算：
491 - aclnnBlockSparseAttention默认确定性实现。495 - aclnnBlockSparseAttention默认确定性实现。
492- 该接口与PyTorch配合使用时，需要保证CANN相关包与PyTorch相关包的版本匹配。496- 该接口与PyTorch配合使用时，需要保证CANN相关包与PyTorch相关包的版本匹配。
493-- qInputLayout当前仅支持"TND"和"BNSD"。497+- qInputLayout当前仅支持"TND"和"BNSD"和"BSND"。
494-- kvInputLayout当前仅支持"TND"和"BNSD"。498+- kvInputLayout当前仅支持"TND"和"BNSD"和"BSND"。
495- 当前query、key、value的InputLayout必须保持一致。499- 当前query、key、value的InputLayout必须保持一致。
496- 输入query、key、value的数据类型必须一致，支持FLOAT16和BFLOAT16。500- 输入query、key、value的数据类型必须一致，支持FLOAT16和BFLOAT16。
497- query、key、value的D轴当前仅支持配置为64或128501- query、key、value的D轴当前仅支持配置为64或128
@@ -502,7 +506,7 @@ aclnnStatus aclnnBlockSparseAttention(
502- actualSeqLengthsOptional与actualSeqLengthsKvOptional当前必须同时配置或同时不配置，仅配置其中之一的行为将被算子拦截。506- actualSeqLengthsOptional与actualSeqLengthsKvOptional当前必须同时配置或同时不配置，仅配置其中之一的行为将被算子拦截。
503- blockTableOptional当前只支持传入nullptr，表示不开启PagedAttention特性。507- blockTableOptional当前只支持传入nullptr，表示不开启PagedAttention特性。
504- innerPrecise必须为0或1或4，其中，<term>Ascend 950PR/Ascend 950DT</term>仅支持配置为4，<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>仅支持配置为0或1。508- innerPrecise必须为0或1或4，其中，<term>Ascend 950PR/Ascend 950DT</term>仅支持配置为4，<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>仅支持配置为0或1。
505-- softmaxLseFlag仅支持配置0或1，分别表示不开启/开启softmaxLse输出。当前，<term>Ascend 950PR/Ascend 950DT</term>仅支持配置为0，<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>支持配置为0或1。509+- softmaxLseFlag仅支持配置0或1，分别表示不开启/开启softmaxLse输出。
506- qSeqlen和kvSeqlen不需要被blockShape整除，支持非对齐场景，实际分块数通过向上取整计算。510- qSeqlen和kvSeqlen不需要被blockShape整除，支持非对齐场景，实际分块数通过向上取整计算。
507- 输入query的headNum为N1，输入key和value的headNum为N2，则N1 >= N2 && N1 % N2 == 0。511- 输入query的headNum为N1，输入key和value的headNum为N2，则N1 >= N2 && N1 % N2 == 0。
508- maskType当前只支持输入0，表示不加mask。512- maskType当前只支持输入0，表示不加mask。
Mattention/block_sparse_attention/docs/aclnnBlockSparseAttentionV2.md+14-10
@@ -32,8 +32,8 @@
32 - D：表示隐藏层最小的单元尺寸，需满足D=H/N（Head-Dim）32 - D：表示隐藏层最小的单元尺寸，需满足D=H/N（Head-Dim）
33 33 
34 当前支持的布局：34 当前支持的布局：
35- - qInputLayout: "TND" "BNSD"35+ - qInputLayout: "TND" "BNSD" "BSND"
36- - kvInputLayout: "TND" "BNSD"36+ - kvInputLayout: "TND" "BNSD" "BSND"
37 37 
38- **FP8特性说明（仅<term>Ascend 950PR/Ascend 950DT</term>支持）**：38- **FP8特性说明（仅<term>Ascend 950PR/Ascend 950DT</term>支持）**：
39 39
@@ -135,6 +135,7 @@ aclnnStatus aclnnBlockSparseAttentionV2(
135 <td>支持的shape为：135 <td>支持的shape为：
136 <ul><li>TND: [totalQTokens, headNum, headDim]。</li>136 <ul><li>TND: [totalQTokens, headNum, headDim]。</li>
137 <li>BNSD: [batch, headNum, maxQSeqLength, headDim]。</li></ul>137 <li>BNSD: [batch, headNum, maxQSeqLength, headDim]。</li></ul>
138+ <li>BSND: [batch, maxQSeqLength, headNum, headDim]。</li></ul>
138 </td>139 </td>
139 <td>FLOAT16、BFLOAT16</td>140 <td>FLOAT16、BFLOAT16</td>
140 <td>ND</td>141 <td>ND</td>
@@ -149,6 +150,7 @@ aclnnStatus aclnnBlockSparseAttentionV2(
149 <ul>150 <ul>
150 <li>TND: [totalKTokens, numKeyValueHeads, headDim]。</li>151 <li>TND: [totalKTokens, numKeyValueHeads, headDim]。</li>
151 <li>BNSD: [batch, numKeyValueHeads, maxKvSeqLength, headDim]。</li>152 <li>BNSD: [batch, numKeyValueHeads, maxKvSeqLength, headDim]。</li>
153+ <li>BSND: [batch, maxKvSeqLength, numKeyValueHeads, headDim]。</li>
152 </ul>154 </ul>
153 </td>155 </td>
154 <td>FLOAT16、BFLOAT16</td>156 <td>FLOAT16、BFLOAT16</td>
@@ -165,6 +167,7 @@ aclnnStatus aclnnBlockSparseAttentionV2(
165 <ul>167 <ul>
166 <li>TND: [totalVTokens, numKeyValueHeads, headDim]。</li>168 <li>TND: [totalVTokens, numKeyValueHeads, headDim]。</li>
167 <li>BNSD: [batch, numKeyValueHeads, maxKvSeqLength, headDim]。</li>169 <li>BNSD: [batch, numKeyValueHeads, maxKvSeqLength, headDim]。</li>
170+ <li>BSND: [batch, maxKvSeqLength, numKeyValueHeads, headDim]。</li>
168 </ul>171 </ul>
169 </td>172 </td>
170 <td>FLOAT16、BFLOAT16</td>173 <td>FLOAT16、BFLOAT16</td>
@@ -228,7 +231,7 @@ aclnnStatus aclnnBlockSparseAttentionV2(
228 可选输入，用于变长序列场景：231 可选输入，用于变长序列场景：
229 <ul>232 <ul>
230 <li>当qInputLayout为"TND"时：该项输入必须配置。</li>233 <li>当qInputLayout为"TND"时：该项输入必须配置。</li>
231- <li>当qInputLayout为"BNSD"时：如配置该项输入，算子内会按该输入指定的实际序列长度进行处理；如不配置该项输入（传入nullptr），算子内会按照query的shape中的S进行处理。</li>234+ <li>当qInputLayout为"BNSD"或"BSND"时：如配置该项输入，算子内会按该输入指定的实际序列长度进行处理；如不配置该项输入（传入nullptr），算子内会按照query的shape中的S进行处理。</li>
232 </ul>235 </ul>
233 </td>236 </td>
234 <td>INT64</td>237 <td>INT64</td>
@@ -244,7 +247,7 @@ aclnnStatus aclnnBlockSparseAttentionV2(
244 可选输入，用于变长序列场景：247 可选输入，用于变长序列场景：
245 <ul>248 <ul>
246 <li>当kvInputLayout为"TND"时：该项输入必须配置。</li>249 <li>当kvInputLayout为"TND"时：该项输入必须配置。</li>
247- <li>当kvInputLayout为"BNSD"时：如配置该项输入，算子内会按该输入指定的实际序列长度进行处理；如不配置该项输入（传入nullptr），算子内会按照key/value的shape中的S进行处理。</li>250+ <li>当kvInputLayout为"BNSD"或"BSND"时：如配置该项输入，算子内会按该输入指定的实际序列长度进行处理；如不配置该项输入（传入nullptr），算子内会按照key/value的shape中的S进行处理。</li>
248 </ul>251 </ul>
249 </td>252 </td>
250 <td>INT64</td>253 <td>INT64</td>
@@ -326,7 +329,7 @@ aclnnStatus aclnnBlockSparseAttentionV2(
326 <td>qInputLayout（char*）</td>329 <td>qInputLayout（char*）</td>
327 <td>输入</td>330 <td>输入</td>
328 <td>代表输入query的数据排布格式。</td>331 <td>代表输入query的数据排布格式。</td>
329- <td>当前仅支持"TND"和"BNSD"，qInputLayout与kvInputLayout需要保持一致。</td>332+ <td>当前仅支持"TND"和"BNSD"和"BSND"，qInputLayout与kvInputLayout需要保持一致。</td>
330 <td>String</td>333 <td>String</td>
331 <td>-</td>334 <td>-</td>
332 <td>-</td>335 <td>-</td>
@@ -336,7 +339,7 @@ aclnnStatus aclnnBlockSparseAttentionV2(
336 <td>kvInputLayout（char*）</td>339 <td>kvInputLayout（char*）</td>
337 <td>输入</td>340 <td>输入</td>
338 <td>代表输入key、value的数据排布格式。</td>341 <td>代表输入key、value的数据排布格式。</td>
339- <td>当前仅支持"TND"和"BNSD"，qInputLayout与kvInputLayout需要保持一致。</td>342+ <td>当前仅支持"TND"和"BNSD"和"BSND"，qInputLayout与kvInputLayout需要保持一致。</td>
340 <td>String</td>343 <td>String</td>
341 <td>-</td>344 <td>-</td>
342 <td>-</td>345 <td>-</td>
@@ -429,7 +432,7 @@ aclnnStatus aclnnBlockSparseAttentionV2(
429 <td>输入</td>432 <td>输入</td>
430 <td>是否使能softmaxLse输出的标志位。</td>433 <td>是否使能softmaxLse输出的标志位。</td>
431 <td>434 <td>
432- 当前只支持传0或1。其中，<term>Ascend 950PR/Ascend 950DT</term>仅支持配置为0，<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>支持配置为0或1435+ 当前只支持传0或1。
433 <ul>436 <ul>
434 <li>0：表示不输出softmaxLse。</li>437 <li>0：表示不输出softmaxLse。</li>
435 <li>1：表示输出softmaxLse，相比不输出softmaxLse可能存在性能损失。</li>438 <li>1：表示输出softmaxLse，相比不输出softmaxLse可能存在性能损失。</li>
@@ -459,6 +462,7 @@ aclnnStatus aclnnBlockSparseAttentionV2(
459 <ul>462 <ul>
460 <li>query为"TND": [totalQTokens, headNum, 1]。</li>463 <li>query为"TND": [totalQTokens, headNum, 1]。</li>
461 <li>query为"BNSD": [batch, headNum, maxQSeqLength, 1]。</li>464 <li>query为"BNSD": [batch, headNum, maxQSeqLength, 1]。</li>
465+ <li>query为"BNSD": [batch, maxQSeqLength, headNum, 1]。</li>
462 </ul>466 </ul>
463 </td>467 </td>
464 <td>FLOAT</td>468 <td>FLOAT</td>
@@ -578,8 +582,8 @@ aclnnStatus aclnnBlockSparseAttentionV2(
578- 确定性计算：582- 确定性计算：
579 - aclnnBlockSparseAttentionV2默认确定性实现。583 - aclnnBlockSparseAttentionV2默认确定性实现。
580- 该接口与PyTorch配合使用时，需要保证CANN相关包与PyTorch相关包的版本匹配。584- 该接口与PyTorch配合使用时，需要保证CANN相关包与PyTorch相关包的版本匹配。
581-- qInputLayout当前仅支持"TND"和"BNSD"。585+- qInputLayout当前仅支持"TND"和"BNSD"和"BSND"。
582-- kvInputLayout当前仅支持"TND"和"BNSD"。586+- kvInputLayout当前仅支持"TND"和"BNSD"和"BSND"。
583- 当前query、key、value的InputLayout必须保持一致。587- 当前query、key、value的InputLayout必须保持一致。
584- 输入query、key、value的数据类型必须一致，支持FLOAT16和BFLOAT16。588- 输入query、key、value的数据类型必须一致，支持FLOAT16和BFLOAT16。
585- query、key、value的D轴当前仅支持配置为64或128589- query、key、value的D轴当前仅支持配置为64或128
@@ -590,7 +594,7 @@ aclnnStatus aclnnBlockSparseAttentionV2(
590- actualSeqLengthsOptional与actualSeqLengthsKvOptional当前必须同时配置或同时不配置，仅配置其中之一的行为将被算子拦截。594- actualSeqLengthsOptional与actualSeqLengthsKvOptional当前必须同时配置或同时不配置，仅配置其中之一的行为将被算子拦截。
591- blockTableOptional当前只支持传入nullptr，表示不开启PagedAttention特性。595- blockTableOptional当前只支持传入nullptr，表示不开启PagedAttention特性。
592- innerPrecise仅支持配置4，表示混合精度运算，在性能与精度上取得一个折中。596- innerPrecise仅支持配置4，表示混合精度运算，在性能与精度上取得一个折中。
593-- softmaxLseFlag仅支持配置0或1，分别表示不开启/开启softmaxLse输出。当前，<term>Ascend 950PR/Ascend 950DT</term>仅支持配置为0，<term>Atlas A2 训练系列产品/Atlas A2 推理系列产品</term>、<term>Atlas A3 训练系列产品/Atlas A3 推理系列产品</term>支持配置为0或1。597+- softmaxLseFlag仅支持配置0或1，分别表示不开启/开启softmaxLse输出。
594- qSeqlen和kvSeqlen不需要被blockShape整除，支持非对齐场景，实际分块数通过向上取整计算。598- qSeqlen和kvSeqlen不需要被blockShape整除，支持非对齐场景，实际分块数通过向上取整计算。
595- 输入query的headNum为N1，输入key和value的headNum为N2，则N1 >= N2 && N1 % N2 == 0。599- 输入query的headNum为N1，输入key和value的headNum为N2，则N1 >= N2 && N1 % N2 == 0。
596- maskType当前只支持输入0，表示不加mask。600- maskType当前只支持输入0，表示不加mask。
Mattention/block_sparse_attention/op_host/block_sparse_attention_tiling.cpp+0-4
@@ -794,10 +794,6 @@ ge::graphStatus BSATiling::ParseAttrs(gert::TilingContext *bsaContext)
794 OP_LOGE(bsaContext->GetNodeName(), "Attr softmaxLseFlag is nullptr.");794 OP_LOGE(bsaContext->GetNodeName(), "Attr softmaxLseFlag is nullptr.");
795 return ge::GRAPH_FAILED;795 return ge::GRAPH_FAILED;
796 } else if (*softmaxLsePtr == LSE_OUT) {796 } else if (*softmaxLsePtr == LSE_OUT) {
797- if (socVer_ == SOC_VER_950_CODE) {
798- OP_LOGE(bsaContext->GetNodeName(), "Attr softmaxLseFlag must be 0 on chip 950.");
799- return ge::GRAPH_FAILED;
800- }
801 softmaxLseFlag_ = true;797 softmaxLseFlag_ = true;
802 } else if (*softmaxLsePtr == LSE_NO_OUT) {798 } else if (*softmaxLsePtr == LSE_NO_OUT) {
803 softmaxLseFlag_ = false;799 softmaxLseFlag_ = false;
Mattention/block_sparse_attention/op_kernel/arch35/block_sparse_attention_kernel_arch35_full_quant.h+33-4
@@ -35,6 +35,7 @@ public:
35 using ElementV = typename BlockMmadPV::ElementB; // fp8_e4m3fn_t35 using ElementV = typename BlockMmadPV::ElementB; // fp8_e4m3fn_t
36 using ElementOTmp = typename BlockMmadPV::ElementC; // float36 using ElementOTmp = typename BlockMmadPV::ElementC; // float
37 using ElementO = typename EpilogueOnlineSoftmax::ElementInput; // half/bloat16_t37 using ElementO = typename EpilogueOnlineSoftmax::ElementInput; // half/bloat16_t
38+ using ElementLse = typename EpilogueRescaleO::ElementLse;
38 using ElementSparseMask = typename EpilogueMask2Idx::ElementSparseMask;39 using ElementSparseMask = typename EpilogueMask2Idx::ElementSparseMask;
39 using ElementSparseIdx = typename EpilogueMask2Idx::ElementSparseIdx;40 using ElementSparseIdx = typename EpilogueMask2Idx::ElementSparseIdx;
40 using ElementSparseCount = typename EpilogueMask2Idx::ElementSparseCount;41 using ElementSparseCount = typename EpilogueMask2Idx::ElementSparseCount;
@@ -45,6 +46,7 @@ public:
45 using LayoutP = layout::RowMajor;46 using LayoutP = layout::RowMajor;
46 using LayoutV = layout::RowMajor;47 using LayoutV = layout::RowMajor;
47 using LayoutO = layout::RowMajor;48 using LayoutO = layout::RowMajor;
49+ using LayoutLse = layout::RowMajor;
48 using LayoutOTmp = layout::RowMajor;50 using LayoutOTmp = layout::RowMajor;
49 using LayoutSparseIdx = layout::RowMajor;51 using LayoutSparseIdx = layout::RowMajor;
50 using LayoutSparseCount = layout::RowMajor;52 using LayoutSparseCount = layout::RowMajor;
@@ -87,6 +89,8 @@ public:
87 gVDequantScale.SetGlobalBuffer((__gm__ float *)params.vDequantScale);89 gVDequantScale.SetGlobalBuffer((__gm__ float *)params.vDequantScale);
88 AscendC::GlobalTensor<ElementO> gO;90 AscendC::GlobalTensor<ElementO> gO;
89 gO.SetGlobalBuffer((__gm__ ElementO *)params.attentionOut);91 gO.SetGlobalBuffer((__gm__ ElementO *)params.attentionOut);
92+ AscendC::GlobalTensor<ElementLse> gLse;
93+ gLse.SetGlobalBuffer((__gm__ ElementLse *)params.lse);
90 AscendC::GlobalTensor<ElementSparseIdx> gSparseIdx;94 AscendC::GlobalTensor<ElementSparseIdx> gSparseIdx;
91 gSparseIdx.SetGlobalBuffer((__gm__ ElementSparseIdx *)params.workSpace);95 gSparseIdx.SetGlobalBuffer((__gm__ ElementSparseIdx *)params.workSpace);
92 AscendC::GlobalTensor<ElementSparseCount> gSparseCount;96 AscendC::GlobalTensor<ElementSparseCount> gSparseCount;
@@ -126,23 +130,34 @@ public:
126 // For BSND: [B, S, N, D], strideB = S * B * D, strideS = N * D, strideN = D130 // For BSND: [B, S, N, D], strideB = S * B * D, strideS = N * D, strideN = D
127 int64_t strideQO = 0;131 int64_t strideQO = 0;
128 int64_t strideKV = 0;132 int64_t strideKV = 0;
133+ int64_t strideLse = 0;
129 int64_t strideQOB = 0; // BNSD/BSND batch_ stride for Q134 int64_t strideQOB = 0; // BNSD/BSND batch_ stride for Q
130 int64_t strideQON = 0; // BNSD/BSND head stride for Q135 int64_t strideQON = 0; // BNSD/BSND head stride for Q
131 int64_t strideQOS = 0; // BNSD/BSND seq stride for Q136 int64_t strideQOS = 0; // BNSD/BSND seq stride for Q
132 int64_t strideKVB = 0; // BNSD/BSND batch_ stride for KV137 int64_t strideKVB = 0; // BNSD/BSND batch_ stride for KV
133 int64_t strideKVN = 0; // BNSD/BSND head stride for KV138 int64_t strideKVN = 0; // BNSD/BSND head stride for KV
134 int64_t strideKVS = 0; // BNSD/BSND seq stride for KV139 int64_t strideKVS = 0; // BNSD/BSND seq stride for KV
140+ int64_t strideLseB = 0;
141+ int64_t strideLseN = 0;
142+ int64_t strideLseS = 0;
135 143 
136 if constexpr (qFormat == Format::TND) {144 if constexpr (qFormat == Format::TND) {
137 strideQO = qHeads_ * embed_;145 strideQO = qHeads_ * embed_;
146+ strideLse = qHeads_;
138 } else if constexpr (qFormat == Format::BNSD) {147 } else if constexpr (qFormat == Format::BNSD) {
139- strideQOB = qHeads_ * qSeqlenAligned_ * embed_; // batch_ stride148+ strideQOB = qHeads_ * qSeqlenAligned_ * embed_; // batch_ stride
140- strideQON = qSeqlenAligned_ * embed_; // head stride149+ strideQON = qSeqlenAligned_ * embed_; // head stride
141- strideQOS = embed_; // seq stride150+ strideQOS = embed_; // seq stride
151+ strideLseB = qHeads_ * qSeqlenAligned_;
152+ strideLseN = qSeqlenAligned_;
153+ strideLseS = 1;
142 } else if constexpr (qFormat == Format::BSND) {154 } else if constexpr (qFormat == Format::BSND) {
143 strideQOB = qSeqlenAligned_ * qHeads_ * embed_; // batch_ stride155 strideQOB = qSeqlenAligned_ * qHeads_ * embed_; // batch_ stride
144 strideQOS = qHeads_ * embed_; // seq stride156 strideQOS = qHeads_ * embed_; // seq stride
145 strideQON = embed_; // head stride157 strideQON = embed_; // head stride
158+ strideLseB = qSeqlenAligned_ * qHeads_;
159+ strideLseS = qHeads_;
160+ strideLseN = 1;
146 }161 }
147 162 
148 if constexpr (kvFormat == Format::TND) {163 if constexpr (kvFormat == Format::TND) {
@@ -163,6 +178,7 @@ public:
163 int64_t kBOffset = 0;178 int64_t kBOffset = 0;
164 int64_t vBOffset = 0;179 int64_t vBOffset = 0;
165 int64_t oBOffset = 0;180 int64_t oBOffset = 0;
181+ int64_t lseBOffset = 0;
166 uint32_t preTotalTaskNum = 0;182 uint32_t preTotalTaskNum = 0;
167 uint32_t curBatch = 0;183 uint32_t curBatch = 0;
168 int64_t qSeqlen = actSeqAval_ ? gActualQseqlen.GetValue(curBatch) : qSeqlenAligned_;184 int64_t qSeqlen = actSeqAval_ ? gActualQseqlen.GetValue(curBatch) : qSeqlenAligned_;
@@ -177,6 +193,7 @@ public:
177 if constexpr (qFormat == Format::TND) {193 if constexpr (qFormat == Format::TND) {
178 qBOffset += qSeqlen * strideQO;194 qBOffset += qSeqlen * strideQO;
179 oBOffset += qSeqlen * strideQO;195 oBOffset += qSeqlen * strideQO;
196+ lseBOffset += qSeqlen * strideLse;
180 }197 }
181 if constexpr (kvFormat == Format::TND) {198 if constexpr (kvFormat == Format::TND) {
182 kBOffset += kvSeqlen * strideKV;199 kBOffset += kvSeqlen * strideKV;
@@ -207,21 +224,27 @@ public:
207 int64_t gmOffsetK = 0;224 int64_t gmOffsetK = 0;
208 int64_t gmOffsetV = 0;225 int64_t gmOffsetV = 0;
209 int64_t gmOffsetO = 0;226 int64_t gmOffsetO = 0;
227+ int64_t gmOffsetLse = 0;
210 int64_t qSOffset = xBlockIdx * blockShapeX_ + qSTileIdxCurXBlock * qBaseTile_;228 int64_t qSOffset = xBlockIdx * blockShapeX_ + qSTileIdxCurXBlock * qBaseTile_;
211 229 
212 if constexpr (qFormat == Format::TND) {230 if constexpr (qFormat == Format::TND) {
213 gmOffsetQ = qBOffset + qSOffset * strideQO + qHeadIdx * embed_;231 gmOffsetQ = qBOffset + qSOffset * strideQO + qHeadIdx * embed_;
214 gmOffsetO = oBOffset + qSOffset * strideQO + qHeadIdx * embed_;232 gmOffsetO = oBOffset + qSOffset * strideQO + qHeadIdx * embed_;
233+ gmOffsetLse = lseBOffset + qSOffset * strideLse + qHeadIdx;
215 } else if constexpr (qFormat == Format::BNSD) {234 } else if constexpr (qFormat == Format::BNSD) {
216 qBOffset = curBatch * strideQOB;235 qBOffset = curBatch * strideQOB;
217 oBOffset = curBatch * strideQOB;236 oBOffset = curBatch * strideQOB;
237+ lseBOffset = curBatch * strideLseB;
218 gmOffsetQ = qBOffset + qHeadIdx * strideQON + qSOffset * strideQOS;238 gmOffsetQ = qBOffset + qHeadIdx * strideQON + qSOffset * strideQOS;
219 gmOffsetO = oBOffset + qHeadIdx * strideQON + qSOffset * strideQOS;239 gmOffsetO = oBOffset + qHeadIdx * strideQON + qSOffset * strideQOS;
240+ gmOffsetLse = lseBOffset + qHeadIdx * strideLseN + qSOffset * strideLseS;
220 } else if constexpr (qFormat == Format::BSND) {241 } else if constexpr (qFormat == Format::BSND) {
221 qBOffset = curBatch * strideQOB;242 qBOffset = curBatch * strideQOB;
222 oBOffset = curBatch * strideQOB;243 oBOffset = curBatch * strideQOB;
244+ lseBOffset = curBatch * strideLseB;
223 gmOffsetQ = qBOffset + qSOffset * strideQOS + qHeadIdx * strideQON;245 gmOffsetQ = qBOffset + qSOffset * strideQOS + qHeadIdx * strideQON;
224 gmOffsetO = oBOffset + qSOffset * strideQOS + qHeadIdx * strideQON;246 gmOffsetO = oBOffset + qSOffset * strideQOS + qHeadIdx * strideQON;
247+ gmOffsetLse = lseBOffset + qSOffset * strideLseS + qHeadIdx * strideLseN;
225 }248 }
226 249 
227 if constexpr (kvFormat == Format::TND) {250 if constexpr (kvFormat == Format::TND) {
@@ -301,16 +324,22 @@ public:
301#endif324#endif
302#ifdef __DAV_VEC__325#ifdef __DAV_VEC__
303 uint32_t oShapeCol = 0;326 uint32_t oShapeCol = 0;
327+ uint32_t lseShapeCol = 0;
304 if constexpr (qFormat == Format::TND) {328 if constexpr (qFormat == Format::TND) {
305 oShapeCol = strideQO;329 oShapeCol = strideQO;
330+ lseShapeCol = strideLse;
306 } else if constexpr (qFormat == Format::BNSD) {331 } else if constexpr (qFormat == Format::BNSD) {
307 oShapeCol = strideQOS;332 oShapeCol = strideQOS;
333+ lseShapeCol = strideLseS;
308 } else if constexpr (qFormat == Format::BSND) {334 } else if constexpr (qFormat == Format::BSND) {
309 oShapeCol = strideQOS;335 oShapeCol = strideQOS;
336+ lseShapeCol = strideLseS;
310 }337 }
311 338 
312 auto gmOLayoutTla = tla::MakeLayout<ElementO, LayoutO>(qBaseTile_, oShapeCol);339 auto gmOLayoutTla = tla::MakeLayout<ElementO, LayoutO>(qBaseTile_, oShapeCol);
313 auto gmOTensorTla = tla::MakeTensor(gO[gmOffsetO], gmOLayoutTla, Arch::PositionGM{});340 auto gmOTensorTla = tla::MakeTensor(gO[gmOffsetO], gmOLayoutTla, Arch::PositionGM{});
341+ auto gmLseLayoutTla = tla::MakeLayout<ElementLse, LayoutLse>(qBaseTile_, lseShapeCol);
342+ auto gmLseTensorTla = tla::MakeTensor(gLse[gmOffsetLse], gmLseLayoutTla, Arch::PositionGM{});
314#endif343#endif
315 for (uint32_t gatheredKvSTileIdx = 0; gatheredKvSTileIdx < kvSLoopNum + PRE_LAUNCH; gatheredKvSTileIdx++) {344 for (uint32_t gatheredKvSTileIdx = 0; gatheredKvSTileIdx < kvSLoopNum + PRE_LAUNCH; gatheredKvSTileIdx++) {
316 if (gatheredKvSTileIdx < kvSLoopNum) {345 if (gatheredKvSTileIdx < kvSLoopNum) {
@@ -406,7 +435,7 @@ public:
406 Arch::CrossCoreFlag mm2ToReFlag(Mm2ToReFlagId);435 Arch::CrossCoreFlag mm2ToReFlag(Mm2ToReFlagId);
407 uint32_t curTileMod = gatheredKvSTileIdxDe % (PRE_LAUNCH + 1);436 uint32_t curTileMod = gatheredKvSTileIdxDe % (PRE_LAUNCH + 1);
408 epilogueRescaleO(437 epilogueRescaleO(
409- gmOTensorTla, actualBlockShapePV,438+ gmOTensorTla, gmLseTensorTla, actualBlockShapePV,
410 curTileMod, gatheredKvSTileIdxDe,439 curTileMod, gatheredKvSTileIdxDe,
411 (gatheredKvSTileIdxDe == 0),440 (gatheredKvSTileIdxDe == 0),
412 (gatheredKvSTileIdxDe == kvSLoopNum - 1),441 (gatheredKvSTileIdxDe == kvSLoopNum - 1),
Mattention/block_sparse_attention/op_kernel/arch35/block_sparse_attention_kernel_arch35_regular.h+34-6
@@ -35,6 +35,7 @@ public:
35 using ElementV = typename BlockMmadPV::ElementB;35 using ElementV = typename BlockMmadPV::ElementB;
36 using ElementOTmp = typename BlockMmadPV::ElementC;36 using ElementOTmp = typename BlockMmadPV::ElementC;
37 using ElementO = typename BlockMmadQK::ElementA;37 using ElementO = typename BlockMmadQK::ElementA;
38+ using ElementLse = typename EpilogueRescaleO::ElementLse;
38 using ElementSparseMask = typename EpilogueMask2Idx::ElementSparseMask;39 using ElementSparseMask = typename EpilogueMask2Idx::ElementSparseMask;
39 using ElementSparseIdx = typename EpilogueMask2Idx::ElementSparseIdx;40 using ElementSparseIdx = typename EpilogueMask2Idx::ElementSparseIdx;
40 using ElementSparseCount = typename EpilogueMask2Idx::ElementSparseCount;41 using ElementSparseCount = typename EpilogueMask2Idx::ElementSparseCount;
@@ -45,6 +46,7 @@ public:
45 using LayoutP = layout::RowMajor;46 using LayoutP = layout::RowMajor;
46 using LayoutV = layout::RowMajor;47 using LayoutV = layout::RowMajor;
47 using LayoutO = layout::RowMajor;48 using LayoutO = layout::RowMajor;
49+ using LayoutLse = layout::RowMajor;
48 using LayoutOTmp = layout::RowMajor;50 using LayoutOTmp = layout::RowMajor;
49 using LayoutSparseIdx = layout::RowMajor;51 using LayoutSparseIdx = layout::RowMajor;
50 using LayoutSparseCount = layout::RowMajor;52 using LayoutSparseCount = layout::RowMajor;
@@ -81,6 +83,8 @@ public:
81 gBlockSparseMask.SetGlobalBuffer((__gm__ uint8_t *)params.blockSparseMask);83 gBlockSparseMask.SetGlobalBuffer((__gm__ uint8_t *)params.blockSparseMask);
82 AscendC::GlobalTensor<ElementO> gO;84 AscendC::GlobalTensor<ElementO> gO;
83 gO.SetGlobalBuffer((__gm__ ElementO *)params.o);85 gO.SetGlobalBuffer((__gm__ ElementO *)params.o);
86+ AscendC::GlobalTensor<ElementLse> gLse;
87+ gLse.SetGlobalBuffer((__gm__ ElementLse *)params.lse);
84 AscendC::GlobalTensor<ElementSparseIdx> gSparseIdx;88 AscendC::GlobalTensor<ElementSparseIdx> gSparseIdx;
85 gSparseIdx.SetGlobalBuffer((__gm__ ElementSparseIdx *)params.workSpace);89 gSparseIdx.SetGlobalBuffer((__gm__ ElementSparseIdx *)params.workSpace);
86 AscendC::GlobalTensor<ElementSparseCount> gSparseCount;90 AscendC::GlobalTensor<ElementSparseCount> gSparseCount;
@@ -120,23 +124,33 @@ public:
120 // For BSND: [B, S, N, D], strideB = S * B * D, strideS = N * D, strideN = D124 // For BSND: [B, S, N, D], strideB = S * B * D, strideS = N * D, strideN = D
121 int64_t strideQO = 0;125 int64_t strideQO = 0;
122 int64_t strideKV = 0;126 int64_t strideKV = 0;
127+ int64_t strideLse = 0;
123 int64_t strideQOB = 0; // BNSD/BSND batch_ stride for Q128 int64_t strideQOB = 0; // BNSD/BSND batch_ stride for Q
124 int64_t strideQON = 0; // BNSD/BSND head stride for Q129 int64_t strideQON = 0; // BNSD/BSND head stride for Q
125 int64_t strideQOS = 0; // BNSD/BSND seq stride for Q130 int64_t strideQOS = 0; // BNSD/BSND seq stride for Q
126 int64_t strideKVB = 0; // BNSD/BSND batch_ stride for KV131 int64_t strideKVB = 0; // BNSD/BSND batch_ stride for KV
127 int64_t strideKVN = 0; // BNSD/BSND head stride for KV132 int64_t strideKVN = 0; // BNSD/BSND head stride for KV
128 int64_t strideKVS = 0; // BNSD/BSND seq stride for KV133 int64_t strideKVS = 0; // BNSD/BSND seq stride for KV
129- 134+ int64_t strideLseB = 0;
135+ int64_t strideLseN = 0;
136+ int64_t strideLseS = 0;
130 if constexpr (qFormat == Format::TND) {137 if constexpr (qFormat == Format::TND) {
131 strideQO = qHeads_ * embed_;138 strideQO = qHeads_ * embed_;
139+ strideLse = qHeads_;
132 } else if constexpr (qFormat == Format::BNSD) {140 } else if constexpr (qFormat == Format::BNSD) {
133- strideQOB = qHeads_ * qSeqlenAligned_ * embed_; // batch_ stride141+ strideQOB = qHeads_ * qSeqlenAligned_ * embed_; // batch_ stride
134- strideQON = qSeqlenAligned_ * embed_; // head stride142+ strideQON = qSeqlenAligned_ * embed_; // head stride
135- strideQOS = embed_; // seq stride143+ strideQOS = embed_; // seq stride
144+ strideLseB = qHeads_ * qSeqlenAligned_;
145+ strideLseN = qSeqlenAligned_;
146+ strideLseS = 1;
136 } else if constexpr (qFormat == Format::BSND) {147 } else if constexpr (qFormat == Format::BSND) {
137 strideQOB = qSeqlenAligned_ * qHeads_ * embed_; // batch_ stride148 strideQOB = qSeqlenAligned_ * qHeads_ * embed_; // batch_ stride
138 strideQOS = qHeads_ * embed_; // seq stride149 strideQOS = qHeads_ * embed_; // seq stride
139 strideQON = embed_; // head stride150 strideQON = embed_; // head stride
151+ strideLseB = qSeqlenAligned_ * qHeads_;
152+ strideLseS = qHeads_;
153+ strideLseN = 1;
140 }154 }
141 155 
142 if constexpr (kvFormat == Format::TND) {156 if constexpr (kvFormat == Format::TND) {
@@ -157,6 +171,7 @@ public:
157 int64_t kBOffset = 0;171 int64_t kBOffset = 0;
158 int64_t vBOffset = 0;172 int64_t vBOffset = 0;
159 int64_t oBOffset = 0;173 int64_t oBOffset = 0;
174+ int64_t lseBOffset = 0;
160 uint32_t preTotalTaskNum = 0;175 uint32_t preTotalTaskNum = 0;
161 uint32_t curBatch = 0;176 uint32_t curBatch = 0;
162 int64_t qSeqlen = actSeqAval_ ? gActualQseqlen.GetValue(curBatch) : qSeqlenAligned_;177 int64_t qSeqlen = actSeqAval_ ? gActualQseqlen.GetValue(curBatch) : qSeqlenAligned_;
@@ -171,6 +186,7 @@ public:
171 if constexpr (qFormat == Format::TND) {186 if constexpr (qFormat == Format::TND) {
172 qBOffset += qSeqlen * strideQO;187 qBOffset += qSeqlen * strideQO;
173 oBOffset += qSeqlen * strideQO;188 oBOffset += qSeqlen * strideQO;
189+ lseBOffset += qSeqlen * strideLse;
174 }190 }
175 if constexpr (kvFormat == Format::TND) {191 if constexpr (kvFormat == Format::TND) {
176 kBOffset += kvSeqlen * strideKV;192 kBOffset += kvSeqlen * strideKV;
@@ -196,26 +212,32 @@ public:
196 int64_t sparseMaskNOffset = qHeadIdx * xBlockNumAligned_ * yBlockNumAligned_;212 int64_t sparseMaskNOffset = qHeadIdx * xBlockNumAligned_ * yBlockNumAligned_;
197 int64_t sparseMaskXOffset = xBlockIdx * yBlockNumAligned_;213 int64_t sparseMaskXOffset = xBlockIdx * yBlockNumAligned_;
198 int64_t gmOffsetSparseMask = sparseMaskBOffset + sparseMaskNOffset + sparseMaskXOffset;214 int64_t gmOffsetSparseMask = sparseMaskBOffset + sparseMaskNOffset + sparseMaskXOffset;
199- // corresponding Q/K/V/O gm offset of cur task215+ // corresponding Q/K/V/O/Lse gm offset of cur task
200 int64_t gmOffsetQ = 0;216 int64_t gmOffsetQ = 0;
201 int64_t gmOffsetK = 0;217 int64_t gmOffsetK = 0;
202 int64_t gmOffsetV = 0;218 int64_t gmOffsetV = 0;
203 int64_t gmOffsetO = 0;219 int64_t gmOffsetO = 0;
220+ int64_t gmOffsetLse = 0;
204 int64_t qSOffset = xBlockIdx * blockShapeX_ + qSTileIdxCurXBlock * qBaseTile_;221 int64_t qSOffset = xBlockIdx * blockShapeX_ + qSTileIdxCurXBlock * qBaseTile_;
205 222 
206 if constexpr (qFormat == Format::TND) {223 if constexpr (qFormat == Format::TND) {
207 gmOffsetQ = qBOffset + qSOffset * strideQO + qHeadIdx * embed_;224 gmOffsetQ = qBOffset + qSOffset * strideQO + qHeadIdx * embed_;
208 gmOffsetO = oBOffset + qSOffset * strideQO + qHeadIdx * embed_;225 gmOffsetO = oBOffset + qSOffset * strideQO + qHeadIdx * embed_;
226+ gmOffsetLse = lseBOffset + qSOffset * strideLse + qHeadIdx;
209 } else if constexpr (qFormat == Format::BNSD) {227 } else if constexpr (qFormat == Format::BNSD) {
210 qBOffset = curBatch * strideQOB;228 qBOffset = curBatch * strideQOB;
211 oBOffset = curBatch * strideQOB;229 oBOffset = curBatch * strideQOB;
230+ lseBOffset = curBatch * strideLseB;
212 gmOffsetQ = qBOffset + qHeadIdx * strideQON + qSOffset * strideQOS;231 gmOffsetQ = qBOffset + qHeadIdx * strideQON + qSOffset * strideQOS;
213 gmOffsetO = oBOffset + qHeadIdx * strideQON + qSOffset * strideQOS;232 gmOffsetO = oBOffset + qHeadIdx * strideQON + qSOffset * strideQOS;
233+ gmOffsetLse = lseBOffset + qHeadIdx * strideLseN + qSOffset * strideLseS;
214 } else if constexpr (qFormat == Format::BSND) {234 } else if constexpr (qFormat == Format::BSND) {
215 qBOffset = curBatch * strideQOB;235 qBOffset = curBatch * strideQOB;
216 oBOffset = curBatch * strideQOB;236 oBOffset = curBatch * strideQOB;
237+ lseBOffset = curBatch * strideLseB;
217 gmOffsetQ = qBOffset + qSOffset * strideQOS + qHeadIdx * strideQON;238 gmOffsetQ = qBOffset + qSOffset * strideQOS + qHeadIdx * strideQON;
218 gmOffsetO = oBOffset + qSOffset * strideQOS + qHeadIdx * strideQON;239 gmOffsetO = oBOffset + qSOffset * strideQOS + qHeadIdx * strideQON;
240+ gmOffsetLse = lseBOffset + qSOffset * strideLseS + qHeadIdx * strideLseN;
219 }241 }
220 242 
221 if constexpr (kvFormat == Format::TND) {243 if constexpr (kvFormat == Format::TND) {
@@ -291,16 +313,22 @@ public:
291#endif313#endif
292#ifdef __DAV_VEC__314#ifdef __DAV_VEC__
293 uint32_t oShapeCol = 0;315 uint32_t oShapeCol = 0;
316+ uint32_t lseShapeCol = 0;
294 if constexpr (qFormat == Format::TND) {317 if constexpr (qFormat == Format::TND) {
295 oShapeCol = strideQO;318 oShapeCol = strideQO;
319+ lseShapeCol = strideLse;
296 } else if constexpr (qFormat == Format::BNSD) {320 } else if constexpr (qFormat == Format::BNSD) {
297 oShapeCol = strideQOS;321 oShapeCol = strideQOS;
322+ lseShapeCol = strideLseS;
298 } else if constexpr (qFormat == Format::BSND) {323 } else if constexpr (qFormat == Format::BSND) {
299 oShapeCol = strideQOS;324 oShapeCol = strideQOS;
325+ lseShapeCol = strideLseS;
300 }326 }
301 327 
302 auto gmOLayoutTla = tla::MakeLayout<ElementO, LayoutO>(qBaseTile_, oShapeCol);328 auto gmOLayoutTla = tla::MakeLayout<ElementO, LayoutO>(qBaseTile_, oShapeCol);
303 auto gmOTensorTla = tla::MakeTensor(gO[gmOffsetO], gmOLayoutTla, Arch::PositionGM{});329 auto gmOTensorTla = tla::MakeTensor(gO[gmOffsetO], gmOLayoutTla, Arch::PositionGM{});
330+ auto gmLseLayoutTla = tla::MakeLayout<ElementLse, LayoutLse>(qBaseTile_, lseShapeCol);
331+ auto gmLseTensorTla = tla::MakeTensor(gLse[gmOffsetLse], gmLseLayoutTla, Arch::PositionGM{});
304#endif332#endif
305 for (uint32_t gatheredKvSTileIdx = 0; gatheredKvSTileIdx < kvSLoopNum + PRE_LAUNCH; gatheredKvSTileIdx++) {333 for (uint32_t gatheredKvSTileIdx = 0; gatheredKvSTileIdx < kvSLoopNum + PRE_LAUNCH; gatheredKvSTileIdx++) {
306 if (gatheredKvSTileIdx < kvSLoopNum) {334 if (gatheredKvSTileIdx < kvSLoopNum) {
@@ -388,7 +416,7 @@ public:
388 Arch::CrossCoreFlag mm2ToReFlag(Mm2ToReFlagId);416 Arch::CrossCoreFlag mm2ToReFlag(Mm2ToReFlagId);
389 uint32_t curTileMod = gatheredKvSTileIdxDe % (PRE_LAUNCH + 1);417 uint32_t curTileMod = gatheredKvSTileIdxDe % (PRE_LAUNCH + 1);
390 epilogueRescaleO(418 epilogueRescaleO(
391- gmOTensorTla, actualBlockShapePV,419+ gmOTensorTla, gmLseTensorTla, actualBlockShapePV,
392 curTileMod, gatheredKvSTileIdxDe,420 curTileMod, gatheredKvSTileIdxDe,
393 (gatheredKvSTileIdxDe == 0),421 (gatheredKvSTileIdxDe == 0),
394 (gatheredKvSTileIdxDe == kvSLoopNum - 1),422 (gatheredKvSTileIdxDe == kvSLoopNum - 1),
Mattention/block_sparse_attention/op_kernel/arch35/kernel_utils.hpp+6-4
@@ -46,6 +46,7 @@ struct BsaKernelParamsArch35 {
46 GM_ADDR actualKvseqlen;46 GM_ADDR actualKvseqlen;
47 GM_ADDR blockSparseMask;47 GM_ADDR blockSparseMask;
48 GM_ADDR o;48 GM_ADDR o;
49+ GM_ADDR lse;
49 GM_ADDR workSpace;50 GM_ADDR workSpace;
50 GM_ADDR tiling;51 GM_ADDR tiling;
51 52 
@@ -55,10 +56,10 @@ struct BsaKernelParamsArch35 {
55 __aicore__ inline56 __aicore__ inline
56 BsaKernelParamsArch35(GM_ADDR q_, GM_ADDR k_, GM_ADDR v_, GM_ADDR mask_, GM_ADDR blockTables_,57 BsaKernelParamsArch35(GM_ADDR q_, GM_ADDR k_, GM_ADDR v_, GM_ADDR mask_, GM_ADDR blockTables_,
57 GM_ADDR actualQseqlen_, GM_ADDR actualKvseqlen_, GM_ADDR blockSparseMask_, GM_ADDR o_,58 GM_ADDR actualQseqlen_, GM_ADDR actualKvseqlen_, GM_ADDR blockSparseMask_, GM_ADDR o_,
58- GM_ADDR workSpace_, GM_ADDR tiling_)59+ GM_ADDR workSpace_, GM_ADDR lse_, GM_ADDR tiling_)
59 : q(q_), k(k_), v(v_), mask(mask_), blockTables(blockTables_), actualQseqlen(actualQseqlen_),60 : q(q_), k(k_), v(v_), mask(mask_), blockTables(blockTables_), actualQseqlen(actualQseqlen_),
60 actualKvseqlen(actualKvseqlen_), blockSparseMask(blockSparseMask_), o(o_),61 actualKvseqlen(actualKvseqlen_), blockSparseMask(blockSparseMask_), o(o_),
61- workSpace(workSpace_), tiling(tiling_) {}62+ workSpace(workSpace_), lse(lse_), tiling(tiling_) {}
62};63};
63 64 
64struct BsaFullQuantKernelParamsArch35 {65struct BsaFullQuantKernelParamsArch35 {
@@ -74,6 +75,7 @@ struct BsaFullQuantKernelParamsArch35 {
74 GM_ADDR kDequantScale;75 GM_ADDR kDequantScale;
75 GM_ADDR vDequantScale;76 GM_ADDR vDequantScale;
76 GM_ADDR attentionOut;77 GM_ADDR attentionOut;
78+ GM_ADDR lse;
77 GM_ADDR workSpace;79 GM_ADDR workSpace;
78 GM_ADDR tiling;80 GM_ADDR tiling;
79 81 
@@ -84,11 +86,11 @@ struct BsaFullQuantKernelParamsArch35 {
84 GM_ADDR actualSeqLengths_, GM_ADDR actualSeqLengthsKv_,86 GM_ADDR actualSeqLengths_, GM_ADDR actualSeqLengthsKv_,
85 GM_ADDR qDequantScale_, GM_ADDR kDequantScale_,87 GM_ADDR qDequantScale_, GM_ADDR kDequantScale_,
86 GM_ADDR vDequantScale_, GM_ADDR attentionOut_,88 GM_ADDR vDequantScale_, GM_ADDR attentionOut_,
87- GM_ADDR workSpace_, GM_ADDR tiling_)89+ GM_ADDR workSpace_, GM_ADDR lse_, GM_ADDR tiling_)
88 : query(query_), key(key_), value(value_), blockSparseMask(blockSparseMask_), attenMask(attenMask_),90 : query(query_), key(key_), value(value_), blockSparseMask(blockSparseMask_), attenMask(attenMask_),
89 blockTable(blockTable_), actualSeqLengths(actualSeqLengths_), actualSeqLengthsKv(actualSeqLengthsKv_),91 blockTable(blockTable_), actualSeqLengths(actualSeqLengths_), actualSeqLengthsKv(actualSeqLengthsKv_),
90 qDequantScale(qDequantScale_), kDequantScale(kDequantScale_), vDequantScale(vDequantScale_),92 qDequantScale(qDequantScale_), kDequantScale(kDequantScale_), vDequantScale(vDequantScale_),
91- attentionOut(attentionOut_), workSpace(workSpace_), tiling(tiling_) {}93+ attentionOut(attentionOut_), workSpace(workSpace_), lse(lse_), tiling(tiling_) {}
92};94};
93 95 
94__aicore__ inline96__aicore__ inline
Mattention/block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue_rescale_o_arch35_reg_high_prec.hpp+132-21
@@ -32,10 +32,11 @@ template <
32 class ElementOTmp_,32 class ElementOTmp_,
33 class ElementS_,33 class ElementS_,
34 class TileCopy_,34 class TileCopy_,
35- class OTmpSrcPos_ // the src TPosition of pv res, viable configurations: GM/L0C35+ class OTmpSrcPos_, // the src TPosition of pv res, viable configurations: GM/L0C
36->36+ LseMode LSE_MODE_,
37+ LseFormat LSE_FORMAT_>
37class BlockEpilogue<38class BlockEpilogue<
38- EpilogueAtlasA5BsaRescaleO,39+ EpilogueAtlasA5BsaRescaleO<LSE_MODE_, LSE_FORMAT_>,
39 ElementO_,40 ElementO_,
40 ElementOTmp_,41 ElementOTmp_,
41 ElementS_,42 ElementS_,
@@ -43,10 +44,11 @@ class BlockEpilogue<
43 OTmpSrcPos_>44 OTmpSrcPos_>
44{45{
45public:46public:
46- using DispatchPolicy = EpilogueAtlasA5BsaRescaleO;47+ using DispatchPolicy = EpilogueAtlasA5BsaRescaleO<LSE_MODE_, LSE_FORMAT_>;
47 using ArchTag = typename DispatchPolicy::ArchTag;48 using ArchTag = typename DispatchPolicy::ArchTag;
48 using ElementO = ElementO_;49 using ElementO = ElementO_;
49 using ElementOTmp = ElementOTmp_;50 using ElementOTmp = ElementOTmp_;
51+ using ElementLse = float;
50 using SMDtype = ElementS_;52 using SMDtype = ElementS_;
51 using TileCopy = TileCopy_;53 using TileCopy = TileCopy_;
52 using OTmpSrcPos = OTmpSrcPos_;54 using OTmpSrcPos = OTmpSrcPos_;
@@ -70,6 +72,7 @@ public:
70 constexpr uint32_t DM_UB_TENSOR_OFFSET = GM_UB_TENSOR_OFFSET + 64 * sizeof(float);72 constexpr uint32_t DM_UB_TENSOR_OFFSET = GM_UB_TENSOR_OFFSET + 64 * sizeof(float);
71 constexpr uint32_t LL_UB_TENSOR_OFFSET = DM_UB_TENSOR_OFFSET + 3 * 64 * sizeof(float);73 constexpr uint32_t LL_UB_TENSOR_OFFSET = DM_UB_TENSOR_OFFSET + 3 * 64 * sizeof(float);
72 constexpr uint32_t GL_UB_TENSOR_OFFSET = LL_UB_TENSOR_OFFSET + 64 * sizeof(float);74 constexpr uint32_t GL_UB_TENSOR_OFFSET = LL_UB_TENSOR_OFFSET + 64 * sizeof(float);
75+ constexpr uint32_t LSE_UB_TENSOR_OFFSET = GL_UB_TENSOR_OFFSET + 64 * sizeof(float);
73 76 
74 for (uint32_t i = 0; i < UB_OTMP_BUF_STAGES; i++) {77 for (uint32_t i = 0; i < UB_OTMP_BUF_STAGES; i++) {
75 loUbTensor[i] = resource.ubBuf.template GetBufferByByte<ElementOTmp>(78 loUbTensor[i] = resource.ubBuf.template GetBufferByByte<ElementOTmp>(
@@ -79,6 +82,8 @@ public:
79 goUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementO>(GO_UB_TENSOR_OFFSET);82 goUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementO>(GO_UB_TENSOR_OFFSET);
80 glUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);83 glUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
81 dmUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);84 dmUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
85+ gmUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
86+ lseUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(LSE_UB_TENSOR_OFFSET);
82 }87 }
83 88 
84 __aicore__ inline89 __aicore__ inline
@@ -106,9 +111,32 @@ public:
106 }111 }
107 }112 }
108 113 
109- template <class TensorDst>114+ template <class TensorLseGm, class TensorLseUb>
110 __aicore__ inline115 __aicore__ inline
111- void SubCoreCompute(TensorDst &gOTensorTlaTile,116+ void CopyUbToGmLse(TensorLseGm const &gLseTensorTlaTile, TensorLseUb const &ubLseTensorTla)
117+ {
118+ AscendC::DataCopyExtParams repeatParams;
119+ if constexpr ((DispatchPolicy::LSE_FORMAT == LseFormat::TN1) ||
120+ (DispatchPolicy::LSE_FORMAT == LseFormat::BSN1)) {
121+ repeatParams.blockCount = tla::get<0>(ubLseTensorTla.shape());
122+ repeatParams.blockLen = sizeof(float);
123+ repeatParams.srcStride = 0;
124+ repeatParams.dstStride = (tla::get<0>(gLseTensorTlaTile.stride()) - 1) * sizeof(float);
125+ } else if constexpr (DispatchPolicy::LSE_FORMAT == LseFormat::BNS1) {
126+ repeatParams.blockCount = 1;
127+ repeatParams.blockLen = tla::get<0>(ubLseTensorTla.shape()) * sizeof(float);
128+ repeatParams.srcStride = 0;
129+ repeatParams.dstStride = 0;
130+ }
131+ auto dstOffset = gLseTensorTlaTile.layout()(gLseTensorTlaTile.coord());
132+ auto srcOffset = ubLseTensorTla.layout()(ubLseTensorTla.coord());
133+ AscendC::DataCopyPad(gLseTensorTlaTile.data()[dstOffset], ubLseTensorTla.data()[srcOffset], repeatParams);
134+ }
135+ 
136+ template <class TensorO, class TensorLse>
137+ __aicore__ inline
138+ void SubCoreCompute(TensorO &gOTensorTlaTile,
139+ TensorLse &gLseTensorTlaTile,
112 uint32_t curTileMod,140 uint32_t curTileMod,
113 uint32_t ubOTmpBufId,141 uint32_t ubOTmpBufId,
114 bool isFirstKvSTile,142 bool isFirstKvSTile,
@@ -126,9 +154,11 @@ public:
126 154 
127 __ubuf__ ElementOTmp *goUb = (__ubuf__ ElementOTmp *) goUbTensor32.GetPhyAddr();155 __ubuf__ ElementOTmp *goUb = (__ubuf__ ElementOTmp *) goUbTensor32.GetPhyAddr();
128 __ubuf__ ElementOTmp *loUb = (__ubuf__ ElementOTmp *) loUbTensor[ubOTmpBufId].GetPhyAddr();156 __ubuf__ ElementOTmp *loUb = (__ubuf__ ElementOTmp *) loUbTensor[ubOTmpBufId].GetPhyAddr();
129- __ubuf__ ElementOTmp *glUb = ( __ubuf__ ElementOTmp *) glUbTensor32.GetPhyAddr();157+ __ubuf__ ElementOTmp *glUb = (__ubuf__ ElementOTmp *) glUbTensor32.GetPhyAddr();
130 __ubuf__ ElementOTmp *dmUb =158 __ubuf__ ElementOTmp *dmUb =
131 (__ubuf__ ElementOTmp *) dmUbTensor32[curTileMod * DM_UB_GLOBAL_ELEM_NUM].GetPhyAddr();159 (__ubuf__ ElementOTmp *) dmUbTensor32[curTileMod * DM_UB_GLOBAL_ELEM_NUM].GetPhyAddr();
160+ __ubuf__ float *gmUb = (__ubuf__ float *) gmUbTensor32.GetPhyAddr();
161+ __ubuf__ float *lseUb = (__ubuf__ float *) lseUbTensor32.GetPhyAddr();
132 162
133 WaitCrossCoreSync<4, PIPE_V>(mm2ToReFlag);163 WaitCrossCoreSync<4, PIPE_V>(mm2ToReFlag);
134 AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);164 AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
@@ -145,8 +175,6 @@ public:
145 DivFuncLastAndFirst<DRegSplitStages::TWO>(175 DivFuncLastAndFirst<DRegSplitStages::TWO>(
146 goUb, loUb, glUb, rowNumCurSubCore, colStrideCurSubCore, colTail, vlElemNum);176 goUb, loUb, glUb, rowNumCurSubCore, colStrideCurSubCore, colTail, vlElemNum);
147 }177 }
148- // DivFuncLastAndFirst<ElementOTmp>(
149- // goUb, loUb, glUb, rowNumCurSubCore, colStrideCurSubCore, colFullLoop, colTail, vlElemNum);
150 }178 }
151 } else if (!isLastKvSTile) {179 } else if (!isLastKvSTile) {
152 if (dStages == 1) {180 if (dStages == 1) {
@@ -156,8 +184,6 @@ public:
156 RescaleFunc<DRegSplitStages::TWO>(184 RescaleFunc<DRegSplitStages::TWO>(
157 goUb, loUb, dmUb, rowNumCurSubCore, colStrideCurSubCore, colTail, vlElemNum);185 goUb, loUb, dmUb, rowNumCurSubCore, colStrideCurSubCore, colTail, vlElemNum);
158 }186 }
159- // RescaleFunc<ElementOTmp>(
160- // goUb, loUb, dmUb, rowNumCurSubCore, colStrideCurSubCore, colFullLoop, colTail, vlElemNum);
161 } else {187 } else {
162 if (dStages == 1) {188 if (dStages == 1) {
163 RescaleFuncLastNotFirst<DRegSplitStages::ONE>(189 RescaleFuncLastNotFirst<DRegSplitStages::ONE>(
@@ -166,12 +192,31 @@ public:
166 RescaleFuncLastNotFirst<DRegSplitStages::TWO>(192 RescaleFuncLastNotFirst<DRegSplitStages::TWO>(
167 goUb, loUb, dmUb, glUb, rowNumCurSubCore, colStrideCurSubCore, colTail, vlElemNum);193 goUb, loUb, dmUb, glUb, rowNumCurSubCore, colStrideCurSubCore, colTail, vlElemNum);
168 }194 }
169- // RescaleFuncLastNotFirst<ElementOTmp>(
170- // goUb, loUb, dmUb, glUb, rowNumCurSubCore, colStrideCurSubCore, colFullLoop, colTail, vlElemNum);
171 }195 }
172 // release lo buf196 // release lo buf
173 SetCrossCoreSync<4, PIPE_V>(mm2ToReFlag);197 SetCrossCoreSync<4, PIPE_V>(mm2ToReFlag);
174 if (isLastKvSTile) {198 if (isLastKvSTile) {
199+ if constexpr (DispatchPolicy::LSE_MODE == LseMode::OUT_ONLY) {
200+ uint32_t colNumLseUb = 0;
201+ uint32_t colStrideLseUb = 0;
202+ if constexpr ((DispatchPolicy::LSE_FORMAT == LseFormat::TN1) ||
203+ (DispatchPolicy::LSE_FORMAT == LseFormat::BSN1)) {
204+ LogSumExpFuncQSAxisIncontinuous(gmUb, glUb, lseUb, rowNumCurSubCore);
205+ colNumLseUb = 8;
206+ colStrideLseUb = 8;
207+ } else if constexpr (DispatchPolicy::LSE_FORMAT == LseFormat::BNS1) {
208+ LogSumExpFuncQSAxisContinuous(gmUb, glUb, lseUb, rowNumCurSubCore);
209+ colNumLseUb = 1;
210+ colStrideLseUb = 1;
211+ }
212+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID1);
213+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID1);
214+ auto ubLseLayoutTla = tla::MakeLayout(
215+ tla::MakeShape(rowNumCurSubCore, colNumLseUb),
216+ tla::MakeStride(colStrideLseUb, tla::Int<1>{}));
217+ auto ubLseTensorTla = tla::MakeTensor(lseUbTensor32, ubLseLayoutTla, Arch::PositionUB{});
218+ CopyUbToGmLse(gLseTensorTlaTile, ubLseTensorTla);
219+ }
175 AscendC::PipeBarrier<PIPE_V>();220 AscendC::PipeBarrier<PIPE_V>();
176 if (isFullQuantFp8) {221 if (isFullQuantFp8) {
177 if (dStages == 1) {222 if (dStages == 1) {
@@ -227,7 +272,7 @@ public:
227 RegTensor<float> mulVreg;272 RegTensor<float> mulVreg;
228 RegTensor<float> goCurVreg;273 RegTensor<float> goCurVreg;
229 MaskReg pregTail = UpdateMask<float>(colTail);274 MaskReg pregTail = UpdateMask<float>(colTail);
230- for (uint32_t i = 0; i < row; i++) {275+ for (uint16_t i = 0; i < row; i++) {
231 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(dmVreg, dmUb + i);276 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(dmVreg, dmUb + i);
232 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goPreVreg, goUb + i * colStride);277 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goPreVreg, goUb + i * colStride);
233 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(loVreg, loUb + i * colStride);278 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(loVreg, loUb + i * colStride);
@@ -256,7 +301,7 @@ public:
256 RegTensor<float> goCurVreg1;301 RegTensor<float> goCurVreg1;
257 MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();302 MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
258 MaskReg pregTail = UpdateMask<float>(colTail);303 MaskReg pregTail = UpdateMask<float>(colTail);
259- for (uint32_t i = 0; i < row; i++) {304+ for (uint16_t i = 0; i < row; i++) {
260 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(dmVreg, dmUb + i);305 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(dmVreg, dmUb + i);
261 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goPreVreg0, goUb + i * colStride);306 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goPreVreg0, goUb + i * colStride);
262 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goPreVreg1, goUb + i * colStride + vlElemNum);307 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goPreVreg1, goUb + i * colStride + vlElemNum);
@@ -298,7 +343,7 @@ public:
298 RegTensor<float> glVreg;343 RegTensor<float> glVreg;
299 RegTensor<float> divVreg;344 RegTensor<float> divVreg;
300 MaskReg pregTail = UpdateMask<float>(colTail);345 MaskReg pregTail = UpdateMask<float>(colTail);
301- for (uint32_t i = 0; i < row; i++) {346+ for (uint16_t i = 0; i < row; i++) {
302 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(dmVreg, dmUb + i);347 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(dmVreg, dmUb + i);
303 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(glVreg, glUb + i);348 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(glVreg, glUb + i);
304 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goPreVreg, goUb + i * colStride);349 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goPreVreg, goUb + i * colStride);
@@ -332,7 +377,7 @@ public:
332 RegTensor<float> divVreg1;377 RegTensor<float> divVreg1;
333 MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();378 MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
334 MaskReg pregTail = UpdateMask<float>(colTail);379 MaskReg pregTail = UpdateMask<float>(colTail);
335- for (uint32_t i = 0; i < row; i++) {380+ for (uint16_t i = 0; i < row; i++) {
336 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(dmVreg, dmUb + i);381 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(dmVreg, dmUb + i);
337 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(glVreg, glUb + i);382 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(glVreg, glUb + i);
338 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goPreVreg0, goUb + i * colStride);383 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goPreVreg0, goUb + i * colStride);
@@ -369,7 +414,7 @@ public:
369 RegTensor<float> glVreg;414 RegTensor<float> glVreg;
370 RegTensor<float> divVreg;415 RegTensor<float> divVreg;
371 MaskReg pregTail = UpdateMask<float>(colTail);416 MaskReg pregTail = UpdateMask<float>(colTail);
372- for (uint32_t i = 0; i < row; i++) {417+ for (uint16_t i = 0; i < row; i++) {
373 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(glVreg, glUb + i);418 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(glVreg, glUb + i);
374 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goCurVreg, loUb + i * colStride);419 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goCurVreg, loUb + i * colStride);
375 Div(divVreg, goCurVreg, glVreg, pregTail);420 Div(divVreg, goCurVreg, glVreg, pregTail);
@@ -391,7 +436,7 @@ public:
391 RegTensor<float> divVreg1;436 RegTensor<float> divVreg1;
392 MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();437 MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
393 MaskReg pregTail = UpdateMask<float>(colTail);438 MaskReg pregTail = UpdateMask<float>(colTail);
394- for (uint32_t i = 0; i < row; i++) {439+ for (uint16_t i = 0; i < row; i++) {
395 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(glVreg, glUb + i);440 LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(glVreg, glUb + i);
396 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goCurVreg0, loUb + i * colStride);441 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goCurVreg0, loUb + i * colStride);
397 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goCurVreg1, loUb + i * colStride + vlElemNum);442 LoadAlign<ElementOTmp, LoadDist::DIST_NORM>(goCurVreg1, loUb + i * colStride + vlElemNum);
@@ -447,9 +492,70 @@ public:
447 }492 }
448 }493 }
449 494 
450- template <class TensorDst>495+ // When qS can be copied from UB to GM continuously,
496+ // qS would be stored from reg to UB continuously.
497+ __simd_vf__ inline void LogSumExpFuncQSAxisContinuous(
498+ __ubuf__ float *gmUb, __ubuf__ float *glUb, __ubuf__ float *lseUb,
499+ uint32_t row)
500+ {
501+ // This vf works only when the rowNum in each AIV does not exceed 64
502+ using namespace AscendC::MicroAPI;
503+ RegTensor<float> gmVreg;
504+ RegTensor<float> glVreg;
505+ RegTensor<float> logGlVreg;
506+ RegTensor<float> lseVreg;
507+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
508+ MaskReg pregTail = UpdateMask<float>(row);
509+ static constexpr LnSpecificMode mode = {MaskMergeMode::ZEROING, AscendC::LnAlgo::PRECISION_1ULP_FTZ_FALSE};
510+ 
511+ LoadAlign<float, LoadDist::DIST_NORM>(glVreg, glUb);
512+ LoadAlign<float, LoadDist::DIST_NORM>(gmVreg, gmUb);
513+ Ln<float, &mode>(logGlVreg, glVreg, pregTail);
514+ Add(lseVreg, logGlVreg, gmVreg, pregTail);
515+ StoreAlign<float, StoreDist::DIST_NORM_B32>(lseUb, lseVreg, pregTail);
516+ }
517+ 
518+ // When qS cannot be copied from UB to GM continuously,
519+ // lse would be broadcasted and then stored from reg to UB continuously.
520+ __simd_vf__ inline void LogSumExpFuncQSAxisIncontinuous(
521+ __ubuf__ float *gmUb, __ubuf__ float *glUb, __ubuf__ float *lseUb,
522+ uint32_t row)
523+ {
524+ using namespace AscendC::MicroAPI;
525+ RegTensor<float> gmRowwiseVreg0;
526+ RegTensor<float> gmRowwiseVreg1;
527+ RegTensor<float> glRowwiseVreg0;
528+ RegTensor<float> glRowwiseVreg1;
529+ RegTensor<float> logGlRowwiseVreg0;
530+ RegTensor<float> logGlRowwiseVreg1;
531+ RegTensor<float> lseRowwiseVreg0;
532+ RegTensor<float> lseRowwiseVreg1;
533+ UnalignReg rowwiseUreg0;
534+ UnalignReg rowwiseUreg1;
535+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
536+ static constexpr LnSpecificMode mode = {MaskMergeMode::ZEROING, AscendC::LnAlgo::PRECISION_1ULP_FTZ_FALSE};
537+ static constexpr uint32_t postUpdateStride = 32 / sizeof(float);
538+ 
539+ for (uint16_t i = 0; i < row; i+=2) {
540+ LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(glRowwiseVreg0, glUb + i);
541+ LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(glRowwiseVreg1, glUb + (i + 1));
542+ LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(gmRowwiseVreg0, gmUb + i);
543+ LoadAlign<ElementOTmp, LoadDist::DIST_BRC_B32>(gmRowwiseVreg1, gmUb + (i + 1));
544+ Ln<float, &mode>(logGlRowwiseVreg0, glRowwiseVreg0, pregFull);
545+ Ln<float, &mode>(logGlRowwiseVreg1, glRowwiseVreg1, pregFull);
546+ Add(lseRowwiseVreg0, logGlRowwiseVreg0, gmRowwiseVreg0, pregFull);
547+ Add(lseRowwiseVreg1, logGlRowwiseVreg1, gmRowwiseVreg1, pregFull);
548+ StoreUnAlign<float, PostLiteral::POST_MODE_UPDATE>(lseUb, lseRowwiseVreg0, rowwiseUreg0, postUpdateStride);
549+ StoreUnAlign<float, PostLiteral::POST_MODE_UPDATE>(lseUb, lseRowwiseVreg1, rowwiseUreg1, postUpdateStride);
550+ }
551+ StoreUnAlignPost<float, PostLiteral::POST_MODE_UPDATE>(lseUb, rowwiseUreg0, postUpdateStride);
552+ StoreUnAlignPost<float, PostLiteral::POST_MODE_UPDATE>(lseUb, rowwiseUreg1, postUpdateStride);
553+ }
554+ 
555+ template <class TensorO, class TensorLse>
451 __aicore__ inline556 __aicore__ inline
452- void operator()(TensorDst &gOTensor,557+ void operator()(TensorO &gOTensor,
558+ TensorLse &gLseTensor,
453 GemmCoord actualOriShape,559 GemmCoord actualOriShape,
454 uint32_t curTileMod,560 uint32_t curTileMod,
455 uint32_t gatheredKvSTileIdx,561 uint32_t gatheredKvSTileIdx,
@@ -475,11 +581,14 @@ public:
475 581 
476 auto gOTensorTlaTile = GetTile(gOTensor,582 auto gOTensorTlaTile = GetTile(gOTensor,
477 tla::MakeCoord(rowOffsetCurSubCore, 0), tla::MakeShape(rowNumCurSubCore, colNumCurSubCore));583 tla::MakeCoord(rowOffsetCurSubCore, 0), tla::MakeShape(rowNumCurSubCore, colNumCurSubCore));
584+ auto gLseTensorTlaTile = GetTile(gLseTensor,
585+ tla::MakeCoord(rowOffsetCurSubCore, 0), tla::MakeShape(rowNumCurSubCore, 1));
478 uint32_t ubOTmpBufId = gatheredKvSTileIdx % UB_OTMP_BUF_STAGES;586 uint32_t ubOTmpBufId = gatheredKvSTileIdx % UB_OTMP_BUF_STAGES;
479 587 
480 if (rowNumCurSubCore > 0) {588 if (rowNumCurSubCore > 0) {
481 SubCoreCompute(589 SubCoreCompute(
482 gOTensorTlaTile,590 gOTensorTlaTile,
591+ gLseTensorTlaTile,
483 curTileMod,592 curTileMod,
484 ubOTmpBufId,593 ubOTmpBufId,
485 isFirstKvSTile,594 isFirstKvSTile,
@@ -497,10 +606,12 @@ private:
497 AscendC::LocalTensor<ElementOTmp> loUbTensor[UB_OTMP_BUF_STAGES];606 AscendC::LocalTensor<ElementOTmp> loUbTensor[UB_OTMP_BUF_STAGES];
498 AscendC::LocalTensor<SMDtype> dmUbTensor16;607 AscendC::LocalTensor<SMDtype> dmUbTensor16;
499 AscendC::LocalTensor<SMDtype> glUbTensor16;608 AscendC::LocalTensor<SMDtype> glUbTensor16;
609+ AscendC::LocalTensor<float> gmUbTensor32;
500 AscendC::LocalTensor<float> dmUbTensor32;610 AscendC::LocalTensor<float> dmUbTensor32;
501 AscendC::LocalTensor<float> glUbTensor32;611 AscendC::LocalTensor<float> glUbTensor32;
502 AscendC::LocalTensor<ElementO> goUbTensor16;612 AscendC::LocalTensor<ElementO> goUbTensor16;
503 AscendC::LocalTensor<ElementOTmp> goUbTensor32;613 AscendC::LocalTensor<ElementOTmp> goUbTensor32;
614+ AscendC::LocalTensor<float> lseUbTensor32;
504 615 
505 CopyUbToGmO copyUbToGmO;616 CopyUbToGmO copyUbToGmO;
506 617 
Mattention/block_sparse_attention/op_kernel/attn_infra/epilogue/bsa_epilogue_dispatch_policy.hpp+6-0
@@ -18,6 +18,7 @@ namespace NpuArch::Epilogue
18{18{
19 19 
20enum class LseMode {NONE = 0, OUT_ONLY = 1};20enum class LseMode {NONE = 0, OUT_ONLY = 1};
21+enum class LseFormat {TN1 = 0, BNS1 = 1, BSN1 = 2};
21// For AtlasA2, FA Infer online Softmax22// For AtlasA2, FA Infer online Softmax
22template <LseMode LSE_MODE_, typename SM_DTYPE_>23template <LseMode LSE_MODE_, typename SM_DTYPE_>
23struct EpilogueAtlasA2OnlineSoftmax {24struct EpilogueAtlasA2OnlineSoftmax {
@@ -44,8 +45,13 @@ struct EpilogueOnlineSoftmaxBsa {
44 using ArchTag = Arch::AtlasA5;45 using ArchTag = Arch::AtlasA5;
45};46};
46 47 
48+template <
49+ LseMode LSE_MODE_ = LseMode::NONE,
50+ LseFormat LSE_FORMAT_ = LseFormat::BNS1>
47struct EpilogueAtlasA5BsaRescaleO {51struct EpilogueAtlasA5BsaRescaleO {
48 using ArchTag = Arch::AtlasA5;52 using ArchTag = Arch::AtlasA5;
53+ static constexpr LseMode LSE_MODE = LSE_MODE_;
54+ static constexpr LseFormat LSE_FORMAT = LSE_FORMAT_;
49};55};
50 56 
51} // namespace NpuArch::Epilogue57} // namespace NpuArch::Epilogue
Mattention/block_sparse_attention/op_kernel/block_sparse_attention.cpp+132-35
@@ -98,85 +98,182 @@ extern "C" __global__ __aicore__ void block_sparse_attention(
98 #endif98 #endif
99#endif99#endif
100#if (__CCE_AICORE__ == 310)100#if (__CCE_AICORE__ == 310)
101- // 非量化FP16/BF16场景101+ // 非量化+混合精度softmax
102 TILING_KEY_IS(QF16_KVF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY);102 TILING_KEY_IS(QF16_KVF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY);
103 TILING_KEY_IS(QBF16_KVBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY);103 TILING_KEY_IS(QBF16_KVBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY);
104 TILING_KEY_IS(QF16_KVF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY);104 TILING_KEY_IS(QF16_KVF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY);
105 TILING_KEY_IS(QBF16_KVBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY);105 TILING_KEY_IS(QBF16_KVBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY);
106 TILING_KEY_IS(QF16_KVF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_KEY);106 TILING_KEY_IS(QF16_KVF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_KEY);
107 TILING_KEY_IS(QBF16_KVBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_KEY);107 TILING_KEY_IS(QBF16_KVBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_KEY);
108- 108+ // 非量化+混合精度softmax+lse
109- 109+ TILING_KEY_IS(QF16_KVF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY);
110- // 全量化FP8场景110+ TILING_KEY_IS(QBF16_KVBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY);
111+ TILING_KEY_IS(QF16_KVF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY);
112+ TILING_KEY_IS(QBF16_KVBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY);
113+ TILING_KEY_IS(QF16_KVF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY);
114+ TILING_KEY_IS(QBF16_KVBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY);
115+ // 全量化FP8+混合精度softmax
111 TILING_KEY_IS(QF8_KVF8_OF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY);116 TILING_KEY_IS(QF8_KVF8_OF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY);
112 TILING_KEY_IS(QF8_KVF8_OF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY);117 TILING_KEY_IS(QF8_KVF8_OF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY);
113 TILING_KEY_IS(QF8_KVF8_OF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_KEY);118 TILING_KEY_IS(QF8_KVF8_OF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_KEY);
114 TILING_KEY_IS(QF8_KVF8_OBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY);119 TILING_KEY_IS(QF8_KVF8_OBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY);
115 TILING_KEY_IS(QF8_KVF8_OBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY);120 TILING_KEY_IS(QF8_KVF8_OBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY);
116 TILING_KEY_IS(QF8_KVF8_OBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_KEY);121 TILING_KEY_IS(QF8_KVF8_OBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_KEY);
122+ // 全量化FP8+混合精度softmax+lse
123+ TILING_KEY_IS(QF8_KVF8_OF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY);
124+ TILING_KEY_IS(QF8_KVF8_OF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY);
125+ TILING_KEY_IS(QF8_KVF8_OF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY);
126+ TILING_KEY_IS(QF8_KVF8_OBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY);
127+ TILING_KEY_IS(QF8_KVF8_OBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY);
128+ TILING_KEY_IS(QF8_KVF8_OBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY);
117 129 
118- // 非量化FP16/BF16场景
119 #if TILING_KEY_VAR == QF16_KVF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY130 #if TILING_KEY_VAR == QF16_KVF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY
120 BsaInferIntfRegular<131 BsaInferIntfRegular<
121- half, half, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND>(132+ half, half, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND,
133+ Epilogue::LseMode::NONE, Epilogue::LseFormat::TN1>(
122 query, key, value, mask, blockTable, attentionOut,134 query, key, value, mask, blockTable, attentionOut,
123- actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, tiling);135+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, softmaxLse, tiling);
124 #elif TILING_KEY_VAR == QF16_KVF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY136 #elif TILING_KEY_VAR == QF16_KVF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY
125 BsaInferIntfRegular<137 BsaInferIntfRegular<
126- half, half, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD>(138+ half, half, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD,
139+ Epilogue::LseMode::NONE, Epilogue::LseFormat::BNS1>(
127 query, key, value, mask, blockTable, attentionOut,140 query, key, value, mask, blockTable, attentionOut,
128- actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, tiling);141+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, softmaxLse, tiling);
129- #elif TILING_KEY_VAR == QBF16_KVBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY
130- BsaInferIntfRegular<
131- bfloat16_t, bfloat16_t, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND>(
132- query, key, value, mask, blockTable, attentionOut,
133- actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, tiling);
134- #elif TILING_KEY_VAR == QBF16_KVBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY
135- BsaInferIntfRegular<
136- bfloat16_t, bfloat16_t, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD>(
137- query, key, value, mask, blockTable, attentionOut,
138- actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, tiling);
139 #elif TILING_KEY_VAR == QF16_KVF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_KEY142 #elif TILING_KEY_VAR == QF16_KVF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_KEY
140 BsaInferIntfRegular<143 BsaInferIntfRegular<
141- half, half, float, BsaKernelArch35::Format::BSND, BsaKernelArch35::Format::BSND>(144+ half, half, float, BsaKernelArch35::Format::BSND, BsaKernelArch35::Format::BSND,
145+ Epilogue::LseMode::NONE, Epilogue::LseFormat::BSN1>(
142 query, key, value, mask, blockTable, attentionOut,146 query, key, value, mask, blockTable, attentionOut,
143- actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, tiling);147+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, softmaxLse, tiling);
148+ #elif TILING_KEY_VAR == QBF16_KVBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY
149+ BsaInferIntfRegular<
150+ bfloat16_t, bfloat16_t, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND,
151+ Epilogue::LseMode::NONE, Epilogue::LseFormat::TN1>(
152+ query, key, value, mask, blockTable, attentionOut,
153+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, softmaxLse, tiling);
154+ #elif TILING_KEY_VAR == QBF16_KVBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY
155+ BsaInferIntfRegular<
156+ bfloat16_t, bfloat16_t, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD,
157+ Epilogue::LseMode::NONE, Epilogue::LseFormat::BNS1>(
158+ query, key, value, mask, blockTable, attentionOut,
159+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, softmaxLse, tiling);
144 #elif TILING_KEY_VAR == QBF16_KVBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_KEY160 #elif TILING_KEY_VAR == QBF16_KVBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_KEY
145 BsaInferIntfRegular<161 BsaInferIntfRegular<
146- bfloat16_t, bfloat16_t, float, BsaKernelArch35::Format::BSND, BsaKernelArch35::Format::BSND>(162+ bfloat16_t, bfloat16_t, float, BsaKernelArch35::Format::BSND, BsaKernelArch35::Format::BSND,
163+ Epilogue::LseMode::NONE, Epilogue::LseFormat::BSN1>(
147 query, key, value, mask, blockTable, attentionOut,164 query, key, value, mask, blockTable, attentionOut,
148- actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, tiling);165+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, softmaxLse, tiling);
166+ // lse
167+ #elif TILING_KEY_VAR == QF16_KVF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY
168+ BsaInferIntfRegular<
169+ half, half, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND,
170+ Epilogue::LseMode::OUT_ONLY, Epilogue::LseFormat::TN1>(
171+ query, key, value, mask, blockTable, attentionOut,
172+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, softmaxLse, tiling);
173+ #elif TILING_KEY_VAR == QBF16_KVBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY
174+ BsaInferIntfRegular<
175+ bfloat16_t, bfloat16_t, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND,
176+ Epilogue::LseMode::OUT_ONLY, Epilogue::LseFormat::TN1>(
177+ query, key, value, mask, blockTable, attentionOut,
178+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, softmaxLse, tiling);
179+ #elif TILING_KEY_VAR == QF16_KVF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY
180+ BsaInferIntfRegular<
181+ half, half, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD,
182+ Epilogue::LseMode::OUT_ONLY, Epilogue::LseFormat::BNS1>(
183+ query, key, value, mask, blockTable, attentionOut,
184+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, softmaxLse, tiling);
185+ #elif TILING_KEY_VAR == QBF16_KVBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY
186+ BsaInferIntfRegular<
187+ bfloat16_t, bfloat16_t, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD,
188+ Epilogue::LseMode::OUT_ONLY, Epilogue::LseFormat::BNS1>(
189+ query, key, value, mask, blockTable, attentionOut,
190+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, softmaxLse, tiling);
191+ #elif TILING_KEY_VAR == QF16_KVF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY
192+ BsaInferIntfRegular<
193+ half, half, float, BsaKernelArch35::Format::BSND, BsaKernelArch35::Format::BSND,
194+ Epilogue::LseMode::OUT_ONLY, Epilogue::LseFormat::BSN1>(
195+ query, key, value, mask, blockTable, attentionOut,
196+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, softmaxLse, tiling);
197+ #elif TILING_KEY_VAR == QBF16_KVBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY
198+ BsaInferIntfRegular<
199+ bfloat16_t, bfloat16_t, float, BsaKernelArch35::Format::BSND, BsaKernelArch35::Format::BSND,
200+ Epilogue::LseMode::OUT_ONLY, Epilogue::LseFormat::BSN1>(
201+ query, key, value, mask, blockTable, attentionOut,
202+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, softmaxLse, tiling);
149 // 全量化FP8场景203 // 全量化FP8场景
150 #elif TILING_KEY_VAR == QF8_KVF8_OF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY204 #elif TILING_KEY_VAR == QF8_KVF8_OF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY
151 BsaInferInterfaceFullQuant<205 BsaInferInterfaceFullQuant<
152- fp8_e4m3fn_t, half, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND>(206+ fp8_e4m3fn_t, half, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND,
207+ Epilogue::LseMode::NONE, Epilogue::LseFormat::TN1>(
153 query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,208 query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,
154- qDequantScale, kDequantScale, vDequantScale, attentionOut, user, tiling);209+ qDequantScale, kDequantScale, vDequantScale, attentionOut, user, softmaxLse, tiling);
155 #elif TILING_KEY_VAR == QF8_KVF8_OF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY210 #elif TILING_KEY_VAR == QF8_KVF8_OF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY
156 BsaInferInterfaceFullQuant<211 BsaInferInterfaceFullQuant<
157- fp8_e4m3fn_t, half, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD>(212+ fp8_e4m3fn_t, half, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD,
213+ Epilogue::LseMode::NONE, Epilogue::LseFormat::BNS1>(
158 query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,214 query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,
159- qDequantScale, kDequantScale, vDequantScale, attentionOut, user, tiling);215+ qDequantScale, kDequantScale, vDequantScale, attentionOut, user, softmaxLse, tiling);
160 #elif TILING_KEY_VAR == QF8_KVF8_OF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_KEY216 #elif TILING_KEY_VAR == QF8_KVF8_OF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_KEY
161 BsaInferInterfaceFullQuant<217 BsaInferInterfaceFullQuant<
162- fp8_e4m3fn_t, half, float, BsaKernelArch35::Format::BSND, BsaKernelArch35::Format::BSND>(218+ fp8_e4m3fn_t, half, float, BsaKernelArch35::Format::BSND, BsaKernelArch35::Format::BSND,
219+ Epilogue::LseMode::NONE, Epilogue::LseFormat::BSN1>(
163 query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,220 query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,
164- qDequantScale, kDequantScale, vDequantScale, attentionOut, user, tiling);221+ qDequantScale, kDequantScale, vDequantScale, attentionOut, user, softmaxLse, tiling);
165 #elif TILING_KEY_VAR == QF8_KVF8_OBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY222 #elif TILING_KEY_VAR == QF8_KVF8_OBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY
166 BsaInferInterfaceFullQuant<223 BsaInferInterfaceFullQuant<
167- fp8_e4m3fn_t, bfloat16_t, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND>(224+ fp8_e4m3fn_t, bfloat16_t, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND,
225+ Epilogue::LseMode::NONE, Epilogue::LseFormat::TN1>(
168 query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,226 query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,
169- qDequantScale, kDequantScale, vDequantScale, attentionOut, user, tiling);227+ qDequantScale, kDequantScale, vDequantScale, attentionOut, user, softmaxLse, tiling);
170 #elif TILING_KEY_VAR == QF8_KVF8_OBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY228 #elif TILING_KEY_VAR == QF8_KVF8_OBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY
171 BsaInferInterfaceFullQuant<229 BsaInferInterfaceFullQuant<
172- fp8_e4m3fn_t, bfloat16_t, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD>(230+ fp8_e4m3fn_t, bfloat16_t, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD,
231+ Epilogue::LseMode::NONE, Epilogue::LseFormat::BNS1>(
173 query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,232 query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,
174- qDequantScale, kDequantScale, vDequantScale, attentionOut, user, tiling);233+ qDequantScale, kDequantScale, vDequantScale, attentionOut, user, softmaxLse, tiling);
175 #elif TILING_KEY_VAR == QF8_KVF8_OBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_KEY234 #elif TILING_KEY_VAR == QF8_KVF8_OBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_KEY
176 BsaInferInterfaceFullQuant<235 BsaInferInterfaceFullQuant<
177- fp8_e4m3fn_t, bfloat16_t, float, BsaKernelArch35::Format::BSND, BsaKernelArch35::Format::BSND>(236+ fp8_e4m3fn_t, bfloat16_t, float, BsaKernelArch35::Format::BSND, BsaKernelArch35::Format::BSND,
237+ Epilogue::LseMode::NONE, Epilogue::LseFormat::BSN1>(
178 query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,238 query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,
179- qDequantScale, kDequantScale, vDequantScale, attentionOut, user, tiling);239+ qDequantScale, kDequantScale, vDequantScale, attentionOut, user, softmaxLse, tiling);
240+ // 全量化FP8+lse
241+ #elif TILING_KEY_VAR == QF8_KVF8_OF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY
242+ BsaInferInterfaceFullQuant<
243+ fp8_e4m3fn_t, half, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND,
244+ Epilogue::LseMode::OUT_ONLY, Epilogue::LseFormat::TN1>(
245+ query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,
246+ qDequantScale, kDequantScale, vDequantScale, attentionOut, user, softmaxLse, tiling);
247+ #elif TILING_KEY_VAR == QF8_KVF8_OF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY
248+ BsaInferInterfaceFullQuant<
249+ fp8_e4m3fn_t, half, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD,
250+ Epilogue::LseMode::OUT_ONLY, Epilogue::LseFormat::BNS1>(
251+ query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,
252+ qDequantScale, kDequantScale, vDequantScale, attentionOut, user, softmaxLse, tiling);
253+ #elif TILING_KEY_VAR == QF8_KVF8_OF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY
254+ BsaInferInterfaceFullQuant<
255+ fp8_e4m3fn_t, half, float, BsaKernelArch35::Format::BSND, BsaKernelArch35::Format::BSND,
256+ Epilogue::LseMode::OUT_ONLY, Epilogue::LseFormat::BSN1>(
257+ query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,
258+ qDequantScale, kDequantScale, vDequantScale, attentionOut, user, softmaxLse, tiling);
259+ #elif TILING_KEY_VAR == QF8_KVF8_OBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY
260+ BsaInferInterfaceFullQuant<
261+ fp8_e4m3fn_t, bfloat16_t, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND,
262+ Epilogue::LseMode::OUT_ONLY, Epilogue::LseFormat::TN1>(
263+ query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,
264+ qDequantScale, kDequantScale, vDequantScale, attentionOut, user, softmaxLse, tiling);
265+ #elif TILING_KEY_VAR == QF8_KVF8_OBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY
266+ BsaInferInterfaceFullQuant<
267+ fp8_e4m3fn_t, bfloat16_t, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD,
268+ Epilogue::LseMode::OUT_ONLY, Epilogue::LseFormat::BNS1>(
269+ query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,
270+ qDequantScale, kDequantScale, vDequantScale, attentionOut, user, softmaxLse, tiling);
271+ #elif TILING_KEY_VAR == QF8_KVF8_OBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY
272+ BsaInferInterfaceFullQuant<
273+ fp8_e4m3fn_t, bfloat16_t, float, BsaKernelArch35::Format::BSND, BsaKernelArch35::Format::BSND,
274+ Epilogue::LseMode::OUT_ONLY, Epilogue::LseFormat::BSN1>(
275+ query, key, value, blockSparseMask, mask, blockTable, actualSeqLengths, actualSeqLengthsKv,
276+ qDequantScale, kDequantScale, vDequantScale, attentionOut, user, softmaxLse, tiling);
180 #endif277 #endif
181#endif278#endif
182 }279 }
Mattention/block_sparse_attention/op_kernel/block_sparse_attention_kernel_interface.cpp+14-7
@@ -119,11 +119,14 @@ __global__ __aicore__ void BlockSparseAttentionInfer(
119 119 
120using namespace BsaKernelArch35;120using namespace BsaKernelArch35;
121 121 
122-template <class InDtype, class SMDtype, class REDtype, Format qFormat, Format kvFormat>122+template <
123+ class InDtype, class SMDtype, class REDtype,
124+ Format qFormat, Format kvFormat,
125+ Epilogue::LseMode lseMode, Epilogue::LseFormat lseFormat>
123__global__ __aicore__ void BsaInferIntfRegular(126__global__ __aicore__ void BsaInferIntfRegular(
124 GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR mask, GM_ADDR blockTables,127 GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR mask, GM_ADDR blockTables,
125 GM_ADDR o, GM_ADDR actualQseqlen, GM_ADDR actualKvseqlen,128 GM_ADDR o, GM_ADDR actualQseqlen, GM_ADDR actualKvseqlen,
126- GM_ADDR blockSparseMask, GM_ADDR workspace,129+ GM_ADDR blockSparseMask, GM_ADDR workspace, GM_ADDR lse,
127 GM_ADDR tiling130 GM_ADDR tiling
128) {131) {
129 using ArchTag = Arch::AtlasA5;132 using ArchTag = Arch::AtlasA5;
@@ -178,7 +181,7 @@ __global__ __aicore__ void BsaInferIntfRegular(
178 using BlockMmadPV = Gemm::Block::BlockMmadTla<181 using BlockMmadPV = Gemm::Block::BlockMmadTla<
179 DispatchPolicyPV, L1TileShapePV, L0TileShapePV, ElementP, ElementV, ElementOTmp, void, TileCopyPV>;182 DispatchPolicyPV, L1TileShapePV, L0TileShapePV, ElementP, ElementV, ElementOTmp, void, TileCopyPV>;
180 // rescale O183 // rescale O
181- using DispatchPolicyRescaleO = Epilogue::EpilogueAtlasA5BsaRescaleO;184+ using DispatchPolicyRescaleO = Epilogue::EpilogueAtlasA5BsaRescaleO<lseMode, lseFormat>;
182 using TileCopyRescaleO = Epilogue::Tile::TileCopyRescaleO<185 using TileCopyRescaleO = Epilogue::Tile::TileCopyRescaleO<
183 ArchTag, ElementO, LayoutO, LayoutOTmp>;186 ArchTag, ElementO, LayoutO, LayoutOTmp>;
184 using EpilogueRescaleO = Epilogue::Block::BlockEpilogue<187 using EpilogueRescaleO = Epilogue::Block::BlockEpilogue<
@@ -187,17 +190,20 @@ __global__ __aicore__ void BsaInferIntfRegular(
187 using BsaRegularKernelArch35 = BsaRegularKernelArch35<190 using BsaRegularKernelArch35 = BsaRegularKernelArch35<
188 EpilogueMask2Idx, BlockMmadQK, EpilogueOnlineSoftmax, BlockMmadPV, EpilogueRescaleO, qFormat, kvFormat>;191 EpilogueMask2Idx, BlockMmadQK, EpilogueOnlineSoftmax, BlockMmadPV, EpilogueRescaleO, qFormat, kvFormat>;
189 BsaKernelParamsArch35 params{q, k, v, mask, blockTables,192 BsaKernelParamsArch35 params{q, k, v, mask, blockTables,
190- actualQseqlen, actualKvseqlen, blockSparseMask, o, workspace, tiling};193+ actualQseqlen, actualKvseqlen, blockSparseMask, o, workspace, lse, tiling};
191 BsaRegularKernelArch35 bsaRegularKernelArch35;194 BsaRegularKernelArch35 bsaRegularKernelArch35;
192 bsaRegularKernelArch35(params);195 bsaRegularKernelArch35(params);
193}196}
194 197 
195-template <class InDtype, class SMDtype, class REDtype, Format qFormat, Format kvFormat>198+template <
199+ class InDtype, class SMDtype, class REDtype,
200+ Format qFormat, Format kvFormat,
201+ Epilogue::LseMode lseMode, Epilogue::LseFormat lseFormat>
196__global__ __aicore__ void BsaInferInterfaceFullQuant(202__global__ __aicore__ void BsaInferInterfaceFullQuant(
197 GM_ADDR query, GM_ADDR key, GM_ADDR value, GM_ADDR blockSparseMask, GM_ADDR attenMask,203 GM_ADDR query, GM_ADDR key, GM_ADDR value, GM_ADDR blockSparseMask, GM_ADDR attenMask,
198 GM_ADDR blockTable, GM_ADDR actualSeqLengths, GM_ADDR actualSeqLengthsKv,204 GM_ADDR blockTable, GM_ADDR actualSeqLengths, GM_ADDR actualSeqLengthsKv,
199 GM_ADDR qDequantScale, GM_ADDR kDequantScale, GM_ADDR vDequantScale, GM_ADDR attentionOut,205 GM_ADDR qDequantScale, GM_ADDR kDequantScale, GM_ADDR vDequantScale, GM_ADDR attentionOut,
200- GM_ADDR workspace, GM_ADDR tiling)206+ GM_ADDR workspace, GM_ADDR lse, GM_ADDR tiling)
201{207{
202 using ArchTag = Arch::AtlasA5;208 using ArchTag = Arch::AtlasA5;
203 using ElementSparseMask = uint8_t;209 using ElementSparseMask = uint8_t;
@@ -254,7 +260,7 @@ __global__ __aicore__ void BsaInferInterfaceFullQuant(
254 using BlockMmadPV = Gemm::Block::BlockMmadTla<DispatchPolicyPV, L1TileShapePV, L0TileShapePV, ElementP, ElementV,260 using BlockMmadPV = Gemm::Block::BlockMmadTla<DispatchPolicyPV, L1TileShapePV, L0TileShapePV, ElementP, ElementV,
255 ElementOTmp, void, TileCopyPV>;261 ElementOTmp, void, TileCopyPV>;
256 // rescale o262 // rescale o
257- using DispatchPolicyRescaleO = Epilogue::EpilogueAtlasA5BsaRescaleO;263+ using DispatchPolicyRescaleO = Epilogue::EpilogueAtlasA5BsaRescaleO<lseMode, lseFormat>;
258 using TileCopyRescaleO = Epilogue::Tile::TileCopyRescaleO<ArchTag, ElementO, LayoutO, LayoutOTmp>;264 using TileCopyRescaleO = Epilogue::Tile::TileCopyRescaleO<ArchTag, ElementO, LayoutO, LayoutOTmp>;
259 using EpilogueRescaleO = Epilogue::Block::BlockEpilogue<DispatchPolicyRescaleO, ElementO, ElementOTmp, ElementS,265 using EpilogueRescaleO = Epilogue::Block::BlockEpilogue<DispatchPolicyRescaleO, ElementO, ElementOTmp, ElementS,
260 TileCopyRescaleO, Arch::PositionL0C>;266 TileCopyRescaleO, Arch::PositionL0C>;
@@ -270,6 +276,7 @@ __global__ __aicore__ void BsaInferInterfaceFullQuant(
270 kDequantScale,276 kDequantScale,
271 vDequantScale,277 vDequantScale,
272 attentionOut,278 attentionOut,
279+ lse,
273 workspace,280 workspace,
274 tiling};281 tiling};
275 using BsaFullQuantKernelArch35 = BsaFullQuantKernelArch35<EpilogueMask2Idx, BlockMmadQK, EpilogueOnlineSoftmax,282 using BsaFullQuantKernelArch35 = BsaFullQuantKernelArch35<EpilogueMask2Idx, BlockMmadQK, EpilogueOnlineSoftmax,
Mattention/block_sparse_attention/op_kernel/block_sparse_attention_tilingkey.h+16-3
@@ -68,21 +68,34 @@
68#endif68#endif
69 69 
70#if (__CCE_AICORE__ == 310)70#if (__CCE_AICORE__ == 310)
71-// 非量化FP16/BF16场景71+// 非量化+混合精度softmax
72#define QF16_KVF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY 905000003040000272#define QF16_KVF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY 9050000030400002
73#define QBF16_KVBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY 905000003042222273#define QBF16_KVBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY 9050000030422222
74#define QF16_KVF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY 905000005040000374#define QF16_KVF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY 9050000050400003
75#define QBF16_KVBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY 905000005042222375#define QBF16_KVBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY 9050000050422223
76#define QF16_KVF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_KEY 905000006040000476#define QF16_KVF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_KEY 9050000060400004
77#define QBF16_KVBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_KEY 905000006042222477#define QBF16_KVBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_KEY 9050000060422224
78- 78+// 非量化+混合精度softmax+lse
79-// 全量化FP8场景79+#define QF16_KVF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY 9050000130400002
80+#define QBF16_KVBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY 9050000130422222
81+#define QF16_KVF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY 9050000150400003
82+#define QBF16_KVBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY 9050000150422223
83+#define QF16_KVF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY 9050000160400004
84+#define QBF16_KVBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY 9050000160422224
85+// 全量化FP8+混合精度softmax
80#define QF8_KVF8_OF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY 905000003040001286#define QF8_KVF8_OF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY 9050000030400012
81#define QF8_KVF8_OF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY 905000005040001387#define QF8_KVF8_OF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY 9050000050400013
82#define QF8_KVF8_OF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_KEY 905000006040001488#define QF8_KVF8_OF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_KEY 9050000060400014
83#define QF8_KVF8_OBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY 905000003040002289#define QF8_KVF8_OBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY 9050000030400022
84#define QF8_KVF8_OBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY 905000005040002390#define QF8_KVF8_OBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY 9050000050400023
85#define QF8_KVF8_OBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_KEY 905000006040002491#define QF8_KVF8_OBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_KEY 9050000060400024
92+// 全量化FP8+混合精度softmax+lse
93+#define QF8_KVF8_OF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY 9050000130400012
94+#define QF8_KVF8_OF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY 9050000150400013
95+#define QF8_KVF8_OF16_QBSND_KVBSND_NOCACHE_SMF16_REF32_NOMASK_LSE_OUT_KEY 9050000160400014
96+#define QF8_KVF8_OBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY 9050000130400022
97+#define QF8_KVF8_OBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY 9050000150400023
98+#define QF8_KVF8_OBF16_QBSND_KVBSND_NOCACHE_SMBF16_REF32_NOMASK_LSE_OUT_KEY 9050000160400024
86 99 
87#endif100#endif
88#endif // BLOCK_SPARSE_ATTENTION_TILINGKEY_H_101#endif // BLOCK_SPARSE_ATTENTION_TILINGKEY_H_