Star106
140
代码
代码
Issues111
Pull Requests20
项目讨论
Wiki
分析
Star106
140
已合并
[Feature][op][eqbsa-02] add eqbsa operator, attention infra, and enable build #514
lanwangli创建于 20 天前
[Feature][op][eqbsa-02] add eqbsa operator, attention infra, and enable build #514
已合并
lanwangli创建于 20 天前
从lanwangli/MindIE-SD_0813eqbsa-op-and-build合入到Ascend/MindIE-SDdev
共 83 个文件变更+22792-1
Mbuild/build_ops.sh+2-1
@@ -36,7 +36,8 @@ if [ ! -f ${msopgen} ]; then
36 echo "${msopgen} not exists"36 echo "${msopgen} not exists"
37 exit 137 exit 1
38fi38fi
39-ascendc_ops=${ASCEND_OP_NAME:-'laser_attention;la_preprocess;ada_block_sparse_attention;sparse_block_estimate;norm_rope_concat;quant_flash_attn;quant_flash_attn_metadata;fused_infer_attention_score;mul_add'}39+ 
40+ascendc_ops=${ASCEND_OP_NAME:-'laser_attention;la_preprocess;ada_block_sparse_attention;sparse_block_estimate;norm_rope_concat;quant_flash_attn;quant_flash_attn_metadata;fused_infer_attention_score;eagle_quant_block_sparse_attention;mul_add'}
40 41 
41# ascend950 backend requires CANN 9.0+; remove ascend950 when CANN < 9.042# ascend950 backend requires CANN 9.0+; remove ascend950 when CANN < 9.0
42default_compute_unit='ascend910;ascend910b;ascend910_93;ascend950'43default_compute_unit='ascend910;ascend910b;ascend910_93;ascend950'
Acsrc/ops/eagle_quant_block_sparse_attention/op_api/aclnn_eagle_quant_block_sparse_attention.cpp+332-0
@@ -0,0 +1,332 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#include "aclnn_eagle_quant_block_sparse_attention.h"
12+ 
13+#include "eagle_quant_block_sparse_attention.h"
14+#include "aclnn_kernels/contiguous.h"
15+#include "opdev/make_op_executor.h"
16+#include "opdev/op_dfx.h"
17+#include "opdev/tensor_view_utils.h"
18+#include "opdev/common_types.h"
19+#include "opdev/op_errno.h"
20+#include "opdev/op_executor.h"
21+#include <acl/acl.h>
22+#include <algorithm>
23+#include <unordered_map>
24+#include <string>
25+ 
26+using namespace op;
27+ 
28+#ifdef __cplusplus
29+extern "C" {
30+#endif
31+ 
32+namespace {
33+ 
34+static constexpr uint64_t LSE_OUT = 1;
35+ 
36+static bool CheckDataType(const aclTensor *query,
37+ const aclTensor *key,
38+ const aclTensor *value)
39+{
40+ const DataType qDtype = query->GetDataType();
41+ const DataType kDtype = key->GetDataType();
42+ const DataType vDtype = value->GetDataType();
43+ 
44+ 
45+ static const std::unordered_map<DataType, std::vector<DataType>> validKvType = {
46+ {DataType::DT_FLOAT16, {DataType::DT_FLOAT16}},
47+ {DataType::DT_BF16, {DataType::DT_BF16}},
48+ {DataType::DT_INT8, {DataType::DT_INT8, DataType::DT_FLOAT8_E4M3FN}}
49+ };
50+ 
51+ auto iter = validKvType.find(qDtype);
52+ if (iter == validKvType.end()) {
53+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Unsupported query datatype %d.", static_cast<int>(qDtype));
54+ return false;
55+ }
56+
57+ if (std::find(iter->second.begin(), iter->second.end(), kDtype) == iter->second.end() ||
58+ std::find(iter->second.begin(), iter->second.end(), vDtype) == iter->second.end()) {
59+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Key/Value datatype mismatch with query.");
60+ return false;
61+ }
62+ 
63+ return true;
64+}
65+ 
66+ 
67+static aclnnStatus CheckMandatoryTensors(const aclTensor *query,
68+ const aclTensor *key,
69+ const aclTensor *value)
70+{
71+ CHECK_RET(query != nullptr, ACLNN_ERR_PARAM_NULLPTR);
72+ CHECK_RET(key != nullptr, ACLNN_ERR_PARAM_NULLPTR);
73+ CHECK_RET(value != nullptr, ACLNN_ERR_PARAM_NULLPTR);
74+ return ACLNN_SUCCESS;
75+}
76+ 
77+static aclnnStatus ParseblockShapeOptional(const aclIntArray *blockShapeOptional)
78+{
79+ if (blockShapeOptional == nullptr) {
80+ OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "blockShapeOptional is null.");
81+ return ACLNN_ERR_PARAM_NULLPTR;
82+ }
83+ 
84+ uint64_t size = blockShapeOptional->Size();
85+ if (size < 2) {
86+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "blockShapeOptional must contain at least two elements [x, y].");
87+ return ACLNN_ERR_PARAM_INVALID;
88+ }
89+ 
90+ const int64_t *data = blockShapeOptional->GetData();
91+ if (data == nullptr) {
92+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "blockShapeOptional data is null.");
93+ return ACLNN_ERR_PARAM_INVALID;
94+ }
95+ 
96+ if (data[0] <= 0 || data[1] <= 0) {
97+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "blockShapeOptional values must be positive, got [%ld, %ld].", data[0], data[1]);
98+ return ACLNN_ERR_PARAM_INVALID;
99+ }
100+ 
101+ return ACLNN_SUCCESS;
102+}
103+ 
104+static aclnnStatus ValidateParams(const aclTensor *query,
105+ const aclTensor *key,
106+ const aclTensor *value,
107+ char *qInputLayout,
108+ char *kvInputLayout,
109+ const aclIntArray *blockShapeOptional)
110+{
111+ CHECK_RET(CheckMandatoryTensors(query, key, value) == ACLNN_SUCCESS,
112+ ACLNN_ERR_PARAM_NULLPTR);
113+ 
114+ if (!CheckDataType(query, key, value)) {
115+ return ACLNN_ERR_PARAM_INVALID;
116+ }
117+ 
118+ if (qInputLayout == nullptr || kvInputLayout == nullptr) {
119+ OP_LOGE(ACLNN_ERR_PARAM_NULLPTR, "Input layout strings are null.");
120+ return ACLNN_ERR_PARAM_NULLPTR;
121+ }
122+ std::string qLayout(qInputLayout);
123+ std::string kvLayout(kvInputLayout);
124+
125+ // 验证Q layout
126+ if (qLayout != "TND" && qLayout != "BNSD") {
127+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "qInputLayout only supports TND or BNSD, got %s.", qLayout.c_str());
128+ return ACLNN_ERR_PARAM_INVALID;
129+ }
130+
131+ // 验证KV layout
132+ if (kvLayout != "TND" && kvLayout != "BNSD") {
133+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "kvInputLayout only supports TND or BNSD, got %s.", kvLayout.c_str());
134+ return ACLNN_ERR_PARAM_INVALID;
135+ }
136+
137+ // 验证Q和KV格式一致性：如果其中一个是BNSD，另一个也必须是BNSD
138+ bool qIsBNSD = (qLayout == "BNSD");
139+ bool kvIsBNSD = (kvLayout == "BNSD");
140+ if (qIsBNSD != kvIsBNSD) {
141+ OP_LOGE(ACLNN_ERR_PARAM_INVALID,
142+ "Q and KV layouts must match: if one is BNSD, the other must also be BNSD. "
143+ "Q layout: %s, KV layout: %s", qLayout.c_str(), kvLayout.c_str());
144+ return ACLNN_ERR_PARAM_INVALID;
145+ }
146+ 
147+ return ParseblockShapeOptional(blockShapeOptional);
148+}
149+ 
150+static aclnnStatus MakeContiguous(const aclTensor *&query,
151+ const aclTensor *&key,
152+ const aclTensor *&value,
153+ const aclTensor *&blockSparseMaskOptional,
154+ const aclTensor *&attenMaskOptional,
155+ const aclTensor *&blockTableOptional,
156+ const aclTensor *&queryScaleOptional,
157+ const aclTensor *&keyScaleOptional,
158+ const aclTensor *&valueScaleOptional,
159+ aclOpExecutor *executor)
160+{
161+ query = l0op::Contiguous(query, executor);
162+ CHECK_RET(query != nullptr, ACLNN_ERR_PARAM_NULLPTR);
163+ 
164+ key = l0op::Contiguous(key, executor);
165+ CHECK_RET(key != nullptr, ACLNN_ERR_PARAM_NULLPTR);
166+ 
167+ value = l0op::Contiguous(value, executor);
168+ CHECK_RET(value != nullptr, ACLNN_ERR_PARAM_NULLPTR);
169+ //新增blockSparseMaskOptional非空校验且必须为四维
170+
171+ if (blockSparseMaskOptional != nullptr) {
172+ blockSparseMaskOptional = l0op::Contiguous(blockSparseMaskOptional, executor);
173+ CHECK_RET(blockSparseMaskOptional != nullptr, ACLNN_ERR_INNER_NULLPTR);
174+ if (blockSparseMaskOptional->GetStorageShape().GetDimNum() != 4) {
175+ OP_LOGE(ACLNN_ERR_PARAM_INVALID,
176+ "blockSparseMask must be 4D tensor");
177+ return ACLNN_ERR_PARAM_INVALID;
178+ }
179+ //CHECK_RET(blockSparseMaskOptional->GetStorageShape().GetDim() == 4, "blockSparseMask must be 4D tensor");
180+ }
181+ 
182+ if (attenMaskOptional != nullptr) {
183+ attenMaskOptional = l0op::Contiguous(attenMaskOptional, executor);
184+ CHECK_RET(attenMaskOptional != nullptr, ACLNN_ERR_INNER_NULLPTR);
185+ }
186+ 
187+ if (blockTableOptional != nullptr) {
188+ blockTableOptional = l0op::Contiguous(blockTableOptional, executor);
189+ CHECK_RET(blockTableOptional != nullptr, ACLNN_ERR_INNER_NULLPTR);
190+ }
191+ 
192+ if (queryScaleOptional != nullptr) {
193+ queryScaleOptional = l0op::Contiguous(queryScaleOptional, executor);
194+ CHECK_RET(queryScaleOptional != nullptr, ACLNN_ERR_INNER_NULLPTR);
195+ }
196+ 
197+ if (keyScaleOptional != nullptr) {
198+ keyScaleOptional = l0op::Contiguous(keyScaleOptional, executor);
199+ CHECK_RET(keyScaleOptional != nullptr, ACLNN_ERR_INNER_NULLPTR);
200+ }
201+ 
202+ if (valueScaleOptional != nullptr) {
203+ valueScaleOptional = l0op::Contiguous(valueScaleOptional, executor);
204+ CHECK_RET(valueScaleOptional != nullptr, ACLNN_ERR_INNER_NULLPTR);
205+ }
206+ 
207+ return ACLNN_SUCCESS;
208+}
209+ 
210+static aclnnStatus ValidateAdditionalParams(int64_t innerPrecise,
211+ const aclTensor *attentionOut,
212+ uint64_t *workspaceSize,
213+ aclOpExecutor **executor)
214+{
215+ if (innerPrecise != 0 && innerPrecise != 1 && innerPrecise != 4) {
216+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "innerPrecise must be 0 or 1 or 4, got %ld.",
217+ innerPrecise);
218+ return ACLNN_ERR_PARAM_INVALID;
219+ }
220+
221+ CHECK_RET(attentionOut != nullptr, ACLNN_ERR_PARAM_NULLPTR);
222+ CHECK_RET(workspaceSize != nullptr, ACLNN_ERR_PARAM_NULLPTR);
223+ CHECK_RET(executor != nullptr, ACLNN_ERR_PARAM_NULLPTR);
224+
225+ return ACLNN_SUCCESS;
226+}
227+ 
228+static string ConvertLayoutString(char *layoutStr)
229+{
230+ return op::ToString(layoutStr).GetString();
231+}
232+ 
233+} // namespace
234+ 
235+__attribute__((visibility("default"))) aclnnStatus aclnnEagleQuantBlockSparseAttentionGetWorkspaceSize(
236+ const aclTensor *query,
237+ const aclTensor *key,
238+ const aclTensor *value,
239+ const aclTensor *blockSparseMask,
240+ const aclTensor *attenMaskOptional,
241+ const aclIntArray *blockShape,
242+ const aclIntArray *actualSeqLengthsOptional,
243+ const aclIntArray *actualSeqLengthsKvOptional,
244+ const aclTensor *blockTableOptional,
245+ const aclTensor *queryScaleOptional,
246+ const aclTensor *keyScaleOptional,
247+ const aclTensor *valueScaleOptional,
248+ char *qInputLayout,
249+ char *kvInputLayout,
250+ int64_t numKeyValueHeads,
251+ int64_t maskType,
252+ double scaleValue,
253+ int64_t innerPrecise,
254+ int64_t blockSize,
255+ int64_t preTokens,
256+ int64_t nextTokens,
257+ int64_t softmaxLseFlag,
258+ aclTensor *attentionOut,
259+ aclTensor *softmaxLseOptional,
260+ uint64_t *workspaceSize,
261+ aclOpExecutor **executor)
262+{
263+ aclnnStatus ret = ValidateParams(query, key, value,
264+ qInputLayout, kvInputLayout, blockShape);
265+ if (ret != ACLNN_SUCCESS) {
266+ return ret;
267+ }
268+ 
269+ ret = ValidateAdditionalParams(innerPrecise, attentionOut, workspaceSize, executor);
270+ if (ret != ACLNN_SUCCESS) {
271+ return ret;
272+ }
273+ //去掉了idx ,idxnums
274+ L2_DFX_PHASE_1(aclnnEagleQuantBlockSparseAttention,
275+ DFX_IN(query, key, value, blockSparseMask, attenMaskOptional, blockShape, actualSeqLengthsOptional,
276+ queryScaleOptional, keyScaleOptional, valueScaleOptional,
277+ actualSeqLengthsKvOptional, blockTableOptional, qInputLayout, qInputLayout, numKeyValueHeads,
278+ maskType, scaleValue, innerPrecise, blockSize, preTokens, nextTokens, softmaxLseFlag),
279+ DFX_OUT(attentionOut, softmaxLseOptional));
280+ 
281+ auto uniqueExecutor = CREATE_EXECUTOR();
282+ CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_NULLPTR);
283+ auto *executorImpl = uniqueExecutor.get();
284+ //新增blockSparseMaskOptional参数
285+ ret = MakeContiguous(query, key, value, blockSparseMask, attenMaskOptional, blockTableOptional,
286+ queryScaleOptional, keyScaleOptional, valueScaleOptional,
287+ executorImpl);
288+ if (ret != ACLNN_SUCCESS) {
289+ return ret;
290+ }
291+
292+ string qInputLayoutStr = ConvertLayoutString(qInputLayout);
293+ string kvInputLayoutStr = ConvertLayoutString(kvInputLayout);
294+ //新增blockSparseMaskOptional参数
295+ auto outputs = l0op::EagleQuantBlockSparseAttention(query, key, value,
296+ blockSparseMask, attenMaskOptional, blockShape, actualSeqLengthsOptional,
297+ actualSeqLengthsKvOptional, blockTableOptional,
298+ queryScaleOptional, keyScaleOptional, valueScaleOptional,
299+ qInputLayoutStr.c_str(), kvInputLayoutStr.c_str(), numKeyValueHeads,
300+ maskType, scaleValue, innerPrecise, blockSize, preTokens, nextTokens, softmaxLseFlag, executorImpl,
301+ attentionOut->GetDataType());
302+ if (outputs[0] == nullptr || outputs[1] == nullptr) {
303+ OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "EagleQuantBlockSparseAttention returned nullptr outputs.");
304+ return ACLNN_ERR_INNER_NULLPTR;
305+ }
306+ 
307+ auto viewCopyResult = l0op::ViewCopy(outputs[0], attentionOut, executorImpl);
308+ CHECK_RET(viewCopyResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
309+ if (softmaxLseFlag == LSE_OUT) {
310+ auto viewCopyLseResult = l0op::ViewCopy(outputs[1], softmaxLseOptional, executorImpl);
311+ CHECK_RET(viewCopyLseResult != nullptr, ACLNN_ERR_INNER_NULLPTR);
312+ }
313+ 
314+ *workspaceSize = executorImpl->GetWorkspaceSize();
315+ uniqueExecutor.ReleaseTo(executor);
316+ return ACLNN_SUCCESS;
317+}
318+ 
319+__attribute__((visibility("default"))) aclnnStatus aclnnEagleQuantBlockSparseAttention(
320+ void *workspace,
321+ uint64_t workspaceSize,
322+ aclOpExecutor *executor,
323+ aclrtStream stream)
324+{
325+ L2_DFX_PHASE_2(aclnnEagleQuantBlockSparseAttention);
326+ return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
327+}
328+ 
329+#ifdef __cplusplus
330+}
331+#endif
332+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_api/aclnn_eagle_quant_block_sparse_attention.h+59-0
@@ -0,0 +1,59 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef ACLNN_EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_H_
12+#define ACLNN_EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_H_
13+ 
14+#include "aclnn/acl_meta.h"
15+ 
16+#ifdef __cplusplus
17+extern "C" {
18+#endif
19+ 
20+__attribute__((visibility("default"))) aclnnStatus aclnnEagleQuantBlockSparseAttentionGetWorkspaceSize(
21+ const aclTensor *query,
22+ const aclTensor *key,
23+ const aclTensor *value,
24+ const aclTensor *blockSparseMaskOptional,
25+ const aclTensor *attenMaskOptional,
26+ const aclIntArray *blockShapeOptional,
27+ const aclIntArray *actualSeqLengthsOptional,
28+ const aclIntArray *actualSeqLengthsKvOptional,
29+ const aclTensor *blockTableOptional,
30+ const aclTensor *queryScaleOptional,
31+ const aclTensor *keyScaleOptional,
32+ const aclTensor *valueSclaleOptional,
33+ char *qInputLayout,
34+ char *kvInputLayout,
35+ int64_t numKeyValueHeads,
36+ int64_t maskType,
37+ double scaleValue,
38+ int64_t innerPrecise,
39+ int64_t blockSize,
40+ int64_t preTokens,
41+ int64_t nextTokens,
42+ int64_t softmaxLseFlag,
43+ aclTensor *attentionOut,
44+ aclTensor *softmaxLseOptional,
45+ uint64_t *workspaceSize,
46+ aclOpExecutor **executor);
47+ 
48+__attribute__((visibility("default"))) aclnnStatus aclnnEagleQuantBlockSparseAttention(
49+ void *workspace,
50+ uint64_t workspaceSize,
51+ aclOpExecutor *executor,
52+ aclrtStream stream);
53+ 
54+#ifdef __cplusplus
55+}
56+#endif
57+ 
58+#endif // ACLNN_EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_H_
59+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_api/eagle_quant_block_sparse_attention.cpp+125-0
@@ -0,0 +1,125 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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 eagle_quant_block_sparse_attention.cpp
13+ * \brief
14+ */
15+ 
16+ 
17+#include "eagle_quant_block_sparse_attention.h"
18+ 
19+#include "opdev/make_op_executor.h"
20+#include "opdev/op_dfx.h"
21+ 
22+using namespace op;
23+ 
24+namespace l0op {
25+ 
26+OP_TYPE_REGISTER(EagleQuantBlockSparseAttention);
27+ 
28+static const aclTensor *ConvertIntArrayToTensor(const aclIntArray *intArray,
29+ aclOpExecutor *executor,
30+ DataType dtype)
31+{
32+ if (intArray != nullptr) {
33+ const aclTensor *tensor = executor->ConvertToTensor(intArray, dtype);
34+ auto mutableTensor = const_cast<aclTensor *>(tensor);
35+ mutableTensor->SetStorageFormat(Format::FORMAT_ND);
36+ mutableTensor->SetViewFormat(Format::FORMAT_ND);
37+ mutableTensor->SetOriginalFormat(Format::FORMAT_ND);
38+ return tensor;
39+ }
40+ return executor->AllocTensor(dtype, Format::FORMAT_ND, Format::FORMAT_ND);
41+}
42+ 
43+const std::array<const aclTensor *, 2> EagleQuantBlockSparseAttention(
44+ const aclTensor *query,
45+ const aclTensor *key,
46+ const aclTensor *value,
47+ const aclTensor *blockSparseMaskOptional,
48+ const aclTensor *attenMaskOptional,
49+ const aclIntArray *blockShapeOptional,
50+ const aclIntArray *actualSeqLengthsOptional,
51+ const aclIntArray *actualSeqLengthsKvOptional,
52+ const aclTensor *blockTableOptional,
53+ const aclTensor *queryScaleOptional,
54+ const aclTensor *keyScaleOptional,
55+ const aclTensor *valueScaleOptional,
56+ const char *qInputLayout,
57+ const char *kvInputLayout,
58+ int64_t numKeyValueHeads,
59+ int64_t maskType,
60+ double scaleValue,
61+ int64_t innerPrecise,
62+ int64_t blockSize,
63+ int64_t preTokens,
64+ int64_t nextTokens,
65+ int64_t softmaxLseFlag,
66+ aclOpExecutor *executor,
67+ DataType outDType
68+ )
69+{
70+ const char *safeKvInputLayout = (kvInputLayout != nullptr) ? kvInputLayout : qInputLayout;
71+
72+ L0_DFX(EagleQuantBlockSparseAttention, query, key, value, blockSparseMaskOptional,
73+ attenMaskOptional, blockShapeOptional, actualSeqLengthsOptional, actualSeqLengthsKvOptional,
74+ blockTableOptional, queryScaleOptional, keyScaleOptional, valueScaleOptional,
75+ qInputLayout, safeKvInputLayout, numKeyValueHeads,
76+ maskType, scaleValue, innerPrecise, blockSize, preTokens, nextTokens, softmaxLseFlag);
77+ 
78+ const aclTensor *blockShapeOptionalTensor = nullptr;
79+ if (blockShapeOptional) {
80+ blockShapeOptionalTensor = ConvertIntArrayToTensor(blockShapeOptional, executor, DataType::DT_INT64);
81+ }
82+ const aclTensor *actualSeqTensor = nullptr;
83+ if (actualSeqLengthsOptional) {
84+ actualSeqTensor = ConvertIntArrayToTensor(actualSeqLengthsOptional, executor, DataType::DT_INT64);
85+ }
86+ const aclTensor *actualSeqKvTensor = nullptr;
87+ if (actualSeqLengthsKvOptional) {
88+ actualSeqKvTensor = ConvertIntArrayToTensor(actualSeqLengthsKvOptional, executor, DataType::DT_INT64);
89+ }
90+ 
91+ auto attentionOutTensor = executor->AllocTensor(outDType, Format::FORMAT_ND, Format::FORMAT_ND);
92+ auto softmaxLseTensor = executor->AllocTensor(DataType::DT_FLOAT, Format::FORMAT_ND, Format::FORMAT_ND);
93+ 
94+ // scaleValue is already float type, no need for cast
95+ auto ret = INFER_SHAPE(EagleQuantBlockSparseAttention,
96+ OP_INPUT(query, key, value, blockSparseMaskOptional, attenMaskOptional,
97+ blockShapeOptionalTensor, actualSeqTensor, actualSeqKvTensor, blockTableOptional,
98+ queryScaleOptional, keyScaleOptional, valueScaleOptional),
99+ OP_OUTPUT(attentionOutTensor, softmaxLseTensor),
100+ OP_ATTR(qInputLayout, safeKvInputLayout,
101+ static_cast<int64_t>(numKeyValueHeads), static_cast<int64_t>(maskType),
102+ static_cast<float>(scaleValue), static_cast<int64_t>(innerPrecise),
103+ static_cast<int64_t>(blockSize), static_cast<uint32_t>(preTokens),
104+ static_cast<int64_t>(nextTokens), static_cast<int64_t>(softmaxLseFlag)));
105+ if (ret != ACLNN_SUCCESS) {
106+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "EagleQuantBlockSparseAttention infer shape failed, scaleValue: %f.", scaleValue);
107+ return {nullptr, nullptr};
108+ }
109+
110+ ADD_TO_LAUNCHER_LIST_AICORE(EagleQuantBlockSparseAttention,
111+ OP_INPUT(query, key, value, blockSparseMaskOptional, attenMaskOptional,
112+ blockShapeOptionalTensor, actualSeqTensor, actualSeqKvTensor,
113+ blockTableOptional, queryScaleOptional, keyScaleOptional, valueScaleOptional),
114+ OP_OUTPUT(attentionOutTensor, softmaxLseTensor),
115+ OP_ATTR(qInputLayout, safeKvInputLayout, static_cast<int64_t>(numKeyValueHeads),
116+ static_cast<int64_t>(maskType), static_cast<float>(scaleValue),
117+ static_cast<int64_t>(innerPrecise), static_cast<int64_t>(blockSize),
118+ static_cast<int64_t>(preTokens), static_cast<int64_t>(nextTokens),
119+ static_cast<int64_t>(softmaxLseFlag)));
120+ 
121+ return {attentionOutTensor, softmaxLseTensor};
122+}
123+} // namespace l0op
124+ 
125+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_api/eagle_quant_block_sparse_attention.h+49-0
@@ -0,0 +1,49 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_H_
12+#define EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_H_
13+ 
14+#include <array>
15+#include "opdev/op_executor.h"
16+#include "opdev/common_types.h"
17+ 
18+namespace l0op {
19+ 
20+const std::array<const aclTensor *, 2> EagleQuantBlockSparseAttention(
21+ const aclTensor *query,
22+ const aclTensor *key,
23+ const aclTensor *value,
24+ const aclTensor *blockSparseMaskOptional,
25+ const aclTensor *attenMaskOptional,
26+ const aclIntArray *blockShapeOptional,
27+ const aclIntArray *actualSeqLengthsOptional,
28+ const aclIntArray *actualSeqLengthsKvOptional,
29+ const aclTensor *blockTableOptional,
30+ const aclTensor *queryScaleOptional,
31+ const aclTensor *keyScaleOptional,
32+ const aclTensor *valueScaleOptional,
33+ const char *qInputLayout,
34+ const char *kvInputLayout,
35+ int64_t numKeyValueHeads,
36+ int64_t maskType,
37+ double scaleValue,
38+ int64_t innerPrecise,
39+ int64_t blockSize,
40+ int64_t preTokens,
41+ int64_t nextTokens,
42+ int64_t softmaxLseFlag,
43+ aclOpExecutor *executor,
44+ op::DataType outDType);
45+ 
46+} // namespace l0op
47+ 
48+#endif // EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_H_
49+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_host/CMakeLists.txt+58-0
@@ -0,0 +1,58 @@
1+# Copyright (c) Huawei Technologies Co., Ltd. 2026-2026. All rights reserved.
2+# MindIE is licensed under Mulan PSL v2.
3+# You can use this software according to the terms and conditions of the Mulan PSL v2.
4+# You may obtain a copy of Mulan PSL v2 at:
5+# http://license.coscl.org.cn/MulanPSL2
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
7+# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
8+# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
9+# See the Mulan PSL v2 for more details.
10+ 
11+add_ops_compile_options(
12+ OP_NAME EagleQuantBlockSparseAttention
13+ OPTIONS --cce-auto-sync=off
14+ -Wno-deprecated-declarations
15+ -fpermissive
16+)
17+ 
18+# def.cpp is routed to the "inner" aclnn generator so that opbuild produces
19+# aclnnInnerEagleQuantBlockSparseAttention and does NOT collide with the
20+# hand-written public aclnnEagleQuantBlockSparseAttention provided via op_api.
21+target_sources(op_host_aclnnInner PRIVATE
22+ eagle_quant_block_sparse_attention_def.cpp
23+)
24+ 
25+target_sources(optiling PRIVATE
26+ eagle_quant_block_sparse_attention_tiling.cpp
27+)
28+ 
29+if (NOT BUILD_OPEN_PROJECT)
30+ target_sources(opmaster_ct PRIVATE
31+ eagle_quant_block_sparse_attention_tiling.cpp
32+ )
33+endif ()
34+ 
35+target_sources(opsproto PRIVATE
36+ eagle_quant_block_sparse_attention_infershape.cpp
37+)
38+ 
39+target_sources(opapi PRIVATE
40+ ${CMAKE_CURRENT_SOURCE_DIR}/../op_api/aclnn_eagle_quant_block_sparse_attention.cpp
41+ ${CMAKE_CURRENT_SOURCE_DIR}/../op_api/eagle_quant_block_sparse_attention.cpp
42+)
43+ 
44+target_include_directories(optiling PRIVATE
45+ ${CMAKE_CURRENT_SOURCE_DIR}
46+ ${CMAKE_CURRENT_SOURCE_DIR}/..
47+ ${CMAKE_CURRENT_SOURCE_DIR}/../../utils/inc
48+)
49+ 
50+target_include_directories(opsproto PRIVATE
51+ ${CMAKE_CURRENT_SOURCE_DIR}/..
52+ ${CMAKE_CURRENT_SOURCE_DIR}/../../utils/inc
53+)
54+ 
55+target_include_directories(opapi PRIVATE
56+ ${CMAKE_CURRENT_SOURCE_DIR}/../op_api
57+ ${CMAKE_CURRENT_SOURCE_DIR}/../../utils/inc
58+)
Acsrc/ops/eagle_quant_block_sparse_attention/op_host/OWNERS+59-0
@@ -0,0 +1,59 @@
1+approvers:
2+- andylhy
3+- xiaolong_han
4+- huxiaobang
5+- startzgf168
6+- huangwei791
7+- wangjun-qt
8+- song-jionghui
9+- yang-binrong
10+- jarr0d
11+- jiang-lirui
12+- huangli70
13+- fzzach
14+- zhangtj0209
15+ 
16+reviewers:
17+- li-xulong
18+- Allan_Yu
19+- mabing726
20+- miao-fangzheng
21+- chengsheng304064
22+- wangfei6
23+- shawn-hu
24+- li-shengxian
25+- xig514
26+- hilihli
27+- chen-vvjob
28+- renkyk
29+- yang-binrong
30+- song-jionghui
31+- wang-zhe123456789
32+- huangli70
33+- huangwei791
34+- shunqi
35+- wangjun-qt
36+- jiang-lirui
37+- realmadrid1016
38+- monologue815
39+- zhangtj0209
40+- GodantShen
41+- Liexss
42+- fzzach
43+- zhanglei_hw
44+- songkai-huawei
45+- yangxu19921206
46+- crystalhu
47+- yu-xinjie62
48+- zhu-yijun-Julius
49+ 
50+files:
51+ "aclnn_eagle_quant_block_sparse_attention.h":
52+ approvers:
53+ - yinqiande
54+ - miao-fangzheng
55+ "eagle_quant_block_sparse_attention_def.cpp":
56+ approvers:
57+ - yinqiande
58+ - miao-fangzheng
59+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_host/eagle_quant_block_sparse_attention_def.cpp+110-0
@@ -0,0 +1,110 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#include "register/op_def_registry.h"
12+ 
13+namespace ops {
14+ 
15+class EagleQuantBlockSparseAttention : public OpDef {
16+public:
17+ explicit EagleQuantBlockSparseAttention(const char* name) : OpDef(name)
18+ {
19+ this->Input("query")
20+ .ParamType(REQUIRED)
21+ .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8,
22+ ge::DT_FLOAT16, ge::DT_BF16, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8})
23+ .FormatList({ge::FORMAT_ND});
24+ this->Input("key")
25+ .ParamType(REQUIRED)
26+ .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8,
27+ ge::DT_FLOAT16, ge::DT_BF16, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8})
28+ .FormatList({ge::FORMAT_ND});
29+ this->Input("value")
30+ .ParamType(REQUIRED)
31+ .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT8_E4M3FN, ge::DT_INT8,
32+ ge::DT_FLOAT8_E4M3FN, ge::DT_INT8, ge::DT_FLOAT16, ge::DT_BF16,
33+ ge::DT_FLOAT8_E4M3FN, ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_INT8, ge::DT_FLOAT8_E4M3FN, ge::DT_FLOAT8_E4M3FN})
34+ .FormatList({ge::FORMAT_ND});
35+ this->Input("blockSparseMask")
36+ .ParamType(OPTIONAL)
37+ .DataType({ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8, ge::DT_INT8,
38+ ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_BOOL, ge::DT_BOOL})
39+ .FormatList({ge::FORMAT_ND});
40+ this->Input("attenMask")
41+ .ParamType(OPTIONAL)
42+ .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16,
43+ ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_FLOAT16})
44+ .FormatList({ge::FORMAT_ND});
45+ this->Input("blockShape")
46+ .ParamType(OPTIONAL)
47+ .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64,
48+ ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
49+ .FormatList({ge::FORMAT_ND});
50+ this->Input("actualSeqLengths")
51+ .ParamType(OPTIONAL)
52+ .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64,
53+ ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
54+ .FormatList({ge::FORMAT_ND});
55+ this->Input("actualSeqLengthsKv")
56+ .ParamType(OPTIONAL)
57+ .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64,
58+ ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
59+ .FormatList({ge::FORMAT_ND});
60+ this->Input("blockTable")
61+ .ParamType(OPTIONAL)
62+ .DataType({ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32,
63+ ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32, ge::DT_INT32})
64+ .FormatList({ge::FORMAT_ND});
65+ this->Input("queryScale")
66+ .ParamType(OPTIONAL)
67+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT,
68+ ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
69+ .FormatList({ge::FORMAT_ND});
70+ this->Input("keyScale")
71+ .ParamType(OPTIONAL)
72+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT,
73+ ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
74+ .FormatList({ge::FORMAT_ND});
75+ this->Input("valueScale")
76+ .ParamType(OPTIONAL)
77+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT,
78+ ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
79+ .FormatList({ge::FORMAT_ND});
80+ this->Output("attentionOut")
81+ .ParamType(REQUIRED)
82+ .DataType({ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_BF16,
83+ ge::DT_FLOAT16, ge::DT_BF16, ge::DT_FLOAT16, ge::DT_FLOAT16, ge::DT_BF16, ge::DT_BF16, ge::DT_BF16, ge::DT_FLOAT16})
84+ .FormatList({ge::FORMAT_ND});
85+ this->Output("softmaxLse")
86+ .ParamType(OPTIONAL)
87+ .DataType({ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT,
88+ ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT, ge::DT_FLOAT})
89+ .FormatList({ge::FORMAT_ND});
90+
91+ this->Attr("qInputLayout").AttrType(OPTIONAL).String("TND");
92+ this->Attr("kvInputLayout").AttrType(OPTIONAL).String("TND");
93+ this->Attr("numKeyValueHeads").AttrType(OPTIONAL).Int(1);
94+ this->Attr("maskType").AttrType(OPTIONAL).Int(0);
95+ this->Attr("scaleValue").AttrType(OPTIONAL).Float(0.0);
96+ this->Attr("innerPrecise").AttrType(OPTIONAL).Int(0);
97+ this->Attr("blockSize").AttrType(OPTIONAL).Int(0);
98+ this->Attr("preTokens").AttrType(OPTIONAL).Int(2147483647);
99+ this->Attr("nextTokens").AttrType(OPTIONAL).Int(2147483647);
100+ this->Attr("softmaxLseFlag").AttrType(OPTIONAL).Int(0);
101+ 
102+ this->AICore().AddConfig("ascend910b");
103+ this->AICore().AddConfig("ascend910_93");
104+ this->AICore().AddConfig("ascend950");
105+ }
106+};
107+ 
108+OP_ADD(EagleQuantBlockSparseAttention);
109+ 
110+} // namespace ops
Acsrc/ops/eagle_quant_block_sparse_attention/op_host/eagle_quant_block_sparse_attention_infershape.cpp+224-0
@@ -0,0 +1,224 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#include <graph/utils/type_utils.h>
12+#include <register/op_impl_registry.h>
13+#include "log/log.h"
14+#include "log/error_code.h"
15+ 
16+using namespace ge;
17+ 
18+namespace ops {
19+ 
20+static constexpr uint32_t QUERY_INDEX = 0;
21+static constexpr uint32_t KEY_INDEX = 1;
22+static constexpr uint32_t VALUE_INDEX = 2;
23+static constexpr uint32_t ATTENTION_OUT_INDEX = 0;
24+static constexpr uint32_t SOFTMAX_LSE_INDEX = 1;
25+ 
26+static constexpr uint32_t ATTR_Q_INPUT_LAYOUT_INDEX = 0;
27+static constexpr uint32_t ATTR_KV_INPUT_LAYOUT_INDEX = 1;
28+static constexpr uint32_t ATTR_NUM_KV_HEADS_INDEX = 2;
29+ 
30+static constexpr uint32_t DIM_BSH = 3;
31+static constexpr uint32_t DIM_TND = 3;
32+static constexpr uint32_t DIM_BNSD = 4;
33+ 
34+static constexpr uint32_t BSH_DIM_B = 0;
35+static constexpr uint32_t BSH_DIM_S = 1;
36+static constexpr uint32_t BSH_DIM_H = 2;
37+ 
38+static constexpr uint32_t TND_DIM_T = 0;
39+static constexpr uint32_t TND_DIM_N = 1;
40+static constexpr uint32_t TND_DIM_D = 2;
41+static constexpr uint32_t TND_DIM_NUM = 3;
42+ 
43+static constexpr uint32_t BNSD_DIM_B = 0;
44+static constexpr uint32_t BNSD_DIM_N = 1;
45+static constexpr uint32_t BNSD_DIM_S = 2;
46+static constexpr uint32_t BNSD_DIM_D = 3;
47+static constexpr uint32_t BNSD_DIM_NUM = 4;
48+static constexpr uint32_t LSE_DIM_D = 1;
49+ 
50+static constexpr int32_t UNKNOWN_DIMS = -2;
51+ 
52+static ge::graphStatus InferShapeEagleQuantBlockSparseAttention(gert::InferShapeContext *context)
53+{
54+ if (context == nullptr) {
55+ OP_LOGE("EagleQuantBlockSparseAttention", "context is nullptr!");
56+ return ge::GRAPH_FAILED;
57+ }
58+
59+ OP_LOGD(context->GetNodeName(), "Enter EagleQuantBlockSparseAttention InferShape impl.");
60+
61+ // 获取Query shape
62+ const gert::Shape *queryShape = context->GetInputShape(QUERY_INDEX);
63+ OP_CHECK_NULL_WITH_CONTEXT(context, queryShape);
64+
65+ // 获取Key shape
66+ const gert::Shape *keyShape = context->GetInputShape(KEY_INDEX);
67+ OP_CHECK_NULL_WITH_CONTEXT(context, keyShape);
68+
69+ // 获取Value shape
70+ const gert::Shape *valueShape = context->GetInputShape(VALUE_INDEX);
71+ OP_CHECK_NULL_WITH_CONTEXT(context, valueShape);
72+
73+ // 获取输出shape
74+ gert::Shape *attentionOutShape = context->GetOutputShape(ATTENTION_OUT_INDEX);
75+ OP_CHECK_NULL_WITH_CONTEXT(context, attentionOutShape);
76+
77+ gert::Shape *softmaxLseShape = context->GetOutputShape(SOFTMAX_LSE_INDEX);
78+ OP_CHECK_NULL_WITH_CONTEXT(context, softmaxLseShape);
79+
80+ // 获取属性
81+ auto attrs = context->GetAttrs();
82+ OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
83+
84+ const char *qInputLayoutPtr = attrs->GetAttrPointer<char>(ATTR_Q_INPUT_LAYOUT_INDEX);
85+ OP_CHECK_NULL_WITH_CONTEXT(context, qInputLayoutPtr);
86+
87+ const char *kvInputLayoutPtr = attrs->GetAttrPointer<char>(ATTR_KV_INPUT_LAYOUT_INDEX);
88+ OP_CHECK_NULL_WITH_CONTEXT(context, kvInputLayoutPtr);
89+
90+ const int64_t *numKvHeadsPtr = attrs->GetInt(ATTR_NUM_KV_HEADS_INDEX);
91+ OP_CHECK_NULL_WITH_CONTEXT(context, numKvHeadsPtr);
92+
93+ // UNKNOWN DIM处理
94+ if ((queryShape->GetDimNum() == 1 && queryShape->GetDim(0) == UNKNOWN_DIMS) ||
95+ (keyShape->GetDimNum() == 1 && keyShape->GetDim(0) == UNKNOWN_DIMS) ||
96+ (valueShape->GetDimNum() == 1 && valueShape->GetDim(0) == UNKNOWN_DIMS)) {
97+ attentionOutShape->SetDimNum(1);
98+ (*attentionOutShape)[0] = UNKNOWN_DIMS;
99+ softmaxLseShape->SetDimNum(1);
100+ (*softmaxLseShape)[0] = UNKNOWN_DIMS;
101+ return ge::GRAPH_SUCCESS;
102+ }
103+
104+ // 设置AttentionOut shape (与Query shape相同)
105+ *attentionOutShape = *queryShape;
106+
107+ // 验证Q layout和KV layout
108+ std::string qLayout(qInputLayoutPtr);
109+ std::string kvLayout(kvInputLayoutPtr);
110+
111+ // 验证Q layout (只支持TND和BNSD)
112+ if (qLayout == "TND") {
113+ if (queryShape->GetDimNum() != DIM_TND) {
114+ OP_LOGE(context->GetNodeName(), "Layout TND, queryDims(%zu) must be 3!", queryShape->GetDimNum());
115+ return ge::GRAPH_FAILED;
116+ }
117+ } else if (qLayout == "BNSD") {
118+ if (queryShape->GetDimNum() != DIM_BNSD) {
119+ OP_LOGE(context->GetNodeName(), "Layout BNSD, queryDims(%zu) must be 4!", queryShape->GetDimNum());
120+ return ge::GRAPH_FAILED;
121+ }
122+ } else {
123+ OP_LOGE(context->GetNodeName(), "Unsupported Q layout: %s. Only TND and BNSD are supported.", qInputLayoutPtr);
124+ return ge::GRAPH_FAILED;
125+ }
126+
127+ // 验证Q和KV格式一致性：如果其中一个是BNSD，另一个也必须是BNSD
128+ bool isQBNSD = (qLayout == "BNSD");
129+ bool isKvBNSD = (kvLayout == "BNSD");
130+
131+ if (isQBNSD != isKvBNSD) {
132+ OP_LOGE(context->GetNodeName(),
133+ "Q and KV layouts must match: if one is BNSD, the other must also be BNSD. "
134+ "Q layout: %s, KV layout: %s", qLayout.c_str(), kvLayout.c_str());
135+ return ge::GRAPH_FAILED;
136+ }
137+
138+ // 验证KV layout
139+ if (kvLayout == "TND") {
140+ if (keyShape->GetDimNum() != DIM_TND || valueShape->GetDimNum() != DIM_TND) {
141+ OP_LOGE(context->GetNodeName(), "Layout TND, KV dims must be 3!");
142+ return ge::GRAPH_FAILED;
143+ }
144+
145+ // TND格式: [T, N, D]
146+ int64_t kvN = keyShape->GetDim(TND_DIM_N);
147+ int64_t kvD = keyShape->GetDim(TND_DIM_D);
148+
149+ if (*numKvHeadsPtr != 0 && kvN != *numKvHeadsPtr) {
150+ OP_LOGE(context->GetNodeName(), "KV heads mismatch in TND format: %ld != %ld", kvN, *numKvHeadsPtr);
151+ return ge::GRAPH_FAILED;
152+ }
153+
154+ if (valueShape->GetDim(TND_DIM_D) != kvD) {
155+ OP_LOGE(context->GetNodeName(), "K and V head dimension mismatch in TND format");
156+ return ge::GRAPH_FAILED;
157+ }
158+ } else if (kvLayout == "BNSD") {
159+ if (keyShape->GetDimNum() != DIM_BNSD || valueShape->GetDimNum() != DIM_BNSD) {
160+ OP_LOGE(context->GetNodeName(), "Layout BNSD, KV dims must be 4!");
161+ return ge::GRAPH_FAILED;
162+ }
163+
164+ // BNSD格式: [B, N, S, D]
165+ int64_t kvB = keyShape->GetDim(BNSD_DIM_B);
166+ int64_t kvN = keyShape->GetDim(BNSD_DIM_N);
167+ int64_t kvD = keyShape->GetDim(BNSD_DIM_D);
168+
169+ if (*numKvHeadsPtr != 0 && kvN != *numKvHeadsPtr) {
170+ OP_LOGE(context->GetNodeName(), "KV heads mismatch in BNSD format: %ld != %ld", kvN, *numKvHeadsPtr);
171+ return ge::GRAPH_FAILED;
172+ }
173+
174+ if (valueShape->GetDim(BNSD_DIM_D) != kvD) {
175+ OP_LOGE(context->GetNodeName(), "K and V head dimension mismatch in BNSD format");
176+ return ge::GRAPH_FAILED;
177+ }
178+ } else {
179+ OP_LOGE(context->GetNodeName(), "Unsupported KV layout: %s. Only TND format is supported.", kvInputLayoutPtr);
180+ return ge::GRAPH_FAILED;
181+ }
182+
183+ // 设置SoftmaxLse shape (如果需要)
184+ // SoftmaxLse shape通常是 [batch, num_heads, q_seqlen, 1] 或类似维度
185+ if (qLayout == "TND") {
186+ // TND格式
187+ softmaxLseShape->SetDimNum(TND_DIM_NUM);
188+ (*softmaxLseShape)[TND_DIM_T] = queryShape->GetDim(TND_DIM_T);
189+ (*softmaxLseShape)[TND_DIM_N] = queryShape->GetDim(TND_DIM_N);
190+ (*softmaxLseShape)[TND_DIM_D] = LSE_DIM_D;
191+ } else if (qLayout == "BNSD") {
192+ // BNSD格式
193+ softmaxLseShape->SetDimNum(BNSD_DIM_NUM);
194+ (*softmaxLseShape)[BNSD_DIM_B] = queryShape->GetDim(BNSD_DIM_B);
195+ (*softmaxLseShape)[BNSD_DIM_N] = queryShape->GetDim(BNSD_DIM_N);
196+ (*softmaxLseShape)[BNSD_DIM_S] = queryShape->GetDim(BNSD_DIM_S);
197+ (*softmaxLseShape)[BNSD_DIM_D] = LSE_DIM_D;
198+ } else {
199+ OP_LOGE(context->GetNodeName(), "Unexpected Q layout in softmaxLse shape calculation: %s", qInputLayoutPtr);
200+ return ge::GRAPH_FAILED;
201+ }
202+
203+ OP_LOGD(context->GetNodeName(), "EagleQuantBlockSparseAttention InferShape success.");
204+ return ge::GRAPH_SUCCESS;
205+}
206+ 
207+ge::graphStatus InferDataTypeEagleQuantBlockSparseAttention(gert::InferDataTypeContext *context)
208+{
209+ if (context == nullptr) {
210+ return ge::GRAPH_FAILED;
211+ }
212+ auto dtype = context->GetInputDataType(0);
213+ context->SetOutputDataType(0, dtype);
214+ context->SetOutputDataType(1, DT_FLOAT);
215+ 
216+ return GRAPH_SUCCESS;
217+}
218+ 
219+IMPL_OP_INFERSHAPE(EagleQuantBlockSparseAttention)
220+ .InferShape(InferShapeEagleQuantBlockSparseAttention)
221+ .InferDataType(InferDataTypeEagleQuantBlockSparseAttention);
222+ 
223+} // namespace ops
224+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_host/eagle_quant_block_sparse_attention_tiling.cpp+1018-0
@@ -0,0 +1,1018 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#include "eagle_quant_block_sparse_attention_tiling.h"
12+#include <cmath>
13+#include <cstring>
14+#include "log/log.h"
15+ 
16+#include <cstdint>
17+#include <string>
18+#include <unordered_map>
19+#include "err/ops_err.h"
20+#include "graph/types.h"
21+#include "graph/tensor.h"
22+#include "tiling/platform/platform_ascendc.h"
23+#include "tiling/tiling_base.h" // provides ASCENDC_EXTERN_C
24+ 
25+using namespace ge;
26+using namespace std;
27+ 
28+constexpr int TND_DIM_T = 0;
29+constexpr int TND_DIM_N = 1;
30+constexpr int TND_DIM_D = 2;
31+constexpr int TND_DIM_NUM = 3;
32+ 
33+constexpr int BNSD_DIM_B = 0;
34+constexpr int BNSD_DIM_N = 1;
35+constexpr int BNSD_DIM_S = 2;
36+constexpr int BNSD_DIM_D = 3;
37+constexpr int BNSD_DIM_NUM = 4;
38+ 
39+constexpr int BSH_DIM_B = 0;
40+constexpr int BSH_DIM_S = 1;
41+constexpr int BSH_DIM_H = 2;
42+ 
43+constexpr int QUERY_INDEX = 0;
44+constexpr int KEY_INDEX = 1;
45+constexpr int VALUE_INDEX = 2;
46+constexpr int BLOCK_SPARSE_MASK_INDEX = 3;
47+constexpr int ATTEN_MASK_INDEX = 4;
48+constexpr int BLOCK_SHAPE_INDEX = 5;
49+constexpr int ACTUAL_SEQ_LENGTHS_INDEX = 6;
50+constexpr int ACTUAL_SEQ_LENGTHS_KV_INDEX = 7;
51+constexpr int BLOCK_TABLE_INDEX = 8;
52+constexpr int SOFTMAX_LSE_INDEX = 10;
53+constexpr int MAX_BLOCK_NUM_INDEX = 2;
54+ 
55+ 
56+constexpr int Q_INPUT_LAYOUT_INDEX = 0;
57+constexpr int KV_INPUT_LAYOUT_INDEX = 1;
58+constexpr int NUM_KEY_VALUE_HEADS_INDEX = 2;
59+constexpr int MASK_TYPE_INDEX = 3;
60+constexpr int SCALE_VALUE_INDEX = 4;
61+constexpr int INNER_PRECISE_INDEX = 5;
62+constexpr int BLOCK_SIZE_INDEX = 6;
63+constexpr int PRE_TOKENS_INDEX = 7;
64+constexpr int NEXT_TOKENS_INDEX = 8;
65+constexpr int SOFTMAX_LSE_FLAG_INDEX = 9;
66+ 
67+constexpr int VALID_EMBEDDING_SIZE_64 = 64;
68+constexpr int VALID_EMBEDDING_SIZE_128 = 128;
69+ 
70+constexpr int LSE_NO_OUT = 0;
71+constexpr int LSE_OUT = 1;
72+ 
73+namespace optiling {
74+ 
75+constexpr uint32_t BASIC_BLOCK_SIZE = 128;
76+constexpr uint32_t WORKSPACE_BLOCK_SIZE_DB = 131072;
77+constexpr uint32_t NUM3 = 3;
78+constexpr uint32_t SOC_VER_950_CODE = 4;
79+constexpr uint32_t INF_WINDOW_SIZE_PRE_NEXT = 2147483647;
80+constexpr uint32_t SPARSE_PATTERN_MODE_MASK = 0;
81+constexpr uint32_t SPARSE_PATTERN_MODE_TABLE = 1;
82+ 
83+constexpr uint32_t TILE_SIZE_128 = 128;
84+constexpr uint32_t TILE_SIZE_256 = 256;
85+constexpr uint32_t TILE_SIZE_512 = 512;
86+constexpr uint32_t D_SIZE_128 = 128;
87+constexpr uint32_t D_SIZE_256 = 256;
88+ 
89+constexpr uint32_t SOLO_BUF = 1;
90+constexpr uint32_t DUO_BUF = 2;
91+constexpr uint32_t TRIO_BUF = 3;
92+ 
93+std::unordered_map<std::string, std::string> inputLayoutMapQ2Kv = {
94+ {"TND", "TND"},
95+ {"BNSD", "BNSD"}
96+};
97+ 
98+static inline uint32_t CeilDiv(uint32_t n1, uint32_t n2)
99+{
100+ if (n1 == 0) {
101+ return 0;
102+ }
103+ return (n2 != 0) ? ((n1 + n2 - 1) / n2) : n1;
104+}
105+ 
106+static inline uint32_t GetQBlocks(int32_t qseqlen, int32_t x)
107+{
108+ uint32_t qBlocksInX = CeilDiv(x, BASIC_BLOCK_SIZE);
109+ uint32_t completeXBlocks = x != 0 ? qseqlen / x : qseqlen / BASIC_BLOCK_SIZE;
110+ uint32_t remainingSeqlen = x != 0 ? qseqlen - completeXBlocks * x : qseqlen % BASIC_BLOCK_SIZE;
111+ uint32_t remainingBlocks = CeilDiv(remainingSeqlen, BASIC_BLOCK_SIZE);
112+ return qBlocksInX * completeXBlocks + remainingBlocks;
113+}
114+ 
115+static inline uint32_t GetQNBlockTile()
116+{
117+ uint32_t qNBlockTile = 1;
118+ return qNBlockTile;
119+}
120+ 
121+ge::graphStatus BSATiling::GetNpuInfo(gert::TilingContext *bsaContext)
122+{
123+ auto ascendcPlatform = platform_ascendc::PlatformAscendC(bsaContext->GetPlatformInfo());
124+
125+ ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize_);
126+ libapiSize_ = ascendcPlatform.GetLibApiWorkSpaceSize();
127+ aivNum_ = ascendcPlatform.GetCoreNumAiv();
128+ aicNum_ = ascendcPlatform.GetCoreNumAic();
129+ socVer_ = static_cast<uint32_t>(ascendcPlatform.GetSocVersion());
130+ return ge::GRAPH_SUCCESS;
131+}
132+ 
133+ge::graphStatus BSATiling::ValidateTNDSeqlenSum(gert::TilingContext *bsaContext)
134+{
135+ // 只在TND格式时进行校验
136+ if (qInputLayout_ != RFAQInputLayout::TND_Q || kvCacheLayout_ != RFAKvCacheLayout::TND) {
137+ return ge::GRAPH_SUCCESS;
138+ }
139+
140+ // 计算所有batch的qseqlen之和
141+ int64_t sumQSeqlen = 0;
142+ int64_t sumKvSeqlen = 0;
143+ 
144+ for (uint32_t i = 0; i < batch_; i++) {
145+ sumQSeqlen += qSeqLenList_[i];
146+ sumKvSeqlen += kvSeqLenList_[i];
147+ }
148+
149+ // 校验qseqlen之和是否等于Q的T
150+ if (sumQSeqlen != totalTokensT_) {
151+ OP_LOGE(bsaContext->GetNodeName(),
152+ "TND format validation failed: sum of qseqlen across all batches (%ld) != Q T (%ld)",
153+ sumQSeqlen, totalTokensT_);
154+ return ge::GRAPH_FAILED;
155+ }
156+
157+ // 校验kvseqlen之和是否等于KV的T
158+ if (sumKvSeqlen != totalTokensKv_) {
159+ OP_LOGE(bsaContext->GetNodeName(),
160+ "TND format validation failed: sum of kvseqlen across all batches (%ld) != KV T (%ld)",
161+ sumKvSeqlen, totalTokensKv_);
162+ return ge::GRAPH_FAILED;
163+ }
164+
165+ return ge::GRAPH_SUCCESS;
166+}
167+ 
168+ge::graphStatus BSATiling::GetInputLayout(gert::TilingContext *bsaContext)
169+{
170+ auto attrs = bsaContext->GetAttrs();
171+ OP_CHECK_NULL_WITH_CONTEXT(bsaContext, attrs);
172+ auto attrLayoutQ = attrs->GetAttrPointer<char>(Q_INPUT_LAYOUT_INDEX);
173+ auto attrLayoutKv = attrs->GetAttrPointer<char>(KV_INPUT_LAYOUT_INDEX);
174+ if (attrLayoutQ == nullptr || attrLayoutKv == nullptr) {
175+ OP_LOGE(bsaContext->GetNodeName(), "qInputLayout, kvInputLayout must be provided.");
176+ return ge::GRAPH_FAILED;
177+ }
178+ std::string qLayout(attrLayoutQ);
179+ std::string kvLayout(attrLayoutKv);
180+ auto it = inputLayoutMapQ2Kv.find(qLayout);
181+ auto itRev = inputLayoutMapQ2Kv.find(kvLayout);
182+ if (it == inputLayoutMapQ2Kv.end() || itRev == inputLayoutMapQ2Kv.end()) {
183+ OP_LOGE(bsaContext->GetNodeName(), "The inputLayout attrs only support TND/BNSD.");
184+ return ge::GRAPH_FAILED;
185+ } else if (it->second != kvLayout) {
186+ OP_LOGE(bsaContext->GetNodeName(), "QInputLayout and kvInputLayout must be the same.");
187+ return ge::GRAPH_FAILED;
188+ } else if (it->first == "TND") {
189+ qInputLayout_ = RFAQInputLayout::TND_Q;
190+ kvCacheLayout_ = RFAKvCacheLayout::TND;
191+ } else if (it->first == "BNSD") {
192+ qInputLayout_ = RFAQInputLayout::BNSD_Q;
193+ kvCacheLayout_ = RFAKvCacheLayout::BNSD;
194+ }
195+ return ge::GRAPH_SUCCESS;
196+}
197+ 
198+ge::graphStatus BSATiling::CheckQKVDtype(gert::TilingContext *bsaContext)
199+{
200+ auto qInputDesc = bsaContext->GetInputDesc(QUERY_INDEX);
201+ auto kInputDesc = bsaContext->GetInputDesc(KEY_INDEX);
202+ auto vInputDesc = bsaContext->GetInputDesc(VALUE_INDEX);
203+ OP_CHECK_NULL_WITH_CONTEXT(bsaContext, qInputDesc);
204+ OP_CHECK_NULL_WITH_CONTEXT(bsaContext, kInputDesc);
205+ OP_CHECK_NULL_WITH_CONTEXT(bsaContext, vInputDesc);
206+ dataType_ = qInputDesc->GetDataType();
207+ auto kDataType = kInputDesc->GetDataType();
208+ vDataType_ = vInputDesc->GetDataType();
209+ 
210+ // 原有FP16/BF16校验保留
211+ if (dataType_ != ge::DT_FLOAT16 && dataType_ != ge::DT_BF16 && dataType_ != ge::DT_INT8) {
212+ OP_LOGE(bsaContext->GetNodeName(),
213+ "The supported dtype of query/key is float16, bfloat16 or int8.");
214+ return ge::GRAPH_FAILED;
215+ }
216+
217+ // 950 混合精度特判：Q/K=int8, V=FLOAT8_E4M3FN
218+ if (socVer_ == SOC_VER_950_CODE &&
219+ dataType_ == ge::DT_INT8 &&
220+ kDataType == ge::DT_INT8 ) {
221+ return ge::GRAPH_SUCCESS;
222+ } else if (dataType_ != kDataType || dataType_ != vDataType_) {
223+ // 非混合精度场景保持原有严格一致性校验
224+ OP_LOGE(bsaContext->GetNodeName(),
225+ "Tensor query/key/value must have consistent dtype, or use 950 mixed prec (Q/K=int8, V=fp8).");
226+ return ge::GRAPH_FAILED;
227+ }
228+ 
229+ return ge::GRAPH_SUCCESS;
230+}
231+ 
232+ge::graphStatus BSATiling::CheckQKVDimVal(
233+ gert::TilingContext *bsaContext, uint32_t kHeads, uint32_t vHeads, uint32_t kHeadDim, uint32_t vHeadDim)
234+{
235+ if (embeddingSize_ != kHeadDim || embeddingSize_ != vHeadDim) {
236+ OP_LOGE(bsaContext->GetNodeName(), "Tensor query/key/value must have consistent headDim with each other, "
237+ "but got qHeadDim %u, kHeadDim %u, vHeadDim %u.", embeddingSize_, kHeadDim, vHeadDim);
238+ return ge::GRAPH_FAILED;
239+ }
240+ if (kvHeads_ != kHeads || kvHeads_ != vHeads) {
241+ OP_LOGE(bsaContext->GetNodeName(),
242+ "Tensor key/value must have consistent headNum with each other and attr kvHeads, "
243+ "but got kHeads %u, vHeads %u, kvHeads(attr) %u.", kHeads, vHeads, kvHeads_);
244+ return ge::GRAPH_FAILED;
245+ }
246+ // temporary regulations of D
247+ if ((embeddingSize_ != VALID_EMBEDDING_SIZE_128) && (embeddingSize_ != VALID_EMBEDDING_SIZE_64)) {
248+ OP_LOGE(bsaContext->GetNodeName(),
249+ "The supported headDim so far is 64 or 128, but got %u.", embeddingSize_);
250+ return ge::GRAPH_FAILED;
251+ }
252+ return ge::GRAPH_SUCCESS;
253+}
254+ 
255+ge::graphStatus BSATiling::ParseQKVInTND(gert::TilingContext *bsaContext)
256+{
257+ const auto *queryShape = bsaContext->GetInputShape(QUERY_INDEX);
258+ const auto *keyShape = bsaContext->GetInputShape(KEY_INDEX);
259+ const auto *valueShape = bsaContext->GetInputShape(VALUE_INDEX);
260+ if (queryShape == nullptr || keyShape == nullptr || valueShape == nullptr) {
261+ OP_LOGE(bsaContext->GetNodeName(), "Query/Key/Value shape is null");
262+ return ge::GRAPH_FAILED;
263+ }
264+ if (queryShape->GetStorageShape().GetDimNum() != TND_DIM_NUM ||
265+ keyShape->GetStorageShape().GetDimNum() != TND_DIM_NUM ||
266+ valueShape->GetStorageShape().GetDimNum() != TND_DIM_NUM) {
267+ OP_LOGE(bsaContext->GetNodeName(), "Tensor query/key/value must have 3 dimensions when layout is 'TND'.");
268+ return ge::GRAPH_FAILED;
269+ }
270+ totalTokensT_ = queryShape->GetStorageShape().GetDim(TND_DIM_T);
271+ totalTokensKv_ = keyShape->GetStorageShape().GetDim(TND_DIM_T);
272+ numHeads_ = static_cast<uint32_t>(queryShape->GetStorageShape().GetDim(TND_DIM_N));
273+ auto kHeads = static_cast<uint32_t>(keyShape->GetStorageShape().GetDim(TND_DIM_N));
274+ auto vHeads = static_cast<uint32_t>(valueShape->GetStorageShape().GetDim(TND_DIM_N));
275+ embeddingSize_ = static_cast<uint32_t>(queryShape->GetStorageShape().GetDim(TND_DIM_D));
276+ auto embeddingSizeK = static_cast<uint32_t>(keyShape->GetStorageShape().GetDim(TND_DIM_D));
277+ auto embeddingSizeV = static_cast<uint32_t>(valueShape->GetStorageShape().GetDim(TND_DIM_D));
278+ if (CheckQKVDimVal(bsaContext, kHeads, vHeads, embeddingSizeK, embeddingSizeV) != ge::GRAPH_SUCCESS) {
279+ OP_LOGE(bsaContext->GetNodeName(), "Check query/key/value dim values failed.");
280+ return ge::GRAPH_FAILED;
281+ }
282+ return ge::GRAPH_SUCCESS;
283+}
284+ 
285+ge::graphStatus BSATiling::ParseQKVInBNSD(gert::TilingContext *bsaContext)
286+{
287+ const auto *queryShape = bsaContext->GetInputShape(QUERY_INDEX);
288+ const auto *keyShape = bsaContext->GetInputShape(KEY_INDEX);
289+ const auto *valueShape = bsaContext->GetInputShape(VALUE_INDEX);
290+ if (queryShape == nullptr || keyShape == nullptr || valueShape == nullptr) {
291+ OP_LOGE(bsaContext->GetNodeName(), "Query/Key/Value shape is null");
292+ return ge::GRAPH_FAILED;
293+ }
294+ if (queryShape->GetStorageShape().GetDimNum() != BNSD_DIM_NUM ||
295+ keyShape->GetStorageShape().GetDimNum() != BNSD_DIM_NUM ||
296+ valueShape->GetStorageShape().GetDimNum() != BNSD_DIM_NUM) {
297+ OP_LOGE(bsaContext->GetNodeName(), "Tensor query/key/value must have 4 dimensions when layout is 'BNSD'.");
298+ return ge::GRAPH_FAILED;
299+ }
300+ batch_ = static_cast<uint32_t>(queryShape->GetStorageShape().GetDim(BNSD_DIM_B));
301+ auto kBatch = static_cast<uint32_t>(keyShape->GetStorageShape().GetDim(BNSD_DIM_B));
302+ auto vBatch = static_cast<uint32_t>(valueShape->GetStorageShape().GetDim(BNSD_DIM_B));
303+ numHeads_ = static_cast<uint32_t>(queryShape->GetStorageShape().GetDim(BNSD_DIM_N));
304+ auto kHeads = static_cast<uint32_t>(keyShape->GetStorageShape().GetDim(BNSD_DIM_N));
305+ auto vHeads = static_cast<uint32_t>(valueShape->GetStorageShape().GetDim(BNSD_DIM_N));
306+ embeddingSize_ = static_cast<uint32_t>(queryShape->GetStorageShape().GetDim(BNSD_DIM_D));
307+ auto embeddingSizeK = static_cast<uint32_t>(keyShape->GetStorageShape().GetDim(BNSD_DIM_D));
308+ auto embeddingSizeV = static_cast<uint32_t>(valueShape->GetStorageShape().GetDim(BNSD_DIM_D));
309+ maxQSeqlen_ = static_cast<uint32_t>(queryShape->GetStorageShape().GetDim(BNSD_DIM_S));
310+ maxKvSeqlen_ = static_cast<uint32_t>(keyShape->GetStorageShape().GetDim(BNSD_DIM_S));
311+ if (batch_ != kBatch || batch_ != vBatch) {
312+ OP_LOGE(bsaContext->GetNodeName(), "Tensor query/key/value must have consistent batch with each other, "
313+ "but got qBatch %u, kBatch %u, vBatch %u.", batch_, kBatch, vBatch);
314+ return ge::GRAPH_FAILED;
315+ }
316+ if (CheckQKVDimVal(bsaContext, kHeads, vHeads, embeddingSizeK, embeddingSizeV) != ge::GRAPH_SUCCESS) {
317+ OP_LOGE(bsaContext->GetNodeName(), "Check query/key/value dim values failed.");
318+ return ge::GRAPH_FAILED;
319+ }
320+ return ge::GRAPH_SUCCESS;
321+}
322+ 
323+ge::graphStatus BSATiling::ParseRequiredTensors(gert::TilingContext *bsaContext)
324+{
325+ if (CheckQKVDtype(bsaContext) != ge::GRAPH_SUCCESS) {
326+ return ge::GRAPH_FAILED;
327+ }
328+ ge::graphStatus ret = ge::GRAPH_SUCCESS;
329+ if (qInputLayout_ == RFAQInputLayout::TND_Q) {
330+ ret = ParseQKVInTND(bsaContext);
331+ } else if (qInputLayout_ == RFAQInputLayout::BNSD_Q) {
332+ ret = ParseQKVInBNSD(bsaContext);
333+ }
334+ return ret;
335+}
336+ 
337+ge::graphStatus BSATiling::ParseSeqlensInTND(gert::TilingContext *bsaContext)
338+{
339+ const auto *actualSeqLengths = bsaContext->GetOptionalInputTensor(ACTUAL_SEQ_LENGTHS_INDEX);
340+ const auto *actualSeqLengthsKv = bsaContext->GetOptionalInputTensor(ACTUAL_SEQ_LENGTHS_KV_INDEX);
341+ if (actualSeqLengths == nullptr || actualSeqLengthsKv == nullptr) {
342+ OP_LOGE(bsaContext->GetNodeName(),
343+ "ActualSeqLengths/actualSeqLengthsKv must be provided when corresponding layout is 'TND'.");
344+ return ge::GRAPH_FAILED;
345+ }
346+ batch_ = static_cast<uint32_t>(actualSeqLengths->GetShapeSize());
347+ auto batchKvS = static_cast<uint32_t>(actualSeqLengthsKv->GetShapeSize());
348+ if (batch_ != batchKvS) {
349+ OP_LOGE(bsaContext->GetNodeName(),
350+ "ActualSeqLengths & actualSeqLengthsKv must have consistent batch size, "
351+ "but got batch in actualSeqLengths: %u, batch in actualSeqLengthsKv: %u", batch_, batchKvS);
352+ return ge::GRAPH_FAILED;
353+ }
354+ qSeqLenList_ = actualSeqLengths->GetData<int64_t>();
355+ kvSeqLenList_ = actualSeqLengthsKv->GetData<int64_t>();
356+ useUniformQSeqlen_ = false;
357+ useUniformKvSeqlen_ = false;
358+ return ge::GRAPH_SUCCESS;
359+}
360+ 
361+ge::graphStatus BSATiling::ParseSeqlensInBNSD(gert::TilingContext *bsaContext)
362+{
363+ const auto *actualSeqLengths = bsaContext->GetOptionalInputTensor(ACTUAL_SEQ_LENGTHS_INDEX);
364+ const auto *actualSeqLengthsKv = bsaContext->GetOptionalInputTensor(ACTUAL_SEQ_LENGTHS_KV_INDEX);
365+ if (actualSeqLengths != nullptr && actualSeqLengthsKv != nullptr) {
366+ uint32_t batchQS = static_cast<uint32_t>(actualSeqLengths->GetShapeSize());
367+ uint32_t batchKvS = static_cast<uint32_t>(actualSeqLengthsKv->GetShapeSize());
368+ if (batch_ != batchKvS || batch_ != batchQS) {
369+ OP_LOGE(bsaContext->GetNodeName(),
370+ "ActualSeqLengths & actualSeqLengthsKv must have consistent batch size with each other and context,"
371+ "but got batch in actualSeqLengths: %u, batch in actualSeqLengthsKv: %u, "
372+ "batch in context(from query dim0): %u",
373+ batchQS, batchKvS, batch_);
374+ return ge::GRAPH_FAILED;
375+ }
376+ qSeqLenList_ = actualSeqLengths->GetData<int64_t>();
377+ kvSeqLenList_ = actualSeqLengthsKv->GetData<int64_t>();
378+ useUniformQSeqlen_ = false;
379+ useUniformKvSeqlen_ = false;
380+ } else if (actualSeqLengths == nullptr && actualSeqLengthsKv == nullptr) {
381+ useUniformQSeqlen_ = true;
382+ useUniformKvSeqlen_ = true;
383+ } else {
384+ OP_LOGE(bsaContext->GetNodeName(),
385+ "ActualSeqLengths & actualSeqLengthsKv must be either both provided, or neither provided.");
386+ return ge::GRAPH_FAILED;
387+ }
388+ return ge::GRAPH_SUCCESS;
389+}
390+ 
391+ge::graphStatus BSATiling::ParseSeqlens(gert::TilingContext *bsaContext)
392+{
393+ ge::graphStatus ret = ge::GRAPH_SUCCESS;
394+ if (qInputLayout_ == RFAQInputLayout::TND_Q) {
395+ ret = ParseSeqlensInTND(bsaContext);
396+ } else if (qInputLayout_ == RFAQInputLayout::BNSD_Q) {
397+ ret = ParseSeqlensInBNSD(bsaContext);
398+ }
399+ return ret;
400+}
401+ 
402+ge::graphStatus BSATiling::CheckSparsePattern(gert::TilingContext *bsaContext, const int64_t defaultShape)
403+{
404+ const auto *blockSparseMaskDesc = bsaContext->GetInputDesc(BLOCK_SPARSE_MASK_INDEX);
405+ const auto *blockSparseMaskShape = bsaContext->GetInputShape(BLOCK_SPARSE_MASK_INDEX);
406+ OP_CHECK_NULL_WITH_CONTEXT(bsaContext, blockSparseMaskDesc);
407+ OP_CHECK_NULL_WITH_CONTEXT(bsaContext, blockSparseMaskShape);
408+ auto sparsePatternDtype = blockSparseMaskDesc->GetDataType();
409+ if (sparsePatternDtype == ge::DT_INT8 || sparsePatternDtype == ge::DT_BOOL ) {
410+ sparsePatternMode_ = SPARSE_PATTERN_MODE_MASK;
411+ } else if (sparsePatternDtype == ge::DT_INT32) {
412+ sparsePatternMode_ = SPARSE_PATTERN_MODE_TABLE;
413+ } else {
414+ OP_LOGE(bsaContext->GetNodeName(),
415+ "BlockSparseMask only supports int8 mask or int32 sparse table, but got dtype %d.",
416+ static_cast<int32_t>(sparsePatternDtype));
417+ return ge::GRAPH_FAILED;
418+ }
419+ if (blockShapeX_ <= 0 || blockShapeY_ <= 0) {
420+ OP_LOGE(bsaContext->GetNodeName(), "BlockShape elems must be greater than 0, "
421+ "but got elem0: %ld, elem1: %ld.", blockShapeX_, blockShapeY_);
422+ return ge::GRAPH_FAILED;
423+ }
424+ // temporary regulation of blockShapeY
425+ if (blockShapeY_ % defaultShape != 0) {
426+ OP_LOGE(bsaContext->GetNodeName(), "BlockShape elem1 must be a multiple of 128 so far, "
427+ "but got elem1: %ld.", blockShapeY_);
428+ return ge::GRAPH_FAILED;
429+ }
430+ if (blockSparseMaskShape->GetStorageShape().GetDimNum() != 4) {
431+ printf("Current DimNum%d\n",blockSparseMaskShape->GetStorageShape().GetDimNum());
432+ //OP_LOGE(bsaContext->GetNodeName(), "BlockSparseMask must have 4 dims.");
433+ //return ge::GRAPH_FAILED;
434+ }
435+ uint32_t bsmBatch = blockSparseMaskShape->GetStorageShape().GetDim(0); // batch_size
436+ uint32_t bsmNumHead = blockSparseMaskShape->GetStorageShape().GetDim(1); // num_heads
437+ if (bsmBatch != batch_) {
438+ OP_LOGE(bsaContext->GetNodeName(), "BlockSparseMask must have consistent batch with context,"
439+ "but got BlockSparseMask batch(dim0): %u, context batch: %u.", bsmBatch, batch_);
440+ return ge::GRAPH_FAILED;
441+ }
442+ if (bsmNumHead != numHeads_) {
443+ OP_LOGE(bsaContext->GetNodeName(), "BlockSparseMask must have consistent numHeads with context,"
444+ "but got BlockSparseMask numHeads(dim1): %u, context numHeads: %u.", bsmNumHead, numHeads_);
445+ return ge::GRAPH_FAILED;
446+ }
447+ return ge::GRAPH_SUCCESS;
448+}
449+ 
450+ge::graphStatus BSATiling::ParseSparsePattern(gert::TilingContext *bsaContext)
451+{
452+ constexpr int64_t DEFAULT_BLOCK_SHAPE = 128;
453+ blockShapeX_ = DEFAULT_BLOCK_SHAPE;
454+ blockShapeY_ = DEFAULT_BLOCK_SHAPE;
455+ const auto *blockSparseMaskTensor = bsaContext->GetOptionalInputTensor(BLOCK_SPARSE_MASK_INDEX);
456+ const auto *blockShapeTensor = bsaContext->GetOptionalInputTensor(BLOCK_SHAPE_INDEX);
457+
458+ if (blockSparseMaskTensor == nullptr) {
459+ OP_LOGE(bsaContext->GetNodeName(), "BlockSparseMask should be provided so far.");
460+ return ge::GRAPH_FAILED;
461+ }
462+ if (blockShapeTensor != nullptr) {
463+ uint32_t blockShapeElemNum = static_cast<uint32_t>(blockShapeTensor->GetShapeSize());
464+ if (blockShapeElemNum != 2) {
465+ OP_LOGE(bsaContext->GetNodeName(), "BlockShape elem num must be 2.");
466+ return ge::GRAPH_FAILED;
467+ }
468+ blockShapeList = blockShapeTensor->GetData<int64_t>();
469+ if (blockShapeList != nullptr) {
470+ blockShapeX_ = blockShapeList[0];
471+ blockShapeY_ = blockShapeList[1];
472+ }
473+ }
474+ if (CheckSparsePattern(bsaContext, DEFAULT_BLOCK_SHAPE) != ge::GRAPH_SUCCESS) {
475+ return ge::GRAPH_FAILED;
476+ }
477+ return ge::GRAPH_SUCCESS;
478+}
479+ 
480+ge::graphStatus BSATiling::ParseAttenMask(gert::TilingContext *bsaContext)
481+{
482+ const auto *attenMaskTensor = bsaContext->GetOptionalInputTensor(ATTEN_MASK_INDEX);
483+ if (attenMaskTensor != nullptr) {
484+ OP_LOGE(bsaContext->GetNodeName(), "AttenMask is NOT YET supported.");
485+ return ge::GRAPH_FAILED;
486+ }
487+ return ge::GRAPH_SUCCESS;
488+}
489+ 
490+ge::graphStatus BSATiling::ParseBlockTable(gert::TilingContext *bsaContext)
491+{
492+ const auto *blockTableTensor = bsaContext->GetOptionalInputTensor(BLOCK_TABLE_INDEX);
493+ if (blockTableTensor != nullptr) {
494+ OP_LOGE(bsaContext->GetNodeName(),
495+ "Paged cache is NOT YET supported, therefore blockTable should be nullptr.");
496+ return ge::GRAPH_FAILED;
497+ }
498+ return ge::GRAPH_SUCCESS;
499+}
500+ 
501+ge::graphStatus BSATiling::ParseOptionalTensors(gert::TilingContext *bsaContext)
502+{
503+ if (ParseSeqlens(bsaContext) != ge::GRAPH_SUCCESS ||
504+ ParseSparsePattern(bsaContext) != ge::GRAPH_SUCCESS ||
505+ ParseAttenMask(bsaContext) != ge::GRAPH_SUCCESS ||
506+ ParseBlockTable(bsaContext) != ge::GRAPH_SUCCESS) {
507+ return ge::GRAPH_FAILED;
508+ }
509+ return ge::GRAPH_SUCCESS;
510+}
511+ 
512+ge::graphStatus BSATiling::ParseAttrs(gert::TilingContext *bsaContext)
513+{
514+ auto attrs = bsaContext->GetAttrs();
515+ OP_CHECK_NULL_WITH_CONTEXT(bsaContext, attrs);
516+ if (attrs->GetAttrPointer<uint32_t>(NUM_KEY_VALUE_HEADS_INDEX) == nullptr) {
517+ OP_LOGE(bsaContext->GetNodeName(), "numKeyValueHeads is null");
518+ return ge::GRAPH_FAILED;
519+ }
520+ kvHeads_ = *attrs->GetAttrPointer<uint32_t>(NUM_KEY_VALUE_HEADS_INDEX);
521+
522+ if (attrs->GetAttrPointer<float>(SCALE_VALUE_INDEX) == nullptr) {
523+ scaleValue_ = 1.0f / std::sqrt(static_cast<float>(embeddingSize_));
524+ } else {
525+ scaleValue_ = *attrs->GetAttrPointer<float>(SCALE_VALUE_INDEX);
526+ }
527+
528+ if (attrs->GetAttrPointer<uint32_t>(MASK_TYPE_INDEX) != nullptr) {
529+ maskType_ = *attrs->GetAttrPointer<uint32_t>(MASK_TYPE_INDEX);
530+ }
531+
532+ // 获取innerPrecise参数
533+ if (attrs->GetAttrPointer<uint32_t>(INNER_PRECISE_INDEX) != nullptr) {
534+ innerPrecise_ = *attrs->GetAttrPointer<uint32_t>(INNER_PRECISE_INDEX);
535+ }
536+ if (socVer_ == SOC_VER_950_CODE) {
537+ if (innerPrecise_ != BsaInnerCalcPrec::LOW_HIGH_MIXED) {
538+ OP_LOGE(bsaContext->GetNodeName(), "On chip 950, only innerPrec = 4 is supported, "
539+ "but got %u.", innerPrecise_);
540+ return ge::GRAPH_FAILED;
541+ }
542+ } else {
543+ auto qInputDesc = bsaContext->GetInputDesc(QUERY_INDEX);
544+ OP_CHECK_NULL_WITH_CONTEXT(bsaContext, qInputDesc);
545+ auto dtypeQ = qInputDesc->GetDataType();
546+ if (innerPrecise_ != BsaInnerCalcPrec::ALL_HIGH && innerPrecise_ != BsaInnerCalcPrec::ALL_LOW) {
547+ OP_LOGE(bsaContext->GetNodeName(), "On chip 910 & 910_93, only innerPrec = 0 or 1 is supported, "
548+ "but got %u.", innerPrecise_);
549+ return ge::GRAPH_FAILED;
550+ } else if (innerPrecise_ == BsaInnerCalcPrec::ALL_LOW && dtypeQ == ge::DT_BF16) {
551+ OP_LOGE(bsaContext->GetNodeName(), "On chip 910 & 910_93, when query dtype is bfloat16, "
552+ "only innerPrec = 0 is supported, but got %u.", innerPrecise_);
553+ return ge::GRAPH_FAILED;
554+ }
555+ }
556+ // reserved yet non-configurable attrs
557+ int64_t blockSize = *attrs->GetAttrPointer<int64_t>(BLOCK_SIZE_INDEX);
558+ if (blockSize != 0) {
559+ OP_LOGE(bsaContext->GetNodeName(), "Since paged cache is not yet supported, "
560+ "blocksize must be 0, but got %ld.", blockSize);
561+ return ge::GRAPH_FAILED;
562+ }
563+ int64_t preTokens = *attrs->GetAttrPointer<int64_t>(PRE_TOKENS_INDEX);
564+ int64_t nextTokens = *attrs->GetAttrPointer<int64_t>(NEXT_TOKENS_INDEX);
565+ if (preTokens != INF_WINDOW_SIZE_PRE_NEXT || nextTokens != INF_WINDOW_SIZE_PRE_NEXT) {
566+ OP_LOGE(bsaContext->GetNodeName(), "Since windowed atten mask is not yet supported, "
567+ "preTokens & nextTokens must be 2147483647, but got %ld, %ld.", preTokens, nextTokens);
568+ return ge::GRAPH_FAILED;
569+ }
570+ auto softmaxLsePtr = attrs->GetAttrPointer<int64_t>(SOFTMAX_LSE_FLAG_INDEX);
571+ if (softmaxLsePtr == nullptr) {
572+ OP_LOGE(bsaContext->GetNodeName(), "Attr softmaxLseFlag is nullptr.");
573+ return ge::GRAPH_FAILED;
574+ } else if (*softmaxLsePtr == LSE_OUT) {
575+ if (socVer_ == SOC_VER_950_CODE) {
576+ OP_LOGE(bsaContext->GetNodeName(), "Attr softmaxLseFlag must be 0 on chip 950.");
577+ return ge::GRAPH_FAILED;
578+ }
579+ softmaxLseFlag_ = true;
580+ } else if (*softmaxLsePtr == LSE_NO_OUT) {
581+ softmaxLseFlag_ = false;
582+ } else {
583+ OP_LOGE(bsaContext->GetNodeName(), "Attr softmaxLseFlag must be 0 or 1, but got: %ld.", *softmaxLsePtr);
584+ return ge::GRAPH_FAILED;
585+ }
586+
587+ return ge::GRAPH_SUCCESS;
588+}
589+ 
590+void BSATiling::CalculateBatchTaskSplit(int64_t qSeqlen, uint32_t groupSize,
591+ uint32_t &curTaskNum, uint32_t &curQBlockNum)
592+{
593+ uint32_t curQBlockTile = GetQNBlockTile();
594+ uint32_t qNBlockNumPerGroup = CeilDiv(groupSize, curQBlockTile);
595+ uint32_t curQNBlockNum = qNBlockNumPerGroup * kvHeads_;
596+ curTaskNum = GetQBlocks(qSeqlen, blockShapeX_) * curQNBlockNum;
597+ curQBlockNum = CeilDiv(qSeqlen, blockShapeX_) * numHeads_;
598+}
599+ 
600+uint32_t BSATiling::GetCurQSTileNum950(int64_t curQSeqlen)
601+{
602+ uint32_t fullXBlockNum = curQSeqlen / blockShapeX_;
603+ uint32_t tailXBlockSize = curQSeqlen % blockShapeX_;
604+ uint32_t qSTileNumPerFullXBlock = (blockShapeX_ + qBaseTile_ - 1) / qBaseTile_;
605+ uint32_t qSTileNumTailXBlock = (tailXBlockSize + qBaseTile_ - 1) / qBaseTile_;
606+ uint32_t curQSTileNum = qSTileNumPerFullXBlock * fullXBlockNum + qSTileNumTailXBlock;
607+ return curQSTileNum;
608+}
609+ 
610+void BSATiling::CalcBaseTileTilingParams950()
611+{
612+ qBaseTile_ = (blockShapeX_ > TILE_SIZE_128) ? TILE_SIZE_128 : blockShapeX_;
613+ if (innerPrecise_ == BsaInnerCalcPrec::LOW_HIGH_MIXED && embeddingSize_ <= D_SIZE_128) {
614+ kvBaseTile_ = TILE_SIZE_512;
615+ } else {
616+ kvBaseTile_ = TILE_SIZE_128;
617+ }
618+}
619+ 
620+void BSATiling::CalcSplitCoreTilingParams950()
621+{
622+ CalcBaseTileTilingParams950();
623+ for (uint32_t bIdx = 0; bIdx < batch_; bIdx++) {
624+ int64_t curQSeqlen = useUniformQSeqlen_ ? maxQSeqlen_ : qSeqLenList_[bIdx];
625+ int64_t curKvSeqlen = useUniformKvSeqlen_ ? maxKvSeqlen_ : kvSeqLenList_[bIdx];
626+ uint32_t curQSTileNum = GetCurQSTileNum950(curQSeqlen);
627+ uint32_t curBatchTaskNum = curQSTileNum * numHeads_;
628+ totalTaskNum_ += curBatchTaskNum;
629+ if (bIdx == 0) {
630+ firstBatchTaskNum_ = curBatchTaskNum;
631+ }
632+ maxQSeqlen_ = (curQSeqlen > maxQSeqlen_) ? curQSeqlen : maxQSeqlen_;
633+ maxKvSeqlen_ = (curKvSeqlen > maxKvSeqlen_) ? curKvSeqlen : maxKvSeqlen_;
634+ }
635+ blockDim_ = aicNum_;
636+ // mask2idx split core
637+ xBlockNumAligned_ = (maxQSeqlen_ + blockShapeX_ - 1) / blockShapeX_;
638+ yBlockNumAligned_ = (maxKvSeqlen_ + blockShapeY_ - 1) / blockShapeY_;
639+ uint32_t totalRowNumBlockMask = batch_ * numHeads_ * xBlockNumAligned_;
640+ avgRowPerSubCore_ = (totalRowNumBlockMask + aivNum_ - 1) / aivNum_;
641+ preActiveSubCoreNum_ = (totalRowNumBlockMask + avgRowPerSubCore_ - 1) / avgRowPerSubCore_;
642+}
643+ 
644+ge::graphStatus BSATiling::CalculateTaskSplit(gert::TilingContext *bsaContext)
645+{
646+ // 计算总的Q块数量和最大KV块数量
647+ totalQBlocks_ = 0;
648+ 
649+ if (kvHeads_ == 0) {
650+ OP_LOGE(bsaContext->GetNodeName(), "kvHeads_ is 0, cannot calculate groupSize");
651+ return ge::GRAPH_FAILED;
652+ }
653+ if (batch_ == 0) {
654+ OP_LOGE(bsaContext->GetNodeName(), "batch_ is 0 in CalculateTaskSplit");
655+ return ge::GRAPH_FAILED;
656+ }
657+
658+ // 根据useUniformQSeqlen_标志位决定分核时使用actualSeqLengths数组还是maxQSeqlen_
659+ uint32_t groupSize = numHeads_ / kvHeads_;
660+
661+ // 遍历每个batch进行分核计算
662+ for (auto i = 0; i < batch_; i++) {
663+ // 根据useUniformQSeqlen_标志位决定使用actualSeqLengths数组还是maxQSeqlen_
664+ int64_t qSeqlen;
665+ if (useUniformQSeqlen_) {
666+ // BNSD格式下actualSeqLengths为nullptr，使用maxQSeqlen_作为统一的qseqlen值
667+ qSeqlen = static_cast<int64_t>(maxQSeqlen_);
668+ } else {
669+ // 使用actualSeqLengths数组（TND格式或BNSD格式但提供了actualSeqLengths）
670+ if (qSeqLenList_ == nullptr) {
671+ OP_LOGE(bsaContext->GetNodeName(), "qSeqLenList_ is nullptr, cannot calculate task split");
672+ return ge::GRAPH_FAILED;
673+ }
674+ qSeqlen = qSeqLenList_[i];
675+ }
676+ 
677+ uint32_t curTaskNum = 0;
678+ uint32_t curQBlockNum = 0;
679+ CalculateBatchTaskSplit(qSeqlen, groupSize, curTaskNum, curQBlockNum);
680+
681+ if (i == 0) {
682+ firstBatchTaskNum_ = curTaskNum;
683+ firstQBlockNum_ = curQBlockNum;
684+ }
685+ totalTaskNum_ += curTaskNum;
686+ totalQBlocks_ += curQBlockNum;
687+ }
688+ blockDim_ = std::min(aicNum_, totalTaskNum_);
689+ return ge::GRAPH_SUCCESS;
690+}
691+ 
692+void BSATiling::CalcWorkspaceTilingParams950(gert::TilingContext *bsaContext)
693+{
694+ selectIdxSize_ = (sparsePatternMode_ == SPARSE_PATTERN_MODE_TABLE) ?
695+ 0 : batch_ * numHeads_ * xBlockNumAligned_ * yBlockNumAligned_ * sizeof(int32_t);
696+ selectNumIdxSize_ = batch_ * numHeads_ * xBlockNumAligned_ * sizeof(int32_t);
697+ workSpaceSize_ = libapiSize_ + selectIdxSize_ + selectNumIdxSize_;
698+ bsaContext->GetWorkspaceSizes(1)[0] = workSpaceSize_;
699+}
700+ 
701+ge::graphStatus BSATiling::CalculateWorkSpace(gert::TilingContext *bsaContext)
702+{
703+ if (blockDim_ == 0) {
704+ OP_LOGE(bsaContext->GetNodeName(), "blockDim is 0");
705+ return ge::GRAPH_FAILED;
706+ }
707+ 
708+ const auto *blockSparseMaskShape = bsaContext->GetInputShape(BLOCK_SPARSE_MASK_INDEX);
709+ maxKvBlockNum_ = blockSparseMaskShape->GetStorageShape().GetDim(3);
710+ maxQBlockNum_ = blockSparseMaskShape->GetStorageShape().GetDim(2);
711+ selectIdxSize_ = (sparsePatternMode_ == SPARSE_PATTERN_MODE_TABLE) ?
712+ 0 : CeilDiv(blockShapeX_, 128) * CeilDiv(maxKvBlockNum_, 32) * 32 * sizeof(uint32_t) * batch_ * numHeads_ * maxQBlockNum_;
713+ selectNumIdxSize_ = CeilDiv(blockShapeX_, 128) * sizeof(uint32_t) * 32 * batch_ * numHeads_ * maxQBlockNum_;
714+ int32_t syncSize_ = sizeof(uint32_t) * 256;
715+
716+ mm1OutSize_ = blockDim_ * WORKSPACE_BLOCK_SIZE_DB * sizeof(float) * NUM3;
717+ smOnlineOutSize_ = blockDim_ * WORKSPACE_BLOCK_SIZE_DB * sizeof(uint16_t) * NUM3;
718+ mm2OutSize_ = blockDim_ * WORKSPACE_BLOCK_SIZE_DB * sizeof(float) * NUM3;
719+ updateSize_ = blockDim_ * WORKSPACE_BLOCK_SIZE_DB * sizeof(float) * NUM3;
720+
721+ workSpaceSize_ = libapiSize_ + mm1OutSize_ + smOnlineOutSize_ + mm2OutSize_ + updateSize_ + selectNumIdxSize_ + selectIdxSize_ + syncSize_;
722+ bsaContext->GetWorkspaceSizes(1)[0] = workSpaceSize_;
723+ uint32_t totalTaskNumMask = batch_ * numHeads_ * maxQBlockNum_;
724+ avgRowNumPerSubCore_ = CeilDiv(totalTaskNumMask, blockDim_ * 2);
725+ preActivateSubCoreNum_ = CeilDiv(totalTaskNumMask, avgRowNumPerSubCore_);
726+
727+ return ge::GRAPH_SUCCESS;
728+}
729+ 
730+void BSATiling::CalcMatmulPhaseL1TileInfo950()
731+{
732+ uint32_t qBaseTileAligned128 = (qBaseTile_ + TILE_SIZE_128 - 1) / TILE_SIZE_128 * TILE_SIZE_128;
733+ uint32_t embeddingSizeAligned128 = (embeddingSize_ + TILE_SIZE_128 - 1) / TILE_SIZE_128 * TILE_SIZE_128;
734+ uint32_t kvBaseTileAligned128 = (kvBaseTile_ + TILE_SIZE_128 - 1) / TILE_SIZE_128 * TILE_SIZE_128;
735+ 
736+ // 基本原则，Q的基块常驻在L1
737+ mm1L1TileM_ = qBaseTileAligned128;
738+ mm1L1TileKLeft_ = embeddingSizeAligned128;
739+ qL1BufNum_ = SOLO_BUF;
740+ if (embeddingSizeAligned128 == D_SIZE_256) {
741+ // K矩阵开启2buf，D按256分割，S2按128分割
742+ // 可以证明，当Q常驻在L1时，无论D和kvBaseTile_为多少，均不会有QK的MTE2重复搬运
743+ mm1L1TileN_ = TILE_SIZE_128;
744+ mm1L1TileKRight_ = TILE_SIZE_256;
745+ kL1BufNum_ = DUO_BUF;
746+ // V矩阵D按256分割，kvBaseTile_不分割，指令提前于核间同步下发，是否开启db取决于kvBaseTile_的大小
747+ mm2L1TileN_ = TILE_SIZE_256;
748+ mm2L1TileKLeft_ = kvBaseTileAligned128;
749+ vL1BufNum_ = TILE_SIZE_256 / kvBaseTileAligned128;
750+ // P矩阵在950上会常驻L1，由于基块的prelaunch为2，因此最好有3 buf，以免基块间流水阻塞
751+ mm2L1TileM_ = qBaseTileAligned128;
752+ mm2L1TileKRight_ = kvBaseTileAligned128;
753+ pL1BufNum_ = TRIO_BUF;
754+ } else if (embeddingSizeAligned128 == D_SIZE_128) {
755+ // K矩阵开启2buf，D按128分割，S2按512分割
756+ mm1L1TileN_ = TILE_SIZE_512;
757+ mm1L1TileKRight_ = TILE_SIZE_128;
758+ kL1BufNum_ = DUO_BUF;
759+ // V矩阵开启db，D按128分割，kvBaseTile_不分割，指令同样提前于核间同步下发
760+ // 如果kvBaseTile_进一步增大，考虑关闭db，使得kvBaseTile_不分割
761+ mm2L1TileN_ = TILE_SIZE_128;
762+ mm2L1TileKLeft_ = kvBaseTileAligned128;
763+ vL1BufNum_ = DUO_BUF;
764+ // P矩阵在950上会常驻L1，由于基块的prelaunch为2，因此最好有3 buf，以免基块间流水阻塞
765+ mm2L1TileM_ = qBaseTileAligned128;
766+ mm2L1TileKRight_ = kvBaseTileAligned128;
767+ pL1BufNum_ = TRIO_BUF;
768+ }
769+}
770+ 
771+ge::graphStatus BSATiling::FillTilingData(gert::TilingContext *bsaContext)
772+{
773+ if (tilingData_ == nullptr) {
774+ return ge::GRAPH_FAILED;
775+ }
776+ tilingData_->set_numHeads(numHeads_);
777+ tilingData_->set_embeddingSize(embeddingSize_);
778+ tilingData_->set_blockSize(blockSize_);
779+ tilingData_->set_kvHeads(kvHeads_);
780+ tilingData_->set_batch(batch_);
781+ tilingData_->set_maxNumBlocksPerBatch(maxNumBlocksPerBatch_);
782+ tilingData_->set_firstBatchTaskNum(firstBatchTaskNum_);
783+ tilingData_->set_totalTaskNum(totalTaskNum_);
784+ tilingData_->set_maskType(maskType_);
785+
786+ tilingData_->set_blockShapeX(blockShapeX_);
787+ tilingData_->set_blockShapeY(blockShapeY_);
788+
789+ tilingData_->set_firstQBlockNum(firstQBlockNum_);
790+ tilingData_->set_totalQBlocks(totalQBlocks_);
791+ tilingData_->set_maxKvBlockNum(maxKvBlockNum_);
792+ tilingData_->set_maxQBlockNum(maxQBlockNum_);
793+ tilingData_->set_avgRowNumPerSubCore(avgRowNumPerSubCore_);
794+ tilingData_->set_preActivateSubCoreNum(preActivateSubCoreNum_);
795+
796+ tilingData_->set_kvCacheLayout(static_cast<uint32_t>(kvCacheLayout_));
797+ tilingData_->set_queryLayout(static_cast<uint32_t>(qInputLayout_));
798+ tilingData_->set_maxQSeqlen(maxQSeqlen_);
799+ tilingData_->set_maxKvSeqlen(maxKvSeqlen_);
800+ tilingData_->set_sparsePatternMode(sparsePatternMode_);
801+ // BNSD格式下当actualSeqLengths为nullptr时，使用maxQSeqlen和maxKvSeqlen作为统一值
802+ tilingData_->set_useUniformQSeqlen(useUniformQSeqlen_ ? 1 : 0);
803+ tilingData_->set_useUniformKvSeqlen(useUniformKvSeqlen_ ? 1 : 0);
804+
805+ // 生成tilingKey（按照开发规范：在tiling层生成）
806+ uint64_t tilingKey = GenerateTilingKey(bsaContext);
807+ tilingData_->set_tilingKey(tilingKey);
808+ bsaContext->SetTilingKey(tilingKey);
809+ bsaContext->SetBlockDim(blockDim_);
810+
811+ tilingData_->set_mm1OutSize(mm1OutSize_);
812+ tilingData_->set_smOnlineOutSize(smOnlineOutSize_);
813+ tilingData_->set_mm2OutSize(mm2OutSize_);
814+ tilingData_->set_updateSize(updateSize_);
815+ tilingData_->set_workSpaceSize(workSpaceSize_);
816+ tilingData_->set_scaleValue(scaleValue_);
817+ tilingData_->set_selectNumIdxSize(selectNumIdxSize_);
818+ tilingData_->set_selectIdxSize(selectIdxSize_);
819+ // fill 950 mask2idx tile info
820+ tilingData_->BsaMask2IdxTileInfo.set_xBlockNumAligned(xBlockNumAligned_);
821+ tilingData_->BsaMask2IdxTileInfo.set_yBlockNumAligned(yBlockNumAligned_);
822+ tilingData_->BsaMask2IdxTileInfo.set_avgRowPerSubCore(avgRowPerSubCore_);
823+ tilingData_->BsaMask2IdxTileInfo.set_preActiveSubCoreNum(preActiveSubCoreNum_);
824+ // fill 950 base tile info
825+ tilingData_->BsaBaseTileInfo.set_qBaseTile(qBaseTile_);
826+ tilingData_->BsaBaseTileInfo.set_kvBaseTile(kvBaseTile_);
827+ // fill 950 matmul phase L1 tile info
828+ tilingData_->BsaMmPhaseL1TileInfo.set_mm1L1TileM(mm1L1TileM_);
829+ tilingData_->BsaMmPhaseL1TileInfo.set_mm1L1TileN(mm1L1TileN_);
830+ tilingData_->BsaMmPhaseL1TileInfo.set_mm1L1TileKLeft(mm1L1TileKLeft_);
831+ tilingData_->BsaMmPhaseL1TileInfo.set_mm1L1TileKRight(mm1L1TileKRight_);
832+ tilingData_->BsaMmPhaseL1TileInfo.set_mm2L1TileM(mm2L1TileM_);
833+ tilingData_->BsaMmPhaseL1TileInfo.set_mm2L1TileN(mm2L1TileN_);
834+ tilingData_->BsaMmPhaseL1TileInfo.set_mm2L1TileKLeft(mm2L1TileKLeft_);
835+ tilingData_->BsaMmPhaseL1TileInfo.set_mm2L1TileKRight(mm2L1TileKRight_);
836+ tilingData_->BsaMmPhaseL1TileInfo.set_qL1BufNum(qL1BufNum_);
837+ tilingData_->BsaMmPhaseL1TileInfo.set_kL1BufNum(kL1BufNum_);
838+ tilingData_->BsaMmPhaseL1TileInfo.set_vL1BufNum(vL1BufNum_);
839+ tilingData_->BsaMmPhaseL1TileInfo.set_pL1BufNum(pL1BufNum_);
840+ return ge::GRAPH_SUCCESS;
841+}
842+ 
843+uint64_t BSATiling::GenerateTilingKey(gert::TilingContext *bsaContext)
844+{
845+ /**
846+ * 64位整数，使用十进制位域表示：
847+ * AAAABBBBCCCCDDDDEEEE
848+ * - 位0-1: Q Layout（个位） 2=TND, 3=BNSD
849+ * - 位2-4: Mask Type（千位） 0=NoMask, 3=CausalMask
850+ * - 位5-7: Softmax Precision（十万位） 0=Float, 1=Half
851+ * - 位8-10: PagedCache Flag（千万位） 0=NoCache, 1=WithCache
852+ * - 位11-13: KV Layout（十亿位） 00=TND, 20=BNSD
853+ * - 位14-15: Data Type（百亿位） 00=FP16, 22=BF16
854+ * - 位16-18: Operator Category（千万亿位） 900=EagleQuantBlockSparseAttention910, 905=EagleQuantBlockSparseAttention950
855+ *
856+ * 示例：
857+ * - FP16, TND, TND, NoCache, Half, NoMask = 9000000030100002
858+ * - FP16, TND, TND, NoCache, Float, NoMask = 9000000030000002
859+ */
860+
861+ uint64_t tilingKey = 9000000000000000ULL; // RFA基础值（Operator Category = 900）
862+ if (socVer_ == SOC_VER_950_CODE) {
863+ tilingKey = 9050000000000000ULL;
864+ }
865+
866+ // [位14-15] Data Type（百亿位）
867+ if (dataType_ == ge::DT_FLOAT16) {
868+ tilingKey += 0; // 00 for FP16
869+ } else if (dataType_ == ge::DT_BF16) {
870+ tilingKey += 22220ULL; // 22 for BF16 -> 9000000030000002 + 22220 = 9000000030022222
871+ } else if (dataType_ == ge::DT_INT8) {
872+ auto oOutputDesc = bsaContext->GetOutputDesc(0);
873+ OP_CHECK_NULL_WITH_CONTEXT(bsaContext, oOutputDesc);
874+ auto dtypeO = oOutputDesc->GetDataType();
875+ if (dtypeO == ge::DT_FLOAT16) {
876+ tilingKey += 44440ULL; // 11 for INT8 output with FP16
877+ } else if (dtypeO == ge::DT_BF16) {
878+ tilingKey += 55550ULL; // 33 for INT8 output with BF16
879+ }
880+ }
881+
882+ // [位11-13] KV Layout（十亿位）
883+ if (kvCacheLayout_ == RFAKvCacheLayout::TND) {
884+ tilingKey += 30000000ULL; // 00 for TND
885+ } else if (kvCacheLayout_ == RFAKvCacheLayout::BNSD) {
886+ tilingKey += 50000000ULL; // 20 for BNSD
887+ }
888+
889+ // [位8-10] PagedCache Flag（千万位）
890+ bool hasPagedCache = (bsaContext->GetOptionalInputTensor(BLOCK_TABLE_INDEX) != nullptr);
891+ if (hasPagedCache) {
892+ tilingKey += 1000000ULL; // 1 for WithCache
893+ }
894+
895+ // [位5-7] Softmax Precision（十万位）
896+ if (innerPrecise_ == 1) {
897+ tilingKey += 100000ULL; // 1 for Half (FP16) Softmax
898+ } else if (innerPrecise_ == 4) {
899+ tilingKey += 400000ULL; // 4 for low prec online softmax & high prec rescale O
900+ }
901+ // innerPrecise_ == 0: 0 for Float Softmax（默认值）
902+
903+ // [位2-4] Mask Type（千位）
904+ if (maskType_ == 3) { // Causal mask
905+ tilingKey += 3000ULL;
906+ }
907+ // maskType_ == 0: 0 for NoMask（默认值）
908+
909+ // [位0-1] Q Layout（个位）
910+ if (qInputLayout_ == RFAQInputLayout::TND_Q) {
911+ tilingKey += 2; // 2 for TND
912+ } else if (qInputLayout_ == RFAQInputLayout::BNSD_Q) {
913+ tilingKey += 3; // 3 for BNSD
914+ }
915+ 
916+ // Softmax LSE（亿位）
917+ if (softmaxLseFlag_) {
918+ tilingKey += 100000000ULL; // 1 for lse out
919+ }
920+ 
921+ constexpr uint64_t V_DTYPE_OFFSET = 10000000000ULL; // 10^10
922+ if (vDataType_ == ge::DT_FLOAT8_E4M3FN || vDataType_ == ge::DT_INT8) {
923+ tilingKey += V_DTYPE_OFFSET * 1ULL;
924+ } else if (vDataType_ == ge::DT_FLOAT8_E5M2) {
925+ tilingKey += V_DTYPE_OFFSET * 2ULL;
926+ }
927+ // vDataType 与 Q/K 相同时，加 0，保持兼容
928+
929+
930+ return tilingKey;
931+}
932+ 
933+ge::graphStatus BSATiling::GetBsaTiling(gert::TilingContext *bsaContext,
934+ EagleQuantBlockSparseAttentionTilingData &tilingData)
935+{
936+ tilingData_ = &tilingData;
937+ ge::graphStatus ret = GetNpuInfo(bsaContext);
938+ if (ret != ge::GRAPH_SUCCESS) {
939+ OP_LOGE(bsaContext->GetNodeName(), "GetNpuInfo failed");
940+ return ret;
941+ }
942+ if (GetInputLayout(bsaContext) != ge::GRAPH_SUCCESS ||
943+ ParseAttrs(bsaContext) != ge::GRAPH_SUCCESS ||
944+ ParseRequiredTensors(bsaContext) != ge::GRAPH_SUCCESS ||
945+ ParseOptionalTensors(bsaContext) != ge::GRAPH_SUCCESS) {
946+ ret = ge::GRAPH_FAILED;
947+ return ret;
948+ }
949+ // 校验TND格式下qseqlen和kvseqlen之和是否分别等于Q和KV的T
950+ ret = ValidateTNDSeqlenSum(bsaContext);
951+ if (ret != ge::GRAPH_SUCCESS) {
952+ OP_LOGE(bsaContext->GetNodeName(), "ValidateTNDSeqlenSum failed");
953+ return ret;
954+ }
955+
956+ if (socVer_ == SOC_VER_950_CODE) {
957+ CalcSplitCoreTilingParams950();
958+ CalcMatmulPhaseL1TileInfo950();
959+ CalcWorkspaceTilingParams950(bsaContext);
960+ } else {
961+ ret = CalculateTaskSplit(bsaContext);
962+ if (ret != ge::GRAPH_SUCCESS) {
963+ OP_LOGE(bsaContext->GetNodeName(), "CalculateTaskSplit failed");
964+ return ret;
965+ }
966+ ret = CalculateWorkSpace(bsaContext);
967+ if (ret != ge::GRAPH_SUCCESS) {
968+ OP_LOGE(bsaContext->GetNodeName(), "CalculateWorkSpace failed");
969+ return ret;
970+ }
971+ }
972+ 
973+ ret = FillTilingData(bsaContext);
974+ if (ret != ge::GRAPH_SUCCESS) {
975+ OP_LOGE(bsaContext->GetNodeName(), "FillTilingData failed");
976+ return ret;
977+ }
978+ return ge::GRAPH_SUCCESS;
979+}
980+ 
981+ge::graphStatus BSATiling::BsaSetTilingData(gert::TilingContext *context,
982+ EagleQuantBlockSparseAttentionTilingData &tilingData)
983+{
984+ OP_CHECK_IF(context->GetRawTilingData() == nullptr,
985+ OPS_REPORT_VECTOR_INNER_ERR("EagleQuantBlockSparseAttention",
986+ "RawTilingData got from GE context is nullptr."), return ge::GRAPH_FAILED);
987+ tilingData.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
988+ context->GetRawTilingData()->SetDataSize(tilingData.GetDataSize());
989+ return ge::GRAPH_SUCCESS;
990+}
991+ 
992+ASCENDC_EXTERN_C ge::graphStatus TilingEagleQuantBlockSparseAttention(gert::TilingContext* context)
993+{
994+ OP_CHECK_IF(context == nullptr, OPS_REPORT_VECTOR_INNER_ERR("EagleQuantBlockSparseAttention",
995+ "Context is nullptr."), return ge::GRAPH_FAILED);
996+ EagleQuantBlockSparseAttentionTilingData tilingData;
997+ BSATiling bsaTiling;
998+ if (bsaTiling.GetBsaTiling(context, tilingData) == ge::GRAPH_SUCCESS) {
999+ bsaTiling.BsaSetTilingData(context, tilingData);
1000+ return ge::GRAPH_SUCCESS;
1001+ } else {
1002+ OP_LOGE(context->GetNodeName(), "GetBsaTiling failed");
1003+ return ge::GRAPH_FAILED;
1004+ }
1005+}
1006+ 
1007+ASCENDC_EXTERN_C ge::graphStatus TilingPrepareForEagleQuantBlockSparseAttention(gert::TilingParseContext* context)
1008+{
1009+ (void) context;
1010+ return ge::GRAPH_SUCCESS;
1011+}
1012+ 
1013+IMPL_OP_OPTILING(EagleQuantBlockSparseAttention)
1014+ .Tiling(TilingEagleQuantBlockSparseAttention)
1015+ .TilingInputsDataDependency({5, 6, 7}, {gert::TilingPlacement::TILING_ON_HOST, gert::TilingPlacement::TILING_ON_AICPU})
1016+ .TilingParse<EagleQuantBlockSparseAttentionCompileInfo>(TilingPrepareForEagleQuantBlockSparseAttention); // 向框架注册入口函数;
1017+ 
1018+} // namespace optiling
Acsrc/ops/eagle_quant_block_sparse_attention/op_host/eagle_quant_block_sparse_attention_tiling.h+286-0
@@ -0,0 +1,286 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_TILING_H
12+#define EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_TILING_H
13+ 
14+#include <cstdint>
15+#include "register/tilingdata_base.h"
16+#include "tiling/platform/platform_ascendc.h"
17+#include "tiling/tiling_api.h"
18+#include "register/op_def_registry.h"
19+ 
20+namespace optiling {
21+// mask 2 idx tile info
22+BEGIN_TILING_DATA_DEF(BsaMask2IdxTiling)
23+TILING_DATA_FIELD_DEF(uint32_t, xBlockNumAligned);
24+TILING_DATA_FIELD_DEF(uint32_t, yBlockNumAligned);
25+TILING_DATA_FIELD_DEF(uint32_t, avgRowPerSubCore);
26+TILING_DATA_FIELD_DEF(uint32_t, preActiveSubCoreNum);
27+END_TILING_DATA_DEF;
28+REGISTER_TILING_DATA_CLASS(BsaMask2IdxTilingOp, BsaMask2IdxTiling)
29+// attention tile info
30+BEGIN_TILING_DATA_DEF(BsaBaseTiling)
31+TILING_DATA_FIELD_DEF(uint32_t, qBaseTile);
32+TILING_DATA_FIELD_DEF(uint32_t, kvBaseTile);
33+END_TILING_DATA_DEF;
34+REGISTER_TILING_DATA_CLASS(BsaBaseTilingOp, BsaBaseTiling)
35+// matmul phase L1 tile info
36+BEGIN_TILING_DATA_DEF(BsaMmPhaseL1Tiling)
37+TILING_DATA_FIELD_DEF(uint32_t, mm1L1TileM);
38+TILING_DATA_FIELD_DEF(uint32_t, mm1L1TileN);
39+TILING_DATA_FIELD_DEF(uint32_t, mm1L1TileKLeft);
40+TILING_DATA_FIELD_DEF(uint32_t, mm1L1TileKRight);
41+TILING_DATA_FIELD_DEF(uint32_t, mm2L1TileM);
42+TILING_DATA_FIELD_DEF(uint32_t, mm2L1TileN);
43+TILING_DATA_FIELD_DEF(uint32_t, mm2L1TileKLeft);
44+TILING_DATA_FIELD_DEF(uint32_t, mm2L1TileKRight);
45+TILING_DATA_FIELD_DEF(uint32_t, qL1BufNum);
46+TILING_DATA_FIELD_DEF(uint32_t, kL1BufNum);
47+TILING_DATA_FIELD_DEF(uint32_t, vL1BufNum);
48+TILING_DATA_FIELD_DEF(uint32_t, pL1BufNum);
49+END_TILING_DATA_DEF;
50+REGISTER_TILING_DATA_CLASS(BsaMmPhaseL1TilingOp, BsaMmPhaseL1Tiling)
51+// EagleQuantBlockSparseAttention Tiling数据定义
52+BEGIN_TILING_DATA_DEF(EagleQuantBlockSparseAttentionTilingData)
53+// 基础参数
54+TILING_DATA_FIELD_DEF(uint32_t, batch);
55+TILING_DATA_FIELD_DEF(uint32_t, numHeads);
56+TILING_DATA_FIELD_DEF(uint32_t, kvHeads);
57+TILING_DATA_FIELD_DEF(uint32_t, embeddingSize);
58+TILING_DATA_FIELD_DEF(uint32_t, blockSize);
59+TILING_DATA_FIELD_DEF(uint32_t, maxNumBlocksPerBatch);
60+TILING_DATA_FIELD_DEF(uint32_t, firstBatchTaskNum);
61+TILING_DATA_FIELD_DEF(uint32_t, totalTaskNum);
62+TILING_DATA_FIELD_DEF(uint32_t, maskType);
63+TILING_DATA_FIELD_DEF(float, scaleValue);
64+TILING_DATA_FIELD_DEF(uint32_t, totalQBlocks); // T: 所有batch中Q方向切块的总数
65+TILING_DATA_FIELD_DEF(uint32_t, firstQBlockNum); // T: 所有batch中Q方向切块的总数
66+ 
67+// 稀疏分块参数 (blockShape)
68+TILING_DATA_FIELD_DEF(uint64_t, blockShapeX); // block的x维度(Q方向)
69+TILING_DATA_FIELD_DEF(uint64_t, blockShapeY); // block的y维度(KV方向)
70+ 
71+// selectIdx相关参数
72+TILING_DATA_FIELD_DEF(uint32_t, maxKvBlockNum); // 最大KV块数量（selectIdx的最后一维）
73+TILING_DATA_FIELD_DEF(uint32_t, maxQBlockNum); // 最大KV块数量（selectIdx的最后一维）
74+TILING_DATA_FIELD_DEF(uint32_t, avgRowNumPerSubCore);
75+TILING_DATA_FIELD_DEF(uint32_t, preActivateSubCoreNum);
76+ 
77+ 
78+// query Layout: 0=TND, 1=BNSD
79+TILING_DATA_FIELD_DEF(uint32_t, queryLayout);
80+ 
81+// KV Cache Layout: 0=TND, 1=BNSD
82+TILING_DATA_FIELD_DEF(uint32_t, kvCacheLayout);
83+ 
84+// BNSD格式的最大序列长度（用于计算stride）
85+// 当actualSeqLengths为nullptr时，maxQSeqlen也用作统一的qseqlen值
86+TILING_DATA_FIELD_DEF(uint32_t, maxQSeqlen); // BNSD格式Q的第三维（S维度），或统一的qseqlen值
87+// 当actualSeqLengthsKv为nullptr时，maxKvSeqlen也用作统一的kvseqlen值
88+TILING_DATA_FIELD_DEF(uint32_t, maxKvSeqlen); // BNSD格式KV的第三维（S维度），或统一的kvseqlen值
89+TILING_DATA_FIELD_DEF(uint32_t, useUniformQSeqlen); // 是否使用统一的qseqlen值（1=是，0=否）
90+TILING_DATA_FIELD_DEF(uint32_t, useUniformKvSeqlen); // 是否使用统一的kvseqlen值（1=是，0=否）
91+ 
92+// TilingKey for kernel dispatch (生成在tiling层)
93+TILING_DATA_FIELD_DEF(uint64_t, tilingKey);
94+TILING_DATA_FIELD_DEF(uint64_t, selectNumIdxSize);
95+TILING_DATA_FIELD_DEF(uint64_t, selectIdxSize);
96+// Workspace大小
97+TILING_DATA_FIELD_DEF(uint64_t, mm1OutSize);
98+TILING_DATA_FIELD_DEF(uint64_t, smOnlineOutSize);
99+TILING_DATA_FIELD_DEF(uint64_t, mm2OutSize);
100+TILING_DATA_FIELD_DEF(uint64_t, updateSize);
101+TILING_DATA_FIELD_DEF(uint64_t, workSpaceSize);
102+ 
103+TILING_DATA_FIELD_DEF_STRUCT(BsaMask2IdxTiling, BsaMask2IdxTileInfo);
104+TILING_DATA_FIELD_DEF_STRUCT(BsaBaseTiling, BsaBaseTileInfo);
105+TILING_DATA_FIELD_DEF_STRUCT(BsaMmPhaseL1Tiling, BsaMmPhaseL1TileInfo);
106+TILING_DATA_FIELD_DEF(uint32_t, sparsePatternMode);
107+END_TILING_DATA_DEF;
108+REGISTER_TILING_DATA_CLASS(EagleQuantBlockSparseAttention, EagleQuantBlockSparseAttentionTilingData)
109+ 
110+// EagleQuantBlockSparseAttention编译信息
111+struct EagleQuantBlockSparseAttentionCompileInfo {
112+ uint32_t inputDataByte = 2;
113+ ge::DataType inputDataType;
114+
115+ uint32_t coreNum = 0;
116+ uint32_t aivNum = 0;
117+ uint32_t aicNum = 0;
118+ uint64_t ubSize = 0;
119+ uint64_t l1Size = 0;
120+ uint64_t sysWorkspaceSize = 0;
121+ platform_ascendc::SocVersion socVersion;
122+};
123+ 
124+// 输入参数信息
125+struct RequiredParaInfo {
126+ const gert::CompileTimeTensorDesc *desc;
127+ const gert::StorageShape *shape;
128+};
129+ 
130+struct OptionalParaInfo {
131+ const gert::CompileTimeTensorDesc *desc;
132+ const gert::Tensor *tensor;
133+};
134+ 
135+// KVCache Layout枚举
136+enum RFAKvCacheLayout : uint32_t {
137+ TND = 0, // [T, N, D] format
138+ BNSD = 1 // [B, N, S, D] format
139+};
140+ 
141+// Q Input Layout枚举
142+enum RFAQInputLayout : uint32_t {
143+ TND_Q = 0, // [T, N, D] format
144+ BNSD_Q = 1 // [B, N, S, D] format
145+};
146+ 
147+// inner prec 枚举
148+enum BsaInnerCalcPrec : uint32_t {
149+ ALL_HIGH = 0,
150+ ALL_LOW = 1,
151+ LOW_HIGH_MIXED = 4
152+};
153+ 
154+// Tiling类
155+class BSATiling {
156+public:
157+ BSATiling() = default;
158+ ~BSATiling() = default;
159+
160+ ge::graphStatus GetBsaTiling(gert::TilingContext *bsaContext,
161+ EagleQuantBlockSparseAttentionTilingData &tilingData);
162+ ge::graphStatus BsaSetTilingData(gert::TilingContext *context,
163+ EagleQuantBlockSparseAttentionTilingData &tilingData);
164+ 
165+private:
166+ ge::graphStatus GetNpuInfo(gert::TilingContext *bsaContext);
167+ ge::graphStatus ParseAttrs(gert::TilingContext *bsaContext);
168+ ge::graphStatus GetInputLayout(gert::TilingContext *bsaContext);
169+ ge::graphStatus ParseRequiredTensors(gert::TilingContext *bsaContext);
170+ ge::graphStatus ParseOptionalTensors(gert::TilingContext *bsaContext);
171+ ge::graphStatus CheckQKVDtype(gert::TilingContext *bsaContext);
172+ ge::graphStatus CheckQKVDimVal(gert::TilingContext *bsaContext,
173+ uint32_t kHeads, uint32_t vHeads, uint32_t kHeadDim, uint32_t vHeadDim);
174+ ge::graphStatus ParseQKVInTND(gert::TilingContext *bsaContext);
175+ ge::graphStatus ParseQKVInBNSD(gert::TilingContext *bsaContext);
176+ ge::graphStatus ParseSeqlensInTND(gert::TilingContext *bsaContext);
177+ ge::graphStatus ParseSeqlensInBNSD(gert::TilingContext *bsaContext);
178+ ge::graphStatus ParseSeqlens(gert::TilingContext *bsaContext);
179+ ge::graphStatus ParseSparsePattern(gert::TilingContext *bsaContext);
180+ ge::graphStatus ParseAttenMask(gert::TilingContext *bsaContext);
181+ ge::graphStatus ParseBlockTable(gert::TilingContext *bsaContext);
182+ ge::graphStatus CheckSparsePattern(gert::TilingContext *bsaContext, const int64_t defaultShape);
183+ ge::graphStatus ValidateTNDSeqlenSum(gert::TilingContext *bsaContext);
184+ // 950 exclusive
185+ uint32_t GetCurQSTileNum950(int64_t curQSeqlen);
186+ void CalcBaseTileTilingParams950();
187+ void CalcSplitCoreTilingParams950();
188+ void CalcWorkspaceTilingParams950(gert::TilingContext *bsaContext);
189+ void CalcMatmulPhaseL1TileInfo950();
190+ // 910 exclusive
191+ ge::graphStatus CalculateTaskSplit(gert::TilingContext *bsaContext);
192+ ge::graphStatus CalculateWorkSpace(gert::TilingContext *bsaContext);
193+ // shared
194+ void CalculateBatchTaskSplit(int64_t qSeqlen, uint32_t groupSize,
195+ uint32_t &curTaskNum, uint32_t &curQBlockNum);
196+ ge::graphStatus FillTilingData(gert::TilingContext *bsaContext);
197+ uint64_t GenerateTilingKey(gert::TilingContext *bsaContext);
198+
199+private:
200+ uint32_t batch_ = 0;
201+ uint32_t qSeqlen_ = 0;
202+ uint32_t kvSeqlen_ = 0;
203+ uint32_t numHeads_ = 0;
204+ uint32_t kvHeads_ = 0;
205+ uint32_t embeddingSize_ = 0;
206+ uint32_t blockSize_ = 128;
207+ int64_t blockShapeX_ = 0; // block的x维度
208+ int64_t blockShapeY_ = 0; // block的y维度
209+ float scaleValue_ = 0.0f;
210+ uint32_t maskType_ = 0;
211+ uint32_t sparsePatternMode_ = 0;
212+ uint32_t innerPrecise_ = 1; // 0=float32 softmax, 1=fp16 softmax
213+ bool softmaxLseFlag_ = false;
214+
215+ uint32_t totalQBlocks_ = 0;
216+ uint32_t maxKvBlockNum_ = 0;
217+ uint32_t maxQBlockNum_ = 0;
218+ uint32_t avgRowNumPerSubCore_ = 0;
219+ uint32_t preActivateSubCoreNum_ = 0;
220+ uint32_t firstQBlockNum_ = 0;
221+ uint32_t firstBatchTaskNum_ = 0;
222+ uint32_t totalTaskNum_ = 0;
223+ uint32_t maxNumBlocksPerBatch_ = 0;
224+ const int64_t *qSeqLenList_ = nullptr;
225+ const int64_t *kvSeqLenList_ = nullptr;
226+ const int64_t *blockShapeList = nullptr;
227+ bool useUniformQSeqlen_ = false; // 是否使用统一的qseqlen值（使用maxQSeqlen_）
228+ bool useUniformKvSeqlen_ = false; // 是否使用统一的kvseqlen值（使用maxKvSeqlen_）
229+ 
230+ uint64_t mm1OutSize_ = 0;
231+ uint64_t smOnlineOutSize_ = 0;
232+ uint64_t mm2OutSize_ = 0;
233+ uint64_t updateSize_ = 0;
234+ uint64_t selectNumIdxSize_ = 0;
235+ uint64_t selectIdxSize_ = 0;
236+
237+ RFAKvCacheLayout kvCacheLayout_ = RFAKvCacheLayout::TND;
238+ RFAQInputLayout qInputLayout_ = RFAQInputLayout::TND_Q;
239+
240+ uint32_t blockDim_ = 20;
241+ uint32_t aivNum_ = 0;
242+ uint32_t aicNum_ = 0;
243+ uint32_t socVer_ = 0;
244+ uint64_t ubSize_ = 0;
245+ uint64_t workSpaceSize_ = 0;
246+ uint64_t libapiSize_ = 0;
247+
248+ uint32_t maxQSeqlen_ = 0; // BNSD格式Q的第三维（S维度）
249+ uint32_t maxKvSeqlen_ = 0; // BNSD格式KV的第三维（S维度）
250+ int64_t totalTokensT_ = 0; // TND格式Q的第一维（T维度，总token数）
251+ int64_t totalTokensKv_ = 0; // TND格式KV的第一维（T维度，总token数
252+ 
253+ // mask2idx tile info
254+ uint32_t xBlockNumAligned_;
255+ uint32_t yBlockNumAligned_;
256+ uint32_t avgRowPerSubCore_;
257+ uint32_t preActiveSubCoreNum_;
258+ // base tile info
259+ uint32_t qBaseTile_;
260+ uint32_t kvBaseTile_;
261+ // L1 tile info
262+ // further splits the base tiles
263+ uint32_t mm1L1TileM_;
264+ uint32_t mm1L1TileN_;
265+ uint32_t mm1L1TileKLeft_;
266+ uint32_t mm1L1TileKRight_;
267+ uint32_t mm2L1TileM_;
268+ uint32_t mm2L1TileN_;
269+ uint32_t mm2L1TileKLeft_;
270+ uint32_t mm2L1TileKRight_;
271+ uint32_t qL1BufNum_;
272+ uint32_t kL1BufNum_;
273+ uint32_t vL1BufNum_;
274+ uint32_t pL1BufNum_;
275+
276+ ge::DataType dataType_ = ge::DT_FLOAT16;
277+ // 在 dataType_ 声明附近添加：
278+ ge::DataType vDataType_ = ge::DT_FLOAT16; // 新增：记录V的独立dtype
279+ 
280+ EagleQuantBlockSparseAttentionTilingData *tilingData_ = nullptr;
281+};
282+ 
283+} // namespace optiling
284+ 
285+#endif // EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_TILING_H
286+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/arch35/eagle_quant_block_sparse_attention_kernel_arch35_qmode1.h+799-0
@@ -0,0 +1,799 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef BSA_ARCH35_KERNEL_QMODE1
12+#define BSA_ARCH35_KERNEL_QMODE1
13+ 
14+#include "eagle_quant_block_sparse_attention_kernel_arch35_regular.h"
15+ 
16+using namespace NpuArch;
17+using namespace tla;
18+ 
19+namespace BsaKernelArch35 {
20+ 
21+template <
22+ class EpilogueMask2Idx,
23+ class QKL1TileShape,
24+ class QKL0TileShape,
25+ class ElementS_,
26+ class QKBias,
27+ class QKTileCopy,
28+ class QKTileMmad,
29+ class EpilogueOnlineSoftmax,
30+ class PVL1TileShape,
31+ class PVL0TileShape,
32+ class ElementOTmp_,
33+ class PVBias,
34+ class PVTileCopy,
35+ class PVTileMmad,
36+ class EpilogueRescaleO,
37+ Format qFormat,
38+ Format kvFormat>
39+class BsaRegularKernelArch35<
40+ EpilogueMask2Idx,
41+ Gemm::Block::BlockMmadTla<Gemm::MmadAtlasA5BsaQK, QKL1TileShape, QKL0TileShape,
42+ int8_t, int8_t, ElementS_, QKBias, QKTileCopy, QKTileMmad>,
43+ EpilogueOnlineSoftmax,
44+ Gemm::Block::BlockMmadTla<Gemm::MmadAtlasA5BsaPV, PVL1TileShape, PVL0TileShape,
45+ float8_e4m3_t, float8_e4m3_t, ElementOTmp_, PVBias, PVTileCopy, PVTileMmad>,
46+ EpilogueRescaleO,
47+ qFormat,
48+ kvFormat> {
49+public:
50+ using BlockMmadQK = Gemm::Block::BlockMmadTla<Gemm::MmadAtlasA5BsaQK, QKL1TileShape, QKL0TileShape,
51+ int8_t, int8_t, ElementS_, QKBias, QKTileCopy, QKTileMmad>;
52+ using BlockMmadPV = Gemm::Block::BlockMmadTla<Gemm::MmadAtlasA5BsaPV, PVL1TileShape, PVL0TileShape,
53+ float8_e4m3_t, float8_e4m3_t, ElementOTmp_, PVBias, PVTileCopy, PVTileMmad>;
54+ using ArchTag = typename BlockMmadPV::ArchTag;
55+
56+ using ElementQ = typename BlockMmadQK::ElementA;
57+ using ElementK = typename BlockMmadQK::ElementB;
58+ using ElementS = typename EpilogueOnlineSoftmax::ElementInput;
59+ using ElementP = typename BlockMmadPV::ElementA;
60+ using ElementV = typename BlockMmadPV::ElementB;
61+ using ElementOTmp = typename BlockMmadPV::ElementC;
62+ using ElementO = typename EpilogueRescaleO::ElementO;
63+ using ElementSparseMask = typename EpilogueMask2Idx::ElementSparseMask;
64+ using ElementSparseIdx = typename EpilogueMask2Idx::ElementSparseIdx;
65+ using ElementSparseCount = typename EpilogueMask2Idx::ElementSparseCount;
66+ 
67+ using LayoutQ = layout::ColumnMajor;
68+ using LayoutK = layout::RowMajor;
69+ //qmode为1的时候设置为RowMajor，原本是Q*KT，现在改成了(K * QT)T，所以LayoutS是RowMajor
70+ using LayoutS = layout::RowMajor;
71+ using LayoutP = layout::RowMajor;
72+ using LayoutV = layout::RowMajor;
73+ using LayoutO = layout::RowMajor;
74+ using LayoutOTmp = layout::RowMajor;
75+ using LayoutSparseIdx = layout::RowMajor;
76+ using LayoutSparseCount = layout::RowMajor;
77+ 
78+ using LayoutTagL1P = typename BlockMmadPV::LayoutTagL1A;
79+ 
80+ static constexpr uint32_t PRE_LAUNCH = 2;
81+ static constexpr uint32_t MAX_CROSS_CORE_BUF_STAGES = PRE_LAUNCH + 1;
82+ static constexpr uint32_t UB_S_OTMP_BUF_STAGES = 2;
83+ static constexpr uint32_t QUANT_BLOCK_SIZE = 64;
84+ 
85+ // Methods
86+ __aicore__ inline
87+ BsaRegularKernelArch35() {}
88+ 
89+ __aicore__ inline
90+ int GetQkStaticLoopNum()
91+ {
92+ return ((qBaseTile_ == BlockMmadQK::L0_TILE_M) &&
93+ (embed_ == BlockMmadQK::L0_TILE_K) &&
94+ ((kvBaseTile_ == BlockMmadQK::L0_TILE_N) ||
95+ (kvBaseTile_ == BlockMmadQK::L0_TILE_N * 2U))) ?
96+ 1 : Gemm::Block::Arch35MmadOpt::DYNAMIC_LOOP;
97+ }
98+ 
99+ __aicore__ inline
100+ int GetPvStaticNLoopNum()
101+ {
102+ if (qBaseTile_ != BlockMmadPV::L0_TILE_M) {
103+ return Gemm::Block::Arch35MmadOpt::DYNAMIC_LOOP;
104+ }
105+ if (embed_ == BlockMmadPV::L0_TILE_N) {
106+ return 1;
107+ }
108+ if (embed_ == BlockMmadPV::L0_TILE_N * 2U) {
109+ return 2;
110+ }
111+ return Gemm::Block::Arch35MmadOpt::DYNAMIC_LOOP;
112+ }
113+ 
114+ template<int quant_mode=1>
115+ __aicore__ inline
116+ void operator()(BsaKernelParamsArch35 const &params)
117+ {
118+ static_assert(quant_mode == 1);
119+ __gm__ EagleQuantBlockSparseAttentionTilingData *bsaTilingData =
120+ reinterpret_cast<__gm__ EagleQuantBlockSparseAttentionTilingData *>(params.tiling);
121+ FetchBaseShapeInfo(bsaTilingData);
122+ CalcOnChipBufTileInfo(bsaTilingData);
123+ // global buffers
124+ AscendC::GlobalTensor<ElementQ> gQ;
125+ gQ.SetGlobalBuffer((__gm__ ElementQ *)params.q);
126+ AscendC::GlobalTensor<ElementK> gK;
127+ gK.SetGlobalBuffer((__gm__ ElementK *)params.k);
128+ AscendC::GlobalTensor<ElementV> gV;
129+ gV.SetGlobalBuffer((__gm__ ElementV *)params.v);
130+ AscendC::GlobalTensor<int64_t> gActualQseqlen;
131+ gActualQseqlen.SetGlobalBuffer((__gm__ int64_t *)params.actualQseqlen);
132+ AscendC::GlobalTensor<int64_t> gActualKvseqlen;
133+ gActualKvseqlen.SetGlobalBuffer((__gm__ int64_t *)params.actualKvseqlen);
134+ AscendC::GlobalTensor<uint8_t> gBlockSparseMask;
135+ gBlockSparseMask.SetGlobalBuffer((__gm__ uint8_t *)params.blockSparseMask);
136+ AscendC::GlobalTensor<ElementO> gO;
137+ gO.SetGlobalBuffer((__gm__ ElementO *)params.o);
138+ AscendC::GlobalTensor<ElementSparseIdx> gSparseIdx;
139+ if (bsaTilingData->sparsePatternMode == SPARSE_PATTERN_MODE_TABLE) {
140+ gSparseIdx.SetGlobalBuffer((__gm__ ElementSparseIdx *)params.blockSparseMask);
141+ } else {
142+ gSparseIdx.SetGlobalBuffer((__gm__ ElementSparseIdx *)params.workSpace);
143+ }
144+ AscendC::GlobalTensor<ElementSparseCount> gSparseCount;
145+ gSparseCount.SetGlobalBuffer((__gm__ ElementSparseCount *)(params.workSpace + sparseIdxSize_));
146+ AscendC::GlobalTensor<float> queryScale;
147+ queryScale.SetGlobalBuffer((__gm__ float *)(params.query_scale));
148+ AscendC::GlobalTensor<float> keyScale;
149+ keyScale.SetGlobalBuffer((__gm__ float *)(params.key_scale));
150+ AscendC::GlobalTensor<float> valueScale;
151+ valueScale.SetGlobalBuffer((__gm__ float *)(params.value_scale));
152+
153+ // cross core data move dst buffers
154+ AscendC::LocalTensor<ElementP> l1PTensor[MAX_CROSS_CORE_BUF_STAGES];
155+ AscendC::LocalTensor<ElementS> ubSTensor[UB_S_OTMP_BUF_STAGES];
156+ AscendC::LocalTensor<ElementOTmp> ubOTmpTensor[UB_S_OTMP_BUF_STAGES];
157+ InitCrossCoreDstBuf(l1PTensor, ubSTensor, ubOTmpTensor);
158+ // core idx
159+ uint32_t coreIdx = AscendC::GetBlockIdx();
160+ uint32_t coreNum = AscendC::GetBlockNum();
161+ // set reverse sync flags
162+ InitSyncFlags<4, 4, 4>();
163+#ifdef __DAV_VEC__
164+ uint32_t totalRowNumBlockMask = batch_ * qHeads_ * xBlockNumAligned_;
165+ if (bsaTilingData->sparsePatternMode == SPARSE_PATTERN_MODE_TABLE) {
166+ //AscendC::printf("good job\n");
167+ SparseTable2Count(resource, gSparseIdx, gSparseCount,
168+ totalRowNumBlockMask, yBlockNumAligned_, avgRowPerSubCore_, preActiveSubCoreNum_);
169+ } else {
170+ EpilogueMask2Idx epilogueMask2Idx(resource);
171+ epilogueMask2Idx(
172+ gBlockSparseMask, gSparseIdx, gSparseCount,
173+ totalRowNumBlockMask, yBlockNumAligned_, avgRowPerSubCore_, preActiveSubCoreNum_);
174+ }
175+#endif
176+ AscendC::SyncAll<false>();
177+#ifdef __DAV_CUBE__
178+ coreIdx = AscendC::GetBlockIdx();
179+ BlockMmadQK blockMmadQK(resource, mm1L1TileHelper_);
180+ BlockMmadPV blockMmadPV(resource, mm2L1AddrStart_, mm2L1TileHelper_);
181+#endif
182+#ifdef __DAV_VEC__
183+ coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum();
184+ EpilogueOnlineSoftmax epilogueOnlineSoftmax(resource, scaleValue_, blockShapeY_);
185+ EpilogueRescaleO epilogueRescaleO(resource, embed_);
186+ auto scaleTensor = resource.ubBuf.template GetBufferByByte<float>(7 * 32 * 1024 + 4096 * 2);
187+#endif
188+ uint32_t qSTileNumPerFullXBlock = CeilDiv(blockShapeX_, qBaseTile_);
189+ // Calculate strides based on layout
190+ // For TND: [T, N, D], stride = N * D
191+ // For BNSD: [B, N, S, D], strideB = N * S * D, strideN = S * D, strideS = D
192+ int64_t strideQO = 0;
193+ int64_t strideKV = 0;
194+ int64_t strideQOB = 0; // BNSD batch_ stride for Q
195+ int64_t strideQON = 0; // BNSD head stride for Q
196+ int64_t strideQOS = 0; // BNSD seq stride for Q
197+ int64_t strideKVB = 0; // BNSD batch_ stride for KV
198+ int64_t strideKVN = 0; // BNSD head stride for KV
199+ int64_t strideKVS = 0; // BNSD seq stride for KV
200+ if constexpr (qFormat == Format::BNSD) {
201+ strideQOB = qHeads_ * qSeqlenAligned_ * embed_; // batch_ stride
202+ strideQON = qSeqlenAligned_ * embed_; // head stride
203+ strideQOS = embed_; // seq stride
204+ } else if constexpr (qFormat == Format::TND) {
205+ strideQO = qHeads_ * embed_;
206+ }
207+ if constexpr (kvFormat == Format::BNSD) {
208+ strideKVB = kvHeads_ * kvSeqlenAligned_ * embed_; // batch_ stride
209+ strideKVN = kvSeqlenAligned_ * embed_; // head stride
210+ strideKVS = embed_; // seq stride
211+ } else if constexpr (kvFormat == Format::TND) {
212+ strideKV = kvHeads_ * embed_;
213+ }
214+ uint32_t embedRound = RoundUp(embed_, 16);
215+ uint32_t groupSize = qHeads_ / kvHeads_;
216+ int qkStaticLoopNum = GetQkStaticLoopNum();
217+ int pvStaticNLoopNum = GetPvStaticNLoopNum();
218+ int64_t qBOffset = 0;
219+ int64_t kBOffset = 0;
220+ int64_t vBOffset = 0;
221+ int64_t oBOffset = 0;
222+ uint32_t preTotalTaskNum = 0;
223+ uint32_t curBatch = 0;
224+ int64_t qSeqlen = actSeqAval_ ? gActualQseqlen.GetValue(curBatch) : qSeqlenAligned_;
225+ int64_t kvSeqlen = actSeqAval_ ? gActualKvseqlen.GetValue(curBatch) : kvSeqlenAligned_;
226+ uint32_t curQSTileNum = GetCurQSTileNum(qSeqlen, blockShapeX_, qBaseTile_);
227+ uint32_t curTotalTaskNum = firstBatchTaskNum_;
228+ // Go through each task
229+ for (uint32_t taskIdx = coreIdx; taskIdx < totalTaskNum_; taskIdx += coreNum) {
230+ while (taskIdx >= curTotalTaskNum) {
231+ ++curBatch;
232+ preTotalTaskNum = curTotalTaskNum;
233+ if constexpr (qFormat == Format::TND) {
234+ qBOffset += qSeqlen * strideQO;
235+ oBOffset += qSeqlen * strideQO;
236+ }
237+ if constexpr (kvFormat == Format::TND) {
238+ kBOffset += kvSeqlen * strideKV;
239+ vBOffset += kvSeqlen * strideKV;
240+ }
241+ qSeqlen = actSeqAval_ ? gActualQseqlen.GetValue(curBatch) : qSeqlenAligned_;
242+ kvSeqlen = actSeqAval_ ? gActualKvseqlen.GetValue(curBatch) : kvSeqlenAligned_;
243+ curQSTileNum = GetCurQSTileNum(qSeqlen, blockShapeX_, qBaseTile_);
244+ curTotalTaskNum += curQSTileNum * qHeads_;
245+ }
246+ uint32_t taskIdxCurBatch = taskIdx - preTotalTaskNum;
247+ uint32_t qHeadIdx = taskIdxCurBatch / curQSTileNum;
248+ uint32_t kvHeadIdx = qHeadIdx / groupSize;
249+ uint32_t qSTileIdx = taskIdxCurBatch - qHeadIdx * curQSTileNum;
250+ // corresponding xBlock index of cur task
251+ uint32_t xBlockIdx = qSTileIdx / qSTileNumPerFullXBlock;
252+ // the q base tile index within the corrsponding xBlock of cur task
253+ uint32_t qSTileIdxCurXBlock = qSTileIdx - xBlockIdx * qSTileNumPerFullXBlock;
254+ // corresponding head index of cur task
255+ // corresponding blockSparseMask gm offset of cur task
256+ // gmBlockSparseMask has the shape [B, qN, xBlockNumAligned, yBlockNumAligned]
257+ int64_t sparseMaskBOffset = curBatch * qHeads_ * xBlockNumAligned_ * yBlockNumAligned_;
258+ int64_t sparseMaskNOffset = qHeadIdx * xBlockNumAligned_ * yBlockNumAligned_;
259+ int64_t sparseMaskXOffset = xBlockIdx * yBlockNumAligned_;
260+ int64_t gmOffsetSparseMask = sparseMaskBOffset + sparseMaskNOffset + sparseMaskXOffset;
261+ // corresponding Q/K/V/O gm offset of cur task
262+ int64_t gmOffsetQ = 0;
263+ int64_t gmOffsetK = 0;
264+ int64_t gmOffsetV = 0;
265+ int64_t gmOffsetO = 0;
266+ int64_t qSOffset = xBlockIdx * blockShapeX_ + qSTileIdxCurXBlock * qBaseTile_;
267+ if constexpr (qFormat == Format::BNSD) {
268+ qBOffset = curBatch * strideQOB;
269+ oBOffset = curBatch * strideQOB;
270+ gmOffsetQ = qBOffset + qHeadIdx * strideQON + qSOffset * strideQOS;
271+ gmOffsetO = oBOffset + qHeadIdx * strideQON + qSOffset * strideQOS;
272+ } else if constexpr (qFormat == Format::TND) {
273+ gmOffsetQ = qBOffset + qSOffset * strideQO + qHeadIdx * embed_;
274+ gmOffsetO = oBOffset + qSOffset * strideQO + qHeadIdx * embed_;
275+ }
276+ if constexpr (kvFormat == Format::BNSD) {
277+ kBOffset = curBatch * strideKVB;
278+ vBOffset = curBatch * strideKVB;
279+ gmOffsetK = kBOffset + kvHeadIdx * strideKVN;
280+ gmOffsetV = vBOffset + kvHeadIdx * strideKVN;
281+ } else if constexpr (kvFormat == Format::TND) {
282+ gmOffsetK = kBOffset + kvHeadIdx * embed_;
283+ gmOffsetV = vBOffset + kvHeadIdx * embed_;
284+ }
285+ // the actual x block num of cur batch_, calc by actual qseqlen
286+ uint32_t xBlockNumAval = static_cast<uint32_t>(CeilDiv(qSeqlen, static_cast<int64_t>(blockShapeX_)));
287+ uint32_t xBlockSize = (xBlockIdx == xBlockNumAval - 1) ?
288+ (qSeqlen - xBlockIdx * blockShapeX_) : blockShapeX_;
289+ uint32_t qSTileNumCurXBlock = CeilDiv(xBlockSize, qBaseTile_);
290+ uint32_t qSTileSizeAct = (qSTileIdxCurXBlock == qSTileNumCurXBlock - 1) ?
291+ (xBlockSize - qSTileIdxCurXBlock * qBaseTile_) : qBaseTile_;
292+ // calc the gathered kvS from sparse mask
293+ uint32_t gmOffsetSparseCount =
294+ curBatch * qHeads_ * xBlockNumAligned_ + qHeadIdx * xBlockNumAligned_ + xBlockIdx;
295+ uint32_t yBlockNumRsvd = gSparseCount.GetValue(gmOffsetSparseCount);
296+ if (yBlockNumRsvd == 0) {
297+ continue;
298+ }
299+ uint32_t gmOffsetSparseIdx = gmOffsetSparseCount * yBlockNumAligned_;
300+ uint32_t lastIdxOffset = gmOffsetSparseIdx + yBlockNumRsvd - 1;
301+ uint32_t lastSparseIdx = gSparseIdx.GetValue(lastIdxOffset);
302+ 
303+ uint32_t yBlockNumAval = static_cast<uint32_t>(CeilDiv(kvSeqlen, static_cast<int64_t>(blockShapeY_)));
304+ uint32_t lastYBlockSize = (lastSparseIdx == yBlockNumAval - 1) ?
305+ kvSeqlen - lastSparseIdx * blockShapeY_ : blockShapeY_;
306+ int64_t gatheredKvSeqlen = (yBlockNumRsvd - 1) * blockShapeY_ + lastYBlockSize;
307+ // the rowNum of cur task
308+ // no qS*qN combination even in GQA/MQA senario, since each qN has a different sparse pattern
309+ uint32_t rowNum = qSTileSizeAct;
310+ uint32_t rowNumRound = RoundUp(rowNum, std::is_same_v<LayoutQ, layout::RowMajor> ? 16 : 64);
311+ uint32_t kvSLoopNum = static_cast<uint32_t>(CeilDiv(gatheredKvSeqlen, static_cast<int64_t>(kvBaseTile_)));
312+ uint32_t kvSTileSizeAct = kvBaseTile_;
313+#ifdef __DAV_CUBE__
314+ uint32_t kvShapeCol = 0;
315+ uint32_t qShapeCol = 0;
316+ if constexpr (qFormat == Format::BNSD) {
317+ qShapeCol = strideQOS;
318+ } else if constexpr (qFormat == Format::TND) {
319+ qShapeCol = strideQO;
320+ }
321+ if constexpr (kvFormat == Format::BNSD) {
322+ kvShapeCol = strideKVS;
323+ } else if constexpr (kvFormat == Format::TND) {
324+ kvShapeCol = strideKV;
325+ }
326+ 
327+ GemmCoord actualBlockShapeQ{rowNum, embed_, 0};
328+ if constexpr (std::is_same_v<LayoutQ, layout::RowMajor>) {
329+ auto gmQLayoutTla = tla::MakeLayout<ElementQ, LayoutQ>(qBaseTile_, qShapeCol);
330+ auto gmQTensorTla = tla::MakeTensor(gQ[gmOffsetQ], gmQLayoutTla, Arch::PositionGM{});
331+ blockMmadQK.loadQGM(gmQTensorTla, actualBlockShapeQ);
332+ } else {
333+ auto gmQLayoutTla = tla::MakeLayout<ElementQ, LayoutQ>(qShapeCol, qBaseTile_);
334+ auto gmQTensorTla = tla::MakeTensor(gQ[gmOffsetQ], gmQLayoutTla, Arch::PositionGM{});
335+ blockMmadQK.loadQGM(gmQTensorTla, actualBlockShapeQ);
336+ }
337+ 
338+ auto gmKLayoutTla = tla::MakeLayout<ElementK, LayoutK>(kvSeqlenAligned_, kvShapeCol);
339+ auto gmKTensorTla = tla::MakeTensor(gK[gmOffsetK], gmKLayoutTla, Arch::PositionGM{});
340+ 
341+ auto gmVLayoutTla = tla::MakeLayout<ElementV, LayoutV>(kvBaseTile_, kvShapeCol);
342+ auto gmVTensorTla = tla::MakeTensor(gV[gmOffsetV], gmVLayoutTla, Arch::PositionGM{});
343+#endif
344+#ifdef __DAV_VEC__
345+ uint32_t oShapeCol = 0;
346+ if constexpr (qFormat == Format::BNSD) {
347+ oShapeCol = strideQOS;
348+ } else if constexpr (qFormat == Format::TND) {
349+ oShapeCol = strideQO;
350+ }
351+#endif
352+ for (uint32_t gatheredKvSTileIdx = 0; gatheredKvSTileIdx < kvSLoopNum + PRE_LAUNCH; gatheredKvSTileIdx++) {
353+ if (gatheredKvSTileIdx < kvSLoopNum) {
354+ if (gatheredKvSTileIdx == kvSLoopNum - 1) {
355+ kvSTileSizeAct = gatheredKvSeqlen - gatheredKvSTileIdx * kvBaseTile_;
356+ } else {
357+ kvSTileSizeAct = kvBaseTile_;
358+ }
359+ // QK
360+ GemmCoord actualBlockShapeQK{rowNum, kvSTileSizeAct, embed_};
361+ uint32_t ubSBufId = gatheredKvSTileIdx % UB_S_OTMP_BUF_STAGES;
362+ auto ubSLayoutTla = tla::MakeLayout<ElementS, LayoutS>(RoundUp(kvSTileSizeAct, 16)*rowNumRound >> 5, 32);
363+ auto ubSTensorTla = tla::MakeTensor(ubSTensor[ubSBufId],
364+ ubSLayoutTla, Arch::PositionUB{});
365+ uint32_t Mm1ToSmFlagId = ubSBufId;
366+ Arch::CrossCoreFlag mm1ToSmFlag(Mm1ToSmFlagId);
367+#ifdef __DAV_CUBE__
368+ uint64_t prefixSumL0AStages = CalcCrossMm1Mm2PrefixSumL0ABStages(
369+ gatheredKvSTileIdx, mm1L0ATotalStages_, mm2L0ATotalStages_, kvSLoopNum, true);
370+ uint64_t prefixSumL0BStages = CalcCrossMm1Mm2PrefixSumL0ABStages(
371+ gatheredKvSTileIdx, mm1L0BTotalStages_, mm2L0BTotalStages_, kvSLoopNum, true);
372+ if (qkStaticLoopNum == 1) {
373+ blockMmadQK.template operator()<quant_mode, 1, 1>(
374+ gmKTensorTla, ubSTensorTla, gSparseIdx[gmOffsetSparseIdx],
375+ actualBlockShapeQK,
376+ gatheredKvSTileIdx, kvSeqlen,
377+ kvBaseTile_, blockShapeY_, yBlockNumAval, yBlockNumRsvd,
378+ prefixSumL0AStages, prefixSumL0BStages,
379+ mm1ToSmFlag, scaleValue_);
380+ } else {
381+ blockMmadQK.template operator()<quant_mode, -1, -1>(
382+ gmKTensorTla, ubSTensorTla, gSparseIdx[gmOffsetSparseIdx],
383+ actualBlockShapeQK,
384+ gatheredKvSTileIdx, kvSeqlen,
385+ kvBaseTile_, blockShapeY_, yBlockNumAval, yBlockNumRsvd,
386+ prefixSumL0AStages, prefixSumL0BStages,
387+ mm1ToSmFlag, scaleValue_);
388+ }
389+ if (gatheredKvSTileIdx == kvSLoopNum - 1)
390+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
391+#endif
392+ // SM
393+ uint32_t l1PBufId = gatheredKvSTileIdx % pL1BufNum_;
394+ uint32_t smToMm2FlagId = l1PBufId + UB_S_OTMP_BUF_STAGES;
395+ Arch::CrossCoreFlag smToMm2Flag(smToMm2FlagId);
396+ auto l1PLayoutTla = tla::MakeLayout<ElementP, NpuArch::layout::zN>(rowNumRound, kvSTileSizeAct);
397+ auto l1PTensorTla = tla::MakeTensor(l1PTensor[l1PBufId],
398+ l1PLayoutTla, Arch::PositionL1{});
399+#ifdef __DAV_VEC__
400+ uint32_t qsOffset = 0;
401+ uint32_t ksOffset = 0;
402+ uint32_t vsOffset = curBatch * kvHeads_ * embed_ + kvHeadIdx * embed_;
403+ if (gatheredKvSTileIdx == 0) {
404+ AscendC::DataCopy(scaleTensor, valueScale[vsOffset], embed_);
405+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID7);
406+ }
407+ if constexpr (qFormat == Format::BNSD) {
408+ uint32_t qScaleSeqNum = CeilDiv(qSeqlenAligned_,QUANT_BLOCK_SIZE);
409+ qsOffset = curBatch * qHeads_ * qScaleSeqNum + qHeadIdx * qScaleSeqNum + CeilDiv(qSOffset,QUANT_BLOCK_SIZE);
410+ } else if constexpr (qFormat == Format::TND) {
411+ qsOffset = CeilDiv((qBOffset / embed_), QUANT_BLOCK_SIZE) + CeilDiv(qSOffset, QUANT_BLOCK_SIZE) * qHeads_ + qHeadIdx;
412+ }
413+ if constexpr (kvFormat == Format::BNSD) {
414+ uint32_t kScaleSeqNum = CeilDiv(kvSeqlenAligned_, QUANT_BLOCK_SIZE);
415+ ksOffset = curBatch * kvHeads_ * kScaleSeqNum + kvHeadIdx * kScaleSeqNum;
416+ } else if constexpr (kvFormat == Format::TND) {
417+ ksOffset = CeilDiv(kBOffset / embed_, QUANT_BLOCK_SIZE) + kvHeadIdx;
418+ }
419+ uint32_t qScaleStride = (qFormat == Format::BNSD) ? 1 : qHeads_;
420+ uint32_t kScaleStride = (kvFormat == Format::BNSD) ? 1 : kvHeads_;
421+ auto gmQSLayoutTla = tla::MakeLayout<float, layout::RowMajor>(
422+ RoundUp(rowNum, QUANT_BLOCK_SIZE) / QUANT_BLOCK_SIZE, qScaleStride);
423+ auto gmQSTensorTla = tla::MakeTensor(queryScale[qsOffset],
424+ gmQSLayoutTla, Arch::PositionGM{});
425+ auto gmKSLayoutTla = tla::MakeLayout<float, layout::RowMajor>(
426+ RoundUp(kvSeqlen, QUANT_BLOCK_SIZE) / QUANT_BLOCK_SIZE, kScaleStride);
427+ auto gmKSTensorTla = tla::MakeTensor(keyScale[ksOffset],
428+ gmKSLayoutTla, Arch::PositionGM{});
429+ int32_t sparseTableStartOffset = gmOffsetSparseIdx + (gatheredKvSTileIdx) * (kvBaseTile_/blockShapeY_);
430+ epilogueOnlineSoftmax.template operator()<quant_mode>(
431+ l1PTensorTla,
432+ gmQSTensorTla,
433+ gmKSTensorTla,
434+ actualBlockShapeQK,
435+ (gatheredKvSTileIdx == 0),
436+ ubSBufId,
437+ l1PBufId,
438+ mm1ToSmFlag,
439+ smToMm2Flag,
440+ gSparseIdx[sparseTableStartOffset]
441+ );
442+ if (gatheredKvSTileIdx == 0) {
443+ static constexpr float fp8MaxReciprocal = 1.0f / 448.0f; // exp((half)ln(448))
444+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID7);
445+ AscendC::Muls(scaleTensor[embed_], scaleTensor, fp8MaxReciprocal, embed_);
446+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID7);
447+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID7);
448+ }
449+ AscendC::PipeBarrier<PIPE_V>();
450+#endif
451+ }
452+ if (gatheredKvSTileIdx >= PRE_LAUNCH) {
453+ uint32_t gatheredKvSTileIdxDe = gatheredKvSTileIdx - PRE_LAUNCH;
454+ if (gatheredKvSTileIdxDe == kvSLoopNum - 1) {
455+ kvSTileSizeAct = gatheredKvSeqlen - gatheredKvSTileIdxDe * kvBaseTile_;
456+ } else {
457+ kvSTileSizeAct = kvBaseTile_;
458+ }
459+ // PV
460+ GemmCoord actualBlockShapePV{rowNumRound, embed_, kvSTileSizeAct};
461+ uint32_t ubOTmpBufId = gatheredKvSTileIdxDe % UB_S_OTMP_BUF_STAGES;
462+ // 核间同步flagId规律：每份dst对应一个id，从小到大顺序依次为qk->sm, sm->pv, pv->re
463+ uint32_t Mm2ToReFlagId = ubOTmpBufId + UB_S_OTMP_BUF_STAGES + pL1BufNum_;
464+#ifdef __DAV_CUBE__
465+ uint32_t l1PBufId = gatheredKvSTileIdxDe % pL1BufNum_;
466+ auto ubOTmpLayoutTla = tla::MakeLayout<ElementOTmp, LayoutOTmp>(rowNumRound, embedRound);
467+ auto ubOTmpTensorTla = tla::MakeTensor(ubOTmpTensor[ubOTmpBufId],
468+ ubOTmpLayoutTla, Arch::PositionUB{});
469+ uint32_t smToMm2FlagId = l1PBufId + UB_S_OTMP_BUF_STAGES;
470+
471+ Arch::CrossCoreFlag smToMm2Flag(smToMm2FlagId);
472+ Arch::CrossCoreFlag mm2ToReFlag(Mm2ToReFlagId);
473+ uint64_t prefixSumL0AStages = CalcCrossMm1Mm2PrefixSumL0ABStages(
474+ gatheredKvSTileIdxDe, mm1L0ATotalStages_, mm2L0ATotalStages_, kvSLoopNum, false);
475+ uint64_t prefixSumL0BStages = CalcCrossMm1Mm2PrefixSumL0ABStages(
476+ gatheredKvSTileIdxDe, mm1L0BTotalStages_, mm2L0BTotalStages_, kvSLoopNum, false);
477+ uint32_t vsOffset = curBatch * kvHeads_ * embed_ + kvHeadIdx * embed_;
478+ auto gmVSLayoutTla = tla::MakeLayout<float, layout::RowMajor>(
479+ 1, embed_);
480+ auto gmVSTensorTla = tla::MakeTensor(valueScale[vsOffset],
481+ gmVSLayoutTla, Arch::PositionGM{});
482+ if (pvStaticNLoopNum == 1) {
483+ blockMmadPV.template operator()<quant_mode, 1, 1>(
484+ gmVTensorTla, ubOTmpTensorTla, gmVSTensorTla, gSparseIdx[gmOffsetSparseIdx],
485+ actualBlockShapePV,
486+ gatheredKvSTileIdxDe, kvSeqlen,
487+ kvBaseTile_, blockShapeY_, yBlockNumAval, yBlockNumRsvd,
488+ prefixSumL0AStages, prefixSumL0BStages,
489+ smToMm2Flag, mm2ToReFlag);
490+ } else if (pvStaticNLoopNum == 2) {
491+ blockMmadPV.template operator()<quant_mode, 1, 2>(
492+ gmVTensorTla, ubOTmpTensorTla, gmVSTensorTla, gSparseIdx[gmOffsetSparseIdx],
493+ actualBlockShapePV,
494+ gatheredKvSTileIdxDe, kvSeqlen,
495+ kvBaseTile_, blockShapeY_, yBlockNumAval, yBlockNumRsvd,
496+ prefixSumL0AStages, prefixSumL0BStages,
497+ smToMm2Flag, mm2ToReFlag);
498+ } else {
499+ blockMmadPV.template operator()<quant_mode, -1, -1>(
500+ gmVTensorTla, ubOTmpTensorTla, gmVSTensorTla, gSparseIdx[gmOffsetSparseIdx],
501+ actualBlockShapePV,
502+ gatheredKvSTileIdxDe, kvSeqlen,
503+ kvBaseTile_, blockShapeY_, yBlockNumAval, yBlockNumRsvd,
504+ prefixSumL0AStages, prefixSumL0BStages,
505+ smToMm2Flag, mm2ToReFlag);
506+ }
507+#endif
508+#ifdef __DAV_VEC__
509+ // rescale O
510+ Arch::CrossCoreFlag mm2ToReFlag(Mm2ToReFlagId);
511+ uint32_t curTileMod = gatheredKvSTileIdxDe % (PRE_LAUNCH + 1);
512+ epilogueRescaleO(
513+ gO[gmOffsetO], (GemmCoord){rowNum, embed_, kvSTileSizeAct},
514+ curTileMod, gatheredKvSTileIdxDe,
515+ (gatheredKvSTileIdxDe == 0),
516+ (gatheredKvSTileIdxDe == kvSLoopNum - 1),
517+ mm2ToReFlag);
518+#endif
519+ }
520+ }
521+ }
522+ // release reverse sync flags
523+ ReleaseSyncFlags<4, 4, 4>();
524+ }
525+ 
526+ __aicore__ inline
527+ void FetchBaseShapeInfo(__gm__ EagleQuantBlockSparseAttentionTilingData *bsaTilingData)
528+ {
529+ batch_ = bsaTilingData->batch;
530+ qHeads_ = bsaTilingData->numHeads;
531+ kvHeads_ = bsaTilingData->kvHeads;
532+ embed_ = bsaTilingData->embeddingSize;
533+ firstBatchTaskNum_ = bsaTilingData->firstBatchTaskNum;
534+ totalTaskNum_ = bsaTilingData->totalTaskNum;
535+ blockShapeX_ = bsaTilingData->blockShapeX;
536+ blockShapeY_ = bsaTilingData->blockShapeY;
537+ scaleValue_ = bsaTilingData->scaleValue;
538+ // mask2idx tile info
539+ xBlockNumAligned_ = bsaTilingData->BsaMask2IdxTileInfo.xBlockNumAligned;
540+ yBlockNumAligned_ = bsaTilingData->BsaMask2IdxTileInfo.yBlockNumAligned;
541+ avgRowPerSubCore_ = bsaTilingData->BsaMask2IdxTileInfo.avgRowPerSubCore;
542+ preActiveSubCoreNum_ = bsaTilingData->BsaMask2IdxTileInfo.preActiveSubCoreNum;
543+ // base tile info
544+ qBaseTile_ = bsaTilingData->BsaBaseTileInfo.qBaseTile;
545+ kvBaseTile_ = bsaTilingData->BsaBaseTileInfo.kvBaseTile;
546+ // whether actual seqlen is provided
547+ actSeqAval_ = (!bsaTilingData->useUniformQSeqlen) && (!bsaTilingData->useUniformKvSeqlen);
548+ sparseIdxSize_ = bsaTilingData->selectIdxSize;
549+ // aligned seqlen q & kv
550+ qSeqlenAligned_ = bsaTilingData->maxQSeqlen;
551+ kvSeqlenAligned_ = bsaTilingData->maxKvSeqlen;
552+ }
553+ 
554+ __aicore__ inline
555+ void CalcOnChipBufTileInfo(__gm__ EagleQuantBlockSparseAttentionTilingData *bsaTilingData)
556+ {
557+ mm1L1TileM_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm1L1TileM;
558+ mm1L1TileN_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm1L1TileN;
559+ mm1L1TileKLeft_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm1L1TileKLeft;
560+ mm1L1TileKRight_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm1L1TileKRight;
561+ mm2L1TileM_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm2L1TileM;
562+ mm2L1TileN_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm2L1TileN;
563+ mm2L1TileKLeft_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm2L1TileKLeft;
564+ mm2L1TileKRight_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm2L1TileKRight;
565+ qL1BufNum_ = bsaTilingData->BsaMmPhaseL1TileInfo.qL1BufNum;
566+ kL1BufNum_ = bsaTilingData->BsaMmPhaseL1TileInfo.kL1BufNum;
567+ vL1BufNum_ = bsaTilingData->BsaMmPhaseL1TileInfo.vL1BufNum;
568+ pL1BufNum_ = bsaTilingData->BsaMmPhaseL1TileInfo.pL1BufNum;
569+ Gemm::Block::Mm1L1TileHelper mm1L1TileHelper(mm1L1TileM_, mm1L1TileN_, mm1L1TileKLeft_, mm1L1TileKRight_,
570+ qL1BufNum_, kL1BufNum_);
571+ mm1L1TileHelper_ = mm1L1TileHelper;
572+ Gemm::Block::Mm2L1TileHelper mm2L1TileHelper(mm2L1TileM_, mm2L1TileN_, mm2L1TileKLeft_, mm2L1TileKRight_,
573+ pL1BufNum_, vL1BufNum_);
574+ mm2L1TileHelper_ = mm2L1TileHelper;
575+ mm2L1AddrStart_ = mm1L1TileM_ * mm1L1TileKLeft_ * qL1BufNum_ * sizeof(ElementQ) +
576+ mm1L1TileKRight_ * mm1L1TileN_ * kL1BufNum_ * sizeof(ElementK);
577+ mm1L0ATotalStages_ = (qBaseTile_ / BlockMmadQK::L0_TILE_M) * (embed_ / BlockMmadQK::L0_TILE_K);
578+ mm1L0BTotalStages_ = (kvBaseTile_ / BlockMmadQK::L0_TILE_N) * (embed_ / BlockMmadQK::L0_TILE_K);
579+ mm2L0ATotalStages_ = (qBaseTile_ / BlockMmadPV::L0_TILE_M) * (kvBaseTile_ / BlockMmadPV::L0_TILE_K);
580+ mm2L0BTotalStages_ = (kvBaseTile_ / BlockMmadPV::L0_TILE_K) * (embed_ / BlockMmadPV::L0_TILE_N);
581+ }
582+ 
583+ __aicore__ inline
584+ uint64_t CalcCrossMm1Mm2PrefixSumL0ABStages(
585+ uint32_t gatheredKvSTileIdx, uint32_t singleMm1L0Stages,
586+ uint32_t singleMm2L0Stages, uint32_t kvSLoopNum,
587+ bool isCurPhaseMm1)
588+ {
589+ uint64_t prefixSumStages;
590+ if (isCurPhaseMm1) {
591+ prefixSumStages = (gatheredKvSTileIdx <= PRE_LAUNCH) ?
592+ gatheredKvSTileIdx * singleMm1L0Stages :
593+ gatheredKvSTileIdx * singleMm1L0Stages + (gatheredKvSTileIdx - PRE_LAUNCH) * singleMm2L0Stages;
594+ } else {
595+ prefixSumStages = (gatheredKvSTileIdx < kvSLoopNum - PRE_LAUNCH) ?
596+ (gatheredKvSTileIdx + 1 + PRE_LAUNCH) * singleMm1L0Stages + gatheredKvSTileIdx * singleMm2L0Stages:
597+ kvSLoopNum * singleMm1L0Stages + gatheredKvSTileIdx * singleMm2L0Stages;
598+ }
599+ return prefixSumStages;
600+ }
601+ 
602+ __aicore__ inline
603+ void InitCrossCoreDstBuf(
604+ AscendC::LocalTensor<ElementP> (&l1PTensor)[MAX_CROSS_CORE_BUF_STAGES],
605+ AscendC::LocalTensor<ElementS> (&ubSTensor)[UB_S_OTMP_BUF_STAGES],
606+ AscendC::LocalTensor<ElementOTmp> (&ubOTmpTensor)[UB_S_OTMP_BUF_STAGES])
607+ {
608+ for (uint32_t i = 0; i < pL1BufNum_; i++) {
609+ l1PTensor[i] = resource.l1Buf.template GetBufferByByte<ElementP>(
610+ mm2L1AddrStart_ + mm2L1TileM_ * mm2L1TileKLeft_ * sizeof(ElementP) * i);
611+ }
612+ uint32_t rowNumPerSubCore = EpilogueOnlineSoftmax::SM_ROW_MAX_ELEM_NUM;
613+ uint32_t colNumPerSubCore = EpilogueOnlineSoftmax::SM_COL_MAX_ELEM_NUM;
614+ uint32_t rescaleCol = EpilogueRescaleO::RESCALE_COL_MAX_ELEM_NUM;
615+ for (uint32_t i = 0; i < UB_S_OTMP_BUF_STAGES; i++) {
616+ ubSTensor[i] = resource.ubBuf.template GetBufferByByte<ElementS>(
617+ rowNumPerSubCore * colNumPerSubCore * sizeof(ElementS) * i);
618+ ubOTmpTensor[i] = resource.ubBuf.template GetBufferByByte<ElementOTmp>(
619+ rowNumPerSubCore * colNumPerSubCore * sizeof(ElementS) * UB_S_OTMP_BUF_STAGES +
620+ rowNumPerSubCore * rescaleCol * sizeof(ElementOTmp) * i);
621+ }
622+ }
623+ 
624+ template <uint32_t MM1_SM_MODE, uint32_t MM2_RE_MODE, uint32_t SM_MM2_MODE>
625+ __aicore__ inline
626+ void InitSyncFlags()
627+ {
628+#ifdef __DAV_CUBE__
629+ // same core sync between pipes
630+ // Query
631+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
632+ // Key
633+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID1);
634+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
635+ // Value
636+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID3);
637+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID4);
638+ // L0A
639+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
640+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID1);
641+ // L0B
642+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID2);
643+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID3);
644+ // L0C
645+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
646+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID1);
647+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID2);
648+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID3);
649+ // VSDB (L1 Value Scale Double Buffer)
650+ AscendC::SetFlag<AscendC::HardEvent::MTE2_FIX>(EVENT_ID0);
651+ AscendC::SetFlag<AscendC::HardEvent::MTE2_FIX>(EVENT_ID1);
652+ AscendC::SetFlag<AscendC::HardEvent::FIX_MTE2>(EVENT_ID0);
653+ // cross core sync
654+ if constexpr (SM_MM2_MODE == 4U) {
655+ AscendC::CrossCoreSetFlag<SM_MM2_MODE, PIPE_MTE1>(2);
656+ AscendC::CrossCoreSetFlag<SM_MM2_MODE, PIPE_MTE1>(18);
657+ AscendC::CrossCoreSetFlag<SM_MM2_MODE, PIPE_MTE1>(3);
658+ AscendC::CrossCoreSetFlag<SM_MM2_MODE, PIPE_MTE1>(19);
659+ AscendC::CrossCoreSetFlag<SM_MM2_MODE, PIPE_MTE1>(4);
660+ AscendC::CrossCoreSetFlag<SM_MM2_MODE, PIPE_MTE1>(20);
661+ }
662+#endif
663+#ifdef __DAV_VEC__
664+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
665+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID1);
666+ // mask2index
667+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
668+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID1);
669+ // softmax
670+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID2);
671+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);
672+ // rescale
673+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
674+ 
675+ AscendC::SetFlag<AscendC::HardEvent::S_MTE3>(EVENT_ID0);
676+ AscendC::SetFlag<AscendC::HardEvent::MTE3_S>(EVENT_ID0);
677+ AscendC::SetFlag<AscendC::HardEvent::MTE3_S>(EVENT_ID1);
678+ if constexpr (MM1_SM_MODE == 4U) {
679+ AscendC::CrossCoreSetFlag<MM1_SM_MODE, PIPE_V>(0);
680+ AscendC::CrossCoreSetFlag<MM1_SM_MODE, PIPE_V>(1);
681+ }
682+ if constexpr (MM2_RE_MODE == 4U) {
683+ AscendC::CrossCoreSetFlag<MM2_RE_MODE, PIPE_V>(5);
684+ AscendC::CrossCoreSetFlag<MM2_RE_MODE, PIPE_V>(6);
685+ }
686+#endif
687+ }
688+ 
689+ template <uint32_t MM1_SM_MODE, uint32_t MM2_RE_MODE, uint32_t SM_MM2_MODE>
690+ __aicore__ inline
691+ void ReleaseSyncFlags()
692+ {
693+#ifdef __DAV_CUBE__
694+ // same core sync between pipes
695+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
696+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID1);
697+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
698+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID3);
699+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID4);
700+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
701+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID1);
702+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID2);
703+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID3);
704+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
705+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID1);
706+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID2);
707+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID3);
708+ // VSDB (L1 Value Scale Double Buffer)
709+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_FIX>(EVENT_ID0);
710+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_FIX>(EVENT_ID1);
711+ AscendC::WaitFlag<AscendC::HardEvent::FIX_MTE2>(EVENT_ID0);
712+ if constexpr (MM1_SM_MODE == 4U) {
713+ AscendC::CrossCoreWaitFlag<MM1_SM_MODE, PIPE_FIX>(0);
714+ AscendC::CrossCoreWaitFlag<MM1_SM_MODE, PIPE_FIX>(1);
715+ AscendC::CrossCoreWaitFlag<MM1_SM_MODE, PIPE_FIX>(16);
716+ AscendC::CrossCoreWaitFlag<MM1_SM_MODE, PIPE_FIX>(17);
717+ }
718+ if constexpr (MM2_RE_MODE == 4U) {
719+ AscendC::CrossCoreWaitFlag<MM2_RE_MODE, PIPE_FIX>(5);
720+ AscendC::CrossCoreWaitFlag<MM2_RE_MODE, PIPE_FIX>(21);
721+ AscendC::CrossCoreWaitFlag<MM2_RE_MODE, PIPE_FIX>(6);
722+ AscendC::CrossCoreWaitFlag<MM2_RE_MODE, PIPE_FIX>(22);
723+ }
724+#endif
725+#ifdef __DAV_VEC__
726+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
727+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID1);
728+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
729+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID1);
730+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID2);
731+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);
732+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
733+ AscendC::WaitFlag<AscendC::HardEvent::S_MTE3>(EVENT_ID0);
734+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_S>(EVENT_ID0);
735+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_S>(EVENT_ID1);
736+ if constexpr (SM_MM2_MODE == 4U) {
737+ AscendC::CrossCoreWaitFlag<SM_MM2_MODE, PIPE_MTE3>(2);
738+ AscendC::CrossCoreWaitFlag<SM_MM2_MODE, PIPE_MTE3>(3);
739+ AscendC::CrossCoreWaitFlag<SM_MM2_MODE, PIPE_MTE3>(4);
740+ }
741+#endif
742+ AscendC::PipeBarrier<PIPE_ALL>();
743+ }
744+ 
745+private:
746+ Arch::Resource<ArchTag> resource;
747+ /*
748+ tiling info, which are const in each kernel launch
749+ */
750+ // basic shape info
751+ uint32_t batch_;
752+ uint32_t qHeads_;
753+ uint32_t kvHeads_;
754+ uint32_t embed_;
755+ uint32_t firstBatchTaskNum_;
756+ uint32_t totalTaskNum_;
757+ uint32_t blockShapeX_;
758+ uint32_t blockShapeY_;
759+ float scaleValue_;
760+ // mask2idx tile info
761+ uint32_t xBlockNumAligned_;
762+ uint32_t yBlockNumAligned_;
763+ uint32_t avgRowPerSubCore_;
764+ uint32_t preActiveSubCoreNum_;
765+ // base tile info
766+ uint32_t qBaseTile_;
767+ uint32_t kvBaseTile_;
768+ // whether actual seqlen is provided
769+ uint32_t actSeqAval_;
770+ // workspace size
771+ uint64_t sparseIdxSize_;
772+ // aligned seqlen q & kv
773+ int64_t qSeqlenAligned_;
774+ int64_t kvSeqlenAligned_;
775+ // L1 tile info
776+ uint32_t mm1L1TileM_;
777+ uint32_t mm1L1TileN_;
778+ uint32_t mm1L1TileKLeft_;
779+ uint32_t mm1L1TileKRight_;
780+ uint32_t mm2L1TileM_;
781+ uint32_t mm2L1TileN_;
782+ uint32_t mm2L1TileKLeft_;
783+ uint32_t mm2L1TileKRight_;
784+ uint32_t qL1BufNum_;
785+ uint32_t kL1BufNum_;
786+ uint32_t vL1BufNum_;
787+ uint32_t pL1BufNum_;
788+ uint32_t mm1L0ATotalStages_;
789+ uint32_t mm1L0BTotalStages_;
790+ uint32_t mm2L0ATotalStages_;
791+ uint32_t mm2L0BTotalStages_;
792+ uint32_t mm2L1AddrStart_ = 0;
793+ Gemm::Block::Mm1L1TileHelper mm1L1TileHelper_;
794+ Gemm::Block::Mm2L1TileHelper mm2L1TileHelper_;
795+};
796+ 
797+}
798+ 
799+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/arch35/eagle_quant_block_sparse_attention_kernel_arch35_regular.h+722-0
@@ -0,0 +1,722 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef BSA_ARCH35_KERNEL_REGULAR
12+#define BSA_ARCH35_KERNEL_REGULAR
13+ 
14+#include "../arch35/kernel_utils.hpp"
15+ 
16+ 
17+using namespace NpuArch;
18+using namespace tla;
19+ 
20+namespace BsaKernelArch35 {
21+ 
22+template <
23+ class EpilogueMask2Idx,
24+ class BlockMmadQK,
25+ class EpilogueOnlineSoftmax,
26+ class BlockMmadPV,
27+ class EpilogueRescaleO,
28+ Format qFormat,
29+ Format kvFormat>
30+class BsaRegularKernelArch35 {
31+public:
32+ using ArchTag = typename BlockMmadPV::ArchTag;
33+
34+ using ElementQ = typename BlockMmadQK::ElementA;
35+ using ElementK = typename BlockMmadQK::ElementB;
36+ using ElementS = typename EpilogueOnlineSoftmax::ElementInput;
37+ using ElementP = typename BlockMmadPV::ElementA;
38+ using ElementV = typename BlockMmadPV::ElementB;
39+ using ElementOTmp = typename BlockMmadPV::ElementC;
40+ using ElementO = typename EpilogueRescaleO::ElementO;
41+ using ElementSparseMask = typename EpilogueMask2Idx::ElementSparseMask;
42+ using ElementSparseIdx = typename EpilogueMask2Idx::ElementSparseIdx;
43+ using ElementSparseCount = typename EpilogueMask2Idx::ElementSparseCount;
44+ 
45+ using LayoutQ = layout::RowMajor;
46+ using LayoutK = layout::ColumnMajor;
47+ using LayoutS = layout::RowMajor;
48+ using LayoutP = layout::RowMajor;
49+ using LayoutV = layout::RowMajor;
50+ using LayoutO = layout::RowMajor;
51+ using LayoutOTmp = layout::RowMajor;
52+ using LayoutSparseIdx = layout::RowMajor;
53+ using LayoutSparseCount = layout::RowMajor;
54+ 
55+ using LayoutTagL1P = typename BlockMmadPV::LayoutTagL1A;
56+ 
57+ static constexpr uint32_t PRE_LAUNCH = 2;
58+ static constexpr uint32_t MAX_CROSS_CORE_BUF_STAGES = PRE_LAUNCH + 1;
59+ static constexpr uint32_t UB_S_OTMP_BUF_STAGES = 2;
60+ // Methods
61+ __aicore__ inline
62+ BsaRegularKernelArch35() {}
63+ 
64+ __aicore__ inline
65+ int GetQkStaticLoopNum()
66+ {
67+ return ((qBaseTile_ == BlockMmadQK::L0_TILE_M) &&
68+ (embed_ == BlockMmadQK::L0_TILE_K) &&
69+ ((kvBaseTile_ == BlockMmadQK::L0_TILE_N) ||
70+ (kvBaseTile_ == BlockMmadQK::L0_TILE_N * 2U))) ?
71+ 1 : Gemm::Block::Arch35MmadOpt::DYNAMIC_LOOP;
72+ }
73+ 
74+ __aicore__ inline
75+ int GetPvStaticNLoopNum()
76+ {
77+ if (qBaseTile_ != BlockMmadPV::L0_TILE_M) {
78+ return Gemm::Block::Arch35MmadOpt::DYNAMIC_LOOP;
79+ }
80+ if (embed_ == BlockMmadPV::L0_TILE_N) {
81+ return 1;
82+ }
83+ if (embed_ == BlockMmadPV::L0_TILE_N * 2U) {
84+ return 2;
85+ }
86+ return Gemm::Block::Arch35MmadOpt::DYNAMIC_LOOP;
87+ }
88+ 
89+ template<int quant_mode=0>
90+ __aicore__ inline
91+ void operator()(BsaKernelParamsArch35 const &params)
92+ {
93+ static_assert(quant_mode == 0);
94+ __gm__ EagleQuantBlockSparseAttentionTilingData *bsaTilingData =
95+ reinterpret_cast<__gm__ EagleQuantBlockSparseAttentionTilingData *>(params.tiling);
96+ FetchBaseShapeInfo(bsaTilingData);
97+ CalcOnChipBufTileInfo(bsaTilingData);
98+ // global buffers
99+ AscendC::GlobalTensor<ElementQ> gQ;
100+ gQ.SetGlobalBuffer((__gm__ ElementQ *)params.q);
101+ AscendC::GlobalTensor<ElementK> gK;
102+ gK.SetGlobalBuffer((__gm__ ElementK *)params.k);
103+ AscendC::GlobalTensor<ElementV> gV;
104+ gV.SetGlobalBuffer((__gm__ ElementV *)params.v);
105+ AscendC::GlobalTensor<int64_t> gActualQseqlen;
106+ gActualQseqlen.SetGlobalBuffer((__gm__ int64_t *)params.actualQseqlen);
107+ AscendC::GlobalTensor<int64_t> gActualKvseqlen;
108+ gActualKvseqlen.SetGlobalBuffer((__gm__ int64_t *)params.actualKvseqlen);
109+ AscendC::GlobalTensor<uint8_t> gBlockSparseMask;
110+ gBlockSparseMask.SetGlobalBuffer((__gm__ uint8_t *)params.blockSparseMask);
111+ AscendC::GlobalTensor<ElementO> gO;
112+ gO.SetGlobalBuffer((__gm__ ElementO *)params.o);
113+ AscendC::GlobalTensor<ElementSparseIdx> gSparseIdx;
114+ if (bsaTilingData->sparsePatternMode == SPARSE_PATTERN_MODE_TABLE) {
115+ gSparseIdx.SetGlobalBuffer((__gm__ ElementSparseIdx *)params.blockSparseMask);
116+ } else {
117+ gSparseIdx.SetGlobalBuffer((__gm__ ElementSparseIdx *)params.workSpace);
118+ }
119+ AscendC::GlobalTensor<ElementSparseCount> gSparseCount;
120+ gSparseCount.SetGlobalBuffer((__gm__ ElementSparseCount *)(params.workSpace + sparseIdxSize_));
121+ 
122+ // cross core data move dst buffers
123+ AscendC::LocalTensor<ElementP> l1PTensor[MAX_CROSS_CORE_BUF_STAGES];
124+ AscendC::LocalTensor<ElementS> ubSTensor[UB_S_OTMP_BUF_STAGES];
125+ AscendC::LocalTensor<ElementOTmp> ubOTmpTensor[UB_S_OTMP_BUF_STAGES];
126+ InitCrossCoreDstBuf(l1PTensor, ubSTensor, ubOTmpTensor);
127+ // core idx
128+ uint32_t coreIdx = AscendC::GetBlockIdx();
129+ uint32_t coreNum = AscendC::GetBlockNum();
130+ // set reverse sync flags
131+ InitSyncFlags<4, 4, 4>();
132+#ifdef __DAV_VEC__
133+ uint32_t totalRowNumBlockMask = batch_ * qHeads_ * xBlockNumAligned_;
134+ if (bsaTilingData->sparsePatternMode == SPARSE_PATTERN_MODE_TABLE) {
135+ SparseTable2Count(resource, gSparseIdx, gSparseCount,
136+ totalRowNumBlockMask, yBlockNumAligned_, avgRowPerSubCore_, preActiveSubCoreNum_);
137+ } else {
138+ EpilogueMask2Idx epilogueMask2Idx(resource);
139+ epilogueMask2Idx(
140+ gBlockSparseMask, gSparseIdx, gSparseCount,
141+ totalRowNumBlockMask, yBlockNumAligned_, avgRowPerSubCore_, preActiveSubCoreNum_);
142+ }
143+#endif
144+ AscendC::SyncAll<false>();
145+#ifdef __DAV_CUBE__
146+ coreIdx = AscendC::GetBlockIdx();
147+ BlockMmadQK blockMmadQK(resource, mm1L1TileHelper_);
148+ BlockMmadPV blockMmadPV(resource, mm2L1AddrStart_, mm2L1TileHelper_);
149+#endif
150+#ifdef __DAV_VEC__
151+ coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum();
152+ EpilogueOnlineSoftmax epilogueOnlineSoftmax(resource, scaleValue_);
153+ EpilogueRescaleO epilogueRescaleO(resource);
154+#endif
155+ uint32_t qSTileNumPerFullXBlock = CeilDiv(blockShapeX_, qBaseTile_);
156+ // Calculate strides based on layout
157+ // For TND: [T, N, D], stride = N * D
158+ // For BNSD: [B, N, S, D], strideB = N * S * D, strideN = S * D, strideS = D
159+ int64_t strideQO = 0;
160+ int64_t strideKV = 0;
161+ int64_t strideQOB = 0; // BNSD batch_ stride for Q
162+ int64_t strideQON = 0; // BNSD head stride for Q
163+ int64_t strideQOS = 0; // BNSD seq stride for Q
164+ int64_t strideKVB = 0; // BNSD batch_ stride for KV
165+ int64_t strideKVN = 0; // BNSD head stride for KV
166+ int64_t strideKVS = 0; // BNSD seq stride for KV
167+ if constexpr (qFormat == Format::BNSD) {
168+ strideQOB = qHeads_ * qSeqlenAligned_ * embed_; // batch_ stride
169+ strideQON = qSeqlenAligned_ * embed_; // head stride
170+ strideQOS = embed_; // seq stride
171+ } else if constexpr (qFormat == Format::TND) {
172+ strideQO = qHeads_ * embed_;
173+ }
174+ if constexpr (kvFormat == Format::BNSD) {
175+ strideKVB = kvHeads_ * kvSeqlenAligned_ * embed_; // batch_ stride
176+ strideKVN = kvSeqlenAligned_ * embed_; // head stride
177+ strideKVS = embed_; // seq stride
178+ } else if constexpr (kvFormat == Format::TND) {
179+ strideKV = kvHeads_ * embed_;
180+ }
181+ uint32_t embedRound = RoundUp(embed_, 16);
182+ uint32_t groupSize = qHeads_ / kvHeads_;
183+ int qkStaticLoopNum = GetQkStaticLoopNum();
184+ int pvStaticNLoopNum = GetPvStaticNLoopNum();
185+ int64_t qBOffset = 0;
186+ int64_t kBOffset = 0;
187+ int64_t vBOffset = 0;
188+ int64_t oBOffset = 0;
189+ uint32_t preTotalTaskNum = 0;
190+ uint32_t curBatch = 0;
191+ int64_t qSeqlen = actSeqAval_ ? gActualQseqlen.GetValue(curBatch) : qSeqlenAligned_;
192+ int64_t kvSeqlen = actSeqAval_ ? gActualKvseqlen.GetValue(curBatch) : kvSeqlenAligned_;
193+ uint32_t curQSTileNum = GetCurQSTileNum(qSeqlen, blockShapeX_, qBaseTile_);
194+ uint32_t curTotalTaskNum = firstBatchTaskNum_;
195+ // Go through each task
196+ for (uint32_t taskIdx = coreIdx; taskIdx < totalTaskNum_; taskIdx += coreNum) {
197+ while (taskIdx >= curTotalTaskNum) {
198+ ++curBatch;
199+ preTotalTaskNum = curTotalTaskNum;
200+ if constexpr (qFormat == Format::TND) {
201+ qBOffset += qSeqlen * strideQO;
202+ oBOffset += qSeqlen * strideQO;
203+ }
204+ if constexpr (kvFormat == Format::TND) {
205+ kBOffset += kvSeqlen * strideKV;
206+ vBOffset += kvSeqlen * strideKV;
207+ }
208+ qSeqlen = actSeqAval_ ? gActualQseqlen.GetValue(curBatch) : qSeqlenAligned_;
209+ kvSeqlen = actSeqAval_ ? gActualKvseqlen.GetValue(curBatch) : kvSeqlenAligned_;
210+ curQSTileNum = GetCurQSTileNum(qSeqlen, blockShapeX_, qBaseTile_);
211+ curTotalTaskNum += curQSTileNum * qHeads_;
212+ }
213+ uint32_t taskIdxCurBatch = taskIdx - preTotalTaskNum;
214+ uint32_t qHeadIdx = taskIdxCurBatch / curQSTileNum;
215+ uint32_t kvHeadIdx = qHeadIdx / groupSize;
216+ uint32_t qSTileIdx = taskIdxCurBatch - qHeadIdx * curQSTileNum;
217+ // corresponding xBlock index of cur task
218+ uint32_t xBlockIdx = qSTileIdx / qSTileNumPerFullXBlock;
219+ // the q base tile index within the corrsponding xBlock of cur task
220+ uint32_t qSTileIdxCurXBlock = qSTileIdx - xBlockIdx * qSTileNumPerFullXBlock;
221+ // corresponding head index of cur task
222+ // corresponding blockSparseMask gm offset of cur task
223+ // gmBlockSparseMask has the shape [B, qN, xBlockNumAligned, yBlockNumAligned]
224+ int64_t sparseMaskBOffset = curBatch * qHeads_ * xBlockNumAligned_ * yBlockNumAligned_;
225+ int64_t sparseMaskNOffset = qHeadIdx * xBlockNumAligned_ * yBlockNumAligned_;
226+ int64_t sparseMaskXOffset = xBlockIdx * yBlockNumAligned_;
227+ int64_t gmOffsetSparseMask = sparseMaskBOffset + sparseMaskNOffset + sparseMaskXOffset;
228+ // corresponding Q/K/V/O gm offset of cur task
229+ int64_t gmOffsetQ = 0;
230+ int64_t gmOffsetK = 0;
231+ int64_t gmOffsetV = 0;
232+ int64_t gmOffsetO = 0;
233+ int64_t qSOffset = xBlockIdx * blockShapeX_ + qSTileIdxCurXBlock * qBaseTile_;
234+ if constexpr (qFormat == Format::BNSD) {
235+ qBOffset = curBatch * strideQOB;
236+ oBOffset = curBatch * strideQOB;
237+ gmOffsetQ = qBOffset + qHeadIdx * strideQON + qSOffset * strideQOS;
238+ gmOffsetO = oBOffset + qHeadIdx * strideQON + qSOffset * strideQOS;
239+ } else if constexpr (qFormat == Format::TND) {
240+ gmOffsetQ = qBOffset + qSOffset * strideQO + qHeadIdx * embed_;
241+ gmOffsetO = oBOffset + qSOffset * strideQO + qHeadIdx * embed_;
242+ }
243+ if constexpr (kvFormat == Format::BNSD) {
244+ kBOffset = curBatch * strideKVB;
245+ vBOffset = curBatch * strideKVB;
246+ gmOffsetK = kBOffset + kvHeadIdx * strideKVN;
247+ gmOffsetV = vBOffset + kvHeadIdx * strideKVN;
248+ } else if constexpr (kvFormat == Format::TND) {
249+ gmOffsetK = kBOffset + kvHeadIdx * embed_;
250+ gmOffsetV = vBOffset + kvHeadIdx * embed_;
251+ }
252+ // the actual x block num of cur batch_, calc by actual qseqlen
253+ uint32_t xBlockNumAval = static_cast<uint32_t>(CeilDiv(qSeqlen, static_cast<int64_t>(blockShapeX_)));
254+ uint32_t xBlockSize = (xBlockIdx == xBlockNumAval - 1) ?
255+ (qSeqlen - xBlockIdx * blockShapeX_) : blockShapeX_;
256+ uint32_t qSTileNumCurXBlock = CeilDiv(xBlockSize, qBaseTile_);
257+ uint32_t qSTileSizeAct = (qSTileIdxCurXBlock == qSTileNumCurXBlock - 1) ?
258+ (xBlockSize - qSTileIdxCurXBlock * qBaseTile_) : qBaseTile_;
259+ // calc the gathered kvS from sparse mask
260+ uint32_t gmOffsetSparseCount =
261+ curBatch * qHeads_ * xBlockNumAligned_ + qHeadIdx * xBlockNumAligned_ + xBlockIdx;
262+ uint32_t yBlockNumRsvd = gSparseCount.GetValue(gmOffsetSparseCount);
263+ if (yBlockNumRsvd == 0) {
264+ continue;
265+ }
266+ uint32_t gmOffsetSparseIdx = gmOffsetSparseCount * yBlockNumAligned_;
267+ uint32_t lastIdxOffset = gmOffsetSparseIdx + yBlockNumRsvd - 1;
268+ uint32_t lastSparseIdx = gSparseIdx.GetValue(lastIdxOffset);
269+ 
270+ uint32_t yBlockNumAval = static_cast<uint32_t>(CeilDiv(kvSeqlen, static_cast<int64_t>(blockShapeY_)));
271+ uint32_t lastYBlockSize = (lastSparseIdx == yBlockNumAval - 1) ?
272+ kvSeqlen - lastSparseIdx * blockShapeY_ : blockShapeY_;
273+ int64_t gatheredKvSeqlen = (yBlockNumRsvd - 1) * blockShapeY_ + lastYBlockSize;
274+ // the rowNum of cur task
275+ // no qS*qN combination even in GQA/MQA senario, since each qN has a different sparse pattern
276+ uint32_t rowNum = qSTileSizeAct;
277+ uint32_t rowNumRound = RoundUp(rowNum, 16);
278+ uint32_t kvSLoopNum = static_cast<uint32_t>(CeilDiv(gatheredKvSeqlen, static_cast<int64_t>(kvBaseTile_)));
279+ uint32_t kvSTileSizeAct = kvBaseTile_;
280+#ifdef __DAV_CUBE__
281+ uint32_t kvShapeCol = 0;
282+ uint32_t qShapeCol = 0;
283+ if constexpr (qFormat == Format::BNSD) {
284+ qShapeCol = strideQOS;
285+ } else if constexpr (qFormat == Format::TND) {
286+ qShapeCol = strideQO;
287+ }
288+ if constexpr (kvFormat == Format::BNSD) {
289+ kvShapeCol = strideKVS;
290+ } else if constexpr (kvFormat == Format::TND) {
291+ kvShapeCol = strideKV;
292+ }
293+ 
294+ GemmCoord actualBlockShapeQ{rowNum, embed_, 0};
295+ if constexpr (std::is_same_v<LayoutQ, layout::RowMajor>) {
296+ auto gmQLayoutTla = tla::MakeLayout<ElementQ, LayoutQ>(qBaseTile_, qShapeCol);
297+ auto gmQTensorTla = tla::MakeTensor(gQ[gmOffsetQ], gmQLayoutTla, Arch::PositionGM{});
298+ blockMmadQK.loadQGM(gmQTensorTla, actualBlockShapeQ);
299+ } else {
300+ auto gmQLayoutTla = tla::MakeLayout<ElementQ, LayoutQ>(qShapeCol, qBaseTile_);
301+ auto gmQTensorTla = tla::MakeTensor(gQ[gmOffsetQ], gmQLayoutTla, Arch::PositionGM{});
302+ blockMmadQK.loadQGM(gmQTensorTla, actualBlockShapeQ);
303+ }
304+ 
305+ auto gmKLayoutTla = tla::MakeLayout<ElementK, LayoutK>(kvShapeCol, kvBaseTile_);
306+ auto gmKTensorTla = tla::MakeTensor(gK[gmOffsetK], gmKLayoutTla, Arch::PositionGM{});
307+ 
308+ auto gmVLayoutTla = tla::MakeLayout<ElementV, LayoutV>(kvBaseTile_, kvShapeCol);
309+ auto gmVTensorTla = tla::MakeTensor(gV[gmOffsetV], gmVLayoutTla, Arch::PositionGM{});
310+#endif
311+#ifdef __DAV_VEC__
312+ uint32_t oShapeCol = 0;
313+ if constexpr (qFormat == Format::BNSD) {
314+ oShapeCol = strideQOS;
315+ } else if constexpr (qFormat == Format::TND) {
316+ oShapeCol = strideQO;
317+ }
318+ auto gmOLayoutTla = tla::MakeLayout<ElementO, LayoutO>(qBaseTile_, oShapeCol);
319+ auto gmOTensorTla = tla::MakeTensor(gO[gmOffsetO], gmOLayoutTla, Arch::PositionGM{});
320+#endif
321+ for (uint32_t gatheredKvSTileIdx = 0; gatheredKvSTileIdx < kvSLoopNum + PRE_LAUNCH; gatheredKvSTileIdx++) {
322+ if (gatheredKvSTileIdx < kvSLoopNum) {
323+ if (gatheredKvSTileIdx == kvSLoopNum - 1) {
324+ kvSTileSizeAct = gatheredKvSeqlen - gatheredKvSTileIdx * kvBaseTile_;
325+ } else {
326+ kvSTileSizeAct = kvBaseTile_;
327+ }
328+ // QK
329+ GemmCoord actualBlockShapeQK{rowNum, kvSTileSizeAct, embed_};
330+ uint32_t ubSBufId = gatheredKvSTileIdx % UB_S_OTMP_BUF_STAGES;
331+ auto ubSLayoutTla = tla::MakeLayout<ElementS, LayoutS>(rowNumRound, RoundUp(kvSTileSizeAct, 16));
332+ auto ubSTensorTla = tla::MakeTensor(ubSTensor[ubSBufId],
333+ ubSLayoutTla, Arch::PositionUB{});
334+ uint32_t Mm1ToSmFlagId = ubSBufId;
335+ Arch::CrossCoreFlag mm1ToSmFlag(Mm1ToSmFlagId);
336+#ifdef __DAV_CUBE__
337+ uint64_t prefixSumL0AStages = CalcCrossMm1Mm2PrefixSumL0ABStages(
338+ gatheredKvSTileIdx, mm1L0ATotalStages_, mm2L0ATotalStages_, kvSLoopNum, true);
339+ uint64_t prefixSumL0BStages = CalcCrossMm1Mm2PrefixSumL0ABStages(
340+ gatheredKvSTileIdx, mm1L0BTotalStages_, mm2L0BTotalStages_, kvSLoopNum, true);
341+ if (qkStaticLoopNum == 1) {
342+ blockMmadQK.template operator()<0, 1, 1>(
343+ gmKTensorTla, ubSTensorTla, gSparseIdx[gmOffsetSparseIdx],
344+ actualBlockShapeQK,
345+ gatheredKvSTileIdx, kvSeqlen,
346+ kvBaseTile_, blockShapeY_, yBlockNumAval, yBlockNumRsvd,
347+ prefixSumL0AStages, prefixSumL0BStages,
348+ mm1ToSmFlag, 1.0f);
349+ } else {
350+ blockMmadQK.template operator()<0, -1, -1>(
351+ gmKTensorTla, ubSTensorTla, gSparseIdx[gmOffsetSparseIdx],
352+ actualBlockShapeQK,
353+ gatheredKvSTileIdx, kvSeqlen,
354+ kvBaseTile_, blockShapeY_, yBlockNumAval, yBlockNumRsvd,
355+ prefixSumL0AStages, prefixSumL0BStages,
356+ mm1ToSmFlag, 1.0f);
357+ }
358+ if (gatheredKvSTileIdx == kvSLoopNum - 1)
359+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
360+#endif
361+ // SM
362+ uint32_t l1PBufId = gatheredKvSTileIdx % pL1BufNum_;
363+ uint32_t smToMm2FlagId = l1PBufId + UB_S_OTMP_BUF_STAGES;
364+ Arch::CrossCoreFlag smToMm2Flag(smToMm2FlagId);
365+ auto l1PLayoutTla = tla::MakeLayout<ElementP, NpuArch::layout::zN>(rowNum, kvSTileSizeAct);
366+ auto l1PTensorTla = tla::MakeTensor(l1PTensor[l1PBufId],
367+ l1PLayoutTla, Arch::PositionL1{});
368+#ifdef __DAV_VEC__
369+ epilogueOnlineSoftmax(
370+ l1PTensorTla,
371+ actualBlockShapeQK,
372+ (gatheredKvSTileIdx == 0),
373+ ubSBufId,
374+ l1PBufId,
375+ mm1ToSmFlag,
376+ smToMm2Flag);
377+#endif
378+ }
379+ if (gatheredKvSTileIdx >= PRE_LAUNCH) {
380+ uint32_t gatheredKvSTileIdxDe = gatheredKvSTileIdx - PRE_LAUNCH;
381+ if (gatheredKvSTileIdxDe == kvSLoopNum - 1) {
382+ kvSTileSizeAct = gatheredKvSeqlen - gatheredKvSTileIdxDe * kvBaseTile_;
383+ } else {
384+ kvSTileSizeAct = kvBaseTile_;
385+ }
386+ // PV
387+ GemmCoord actualBlockShapePV{rowNum, embed_, kvSTileSizeAct};
388+ uint32_t ubOTmpBufId = gatheredKvSTileIdxDe % UB_S_OTMP_BUF_STAGES;
389+ // 核间同步flagId规律：每份dst对应一个id，从小到大顺序依次为qk->sm, sm->pv, pv->re
390+ uint32_t Mm2ToReFlagId = ubOTmpBufId + UB_S_OTMP_BUF_STAGES + pL1BufNum_;
391+#ifdef __DAV_CUBE__
392+ uint32_t l1PBufId = gatheredKvSTileIdxDe % pL1BufNum_;
393+ auto ubOTmpLayoutTla = tla::MakeLayout<ElementOTmp, LayoutOTmp>(rowNumRound, embedRound);
394+ auto ubOTmpTensorTla = tla::MakeTensor(ubOTmpTensor[ubOTmpBufId],
395+ ubOTmpLayoutTla, Arch::PositionUB{});
396+ uint32_t smToMm2FlagId = l1PBufId + UB_S_OTMP_BUF_STAGES;
397+
398+ Arch::CrossCoreFlag smToMm2Flag(smToMm2FlagId);
399+ Arch::CrossCoreFlag mm2ToReFlag(Mm2ToReFlagId);
400+ uint64_t prefixSumL0AStages = CalcCrossMm1Mm2PrefixSumL0ABStages(
401+ gatheredKvSTileIdxDe, mm1L0ATotalStages_, mm2L0ATotalStages_, kvSLoopNum, false);
402+ uint64_t prefixSumL0BStages = CalcCrossMm1Mm2PrefixSumL0ABStages(
403+ gatheredKvSTileIdxDe, mm1L0BTotalStages_, mm2L0BTotalStages_, kvSLoopNum, false);
404+ if (pvStaticNLoopNum == 1) {
405+ blockMmadPV.template operator()<1, 1>(
406+ gmVTensorTla, ubOTmpTensorTla, gSparseIdx[gmOffsetSparseIdx],
407+ actualBlockShapePV,
408+ gatheredKvSTileIdxDe, kvSeqlen,
409+ kvBaseTile_, blockShapeY_, yBlockNumAval, yBlockNumRsvd,
410+ prefixSumL0AStages, prefixSumL0BStages,
411+ smToMm2Flag, mm2ToReFlag);
412+ } else if (pvStaticNLoopNum == 2) {
413+ blockMmadPV.template operator()<1, 2>(
414+ gmVTensorTla, ubOTmpTensorTla, gSparseIdx[gmOffsetSparseIdx],
415+ actualBlockShapePV,
416+ gatheredKvSTileIdxDe, kvSeqlen,
417+ kvBaseTile_, blockShapeY_, yBlockNumAval, yBlockNumRsvd,
418+ prefixSumL0AStages, prefixSumL0BStages,
419+ smToMm2Flag, mm2ToReFlag);
420+ } else {
421+ blockMmadPV.template operator()<-1, -1>(
422+ gmVTensorTla, ubOTmpTensorTla, gSparseIdx[gmOffsetSparseIdx],
423+ actualBlockShapePV,
424+ gatheredKvSTileIdxDe, kvSeqlen,
425+ kvBaseTile_, blockShapeY_, yBlockNumAval, yBlockNumRsvd,
426+ prefixSumL0AStages, prefixSumL0BStages,
427+ smToMm2Flag, mm2ToReFlag);
428+ }
429+#endif
430+#ifdef __DAV_VEC__
431+ // rescale O
432+ Arch::CrossCoreFlag mm2ToReFlag(Mm2ToReFlagId);
433+ uint32_t curTileMod = gatheredKvSTileIdxDe % (PRE_LAUNCH + 1);
434+ epilogueRescaleO(
435+ gO[gmOffsetO], actualBlockShapePV,
436+ curTileMod, gatheredKvSTileIdxDe,
437+ (gatheredKvSTileIdxDe == 0),
438+ (gatheredKvSTileIdxDe == kvSLoopNum - 1),
439+ mm2ToReFlag);
440+#endif
441+ }
442+ }
443+ }
444+ // release reverse sync flags
445+ ReleaseSyncFlags<4, 4, 4>();
446+ }
447+ 
448+ __aicore__ inline
449+ void FetchBaseShapeInfo(__gm__ EagleQuantBlockSparseAttentionTilingData *bsaTilingData)
450+ {
451+ batch_ = bsaTilingData->batch;
452+ qHeads_ = bsaTilingData->numHeads;
453+ kvHeads_ = bsaTilingData->kvHeads;
454+ embed_ = bsaTilingData->embeddingSize;
455+ firstBatchTaskNum_ = bsaTilingData->firstBatchTaskNum;
456+ totalTaskNum_ = bsaTilingData->totalTaskNum;
457+ blockShapeX_ = bsaTilingData->blockShapeX;
458+ blockShapeY_ = bsaTilingData->blockShapeY;
459+ scaleValue_ = bsaTilingData->scaleValue;
460+ // mask2idx tile info
461+ xBlockNumAligned_ = bsaTilingData->BsaMask2IdxTileInfo.xBlockNumAligned;
462+ yBlockNumAligned_ = bsaTilingData->BsaMask2IdxTileInfo.yBlockNumAligned;
463+ avgRowPerSubCore_ = bsaTilingData->BsaMask2IdxTileInfo.avgRowPerSubCore;
464+ preActiveSubCoreNum_ = bsaTilingData->BsaMask2IdxTileInfo.preActiveSubCoreNum;
465+ // base tile info
466+ qBaseTile_ = bsaTilingData->BsaBaseTileInfo.qBaseTile;
467+ kvBaseTile_ = bsaTilingData->BsaBaseTileInfo.kvBaseTile;
468+ // whether actual seqlen is provided
469+ actSeqAval_ = (!bsaTilingData->useUniformQSeqlen) && (!bsaTilingData->useUniformKvSeqlen);
470+ sparseIdxSize_ = bsaTilingData->selectIdxSize;
471+ // aligned seqlen q & kv
472+ qSeqlenAligned_ = bsaTilingData->maxQSeqlen;
473+ kvSeqlenAligned_ = bsaTilingData->maxKvSeqlen;
474+ }
475+ 
476+ __aicore__ inline
477+ void CalcOnChipBufTileInfo(__gm__ EagleQuantBlockSparseAttentionTilingData *bsaTilingData)
478+ {
479+ mm1L1TileM_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm1L1TileM;
480+ mm1L1TileN_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm1L1TileN;
481+ mm1L1TileKLeft_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm1L1TileKLeft;
482+ mm1L1TileKRight_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm1L1TileKRight;
483+ mm2L1TileM_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm2L1TileM;
484+ mm2L1TileN_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm2L1TileN;
485+ mm2L1TileKLeft_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm2L1TileKLeft;
486+ mm2L1TileKRight_ = bsaTilingData->BsaMmPhaseL1TileInfo.mm2L1TileKRight;
487+ qL1BufNum_ = bsaTilingData->BsaMmPhaseL1TileInfo.qL1BufNum;
488+ kL1BufNum_ = bsaTilingData->BsaMmPhaseL1TileInfo.kL1BufNum;
489+ vL1BufNum_ = bsaTilingData->BsaMmPhaseL1TileInfo.vL1BufNum;
490+ pL1BufNum_ = bsaTilingData->BsaMmPhaseL1TileInfo.pL1BufNum;
491+ Gemm::Block::Mm1L1TileHelper mm1L1TileHelper(mm1L1TileM_, mm1L1TileN_, mm1L1TileKLeft_, mm1L1TileKRight_,
492+ qL1BufNum_, kL1BufNum_);
493+ mm1L1TileHelper_ = mm1L1TileHelper;
494+ Gemm::Block::Mm2L1TileHelper mm2L1TileHelper(mm2L1TileM_, mm2L1TileN_, mm2L1TileKLeft_, mm2L1TileKRight_,
495+ pL1BufNum_, vL1BufNum_);
496+ mm2L1TileHelper_ = mm2L1TileHelper;
497+ mm2L1AddrStart_ = mm1L1TileM_ * mm1L1TileKLeft_ * qL1BufNum_ * sizeof(ElementQ) +
498+ mm1L1TileKRight_ * mm1L1TileN_ * kL1BufNum_ * sizeof(ElementK);
499+ mm1L0ATotalStages_ = (qBaseTile_ / BlockMmadQK::L0_TILE_M) * (embed_ / BlockMmadQK::L0_TILE_K);
500+ mm1L0BTotalStages_ = (kvBaseTile_ / BlockMmadQK::L0_TILE_N) * (embed_ / BlockMmadQK::L0_TILE_K);
501+ mm2L0ATotalStages_ = (qBaseTile_ / BlockMmadPV::L0_TILE_M) * (kvBaseTile_ / BlockMmadPV::L0_TILE_K);
502+ mm2L0BTotalStages_ = (kvBaseTile_ / BlockMmadPV::L0_TILE_K) * (embed_ / BlockMmadPV::L0_TILE_N);
503+ }
504+ 
505+ __aicore__ inline
506+ uint64_t CalcCrossMm1Mm2PrefixSumL0ABStages(
507+ uint32_t gatheredKvSTileIdx, uint32_t singleMm1L0Stages,
508+ uint32_t singleMm2L0Stages, uint32_t kvSLoopNum,
509+ bool isCurPhaseMm1)
510+ {
511+ uint64_t prefixSumStages;
512+ if (isCurPhaseMm1) {
513+ prefixSumStages = (gatheredKvSTileIdx <= PRE_LAUNCH) ?
514+ gatheredKvSTileIdx * singleMm1L0Stages :
515+ gatheredKvSTileIdx * singleMm1L0Stages + (gatheredKvSTileIdx - PRE_LAUNCH) * singleMm2L0Stages;
516+ } else {
517+ prefixSumStages = (gatheredKvSTileIdx < kvSLoopNum - PRE_LAUNCH) ?
518+ (gatheredKvSTileIdx + 1 + PRE_LAUNCH) * singleMm1L0Stages + gatheredKvSTileIdx * singleMm2L0Stages:
519+ kvSLoopNum * singleMm1L0Stages + gatheredKvSTileIdx * singleMm2L0Stages;
520+ }
521+ return prefixSumStages;
522+ }
523+ 
524+ __aicore__ inline
525+ void InitCrossCoreDstBuf(
526+ AscendC::LocalTensor<ElementP> (&l1PTensor)[MAX_CROSS_CORE_BUF_STAGES],
527+ AscendC::LocalTensor<ElementS> (&ubSTensor)[UB_S_OTMP_BUF_STAGES],
528+ AscendC::LocalTensor<ElementOTmp> (&ubOTmpTensor)[UB_S_OTMP_BUF_STAGES])
529+ {
530+ for (uint32_t i = 0; i < pL1BufNum_; i++) {
531+ l1PTensor[i] = resource.l1Buf.template GetBufferByByte<ElementP>(
532+ mm2L1AddrStart_ + mm2L1TileM_ * mm2L1TileKLeft_ * sizeof(ElementP) * i);
533+ }
534+ uint32_t rowNumPerSubCore = EpilogueOnlineSoftmax::SM_ROW_MAX_ELEM_NUM;
535+ uint32_t colNumPerSubCore = EpilogueOnlineSoftmax::SM_COL_MAX_ELEM_NUM;
536+ uint32_t rescaleCol = EpilogueRescaleO::RESCALE_COL_MAX_ELEM_NUM;
537+ for (uint32_t i = 0; i < UB_S_OTMP_BUF_STAGES; i++) {
538+ ubSTensor[i] = resource.ubBuf.template GetBufferByByte<ElementS>(
539+ rowNumPerSubCore * colNumPerSubCore * sizeof(ElementS) * i);
540+ ubOTmpTensor[i] = resource.ubBuf.template GetBufferByByte<ElementOTmp>(
541+ rowNumPerSubCore * colNumPerSubCore * sizeof(ElementS) * UB_S_OTMP_BUF_STAGES +
542+ rowNumPerSubCore * colNumPerSubCore * sizeof(ElementS) * UB_S_OTMP_BUF_STAGES +
543+ rowNumPerSubCore * rescaleCol * sizeof(ElementOTmp) * i);
544+ }
545+ }
546+ 
547+ template <uint32_t MM1_SM_MODE, uint32_t MM2_RE_MODE, uint32_t SM_MM2_MODE>
548+ __aicore__ inline
549+ void InitSyncFlags()
550+ {
551+#ifdef __DAV_CUBE__
552+ // same core sync between pipes
553+ // Query
554+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
555+ // Key
556+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID1);
557+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
558+ // Value
559+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID3);
560+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID4);
561+ // L0A
562+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
563+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID1);
564+ // L0B
565+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID2);
566+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID3);
567+ // L0C
568+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
569+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID1);
570+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID2);
571+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID3);
572+ // VSDB (L1 Value Scale Double Buffer)
573+ AscendC::SetFlag<AscendC::HardEvent::MTE2_FIX>(EVENT_ID0);
574+ AscendC::SetFlag<AscendC::HardEvent::MTE2_FIX>(EVENT_ID1);
575+ AscendC::SetFlag<AscendC::HardEvent::FIX_MTE2>(EVENT_ID0);
576+ // cross core sync
577+ if constexpr (SM_MM2_MODE == 4U) {
578+ AscendC::CrossCoreSetFlag<SM_MM2_MODE, PIPE_MTE1>(2);
579+ AscendC::CrossCoreSetFlag<SM_MM2_MODE, PIPE_MTE1>(18);
580+ AscendC::CrossCoreSetFlag<SM_MM2_MODE, PIPE_MTE1>(3);
581+ AscendC::CrossCoreSetFlag<SM_MM2_MODE, PIPE_MTE1>(19);
582+ AscendC::CrossCoreSetFlag<SM_MM2_MODE, PIPE_MTE1>(4);
583+ AscendC::CrossCoreSetFlag<SM_MM2_MODE, PIPE_MTE1>(20);
584+ }
585+#endif
586+#ifdef __DAV_VEC__
587+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
588+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID1);
589+ // mask2index
590+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
591+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID1);
592+ // softmax
593+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID2);
594+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);
595+ // rescale
596+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
597+ 
598+ AscendC::SetFlag<AscendC::HardEvent::S_MTE3>(EVENT_ID0);
599+ AscendC::SetFlag<AscendC::HardEvent::MTE3_S>(EVENT_ID0);
600+ AscendC::SetFlag<AscendC::HardEvent::MTE3_S>(EVENT_ID1);
601+ if constexpr (MM1_SM_MODE == 4U) {
602+ AscendC::CrossCoreSetFlag<MM1_SM_MODE, PIPE_V>(0);
603+ AscendC::CrossCoreSetFlag<MM1_SM_MODE, PIPE_V>(1);
604+ }
605+ if constexpr (MM2_RE_MODE == 4U) {
606+ AscendC::CrossCoreSetFlag<MM2_RE_MODE, PIPE_V>(5);
607+ AscendC::CrossCoreSetFlag<MM2_RE_MODE, PIPE_V>(6);
608+ }
609+#endif
610+ }
611+ 
612+ template <uint32_t MM1_SM_MODE, uint32_t MM2_RE_MODE, uint32_t SM_MM2_MODE>
613+ __aicore__ inline
614+ void ReleaseSyncFlags()
615+ {
616+#ifdef __DAV_CUBE__
617+ // same core sync between pipes
618+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
619+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID1);
620+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
621+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID3);
622+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID4);
623+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
624+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID1);
625+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID2);
626+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID3);
627+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
628+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID1);
629+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID2);
630+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID3);
631+ // VSDB (L1 Value Scale Double Buffer)
632+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_FIX>(EVENT_ID0);
633+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_FIX>(EVENT_ID1);
634+ AscendC::WaitFlag<AscendC::HardEvent::FIX_MTE2>(EVENT_ID0);
635+ if constexpr (MM1_SM_MODE == 4U) {
636+ AscendC::CrossCoreWaitFlag<MM1_SM_MODE, PIPE_FIX>(0);
637+ AscendC::CrossCoreWaitFlag<MM1_SM_MODE, PIPE_FIX>(1);
638+ AscendC::CrossCoreWaitFlag<MM1_SM_MODE, PIPE_FIX>(16);
639+ AscendC::CrossCoreWaitFlag<MM1_SM_MODE, PIPE_FIX>(17);
640+ }
641+ if constexpr (MM2_RE_MODE == 4U) {
642+ AscendC::CrossCoreWaitFlag<MM2_RE_MODE, PIPE_FIX>(5);
643+ AscendC::CrossCoreWaitFlag<MM2_RE_MODE, PIPE_FIX>(21);
644+ AscendC::CrossCoreWaitFlag<MM2_RE_MODE, PIPE_FIX>(6);
645+ AscendC::CrossCoreWaitFlag<MM2_RE_MODE, PIPE_FIX>(22);
646+ }
647+#endif
648+#ifdef __DAV_VEC__
649+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
650+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID1);
651+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
652+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID1);
653+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID2);
654+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID3);
655+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
656+ AscendC::WaitFlag<AscendC::HardEvent::S_MTE3>(EVENT_ID0);
657+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_S>(EVENT_ID0);
658+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_S>(EVENT_ID1);
659+ if constexpr (SM_MM2_MODE == 4U) {
660+ AscendC::CrossCoreWaitFlag<SM_MM2_MODE, PIPE_MTE3>(2);
661+ AscendC::CrossCoreWaitFlag<SM_MM2_MODE, PIPE_MTE3>(3);
662+ AscendC::CrossCoreWaitFlag<SM_MM2_MODE, PIPE_MTE3>(4);
663+ }
664+#endif
665+ AscendC::PipeBarrier<PIPE_ALL>();
666+ }
667+ 
668+private:
669+ Arch::Resource<ArchTag> resource;
670+ /*
671+ tiling info, which are const in each kernel launch
672+ */
673+ // basic shape info
674+ uint32_t batch_;
675+ uint32_t qHeads_;
676+ uint32_t kvHeads_;
677+ uint32_t embed_;
678+ uint32_t firstBatchTaskNum_;
679+ uint32_t totalTaskNum_;
680+ uint32_t blockShapeX_;
681+ uint32_t blockShapeY_;
682+ float scaleValue_;
683+ // mask2idx tile info
684+ uint32_t xBlockNumAligned_;
685+ uint32_t yBlockNumAligned_;
686+ uint32_t avgRowPerSubCore_;
687+ uint32_t preActiveSubCoreNum_;
688+ // base tile info
689+ uint32_t qBaseTile_;
690+ uint32_t kvBaseTile_;
691+ // whether actual seqlen is provided
692+ uint32_t actSeqAval_;
693+ // workspace size
694+ uint64_t sparseIdxSize_;
695+ // aligned seqlen q & kv
696+ int64_t qSeqlenAligned_;
697+ int64_t kvSeqlenAligned_;
698+ // L1 tile info
699+ uint32_t mm1L1TileM_;
700+ uint32_t mm1L1TileN_;
701+ uint32_t mm1L1TileKLeft_;
702+ uint32_t mm1L1TileKRight_;
703+ uint32_t mm2L1TileM_;
704+ uint32_t mm2L1TileN_;
705+ uint32_t mm2L1TileKLeft_;
706+ uint32_t mm2L1TileKRight_;
707+ uint32_t qL1BufNum_;
708+ uint32_t kL1BufNum_;
709+ uint32_t vL1BufNum_;
710+ uint32_t pL1BufNum_;
711+ uint32_t mm1L0ATotalStages_;
712+ uint32_t mm1L0BTotalStages_;
713+ uint32_t mm2L0ATotalStages_;
714+ uint32_t mm2L0BTotalStages_;
715+ uint32_t mm2L1AddrStart_ = 0;
716+ Gemm::Block::Mm1L1TileHelper mm1L1TileHelper_;
717+ Gemm::Block::Mm2L1TileHelper mm2L1TileHelper_;
718+};
719+ 
720+}
721+ 
722+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/arch35/kernel_utils.hpp+230-0
@@ -0,0 +1,230 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef BSA_ARCH35_KERNEL_UTILS
11+#define BSA_ARCH35_KERNEL_UTILS
12+ 
13+#include "../attn_infra/base_defs.hpp"
14+#include "../attn_infra/arch/arch.hpp"
15+#include "../attn_infra/layout/layout.hpp"
16+ 
17+#include "../attn_infra/gemm/block/block_mmad.hpp"
18+#include "../attn_infra/gemm/dispatch_policy.hpp"
19+#include "../attn_infra/gemm/gemm_type.hpp"
20+ 
21+#include "../attn_infra/arch/cross_core_sync.hpp"
22+#include "../attn_infra/arch/resource.hpp"
23+#include "../attn_infra/epilogue/block/block_epilogue.hpp"
24+#include "../attn_infra/epilogue/dispatch_policy.hpp"
25+#include "../tla/tensor.hpp"
26+#include "../tla/layout.hpp"
27+#include "kernel_operator.h"
28+#include "lib/matmul_intf.h"
29+#include "kernel_tiling/kernel_tiling.h"
30+ 
31+namespace BsaKernelArch35 {
32+ 
33+enum class Format {
34+ TND = 0,
35+ BNSD = 1
36+};
37+ 
38+struct BsaKernelParamsArch35 {
39+ GM_ADDR q;
40+ GM_ADDR k;
41+ GM_ADDR v;
42+ GM_ADDR mask;
43+ GM_ADDR blockTables;
44+ GM_ADDR query_scale;
45+ GM_ADDR key_scale;
46+ GM_ADDR value_scale;
47+ GM_ADDR actualQseqlen;
48+ GM_ADDR actualKvseqlen;
49+ GM_ADDR blockSparseMask;
50+ GM_ADDR o;
51+ GM_ADDR workSpace;
52+ GM_ADDR tiling;
53+ 
54+ // Methods
55+ __aicore__ inline
56+ BsaKernelParamsArch35() {}
57+ __aicore__ inline
58+ BsaKernelParamsArch35(GM_ADDR q_, GM_ADDR k_, GM_ADDR v_, GM_ADDR mask_, GM_ADDR blockTables_,
59+ GM_ADDR query_scale_, GM_ADDR key_scale_, GM_ADDR value_scale_,
60+ GM_ADDR actualQseqlen_, GM_ADDR actualKvseqlen_, GM_ADDR blockSparseMask_, GM_ADDR o_,
61+ GM_ADDR workSpace_, GM_ADDR tiling_)
62+ : q(q_), k(k_), v(v_), mask(mask_), blockTables(blockTables_), actualQseqlen(actualQseqlen_),
63+ query_scale(query_scale_), key_scale(key_scale_), value_scale(value_scale_),
64+ actualKvseqlen(actualKvseqlen_), blockSparseMask(blockSparseMask_), o(o_),
65+ workSpace(workSpace_), tiling(tiling_) {}
66+};
67+ 
68+constexpr uint32_t SPARSE_PATTERN_MODE_TABLE = 1;
69+ 
70+ 
71+template<class ArchTag>
72+__aicore__ inline void SparseTable2Count(
73+ NpuArch::Arch::Resource<ArchTag> &resource,
74+ AscendC::GlobalTensor<int32_t> sparseTableGM,
75+ AscendC::GlobalTensor<int32_t> sparseCountGM,
76+ uint32_t totalRowNumBlockMask,
77+ uint32_t yBlockNumAligned,
78+ uint32_t avgRowPerSubCore,
79+ uint32_t preActiveSubCoreNum)
80+{
81+ static constexpr uint32_t PRE_ROW_TILE = 128;
82+ static constexpr uint32_t PRE_COL_TILE = 128;
83+ static constexpr uint32_t PRE_ELEM_NUM_PER_LOOP = PRE_ROW_TILE * PRE_COL_TILE;
84+ static constexpr uint32_t TABLE_IN_INT32 = 0;
85+ static constexpr uint32_t TABLE_VALID_BIT = TABLE_IN_INT32 + PRE_ELEM_NUM_PER_LOOP * sizeof(int32_t);
86+ static constexpr uint32_t TABLE_ONE_VALUE = TABLE_VALID_BIT + PRE_COL_TILE * sizeof(uint8_t);
87+ static constexpr uint32_t TABLE_VALID_VALUE = TABLE_ONE_VALUE + PRE_COL_TILE * sizeof(float);
88+ static constexpr uint32_t TABLE_TILE_COUNT = TABLE_VALID_VALUE + PRE_ELEM_NUM_PER_LOOP * sizeof(float);
89+ static constexpr uint32_t TABLE_COUNT_FLOAT = TABLE_TILE_COUNT + PRE_ROW_TILE * sizeof(float);
90+ static constexpr uint32_t RSVD_SPARSE_COUNT = TABLE_COUNT_FLOAT + PRE_ROW_TILE * sizeof(float);
91+ // Pass Sum an explicit temp buffer so it will not pop hidden UB space outside this layout.
92+ static constexpr uint32_t TABLE_SUM_TMP = RSVD_SPARSE_COUNT + PRE_ROW_TILE * sizeof(int32_t);
93+ static constexpr uint32_t TABLE_SUM_TMP_SIZE = PRE_ROW_TILE * PRE_COL_TILE * sizeof(float);
94+ static constexpr uint32_t SPARSE_TABLE_UB_SIZE = TABLE_SUM_TMP + TABLE_SUM_TMP_SIZE;
95+ static constexpr uint32_t SPARSE_TABLE_UB_LIMIT = ArchTag::UB_SIZE - 8U * 1024U;
96+ static constexpr int32_t SPARSE_TABLE_END = -1;
97+ static constexpr int32_t SPARSE_TABLE_VALID_MIN = 0;
98+ static_assert(SPARSE_TABLE_UB_SIZE <= SPARSE_TABLE_UB_LIMIT,
99+ "SparseTable2Count UB layout exceeds the usable arch UB size.");
100+ 
101+ AscendC::LocalTensor<int32_t> sparseTableUb =
102+ resource.ubBuf.template GetBufferByByte<int32_t>(TABLE_IN_INT32);//128*128*4Bytes
103+ AscendC::LocalTensor<uint8_t> validMaskUb =
104+ resource.ubBuf.template GetBufferByByte<uint8_t>(TABLE_VALID_BIT);//128Bytes
105+ AscendC::LocalTensor<float> oneValueUb =
106+ resource.ubBuf.template GetBufferByByte<float>(TABLE_ONE_VALUE);//128*4Bytes
107+ AscendC::LocalTensor<float> validValueUb =
108+ resource.ubBuf.template GetBufferByByte<float>(TABLE_VALID_VALUE);//128*128*4Bytes
109+ AscendC::LocalTensor<float> tileCountUb =
110+ resource.ubBuf.template GetBufferByByte<float>(TABLE_TILE_COUNT);//128*4Bytes
111+ AscendC::LocalTensor<float> countFloatUb =
112+ resource.ubBuf.template GetBufferByByte<float>(TABLE_COUNT_FLOAT);//128*4Bytes
113+ AscendC::LocalTensor<int32_t> sparseCountUb =
114+ resource.ubBuf.template GetBufferByByte<int32_t>(RSVD_SPARSE_COUNT);//128*4Bytes
115+ AscendC::LocalTensor<uint8_t> sumTmpUb =
116+ resource.ubBuf.template GetBufferByByte<uint8_t>(TABLE_SUM_TMP);//128*128*4Bytes
117+ 
118+ uint32_t subCoreIdx = AscendC::GetBlockIdx();
119+ uint64_t curSubCoreRowOffset = static_cast<uint64_t>(subCoreIdx) * avgRowPerSubCore;
120+ uint32_t actDealtRow = (subCoreIdx == preActiveSubCoreNum - 1) ?
121+ static_cast<uint32_t>(totalRowNumBlockMask - curSubCoreRowOffset) : avgRowPerSubCore;
122+ if (subCoreIdx >= preActiveSubCoreNum) {
123+ return;
124+ }
125+ 
126+ AscendC::Duplicate(oneValueUb, static_cast<float>(1.0), PRE_COL_TILE);
127+ AscendC::PipeBarrier<PIPE_V>();
128+ uint32_t rowLoop = (actDealtRow + PRE_ROW_TILE - 1) / PRE_ROW_TILE;
129+ for (uint32_t i = 0; i < rowLoop; i++) {
130+ uint32_t curLoopRowOffset = i * PRE_ROW_TILE;
131+ uint32_t actDealtRowCurLoop =
132+ (i == rowLoop - 1) ? (actDealtRow - curLoopRowOffset) : PRE_ROW_TILE;
133+ 
134+ AscendC::Duplicate(countFloatUb, static_cast<float>(0.0), actDealtRowCurLoop);
135+ AscendC::PipeBarrier<PIPE_V>();
136+ uint32_t colLoop = (yBlockNumAligned + PRE_COL_TILE - 1) / PRE_COL_TILE;
137+ for (uint32_t j = 0; j < colLoop; j++) {
138+ uint32_t curLoopColOffset = j * PRE_COL_TILE;
139+ uint32_t actDealtColCurLoop =
140+ (j == colLoop - 1) ? (yBlockNumAligned - curLoopColOffset) : PRE_COL_TILE;
141+ // MTE copy only needs 32B alignment; CompareScalar keeps the 256B count alignment.
142+ uint32_t actDealtColCurLoopCopyAlign = ((actDealtColCurLoop + 7) / 8) * 8;
143+ uint32_t actDealtColCurLoopCompareAlign = ((actDealtColCurLoop + 63) / 64) * 64;
144+ uint32_t copyRightPadding = actDealtColCurLoopCopyAlign - actDealtColCurLoop;
145+ uint64_t sparseTableOffset =
146+ (curSubCoreRowOffset + curLoopRowOffset) * yBlockNumAligned + curLoopColOffset;
147+ 
148+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(0);
149+ if (actDealtColCurLoopCompareAlign > actDealtColCurLoopCopyAlign) {
150+ // Pre-fill the compare-only tail. This avoids large MTE right-padding configs.
151+ AscendC::Duplicate(sparseTableUb, SPARSE_TABLE_END, actDealtRowCurLoop * PRE_COL_TILE);
152+ AscendC::PipeBarrier<PIPE_V>();
153+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(0);
154+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(0);
155+ }
156+ // Load one int32 sparse-table tile. MTE pads only to 32B; any wider compare tail is already -1.
157+ AscendC::DataCopyPad(
158+ sparseTableUb,
159+ sparseTableGM[sparseTableOffset],
160+ AscendC::DataCopyExtParams(
161+ actDealtRowCurLoop, //BlockCount 指定该指令包含的连续传输数据块的个数
162+ actDealtColCurLoop * sizeof(int32_t), //指定该指令每个连续传输数据块长度，单位为字节
163+ (yBlockNumAligned - actDealtColCurLoop) * sizeof(int32_t), //源操作数相邻连续数据块间隔（从前一个尾到后一个头），单位为字节
164+ (PRE_COL_TILE - actDealtColCurLoopCopyAlign) * sizeof(int32_t) / 32,//目的操作数相邻连续数据块间隔（从前一个尾到后一个头），单位为32B
165+ 0),
166+ //isPad表示要填充 leftPadding，左侧需要填充的元素个数，字节数不要超过32B，rightPadding，右侧要填充元素个数，不超过32B，PaddingValue是要填充的数值
167+ AscendC::DataCopyPadExtParams<int32_t>(
168+ true, 0, copyRightPadding, SPARSE_TABLE_END));
169+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(0);
170+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(0);
171+ for (uint32_t row = 0; row < actDealtRowCurLoop; row++) {
172+ // Sparse table rows are terminated by -1; valid table indices are non-negative.
173+ AscendC::CompareScalar(
174+ validMaskUb,//输出，目的操作数
175+ sparseTableUb[row * PRE_COL_TILE],//源操作数0
176+ SPARSE_TABLE_VALID_MIN,
177+ AscendC::CMPMODE::GE,
178+ actDealtColCurLoopCompareAlign);//calCount，输入数据元素个数，设置CalCount时，需要保证calCount个元素所占空间256字节对齐
179+ AscendC::PipeBarrier<PIPE_V>();
180+ //Mask数值为0,从src0选取，否则从src1中选取
181+ AscendC::Select(
182+ validValueUb[row * PRE_COL_TILE],//输出
183+ validMaskUb,//mask输入
184+ oneValueUb,//源操作数0
185+ static_cast<float>(0.0),//源操作数1
186+ AscendC::SELMODE::VSEL_TENSOR_SCALAR_MODE, //selMode
187+ actDealtColCurLoopCompareAlign); //calCount
188+ AscendC::PipeBarrier<PIPE_V>();
189+ }
190+ AscendC::Sum(
191+ tileCountUb,
192+ validValueUb,
193+ sumTmpUb,
194+ AscendC::SumParams{actDealtRowCurLoop, PRE_COL_TILE, actDealtColCurLoop});
195+ AscendC::PipeBarrier<PIPE_V>();
196+ AscendC::Add(countFloatUb, countFloatUb, tileCountUb, actDealtRowCurLoop);
197+ AscendC::PipeBarrier<PIPE_V>();
198+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(0);
199+ }
200+ AscendC::Cast(
201+ sparseCountUb,
202+ countFloatUb,
203+ AscendC::RoundMode::CAST_ROUND,
204+ actDealtRowCurLoop);
205+ AscendC::PipeBarrier<PIPE_V>();
206+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(0);
207+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(0);
208+ uint64_t sparseCountOffset = curSubCoreRowOffset + curLoopRowOffset;
209+ AscendC::DataCopyPad(
210+ sparseCountGM[sparseCountOffset],
211+ sparseCountUb,
212+ AscendC::DataCopyExtParams(1, actDealtRowCurLoop * sizeof(int32_t), 0, 0, 0));
213+ }
214+ 
215+}
216+ 
217+__aicore__ inline
218+uint32_t GetCurQSTileNum(int64_t curQSeqlen, uint32_t blockShapeX, uint32_t qBaseTile)
219+{
220+ uint32_t fullXBlockNum = curQSeqlen / blockShapeX;
221+ uint32_t tailXBlockSize = curQSeqlen % blockShapeX;
222+ uint32_t qSTileNumPerFullXBlock = (blockShapeX + qBaseTile - 1) / qBaseTile;
223+ uint32_t qSTileNumTailXBlock = (tailXBlockSize + qBaseTile - 1) / qBaseTile;
224+ uint32_t curQSTileNum = qSTileNumPerFullXBlock * fullXBlockNum + qSTileNumTailXBlock;
225+ return curQSTileNum;
226+}
227+ 
228+}
229+ 
230+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/arch/arch.hpp+51-0
@@ -0,0 +1,51 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef ARCH_ARCH_HPP
12+#define ARCH_ARCH_HPP
13+ 
14+#include "../../attn_infra/base_defs.hpp"
15+ 
16+namespace NpuArch::Arch
17+{
18+ 
19+struct AtlasA2 {
20+ static constexpr uint32_t BIAS_SIZE = 1024;
21+ static constexpr uint32_t FIXBUF_SIZE = 7U * 1024U;
22+ static constexpr uint32_t UB_SIZE = 192U * 1024U;
23+ static constexpr uint32_t L1_SIZE = 512U * 1024U;
24+ static constexpr uint32_t L0A_SIZE = 64U * 1024U;
25+ static constexpr uint32_t L0B_SIZE = 64U * 1024U;
26+ static constexpr uint32_t L0C_SIZE = 128U * 1024U;
27+};
28+ 
29+struct AtlasA5 {
30+ static constexpr uint32_t BIAS_SIZE = 4U * 1024U;
31+ static constexpr uint32_t FIXBUF_SIZE = 16U * 1024U;
32+ static constexpr uint32_t UB_SIZE = 256U * 1024U;
33+ static constexpr uint32_t L1_SIZE = 512U * 1024U;
34+ static constexpr uint32_t L0A_SIZE = 64U * 1024U;
35+ static constexpr uint32_t L0B_SIZE = 64U * 1024U;
36+ static constexpr uint32_t L0C_SIZE = 256U * 1024U;
37+};
38+ 
39+template <AscendC::TPosition POS>
40+using PositionType = std::integral_constant<AscendC::TPosition, POS>;
41+ 
42+using PositionGM = PositionType<AscendC::TPosition::GM>;
43+using PositionL1 = PositionType<AscendC::TPosition::A1>;
44+using PositionL0A = PositionType<AscendC::TPosition::A2>;
45+using PositionL0B = PositionType<AscendC::TPosition::B2>;
46+using PositionL0C = PositionType<AscendC::TPosition::CO1>;
47+using PositionUB = PositionType<AscendC::TPosition::VECCALC>;
48+ 
49+} // namespace NpuArch::Arch
50+ 
51+#endif // ARCH_ARCH_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/arch/cross_core_sync.hpp+116-0
@@ -0,0 +1,116 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef ARCH_CROSS_CORE_SYNC_HPP
12+#define ARCH_CROSS_CORE_SYNC_HPP
13+ 
14+#include "../../attn_infra/base_defs.hpp"
15+ 
16+namespace NpuArch::Arch
17+{
18+ 
19+constexpr uint32_t MAX_REVERSE_DEPTH = 16;
20+ 
21+using FlagID = uint16_t;
22+constexpr FlagID AIV_INTER_BLOCK_BARRIER = 8;
23+constexpr FlagID AIC_INTER_BLOCK_BARRIER = 9;
24+constexpr FlagID AIV_INTER_SUBBLOCK_BARRIER = 10;
25+constexpr FlagID FFTS_MAX_FLAG = 7;
26+ 
27+struct CrossCoreFlag {
28+ __aicore__ inline
29+ CrossCoreFlag() : id(0) {}
30+ 
31+ __aicore__ inline
32+ CrossCoreFlag(FlagID id) : id(id) {}
33+ 
34+ FlagID id;
35+};
36+ 
37+template <uint32_t REVERSE_DEPTH_ = MAX_REVERSE_DEPTH>
38+struct CrossCoreFlagWithReverse {
39+ __aicore__ inline
40+ CrossCoreFlagWithReverse() : id(0), reverseId(0) {}
41+ 
42+ __aicore__ inline
43+ CrossCoreFlagWithReverse(FlagID id, FlagID reverseId) : id(id), reverseId(reverseId) {}
44+ 
45+ FlagID id;
46+ FlagID reverseId;
47+ uint32_t count{ 0 };
48+};
49+ 
50+template <uint8_t MODE, int32_t CORE_TYPE>
51+struct BarrierFlag {
52+ static_assert(MODE != MODE, "Unsupported cross core barrier flag, can not find the specialization.");
53+};
54+ 
55+template <>
56+struct BarrierFlag<0x0, AscendC::AIV> {
57+ static constexpr FlagID ID = AIV_INTER_BLOCK_BARRIER;
58+};
59+ 
60+template <>
61+struct BarrierFlag<0x0, AscendC::AIC> {
62+ static constexpr FlagID ID = AIC_INTER_BLOCK_BARRIER;
63+};
64+ 
65+template <>
66+struct BarrierFlag<0x1, AscendC::AIV> {
67+ static constexpr FlagID ID = AIV_INTER_SUBBLOCK_BARRIER;
68+};
69+ 
70+template <uint8_t MODE, pipe_t PIPE>
71+__aicore__ inline
72+void CrossCoreBarrier()
73+{
74+ constexpr FlagID flagId = BarrierFlag<MODE, g_coreType>::ID;
75+ AscendC::CrossCoreSetFlag<MODE, PIPE>(flagId);
76+ AscendC::CrossCoreWaitFlag(flagId);
77+}
78+ 
79+template <uint8_t MODE, pipe_t PIPE>
80+__aicore__ inline
81+void CrossCoreSetFlag(CrossCoreFlag &flag)
82+{
83+ AscendC::CrossCoreSetFlag<MODE, PIPE>(flag.id);
84+}
85+ 
86+template <uint8_t MODE = 0, pipe_t PIPE = PIPE_S>
87+__aicore__ inline void CrossCoreWaitFlag(CrossCoreFlag &flag)
88+{
89+ AscendC::CrossCoreWaitFlag<MODE, PIPE>(flag.id);
90+}
91+ 
92+template <uint8_t MODE, pipe_t PIPE, uint32_t REVERSE_DEPTH>
93+__aicore__ inline
94+void CrossCoreSetFlagWithReverse(CrossCoreFlagWithReverse<REVERSE_DEPTH> &flag)
95+{
96+ AscendC::CrossCoreSetFlag<MODE, PIPE>(flag.id);
97+ if (++flag.count >= REVERSE_DEPTH) {
98+ AscendC::CrossCoreWaitFlag(flag.reverseId);
99+ flag.count = 0;
100+ }
101+}
102+ 
103+template <uint8_t MODE, pipe_t PIPE, uint32_t REVERSE_DEPTH>
104+__aicore__ inline
105+void CrossCoreWaitFlagWithReverse(CrossCoreFlagWithReverse<REVERSE_DEPTH> &flag)
106+{
107+ AscendC::CrossCoreWaitFlag(flag.id);
108+ if (++flag.count >= REVERSE_DEPTH) {
109+ AscendC::CrossCoreSetFlag<MODE, PIPE>(flag.reverseId);
110+ flag.count = 0;
111+ }
112+}
113+ 
114+} // namespace NpuArch::Arch
115+ 
116+#endif // ARCH_CROSS_CORE_SYNC_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/arch/local_tensor_buffer.hpp+343-0
@@ -0,0 +1,343 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef INCLUDE_ARCH_MEMORY_H
12+#define INCLUDE_ARCH_MEMORY_H
13+ 
14+#include "../../attn_infra/base_defs.hpp"
15+#include "../../attn_infra/arch/arch.hpp"
16+ 
17+namespace NpuArch::Arch
18+{
19+ 
20+struct LocalTensorBufferBase {
21+public:
22+ template <class Element = half>
23+ __aicore__ inline
24+ AscendC::LocalTensor<Element> GetBufferByByte(const uint32_t offset) const
25+ {
26+ return tensor[offset].template ReinterpretCast<Element>();
27+ }
28+ 
29+protected:
30+ __aicore__ inline
31+ LocalTensorBufferBase() = default;
32+ 
33+ AscendC::LocalTensor<uint8_t> tensor;
34+};
35+ 
36+template <
37+ class ArchTag,
38+ AscendC::TPosition Position
39+>
40+struct LocalTensorBuffer {
41+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported local tensor buffer, can not find the specialization.");
42+};
43+ 
44+/// Partial specialization for TPosition::A1
45+template <>
46+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::A1> : LocalTensorBufferBase {
47+public:
48+ static constexpr AscendC::TPosition Position = AscendC::TPosition::A1;
49+ 
50+ __aicore__ inline
51+ LocalTensorBuffer()
52+ {
53+ tensor = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::A1, 0, Arch::AtlasA5::L1_SIZE);
54+ }
55+};
56+ 
57+///////////////////////////////////////////////////////////
58+ 
59+/// Partial specialization for TPosition::A2
60+template <>
61+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::A2> : LocalTensorBufferBase {
62+public:
63+ static constexpr AscendC::TPosition Position = AscendC::TPosition::A2;
64+ 
65+ __aicore__ inline
66+ LocalTensorBuffer()
67+ {
68+ tensor = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::A2, 0, Arch::AtlasA5::L0A_SIZE);
69+ }
70+};
71+ 
72+///////////////////////////////////////////////////////////
73+ 
74+/// Partial specialization for TPosition::B1
75+template <>
76+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::B1> : LocalTensorBufferBase {
77+public:
78+ static constexpr AscendC::TPosition Position = AscendC::TPosition::B1;
79+ 
80+ __aicore__ inline
81+ LocalTensorBuffer()
82+ {
83+ tensor = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::B1, 0, Arch::AtlasA5::L1_SIZE);
84+ }
85+};
86+ 
87+///////////////////////////////////////////////////////////
88+ 
89+/// Partial specialization for TPosition::B2
90+template <>
91+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::B2> : LocalTensorBufferBase {
92+public:
93+ static constexpr AscendC::TPosition Position = AscendC::TPosition::B2;
94+ 
95+ __aicore__ inline
96+ LocalTensorBuffer()
97+ {
98+ tensor = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::B2, 0, Arch::AtlasA5::L0B_SIZE);
99+ }
100+};
101+ 
102+///////////////////////////////////////////////////////////
103+ 
104+/// Partial specialization for TPosition::C1
105+template <>
106+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::C1> : LocalTensorBufferBase {
107+public:
108+ static constexpr AscendC::TPosition Position = AscendC::TPosition::C1;
109+ 
110+ __aicore__ inline
111+ LocalTensorBuffer()
112+ {
113+ tensor = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::C1, 0, Arch::AtlasA5::L1_SIZE);
114+ }
115+};
116+ 
117+///////////////////////////////////////////////////////////
118+ 
119+/// Partial specialization for TPosition::C2
120+template <>
121+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::C2> : LocalTensorBufferBase {
122+public:
123+ static constexpr AscendC::TPosition Position = AscendC::TPosition::C2;
124+ 
125+ __aicore__ inline
126+ LocalTensorBuffer()
127+ {
128+ tensor = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::C2, 0, Arch::AtlasA5::BIAS_SIZE);
129+ }
130+};
131+ 
132+///////////////////////////////////////////////////////////
133+ 
134+/// Partial specialization for TPosition::CO1
135+template <>
136+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::CO1> : LocalTensorBufferBase {
137+public:
138+ static constexpr AscendC::TPosition Position = AscendC::TPosition::CO1;
139+ 
140+ __aicore__ inline
141+ LocalTensorBuffer()
142+ {
143+ tensor = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::CO1, 0, Arch::AtlasA5::L0C_SIZE);
144+ }
145+};
146+ 
147+///////////////////////////////////////////////////////////
148+ 
149+/// Partial specialization for TPosition::C2PIPE2GM
150+template <>
151+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::C2PIPE2GM> : LocalTensorBufferBase {
152+public:
153+ static constexpr AscendC::TPosition Position = AscendC::TPosition::C2PIPE2GM;
154+ 
155+ __aicore__ inline
156+ LocalTensorBuffer()
157+ {
158+ tensor = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::C2PIPE2GM, 0, Arch::AtlasA5::FIXBUF_SIZE);
159+ }
160+};
161+ 
162+///////////////////////////////////////////////////////////
163+ 
164+/// Partial specialization for TPosition::VECIN
165+template <>
166+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::VECIN> : LocalTensorBufferBase {
167+public:
168+ static constexpr AscendC::TPosition Position = AscendC::TPosition::VECIN;
169+ 
170+ __aicore__ inline
171+ LocalTensorBuffer()
172+ {
173+ tensor = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::VECIN, 0, Arch::AtlasA5::UB_SIZE);
174+ }
175+};
176+ 
177+///////////////////////////////////////////////////////////
178+ 
179+/// Partial specialization for TPosition::VECOUT
180+template <>
181+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::VECOUT> : LocalTensorBufferBase {
182+public:
183+ static constexpr AscendC::TPosition Position = AscendC::TPosition::VECOUT;
184+ 
185+ __aicore__ inline
186+ LocalTensorBuffer()
187+ {
188+ tensor = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::VECOUT, 0, Arch::AtlasA5::UB_SIZE);
189+ }
190+};
191+ 
192+///////////////////////////////////////////////////////////
193+ 
194+/// Partial specialization for TPosition::VECCALC
195+template <>
196+struct LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::VECCALC> : LocalTensorBufferBase {
197+public:
198+ static constexpr AscendC::TPosition Position = AscendC::TPosition::VECCALC;
199+ 
200+ __aicore__ inline
201+ LocalTensorBuffer()
202+ {
203+ tensor = AscendC::LocalTensor<uint8_t>(AscendC::TPosition::VECCALC, 0, Arch::AtlasA5::UB_SIZE);
204+ }
205+};
206+ 
207+ 
208+/// Partial specialization for TPosition::A1
209+template <>
210+struct LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::A1> : LocalTensorBufferBase {
211+public:
212+ static constexpr AscendC::TPosition Position = AscendC::TPosition::A1;
213+ 
214+ __aicore__ inline
215+ LocalTensorBuffer()
216+ {
217+ AscendC::TBuf<AscendC::TPosition::A1> tbufA1;
218+ GetTPipePtr()->InitBuffer(tbufA1, Arch::AtlasA2::L1_SIZE);
219+ tensor = tbufA1.Get<uint8_t>();
220+ }
221+};
222+ 
223+///////////////////////////////////////////////////////////
224+ 
225+/// Partial specialization for TPosition::A2
226+template <>
227+struct LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::A2> : LocalTensorBufferBase {
228+public:
229+ static constexpr AscendC::TPosition Position = AscendC::TPosition::A2;
230+ 
231+ __aicore__ inline
232+ LocalTensorBuffer()
233+ {
234+ AscendC::TBuf<AscendC::TPosition::A2> tbufA2;
235+ GetTPipePtr()->InitBuffer(tbufA2, Arch::AtlasA2::L0A_SIZE);
236+ tensor = tbufA2.Get<uint8_t>();
237+ }
238+};
239+ 
240+///////////////////////////////////////////////////////////
241+ 
242+/// Partial specialization for TPosition::B1
243+template <>
244+struct LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::B1> : LocalTensorBufferBase {
245+public:
246+ static constexpr AscendC::TPosition Position = AscendC::TPosition::B1;
247+ 
248+ __aicore__ inline
249+ LocalTensorBuffer()
250+ {
251+ AscendC::TBuf<AscendC::TPosition::B1> tbufB1;
252+ GetTPipePtr()->InitBuffer(tbufB1, Arch::AtlasA2::L1_SIZE);
253+ tensor = tbufB1.Get<uint8_t>();
254+ }
255+};
256+ 
257+///////////////////////////////////////////////////////////
258+ 
259+/// Partial specialization for AtlasA2, TPosition::B2
260+template <>
261+struct LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::B2> : LocalTensorBufferBase {
262+public:
263+ static constexpr AscendC::TPosition Position = AscendC::TPosition::B2;
264+ 
265+ __aicore__ inline
266+ LocalTensorBuffer()
267+ {
268+ AscendC::TBuf<AscendC::TPosition::B2> tbufB2;
269+ GetTPipePtr()->InitBuffer(tbufB2, Arch::AtlasA2::L0B_SIZE);
270+ tensor = tbufB2.Get<uint8_t>();
271+ }
272+};
273+ 
274+///////////////////////////////////////////////////////////
275+ 
276+/// Partial specialization for AtlasA2, TPosition::C1
277+template <>
278+struct LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::C1> : LocalTensorBufferBase {
279+public:
280+ using ArchTag = Arch::AtlasA2;
281+ static constexpr AscendC::TPosition Position = AscendC::TPosition::C1;
282+ 
283+ __aicore__ inline
284+ LocalTensorBuffer()
285+ {
286+ AscendC::TBuf<AscendC::TPosition::C1> tbufC1;
287+ GetTPipePtr()->InitBuffer(tbufC1, Arch::AtlasA2::L1_SIZE);
288+ tensor = tbufC1.Get<uint8_t>();
289+ }
290+};
291+ 
292+///////////////////////////////////////////////////////////
293+ 
294+/// Partial specialization for AtlasA2, TPosition::C2
295+template <>
296+struct LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::C2> : LocalTensorBufferBase {
297+public:
298+ using ArchTag = Arch::AtlasA2;
299+ static constexpr AscendC::TPosition Position = AscendC::TPosition::C2;
300+ 
301+ __aicore__ inline
302+ LocalTensorBuffer()
303+ {
304+ AscendC::TBuf<AscendC::TPosition::C2> tbufC2;
305+ GetTPipePtr()->InitBuffer(tbufC2, Arch::AtlasA2::BIAS_SIZE);
306+ tensor = tbufC2.Get<uint8_t>();
307+ }
308+};
309+ 
310+///////////////////////////////////////////////////////////
311+ 
312+/// Partial specialization for TPosition::CO1
313+template <>
314+struct LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::CO1> : LocalTensorBufferBase {
315+public:
316+ static constexpr AscendC::TPosition Position = AscendC::TPosition::CO1;
317+ 
318+ __aicore__ inline
319+ LocalTensorBuffer()
320+ {
321+ AscendC::TBuf<AscendC::TPosition::CO1> tbufCO1;
322+ GetTPipePtr()->InitBuffer(tbufCO1, Arch::AtlasA2::L0C_SIZE);
323+ tensor = tbufCO1.Get<uint8_t>();
324+ }
325+};
326+ 
327+/// Partial specialization for TPosition::VECCALC
328+template <>
329+struct LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::VECCALC> : LocalTensorBufferBase {
330+public:
331+ static constexpr AscendC::TPosition Position = AscendC::TPosition::VECCALC;
332+ 
333+ __aicore__ inline
334+ LocalTensorBuffer()
335+ {
336+ AscendC::TBuf<AscendC::TPosition::VECCALC> tbufVECCALC;
337+ GetTPipePtr()->InitBuffer(tbufVECCALC, Arch::AtlasA2::UB_SIZE);
338+ tensor = tbufVECCALC.Get<uint8_t>();
339+ }
340+};
341+} // namespace NpuArch::Arch
342+ 
343+#endif // INCLUDE_ARCH_MEMORY_H
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/arch/resource.hpp+69-0
@@ -0,0 +1,69 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef INCLUDE_ARCH_RESOURCE_HPP
12+#define INCLUDE_ARCH_RESOURCE_HPP
13+ 
14+#include "../../attn_infra/base_defs.hpp"
15+#include "../../attn_infra/arch/local_tensor_buffer.hpp"
16+ 
17+namespace NpuArch::Arch
18+{
19+ 
20+template<class ArchTag>
21+struct Resource
22+{};
23+ 
24+template<>
25+struct Resource<Arch::AtlasA5> {
26+public:
27+ LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::A1> l1Buf;
28+ LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::A2> l0ABuf;
29+ LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::B2> l0BBuf;
30+ LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::C2> btBuf;
31+ LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::CO1> l0CBuf;
32+ LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::VECCALC> ubBuf;
33+ LocalTensorBuffer<Arch::AtlasA5, AscendC::TPosition::C2PIPE2GM> fpBuf;
34+ 
35+ __aicore__ inline
36+ Resource()
37+ {
38+ AscendC::InitSocState();
39+ }
40+ 
41+ __aicore__ inline
42+ ~Resource()
43+ {
44+ AscendC::InitSocState();
45+ }
46+};
47+ 
48+template<>
49+struct Resource<Arch::AtlasA2> {
50+public:
51+ AscendC::TPipe pipe;
52+ 
53+ LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::A1> l1Buf;
54+ LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::A2> l0ABuf;
55+ LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::B2> l0BBuf;
56+ LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::C2> btBuf;
57+ LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::CO1> l0CBuf;
58+ LocalTensorBuffer<Arch::AtlasA2, AscendC::TPosition::VECCALC> ubBuf;
59+ 
60+ __aicore__ inline
61+ Resource()
62+ {
63+ pipe.Destroy();
64+ }
65+};
66+ 
67+} // namespace NpuArch::Arch
68+ 
69+#endif // INCLUDE_ARCH_RESOURCE_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/base_defs.hpp+59-0
@@ -0,0 +1,59 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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 base_defs.hpp
13+ * \brief
14+ */
15+ 
16+#ifndef HPP_HPP
17+#define HPP_HPP
18+ 
19+#include <cstdint>
20+ 
21+#include <kernel_operator.h>
22+ 
23+#include "../attn_infra/detail/alignment.hpp"
24+#include "../attn_infra/detail/dependent_false.hpp"
25+#include "../attn_infra/detail/macros.hpp"
26+#ifndef DUMP_TENSOR
27+#define DEBUG_PRINT(str) do { AscendC::printf("[DEBUG] Line: %d %s\n", __LINE__, str); } while(0)
28+#define DUMP_TENSOR(mmResUb, m, n ) \
29+ do { AscendC::printf("line %d, tensor: %s\n", __LINE__, #mmResUb); \
30+ uint32_t array_ ## __LINE__[] = {static_cast<uint32_t>(m), static_cast<uint32_t>(n)}; \
31+ AscendC::DumpTensor(mmResUb, 2, (m)*(n)); \
32+ } while(0)
33+#endif
34+ 
35+namespace NpuArch {
36+ 
37+constexpr uint32_t BYTE_PER_C0 = 32;
38+constexpr uint32_t BYTE_PER_C2 = 64;
39+constexpr uint32_t C0_NUM_PER_FRACTAL = 16;
40+constexpr uint32_t BYTE_PER_FRACTAL = BYTE_PER_C0 * C0_NUM_PER_FRACTAL;
41+ 
42+constexpr uint32_t BYTE_PER_BLK = 32;
43+constexpr uint32_t BLK_NUM_PER_VECTOR_FRACTAL = 8;
44+constexpr uint32_t BYTE_PER_VECTOR_FRACTAL = BYTE_PER_BLK * BLK_NUM_PER_VECTOR_FRACTAL;
45+ 
46+constexpr uint64_t L2_OFFSET = 0;
47+constexpr uint32_t STRIDE_LIMIT = 65536;
48+ 
49+constexpr uint32_t BYTE_PER_BLK_FP = 128; /// datablock size of A1->C2PiPE2GM
50+ 
51+constexpr uint32_t MX_SCALE_COPY_GROUP_NUM = 2;
52+constexpr uint32_t MX_SCALE_GROUP_NUM = 32;
53+constexpr uint32_t MX_BASEK_FACTOR = 64;
54+ 
55+class EmptyClass {};
56+ 
57+} // namespace NpuArch
58+ 
59+#endif // HPP_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/coord.hpp+442-0
@@ -0,0 +1,442 @@
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 coord.hpp
13+ * \brief
14+ */
15+ 
16+#ifndef COORD_HPP
17+#define COORD_HPP
18+ 
19+#include "../attn_infra/base_defs.hpp"
20+ 
21+namespace NpuArch {
22+ 
23+/// Statically-sized array specifying Coords within a tensor
24+template <
25+ int RANK_, ///< Logical rank of coordinate
26+ class Index_ = uint32_t, ///< Index type used for each dimension
27+ class LongIndex_ = int64_t ///< Long index type used for linear offsets
28+>
29+struct Coord {
30+public:
31+ // Number of elements in Coord
32+ static const int RANK = RANK_;
33+ 
34+ // Index typen used to store elements
35+ using Index = Index_;
36+ 
37+ // Type used to represent linear offsets
38+ using LongIndex = LongIndex_;
39+ 
40+ // Default ctor initializes uniformly
41+ HOST_DEVICE constexpr
42+ explicit Coord(Index value = Index(0))
43+ {
44+ for (int i = 0; i < RANK; ++i) {
45+ idx[i] = value;
46+ }
47+ }
48+ 
49+ // Constructs from an array of integers
50+ HOST_DEVICE constexpr
51+ Coord(Index const (&idx_)[RANK])
52+ {
53+ for (int i = 0; i < RANK; ++i) {
54+ idx[i] = idx_[i];
55+ }
56+ }
57+ 
58+ HOST_DEVICE
59+ int Argmin() const
60+ {
61+ return ArgminImpl<1>(0);
62+ }
63+ 
64+ // Returns the index of the dimension with greatest value
65+ HOST_DEVICE
66+ int Argmax() const
67+ {
68+ return ArgmaxImpl<1>(0);
69+ }
70+ 
71+ // Returns true if Coord is non-zero
72+ HOST_DEVICE
73+ explicit operator bool() const
74+ {
75+ return AnyImpl<0>();
76+ }
77+ 
78+ // Return true if Coord is uniformly zero.
79+ HOST_DEVICE
80+ bool operator!() const
81+ {
82+ return !AnyImpl<0>();
83+ }
84+ 
85+ // Element-wise addition
86+ HOST_DEVICE
87+ Coord operator+(Coord const &b) const
88+ {
89+ Coord c;
90+ AddCoordImpl<0>(c, b);
91+ return c;
92+ }
93+ 
94+ // Add a scalar to each element
95+ HOST_DEVICE
96+ Coord operator+(const Index val) const
97+ {
98+ Coord c;
99+ AddScalarImpl<0>(c, val);
100+ return c;
101+ }
102+ 
103+ // Element-wise subtraction
104+ HOST_DEVICE
105+ Coord operator-(Coord const &b) const
106+ {
107+ Coord c;
108+ SubCoordImpl<0>(c, b);
109+ return c;
110+ }
111+ 
112+ // Subtract a scalar from each element
113+ HOST_DEVICE
114+ Coord operator-(Index const val) const
115+ {
116+ Coord c;
117+ SubScalarImpl<0>(c, val);
118+ return c;
119+ }
120+ 
121+ // Element-wise multiply
122+ HOST_DEVICE
123+ Coord operator*(Coord const &b) const
124+ {
125+ Coord c;
126+ MulCoordImpl<0>(c, b);
127+ return c;
128+ }
129+ 
130+ // Element-wise division
131+ HOST_DEVICE
132+ Coord operator/(Coord const &b) const
133+ {
134+ Coord c;
135+ DivCoordImpl<0>(c, b);
136+ return c;
137+ }
138+ 
139+ // Element-wise mod
140+ HOST_DEVICE
141+ Coord operator%(Coord const &b) const
142+ {
143+ Coord c;
144+ ModCoordImpl<0>(c, b);
145+ return c;
146+ }
147+ 
148+ // In-place addition
149+ HOST_DEVICE
150+ Coord &operator+=(Coord const &b)
151+ {
152+ PlusEqualImpl<0>(b);
153+ return *this;
154+ }
155+ 
156+ // In-place equal
157+ HOST_DEVICE
158+ bool operator==(Coord const &b) const
159+ {
160+ return EqualCoordImpl<0>(b);
161+ }
162+ 
163+ // In-place equal
164+ HOST_DEVICE
165+ bool operator==(Index const val) const
166+ {
167+ return EqualScalarImpl<0>(val);
168+ }
169+ 
170+ // Member acces operator
171+ HOST_DEVICE
172+ Index &operator[](int dim)
173+ {
174+ return idx[dim];
175+ }
176+ 
177+ // Member access operator
178+ HOST_DEVICE
179+ Index const &operator[](int dim) const
180+ {
181+ return idx[dim];
182+ }
183+ 
184+ // Gets the index of a given Coord element
185+ template <int DIM>
186+ HOST_DEVICE
187+ Index &At()
188+ {
189+ return idx[DIM];
190+ }
191+ 
192+ // Access via index; may limit unrolling potential
193+ HOST_DEVICE
194+ Index &At(int dim)
195+ {
196+ return idx[dim];
197+ }
198+ 
199+ // Gets the index of a given Coord element
200+ template <int DIM>
201+ HOST_DEVICE
202+ Index const &At() const
203+ {
204+ return idx[DIM];
205+ }
206+ 
207+ // Access via index; may limit unrolling potential
208+ HOST_DEVICE
209+ Index const &At(int dim) const
210+ {
211+ return idx[dim];
212+ }
213+ 
214+ template <int... Is>
215+ HOST_DEVICE
216+ auto GetCoordByAxis() const
217+ {
218+ Index idx_[sizeof...(Is)]{idx[Is]...};
219+ return Coord<sizeof...(Is), Index, LongIndex>{idx_};
220+ }
221+ 
222+ HOST_DEVICE
223+ static Coord Min(Coord const &a, Coord const &b)
224+ {
225+ Coord res;
226+ for (int i = 0; i < RANK; ++i) {
227+ res[i] = a[i] < b[i] ? a[i] : b[i];
228+ }
229+ return res;
230+ }
231+ 
232+private:
233+ template <int N>
234+ HOST_DEVICE
235+ int ArgminImpl(int i) const
236+ {
237+ if constexpr (N == RANK) {
238+ return i;
239+ }
240+ else {
241+ return ArgminImpl<N + 1>(idx[N] < idx[i] ? N : i);
242+ }
243+ }
244+ 
245+ template <int N>
246+ HOST_DEVICE
247+ int ArgmaxImpl(int i) const
248+ {
249+ if constexpr (N == RANK) {
250+ return i;
251+ }
252+ else {
253+ return ArgmaxImpl<N + 1>(idx[N] > idx[i] ? N : i);
254+ }
255+ }
256+
257+ template <int N>
258+ HOST_DEVICE
259+ bool AnyImpl() const
260+ {
261+ if constexpr (N == RANK) {
262+ return false;
263+ }
264+ else {
265+ return idx[N] || AnyImpl<N + 1>();
266+ }
267+ }
268+ 
269+ template <int N>
270+ HOST_DEVICE
271+ void AddCoordImpl(Coord &c, Coord const &b) const
272+ {
273+ if constexpr (N < RANK) {
274+ c.idx[N] = idx[N] + b.idx[N];
275+ AddCoordImpl<N + 1>(c, b);
276+ }
277+ }
278+ 
279+ template <int N>
280+ HOST_DEVICE
281+ void AddScalarImpl(Coord &c, Index const val) const
282+ {
283+ if constexpr (N < RANK) {
284+ c.idx[N] = idx[N] + val;
285+ AddScalarImpl<N + 1>(c, val);
286+ }
287+ }
288+ 
289+ template <int N>
290+ HOST_DEVICE
291+ void SubCoordImpl(Coord &c, Coord const &b) const
292+ {
293+ if constexpr (N < RANK) {
294+ c.idx[N] = idx[N] - b.idx[N];
295+ SubCoordImpl<N + 1>(c, b);
296+ }
297+ }
298+ 
299+ template <int N>
300+ HOST_DEVICE
301+ void SubScalarImpl(Coord &c, Index const val) const
302+ {
303+ if constexpr (N < RANK) {
304+ c.idx[N] = idx[N] - val;
305+ SubScalarImpl<N + 1>(c, val);
306+ }
307+ }
308+ 
309+ template <int N>
310+ HOST_DEVICE
311+ void MulCoordImpl(Coord &c, Coord const &b) const
312+ {
313+ if constexpr (N < RANK) {
314+ c.idx[N] = idx[N] * b.idx[N];
315+ MulCoordImpl<N + 1>(c, b);
316+ }
317+ }
318+ 
319+ template <int N>
320+ HOST_DEVICE
321+ void DivCoordImpl(Coord &c, Coord const &b) const
322+ {
323+ if constexpr (N < RANK) {
324+ c.idx[N] = idx[N] / b.idx[N];
325+ DivCoordImpl<N + 1>(c, b);
326+ }
327+ }
328+ 
329+ template <int N>
330+ HOST_DEVICE
331+ void ModCoordImpl(Coord &c, Coord const &b) const
332+ {
333+ if constexpr (N < RANK) {
334+ c.idx[N] = idx[N] % b.idx[N];
335+ ModCoordImpl<N + 1>(c, b);
336+ }
337+ }
338+ 
339+ template <int N>
340+ HOST_DEVICE
341+ void PlusEqualImpl(Coord const &b)
342+ {
343+ if constexpr (N < RANK) {
344+ idx[N] += b.idx[N];
345+ PlusEqualImpl<N + 1>(b);
346+ }
347+ }
348+ 
349+ template <int N>
350+ HOST_DEVICE
351+ bool EqualCoordImpl(Coord const &b) const
352+ {
353+ if constexpr (N == RANK) {
354+ return true;
355+ }
356+ else {
357+ return idx[N] == b.idx[N] && EqualCoordImpl<N + 1>(b);
358+ }
359+ }
360+ 
361+ template <int N>
362+ HOST_DEVICE
363+ bool EqualScalarImpl(Index const val) const
364+ {
365+ if constexpr (N == RANK) {
366+ return true;
367+ }
368+ else {
369+ return idx[N] == val && EqualScalarImpl<N + 1>(val);
370+ }
371+ }
372+ 
373+ // Indices
374+ Index idx[RANK];
375+};
376+ 
377+// Helper to make a 1-element coordinate
378+template <class T>
379+HOST_DEVICE constexpr
380+Coord<1, T> MakeCoord(T dim0)
381+{
382+ T values[1] = {dim0};
383+ return Coord<1, T>(values);
384+}
385+ 
386+/// Helper to make a 2-element coordinate
387+template <class T>
388+HOST_DEVICE constexpr
389+Coord<2, T> MakeCoord(T dim0, T dim1)
390+{
391+ T values[2] = {dim0, dim1};
392+ return Coord<2, T>(values);
393+}
394+ 
395+/// Helper to make a 3-element coordinate
396+template <class T>
397+HOST_DEVICE constexpr
398+Coord<3, T> MakeCoord(T dim0, T dim1, T dim2)
399+{
400+ T values[3] = {dim0, dim1, dim2};
401+ return Coord<3, T>(values);
402+}
403+ 
404+/// Helper to make a 4-element coordinate
405+template <class T>
406+HOST_DEVICE constexpr
407+Coord<4, T> MakeCoord(T dim0, T dim1, T dim2, T dim3)
408+{
409+ T values[4] = {dim0, dim1, dim2, dim3};
410+ return Coord<4, T>(values);
411+}
412+ 
413+/// Helper to make a 5-element coordinate
414+template <class T>
415+HOST_DEVICE constexpr
416+Coord<5, T> MakeCoord(T dim0, T dim1, T dim2, T dim3, T dim4)
417+{
418+ T values[5] = {dim0, dim1, dim2, dim3, dim4};
419+ return Coord<5, T>(values);
420+}
421+ 
422+/// Helper to make a 6-element coordinate
423+template <class T>
424+HOST_DEVICE constexpr
425+Coord<6, T> MakeCoord(T dim0, T dim1, T dim2, T dim3, T dim4, T dim5)
426+{
427+ T values[6] = {dim0, dim1, dim2, dim3, dim4, dim5};
428+ return Coord<6, T>(values);
429+}
430+ 
431+/// Helper to make a 7-element coordinate
432+template <class T>
433+HOST_DEVICE constexpr
434+Coord<7, T> MakeCoord(T dim0, T dim1, T dim2, T dim3, T dim4, T dim5, T dim6)
435+{
436+ T values[7] = {dim0, dim1, dim2, dim3, dim4, dim5, dim6};
437+ return Coord<7, T>(values);
438+}
439+ 
440+} // namespace NpuArch
441+ 
442+#endif // COORD_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/detail/alignment.hpp+129-0
@@ -0,0 +1,129 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef ALIGNMENT_HPP
12+#define ALIGNMENT_HPP
13+ 
14+#include <limits>
15+ 
16+#include "../../attn_infra/detail/macros.hpp"
17+#include "../../tla/numeric/integral_constant.hpp"
18+ 
19+template <uint32_t ALIGN, typename T>
20+HOST_DEVICE
21+constexpr T RoundDown(const T val)
22+{
23+ static_assert(ALIGN != 0U, "ALIGN must not be 0");
24+ return val / ALIGN * ALIGN;
25+}
26+ 
27+template <class T, class U>
28+HOST_DEVICE
29+constexpr auto RoundDown(T const &val, U const &align)
30+{
31+ if constexpr (tla::is_static<T>::value && tla::is_static<U>::value) { // Int, Int
32+ constexpr uint32_t res = T::value / U::value * U::value;
33+ return tla::Int<res>{};
34+ } else if constexpr (tla::is_static<T>::value) { // Int, int
35+ return T::value / align * align;
36+ } else if constexpr (tla::is_static<U>::value) { // int, Int
37+ return val / U::value * U::value;
38+ } else { // int, int
39+ return val / align * align;
40+ }
41+}
42+ 
43+template <uint32_t ALIGN, typename T = uint32_t>
44+HOST_DEVICE
45+constexpr T RoundUp(const T val)
46+{
47+ static_assert(ALIGN != 0U, "ALIGN must not be 0");
48+ T align = ALIGN;
49+ if (val + align - 1 < val) {
50+ return val;
51+ }
52+ return (val + align - 1) / align * align;
53+}
54+ 
55+template <class T, class U>
56+HOST_DEVICE
57+constexpr auto RoundUp(T const &val, U const &align)
58+{
59+ if constexpr (tla::is_static<T>::value && tla::is_static<U>::value) { // Int, Int
60+ constexpr uint32_t res = (T::value + U::value - 1) / U::value * U::value;
61+ return tla::Int<res>{};
62+ } else if constexpr (tla::is_static<T>::value) { // Int, int
63+ return (T::value + align - 1) / align * align;
64+ } else if constexpr (tla::is_static<U>::value) { // int, Int
65+ return (val + U::value - 1) / U::value * U::value;
66+ } else { // int, int
67+ return (val + align - 1) / align * align;
68+ }
69+}
70+ 
71+template <uint32_t DIVISOR, typename T = uint32_t>
72+HOST_DEVICE
73+constexpr T CeilDiv(const T dividend)
74+{
75+ static_assert(DIVISOR != 0U, "DIVISOR must not be 0");
76+ T divisor = DIVISOR;
77+ if (dividend + divisor - 1 < dividend) {
78+ return dividend;
79+ }
80+ return (dividend + divisor - 1) / divisor;
81+}
82+ 
83+template <class T>
84+HOST_DEVICE
85+constexpr T CeilDiv(const T dividend, const T divisor)
86+{
87+ if (divisor == 0 || dividend + divisor - 1 < dividend) {
88+ return std::numeric_limits<T>::max();
89+ }
90+ return (dividend + divisor - 1) / divisor;
91+}
92+ 
93+template <class T, class U>
94+HOST_DEVICE
95+constexpr auto CeilDiv(T const &dividend, U const &divisor)
96+{
97+ if constexpr (tla::is_static<T>::value && tla::is_static<U>::value) { // Int, Int
98+ constexpr uint32_t res = (T::value + U::value - 1) / U::value;
99+ return tla::Int<res>{};
100+ } else if constexpr (tla::is_static<T>::value) { // Int, int
101+ return (T::value + divisor - 1) / divisor;
102+ } else if constexpr (tla::is_static<U>::value) { // int, Int
103+ return (dividend + U::value - 1) / U::value;
104+ } else { // int, int
105+ return (dividend + divisor - 1) / divisor;
106+ }
107+}
108+ 
109+template <class T, class U>
110+HOST_DEVICE
111+constexpr auto Max(T const &a, U const &b) {
112+ if (a > b) {
113+ return a;
114+ } else {
115+ return b;
116+ }
117+}
118+ 
119+template <class T, class U>
120+HOST_DEVICE
121+constexpr auto Min(T const &a, U const &b) {
122+ if (a < b) {
123+ return a;
124+ } else {
125+ return b;
126+ }
127+}
128+ 
129+#endif // ALIGNMENT_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/detail/dependent_false.hpp+20-0
@@ -0,0 +1,20 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef DETAIL_DEPENDENT_FALSE_HPP
12+#define DETAIL_DEPENDENT_FALSE_HPP
13+ 
14+template <bool VALUE, class... Args>
15+constexpr bool DEPENDENT_BOOL_VALUE = VALUE;
16+ 
17+template <class... Args>
18+constexpr bool DEPENDENT_FALSE = DEPENDENT_BOOL_VALUE<false, Args...>;
19+ 
20+#endif // DETAIL_DEPENDENT_FALSE_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/detail/macros.hpp+16-0
@@ -0,0 +1,16 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef DETAIL_MACROS_HPP
12+#define DETAIL_MACROS_HPP
13+ 
14+#define HOST_DEVICE __forceinline__ [host, aicore]
15+ 
16+#endif // DETAIL_MACROS_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/detail/tag_to_layout.hpp+106-0
@@ -0,0 +1,106 @@
1+/**
2+ * Copyright (c) 2025 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+#ifndef DETAIL_TAG_TO_LAYOUT_HPP
12+#define DETAIL_TAG_TO_LAYOUT_HPP
13+ 
14+#include "../../attn_infra/layout/layout.hpp"
15+#include "../../tla/layout.hpp"
16+ 
17+////////////////////////////////////////////////////////////////////////////////////////////////////
18+ 
19+namespace NpuArch::detail {
20+////////////////////////////////////////////////////////////////////////////////////////////////////
21+// For each NpuArch::layout, provides its corresponding tla layout types
22+template <class Element, class LayoutTag>
23+struct TagToLayout {
24+ using type = LayoutTag;
25+};
26+ 
27+template <class Element>
28+struct TagToLayout<Element, layout::RowMajor> {
29+ using type = tla::Layout<tla::Shape<uint32_t, uint32_t>, tla::Stride<int64_t, tla::Int<1>>>;
30+};
31+ 
32+template <class Element>
33+struct TagToLayout<Element, layout::ColumnMajor> {
34+ using type = tla::Layout<tla::Shape<uint32_t, uint32_t>, tla::Stride<tla::Int<1>, int64_t>>;
35+};
36+ 
37+template <class Element>
38+struct TagToLayout<Element, layout::VectorLayout> {
39+ using type = tla::Layout<tla::Shape<uint32_t>, tla::Stride<tla::Int<1>>>;
40+};
41+ 
42+template <class Element>
43+struct TagToLayout<Element, layout::zN> {
44+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
45+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
46+ using type = tla::Layout<
47+ tla::Shape<tla::Shape<tla::Int<C0_NUM_PER_FRACTAL>, uint32_t>, tla::Shape<tla::Int<ELE_NUM_PER_C0>, uint32_t>>,
48+ tla::Stride<tla::Stride<tla::Int<ELE_NUM_PER_C0>, tla::Int<ELE_NUM_PER_FRACTAL>>,
49+ tla::Stride<tla::Int<1>, int64_t>>>;
50+};
51+ 
52+template <class Element>
53+struct TagToLayout<Element, layout::zZ> {
54+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
55+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
56+ using type = tla::Layout<
57+ tla::Shape<tla::Shape<tla::Int<C0_NUM_PER_FRACTAL>, uint32_t>, tla::Shape<tla::Int<ELE_NUM_PER_C0>, uint32_t>>,
58+ tla::Stride<tla::Stride<tla::Int<ELE_NUM_PER_C0>, int64_t>,
59+ tla::Stride<tla::Int<1>, tla::Int<ELE_NUM_PER_FRACTAL>>>>;
60+};
61+ 
62+template <class Element>
63+struct TagToLayout<Element, layout::nZ> {
64+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
65+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
66+ using type = tla::Layout<
67+ tla::Shape<tla::Shape<tla::Int<ELE_NUM_PER_C0>, uint32_t>, tla::Shape<tla::Int<C0_NUM_PER_FRACTAL>, uint32_t>>,
68+ tla::Stride<tla::Stride<tla::Int<1>, int64_t>,
69+ tla::Stride<tla::Int<ELE_NUM_PER_C0>, tla::Int<ELE_NUM_PER_FRACTAL>>>>;
70+};
71+ 
72+template <>
73+struct TagToLayout<AscendC::fp8_e8m0_t, layout::zZ> {
74+ static constexpr uint32_t ELE_NUM_PER_C0 = 2;
75+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = 32;
76+ using type = tla::Layout<
77+ tla::Shape<tla::Shape<tla::Int<C0_NUM_PER_FRACTAL>, uint32_t>, tla::Shape<tla::Int<ELE_NUM_PER_C0>, uint32_t>>,
78+ tla::Stride<tla::Stride<tla::Int<ELE_NUM_PER_C0>, int64_t>,
79+ tla::Stride<tla::Int<1>, tla::Int<ELE_NUM_PER_FRACTAL>>>>;
80+};
81+ 
82+template <>
83+struct TagToLayout<AscendC::fp8_e8m0_t, layout::nN> {
84+ static constexpr uint32_t ELE_NUM_PER_C0 = 2;
85+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = 32;
86+ using type = tla::Layout<
87+ tla::Shape<tla::Shape<tla::Int<ELE_NUM_PER_C0>, uint32_t>, tla::Shape<tla::Int<C0_NUM_PER_FRACTAL>, uint32_t>>,
88+ tla::Stride<tla::Stride<tla::Int<1>, tla::Int<ELE_NUM_PER_FRACTAL>>,
89+ tla::Stride<tla::Int<ELE_NUM_PER_C0>, int64_t>>>;
90+};
91+ 
92+// Convenience aliases
93+template <class Element, class LayoutTag>
94+using TagToLayout_t = typename TagToLayout<Element, LayoutTag>::type;
95+ 
96+constexpr uint32_t ELE_NUM_PER_FRACTAL_L0C = 256;
97+using LayoutL0C = tla::Layout<
98+ tla::Shape<tla::Shape<tla::Int<C0_NUM_PER_FRACTAL>, uint32_t>, tla::Shape<tla::Int<C0_NUM_PER_FRACTAL>, uint32_t>>,
99+ tla::Stride<tla::Stride<tla::Int<C0_NUM_PER_FRACTAL>, tla::Int<ELE_NUM_PER_FRACTAL_L0C>>,
100+ tla::Stride<tla::Int<1>, int64_t>>>;
101+ 
102+////////////////////////////////////////////////////////////////////////////////////////////////////
103+ 
104+} // namespace NpuArch::detail
105+ 
106+#endif // CATLASS_DETAIL_TAG_TO_LAYOUT_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue.hpp+39-0
@@ -0,0 +1,39 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_HPP
12+#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+ 
16+namespace NpuArch::Epilogue::Block {
17+ 
18+template <
19+ class DispatchPolicy,
20+ class... Args
21+>
22+class BlockEpilogue {
23+ static_assert(DEPENDENT_FALSE<DispatchPolicy>, "Could not find an epilogue specialization");
24+};
25+ 
26+} // namespace NpuArch::Epilogue::Block
27+#if (__CCE_AICORE__ == 220)
28+#include "../../../attn_infra/epilogue/block/block_epilogue_online_softmax.hpp"
29+#include "../../../attn_infra/epilogue/block/block_epilogue_online_softmax_low_prec.hpp"
30+#include "../../../attn_infra/epilogue/block/block_epilogue_rescale_o.hpp"
31+#include "../../../attn_infra/epilogue/block/block_epilogue_rescale_o_low_prec.hpp"
32+#endif
33+#if (__CCE_AICORE__ == 310)
34+#include "../../../attn_infra/epilogue/block/block_epilogue_mask2idx_arch35.hpp"
35+#include "../../../attn_infra/epilogue/block/block_epilogue_rescale_o_arch35_reg_high_prec.hpp"
36+#include "../../../attn_infra/epilogue/block/block_epilogue_online_softmax_arch35_reg_low_prec.hpp"
37+#include "../../../attn_infra/epilogue/block/block_epilogue_online_softmax_arch35_reg_low_prec_bf16.hpp"
38+#endif
39+#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue_mask2idx_arch35.hpp+199-0
@@ -0,0 +1,199 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_MASK2IDX_ARCH35
12+#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_MASK2IDX_ARCH35
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/resource.hpp"
16+#include "../../../attn_infra/epilogue/dispatch_policy.hpp"
17+#include "../../../attn_infra/gemm_coord.hpp"
18+#include "../../../attn_infra/matrix_coord.hpp"
19+ 
20+namespace NpuArch::Epilogue::Block {
21+ 
22+template <
23+ class ElementSparseMask_,
24+ class ElementSparseIdx_,
25+ class ElementSparseCount_
26+>
27+class BlockEpilogue <
28+ EpilogueBsaMask2Idx,
29+ ElementSparseMask_,
30+ ElementSparseIdx_,
31+ ElementSparseCount_>
32+{
33+public:
34+ using DispatchPolicy = EpilogueBsaMask2Idx;
35+ using ArchTag = typename DispatchPolicy::ArchTag;
36+ using ElementSparseMask = ElementSparseMask_;
37+ using ElementSparseIdx = ElementSparseIdx_;
38+ using ElementSparseCount = ElementSparseCount_;
39+ 
40+ static constexpr uint32_t IO_STAGES = DispatchPolicy::IO_STAGES;
41+ static constexpr uint32_t PRE_ROW_TILE = 128;
42+ static constexpr uint32_t PRE_COL_TILE = 64;
43+ static constexpr uint32_t PRE_ELEM_NUM_PER_LOOP = PRE_ROW_TILE * PRE_COL_TILE;
44+ 
45+ __aicore__ inline
46+ BlockEpilogue(Arch::Resource<ArchTag> &resource)
47+ {
48+ constexpr uint32_t MASK_PAT_IN_UINT8 = 0;
49+ constexpr uint32_t MASK_PAT_IN_FP16 = 2 * PRE_ELEM_NUM_PER_LOOP;
50+ constexpr uint32_t MASK_PAT_IN_FP32 = 4 * PRE_ELEM_NUM_PER_LOOP; // 2(db)+2
51+ constexpr uint32_t MASK_PAT_IN_BIT = 8 * PRE_ELEM_NUM_PER_LOOP; // 1+2+4
52+ constexpr uint32_t MASK_IDX = 9 * PRE_ELEM_NUM_PER_LOOP; // 1+2+4+1
53+ constexpr uint32_t RSVD_SPARSE_IDX = 13 * PRE_ELEM_NUM_PER_LOOP; // 1+2+4+1+4
54+ constexpr uint32_t RSVD_SPARSE_COUNT = 21 * PRE_ELEM_NUM_PER_LOOP; // 1+2+4+1+4+8(db)
55+ 
56+ for (uint32_t i = 0; i < IO_STAGES; i++) {
57+ maskPatUb8[i] = resource.ubBuf.template GetBufferByByte<ElementSparseMask>(
58+ MASK_PAT_IN_UINT8 + i * PRE_ELEM_NUM_PER_LOOP * sizeof(ElementSparseMask));
59+ sparseIdxUb[i] = resource.ubBuf.template GetBufferByByte<ElementSparseIdx>(
60+ RSVD_SPARSE_IDX + i * PRE_ELEM_NUM_PER_LOOP * sizeof(ElementSparseIdx));
61+ sparseCountUb[i] = resource.ubBuf.template GetBufferByByte<ElementSparseCount>(
62+ RSVD_SPARSE_COUNT + i * PRE_ROW_TILE * sizeof(ElementSparseCount));
63+ }
64+ maskPatUb16 = resource.ubBuf.template GetBufferByByte<half>(MASK_PAT_IN_FP16);
65+ maskPatUb32 = resource.ubBuf.template GetBufferByByte<float>(MASK_PAT_IN_FP32);
66+ maskPatInBitUb8 = resource.ubBuf.template GetBufferByByte<uint8_t>(MASK_PAT_IN_BIT);
67+ maskPatInBitUb32 = resource.ubBuf.template GetBufferByByte<uint32_t>(MASK_PAT_IN_BIT);
68+ maskIdxUb = resource.ubBuf.template GetBufferByByte<ElementSparseIdx>(MASK_IDX);
69+ }
70+ 
71+ __aicore__ inline
72+ void operator()(
73+ AscendC::GlobalTensor<ElementSparseMask> gSparseMask,
74+ AscendC::GlobalTensor<ElementSparseIdx> gSparseIdx,
75+ AscendC::GlobalTensor<ElementSparseCount> gSparseCount,
76+ uint32_t totalRowNumBlockMask,
77+ uint32_t yBlockNumAligned,
78+ uint32_t avgRowPerSubCore,
79+ uint32_t preActiveSubCoreNum)
80+ {
81+ uint32_t subCoreIdx = AscendC::GetBlockIdx();
82+ uint32_t curSubCoreRowOffset = subCoreIdx * avgRowPerSubCore;
83+ uint32_t actDealtRow = (subCoreIdx == preActiveSubCoreNum - 1) ?
84+ (totalRowNumBlockMask - curSubCoreRowOffset) : avgRowPerSubCore;
85+
86+ if (subCoreIdx < preActiveSubCoreNum) {
87+ uint32_t rowLoop = CeilDiv(actDealtRow, PRE_ROW_TILE);
88+ uint32_t colLoop = CeilDiv(yBlockNumAligned, PRE_COL_TILE);
89+ uint32_t IdxPingPongFlag = 0;
90+ uint32_t CountPingPongFlag = 0;
91+ for (uint32_t i = 0; i < rowLoop; i++) {
92+ uint32_t curLoopRowOffset = i * PRE_ROW_TILE;
93+ uint32_t globalRowOffset = curSubCoreRowOffset + curLoopRowOffset;
94+ uint32_t actDealtRowCurLoop = (i == rowLoop - 1) ? (actDealtRow - curLoopRowOffset) : PRE_ROW_TILE;
95+ 
96+ uint64_t rsvdCountPerRow[PRE_ROW_TILE] = {0};
97+ uint64_t rsvdCountPerRowCurColLoop[PRE_ROW_TILE] = {0};
98+ for (uint32_t j = 0; j < colLoop; j++) {
99+ uint32_t curLoopColOffset = j * PRE_COL_TILE;
100+ uint32_t actDealtColCurLoop =
101+ (j == colLoop - 1) ? (yBlockNumAligned - curLoopColOffset) : PRE_COL_TILE;
102+ uint32_t actDealtColCurLoop32 = CeilDiv(actDealtColCurLoop, 32) * 32;
103+ AscendC::CreateVecIndex(
104+ maskIdxUb, static_cast<int32_t>(curLoopColOffset), actDealtColCurLoop);
105+ AscendC::PipeBarrier<PIPE_V>();
106+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(IdxPingPongFlag);
107+ AscendC::DataCopyPad(
108+ maskPatUb8[IdxPingPongFlag],
109+ gSparseMask[globalRowOffset * yBlockNumAligned + curLoopColOffset],
110+ AscendC::DataCopyExtParams(
111+ actDealtRowCurLoop,
112+ actDealtColCurLoop * sizeof(ElementSparseMask),
113+ (yBlockNumAligned - actDealtColCurLoop) * sizeof(ElementSparseMask),
114+ 0, 0),
115+ AscendC::DataCopyPadExtParams<ElementSparseMask>(
116+ true, 0, (actDealtColCurLoop32 - actDealtColCurLoop), 0));
117+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(0);
118+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(0);
119+ AscendC::Cast(
120+ maskPatUb16, maskPatUb8[IdxPingPongFlag],
121+ AscendC::RoundMode::CAST_NONE, actDealtRowCurLoop * actDealtColCurLoop32);
122+ AscendC::PipeBarrier<PIPE_V>();
123+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(IdxPingPongFlag);
124+ 
125+ AscendC::Cast(
126+ maskPatUb32, maskPatUb16,
127+ AscendC::RoundMode::CAST_NONE, actDealtRowCurLoop * actDealtColCurLoop32);
128+ AscendC::PipeBarrier<PIPE_V>();
129+ for (uint32_t k = 0; k < actDealtRowCurLoop; k++) {
130+ AscendC::CompareScalar(
131+ maskPatInBitUb8[k * PRE_COL_TILE],
132+ maskPatUb32[k * actDealtColCurLoop32],
133+ static_cast<float>(1.0), AscendC::CMPMODE::GE, actDealtColCurLoop);
134+ AscendC::PipeBarrier<PIPE_V>();
135+ if (k == 0) {
136+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(IdxPingPongFlag);
137+ }
138+ AscendC::GatherMask(
139+ sparseIdxUb[IdxPingPongFlag][k * PRE_COL_TILE],
140+ maskIdxUb,
141+ maskPatInBitUb32[k * PRE_COL_TILE / 4],
142+ true, actDealtColCurLoop, {1, 1, 0, 0}, rsvdCountPerRowCurColLoop[k]);
143+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(0);
144+ 
145+ 
146+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(0);
147+ if (k == 0) {
148+ AscendC::WaitFlag<AscendC::HardEvent::S_MTE3>(0);
149+ }
150+ AscendC::DataCopyPad(
151+ gSparseIdx[
152+ globalRowOffset * yBlockNumAligned + k * yBlockNumAligned +
153+ rsvdCountPerRow[k]],
154+ sparseIdxUb[IdxPingPongFlag][k * PRE_COL_TILE],
155+ AscendC::DataCopyExtParams(
156+ 1, rsvdCountPerRowCurColLoop[k] * sizeof(ElementSparseIdx), 0, 0, 0));
157+ if (k == actDealtRowCurLoop - 1) {
158+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(IdxPingPongFlag);
159+ }
160+ AscendC::SetFlag<AscendC::HardEvent::MTE3_S>(0);
161+ 
162+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_S>(0);
163+ rsvdCountPerRow[k] += rsvdCountPerRowCurColLoop[k];
164+ if (k == actDealtRowCurLoop - 1) {
165+ AscendC::SetFlag<AscendC::HardEvent::S_MTE3>(0);
166+ }
167+ }
168+ IdxPingPongFlag = 1 - IdxPingPongFlag;
169+ }
170+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_S>(CountPingPongFlag);
171+ for (uint32_t k = 0; k < actDealtRowCurLoop; k++) {
172+ sparseCountUb[CountPingPongFlag].SetValue(k, static_cast<int32_t>(rsvdCountPerRow[k]));
173+ }
174+ AscendC::SetFlag<AscendC::HardEvent::S_MTE3>(CountPingPongFlag + 2);
175+ AscendC::WaitFlag<AscendC::HardEvent::S_MTE3>(CountPingPongFlag + 2);
176+ AscendC::DataCopyPad(
177+ gSparseCount[globalRowOffset],
178+ sparseCountUb[CountPingPongFlag],
179+ AscendC::DataCopyExtParams(1, actDealtRowCurLoop * sizeof(ElementSparseCount), 0, 0, 0));
180+ AscendC::SetFlag<AscendC::HardEvent::MTE3_S>(CountPingPongFlag);
181+ 
182+ CountPingPongFlag = 1 - CountPingPongFlag;
183+ }
184+ }
185+ }
186+private:
187+ AscendC::LocalTensor<uint8_t> maskPatUb8[IO_STAGES];
188+ AscendC::LocalTensor<int32_t> sparseCountUb[IO_STAGES];
189+ AscendC::LocalTensor<uint8_t> maskPatInBitUb8;
190+ AscendC::LocalTensor<uint32_t> maskPatInBitUb32;
191+
192+ AscendC::LocalTensor<int32_t> maskIdxUb;
193+ AscendC::LocalTensor<int32_t> sparseIdxUb[IO_STAGES];
194+ AscendC::LocalTensor<half> maskPatUb16;
195+ AscendC::LocalTensor<float> maskPatUb32;
196+};
197+ 
198+} // namespace NpuArch::Epilogue::Block
199+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue_online_softmax.hpp+1165-0
@@ -0,0 +1,1165 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_HPP
12+#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/cross_core_sync.hpp"
16+#include "../../../attn_infra/arch/resource.hpp"
17+#include "../../../attn_infra/epilogue/dispatch_policy.hpp"
18+#include "../../../attn_infra/epilogue/tile_common/tile_copy.hpp"
19+#include "../../../attn_infra/gemm_coord.hpp"
20+#include "../../../attn_infra/matrix_coord.hpp"
21+ 
22+namespace NpuArch::Epilogue::Block {
23+ 
24+template <
25+ class OutputType_,
26+ class InputType_,
27+ class MaskType_,
28+ LseMode LSE_MODE_>
29+class BlockEpilogue<
30+ EpilogueAtlasA2OnlineSoftmax<LSE_MODE_, float>,
31+ OutputType_,
32+ InputType_,
33+ MaskType_>
34+{
35+public:
36+ using DispatchPolicy = EpilogueAtlasA2OnlineSoftmax<LSE_MODE_, float>;
37+ using ArchTag = typename DispatchPolicy::ArchTag;
38+ using ElementOutput = typename OutputType_::Element;
39+ using ElementInput = typename InputType_::Element;
40+ using ElementMask = typename MaskType_::Element;
41+ 
42+ using LayoutOutput = typename OutputType_::Layout;
43+ using LayoutInput = typename InputType_::Layout;
44+ using LayoutMask = typename MaskType_::Layout;
45+ 
46+ static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE;
47+ 
48+ static constexpr uint32_t BLOCK_SIZE_IN_BYTE = 32;
49+ static constexpr uint32_t REPEAT_SIZE_IN_BYTE = 256;
50+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
51+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
52+ static constexpr uint32_t HALF_VECTOR_SIZE = 128;
53+ static constexpr uint32_t BLOCK_SIZE = 16;
54+ static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024;
55+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384;
56+ static constexpr uint32_t VECTOR_SIZE = 128;
57+ static constexpr uint32_t MAX_UB_S_ELEM_NUM = 8192;
58+ 
59+ static constexpr uint32_t REDUCE_UB_SIZE = 1024;
60+ static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32;
61+ static constexpr uint32_t ROW_OPS_SPEC_MASK_16 = 16;
62+ static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4;
63+ static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256;
64+ static constexpr int64_t UB_FLOAT_LINE_SIZE = 64;
65+ 
66+ __aicore__ inline
67+ BlockEpilogue(Arch::Resource<ArchTag> &resource, float scaleValue_)
68+ {
69+ // Allocate UB space
70+ constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;
71+ constexpr uint32_t LP_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
72+ constexpr uint32_t MASK_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
73+ constexpr uint32_t MASK32_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
74+ 
75+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE;
76+ constexpr uint32_t LM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 8 * UB_UINT8_VECTOR_SIZE;
77+ 
78+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE;
79+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE;
80+ constexpr uint32_t LL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 11 * UB_UINT8_VECTOR_SIZE;
81+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
82+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE;
83+ 
84+ constexpr uint32_t MASK16_UB_TENSOR_OFFSET = 11 * UB_UINT8_BLOCK_SIZE;
85+ 
86+ scaleValue = scaleValue_;
87+ lsUbTensor = resource.ubBuf.template GetBufferByByte<float>(LS_UB_TENSOR_OFFSET);
88+ lpUbTensor = resource.ubBuf.template GetBufferByByte<ElementOutput>(LP_UB_TENSOR_OFFSET);
89+ maskUbTensor = resource.ubBuf.template GetBufferByByte<ElementMask>(MASK_UB_TENSOR_OFFSET);
90+ maskUbTensor16 = resource.ubBuf.template GetBufferByByte<half>(MASK16_UB_TENSOR_OFFSET);
91+ maskUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(MASK32_UB_TENSOR_OFFSET);
92+ lmUbTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_TENSOR_OFFSET);
93+ hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
94+ gmUbTensor = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
95+ dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
96+ llUbTensor = resource.ubBuf.template GetBufferByByte<float>(LL_UB_TENSOR_OFFSET);
97+ tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
98+ glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
99+ }
100+ 
101+ __aicore__ inline
102+ ~BlockEpilogue() {}
103+ 
104+ template <typename T>
105+ __aicore__ inline T Min(T a, T b)
106+ {
107+ return (a > b) ? b : a;
108+ }
109+ 
110+ __aicore__ inline
111+ void SetVecMask(int32_t len)
112+ {
113+ uint64_t mask = 0;
114+ uint64_t one = 1;
115+ uint64_t temp = len % FLOAT_VECTOR_SIZE;
116+ for (int64_t i = 0; i < temp; i++) {
117+ mask |= one << i;
118+ }
119+ 
120+ if (len == VECTOR_SIZE || len == 0) {
121+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
122+ } else if (len >= FLOAT_VECTOR_SIZE) {
123+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
124+ } else {
125+ AscendC::SetVectorMask<int8_t>(0x0, mask);
126+ }
127+ }
128+ 
129+ __aicore__ inline
130+ void SetBlockReduceMask(int32_t len)
131+ {
132+ if (len > 8 || len < 1) {
133+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
134+ return;
135+ }
136+ uint64_t subMask = ((uint64_t)1 << len) - 1;
137+ uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask + (subMask << 56) +
138+ (subMask << 40) + (subMask << 24) + (subMask << 8);
139+ AscendC::SetVectorMask<int8_t>(maskValue, maskValue);
140+ }
141+ 
142+ __aicore__ inline
143+ void RowsumSPECTILE1024(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowsumUb,
144+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
145+ uint32_t numElemsAligned)
146+ {
147+ AscendC::BlockReduceSum<float, false>(
148+ tvUbTensor,
149+ srcUb,
150+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE,
151+ AscendC::MASK_PLACEHOLDER, // (uint64_t)0
152+ 1,
153+ 1,
154+ 8);
155+ AscendC::PipeBarrier<PIPE_V>();
156+ 
157+ AscendC::BlockReduceSum<float, false>(
158+ tvUbTensor[REDUCE_UB_SIZE],
159+ tvUbTensor,
160+ numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE,
161+ AscendC::MASK_PLACEHOLDER, // (uint64_t)0
162+ 1,
163+ 1,
164+ 8);
165+ AscendC::PipeBarrier<PIPE_V>();
166+ 
167+ SetVecMask(ROW_OPS_SPEC_MASK_16);
168+ AscendC::WholeReduceSum<float, false>(
169+ rowsumUb,
170+ tvUbTensor[REDUCE_UB_SIZE],
171+ AscendC::MASK_PLACEHOLDER, // (uint64_t)0
172+ numRowsRound,
173+ 1,
174+ 1,
175+ 2);
176+ AscendC::PipeBarrier<PIPE_V>();
177+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
178+ }
179+ 
180+ __aicore__ inline
181+ void RowsumSPECTILE512(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowsumUb,
182+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
183+ uint32_t numElemsAligned)
184+ {
185+ AscendC::BlockReduceSum<float, false>(
186+ tvUbTensor,
187+ srcUb,
188+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE,
189+ 0, 1, 1, 8);
190+ AscendC::PipeBarrier<PIPE_V>();
191+ 
192+ AscendC::BlockReduceSum<float, false>(
193+ tvUbTensor[REDUCE_UB_SIZE],
194+ tvUbTensor,
195+ numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE,
196+ 0, 1, 1, 8);
197+ AscendC::PipeBarrier<PIPE_V>();
198+ AscendC::BlockReduceSum<float, false>(
199+ rowsumUb,
200+ tvUbTensor[REDUCE_UB_SIZE],
201+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE,
202+ 0, 1, 1, 8);
203+ AscendC::PipeBarrier<PIPE_V>();
204+ }
205+ 
206+ __aicore__ inline
207+ void RowsumSPECTILE256(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowsumUb,
208+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
209+ uint32_t numElemsAligned)
210+ {
211+ AscendC::BlockReduceSum<float, false>(
212+ tvUbTensor,
213+ srcUb,
214+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE,
215+ 0, 1, 1, 8);
216+ AscendC::PipeBarrier<PIPE_V>();
217+ SetVecMask(ROW_OPS_SPEC_MASK_32);
218+ AscendC::BlockReduceSum<float, false>(
219+ tvUbTensor[REDUCE_UB_SIZE],
220+ tvUbTensor,
221+ numRowsRound,
222+ 0, 1, 1, 4);
223+ AscendC::PipeBarrier<PIPE_V>();
224+ SetBlockReduceMask(ROW_OPS_SPEC_MASK_4);
225+ AscendC::BlockReduceSum<float, false>(
226+ rowsumUb,
227+ tvUbTensor[REDUCE_UB_SIZE],
228+ CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
229+ 0, 1, 1, 8);
230+ AscendC::PipeBarrier<PIPE_V>();
231+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
232+ }
233+ 
234+ __aicore__ inline
235+ void RowsumTAILTILE(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowsumUb,
236+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
237+ uint32_t numElemsAligned)
238+ {
239+ if (numElems >= FLOAT_VECTOR_SIZE) {
240+ AscendC::BlockReduceSum<float, false>(
241+ tvUbTensor,
242+ srcUb,
243+ numRowsRound,
244+ 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
245+ AscendC::PipeBarrier<PIPE_V>();
246+ AscendC::BlockReduceSum<float, false>(
247+ rowsumUb,
248+ tvUbTensor,
249+ CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
250+ 0, 1, 1, 8);
251+ AscendC::PipeBarrier<PIPE_V>();
252+ for (uint64_t rowSumIdx = 1; rowSumIdx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowSumIdx) {
253+ AscendC::BlockReduceSum<float, false>(
254+ tvUbTensor,
255+ srcUb[rowSumIdx * FLOAT_VECTOR_SIZE],
256+ numRowsRound,
257+ 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
258+ AscendC::PipeBarrier<PIPE_V>();
259+ AscendC::BlockReduceSum<float, false>(
260+ tvUbTensor[REDUCE_UB_SIZE],
261+ tvUbTensor,
262+ CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
263+ 0, 1, 1, 8);
264+ AscendC::PipeBarrier<PIPE_V>();
265+ SetVecMask(numRowsRound);
266+ AscendC::Add<float, false>(
267+ rowsumUb,
268+ rowsumUb,
269+ tvUbTensor[REDUCE_UB_SIZE],
270+ (uint64_t)0,
271+ 1,
272+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
273+ AscendC::PipeBarrier<PIPE_V>();
274+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
275+ }
276+ }
277+ if (numElems % FLOAT_VECTOR_SIZE > 0) {
278+ SetVecMask(numElems % FLOAT_VECTOR_SIZE);
279+ AscendC::BlockReduceSum<float, false>(
280+ tvUbTensor,
281+ srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
282+ numRowsRound,
283+ 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
284+ AscendC::PipeBarrier<PIPE_V>();
285+ SetBlockReduceMask(CeilDiv(numElems % FLOAT_VECTOR_SIZE, FLOAT_BLOCK_SIZE));
286+ if (numElems < FLOAT_VECTOR_SIZE) {
287+ AscendC::BlockReduceSum<float, false>(
288+ rowsumUb,
289+ tvUbTensor,
290+ CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
291+ 0, 1, 1, 8);
292+ AscendC::PipeBarrier<PIPE_V>();
293+ } else {
294+ AscendC::BlockReduceSum<float, false>(
295+ tvUbTensor[REDUCE_UB_SIZE],
296+ tvUbTensor,
297+ CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
298+ 0, 1, 1, 8);
299+ AscendC::PipeBarrier<PIPE_V>();
300+ SetVecMask(numRowsRound);
301+ AscendC::Add<float, false>(
302+ rowsumUb,
303+ rowsumUb,
304+ tvUbTensor[REDUCE_UB_SIZE],
305+ (uint64_t)0,
306+ 1,
307+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
308+ AscendC::PipeBarrier<PIPE_V>();
309+ }
310+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
311+ }
312+ }
313+ 
314+ __aicore__ inline
315+ void RowmaxSPECTILE1024(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowmaxUb,
316+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
317+ uint32_t numElemsAligned)
318+ {
319+ AscendC::BlockReduceMax<float, false>(
320+ tvUbTensor,
321+ srcUb,
322+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE,
323+ AscendC::MASK_PLACEHOLDER, // (uint64_t)0
324+ 1,
325+ 1,
326+ 8);
327+ AscendC::PipeBarrier<PIPE_V>();
328+ 
329+ AscendC::BlockReduceMax<float, false>(
330+ tvUbTensor[REDUCE_UB_SIZE],
331+ tvUbTensor,
332+ numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE,
333+ AscendC::MASK_PLACEHOLDER, // (uint64_t)0
334+ 1,
335+ 1,
336+ 8);
337+ AscendC::PipeBarrier<PIPE_V>();
338+ 
339+ SetVecMask(ROW_OPS_SPEC_MASK_16);
340+ AscendC::WholeReduceMax<float, false>(
341+ rowmaxUb,
342+ tvUbTensor[REDUCE_UB_SIZE],
343+ AscendC::MASK_PLACEHOLDER, // (uint64_t)0
344+ numRowsRound,
345+ 1,
346+ 1,
347+ 2,
348+ AscendC::ReduceOrder::ORDER_ONLY_VALUE);
349+ AscendC::PipeBarrier<PIPE_V>();
350+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
351+ }
352+ 
353+ __aicore__ inline
354+ void RowmaxSPECTILE512(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowmaxUb,
355+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
356+ uint32_t numElemsAligned)
357+ {
358+ AscendC::BlockReduceMax<float, false>(
359+ tvUbTensor,
360+ srcUb,
361+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE,
362+ 0, 1, 1, 8);
363+ AscendC::PipeBarrier<PIPE_V>();
364+ AscendC::BlockReduceMax<float, false>(
365+ tvUbTensor[REDUCE_UB_SIZE],
366+ tvUbTensor,
367+ numRowsRound * numElemsAligned / FLOAT_BLOCK_SIZE / FLOAT_VECTOR_SIZE,
368+ 0, 1, 1, 8);
369+ AscendC::PipeBarrier<PIPE_V>();
370+ AscendC::BlockReduceMax<float, false>(
371+ rowmaxUb,
372+ tvUbTensor[REDUCE_UB_SIZE],
373+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE / FLOAT_VECTOR_SIZE,
374+ 0, 1, 1, 8);
375+ AscendC::PipeBarrier<PIPE_V>();
376+ }
377+ 
378+ __aicore__ inline
379+ void RowmaxSPECTILE256(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowmaxUb,
380+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
381+ uint32_t numElemsAligned)
382+ {
383+ AscendC::BlockReduceMax<float, false>(
384+ tvUbTensor,
385+ srcUb,
386+ numRowsRound * numElemsAligned / FLOAT_VECTOR_SIZE,
387+ 0, 1, 1, 8);
388+ AscendC::PipeBarrier<PIPE_V>();
389+ SetVecMask(ROW_OPS_SPEC_MASK_32);
390+ AscendC::BlockReduceMax<float, false>(
391+ tvUbTensor[REDUCE_UB_SIZE],
392+ tvUbTensor,
393+ numRowsRound,
394+ 0, 1, 1, 4);
395+ AscendC::PipeBarrier<PIPE_V>();
396+ SetBlockReduceMask(ROW_OPS_SPEC_MASK_4);
397+ AscendC::BlockReduceMax<float, false>(
398+ rowmaxUb,
399+ tvUbTensor[REDUCE_UB_SIZE],
400+ CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
401+ 0, 1, 1, 8);
402+ AscendC::PipeBarrier<PIPE_V>();
403+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
404+ }
405+ 
406+ __aicore__ inline
407+ void RowmaxTAILTILE(const AscendC::LocalTensor<float> &srcUb, const AscendC::LocalTensor<float> &rowmaxUb,
408+ const AscendC::LocalTensor<float> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
409+ uint32_t numElemsAligned)
410+ {
411+ if (numElems >= FLOAT_VECTOR_SIZE) {
412+ AscendC::BlockReduceMax<float, false>(
413+ tvUbTensor,
414+ srcUb,
415+ numRowsRound,
416+ 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
417+ AscendC::PipeBarrier<PIPE_V>();
418+ AscendC::BlockReduceMax<float, false>(
419+ rowmaxUb,
420+ tvUbTensor,
421+ CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
422+ 0, 1, 1, 8);
423+ AscendC::PipeBarrier<PIPE_V>();
424+ for (uint64_t rowmax_idx = 1; rowmax_idx < (uint64_t)numElems / FLOAT_VECTOR_SIZE; ++rowmax_idx) {
425+ AscendC::BlockReduceMax<float, false>(
426+ tvUbTensor,
427+ srcUb[rowmax_idx * FLOAT_VECTOR_SIZE],
428+ numRowsRound,
429+ 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
430+ AscendC::PipeBarrier<PIPE_V>();
431+ AscendC::BlockReduceMax<float, false>(
432+ tvUbTensor[REDUCE_UB_SIZE],
433+ tvUbTensor,
434+ CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
435+ 0, 1, 1, 8);
436+ AscendC::PipeBarrier<PIPE_V>();
437+ SetVecMask(numRowsRound);
438+ AscendC::Max<float, false>(rowmaxUb,
439+ rowmaxUb,
440+ tvUbTensor[REDUCE_UB_SIZE],
441+ (uint64_t)0,
442+ 1,
443+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
444+ AscendC::PipeBarrier<PIPE_V>();
445+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
446+ }
447+ }
448+ if (numElems % FLOAT_VECTOR_SIZE > 0) {
449+ SetVecMask(numElems % FLOAT_VECTOR_SIZE);
450+ AscendC::BlockReduceMax<float, false>(
451+ tvUbTensor,
452+ srcUb[numElems / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
453+ numRowsRound,
454+ 0, 1, 1, numElemsAligned / FLOAT_BLOCK_SIZE);
455+ AscendC::PipeBarrier<PIPE_V>();
456+ SetBlockReduceMask(CeilDiv(numElems % FLOAT_VECTOR_SIZE, FLOAT_BLOCK_SIZE));
457+ if (numElems < FLOAT_VECTOR_SIZE) {
458+ AscendC::BlockReduceMax<float, false>(rowmaxUb,
459+ tvUbTensor,
460+ CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
461+ 0, 1, 1, 8);
462+ AscendC::PipeBarrier<PIPE_V>();
463+ } else {
464+ AscendC::BlockReduceMax<float, false>(tvUbTensor[REDUCE_UB_SIZE],
465+ tvUbTensor,
466+ CeilDiv(numRowsRound * FLOAT_BLOCK_SIZE, FLOAT_VECTOR_SIZE),
467+ 0, 1, 1, 8);
468+ AscendC::PipeBarrier<PIPE_V>();
469+ SetVecMask(numRowsRound);
470+ AscendC::Max<float, false>(rowmaxUb,
471+ rowmaxUb,
472+ tvUbTensor[REDUCE_UB_SIZE],
473+ (uint64_t)0,
474+ 1,
475+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
476+ AscendC::PipeBarrier<PIPE_V>();
477+ }
478+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
479+ }
480+ }
481+ 
482+ __aicore__ inline
483+ void CopySGmToUb(
484+ AscendC::GlobalTensor<ElementInput> gInput,
485+ uint32_t sUbOffset,
486+ uint32_t rowNumCurLoop,
487+ uint32_t columnNumRound,
488+ uint32_t columnNumPad)
489+ {
490+ AscendC::DataCopy(
491+ lsUbTensor[sUbOffset],
492+ gInput,
493+ AscendC::DataCopyParams(
494+ rowNumCurLoop, columnNumRound / FLOAT_BLOCK_SIZE,
495+ (columnNumPad - columnNumRound) / FLOAT_BLOCK_SIZE, 0));
496+ }
497+ 
498+ __aicore__ inline
499+ void CopyMaskGmToUb(
500+ AscendC::GlobalTensor<ElementMask> gMask,
501+ uint32_t columnNum, uint32_t columnNumRound,
502+ uint32_t maskStride, uint32_t tokenNumPerHead,
503+ uint32_t proTokenIdx, uint32_t proTokenNum,
504+ uint32_t integralHeadNum, uint32_t epiTokenNum)
505+ {
506+ uint32_t innerUbRowOffset = 0;
507+ if (proTokenNum != 0) {
508+ AscendC::DataCopyPad(
509+ maskUbTensor[innerUbRowOffset], gMask[proTokenIdx * maskStride],
510+ AscendC::DataCopyExtParams(
511+ proTokenNum, columnNum * sizeof(ElementMask),
512+ (maskStride - columnNum) * sizeof(ElementMask), 0, 0),
513+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
514+ innerUbRowOffset += proTokenNum * columnNumRound;
515+ }
516+ for (uint32_t headIdx = 0; headIdx < integralHeadNum; headIdx++) {
517+ AscendC::DataCopyPad(
518+ maskUbTensor[innerUbRowOffset], gMask,
519+ AscendC::DataCopyExtParams(
520+ tokenNumPerHead, columnNum * sizeof(ElementMask),
521+ (maskStride - columnNum) * sizeof(ElementMask), 0, 0),
522+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
523+ innerUbRowOffset += tokenNumPerHead * columnNumRound;
524+ }
525+ if (epiTokenNum != 0) {
526+ AscendC::DataCopyPad(
527+ maskUbTensor[innerUbRowOffset], gMask,
528+ AscendC::DataCopyExtParams(
529+ epiTokenNum, columnNum * sizeof(ElementMask),
530+ (maskStride - columnNum) * sizeof(ElementMask), 0, 0),
531+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
532+ }
533+ }
534+ 
535+ __aicore__ inline
536+ void ScaleS(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound)
537+ {
538+ AscendC::Muls<float, false>(
539+ lsUbTensor[sUbOffset],
540+ lsUbTensor[sUbOffset],
541+ scaleValue,
542+ (uint64_t)0,
543+ CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE),
544+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
545+ AscendC::PipeBarrier<PIPE_V>();
546+ }
547+ 
548+ template<typename ElementMaskDst, typename ElementMaskSrc>
549+ __aicore__ inline
550+ void UpCastMask(
551+ const AscendC::LocalTensor<ElementMaskDst> &maskUbTensorDst,
552+ const AscendC::LocalTensor<ElementMaskSrc> &maskUbTensorSrc,
553+ uint32_t rowNumCurLoop,
554+ uint32_t columnNumRound)
555+ {
556+ AscendC::Cast<ElementMaskDst, ElementMaskSrc, false>(
557+ maskUbTensorDst, maskUbTensorSrc, AscendC::RoundMode::CAST_NONE, (uint64_t)0,
558+ CeilDiv(rowNumCurLoop * columnNumRound, (uint32_t)(REPEAT_SIZE_IN_BYTE / sizeof(ElementMaskDst))),
559+ AscendC::UnaryRepeatParams(1, 1, 8, 4));
560+ AscendC::PipeBarrier<PIPE_V>();
561+ }
562+ 
563+ __aicore__ inline
564+ void ApplyMask(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound, uint32_t maskColumnRound,
565+ uint32_t addMaskUbOffset)
566+ {
567+ AscendC::Muls<float, false>(
568+ maskUbTensor32,
569+ maskUbTensor32,
570+ (float)-3e38,
571+ (uint64_t)0,
572+ CeilDiv(rowNumCurLoop * maskColumnRound, FLOAT_VECTOR_SIZE),
573+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
574+ AscendC::PipeBarrier<PIPE_V>();
575+ if (maskColumnRound == columnNumRound) {
576+ AscendC::Add<float, false>(
577+ lsUbTensor[sUbOffset],
578+ lsUbTensor[sUbOffset],
579+ maskUbTensor32,
580+ (uint64_t)0,
581+ CeilDiv(rowNumCurLoop * maskColumnRound, FLOAT_VECTOR_SIZE),
582+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
583+ } else {
584+ uint32_t loop = maskColumnRound / FLOAT_VECTOR_SIZE;
585+ for (uint32_t i = 0; i < loop; i++) {
586+ AscendC::Add<float, false>(lsUbTensor[sUbOffset][addMaskUbOffset + i * FLOAT_VECTOR_SIZE],
587+ lsUbTensor[sUbOffset][addMaskUbOffset + i * FLOAT_VECTOR_SIZE],
588+ maskUbTensor32[i * FLOAT_VECTOR_SIZE],
589+ (uint64_t)0,
590+ rowNumCurLoop,
591+ AscendC::BinaryRepeatParams(
592+ 1, 1, 1,
593+ columnNumRound / FLOAT_BLOCK_SIZE,
594+ columnNumRound / FLOAT_BLOCK_SIZE,
595+ maskColumnRound / FLOAT_BLOCK_SIZE));
596+ }
597+ if (maskColumnRound % FLOAT_VECTOR_SIZE > 0) {
598+ SetVecMask(maskColumnRound % FLOAT_VECTOR_SIZE);
599+ AscendC::Add<float, false>(lsUbTensor[sUbOffset][addMaskUbOffset + loop * FLOAT_VECTOR_SIZE],
600+ lsUbTensor[sUbOffset][addMaskUbOffset + loop * FLOAT_VECTOR_SIZE],
601+ maskUbTensor32[loop * FLOAT_VECTOR_SIZE],
602+ (uint64_t)0,
603+ rowNumCurLoop,
604+ AscendC::BinaryRepeatParams(
605+ 1, 1, 1,
606+ columnNumRound / FLOAT_BLOCK_SIZE,
607+ columnNumRound / FLOAT_BLOCK_SIZE,
608+ maskColumnRound / FLOAT_BLOCK_SIZE));
609+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
610+ }
611+ }
612+ AscendC::PipeBarrier<PIPE_V>();
613+ }
614+ 
615+ __aicore__ inline
616+ void CalcLocalRowMax(uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
617+ uint32_t rowOffset)
618+ {
619+ if (columnNum == 1024) {
620+ RowmaxSPECTILE1024(
621+ lsUbTensor[sUbOffset],
622+ lmUbTensor[rowOffset],
623+ tvUbTensor,
624+ rowNumCurLoopRound,
625+ columnNum,
626+ columnNumRound);
627+ } else if (columnNum == 512) {
628+ RowmaxSPECTILE512(
629+ lsUbTensor[sUbOffset],
630+ lmUbTensor[rowOffset],
631+ tvUbTensor,
632+ rowNumCurLoopRound,
633+ columnNum,
634+ columnNumRound);
635+ } else if (columnNum == 256) {
636+ RowmaxSPECTILE256(
637+ lsUbTensor[sUbOffset],
638+ lmUbTensor[rowOffset],
639+ tvUbTensor,
640+ rowNumCurLoopRound,
641+ columnNum,
642+ columnNumRound);
643+ } else {
644+ RowmaxTAILTILE(
645+ lsUbTensor[sUbOffset],
646+ lmUbTensor[rowOffset],
647+ tvUbTensor,
648+ rowNumCurLoopRound,
649+ columnNum,
650+ columnNumRound);
651+ }
652+ }
653+ 
654+ __aicore__ inline
655+ void UpdateGlobalRowMax(uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum,
656+ uint32_t columnNumRound, uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile)
657+ {
658+ if (isFirstStackTile) {
659+ AscendC::DataCopy(
660+ hmUbTensor[rowOffset],
661+ lmUbTensor[rowOffset],
662+ AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0));
663+ AscendC::PipeBarrier<PIPE_V>();
664+ } else {
665+ SetVecMask(rowNumCurLoop);
666+ // *** hm = vmax(lm, gm)
667+ AscendC::Max<float, false>(
668+ hmUbTensor[rowOffset],
669+ lmUbTensor[rowOffset],
670+ gmUbTensor[rowOffset],
671+ (uint64_t)0,
672+ 1,
673+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
674+ AscendC::PipeBarrier<PIPE_V>();
675+ // *** dm = gm - hm
676+ AscendC::Sub<float, false>(
677+ dmUbTensor[dmUbOffsetCurCycle],
678+ gmUbTensor[rowOffset],
679+ hmUbTensor[rowOffset],
680+ (uint64_t)0,
681+ 1,
682+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
683+ AscendC::PipeBarrier<PIPE_V>();
684+ // *** dm = exp(dm)
685+ AscendC::Exp<float, false>(
686+ dmUbTensor[dmUbOffsetCurCycle],
687+ dmUbTensor[dmUbOffsetCurCycle],
688+ (uint64_t)0,
689+ 1,
690+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
691+ }
692+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
693+ AscendC::PipeBarrier<PIPE_V>();
694+ // *** gm = hm
695+ AscendC::DataCopy(
696+ gmUbTensor[rowOffset],
697+ hmUbTensor[rowOffset],
698+ AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0));
699+ AscendC::PipeBarrier<PIPE_V>();
700+ }
701+ 
702+ __aicore__ inline
703+ void CalcExp(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum,
704+ uint32_t columnNumRound, uint32_t rowOffset)
705+ {
706+ // *** hm_block = expand_to_block(hm), 存放于 tv
707+ AscendC::Brcb(
708+ tvUbTensor.template ReinterpretCast<uint32_t>(),
709+ hmUbTensor[rowOffset].template ReinterpretCast<uint32_t>(),
710+ rowNumCurLoopRound / FLOAT_BLOCK_SIZE,
711+ AscendC::BrcbRepeatParams(1, 8));
712+ AscendC::PipeBarrier<PIPE_V>();
713+ // *** ls = ls - hm_block
714+ for (uint32_t subIdx = 0; subIdx < columnNum / FLOAT_VECTOR_SIZE; ++subIdx) {
715+ AscendC::Sub<float, false>(
716+ lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE],
717+ lsUbTensor[sUbOffset][subIdx * FLOAT_VECTOR_SIZE],
718+ tvUbTensor,
719+ (uint64_t)0,
720+ rowNumCurLoop,
721+ AscendC::BinaryRepeatParams(
722+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
723+ }
724+ if (columnNum % FLOAT_VECTOR_SIZE > 0) {
725+ SetVecMask(columnNum % FLOAT_VECTOR_SIZE);
726+ AscendC::Sub<float, false>(
727+ lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
728+ lsUbTensor[sUbOffset][columnNum / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
729+ tvUbTensor,
730+ (uint64_t)0,
731+ rowNumCurLoop,
732+ AscendC::BinaryRepeatParams(
733+ 1, 1, 0, columnNumRound / FLOAT_BLOCK_SIZE, columnNumRound / FLOAT_BLOCK_SIZE, 1));
734+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
735+ }
736+ AscendC::PipeBarrier<PIPE_V>();
737+ // *** ls = exp(ls)
738+ AscendC::Exp<float, false>(
739+ lsUbTensor[sUbOffset],
740+ lsUbTensor[sUbOffset],
741+ (uint64_t)0,
742+ CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE),
743+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
744+ AscendC::PipeBarrier<PIPE_V>();
745+ }
746+ 
747+ __aicore__ inline
748+ void CalcLocalRowSum(uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
749+ uint32_t rowOffset)
750+ {
751+ // *** ll = rowsum(ls32)
752+ if (columnNum == 1024) {
753+ RowsumSPECTILE1024(
754+ lsUbTensor[sUbOffset],
755+ llUbTensor[rowOffset],
756+ tvUbTensor,
757+ rowNumCurLoopRound,
758+ columnNum,
759+ columnNumRound);
760+ } else if (columnNum == 512) {
761+ RowsumSPECTILE512(
762+ lsUbTensor[sUbOffset],
763+ llUbTensor[rowOffset],
764+ tvUbTensor,
765+ rowNumCurLoopRound,
766+ columnNum,
767+ columnNumRound);
768+ } else if (columnNum == 256) {
769+ RowsumSPECTILE256(
770+ lsUbTensor[sUbOffset],
771+ llUbTensor[rowOffset],
772+ tvUbTensor,
773+ rowNumCurLoopRound,
774+ columnNum,
775+ columnNumRound);
776+ } else {
777+ RowsumTAILTILE(
778+ lsUbTensor[sUbOffset],
779+ llUbTensor[rowOffset],
780+ tvUbTensor,
781+ rowNumCurLoopRound,
782+ columnNum,
783+ columnNumRound);
784+ }
785+ }
786+ 
787+ __aicore__ inline
788+ void UpdateGlobalRowSum(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound,
789+ uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile)
790+ {
791+ if (isFirstStackTile) {
792+ // *** gl = ll
793+ AscendC::DataCopy(
794+ glUbTensor[rowOffset],
795+ llUbTensor[rowOffset],
796+ AscendC::DataCopyParams(1, rowNumCurLoopRound / FLOAT_BLOCK_SIZE, 0, 0));
797+ AscendC::PipeBarrier<PIPE_V>();
798+ } else {
799+ SetVecMask(rowNumCurLoop);
800+ // *** gl = dm * gl
801+ AscendC::Mul<float, false>(
802+ glUbTensor[rowOffset],
803+ dmUbTensor[dmUbOffsetCurCycle],
804+ glUbTensor[rowOffset],
805+ (uint64_t)0,
806+ 1,
807+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
808+ AscendC::PipeBarrier<PIPE_V>();
809+ // *** gl = ll + gl
810+ AscendC::Add<float, false>(
811+ glUbTensor[rowOffset],
812+ glUbTensor[rowOffset],
813+ llUbTensor[rowOffset],
814+ (uint64_t)0,
815+ 1,
816+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
817+ AscendC::PipeBarrier<PIPE_V>();
818+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
819+ }
820+ }
821+ 
822+ __aicore__ inline
823+ void DownCastP(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound)
824+ {
825+ // *** lp = castfp32to16(ls)
826+ if (std::is_same<ElementOutput, bfloat16_t>::value) {
827+ AscendC::Cast<ElementOutput, float, false>(
828+ lpUbTensor[sUbOffset],
829+ lsUbTensor[sUbOffset],
830+ AscendC::RoundMode::CAST_RINT,
831+ (uint64_t)0,
832+ CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE),
833+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
834+ } else {
835+ AscendC::Cast<ElementOutput, float, false>(
836+ lpUbTensor[sUbOffset],
837+ lsUbTensor[sUbOffset],
838+ AscendC::RoundMode::CAST_NONE,
839+ (uint64_t)0,
840+ CeilDiv(rowNumCurLoop * columnNumRound, FLOAT_VECTOR_SIZE),
841+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
842+ }
843+ }
844+ 
845+ __aicore__ inline
846+ void CopyPUbToGm(AscendC::GlobalTensor<ElementOutput> gOutput, uint32_t sUbOffset, uint32_t rowNumCurLoop,
847+ uint32_t columnNumRound, uint32_t columnNumPad)
848+ {
849+ AscendC::DataCopy(
850+ gOutput,
851+ lpUbTensor[sUbOffset],
852+ AscendC::DataCopyParams(
853+ rowNumCurLoop, columnNumRound / BLOCK_SIZE, 0, (columnNumPad - columnNumRound) / BLOCK_SIZE));
854+ }
855+ 
856+ template <bool doTriUMask>
857+ __aicore__ inline
858+ void SubCoreCompute(
859+ AscendC::GlobalTensor<ElementOutput> gOutput, const LayoutOutput &layoutOutput,
860+ uint32_t rowOffset, uint32_t isFirstStackTile, uint32_t isLastNoMaskStackTile,
861+ uint32_t isFirstRowLoop, uint32_t isLastRowLoop,
862+ uint32_t columnNumRound, uint32_t pingpongFlag,
863+ uint32_t curStackTileMod, Arch::CrossCoreFlag softmaxFlag)
864+ {
865+ uint32_t rowNumCurLoop = layoutOutput.shape(0);
866+ uint32_t rowNumCurLoopRound = RoundUp(rowNumCurLoop, FLOAT_BLOCK_SIZE);
867+ uint32_t columnNum = layoutOutput.shape(1);
868+ uint32_t columnNumPad = layoutOutput.stride(0);
869+ uint32_t sUbOffset = pingpongFlag * MAX_UB_S_ELEM_NUM;
870+ uint32_t dmUbOffsetCurCycle = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffset;
871+ 
872+ if constexpr (LSE_MODE_ == LseMode::OUT_ONLY) {
873+ // wait for lse from ub to gm
874+ // In lse out-only mode, tv is used in the last stack tile to transport lse
875+ if (isFirstStackTile && isFirstRowLoop) {
876+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
877+ }
878+ }
879+ CalcLocalRowMax(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
880+ UpdateGlobalRowMax(
881+ rowNumCurLoop, rowNumCurLoopRound,
882+ columnNum, columnNumRound,
883+ dmUbOffsetCurCycle,
884+ rowOffset,
885+ isFirstStackTile);
886+ 
887+ CalcExp(sUbOffset, rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
888+ if constexpr (!doTriUMask) {
889+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(pingpongFlag);
890+ }
891+ 
892+ DownCastP(sUbOffset, rowNumCurLoop, columnNumRound);
893+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
894+ 
895+ CalcLocalRowSum(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
896+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(pingpongFlag);
897+ 
898+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(pingpongFlag);
899+ CopyPUbToGm(gOutput, sUbOffset, rowNumCurLoop, columnNumRound, columnNumPad);
900+ if constexpr (!doTriUMask) {
901+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(pingpongFlag);
902+ if (isLastNoMaskStackTile && isLastRowLoop) {
903+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
904+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
905+ }
906+ } else {
907+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
908+ }
909+ if (isLastRowLoop) {
910+ NpuArch::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxFlag);
911+ }
912+ UpdateGlobalRowSum(
913+ sUbOffset, rowNumCurLoop, rowNumCurLoopRound, dmUbOffsetCurCycle, rowOffset, isFirstStackTile);
914+ }
915+ 
916+ __aicore__ inline
917+ void operator()(AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
918+ const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, GemmCoord actualBlockShape,
919+ uint32_t isFirstStackTile, uint32_t isLastNoMaskStackTile,
920+ uint32_t qSBlockSize, uint32_t qNBlockSize, uint32_t curStackTileMod, Arch::CrossCoreFlag softmaxFlag)
921+ {
922+ uint32_t rowNum = actualBlockShape.m();
923+ uint32_t columnNum = actualBlockShape.n();
924+ uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE);
925+ uint32_t columnNumPad = layoutInput.stride(0);
926+ 
927+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
928+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
929+ 
930+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
931+ uint32_t qNThisSubBlock = (qNBlockSize == 1) ?
932+ 0 : (subBlockIdx == 1) ? (qNBlockSize - qNSplitSubBlock) : qNSplitSubBlock;
933+ uint32_t rowSplitSubBlock = (qNBlockSize == 1) ?
934+ (qSBlockSize / 2) : (qSBlockSize * qNSplitSubBlock);
935+ uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock;
936+ uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock;
937+ uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound;
938+ uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE);
939+ rowNumTile = AscendC::Std::min(rowNumTile, FLOAT_VECTOR_SIZE);
940+ uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile);
941+ uint32_t preLoad = 1;
942+ if (rowActualThisSubBlock == 0) {
943+ NpuArch::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxFlag);
944+ return;
945+ }
946+ 
947+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) {
948+ if (rowLoopIdx < rowLoopNum) {
949+ uint32_t pingpongFlag = rowLoopIdx % 2;
950+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
951+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
952+ uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ?
953+ (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
954+ 
955+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
956+ auto gInputCurLoop = gInput[offsetInput];
957+ 
958+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(pingpongFlag);
959+ CopySGmToUb(
960+ gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad);
961+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
962+ }
963+ if (rowLoopIdx >= preLoad) {
964+ uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad;
965+ uint32_t pingpongFlag = delayedRowLoopIdx % 2;
966+ uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile;
967+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
968+ uint32_t rowNumCurLoop =
969+ (delayedRowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
970+ 
971+ int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
972+ auto gOutputCurLoop = gOutput[offsetOutput];
973+ auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum));
974+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
975+ ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound);
976+ SubCoreCompute<false>(
977+ gOutputCurLoop,
978+ layoutOutputCurLoop,
979+ rowOffsetCurLoop,
980+ isFirstStackTile,
981+ isLastNoMaskStackTile,
982+ (delayedRowLoopIdx == 0),
983+ (delayedRowLoopIdx == rowLoopNum - 1),
984+ columnNumRound,
985+ pingpongFlag,
986+ curStackTileMod,
987+ softmaxFlag);
988+ }
989+ }
990+ }
991+ 
992+ __aicore__ inline
993+ void operator()(AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<ElementInput> gInput,
994+ AscendC::GlobalTensor<ElementMask> gMask, const LayoutOutput &layoutOutput, const LayoutInput &layoutInput,
995+ const LayoutInput &layoutMask, GemmCoord actualBlockShape, uint32_t isFirstStackTile, uint32_t qSBlockSize,
996+ uint32_t qNBlockSize, uint32_t curStackTileMod, Arch::CrossCoreFlag qkReady, uint32_t triUp, uint32_t triDown,
997+ uint32_t kvSStartIdx, uint32_t kvSEndIdx)
998+ {
999+ uint32_t rowNum = actualBlockShape.m();
1000+ uint32_t columnNum = actualBlockShape.n();
1001+ uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE_IN_BYTE);
1002+ uint32_t columnNumPad = layoutInput.stride(0);
1003+ uint32_t maskStride = layoutMask.stride(0);
1004+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
1005+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
1006+ 
1007+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
1008+ uint32_t qNThisSubBlock = (qNBlockSize == 1) ?
1009+ 0 : (subBlockIdx == 1) ? (qNBlockSize - qNSplitSubBlock) : qNSplitSubBlock;
1010+ uint32_t rowSplitSubBlock = (qNBlockSize == 1) ?
1011+ (qSBlockSize / 2) : (qSBlockSize * qNSplitSubBlock);
1012+ uint32_t rowActualThisSubBlock = (subBlockIdx == 1) ?
1013+ (rowNum - rowSplitSubBlock) : rowSplitSubBlock;
1014+ uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock;
1015+ 
1016+ uint32_t tokenNumPerHeadThisSubBlock = Min(qSBlockSize, rowActualThisSubBlock);
1017+ uint32_t maskOffsetThisSubBlock = (qNBlockSize == 1) ?
1018+ rowOffsetThisSubBlock : 0;
1019+ 
1020+ // calc mask shift in gm
1021+ uint32_t gmOffsetMaskRow;
1022+ uint32_t gmOffsetMaskColumn;
1023+ uint32_t maskColumn;
1024+ uint32_t addMaskUbOffset;
1025+ if (triUp >= kvSStartIdx) {
1026+ uint32_t triUpRoundDown = RoundDown(triUp, BLOCK_SIZE_IN_BYTE);
1027+ gmOffsetMaskRow = triUp - triUpRoundDown;
1028+ gmOffsetMaskColumn = 0;
1029+ maskColumn = kvSEndIdx - triUpRoundDown;
1030+ addMaskUbOffset = triUpRoundDown - kvSStartIdx;
1031+ } else {
1032+ gmOffsetMaskRow = 0;
1033+ gmOffsetMaskColumn = kvSStartIdx - triUp;
1034+ maskColumn = columnNum;
1035+ addMaskUbOffset = 0;
1036+ }
1037+ uint32_t maskColumnRound = RoundUp(maskColumn, BLOCK_SIZE_IN_BYTE);
1038+ 
1039+ int64_t offsetMask =
1040+ layoutMask.GetOffset(MatrixCoord(gmOffsetMaskRow + maskOffsetThisSubBlock, gmOffsetMaskColumn));
1041+ auto gMaskThisSubBlock = gMask[offsetMask];
1042+ auto layoutMaskThisSubBlock = layoutMask;
1043+ 
1044+ uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound;
1045+ uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE);
1046+ rowNumTile = AscendC::Std::min(rowNumTile, FLOAT_VECTOR_SIZE);
1047+ uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile);
1048+ uint32_t preLoad = 1;
1049+ 
1050+ if (rowActualThisSubBlock == 0) {
1051+ Arch::CrossCoreWaitFlag(qkReady);
1052+ return;
1053+ }
1054+ 
1055+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum + preLoad; rowLoopIdx++) {
1056+ if (rowLoopIdx < rowLoopNum) {
1057+ uint32_t pingpongFlag = rowLoopIdx % 2;
1058+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
1059+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
1060+ uint32_t rowNumCurLoop = (rowLoopIdx == rowLoopNum - 1) ?
1061+ (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
1062+ // loop 0 mask load before cross core sync
1063+ if (rowLoopIdx == 0) {
1064+ // the token idx of the start token of the prologue part
1065+ uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock;
1066+ // the token num of the prologue part
1067+ uint32_t proTokenNum =
1068+ Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock;
1069+ // the token num of the epilogue part
1070+ uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock;
1071+ // the number of integral heads within a cycle
1072+ uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock;
1073+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
1074+ CopyMaskGmToUb(
1075+ gMaskThisSubBlock,
1076+ maskColumn, maskColumnRound, maskStride,
1077+ tokenNumPerHeadThisSubBlock,
1078+ proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum);
1079+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
1080+ Arch::CrossCoreWaitFlag(qkReady);
1081+ }
1082+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
1083+ auto gInputCurLoop = gInput[offsetInput];
1084+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(pingpongFlag);
1085+ CopySGmToUb(
1086+ gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad);
1087+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
1088+ }
1089+ if (rowLoopIdx >= preLoad) {
1090+ uint32_t delayedRowLoopIdx = rowLoopIdx - preLoad;
1091+ uint32_t pingpongFlag = delayedRowLoopIdx % 2;
1092+ uint32_t rowOffsetCurLoop = delayedRowLoopIdx * rowNumTile;
1093+ uint32_t rowNumCurLoop = (delayedRowLoopIdx == rowLoopNum - 1) ?
1094+ (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
1095+ 
1096+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
1097+ UpCastMask<half, ElementMask>(maskUbTensor16, maskUbTensor, rowNumCurLoop, columnNumRound);
1098+ UpCastMask<float, half>(maskUbTensor32, maskUbTensor16, rowNumCurLoop, columnNumRound);
1099+
1100+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(pingpongFlag);
1101+ ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound);
1102+ ApplyMask(
1103+ (pingpongFlag * MAX_UB_S_ELEM_NUM),
1104+ rowNumCurLoop, columnNumRound,
1105+ maskColumnRound, addMaskUbOffset);
1106+ // next loop mask load
1107+ if (rowLoopIdx < rowLoopNum) {
1108+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
1109+ uint32_t rowNumCurLoop =
1110+ (rowLoopIdx == rowLoopNum - 1) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
1111+ // the token idx of the start token of the prologue part
1112+ uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock;
1113+ // the token num of the prologue part
1114+ uint32_t proTokenNum =
1115+ Min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) % tokenNumPerHeadThisSubBlock;
1116+ // the number of integral heads within a cycle
1117+ uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock;
1118+ // the token num of the epilogue part
1119+ uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock;
1120+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID0);
1121+ CopyMaskGmToUb(
1122+ gMaskThisSubBlock,
1123+ maskColumn, maskColumnRound, maskStride,
1124+ tokenNumPerHeadThisSubBlock,
1125+ proTokenIdx, proTokenNum, integralHeadNum, epiTokenNum);
1126+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID2);
1127+ }
1128+ // online softmax vectorized compute
1129+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
1130+ int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
1131+ auto gOutputCurLoop = gOutput[offsetOutput];
1132+ auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum));
1133+ SubCoreCompute<true>(
1134+ gOutputCurLoop,
1135+ layoutOutputCurLoop,
1136+ rowOffsetCurLoop,
1137+ isFirstStackTile,
1138+ 0,
1139+ (delayedRowLoopIdx == 0),
1140+ (delayedRowLoopIdx == rowLoopNum - 1),
1141+ columnNumRound,
1142+ pingpongFlag,
1143+ curStackTileMod);
1144+ }
1145+ }
1146+ }
1147+ 
1148+private:
1149+ float scaleValue;
1150+ AscendC::LocalTensor<float> lsUbTensor;
1151+ AscendC::LocalTensor<ElementOutput> lpUbTensor;
1152+ AscendC::LocalTensor<ElementMask> maskUbTensor;
1153+ AscendC::LocalTensor<half> maskUbTensor16;
1154+ AscendC::LocalTensor<float> maskUbTensor32;
1155+ AscendC::LocalTensor<float> lmUbTensor;
1156+ AscendC::LocalTensor<float> hmUbTensor;
1157+ AscendC::LocalTensor<float> gmUbTensor;
1158+ AscendC::LocalTensor<float> dmUbTensor;
1159+ AscendC::LocalTensor<float> llUbTensor;
1160+ AscendC::LocalTensor<float> tvUbTensor;
1161+ AscendC::LocalTensor<float> glUbTensor;
1162+};
1163+}
1164+ 
1165+#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue_online_softmax_arch35_reg_low_prec.hpp+528-0
@@ -0,0 +1,528 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_ARCH35_REG_LOW_PREC_HPP
12+#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_ARCH35_REG_LOW_PREC_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/resource.hpp"
16+#include "../../../attn_infra/epilogue/dispatch_policy.hpp"
17+#include "../../../attn_infra/epilogue/tile_common/tile_copy.hpp"
18+#include "../../../attn_infra/gemm_coord.hpp"
19+#include "../../../attn_infra/matrix_coord.hpp"
20+#include "../../../tla/tensor.hpp"
21+#include "../../../tla/layout.hpp"
22+using AscendC::printf;
23+#define VF_MIN(a, b) ((b) + (uint16_t((a) - (b)) & -(uint16_t((a) - (b)) >> 15)))
24+ 
25+namespace NpuArch::Epilogue::Block {
26+ 
27+template <typename DefaultType>
28+struct ElementPTmpTypeLookup {
29+ using type = DefaultType;
30+};
31+ 
32+template <>
33+struct ElementPTmpTypeLookup<float8_e4m3_t> {
34+ using type = half;
35+};
36+ 
37+enum class KvBaseTileRegSplitStages {
38+ ONE,
39+ TWO
40+};
41+ 
42+template <
43+ class OutputType_,
44+ class LayoutS_>
45+class BlockEpilogue<
46+ EpilogueOnlineSoftmaxBsa,
47+ OutputType_,
48+ Gemm::GemmType<half, LayoutS_>>
49+{
50+public:
51+ using DispatchPolicy = EpilogueOnlineSoftmaxBsa;
52+ using ArchTag = typename DispatchPolicy::ArchTag;
53+ using ElementOutput = typename OutputType_::Element;
54+ using ElementOutputTmp = typename ElementPTmpTypeLookup<ElementOutput>::type;
55+ 
56+ 
57+ using ElementInput = half;
58+ using LayoutOutput = typename OutputType_::Layout;
59+ using LayoutInput = LayoutS_;
60+ 
61+ static constexpr uint32_t BLOCK_SIZE_IN_BYTE = 32;
62+ static constexpr uint32_t REPEAT_SIZE_IN_BYTE = 256;
63+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
64+ static constexpr uint32_t FP8_BLOCK_SIZE = 32;
65+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
66+ static constexpr uint32_t HALF_VECTOR_SIZE = 128;
67+ static constexpr uint32_t BLOCK_SIZE = 16;
68+ static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024;
69+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 32768;
70+ static constexpr uint32_t VECTOR_SIZE = 128;
71+ static constexpr uint32_t MAX_UB_S_ELEM_NUM = 32768;
72+ static constexpr uint32_t DM_UB_GLOBAL_ELEM_NUM = 64 * 2; //! 2* for 11 22 ... 32,32
73+ static constexpr uint32_t ELE_NUM_PER_C0 = 16;
74+ 
75+ static constexpr uint32_t REDUCE_UB_SIZE = 1024;
76+ static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32;
77+ static constexpr uint32_t ROW_OPS_SPEC_MASK_8 = 8;
78+ static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4;
79+ static constexpr uint32_t ROW_OPS_SPEC_MASK_2 = 2;
80+ static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256;
81+ static constexpr int64_t UB_FLOAT_LINE_SIZE = 64;
82+ 
83+ static constexpr uint32_t SPLIT_COL_IDX_2 = 2;
84+ static constexpr uint32_t SPLIT_COL_IDX_3 = 3;
85+ static constexpr ElementInput MIN_VALUE = -65504.0f;
86+ static constexpr uint32_t FP8_REP_SIZE = 128 *2;
87+ static constexpr uint32_t HALF_REP_SIZE = 128;
88+ static constexpr uint32_t FLOAT_REP_SIZE = 64;
89+ static constexpr uint32_t BLOCK_REP_SIZE = 8;
90+ static constexpr uint32_t REPEAT_STRIDE = 1;
91+ static constexpr uint32_t C0_NUM_PER_FRACTAL = 16;
92+ static constexpr uint32_t SM_ROW_MAX_ELEM_NUM = 64;
93+ static constexpr uint32_t SM_COL_MAX_ELEM_NUM = 512;
94+ static constexpr uint32_t SM_VREG_SIZE = 256 / sizeof(ElementInput);
95+ static constexpr uint32_t QUANT_MODE1_SCALE_BYTES = 1024;
96+ 
97+ __aicore__ inline
98+ BlockEpilogue(Arch::Resource<ArchTag> &resource, float scaleValue_, uint32_t blockSizeY_ = 128)
99+ {
100+ // Allocate UB space
101+ constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;
102+ constexpr uint32_t LP_UB_TENSOR_OFFSET = LS_UB_TENSOR_OFFSET;
103+ 
104+ constexpr uint32_t LM_UB_TENSOR_OFFSET = 7 * UB_UINT8_BLOCK_SIZE;
105+ constexpr uint32_t GM_UB_TENSOR_OFFSET = LM_UB_TENSOR_OFFSET + 128 * sizeof(float);
106+ constexpr uint32_t DM_UB_TENSOR_OFFSET = GM_UB_TENSOR_OFFSET + 128 * sizeof(float);
107+ constexpr uint32_t LL_UB_TENSOR_OFFSET = DM_UB_TENSOR_OFFSET + 3 * 128 * sizeof(float);
108+ constexpr uint32_t GL_UB_TENSOR_OFFSET = LL_UB_TENSOR_OFFSET + 128 * sizeof(float);
109+ constexpr uint32_t SCALE_UB_TENSOR_OFFSET = GL_UB_TENSOR_OFFSET + 128 * sizeof(float);
110+ 
111+ subBlockIdx_ = AscendC::GetSubBlockIdx();
112+ 
113+ scaleValue = scaleValue_;
114+ blockSizeK = blockSizeY_;
115+ lsUbTensor = resource.ubBuf.template GetBufferByByte<ElementInput>(LS_UB_TENSOR_OFFSET);
116+ lpUbTensor = resource.ubBuf.template GetBufferByByte<uint16_t>(LP_UB_TENSOR_OFFSET);
117+ gmUbTensor = resource.ubBuf.template GetBufferByByte<ElementInput>(GM_UB_TENSOR_OFFSET);
118+ glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
119+ dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
120+ lmUbTensor = resource.ubBuf.template GetBufferByByte<ElementInput>(LM_UB_TENSOR_OFFSET);
121+ llUbTensor = resource.ubBuf.template GetBufferByByte<ElementInput>(LL_UB_TENSOR_OFFSET);
122+ scaleTensor = resource.ubBuf.template GetBufferByByte<ElementInput>(SCALE_UB_TENSOR_OFFSET);
123+ tmpTensor = resource.ubBuf.template GetBufferByByte<uint8_t>(LM_UB_TENSOR_OFFSET + 4096U);
124+ AscendC::SetFlag<AscendC::HardEvent::V_S>(EVENT_ID0);
125+ }
126+ 
127+ __aicore__ inline
128+ ~BlockEpilogue()
129+ {
130+ AscendC::WaitFlag<AscendC::HardEvent::V_S>(EVENT_ID0);
131+ }
132+ 
133+ template <class TensorDst, class TensorSrc>
134+ __aicore__ inline
135+ void CopyPUbToPL1(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t m, uint32_t nRound, uint32_t actRowsCube)
136+ {
137+ constexpr uint32_t BlockElements = BLOCK_SIZE_IN_BYTE / sizeof(ElementOutput);
138+ if constexpr (sizeof(ElementOutput) == 2) {
139+ const uint32_t blockCount = tla::get<1, 1>(srcTensor.shape());
140+ const uint32_t blockLen = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape());
141+ const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
142+ constexpr int32_t C0_SIZE = BLOCK_SIZE_IN_BYTE / sizeof(typename TensorDst::Element);
143+ AscendC::DataCopyParams repeatParams;
144+ 
145+ repeatParams.blockCount = blockCount;
146+ repeatParams.blockLen = m;
147+ repeatParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / C0_SIZE - m;
148+ repeatParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / C0_SIZE - m;
149+ 
150+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
151+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
152+ 
153+ AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams);
154+ } else if constexpr (sizeof(ElementOutput) == 1) { // dn 2 Zn
155+ auto mRound = RoundUp(m, BlockElements);
156+ auto dstOffset = subBlockIdx_ == 0 ? 0 : (actRowsCube - m) * nRound;
157+ AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data(), mRound * nRound);
158+ }
159+ }
160+ 
161+ template <uint32_t MODE, pipe_t PIPE>
162+ __aicore__ inline
163+ void SetCrossCoreSync(Arch::CrossCoreFlag &crossCoreFlag)
164+ {
165+ if constexpr (MODE == 4U) {
166+ Arch::CrossCoreSetFlag<MODE, PIPE>(crossCoreFlag);
167+ }
168+ }
169+ 
170+ template <uint32_t MODE, pipe_t PIPE>
171+ __aicore__ inline
172+ void WaitCrossCoreSync(Arch::CrossCoreFlag &crossCoreFlag)
173+ {
174+ if constexpr (MODE == 4U) {
175+ Arch::CrossCoreWaitFlag<MODE, PIPE>(crossCoreFlag);
176+ }
177+ }
178+ 
179+ template <class TensorP>
180+ __aicore__ inline
181+ void operator()(TensorP &l1PTensorTla, GemmCoord actualBlockShape,
182+ uint32_t isFirstKvSTile, uint32_t ubSBufId, uint32_t l1PBufId,
183+ Arch::CrossCoreFlag mm1ToSmFlag, Arch::CrossCoreFlag smToMm2Flag){
184+
185+ struct EmptyTensor {};
186+ EmptyTensor dummyQS;
187+ EmptyTensor dummyKS;
188+ operator()<0>(l1PTensorTla, dummyQS, dummyKS,
189+ actualBlockShape,
190+ isFirstKvSTile, ubSBufId, l1PBufId,
191+ mm1ToSmFlag, smToMm2Flag, AscendC::GlobalTensor<int32_t>());
192+ }
193+
194+ template <int quant_mode, class TensorP, class TensorQS, class TensorKS, typename ElementIndex>
195+ __aicore__ inline
196+ void operator()(TensorP &l1PTensorTla, TensorQS &gmQSTensorTla, TensorKS &gmKSTensorTla,
197+ GemmCoord actualBlockShape,
198+ uint32_t isFirstKvSTile, uint32_t ubSBufId, uint32_t l1PBufId,
199+ Arch::CrossCoreFlag mm1ToSmFlag, Arch::CrossCoreFlag smToMm2Flag,
200+ const AscendC::GlobalTensor<ElementIndex>& sparseIndex)
201+ {
202+ static_assert(quant_mode == 0 || quant_mode == 1);
203+ constexpr int16_t vlSize = static_cast<int16_t>(AscendC::GetVecLen() / sizeof(ElementInput));
204+ uint32_t m;
205+ uint16_t mRound;
206+ uint32_t mCopyOffset = 0;
207+ const uint32_t mTot = actualBlockShape.m();
208+ if constexpr (quant_mode == 0) {
209+ mCopyOffset = RoundUp(actualBlockShape.m(), 8) / 2;
210+ m = actualBlockShape.m() < mCopyOffset ? actualBlockShape.m() : mCopyOffset;
211+ m = subBlockIdx_ == 0 ? m : actualBlockShape.m() - m;
212+ mRound = RoundUp(m, C0_NUM_PER_FRACTAL);
213+ } else {
214+ if (mTot <= FP8_BLOCK_SIZE) {
215+ m = (subBlockIdx_ == 0) ? mTot : 0;
216+ } else {
217+ uint32_t mhalf = (mTot + 1) / 2;
218+ uint32_t mAlign = (mhalf > FP8_BLOCK_SIZE) ? RoundUp(mhalf, FP8_BLOCK_SIZE) : FP8_BLOCK_SIZE;
219+ m = (subBlockIdx_ == 0) ? mAlign : (mTot - mAlign);
220+ mCopyOffset = (subBlockIdx_ == 0) ? 0 : mAlign;
221+ }
222+ mRound = RoundUp(m, FP8_BLOCK_SIZE);
223+ }
224+ 
225+ if (m == 0) {
226+ WaitCrossCoreSync<4, PIPE_V>(mm1ToSmFlag);
227+ SetCrossCoreSync<4, PIPE_V>(mm1ToSmFlag);
228+ WaitCrossCoreSync<4, PIPE_MTE3>(smToMm2Flag);
229+ SetCrossCoreSync<4, PIPE_MTE3>(smToMm2Flag);
230+ return;
231+ }
232+
233+ uint32_t n = actualBlockShape.n();
234+ uint16_t nRound = RoundUp(n, 16);
235+ auto paddingSize = (nRound - n) * FP8_BLOCK_SIZE;
236+ 
237+ // wait QK Fixpipe finsh
238+ WaitCrossCoreSync<4, PIPE_V>(mm1ToSmFlag);
239+ constexpr uint32_t QUANT_BLOCK_SIZE = 64;
240+ if constexpr (quant_mode == 1){
241+ auto qlayout = gmQSTensorTla.layout();
242+ auto klayout = gmKSTensorTla.layout();
243+ uint32_t blockNum = AscendC::CeilDivision(n, blockSizeK);
244+ uint32_t subNLoops = AscendC::CeilDivision(blockSizeK, QUANT_BLOCK_SIZE);
245+
246+ auto qOffset = qlayout(tla::MakeCoord(mTot <= QUANT_BLOCK_SIZE ? 0 : subBlockIdx_, static_cast<uint32_t>(0)));
247+ float qscale = gmQSTensorTla.data().GetValue(qOffset);
248+ AscendC::WaitFlag<AscendC::HardEvent::V_S>(EVENT_ID0);
249+ for(uint32_t i = 0; i < blockNum; i++){
250+ uint32_t index = sparseIndex.GetValue(i);
251+ auto jloops = n - i * blockSizeK <= QUANT_BLOCK_SIZE ? 1 : subNLoops;
252+ for(uint32_t j = 0; j < jloops; j++) {
253+ float ks = gmKSTensorTla.data().GetValue(index * subNLoops + j);
254+ auto s = ks * qscale; // scaleValue on fixpipe
255+ scaleTensor.SetValue(i*subNLoops + j, static_cast<ElementInput>(s));
256+ }
257+ }
258+ AscendC::SetFlag<AscendC::HardEvent::S_V>(EVENT_ID0);
259+ AscendC::WaitFlag<AscendC::HardEvent::S_V>(EVENT_ID0);
260+ }
261+
262+ ElementInput minValue = -60000.0f;
263+ if (isFirstKvSTile) {
264+ AscendC::Duplicate(gmUbTensor, minValue, 256); // max
265+ AscendC::Duplicate(glUbTensor, 0.0f, 128); // sum
266+ }
267+ 
268+ if (unlikely(n < nRound)) {
269+ AscendC::Duplicate(lsUbTensor[ubSBufId * MAX_UB_S_ELEM_NUM + n * FP8_BLOCK_SIZE], minValue, paddingSize);
270+ if (mRound == 64) {AscendC::Duplicate(lsUbTensor[ubSBufId * MAX_UB_S_ELEM_NUM + (n + nRound) * FP8_BLOCK_SIZE], minValue, paddingSize);}
271+ }
272+ 
273+ AscendC::PipeBarrier<PIPE_V>();
274+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(ubSBufId + 2);// alloc P
275+ auto lpUbOutTensor = lpUbTensor[ubSBufId * MAX_UB_S_ELEM_NUM].template ReinterpretCast<float8_e4m3_t>();
276+ DnSoftmaxQuantBlock<float8_e4m3_t, half, FP8_BLOCK_SIZE, true>(lpUbOutTensor, lsUbTensor[ubSBufId * MAX_UB_S_ELEM_NUM],
277+ gmUbTensor, gmUbTensor,
278+ dmUbTensor[l1PBufId * DM_UB_GLOBAL_ELEM_NUM],
279+ glUbTensor, glUbTensor, scaleTensor,
280+ tmpTensor, nRound);
281+ 
282+ if (mRound == 64) {
283+ DnSoftmaxQuantBlock<float8_e4m3_t, half, FP8_BLOCK_SIZE, true>(lpUbOutTensor[FP8_BLOCK_SIZE * nRound], lsUbTensor[ubSBufId * MAX_UB_S_ELEM_NUM + FP8_BLOCK_SIZE * nRound],
284+ gmUbTensor[128], gmUbTensor[128],
285+ dmUbTensor[l1PBufId * DM_UB_GLOBAL_ELEM_NUM + 64],
286+ glUbTensor[64], glUbTensor[64], scaleTensor,
287+ tmpTensor, nRound);
288+ }
289+ if constexpr (quant_mode == 1) {
290+ AscendC::SetFlag<AscendC::HardEvent::V_S>(EVENT_ID0);
291+ }
292+
293+ if (unlikely(n < nRound)) {
294+ AscendC::PipeBarrier<PIPE_V>();
295+ auto lpPad = lpUbOutTensor[n * FP8_BLOCK_SIZE].template ReinterpretCast<int8_t>();
296+ AscendC::Duplicate(lpPad, (int8_t)0, paddingSize);
297+ if (mRound == 64) {AscendC::Duplicate(lpPad[nRound * FP8_BLOCK_SIZE], (int8_t)0, paddingSize);}
298+ }
299+ 
300+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(ubSBufId); // P enque
301+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(ubSBufId); // P deque
302+ 
303+ auto ubPLayoutTla = tla::MakeLayout<ElementOutput, LayoutOutput>(mRound, nRound);
304+ 
305+ auto ubPTensorTla = tla::MakeTensor(lpUbOutTensor, ubPLayoutTla, Arch::PositionUB{});
306+ auto ubPTensorTlaTile = GetTile(ubPTensorTla, tla::MakeCoord(0, 0), tla::MakeShape(m, n));
307+ auto l1PTensorTlaTile = GetTile(l1PTensorTla, tla::MakeCoord(subBlockIdx_ * mCopyOffset, 0), tla::MakeShape(m, n));
308+ WaitCrossCoreSync<4, PIPE_MTE3>(smToMm2Flag);
309+ 
310+ CopyPUbToPL1(l1PTensorTlaTile, ubPTensorTlaTile, m, nRound, actualBlockShape.m());
311+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(ubSBufId + 2); //Free P
312+ SetCrossCoreSync<4, PIPE_MTE3>(mm1ToSmFlag);
313+ SetCrossCoreSync<4, PIPE_MTE3>(smToMm2Flag);
314+ }
315+ 
316+private:
317+ float scaleValue;
318+ AscendC::LocalTensor<ElementInput> lsUbTensor;
319+ AscendC::LocalTensor<uint16_t> lpUbTensor;
320+ AscendC::LocalTensor<ElementInput> gmUbTensor;
321+ AscendC::LocalTensor<float> glUbTensor;
322+ AscendC::LocalTensor<float> dmUbTensor;
323+ AscendC::LocalTensor<ElementInput> lmUbTensor;
324+ AscendC::LocalTensor<ElementInput> llUbTensor;
325+ AscendC::LocalTensor<ElementInput> scaleTensor;
326+ 
327+ AscendC::LocalTensor<uint8_t> tmpTensor;
328+ uint32_t subBlockIdx_, blockSizeK;
329+ 
330+ template <typename T2, typename T, uint16_t m = 32, bool DEQ = true, uint16_t MAX_COLS = 256, uint16_t DEQ_BLK = 64>
331+ __aicore__ inline void DnSoftmaxQuantBlock(
332+ const AscendC::LocalTensor<T2>& dstTensor,
333+ const AscendC::LocalTensor<T>& srcDnTensor,
334+ const AscendC::LocalTensor<T>& maxTensor,
335+ const AscendC::LocalTensor<T>& inMaxTensor,
336+ const AscendC::LocalTensor<float>& expMaxTensor,
337+ const AscendC::LocalTensor<float>& expSumTensor,
338+ const AscendC::LocalTensor<float>& inSumTensor,
339+ const AscendC::LocalTensor<T>& scaleTensor,
340+ const AscendC::LocalTensor<uint8_t>& tmpTensor,
341+ const uint32_t n)
342+ {
343+ using namespace AscendC::MicroAPI;
344+
345+ constexpr static CastTrait castTraitFp32Zero = {
346+ RegLayout::ZERO, SatMode::UNKNOWN, MaskMergeMode::ZEROING, AscendC::RoundMode::UNKNOWN,
347+ };
348+ constexpr static CastTrait castTraitFp32One = {
349+ RegLayout::ONE, SatMode::UNKNOWN, MaskMergeMode::ZEROING, AscendC::RoundMode::UNKNOWN,
350+ };
351+ 
352+ __ubuf__ T* srcUb = (__ubuf__ T*)srcDnTensor.GetPhyAddr();
353+ __ubuf__ T2* dstUb = (__ubuf__ T2*)dstTensor.GetPhyAddr();
354+ 
355+ __ubuf__ T* inMaxUb = (__ubuf__ T*)inMaxTensor.GetPhyAddr();
356+ __ubuf__ T* maxUb = (__ubuf__ T*)maxTensor.GetPhyAddr();
357+ 
358+ __ubuf__ float* expMaxUb = (__ubuf__ float*)expMaxTensor.GetPhyAddr();
359+ __ubuf__ float* inSumUb = (__ubuf__ float*)inSumTensor.GetPhyAddr();
360+ __ubuf__ float* expSumUb = (__ubuf__ float*)expSumTensor.GetPhyAddr();
361+ 
362+ __ubuf__ float* tmpMaxUbStart = (__ubuf__ float*)tmpTensor.GetPhyAddr();
363+ __ubuf__ T* scaleUb = (__ubuf__ T*) scaleTensor.GetPhyAddr();
364+
365+ T minValue = -65504.0f;
366+ constexpr uint16_t VL = 256;
367+ constexpr uint16_t UNROLL = 4;
368+ constexpr uint16_t VL_ELE_B16 = 128;
369+ constexpr uint16_t R = VL / sizeof(T) / m; // R = 4
370+ 
371+ __VEC_SCOPE__
372+ {
373+ RegTensor<T> vregMax0, vregMax1, vregMax2, vregMax3;
374+ RegTensor<T> vregAcc0, vregAcc1, vregAcc2, vregAcc3;
375+ RegTensor<T> vregSrc0, vregSrc1, vregSrc2, vregSrc3;
376+ RegTensor<T> vregScale;
377+ RegTensor<int8_t> vregZero0, vregZero1, vregFp8High0, vregFp8High1;
378+ RegTensor<T2> vregDst0, vregDst1;
379+
380+ RegTensor<T> vregInMax, vregExpMax16;
381+ RegTensor<T> vregGlobalMax;
382+ RegTensor<float> vregExpMax32;
383+ RegTensor<float> vregInExpSum;
384+ RegTensor<float> vregExpSum32;
385+ RegTensor<T> vregTmp0, vregTmp1;
386+ 
387+ RegTensor<T> vregLocalReduce;
388+ 
389+ MaskReg pregAll = CreateMask<T, MaskPattern::ALL>();
390+ MaskReg pregAll32 = CreateMask<float, MaskPattern::ALL>();
391+ uint32_t tmpM = m;
392+ 
393+ MaskReg pregM = UpdateMask<T>(tmpM);
394+ MaskReg preg64 = CreateMask<T, MaskPattern::VL64>();
395+ MaskReg preg32 = CreateMask<T, MaskPattern::VL32>();
396+ MaskReg pregAllB8 = CreateMask<T2, MaskPattern::ALL>();
397+ MaskReg pregHalfFp8 = CreateMask<T2, MaskPattern::VL128>();
398+ 
399+ Duplicate<T, T>(vregMax0, minValue);
400+ Duplicate<T, T>(vregMax1, minValue);
401+ Duplicate<T, T>(vregMax2, minValue);
402+ Duplicate<T, T>(vregMax3, minValue);
403+ 
404+ DataCopy<T, LoadDist::DIST_NORM>(vregInMax, inMaxUb); // 1111 2222 3333 ... 32,32,32,32
405+ DataCopy<float, LoadDist::DIST_NORM>(vregInExpSum, inSumUb); // 11 22 ... 32,32
406+ 
407+
408+ for (uint16_t k = 0; k < (uint16_t)((n + DEQ_BLK - 1) / DEQ_BLK); ++k) {
409+ uint16_t n2 = VF_MIN(n - k * DEQ_BLK, DEQ_BLK);
410+ DataCopy<T, LoadDist::DIST_BRC_B16>(vregScale, scaleUb + k);
411+
412+ for (uint16_t j = 0; j < (uint16_t)(n2 / R / UNROLL) ; j++) {
413+ uint32_t offset = (k * DEQ_BLK * m) + (j * VL_ELE_B16 * UNROLL);
414+ DataCopy<T, LoadDist::DIST_NORM>(vregSrc0, srcUb + offset + 0);
415+ DataCopy<T, LoadDist::DIST_NORM>(vregSrc1, srcUb + offset + VL_ELE_B16);
416+ DataCopy<T, LoadDist::DIST_NORM>(vregSrc2, srcUb + offset + VL_ELE_B16 *2);
417+ DataCopy<T, LoadDist::DIST_NORM>(vregSrc3, srcUb + offset + VL_ELE_B16 *3);
418+ 
419+ Mul(vregSrc0, vregSrc0, vregScale, pregAll);
420+ Mul(vregSrc1, vregSrc1, vregScale, pregAll);
421+ Mul(vregSrc2, vregSrc2, vregScale, pregAll);
422+ Mul(vregSrc3, vregSrc3, vregScale, pregAll);
423+ 
424+ Max(vregMax0, vregMax0, vregSrc0, pregAll);
425+ StoreAlign(srcUb + offset + VL_ELE_B16*0 ,vregSrc0, pregAll);
426+ Max(vregMax1, vregMax1, vregSrc1, pregAll);
427+ StoreAlign(srcUb + offset + VL_ELE_B16*1 ,vregSrc1, pregAll);
428+ Max(vregMax2, vregMax2, vregSrc2, pregAll);
429+ StoreAlign(srcUb + offset + VL_ELE_B16*2 ,vregSrc2, pregAll);
430+ Max(vregMax3, vregMax3, vregSrc3, pregAll);
431+ StoreAlign(srcUb + offset + VL_ELE_B16*3 ,vregSrc3, pregAll);
432+ }
433+ }
434+ Max(vregMax0, vregMax0, vregMax1, pregAll);
435+ Max(vregMax2, vregMax2, vregMax3, pregAll);
436+ Max(vregMax0, vregMax0, vregMax2, pregAll);
437+
438+ Interleave(vregTmp0, vregTmp1, vregMax0, vregMax0);
439+ Max(vregMax1, vregTmp0, vregTmp1, pregAll); // 11 22, ...., 64,64
440+ Interleave(vregTmp0, vregTmp1, vregMax1, vregMax1);
441+ Max(vregLocalReduce, vregTmp0, vregTmp1, pregAll); // 1111 2222 3333 ... 32,32,32,32
442+
443+ // vregLocalReduce is Local max
444+ Max(vregGlobalMax, vregLocalReduce, vregInMax, pregAll);
445+ Sub(vregExpMax16, vregInMax, vregGlobalMax, pregAll);
446+ Exp(vregExpMax16, vregExpMax16, pregAll);
447+ Cast<float, T, castTraitFp32Zero>(vregExpMax32, vregExpMax16, pregAll); // 11 22 33 .... 32,32
448+ 
449+ DeInterleave(vregTmp0, vregTmp1, vregGlobalMax, vregGlobalMax);
450+ DeInterleave(vregInMax, vregTmp1, vregTmp0, vregTmp0);
451+ 
452+ LocalMemBar<MemType::VEC_STORE, MemType::VEC_LOAD>();
453+ Duplicate<T, T>(vregAcc0, (T)0.0f);
454+ Duplicate<T, T>(vregAcc1, (T)0.0f);
455+ Duplicate<T, T>(vregAcc2, (T)0.0f);
456+ Duplicate<T, T>(vregAcc3, (T)0.0f);
457+ 
458+ for (uint16_t i = 0; i < (uint16_t)(n >> 4); ++i) {
459+ uint32_t offset = i << 9; // 128 * 4
460+ DataCopy<T, LoadDist::DIST_NORM>(vregSrc0, srcUb + offset + 0);
461+ DataCopy<T, LoadDist::DIST_NORM>(vregSrc1, srcUb + offset + VL_ELE_B16);
462+ DataCopy<T, LoadDist::DIST_NORM>(vregSrc2, srcUb + offset + VL_ELE_B16 *2);
463+ DataCopy<T, LoadDist::DIST_NORM>(vregSrc3, srcUb + offset + VL_ELE_B16 *3);
464+ 
465+ Sub(vregSrc0, vregSrc0, vregInMax, pregAll);
466+ Sub(vregSrc1, vregSrc1, vregInMax, pregAll);
467+ Sub(vregSrc2, vregSrc2, vregInMax, pregAll);
468+ Sub(vregSrc3, vregSrc3, vregInMax, pregAll);
469+
470+ Exp(vregSrc0, vregSrc0, pregAll);
471+ Exp(vregSrc1, vregSrc1, pregAll);
472+ Exp(vregSrc2, vregSrc2, pregAll);
473+ Exp(vregSrc3, vregSrc3, pregAll);
474+ 
475+ Add(vregAcc0, vregAcc0, vregSrc0, pregAll);
476+ Add(vregAcc1, vregAcc1, vregSrc1, pregAll);
477+ Add(vregAcc2, vregAcc2, vregSrc2, pregAll);
478+ Add(vregAcc3, vregAcc3, vregSrc3, pregAll);
479+ 
480+ Muls(vregSrc0, vregSrc0, (T)448.0f, pregAll);
481+ Muls(vregSrc1, vregSrc1, (T)448.0f, pregAll);
482+ Muls(vregSrc2, vregSrc2, (T)448.0f, pregAll);
483+ Muls(vregSrc3, vregSrc3, (T)448.0f, pregAll);
484+ 
485+ // rna
486+ Maxs((RegTensor<int16_t>&) vregSrc0, (RegTensor<int16_t>&) vregSrc0, (int16_t)(8128+128), pregAll);
487+ Maxs((RegTensor<int16_t>&) vregSrc1, (RegTensor<int16_t>&) vregSrc1, (int16_t)(8128+128), pregAll);
488+ Maxs((RegTensor<int16_t>&) vregSrc2, (RegTensor<int16_t>&) vregSrc2, (int16_t)(8128+128), pregAll);
489+ Maxs((RegTensor<int16_t>&) vregSrc3, (RegTensor<int16_t>&) vregSrc3, (int16_t)(8128+128), pregAll);
490+ 
491+ Adds((RegTensor<int16_t>&) vregSrc0, (RegTensor<int16_t>&) vregSrc0, (int16_t)(-8128), pregAll);
492+ Adds((RegTensor<int16_t>&) vregSrc1, (RegTensor<int16_t>&) vregSrc1, (int16_t)(-8128), pregAll);
493+ Adds((RegTensor<int16_t>&) vregSrc2, (RegTensor<int16_t>&) vregSrc2, (int16_t)(-8128), pregAll);
494+ Adds((RegTensor<int16_t>&) vregSrc3, (RegTensor<int16_t>&) vregSrc3, (int16_t)(-8128), pregAll);
495+ 
496+ ShiftRights((RegTensor<int16_t>&) vregSrc0, (RegTensor<int16_t>&) vregSrc0, (int16_t)7, pregAll);
497+ ShiftRights((RegTensor<int16_t>&) vregSrc1, (RegTensor<int16_t>&) vregSrc1, (int16_t)7, pregAll);
498+ ShiftRights((RegTensor<int16_t>&) vregSrc2, (RegTensor<int16_t>&) vregSrc2, (int16_t)7, pregAll);
499+ ShiftRights((RegTensor<int16_t>&) vregSrc3, (RegTensor<int16_t>&) vregSrc3, (int16_t)7, pregAll);
500+ 
501+ StoreAlign<T2, StoreDist::DIST_PACK_B16>(dstUb + offset + VL_ELE_B16 * 0, (RegTensor<T2>&) vregSrc0, pregAll);
502+ StoreAlign<T2, StoreDist::DIST_PACK_B16>(dstUb + offset + VL_ELE_B16 * 1, (RegTensor<T2>&) vregSrc1, pregAll);
503+ StoreAlign<T2, StoreDist::DIST_PACK_B16>(dstUb + offset + VL_ELE_B16 * 2, (RegTensor<T2>&) vregSrc2, pregAll);
504+ StoreAlign<T2, StoreDist::DIST_PACK_B16>(dstUb + offset + VL_ELE_B16 * 3, (RegTensor<T2>&) vregSrc3, pregAll); // B16 mask reg need
505+ }
506+ Add(vregAcc0, vregAcc0, vregAcc1, pregAll);
507+ Add(vregAcc2, vregAcc2, vregAcc3, pregAll);
508+ Add(vregAcc0, vregAcc0, vregAcc2, pregAll);
509+ 
510+ Interleave(vregTmp0, vregTmp1, vregAcc0, vregAcc0);
511+ Add(vregAcc1, vregTmp0, vregTmp1, pregAll);
512+ Interleave(vregTmp0, vregTmp1, vregAcc1, vregAcc1);
513+ Add(vregLocalReduce, vregTmp0, vregTmp1, pregAll); // 1111 2222 .... 32,32,32,32
514+ Cast<float, T, castTraitFp32Zero>(vregExpSum32, vregLocalReduce, pregAll); // 11 22 ... 32,32
515+ 
516+ // x_sum = sum(exp_max * in_sum + x_sum)
517+ Mul<float, MaskMergeMode::ZEROING>(vregInExpSum, vregExpMax32, vregInExpSum, pregAll32);
518+ Add<float, MaskMergeMode::ZEROING>(vregInExpSum, vregInExpSum, vregExpSum32, pregAll32);
519+ 
520+ DataCopy<float, StoreDist::DIST_NORM>(expSumUb, vregInExpSum, pregAll32);
521+ DataCopy<float, StoreDist::DIST_NORM>(expMaxUb, vregExpMax32, pregAll32);
522+ DataCopy<T, StoreDist::DIST_NORM>(maxUb, vregGlobalMax, pregAll);
523+ }
524+ }
525+};
526+}
527+ 
528+#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_BSA_LOW_PREC_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue_online_softmax_arch35_reg_low_prec_bf16.hpp+684-0
@@ -0,0 +1,684 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_ARCH35_REG_LOW_PREC_BF16_HPP
12+#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_ARCH35_REG_LOW_PREC_BF16_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/resource.hpp"
16+#include "../../../attn_infra/epilogue/dispatch_policy.hpp"
17+#include "../../../attn_infra/epilogue/tile_common/tile_copy.hpp"
18+#include "../../../attn_infra/gemm_coord.hpp"
19+#include "../../../attn_infra/matrix_coord.hpp"
20+#include "../../../tla/tensor.hpp"
21+#include "../../../tla/layout.hpp"
22+ 
23+namespace NpuArch::Epilogue::Block {
24+ 
25+enum class KvBaseTileRegSplitStagesBf16 {
26+ ONE,
27+ TWO
28+};
29+ 
30+template <
31+ class OutputType_,
32+ class LayoutS_>
33+class BlockEpilogue<
34+ EpilogueOnlineSoftmaxBsa,
35+ OutputType_,
36+ Gemm::GemmType<bfloat16_t, LayoutS_>>
37+{
38+public:
39+ using DispatchPolicy = EpilogueOnlineSoftmaxBsa;
40+ using ArchTag = typename DispatchPolicy::ArchTag;
41+ using ElementOutput = typename OutputType_::Element;
42+ using ElementInput = bfloat16_t;
43+ 
44+ using LayoutOutput = typename OutputType_::Layout;
45+ using LayoutInput = LayoutS_;
46+ 
47+ static constexpr uint32_t BLOCK_SIZE_IN_BYTE = 32;
48+ static constexpr uint32_t REPEAT_SIZE_IN_BYTE = 256;
49+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
50+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
51+ static constexpr uint32_t HALF_VECTOR_SIZE = 128;
52+ static constexpr uint32_t BLOCK_SIZE = 16;
53+ static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024;
54+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 32768;
55+ static constexpr uint32_t VECTOR_SIZE = 128;
56+ static constexpr uint32_t MAX_UB_S_ELEM_NUM = 16384;
57+ static constexpr uint32_t DM_UB_GLOBAL_ELEM_NUM = 64;
58+ static constexpr uint32_t ELE_NUM_PER_C0 = 16;
59+ static constexpr uint32_t C0_NUM_PER_FRACTAL = 16;
60+ 
61+ static constexpr uint32_t REDUCE_UB_SIZE = 1024;
62+ static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32;
63+ static constexpr uint32_t ROW_OPS_SPEC_MASK_8 = 8;
64+ static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4;
65+ static constexpr uint32_t ROW_OPS_SPEC_MASK_2 = 2;
66+ static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256;
67+ static constexpr int64_t UB_FLOAT_LINE_SIZE = 64;
68+ 
69+ static constexpr uint32_t SPLIT_COL_IDX_2 = 2;
70+ static constexpr uint32_t SPLIT_COL_IDX_3 = 3;
71+ static constexpr uint32_t HALF_REP_SIZE = 128;
72+ static constexpr uint32_t FLOAT_REP_SIZE = 64;
73+ static constexpr uint32_t BLOCK_REP_SIZE = 8;
74+ static constexpr uint32_t REPEAT_STRIDE = 1;
75+ static constexpr uint32_t SM_ROW_MAX_ELEM_NUM = 64;
76+ static constexpr uint32_t SM_COL_MAX_ELEM_NUM = 256;
77+ static constexpr uint32_t SM_VREG_SIZE = 256 / sizeof(ElementInput);
78+ 
79+ __aicore__ inline
80+ BlockEpilogue(Arch::Resource<ArchTag> &resource, float scaleValue_)
81+ {
82+ // Allocate UB space
83+ constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;
84+ constexpr uint32_t LP_UB_TENSOR_OFFSET = 2 * UB_UINT8_BLOCK_SIZE;
85+ 
86+ constexpr uint32_t LM_UB_TENSOR_OFFSET = 7 * UB_UINT8_BLOCK_SIZE;
87+ constexpr uint32_t GM_UB_TENSOR_OFFSET = LM_UB_TENSOR_OFFSET + 64 * sizeof(float);
88+ constexpr uint32_t DM_UB_TENSOR_OFFSET = GM_UB_TENSOR_OFFSET + 64 * sizeof(float);
89+ constexpr uint32_t LL_UB_TENSOR_OFFSET = DM_UB_TENSOR_OFFSET + 3 * 64 * sizeof(float);
90+ constexpr uint32_t GL_UB_TENSOR_OFFSET = LL_UB_TENSOR_OFFSET + 64 * sizeof(float);
91+ 
92+ subBlockIdx_ = AscendC::GetSubBlockIdx();
93+ scaleValue = AscendC::ToBfloat16(scaleValue_);
94+ MIN_VALUE = AscendC::ToBfloat16(-3.389531390315715675e+38);
95+ 
96+ lsUbTensor = resource.ubBuf.template GetBufferByByte<ElementInput>(LS_UB_TENSOR_OFFSET);
97+ lpUbTensor = resource.ubBuf.template GetBufferByByte<ElementOutput>(LP_UB_TENSOR_OFFSET);
98+ gmUbTensor = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
99+ glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
100+ dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
101+ lmUbTensor = resource.ubBuf.template GetBufferByByte<ElementInput>(LM_UB_TENSOR_OFFSET);
102+ llUbTensor = resource.ubBuf.template GetBufferByByte<ElementInput>(LL_UB_TENSOR_OFFSET);
103+ lmUbFloatTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_TENSOR_OFFSET);
104+ llUbFloatTensor = resource.ubBuf.template GetBufferByByte<float>(LL_UB_TENSOR_OFFSET);
105+ }
106+ 
107+ __aicore__ inline
108+ ~BlockEpilogue()
109+ {
110+ }
111+ 
112+ template <class TensorDst, class TensorSrc>
113+ __aicore__ inline
114+ void CopyPUbToPL1(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t m)
115+ {
116+ const uint32_t blockCount = tla::get<1, 1>(srcTensor.shape());
117+ const uint32_t blockLen = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape());
118+ const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
119+ 
120+ AscendC::DataCopyParams repeatParams;
121+ 
122+ repeatParams.blockCount = blockCount;
123+ repeatParams.blockLen = m;
124+ repeatParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / ELE_NUM_PER_C0 - m;
125+ repeatParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / ELE_NUM_PER_C0 - m;
126+ 
127+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
128+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
129+ 
130+ AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams);
131+ }
132+ 
133+ template <uint32_t MODE, pipe_t PIPE>
134+ __aicore__ inline
135+ void SetCrossCoreSync(Arch::CrossCoreFlag &crossCoreFlag)
136+ {
137+ // in mode 4, AIC set for 2 AIVs seperately
138+ if constexpr (MODE == 4U) {
139+ Arch::CrossCoreSetFlag<MODE, PIPE>(crossCoreFlag);
140+ }
141+ }
142+ 
143+ template <uint32_t MODE, pipe_t PIPE>
144+ __aicore__ inline
145+ void WaitCrossCoreSync(Arch::CrossCoreFlag &crossCoreFlag)
146+ {
147+ // in mode 4, AIC wait for 2 AIVs seperately
148+ if constexpr (MODE == 4U) {
149+ Arch::CrossCoreWaitFlag<MODE, PIPE>(crossCoreFlag);
150+ }
151+ }
152+ 
153+ template <class TensorP>
154+ __aicore__ inline
155+ void operator()(TensorP &l1PTensorTla, GemmCoord actualBlockShape,
156+ uint32_t isFirstKvSTile, uint32_t ubSBufId, uint32_t l1PBufId,
157+ Arch::CrossCoreFlag mm1ToSmFlag, Arch::CrossCoreFlag smToMm2Flag,
158+ int actualRownum=0){
159+ 
160+ struct EmptyTensor {};
161+ EmptyTensor dummyQS;
162+ EmptyTensor dummyKS;
163+ operator()<0>(l1PTensorTla, dummyQS, dummyKS,
164+ actualBlockShape,
165+ isFirstKvSTile, ubSBufId, l1PBufId,
166+ mm1ToSmFlag, smToMm2Flag, AscendC::GlobalTensor<int32_t>(),
167+ 0, actualRownum);
168+ }
169+
170+ template <int quant_mode, class TensorP, class TensorQS, class TensorKS, typename ElementIndex>
171+ __aicore__ inline
172+ void operator()(TensorP &l1PTensorTla, TensorQS &gmQSTensorTla, TensorKS &gmKSTensorTla,
173+ GemmCoord actualBlockShape,
174+ uint32_t isFirstKvSTile, uint32_t ubSBufId, uint32_t l1PBufId,
175+ Arch::CrossCoreFlag mm1ToSmFlag, Arch::CrossCoreFlag smToMm2Flag,
176+ AscendC::GlobalTensor<ElementIndex> sparseIndex, int blockNum=0, int actualRownum=0)
177+ {
178+ uint32_t mCopyOffset = (actualRownum > 0) ?
179+ static_cast<uint32_t>(actualRownum) / 2 : RoundUp(actualBlockShape.m(), 8) / 2;
180+ uint32_t m = actualBlockShape.m() < mCopyOffset ? actualBlockShape.m() : mCopyOffset;
181+ m = subBlockIdx_ == 0 ? m : actualBlockShape.m() - m;
182+ if (m == 0) {
183+ WaitCrossCoreSync<4, PIPE_V>(mm1ToSmFlag);
184+ SetCrossCoreSync<4, PIPE_V>(mm1ToSmFlag);
185+ WaitCrossCoreSync<4, PIPE_MTE3>(smToMm2Flag);
186+ SetCrossCoreSync<4, PIPE_MTE3>(smToMm2Flag);
187+ return;
188+ }
189+ uint32_t n = actualBlockShape.n();
190+ uint16_t mRound = RoundUp(m, C0_NUM_PER_FRACTAL);
191+ uint16_t nRound = RoundUp(n, ELE_NUM_PER_C0);
192+ uint32_t blockStride = mRound;
193+ constexpr int16_t vlSize = static_cast<int16_t>(AscendC::GetVecLen() / sizeof(ElementInput));
194+ constexpr int16_t vlFloatSize = static_cast<int16_t>(AscendC::GetVecLen() / sizeof(float));
195+ int16_t nLoops = AscendC::CeilDivision(n, vlSize) - 1;
196+ uint32_t tailN = (n - 1) % vlSize + 1;
197+ int16_t mLoops = AscendC::CeilDivision(m, vlFloatSize) - 1;
198+ uint32_t tailM = (m - 1) % vlFloatSize + 1;
199+ uint32_t nPadding = (tailN + BLOCK_SIZE_IN_BYTE - 1) / BLOCK_SIZE_IN_BYTE * BLOCK_SIZE_IN_BYTE;
200+ __ubuf__ ElementOutput *pAddr = (__ubuf__ ElementOutput*) lpUbTensor[ubSBufId * MAX_UB_S_ELEM_NUM].GetPhyAddr();
201+ __ubuf__ ElementInput *sAddr = (__ubuf__ ElementInput*) lsUbTensor[ubSBufId * MAX_UB_S_ELEM_NUM].GetPhyAddr();
202+ __ubuf__ float *lastMaxAddr = (__ubuf__ float *)gmUbTensor.GetPhyAddr();
203+ __ubuf__ float *lastSumAddr = (__ubuf__ float*) glUbTensor.GetPhyAddr();
204+ __ubuf__ ElementInput *nowMaxAddr = (__ubuf__ ElementInput*) lmUbTensor.GetPhyAddr();
205+ __ubuf__ float *nowMaxFloatAddr = (__ubuf__ float*) lmUbFloatTensor.GetPhyAddr();
206+ __ubuf__ float *nowSumAddr = (__ubuf__ float*) llUbFloatTensor.GetPhyAddr();
207+ __ubuf__ float *expMaxUbAddr = (__ubuf__ float *)dmUbTensor[l1PBufId * DM_UB_GLOBAL_ELEM_NUM].GetPhyAddr();
208+ 
209+ // wait QK Fixpipe finsh
210+ WaitCrossCoreSync<4, PIPE_V>(mm1ToSmFlag);
211+ if (isFirstKvSTile) {
212+ nowMaxFloatAddr = lastMaxAddr;
213+ nowSumAddr = lastSumAddr;
214+ }
215+ uint32_t kvBaseTileRegStages = CeilDiv(n, SM_VREG_SIZE);
216+ if (kvBaseTileRegStages == 1) {
217+ ComputeScaleAndMax<KvBaseTileRegSplitStagesBf16::ONE>(
218+ sAddr, nowMaxFloatAddr, m, tailN, nPadding, scaleValue, nRound);
219+ } else if (kvBaseTileRegStages == 2) {
220+ ComputeScaleAndMax<KvBaseTileRegSplitStagesBf16::TWO>(
221+ sAddr, nowMaxFloatAddr, m, tailN, nPadding, scaleValue, nRound);
222+ }
223+ 
224+ if (!isFirstKvSTile) {
225+ UpdateMax(nowMaxFloatAddr, lastMaxAddr, mLoops, tailM);
226+ }
227+ 
228+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(ubSBufId + 2);
229+ uint32_t tailNOdd = tailN / 2 ;
230+ uint32_t tailNEven = tailNOdd + tailN % 2;
231+ if (kvBaseTileRegStages == 1) {
232+ ComputeExpSubSum16<KvBaseTileRegSplitStagesBf16::ONE>(
233+ pAddr, sAddr, nowMaxFloatAddr, nowSumAddr, m, tailN, blockStride, nRound, tailNOdd, tailNEven);
234+ } else if (kvBaseTileRegStages == 2) {
235+ ComputeExpSubSum16<KvBaseTileRegSplitStagesBf16::TWO>(
236+ pAddr, sAddr, nowMaxFloatAddr, nowSumAddr, m, tailN, blockStride, nRound, tailNOdd, tailNEven);
237+ }
238+ 
239+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(ubSBufId);
240+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(ubSBufId);
241+ SetCrossCoreSync<4, PIPE_V>(mm1ToSmFlag);
242+ 
243+ auto ubPLayoutTla = tla::MakeLayout<ElementOutput, LayoutOutput>(mRound, nRound);
244+ auto ubPTensorTla = tla::MakeTensor(lpUbTensor[ubSBufId * MAX_UB_S_ELEM_NUM],
245+ ubPLayoutTla, Arch::PositionUB{});
246+ auto ubPTensorTlaTile = GetTile(ubPTensorTla,
247+ tla::MakeCoord(0, 0), tla::MakeShape(m, n));
248+ auto l1PTensorTlaTile = GetTile(l1PTensorTla,
249+ tla::MakeCoord(subBlockIdx_ * mCopyOffset, 0), tla::MakeShape(m, n));
250+ WaitCrossCoreSync<4, PIPE_MTE3>(smToMm2Flag);
251+ CopyPUbToPL1(l1PTensorTlaTile, ubPTensorTlaTile, m);
252+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(ubSBufId + 2);
253+ // crossCoreSync after PIPE_MTE1 move
254+ SetCrossCoreSync<4, PIPE_MTE3>(smToMm2Flag);
255+ if (!isFirstKvSTile) {
256+ UpdateExpSumAndExpMax(
257+ lastSumAddr, expMaxUbAddr, lastMaxAddr, nowSumAddr, nowMaxFloatAddr, mLoops, tailM);
258+ }
259+ AscendC::PipeBarrier<PIPE_V>();
260+ }
261+ 
262+private:
263+ ElementInput scaleValue;
264+ AscendC::LocalTensor<ElementInput> lsUbTensor;
265+ AscendC::LocalTensor<ElementOutput> lpUbTensor;
266+ AscendC::LocalTensor<float> gmUbTensor;
267+ AscendC::LocalTensor<float> glUbTensor;
268+ AscendC::LocalTensor<float> dmUbTensor;
269+ AscendC::LocalTensor<ElementInput> lmUbTensor;
270+ AscendC::LocalTensor<ElementInput> llUbTensor;
271+ AscendC::LocalTensor<float> lmUbFloatTensor;
272+ AscendC::LocalTensor<float> llUbFloatTensor;
273+ uint32_t subBlockIdx_;
274+ ElementInput MIN_VALUE;
275+ 
276+ template <KvBaseTileRegSplitStagesBf16 kvBaseTileRegSplitStages>
277+ __simd_vf__ inline void ComputeScaleAndMax(
278+ __ubuf__ ElementInput *srcUb, __ubuf__ float *newMaxUb,
279+ uint16_t m, uint32_t tailN, uint32_t nPadding, ElementInput dScale, uint16_t S2BaseSize)
280+ {
281+ }
282+
283+ template <>
284+ __simd_vf__ inline void ComputeScaleAndMax<KvBaseTileRegSplitStagesBf16::ONE>(
285+ __ubuf__ ElementInput *srcUb, __ubuf__ float *newMaxUb,
286+ uint16_t m, uint32_t tailN, uint32_t nPadding, ElementInput dScale, uint16_t S2BaseSize)
287+ {
288+ using namespace AscendC::MicroAPI;
289+ 
290+ constexpr static CastTrait castTraitZero = {
291+ RegLayout::ZERO,
292+ SatMode::UNKNOWN,
293+ MaskMergeMode::ZEROING,
294+ AscendC::RoundMode::UNKNOWN,
295+ };
296+ constexpr static CastTrait castTraitOne = {
297+ RegLayout::ONE,
298+ SatMode::UNKNOWN,
299+ MaskMergeMode::ZEROING,
300+ AscendC::RoundMode::UNKNOWN,
301+ };
302+ 
303+ RegTensor<ElementInput> minVreg;
304+ RegTensor<ElementInput> srcVreg;
305+ // RegTensor<ElementInput> maxSrcVreg;
306+ RegTensor<ElementInput> maxTmpVreg;
307+ RegTensor<ElementInput> scaleVreg;
308+ RegTensor<float> maxFloatVreg0;
309+ RegTensor<float> maxFloatVreg1;
310+ RegTensor<float> maxTmpFloatVreg;
311+ RegTensor<float> maxTmpFloatVreg0;
312+ RegTensor<float> maxTmpFloatVreg1;
313+ UnalignReg maxUreg;
314+ MaskReg pregCompare;
315+ MaskReg pregFull = CreateMask<ElementInput, MaskPattern::ALL>();
316+ MaskReg pregFloatFull = CreateMask<float, MaskPattern::ALL>();
317+ MaskReg pregTailN = UpdateMask<ElementInput>(tailN);
318+ MaskReg pregFloatTailN = UpdateMask<float>(tailN);
319+ 
320+ Duplicate(minVreg, MIN_VALUE);
321+ Duplicate(scaleVreg, dScale);
322+ for (uint16_t i = 0; i < m; ++i) {
323+ LoadAlign(srcVreg, srcUb + i * S2BaseSize);
324+ Mul(srcVreg, srcVreg, scaleVreg, pregFull);
325+ Select(srcVreg, srcVreg, minVreg, pregTailN);
326+ StoreAlign<ElementInput, StoreDist::DIST_NORM_B16>(
327+ srcUb + i * S2BaseSize, srcVreg, pregTailN);
328+ Cast<float, ElementInput, castTraitZero>(maxFloatVreg0, srcVreg, pregFull);
329+ Cast<float, ElementInput, castTraitOne>(maxFloatVreg1, srcVreg, pregFull);
330+ ReduceMax(maxTmpFloatVreg0, maxFloatVreg0, pregFull);
331+ ReduceMax(maxTmpFloatVreg1, maxFloatVreg1, pregFull);
332+ Max(maxTmpFloatVreg, maxTmpFloatVreg0, maxTmpFloatVreg1, pregFull);
333+ StoreUnAlign<float, PostLiteral::POST_MODE_UPDATE>(newMaxUb, maxTmpFloatVreg, maxUreg, 1);
334+ }
335+ vstas(maxUreg, newMaxUb, 0, POST_UPDATE);
336+ }
337+ 
338+ template <>
339+ __simd_vf__ inline void ComputeScaleAndMax<KvBaseTileRegSplitStagesBf16::TWO>(
340+ __ubuf__ ElementInput *srcUb, __ubuf__ float *newMaxUb,
341+ uint16_t m, uint32_t tailN, uint32_t nPadding, ElementInput dScale, uint16_t S2BaseSize)
342+ {
343+ using namespace AscendC::MicroAPI;
344+ constexpr static CastTrait castTraitZero = {
345+ RegLayout::ZERO,
346+ SatMode::UNKNOWN,
347+ MaskMergeMode::ZEROING,
348+ AscendC::RoundMode::UNKNOWN,
349+ };
350+ constexpr static CastTrait castTraitOne = {
351+ RegLayout::ONE,
352+ SatMode::UNKNOWN,
353+ MaskMergeMode::ZEROING,
354+ AscendC::RoundMode::UNKNOWN,
355+ };
356+ 
357+ RegTensor<ElementInput> minVreg;
358+ RegTensor<ElementInput> srcVreg0;
359+ RegTensor<ElementInput> srcVreg1;
360+ // RegTensor<ElementInput> maxSrcVreg;
361+ RegTensor<ElementInput> maxTmpVreg;
362+ RegTensor<ElementInput> scaleVreg;
363+ RegTensor<float> maxFloatVreg0;
364+ RegTensor<float> maxFloatVreg1;
365+ RegTensor<float> maxTmpFloatVreg;
366+ RegTensor<float> maxTmpFloatVreg0;
367+ RegTensor<float> maxTmpFloatVreg1;
368+ UnalignReg maxUreg;
369+ MaskReg pregCompare;
370+ MaskReg pregFull = CreateMask<ElementInput, MaskPattern::ALL>();
371+ MaskReg pregFloatFull = CreateMask<float, MaskPattern::ALL>();
372+ MaskReg pregTailN = UpdateMask<ElementInput>(tailN);
373+ MaskReg pregFloatTailN = UpdateMask<float>(tailN);
374+ 
375+ Duplicate(minVreg, MIN_VALUE);
376+ Duplicate(scaleVreg, dScale);
377+ for (uint16_t i = 0; i < m; ++i) {
378+ LoadAlign(srcVreg0, srcUb + i * S2BaseSize);
379+ LoadAlign(srcVreg1, srcUb + i * S2BaseSize + HALF_REP_SIZE);
380+ Mul(srcVreg0, srcVreg0, scaleVreg, pregFull);
381+ Mul(srcVreg1, srcVreg1, scaleVreg, pregFull);
382+ StoreAlign<ElementInput, StoreDist::DIST_NORM_B16>(
383+ srcUb + i * S2BaseSize, srcVreg0, pregFull);
384+ StoreAlign<ElementInput, StoreDist::DIST_NORM_B16>(
385+ srcUb + i * S2BaseSize + HALF_REP_SIZE, srcVreg1, pregTailN);
386+ Max<ElementInput, MaskMergeMode::MERGING>(srcVreg0, srcVreg0, srcVreg1, pregTailN);
387+ 
388+ Cast<float, ElementInput, castTraitZero>(maxFloatVreg0, srcVreg0, pregFull);
389+ Cast<float, ElementInput, castTraitOne>(maxFloatVreg1, srcVreg0, pregFull);
390+ ReduceMax(maxTmpFloatVreg0, maxFloatVreg0, pregFull);
391+ ReduceMax(maxTmpFloatVreg1, maxFloatVreg1, pregFull);
392+ Max(maxTmpFloatVreg, maxTmpFloatVreg0, maxTmpFloatVreg1, pregFull);
393+ StoreUnAlign<float, PostLiteral::POST_MODE_UPDATE>(newMaxUb, maxTmpFloatVreg, maxUreg, 1);
394+ }
395+ vstas(maxUreg, newMaxUb, 0, POST_UPDATE);
396+ }
397+ 
398+ template <typename ElementS>
399+ __simd_vf__ inline void CastMax(
400+ __ubuf__ ElementS *nowMaxUb, __ubuf__ float *nowMaxFloatUb, uint16_t mLoops, uint32_t tailM)
401+ {
402+ using namespace AscendC::MicroAPI;
403+ constexpr static CastTrait castTraitZero = {
404+ RegLayout::ZERO,
405+ SatMode::UNKNOWN,
406+ MaskMergeMode::ZEROING,
407+ AscendC::RoundMode::UNKNOWN,
408+ };
409+ 
410+ RegTensor<ElementS> nowMaxVreg;
411+ RegTensor<float> nowMaxFloatVreg;
412+ RegTensor<ElementS> maxVreg;
413+ 
414+ MaskReg pregFull = CreateMask<ElementS, MaskPattern::ALL>();
415+ MaskReg pregFloatFull = CreateMask<float, MaskPattern::ALL>();
416+ MaskReg pregTailM = UpdateMask<ElementS>(tailM);
417+ MaskReg pregFloatTailM = UpdateMask<float>(tailM);
418+ for (uint16_t i = 0; i < mLoops; ++i) {
419+ LoadAlign(nowMaxVreg, nowMaxUb + i * HALF_REP_SIZE);
420+ Cast<float, ElementS, castTraitZero>(nowMaxFloatVreg, nowMaxVreg, pregFull);
421+ StoreAlign<float, StoreDist::DIST_NORM_B32>(
422+ nowMaxFloatUb + i * FLOAT_REP_SIZE, nowMaxFloatVreg, pregFloatFull);
423+ }
424+ LoadAlign(nowMaxVreg, nowMaxUb + mLoops * HALF_REP_SIZE);
425+ Cast<float, ElementS, castTraitZero>(nowMaxFloatVreg, nowMaxVreg, pregFull);
426+ StoreAlign<float, StoreDist::DIST_NORM_B32>(
427+ nowMaxFloatUb + mLoops * FLOAT_REP_SIZE, nowMaxFloatVreg, pregFloatTailM);
428+ }
429+ 
430+ __simd_vf__ inline void UpdateMax(
431+ __ubuf__ float *nowMaxUb, __ubuf__ float *lastMaxUb, uint16_t mLoops, uint32_t tailM)
432+ {
433+ using namespace AscendC::MicroAPI;
434+ 
435+ RegTensor<float> nowMaxVreg;
436+ RegTensor<float> lastMaxFloatVreg;
437+ RegTensor<float> maxVreg;
438+ 
439+ MaskReg pregFloatFull = CreateMask<float, MaskPattern::ALL>();
440+ MaskReg pregFloatTailM = UpdateMask<float>(tailM);
441+ for (uint16_t i = 0; i < mLoops; ++i) {
442+ LoadAlign(lastMaxFloatVreg, lastMaxUb + i * FLOAT_REP_SIZE);
443+ LoadAlign(nowMaxVreg, nowMaxUb + i * FLOAT_REP_SIZE);
444+ Max(maxVreg, nowMaxVreg, lastMaxFloatVreg, pregFloatFull);
445+ StoreAlign<float, StoreDist::DIST_NORM_B32>(nowMaxUb + i * FLOAT_REP_SIZE, maxVreg, pregFloatFull);
446+ }
447+ LoadAlign(lastMaxFloatVreg, lastMaxUb + mLoops * FLOAT_REP_SIZE);
448+ LoadAlign(nowMaxVreg, nowMaxUb + mLoops * FLOAT_REP_SIZE);
449+ Max(maxVreg, nowMaxVreg, lastMaxFloatVreg, pregFloatFull);
450+ StoreAlign<float, StoreDist::DIST_NORM_B32>(nowMaxUb + mLoops * FLOAT_REP_SIZE, maxVreg, pregFloatTailM);
451+ }
452+ 
453+ template <KvBaseTileRegSplitStagesBf16 kvBaseTileRegSplitStages>
454+ __simd_vf__ inline void ComputeExpSubSum16(
455+ __ubuf__ ElementOutput *expUb, __ubuf__ ElementInput *srcUb,
456+ __ubuf__ float *nowMaxUb, __ubuf__ float *expSumUb,
457+ uint16_t m, uint32_t tailN, uint32_t blockStride,
458+ uint16_t S2BaseSize, uint32_t tailNOdd, uint32_t tailNEven)
459+ {
460+ }
461+ 
462+ template <>
463+ __simd_vf__ inline void ComputeExpSubSum16<KvBaseTileRegSplitStagesBf16::ONE>(
464+ __ubuf__ ElementOutput *expUb, __ubuf__ ElementInput *srcUb,
465+ __ubuf__ float *nowMaxUb, __ubuf__ float *expSumUb,
466+ uint16_t m, uint32_t tailN, uint32_t blockStride,
467+ uint16_t S2BaseSize, uint32_t tailNOdd, uint32_t tailNEven)
468+ {
469+ using namespace AscendC::MicroAPI;
470+ constexpr static CastTrait castTraitZero = {
471+ RegLayout::ZERO,
472+ SatMode::UNKNOWN,
473+ MaskMergeMode::ZEROING,
474+ AscendC::RoundMode::UNKNOWN,
475+ };
476+ constexpr static CastTrait castTraitOne = {
477+ RegLayout::ONE,
478+ SatMode::UNKNOWN,
479+ MaskMergeMode::ZEROING,
480+ AscendC::RoundMode::UNKNOWN,
481+ };
482+ 
483+ constexpr static CastTrait castTraitZeroDown = {
484+ RegLayout::ZERO,
485+ SatMode::SAT,
486+ MaskMergeMode::ZEROING,
487+ AscendC::RoundMode::CAST_ROUND,
488+ };
489+ 
490+ constexpr static CastTrait castTraitOneDown = {
491+ RegLayout::ONE,
492+ SatMode::SAT,
493+ MaskMergeMode::ZEROING,
494+ AscendC::RoundMode::CAST_ROUND,
495+ };
496+ 
497+ RegTensor<ElementInput> expVreg;
498+ RegTensor<float> expFloatVreg0;
499+ RegTensor<float> expFloatVreg1;
500+ RegTensor<float> expSumVreg;
501+ RegTensor<float> maxVreg;
502+ 
503+ RegTensor<float> expDstFloatVreg0;
504+ RegTensor<float> expDstFloatVreg1;
505+ RegTensor<ElementInput> expDstVreg;
506+ RegTensor<ElementInput> expDstVreg0;
507+ RegTensor<ElementInput> expDstVreg1;
508+ 
509+ UnalignReg expSumUreg;
510+ 
511+ MaskReg pregFull = CreateMask<ElementInput, MaskPattern::ALL>();
512+ MaskReg pregFloatFull = CreateMask<float, MaskPattern::ALL>();
513+ MaskReg pregTailN = UpdateMask<ElementInput>(tailN);
514+ MaskReg pregtailNOdd= UpdateMask<float>(tailNOdd);
515+ MaskReg pregtailNEven = UpdateMask<float>(tailNEven);
516+ for (uint16_t i = 0; i < m; ++i) {
517+ LoadAlign<float, LoadDist::DIST_BRC_B32>(maxVreg, nowMaxUb + i);
518+ Duplicate(expSumVreg, 0);
519+ LoadAlign(expVreg, srcUb + i * S2BaseSize);
520+ Cast<float, ElementInput, castTraitZero>(expFloatVreg0, expVreg, pregFull);
521+ Cast<float, ElementInput, castTraitOne>(expFloatVreg1, expVreg, pregFull);
522+ FusedExpSub(expDstFloatVreg0, expFloatVreg0, maxVreg, pregtailNEven);
523+ FusedExpSub(expDstFloatVreg1, expFloatVreg1, maxVreg, pregtailNOdd);
524+ Add<float, MaskMergeMode::MERGING>(expSumVreg, expSumVreg, expDstFloatVreg0, pregtailNEven);
525+ Add<float, MaskMergeMode::MERGING>(expSumVreg, expSumVreg, expDstFloatVreg1, pregtailNOdd);
526+ Cast<ElementInput, float, castTraitZeroDown>(expDstVreg0, expDstFloatVreg0, pregFloatFull);
527+ Cast<ElementInput, float, castTraitOneDown>(expDstVreg1, expDstFloatVreg1, pregFloatFull);
528+ Or((RegTensor<uint16_t>&)expDstVreg,
529+ (RegTensor<uint16_t>&)expDstVreg0, (RegTensor<uint16_t>&)expDstVreg1,
530+ pregFull);
531+ StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(
532+ expUb + i * ELE_NUM_PER_C0,
533+ expDstVreg, blockStride, pregTailN);
534+ ReduceSum(expSumVreg, expSumVreg, pregFull);
535+ StoreUnAlign<float, PostLiteral::POST_MODE_UPDATE>(expSumUb, expSumVreg, expSumUreg, 1);
536+ }
537+ vstas(expSumUreg, expSumUb, 0, POST_UPDATE);
538+ }
539+ 
540+ template <>
541+ __simd_vf__ inline void ComputeExpSubSum16<KvBaseTileRegSplitStagesBf16::TWO>(
542+ __ubuf__ ElementOutput *expUb, __ubuf__ ElementInput *srcUb,
543+ __ubuf__ float *nowMaxUb, __ubuf__ float *expSumUb,
544+ uint16_t m, uint32_t tailN, uint32_t blockStride,
545+ uint16_t S2BaseSize, uint32_t tailNOdd, uint32_t tailNEven)
546+ {
547+ using namespace AscendC::MicroAPI;
548+ constexpr static CastTrait castTraitZero = {
549+ RegLayout::ZERO,
550+ SatMode::UNKNOWN,
551+ MaskMergeMode::ZEROING,
552+ AscendC::RoundMode::UNKNOWN,
553+ };
554+ constexpr static CastTrait castTraitOne = {
555+ RegLayout::ONE,
556+ SatMode::UNKNOWN,
557+ MaskMergeMode::ZEROING,
558+ AscendC::RoundMode::UNKNOWN,
559+ };
560+ 
561+ constexpr static CastTrait castTraitZeroDown = {
562+ RegLayout::ZERO,
563+ SatMode::SAT,
564+ MaskMergeMode::ZEROING,
565+ AscendC::RoundMode::CAST_ROUND,
566+ };
567+ 
568+ constexpr static CastTrait castTraitOneDown = {
569+ RegLayout::ONE,
570+ SatMode::SAT,
571+ MaskMergeMode::ZEROING,
572+ AscendC::RoundMode::CAST_ROUND,
573+ };
574+ 
575+ RegTensor<ElementInput> expVreg0;
576+ RegTensor<ElementInput> expVreg1;
577+ RegTensor<float> expFloatVreg0;
578+ RegTensor<float> expFloatVreg1;
579+ RegTensor<float> expFloatVreg2;
580+ RegTensor<float> expFloatVreg3;
581+ RegTensor<float> expSumVreg;
582+ RegTensor<float> maxVreg;
583+ 
584+ RegTensor<float> expDstFloatVreg0;
585+ RegTensor<float> expDstFloatVreg1;
586+ RegTensor<float> expDstFloatVreg2;
587+ RegTensor<float> expDstFloatVreg3;
588+ RegTensor<ElementInput> expOutVreg0;
589+ RegTensor<ElementInput> expOutVreg1;
590+ RegTensor<ElementInput> expDstVreg0;
591+ RegTensor<ElementInput> expDstVreg1;
592+ RegTensor<ElementInput> expDstVreg2;
593+ RegTensor<ElementInput> expDstVreg3;
594+ 
595+ UnalignReg expSumUreg;
596+ 
597+ MaskReg pregFull = CreateMask<ElementInput, MaskPattern::ALL>();
598+ MaskReg pregFloatFull = CreateMask<float, MaskPattern::ALL>();
599+ MaskReg pregTailN = UpdateMask<ElementInput>(tailN);
600+ MaskReg pregtailNOdd= UpdateMask<float>(tailNOdd);
601+ MaskReg pregtailNEven = UpdateMask<float>(tailNEven);
602+ for (uint16_t i = 0; i < m; ++i) {
603+ LoadAlign<float, LoadDist::DIST_BRC_B32>(maxVreg, nowMaxUb + i);
604+ Duplicate(expSumVreg, 0);
605+ LoadAlign(expVreg0, srcUb + i * S2BaseSize);
606+ LoadAlign(expVreg1, srcUb + i * S2BaseSize + HALF_REP_SIZE);
607+ Cast<float, ElementInput, castTraitZero>(expFloatVreg0, expVreg0, pregFull);
608+ Cast<float, ElementInput, castTraitOne>(expFloatVreg1, expVreg0, pregFull);
609+ Cast<float, ElementInput, castTraitZero>(expFloatVreg2, expVreg1, pregFull);
610+ Cast<float, ElementInput, castTraitOne>(expFloatVreg3, expVreg1, pregFull);
611+ FusedExpSub(expDstFloatVreg0, expFloatVreg0, maxVreg, pregFloatFull);
612+ FusedExpSub(expDstFloatVreg1, expFloatVreg1, maxVreg, pregFloatFull);
613+ FusedExpSub(expDstFloatVreg2, expFloatVreg2, maxVreg, pregtailNEven);
614+ FusedExpSub(expDstFloatVreg3, expFloatVreg3, maxVreg, pregtailNOdd);
615+ Add<float, MaskMergeMode::MERGING>(expSumVreg, expSumVreg, expDstFloatVreg0, pregFloatFull);
616+ Add<float, MaskMergeMode::MERGING>(expSumVreg, expSumVreg, expDstFloatVreg1, pregFloatFull);
617+ Add<float, MaskMergeMode::MERGING>(expSumVreg, expSumVreg, expDstFloatVreg2, pregtailNEven);
618+ Add<float, MaskMergeMode::MERGING>(expSumVreg, expSumVreg, expDstFloatVreg3, pregtailNOdd);
619+ Cast<ElementInput, float, castTraitZeroDown>(expDstVreg0, expDstFloatVreg0, pregFloatFull);
620+ Cast<ElementInput, float, castTraitOneDown>(expDstVreg1, expDstFloatVreg1, pregFloatFull);
621+ Cast<ElementInput, float, castTraitZeroDown>(expDstVreg2, expDstFloatVreg2, pregFloatFull);
622+ Cast<ElementInput, float, castTraitOneDown>(expDstVreg3, expDstFloatVreg3, pregFloatFull);
623+ Or((RegTensor<uint16_t>&)expOutVreg0,
624+ (RegTensor<uint16_t>&)expDstVreg0, (RegTensor<uint16_t>&)expDstVreg1,
625+ pregFull);
626+ Or((RegTensor<uint16_t>&)expOutVreg1,
627+ (RegTensor<uint16_t>&)expDstVreg2, (RegTensor<uint16_t>&)expDstVreg3,
628+ pregFull);
629+ StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(
630+ expUb + i * ELE_NUM_PER_C0,
631+ expOutVreg0, blockStride, pregFull);
632+ StoreAlign<ElementOutput, DataCopyMode::DATA_BLOCK_COPY>(
633+ expUb + i * ELE_NUM_PER_C0 + blockStride * ELE_NUM_PER_C0 * BLOCK_REP_SIZE,
634+ expOutVreg1, blockStride, pregTailN);
635+ 
636+ ReduceSum(expSumVreg, expSumVreg, pregFull);
637+ StoreUnAlign<float, PostLiteral::POST_MODE_UPDATE>(expSumUb, expSumVreg, expSumUreg, 1);
638+ }
639+ vstas(expSumUreg, expSumUb, 0, POST_UPDATE);
640+ }
641+ 
642+ __simd_vf__ inline void UpdateExpSumAndExpMax(__ubuf__ float *sumUb, __ubuf__ float *expMaxUb,
643+ __ubuf__ float *maxUb, __ubuf__ float *expSumUb, __ubuf__ float *nowMaxUb,
644+ uint16_t mLoops, uint32_t tailM)
645+ {
646+ using namespace AscendC::MicroAPI;
647+ 
648+ RegTensor<float> nowMaxFloatVreg;
649+ RegTensor<float> lastMaxVreg;
650+ RegTensor<float> expMaxVreg;
651+ RegTensor<float> lastExpSumVreg;
652+ RegTensor<float> brcExpSumFloatVreg;
653+ RegTensor<float> updateExpSumVreg;
654+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
655+ MaskReg pregTailM = UpdateMask<float>(tailM);
656+ for (int16_t i = 0; i < mLoops; ++i) {
657+ LoadAlign(lastMaxVreg, maxUb + i * FLOAT_REP_SIZE);
658+ LoadAlign(nowMaxFloatVreg, nowMaxUb + i * FLOAT_REP_SIZE);
659+ FusedExpSub(expMaxVreg, lastMaxVreg, nowMaxFloatVreg, pregFull);
660+ StoreAlign<float, StoreDist::DIST_NORM_B32>(expMaxUb + i * FLOAT_REP_SIZE, expMaxVreg, pregFull);
661+ StoreAlign<float, StoreDist::DIST_NORM_B32>(maxUb + i * FLOAT_REP_SIZE, nowMaxFloatVreg, pregFull);
662+ 
663+ LoadAlign(lastExpSumVreg, sumUb + i * FLOAT_REP_SIZE);
664+ LoadAlign(brcExpSumFloatVreg, expSumUb + i * FLOAT_REP_SIZE);
665+ Mul(updateExpSumVreg, expMaxVreg, lastExpSumVreg, pregFull);
666+ Add(updateExpSumVreg, updateExpSumVreg, brcExpSumFloatVreg, pregFull);
667+ StoreAlign<float, StoreDist::DIST_NORM_B32>(sumUb + i * FLOAT_REP_SIZE, updateExpSumVreg, pregFull);
668+ }
669+ LoadAlign(lastMaxVreg, maxUb + mLoops * FLOAT_REP_SIZE);
670+ LoadAlign(nowMaxFloatVreg, nowMaxUb + mLoops * FLOAT_REP_SIZE);
671+ FusedExpSub(expMaxVreg, lastMaxVreg, nowMaxFloatVreg, pregTailM);
672+ StoreAlign<float, StoreDist::DIST_NORM_B32>(expMaxUb + mLoops * FLOAT_REP_SIZE, expMaxVreg, pregTailM);
673+ StoreAlign<float, StoreDist::DIST_NORM_B32>(maxUb + mLoops * FLOAT_REP_SIZE, nowMaxFloatVreg, pregTailM);
674+ 
675+ LoadAlign(lastExpSumVreg, sumUb + mLoops * FLOAT_REP_SIZE);
676+ LoadAlign(brcExpSumFloatVreg, expSumUb + mLoops * FLOAT_REP_SIZE);
677+ Mul(updateExpSumVreg, expMaxVreg, lastExpSumVreg, pregTailM);
678+ Add(updateExpSumVreg, updateExpSumVreg, brcExpSumFloatVreg, pregTailM);
679+ StoreAlign<float, StoreDist::DIST_NORM_B32>(sumUb + mLoops * FLOAT_REP_SIZE, updateExpSumVreg, pregTailM);
680+ }
681+};
682+}
683+ 
684+#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_ARCH35_REG_LOW_PREC_BF16_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue_online_softmax_low_prec.hpp+983-0
@@ -0,0 +1,983 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_LOW_PREC_HPP
12+#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_LOW_PREC_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/cross_core_sync.hpp"
16+#include "../../../attn_infra/arch/resource.hpp"
17+#include "../../../attn_infra/epilogue/dispatch_policy.hpp"
18+#include "../../../attn_infra/epilogue/tile_common/tile_copy.hpp"
19+#include "../../../attn_infra/gemm_coord.hpp"
20+#include "../../../attn_infra/matrix_coord.hpp"
21+ 
22+namespace NpuArch::Epilogue::Block {
23+ 
24+template <
25+ class OutputType_,
26+ class InputType_,
27+ class MaskType_,
28+ LseMode LSE_MODE_>
29+class BlockEpilogue<
30+ EpilogueAtlasA2OnlineSoftmax<LSE_MODE_, half>,
31+ OutputType_,
32+ InputType_,
33+ MaskType_>
34+{
35+public:
36+ using DispatchPolicy = EpilogueAtlasA2OnlineSoftmax<LSE_MODE_, half>;
37+ using ArchTag = typename DispatchPolicy::ArchTag;
38+ using ElementOutput = typename OutputType_::Element;
39+ using ElementInput = typename InputType_::Element;
40+ using ElementMask = typename MaskType_::Element;
41+ 
42+ using LayoutOutput = typename OutputType_::Layout;
43+ using LayoutInput = typename InputType_::Layout;
44+ using LayoutMask = typename MaskType_::Layout;
45+ 
46+ static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE;
47+ 
48+ static constexpr uint32_t BLOCK_SIZE_IN_BYTE = 32;
49+ static constexpr uint32_t REPEAT_SIZE_IN_BYTE = 256;
50+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
51+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
52+ static constexpr uint32_t HALF_VECTOR_SIZE = 128;
53+ static constexpr uint32_t BLOCK_SIZE = 16;
54+ static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024;
55+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384;
56+ static constexpr uint32_t VECTOR_SIZE = 128;
57+ static constexpr uint32_t MAX_UB_S_ELEM_NUM = 16384;
58+ 
59+ static constexpr uint32_t REDUCE_UB_SIZE = 1024;
60+ static constexpr uint32_t ROW_OPS_SPEC_MASK_32 = 32;
61+ static constexpr uint32_t ROW_OPS_SPEC_MASK_8 = 8;
62+ static constexpr uint32_t ROW_OPS_SPEC_MASK_4 = 4;
63+ static constexpr uint32_t ROW_OPS_SPEC_MASK_2 = 2;
64+ static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256;
65+ static constexpr int64_t UB_FLOAT_LINE_SIZE = 64;
66+ 
67+ static constexpr uint32_t SPLIT_COL_IDX_2 = 2;
68+ static constexpr uint32_t SPLIT_COL_IDX_3 = 3;
69+ __aicore__ inline
70+ BlockEpilogue(Arch::Resource<ArchTag> &resource, float scaleValue_)
71+ {
72+ // Allocate UB space
73+ constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;
74+ constexpr uint32_t COMPUTE_UB_TENSOR_OFFSET = 2 * UB_UINT8_BLOCK_SIZE;
75+ constexpr uint32_t LP_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
76+ constexpr uint32_t MASK16_UB_TENSOR_OFFSET = 0;
77+ 
78+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE;
79+ constexpr uint32_t LM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 8 * UB_UINT8_VECTOR_SIZE;
80+ 
81+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE;
82+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE;
83+ constexpr uint32_t LL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 11 * UB_UINT8_VECTOR_SIZE;
84+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
85+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE;
86+ 
87+ constexpr uint32_t MASK_UB_TENSOR_OFFSET = 11 * UB_UINT8_BLOCK_SIZE;
88+ 
89+ scaleValue = static_cast<half>(scaleValue_);
90+ lsUbTensor = resource.ubBuf.template GetBufferByByte<half>(LS_UB_TENSOR_OFFSET);
91+ computeUbTensor = resource.ubBuf.template GetBufferByByte<half>(COMPUTE_UB_TENSOR_OFFSET);
92+ lpUbTensor = resource.ubBuf.template GetBufferByByte<ElementOutput>(LP_UB_TENSOR_OFFSET);
93+ maskUbTensor = resource.ubBuf.template GetBufferByByte<ElementMask>(MASK_UB_TENSOR_OFFSET);
94+ maskUbTensor16 = resource.ubBuf.template GetBufferByByte<half>(MASK16_UB_TENSOR_OFFSET);
95+ lmUbTensor = resource.ubBuf.template GetBufferByByte<half>(LM_UB_TENSOR_OFFSET);
96+ hmUbTensor = resource.ubBuf.template GetBufferByByte<half>(HM_UB_TENSOR_OFFSET);
97+ gmUbTensor = resource.ubBuf.template GetBufferByByte<half>(GM_UB_TENSOR_OFFSET);
98+ dmUbTensor = resource.ubBuf.template GetBufferByByte<half>(DM_UB_TENSOR_OFFSET);
99+ llUbTensor = resource.ubBuf.template GetBufferByByte<half>(LL_UB_TENSOR_OFFSET);
100+ tvUbTensor = resource.ubBuf.template GetBufferByByte<half>(TV_UB_TENSOR_OFFSET);
101+ glUbTensor = resource.ubBuf.template GetBufferByByte<half>(GL_UB_TENSOR_OFFSET);
102+ }
103+ 
104+ __aicore__ inline
105+ ~BlockEpilogue() {}
106+ 
107+ __aicore__ inline
108+ void SetVecMask(int32_t len)
109+ {
110+ const int32_t MAX_MASK_LEN = 128;
111+ const int32_t HALF_MASK_LEN = 64;
112+ if (len >= MAX_MASK_LEN) {
113+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
114+ return;
115+ }
116+ int32_t highMask = len - HALF_MASK_LEN > 0 ? len - HALF_MASK_LEN : 0;
117+ int32_t lowMask = len - HALF_MASK_LEN >= 0 ? HALF_MASK_LEN : len;
118+ if (len < HALF_MASK_LEN) {
119+ AscendC::SetVectorMask<int8_t>(0x0, ((uint64_t)1 << lowMask) - 1);
120+ } else {
121+ AscendC::SetVectorMask<int8_t>(((uint64_t)1 << highMask) - 1, 0xffffffffffffffff);
122+ }
123+ }
124+ 
125+ __aicore__ inline
126+ void SetBlockReduceMask(int32_t len)
127+ {
128+ const int32_t MAX_LEN = 16;
129+ if (len > MAX_LEN) {
130+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
131+ return;
132+ }
133+ uint64_t subMask = (static_cast<uint64_t>(1) << len) - 1;
134+ uint64_t maskValue = (subMask << 48) + (subMask << 32) + (subMask << 16) + subMask;
135+ AscendC::SetVectorMask<int8_t>(maskValue, maskValue);
136+ }
137+ 
138+ __aicore__ inline
139+ void ReduceSumByPair(const AscendC::LocalTensor<half> &srcUb, uint32_t numRowsRound, uint32_t loopCount,
140+ uint32_t columnStrideIndex, uint8_t dataBlockStride, uint8_t repeatStride)
141+ {
142+ for (uint32_t i = 0; i < loopCount; i += columnStrideIndex) {
143+ uint32_t src0Start = i * HALF_VECTOR_SIZE;
144+ uint32_t src1Start = (i + columnStrideIndex / 2) * HALF_VECTOR_SIZE;
145+ AscendC::Add<half, false>(
146+ srcUb[src0Start],
147+ srcUb[src0Start],
148+ srcUb[src1Start],
149+ AscendC::MASK_PLACEHOLDER, // (uint64_t)0
150+ numRowsRound,
151+ AscendC::BinaryRepeatParams(
152+ dataBlockStride,
153+ dataBlockStride,
154+ dataBlockStride,
155+ repeatStride,
156+ repeatStride,
157+ repeatStride));
158+ }
159+ }
160+ 
161+ __aicore__ inline
162+ void RowsumTileFlexible(const AscendC::LocalTensor<half> &srcUb, const AscendC::LocalTensor<half> &rowsumUb,
163+ uint32_t numRowsRound, uint32_t numElemsAligned)
164+ {
165+ // Vector计算单元每个迭代最多处理256Byte数据，因此half低精度场景，每次迭代最多处理256/2=128个元素
166+ uint32_t loopCount = numElemsAligned / HALF_VECTOR_SIZE; // half低精度场景，每行需要numElemsAligned/128次循环处理
167+ // 每个datablock长度32Byte，因此half低精度场景，每个datablock内有32/2=16个元素
168+ uint8_t dataBlockNumPerRow = numElemsAligned / BLOCK_SIZE; // half低精度场景，每行共有numElemsAligned/16个datablock
169+ uint8_t dataBlockStride = 1;
170+ 
171+ // 举例，若numElemsAligned为1024，以128为单位分治求和，1024->512->256->128
172+ for (uint32_t columnStrideIndex = 2; columnStrideIndex <= loopCount; columnStrideIndex *= 2) {
173+ ReduceSumByPair(srcUb, numRowsRound, loopCount, columnStrideIndex, dataBlockStride, dataBlockNumPerRow);
174+ AscendC::PipeBarrier<PIPE_V>();
175+ }
176+ 
177+ //每行128个元素分别规约求和
178+ AscendC::WholeReduceSum<half, false>(
179+ rowsumUb,
180+ srcUb,
181+ AscendC::MASK_PLACEHOLDER, // (uint64_t)0
182+ numRowsRound,
183+ dataBlockStride,
184+ dataBlockStride,
185+ dataBlockNumPerRow);
186+ AscendC::PipeBarrier<PIPE_V>();
187+ }
188+
189+ __aicore__ inline
190+ void RowsumSPECTILE1024(const AscendC::LocalTensor<half> &srcUb, const AscendC::LocalTensor<half> &rowsumUb,
191+ const AscendC::LocalTensor<half> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
192+ uint32_t numElemsAligned)
193+ {
194+ RowsumTileFlexible(srcUb, rowsumUb, numRowsRound, numElemsAligned);
195+ }
196+ 
197+ __aicore__ inline
198+ void RowsumSPECTILE512(const AscendC::LocalTensor<half> &srcUb, const AscendC::LocalTensor<half> &rowsumUb,
199+ const AscendC::LocalTensor<half> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
200+ uint32_t numElemsAligned)
201+ {
202+ RowsumTileFlexible(srcUb, rowsumUb, numRowsRound, numElemsAligned);
203+ }
204+ 
205+ __aicore__ inline
206+ void RowsumTAILTILE(const AscendC::LocalTensor<half> &srcUb, const AscendC::LocalTensor<half> &rowsumUb,
207+ const AscendC::LocalTensor<half> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
208+ uint32_t numElemsAligned)
209+ {
210+ if (numElems <= HALF_VECTOR_SIZE) {
211+ SetVecMask(numElems);
212+ AscendC::WholeReduceSum<half, false>(
213+ rowsumUb, srcUb, (int32_t)0, numRowsRound, 1, 1,
214+ numElemsAligned / BLOCK_SIZE);
215+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
216+ } else {
217+ for (uint32_t vmaxIdx = 1; vmaxIdx < numElems / HALF_VECTOR_SIZE; vmaxIdx++) {
218+ AscendC::Add<half, false>(
219+ srcUb,
220+ srcUb,
221+ srcUb[vmaxIdx * HALF_VECTOR_SIZE],
222+ (uint64_t)0,
223+ numRowsRound,
224+ AscendC::BinaryRepeatParams(
225+ 1, 1, 1,
226+ numElemsAligned / BLOCK_SIZE,
227+ numElemsAligned / BLOCK_SIZE,
228+ numElemsAligned / BLOCK_SIZE));
229+ }
230+ AscendC::PipeBarrier<PIPE_V>();
231+ if (numElems % HALF_VECTOR_SIZE > 0) {
232+ SetVecMask(numElems % HALF_VECTOR_SIZE);
233+ AscendC::Add<half, false>(
234+ srcUb,
235+ srcUb,
236+ srcUb[numElems / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE],
237+ (uint64_t)0,
238+ numRowsRound,
239+ AscendC::BinaryRepeatParams(
240+ 1, 1, 1,
241+ numElemsAligned / BLOCK_SIZE,
242+ numElemsAligned / BLOCK_SIZE,
243+ numElemsAligned / BLOCK_SIZE));
244+ AscendC::PipeBarrier<PIPE_V>();
245+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
246+ }
247+ AscendC::WholeReduceSum<half, false>(
248+ rowsumUb, srcUb, (int32_t)0, numRowsRound, 1, 1,
249+ numElemsAligned / BLOCK_SIZE);
250+ }
251+ AscendC::PipeBarrier<PIPE_V>();
252+ }
253+ 
254+ __aicore__ inline
255+ void ReduceMaxByPair(const AscendC::LocalTensor<half> &dstUb, const AscendC::LocalTensor<half> &srcUb,
256+ uint32_t numRowsRound, uint32_t loopCount, uint32_t columnStrideIndex,
257+ uint8_t dataBlockStride, uint8_t repeatStride)
258+ {
259+ for (uint32_t i = 0; i < loopCount; i += columnStrideIndex) {
260+ uint32_t src0Start = i * HALF_VECTOR_SIZE;
261+ uint32_t src1Start = (i + columnStrideIndex / 2) * HALF_VECTOR_SIZE;
262+ AscendC::Max<half, false>(
263+ dstUb[src0Start],
264+ srcUb[src0Start],
265+ srcUb[src1Start],
266+ AscendC::MASK_PLACEHOLDER, // (uint64_t)0
267+ numRowsRound,
268+ AscendC::BinaryRepeatParams(
269+ dataBlockStride,
270+ dataBlockStride,
271+ dataBlockStride,
272+ repeatStride,
273+ repeatStride,
274+ repeatStride));
275+ }
276+ }
277+ 
278+ __aicore__ inline
279+ void RowmaxTileFlexible(const AscendC::LocalTensor<half> &srcUb, const AscendC::LocalTensor<half> &rowmaxUb,
280+ uint32_t numRowsRound, uint32_t numElemsAligned)
281+ {
282+ // Vector计算单元每个迭代最多处理256Byte数据，因此half低精度场景，每次迭代最多处理256/2=128个元素
283+ uint32_t loopCount = numElemsAligned / HALF_VECTOR_SIZE; // half低精度场景，每行需要numElemsAligned/128次循环处理
284+ // 每个datablock长度32Byte，因此half低精度场景，每个datablock内有32/2=16个元素
285+ uint8_t dataBlockNumPerRow = numElemsAligned / BLOCK_SIZE; // half低精度场景，每行共有numElemsAligned/16个datablock
286+ uint8_t dataBlockStride = 1;
287+ 
288+ // 举例，若numElemsAligned为1024，以128为单位分治求最大值，1024->512->256->128
289+ uint32_t columnStrideIndex = 2;
290+ // 后续Rowsum计算还会使用到srcUb，因此第一轮分治使用lsUbTensor作为目的操作数，srcUb作为源操作数
291+ ReduceMaxByPair(lsUbTensor, srcUb, numRowsRound, loopCount, columnStrideIndex, dataBlockStride, dataBlockNumPerRow);
292+ AscendC::PipeBarrier<PIPE_V>();
293+ columnStrideIndex *= 2;
294+ for (; columnStrideIndex <= loopCount; columnStrideIndex *= 2) {
295+ ReduceMaxByPair(lsUbTensor, lsUbTensor, numRowsRound, loopCount, columnStrideIndex, dataBlockStride, dataBlockNumPerRow);
296+ AscendC::PipeBarrier<PIPE_V>();
297+ }
298+ 
299+ //每行128个元素分别规约求最大值
300+ AscendC::WholeReduceMax<half, false>(
301+ rowmaxUb,
302+ lsUbTensor,
303+ AscendC::MASK_PLACEHOLDER, // (uint64_t)0
304+ numRowsRound,
305+ dataBlockStride,
306+ dataBlockStride,
307+ dataBlockNumPerRow,
308+ AscendC::ReduceOrder::ORDER_ONLY_VALUE);
309+ AscendC::PipeBarrier<PIPE_V>();
310+ }
311+ 
312+ __aicore__ inline
313+ void RowmaxSPECTILE1024(const AscendC::LocalTensor<half> &srcUb, const AscendC::LocalTensor<half> &rowmaxUb,
314+ const AscendC::LocalTensor<half> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
315+ uint32_t numElemsAligned)
316+ {
317+ RowmaxTileFlexible(srcUb, rowmaxUb, numRowsRound, numElemsAligned);
318+ }
319+ 
320+ __aicore__ inline
321+ void RowmaxSPECTILE512(const AscendC::LocalTensor<half> &srcUb, const AscendC::LocalTensor<half> &rowmaxUb,
322+ const AscendC::LocalTensor<half> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
323+ uint32_t numElemsAligned)
324+ {
325+ RowmaxTileFlexible(srcUb, rowmaxUb, numRowsRound, numElemsAligned);
326+ }
327+ 
328+ __aicore__ inline
329+ void RowmaxTAILTILE(const AscendC::LocalTensor<half> &srcUb, const AscendC::LocalTensor<half> &rowmaxUb,
330+ const AscendC::LocalTensor<half> &tvUbTensor, uint32_t numRowsRound, uint32_t numElems,
331+ uint32_t numElemsAligned)
332+ {
333+ if (numElems <= HALF_VECTOR_SIZE) {
334+ SetVecMask(numElems);
335+ AscendC::WholeReduceMax<half, false>(
336+ rowmaxUb, srcUb, (int32_t)0, numRowsRound, 1, 1,
337+ numElemsAligned / BLOCK_SIZE, AscendC::ReduceOrder::ORDER_ONLY_VALUE);
338+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
339+ } else {
340+ AscendC::DataCopy(
341+ lsUbTensor,
342+ srcUb,
343+ AscendC::DataCopyParams(
344+ numRowsRound,
345+ HALF_VECTOR_SIZE / BLOCK_SIZE,
346+ (numElemsAligned - HALF_VECTOR_SIZE) / BLOCK_SIZE,
347+ (numElemsAligned - HALF_VECTOR_SIZE) / BLOCK_SIZE));
348+ AscendC::PipeBarrier<PIPE_V>();
349+ for (uint32_t vmaxIdx = 1; vmaxIdx < numElems / HALF_VECTOR_SIZE; vmaxIdx++) {
350+ AscendC::Max<half, false>(
351+ lsUbTensor,
352+ lsUbTensor,
353+ srcUb[vmaxIdx * HALF_VECTOR_SIZE],
354+ (uint64_t)0,
355+ numRowsRound,
356+ AscendC::BinaryRepeatParams(
357+ 1, 1, 1,
358+ numElemsAligned / BLOCK_SIZE,
359+ numElemsAligned / BLOCK_SIZE,
360+ numElemsAligned / BLOCK_SIZE));
361+ }
362+ AscendC::PipeBarrier<PIPE_V>();
363+ if (numElems % HALF_VECTOR_SIZE > 0) {
364+ SetVecMask(numElems % HALF_VECTOR_SIZE);
365+ AscendC::Max<half, false>(
366+ lsUbTensor,
367+ lsUbTensor,
368+ srcUb[numElems / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE],
369+ (uint64_t)0,
370+ numRowsRound,
371+ AscendC::BinaryRepeatParams(
372+ 1, 1, 1,
373+ numElemsAligned / BLOCK_SIZE,
374+ numElemsAligned / BLOCK_SIZE,
375+ numElemsAligned / BLOCK_SIZE));
376+ AscendC::PipeBarrier<PIPE_V>();
377+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
378+ }
379+ AscendC::WholeReduceMax<half, false>(
380+ rowmaxUb, lsUbTensor, (int32_t)0, numRowsRound, 1, 1,
381+ numElemsAligned / BLOCK_SIZE, AscendC::ReduceOrder::ORDER_ONLY_VALUE);
382+ }
383+ AscendC::PipeBarrier<PIPE_V>();
384+ }
385+ 
386+ __aicore__ inline
387+ void CopySGmToUb(AscendC::GlobalTensor<half> gInput, uint32_t sUbOffset, uint32_t rowNumCurLoop,
388+ uint32_t columnNumRound, uint32_t columnNumPad)
389+ {
390+ // input S
391+ AscendC::DataCopy(
392+ lsUbTensor,
393+ gInput,
394+ AscendC::DataCopyParams(rowNumCurLoop,
395+ columnNumRound / BLOCK_SIZE,
396+ (columnNumPad - columnNumRound) / BLOCK_SIZE,
397+ 0));
398+ }
399+ 
400+ __aicore__ inline
401+ void CopyMaskGmToUb(AscendC::GlobalTensor<ElementMask> gMask, uint32_t columnNum, uint32_t columnNumRound,
402+ uint32_t maskStride, uint32_t tokenNumPerHead, uint32_t proTokenIdx, uint32_t proTokenNum,
403+ uint32_t integralHeadNum, uint32_t epiTokenNum)
404+ {
405+ uint32_t innerUbRowOffset = 0;
406+ if (proTokenNum != 0U) {
407+ AscendC::DataCopyPad(
408+ maskUbTensor[innerUbRowOffset],
409+ gMask[proTokenIdx * maskStride],
410+ AscendC::DataCopyExtParams(
411+ proTokenNum, columnNum * sizeof(ElementMask),
412+ (maskStride - columnNum) * sizeof(ElementMask), 0, 0),
413+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
414+ innerUbRowOffset += proTokenNum * columnNumRound;
415+ }
416+ for (uint32_t headIdx = 0; headIdx < integralHeadNum; headIdx++) {
417+ AscendC::DataCopyPad(
418+ maskUbTensor[innerUbRowOffset],
419+ gMask,
420+ AscendC::DataCopyExtParams(
421+ tokenNumPerHead, columnNum * sizeof(ElementMask),
422+ (maskStride - columnNum) * sizeof(ElementMask), 0, 0),
423+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
424+ innerUbRowOffset += tokenNumPerHead * columnNumRound;
425+ }
426+ if (epiTokenNum != 0) {
427+ AscendC::DataCopyPad(
428+ maskUbTensor[innerUbRowOffset],
429+ gMask,
430+ AscendC::DataCopyExtParams(
431+ epiTokenNum, columnNum * sizeof(ElementMask),
432+ (maskStride - columnNum) * sizeof(ElementMask), 0, 0),
433+ AscendC::DataCopyPadExtParams<ElementMask>(false, 0, 0, 0));
434+ }
435+ }
436+ 
437+ __aicore__ inline
438+ void ScaleS(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound)
439+ {
440+ // *** ls = scaleValue * ls
441+ AscendC::Muls<half, false>(
442+ computeUbTensor,
443+ lsUbTensor,
444+ scaleValue,
445+ (uint64_t)0,
446+ (rowNumCurLoop * columnNumRound + HALF_VECTOR_SIZE - 1) / HALF_VECTOR_SIZE,
447+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
448+ AscendC::PipeBarrier<PIPE_V>();
449+ }
450+ 
451+ template<typename ElementMaskDst, typename ElementMaskSrc>
452+ __aicore__ inline
453+ void UpCastMask(
454+ const AscendC::LocalTensor<ElementMaskDst> &maskUbTensorDst,
455+ const AscendC::LocalTensor<ElementMaskSrc> &maskUbTensorSrc,
456+ uint32_t rowNumCurLoop,
457+ uint32_t columnNumRound)
458+ {
459+ AscendC::Cast<ElementMaskDst, ElementMaskSrc, false>(
460+ maskUbTensorDst, maskUbTensorSrc, AscendC::RoundMode::CAST_NONE, (uint64_t)0,
461+ CeilDiv(rowNumCurLoop * columnNumRound, (uint32_t)(REPEAT_SIZE_IN_BYTE / sizeof(ElementMaskDst))),
462+ AscendC::UnaryRepeatParams(1, 1, 8, 4));
463+ AscendC::PipeBarrier<PIPE_V>();
464+ }
465+ 
466+ __aicore__ inline
467+ void ApplyMask(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound, uint32_t maskColumnRound,
468+ uint32_t addMaskUbOffset)
469+ {
470+ AscendC::Muls<half, false>(
471+ maskUbTensor16,
472+ maskUbTensor16,
473+ (half)-6e4, // -65504
474+ (uint64_t)0,
475+ (rowNumCurLoop * maskColumnRound + HALF_VECTOR_SIZE - 1) / HALF_VECTOR_SIZE,
476+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
477+ AscendC::PipeBarrier<PIPE_V>();
478+ if (maskColumnRound == columnNumRound) {
479+ AscendC::Add<half, false>(
480+ computeUbTensor,
481+ computeUbTensor,
482+ maskUbTensor16,
483+ (uint64_t)0,
484+ (rowNumCurLoop * maskColumnRound + HALF_VECTOR_SIZE - 1) / HALF_VECTOR_SIZE,
485+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
486+ } else {
487+ uint32_t loop = maskColumnRound / HALF_VECTOR_SIZE;
488+ for (uint32_t i = 0; i < loop; i++) {
489+ AscendC::Add<half, false>(
490+ computeUbTensor[addMaskUbOffset + i * HALF_VECTOR_SIZE],
491+ computeUbTensor[addMaskUbOffset + i * HALF_VECTOR_SIZE],
492+ maskUbTensor16[i * HALF_VECTOR_SIZE],
493+ (uint64_t)0,
494+ rowNumCurLoop,
495+ AscendC::BinaryRepeatParams(1,
496+ 1,
497+ 1,
498+ columnNumRound / BLOCK_SIZE,
499+ columnNumRound / BLOCK_SIZE,
500+ maskColumnRound / BLOCK_SIZE));
501+ }
502+ if (maskColumnRound % HALF_VECTOR_SIZE > 0) {
503+ SetVecMask(maskColumnRound % HALF_VECTOR_SIZE);
504+ AscendC::Add<half, false>(
505+ computeUbTensor[addMaskUbOffset + loop * HALF_VECTOR_SIZE],
506+ computeUbTensor[addMaskUbOffset + loop * HALF_VECTOR_SIZE],
507+ maskUbTensor16[loop * HALF_VECTOR_SIZE],
508+ (uint64_t)0,
509+ rowNumCurLoop,
510+ AscendC::BinaryRepeatParams(1,
511+ 1,
512+ 1,
513+ columnNumRound / BLOCK_SIZE,
514+ columnNumRound / BLOCK_SIZE,
515+ maskColumnRound / BLOCK_SIZE));
516+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
517+ }
518+ }
519+ AscendC::PipeBarrier<PIPE_V>();
520+ }
521+ 
522+ __aicore__ inline
523+ void CalcLocalRowMax(uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
524+ uint32_t rowOffset)
525+ {
526+ if (columnNum == 1024U) {
527+ RowmaxSPECTILE1024(
528+ computeUbTensor,
529+ lmUbTensor[rowOffset],
530+ tvUbTensor,
531+ rowNumCurLoopRound,
532+ columnNum,
533+ columnNumRound);
534+ } else if (columnNum == 512U) {
535+ RowmaxSPECTILE512(
536+ computeUbTensor,
537+ lmUbTensor[rowOffset],
538+ tvUbTensor,
539+ rowNumCurLoopRound,
540+ columnNum,
541+ columnNumRound);
542+ } else {
543+ RowmaxTAILTILE(
544+ computeUbTensor,
545+ lmUbTensor[rowOffset],
546+ tvUbTensor,
547+ rowNumCurLoopRound,
548+ columnNum,
549+ columnNumRound);
550+ }
551+ }
552+ 
553+ __aicore__ inline
554+ void UpdateGlobalRowMax(uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum,
555+ uint32_t columnNumRound, uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile)
556+ {
557+ if (isFirstStackTile) {
558+ AscendC::DataCopy(
559+ hmUbTensor[rowOffset],
560+ lmUbTensor[rowOffset],
561+ AscendC::DataCopyParams(1, rowNumCurLoopRound / BLOCK_SIZE, 0, 0));
562+ AscendC::PipeBarrier<PIPE_V>();
563+ } else {
564+ SetVecMask(rowNumCurLoop);
565+ // *** hm = vmax(lm, gm)
566+ AscendC::Max<half, false>(
567+ hmUbTensor[rowOffset],
568+ lmUbTensor[rowOffset],
569+ gmUbTensor[rowOffset],
570+ (uint64_t)0,
571+ 1,
572+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
573+ 
574+ AscendC::PipeBarrier<PIPE_V>();
575+ // *** dm = gm - hm
576+ AscendC::Sub<half, false>(
577+ dmUbTensor[dmUbOffsetCurCycle],
578+ gmUbTensor[rowOffset],
579+ hmUbTensor[rowOffset],
580+ (uint64_t)0,
581+ 1,
582+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
583+ AscendC::PipeBarrier<PIPE_V>();
584+ // *** dm = exp(dm)
585+ AscendC::Exp<half, false>(dmUbTensor[dmUbOffsetCurCycle],
586+ dmUbTensor[dmUbOffsetCurCycle],
587+ (uint64_t)0,
588+ 1,
589+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
590+ }
591+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
592+ AscendC::PipeBarrier<PIPE_V>();
593+ // *** gm = hm
594+ AscendC::DataCopy(gmUbTensor[rowOffset],
595+ hmUbTensor[rowOffset],
596+ AscendC::DataCopyParams(1, rowNumCurLoopRound / BLOCK_SIZE, 0, 0));
597+ AscendC::PipeBarrier<PIPE_V>();
598+ }
599+ 
600+ __aicore__ inline
601+ void CalcExp(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound, uint32_t columnNum,
602+ uint32_t columnNumRound, uint32_t rowOffset)
603+ {
604+ // *** hm_block = expand_to_block(hm), 存放于 tv
605+ AscendC::Brcb(
606+ tvUbTensor.template ReinterpretCast<uint16_t>(),
607+ hmUbTensor[rowOffset].template ReinterpretCast<uint16_t>(),
608+ rowNumCurLoopRound / FLOAT_BLOCK_SIZE,
609+ AscendC::BrcbRepeatParams(1, 8));
610+ AscendC::PipeBarrier<PIPE_V>();
611+ // *** ls = ls - hm_block
612+ for (uint32_t subIdx = 0; subIdx < columnNum / HALF_VECTOR_SIZE; ++subIdx) {
613+ AscendC::Sub<half, false>(
614+ computeUbTensor[subIdx * HALF_VECTOR_SIZE],
615+ computeUbTensor[subIdx * HALF_VECTOR_SIZE],
616+ tvUbTensor,
617+ (uint64_t)0,
618+ rowNumCurLoop,
619+ AscendC::BinaryRepeatParams(
620+ 1, 1, 0, columnNumRound / BLOCK_SIZE, columnNumRound / BLOCK_SIZE, 1));
621+ }
622+ if (columnNum % HALF_VECTOR_SIZE > 0) {
623+ SetVecMask(columnNum % HALF_VECTOR_SIZE);
624+ AscendC::Sub<half, false>(
625+ computeUbTensor[columnNum / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE],
626+ computeUbTensor[columnNum / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE],
627+ tvUbTensor,
628+ (uint64_t)0,
629+ rowNumCurLoop,
630+ AscendC::BinaryRepeatParams(
631+ 1, 1, 0, columnNumRound / BLOCK_SIZE, columnNumRound / BLOCK_SIZE, 1));
632+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
633+ }
634+ AscendC::PipeBarrier<PIPE_V>();
635+ // *** ls = exp(ls)
636+ AscendC::Exp<half, false>(
637+ computeUbTensor,
638+ computeUbTensor,
639+ (uint64_t)0,
640+ (rowNumCurLoop * columnNumRound + HALF_VECTOR_SIZE - 1) / HALF_VECTOR_SIZE,
641+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
642+ AscendC::PipeBarrier<PIPE_V>();
643+ }
644+ 
645+ __aicore__ inline
646+ void CalcLocalRowSum(uint32_t sUbOffset, uint32_t rowNumCurLoopRound, uint32_t columnNum, uint32_t columnNumRound,
647+ uint32_t rowOffset)
648+ {
649+ // *** ll = rowsum(ls32)
650+ if (columnNum == 1024U) {
651+ RowsumSPECTILE1024(
652+ computeUbTensor,
653+ llUbTensor[rowOffset],
654+ tvUbTensor,
655+ rowNumCurLoopRound,
656+ columnNum,
657+ columnNumRound);
658+ } else if (columnNum == 512U) {
659+ RowsumSPECTILE512(
660+ computeUbTensor,
661+ llUbTensor[rowOffset],
662+ tvUbTensor,
663+ rowNumCurLoopRound,
664+ columnNum,
665+ columnNumRound);
666+ } else {
667+ RowsumTAILTILE(
668+ computeUbTensor,
669+ llUbTensor[rowOffset],
670+ tvUbTensor,
671+ rowNumCurLoopRound,
672+ columnNum,
673+ columnNumRound);
674+ }
675+ }
676+ 
677+ __aicore__ inline
678+ void UpdateGlobalRowSum(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t rowNumCurLoopRound,
679+ uint32_t dmUbOffsetCurCycle, uint32_t rowOffset, uint32_t isFirstStackTile)
680+ {
681+ if (isFirstStackTile) {
682+ // *** gl = ll
683+ AscendC::DataCopy(
684+ glUbTensor[rowOffset],
685+ llUbTensor[rowOffset],
686+ AscendC::DataCopyParams(1, rowNumCurLoopRound / BLOCK_SIZE, 0, 0));
687+ AscendC::PipeBarrier<PIPE_V>();
688+ } else {
689+ SetVecMask(rowNumCurLoop);
690+ // *** gl = dm * gl
691+ AscendC::Mul<half, false>(
692+ glUbTensor[rowOffset],
693+ dmUbTensor[dmUbOffsetCurCycle],
694+ glUbTensor[rowOffset],
695+ (uint64_t)0,
696+ 1,
697+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
698+ AscendC::PipeBarrier<PIPE_V>();
699+ // *** gl = ll + gl
700+ AscendC::Add<half, false>(
701+ glUbTensor[rowOffset],
702+ glUbTensor[rowOffset],
703+ llUbTensor[rowOffset],
704+ (uint64_t)0,
705+ 1,
706+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
707+ AscendC::PipeBarrier<PIPE_V>();
708+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
709+ }
710+ }
711+ 
712+ __aicore__ inline
713+ void MoveP(uint32_t sUbOffset, uint32_t rowNumCurLoop, uint32_t columnNumRound)
714+ {
715+ AscendC::DataCopyParams repeatParams;
716+ repeatParams.blockCount = 1;
717+ repeatParams.srcStride = 0;
718+ repeatParams.blockLen = CeilDiv(rowNumCurLoop * columnNumRound, BLOCK_SIZE);
719+ AscendC::DataCopy<half>(lpUbTensor[sUbOffset], computeUbTensor, repeatParams);
720+ AscendC::PipeBarrier<PIPE_V>();
721+ }
722+ 
723+ __aicore__ inline
724+ void CopyPUbToGm(AscendC::GlobalTensor<ElementOutput> gOutput, uint32_t sUbOffset, uint32_t rowNumCurLoop,
725+ uint32_t columnNumRound, uint32_t columnNumPad)
726+ {
727+ AscendC::DataCopy(gOutput,
728+ lpUbTensor[sUbOffset],
729+ AscendC::DataCopyParams(
730+ rowNumCurLoop, columnNumRound / BLOCK_SIZE, 0, (columnNumPad - columnNumRound) / BLOCK_SIZE));
731+ }
732+ 
733+ __aicore__ inline
734+ void SubCoreCompute(
735+ AscendC::GlobalTensor<ElementOutput> gOutput, const LayoutOutput &layoutOutput,
736+ uint32_t rowOffset, uint32_t isFirstStackTile, uint32_t isFirstRowLoop,
737+ uint32_t columnNumRound, uint32_t pingpongFlag,
738+ uint32_t curStackTileMod, Arch::CrossCoreFlag softmaxFlag, uint32_t isLastLoop)
739+ {
740+ uint32_t rowNumCurLoop = layoutOutput.shape(0);
741+ uint32_t rowNumCurLoopRound = RoundUp(rowNumCurLoop, BLOCK_SIZE);
742+ uint32_t columnNum = layoutOutput.shape(1);
743+ uint32_t columnNumPad = layoutOutput.stride(0);
744+ uint32_t sUbOffset = pingpongFlag * MAX_UB_S_ELEM_NUM;
745+ uint32_t dmUbOffsetCurCycle = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffset;
746+ 
747+ if constexpr (LSE_MODE_ == LseMode::OUT_ONLY) {
748+ // wait for lse from ub to gm (low pre)
749+ // In lse out-only mode, tv is used in the last stack tile to transport lse
750+ if (isFirstStackTile && isFirstRowLoop) {
751+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
752+ }
753+ }
754+ CalcLocalRowMax(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
755+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
756+ UpdateGlobalRowMax(rowNumCurLoop,
757+ rowNumCurLoopRound,
758+ columnNum,
759+ columnNumRound,
760+ dmUbOffsetCurCycle,
761+ rowOffset,
762+ isFirstStackTile);
763+ CalcExp(sUbOffset, rowNumCurLoop, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
764+ 
765+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(pingpongFlag);
766+ MoveP(sUbOffset, rowNumCurLoop, columnNumRound);
767+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
768+ 
769+ CalcLocalRowSum(sUbOffset, rowNumCurLoopRound, columnNum, columnNumRound, rowOffset);
770+ 
771+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
772+ CopyPUbToGm(gOutput, sUbOffset, rowNumCurLoop, columnNumRound, columnNumPad);
773+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(pingpongFlag);
774+ if (isLastLoop) {
775+ NpuArch::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxFlag);
776+ }
777+ UpdateGlobalRowSum(
778+ sUbOffset, rowNumCurLoop, rowNumCurLoopRound, dmUbOffsetCurCycle, rowOffset, isFirstStackTile);
779+ }
780+ 
781+ __aicore__ inline
782+ void operator()(AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<half> gInput,
783+ const LayoutOutput &layoutOutput, const LayoutInput &layoutInput, GemmCoord actualBlockShape,
784+ uint32_t isFirstStackTile, uint32_t isLastNoMaskStackTile,
785+ uint32_t qSBlockSize, uint32_t qNBlockSize, uint32_t curStackTileMod, Arch::CrossCoreFlag softmaxFlag)
786+ {
787+ uint32_t rowNum = actualBlockShape.m();
788+ uint32_t columnNum = actualBlockShape.n();
789+ uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE);
790+ uint32_t columnNumPad = layoutInput.stride(0);
791+ 
792+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
793+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
794+ 
795+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
796+ uint32_t qNThisSubBlock = (qNBlockSize == 1U) ?
797+ 0 : (subBlockIdx == 1U) ? (qNBlockSize - qNSplitSubBlock) : qNSplitSubBlock;
798+ uint32_t rowSplitSubBlock = (qNBlockSize == 1U) ? (qSBlockSize / 2U) : (qSBlockSize * qNSplitSubBlock);
799+ uint32_t rowActualThisSubBlock = (subBlockIdx == 1U) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock;
800+ uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock;
801+ uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound;
802+ uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, BLOCK_SIZE);
803+ rowNumTile = AscendC::Std::min(rowNumTile, HALF_VECTOR_SIZE);
804+ uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile);
805+ 
806+ if (rowActualThisSubBlock == 0) {
807+ NpuArch::Arch::CrossCoreSetFlag<0x2, PIPE_MTE3>(softmaxFlag);
808+ return;
809+ }
810+ 
811+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum; rowLoopIdx++) {
812+ uint32_t pingpongFlag = rowLoopIdx % 2U;
813+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
814+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
815+ uint32_t rowNumCurLoop =
816+ (rowLoopIdx == rowLoopNum - 1U) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
817+ 
818+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
819+ auto gInputCurLoop = gInput[offsetInput];
820+ 
821+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
822+ CopySGmToUb(
823+ gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad);
824+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
825+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
826+ ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound);
827+ 
828+ int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
829+ auto gOutputCurLoop = gOutput[offsetOutput];
830+ auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum));
831+ SubCoreCompute(
832+ gOutputCurLoop,
833+ layoutOutputCurLoop,
834+ rowOffsetCurLoop,
835+ isFirstStackTile,
836+ (rowLoopIdx == 0U),
837+ columnNumRound,
838+ pingpongFlag,
839+ curStackTileMod,
840+ softmaxFlag,
841+ (rowLoopIdx == rowLoopNum - 1));
842+ }
843+ }
844+ 
845+ __aicore__ inline
846+ void operator()(AscendC::GlobalTensor<ElementOutput> gOutput, AscendC::GlobalTensor<half> gInput,
847+ AscendC::GlobalTensor<ElementMask> gMask, const LayoutOutput &layoutOutput, const LayoutInput &layoutInput,
848+ const LayoutInput &layoutMask, GemmCoord actualBlockShape, uint32_t isFirstStackTile, uint32_t qSBlockSize,
849+ uint32_t qNBlockSize, uint32_t curStackTileMod, Arch::CrossCoreFlag qkReady, uint32_t triUp, uint32_t triDown,
850+ uint32_t kvSStartIdx, uint32_t kvSEndIdx, Arch::CrossCoreFlag softmaxFlag)
851+ {
852+ uint32_t rowNum = actualBlockShape.m();
853+ uint32_t columnNum = actualBlockShape.n();
854+ uint32_t columnNumRound = RoundUp(columnNum, BLOCK_SIZE);
855+ uint32_t columnNumPad = layoutInput.stride(0);
856+ uint32_t maskStride = layoutMask.stride(0);
857+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
858+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
859+ 
860+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
861+ uint32_t qNThisSubBlock = (qNBlockSize == 1U) ?
862+ 0 : (subBlockIdx == 1U) ? (qNBlockSize - qNSplitSubBlock) : qNSplitSubBlock;
863+ uint32_t rowSplitSubBlock = (qNBlockSize == 1U) ? (qSBlockSize / 2U) : (qSBlockSize * qNSplitSubBlock);
864+ uint32_t rowActualThisSubBlock = (subBlockIdx == 1U) ? (rowNum - rowSplitSubBlock) : rowSplitSubBlock;
865+ uint32_t rowOffsetThisSubBlock = subBlockIdx * rowSplitSubBlock;
866+ 
867+ uint32_t tokenNumPerHeadThisSubBlock = AscendC::Std::min(qSBlockSize, rowActualThisSubBlock);
868+ 
869+ uint32_t maskOffsetThisSubBlock = (qNBlockSize == 1U) ? rowOffsetThisSubBlock : 0;
870+ 
871+ uint32_t gmOffsetMaskRow;
872+ uint32_t gmOffsetMaskColumn;
873+ uint32_t maskColumn;
874+ uint32_t addMaskUbOffset;
875+ if (triUp >= kvSStartIdx) {
876+ uint32_t triUpRoundDown = RoundDown(triUp, BLOCK_SIZE);
877+ gmOffsetMaskRow = triUp - triUpRoundDown;
878+ gmOffsetMaskColumn = 0U;
879+ maskColumn = kvSEndIdx - triUpRoundDown;
880+ addMaskUbOffset = triUpRoundDown - kvSStartIdx;
881+ } else {
882+ gmOffsetMaskRow = 0U;
883+ gmOffsetMaskColumn = kvSStartIdx - triUp;
884+ maskColumn = columnNum;
885+ addMaskUbOffset = 0U;
886+ }
887+ uint32_t maskColumnRound = RoundUp(maskColumn, BLOCK_SIZE);
888+ 
889+ int64_t offsetMask =
890+ layoutMask.GetOffset(MatrixCoord(gmOffsetMaskRow + maskOffsetThisSubBlock, gmOffsetMaskColumn));
891+ auto gMaskThisSubBlock = gMask[offsetMask];
892+ auto layoutMaskThisSubBlock = layoutMask;
893+ 
894+ uint32_t maxRowNumPerLoop = MAX_UB_S_ELEM_NUM / columnNumRound;
895+ uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, BLOCK_SIZE);
896+ rowNumTile = AscendC::Std::min(rowNumTile, HALF_VECTOR_SIZE);
897+ uint32_t rowLoopNum = CeilDiv(rowActualThisSubBlock, rowNumTile);
898+ 
899+ if (rowActualThisSubBlock == 0U) {
900+ Arch::CrossCoreWaitFlag(qkReady);
901+ return;
902+ }
903+ Arch::CrossCoreWaitFlag(qkReady);
904+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoopNum; rowLoopIdx++) {
905+ uint32_t pingpongFlag = rowLoopIdx % 2U;
906+ uint32_t rowOffsetCurLoop = rowLoopIdx * rowNumTile;
907+ uint32_t rowOffsetIoGm = rowOffsetCurLoop + rowOffsetThisSubBlock;
908+ uint32_t rowNumCurLoop =
909+ (rowLoopIdx == rowLoopNum - 1U) ? (rowActualThisSubBlock - rowOffsetCurLoop) : rowNumTile;
910+ 
911+ uint32_t proTokenIdx = rowOffsetCurLoop % tokenNumPerHeadThisSubBlock;
912+ uint32_t proTokenNum = AscendC::Std::min(rowNumCurLoop, (tokenNumPerHeadThisSubBlock - proTokenIdx)) %
913+ tokenNumPerHeadThisSubBlock;
914+ uint32_t integralHeadNum = (rowNumCurLoop - proTokenNum) / tokenNumPerHeadThisSubBlock;
915+ uint32_t epiTokenNum = rowNumCurLoop - proTokenNum - integralHeadNum * tokenNumPerHeadThisSubBlock;
916+ 
917+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
918+ auto gInputCurLoop = gInput[offsetInput];
919+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
920+ CopySGmToUb(
921+ gInputCurLoop, (pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound, columnNumPad);
922+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
923+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
924+ ScaleS((pingpongFlag * MAX_UB_S_ELEM_NUM), rowNumCurLoop, columnNumRound);
925+
926+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
927+ CopyMaskGmToUb(
928+ gMaskThisSubBlock,
929+ maskColumn,
930+ maskColumnRound,
931+ maskStride,
932+ tokenNumPerHeadThisSubBlock,
933+ proTokenIdx,
934+ proTokenNum,
935+ integralHeadNum,
936+ epiTokenNum);
937+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1);
938+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1);
939+ UpCastMask<half, ElementMask>(maskUbTensor16, maskUbTensor, rowNumCurLoop, columnNumRound);
940+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
941+ ApplyMask(
942+ (pingpongFlag * MAX_UB_S_ELEM_NUM),
943+ rowNumCurLoop,
944+ columnNumRound,
945+ maskColumnRound,
946+ addMaskUbOffset);
947+ 
948+ // online softmax vectorized compute
949+ int64_t offsetOutput = layoutOutput.GetOffset(MatrixCoord(rowOffsetIoGm, 0));
950+ auto gOutputCurLoop = gOutput[offsetOutput];
951+ auto layoutOutputCurLoop = layoutOutput.GetTileLayout(MatrixCoord(rowNumCurLoop, columnNum));
952+ SubCoreCompute(
953+ gOutputCurLoop,
954+ layoutOutputCurLoop,
955+ rowOffsetCurLoop,
956+ isFirstStackTile,
957+ (rowLoopIdx == 0),
958+ columnNumRound,
959+ pingpongFlag,
960+ curStackTileMod,
961+ softmaxFlag,
962+ 0);
963+ }
964+ }
965+ 
966+private:
967+ half scaleValue;
968+ AscendC::LocalTensor<half> lsUbTensor;
969+ AscendC::LocalTensor<half> computeUbTensor;
970+ AscendC::LocalTensor<ElementOutput> lpUbTensor;
971+ AscendC::LocalTensor<ElementMask> maskUbTensor;
972+ AscendC::LocalTensor<half> maskUbTensor16;
973+ AscendC::LocalTensor<half> lmUbTensor;
974+ AscendC::LocalTensor<half> hmUbTensor;
975+ AscendC::LocalTensor<half> gmUbTensor;
976+ AscendC::LocalTensor<half> dmUbTensor;
977+ AscendC::LocalTensor<half> llUbTensor;
978+ AscendC::LocalTensor<half> tvUbTensor;
979+ AscendC::LocalTensor<half> glUbTensor;
980+};
981+}
982+ 
983+#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_ONLINE_SOFTMAX_LOW_PREC_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue_rescale_o.hpp+461-0
@@ -0,0 +1,461 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_HPP
12+#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/resource.hpp"
16+#include "../../../attn_infra/epilogue/dispatch_policy.hpp"
17+#include "../../../attn_infra/epilogue/tile_common/tile_copy.hpp"
18+#include "../../../attn_infra/gemm_coord.hpp"
19+#include "../../../attn_infra/matrix_coord.hpp"
20+ 
21+namespace NpuArch::Epilogue::Block {
22+ 
23+template <
24+ class OutputType_,
25+ class InputType_,
26+ class UpdateType_,
27+ class LseType_,
28+ LseMode LSE_MODE_>
29+class BlockEpilogue<
30+ EpilogueAtlasA2RescaleO<LSE_MODE_, float>,
31+ OutputType_,
32+ InputType_,
33+ UpdateType_,
34+ LseType_>
35+{
36+public:
37+ // Type aliases
38+ using DispatchPolicy = EpilogueAtlasA2RescaleO<LSE_MODE_, float>;
39+ using ArchTag = typename DispatchPolicy::ArchTag;
40+ 
41+ using ElementOutput = typename OutputType_::Element;
42+ using ElementInput = typename InputType_::Element;
43+ using ElementUpdate = typename UpdateType_::Element;
44+ using ElementLse = typename LseType_::Element;
45+ 
46+ using LayoutOutput = typename OutputType_::Layout;
47+ using LayoutInput = typename InputType_::Layout;
48+ using LayoutUpdate = typename UpdateType_::Layout;
49+ using LayoutLse = typename LseType_::Layout;
50+ 
51+ static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE;
52+ 
53+ static constexpr uint32_t HALF_ELENUM_PER_BLK = 16;
54+ static constexpr uint32_t BLOCK_SIZE = 16;
55+ static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128;
56+ static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64;
57+ static constexpr uint32_t HALF_ELENUM_PER_LINE = 256;
58+ static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128;
59+ static constexpr uint32_t MULTIPLIER = 2;
60+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
61+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
62+ static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024;
63+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384;
64+ static constexpr uint32_t HALF_DM_UB_SIZE = 64;
65+ static constexpr uint32_t HALF_LL_UB_SIZE = 256;
66+ static constexpr uint32_t VECTOR_SIZE = 128;
67+ static constexpr uint32_t NUM4 = 4;
68+ static constexpr uint32_t MAX_UB_O_ELEM_NUM = 8192;
69+ static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256;
70+ static constexpr uint32_t SIZE_OF_16BIT = 2;
71+ 
72+ __aicore__ inline
73+ BlockEpilogue(Arch::Resource<ArchTag> &resource)
74+ {
75+ // Allocate UB space
76+ constexpr uint32_t LO_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE;
77+ constexpr uint32_t GO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE;
78+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE;
79+ 
80+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE;
81+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE;
82+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
83+ constexpr uint32_t LSE_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
84+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE;
85+ 
86+ loUbTensor = resource.ubBuf.template GetBufferByByte<float>(LO_UB_TENSOR_OFFSET);
87+ dmUbTensor = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
88+ glUbTensor = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
89+ tvUbTensor = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
90+ goUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementOutput>(GO_UB_TENSOR_OFFSET);
91+ goUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GO_UB_TENSOR_OFFSET);
92+ hmUbTensor = resource.ubBuf.template GetBufferByByte<float>(HM_UB_TENSOR_OFFSET);
93+ gmUbTensor = resource.ubBuf.template GetBufferByByte<float>(GM_UB_TENSOR_OFFSET);
94+ lse32_ubuf_tensor = resource.ubBuf.template GetBufferByByte<float>(LSE_UB_TENSOR_OFFSET);
95+ }
96+ 
97+ __aicore__ inline
98+ ~BlockEpilogue() {}
99+ 
100+ __aicore__ inline
101+ void SetMask(int32_t len)
102+ {
103+ uint64_t mask = 0;
104+ uint64_t one = 1;
105+ uint64_t temp = static_cast<uint64_t>(len) % static_cast<uint64_t>(FLOAT_VECTOR_SIZE);
106+ for (uint64_t i = 0; i < temp; i++) {
107+ mask |= one << i;
108+ }
109+ 
110+ if (len == VECTOR_SIZE) {
111+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
112+ } else if (len >= FLOAT_VECTOR_SIZE) {
113+ AscendC::SetVectorMask<int8_t>(mask, (uint64_t)-1);
114+ } else {
115+ AscendC::SetVectorMask<int8_t>(0x0, mask);
116+ }
117+ }
118+ 
119+ __aicore__ inline
120+ void CopyOToGm(AscendC::GlobalTensor<ElementOutput> gOutput, uint32_t proTokenIdx, uint32_t proTokenNum,
121+ uint32_t epiTokenNum, uint32_t integralHeadNum, uint32_t qSThisSubBlock, uint32_t embed, uint32_t oHiddenSize)
122+ {
123+ uint32_t innerOGmOffset = 0;
124+ uint32_t innerGOUbOffset = 0;
125+ if (proTokenNum != 0U) {
126+ AscendC::DataCopyPad(
127+ gOutput[innerOGmOffset + proTokenIdx * oHiddenSize],
128+ goUbTensor16[innerGOUbOffset],
129+ AscendC::DataCopyExtParams(
130+ proTokenNum, embed * SIZE_OF_16BIT, 0, (oHiddenSize - embed) * SIZE_OF_16BIT, 0));
131+ innerOGmOffset += embed;
132+ innerGOUbOffset += proTokenNum * embed;
133+ }
134+ for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) {
135+ AscendC::DataCopyPad(
136+ gOutput[innerOGmOffset],
137+ goUbTensor16[innerGOUbOffset],
138+ AscendC::DataCopyExtParams(
139+ qSThisSubBlock, embed * SIZE_OF_16BIT, 0, (oHiddenSize - embed) * SIZE_OF_16BIT, 0));
140+ innerOGmOffset += embed;
141+ innerGOUbOffset += qSThisSubBlock * embed;
142+ }
143+ if (epiTokenNum != 0U) {
144+ AscendC::DataCopyPad(
145+ gOutput[innerOGmOffset],
146+ goUbTensor16[innerGOUbOffset],
147+ AscendC::DataCopyExtParams(
148+ epiTokenNum, embed * SIZE_OF_16BIT, 0, (oHiddenSize - embed) * SIZE_OF_16BIT, 0));
149+ }
150+ }
151+ 
152+ __aicore__ inline
153+ void SubCoreCompute(
154+ AscendC::GlobalTensor<ElementOutput> gOutput,
155+ AscendC::GlobalTensor<ElementInput> gInput,
156+ AscendC::GlobalTensor<ElementUpdate> gUpdate,
157+ AscendC::GlobalTensor<ElementLse> gLse,
158+ const LayoutOutput &layoutOutput,
159+ const LayoutInput &layoutInput,
160+ const LayoutUpdate &layoutUpdate,
161+ const LayoutLse &layoutLse,
162+ uint32_t qNThisSubBlock, uint32_t qSThisSubBlock, uint32_t totalRowNum,
163+ uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod,
164+ uint32_t needRowLoop, uint32_t isLastRowLoop, uint32_t rowOffsetLoop,
165+ uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t epiTokenNum, uint32_t integralHeadNum)
166+ {
167+ uint32_t curRowNum = layoutInput.shape(0);
168+ uint32_t embed = layoutInput.shape(1);
169+ uint32_t embedRound = layoutInput.stride(0);
170+ uint32_t curRowNumRound = RoundUp(curRowNum, FLOAT_BLOCK_SIZE);
171+ uint32_t qSBlockSize = layoutOutput.shape(0);
172+ uint32_t oHiddenSize = layoutOutput.shape(1);
173+ uint32_t stride = layoutLse.shape(1); // stride for lse copy out
174+ uint32_t dmUbOffsetCurStackTile = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffsetLoop;
175+ 
176+ if (!isFirstStackTile) {
177+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
178+ AscendC::DataCopy(
179+ loUbTensor, gInput, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
180+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
181+ }
182+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID6);
183+ if (!isFirstStackTile) {
184+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
185+ AscendC::Brcb(tvUbTensor.ReinterpretCast<uint32_t>(),
186+ dmUbTensor[dmUbOffsetCurStackTile].ReinterpretCast<uint32_t>(),
187+ curRowNumRound / FLOAT_BLOCK_SIZE,
188+ AscendC::BrcbRepeatParams(1, 8));
189+ AscendC::PipeBarrier<PIPE_V>();
190+ if (needRowLoop) {
191+ AscendC::DataCopy(
192+ goUbTensor32, gUpdate,
193+ AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
194+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1);
195+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1);
196+ }
197+ // *** go = go * dm_block
198+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
199+ for (uint32_t vmul_idx = 0; vmul_idx < embed / FLOAT_VECTOR_SIZE; ++vmul_idx) {
200+ AscendC::Mul<float, false>(
201+ goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE],
202+ goUbTensor32[vmul_idx * FLOAT_VECTOR_SIZE],
203+ tvUbTensor,
204+ (uint64_t)0,
205+ curRowNum,
206+ AscendC::BinaryRepeatParams(
207+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
208+ }
209+ if (embed % FLOAT_VECTOR_SIZE > 0) {
210+ SetMask(embed % FLOAT_VECTOR_SIZE);
211+ AscendC::Mul<float, false>(
212+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
213+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
214+ tvUbTensor,
215+ (uint64_t)0,
216+ curRowNum,
217+ AscendC::BinaryRepeatParams(
218+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
219+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
220+ }
221+ AscendC::PipeBarrier<PIPE_V>();
222+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
223+ // *** go = lo + go
224+ AscendC::Add<float, false>(
225+ goUbTensor32,
226+ goUbTensor32,
227+ loUbTensor,
228+ (uint64_t)0,
229+ (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
230+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
231+ AscendC::PipeBarrier<PIPE_V>();
232+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
233+ } else {
234+ // *** go = lo
235+ AscendC::DataCopy(
236+ goUbTensor32, gInput, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
237+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
238+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
239+ }
240+ 
241+ if (isLastStackTile) {
242+ // *** gl_block = expand_to_block(gl), 存放于 tv
243+ AscendC::Brcb(
244+ tvUbTensor.ReinterpretCast<uint32_t>(),
245+ glUbTensor.ReinterpretCast<uint32_t>()[rowOffsetLoop],
246+ curRowNumRound / FLOAT_BLOCK_SIZE,
247+ AscendC::BrcbRepeatParams(1, 8));
248+ AscendC::PipeBarrier<PIPE_V>();
249+ // *** go = go / gl_block
250+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
251+ for (uint32_t vdiv_idx = 0; vdiv_idx < embed / FLOAT_VECTOR_SIZE; ++vdiv_idx) {
252+ AscendC::Div<float, false>(
253+ goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE],
254+ goUbTensor32[vdiv_idx * FLOAT_VECTOR_SIZE],
255+ tvUbTensor,
256+ (uint64_t)0,
257+ curRowNum,
258+ AscendC::BinaryRepeatParams(
259+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
260+ }
261+ if (embed % FLOAT_VECTOR_SIZE > 0) {
262+ SetMask(embed % FLOAT_VECTOR_SIZE);
263+ AscendC::Div<float, false>(
264+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
265+ goUbTensor32[embed / FLOAT_VECTOR_SIZE * FLOAT_VECTOR_SIZE],
266+ tvUbTensor,
267+ (uint64_t)0,
268+ curRowNum,
269+ AscendC::BinaryRepeatParams(
270+ 1, 1, 0, embedRound / FLOAT_BLOCK_SIZE, embedRound / FLOAT_BLOCK_SIZE, 1));
271+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
272+ }
273+ AscendC::PipeBarrier<PIPE_V>();
274+ 
275+ // *** go = castfp32to16(go)
276+ if (std::is_same<ElementOutput, bfloat16_t>::value) {
277+ AscendC::Cast<ElementOutput, float, false>(
278+ goUbTensor16, goUbTensor32,
279+ AscendC::RoundMode::CAST_RINT, (uint64_t)0,
280+ (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
281+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
282+ } else {
283+ AscendC::Cast<ElementOutput, float, false>(
284+ goUbTensor16, goUbTensor32,
285+ AscendC::RoundMode::CAST_NONE, (uint64_t)0,
286+ (curRowNum * embedRound + FLOAT_VECTOR_SIZE - 1) / FLOAT_VECTOR_SIZE,
287+ AscendC::UnaryRepeatParams(1, 1, 4, 8));
288+ }
289+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
290+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
291+ 
292+ // ***move O to GM
293+ CopyOToGm(
294+ gOutput, proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, qSThisSubBlock, embed, oHiddenSize);
295+ if constexpr(LSE_MODE == LseMode::OUT_ONLY) {
296+ if (isLastRowLoop) {
297+ AscendC::PipeBarrier<PIPE_V>();
298+ AscendC::Ln<float, false>(
299+ lse32_ubuf_tensor,
300+ glUbTensor,
301+ (uint64_t)0,
302+ CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE),
303+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
304+ AscendC::PipeBarrier<PIPE_V>();
305+ AscendC::Add<float, false>(
306+ lse32_ubuf_tensor,
307+ lse32_ubuf_tensor,
308+ gmUbTensor,
309+ (uint64_t)0,
310+ CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE),
311+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
312+ AscendC::PipeBarrier<PIPE_V>();
313+ AscendC::Brcb(
314+ tvUbTensor.ReinterpretCast<uint32_t>(),
315+ lse32_ubuf_tensor.ReinterpretCast<uint32_t>(),
316+ CeilDiv(totalRowNum, FLOAT_BLOCK_SIZE),
317+ AscendC::BrcbRepeatParams(1, 8));
318+ AscendC::PipeBarrier<PIPE_V>();
319+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID4);
320+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID4);
321+ if (qNThisSubBlock == 0U) {
322+ AscendC::DataCopyPad(
323+ gLse, tvUbTensor,
324+ AscendC::DataCopyExtParams(totalRowNum, sizeof(float), 0, (stride - 1) * sizeof(float), 0));
325+ } else {
326+ for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) {
327+ AscendC::DataCopyPad(
328+ gLse[qNIdx],
329+ tvUbTensor[qNIdx * qSBlockSize * FLOAT_BLOCK_SIZE],
330+ AscendC::DataCopyExtParams(
331+ qSBlockSize, sizeof(float), 0, (stride - 1) * sizeof(float), 0));
332+ }
333+ }
334+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
335+ }
336+ }
337+ } else if (needRowLoop) {
338+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID5);
339+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID5);
340+ AscendC::DataCopy(
341+ gUpdate, goUbTensor32, AscendC::DataCopyParams(1, curRowNum * embedRound / FLOAT_BLOCK_SIZE, 0, 0));
342+ }
343+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID6);
344+ }
345+ 
346+ __aicore__ inline
347+ void operator()(
348+ AscendC::GlobalTensor<ElementOutput> gOutput,
349+ AscendC::GlobalTensor<ElementInput> gInput,
350+ AscendC::GlobalTensor<ElementUpdate> gUpdate,
351+ AscendC::GlobalTensor<ElementLse> gLse,
352+ const LayoutOutput &layoutOutput,
353+ const LayoutInput &layoutInput,
354+ const LayoutUpdate &layoutUpdate,
355+ const LayoutLse &layoutLse,
356+ GemmCoord actualBlockShape,
357+ uint32_t qSBlockSize, uint32_t qNBlockSize,
358+ uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod)
359+ {
360+ uint32_t rowNum = actualBlockShape.m();
361+ uint32_t embed = actualBlockShape.n();
362+ uint32_t maxRowNumPerLoop = MAX_UB_O_ELEM_NUM / embed;
363+ uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, FLOAT_BLOCK_SIZE);
364+ 
365+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
366+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
367+ 
368+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
369+ uint32_t qNThisSubBlock = (qNBlockSize == 1U) ? 0
370+ : (subBlockIdx == 1U) ? (qNBlockSize - qNSplitSubBlock)
371+ : qNSplitSubBlock;
372+ uint32_t inRowSplitSubBlock =
373+ (qNBlockSize == 1U) ? (qSBlockSize / subBlockNum) : (qSBlockSize * qNSplitSubBlock);
374+ uint32_t inRowActualThisSubBlock = (subBlockIdx == 1U) ? (rowNum - inRowSplitSubBlock) : inRowSplitSubBlock;
375+ uint32_t inRowOffsetThisSubBlock = subBlockIdx * inRowSplitSubBlock;
376+ uint32_t outRowOffsetThisSubBlock = (qNBlockSize == 1U) ? inRowOffsetThisSubBlock : 0;
377+ uint32_t outColOffsetThisSubBlock = (qNBlockSize == 1U) ? 0 : subBlockIdx * qNSplitSubBlock * embed;
378+ uint32_t qSThisSubBlock = (qNBlockSize == 1U) ? inRowActualThisSubBlock : qSBlockSize;
379+ int64_t outOffsetSubBlock =
380+ layoutOutput.GetOffset(MatrixCoord(outRowOffsetThisSubBlock, outColOffsetThisSubBlock));
381+ 
382+ uint32_t outLseRowOffsetThisSubBlock = (qNBlockSize == 1U) ? inRowOffsetThisSubBlock : 0;
383+ uint32_t outLseColOffsetThisSubBlock = (qNBlockSize == 1U) ? 0 : subBlockIdx * qNSplitSubBlock;
384+ int64_t offsetLse = layoutLse.GetOffset(MatrixCoord(outLseRowOffsetThisSubBlock, outLseColOffsetThisSubBlock));
385+ auto gLseThisSubBlock = gLse[offsetLse];
386+ 
387+ if (inRowActualThisSubBlock > 0U) {
388+ uint32_t rowLoop = CeilDiv(inRowActualThisSubBlock, rowNumTile);
389+ uint32_t needRowLoop = (rowLoop > 1U) ? 1 : 0;
390+ 
391+ // The rows of each cycle consist of multiple heads with several tokens.
392+ // There are several integral heads, one prologue head, one epilogue head.
393+ uint32_t proTokenIdx = 0; // the token idx of the start token of the prologue part
394+ uint32_t proTokenIdxPre = 0; // the token idx of the start token of the pre prologue part
395+ uint32_t proTokenNum = 0; // the token num of the prologue part
396+ uint32_t epiTokenNum = 0; // the token num of the epilogue part
397+ uint32_t integralHeadNum = 0; // the number of integral heads within a cycle
398+ uint32_t qSRemian = qSThisSubBlock;
399+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoop; rowLoopIdx++) {
400+ uint32_t rowOffsetLoop = rowLoopIdx * rowNumTile;
401+ uint32_t rowOffsetCurLoop = inRowOffsetThisSubBlock + rowOffsetLoop;
402+ uint32_t rowActualCurLoop =
403+ (rowLoopIdx == (rowLoop - 1U)) ? inRowActualThisSubBlock - rowLoopIdx * rowNumTile : rowNumTile;
404+ 
405+ int64_t offsetOutput =
406+ static_cast<int64_t>(rowLoopIdx * rowNumTile / qSThisSubBlock * embed) + outOffsetSubBlock;
407+ auto gOutputCurLoop = gOutput[offsetOutput];
408+ auto layoutOutputCurLoop = layoutOutput;
409+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurLoop, 0));
410+ auto gInputCurLoop = gInput[offsetInput];
411+ auto layoutInputCurLoop = layoutInput.GetTileLayout(MatrixCoord(rowActualCurLoop, embed));
412+ 
413+ int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurLoop, 0));
414+ auto gUpdateCurLoop = gUpdate[offsetUpdate];
415+ auto layoutUpdateCurLoop = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurLoop, embed));
416+ 
417+ proTokenIdx = rowOffsetLoop % qSThisSubBlock;
418+ proTokenNum = AscendC::Std::min(rowActualCurLoop, (qSThisSubBlock - proTokenIdx)) % qSThisSubBlock;
419+ integralHeadNum = (rowActualCurLoop - proTokenNum) / qSThisSubBlock;
420+ epiTokenNum = rowActualCurLoop - proTokenNum - integralHeadNum * qSThisSubBlock;
421+ 
422+ SubCoreCompute(
423+ gOutputCurLoop,
424+ gInputCurLoop,
425+ gUpdateCurLoop,
426+ gLseThisSubBlock,
427+ layoutOutputCurLoop,
428+ layoutInputCurLoop,
429+ layoutUpdateCurLoop,
430+ layoutLse,
431+ qNThisSubBlock,
432+ qSThisSubBlock,
433+ inRowActualThisSubBlock,
434+ isFirstStackTile,
435+ isLastStackTile,
436+ curStackTileMod,
437+ needRowLoop,
438+ (rowLoopIdx == rowLoop - 1U),
439+ rowOffsetLoop,
440+ proTokenIdx,
441+ proTokenNum,
442+ epiTokenNum,
443+ integralHeadNum);
444+ }
445+ }
446+ }
447+ 
448+private:
449+ AscendC::LocalTensor<float> loUbTensor;
450+ AscendC::LocalTensor<float> dmUbTensor;
451+ AscendC::LocalTensor<float> hmUbTensor;
452+ AscendC::LocalTensor<float> glUbTensor;
453+ AscendC::LocalTensor<float> tvUbTensor;
454+ AscendC::LocalTensor<ElementOutput> goUbTensor16;
455+ AscendC::LocalTensor<float> goUbTensor32;
456+ AscendC::LocalTensor<float> gmUbTensor;
457+ AscendC::LocalTensor<float> lse32_ubuf_tensor;
458+};
459+}
460+ 
461+#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue_rescale_o_arch35_reg_high_prec.hpp+479-0
@@ -0,0 +1,479 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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 (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+#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_ARCH35_REG_HIGH_PREC
12+#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_O_ARCH35_REG_HIGH_PREC
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/resource.hpp"
16+#include "../../../attn_infra/epilogue/dispatch_policy.hpp"
17+#include "../../../attn_infra/epilogue/tile_common/tile_copy.hpp"
18+#include "../../../attn_infra/gemm_coord.hpp"
19+#include "../../../attn_infra/matrix_coord.hpp"
20+#include "../../../tla/tensor.hpp"
21+#include "../../../tla/layout.hpp"
22+ 
23+namespace NpuArch::Epilogue::Block {
24+ 
25+template <
26+ class ElementO_,
27+ class ElementOTmp_,
28+ class ElementS_,
29+ class TileCopy_,
30+ class OTmpSrcPos_>
31+class BlockEpilogue<
32+ EpilogueAtlasA5BsaRescaleO,
33+ ElementO_,
34+ ElementOTmp_,
35+ ElementS_,
36+ TileCopy_,
37+ OTmpSrcPos_>
38+{
39+public:
40+ using DispatchPolicy = EpilogueAtlasA5BsaRescaleO;
41+ using ArchTag = typename DispatchPolicy::ArchTag;
42+ using ElementO = ElementO_;
43+ using ElementOTmp = ElementOTmp_;
44+ using SMDtype = ElementS_;
45+ using TileCopy = TileCopy_;
46+ using OTmpSrcPos = OTmpSrcPos_;
47+ 
48+ using CopyUbToGmO = typename TileCopy::CopyUbToGmO;
49+ 
50+ static constexpr uint32_t UB_OTMP_BUF_STAGES = 2;
51+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 32768;
52+ static constexpr uint32_t DM_UB_GLOBAL_ELEM_NUM = 64 * 2; //! 2* for 11 22 ... 32,32
53+ static constexpr uint32_t RESCALE_ROW_MAX_ELEM_NUM = 64;
54+ static constexpr uint32_t RESCALE_COL_MAX_ELEM_NUM = 128;
55+ static constexpr uint32_t RESCALE_VREG_SIZE = 256 / sizeof(ElementOTmp);
56+ static constexpr bool DEQ = true;
57+ 
58+ __aicore__ inline
59+ BlockEpilogue(Arch::Resource<ArchTag> &resource, uint32_t embed_ = 128)
60+ {
61+ constexpr uint32_t LO_UB_TENSOR_OFFSET = 4 * UB_UINT8_BLOCK_SIZE;
62+ constexpr uint32_t GO_UB_TENSOR_OFFSET = 6 * UB_UINT8_BLOCK_SIZE;
63+ constexpr uint32_t LM_UB_TENSOR_OFFSET = 7 * UB_UINT8_BLOCK_SIZE;
64+ constexpr uint32_t GM_UB_TENSOR_OFFSET = LM_UB_TENSOR_OFFSET + 128 * sizeof(float);
65+ constexpr uint32_t DM_UB_TENSOR_OFFSET = GM_UB_TENSOR_OFFSET + 128 * sizeof(float);
66+ constexpr uint32_t LL_UB_TENSOR_OFFSET = DM_UB_TENSOR_OFFSET + 3 * 128 * sizeof(float);
67+ constexpr uint32_t GL_UB_TENSOR_OFFSET = LL_UB_TENSOR_OFFSET + 128 * sizeof(float);
68+ 
69+ for (uint32_t i = 0; i < UB_OTMP_BUF_STAGES; i++) {
70+ loUbTensor[i] = resource.ubBuf.template GetBufferByByte<ElementOTmp>(
71+ LO_UB_TENSOR_OFFSET + i * UB_UINT8_BLOCK_SIZE);
72+ }
73+ goUbTensor32 = resource.ubBuf.template GetBufferByByte<ElementOTmp>(GO_UB_TENSOR_OFFSET);
74+ goUbTensor16 = resource.ubBuf.template GetBufferByByte<ElementO>(GO_UB_TENSOR_OFFSET);
75+ glUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(GL_UB_TENSOR_OFFSET);
76+ dmUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(DM_UB_TENSOR_OFFSET);
77+ scaleTensor = resource.ubBuf.template GetBufferByByte<float>(LM_UB_TENSOR_OFFSET + 4096 * 2 + embed_ * sizeof(float));
78+ }
79+ 
80+ __aicore__ inline
81+ ~BlockEpilogue()
82+ {
83+ }
84+ 
85+ template <uint32_t MODE, pipe_t PIPE>
86+ __aicore__ inline
87+ void SetCrossCoreSync(Arch::CrossCoreFlag &crossCoreFlag)
88+ {
89+ if constexpr (MODE == 4U) {
90+ Arch::CrossCoreSetFlag<MODE, PIPE>(crossCoreFlag);
91+ }
92+ }
93+ 
94+ template <uint32_t MODE, pipe_t PIPE>
95+ __aicore__ inline
96+ void WaitCrossCoreSync(Arch::CrossCoreFlag &crossCoreFlag)
97+ {
98+ if constexpr (MODE == 4U) {
99+ Arch::CrossCoreWaitFlag<MODE, PIPE>(crossCoreFlag);
100+ }
101+ }
102+ 
103+ template <class TensorDst>
104+ __aicore__ inline
105+ void SubCoreCompute(const TensorDst &gOTensor,
106+ uint32_t curTileMod,
107+ uint32_t ubOTmpBufId,
108+ bool isFirstKvSTile,
109+ bool isLastKvSTile,
110+ uint32_t colNumOri,
111+ Arch::CrossCoreFlag mm2ToReFlag)
112+ {
113+ __ubuf__ ElementOTmp *goUb = (__ubuf__ ElementOTmp *) goUbTensor32.GetPhyAddr();
114+ __ubuf__ ElementOTmp *loUb = (__ubuf__ ElementOTmp *) loUbTensor[ubOTmpBufId].GetPhyAddr();
115+ __ubuf__ ElementOTmp *glUb = ( __ubuf__ ElementOTmp *) glUbTensor32.GetPhyAddr();
116+ __ubuf__ ElementOTmp *dmUb = (__ubuf__ ElementOTmp *) dmUbTensor32[curTileMod * DM_UB_GLOBAL_ELEM_NUM].GetPhyAddr();
117+ __ubuf__ ElementOTmp *scaleUb = (__ubuf__ ElementOTmp *)scaleTensor.GetPhyAddr();
118+
119+ WaitCrossCoreSync<4, PIPE_V>(mm2ToReFlag);
120+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
121+ 
122+ if (isFirstKvSTile) {
123+ if (!isLastKvSTile) {
124+ if constexpr (DEQ) {
125+ DeqLocalO<ElementOTmp, 128, DEQ>(goUb, loUb, scaleUb, rowNumCurSubCore); AscendC::PipeBarrier<PIPE_V>();
126+ } else {
127+ AscendC::DataCopy(goUbTensor32, loUbTensor[ubOTmpBufId], rowNumCurSubCore * colNumOri); AscendC::PipeBarrier<PIPE_V>();
128+ }
129+ } else { // go = lo div sum
130+ DivFuncLastAndFirst<ElementOTmp, 128, DEQ>(goUb, loUb, glUb, scaleUb, rowNumCurSubCore);
131+ }
132+ } else if (!isLastKvSTile) {
133+ RescaleFunc<ElementOTmp, 128, DEQ>(goUb, loUb, dmUb, scaleUb, rowNumCurSubCore);
134+ } else {
135+ RescaleFuncLastNotFirst<ElementOTmp, 128, DEQ>(goUb, loUb, dmUb, glUb, scaleUb, rowNumCurSubCore);
136+ }
137+
138+ // release lo buf
139+ SetCrossCoreSync<4, PIPE_V>(mm2ToReFlag);
140+ if (isLastKvSTile) {
141+ AscendC::PipeBarrier<PIPE_V>();
142+ if constexpr (std::is_same<ElementO, bfloat16_t>::value) {
143+ AscendC::Cast(
144+ goUbTensor16, goUbTensor32,
145+ AscendC::RoundMode::CAST_RINT,
146+ rowNumCurSubCore * colNumOri);
147+ } else {
148+ AscendC::Cast(
149+ goUbTensor16, goUbTensor32,
150+ AscendC::RoundMode::CAST_NONE,
151+ rowNumCurSubCore * colNumOri);
152+ }
153+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
154+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
155+ DataCopy(gOTensor[rowOffsetCurSubCore * colNumOri], goUbTensor16, rowNumCurSubCore * colNumOri);
156+ }
157+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
158+ }
159+
160+ template <typename T=float, uint32_t colStride=128, bool DEQ=false>
161+ __simd_vf__ inline void RescaleFunc(__ubuf__ T *goUb, __ubuf__ T *loUb, __ubuf__ T *dmUb, __ubuf__ T *scaleUb,
162+ uint32_t row)
163+ {
164+ using namespace AscendC::MicroAPI;
165+ 
166+ RegTensor<float> dmVreg0, dmVreg1, scaleVreg0, scaleVreg1;
167+ RegTensor<float> goPreVreg0, goPreVreg0_2, goPreVreg1, goPreVreg1_2;
168+ RegTensor<float> loVreg0, loVreg0_2, loVreg1, loVreg1_2;
169+ RegTensor<float> mulVreg0, mulVreg0_2, mulVreg1, mulVreg1_2;
170+ RegTensor<float> goCurVreg0, goCurVreg0_2, goCurVreg1, goCurVreg1_2;
171+ 
172+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
173+ constexpr uint32_t vlElemNum = 64;
174+
175+ uint32_t halfRow = (row + 1) / 2; // = 32
176+
177+ // 循环 32 次，每次处理 row_i 和 row_{i+32}
178+ if constexpr (DEQ) {
179+ LoadAlign<T, LoadDist::DIST_NORM>(scaleVreg0, scaleUb); // col0
180+ LoadAlign<T, LoadDist::DIST_NORM>(scaleVreg1, scaleUb + vlElemNum); //col1
181+ }
182+ 
183+ for (uint16_t i = 0; i < halfRow; i++) {
184+ uint32_t row1_idx = i + halfRow;
185+ 
186+ uint32_t baseOffset0 = i * colStride;
187+ uint32_t secondOffset0 = baseOffset0 + vlElemNum;
188+ uint32_t baseOffset1 = row1_idx * colStride;
189+ uint32_t secondOffset1 = baseOffset1 + vlElemNum;
190+ 
191+ // Load scalars and broadcast
192+ //! 2* for 11 22 ... 32,32
193+ LoadAlign<T, LoadDist::DIST_BRC_B32>(dmVreg0, dmUb + 2*i);
194+ LoadAlign<T, LoadDist::DIST_BRC_B32>(dmVreg1, dmUb + 2*row1_idx);
195+ 
196+ // Load Row i
197+ LoadAlign<T, LoadDist::DIST_NORM>(goPreVreg0, goUb + baseOffset0);
198+ LoadAlign<T, LoadDist::DIST_NORM>(goPreVreg0_2, goUb + secondOffset0);
199+ LoadAlign<T, LoadDist::DIST_NORM>(loVreg0, loUb + baseOffset0);
200+ LoadAlign<T, LoadDist::DIST_NORM>(loVreg0_2, loUb + secondOffset0);
201+ 
202+ // Load Row i+32
203+ LoadAlign<T, LoadDist::DIST_NORM>(goPreVreg1, goUb + baseOffset1);
204+ LoadAlign<T, LoadDist::DIST_NORM>(goPreVreg1_2, goUb + secondOffset1);
205+ LoadAlign<T, LoadDist::DIST_NORM>(loVreg1, loUb + baseOffset1);
206+ LoadAlign<T, LoadDist::DIST_NORM>(loVreg1_2, loUb + secondOffset1);
207+ 
208+ // Multiply (4-issue)
209+ Mul(mulVreg0, goPreVreg0, dmVreg0, pregFull);
210+ Mul(mulVreg0_2, goPreVreg0_2, dmVreg0, pregFull);
211+ Mul(mulVreg1, goPreVreg1, dmVreg1, pregFull);
212+ Mul(mulVreg1_2, goPreVreg1_2, dmVreg1, pregFull);
213+ 
214+ if constexpr (DEQ) {
215+ MulDstAdd(loVreg0, scaleVreg0, mulVreg0, pregFull);
216+ MulDstAdd(loVreg0_2, scaleVreg1, mulVreg0_2, pregFull);
217+ MulDstAdd(loVreg1, scaleVreg0, mulVreg1, pregFull);
218+ MulDstAdd(loVreg1_2, scaleVreg1, mulVreg1_2, pregFull);
219+ } else {
220+ Add(loVreg0, mulVreg0, loVreg0, pregFull);
221+ Add(loVreg0_2, mulVreg0_2, loVreg0_2, pregFull);
222+ Add(loVreg1, mulVreg1, loVreg1, pregFull);
223+ Add(loVreg1_2, mulVreg1_2, loVreg1_2, pregFull);
224+ }
225+ 
226+ // Store (4-issue)
227+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + baseOffset0, loVreg0, pregFull);
228+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + secondOffset0, loVreg0_2, pregFull);
229+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + baseOffset1, loVreg1, pregFull);
230+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + secondOffset1, loVreg1_2, pregFull);
231+ }
232+ }
233+ 
234+ template <typename T=float, uint32_t colStride=128, bool DEQ=false>
235+ __simd_vf__ inline void RescaleFuncLastNotFirst(__ubuf__ T *goUb, __ubuf__ T *loUb,
236+ __ubuf__ T *dmUb, __ubuf__ T *glUb,
237+ __ubuf__ T *scaleUb,
238+ uint32_t row)
239+ {
240+ using namespace AscendC::MicroAPI;
241+
242+ RegTensor<float> dmVreg0, glVreg0, dmVreg1, glVreg1;
243+ RegTensor<float> scaleVreg0, scaleVreg1; // 增加 scale 寄存器
244+ RegTensor<float> goPreVreg0, goPreVreg0_2, goPreVreg1, goPreVreg1_2;
245+ RegTensor<float> loVreg0, loVreg0_2, loVreg1, loVreg1_2;
246+ RegTensor<float> mulVreg0, mulVreg0_2, mulVreg1, mulVreg1_2;
247+ RegTensor<float> goCurVreg0, goCurVreg0_2, goCurVreg1, goCurVreg1_2;
248+ RegTensor<float> divVreg0, divVreg0_2, divVreg1, divVreg1_2;
249+ 
250+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
251+ constexpr uint32_t vlElemNum = 64;
252+
253+ uint32_t halfRow = (row + 1) / 2; // = 32
254+ 
255+ // 若开启反量化，在循环外一次性加载 per-channel scale
256+ if constexpr (DEQ) {
257+ LoadAlign<T, LoadDist::DIST_NORM>(scaleVreg0, scaleUb); // col0
258+ LoadAlign<T, LoadDist::DIST_NORM>(scaleVreg1, scaleUb + vlElemNum); // col1
259+ }
260+
261+ for (uint16_t i = 0; i < halfRow; i++) {
262+ uint32_t row1_idx = i + halfRow;
263+ 
264+ uint32_t baseOffset0 = i * colStride;
265+ uint32_t secondOffset0 = baseOffset0 + vlElemNum;
266+ uint32_t baseOffset1 = row1_idx * colStride;
267+ uint32_t secondOffset1 = baseOffset1 + vlElemNum;
268+ 
269+ // Load scalars (i and i+32)
270+ //! 2* for 11 22 ... 32,32
271+ LoadAlign<T, LoadDist::DIST_BRC_B32>(dmVreg0, dmUb + 2*i);
272+ LoadAlign<T, LoadDist::DIST_BRC_B32>(glVreg0, glUb + 2*i);
273+ LoadAlign<T, LoadDist::DIST_BRC_B32>(dmVreg1, dmUb + 2*row1_idx);
274+ LoadAlign<T, LoadDist::DIST_BRC_B32>(glVreg1, glUb + 2*row1_idx);
275+ 
276+ // Load Data
277+ LoadAlign<T, LoadDist::DIST_NORM>(goPreVreg0, goUb + baseOffset0);
278+ LoadAlign<T, LoadDist::DIST_NORM>(goPreVreg0_2, goUb + secondOffset0);
279+ LoadAlign<T, LoadDist::DIST_NORM>(goPreVreg1, goUb + baseOffset1);
280+ LoadAlign<T, LoadDist::DIST_NORM>(goPreVreg1_2, goUb + secondOffset1);
281+ 
282+ LoadAlign<T, LoadDist::DIST_NORM>(loVreg0, loUb + baseOffset0);
283+ LoadAlign<T, LoadDist::DIST_NORM>(loVreg0_2, loUb + secondOffset0);
284+ LoadAlign<T, LoadDist::DIST_NORM>(loVreg1, loUb + baseOffset1);
285+ LoadAlign<T, LoadDist::DIST_NORM>(loVreg1_2, loUb + secondOffset1);
286+ 
287+ // Muls: goPre * dm
288+ Mul(mulVreg0, goPreVreg0, dmVreg0, pregFull);
289+ Mul(mulVreg0_2, goPreVreg0_2, dmVreg0, pregFull);
290+ Mul(mulVreg1, goPreVreg1, dmVreg1, pregFull);
291+ Mul(mulVreg1_2, goPreVreg1_2, dmVreg1, pregFull);
292+ 
293+ // Dequantization: Local * scale + mulvreg
294+ // dstReg与srcReg0相乘后与srcReg1相加
295+ if constexpr (DEQ) {
296+ MulDstAdd(loVreg0, scaleVreg0, mulVreg0, pregFull);
297+ MulDstAdd(loVreg0_2, scaleVreg1, mulVreg0_2, pregFull);
298+ MulDstAdd(loVreg1, scaleVreg0, mulVreg1, pregFull);
299+ MulDstAdd(loVreg1_2, scaleVreg1, mulVreg1_2, pregFull);
300+ } else {
301+ Add(loVreg0, mulVreg0, loVreg0, pregFull);
302+ Add(loVreg0_2, mulVreg0_2, loVreg0_2, pregFull);
303+ Add(loVreg1, mulVreg1, loVreg1, pregFull);
304+ Add(loVreg1_2, mulVreg1_2, loVreg1_2, pregFull);
305+ }
306+ 
307+ Div(divVreg0, loVreg0, glVreg0, pregFull);
308+ Div(divVreg0_2, loVreg0_2, glVreg0, pregFull);
309+ Div(divVreg1, loVreg1, glVreg1, pregFull);
310+ Div(divVreg1_2, loVreg1_2, glVreg1, pregFull);
311+ 
312+ // Stores
313+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + baseOffset0, divVreg0, pregFull);
314+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + secondOffset0, divVreg0_2, pregFull);
315+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + baseOffset1, divVreg1, pregFull);
316+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + secondOffset1, divVreg1_2, pregFull);
317+ }
318+ }
319+ 
320+ template <typename T=float, uint32_t colStride=128, bool DEQ=false>
321+ __simd_vf__ inline void DivFuncLastAndFirst(__ubuf__ T *goUb, __ubuf__ T *loUb, __ubuf__ T *glUb,
322+ __ubuf__ T *scaleUb, // 保持接口一致性
323+ uint32_t row)
324+ {
325+ using namespace AscendC::MicroAPI;
326+
327+ RegTensor<float> goCurVreg0, goCurVreg0_2, goCurVreg1, goCurVreg1_2;
328+ RegTensor<float> glVreg0, glVreg1;
329+ RegTensor<float> scaleVreg0, scaleVreg1; // 仅做声明备用
330+ RegTensor<float> divVreg0, divVreg0_2, divVreg1, divVreg1_2;
331+
332+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
333+ constexpr uint32_t vlElemNum = 64;
334+
335+ if constexpr (colStride==128) {
336+ uint32_t halfRow = (row + 1) / 2; // = 32
337+
338+ if constexpr (DEQ) {
339+ LoadAlign<T, LoadDist::DIST_NORM>(scaleVreg0, scaleUb);
340+ LoadAlign<T, LoadDist::DIST_NORM>(scaleVreg1, scaleUb + vlElemNum);
341+ }
342+ for (uint16_t i = 0; i < halfRow; i++) {
343+ uint32_t row1_idx = i + halfRow;
344+ 
345+ uint32_t baseOffset0 = i * colStride;
346+ uint32_t secondOffset0 = baseOffset0 + vlElemNum;
347+ uint32_t baseOffset1 = row1_idx * colStride;
348+ uint32_t secondOffset1 = baseOffset1 + vlElemNum;
349+ //! 2* for 11 22 ... 32,32
350+ LoadAlign<T, LoadDist::DIST_BRC_B32>(glVreg0, glUb + 2*i);
351+ LoadAlign<T, LoadDist::DIST_BRC_B32>(glVreg1, glUb + 2*row1_idx);
352+ 
353+ LoadAlign<T, LoadDist::DIST_NORM>(goCurVreg0, loUb + baseOffset0);
354+ LoadAlign<T, LoadDist::DIST_NORM>(goCurVreg0_2, loUb + secondOffset0);
355+ LoadAlign<T, LoadDist::DIST_NORM>(goCurVreg1, loUb + baseOffset1);
356+ LoadAlign<T, LoadDist::DIST_NORM>(goCurVreg1_2, loUb + secondOffset1);
357+ 
358+ if constexpr (DEQ) {
359+ Mul(goCurVreg0, goCurVreg0, scaleVreg0, pregFull);
360+ Mul(goCurVreg0_2, goCurVreg0_2, scaleVreg1, pregFull);
361+ Mul(goCurVreg1, goCurVreg1, scaleVreg0, pregFull);
362+ Mul(goCurVreg1_2, goCurVreg1_2, scaleVreg1, pregFull);
363+ }
364+ 
365+ Div(divVreg0, goCurVreg0, glVreg0, pregFull);
366+ Div(divVreg0_2, goCurVreg0_2, glVreg0, pregFull);
367+ Div(divVreg1, goCurVreg1, glVreg1, pregFull);
368+ Div(divVreg1_2, goCurVreg1_2, glVreg1, pregFull);
369+ 
370+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + baseOffset0, divVreg0, pregFull);
371+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + secondOffset0, divVreg0_2, pregFull);
372+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + baseOffset1, divVreg1, pregFull);
373+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + secondOffset1, divVreg1_2, pregFull);
374+ }
375+ }
376+ }
377+ 
378+ template <typename T=float, uint32_t colStride=128, bool DEQ=false>
379+ __simd_vf__ inline void DeqLocalO(__ubuf__ T *goUb, __ubuf__ T *loUb,
380+ __ubuf__ T *scaleUb,
381+ uint32_t row)
382+ {
383+ using namespace AscendC::MicroAPI;
384+
385+ RegTensor<float> goCurVreg0, goCurVreg0_2, goCurVreg1, goCurVreg1_2;
386+ RegTensor<float> scaleVreg0, scaleVreg1;
387+
388+ MaskReg pregFull = CreateMask<float, MaskPattern::ALL>();
389+ constexpr uint32_t vlElemNum = 64;
390+
391+ if constexpr (colStride==128) {
392+ uint32_t halfRow = (row + 1) / 2; // = 32
393+
394+ if constexpr (DEQ) {
395+ LoadAlign<T, LoadDist::DIST_NORM>(scaleVreg0, scaleUb);
396+ LoadAlign<T, LoadDist::DIST_NORM>(scaleVreg1, scaleUb + vlElemNum);
397+ }
398+ for (uint16_t i = 0; i < halfRow; i++) {
399+ uint32_t row1_idx = i + halfRow;
400+ 
401+ uint32_t baseOffset0 = i * colStride;
402+ uint32_t secondOffset0 = baseOffset0 + vlElemNum;
403+ uint32_t baseOffset1 = row1_idx * colStride;
404+ uint32_t secondOffset1 = baseOffset1 + vlElemNum;
405+ 
406+ LoadAlign<T, LoadDist::DIST_NORM>(goCurVreg0, loUb + baseOffset0);
407+ LoadAlign<T, LoadDist::DIST_NORM>(goCurVreg0_2, loUb + secondOffset0);
408+ LoadAlign<T, LoadDist::DIST_NORM>(goCurVreg1, loUb + baseOffset1);
409+ LoadAlign<T, LoadDist::DIST_NORM>(goCurVreg1_2, loUb + secondOffset1);
410+ 
411+ if constexpr (DEQ) {
412+ Mul(goCurVreg0, goCurVreg0, scaleVreg0, pregFull);
413+ Mul(goCurVreg0_2, goCurVreg0_2, scaleVreg1, pregFull);
414+ Mul(goCurVreg1, goCurVreg1, scaleVreg0, pregFull);
415+ Mul(goCurVreg1_2, goCurVreg1_2, scaleVreg1, pregFull);
416+ }
417+ 
418+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + baseOffset0, goCurVreg0, pregFull);
419+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + secondOffset0, goCurVreg0_2, pregFull);
420+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + baseOffset1, goCurVreg1, pregFull);
421+ StoreAlign<T, StoreDist::DIST_NORM_B32>(goUb + secondOffset1, goCurVreg1_2, pregFull);
422+ }
423+ }
424+ }
425+ 
426+ template <class TensorDst>
427+ __aicore__ inline
428+ void operator()(const TensorDst &gOTensor,
429+ GemmCoord actualOriShape,
430+ uint32_t curTileMod,
431+ uint32_t gatheredKvSTileIdx,
432+ bool isFirstKvSTile,
433+ bool isLastKvSTile,
434+ Arch::CrossCoreFlag mm2ToReFlag)
435+ {
436+ uint32_t rowNumOri = actualOriShape[0];
437+ uint32_t colNumOri = actualOriShape[1];
438+ constexpr uint32_t FP8_BLOCK_SIZE = 32;
439+ if (rowNumOri <= FP8_BLOCK_SIZE) {
440+ rowNumCurSubCore = (subBlockIdx == 0) ? rowNumOri : 0;
441+ } else {
442+ uint32_t mhalf = (rowNumOri + subBlockNum -1) / subBlockNum;
443+ uint32_t mAlign = (mhalf > FP8_BLOCK_SIZE) ? RoundUp(mhalf, FP8_BLOCK_SIZE) : FP8_BLOCK_SIZE;
444+ rowNumCurSubCore = (subBlockIdx == 0) ? mAlign : (rowNumOri - mAlign);
445+ }
446+ rowOffsetCurSubCore = subBlockIdx == 0 ? 0: rowNumOri - rowNumCurSubCore;
447+ uint32_t ubOTmpBufId = gatheredKvSTileIdx % UB_OTMP_BUF_STAGES;
448+ 
449+ if (rowNumCurSubCore > 0) {
450+ SubCoreCompute(
451+ gOTensor,
452+ curTileMod,
453+ ubOTmpBufId,
454+ isFirstKvSTile,
455+ isLastKvSTile,
456+ colNumOri,
457+ mm2ToReFlag);
458+ } else {
459+ Arch::CrossCoreWaitFlag<4, PIPE_V>(mm2ToReFlag);
460+ Arch::CrossCoreSetFlag<4, PIPE_V>(mm2ToReFlag);
461+ }
462+ }
463+private:
464+ AscendC::LocalTensor<ElementOTmp> loUbTensor[UB_OTMP_BUF_STAGES];
465+ AscendC::LocalTensor<SMDtype> dmUbTensor16;
466+ AscendC::LocalTensor<SMDtype> glUbTensor16;
467+ AscendC::LocalTensor<float> dmUbTensor32;
468+ AscendC::LocalTensor<float> glUbTensor32;
469+ AscendC::LocalTensor<float> scaleTensor;
470+ AscendC::LocalTensor<ElementO> goUbTensor16;
471+ AscendC::LocalTensor<ElementOTmp> goUbTensor32;
472+ 
473+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
474+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
475+ uint32_t rowOffsetCurSubCore = 0;
476+ uint32_t rowNumCurSubCore = 0;
477+};
478+}
479+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/block/block_epilogue_rescale_o_low_prec.hpp+463-0
@@ -0,0 +1,463 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_LOW_PREC_O_HPP
12+#define EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_LOW_PREC_O_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/resource.hpp"
16+#include "../../../attn_infra/epilogue/dispatch_policy.hpp"
17+#include "../../../attn_infra/epilogue/tile_common/tile_copy.hpp"
18+#include "../../../attn_infra/gemm_coord.hpp"
19+#include "../../../attn_infra/matrix_coord.hpp"
20+ 
21+namespace NpuArch::Epilogue::Block {
22+ 
23+template <
24+ class OutputType_,
25+ class InputType_,
26+ class UpdateType_,
27+ class LseType_,
28+ LseMode LSE_MODE_>
29+class BlockEpilogue<
30+ EpilogueAtlasA2RescaleO<LSE_MODE_, half>,
31+ OutputType_,
32+ InputType_,
33+ UpdateType_,
34+ LseType_>
35+{
36+public:
37+ // Type aliases
38+ using DispatchPolicy = EpilogueAtlasA2RescaleO<LSE_MODE_, half>;
39+ using ArchTag = typename DispatchPolicy::ArchTag;
40+ 
41+ using ElementOutput = typename OutputType_::Element;
42+ using ElementInput = typename InputType_::Element;
43+ using ElementUpdate = typename UpdateType_::Element;
44+ using ElementLse = typename LseType_::Element;
45+ 
46+ using LayoutOutput = typename OutputType_::Layout;
47+ using LayoutInput = typename InputType_::Layout;
48+ using LayoutUpdate = typename UpdateType_::Layout;
49+ using LayoutLse = typename LseType_::Layout;
50+ 
51+ static constexpr LseMode LSE_MODE = DispatchPolicy::LSE_MODE;
52+ 
53+ static constexpr uint32_t HALF_ELENUM_PER_BLK = 16;
54+ static constexpr uint32_t BLOCK_SIZE = 16;
55+ static constexpr uint32_t HALF_ELENUM_PER_VECCALC = 128;
56+ static constexpr uint32_t FLOAT_ELENUM_PER_VECCALC = 64;
57+ static constexpr uint32_t HALF_ELENUM_PER_LINE = 256;
58+ static constexpr uint32_t FLOAT_ELENUM_PER_LINE = 128;
59+ static constexpr uint32_t MULTIPLIER = 2;
60+ static constexpr uint32_t FLOAT_BLOCK_SIZE = 8;
61+ static constexpr uint32_t HALF_BLOCK_SIZE = 16;
62+ static constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
63+ static constexpr uint32_t HALF_VECTOR_SIZE = 128;
64+ static constexpr uint32_t UB_UINT8_VECTOR_SIZE = 1024;
65+ static constexpr uint32_t UB_UINT8_BLOCK_SIZE = 16384;
66+ static constexpr uint32_t HALF_DM_UB_SIZE = 64;
67+ static constexpr uint32_t HALF_LL_UB_SIZE = 256;
68+ static constexpr uint32_t VECTOR_SIZE = 128;
69+ static constexpr uint32_t NUM4 = 4;
70+ static constexpr uint32_t MAX_UB_O_ELEM_NUM = 8192;
71+ static constexpr uint32_t MAX_ROW_NUM_SUB_CORE = 256;
72+ static constexpr uint32_t SIZE_OF_16BIT = 2;
73+ 
74+ __aicore__ inline
75+ BlockEpilogue(Arch::Resource<ArchTag> &resource)
76+ {
77+ // Allocate UB space
78+ constexpr uint32_t LO_UB_TENSOR_OFFSET = 8 * UB_UINT8_BLOCK_SIZE;
79+ constexpr uint32_t GO_UB_TENSOR_OFFSET = 9 * UB_UINT8_BLOCK_SIZE;
80+ constexpr uint32_t TV_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE;
81+ 
82+ constexpr uint32_t HM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 9 * UB_UINT8_VECTOR_SIZE;
83+ constexpr uint32_t GM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE;
84+ constexpr uint32_t LSE32_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 10 * UB_UINT8_VECTOR_SIZE;
85+ constexpr uint32_t GL_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
86+ constexpr uint32_t LSE16_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 12 * UB_UINT8_VECTOR_SIZE;
87+ constexpr uint32_t DM_UB_TENSOR_OFFSET = 10 * UB_UINT8_BLOCK_SIZE + 13 * UB_UINT8_VECTOR_SIZE;
88+ 
89+ loUbTensor = resource.ubBuf.template GetBufferByByte<half>(LO_UB_TENSOR_OFFSET);
90+ dmUbTensor = resource.ubBuf.template GetBufferByByte<half>(DM_UB_TENSOR_OFFSET);
91+ glUbTensor = resource.ubBuf.template GetBufferByByte<half>(GL_UB_TENSOR_OFFSET);
92+ tvUbTensor = resource.ubBuf.template GetBufferByByte<half>(TV_UB_TENSOR_OFFSET);
93+ tvUbTensor32 = resource.ubBuf.template GetBufferByByte<float>(TV_UB_TENSOR_OFFSET);
94+ goUbTensor = resource.ubBuf.template GetBufferByByte<ElementOutput>(GO_UB_TENSOR_OFFSET);
95+ hmUbTensor = resource.ubBuf.template GetBufferByByte<half>(HM_UB_TENSOR_OFFSET);
96+ gmUbTensor = resource.ubBuf.template GetBufferByByte<half>(GM_UB_TENSOR_OFFSET);
97+ lse16_ubuf_tensor = resource.ubBuf.template GetBufferByByte<half>(LSE16_UB_TENSOR_OFFSET);
98+ lse32_ubuf_tensor = resource.ubBuf.template GetBufferByByte<float>(LSE32_UB_TENSOR_OFFSET);
99+ }
100+ 
101+ __aicore__ inline
102+ ~BlockEpilogue() {}
103+ 
104+ __aicore__ inline
105+ void SetMask(int32_t len)
106+ {
107+ const int32_t MAX_MASK_LEN = 128;
108+ const int32_t HALF_MASK_LEN = 64;
109+ if (len >= MAX_MASK_LEN) {
110+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
111+ return;
112+ }
113+ int32_t highMask = len - HALF_MASK_LEN > 0 ? len - HALF_MASK_LEN : 0;
114+ int32_t lowMask = len - HALF_MASK_LEN >= 0 ? HALF_MASK_LEN : len;
115+ if (len < HALF_MASK_LEN) {
116+ AscendC::SetVectorMask<int8_t>(0x0, ((uint64_t)1 << lowMask) - 1);
117+ } else {
118+ AscendC::SetVectorMask<int8_t>(((uint64_t)1 << highMask) - 1, 0xffffffffffffffff);
119+ }
120+ }
121+ 
122+ __aicore__ inline
123+ void CopyOToGm(AscendC::GlobalTensor<ElementOutput> gOutput, uint32_t proTokenIdx, uint32_t proTokenNum,
124+ uint32_t epiTokenNum, uint32_t integralHeadNum, uint32_t qSThisSubBlock, uint32_t embed, uint32_t oHiddenSize)
125+ {
126+ uint32_t innerOGmOffset = 0;
127+ uint32_t innerGOUbOffset = 0;
128+ if (proTokenNum != 0U) {
129+ AscendC::DataCopyPad(
130+ gOutput[innerOGmOffset + proTokenIdx * oHiddenSize],
131+ goUbTensor[innerGOUbOffset],
132+ AscendC::DataCopyExtParams(
133+ proTokenNum, embed * SIZE_OF_16BIT, 0, (oHiddenSize - embed) * SIZE_OF_16BIT, 0));
134+ innerOGmOffset += embed;
135+ innerGOUbOffset += proTokenNum * embed;
136+ }
137+ for (uint32_t qN_idx = 0; qN_idx < integralHeadNum; qN_idx++) {
138+ AscendC::DataCopyPad(
139+ gOutput[innerOGmOffset],
140+ goUbTensor[innerGOUbOffset],
141+ AscendC::DataCopyExtParams(
142+ qSThisSubBlock, embed * SIZE_OF_16BIT, 0, (oHiddenSize - embed) * SIZE_OF_16BIT, 0));
143+ innerOGmOffset += embed;
144+ innerGOUbOffset += qSThisSubBlock * embed;
145+ }
146+ if (epiTokenNum != 0U) {
147+ AscendC::DataCopyPad(
148+ gOutput[innerOGmOffset],
149+ goUbTensor[innerGOUbOffset],
150+ AscendC::DataCopyExtParams(
151+ epiTokenNum, embed * SIZE_OF_16BIT, 0, (oHiddenSize - embed) * SIZE_OF_16BIT, 0));
152+ }
153+ }
154+ 
155+ __aicore__ inline
156+ void SubCoreCompute(
157+ AscendC::GlobalTensor<ElementOutput> gOutput,
158+ AscendC::GlobalTensor<ElementInput> gInput,
159+ AscendC::GlobalTensor<ElementUpdate> gUpdate,
160+ AscendC::GlobalTensor<ElementLse> gLse,
161+ const LayoutOutput &layoutOutput,
162+ const LayoutInput &layoutInput,
163+ const LayoutUpdate &layoutUpdate,
164+ const LayoutLse &layoutLse,
165+ uint32_t qNThisSubBlock, uint32_t qSThisSubBlock, uint32_t totalRowNum,
166+ uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod,
167+ uint32_t needRowLoop, uint32_t isLastRowLoop, uint32_t rowOffsetLoop,
168+ uint32_t proTokenIdx, uint32_t proTokenNum, uint32_t epiTokenNum, uint32_t integralHeadNum)
169+ {
170+ uint32_t curRowNum = layoutInput.shape(0);
171+ uint32_t embed = layoutInput.shape(1);
172+ uint32_t embedRound = layoutInput.stride(0);
173+ uint32_t curRowNumRound = RoundUp(curRowNum, HALF_BLOCK_SIZE);
174+ uint32_t qSBlockSize = layoutOutput.shape(0);
175+ uint32_t oHiddenSize = layoutOutput.shape(1);
176+ uint32_t stride = layoutLse.shape(1);
177+ uint32_t dmUbOffsetCurStackTile = curStackTileMod * MAX_ROW_NUM_SUB_CORE + rowOffsetLoop;
178+ 
179+ if (!isFirstStackTile) {
180+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
181+ AscendC::DataCopy(
182+ loUbTensor, gInput, AscendC::DataCopyParams(1, curRowNum * embedRound / HALF_BLOCK_SIZE, 0, 0));
183+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
184+ }
185+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID6);
186+ if (!isFirstStackTile) {
187+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
188+ AscendC::Brcb(
189+ tvUbTensor.ReinterpretCast<uint16_t>(),
190+ dmUbTensor[dmUbOffsetCurStackTile].ReinterpretCast<uint16_t>(),
191+ curRowNumRound / FLOAT_BLOCK_SIZE,
192+ AscendC::BrcbRepeatParams(1, 8));
193+ AscendC::PipeBarrier<PIPE_V>();
194+ if (needRowLoop) {
195+ AscendC::DataCopy(
196+ goUbTensor, gUpdate,
197+ AscendC::DataCopyParams(1, curRowNum * embedRound / HALF_BLOCK_SIZE, 0, 0));
198+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1);
199+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID1);
200+ }
201+ // *** go = go * dm_block
202+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
203+ for (uint32_t vmul_idx = 0; vmul_idx < embed / HALF_VECTOR_SIZE; ++vmul_idx) {
204+ AscendC::Mul<half, false>(
205+ goUbTensor[vmul_idx * HALF_VECTOR_SIZE],
206+ goUbTensor[vmul_idx * HALF_VECTOR_SIZE],
207+ tvUbTensor,
208+ (uint64_t)0,
209+ curRowNum,
210+ AscendC::BinaryRepeatParams(
211+ 1, 1, 0, embedRound / HALF_BLOCK_SIZE, embedRound / HALF_BLOCK_SIZE, 1));
212+ }
213+ if (embed % HALF_VECTOR_SIZE > 0) {
214+ SetMask(embed % HALF_VECTOR_SIZE);
215+ AscendC::Mul<half, false>(
216+ goUbTensor[embed / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE],
217+ goUbTensor[embed / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE],
218+ tvUbTensor,
219+ (uint64_t)0,
220+ curRowNum,
221+ AscendC::BinaryRepeatParams(
222+ 1, 1, 0, embedRound / HALF_BLOCK_SIZE, embedRound / HALF_BLOCK_SIZE, 1));
223+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
224+ }
225+ AscendC::PipeBarrier<PIPE_V>();
226+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
227+ // *** go = lo + go
228+ AscendC::Add<half, false>(
229+ goUbTensor,
230+ goUbTensor,
231+ loUbTensor,
232+ (uint64_t)0,
233+ (curRowNum * embedRound + HALF_VECTOR_SIZE - 1) / HALF_VECTOR_SIZE,
234+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
235+ AscendC::PipeBarrier<PIPE_V>();
236+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
237+ } else {
238+ // *** go = lo
239+ AscendC::DataCopy(
240+ goUbTensor, gInput, AscendC::DataCopyParams(1, curRowNum * embedRound / HALF_BLOCK_SIZE, 0, 0));
241+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
242+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(EVENT_ID0);
243+ }
244+ 
245+ if (isLastStackTile) {
246+ // *** gl_block = expand_to_block(gl), 存放于 tv
247+ AscendC::Brcb(
248+ tvUbTensor.ReinterpretCast<uint16_t>(),
249+ glUbTensor.ReinterpretCast<uint16_t>()[rowOffsetLoop],
250+ curRowNumRound / FLOAT_BLOCK_SIZE,
251+ AscendC::BrcbRepeatParams(1, 8));
252+ AscendC::PipeBarrier<PIPE_V>();
253+ // *** go = go / gl_block
254+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
255+ for (uint32_t vdiv_idx = 0; vdiv_idx < embed / HALF_VECTOR_SIZE; ++vdiv_idx) {
256+ AscendC::Div<half, false>(
257+ goUbTensor[vdiv_idx * HALF_VECTOR_SIZE],
258+ goUbTensor[vdiv_idx * HALF_VECTOR_SIZE],
259+ tvUbTensor,
260+ (uint64_t)0,
261+ curRowNum,
262+ AscendC::BinaryRepeatParams(
263+ 1, 1, 0, embedRound / HALF_BLOCK_SIZE, embedRound / HALF_BLOCK_SIZE, 1));
264+ }
265+ if (embed % HALF_VECTOR_SIZE > 0) {
266+ SetMask(embed % HALF_VECTOR_SIZE);
267+ AscendC::Div<half, false>(
268+ goUbTensor[embed / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE],
269+ goUbTensor[embed / HALF_VECTOR_SIZE * HALF_VECTOR_SIZE],
270+ tvUbTensor,
271+ (uint64_t)0,
272+ curRowNum,
273+ AscendC::BinaryRepeatParams(
274+ 1, 1, 0, embedRound / HALF_BLOCK_SIZE, embedRound / HALF_BLOCK_SIZE, 1));
275+ AscendC::SetVectorMask<int8_t>((uint64_t)-1, (uint64_t)-1);
276+ }
277+ AscendC::PipeBarrier<PIPE_V>();
278+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
279+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID0);
280+ 
281+ // ***move O to GM
282+ CopyOToGm(
283+ gOutput, proTokenIdx, proTokenNum, epiTokenNum, integralHeadNum, qSThisSubBlock, embed, oHiddenSize);
284+ if constexpr (LSE_MODE == LseMode::OUT_ONLY) {
285+ if (isLastRowLoop) {
286+ AscendC::PipeBarrier<PIPE_V>();
287+ AscendC::Ln<half, false>(
288+ lse16_ubuf_tensor,
289+ glUbTensor,
290+ (uint64_t)0,
291+ CeilDiv(totalRowNum, HALF_VECTOR_SIZE),
292+ AscendC::UnaryRepeatParams(1, 1, 8, 8));
293+ AscendC::PipeBarrier<PIPE_V>();
294+ AscendC::Add<half, false>(
295+ lse16_ubuf_tensor,
296+ lse16_ubuf_tensor,
297+ gmUbTensor,
298+ (uint64_t)0,
299+ CeilDiv(totalRowNum, HALF_VECTOR_SIZE),
300+ AscendC::BinaryRepeatParams(1, 1, 1, 8, 8, 8));
301+ AscendC::PipeBarrier<PIPE_V>();
302+ AscendC::Cast<float, half, false>(
303+ lse32_ubuf_tensor,
304+ lse16_ubuf_tensor,
305+ AscendC::RoundMode::CAST_NONE,
306+ (uint64_t)0,
307+ CeilDiv(totalRowNum, FLOAT_VECTOR_SIZE),
308+ AscendC::UnaryRepeatParams(1, 1, 8, 4));
309+ AscendC::PipeBarrier<PIPE_V>();
310+ 
311+ // *** lse_block = expand_to_block(lse), 存放于 tv
312+ AscendC::Brcb(
313+ tvUbTensor32.ReinterpretCast<uint32_t>(),
314+ lse32_ubuf_tensor.ReinterpretCast<uint32_t>(),
315+ CeilDiv(totalRowNum, FLOAT_BLOCK_SIZE),
316+ AscendC::BrcbRepeatParams(1, 8));
317+ AscendC::PipeBarrier<PIPE_V>();
318+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID4);
319+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID4);
320+
321+ if (qNThisSubBlock == 0U) {
322+ AscendC::DataCopyPad(
323+ gLse, tvUbTensor32,
324+ AscendC::DataCopyExtParams(
325+ totalRowNum, sizeof(float), 0, (stride - 1) * sizeof(float), 0));
326+ } else {
327+ for (uint32_t qNIdx = 0; qNIdx < qNThisSubBlock; qNIdx++) {
328+ AscendC::DataCopyPad(
329+ gLse[qNIdx],
330+ tvUbTensor32[qNIdx * qSBlockSize * FLOAT_BLOCK_SIZE],
331+ AscendC::DataCopyExtParams(
332+ qSBlockSize, sizeof(float), 0, (stride - 1) * sizeof(float), 0));
333+ }
334+ }
335+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
336+ }
337+ }
338+ } else if (needRowLoop) {
339+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID5);
340+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(EVENT_ID5);
341+ AscendC::DataCopy(
342+ gUpdate, goUbTensor, AscendC::DataCopyParams(1, curRowNum * embedRound / HALF_BLOCK_SIZE, 0, 0));
343+ }
344+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID6);
345+ }
346+ 
347+ __aicore__ inline
348+ void operator()(
349+ AscendC::GlobalTensor<ElementOutput> gOutput,
350+ AscendC::GlobalTensor<ElementInput> gInput,
351+ AscendC::GlobalTensor<ElementUpdate> gUpdate,
352+ AscendC::GlobalTensor<ElementLse> gLse,
353+ const LayoutOutput &layoutOutput,
354+ const LayoutInput &layoutInput,
355+ const LayoutUpdate &layoutUpdate,
356+ const LayoutLse &layoutLse,
357+ GemmCoord actualBlockShape,
358+ uint32_t qSBlockSize, uint32_t qNBlockSize,
359+ uint32_t isFirstStackTile, uint32_t isLastStackTile, uint32_t curStackTileMod)
360+ {
361+ uint32_t rowNum = actualBlockShape.m();
362+ uint32_t embed = actualBlockShape.n();
363+ uint32_t maxRowNumPerLoop = MAX_UB_O_ELEM_NUM / embed;
364+ uint32_t rowNumTile = RoundDown(maxRowNumPerLoop, HALF_BLOCK_SIZE);
365+ 
366+ uint32_t subBlockIdx = AscendC::GetSubBlockIdx();
367+ uint32_t subBlockNum = AscendC::GetSubBlockNum();
368+ 
369+ uint32_t qNSplitSubBlock = qNBlockSize / subBlockNum;
370+ uint32_t qNThisSubBlock = (qNBlockSize == 1U) ? 0
371+ : (subBlockIdx == 1U) ? (qNBlockSize - qNSplitSubBlock)
372+ : qNSplitSubBlock;
373+ uint32_t inRowSplitSubBlock =
374+ (qNBlockSize == 1U) ? (qSBlockSize / subBlockNum) : (qSBlockSize * qNSplitSubBlock);
375+ uint32_t inRowActualThisSubBlock = (subBlockIdx == 1U) ? (rowNum - inRowSplitSubBlock) : inRowSplitSubBlock;
376+ uint32_t inRowOffsetThisSubBlock = subBlockIdx * inRowSplitSubBlock;
377+ uint32_t outRowOffsetThisSubBlock = (qNBlockSize == 1U) ? inRowOffsetThisSubBlock : 0;
378+ uint32_t outColOffsetThisSubBlock = (qNBlockSize == 1U) ? 0 : subBlockIdx * qNSplitSubBlock * embed;
379+ uint32_t qSThisSubBlock = (qNBlockSize == 1U) ? inRowActualThisSubBlock : qSBlockSize;
380+ int64_t outOffsetSubBlock =
381+ layoutOutput.GetOffset(MatrixCoord(outRowOffsetThisSubBlock, outColOffsetThisSubBlock));
382+
383+ uint32_t outLseRowOffsetThisSubBlock = (qNBlockSize == 1U) ? inRowOffsetThisSubBlock : 0;
384+ uint32_t outLseColOffsetThisSubBlock = (qNBlockSize == 1U) ? 0 : subBlockIdx * qNSplitSubBlock;
385+ int64_t offsetLse = layoutLse.GetOffset(MatrixCoord(outLseRowOffsetThisSubBlock, outLseColOffsetThisSubBlock));
386+ auto gLseThisSubBlock = gLse[offsetLse];
387+
388+ if (inRowActualThisSubBlock > 0U) {
389+ uint32_t rowLoop = CeilDiv(inRowActualThisSubBlock, rowNumTile);
390+ uint32_t needRowLoop = (rowLoop > 1U) ? 1 : 0;
391+ 
392+ // The rows of each cycle consist of multiple heads with several tokens.
393+ // There are several integral heads, one prologue head, one epilogue head.
394+ uint32_t proTokenIdx = 0; // the token idx of the start token of the prologue part
395+ uint32_t proTokenIdxPre = 0; // the token idx of the start token of the pre prologue part
396+ uint32_t proTokenNum = 0; // the token num of the prologue part
397+ uint32_t epiTokenNum = 0; // the token num of the epilogue part
398+ uint32_t integralHeadNum = 0; // the number of integral heads within a cycle
399+ uint32_t qSRemian = qSThisSubBlock;
400+ for (uint32_t rowLoopIdx = 0; rowLoopIdx < rowLoop; rowLoopIdx++) {
401+ uint32_t rowOffsetLoop = rowLoopIdx * rowNumTile;
402+ uint32_t rowOffsetCurLoop = inRowOffsetThisSubBlock + rowOffsetLoop;
403+ uint32_t rowActualCurLoop =
404+ (rowLoopIdx == (rowLoop - 1U)) ? inRowActualThisSubBlock - rowLoopIdx * rowNumTile : rowNumTile;
405+ 
406+ int64_t offsetOutput =
407+ static_cast<int64_t>(rowLoopIdx * rowNumTile / qSThisSubBlock * embed) + outOffsetSubBlock;
408+ auto gOutputCurLoop = gOutput[offsetOutput];
409+ auto layoutOutputCurLoop = layoutOutput;
410+ int64_t offsetInput = layoutInput.GetOffset(MatrixCoord(rowOffsetCurLoop, 0));
411+ auto gInputCurLoop = gInput[offsetInput];
412+ auto layoutInputCurLoop = layoutInput.GetTileLayout(MatrixCoord(rowActualCurLoop, embed));
413+ 
414+ int64_t offsetUpdate = layoutUpdate.GetOffset(MatrixCoord(rowOffsetCurLoop, 0));
415+ auto gUpdateCurLoop = gUpdate[offsetUpdate];
416+ auto layoutUpdateCurLoop = layoutUpdate.GetTileLayout(MatrixCoord(rowActualCurLoop, embed));
417+ 
418+ proTokenIdx = rowOffsetLoop % qSThisSubBlock;
419+ proTokenNum = AscendC::Std::min(rowActualCurLoop, (qSThisSubBlock - proTokenIdx)) % qSThisSubBlock;
420+ integralHeadNum = (rowActualCurLoop - proTokenNum) / qSThisSubBlock;
421+ epiTokenNum = rowActualCurLoop - proTokenNum - integralHeadNum * qSThisSubBlock;
422+ 
423+ SubCoreCompute(
424+ gOutputCurLoop,
425+ gInputCurLoop,
426+ gUpdateCurLoop,
427+ gLseThisSubBlock,
428+ layoutOutputCurLoop,
429+ layoutInputCurLoop,
430+ layoutUpdateCurLoop,
431+ layoutLse,
432+ qNThisSubBlock,
433+ qSThisSubBlock,
434+ inRowActualThisSubBlock,
435+ isFirstStackTile,
436+ isLastStackTile,
437+ curStackTileMod,
438+ needRowLoop,
439+ (rowLoopIdx == rowLoop - 1U),
440+ rowOffsetLoop,
441+ proTokenIdx,
442+ proTokenNum,
443+ epiTokenNum,
444+ integralHeadNum);
445+ }
446+ }
447+ }
448+ 
449+private:
450+ AscendC::LocalTensor<half> loUbTensor;
451+ AscendC::LocalTensor<half> dmUbTensor;
452+ AscendC::LocalTensor<half> hmUbTensor;
453+ AscendC::LocalTensor<half> glUbTensor;
454+ AscendC::LocalTensor<half> tvUbTensor;
455+ AscendC::LocalTensor<float> tvUbTensor32;
456+ AscendC::LocalTensor<ElementOutput> goUbTensor;
457+ AscendC::LocalTensor<half> gmUbTensor;
458+ AscendC::LocalTensor<half> lse16_ubuf_tensor;
459+ AscendC::LocalTensor<float> lse32_ubuf_tensor;
460+};
461+}
462+ 
463+#endif // EPILOGUE_BLOCK_BLOCK_EPILOGUE_RESCALE_LOW_PREC_O_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/dispatch_policy.hpp+53-0
@@ -0,0 +1,53 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_DISPATCH_POLICY_HPP
12+#define EPILOGUE_DISPATCH_POLICY_HPP
13+ 
14+#include "../../attn_infra/base_defs.hpp"
15+#include "../../attn_infra/arch/arch.hpp"
16+ 
17+namespace NpuArch::Epilogue
18+{
19+ 
20+enum class LseMode {NONE = 0, OUT_ONLY = 1};
21+// For AtlasA2, FA Infer online Softmax
22+template <LseMode LSE_MODE_, typename SM_DTYPE_>
23+struct EpilogueAtlasA2OnlineSoftmax {
24+ using ArchTag = Arch::AtlasA2;
25+ using IntermPrec = SM_DTYPE_;
26+ static constexpr LseMode LSE_MODE = LSE_MODE_;
27+};
28+ 
29+// For AtlasA2, FA Infer RescaleO
30+template <LseMode LSE_MODE_, typename SM_DTYPE_>
31+struct EpilogueAtlasA2RescaleO {
32+ using ArchTag = Arch::AtlasA2;
33+ using IntermPrec = SM_DTYPE_;
34+ static constexpr LseMode LSE_MODE = LSE_MODE_;
35+};
36+ 
37+// For AtlasA5
38+struct EpilogueBsaMask2Idx {
39+ static constexpr uint32_t IO_STAGES = 2;
40+ using ArchTag = Arch::AtlasA5;
41+};
42+ 
43+struct EpilogueOnlineSoftmaxBsa {
44+ using ArchTag = Arch::AtlasA5;
45+};
46+ 
47+struct EpilogueAtlasA5BsaRescaleO {
48+ using ArchTag = Arch::AtlasA5;
49+};
50+ 
51+} // namespace NpuArch::Epilogue
52+ 
53+#endif // EPILOGUE_DISPATCH_POLICY_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/tile_common/copy_gm_to_ub.hpp+188-0
@@ -0,0 +1,188 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_TILE_TILE_COPY_GM_TO_UB_HPP
12+#define EPILOGUE_TILE_TILE_COPY_GM_TO_UB_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/layout/layout.hpp"
17+#include "../../../attn_infra/gemm/gemm_type.hpp"
18+ 
19+namespace NpuArch::Epilogue::Tile
20+{
21+ 
22+template <
23+ class ArchTag,
24+ class GmType
25+>
26+struct CopyGm2Ub {
27+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy gm to ub, can not find the specialization.");
28+};
29+ 
30+template <typename Element>
31+struct CopyGm2Ub<Arch::AtlasA2, Gemm::GemmType<Element, layout::RowMajor>> {
32+ using LayoutSrc = layout::RowMajor;
33+ using LayoutDst = layout::RowMajor;
34+ 
35+ static constexpr uint32_t ELE_NUM_PER_BLK = static_cast<uint32_t>(BYTE_PER_BLK) / static_cast<uint32_t>(sizeof(Element));
36+ 
37+ __aicore__ inline
38+ CopyGm2Ub() = default;
39+ 
40+ __aicore__ inline
41+ void operator()(
42+ AscendC::LocalTensor<Element> const &dstTensor,
43+ AscendC::GlobalTensor<Element> const &srcTensor,
44+ layout::RowMajor const &layoutDst,
45+ layout::RowMajor const &layoutSrc)
46+ {
47+ AscendC::DataCopyExtParams dataCopyParams(
48+ layoutSrc.shape(0),
49+ layoutSrc.shape(1) * sizeof(Element),
50+ (layoutSrc.stride(0) - layoutSrc.shape(1)) * sizeof(Element),
51+ (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK,
52+ 0
53+ );
54+ AscendC::DataCopyPadExtParams<Element> padParams(false, 0, 0, 0);
55+ AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams);
56+ };
57+};
58+ 
59+template <typename Element>
60+struct CopyGm2Ub<Arch::AtlasA2, Gemm::GemmType<Element, layout::VectorLayout>> {
61+ using LayoutSrc = layout::VectorLayout;
62+ using LayoutDst = layout::VectorLayout;
63+ 
64+ static constexpr uint32_t ELE_NUM_PER_BLK = static_cast<uint32_t>(BYTE_PER_BLK) / static_cast<uint32_t>(sizeof(Element));
65+ 
66+ __aicore__ inline
67+ CopyGm2Ub() = default;
68+ 
69+ __aicore__ inline
70+ void operator()(
71+ AscendC::LocalTensor<Element> const &dstTensor,
72+ AscendC::GlobalTensor<Element> const &srcTensor,
73+ layout::VectorLayout const &layoutDst,
74+ layout::VectorLayout const &layoutSrc)
75+ {
76+ AscendC::DataCopyExtParams dataCopyParams(
77+ 1,
78+ layoutSrc.shape(0) * sizeof(Element),
79+ 0,
80+ 0,
81+ 0
82+ );
83+ AscendC::DataCopyPadExtParams<Element> padParams(false, 0, 0, 0);
84+ AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams);
85+ };
86+};
87+ 
88+/// @brief This copy instruction used to copy per token scale from GM to UB.
89+/// Copy the scale of shape (m,1) on GM to the first column of shape (m,n) on UB,
90+/// and pad the first block of each row (i.e. pad to shape (m,8) when element type is float).
91+/// @tparam ArchTag: Architecture tag.
92+/// @tparam GmType: Type of data on GM.
93+template <
94+ class ArchTag,
95+ class GmType
96+>
97+struct CopyPerTokenScale2Ub {
98+ static_assert(std::is_same_v<typename GmType::Layout, layout::ColumnMajor>,
99+ "Unsupported layout for CopyPerTokenScale2Ub.");
100+ 
101+ using Element = typename GmType::Element;
102+ using LayoutSrc = typename GmType::Layout;
103+ using LayoutDst = layout::RowMajor;
104+ 
105+ static constexpr uint32_t ELE_NUM_PER_BLK = static_cast<uint32_t>(BYTE_PER_BLK) / static_cast<uint32_t>(sizeof(Element));
106+ 
107+ __aicore__ inline
108+ CopyPerTokenScale2Ub() = default;
109+ 
110+ __aicore__ inline
111+ void operator()(
112+ AscendC::LocalTensor<Element> const &dstTensor,
113+ AscendC::GlobalTensor<Element> const &srcTensor,
114+ LayoutDst const &layoutDst,
115+ LayoutSrc const &layoutSrc)
116+ {
117+ AscendC::DataCopyExtParams dataCopyParams;
118+ AscendC::DataCopyPadExtParams<Element> padParams;
119+ 
120+ dataCopyParams.blockCount = layoutSrc.shape(0);
121+ dataCopyParams.blockLen = layoutSrc.shape(1) * sizeof(Element); // per token scale has only one column
122+ dataCopyParams.srcStride = 0;
123+ dataCopyParams.dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK;
124+ // Pad the data to the complete block
125+ padParams.isPad = true;
126+ padParams.leftPadding = 0;
127+ padParams.rightPadding = 0;
128+ 
129+ AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams, padParams);
130+ }
131+};
132+ 
133+template <
134+ class ArchTag,
135+ class GmType
136+>
137+struct CopyGm2UbAligned {
138+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy gm to ub aligned, can not find the specialization.");
139+};
140+ 
141+template <typename Element>
142+struct CopyGm2UbAligned<Arch::AtlasA2, Gemm::GemmType<Element, layout::RowMajor>> {
143+ using LayoutSrc = layout::RowMajor;
144+ using LayoutDst = layout::RowMajor;
145+ 
146+ static constexpr uint32_t ELE_NUM_PER_BLK = static_cast<uint32_t>(BYTE_PER_BLK) / static_cast<uint32_t>(sizeof(Element));
147+ static constexpr uint32_t BLOCK_LEN_LIMIT = 65536;
148+ static constexpr uint32_t MAX_REPEAT = 4095;
149+ static constexpr uint32_t STRIDE_LIMIT = 65536;
150+ 
151+ __aicore__ inline
152+ CopyGm2UbAligned() = default;
153+ 
154+ __aicore__ inline
155+ void operator()(
156+ AscendC::LocalTensor<Element> const &dstTensor,
157+ AscendC::GlobalTensor<Element> const &srcTensor,
158+ layout::RowMajor const &layoutDst,
159+ layout::RowMajor const &layoutSrc)
160+ {
161+ uint32_t rows = layoutSrc.shape(0);
162+ uint32_t cols = layoutSrc.shape(1);
163+ uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK;
164+ uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK;
165+ 
166+ if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) {
167+ DataCopy(dstTensor, srcTensor, rows * cols);
168+ } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) {
169+ uint32_t rLoops = CeilDiv(rows, MAX_REPEAT);
170+ for (uint32_t i = 0; i < rLoops; ++i) {
171+ uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT;
172+ AscendC::DataCopyParams dataCopyParams(
173+ rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride
174+ );
175+ DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)],
176+ srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams);
177+ }
178+ } else {
179+ for (uint32_t i = 0; i < rows; ++i) {
180+ DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols);
181+ }
182+ }
183+ };
184+};
185+ 
186+} // NpuArch::Epilogue::Tile
187+ 
188+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/tile_common/copy_gm_to_ub_tla.hpp+113-0
@@ -0,0 +1,113 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_TILE_COPY_GM_TO_UB_TLA_HPP
12+#define EPILOGUE_TILE_COPY_GM_TO_UB_TLA_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../tla/tensor.hpp"
17+#include "../../../tla/layout.hpp"
18+ 
19+namespace NpuArch::Epilogue::Tile
20+{
21+ 
22+template <
23+ class ArchTag,
24+ class TensorSrc,
25+ class TensorDst,
26+ class Enable = void
27+>
28+struct CopyGm2UbTla {
29+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported CopyGm2UbTla, can not find the specialization.");
30+};
31+ 
32+/// Partial specialization for AtlasA2, RowMajor in and RowMajor out.
33+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
34+struct CopyGm2UbTla<Arch::AtlasA2,
35+ tla::Tensor<AscendC::GlobalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::GM>,
36+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::VECCALC>,
37+ std::enable_if_t<tla::detail::isRowMajor<LayoutSrc>::value &&
38+ tla::detail::isRowMajor<LayoutDst>::value>> {
39+ static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(ElementSrc);
40+ 
41+ // Methods
42+ 
43+ __aicore__ inline
44+ CopyGm2UbTla() = default;
45+ 
46+ template <class TensorDst, class TensorSrc>
47+ __aicore__ inline
48+ void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor)
49+ {
50+ static_assert(tla::detail::isRowMajor<typename TensorSrc::Layout>::value &&
51+ tla::detail::isRowMajor<typename TensorDst::Layout>::value &&
52+ TensorSrc::position == AscendC::TPosition::GM &&
53+ TensorDst::position == AscendC::TPosition::VECCALC,
54+ "The input parameters do not match. TensorSrc must be GM and RowMajor, "
55+ "while TensorDst must be UB and RowMajor");
56+ 
57+ AscendC::DataCopyExtParams dataCopyParams(
58+ tla::get<0>(srcTensor.shape()),
59+ tla::get<1>(srcTensor.shape()) * sizeof(ElementSrc),
60+ (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) * sizeof(ElementSrc),
61+ (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) / ELE_NUM_PER_BLK,
62+ 0
63+ );
64+ AscendC::DataCopyPadExtParams<ElementSrc> padParams(false, 0, 0, 0);
65+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
66+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
67+ AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams, padParams);
68+ };
69+};
70+ 
71+/// Partial specialization for AtlasA5, RowMajor in and RowMajor out.
72+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
73+struct CopyGm2UbTla<Arch::AtlasA5,
74+ tla::Tensor<AscendC::GlobalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::GM>,
75+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::VECCALC>,
76+ std::enable_if_t<tla::detail::isRowMajor<LayoutSrc>::value &&
77+ tla::detail::isRowMajor<LayoutDst>::value>> {
78+ static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(ElementSrc);
79+ 
80+ // Methods
81+ 
82+ __aicore__ inline
83+ CopyGm2UbTla() = default;
84+ 
85+ template <class TensorDst, class TensorSrc>
86+ __aicore__ inline
87+ void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor)
88+ {
89+ static_assert(tla::detail::isRowMajor<typename TensorSrc::Layout>::value &&
90+ tla::detail::isRowMajor<typename TensorDst::Layout>::value &&
91+ TensorSrc::position == AscendC::TPosition::GM &&
92+ TensorDst::position == AscendC::TPosition::VECCALC,
93+ "The input parameters do not match. TensorSrc must be GM and RowMajor, "
94+ "while TensorDst must be UB and RowMajor");
95+ 
96+ AscendC::DataCopyExtParams dataCopyParams(
97+ tla::get<0>(srcTensor.shape()),
98+ tla::get<1>(srcTensor.shape()) * sizeof(ElementSrc),
99+ (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) * sizeof(ElementSrc),
100+ (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) / ELE_NUM_PER_BLK,
101+ 0
102+ );
103+ AscendC::DataCopyPadExtParams<ElementSrc> padParams(false, 0, 0, 0);
104+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
105+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
106+ AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams, padParams);
107+ };
108+};
109+ 
110+}
111+ 
112+#endif // EPILOGUE_TILE_COPY_GM_TO_UB_TLA_HPP
113+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/tile_common/copy_ub_to_gm.hpp+144-0
@@ -0,0 +1,144 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_TILE_TILE_COPY_UB_TO_GM_HPP
12+#define EPILOGUE_TILE_TILE_COPY_UB_TO_GM_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/layout/layout.hpp"
17+#include "../../../attn_infra/gemm/gemm_type.hpp"
18+ 
19+namespace NpuArch::Epilogue::Tile
20+{
21+ 
22+template <
23+ class ArchTag,
24+ class GmType
25+>
26+struct CopyUb2Gm {
27+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy ub to gm, can not find the specialization.");
28+};
29+ 
30+template <typename Element>
31+struct CopyUb2Gm<Arch::AtlasA2, Gemm::GemmType<Element, layout::RowMajor>> {
32+ using LayoutDst = layout::RowMajor;
33+ using LayoutSrc = layout::RowMajor;
34+ 
35+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
36+ 
37+ __aicore__ inline
38+ CopyUb2Gm() = default;
39+ 
40+ __aicore__ inline
41+ void operator()(
42+ AscendC::GlobalTensor<Element> const &dstTensor,
43+ AscendC::LocalTensor<Element> const &srcTensor,
44+ layout::RowMajor const &layoutDst,
45+ layout::RowMajor const &layoutSrc)
46+ {
47+ AscendC::DataCopyExtParams dataCopyParams(
48+ layoutDst.shape(0),
49+ layoutDst.shape(1) * sizeof(Element),
50+ (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_C0,
51+ (layoutDst.stride(0) - layoutDst.shape(1)) * sizeof(Element),
52+ 0
53+ );
54+ AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams);
55+ }
56+};
57+ 
58+ 
59+// new add vectorlayout version
60+template <typename Element>
61+struct CopyUb2Gm<Arch::AtlasA2, Gemm::GemmType<Element, layout::VectorLayout>> {
62+ using LayoutSrc = layout::VectorLayout;
63+ using LayoutDst = layout::VectorLayout;
64+ 
65+ static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element);
66+ 
67+ __aicore__ inline
68+ CopyUb2Gm() = default;
69+ 
70+ __aicore__ inline
71+ void operator()(
72+ AscendC::GlobalTensor<Element> const &dstTensor,
73+ AscendC::LocalTensor<Element> const &srcTensor,
74+ layout::VectorLayout const &layoutDst,
75+ layout::VectorLayout const &layoutSrc)
76+ {
77+ AscendC::DataCopyExtParams dataCopyParams(
78+ 1,
79+ layoutDst.shape(0) * sizeof(Element),
80+ 0,
81+ 0,
82+ 0
83+ );
84+ AscendC::DataCopyPad(dstTensor, srcTensor, dataCopyParams);
85+ };
86+};
87+ 
88+ 
89+template <
90+ class ArchTag,
91+ class GmType
92+>
93+struct CopyUb2GmAligned {
94+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy ub to gm aligned, can not find the specialization.");
95+};
96+ 
97+template <typename Element>
98+struct CopyUb2GmAligned<Arch::AtlasA2, Gemm::GemmType<Element, layout::RowMajor>> {
99+ using LayoutSrc = layout::RowMajor;
100+ using LayoutDst = layout::RowMajor;
101+ 
102+ static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element);
103+ static constexpr uint32_t BLOCK_LEN_LIMIT = 65536;
104+ static constexpr uint32_t MAX_REPEAT = 4095;
105+ static constexpr uint32_t STRIDE_LIMIT = 65536;
106+ 
107+ __aicore__ inline
108+ CopyUb2GmAligned() = default;
109+ 
110+ __aicore__ inline
111+ void operator()(
112+ AscendC::GlobalTensor<Element> const &dstTensor,
113+ AscendC::LocalTensor<Element> const &srcTensor,
114+ layout::RowMajor const &layoutDst,
115+ layout::RowMajor const &layoutSrc)
116+ {
117+ uint32_t rows = layoutDst.shape(0);
118+ uint32_t cols = layoutDst.shape(1);
119+ uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK;
120+ uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK;
121+ 
122+ if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) {
123+ DataCopy(dstTensor, srcTensor, rows * cols);
124+ } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) {
125+ uint32_t rLoops = CeilDiv(rows, MAX_REPEAT);
126+ for (uint32_t i = 0; i < rLoops; ++i) {
127+ uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT;
128+ AscendC::DataCopyParams dataCopyParams(
129+ rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride
130+ );
131+ DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)],
132+ srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams);
133+ }
134+ } else {
135+ for (uint32_t i = 0; i < rows; ++i) {
136+ DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols);
137+ }
138+ }
139+ };
140+};
141+ 
142+} // NpuArch::Epilogue::Tile
143+ 
144+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/tile_common/copy_ub_to_gm_tla.hpp+110-0
@@ -0,0 +1,110 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_TILE_COPY_UB_TO_GM_TLA_HPP
12+#define EPILOGUE_TILE_COPY_UB_TO_GM_TLA_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../tla/tensor.hpp"
17+#include "../../../tla/layout.hpp"
18+ 
19+namespace NpuArch::Epilogue::Tile {
20+ 
21+template <
22+ class ArchTag,
23+ class TensorSrc,
24+ class TensorDst,
25+ class Enable = void
26+>
27+struct CopyUb2GmTla {
28+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported CopyUb2GmTla, can not find the specialization.");
29+};
30+ 
31+/// Partial specialization for AtlasA2, RowMajor in and RowMajor out.
32+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
33+struct CopyUb2GmTla<Arch::AtlasA2,
34+ tla::Tensor<AscendC::LocalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::VECCALC>,
35+ tla::Tensor<AscendC::GlobalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::GM>,
36+ std::enable_if_t<tla::detail::isRowMajor<LayoutSrc>::value &&
37+ tla::detail::isRowMajor<LayoutDst>::value>> {
38+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
39+ 
40+ // Methods
41+ 
42+ __aicore__ inline
43+ CopyUb2GmTla() = default;
44+ 
45+ template <class TensorDst, class TensorSrc>
46+ __aicore__ inline
47+ void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor)
48+ {
49+ static_assert(tla::detail::isRowMajor<typename TensorSrc::Layout>::value &&
50+ tla::detail::isRowMajor<typename TensorDst::Layout>::value &&
51+ TensorSrc::position == AscendC::TPosition::VECCALC &&
52+ TensorDst::position == AscendC::TPosition::GM,
53+ "The input parameters do not match. TensorSrc must be UB and RowMajor, "
54+ "while TensorDst must be GM and RowMajor");
55+ 
56+ AscendC::DataCopyExtParams dataCopyParams(
57+ tla::get<0>(dstTensor.shape()),
58+ tla::get<1>(dstTensor.shape()) * sizeof(ElementSrc),
59+ (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) / ELE_NUM_PER_C0,
60+ (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) * sizeof(ElementSrc),
61+ 0
62+ );
63+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
64+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
65+ AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams);
66+ };
67+};
68+ 
69+/// Partial specialization for AtlasA5, RowMajor in and RowMajor out.
70+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
71+struct CopyUb2GmTla<Arch::AtlasA5,
72+ tla::Tensor<AscendC::LocalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::VECCALC>,
73+ tla::Tensor<AscendC::GlobalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::GM>,
74+ std::enable_if_t<tla::detail::isRowMajor<LayoutSrc>::value &&
75+ tla::detail::isRowMajor<LayoutDst>::value>> {
76+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
77+ 
78+ // Methods
79+ 
80+ __aicore__ inline
81+ CopyUb2GmTla() = default;
82+ 
83+ template <class TensorDst, class TensorSrc>
84+ __aicore__ inline
85+ void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor)
86+ {
87+ static_assert(tla::detail::isRowMajor<typename TensorSrc::Layout>::value &&
88+ tla::detail::isRowMajor<typename TensorDst::Layout>::value &&
89+ TensorSrc::position == AscendC::TPosition::VECCALC &&
90+ TensorDst::position == AscendC::TPosition::GM,
91+ "The input parameters do not match. TensorSrc must be UB and RowMajor, "
92+ "while TensorDst must be GM and RowMajor");
93+ 
94+ AscendC::DataCopyExtParams dataCopyParams(
95+ tla::get<0>(dstTensor.shape()),
96+ tla::get<1>(dstTensor.shape()) * sizeof(ElementSrc),
97+ (tla::get<0>(srcTensor.stride()) - tla::get<1>(srcTensor.shape())) / ELE_NUM_PER_C0,
98+ (tla::get<0>(dstTensor.stride()) - tla::get<1>(dstTensor.shape())) * sizeof(ElementSrc),
99+ 0
100+ );
101+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
102+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
103+ AscendC::DataCopyPad(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], dataCopyParams);
104+ };
105+};
106+ 
107+}
108+ 
109+#endif // EPILOGUE_TILE_COPY_UB_TO_GM_TLA_HPP
110+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/epilogue/tile_common/tile_copy.hpp+130-0
@@ -0,0 +1,130 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EPILOGUE_TILE_TILE_COPY_HPP
12+#define EPILOGUE_TILE_TILE_COPY_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/detail/tag_to_layout.hpp"
17+#include "../../../attn_infra/epilogue/tile_common/copy_gm_to_ub.hpp"
18+#include "../../../attn_infra/epilogue/tile_common/copy_ub_to_gm.hpp"
19+#include "../../../attn_infra/epilogue/tile_common/copy_gm_to_ub_tla.hpp"
20+#include "../../../attn_infra/epilogue/tile_common/copy_ub_to_gm_tla.hpp"
21+#include "../../../tla/tensor.hpp"
22+ 
23+namespace NpuArch::Epilogue::Tile
24+{
25+ 
26+template <
27+ /// Tag indicating architecture
28+ class ArchTag,
29+ class... Args
30+>
31+struct TileCopy {
32+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported tile_common copy, can not find the specialization.");
33+};
34+ 
35+template <
36+ class ArchTag,
37+ /// GemmType for C matrix operand
38+ class CType,
39+ /// GemmType for X matrix operand
40+ class XType,
41+ /// GemmType for D matrix operand
42+ class DType
43+>
44+struct TileCopy<ArchTag, CType, XType, DType> {
45+ using ElementC = typename CType::Element;
46+ using ElementX = typename XType::Element;
47+ using ElementD = typename DType::Element;
48+ 
49+ using CopyGmToUbC = CopyGm2Ub<ArchTag, CType>;
50+ using CopyGmToUbX = CopyGm2Ub<ArchTag, XType>;
51+ using CopyUbToGmD = CopyUb2Gm<ArchTag, DType>;
52+};
53+ 
54+template <
55+ class ArchTag,
56+ class CType,
57+ class XType,
58+ class YType,
59+ class DType
60+>
61+struct TileCopy<ArchTag, CType, XType, YType, DType> {
62+ using ElementC = typename CType::Element;
63+ using ElementX = typename XType::Element;
64+ using ElementY = typename YType::Element;
65+ using ElementD = typename DType::Element;
66+ 
67+ using CopyGmToUbC = CopyGm2Ub<ArchTag, CType>;
68+ using CopyGmToUbX = CopyGm2Ub<ArchTag, XType>;
69+ using CopyGmToUbY = CopyGm2Ub<ArchTag, YType>;
70+ using CopyUbToGmD = CopyUb2Gm<ArchTag, DType>;
71+};
72+ 
73+template <
74+ class ArchTag,
75+ class CType,
76+ class XType,
77+ class YType,
78+ class DType
79+>
80+struct TileCopyBf16 {
81+ using ElementC = typename CType::Element;
82+ using ElementX = bfloat16_t;
83+ using ElementY = bfloat16_t;
84+ using ElementD = bfloat16_t;
85+ 
86+ using CopyGmToUbC = CopyGm2Ub<ArchTag, CType>;
87+ using CopyGmToUbX = CopyGm2Ub<ArchTag, Gemm::GemmType<bfloat16_t, typename XType::Layout>>;
88+ using CopyGmToUbY = CopyGm2Ub<ArchTag, Gemm::GemmType<bfloat16_t, typename YType::Layout>>;
89+ using CopyUbToGmD = CopyUb2Gm<ArchTag, Gemm::GemmType<bfloat16_t, typename DType::Layout>>;
90+};
91+ 
92+template <
93+ class ArchTag,
94+ class CType,
95+ class ScaleType,
96+ class PerTokenScaleType,
97+ class DType
98+>
99+struct TileCopyPerTokenDequant {
100+ using ElementC = typename CType::Element;
101+ using ElementScale = typename ScaleType::Element;
102+ using ElementPerTokenScale = typename PerTokenScaleType::Element;
103+ using ElementD = typename DType::Element;
104+ 
105+ using CopyGmToUbC = CopyGm2Ub<ArchTag, CType>;
106+ using CopyGmToUbScale = CopyGm2Ub<ArchTag, ScaleType>;
107+ using CopyGmToUbPerTokenScale = CopyPerTokenScale2Ub<ArchTag, PerTokenScaleType>;
108+ using CopyUbToGmD = CopyUb2Gm<ArchTag, DType>;
109+};
110+ 
111+template <
112+ class ArchTag,
113+ class ElementO_,
114+ class LayoutTagO_,
115+ class LayoutTagOTmp_
116+>
117+struct TileCopyRescaleO{
118+ using ElementO = ElementO_;
119+ using LayoutTagO = LayoutTagO_;
120+ using LayoutTagOTmp = LayoutTagOTmp_;
121+ using LayoutO = detail::TagToLayout_t<ElementO, LayoutTagO>;
122+
123+ using TensorUbO = tla::Tensor<AscendC::LocalTensor<ElementO>, LayoutO, tla::Coord<tla::_0, tla::_0>, AscendC::TPosition::VECCALC>;
124+ using TensorGmO = tla::Tensor<AscendC::GlobalTensor<ElementO>, LayoutO, tla::Coord<tla::_0, tla::_0>, AscendC::TPosition::GM>;
125+ 
126+ using CopyUbToGmO = Tile::CopyUb2GmTla<ArchTag, TensorUbO, TensorGmO>;
127+};
128+}
129+ 
130+#endif // EPILOGUE_TILE_TILE_COPY_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/block/block_mmad.hpp+66-0
@@ -0,0 +1,66 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_BLOCK_BLOCK_MMAD_HPP
12+#define GEMM_BLOCK_BLOCK_MMAD_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/gemm/tile_common/tile_copy.hpp"
16+#include "../../../attn_infra/gemm/tile_common/tile_mmad.hpp"
17+ 
18+namespace NpuArch::Gemm::Block {
19+ 
20+#if (__CCE_AICORE__ == 220)
21+template <
22+ class DispatchPolicy,
23+ class L1TileShape,
24+ class L0TileShape,
25+ class AType,
26+ class BType,
27+ class CType,
28+ class BiasType = void,
29+ class TileCopy = Gemm::Tile::TileCopy<typename DispatchPolicy::ArchTag, AType, BType, CType, BiasType>,
30+ class TileMmad = Gemm::Tile::TileMmad<typename DispatchPolicy::ArchTag, AType, BType, BiasType>
31+>
32+struct BlockMmad {
33+ static_assert(DEPENDENT_FALSE<DispatchPolicy>, "BlockMmad is not implemented for this DispatchPolicy");
34+};
35+#endif
36+ 
37+#if (__CCE_AICORE__ == 310)
38+template <
39+ class DispatchPolicy,
40+ class L1TileShape,
41+ class L0TileShape,
42+ class ElementA,
43+ class ElementB,
44+ class ElementC,
45+ class ElementBias = void,
46+ class TileCopy = Gemm::Tile::PackedTileCopyTla<typename DispatchPolicy::ArchTag, ElementA, layout::RowMajor,
47+ ElementB, layout::RowMajor, ElementC, layout::RowMajor, ElementBias>,
48+ class TileMmad =
49+ Gemm::Tile::TileMmadTla<typename DispatchPolicy::ArchTag, ElementA, typename TileCopy::LayoutTagL1A>
50+>
51+struct BlockMmadTla {
52+ static_assert(DEPENDENT_FALSE<DispatchPolicy>, "BlockMmadTla is not implemented for this DispatchPolicy");
53+};
54+#endif
55+ 
56+} // namespace NpuArch::Gemm::Block
57+#if (__CCE_AICORE__ == 220)
58+#include "../../../attn_infra/gemm/block/block_mmad_qk.hpp"
59+#include "../../../attn_infra/gemm/block/block_mmad_pv.hpp"
60+#endif
61+#if (__CCE_AICORE__ == 310)
62+#include "../../../attn_infra/gemm/block/block_mmad_pv_arch35_ABf16_C_to_UB.hpp"
63+#include "../../../attn_infra/gemm/block/block_mmad_qk_arch35_ABf16_C_to_UB.hpp"
64+#include "../../../attn_infra/gemm/block/block_mmad_qk_arch35_ABint8_C_to_UB.hpp"
65+#endif
66+#endif // GEMM_BLOCK_BLOCK_MMAD_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/block/block_mmad_arch35_opt.hpp+80-0
@@ -0,0 +1,80 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_BLOCK_MMAD_ARCH35_OPT_HPP
12+#define GEMM_BLOCK_MMAD_ARCH35_OPT_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+ 
16+namespace NpuArch::Gemm::Block {
17+ 
18+struct Arch35MmadOpt {
19+ static constexpr int DYNAMIC_LOOP = -1;
20+ 
21+ __aicore__ inline static bool UseM1N12(uint32_t mLoopNum, uint32_t nLoopNum)
22+ {
23+ return (mLoopNum == 1U) && (nLoopNum > 0U) && (nLoopNum <= 2U);
24+ }
25+ 
26+ __aicore__ inline static uint32_t LoopNum(uint32_t total, uint32_t tile)
27+ {
28+ return CeilDiv(total, tile);
29+ }
30+ 
31+ template <int STATIC_LOOP_NUM>
32+ __aicore__ inline static uint32_t LoopBound(uint32_t loopNum)
33+ {
34+ static_assert(STATIC_LOOP_NUM == DYNAMIC_LOOP || STATIC_LOOP_NUM > 0,
35+ "STATIC_LOOP_NUM must be -1 or positive");
36+ if constexpr (STATIC_LOOP_NUM == DYNAMIC_LOOP) {
37+ return loopNum;
38+ } else {
39+ return static_cast<uint32_t>(STATIC_LOOP_NUM);
40+ }
41+ }
42+ 
43+ __aicore__ inline static uint32_t GetCurLoopCounter(uint32_t outerLoopItr, uint32_t loopNum)
44+ {
45+ return outerLoopItr * loopNum;
46+ }
47+ 
48+ __aicore__ inline static uint32_t MainLoopNum(uint32_t total, uint32_t tile)
49+ {
50+ uint32_t loopNum = LoopNum(total, tile);
51+ return (loopNum > 0U) ? (loopNum - 1U) : 0U;
52+ }
53+ 
54+ __aicore__ inline static uint32_t FinalTileSize(uint32_t total, uint32_t tile, uint32_t mainLoopNum)
55+ {
56+ return total - mainLoopNum * tile;
57+ }
58+ 
59+ template <uint32_t STAGES>
60+ __aicore__ inline static uint32_t StageId(uint32_t counter)
61+ {
62+ static_assert(STAGES > 0U, "STAGES must not be 0");
63+ if constexpr ((STAGES & (STAGES - 1U)) == 0U) {
64+ return counter & (STAGES - 1U);
65+ } else {
66+ return counter % STAGES;
67+ }
68+ }
69+ 
70+ template <uint32_t ALIGN>
71+ __aicore__ inline static uint32_t AlignUpPow2(uint32_t value)
72+ {
73+ static_assert(ALIGN > 0U && ((ALIGN & (ALIGN - 1U)) == 0U), "ALIGN must be power of 2");
74+ return (value + ALIGN - 1U) & ~(ALIGN - 1U);
75+ }
76+};
77+ 
78+} // namespace NpuArch::Gemm::Block
79+ 
80+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/block/block_mmad_pv.hpp+289-0
@@ -0,0 +1,289 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_BLOCK_MMAD_SFAI_PV_HPP
12+#define GEMM_BLOCK_MMAD_SFAI_PV_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/resource.hpp"
16+#include "../../../attn_infra/coord.hpp"
17+#include "../../../attn_infra/arch/cross_core_sync.hpp"
18+#include "../../../attn_infra/gemm/dispatch_policy.hpp"
19+#include "../../../attn_infra/gemm/helper.hpp"
20+#include "../../../attn_infra/gemm_coord.hpp"
21+#include "../../../attn_infra/gemm/tile_common/tile_copy.hpp"
22+#include "../../../attn_infra/gemm/tile_common/tile_mmad.hpp"
23+ 
24+////////////////////////////////////////////////////////////////////
25+ 
26+namespace NpuArch::Gemm::Block {
27+////////////////////////////////////////////////////////////////////
28+ 
29+template <
30+ bool PAGED_CACHE_FLAG_,
31+ bool ENABLE_UNIT_FLAG_,
32+ class L1TileShape_,
33+ class L0TileShape_,
34+ class AType_,
35+ class BType_,
36+ class CType_,
37+ class BiasType_,
38+ class TileCopy_,
39+ class TileMmad_>
40+struct BlockMmad<
41+ MmadAtlasA2SFAIPV<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>,
42+ L1TileShape_,
43+ L0TileShape_,
44+ AType_,
45+ BType_,
46+ CType_,
47+ BiasType_,
48+ TileCopy_,
49+ TileMmad_> {
50+public:
51+ // Type Aliases
52+ using DispatchPolicy = MmadAtlasA2SFAIPV<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>;
53+ using ArchTag = typename DispatchPolicy::ArchTag;
54+ using L1TileShape = L1TileShape_;
55+ using L0TileShape = L0TileShape_;
56+ using ElementA = typename AType_::Element;
57+ using LayoutA = typename AType_::Layout;
58+ using ElementB = typename BType_::Element;
59+ using LayoutB = typename BType_::Layout;
60+ using ElementC = typename CType_::Element;
61+ using LayoutC = typename CType_::Layout;
62+ using TileMmad = TileMmad_;
63+ using CopyGmToL1A = typename TileCopy_::CopyGmToL1A;
64+ using CopyGmToL1B = typename TileCopy_::CopyGmToL1B;
65+ using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A;
66+ using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B;
67+ using CopyL0CToGm = typename TileCopy_::CopyL0CToGm;
68+ using ElementAccumulator =
69+ typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
70+ using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc;
71+ using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc;
72+ using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst;
73+ using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst;
74+ using LayoutCInL0 = layout::zN;
75+ 
76+ using L1AAlignHelper = Gemm::helper::L1AlignHelper<ElementA, LayoutA>;
77+ using L1BAlignHelper = Gemm::helper::L1AlignHelper<ElementB, LayoutB>;
78+ 
79+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
80+ static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA);
81+ static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB);
82+ static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE;
83+ static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE;
84+ static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE;
85+ static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES;
86+ static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES;
87+ static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES;
88+ 
89+ // Check LayoutC
90+ static_assert(std::is_same_v<LayoutC, layout::RowMajor>, "LayoutC only support RowMajor yet!");
91+ 
92+ /// Construct
93+ __aicore__ inline
94+ BlockMmad(Arch::Resource<ArchTag> &resource, uint32_t l1BufAddrStart = 0)
95+ {
96+ // Allocate L1 memory space
97+ l1BTensor = resource.l1Buf.template GetBufferByByte<ElementB>(l1BufAddrStart + L1A_SIZE * 2);
98+ for (uint32_t i = 0; i < STAGES; i++) {
99+ l1ATensor[i] = resource.l1Buf.template GetBufferByByte<ElementA>(l1BufAddrStart + L1A_SIZE * i);
100+ l0ATensor[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_PINGPONG_BUF_SIZE * i);
101+ l0BTensor[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_PINGPONG_BUF_SIZE * i);
102+ l0CTensor[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(L0C_PINGPONG_BUF_SIZE * i);
103+ }
104+ }
105+ 
106+ /// Destructor
107+ __aicore__ inline
108+ ~BlockMmad() {}
109+ 
110+ __aicore__ inline
111+ void getBlockShape(
112+ GemmCoord &actualShape, uint32_t &nowNIdx, uint32_t &nLoop, uint32_t &stackSeqTile, uint32_t &blockSize)
113+ {
114+ uint32_t nSplitSize = blockSize;
115+ if (nowNIdx == nLoop - 1) {
116+ nSplitSize = stackSeqTile - nowNIdx * blockSize;
117+ }
118+ actualShape[2] = nSplitSize;
119+ }
120+ 
121+ __aicore__ inline
122+ void operator()(AscendC::GlobalTensor<ElementA> gA,
123+ AscendC::GlobalTensor<ElementB> gB,
124+ AscendC::GlobalTensor<ElementC> gC,
125+ AscendC::GlobalTensor<int32_t> gBlockTable,
126+ AscendC::GlobalTensor<int32_t> gSelectIdx,
127+ LayoutA layoutA, LayoutB layoutB, LayoutC layoutC,GemmCoord actualOriShape,
128+ uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t kvSeqlen, uint32_t strideKV,
129+ uint32_t blockStackNum, Arch::CrossCoreFlag softmaxFlag,
130+ uint32_t &y, uint32_t &selectNum, uint32_t &kvYBlockNum)
131+ {
132+ uint32_t rowNum = actualOriShape[0];
133+ uint32_t embed = actualOriShape[1];
134+ uint32_t stackSeqTile = actualOriShape[2];
135+ GemmCoord actualShape{rowNum, embed, 0};
136+ 
137+ // load V
138+ LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout<ElementB>(stackSeqTile, embed);
139+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID4);
140+ uint32_t nL1Loop = CeilDiv<L1TileShape::N>(stackSeqTile);
141+ 
142+ for (uint32_t blockStackIdx = 0; blockStackIdx < nL1Loop; ++blockStackIdx) {
143+ uint32_t nowNIdx = nIdx + blockStackIdx;
144+ getBlockShape(actualShape, blockStackIdx, nL1Loop, stackSeqTile, blockSize);
145+ uint32_t kActual = actualShape.k();
146+ uint32_t nActual = actualShape.n();
147+ 
148+ uint32_t processSize = 0;
149+ uint32_t nBlockOffset = nowNIdx * blockSize;
150+ uint32_t currentSelectYIdx = nBlockOffset / y;
151+ uint32_t currentYoffset = nBlockOffset % y;
152+ uint32_t currentYIdx = gSelectIdx.GetValue(currentSelectYIdx);
153+ uint32_t offsetInKV = currentYIdx * y + currentYoffset;
154+ 
155+ while (processSize < kActual && currentSelectYIdx < selectNum && currentYIdx < kvYBlockNum && offsetInKV < kvSeqlen) {
156+ uint32_t yAcutal = (currentSelectYIdx == selectNum - 1 && currentYIdx == kvYBlockNum - 1 && kvSeqlen % y != 0) ?
157+ (kvSeqlen - y * currentYIdx) : y;
158+ uint32_t remainingInYBlock = yAcutal - currentYoffset;
159+ uint32_t remainingInNBlock = kActual - processSize;
160+ 
161+ uint32_t actualYSize = min(remainingInNBlock, remainingInYBlock);
162+ if (actualYSize == 0) {
163+ break;
164+ }
165+ 
166+ auto layoutBTile = layoutB.GetTileLayout(MakeCoord(actualYSize, nActual));
167+ MatrixCoord l1BTileCoord{blockStackIdx * blockSize + processSize, 0};
168+ auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BTileCoord)];
169+
170+ copyGmToL1B(l1BTile, gB[offsetInKV * strideKV], layoutBInL1, layoutBTile);
171+
172+ processSize += actualYSize;
173+ currentYoffset += actualYSize;
174+ offsetInKV += actualYSize;
175+ 
176+ if (currentYoffset >= yAcutal) {
177+ currentSelectYIdx++;
178+ if (currentSelectYIdx >= selectNum) {
179+ break;
180+ }
181+ currentYoffset = 0;
182+ currentYIdx = gSelectIdx.GetValue(currentSelectYIdx);
183+ offsetInKV = currentYIdx * y;
184+ }
185+ }
186+ }
187+ 
188+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
189+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
190+ 
191+ Arch::CrossCoreWaitFlag(softmaxFlag);
192+ 
193+ uint32_t mL1Loop = CeilDiv<L1TileShape::M>(rowNum);
194+ uint32_t kL1Loop = CeilDiv<L1TileShape::K>(stackSeqTile);
195+ for (uint32_t mL1Idx = 0; mL1Idx < mL1Loop; mL1Idx++) {
196+ uint32_t mL1Actual = (mL1Idx < mL1Loop - 1) ? L1TileShape::M : (rowNum - mL1Idx * L1TileShape::M);
197+ uint32_t mRound = RoundUp<L1AAlignHelper::M_ALIGNED>(mL1Actual);
198+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CPingPongFlag);
199+ for (uint32_t kL1Idx = 0; kL1Idx < kL1Loop; kL1Idx++) {
200+ uint32_t kL1Actual = (kL1Idx < kL1Loop - 1) ? L1TileShape::K : (stackSeqTile - kL1Idx * L1TileShape::K);
201+ 
202+ // load P
203+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1PPingPongFlag);
204+ MatrixCoord gmATileCoord{mL1Idx * L1TileShape::M, kL1Idx * L1TileShape::K};
205+ auto gmTileA = gA[layoutA.GetOffset(gmATileCoord)];
206+ auto layoutTileA = layoutA.GetTileLayout(MakeCoord(mL1Actual, kL1Actual));
207+ LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout<ElementA>(mL1Actual, kL1Actual);
208+ copyGmToL1A(l1ATensor[l1PPingPongFlag], gmTileA, layoutAInL1, layoutTileA);
209+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1PPingPongFlag);
210+ 
211+ uint32_t kL0Loop = CeilDiv<L0TileShape::K>(kL1Actual);
212+ for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) {
213+ uint32_t kL0Actual =
214+ (kL0Idx < kL0Loop - 1) ? L0TileShape::K : (kL1Actual - kL0Idx * L0TileShape::K);
215+ 
216+ LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout<ElementA>(mL1Actual, kL0Actual);
217+ MatrixCoord l1ATileCoord{0, kL0Idx * L0TileShape::K};
218+ auto l1ATile = l1ATensor[l1PPingPongFlag][layoutAInL1.GetOffset(l1ATileCoord)];
219+ 
220+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag);
221+ if (kL0Idx == 0) {
222+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1PPingPongFlag);
223+ }
224+ copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATile, layoutAInL0, layoutAInL1);
225+ if (kL0Idx == kL0Loop - 1) {
226+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1PPingPongFlag);
227+ }
228+ 
229+ LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout<ElementB>(kL0Actual, embed);
230+ MatrixCoord l1BTileCoord{kL1Idx * L1TileShape::K + kL0Idx * L0TileShape::K, 0};
231+ auto l1BTile = l1BTensor[layoutBInL1.GetOffset(l1BTileCoord)];
232+ 
233+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag + 2);
234+ copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTile, layoutBInL0, layoutBInL1);
235+ 
236+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
237+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
238+ bool initMmad = kL1Idx == 0 && kL0Idx == 0;
239+ tileMmad(l0CTensor[l0CPingPongFlag],
240+ l0ATensor[l0ABPingPongFlag],
241+ l0BTensor[l0ABPingPongFlag],
242+ mRound,
243+ embed,
244+ kL0Actual,
245+ initMmad);
246+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag);
247+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag + 2);
248+ l0ABPingPongFlag = 1 - l0ABPingPongFlag;
249+ }
250+ l1PPingPongFlag = 1 - l1PPingPongFlag;
251+ }
252+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
253+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
254+ MatrixCoord gmCTileCoord{mL1Idx * L0TileShape::M, 0};
255+ LayoutC layoutCTile = layoutC.GetTileLayout(MakeCoord(mL1Actual, embed));
256+ auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mL1Actual, embed));
257+ copyL0CToGm(gC[layoutC.GetOffset(gmCTileCoord)], l0CTensor[l0CPingPongFlag], layoutCTile, layoutInL0C);
258+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CPingPongFlag);
259+ l0CPingPongFlag = 1 - l0CPingPongFlag;
260+ }
261+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID4);
262+ }
263+ 
264+protected:
265+ /// Data members
266+ AscendC::LocalTensor<ElementA> l1ATensor[STAGES];
267+ AscendC::LocalTensor<ElementB> l1BTensor;
268+ AscendC::LocalTensor<ElementA> l0ATensor[STAGES];
269+ AscendC::LocalTensor<ElementB> l0BTensor[STAGES];
270+ AscendC::LocalTensor<ElementAccumulator> l0CTensor[STAGES];
271+ 
272+ TileMmad tileMmad;
273+ CopyGmToL1A copyGmToL1A;
274+ CopyGmToL1B copyGmToL1B;
275+ CopyL1ToL0A copyL1ToL0A;
276+ CopyL1ToL0B copyL1ToL0B;
277+ CopyL0CToGm copyL0CToGm;
278+ 
279+ uint32_t l1PPingPongFlag = 0;
280+ uint32_t l0CPingPongFlag = 0;
281+ uint32_t l0ABPingPongFlag = 0;
282+};
283+ 
284+////////////////////////////////////////////////////////////////////
285+ 
286+} // namespace NpuArch::Gemm::Block
287+ 
288+#endif // GEMM_BLOCK_MMAD_SFAI_PV_HPP
289+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/block/block_mmad_pv_arch35_ABf16_C_to_UB.hpp+462-0
@@ -0,0 +1,462 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+ * @brief matmul implementation for single p&v base tile
13+ * This implementation is designed for the following senario:
14+ * A full p base tile is loaded to L1 from UB, no workspace transit
15+ * A full v base tile is loaded to L1 from GM，relavent instructions launched before p base tile crossCore wait
16+ * A full p*v base tile is loaded to UB from l0C, no workspace transit
17+ */
18+#ifndef GEMM_BLOCK_PV_ARCH35_ABF16_C2UB_HPP
19+#define GEMM_BLOCK_PV_ARCH35_ABF16_C2UB_HPP
20+ 
21+#include "../../../attn_infra/base_defs.hpp"
22+#include "../../../attn_infra/arch/resource.hpp"
23+#include "../../../attn_infra/arch/cross_core_sync.hpp"
24+#include "../../../attn_infra/coord.hpp"
25+#include "../../../attn_infra/gemm/dispatch_policy.hpp"
26+#include "../../../attn_infra/gemm/helper.hpp"
27+#include "../../../attn_infra/gemm_coord.hpp"
28+#include "../../../attn_infra/gemm/block/block_mmad_arch35_opt.hpp"
29+#include "../../../attn_infra/gemm/tile_common/tile_copy.hpp"
30+#include "../../../attn_infra/gemm/tile_common/tile_mmad.hpp"
31+#include "../../../tla/layout.hpp"
32+#include "../../../tla/tensor.hpp"
33+ 
34+////////////////////////////////////////////////////////////////////
35+ 
36+namespace NpuArch::Gemm::Block {
37+////////////////////////////////////////////////////////////////////
38+ 
39+struct Mm2L1TileHelper {
40+ uint32_t mm2L1TileM;
41+ uint32_t mm2L1TileN;
42+ uint32_t mm2L1TileKLeft;
43+ uint32_t mm2L1TileKRight;
44+ uint32_t pL1BufNum;
45+ uint32_t vL1BufNum;
46+ 
47+ __aicore__ inline
48+ Mm2L1TileHelper() {}
49+ 
50+ __aicore__ inline
51+ Mm2L1TileHelper(
52+ uint32_t m,
53+ uint32_t n,
54+ uint32_t kl,
55+ uint32_t kr,
56+ uint32_t pbn,
57+ uint32_t vbn) :
58+ mm2L1TileM(m),
59+ mm2L1TileN(n),
60+ mm2L1TileKLeft(kl),
61+ mm2L1TileKRight(kr),
62+ pL1BufNum(pbn),
63+ vL1BufNum(vbn) {}
64+};
65+ 
66+template <
67+ class L1TileShape_,
68+ class L0TileShape_,
69+ class ElementA_,
70+ class ElementB_,
71+ class ElementC_,
72+ class ElementBias_,
73+ class TileCopy_,
74+ class TileMmad_>
75+struct BlockMmadTla<
76+ MmadAtlasA5BsaPV,
77+ L1TileShape_,
78+ L0TileShape_,
79+ ElementA_,
80+ ElementB_,
81+ ElementC_,
82+ ElementBias_,
83+ TileCopy_,
84+ TileMmad_>
85+{
86+public:
87+ using DispatchPolicy = MmadAtlasA5BsaPV;
88+ using ArchTag = typename DispatchPolicy::ArchTag;
89+ using TileCopy = TileCopy_;
90+ using ElementA = ElementA_;
91+ using ElementB = ElementB_;
92+ using ElementC = ElementC_;
93+ 
94+ using TileMmad = TileMmad_;
95+ 
96+ using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A;
97+ using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B;
98+ 
99+ using ElementAccumulator = typename TileCopy::ElementAccumulator;
100+ 
101+ using LayoutTagL1A = typename TileCopy::LayoutTagL1A;
102+ using LayoutTagL1B = typename TileCopy::LayoutTagL1B;
103+ using LayoutTagL0A = typename TileCopy::LayoutTagL0A;
104+ using LayoutTagL0B = typename TileCopy::LayoutTagL0B;
105+ 
106+ static constexpr uint32_t L0_STAGES = DispatchPolicy::L0_STAGES;
107+ static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape_{});
108+ static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape_{});
109+ static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape_{});
110+ static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / L0_STAGES;
111+ static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / L0_STAGES;
112+ static constexpr uint32_t L0C_HALF_BUF_SIZE = ArchTag::L0C_SIZE / 2;
113+ static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_HALF_BUF_SIZE / L0_STAGES;
114+ 
115+ static constexpr uint32_t MAX_L1_STAGES = 3; // 编译期常量，为静态L1Tensor数组开辟准备。取一个buffer份数的极大值
116+ static constexpr uint32_t V0_V1_FLAG_ID_OFFSET = 16; // 核间同步mode4，AIC侧需要两个flagId分别对应两个AIV
117+ static constexpr uint32_t VSDB_BUF_NUM = 2;
118+ static constexpr uint32_t VSDB_BUF_SIZE = 256 * sizeof(uint64_t);
119+ static constexpr bool enFixpipeDequant = false;
120+ 
121+ __aicore__ inline
122+ BlockMmadTla(Arch::Resource<ArchTag> &resource, uint32_t l1BufAddrStart, Mm2L1TileHelper &mm2L1TileHelper)
123+ {
124+ l1ATileM = mm2L1TileHelper.mm2L1TileM;
125+ l1BTileN = mm2L1TileHelper.mm2L1TileN;
126+ l1ATileK = mm2L1TileHelper.mm2L1TileKLeft;
127+ l1BTileK = mm2L1TileHelper.mm2L1TileKRight;
128+ l1ABufNum = mm2L1TileHelper.pL1BufNum;
129+ l1BBufNum = mm2L1TileHelper.vL1BufNum;
130+ for (uint32_t i = 0; i < l1ABufNum; i++) {
131+ l1ATensor[i] = resource.l1Buf.template GetBufferByByte<ElementA>(
132+ l1BufAddrStart + l1ATileM * l1ATileK * sizeof(ElementA) * i);
133+ }
134+ for (uint32_t i = 0; i < l1BBufNum; i++) {
135+ l1BTensor[i] = resource.l1Buf.template GetBufferByByte<ElementB>(
136+ l1BufAddrStart + l1ATileM * l1ATileK * sizeof(ElementA) * l1ABufNum +
137+ l1BTileK * l1BTileN * sizeof(ElementB) * i);
138+ }
139+ for (uint32_t i = 0; i < L0_STAGES; i++) {
140+ l0ATensor[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(
141+ L0A_PINGPONG_BUF_SIZE * i);
142+ l0BTensor[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(
143+ L0B_PINGPONG_BUF_SIZE * i);
144+ l0CTensor[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(
145+ L0C_HALF_BUF_SIZE + L0C_PINGPONG_BUF_SIZE * i);
146+ }
147+ uint32_t vsdbL1AddrStart = ArchTag::L1_SIZE - VSDB_BUF_NUM * VSDB_BUF_SIZE;
148+ for (uint32_t i = 0; i < VSDB_BUF_NUM; i++) {
149+ l1VSDBTensor[i] = resource.l1Buf.template GetBufferByByte<uint64_t>(
150+ vsdbL1AddrStart + VSDB_BUF_SIZE * i);
151+ }
152+ }
153+ 
154+ /// Destructor
155+ __aicore__ inline
156+ ~BlockMmadTla() {}
157+ 
158+ template <uint32_t MODE, pipe_t PIPE>
159+ __aicore__ inline
160+ void SetCrossCoreSync(Arch::CrossCoreFlag &crossCoreFlag)
161+ {
162+ // in mode 4, AIC set for 2 AIVs seperately
163+ if constexpr (MODE == 4U) {
164+ uint16_t flagIdV0 = crossCoreFlag.id;
165+ uint16_t flagIdV1 = flagIdV0 + V0_V1_FLAG_ID_OFFSET;
166+ Arch::CrossCoreFlag crossCoreFlagV1(flagIdV1);
167+ Arch::CrossCoreSetFlag<MODE, PIPE>(crossCoreFlag);
168+ Arch::CrossCoreSetFlag<MODE, PIPE>(crossCoreFlagV1);
169+ }
170+ }
171+ 
172+ template <uint32_t MODE, pipe_t PIPE>
173+ __aicore__ inline
174+ void WaitCrossCoreSync(Arch::CrossCoreFlag &crossCoreFlag)
175+ {
176+ // in mode 4, AIC wait for 2 AIVs seperately
177+ if constexpr (MODE == 4U) {
178+ uint16_t flagIdV0 = crossCoreFlag.id;
179+ uint16_t flagIdV1 = flagIdV0 + V0_V1_FLAG_ID_OFFSET;
180+ Arch::CrossCoreFlag crossCoreFlagV1(flagIdV1);
181+ Arch::CrossCoreWaitFlag<MODE, PIPE>(crossCoreFlag);
182+ Arch::CrossCoreWaitFlag<MODE, PIPE>(crossCoreFlagV1);
183+ }
184+ }
185+
186+ __aicore__ inline
187+ uint32_t GetCurLoopCounter(uint32_t outterLoopItr, uint32_t curLoopNum, uint32_t curLoopItr)
188+ {
189+ return outterLoopItr * curLoopNum + curLoopItr;
190+ }
191+
192+ template <class TensorB, class TensorL1B>
193+ __aicore__ inline
194+ void SparseVBaseTileL1FullLoad(TensorB &gBTensor, TensorL1B &l1BTensorTla,
195+ AscendC::GlobalTensor<int32_t> gSparseBlockIdx,
196+ uint32_t gatheredKvSTileIdx, uint32_t kvSeqlen,
197+ uint32_t kvSBaseTile, uint32_t blockShapeY,
198+ uint32_t yBlockNumAval, uint32_t yBlockNumRsvd,
199+ uint32_t curBaseTileSize, uint32_t embed)
200+ {
201+ using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B<TensorB>;
202+ CopyGmToL1B copyGmToL1B;
203+ uint32_t baseTileStartOffset = gatheredKvSTileIdx * kvSBaseTile;
204+ uint32_t baseTileEndOffset = baseTileStartOffset + curBaseTileSize;
205+ // 稀疏情况下对实际选中的部分gather后进行基本块切分
206+ // 当前处理的Yblock在gather的序列中的起始偏移，初始值为当前基块的起始偏移
207+ uint32_t gatheredStartOffset = baseTileStartOffset;
208+ // 当前处理的Yblock gather后的下标，初始值为基本块起始偏移对应的按Y方向稀疏block的gather后起始下标
209+ uint32_t gatheredYBlockIdx = gatheredStartOffset / blockShapeY;
210+ // 当前基本块起始偏移对应的按Y方向稀疏后的block内起始偏移
211+ uint32_t yBlockInnerStartOffset = gatheredStartOffset % blockShapeY;
212+ // 当前处理的Yblock原始的下标，初始值为基本块对应的按Y方向稀疏block的原始起始下标
213+ uint32_t oriYBlockIdx = gSparseBlockIdx.GetValue(gatheredYBlockIdx);
214+ // 当前处理的Yblock起始位置在原始序列中的偏移，初始值为基本块在原始序列中的起始偏移
215+ uint32_t oriStartOffset = oriYBlockIdx * blockShapeY + yBlockInnerStartOffset;
216+ // 逐稀疏block搬移填充基本块过程中，已处理的累积序列长度
217+ uint32_t dealtLenAccum = 0;
218+ 
219+ while (dealtLenAccum < curBaseTileSize && gatheredYBlockIdx < yBlockNumRsvd &&
220+ oriYBlockIdx < yBlockNumAval && oriStartOffset < kvSeqlen) {
221+ uint32_t curYBlockSize = blockShapeY;
222+ if (oriYBlockIdx == yBlockNumAval - 1) {
223+ curYBlockSize = kvSeqlen - oriYBlockIdx * blockShapeY;
224+ }
225+ uint32_t gatheredEndOffset =
226+ min(gatheredYBlockIdx * blockShapeY + curYBlockSize,
227+ baseTileEndOffset);
228+ // 当前循环处理的序列长度
229+ uint32_t curDealtLen = gatheredEndOffset - gatheredStartOffset;
230+ if (curDealtLen == 0) {
231+ break;
232+ }
233+ auto l1BTensorTlaTile = GetTile(l1BTensorTla,
234+ tla::MakeCoord(dealtLenAccum, 0), tla::MakeShape(curDealtLen, embed));
235+ auto gBTensorTlaTile = GetTile(gBTensor,
236+ tla::MakeCoord(oriStartOffset, 0), tla::MakeShape(curDealtLen, embed));
237+ copyGmToL1B(l1BTensorTlaTile, gBTensorTlaTile);
238+ // 为下一次循环刷新循环变量
239+ dealtLenAccum += curDealtLen;
240+ gatheredStartOffset += curDealtLen;
241+ gatheredYBlockIdx = gatheredStartOffset / blockShapeY;
242+ yBlockInnerStartOffset = gatheredStartOffset % blockShapeY;
243+ if (dealtLenAccum < curBaseTileSize) {
244+ // 防止最后一次循环之后对GM的访问越界引发硬件error
245+ oriYBlockIdx = gSparseBlockIdx.GetValue(gatheredYBlockIdx);
246+ oriStartOffset = oriYBlockIdx * blockShapeY + yBlockInnerStartOffset;
247+ }
248+ }
249+ }
250+ 
251+
252+ 
253+ template <
254+ int staticML0LoopNum = Arch35MmadOpt::DYNAMIC_LOOP,
255+ int staticNL0LoopNum = Arch35MmadOpt::DYNAMIC_LOOP,
256+ class TensorB,
257+ class TensorC>
258+ __aicore__ inline
259+ void operator()(TensorB &gBTensor, TensorC &ubCTensor,
260+ AscendC::GlobalTensor<int32_t> gSparseBlockIdx,
261+ GemmCoord actualOriShape,
262+ uint32_t gatheredKvSTileIdx, uint32_t kvSeqlen,
263+ uint32_t kvSBaseTile, uint32_t blockShapeY,
264+ uint32_t yBlockNumAval, uint32_t yBlockNumRsvd,
265+ uint64_t prefixSumL0AStages, uint64_t prefixSumL0BStages,
266+ Arch::CrossCoreFlag smToMm2Flag, Arch::CrossCoreFlag mm2ToReFlag)
267+ {
268+ struct EmptyTensor {};
269+ EmptyTensor dummyVS;
270+ this->operator()<0, staticML0LoopNum, staticNL0LoopNum>(
271+ gBTensor, ubCTensor, dummyVS, gSparseBlockIdx,
272+ actualOriShape, gatheredKvSTileIdx, kvSeqlen,
273+ kvSBaseTile, blockShapeY, yBlockNumAval,
274+ yBlockNumRsvd, prefixSumL0AStages, prefixSumL0BStages, smToMm2Flag, mm2ToReFlag);
275+ }
276+ 
277+ 
278+ template <
279+ int quant_mode,
280+ int staticML0LoopNum = Arch35MmadOpt::DYNAMIC_LOOP,
281+ int staticNL0LoopNum = Arch35MmadOpt::DYNAMIC_LOOP,
282+ class TensorB,
283+ class TensorC,
284+ class TensorVS>
285+ __aicore__ inline
286+ void operator()(TensorB &gBTensor, TensorC &ubCTensor, TensorVS &gVSTensor,
287+ AscendC::GlobalTensor<int32_t> gSparseBlockIdx,
288+ GemmCoord actualOriShape,
289+ uint32_t gatheredKvSTileIdx, uint32_t kvSeqlen,
290+ uint32_t kvSBaseTile, uint32_t blockShapeY,
291+ uint32_t yBlockNumAval, uint32_t yBlockNumRsvd,
292+ uint64_t prefixSumL0AStages, uint64_t prefixSumL0BStages,
293+ Arch::CrossCoreFlag smToMm2Flag, Arch::CrossCoreFlag mm2ToReFlag)
294+ {
295+ using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst<TensorC>;
296+ CopyL0CToDst copyL0CToDst;
297+ 
298+ uint32_t rowNum = actualOriShape[0];
299+ uint32_t embed = actualOriShape[1];
300+ uint32_t curBaseTileSize = actualOriShape[2];
301+ 
302+ uint32_t l1BBufId = gatheredKvSTileIdx % l1BBufNum;
303+ uint32_t l1BEventId = l1BBufId + 3;
304+ uint32_t l1ABufId = gatheredKvSTileIdx % l1ABufNum;
305+ 
306+ auto l1BLayoutTla = tla::MakeLayout<ElementB, LayoutTagL1B>(curBaseTileSize, embed);
307+ auto l1BTensorTla = tla::MakeTensor(l1BTensor[l1BBufId], l1BLayoutTla, Arch::PositionL1{});
308+ auto l1ALayoutTla = tla::MakeLayout<ElementA, LayoutTagL1A>(rowNum, curBaseTileSize);
309+ auto l1ATensorTla = tla::MakeTensor(l1ATensor[l1ABufId], l1ALayoutTla, Arch::PositionL1{});
310+ // load V full base tile to L1 before crossCoreSync
311+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventId);
312+ SparseVBaseTileL1FullLoad(
313+ gBTensor, l1BTensorTla, gSparseBlockIdx, gatheredKvSTileIdx, kvSeqlen, kvSBaseTile, blockShapeY,
314+ yBlockNumAval, yBlockNumRsvd, curBaseTileSize, embed);
315+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventId);
316+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventId);
317+ // fwd crossCoreSync from online sm to mm2
318+ WaitCrossCoreSync<4, PIPE_MTE1>(smToMm2Flag);
319+ // P full base tile already on L1
320+ uint32_t mL0LoopNum = Arch35MmadOpt::LoopNum(rowNum, L0_TILE_M);
321+ uint32_t nL0LoopNum = Arch35MmadOpt::LoopNum(embed, L0_TILE_N);
322+ uint32_t kL0LoopNum = Arch35MmadOpt::LoopNum(curBaseTileSize, L0_TILE_K);
323+ uint32_t mFixPAligned8 = Arch35MmadOpt::AlignUpPow2<8>(rowNum);
324+ // while splitting the base tile OTmp to 2 AIVs,
325+ // the order of the elements in each column is expected to be preserved,
326+ // which means a column in l0C cannot be chunked and processed by dualMode FixPipe seperately.
327+ // therefore, FixPipe won't launch until each portion(chunked only by columns, based on nbuffer strategy)
328+ // of the base tile is ready on l0C
329+ uint32_t nLoopCounter = Arch35MmadOpt::GetCurLoopCounter(gatheredKvSTileIdx, nL0LoopNum);
330+ uint32_t nOffset = 0;
331+ for (uint32_t nL0Itr = 0; nL0Itr < Arch35MmadOpt::LoopBound<staticNL0LoopNum>(nL0LoopNum); nL0Itr++) {
332+ uint32_t l0TileNAct = (nL0Itr == nL0LoopNum - 1) ? (embed - nOffset) : L0_TILE_N;
333+ // l0C nbuffer chunked only in n loop
334+ uint32_t l0CLoopCounter = nLoopCounter;
335+ uint32_t l0CBufId = Arch35MmadOpt::StageId<L0_STAGES>(l0CLoopCounter);
336+ // uint32_t l0CEventId = l0CBufId;
337+ uint32_t l0CEventId = l0CBufId + 2;
338+ auto l0CLayoutTla = tla::MakeLayoutL0C(rowNum, l0TileNAct);
339+ auto l0CTensorTla = tla::MakeTensor(l0CTensor[l0CBufId], l0CLayoutTla, Arch::PositionL0C{});
340+ uint32_t mOffset = 0;
341+ uint32_t l0ALoopBase = static_cast<uint32_t>(prefixSumL0AStages);
342+ uint32_t l0BLoopBase = static_cast<uint32_t>(prefixSumL0BStages) + nL0Itr * kL0LoopNum;
343+ for (uint32_t mL0Itr = 0; mL0Itr < Arch35MmadOpt::LoopBound<staticML0LoopNum>(mL0LoopNum); mL0Itr++) {
344+ uint32_t l0TileMAct = (mL0Itr == mL0LoopNum - 1) ? (rowNum - mOffset) : L0_TILE_M;
345+ // uint32_t mLoopCounter = GetCurLoopCounter(gatheredKvSTileIdx, mL0LoopNum, mL0Itr);
346+ // different m chunks will be concated in the same piece of l0C buffer
347+ auto l0CTensorTlaTile = GetTile(l0CTensorTla,
348+ tla::MakeCoord(mOffset, 0), tla::MakeShape(l0TileMAct, l0TileNAct));
349+ uint32_t l0ALoopCounter = l0ALoopBase;
350+ uint32_t l0BLoopCounter = l0BLoopBase;
351+ uint32_t kOffset = 0;
352+ for (uint32_t kL0Itr = 0; kL0Itr < kL0LoopNum; kL0Itr++) {
353+ uint32_t l0TileKAct = (kL0Itr == kL0LoopNum - 1) ?
354+ (curBaseTileSize - kOffset) : L0_TILE_K;
355+ uint32_t l0ABufId = Arch35MmadOpt::StageId<L0_STAGES>(l0ALoopCounter);
356+ // l0ABufId = (mLoopCounter % 2 == 0) ? (1 - l0ABufId) : l0ABufId;
357+ uint32_t l0BBufId = Arch35MmadOpt::StageId<L0_STAGES>(l0BLoopCounter);
358+ uint32_t l0AEventId = l0ABufId;
359+ uint32_t l0BEventId = l0BBufId + 2;
360+ // when L0B buffers wouldn't be reused across the k loop
361+ // redundant L0B load caused by m loop can be avoided
362+ auto l1BTensorTlaTile = GetTile(l1BTensorTla,
363+ tla::MakeCoord(kOffset, nOffset), tla::MakeShape(l0TileKAct, l0TileNAct));
364+ auto l0BLayoutTla = tla::MakeLayout<ElementB, LayoutTagL0B>(l0TileKAct, l0TileNAct);
365+ auto l0BTensorTla = tla::MakeTensor(l0BTensor[l0BBufId], l0BLayoutTla, Arch::PositionL0B{});
366+
367+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventId);
368+ copyL1ToL0B(l0BTensorTla, l1BTensorTlaTile);
369+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0BEventId);
370+ 
371+ if ((mL0Itr == mL0LoopNum - 1) && (nL0Itr == nL0LoopNum - 1) && (kL0Itr == kL0LoopNum - 1)) {
372+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventId);
373+ }
374+ 
375+ auto l1ATensorTlaTile = GetTile(l1ATensorTla,
376+ tla::MakeCoord(mOffset, kOffset), tla::MakeShape(l0TileMAct, l0TileKAct));
377+ auto l0ALayoutTla = tla::MakeLayout<ElementA, LayoutTagL0A>(l0TileMAct, l0TileKAct);
378+ auto l0ATensorTla = tla::MakeTensor(l0ATensor[l0ABufId], l0ALayoutTla, Arch::PositionL0A{});
379+ 
380+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventId);
381+ copyL1ToL0A(l0ATensorTla, l1ATensorTlaTile);
382+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0AEventId);
383+ // reverse crossCoreSync for P
384+ if ((mL0Itr == mL0LoopNum - 1) && (nL0Itr == nL0LoopNum - 1) && (kL0Itr == kL0LoopNum - 1)) {
385+ SetCrossCoreSync<4, PIPE_MTE1>(smToMm2Flag);
386+ }
387+ 
388+ bool initMmad = (kL0Itr == 0);
389+ uint32_t l0TileMAligned = Arch35MmadOpt::AlignUpPow2<16>(l0TileMAct);
390+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0AEventId);
391+
392+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0BEventId);
393+ 
394+ if (mL0Itr == 0 && kL0Itr == 0) {
395+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventId);
396+ }
397+ 
398+ tileMmad(
399+ l0CTensorTlaTile,
400+ l0ATensorTla,
401+ l0BTensorTla,
402+ l0TileMAligned,
403+ l0TileNAct,
404+ l0TileKAct,
405+ initMmad);
406+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventId);
407+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventId);
408+ l0ALoopCounter++;
409+ l0BLoopCounter++;
410+ kOffset += L0_TILE_K;
411+ }
412+ mOffset += L0_TILE_M;
413+ l0ALoopBase += kL0LoopNum;
414+ }
415+ // fixpipe
416+ if (nL0Itr == 0) {
417+ // reverse crossCoreSync, do fixPipe only after ubCTensor is fully released
418+ WaitCrossCoreSync<4, PIPE_FIX>(mm2ToReFlag);
419+ }
420+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CEventId);
421+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CEventId);
422+ // 需要kernel传输ubCTensor的时候确保其shape的m，n是满足32B（8个32位元素）对齐的
423+ // rounded up by 8 and splited in half to each AIV
424+ // valid rows in AIV0: [0, mFixPAligned8 / 2 - 1]
425+ // valid rows in AIV1: [mFixPAligned8 / 2, rowNum - 1]
426+ uint32_t nFixPAligned8 = Arch35MmadOpt::AlignUpPow2<8>(l0TileNAct);
427+ auto ubCTensorTlaTile = GetTile(ubCTensor,
428+ tla::MakeCoord(0, nOffset), tla::MakeShape(mFixPAligned8, nFixPAligned8));
429+ copyL0CToDst(ubCTensorTlaTile, l0CTensorTla);
430+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventId);
431+ nLoopCounter++;
432+ nOffset += L0_TILE_N;
433+ }
434+ // crossCoreSync after all fixPipe move
435+ SetCrossCoreSync<4, PIPE_FIX>(mm2ToReFlag);
436+ }
437+ 
438+protected:
439+ /// Data members
440+ AscendC::LocalTensor<ElementA> l1ATensor[MAX_L1_STAGES];
441+ AscendC::LocalTensor<ElementB> l1BTensor[MAX_L1_STAGES];
442+ AscendC::LocalTensor<uint64_t> l1VSDBTensor[VSDB_BUF_NUM];
443+ AscendC::LocalTensor<ElementA> l0ATensor[L0_STAGES];
444+ AscendC::LocalTensor<ElementB> l0BTensor[L0_STAGES];
445+ AscendC::LocalTensor<ElementAccumulator> l0CTensor[L0_STAGES];
446+ 
447+ TileMmad tileMmad;
448+ CopyL1ToL0A copyL1ToL0A;
449+ CopyL1ToL0B copyL1ToL0B;
450+ 
451+ uint32_t l1ATileM;
452+ uint32_t l1BTileN;
453+ uint32_t l1ATileK;
454+ uint32_t l1BTileK;
455+ uint32_t l1ABufNum;
456+ uint32_t l1BBufNum;
457+ uint32_t l1VSdbSel = 0;
458+};
459+////////////////////////////////////////////////////////////////////
460+ 
461+} // namespace NpuArch::Gemm::Block
462+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/block/block_mmad_qk.hpp+323-0
@@ -0,0 +1,323 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_BLOCK_MMAD_SFAI_QK_HPP
12+#define GEMM_BLOCK_MMAD_SFAI_QK_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/resource.hpp"
16+#include "../../../attn_infra/coord.hpp"
17+#include "../../../attn_infra/gemm/dispatch_policy.hpp"
18+#include "../../../attn_infra/gemm/helper.hpp"
19+#include "../../../attn_infra/gemm_coord.hpp"
20+#include "../../../attn_infra/gemm/tile_common/tile_copy.hpp"
21+#include "../../../attn_infra/gemm/tile_common/tile_mmad.hpp"
22+ 
23+////////////////////////////////////////////////////////////////////
24+ 
25+namespace NpuArch::Gemm::Block {
26+////////////////////////////////////////////////////////////////////
27+ 
28+template <
29+ bool PAGED_CACHE_FLAG_,
30+ bool ENABLE_UNIT_FLAG_,
31+ class L1TileShape_,
32+ class L0TileShape_,
33+ class AType_,
34+ class BType_,
35+ class CType_,
36+ class BiasType_,
37+ class TileCopy_,
38+ class TileMmad_>
39+struct BlockMmad<
40+ MmadAtlasA2SFAIQK<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>,
41+ L1TileShape_,
42+ L0TileShape_,
43+ AType_,
44+ BType_,
45+ CType_,
46+ BiasType_,
47+ TileCopy_,
48+ TileMmad_> {
49+public:
50+ using DispatchPolicy = MmadAtlasA2SFAIQK<PAGED_CACHE_FLAG_, ENABLE_UNIT_FLAG_>;
51+ using ArchTag = typename DispatchPolicy::ArchTag;
52+ 
53+ using L1TileShape = L1TileShape_;
54+ using L0TileShape = L0TileShape_;
55+ 
56+ using ElementA = typename AType_::Element;
57+ using LayoutA = typename AType_::Layout;
58+ using ElementB = typename BType_::Element;
59+ using LayoutB = typename BType_::Layout;
60+ using ElementC = typename CType_::Element;
61+ using LayoutC = typename CType_::Layout;
62+
63+ using TileMmad = TileMmad_;
64+ using CopyGmToL1A = typename TileCopy_::CopyGmToL1A;
65+ using CopyGmToL1B = typename TileCopy_::CopyGmToL1B;
66+ using CopyL1ToL0A = typename TileCopy_::CopyL1ToL0A;
67+ using CopyL1ToL0B = typename TileCopy_::CopyL1ToL0B;
68+ using CopyL0CToGm = typename TileCopy_::CopyL0CToGm;
69+
70+ using ElementAccumulator =
71+ typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
72+
73+ using LayoutAInL1 = typename CopyL1ToL0A::LayoutSrc;
74+ using LayoutBInL1 = typename CopyL1ToL0B::LayoutSrc;
75+ using LayoutAInL0 = typename CopyL1ToL0A::LayoutDst;
76+ using LayoutBInL0 = typename CopyL1ToL0B::LayoutDst;
77+ using LayoutCInL0 = layout::zN;
78+ 
79+ using L1AAlignHelper = Gemm::helper::L1AlignHelper<ElementA, LayoutA>;
80+ using L1BAlignHelper = Gemm::helper::L1AlignHelper<ElementB, LayoutB>;
81+ 
82+ static constexpr uint32_t STAGES = DispatchPolicy::STAGES;
83+
84+ static constexpr uint32_t L1A_SIZE = L1TileShape::M * L1TileShape::K * sizeof(ElementA);
85+ static constexpr uint32_t L1B_SIZE = L1TileShape::N * L1TileShape::K * sizeof(ElementB);
86+
87+ static constexpr uint32_t L0A_SIZE = ArchTag::L0A_SIZE;
88+ static constexpr uint32_t L0B_SIZE = ArchTag::L0B_SIZE;
89+ static constexpr uint32_t L0C_SIZE = ArchTag::L0C_SIZE;
90+
91+ static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = L0A_SIZE / STAGES;
92+ static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = L0B_SIZE / STAGES;
93+ static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_SIZE / STAGES;
94+
95+ static constexpr uint32_t BLOCK_SIZE = 16;
96+ 
97+ static_assert(std::is_same_v<LayoutC, layout::RowMajor>, "LayoutC only support RowMajor yet!");
98+ 
99+ __aicore__ inline
100+ BlockMmad(Arch::Resource<ArchTag> &resource, uint32_t l1BufAddrStart = 0)
101+ {
102+ l1ATensor = resource.l1Buf.template GetBufferByByte<ElementA>(l1BufAddrStart);
103+
104+ for (uint32_t i = 0; i < STAGES; i++) {
105+ l1BTensor[i] = resource.l1Buf.template GetBufferByByte<ElementB>(l1BufAddrStart + L1A_SIZE + L1B_SIZE * i);
106+
107+ l0ATensor[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(L0A_PINGPONG_BUF_SIZE * i);
108+ l0BTensor[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(L0B_PINGPONG_BUF_SIZE * i);
109+ l0CTensor[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(L0C_PINGPONG_BUF_SIZE * i);
110+ }
111+ }
112+ 
113+ __aicore__ inline
114+ ~BlockMmad() {}
115+ 
116+ __aicore__ inline
117+ void loadQGM(AscendC::GlobalTensor<ElementA> gA, LayoutA layoutA, uint32_t rowNum, uint32_t &singleGroupHeads,
118+ uint64_t qStride)
119+ {
120+ uint32_t embed = layoutA.shape(1);
121+ uint32_t rowNumRound = RoundUp<L1AAlignHelper::M_ALIGNED>(rowNum);
122+ uint32_t tokenNumPerGroup = rowNum / singleGroupHeads;
123+
124+ auto layoutSingleANd = layoutA.GetTileLayout(MakeCoord(singleGroupHeads, embed));
125+
126+ LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout<ElementA>(rowNum, embed);
127+
128+ copyGmToL1A(l1ATensor,
129+ gA,
130+ layoutAInL1,
131+ layoutSingleANd,
132+ tokenNumPerGroup,
133+ qStride,
134+ tokenNumPerGroup,
135+ BLOCK_SIZE,
136+ rowNumRound);
137+
138+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID3);
139+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID3);
140+ }
141+ 
142+ __aicore__ inline
143+ void getBlockShape(GemmCoord &actualShape, uint32_t nL1Idx, uint32_t nL1Loop, uint32_t stackSeqTile)
144+ {
145+ uint32_t nSplitSize = L1TileShape::N;
146+ if (nL1Idx == nL1Loop - 1) {
147+ nSplitSize = stackSeqTile - nL1Idx * L1TileShape::N;
148+ }
149+ actualShape[1] = nSplitSize;
150+ }
151+ 
152+ __aicore__ inline
153+ void operator()(AscendC::GlobalTensor<ElementA> gA,
154+ AscendC::GlobalTensor<ElementB> gB,
155+ AscendC::GlobalTensor<ElementC> gC,
156+ AscendC::GlobalTensor<int32_t> gBlockTable,
157+ AscendC::GlobalTensor<int32_t> gSelectIdx,
158+ LayoutA layoutA, LayoutB layoutB, LayoutC layoutC, GemmCoord actualOriShape,
159+ uint32_t &nIdx, uint32_t &nLoop, uint32_t &blockSize, uint32_t strideKV,
160+ uint32_t &y, uint32_t &selectNum, uint32_t &kvYBlockNum, uint32_t &kvSeqlen)
161+ {
162+ uint32_t rowNum = actualOriShape[0];
163+ uint32_t stackSeqTile = actualOriShape[1];
164+ uint32_t embed = actualOriShape[2];
165+ 
166+ GemmCoord actualShape{rowNum, 0, embed};
167+ GemmCoord actualNextShape{rowNum, 0, embed};
168+
169+ uint32_t gBOffset = 0;
170+ uint32_t gBNextOffset = 0;
171+ 
172+ LayoutAInL1 layoutAInL1 = LayoutAInL1::template MakeLayout<ElementA>(rowNum, embed);
173+ 
174+ uint32_t tileNNumPerPaged = blockSize / L1TileShape::N;
175+ uint32_t nL1Loop = CeilDiv<L1TileShape::N>(stackSeqTile);
176+
177+ for (uint32_t nL1Idx = 0; nL1Idx < nL1Loop; ++nL1Idx) {
178+ uint32_t nowNIdx = nIdx + nL1Idx / tileNNumPerPaged;
179+
180+ getBlockShape(actualShape, nL1Idx, nL1Loop, stackSeqTile);
181+ 
182+ uint32_t mActual = actualShape.m();
183+ uint32_t kActual = actualShape.k();
184+ uint32_t nActual = actualShape.n();
185+ uint32_t mRound = RoundUp<L1AAlignHelper::M_ALIGNED>(mActual);
186+ 
187+ LayoutBInL1 layoutBInL1 = LayoutBInL1::template MakeLayout<ElementB>(kActual, nActual);
188+ 
189+ uint32_t processSize = 0;
190+ uint32_t nBlockOffset = nowNIdx * blockSize;
191+ uint32_t currentSelectYIdx = nBlockOffset / y;
192+ uint32_t currentYoffset = nBlockOffset % y;
193+ 
194+ uint32_t currentYIdx = gSelectIdx.GetValue(currentSelectYIdx);
195+ uint32_t offsetInKV = currentYIdx * y + currentYoffset;
196+ 
197+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1KvPingPongFlag);
198+
199+ while (processSize < nActual && currentSelectYIdx < selectNum && currentYIdx < kvYBlockNum && offsetInKV < kvSeqlen) {
200+ uint32_t yAcutal = (currentSelectYIdx == selectNum - 1 && currentYIdx == kvYBlockNum - 1 && kvSeqlen % y != 0) ?
201+ (kvSeqlen - y * currentYIdx) : y;
202+ uint32_t remainingInYBlock = yAcutal - currentYoffset;
203+ uint32_t remainingInNBlock = nActual - processSize;
204+ 
205+ uint32_t actualYSize = min(remainingInNBlock, remainingInYBlock);
206+ if (actualYSize == 0) {
207+ break;
208+ }
209+ 
210+ auto layoutBTile = layoutB.GetTileLayout(MakeCoord(kActual, actualYSize));
211+ 
212+ copyGmToL1B(l1BTensor[l1KvPingPongFlag][processSize], gB[offsetInKV * strideKV], layoutBInL1, layoutBTile);
213+ 
214+ processSize += actualYSize;
215+ currentYoffset += actualYSize;
216+ offsetInKV += actualYSize;
217+ 
218+ if (currentYoffset >= yAcutal) {
219+ currentSelectYIdx++;
220+ if (currentSelectYIdx >= selectNum) {
221+ break;
222+ }
223+ currentYoffset = 0;
224+ currentYIdx = gSelectIdx.GetValue(currentSelectYIdx);
225+ offsetInKV = currentYIdx * y;
226+ }
227+ }
228+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1KvPingPongFlag);
229+ 
230+ uint32_t mL0Loop = CeilDiv<L0TileShape::M>(mActual);
231+ uint32_t kL0Loop = CeilDiv<L0TileShape::K>(kActual);
232+ 
233+ for (uint32_t mL0Idx = 0; mL0Idx < mL0Loop; mL0Idx++) {
234+ uint32_t mL0Actual = (mL0Idx < mL0Loop - 1) ? L0TileShape::M : (mActual - mL0Idx * L0TileShape::M);
235+
236+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CPingPongFlag);
237+
238+ for (uint32_t kL0Idx = 0; kL0Idx < kL0Loop; kL0Idx++) {
239+ uint32_t kL0Actual = (kL0Idx < kL0Loop - 1) ? L0TileShape::K : (kActual - kL0Idx * L0TileShape::K);
240+ 
241+ LayoutAInL0 layoutAInL0 = LayoutAInL0::template MakeLayout<ElementA>(mL0Actual, kL0Actual);
242+ MatrixCoord l1ATileCoord{mL0Idx * L0TileShape::M, kL0Idx * L0TileShape::K};
243+ auto l1ATile = l1ATensor[layoutAInL1.GetOffset(l1ATileCoord)];
244+ 
245+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag);
246+ copyL1ToL0A(l0ATensor[l0ABPingPongFlag], l1ATile, layoutAInL0, layoutAInL1);
247+ 
248+ LayoutBInL0 layoutBInL0 = LayoutBInL0::template MakeLayout<ElementB>(kL0Actual, nActual);
249+ MatrixCoord l1BTileCoord{kL0Idx * L0TileShape::K, 0};
250+ auto l1BTile = l1BTensor[l1KvPingPongFlag][layoutBInL1.GetOffset(l1BTileCoord)];
251+
252+ if ((mL0Idx == 0) && (kL0Idx == 0)) {
253+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1KvPingPongFlag);
254+ }
255+
256+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag + 2);
257+ copyL1ToL0B(l0BTensor[l0ABPingPongFlag], l1BTile, layoutBInL0, layoutBInL1);
258+
259+ if ((mL0Idx == mL0Loop - 1) && (kL0Idx == kL0Loop - 1)) {
260+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1KvPingPongFlag);
261+ }
262+ 
263+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
264+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(EVENT_ID0);
265+
266+ bool initMmad = kL0Idx == 0;
267+ tileMmad(l0CTensor[l0CPingPongFlag],
268+ l0ATensor[l0ABPingPongFlag],
269+ l0BTensor[l0ABPingPongFlag],
270+ mRound,
271+ nActual,
272+ kL0Actual,
273+ initMmad);
274+
275+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag);
276+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0ABPingPongFlag + 2);
277+
278+ l0ABPingPongFlag = 1 - l0ABPingPongFlag;
279+ }
280+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
281+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(EVENT_ID0);
282+
283+ MatrixCoord gmCTileCoord{mL0Idx * L0TileShape::M, nL1Idx * L1TileShape::N};
284+ LayoutC layoutCTile = layoutC.GetTileLayout(MakeCoord(mL0Actual, nActual));
285+ auto layoutInL0C = LayoutCInL0::MakeLayoutInL0C(MakeCoord(mL0Actual, nActual));
286+
287+ copyL0CToGm(gC[layoutC.GetOffset(gmCTileCoord)], l0CTensor[l0CPingPongFlag], layoutCTile, layoutInL0C);
288+
289+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CPingPongFlag);
290+ l0CPingPongFlag = 1 - l0CPingPongFlag;
291+ }
292+
293+ l1KvPingPongFlag = 1 - l1KvPingPongFlag;
294+ }
295+ }
296+ 
297+protected:
298+ AscendC::LocalTensor<ElementA> l1ATensor;
299+ AscendC::LocalTensor<ElementB> l1BTensor[STAGES];
300+
301+ AscendC::LocalTensor<ElementA> l0ATensor[STAGES];
302+ AscendC::LocalTensor<ElementB> l0BTensor[STAGES];
303+ AscendC::LocalTensor<ElementAccumulator> l0CTensor[STAGES];
304+ 
305+ TileMmad tileMmad;
306+ CopyGmToL1A copyGmToL1A;
307+ CopyGmToL1B copyGmToL1B;
308+ CopyL1ToL0A copyL1ToL0A;
309+ CopyL1ToL0B copyL1ToL0B;
310+ CopyL0CToGm copyL0CToGm;
311+ 
312+ uint32_t l1KvPingPongFlag = 0;
313+ uint32_t l0CPingPongFlag = 0;
314+ uint32_t l0ABPingPongFlag = 0;
315+};
316+ 
317+////////////////////////////////////////////////////////////////////
318+ 
319+} // namespace NpuArch::Gemm::Block
320+ 
321+ 
322+#endif // GEMM_BLOCK_MMAD_SFAI_QK_HPP
323+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/block/block_mmad_qk_arch35_ABf16_C_to_UB.hpp+499-0
@@ -0,0 +1,499 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+ * @brief matmul implementation for single q&k^t base tile
13+ * This implementation is designed for the following senario:
14+ * A full q base tile is loaded to L1 from GM at the very beginning,
15+ * and it remains persistent until each k base tile is dealt
16+ * A full q*k^t base tile is loaded to UB from l0C, no workspace transit
17+ */
18+#ifndef GEMM_BLOCK_QK_ARCH35_ABF16_C2UB_HPP
19+#define GEMM_BLOCK_QK_ARCH35_ABF16_C2UB_HPP
20+ 
21+#include "../../../attn_infra/base_defs.hpp"
22+#include "../../../attn_infra/arch/resource.hpp"
23+#include "../../../attn_infra/arch/cross_core_sync.hpp"
24+#include "../../../attn_infra/coord.hpp"
25+#include "../../../attn_infra/gemm/dispatch_policy.hpp"
26+#include "../../../attn_infra/gemm/helper.hpp"
27+#include "../../../attn_infra/gemm_coord.hpp"
28+#include "../../../attn_infra/gemm/block/block_mmad_arch35_opt.hpp"
29+#include "../../../attn_infra/gemm/tile_common/tile_copy.hpp"
30+#include "../../../attn_infra/gemm/tile_common/tile_mmad.hpp"
31+#include "../../../tla/layout.hpp"
32+#include "../../../tla/tensor.hpp"
33+ 
34+////////////////////////////////////////////////////////////////////
35+ 
36+namespace NpuArch::Gemm::Block {
37+////////////////////////////////////////////////////////////////////
38+ 
39+ 
40+#ifndef GEMM_BLOCK_QK_ARCH35_COMMON_HELPER_HPP
41+#define GEMM_BLOCK_QK_ARCH35_COMMON_HELPER_HPP
42+template <class T>
43+struct HasParams {
44+ template <class U>
45+ static char test(typename U::Params*);
46+ template <class U>
47+ static int test(...);
48+ static constexpr bool value = sizeof(test<T>(0)) == sizeof(char);
49+};
50+ 
51+struct Mm1L1TileHelper {
52+ uint32_t mm1L1TileM;
53+ uint32_t mm1L1TileN;
54+ uint32_t mm1L1TileKLeft;
55+ uint32_t mm1L1TileKRight;
56+ uint32_t qL1BufNum;
57+ uint32_t kL1BufNum;
58+ 
59+ __aicore__ inline
60+ Mm1L1TileHelper() {}
61+ 
62+ __aicore__ inline
63+ Mm1L1TileHelper(
64+ uint32_t m,
65+ uint32_t n,
66+ uint32_t kl,
67+ uint32_t kr,
68+ uint32_t pbn,
69+ uint32_t vbn) :
70+ mm1L1TileM(m),
71+ mm1L1TileN(n),
72+ mm1L1TileKLeft(kl),
73+ mm1L1TileKRight(kr),
74+ qL1BufNum(pbn),
75+ kL1BufNum(vbn) {}
76+};
77+#endif
78+ 
79+template <
80+ class L1TileShape_,
81+ class L0TileShape_,
82+ class ElementA_,
83+ class ElementB_,
84+ class ElementC_,
85+ class ElementBias_,
86+ class TileCopy_,
87+ class TileMmad_>
88+struct BlockMmadTla<
89+ MmadAtlasA5BsaQK,
90+ L1TileShape_,
91+ L0TileShape_,
92+ ElementA_,
93+ ElementB_,
94+ ElementC_,
95+ ElementBias_,
96+ TileCopy_,
97+ TileMmad_>
98+{
99+public:
100+ using DispatchPolicy = MmadAtlasA5BsaQK;
101+ using ArchTag = typename DispatchPolicy::ArchTag;
102+ using TileCopy = TileCopy_;
103+ using ElementA = ElementA_;
104+ using ElementB = ElementB_;
105+ using ElementC = ElementC_;
106+ 
107+ using TileMmad = TileMmad_;
108+ 
109+ using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A;
110+ using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B;
111+ 
112+ using ElementAccumulator = typename TileCopy::ElementAccumulator;
113+ 
114+ using LayoutTagL1A = typename TileCopy::LayoutTagL1A;
115+ using LayoutTagL1B = typename TileCopy::LayoutTagL1B;
116+ using LayoutTagL0A = typename TileCopy::LayoutTagL0A;
117+ using LayoutTagL0B = typename TileCopy::LayoutTagL0B;
118+ 
119+ static constexpr uint32_t L0_STAGES = DispatchPolicy::L0_STAGES;
120+ static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape_{});
121+ static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape_{});
122+ static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape_{});
123+ static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / L0_STAGES;
124+ static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / L0_STAGES;
125+ static constexpr uint32_t L0C_HALF_BUF_SIZE = ArchTag::L0C_SIZE / 2;
126+ static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_HALF_BUF_SIZE / L0_STAGES;
127+ 
128+ static constexpr uint32_t MAX_L1_STAGES = 3; // 编译期常量，为静态L1Tensor数组开辟准备。取一个buffer份数的极大值
129+ static constexpr uint32_t V0_V1_FLAG_ID_OFFSET = 16; // 核间同步mode4，AIC侧需要两个flagId分别对应两个AIV
130+ 
131+ __aicore__ inline
132+ BlockMmadTla(Arch::Resource<ArchTag> &resource, Mm1L1TileHelper &mm1L1TileHelper)
133+ {
134+ l1ATileM = mm1L1TileHelper.mm1L1TileM;
135+ l1BTileN = mm1L1TileHelper.mm1L1TileN;
136+ l1ATileK = mm1L1TileHelper.mm1L1TileKLeft;
137+ l1BTileK = mm1L1TileHelper.mm1L1TileKRight;
138+ l1ABufNum = mm1L1TileHelper.qL1BufNum;
139+ l1BBufNum = mm1L1TileHelper.kL1BufNum;
140+ 
141+ for (uint32_t i = 0; i < l1ABufNum; i++) {
142+ l1ATensor[i] = resource.l1Buf.template GetBufferByByte<ElementA>(
143+ l1ATileM * l1ATileK * sizeof(ElementA) * i);
144+ }
145+ for (uint32_t i = 0; i < l1BBufNum; i++) {
146+ l1BTensor[i] = resource.l1Buf.template GetBufferByByte<ElementB>(
147+ l1ATileM * l1ATileK * sizeof(ElementA) * l1ABufNum +
148+ l1BTileK * l1BTileN * sizeof(ElementB) * i);
149+ }
150+ for (uint32_t i = 0; i < L0_STAGES; i++) {
151+ l0ATensor[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(
152+ L0A_PINGPONG_BUF_SIZE * i);
153+ l0BTensor[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(
154+ L0B_PINGPONG_BUF_SIZE * i);
155+ l0CTensor[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(
156+ L0C_PINGPONG_BUF_SIZE * i);
157+ }
158+ }
159+ 
160+ /// Destructor
161+ __aicore__ inline
162+ ~BlockMmadTla() {}
163+ 
164+ template <class TensorA>
165+ __aicore__ inline
166+ void loadQGM(TensorA &gATensor, GemmCoord actualOriShape)
167+ {
168+ using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A<TensorA>;
169+ CopyGmToL1A copyGmToL1A;
170+ uint32_t rowNum = actualOriShape[0];
171+ uint32_t embed = actualOriShape[1];
172+ auto l1ALayoutTla = tla::MakeLayout<ElementA, LayoutTagL1A>(rowNum, embed);
173+ auto l1ATensorTla = tla::MakeTensor(l1ATensor[0], l1ALayoutTla, Arch::PositionL1{});
174+ auto l1ATensorTlaTile = GetTile(l1ATensorTla,
175+ tla::MakeCoord(0, 0), tla::MakeShape(rowNum, embed));
176+ auto gATensorTlaTile = GetTile(gATensor,
177+ tla::MakeCoord(0, 0), tla::MakeShape(rowNum, embed));
178+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
179+ copyGmToL1A(l1ATensorTlaTile, gATensorTlaTile);
180+
181+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
182+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
183+ }
184+ 
185+ template <uint32_t MODE, pipe_t PIPE>
186+ __aicore__ inline
187+ void SetCrossCoreSync(Arch::CrossCoreFlag &crossCoreFlag)
188+ {
189+ // in mode 4, AIC set for 2 AIVs seperately
190+ if constexpr (MODE == 4U) {
191+ uint16_t flagIdV0 = crossCoreFlag.id;
192+ uint16_t flagIdV1 = flagIdV0 + V0_V1_FLAG_ID_OFFSET;
193+ Arch::CrossCoreFlag crossCoreFlagV1(flagIdV1);
194+ Arch::CrossCoreSetFlag<MODE, PIPE>(crossCoreFlag);
195+ Arch::CrossCoreSetFlag<MODE, PIPE>(crossCoreFlagV1);
196+ }
197+ }
198+ 
199+ template <uint32_t MODE, pipe_t PIPE>
200+ __aicore__ inline
201+ void WaitCrossCoreSync(Arch::CrossCoreFlag &crossCoreFlag)
202+ {
203+ // in mode 4, AIC wait for 2 AIVs seperately
204+ if constexpr (MODE == 4U) {
205+ uint16_t flagIdV0 = crossCoreFlag.id;
206+ uint16_t flagIdV1 = flagIdV0 + V0_V1_FLAG_ID_OFFSET;
207+ Arch::CrossCoreFlag crossCoreFlagV1(flagIdV1);
208+ Arch::CrossCoreWaitFlag<MODE, PIPE>(crossCoreFlag);
209+ Arch::CrossCoreWaitFlag<MODE, PIPE>(crossCoreFlagV1);
210+ }
211+ }
212+
213+ __aicore__ inline
214+ uint32_t GetCurLoopCounter(uint32_t outterLoopItr, uint32_t curLoopNum, uint32_t curLoopItr)
215+ {
216+ return outterLoopItr * curLoopNum + curLoopItr;
217+ }
218+
219+ template <class TensorB, class TensorL1B>
220+ __aicore__ inline
221+ void SparseKL1TileNLoad(TensorB &gBTensor, TensorL1B &l1BTensorTla,
222+ AscendC::GlobalTensor<int32_t> gSparseBlockIdx,
223+ uint32_t gatheredKvSTileIdx, uint32_t kvSeqlen,
224+ uint32_t kvSBaseTile, uint32_t blockShapeY,
225+ uint32_t yBlockNumAval, uint32_t yBlockNumRsvd,
226+ uint32_t l1BTileNAct, uint32_t embed,
227+ uint32_t kvSBaseTileInnerOffset)
228+ {
229+ using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B<TensorB>;
230+ CopyGmToL1B copyGmToL1B;
231+ uint32_t baseTileStartOffset = gatheredKvSTileIdx * kvSBaseTile + kvSBaseTileInnerOffset;
232+ uint32_t baseTileEndOffset = baseTileStartOffset + l1BTileNAct;
233+ // 稀疏情况下对实际选中的部分gather后进行基本块切分
234+ // 当前处理的Yblock在gather的序列中的起始偏移，初始值为当前基块的起始偏移
235+ uint32_t gatheredStartOffset = baseTileStartOffset;
236+ // 当前处理的Yblock gather后的下标，初始值为基本块起始偏移对应的按Y方向稀疏block的gather后起始下标
237+ uint32_t gatheredYBlockIdx = gatheredStartOffset / blockShapeY;
238+ // 当前基本块起始偏移对应的按Y方向稀疏后的block内起始偏移
239+ uint32_t yBlockInnerStartOffset = gatheredStartOffset % blockShapeY;
240+ // 当前处理的Yblock原始的下标，初始值为基本块对应的按Y方向稀疏block的原始起始下标
241+ uint32_t oriYBlockIdx = gSparseBlockIdx.GetValue(gatheredYBlockIdx);
242+ // 当前处理的Yblock起始位置在原始序列中的偏移，初始值为基本块在原始序列中的起始偏移
243+ uint32_t oriStartOffset = oriYBlockIdx * blockShapeY + yBlockInnerStartOffset;
244+ // 逐稀疏block搬移填充基本块过程中，已处理的累积序列长度
245+ uint32_t dealtLenAccum = 0;
246+ 
247+ while (dealtLenAccum < l1BTileNAct && gatheredYBlockIdx < yBlockNumRsvd &&
248+ oriYBlockIdx < yBlockNumAval && oriStartOffset < kvSeqlen) {
249+ uint32_t curYBlockSize = blockShapeY;
250+ if (oriYBlockIdx == yBlockNumAval - 1) {
251+ curYBlockSize = kvSeqlen - oriYBlockIdx * blockShapeY;
252+ }
253+ uint32_t gatheredEndOffset =
254+ min(gatheredYBlockIdx * blockShapeY + curYBlockSize,
255+ baseTileEndOffset);
256+ // 当前循环处理的序列长度
257+ uint32_t curDealtLen = gatheredEndOffset - gatheredStartOffset;
258+ if (curDealtLen == 0) {
259+ break;
260+ }
261+ 
262+ auto l1BTensorTlaTile = GetTile(l1BTensorTla,
263+ tla::MakeCoord(0, dealtLenAccum), tla::MakeShape(embed, curDealtLen));
264+ auto gBTensorTlaTile = GetTile(gBTensor,
265+ tla::MakeCoord(0, oriStartOffset), tla::MakeShape(embed, curDealtLen));
266+ copyGmToL1B(l1BTensorTlaTile, gBTensorTlaTile);
267+ // 为下一次循环刷新循环变量
268+ dealtLenAccum += curDealtLen;
269+ gatheredStartOffset += curDealtLen;
270+ gatheredYBlockIdx = gatheredStartOffset / blockShapeY;
271+ yBlockInnerStartOffset = gatheredStartOffset % blockShapeY;
272+ if (dealtLenAccum < l1BTileNAct) {
273+ oriYBlockIdx = gSparseBlockIdx.GetValue(gatheredYBlockIdx);
274+ oriStartOffset = oriYBlockIdx * blockShapeY + yBlockInnerStartOffset;
275+ }
276+ }
277+ }
278+ 
279+ template <class TensorB, class TensorC>
280+ __aicore__ inline
281+ void operator()(TensorB &gBTensor, TensorC &ubCTensor,
282+ AscendC::GlobalTensor<int32_t> gSparseBlockIdx,
283+ GemmCoord actualOriShape,
284+ uint32_t gatheredKvSTileIdx, uint32_t kvSeqlen,
285+ uint32_t kvSBaseTile, uint32_t blockShapeY,
286+ uint32_t yBlockNumAval, uint32_t yBlockNumRsvd,
287+ uint64_t prefixSumL0AStages, uint64_t prefixSumL0BStages,
288+ Arch::CrossCoreFlag mm1ToSmFlag,
289+ float scaleValue = 0.0f){
290+ this->operator()<0>(gBTensor, ubCTensor, gSparseBlockIdx,
291+ actualOriShape, gatheredKvSTileIdx, kvSeqlen,
292+ kvSBaseTile, blockShapeY, yBlockNumAval, yBlockNumRsvd,
293+ prefixSumL0AStages, prefixSumL0BStages,
294+ mm1ToSmFlag, scaleValue);
295+ }
296+ 
297+ template <
298+ int quant_mode,
299+ int staticML0LoopNum = Arch35MmadOpt::DYNAMIC_LOOP,
300+ int staticKL0LoopNum = Arch35MmadOpt::DYNAMIC_LOOP,
301+ class TensorB,
302+ class TensorC>
303+ __aicore__ inline
304+ void operator()(TensorB &gBTensor, TensorC &ubCTensor,
305+ AscendC::GlobalTensor<int32_t> gSparseBlockIdx,
306+ GemmCoord actualOriShape,
307+ uint32_t gatheredKvSTileIdx, uint32_t kvSeqlen,
308+ uint32_t kvSBaseTile, uint32_t blockShapeY,
309+ uint32_t yBlockNumAval, uint32_t yBlockNumRsvd,
310+ uint64_t prefixSumL0AStages, uint64_t prefixSumL0BStages,
311+ Arch::CrossCoreFlag mm1ToSmFlag,
312+ float scaleValue = 0.0f)
313+ {
314+ using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst<TensorC>;
315+ CopyL0CToDst copyL0CToDstSub0;
316+ CopyL0CToDst copyL0CToDstSub1;
317+ /*
318+ if constexpr (HasParams<CopyL0CToDst>::value && quant_mode == 1) {
319+ using DstParams = typename CopyL0CToDst::Params;
320+ CopyL0CToDstSub0 = CopyL0CToDst(DstParams{scaleValue});
321+ CopyL0CToDstSub1 = CopyL0CToDst(DstParams{scaleValue});
322+ }*/
323+ 
324+ uint32_t rowNum = actualOriShape[0];
325+ uint32_t embed = actualOriShape[2];
326+ uint32_t curBaseTileSize = actualOriShape[1];
327+ 
328+ auto l1ALayoutTla = tla::MakeLayout<ElementA, LayoutTagL1A>(rowNum, embed);
329+ auto l1ATensorTla = tla::MakeTensor(l1ATensor[0], l1ALayoutTla, Arch::PositionL1{});
330+ 
331+ // P full base tile already on L1
332+ uint32_t nL1LoopNum = Arch35MmadOpt::LoopNum(curBaseTileSize, l1BTileN);
333+ uint32_t mL0LoopNum = Arch35MmadOpt::LoopNum(rowNum, L0_TILE_M);
334+ uint32_t kL0LoopNum = Arch35MmadOpt::LoopNum(embed, L0_TILE_K);
335+ uint32_t mFixPAligned8 = Arch35MmadOpt::AlignUpPow2<8>(rowNum);
336+ uint32_t mPerSubCore = mFixPAligned8 / 2;
337+ bool l1ToL0ANoRepeatFlag = (mL0LoopNum == 1U) && (kL0LoopNum <= L0_STAGES);
338+ 
339+ // while splitting the base tile S to 2 AIVs,
340+ // the order of the elements in each column is expected to be preserved,
341+ // which means a column in l0C cannot be chunked and processed by dualMode FixPipe seperately.
342+ // therefore, FixPipe won't launch until each portion(chunked only by columns, based on nbuffer strategy)
343+ // of the base tile is ready on l0C
344+ uint32_t nLoopCounterL1 = Arch35MmadOpt::GetCurLoopCounter(gatheredKvSTileIdx, nL1LoopNum);
345+ uint32_t nL1Offset = 0;
346+ for (uint32_t nL1Itr = 0; nL1Itr < nL1LoopNum; nL1Itr++) {
347+ uint32_t l1TileNAct = (nL1Itr == nL1LoopNum - 1) ? (curBaseTileSize - nL1Offset) : l1BTileN;
348+ uint32_t l1BBufId = nLoopCounterL1 % l1BBufNum;
349+ uint32_t l1BEventId = l1BBufId + 1;
350+ auto l1BLayoutTla = tla::MakeLayout<ElementB, LayoutTagL1B>(embed, l1TileNAct);
351+ auto l1BTensorTla = tla::MakeTensor(l1BTensor[l1BBufId], l1BLayoutTla, Arch::PositionL1{});
352+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventId);
353+ SparseKL1TileNLoad(
354+ gBTensor, l1BTensorTla, gSparseBlockIdx, gatheredKvSTileIdx, kvSeqlen, kvSBaseTile, blockShapeY,
355+ yBlockNumAval, yBlockNumRsvd, l1TileNAct, embed, nL1Offset);
356+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventId);
357+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1BEventId);
358+ uint32_t nL0LoopNum = Arch35MmadOpt::LoopNum(l1TileNAct, L0_TILE_N);
359+ uint32_t lNLoopCounterL0 = Arch35MmadOpt::GetCurLoopCounter(nL1Itr, nL0LoopNum);
360+ uint32_t l0CLoopCounter = Arch35MmadOpt::GetCurLoopCounter(nLoopCounterL1, nL0LoopNum);
361+ uint32_t nL0Offset = 0;
362+ for (uint32_t nL0Itr = 0; nL0Itr < nL0LoopNum; nL0Itr++) {
363+ uint32_t l0TileNAct = (nL0Itr == nL0LoopNum - 1) ? (l1TileNAct - nL0Offset) : L0_TILE_N;
364+ // l0C nbuffer chunked only in n loop
365+ uint32_t l0CBufId = Arch35MmadOpt::StageId<L0_STAGES>(l0CLoopCounter);
366+ uint32_t l0CEventId = l0CBufId;
367+ auto l0CLayoutTla = tla::MakeLayoutL0C(rowNum, l0TileNAct);
368+ auto l0CTensorTla = tla::MakeTensor(l0CTensor[l0CBufId], l0CLayoutTla, Arch::PositionL0C{});
369+ uint32_t mOffset = 0;
370+ uint32_t l0ALoopBase = static_cast<uint32_t>(prefixSumL0AStages);
371+ uint32_t l0BLoopBase = static_cast<uint32_t>(prefixSumL0BStages) + lNLoopCounterL0 * kL0LoopNum;
372+ for (uint32_t mL0Itr = 0; mL0Itr < Arch35MmadOpt::LoopBound<staticML0LoopNum>(mL0LoopNum); mL0Itr++) {
373+ uint32_t l0TileMAct = (mL0Itr == mL0LoopNum - 1) ? (rowNum - mOffset) : L0_TILE_M;
374+ // uint32_t mLoopCounter = GetCurLoopCounter(gatheredKvSTileIdx, mL0LoopNum, mL0Itr);
375+ // different m chunks will be concated in the same piece of l0C buffer
376+ auto l0CTensorTlaTile = GetTile(l0CTensorTla,
377+ tla::MakeCoord(mOffset, 0), tla::MakeShape(l0TileMAct, l0TileNAct));
378+ uint32_t l0ALoopCounter = l0ALoopBase;
379+ uint32_t l0BLoopCounter = l0BLoopBase;
380+ uint32_t kOffset = 0;
381+ for (uint32_t kL0Itr = 0; kL0Itr < Arch35MmadOpt::LoopBound<staticKL0LoopNum>(kL0LoopNum); kL0Itr++) {
382+ uint32_t l0TileKAct = (kL0Itr == kL0LoopNum - 1) ?
383+ (embed - kOffset) : L0_TILE_K;
384+ uint32_t l0ABufId = Arch35MmadOpt::StageId<L0_STAGES>(l0ALoopCounter);
385+ uint32_t l0BBufId = Arch35MmadOpt::StageId<L0_STAGES>(l0BLoopCounter);
386+ uint32_t l0AEventId = l0ABufId;
387+ uint32_t l0BEventId = l0BBufId + 2;
388+ // when L0B buffers wouldn't be reused across the k loop
389+ // redundant L0B load caused by m loop can be avoided
390+ auto l1BTensorTlaTile = GetTile(l1BTensorTla, tla::MakeCoord(kOffset, nL0Offset),
391+ tla::MakeShape(l0TileKAct, l0TileNAct));
392+ auto l0BLayoutTla = tla::MakeLayout<ElementB, LayoutTagL0B>(l0TileKAct, l0TileNAct);
393+ auto l0BTensorTla = tla::MakeTensor(l0BTensor[l0BBufId], l0BLayoutTla, Arch::PositionL0B{});
394+
395+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventId);
396+ copyL1ToL0B(l0BTensorTla, l1BTensorTlaTile);
397+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0BEventId);
398+ if ((nL0Itr == nL0LoopNum - 1) && (mL0Itr == mL0LoopNum - 1) && (kL0Itr == kL0LoopNum - 1)) {
399+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1BEventId);
400+ }
401+ 
402+ auto l1ATensorTlaTile = GetTile(l1ATensorTla, tla::MakeCoord(mOffset, kOffset),
403+ tla::MakeShape(l0TileMAct, l0TileKAct));
404+ auto l0ALayoutTla = tla::MakeLayout<ElementA, LayoutTagL0A>(l0TileMAct, l0TileKAct);
405+ auto l0ATensorTla = tla::MakeTensor(l0ATensor[l0ABufId], l0ALayoutTla, Arch::PositionL0A{});
406+ 
407+ if (l1ToL0ANoRepeatFlag && (nL1Itr == 0) && (nL0Itr == 0)) {
408+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventId);
409+ copyL1ToL0A(l0ATensorTla, l1ATensorTlaTile);
410+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0AEventId);
411+ }
412+ 
413+ bool initMmad = (kL0Itr == 0);
414+ uint32_t l0TileMAligned = Arch35MmadOpt::AlignUpPow2<16>(l0TileMAct);
415+ if (l1ToL0ANoRepeatFlag && (nL1Itr == 0) && (nL0Itr == 0)) {
416+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0AEventId);
417+ }
418+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0BEventId);
419+ if (mL0Itr == 0 && kL0Itr == 0) {
420+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventId);
421+ }
422+ tileMmad(
423+ l0CTensorTlaTile,
424+ l0ATensorTla,
425+ l0BTensorTla,
426+ l0TileMAligned,
427+ l0TileNAct,
428+ l0TileKAct,
429+ initMmad);
430+ if (l1ToL0ANoRepeatFlag && (nL1Itr == nL1LoopNum - 1) && (nL0Itr == nL0LoopNum - 1)) {
431+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventId);
432+ }
433+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventId);
434+ l0ALoopCounter++;
435+ l0BLoopCounter++;
436+ kOffset += L0_TILE_K;
437+ }
438+ mOffset += L0_TILE_M;
439+ l0ALoopBase += kL0LoopNum;
440+ }
441+ // fixpipe
442+ if (nL0Itr == 0) {
443+ // reverse crossCoreSync, do fixPipe only after ubCTensor is fully released
444+ WaitCrossCoreSync<4, PIPE_FIX>(mm1ToSmFlag);
445+ }
446+ AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CEventId);
447+ AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CEventId);
448+ // 需要kernel传输ubCTensor的时候确保其shape的m，n是满足32B（8个32位元素）对齐的
449+ // rounded up by 8 and splited in half to each AIV
450+ // valid rows in AIV0: [0, mFixPAligned8 / 2 - 1]
451+ // valid rows in AIV1: [mFixPAligned8 / 2, rowNum - 1]
452+ uint32_t nFixPAligned16 = Arch35MmadOpt::AlignUpPow2<16>(l0TileNAct);
453+ auto ubCTensorTlaTile = GetTile(ubCTensor,
454+ tla::MakeCoord(0, nL0Offset), tla::MakeShape(mPerSubCore, nFixPAligned16));
455+ auto l0CTensorTlaTileSub0 = GetTile(l0CTensorTla,
456+ tla::MakeCoord(0, 0), tla::MakeShape(mPerSubCore, l0TileNAct));
457+ auto l0CTensorTlaTileSub1 = GetTile(l0CTensorTla,
458+ tla::MakeCoord(mPerSubCore, 0), tla::MakeShape(mPerSubCore, l0TileNAct));
459+ copyL0CToDstSub0(ubCTensorTlaTile, l0CTensorTlaTileSub0, false);
460+ copyL0CToDstSub1(ubCTensorTlaTile, l0CTensorTlaTileSub1, true);
461+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventId);
462+ lNLoopCounterL0++;
463+ l0CLoopCounter++;
464+ nL0Offset += L0_TILE_N;
465+ }
466+ nLoopCounterL1++;
467+ nL1Offset += l1BTileN;
468+ }
469+ // crossCoreSync after all fixPipe move
470+ SetCrossCoreSync<4, PIPE_FIX>(mm1ToSmFlag);
471+ }
472+ 
473+protected:
474+ /// Data members
475+ AscendC::LocalTensor<ElementA> l1ATensor[MAX_L1_STAGES];
476+ AscendC::LocalTensor<ElementB> l1BTensor[MAX_L1_STAGES];
477+ AscendC::LocalTensor<ElementA> l0ATensor[L0_STAGES];
478+ AscendC::LocalTensor<ElementB> l0BTensor[L0_STAGES];
479+ AscendC::LocalTensor<ElementAccumulator> l0CTensor[L0_STAGES];
480+ 
481+ TileMmad tileMmad;
482+ CopyL1ToL0A copyL1ToL0A;
483+ CopyL1ToL0B copyL1ToL0B;
484+ 
485+ uint32_t l1ATileM;
486+ uint32_t l1BTileN;
487+ uint32_t l1ATileK;
488+ uint32_t l1BTileK;
489+ uint32_t l1ABufNum;
490+ uint32_t l1BBufNum;
491+ 
492+ uint32_t l1PPingPongFlag = 0;
493+ uint32_t l0CPingPongFlag = 0;
494+ uint32_t l0ABPingPongFlag = 0;
495+};
496+////////////////////////////////////////////////////////////////////
497+ 
498+} // namespace NpuArch::Gemm::Block
499+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/block/block_mmad_qk_arch35_ABint8_C_to_UB.hpp+503-0
@@ -0,0 +1,503 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+ * @brief matmul implementation for single q&k^t base tile
13+ * This implementation is designed for the following senario:
14+ * A full q base tile is loaded to L1 from GM at the very beginning,
15+ * and it remains persistent until each k base tile is dealt
16+ * A full q*k^t base tile is loaded to UB from l0C, no workspace transit
17+ */
18+#ifndef GEMM_BLOCK_QK_ARCH35_ABINT8_C2UB_HPP
19+#define GEMM_BLOCK_QK_ARCH35_ABINT8_C2UB_HPP
20+ 
21+#include "../../../attn_infra/base_defs.hpp"
22+#include "../../../attn_infra/arch/resource.hpp"
23+#include "../../../attn_infra/arch/cross_core_sync.hpp"
24+#include "../../../attn_infra/coord.hpp"
25+#include "../../../attn_infra/gemm/dispatch_policy.hpp"
26+#include "../../../attn_infra/gemm/helper.hpp"
27+#include "../../../attn_infra/gemm_coord.hpp"
28+#include "../../../attn_infra/gemm/block/block_mmad_arch35_opt.hpp"
29+#include "../../../attn_infra/gemm/tile_common/tile_copy.hpp"
30+#include "../../../attn_infra/gemm/tile_common/tile_mmad.hpp"
31+#include "../../../tla/layout.hpp"
32+#include "../../../tla/tensor.hpp"
33+ 
34+////////////////////////////////////////////////////////////////////
35+ 
36+namespace NpuArch::Gemm::Block {
37+////////////////////////////////////////////////////////////////////
38+ 
39+ 
40+#ifndef GEMM_BLOCK_QK_ARCH35_COMMON_HELPER_HPP
41+#define GEMM_BLOCK_QK_ARCH35_COMMON_HELPER_HPP
42+template <class T>
43+struct HasParams {
44+ template <class U>
45+ static char test(typename U::Params*);
46+ template <class U>
47+ static int test(...);
48+ static constexpr bool value = sizeof(test<T>(0)) == sizeof(char);
49+};
50+ 
51+struct Mm1L1TileHelper {
52+ uint32_t mm1L1TileM;
53+ uint32_t mm1L1TileN;
54+ uint32_t mm1L1TileKLeft;
55+ uint32_t mm1L1TileKRight;
56+ uint32_t qL1BufNum;
57+ uint32_t kL1BufNum;
58+ 
59+ __aicore__ inline
60+ Mm1L1TileHelper() {}
61+ 
62+ __aicore__ inline
63+ Mm1L1TileHelper(
64+ uint32_t m,
65+ uint32_t n,
66+ uint32_t kl,
67+ uint32_t kr,
68+ uint32_t pbn,
69+ uint32_t vbn) :
70+ mm1L1TileM(m),
71+ mm1L1TileN(n),
72+ mm1L1TileKLeft(kl),
73+ mm1L1TileKRight(kr),
74+ qL1BufNum(pbn),
75+ kL1BufNum(vbn) {}
76+};
77+#endif
78+ 
79+template <
80+ class L1TileShape_,
81+ class L0TileShape_,
82+ class ElementC_,
83+ class ElementBias_,
84+ class TileCopy_,
85+ class TileMmad_>
86+struct BlockMmadTla<
87+ MmadAtlasA5BsaQK,
88+ L1TileShape_,
89+ L0TileShape_,
90+ int8_t,
91+ int8_t,
92+ ElementC_,
93+ ElementBias_,
94+ TileCopy_,
95+ TileMmad_>
96+{
97+public:
98+ using DispatchPolicy = MmadAtlasA5BsaQK;
99+ using ArchTag = typename DispatchPolicy::ArchTag;
100+ using TileCopy = TileCopy_;
101+ using ElementA = int8_t;
102+ using ElementB = int8_t;
103+ using ElementC = ElementC_;
104+ 
105+ using TileMmad = TileMmad_;
106+ 
107+ using CopyL1ToL0A = typename TileCopy::CopyL1ToL0A;
108+ using CopyL1ToL0B = typename TileCopy::CopyL1ToL0B;
109+ 
110+ using ElementAccumulator = typename TileCopy::ElementAccumulator;
111+ 
112+ using LayoutTagL1A = typename TileCopy::LayoutTagL1A;
113+ using LayoutTagL1B = typename TileCopy::LayoutTagL1B;
114+ using LayoutTagL0A = typename TileCopy::LayoutTagL0A;
115+ using LayoutTagL0B = typename TileCopy::LayoutTagL0B;
116+ 
117+ static constexpr uint32_t L0_STAGES = DispatchPolicy::L0_STAGES;
118+ static constexpr uint32_t L0_TILE_M = tla::get<0>(L0TileShape_{});
119+ static constexpr uint32_t L0_TILE_N = tla::get<1>(L0TileShape_{});
120+ static constexpr uint32_t L0_TILE_K = tla::get<2>(L0TileShape_{});
121+ static constexpr uint32_t L0A_PINGPONG_BUF_SIZE = ArchTag::L0A_SIZE / L0_STAGES;
122+ static constexpr uint32_t L0B_PINGPONG_BUF_SIZE = ArchTag::L0B_SIZE / L0_STAGES;
123+ static constexpr uint32_t L0C_HALF_BUF_SIZE = ArchTag::L0C_SIZE / 2;
124+ static constexpr uint32_t L0C_PINGPONG_BUF_SIZE = L0C_HALF_BUF_SIZE / L0_STAGES;
125+ 
126+ static constexpr uint32_t MAX_L1_STAGES = 3; // 编译期常量，为静态L1Tensor数组开辟准备。取一个buffer份数的极大值
127+ static constexpr uint32_t V0_V1_FLAG_ID_OFFSET = 16; // 核间同步mode4，AIC侧需要两个flagId分别对应两个AIV
128+ 
129+ __aicore__ inline
130+ BlockMmadTla(Arch::Resource<ArchTag> &resource, Mm1L1TileHelper &mm1L1TileHelper)
131+ {
132+ l1ATileM = mm1L1TileHelper.mm1L1TileN;
133+ l1BTileN = mm1L1TileHelper.mm1L1TileM;
134+ l1ATileK = mm1L1TileHelper.mm1L1TileKRight;
135+ l1BTileK = mm1L1TileHelper.mm1L1TileKLeft;
136+ l1ABufNum = mm1L1TileHelper.kL1BufNum;
137+ l1BBufNum = mm1L1TileHelper.qL1BufNum;
138+
139+ for (uint32_t i = 0; i < l1ABufNum; i++) {
140+ l1ATensor[i] = resource.l1Buf.template GetBufferByByte<ElementA>(
141+ l1ATileM * l1ATileK * sizeof(ElementA) * i);
142+ }
143+ for (uint32_t i = 0; i < l1BBufNum; i++) {
144+ l1BTensor[i] = resource.l1Buf.template GetBufferByByte<ElementB>(
145+ l1ATileM * l1ATileK * sizeof(ElementA) * l1ABufNum +
146+ l1BTileK * l1BTileN * sizeof(ElementB) * i);
147+ }
148+ for (uint32_t i = 0; i < L0_STAGES; i++) {
149+ l0ATensor[i] = resource.l0ABuf.template GetBufferByByte<ElementA>(
150+ L0A_PINGPONG_BUF_SIZE * i);
151+ l0BTensor[i] = resource.l0BBuf.template GetBufferByByte<ElementB>(
152+ L0B_PINGPONG_BUF_SIZE * i);
153+ l0CTensor[i] = resource.l0CBuf.template GetBufferByByte<ElementAccumulator>(
154+ L0C_PINGPONG_BUF_SIZE * i);
155+ }
156+ }
157+ 
158+ /// Destructor
159+ __aicore__ inline
160+ ~BlockMmadTla() {}
161+ 
162+ template <class TensorB>
163+ __aicore__ inline
164+ void loadQGM(TensorB &gBTensor, GemmCoord actualOriShape)
165+ {
166+ using CopyGmToL1B = typename TileCopy_::template CopyGmToL1B<TensorB>;
167+ CopyGmToL1B copyGmToL1B;
168+ uint32_t rowNum = actualOriShape[0];
169+ uint32_t embed = actualOriShape[1];
170+ //rowNum需要至少按照32对齐
171+ auto l1BLayoutTla = tla::MakeLayout<ElementB, LayoutTagL1B>(embed, Arch35MmadOpt::AlignUpPow2<32>(rowNum));
172+ auto l1BTensorTla = tla::MakeTensor(l1BTensor[0], l1BLayoutTla, Arch::PositionL1{});
173+ auto l1BTensorTlaTile = GetTile(l1BTensorTla,
174+ tla::MakeCoord(0, 0), tla::MakeShape(embed, rowNum));
175+ auto gBTensorTlaTile = GetTile(gBTensor,
176+ tla::MakeCoord(0, 0), tla::MakeShape(embed, rowNum));
177+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
178+ copyGmToL1B(l1BTensorTlaTile, gBTensorTlaTile);
179+
180+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
181+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(EVENT_ID0);
182+ }
183+ 
184+ template <uint32_t MODE, pipe_t PIPE>
185+ __aicore__ inline
186+ void SetCrossCoreSync(Arch::CrossCoreFlag &crossCoreFlag)
187+ {
188+ // in mode 4, AIC set for 2 AIVs seperately
189+ if constexpr (MODE == 4U) {
190+ uint16_t flagIdV0 = crossCoreFlag.id;
191+ uint16_t flagIdV1 = flagIdV0 + V0_V1_FLAG_ID_OFFSET;
192+ Arch::CrossCoreFlag crossCoreFlagV1(flagIdV1);
193+ Arch::CrossCoreSetFlag<MODE, PIPE>(crossCoreFlag);
194+ Arch::CrossCoreSetFlag<MODE, PIPE>(crossCoreFlagV1);
195+ }
196+ }
197+ 
198+ template <uint32_t MODE, pipe_t PIPE>
199+ __aicore__ inline
200+ void WaitCrossCoreSync(Arch::CrossCoreFlag &crossCoreFlag)
201+ {
202+ // in mode 4, AIC wait for 2 AIVs seperately
203+ if constexpr (MODE == 4U) {
204+ uint16_t flagIdV0 = crossCoreFlag.id;
205+ uint16_t flagIdV1 = flagIdV0 + V0_V1_FLAG_ID_OFFSET;
206+ Arch::CrossCoreFlag crossCoreFlagV1(flagIdV1);
207+ Arch::CrossCoreWaitFlag<MODE, PIPE>(crossCoreFlag);
208+ Arch::CrossCoreWaitFlag<MODE, PIPE>(crossCoreFlagV1);
209+ }
210+ }
211+
212+ __aicore__ inline
213+ uint32_t GetCurLoopCounter(uint32_t outterLoopItr, uint32_t curLoopNum, uint32_t curLoopItr)
214+ {
215+ return outterLoopItr * curLoopNum + curLoopItr;
216+ }
217+
218+ template <class TensorA, class TensorL1A>
219+ __aicore__ inline
220+ void SparseKL1TileNLoad(TensorA &gATensor, TensorL1A &l1ATensorTla,
221+ AscendC::GlobalTensor<int32_t> gSparseBlockIdx,
222+ uint32_t gatheredKvSTileIdx, uint32_t kvSeqlen,
223+ uint32_t kvSBaseTile, uint32_t blockShapeY,
224+ uint32_t yBlockNumAval, uint32_t yBlockNumRsvd,
225+ uint32_t l1BTileNAct, uint32_t embed,
226+ uint32_t kvSBaseTileInnerOffset)
227+ {
228+ using CopyGmToL1A = typename TileCopy_::template CopyGmToL1A<TensorA>;
229+ CopyGmToL1A copyGmToL1A;
230+ uint32_t baseTileStartOffset = gatheredKvSTileIdx * kvSBaseTile + kvSBaseTileInnerOffset;
231+ uint32_t baseTileEndOffset = baseTileStartOffset + l1BTileNAct;
232+ // 稀疏情况下对实际选中的部分gather后进行基本块切分
233+ // 当前处理的Yblock在gather的序列中的起始偏移，初始值为当前基块的起始偏移
234+ uint32_t gatheredStartOffset = baseTileStartOffset;
235+ // 当前处理的Yblock gather后的下标，初始值为基本块起始偏移对应的按Y方向稀疏block的gather后起始下标
236+ uint32_t gatheredYBlockIdx = gatheredStartOffset / blockShapeY;
237+ // 当前基本块起始偏移对应的按Y方向稀疏后的block内起始偏移
238+ uint32_t yBlockInnerStartOffset = gatheredStartOffset % blockShapeY;
239+ // 当前处理的Yblock原始的下标，初始值为基本块对应的按Y方向稀疏block的原始起始下标
240+ uint32_t oriYBlockIdx = gSparseBlockIdx.GetValue(gatheredYBlockIdx);
241+ // 当前处理的Yblock起始位置在原始序列中的偏移，初始值为基本块在原始序列中的起始偏移
242+ uint32_t oriStartOffset = oriYBlockIdx * blockShapeY + yBlockInnerStartOffset;
243+ // 逐稀疏block搬移填充基本块过程中，已处理的累积序列长度
244+ uint32_t dealtLenAccum = 0;
245+ 
246+ while (dealtLenAccum < l1BTileNAct && gatheredYBlockIdx < yBlockNumRsvd &&
247+ oriYBlockIdx < yBlockNumAval && oriStartOffset < kvSeqlen) {
248+ uint32_t curYBlockSize = blockShapeY;
249+ if (oriYBlockIdx == yBlockNumAval - 1) {
250+ curYBlockSize = kvSeqlen - oriYBlockIdx * blockShapeY;
251+ }
252+ uint32_t gatheredEndOffset =
253+ min(gatheredYBlockIdx * blockShapeY + curYBlockSize,
254+ baseTileEndOffset);
255+ // 当前循环处理的序列长度
256+ uint32_t curDealtLen = gatheredEndOffset - gatheredStartOffset;
257+ if (curDealtLen == 0) {
258+ break;
259+ }
260+ 
261+ auto l1ATensorTlaTile = GetTile(l1ATensorTla,
262+ tla::MakeCoord(dealtLenAccum, 0), tla::MakeShape(curDealtLen, embed));
263+ auto gATensorTlaTile = GetTile(gATensor,
264+ tla::MakeCoord(oriStartOffset, 0), tla::MakeShape(curDealtLen, embed));
265+ copyGmToL1A(l1ATensorTlaTile, gATensorTlaTile);
266+ // 为下一次循环刷新循环变量
267+ dealtLenAccum += curDealtLen;
268+ gatheredStartOffset += curDealtLen;
269+ gatheredYBlockIdx = gatheredStartOffset / blockShapeY;
270+ yBlockInnerStartOffset = gatheredStartOffset % blockShapeY;
271+ if (dealtLenAccum < l1BTileNAct) {
272+ oriYBlockIdx = gSparseBlockIdx.GetValue(gatheredYBlockIdx);
273+ oriStartOffset = oriYBlockIdx * blockShapeY + yBlockInnerStartOffset;
274+ }
275+ }
276+ }
277+ 
278+ 
279+ template <
280+ int quant_mode,
281+ int staticML0LoopNum = Arch35MmadOpt::DYNAMIC_LOOP,
282+ int staticKL0LoopNum = Arch35MmadOpt::DYNAMIC_LOOP,
283+ class TensorA,
284+ class TensorC>
285+ __aicore__ inline
286+ void operator()(TensorA &gATensor, TensorC &ubCTensor,
287+ AscendC::GlobalTensor<int32_t> gSparseBlockIdx,
288+ GemmCoord actualOriShape,
289+ uint32_t gatheredKvSTileIdx, uint32_t kvSeqlen,
290+ uint32_t kvSBaseTile, uint32_t blockShapeY,
291+ uint32_t yBlockNumAval, uint32_t yBlockNumRsvd,
292+ uint64_t prefixSumL0AStages, uint64_t prefixSumL0BStages,
293+ Arch::CrossCoreFlag mm1ToSmFlag,
294+ float scaleValue = 0.0f)
295+ {
296+ using CopyL0CToDst = typename TileCopy_::template CopyL0CToDst<TensorC>;
297+ CopyL0CToDst copyL0CToDstSub0;
298+ CopyL0CToDst copyL0CToDstSub1;
299+ 
300+ uint32_t rowNum = Arch35MmadOpt::AlignUpPow2<32>(actualOriShape[0]);
301+ uint32_t embed = actualOriShape[2];
302+ uint32_t curBaseTileSize = actualOriShape[1];
303+ 
304+ auto l1BLayoutTla = tla::MakeLayout<ElementB, LayoutTagL1B>(embed, rowNum);
305+ auto l1BTensorTla = tla::MakeTensor(l1BTensor[0], l1BLayoutTla, Arch::PositionL1{});
306+ 
307+ // P full base tile already on L1
308+ uint32_t nL1LoopNum = Arch35MmadOpt::LoopNum(curBaseTileSize, l1ATileM);
309+ uint32_t mL0LoopNum = Arch35MmadOpt::LoopNum(rowNum, L0_TILE_M);
310+ uint32_t kL0LoopNum = Arch35MmadOpt::LoopNum(embed, L0_TILE_K);
311+ uint32_t mFixPAligned64 = Arch35MmadOpt::AlignUpPow2<64>(rowNum);
312+ uint32_t mFixPAligned32 = Arch35MmadOpt::AlignUpPow2<32>(rowNum);
313+ 
314+ uint32_t mPerSubCore = mFixPAligned64 / 2;
315+
316+ bool l1ToL0ANoRepeatFlag = (mL0LoopNum == 1U) && (kL0LoopNum <= L0_STAGES);
317+ 
318+ uint32_t nLoopCounterL1 = Arch35MmadOpt::GetCurLoopCounter(gatheredKvSTileIdx, nL1LoopNum);
319+ uint32_t nL1Offset = 0;
320+ //原本是mxkxn的矩阵变成了nxkxm;
321+ for (uint32_t nL1Itr = 0; nL1Itr < nL1LoopNum; nL1Itr++) {
322+ uint32_t l1TileNAct = (nL1Itr == nL1LoopNum - 1) ? (curBaseTileSize - nL1Offset) : l1ATileM;
323+ uint32_t l1ABufId = nLoopCounterL1 % l1ABufNum;
324+ uint32_t l1AEventId = l1ABufId + 1;
325+ auto l1ALayoutTla = tla::MakeLayout<ElementA, LayoutTagL1A>(l1TileNAct, embed);
326+ auto l1ATensorTla = tla::MakeTensor(l1ATensor[l1ABufId], l1ALayoutTla, Arch::PositionL1{});
327+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventId);
328+ SparseKL1TileNLoad(
329+ gATensor, l1ATensorTla, gSparseBlockIdx, gatheredKvSTileIdx, kvSeqlen, kvSBaseTile, blockShapeY,
330+ yBlockNumAval, yBlockNumRsvd, l1TileNAct, embed, nL1Offset);
331+ AscendC::SetFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventId);
332+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_MTE1>(l1AEventId);
333+ uint32_t nL0LoopNum = Arch35MmadOpt::LoopNum(l1TileNAct, L0_TILE_N);
334+ uint32_t lNLoopCounterL0 = Arch35MmadOpt::GetCurLoopCounter(nL1Itr, nL0LoopNum);
335+ uint32_t l0CLoopCounter = Arch35MmadOpt::GetCurLoopCounter(nLoopCounterL1, nL0LoopNum);
336+ uint32_t nL0Offset = 0;
337+ for (uint32_t nL0Itr = 0; nL0Itr < nL0LoopNum; nL0Itr++) {
338+ uint32_t l0TileNAct = (nL0Itr == nL0LoopNum - 1) ? (l1TileNAct - nL0Offset) : L0_TILE_N;
339+ // l0C nbuffer chunked only in n loop
340+ uint32_t l0CBufId = Arch35MmadOpt::StageId<1>(l0CLoopCounter);
341+ uint32_t l0CEventId = l0CBufId;
342+ auto l0CLayoutTla = tla::MakeLayoutL0C(l0TileNAct, mFixPAligned64);
343+ auto l0CTensorTla = tla::MakeTensor(l0CTensor[l0CBufId], l0CLayoutTla, Arch::PositionL0C{});
344+ uint32_t mOffset = 0;
345+ uint32_t l0ALoopBase = static_cast<uint32_t>(prefixSumL0AStages) + lNLoopCounterL0 * kL0LoopNum;
346+ uint32_t l0BLoopBase = static_cast<uint32_t>(prefixSumL0BStages);
347+ for (uint32_t mL0Itr = 0; mL0Itr < Arch35MmadOpt::LoopBound<staticML0LoopNum>(mL0LoopNum); mL0Itr++) {
348+ uint32_t l0TileMAct = (mL0Itr == mL0LoopNum - 1) ? (rowNum - mOffset) : L0_TILE_M;
349+ // uint32_t mLoopCounter = GetCurLoopCounter(gatheredKvSTileIdx, mL0LoopNum, mL0Itr);
350+ // different m chunks will be concated in the same piece of l0C buffer
351+ auto l0CTensorTlaTile = GetTile(l0CTensorTla,
352+ tla::MakeCoord(0, mOffset), tla::MakeShape(l0TileNAct, l0TileMAct));
353+ uint32_t l0ALoopCounter = l0ALoopBase;
354+ uint32_t l0BLoopCounter = l0BLoopBase;
355+ uint32_t kOffset = 0;
356+
357+ for (uint32_t kL0Itr = 0; kL0Itr < Arch35MmadOpt::LoopBound<staticKL0LoopNum>(kL0LoopNum); kL0Itr++) {
358+ uint32_t l0TileKAct = (kL0Itr == kL0LoopNum - 1) ?
359+ (embed - kOffset) : L0_TILE_K;
360+ uint32_t l0ABufId = Arch35MmadOpt::StageId<L0_STAGES>(l0ALoopCounter);
361+ uint32_t l0BBufId = Arch35MmadOpt::StageId<L0_STAGES>(l0BLoopCounter);
362+ uint32_t l0AEventId = l0ABufId;
363+ uint32_t l0BEventId = l0BBufId + 2;
364+ // when L0B buffers wouldn't be reused across the k loop
365+ // redundant L0B load caused by m loop can be avoided
366+ auto l1ATensorTlaTile = GetTile(l1ATensorTla, tla::MakeCoord(nL0Offset, kOffset),
367+ tla::MakeShape(l0TileNAct, l0TileKAct));
368+ auto l0ALayoutTla = tla::MakeLayout<ElementA, LayoutTagL0A>(l0TileNAct, l0TileKAct);
369+ auto l0ATensorTla = tla::MakeTensor(l0ATensor[l0ABufId], l0ALayoutTla, Arch::PositionL0A{});
370+
371+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0AEventId);
372+ copyL1ToL0A(l0ATensorTla, l1ATensorTlaTile);
373+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0AEventId);
374+ if ((nL0Itr == nL0LoopNum - 1) && (mL0Itr == mL0LoopNum - 1) && (kL0Itr == kL0LoopNum - 1)) {
375+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(l1AEventId);
376+ }
377+ 
378+ auto l1BTensorTlaTile = GetTile(l1BTensorTla, tla::MakeCoord(kOffset, mOffset),
379+ tla::MakeShape(l0TileKAct, l0TileMAct));
380+ auto l0BLayoutTla = tla::MakeLayout<ElementB, LayoutTagL0B>(l0TileKAct, l0TileMAct);
381+ auto l0BTensorTla = tla::MakeTensor(l0BTensor[l0BBufId], l0BLayoutTla, Arch::PositionL0B{});
382+ 
383+ if (l1ToL0ANoRepeatFlag && (nL1Itr == 0) && (nL0Itr == 0)) {
384+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(l0BEventId);
385+ copyL1ToL0B(l0BTensorTla, l1BTensorTlaTile);
386+ AscendC::SetFlag<AscendC::HardEvent::MTE1_M>(l0BEventId);
387+ }
388+ 
389+ bool initMmad = (kL0Itr == 0);
390+ //uint32_t l0TileMAligned = Arch35MmadOpt::AlignUpPow2<16>(l0TileMAct);
391+ uint32_t l0TileNAligned = Arch35MmadOpt::AlignUpPow2<16>(l0TileNAct);
392+ if (l1ToL0ANoRepeatFlag && (nL1Itr == 0) && (nL0Itr == 0)) {
393+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0BEventId);
394+ }
395+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_M>(l0AEventId);
396+ if (mL0Itr == 0 && kL0Itr == 0) {
397+ // AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(l0CEventId);
398+ }
399+ tileMmad(
400+ l0CTensorTlaTile,
401+ l0ATensorTla,
402+ l0BTensorTla,
403+ l0TileNAligned,
404+ l0TileMAct,
405+ l0TileKAct,
406+ initMmad,
407+ mL0Itr == mL0LoopNum - 1? 3 : 2);
408+ if (l1ToL0ANoRepeatFlag && (nL1Itr == nL1LoopNum - 1) && (nL0Itr == nL0LoopNum - 1)) {
409+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0BEventId);
410+ }
411+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(l0AEventId);
412+ l0ALoopCounter++;
413+ l0BLoopCounter++;
414+ kOffset += L0_TILE_K;
415+ }
416+ mOffset += L0_TILE_M;
417+ l0ALoopBase += kL0LoopNum;
418+ }
419+ // fixpipe
420+ if (nL0Itr == 0) {
421+ // reverse crossCoreSync, do fixPipe only after ubCTensor is fully released
422+ WaitCrossCoreSync<4, PIPE_FIX>(mm1ToSmFlag);
423+ }
424+ //AscendC::SetFlag<AscendC::HardEvent::M_FIX>(l0CEventId);
425+ // AscendC::WaitFlag<AscendC::HardEvent::M_FIX>(l0CEventId);
426+ 
427+ uint32_t nFixPAligned16 = Arch35MmadOpt::AlignUpPow2<16>(l0TileNAct);
428+
429+ float scale = scaleValue;
430+ int32_t curBaseTileSizeAligned16 = Arch35MmadOpt::AlignUpPow2<16>(curBaseTileSize);
431+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams0, intriParams1;
432+ intriParams0.nSize = 32;
433+ intriParams0.mSize = nFixPAligned16;
434+ intriParams0.srcStride = nFixPAligned16;
435+ intriParams0.dstStride = 32;
436+ intriParams0.params.srcNdStride = 2 * nFixPAligned16;
437+ intriParams0.params.dstNdStride = curBaseTileSizeAligned16 * 32;
438+ intriParams0.params.ndNum = mPerSubCore == 32 ? 1:2;
439+ intriParams0.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t*>(&scale));
440+ intriParams0.quantPre = QuantMode_t::DEQF16;
441+ intriParams0.reluEn = false;
442+ intriParams0.unitFlag = 3;
443+ intriParams0.dualDstCtl = 0;
444+ intriParams0.subBlockId = false;
445+ 
446+ intriParams1.nSize = 32;
447+ intriParams1.mSize = nFixPAligned16;
448+ intriParams1.srcStride = nFixPAligned16;
449+ intriParams1.dstStride = 32;
450+ intriParams1.params.srcNdStride = 2*nFixPAligned16;
451+ intriParams1.params.dstNdStride = curBaseTileSizeAligned16 * 32;
452+ intriParams1.params.ndNum = (mFixPAligned32 == 64 || mFixPAligned32 == 96) ? 1 : 2;
453+ intriParams1.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t*>(&scale));
454+ intriParams1.quantPre = QuantMode_t::DEQF16;
455+ intriParams1.reluEn = false;
456+ intriParams1.unitFlag = 3;
457+ intriParams1.dualDstCtl = 0;
458+ intriParams1.subBlockId = true;
459+ constexpr static AscendC::FixpipeConfig CFG_ROW_MAJOR_UB = {AscendC::CO2Layout::ROW_MAJOR, true};
460+ AscendC::Fixpipe<half, int32_t, CFG_ROW_MAJOR_UB>(ubCTensor.data()[nL0Itr * L0_TILE_N * 32], l0CTensor[l0CBufId], intriParams0);
461+ if (rowNum > 32) {
462+ AscendC::Fixpipe<half, int32_t, CFG_ROW_MAJOR_UB>(ubCTensor.data()[nL0Itr * L0_TILE_N * 32], l0CTensor[l0CBufId][mPerSubCore * nFixPAligned16], intriParams1);
463+ }
464+ 
465+ // AscendC::SetFlag<AscendC::HardEvent::FIX_M>(l0CEventId);
466+ lNLoopCounterL0++;
467+ l0CLoopCounter++;
468+ nL0Offset += L0_TILE_N;
469+ }
470+ nLoopCounterL1++;
471+ nL1Offset += l1ATileM;
472+ }
473+ // crossCoreSync after all fixPipe move
474+ SetCrossCoreSync<4, PIPE_FIX>(mm1ToSmFlag);
475+ }
476+ 
477+protected:
478+ /// Data members
479+ AscendC::LocalTensor<ElementA> l1ATensor[MAX_L1_STAGES];
480+ AscendC::LocalTensor<ElementB> l1BTensor[MAX_L1_STAGES];
481+ AscendC::LocalTensor<ElementA> l0ATensor[L0_STAGES];
482+ AscendC::LocalTensor<ElementB> l0BTensor[L0_STAGES];
483+ AscendC::LocalTensor<ElementAccumulator> l0CTensor[L0_STAGES];
484+ 
485+ TileMmad tileMmad;
486+ CopyL1ToL0A copyL1ToL0A;
487+ CopyL1ToL0B copyL1ToL0B;
488+ 
489+ uint32_t l1ATileM;
490+ uint32_t l1BTileN;
491+ uint32_t l1ATileK;
492+ uint32_t l1BTileK;
493+ uint32_t l1ABufNum;
494+ uint32_t l1BBufNum;
495+ 
496+ uint32_t l1PPingPongFlag = 0;
497+ uint32_t l0CPingPongFlag = 0;
498+ uint32_t l0ABPingPongFlag = 0;
499+};
500+////////////////////////////////////////////////////////////////////
501+ 
502+} // namespace NpuArch::Gemm::Block
503+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/dispatch_policy.hpp+61-0
@@ -0,0 +1,61 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_DISPATCH_POLICY_HPP
12+#define GEMM_DISPATCH_POLICY_HPP
13+ 
14+#include "../../attn_infra/base_defs.hpp"
15+#include "../../attn_infra/arch/arch.hpp"
16+ 
17+namespace NpuArch::Gemm
18+{
19+ 
20+// Block Mmad Policies
21+ 
22+template <bool ASYNC_ = false>
23+struct MmadAtlasA2Base {
24+ using ArchTag = Arch::AtlasA2;
25+ static constexpr uint32_t ASYNC = ASYNC_;
26+};
27+ 
28+template <bool ASYNC_ = false>
29+struct MmadAtlasA5Base {
30+ using ArchTag = Arch::AtlasA5;
31+ static constexpr uint32_t ASYNC = ASYNC_;
32+};
33+ 
34+using MmadAtlasA2 = MmadAtlasA2Base<false>;
35+using MmadAtlasA5 = MmadAtlasA5Base<false>;
36+ 
37+template <bool PAGED_CACHE_FLAG_ = false, bool ENABLE_UNIT_FLAG_ = false>
38+struct MmadAtlasA2SFAIQK : public MmadAtlasA2 {
39+ static constexpr uint32_t STAGES = 2;
40+ static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_;
41+ static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
42+};
43+ 
44+template <bool PAGED_CACHE_FLAG_ = false, bool ENABLE_UNIT_FLAG_ = false>
45+struct MmadAtlasA2SFAIPV : public MmadAtlasA2 {
46+ static constexpr uint32_t STAGES = 2;
47+ static constexpr bool PAGED_CACHE_FLAG = PAGED_CACHE_FLAG_;
48+ static constexpr bool ENABLE_UNIT_FLAG = ENABLE_UNIT_FLAG_;
49+};
50+ 
51+struct MmadAtlasA5BsaQK : public MmadAtlasA5 {
52+ static constexpr uint32_t L0_STAGES = 2;
53+};
54+ 
55+struct MmadAtlasA5BsaPV : public MmadAtlasA5 {
56+ static constexpr uint32_t L0_STAGES = 2;
57+};
58+ 
59+} // namespace NpuArch::Gemm
60+ 
61+#endif // GEMM_DISPATCH_POLICY_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/gemm_type.hpp+28-0
@@ -0,0 +1,28 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_GEMM_TYPE_HPP
12+#define GEMM_GEMM_TYPE_HPP
13+ 
14+#include "../../attn_infra/base_defs.hpp"
15+ 
16+namespace NpuArch::Gemm
17+{
18+template <class Element_, class Layout_, AscendC::TPosition POSITION_ = AscendC::TPosition::GM>
19+struct GemmType
20+{
21+ using Element = Element_;
22+ using Layout = Layout_;
23+ static constexpr AscendC::TPosition POSITION = POSITION_;
24+};
25+ 
26+} // namespace NpuArch::Gemm
27+ 
28+#endif // GEMM_GEMM_TYPE_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/helper.hpp+347-0
@@ -0,0 +1,347 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_HELPER_HPP
12+#define GEMM_HELPER_HPP
13+ 
14+#include "../../attn_infra/arch/arch.hpp"
15+#include "../../attn_infra/base_defs.hpp"
16+#include "../../attn_infra/layout/layout.hpp"
17+#include "../../attn_infra/gemm/gemm_type.hpp"
18+#include "../../tla/layout.hpp"
19+ 
20+namespace NpuArch::Gemm::helper
21+{
22+ 
23+template<class Element, class Layout>
24+struct L1AlignHelper {
25+ static_assert(DEPENDENT_FALSE<Element>, "Unsupported align helper, can not find the specialization.");
26+};
27+ 
28+template<class Element>
29+struct L1AlignHelper<Element, layout::RowMajor> {
30+ static constexpr uint32_t ELE_NUM_PER_C0 = static_cast<uint32_t>(BYTE_PER_C0) / static_cast<uint32_t>(sizeof(Element));
31+ static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL;
32+ static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0;
33+ static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0;
34+};
35+ 
36+template<class Element>
37+struct L1AlignHelper<Element, layout::ColumnMajor> {
38+ static constexpr uint32_t ELE_NUM_PER_C0 = static_cast<uint32_t>(BYTE_PER_C0) / static_cast<uint32_t>(sizeof(Element));
39+ static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0;
40+ static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0;
41+ static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL;
42+};
43+ 
44+template<class Element>
45+struct L1AlignHelper<Element, layout::PaddingRowMajor> {
46+ static constexpr uint32_t ELE_NUM_PER_C0 = static_cast<uint32_t>(BYTE_PER_C0) / static_cast<uint32_t>(sizeof(Element));
47+ static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL;
48+ static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0;
49+ static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0;
50+};
51+ 
52+template<class Element>
53+struct L1AlignHelper<Element, layout::PaddingColumnMajor> {
54+ static constexpr uint32_t ELE_NUM_PER_C0 = static_cast<uint32_t>(BYTE_PER_C0) / static_cast<uint32_t>(sizeof(Element));
55+ static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0;
56+ static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0;
57+ static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL;
58+};
59+ 
60+template<class Element>
61+struct L1AlignHelper<Element, layout::zN> {
62+ static constexpr uint32_t ELE_NUM_PER_C0 = static_cast<uint32_t>(BYTE_PER_C0) / static_cast<uint32_t>(sizeof(Element));
63+ static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL;
64+ static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0;
65+ static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0;
66+};
67+ 
68+template<class Element>
69+struct L1AlignHelper<Element, layout::nZ> {
70+ static constexpr uint32_t ELE_NUM_PER_C0 = static_cast<uint32_t>(BYTE_PER_C0) / static_cast<uint32_t>(sizeof(Element));
71+ static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0;
72+ static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0;
73+ static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL;
74+};
75+ 
76+template<class ElementA, class ElementB>
77+struct ElementAccumulatorSelector {
78+ static_assert(DEPENDENT_FALSE<ElementA>,
79+ "Unsupported element accumulator selector, can not find the specialization.");
80+};
81+ 
82+template<>
83+struct ElementAccumulatorSelector<half, half> {
84+ using ElementAccumulator = float;
85+};
86+ 
87+template<>
88+struct ElementAccumulatorSelector<float, float> {
89+ using ElementAccumulator = float;
90+};
91+ 
92+template<>
93+struct ElementAccumulatorSelector<int8_t, int8_t> {
94+ using ElementAccumulator = int32_t;
95+};
96+ 
97+template<>
98+struct ElementAccumulatorSelector<bfloat16_t, bfloat16_t> {
99+ using ElementAccumulator = float;
100+};
101+ 
102+#if (__CCE_AICORE__ == 310)
103+template<>
104+struct ElementAccumulatorSelector<float8_e4m3_t, float8_e4m3_t> {
105+ using ElementAccumulator = float;
106+};
107+ 
108+template<>
109+struct ElementAccumulatorSelector<float8_e5m2_t, float8_e5m2_t> {
110+ using ElementAccumulator = float;
111+};
112+ 
113+template<>
114+struct ElementAccumulatorSelector<float8_e4m3_t, float8_e5m2_t> {
115+ using ElementAccumulator = float;
116+};
117+ 
118+template<>
119+struct ElementAccumulatorSelector<float8_e5m2_t, float8_e4m3_t> {
120+ using ElementAccumulator = float;
121+};
122+#endif
123+ 
124+// template <class Element_, bool isMx = false>
125+// struct GetL0Element {
126+// using Element = Element_;
127+// };
128+// #if defined(__NPU_ARCH__) && __NPU_ARCH__ == 3101
129+// template <>
130+// struct GetL0Element<float8_e4m3_t, true> {
131+// using Element = AscendC::mx_fp8_e4m3_t;
132+// };
133+ 
134+// template <>
135+// struct GetL0Element<float8_e5m2_t, true> {
136+// using Element = AscendC::mx_fp8_e5m2_t;
137+// };
138+// #endif
139+ 
140+template<class GmAType>
141+struct L1ATypeSelector {
142+ static_assert(DEPENDENT_FALSE<GmAType>,
143+ "Unsupported layout selector, can not find the specialization.");
144+};
145+ 
146+template<class Element>
147+struct L1ATypeSelector<Gemm::GemmType<Element, layout::VectorLayout>> {
148+ using L1AType = Gemm::GemmType<Element, layout::VectorLayout, AscendC::TPosition::A1>;
149+};
150+ 
151+template<class Element>
152+struct L1ATypeSelector<Gemm::GemmType<Element, layout::RowMajor>> {
153+ using L1AType = Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>;
154+};
155+ 
156+template<class Element>
157+struct L1ATypeSelector<Gemm::GemmType<Element, layout::zN>> {
158+ using L1AType = Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>;
159+};
160+ 
161+template<class Element>
162+struct L1ATypeSelector<Gemm::GemmType<Element, layout::PaddingRowMajor>> {
163+ using L1AType = Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>;
164+};
165+ 
166+template<class Element>
167+struct L1ATypeSelector<Gemm::GemmType<Element, layout::ColumnMajor>> {
168+ using L1AType = Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::A1>;
169+};
170+ 
171+template<class Element>
172+struct L1ATypeSelector<Gemm::GemmType<Element, layout::nZ>> {
173+ using L1AType = Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::A1>;
174+};
175+ 
176+template<class Element>
177+struct L1ATypeSelector<Gemm::GemmType<Element, layout::PaddingColumnMajor>> {
178+ using L1AType = Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::A1>;
179+};
180+ 
181+template<class GmBType>
182+struct L1BTypeSelector {
183+ static_assert(DEPENDENT_FALSE<GmBType>,
184+ "Unsupported layout selector, can not find the specialization.");
185+};
186+ 
187+template<class Element>
188+struct L1BTypeSelector<Gemm::GemmType<Element, layout::RowMajor>> {
189+ using L1BType = Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>;
190+};
191+ 
192+template<class Element>
193+struct L1BTypeSelector<Gemm::GemmType<Element, layout::zN>> {
194+ using L1BType = Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>;
195+};
196+ 
197+template<class Element>
198+struct L1BTypeSelector<Gemm::GemmType<Element, layout::PaddingRowMajor>> {
199+ using L1BType = Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>;
200+};
201+ 
202+template<class Element>
203+struct L1BTypeSelector<Gemm::GemmType<Element, layout::ColumnMajor>> {
204+ using L1BType = Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::A1>;
205+};
206+ 
207+template<class Element>
208+struct L1BTypeSelector<Gemm::GemmType<Element, layout::nZ>> {
209+ using L1BType = Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::A1>;
210+};
211+ 
212+template<class Element>
213+struct L1BTypeSelector<Gemm::GemmType<Element, layout::PaddingColumnMajor>> {
214+ using L1BType = Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::A1>;
215+};
216+ 
217+template<class GmBiasType, class ElementAccumulator>
218+struct L1BiasTypeSelector {
219+ static_assert(DEPENDENT_FALSE<GmBiasType>,
220+ "Unsupported layout selector, can not find the specialization.");
221+};
222+ 
223+template<class ArchTag>
224+struct L0ALayoutSelector {
225+ static_assert(DEPENDENT_FALSE<ArchTag>,
226+ "Unsupported layout selector, can not find the specialization.");
227+};
228+ 
229+template<>
230+struct L0ALayoutSelector<Arch::AtlasA2> {
231+ using Layout = layout::zZ;
232+};
233+ 
234+template<>
235+struct L0ALayoutSelector<Arch::AtlasA5> {
236+ using Layout = layout::zN;
237+};
238+ 
239+template<class Element, class Layout, class Enable = void>
240+struct L1AlignHelperTla {
241+ static_assert(DEPENDENT_FALSE<Element>, "Unsupported align helper tla, can not find the specialization.");
242+};
243+ 
244+template<class Element, class Layout>
245+struct L1AlignHelperTla<Element, Layout, std::enable_if_t<tla::detail::isRowMajor<Layout>::value>> {
246+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
247+ static constexpr uint32_t M_ALIGNED = C0_NUM_PER_FRACTAL;
248+ static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0;
249+ static constexpr uint32_t N_ALIGNED = ELE_NUM_PER_C0;
250+};
251+ 
252+template<class Element, class Layout>
253+struct L1AlignHelperTla<Element, Layout, std::enable_if_t<tla::detail::isColumnMajor<Layout>::value>> {
254+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
255+ static constexpr uint32_t M_ALIGNED = ELE_NUM_PER_C0;
256+ static constexpr uint32_t K_ALIGNED = ELE_NUM_PER_C0;
257+ static constexpr uint32_t N_ALIGNED = C0_NUM_PER_FRACTAL;
258+};
259+ 
260+template<class ElementAccumulator>
261+struct L1BiasTypeSelector<void, ElementAccumulator> {
262+ using GMBiasType = void;
263+ using L1BiasType = void;
264+ using L0BiasType = void;
265+};
266+ 
267+template<class Element, class ElementAccumulator>
268+struct L1BiasTypeSelector<Gemm::GemmType<Element, layout::VectorLayout>, ElementAccumulator> {
269+ using GMBiasType = Gemm::GemmType<Element, layout::VectorLayout, AscendC::TPosition::GM>;
270+ using L1BiasType = Gemm::GemmType<Element, layout::VectorLayout, AscendC::TPosition::A1>;
271+ using L0BiasType = Gemm::GemmType<ElementAccumulator, layout::VectorLayout, AscendC::TPosition::C2>;
272+};
273+ 
274+///////////////////////////////////////
275+// new add
276+template<>
277+struct ElementAccumulatorSelector<int32_t, int32_t> {
278+ using ElementAccumulator = int32_t;
279+};
280+ 
281+template<class GmAType, class GmBType>
282+struct L1AndL0TypeSelectorGemm{
283+ static_assert(DEPENDENT_FALSE<GmAType>,
284+ "Unsupported layout selector, can not find the specialization.");
285+ static_assert(DEPENDENT_FALSE<GmBType>,
286+ "Unsupported layout selector, can not find the specialization.");
287+};
288+ 
289+template<class Element>
290+struct L1AndL0TypeSelectorGemm<Gemm::GemmType<Element, layout::RowMajor>, Gemm::GemmType<Element, layout::RowMajor>>{
291+ using L1AType = Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>;
292+ using L1BType = Gemm::GemmType<Element, layout::zZ, AscendC::TPosition::B1>;
293+ using L0AType = Gemm::GemmType<Element, layout::zZ, AscendC::TPosition::A2>;
294+ using L0BType = Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::B2>;
295+};
296+ 
297+template<>
298+struct L1AndL0TypeSelectorGemm<Gemm::GemmType<int8_t, layout::RowMajor>, Gemm::GemmType<int8_t, layout::RowMajor>>{
299+ using L1AType = Gemm::GemmType<int8_t, layout::zN, AscendC::TPosition::A1>;
300+ using L1BType = Gemm::GemmType<int8_t, layout::zN, AscendC::TPosition::B1>;
301+ using L0AType = Gemm::GemmType<int8_t, layout::zZ, AscendC::TPosition::A2>;
302+ using L0BType = Gemm::GemmType<int8_t, layout::nZ, AscendC::TPosition::B2>;
303+};
304+ 
305+template<class Element>
306+struct L1AndL0TypeSelectorGemm<Gemm::GemmType<Element, layout::ColumnMajor>, Gemm::GemmType<Element, layout::ColumnMajor>>{
307+ using L1AType = Gemm::GemmType<Element, layout::nN, AscendC::TPosition::A1>;
308+ using L1BType = Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::B1>;
309+ using L0AType = Gemm::GemmType<Element, layout::zZ, AscendC::TPosition::A2>;
310+ using L0BType = Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::B2>;
311+};
312+ 
313+template<>
314+struct L1AndL0TypeSelectorGemm<Gemm::GemmType<int8_t, layout::ColumnMajor>, Gemm::GemmType<int8_t, layout::ColumnMajor>>{
315+ using L1AType = Gemm::GemmType<int8_t, layout::nZ, AscendC::TPosition::A1>;
316+ using L1BType = Gemm::GemmType<int8_t, layout::nZ, AscendC::TPosition::B1>;
317+ using L0AType = Gemm::GemmType<int8_t, layout::zZ, AscendC::TPosition::A2>;
318+ using L0BType = Gemm::GemmType<int8_t, layout::nZ, AscendC::TPosition::B2>;
319+};
320+ 
321+template<class Element>
322+struct L1AndL0TypeSelectorGemm<Gemm::GemmType<Element, layout::RowMajor>, Gemm::GemmType<Element, layout::ColumnMajor>>{
323+ using L1AType = Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>;
324+ using L1BType = Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::B1>;
325+ using L0AType = Gemm::GemmType<Element, layout::zZ, AscendC::TPosition::A2>;
326+ using L0BType = Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::B2>;
327+};
328+ 
329+template<class Element>
330+struct L1AndL0TypeSelectorGemm<Gemm::GemmType<Element, layout::ColumnMajor>, Gemm::GemmType<Element, layout::RowMajor>>{
331+ using L1AType = Gemm::GemmType<Element, layout::nN, AscendC::TPosition::A1>;
332+ using L1BType = Gemm::GemmType<Element, layout::zZ, AscendC::TPosition::B1>;
333+ using L0AType = Gemm::GemmType<Element, layout::zZ, AscendC::TPosition::A2>;
334+ using L0BType = Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::B2>;
335+};
336+ 
337+template<>
338+struct L1AndL0TypeSelectorGemm<Gemm::GemmType<int8_t, layout::ColumnMajor>, Gemm::GemmType<int8_t, layout::RowMajor>>{
339+ using L1AType = Gemm::GemmType<int8_t, layout::nZ, AscendC::TPosition::A1>;
340+ using L1BType = Gemm::GemmType<int8_t, layout::zN, AscendC::TPosition::B1>;
341+ using L0AType = Gemm::GemmType<int8_t, layout::zZ, AscendC::TPosition::A2>;
342+ using L0BType = Gemm::GemmType<int8_t, layout::nZ, AscendC::TPosition::B2>;
343+};
344+///////////////////////////////////////
345+} // namespace NpuArch::Gemm::helper
346+ 
347+#endif // GEMM_HELPER_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/copy_gm_to_l1_a2.hpp+1068-0
@@ -0,0 +1,1068 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_COPY_GM_TO_L1_HPP
12+#define GEMM_TILE_COPY_GM_TO_L1_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/layout/layout.hpp"
17+#include "../../../attn_infra/gemm/gemm_type.hpp"
18+#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp"
19+#include "../../../tla/tensor.hpp"
20+ 
21+namespace NpuArch::Gemm::Tile {
22+ 
23+template <
24+ class ArchTag,
25+ /// GemmType for matrix operand
26+ class GmType,
27+ class L1Type = void
28+>
29+struct CopyGmToL1 {
30+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy gm to l1, can not find the specialization.");
31+};
32+ 
33+template <
34+ class ArchTag,
35+ /// GemmType for matrix operand
36+ class GmType,
37+ class L1Type = void
38+>
39+struct CopyGmToL1IntervalDataCopy {
40+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy gm to l1, can not find the specialization.");
41+};
42+ 
43+////////////////////////////////////////
44+/// Using the standard strided DataCopy interface to implement nd2nz
45+/// transfer may achieve higher data transfer efficiency when the data block shape is short and wide
46+/// Partial specialization for AtlasA2, half, RowMajor in and zN out.
47+template<>
48+struct CopyGmToL1IntervalDataCopy<Arch::AtlasA2, Gemm::GemmType<half, layout::RowMajor>> {
49+ using LayoutDst = layout::zN;
50+ using LayoutSrc = layout::RowMajor;
51+ using Element = half;
52+ 
53+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
54+ 
55+ // Mehtods
56+ 
57+ __aicore__ inline
58+ CopyGmToL1IntervalDataCopy() {};
59+ 
60+ __aicore__ inline
61+ void operator()(
62+ AscendC::LocalTensor<Element> const &dstTensor,
63+ AscendC::GlobalTensor<Element> const &srcTensor,
64+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
65+ {
66+ for (int i = 0; i < layoutSrc.shape(0); ++i) {
67+ AscendC::DataCopyParams dataCopyParams(
68+ CeilDiv(layoutSrc.shape(1), layoutDst.shape(2)),
69+ layoutDst.shape(2) / ELE_NUM_PER_C0,
70+ 0,
71+ (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0
72+ );
73+ AscendC::DataCopy(dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams);
74+ }
75+ }
76+};
77+ 
78+/// Partial specialization for AtlasA2, half, PaddingRowMajor in and zN out.
79+/// Using the standard strided DataCopy interface to implement nd2nz
80+/// transfer may achieve higher data transfer efficiency when the data block shape is short and wide
81+template<>
82+struct CopyGmToL1IntervalDataCopy<Arch::AtlasA2, Gemm::GemmType<half, layout::PaddingRowMajor>> {
83+ using LayoutDst = layout::zN;
84+ using LayoutSrc = layout::PaddingRowMajor;
85+ using Element = half;
86+ 
87+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
88+ 
89+ // Mehtods
90+ 
91+ __aicore__ inline
92+ CopyGmToL1IntervalDataCopy() {};
93+ 
94+ __aicore__ inline
95+ void operator()(
96+ AscendC::LocalTensor<Element> const &dstTensor,
97+ AscendC::GlobalTensor<Element> const &srcTensor,
98+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
99+ {
100+ for (int i = 0; i < layoutSrc.orgShape(0); ++i) {
101+ AscendC::DataCopyParams dataCopyParams(
102+ CeilDiv(layoutSrc.orgShape(1), layoutDst.shape(2)),
103+ layoutDst.shape(2) / ELE_NUM_PER_C0,
104+ 0,
105+ (layoutDst.stride(3) - layoutDst.shape(2)) / ELE_NUM_PER_C0
106+ );
107+ AscendC::DataCopy(dstTensor[i * layoutDst.shape(2)], srcTensor[i * layoutSrc.stride(0)], dataCopyParams);
108+ }
109+ }
110+};
111+ 
112+/// Partial specialization for AtlasA2, half, ColumnMajor in and zN out.
113+/// Using the standard strided DataCopy interface to implement nd2nz
114+/// transfer may achieve higher data transfer efficiency when the data block shape is tall and narrow
115+template<>
116+struct CopyGmToL1IntervalDataCopy<Arch::AtlasA2, Gemm::GemmType<half, layout::ColumnMajor>> {
117+ using LayoutDst = layout::nZ;
118+ using LayoutSrc = layout::ColumnMajor;
119+ using Element = half;
120+ 
121+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
122+ 
123+ // Mehtods
124+ 
125+ __aicore__ inline
126+ CopyGmToL1IntervalDataCopy() {};
127+ 
128+ __aicore__ inline
129+ void operator()(
130+ AscendC::LocalTensor<Element> const &dstTensor,
131+ AscendC::GlobalTensor<Element> const &srcTensor,
132+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
133+ {
134+ for (int i = 0; i < layoutSrc.shape(1); ++i) {
135+ AscendC::DataCopyParams dataCopyParams(
136+ CeilDiv(layoutSrc.shape(0), layoutDst.shape(0)),
137+ layoutDst.shape(0) / ELE_NUM_PER_C0,
138+ 0,
139+ (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0
140+ );
141+ AscendC::DataCopy(dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(1)], dataCopyParams);
142+ }
143+ }
144+};
145+ 
146+/// Partial specialization for AtlasA2, half, PaddingColumnMajor in and zN out.
147+/// Using the standard strided DataCopy interface to implement nd2nz
148+/// transfer may achieve higher data transfer efficiency when the data block shape is tall and narrow
149+template<>
150+struct CopyGmToL1IntervalDataCopy<Arch::AtlasA2, Gemm::GemmType<half, layout::PaddingColumnMajor>> {
151+ using LayoutDst = layout::nZ;
152+ using LayoutSrc = layout::PaddingColumnMajor;
153+ using Element = half;
154+ 
155+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
156+ 
157+ // Mehtods
158+ 
159+ __aicore__ inline
160+ CopyGmToL1IntervalDataCopy() {};
161+ 
162+ __aicore__ inline
163+ void operator()(
164+ AscendC::LocalTensor<Element> const &dstTensor,
165+ AscendC::GlobalTensor<Element> const &srcTensor,
166+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
167+ {
168+ for (int i = 0; i < layoutSrc.orgShape(1); ++i) {
169+ AscendC::DataCopyParams dataCopyParams(
170+ CeilDiv(layoutSrc.orgShape(0), layoutDst.shape(0)),
171+ layoutDst.shape(0) / ELE_NUM_PER_C0,
172+ 0,
173+ (layoutDst.stride(1) - layoutDst.shape(0)) / ELE_NUM_PER_C0
174+ );
175+ AscendC::DataCopy(dstTensor[i * layoutDst.shape(0)], srcTensor[i * layoutSrc.stride(2)], dataCopyParams);
176+ }
177+ }
178+};
179+ 
180+/// new add gemm
181+template <class ArchTag, class Element>
182+struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::RowMajor>, Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>> {
183+ using LayoutDst = layout::zN;
184+ using LayoutSrc = layout::RowMajor;
185+ 
186+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
187+ 
188+ // Mehtods
189+ 
190+ __aicore__ inline
191+ CopyGmToL1() {};
192+ 
193+ __aicore__ inline
194+ void operator()(
195+ AscendC::LocalTensor<Element> const &dstTensor,
196+ AscendC::GlobalTensor<Element> const &srcTensor,
197+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
198+ {
199+ AscendC::Nd2NzParams intriParams;
200+ 
201+ intriParams.ndNum = 1;
202+ intriParams.dValue = layoutSrc.shape(1);
203+ intriParams.srcNdMatrixStride = 0;
204+ intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0;
205+ intriParams.dstNzMatrixStride = 0;
206+ 
207+ if (layoutSrc.stride(0) < STRIDE_LIMIT) {
208+ intriParams.nValue = layoutSrc.shape(0);
209+ intriParams.srcDValue = layoutSrc.stride(0);
210+ intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0;
211+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
212+ } else {
213+ intriParams.nValue = 1;
214+ intriParams.srcDValue = 0;
215+ intriParams.dstNzNStride = 0;
216+ for (uint32_t i = 0; i < layoutSrc.shape(0); i++) {
217+ AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams);
218+ }
219+ }
220+ }
221+};
222+ 
223+template <class ArchTag, class Element>
224+struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::RowMajor>, Gemm::GemmType<Element, layout::zZ, AscendC::TPosition::B1>> {
225+ using LayoutDst = layout::zZ;
226+ using LayoutSrc = layout::RowMajor;
227+ 
228+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
229+ 
230+ // Mehtods
231+ 
232+ __aicore__ inline
233+ CopyGmToL1() {};
234+ 
235+ __aicore__ inline
236+ void operator()(
237+ AscendC::LocalTensor<Element> const &dstTensor,
238+ AscendC::GlobalTensor<Element> const &srcTensor,
239+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
240+ {
241+ AscendC::Nd2NzParams intriParams;
242+ uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(0);
243+ uint32_t ndNum = layoutSrc.shape(0) / C0_NUM_PER_FRACTAL;
244+ uint32_t remains = layoutSrc.shape(0) % C0_NUM_PER_FRACTAL;
245+ if (srcNdStride < STRIDE_LIMIT) {
246+ if (ndNum) {
247+ intriParams.ndNum = ndNum;
248+ intriParams.nValue = C0_NUM_PER_FRACTAL;
249+ intriParams.dValue = layoutSrc.shape(1);
250+ intriParams.srcNdMatrixStride = srcNdStride;
251+ intriParams.srcDValue = layoutSrc.stride(0);
252+ 
253+ intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0;
254+ intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0;
255+ 
256+ intriParams.dstNzMatrixStride = layoutDst.stride(1);
257+ 
258+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
259+ }
260+ 
261+ if (remains) {
262+ AscendC::Nd2NzParams tailParams;
263+ tailParams.ndNum = 1;
264+ tailParams.nValue = remains;
265+ tailParams.dValue = layoutSrc.shape(1);
266+ tailParams.srcNdMatrixStride = srcNdStride;
267+ tailParams.srcDValue = layoutSrc.stride(0);
268+ 
269+ tailParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0;
270+ tailParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0;
271+ tailParams.dstNzMatrixStride = 0; //`
272+ 
273+ AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(1)], srcTensor[ndNum * srcNdStride], tailParams);
274+ }
275+ } else if (layoutSrc.stride(0) < STRIDE_LIMIT) {
276+ for (uint32_t i = 0; i < ndNum; i++) {
277+ AscendC::Nd2NzParams intriParams;
278+ intriParams.ndNum = 1;
279+ intriParams.nValue = C0_NUM_PER_FRACTAL;
280+ intriParams.dValue = layoutSrc.shape(1);
281+ intriParams.srcNdMatrixStride = 0;
282+ intriParams.srcDValue = layoutSrc.stride(0);
283+ 
284+ intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0;
285+ intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0;
286+ intriParams.dstNzMatrixStride = 0;
287+ 
288+ AscendC::DataCopy(dstTensor[i * layoutDst.stride(1)], srcTensor[i * srcNdStride], intriParams);
289+ }
290+ if (remains) {
291+ AscendC::Nd2NzParams tailParams;
292+ tailParams.ndNum = 1;
293+ tailParams.nValue = remains;
294+ tailParams.dValue = layoutSrc.shape(1);
295+ tailParams.srcNdMatrixStride = 0;
296+ tailParams.srcDValue = layoutSrc.stride(0);
297+ 
298+ tailParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0;
299+ tailParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0;
300+ tailParams.dstNzMatrixStride = 0;
301+ 
302+ AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(1)], srcTensor[ndNum * srcNdStride], tailParams);
303+ }
304+ } else {
305+ for (uint32_t i = 0; i < layoutSrc.shape(0); i++) {
306+ uint32_t idxR0 = i / C0_NUM_PER_FRACTAL;
307+ uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL;
308+ 
309+ AscendC::Nd2NzParams intriParams;
310+ intriParams.ndNum = 1;
311+ intriParams.nValue = 1;
312+ intriParams.dValue = layoutSrc.shape(1);
313+ intriParams.srcNdMatrixStride = 0;
314+ intriParams.srcDValue = 0;
315+ 
316+ intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0;
317+ intriParams.dstNzNStride = 0;
318+ intriParams.dstNzMatrixStride = 0;
319+ 
320+ uint32_t offsetDst = i * idxR0 * layoutDst.stride(1) + idxInR0 * ELE_NUM_PER_C0;
321+ uint32_t offsetSrc = i * layoutSrc.stride(0);
322+ AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams);
323+ }
324+ }
325+ }
326+};
327+ 
328+template <class ArchTag, class Element>
329+struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::ColumnMajor>, Gemm::GemmType<Element, layout::nN, AscendC::TPosition::A1>> {
330+ using LayoutDst = layout::nN;
331+ using LayoutSrc = layout::ColumnMajor;
332+ 
333+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
334+ 
335+ // Mehtods
336+ 
337+ __aicore__ inline
338+ CopyGmToL1() {};
339+ 
340+ __aicore__ inline
341+ void operator()(
342+ AscendC::LocalTensor<Element> const &dstTensor,
343+ AscendC::GlobalTensor<Element> const &srcTensor,
344+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
345+ {
346+ AscendC::Nd2NzParams intriParams;
347+ uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(1);
348+ uint32_t ndNum = layoutSrc.shape(1) / C0_NUM_PER_FRACTAL;
349+ uint32_t remains = layoutSrc.shape(1) % C0_NUM_PER_FRACTAL;
350+ if (srcNdStride < STRIDE_LIMIT) {
351+ if (ndNum) {
352+ intriParams.ndNum = ndNum;
353+ intriParams.nValue = C0_NUM_PER_FRACTAL;
354+ intriParams.dValue = layoutSrc.shape(0);
355+ intriParams.srcNdMatrixStride = srcNdStride;
356+ intriParams.srcDValue = layoutSrc.stride(1);
357+ 
358+ intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
359+ intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0;
360+ 
361+ intriParams.dstNzMatrixStride = layoutDst.stride(3);
362+ 
363+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
364+ }
365+ 
366+ if (remains) {
367+ AscendC::Nd2NzParams tailParams;
368+ tailParams.ndNum = 1;
369+ tailParams.nValue = remains;
370+ tailParams.dValue = layoutSrc.shape(0);
371+ tailParams.srcNdMatrixStride = srcNdStride;
372+ tailParams.srcDValue = layoutSrc.stride(1);
373+ 
374+ tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
375+ tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0;
376+ tailParams.dstNzMatrixStride = 0;
377+ 
378+ AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams);
379+ }
380+ } else if (layoutSrc.stride(1) < STRIDE_LIMIT) {
381+ for (uint32_t i = 0; i < ndNum; i++) {
382+ AscendC::Nd2NzParams intriParams;
383+ intriParams.ndNum = 1;
384+ intriParams.nValue = C0_NUM_PER_FRACTAL;
385+ intriParams.dValue = layoutSrc.shape(0);
386+ intriParams.srcNdMatrixStride = 0;
387+ intriParams.srcDValue = layoutSrc.stride(1);
388+ 
389+ intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
390+ intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0;
391+ intriParams.dstNzMatrixStride = 0;
392+ 
393+ AscendC::DataCopy(dstTensor[i * layoutDst.stride(3)], srcTensor[i * srcNdStride], intriParams);
394+ }
395+ if (remains) {
396+ AscendC::Nd2NzParams tailParams;
397+ tailParams.ndNum = 1;
398+ tailParams.nValue = remains;
399+ tailParams.dValue = layoutSrc.shape(0);
400+ tailParams.srcNdMatrixStride = 0;
401+ tailParams.srcDValue = layoutSrc.stride(1);
402+ 
403+ tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
404+ tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0;
405+ tailParams.dstNzMatrixStride = 0;
406+ 
407+ AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams);
408+ }
409+ } else {
410+ for (uint32_t i = 0; i < layoutSrc.shape(1); i++) {
411+ uint32_t idxR0 = i / C0_NUM_PER_FRACTAL;
412+ uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL;
413+ 
414+ AscendC::Nd2NzParams intriParams;
415+ intriParams.ndNum = 1;
416+ intriParams.nValue = 1;
417+ intriParams.dValue = layoutSrc.shape(0);
418+ intriParams.srcNdMatrixStride = 0;
419+ intriParams.srcDValue = 0;
420+ 
421+ intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
422+ intriParams.dstNzNStride = 0;
423+ intriParams.dstNzMatrixStride = 0;
424+ 
425+ uint32_t offsetDst = i * idxR0 * layoutDst.stride(3) + idxInR0 * ELE_NUM_PER_C0;
426+ uint32_t offsetSrc = i * layoutSrc.stride(1);
427+ AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams);
428+ }
429+ }
430+ }
431+};
432+ 
433+template <class ArchTag, class Element>
434+struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::ColumnMajor>, Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::B1>> {
435+ using LayoutDst = layout::nZ;
436+ using LayoutSrc = layout::ColumnMajor;
437+ 
438+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
439+ 
440+ // Mehtods
441+ 
442+ __aicore__ inline
443+ CopyGmToL1() {};
444+ 
445+ __aicore__ inline
446+ void operator()(
447+ AscendC::LocalTensor<Element> const &dstTensor,
448+ AscendC::GlobalTensor<Element> const &srcTensor,
449+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
450+ {
451+ AscendC::Nd2NzParams intriParams;
452+ 
453+ intriParams.ndNum = 1;
454+ intriParams.dValue = layoutSrc.shape(0);
455+ intriParams.srcNdMatrixStride = 0;
456+ intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
457+ intriParams.dstNzMatrixStride = 0;
458+ 
459+ if (layoutSrc.stride(1) < STRIDE_LIMIT) {
460+ intriParams.nValue = layoutSrc.shape(1);
461+ intriParams.srcDValue = layoutSrc.stride(1);
462+ intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0;
463+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
464+ } else {
465+ intriParams.nValue = 1;
466+ intriParams.srcDValue = 0;
467+ intriParams.dstNzNStride = 0;
468+ for (uint32_t i = 0; i < layoutSrc.shape(1); i++) {
469+ AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams);
470+ }
471+ }
472+ }
473+};
474+ 
475+template <class ArchTag, class Element>
476+struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::ColumnMajor>, Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::A1>> {
477+ using LayoutDst = layout::nZ;
478+ using LayoutSrc = layout::ColumnMajor;
479+ 
480+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
481+ 
482+ // Mehtods
483+ 
484+ __aicore__ inline
485+ CopyGmToL1() {};
486+ 
487+ __aicore__ inline
488+ void operator()(
489+ AscendC::LocalTensor<Element> const &dstTensor,
490+ AscendC::GlobalTensor<Element> const &srcTensor,
491+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
492+ {
493+ AscendC::Nd2NzParams intriParams;
494+ 
495+ intriParams.ndNum = 1;
496+ intriParams.dValue = layoutSrc.shape(0);
497+ intriParams.srcNdMatrixStride = 0;
498+ intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
499+ intriParams.dstNzMatrixStride = 0;
500+ 
501+ if (layoutSrc.stride(1) < STRIDE_LIMIT) {
502+ intriParams.nValue = layoutSrc.shape(1);
503+ intriParams.srcDValue = layoutSrc.stride(1);
504+ intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0;
505+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
506+ } else {
507+ intriParams.nValue = 1;
508+ intriParams.srcDValue = 0;
509+ intriParams.dstNzNStride = 0;
510+ for (uint32_t i = 0; i < layoutSrc.shape(1); i++) {
511+ AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams);
512+ }
513+ }
514+ }
515+};
516+////////////////////////////////////////
517+ 
518+///////////////////////////////////////
519+/// new add gemv, VectorLayout -> zN
520+template <class ArchTag, class Element>
521+struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::VectorLayout>, Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>> {
522+ using LayoutDst = layout::zN;
523+ using LayoutSrc = layout::VectorLayout;
524+ 
525+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
526+ 
527+ // Methods
528+ 
529+ __aicore__ inline
530+ CopyGmToL1() {};
531+ 
532+ __aicore__ inline
533+ void operator()(
534+ AscendC::LocalTensor<Element> const &dstTensor,
535+ AscendC::GlobalTensor<Element> const &srcTensor,
536+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
537+ {
538+ AscendC::Nd2NzParams intriParams;
539+ 
540+ intriParams.ndNum = 1;
541+ intriParams.dValue = layoutSrc.shape(0);
542+ intriParams.srcNdMatrixStride = 0;
543+ intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0;
544+ intriParams.dstNzMatrixStride = 0;
545+ intriParams.nValue = 1;
546+ intriParams.srcDValue = layoutSrc.shape(0);
547+ intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0;
548+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
549+ }
550+};
551+ 
552+ 
553+ 
554+///////////////////////////////////////
555+/// new add gemv, ColumnMajor -> nN
556+template <class ArchTag, class Element>
557+struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::ColumnMajor>, Gemm::GemmType<Element, layout::nN, AscendC::TPosition::B1>> {
558+ using LayoutDst = layout::nN;
559+ using LayoutSrc = layout::ColumnMajor;
560+ 
561+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
562+ 
563+ // Methods
564+ 
565+ __aicore__ inline
566+ CopyGmToL1() {};
567+ 
568+ __aicore__ inline
569+ void operator()(
570+ AscendC::LocalTensor<Element> const &dstTensor,
571+ AscendC::GlobalTensor<Element> const &srcTensor,
572+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
573+ {
574+ AscendC::Nd2NzParams intriParams;
575+ uint32_t srcNdStride = C0_NUM_PER_FRACTAL * layoutSrc.stride(1);
576+ uint32_t ndNum = layoutSrc.shape(1) / C0_NUM_PER_FRACTAL;
577+ uint32_t remains = layoutSrc.shape(1) % C0_NUM_PER_FRACTAL;
578+ if (srcNdStride < STRIDE_LIMIT) {
579+ if (ndNum) {
580+ intriParams.ndNum = ndNum;
581+ intriParams.nValue = C0_NUM_PER_FRACTAL;
582+ intriParams.dValue = layoutSrc.shape(0);
583+ intriParams.srcNdMatrixStride = srcNdStride;
584+ intriParams.srcDValue = layoutSrc.stride(1);
585+ 
586+ intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
587+ intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0;
588+ 
589+ intriParams.dstNzMatrixStride = layoutDst.stride(3);
590+ 
591+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
592+ }
593+ 
594+ if (remains) {
595+ AscendC::Nd2NzParams tailParams;
596+ tailParams.ndNum = 1;
597+ tailParams.nValue = remains;
598+ tailParams.dValue = layoutSrc.shape(0);
599+ tailParams.srcNdMatrixStride = srcNdStride;
600+ tailParams.srcDValue = layoutSrc.stride(1);
601+ 
602+ tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
603+ tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0;
604+ tailParams.dstNzMatrixStride = 0;
605+ 
606+ AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams);
607+ }
608+ } else if (layoutSrc.stride(1) < STRIDE_LIMIT) {
609+ for (uint32_t i = 0; i < ndNum; i++) {
610+ AscendC::Nd2NzParams intriParams;
611+ intriParams.ndNum = 1;
612+ intriParams.nValue = C0_NUM_PER_FRACTAL;
613+ intriParams.dValue = layoutSrc.shape(0);
614+ intriParams.srcNdMatrixStride = 0;
615+ intriParams.srcDValue = layoutSrc.stride(1);
616+ 
617+ intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
618+ intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0;
619+ intriParams.dstNzMatrixStride = 0;
620+ 
621+ AscendC::DataCopy(dstTensor[i * layoutDst.stride(3)], srcTensor[i * srcNdStride], intriParams);
622+ }
623+ if (remains) {
624+ AscendC::Nd2NzParams tailParams;
625+ tailParams.ndNum = 1;
626+ tailParams.nValue = remains;
627+ tailParams.dValue = layoutSrc.shape(0);
628+ tailParams.srcNdMatrixStride = 0;
629+ tailParams.srcDValue = layoutSrc.stride(1);
630+ 
631+ tailParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
632+ tailParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0;
633+ tailParams.dstNzMatrixStride = 0;
634+ 
635+ AscendC::DataCopy(dstTensor[ndNum * layoutDst.stride(3)], srcTensor[ndNum * srcNdStride], tailParams);
636+ }
637+ } else {
638+ for (uint32_t i = 0; i < layoutSrc.shape(1); i++) {
639+ uint32_t idxR0 = i / C0_NUM_PER_FRACTAL;
640+ uint32_t idxInR0 = i % C0_NUM_PER_FRACTAL;
641+ 
642+ AscendC::Nd2NzParams intriParams;
643+ intriParams.ndNum = 1;
644+ intriParams.nValue = 1;
645+ intriParams.dValue = layoutSrc.shape(0);
646+ intriParams.srcNdMatrixStride = 0;
647+ intriParams.srcDValue = 0;
648+ 
649+ intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
650+ intriParams.dstNzNStride = 0;
651+ intriParams.dstNzMatrixStride = 0;
652+ 
653+ uint32_t offsetDst = i * idxR0 * layoutDst.stride(3) + idxInR0 * ELE_NUM_PER_C0;
654+ uint32_t offsetSrc = i * layoutSrc.stride(1);
655+ AscendC::DataCopy(dstTensor[offsetDst], srcTensor[offsetSrc], intriParams);
656+ }
657+ }
658+ }
659+};
660+ 
661+template <class ArchTag, class Element>
662+struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::RowMajor>, Gemm::GemmType<Element, layout::zN, AscendC::TPosition::B1>> {
663+ using LayoutDst = layout::zN;
664+ using LayoutSrc = layout::RowMajor;
665+ 
666+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
667+ 
668+ // Methods
669+ 
670+ __aicore__ inline
671+ CopyGmToL1() {};
672+ 
673+ __aicore__ inline
674+ void operator()(
675+ AscendC::LocalTensor<Element> const &dstTensor,
676+ AscendC::GlobalTensor<Element> const &srcTensor,
677+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
678+ {
679+ AscendC::Nd2NzParams intriParams;
680+ 
681+ intriParams.ndNum = 1;
682+ intriParams.dValue = layoutSrc.shape(1);
683+ intriParams.srcNdMatrixStride = 0;
684+ intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0;
685+ intriParams.dstNzMatrixStride = 0;
686+ 
687+ if (layoutSrc.stride(0) < STRIDE_LIMIT) {
688+ intriParams.nValue = layoutSrc.shape(0);
689+ intriParams.srcDValue = layoutSrc.stride(0);
690+ intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0;
691+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
692+ } else {
693+ intriParams.nValue = 1;
694+ intriParams.srcDValue = 0;
695+ intriParams.dstNzNStride = 0;
696+ for (uint32_t i = 0; i < layoutSrc.shape(0); i++) {
697+ AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams);
698+ }
699+ }
700+ }
701+};
702+/////////////////////////////////
703+ 
704+/// Partial specialization for AtlasA2, RowMajor in and zN out.
705+template <class Element>
706+struct CopyGmToL1<Arch::AtlasA2, Gemm::GemmType<Element, layout::RowMajor>> {
707+ using LayoutDst = layout::zN;
708+ using LayoutSrc = layout::RowMajor;
709+ 
710+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
711+ 
712+ // Mehtods
713+ 
714+ __aicore__ inline
715+ CopyGmToL1() {};
716+ 
717+ __aicore__ inline
718+ void operator()(
719+ AscendC::LocalTensor<Element> const &dstTensor,
720+ AscendC::GlobalTensor<Element> const &srcTensor,
721+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
722+ {
723+ AscendC::Nd2NzParams intriParams;
724+ 
725+ intriParams.ndNum = 1;
726+ intriParams.dValue = layoutSrc.shape(1);
727+ intriParams.srcNdMatrixStride = 0;
728+ intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0;
729+ intriParams.dstNzMatrixStride = 0;
730+ 
731+ if (layoutSrc.stride(0) < STRIDE_LIMIT) {
732+ intriParams.nValue = layoutSrc.shape(0);
733+ intriParams.srcDValue = layoutSrc.stride(0);
734+ intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0;
735+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
736+ } else {
737+ intriParams.nValue = 1;
738+ intriParams.srcDValue = 0;
739+ intriParams.dstNzNStride = 0;
740+ for (uint32_t i = 0; i < layoutSrc.shape(0); i++) {
741+ AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(0)], intriParams);
742+ }
743+ }
744+ }
745+ 
746+ // layoutSrc must be the layout of one of the src matrices
747+ __aicore__ inline
748+ void operator()(
749+ AscendC::LocalTensor<Element> const &dstTensor,
750+ AscendC::GlobalTensor<Element> const &srcTensor,
751+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc,
752+ uint32_t ndNum, uint32_t srcNdMatrixStride,
753+ uint32_t dstNzNStride, uint32_t dstNzMatrixStride,
754+ uint32_t dstNzC0Stride)
755+ {
756+ AscendC::Nd2NzParams intriParams;
757+ 
758+ intriParams.nValue = layoutSrc.shape(0);
759+ intriParams.dValue = layoutSrc.shape(1);
760+ intriParams.srcDValue = layoutSrc.stride(0);
761+ intriParams.dstNzNStride = dstNzNStride;
762+ intriParams.dstNzC0Stride = dstNzC0Stride;
763+ if (srcNdMatrixStride < STRIDE_LIMIT) {
764+ intriParams.ndNum = ndNum;
765+ intriParams.srcNdMatrixStride = srcNdMatrixStride;
766+ intriParams.dstNzMatrixStride = dstNzMatrixStride;
767+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
768+ } else {
769+ intriParams.ndNum = 1;
770+ intriParams.srcNdMatrixStride = 0;
771+ intriParams.dstNzMatrixStride = 0;
772+ for (uint32_t i = 0; i < ndNum; i++) {
773+ AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * srcNdMatrixStride], intriParams);
774+ }
775+ }
776+ }
777+};
778+ 
779+/// Partial specialization for AtlasA2, ColumnMajor in and nZ out.
780+template <
781+ class Element
782+>
783+struct CopyGmToL1<Arch::AtlasA2, Gemm::GemmType<Element, layout::ColumnMajor>> {
784+ using LayoutDst = layout::nZ;
785+ using LayoutSrc = layout::ColumnMajor;
786+ 
787+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
788+ 
789+ // Mehtods
790+ 
791+ __aicore__ inline
792+ CopyGmToL1() {};
793+ 
794+ __aicore__ inline
795+ void operator()(
796+ AscendC::LocalTensor<Element> const &dstTensor,
797+ AscendC::GlobalTensor<Element> const &srcTensor,
798+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
799+ {
800+ AscendC::Nd2NzParams intriParams;
801+ 
802+ intriParams.ndNum = 1;
803+ intriParams.dValue = layoutSrc.shape(0);
804+ intriParams.srcNdMatrixStride = 0;
805+ intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
806+ intriParams.dstNzMatrixStride = 0;
807+ 
808+ if (layoutSrc.stride(1) < STRIDE_LIMIT) {
809+ intriParams.nValue = layoutSrc.shape(1);
810+ intriParams.srcDValue = layoutSrc.stride(1);
811+ intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0;
812+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
813+ } else {
814+ intriParams.nValue = 1;
815+ intriParams.srcDValue = 0;
816+ intriParams.dstNzNStride = 0;
817+ for (uint32_t i = 0; i < layoutSrc.shape(1); i++) {
818+ AscendC::DataCopy(dstTensor[i * ELE_NUM_PER_C0], srcTensor[i * layoutSrc.stride(1)], intriParams);
819+ }
820+ }
821+ }
822+};
823+ 
824+/// Partial specialization for zN in and zN out.
825+template <
826+ class ArchTag,
827+ class Element
828+>
829+struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::zN>> {
830+ using LayoutDst = layout::zN;
831+ using LayoutSrc = layout::zN;
832+ 
833+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
834+ 
835+ // Mehtods
836+ 
837+ __aicore__ inline
838+ CopyGmToL1() {};
839+ 
840+ __aicore__ inline
841+ void operator()(
842+ AscendC::LocalTensor<Element> const &dstTensor,
843+ AscendC::GlobalTensor<Element> const &srcTensor,
844+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
845+ {
846+ uint32_t blockCount = CeilDiv<ELE_NUM_PER_C0>(layoutSrc.orgShape(1));
847+ uint32_t blockLen = RoundUp<C0_NUM_PER_FRACTAL>(layoutSrc.orgShape(0));
848+ 
849+ AscendC::DataCopyParams repeatParams;
850+ 
851+ if (layoutSrc.stride(3) / ELE_NUM_PER_C0 < STRIDE_LIMIT) {
852+ repeatParams.blockCount = blockCount;
853+ repeatParams.blockLen = blockLen;
854+ repeatParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_C0 - blockLen;
855+ repeatParams.dstStride = layoutDst.stride(3) / ELE_NUM_PER_C0 - blockLen;
856+ AscendC::DataCopy(dstTensor, srcTensor, repeatParams);
857+ } else {
858+ repeatParams.blockCount = 1;
859+ repeatParams.blockLen = blockLen;
860+ repeatParams.srcStride = 0;
861+ repeatParams.dstStride = 0;
862+ for (uint32_t i = 0; i < blockCount; i++) {
863+ uint64_t dstOffset = i * layoutDst.stride(3);
864+ uint64_t srcOffset = i * layoutSrc.stride(3);
865+ AscendC::DataCopy(dstTensor[dstOffset], srcTensor[srcOffset], repeatParams);
866+ }
867+ }
868+ }
869+};
870+ 
871+/// Partial specialization for nZ in and nZ out.
872+template <
873+ class ArchTag,
874+ class Element
875+>
876+struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::nZ>> {
877+ using LayoutDst = layout::nZ;
878+ using LayoutSrc = layout::nZ;
879+ 
880+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
881+ 
882+ // Mehtods
883+ 
884+ __aicore__ inline
885+ CopyGmToL1() {};
886+ 
887+ __aicore__ inline
888+ void operator()(
889+ AscendC::LocalTensor<Element> const &dstTensor,
890+ AscendC::GlobalTensor<Element> const &srcTensor,
891+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
892+ {
893+ uint32_t blockCount = CeilDiv<ELE_NUM_PER_C0>(layoutSrc.orgShape(0));
894+ uint32_t blockLen = RoundUp<C0_NUM_PER_FRACTAL>(layoutSrc.orgShape(1));
895+ 
896+ AscendC::DataCopyParams repeatParams;
897+ 
898+ if (layoutSrc.stride(1) / ELE_NUM_PER_C0 < STRIDE_LIMIT) {
899+ repeatParams.blockCount = blockCount;
900+ repeatParams.blockLen = blockLen;
901+ repeatParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_C0 - blockLen;
902+ repeatParams.dstStride = layoutDst.stride(1) / ELE_NUM_PER_C0 - blockLen;
903+ AscendC::DataCopy(dstTensor, srcTensor, repeatParams);
904+ } else {
905+ repeatParams.blockCount = 1;
906+ repeatParams.blockLen = blockLen;
907+ repeatParams.srcStride = 0;
908+ repeatParams.dstStride = 0;
909+ for (uint32_t i = 0; i < blockCount; i++) {
910+ uint64_t dstOffset = i * layoutDst.stride(1);
911+ uint64_t srcOffset = i * layoutSrc.stride(1);
912+ AscendC::DataCopy(dstTensor[dstOffset], srcTensor[srcOffset], repeatParams);
913+ }
914+ }
915+ }
916+};
917+ 
918+/// Partial specialization for AtlasA2, PaddingRowMajor in and zN out.
919+template <class Element>
920+struct CopyGmToL1<Arch::AtlasA2, Gemm::GemmType<Element, layout::PaddingRowMajor>> {
921+ using LayoutDst = layout::zN;
922+ using LayoutSrc = layout::PaddingRowMajor;
923+ 
924+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
925+ 
926+ // Mehtods
927+ 
928+ __aicore__ inline
929+ CopyGmToL1() {};
930+ 
931+ __aicore__ inline
932+ void operator()(
933+ AscendC::LocalTensor<Element> const &dstTensor,
934+ AscendC::GlobalTensor<Element> const &srcTensor,
935+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
936+ {
937+ AscendC::Nd2NzParams intriParams;
938+ 
939+ intriParams.ndNum = 1;
940+ intriParams.dValue = layoutSrc.orgShape(1);
941+ intriParams.srcNdMatrixStride = 0;
942+ intriParams.dstNzC0Stride = layoutDst.stride(3) / ELE_NUM_PER_C0;
943+ intriParams.dstNzMatrixStride = 0;
944+ 
945+ intriParams.nValue = layoutSrc.orgShape(0);
946+ intriParams.srcDValue = layoutSrc.stride(0);
947+ intriParams.dstNzNStride = layoutDst.stride(0) / ELE_NUM_PER_C0;
948+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
949+ }
950+};
951+ 
952+/// Partial specialization for AtlasA2, ColumnMajor in and nZ out.
953+template <
954+ class Element
955+>
956+struct CopyGmToL1<Arch::AtlasA2, Gemm::GemmType<Element, layout::PaddingColumnMajor>> {
957+ using LayoutDst = layout::nZ;
958+ using LayoutSrc = layout::PaddingColumnMajor;
959+ 
960+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
961+ 
962+ // Mehtods
963+ 
964+ __aicore__ inline
965+ CopyGmToL1() {};
966+ 
967+ __aicore__ inline
968+ void operator()(
969+ AscendC::LocalTensor<Element> const &dstTensor,
970+ AscendC::GlobalTensor<Element> const &srcTensor,
971+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
972+ {
973+ AscendC::Nd2NzParams intriParams;
974+ 
975+ intriParams.ndNum = 1;
976+ intriParams.dValue = layoutSrc.orgShape(0);
977+ intriParams.srcNdMatrixStride = 0;
978+ intriParams.dstNzC0Stride = layoutDst.stride(1) / ELE_NUM_PER_C0;
979+ intriParams.dstNzMatrixStride = 0;
980+ 
981+ intriParams.nValue = layoutSrc.orgShape(1);
982+ intriParams.srcDValue = layoutSrc.stride(2);
983+ intriParams.dstNzNStride = layoutDst.stride(2) / ELE_NUM_PER_C0;
984+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
985+ }
986+};
987+ 
988+/// Partial specialization for AtlasA2, RowMajor in and RowMajor out.
989+template <class Element>
990+struct CopyGmToL1<Arch::AtlasA2, Gemm::GemmType<Element, layout::RowMajor>,
991+ Gemm::GemmType<Element, layout::RowMajor, AscendC::TPosition::A1>> {
992+ using LayoutDst = layout::RowMajor;
993+ using LayoutSrc = layout::RowMajor;
994+ 
995+ static constexpr uint32_t ELE_NUM_PER_BLK = BYTE_PER_BLK / sizeof(Element);
996+ static constexpr uint32_t BLOCK_LEN_LIMIT = 65536;
997+ static constexpr uint32_t MAX_REPEAT = 4095;
998+ 
999+ // Mehtods
1000+ 
1001+ __aicore__ inline
1002+ CopyGmToL1() {};
1003+ 
1004+ __aicore__ inline
1005+ void operator()(
1006+ AscendC::LocalTensor<Element> const &dstTensor,
1007+ AscendC::GlobalTensor<Element> const &srcTensor,
1008+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
1009+ {
1010+ uint32_t rows = layoutSrc.shape(0);
1011+ uint32_t cols = layoutSrc.shape(1);
1012+ uint32_t srcStride = (layoutSrc.stride(0) - layoutSrc.shape(1)) / ELE_NUM_PER_BLK;
1013+ uint32_t dstStride = (layoutDst.stride(0) - layoutDst.shape(1)) / ELE_NUM_PER_BLK;
1014+ 
1015+ if ((layoutSrc.shape(1) == layoutSrc.stride(0)) && (layoutDst.shape(1) == layoutDst.stride(0))) {
1016+ DataCopy(dstTensor, srcTensor, rows * cols);
1017+ } else if (srcStride < STRIDE_LIMIT && dstStride < STRIDE_LIMIT && (cols / ELE_NUM_PER_BLK) < BLOCK_LEN_LIMIT) {
1018+ uint32_t rLoops = CeilDiv(rows, MAX_REPEAT);
1019+ for (uint32_t i = 0; i < rLoops; ++i) {
1020+ uint32_t rActual = (i < rLoops - 1) ? MAX_REPEAT : rows - i * MAX_REPEAT;
1021+ AscendC::DataCopyParams dataCopyParams(
1022+ rActual, cols / ELE_NUM_PER_BLK, srcStride, dstStride
1023+ );
1024+ DataCopy(dstTensor[i * MAX_REPEAT * layoutDst.stride(0)],
1025+ srcTensor[i * MAX_REPEAT * layoutSrc.stride(0)], dataCopyParams);
1026+ }
1027+ } else {
1028+ for (uint32_t i = 0; i < rows; ++i) {
1029+ DataCopy(dstTensor[i * layoutDst.stride(0)], srcTensor[i * layoutSrc.stride(0)], cols);
1030+ }
1031+ }
1032+ }
1033+};
1034+ 
1035+template <class ArchTag, class Element>
1036+struct CopyGmToL1<ArchTag, Gemm::GemmType<Element, layout::VectorLayout, AscendC::TPosition::GM>,
1037+ Gemm::GemmType<Element, layout::VectorLayout, AscendC::TPosition::A1>> {
1038+ using LayoutDst = layout::VectorLayout;
1039+ using LayoutSrc = layout::VectorLayout;
1040+ 
1041+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
1042+ 
1043+ // Mehtods
1044+ 
1045+ __aicore__ inline
1046+ CopyGmToL1() {};
1047+ 
1048+ __aicore__ inline
1049+ void operator()(
1050+ AscendC::LocalTensor<Element> const &dstTensor,
1051+ AscendC::GlobalTensor<Element> const &srcTensor,
1052+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
1053+ {
1054+ AscendC::DataCopyParams intriParams;
1055+ intriParams.blockCount = 1;
1056+ intriParams.blockLen = layoutDst.shape(0) / ELE_NUM_PER_C0;
1057+ intriParams.srcStride = 0;
1058+ intriParams.dstStride = 0;
1059+ AscendC::DataCopy(dstTensor, srcTensor, intriParams);
1060+ }
1061+};
1062+ 
1063+ 
1064+/////////////////////////////////////////////////////////////////////////////////////////////////////////
1065+ 
1066+} // namespace NpuArch::Gemm::Tile
1067+ 
1068+#endif // GEMM_TILE_COPY_GM_TO_L1_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/copy_gm_to_l1_a5.hpp+221-0
@@ -0,0 +1,221 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_COPY_GM_TO_L1_A5_HPP
12+#define GEMM_TILE_COPY_GM_TO_L1_A5_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp"
17+#include "../../../tla/tensor.hpp"
18+ 
19+namespace NpuArch::Gemm::Tile {
20+ 
21+/// Partial specialization for CopyGmToL1, AtlasA5, RowMajor in and zN out.
22+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
23+struct TileCopyTla<
24+ Arch::AtlasA5,
25+ tla::Tensor<AscendC::GlobalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::GM>,
26+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::A1>,
27+ std::enable_if_t<tla::detail::isRowMajor<LayoutSrc>::value && tla::detail::iszN<ElementDst, LayoutDst>::value>> {
28+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
29+ 
30+ // Methods
31+ 
32+ __aicore__ inline
33+ TileCopyTla() {};
34+ 
35+ template <class TensorDst, class TensorSrc>
36+ __aicore__ inline void operator()(
37+ TensorDst const &dstTensor,
38+ TensorSrc const &srcTensor,
39+ uint32_t ndNum = 1,
40+ uint32_t srcNdMatrixStride = 0,
41+ uint32_t dstNzMatrixStride = 0
42+ )
43+ {
44+ static_assert(
45+ tla::detail::isRowMajor<typename TensorSrc::Layout>::value
46+ && tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
47+ && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1,
48+ "The input parameters do not match. TensorSrc must be GM and RowMajor, while TensorDst must be L1 and zN"
49+ );
50+ 
51+ const uint32_t nValue = tla::get<0>(srcTensor.shape());
52+ const uint32_t dValue = tla::get<1>(srcTensor.shape());
53+ const uint32_t srcDValue = tla::get<0>(srcTensor.stride());
54+ const uint32_t dstInnerStrideRow = tla::get<0, 0>(dstTensor.stride());
55+ const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
56+ 
57+ AscendC::Nd2NzParams intriParams;
58+ 
59+ intriParams.ndNum = ndNum;
60+ intriParams.nValue = nValue;
61+ intriParams.dValue = dValue;
62+ intriParams.srcNdMatrixStride = srcNdMatrixStride;
63+ intriParams.srcDValue = srcDValue;
64+ intriParams.dstNzC0Stride = dstOuterStrideCol / ELE_NUM_PER_C0;
65+ intriParams.dstNzNStride = dstInnerStrideRow / ELE_NUM_PER_C0;
66+ intriParams.dstNzMatrixStride = dstNzMatrixStride;
67+ 
68+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
69+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
70+ 
71+ AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
72+ }
73+};
74+ 
75+/// Partial specialization for CopyGmToL1, AtlasA5, zN in and zN out.
76+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
77+struct TileCopyTla<
78+ Arch::AtlasA5,
79+ tla::Tensor<AscendC::GlobalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::GM>,
80+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::A1>,
81+ std::enable_if_t<tla::detail::iszN<ElementSrc, LayoutSrc>::value && tla::detail::iszN<ElementDst, LayoutDst>::value>> {
82+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
83+ 
84+ // Methods
85+ 
86+ __aicore__ inline
87+ TileCopyTla() {};
88+ 
89+ template <class TensorDst, class TensorSrc>
90+ __aicore__ inline void operator()(
91+ TensorDst const &dstTensor,
92+ TensorSrc const &srcTensor
93+ )
94+ {
95+ static_assert(
96+ tla::detail::iszN<typename TensorSrc::Element, typename TensorSrc::Layout>::value
97+ && tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
98+ && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1,
99+ "The input parameters do not match. TensorSrc must be GM and zN, while TensorDst must be L1 and zN"
100+ );
101+ 
102+ const uint32_t blockCount = tla::get<1, 1>(srcTensor.shape());
103+ const uint32_t blockLen = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape());
104+ 
105+ AscendC::DataCopyParams repeatParams;
106+ 
107+ repeatParams.blockCount = blockCount;
108+ repeatParams.blockLen = blockLen;
109+ repeatParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / ELE_NUM_PER_C0 - blockLen;
110+ repeatParams.dstStride = tla::get<1, 1>(dstTensor.stride()) / ELE_NUM_PER_C0 - blockLen;
111+ 
112+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
113+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
114+ 
115+ AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams);
116+ }
117+};
118+ 
119+/// Partial specialization for CopyGmToL1, AtlasA5, ColumnMajor in and nZ out.
120+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
121+struct TileCopyTla<
122+ Arch::AtlasA5,
123+ tla::Tensor<AscendC::GlobalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::GM>,
124+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::A1>,
125+ std::enable_if_t<tla::detail::isColumnMajor<LayoutSrc>::value && tla::detail::isnZ<ElementDst, LayoutDst>::value>> {
126+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
127+ 
128+ // Methods
129+ 
130+ __aicore__ inline
131+ TileCopyTla() {};
132+ 
133+ template <class TensorDst, class TensorSrc>
134+ __aicore__ inline void operator()(
135+ TensorDst const &dstTensor,
136+ TensorSrc const &srcTensor,
137+ uint32_t ndNum = 1,
138+ uint32_t srcNdMatrixStride = 0,
139+ uint32_t dstNzMatrixStride = 0
140+ )
141+ {
142+ static_assert(
143+ tla::detail::isColumnMajor<typename TensorSrc::Layout>::value
144+ && tla::detail::isnZ<typename TensorDst::Element, typename TensorDst::Layout>::value
145+ && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1,
146+ "The input parameters do not match. TensorSrc must be GM and ColumnMajor, "
147+ "while TensorDst must be L1 and nZ"
148+ );
149+ 
150+ const uint32_t nValue = tla::get<1>(srcTensor.shape());
151+ const uint32_t dValue = tla::get<0>(srcTensor.shape());
152+ const uint32_t srcDValue = tla::get<1>(srcTensor.stride());
153+ const uint32_t dstInnerStrideCol = tla::get<1, 0>(dstTensor.stride());
154+ const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
155+ 
156+ AscendC::Nd2NzParams intriParams;
157+ 
158+ intriParams.ndNum = ndNum;
159+ intriParams.nValue = nValue;
160+ intriParams.dValue = dValue;
161+ intriParams.srcNdMatrixStride = srcNdMatrixStride;
162+ intriParams.srcDValue = srcDValue;
163+ intriParams.dstNzC0Stride = dstOuterStrideRow / ELE_NUM_PER_C0;
164+ intriParams.dstNzNStride = dstInnerStrideCol / ELE_NUM_PER_C0;
165+ intriParams.dstNzMatrixStride = dstNzMatrixStride;
166+ 
167+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
168+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
169+ 
170+ AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
171+ }
172+};
173+ 
174+/// Partial specialization for CopyGmToL1, AtlasA5, nZ in and nZ out.
175+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
176+struct TileCopyTla<
177+ Arch::AtlasA5,
178+ tla::Tensor<AscendC::GlobalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::GM>,
179+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::A1>,
180+ std::enable_if_t<tla::detail::isnZ<ElementSrc, LayoutSrc>::value && tla::detail::isnZ<ElementDst, LayoutDst>::value>> {
181+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
182+ 
183+ // Methods
184+ 
185+ __aicore__ inline
186+ TileCopyTla() {};
187+ 
188+ template <class TensorDst, class TensorSrc>
189+ __aicore__ inline void operator()(
190+ TensorDst const &dstTensor,
191+ TensorSrc const &srcTensor
192+ )
193+ {
194+ static_assert(
195+ tla::detail::isnZ<typename TensorSrc::Element, typename TensorSrc::Layout>::value
196+ && tla::detail::isnZ<typename TensorDst::Element, typename TensorDst::Layout>::value
197+ && TensorSrc::position == AscendC::TPosition::GM && TensorDst::position == AscendC::TPosition::A1,
198+ "The input parameters do not match. TensorSrc must be GM and nZ, "
199+ "while TensorDst must be L1 and nZ"
200+ );
201+ 
202+ const uint32_t blockCount = tla::get<0, 1>(srcTensor.shape());
203+ const uint32_t blockLen = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());
204+ 
205+ AscendC::DataCopyParams repeatParams;
206+ 
207+ repeatParams.blockCount = blockCount;
208+ repeatParams.blockLen = blockLen;
209+ repeatParams.srcStride = tla::get<0, 1>(srcTensor.stride()) / ELE_NUM_PER_C0 - blockLen;
210+ repeatParams.dstStride = tla::get<0, 1>(dstTensor.stride()) / ELE_NUM_PER_C0 - blockLen;
211+ 
212+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
213+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
214+ 
215+ AscendC::DataCopy(dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], repeatParams);
216+ }
217+};
218+ 
219+} // namespace NpuArch::Gemm::Tile
220+ 
221+#endif // GEMM_TILE_COPY_GM_TO_L1_A5_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/copy_gm_to_ub.hpp+22-0
@@ -0,0 +1,22 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_COPY_GM_TO_UB_HPP
12+#define GEMM_TILE_COPY_GM_TO_UB_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp"
17+namespace NpuArch::Gemm::Tile {
18+ 
19+ 
20+} // NpuArch::Gemm::Tile
21+ 
22+#endif // GEMM_TILE_COPY_GM_TO_UB_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/copy_l0c_to_dst.hpp+283-0
@@ -0,0 +1,283 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_COPY_L0C_TO_DST_HPP
12+#define GEMM_TILE_COPY_L0C_TO_DST_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/gemm/gemm_type.hpp"
17+#include "../../../tla/tensor.hpp"
18+ 
19+namespace NpuArch::Gemm::Tile {
20+ 
21+enum class ScaleGranularity {
22+ UNDEFINED = -1,
23+ NO_QUANT = 0,
24+ PER_TENSOR,
25+ PER_CHANNEL,
26+ PER_GROUP
27+};
28+ 
29+template <
30+ class ArchTag,
31+ class ElementSrc,
32+ class ElementDst,
33+ ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT
34+>
35+struct CopyL0CToDstQuantMode {
36+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy l0c to gm, can not find the specialization.");
37+};
38+ 
39+// CopyL0CToGm fp32 to fp32
40+template <class ArchTag>
41+struct CopyL0CToDstQuantMode<
42+ ArchTag,
43+ float, float,
44+ ScaleGranularity::NO_QUANT
45+> {
46+ static constexpr auto VALUE = QuantMode_t::NoQuant;
47+};
48+ 
49+// CopyL0CToGm fp32 to fp32
50+#if defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3101) || (__NPU_ARCH__ == 3510))
51+template <>
52+struct CopyL0CToDstQuantMode<
53+ NpuArch::Arch::AtlasA5,
54+ float, float,
55+ ScaleGranularity::PER_CHANNEL
56+> {
57+ static constexpr auto VALUE = QuantMode_t::VQF322F32_PRE;
58+};
59+#endif
60+ 
61+ 
62+// CopyL0CToGm cast fp32 to fp16
63+template <class ArchTag>
64+struct CopyL0CToDstQuantMode<
65+ ArchTag,
66+ float, half,
67+ ScaleGranularity::NO_QUANT
68+> {
69+ static constexpr auto VALUE = QuantMode_t::F322F16;
70+};
71+ 
72+template <class ArchTag>
73+struct CopyL0CToDstQuantMode<
74+ ArchTag,
75+ float, half,
76+ ScaleGranularity::PER_TENSOR
77+> {
78+ static constexpr auto VALUE = QuantMode_t::QF322F16_PRE;
79+};
80+ 
81+template <class ArchTag>
82+struct CopyL0CToDstQuantMode<
83+ ArchTag,
84+ float, half,
85+ ScaleGranularity::PER_CHANNEL
86+> {
87+ static constexpr auto VALUE = QuantMode_t::VQF322F16_PRE;
88+};
89+ 
90+// CopyL0CToGm cast fp32 to bf16
91+template <class ArchTag>
92+struct CopyL0CToDstQuantMode<
93+ ArchTag,
94+ float, bfloat16_t,
95+ ScaleGranularity::NO_QUANT
96+> {
97+ static constexpr auto VALUE = QuantMode_t::F322BF16;
98+};
99+ 
100+template <class ArchTag>
101+struct CopyL0CToDstQuantMode<
102+ ArchTag,
103+ float, bfloat16_t,
104+ ScaleGranularity::PER_TENSOR
105+> {
106+ static constexpr auto VALUE = QuantMode_t::QF322BF16_PRE;
107+};
108+ 
109+template <class ArchTag>
110+struct CopyL0CToDstQuantMode<
111+ ArchTag,
112+ float, bfloat16_t,
113+ ScaleGranularity::PER_CHANNEL
114+> {
115+ static constexpr auto VALUE = QuantMode_t::VQF322BF16_PRE;
116+};
117+ 
118+// CopyL0CToGm cast float to uint8/int8
119+template <class ArchTag>
120+struct CopyL0CToDstQuantMode<
121+ ArchTag,
122+ float, uint8_t,
123+ ScaleGranularity::PER_TENSOR
124+> {
125+ static constexpr auto VALUE = QuantMode_t::QF322B8_PRE;
126+};
127+ 
128+template <class ArchTag>
129+struct CopyL0CToDstQuantMode<
130+ ArchTag,
131+ float, uint8_t,
132+ ScaleGranularity::PER_CHANNEL
133+> {
134+ static constexpr auto VALUE = QuantMode_t::VQF322B8_PRE;
135+};
136+ 
137+template <class ArchTag>
138+struct CopyL0CToDstQuantMode<
139+ ArchTag,
140+ float, int8_t,
141+ ScaleGranularity::PER_TENSOR
142+> {
143+ static constexpr auto VALUE = QuantMode_t::QF322B8_PRE;
144+};
145+ 
146+template <class ArchTag>
147+struct CopyL0CToDstQuantMode<
148+ ArchTag,
149+ float, int8_t,
150+ ScaleGranularity::PER_CHANNEL
151+> {
152+ static constexpr auto VALUE = QuantMode_t::VQF322B8_PRE;
153+};
154+ 
155+// CopyL0CToGm output int32
156+template <class ArchTag>
157+struct CopyL0CToDstQuantMode<
158+ ArchTag,
159+ int32_t, int32_t,
160+ ScaleGranularity::NO_QUANT
161+> {
162+ static constexpr auto VALUE = QuantMode_t::NoQuant;
163+};
164+ 
165+// CopyL0CToGm cast int32_t to fp16
166+template <class ArchTag>
167+struct CopyL0CToDstQuantMode<
168+ ArchTag,
169+ int32_t, half,
170+ ScaleGranularity::PER_TENSOR
171+> {
172+ static constexpr auto VALUE = QuantMode_t::DEQF16;
173+};
174+ 
175+// CopyL0CToGm cast int32_t to fp16
176+template <class ArchTag>
177+struct CopyL0CToDstQuantMode<
178+ ArchTag,
179+ int32_t, half,
180+ ScaleGranularity::NO_QUANT
181+> {
182+ static constexpr auto VALUE = QuantMode_t::DEQF16;
183+};
184+ 
185+template <class ArchTag>
186+struct CopyL0CToDstQuantMode<
187+ ArchTag,
188+ int32_t, half,
189+ ScaleGranularity::PER_CHANNEL
190+> {
191+ static constexpr auto VALUE = QuantMode_t::VDEQF16;
192+};
193+ 
194+// CopyL0CToGm cast int32 to uint8/int8
195+template <class ArchTag>
196+struct CopyL0CToDstQuantMode<
197+ ArchTag,
198+ int32_t, uint8_t,
199+ ScaleGranularity::PER_TENSOR
200+> {
201+ static constexpr auto VALUE = QuantMode_t::REQ8;
202+};
203+ 
204+template <class ArchTag>
205+struct CopyL0CToDstQuantMode<
206+ ArchTag,
207+ int32_t, uint8_t,
208+ ScaleGranularity::PER_CHANNEL
209+> {
210+ static constexpr auto VALUE = QuantMode_t::VREQ8;
211+};
212+ 
213+template <class ArchTag>
214+struct CopyL0CToDstQuantMode<
215+ ArchTag,
216+ int32_t, int8_t,
217+ ScaleGranularity::PER_TENSOR
218+> {
219+ static constexpr auto VALUE = QuantMode_t::REQ8;
220+};
221+ 
222+template <class ArchTag>
223+struct CopyL0CToDstQuantMode<
224+ ArchTag,
225+ int32_t, int8_t,
226+ ScaleGranularity::PER_CHANNEL
227+> {
228+ static constexpr auto VALUE = QuantMode_t::VREQ8;
229+};
230+ 
231+template <
232+ class ArchTag,
233+ class ElementAccumulator,
234+ class GmType,
235+ ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT,
236+ bool ReluEnable = false
237+>
238+struct CopyL0CToGm {
239+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy l0c to gm, can not find the specialization.");
240+};
241+ 
242+///////////////////////////////////////////CopyL0CToGmTla/////////////////////////////////////////////////
243+// L0C copy mode
244+struct CopyToGM {};
245+struct CopyToL1 {};
246+ 
247+template <
248+ class ArchTag,
249+ class TensorSrc,
250+ class TensorDst,
251+ ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT,
252+ bool ReluEnable = false,
253+ class Enable = void
254+>
255+struct CopyL0CToGmTla {
256+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy l0c to gm, can not find the specialization.");
257+};
258+ 
259+/////////////////////////////////////////////////////////////////////////////////////////////////////////////
260+ 
261+enum class CopyL0CToUBMode {
262+ NO_SPLIT = 0,
263+ SPLIT_M,
264+ SPLIT_N,
265+ RESERVED
266+};
267+ 
268+template <
269+ class ArchTag,
270+ class TensorSrc,
271+ class TensorDst,
272+ CopyL0CToUBMode CopyMode = CopyL0CToUBMode::NO_SPLIT,
273+ ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT,
274+ bool ReluEnable = false,
275+ class Enable = void
276+>
277+struct CopyL0CToUBTla {
278+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy l0c to ub, can not find the specialization.");
279+};
280+ 
281+} // namespace NpuArch::Gemm::Tile
282+ 
283+#endif // GEMM_TILE_COPY_L0C_TO_DST_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/copy_l0c_to_gm_a2.hpp+106-0
@@ -0,0 +1,106 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_COPY_L0C_TO_GM_A2_HPP
12+#define GEMM_TILE_COPY_L0C_TO_GM_A2_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/gemm/tile_common/copy_l0c_to_dst.hpp"
17+#include "../../../attn_infra/gemm/gemm_type.hpp"
18+namespace NpuArch::Gemm::Tile {
19+ 
20+template <
21+ class ElementAccumulator_,
22+ class ElementDst_,
23+ bool ReluEnable_
24+>
25+struct CopyL0CToGm<NpuArch::Arch::AtlasA2,
26+ ElementAccumulator_,
27+ Gemm::GemmType<ElementDst_, layout::RowMajor>,
28+ ScaleGranularity::NO_QUANT,
29+ ReluEnable_>
30+{
31+ using ArchTag = NpuArch::Arch::AtlasA2;
32+ using ElementDst = ElementDst_;
33+ using ElementSrc = ElementAccumulator_;
34+ using LayoutSrc = NpuArch::layout::zN;
35+ using LayoutDst = NpuArch::layout::RowMajor;
36+ static constexpr auto quantPre = CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst,
37+ ScaleGranularity::NO_QUANT>::VALUE;
38+ static constexpr auto reluEn = ReluEnable_;
39+ 
40+ __aicore__ inline
41+ void operator()(AscendC::GlobalTensor<ElementDst> const &dst, AscendC::LocalTensor<ElementSrc> const &src,
42+ LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0)
43+ {
44+ AscendC::FixpipeParamsV220 intriParams;
45+ 
46+ // Fixpipe layout information
47+ intriParams.nSize = dstLayout.shape(1);
48+ intriParams.mSize = dstLayout.shape(0);
49+ intriParams.srcStride = srcLayout.stride(3) / srcLayout.stride(0);
50+ intriParams.dstStride = dstLayout.stride(0);
51+ 
52+ // Fixpipe auxiliary arguments
53+ intriParams.quantPre = quantPre;
54+ intriParams.reluEn = reluEn;
55+ intriParams.unitFlag = unitFlag;
56+ 
57+ // Call AscendC Fixpipe
58+ AscendC::Fixpipe<ElementDst, ElementSrc, AscendC::CFG_ROW_MAJOR>(dst, src, intriParams);
59+ }
60+};
61+ 
62+template <
63+ class ElementAccumulator_,
64+ class ElementDst_,
65+ bool ReluEnable_
66+>
67+struct CopyL0CToGm<NpuArch::Arch::AtlasA2,
68+ ElementAccumulator_,
69+ Gemm::GemmType<ElementDst_, layout::zN>,
70+ ScaleGranularity::NO_QUANT,
71+ ReluEnable_>
72+{
73+ using ArchTag = NpuArch::Arch::AtlasA2;
74+ using ElementDst = ElementDst_;
75+ using ElementSrc = ElementAccumulator_;
76+ using LayoutSrc = NpuArch::layout::zN;
77+ using LayoutDst = NpuArch::layout::zN;
78+ static constexpr auto quantPre = CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst,
79+ ScaleGranularity::NO_QUANT>::VALUE;
80+ static constexpr auto reluEn = ReluEnable_;
81+ 
82+ __aicore__ inline
83+ void operator()(AscendC::GlobalTensor<ElementDst> const &dst, AscendC::LocalTensor<ElementSrc> const &src,
84+ LayoutDst const &dstLayout, LayoutSrc const &srcLayout, uint8_t unitFlag = 0)
85+ {
86+ AscendC::FixpipeParamsV220 intriParams;
87+ 
88+ // Fixpipe layout information
89+ intriParams.nSize = dstLayout.shape(2) * dstLayout.shape(3);
90+ intriParams.mSize = dstLayout.shape(0) * dstLayout.shape(1);
91+ intriParams.srcStride = srcLayout.stride(3) / srcLayout.shape(2);
92+ intriParams.dstStride = dstLayout.stride(3) / (BYTE_PER_C0 / sizeof(ElementDst));
93+ 
94+ // Fixpipe auxiliary arguments
95+ intriParams.quantPre = quantPre;
96+ intriParams.reluEn = reluEn;
97+ intriParams.unitFlag = unitFlag;
98+ 
99+ // Call AscendC Fixpipe
100+ AscendC::Fixpipe<ElementDst, ElementSrc, AscendC::CFG_NZ>(dst, src, intriParams);
101+ }
102+};
103+ 
104+} // namespace NpuArch::Gemm::Tile
105+ 
106+#endif // GEMM_TILE_COPY_L0C_TO_GM_A2_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/copy_l0c_to_ub_a5.hpp+757-0
@@ -0,0 +1,757 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_COPY_L0C_TO_UB_A5_HPP
12+#define GEMM_TILE_COPY_L0C_TO_UB_A5_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/gemm/tile_common/copy_l0c_to_dst.hpp"
17+#include "../../../tla/tensor.hpp"
18+ 
19+#if (__CCE_AICORE__ == 310)
20+constexpr AscendC::FixpipeConfig CFG_ROW_MAJOR_UB = {AscendC::CO2Layout::ROW_MAJOR, true};
21+constexpr AscendC::FixpipeConfig CFG_NZ_UB = {AscendC::CO2Layout::NZ, true};
22+#endif
23+ 
24+namespace NpuArch::Gemm::Tile {
25+ 
26+template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_>
27+struct CopyL0CToUBTla<
28+ NpuArch::Arch::AtlasA5,
29+ TensorSrc_,
30+ tla::Tensor<AscendC::LocalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>,
31+ CopyL0CToUBMode::NO_SPLIT,
32+ ScaleGranularity::NO_QUANT,
33+ ReluEnable_,
34+ std::enable_if_t<tla::detail::isRowMajor<LayoutDst_>::value>> {
35+ using ArchTag = NpuArch::Arch::AtlasA5;
36+ using ElementDst = ElementDst_;
37+ using ElementSrc = typename TensorSrc_::Element;
38+ static constexpr auto quantPre =
39+ CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::NO_QUANT>::VALUE;
40+ static constexpr auto reluEn = ReluEnable_;
41+ 
42+ template <class TensorDst, class TensorSrc>
43+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0)
44+ {
45+ static_assert(
46+ tla::detail::isRowMajor<typename TensorDst::Layout>::value && TensorSrc::position == AscendC::TPosition::CO1
47+ && TensorDst::position == AscendC::TPosition::VECCALC,
48+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"
49+ );
50+ 
51+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;
52+ 
53+ // Fixpipe layout information
54+ intriParams.nSize = tla::get<1>(dstTensor.shape());
55+ intriParams.mSize = tla::get<0>(dstTensor.shape());
56+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
57+ intriParams.dstStride = tla::get<0>(dstTensor.stride());
58+ 
59+ // Fixpipe auxiliary arguments
60+ intriParams.quantPre = quantPre;
61+ intriParams.reluEn = reluEn;
62+ intriParams.unitFlag = unitFlag;
63+ 
64+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
65+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
66+ 
67+ // Call AscendC Fixpipe
68+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_ROW_MAJOR_UB>(
69+ dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
70+ }
71+
72+ template <class TensorDst, class TensorSrc>
73+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, bool subBlockId, uint8_t unitFlag = 0)
74+ {
75+ static_assert(
76+ tla::detail::isRowMajor<typename TensorDst::Layout>::value && TensorSrc::position == AscendC::TPosition::CO1
77+ && TensorDst::position == AscendC::TPosition::VECCALC,
78+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"
79+ );
80+ 
81+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;
82+ 
83+ // Fixpipe layout information
84+ intriParams.nSize = tla::get<1>(dstTensor.shape());
85+ intriParams.mSize = tla::get<0>(dstTensor.shape());
86+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
87+ intriParams.dstStride = tla::get<0>(dstTensor.stride());
88+ 
89+ // Fixpipe auxiliary arguments
90+ intriParams.quantPre = quantPre;
91+ intriParams.reluEn = reluEn;
92+ intriParams.unitFlag = unitFlag;
93+ intriParams.dualDstCtl = 0;
94+ intriParams.subBlockId = subBlockId;
95+ 
96+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
97+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
98+ 
99+ // Call AscendC Fixpipe
100+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_ROW_MAJOR_UB>(
101+ dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
102+ }
103+};
104+ 
105+template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_>
106+struct CopyL0CToUBTla<
107+ NpuArch::Arch::AtlasA5,
108+ TensorSrc_,
109+ tla::Tensor<AscendC::LocalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>,
110+ CopyL0CToUBMode::NO_SPLIT,
111+ ScaleGranularity::PER_TENSOR,
112+ ReluEnable_,
113+ std::enable_if_t<tla::detail::iszN<ElementDst_, LayoutDst_>::value>> {
114+ using ArchTag = NpuArch::Arch::AtlasA5;
115+ using ElementDst = ElementDst_;
116+ using ElementSrc = typename TensorSrc_::Element;
117+ static constexpr auto quantPre =
118+ CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::PER_TENSOR>::VALUE;
119+ static constexpr auto reluEn = ReluEnable_;
120+ 
121+ struct Params {
122+ float scale = 1.0f;
123+ 
124+ __aicore__ inline
125+ Params() = default;
126+ 
127+ __aicore__ inline
128+ Params(float scalar)
129+ {
130+ scale = scalar;
131+ }
132+ };
133+ Params params;
134+ 
135+ __aicore__ inline
136+ CopyL0CToUBTla() = default;
137+ 
138+ __aicore__ inline
139+ CopyL0CToUBTla(Params const &params_) : params(params_) {};
140+ 
141+ template <class TensorDst, class TensorSrc>
142+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0)
143+ {
144+ static_assert(
145+ tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
146+ && TensorSrc::position == AscendC::TPosition::CO1
147+ && TensorDst::position == AscendC::TPosition::VECCALC,
148+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and zN"
149+ );
150+ 
151+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::NZ> intriParams;
152+ 
153+ intriParams.nSize = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
154+ intriParams.mSize = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
155+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape());
156+ intriParams.dstStride = intriParams.mSize * (BYTE_PER_C0 / sizeof(ElementDst));
157+ 
158+ intriParams.quantPre = quantPre;
159+ intriParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t*>(&params.scale));
160+ intriParams.reluEn = reluEn;
161+ intriParams.unitFlag = unitFlag;
162+ 
163+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
164+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
165+ 
166+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_NZ_UB>(
167+ dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
168+ }
169+ 
170+ template <class TensorDst, class TensorSrc>
171+ __aicore__ inline void operator()(
172+ TensorDst const &dstTensor, TensorSrc const &srcTensor, bool subBlockId, uint8_t unitFlag = 0)
173+ {
174+ static_assert(
175+ tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
176+ && TensorSrc::position == AscendC::TPosition::CO1
177+ && TensorDst::position == AscendC::TPosition::VECCALC,
178+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and zN"
179+ );
180+ 
181+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::NZ> intriParams;
182+ 
183+ intriParams.nSize = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
184+ intriParams.mSize = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
185+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape());
186+ intriParams.dstStride = intriParams.mSize * (BYTE_PER_C0 / sizeof(ElementDst));
187+ 
188+ intriParams.quantPre = quantPre;
189+ intriParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t*>(&params.scale));
190+ intriParams.reluEn = reluEn;
191+ intriParams.unitFlag = unitFlag;
192+ intriParams.dualDstCtl = 0;
193+ intriParams.subBlockId = subBlockId;
194+ 
195+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
196+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
197+ 
198+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_NZ_UB>(
199+ dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
200+ }
201+};
202+ 
203+template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_>
204+struct CopyL0CToUBTla<
205+ NpuArch::Arch::AtlasA5,
206+ TensorSrc_,
207+ tla::Tensor<AscendC::LocalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>,
208+ CopyL0CToUBMode::NO_SPLIT,
209+ ScaleGranularity::NO_QUANT,
210+ ReluEnable_,
211+ std::enable_if_t<tla::detail::iszN<ElementDst_, LayoutDst_>::value>> {
212+ using ArchTag = NpuArch::Arch::AtlasA5;
213+ using ElementDst = ElementDst_;
214+ using ElementSrc = typename TensorSrc_::Element;
215+ static constexpr auto quantPre =
216+ CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::NO_QUANT>::VALUE;
217+ static constexpr auto reluEn = ReluEnable_;
218+ 
219+ template <class TensorDst, class TensorSrc>
220+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0)
221+ {
222+ static_assert(
223+ tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
224+ && TensorSrc::position == AscendC::TPosition::CO1
225+ && TensorDst::position == AscendC::TPosition::VECCALC,
226+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and zN"
227+ );
228+ 
229+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::NZ> intriParams;
230+ 
231+ //shape = ((16, ceil_div(rows, 16)), (16, ceil_div(cols, 16)))
232+ //stride = ((16, 256), (1, round_up(rows, 16) * 16))
233+ // zN/NZ Fixpipe consumes the physical fractal extent stored in TLA nested shape.
234+ // 源NZ矩阵在N方向上的大小
235+ intriParams.nSize = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
236+ // 源NZ矩阵在M方向上的大小
237+ intriParams.mSize = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
238+ // 源NZ矩阵中的相邻Z排布的起始地址偏移，单位是C0_SIZE
239+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape());
240+ // 目的NZ矩阵中相邻Z排布的起始地址偏移，单位是元素
241+ intriParams.dstStride = intriParams.mSize * (BYTE_PER_C0 / sizeof(ElementDst));
242+ 
243+ intriParams.quantPre = quantPre;
244+ intriParams.reluEn = reluEn;
245+ intriParams.unitFlag = unitFlag;
246+ 
247+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
248+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
249+ 
250+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_NZ_UB>(
251+ dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
252+ }
253+ 
254+ template <class TensorDst, class TensorSrc>
255+ __aicore__ inline void operator()(
256+ TensorDst const &dstTensor, TensorSrc const &srcTensor, bool subBlockId, uint8_t unitFlag = 0)
257+ {
258+ static_assert(
259+ tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
260+ && TensorSrc::position == AscendC::TPosition::CO1
261+ && TensorDst::position == AscendC::TPosition::VECCALC,
262+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and zN"
263+ );
264+ 
265+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::NZ> intriParams;
266+ 
267+ intriParams.nSize = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
268+ intriParams.mSize = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
269+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape());
270+ intriParams.dstStride = intriParams.mSize * (BYTE_PER_C0 / sizeof(ElementDst));
271+ 
272+ intriParams.quantPre = quantPre;
273+ intriParams.reluEn = reluEn;
274+ intriParams.unitFlag = unitFlag;
275+ intriParams.dualDstCtl = 0;
276+ intriParams.subBlockId = subBlockId;
277+ 
278+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
279+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
280+ 
281+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_NZ_UB>(
282+ dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
283+ }
284+};
285+ 
286+template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_>
287+struct CopyL0CToUBTla<
288+ NpuArch::Arch::AtlasA5,
289+ TensorSrc_,
290+ tla::Tensor<AscendC::LocalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>,
291+ CopyL0CToUBMode::NO_SPLIT,
292+ ScaleGranularity::PER_TENSOR,
293+ ReluEnable_,
294+ std::enable_if_t<tla::detail::isRowMajor<LayoutDst_>::value>> {
295+ using ArchTag = NpuArch::Arch::AtlasA5;
296+ using ElementDst = ElementDst_;
297+ using ElementSrc = typename TensorSrc_::Element;
298+ static constexpr auto quantPre =
299+ CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::PER_TENSOR>::VALUE;
300+ static constexpr auto reluEn = ReluEnable_;
301+ 
302+ struct Params {
303+ float scale = 1.0f;
304+ 
305+ __aicore__ inline
306+ Params() = default;
307+ 
308+ __aicore__ inline
309+ Params(float scalar)
310+ {
311+ scale = scalar;
312+ }
313+ };
314+ Params params;
315+ 
316+ __aicore__ inline
317+ CopyL0CToUBTla() = default;
318+ 
319+ __aicore__ inline
320+ CopyL0CToUBTla(Params const &params_) : params(params_) {};
321+ 
322+ template <class TensorDst, class TensorSrc>
323+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0)
324+ {
325+ static_assert(
326+ tla::detail::isRowMajor<typename TensorDst::Layout>::value && TensorSrc::position == AscendC::TPosition::CO1
327+ && TensorDst::position == AscendC::TPosition::VECCALC,
328+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"
329+ );
330+ 
331+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;
332+ 
333+ // Fixpipe layout information
334+ intriParams.nSize = tla::get<1>(dstTensor.shape());
335+ intriParams.mSize = tla::get<0>(dstTensor.shape());
336+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
337+ intriParams.dstStride = tla::get<0>(dstTensor.stride());
338+ 
339+ // Fixpipe auxiliary arguments
340+ intriParams.quantPre = quantPre;
341+ intriParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t*>(&params.scale));
342+ intriParams.reluEn = reluEn;
343+ intriParams.unitFlag = unitFlag;
344+ 
345+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
346+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
347+ 
348+ // Call AscendC Fixpipe
349+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_ROW_MAJOR_UB>(
350+ dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
351+ }
352+ template <class TensorDst, class TensorSrc>
353+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, bool subBlockId, uint8_t unitFlag = 0)
354+ {
355+ static_assert(
356+ tla::detail::isRowMajor<typename TensorDst::Layout>::value && TensorSrc::position == AscendC::TPosition::CO1
357+ && TensorDst::position == AscendC::TPosition::VECCALC,
358+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"
359+ );
360+ 
361+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;
362+ 
363+ // Fixpipe layout information
364+ intriParams.nSize = tla::get<1>(dstTensor.shape());
365+ intriParams.mSize = tla::get<0>(dstTensor.shape());
366+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
367+ intriParams.dstStride = tla::get<0>(dstTensor.stride());
368+ 
369+ // Fixpipe auxiliary arguments
370+ intriParams.quantPre = quantPre;
371+ intriParams.deqScalar = static_cast<uint64_t>(*reinterpret_cast<int32_t*>(&params.scale));
372+ intriParams.reluEn = reluEn;
373+ intriParams.unitFlag = unitFlag;
374+ intriParams.dualDstCtl = 0;
375+ intriParams.subBlockId = subBlockId;
376+ 
377+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
378+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
379+ 
380+ // Call AscendC Fixpipe
381+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_ROW_MAJOR_UB>(
382+ dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
383+ }
384+};
385+ 
386+template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_>
387+struct CopyL0CToUBTla<
388+ NpuArch::Arch::AtlasA5,
389+ TensorSrc_,
390+ tla::Tensor<AscendC::LocalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>,
391+ CopyL0CToUBMode::SPLIT_M,
392+ ScaleGranularity::NO_QUANT,
393+ ReluEnable_,
394+ std::enable_if_t<tla::detail::isRowMajor<LayoutDst_>::value>> {
395+ using ArchTag = NpuArch::Arch::AtlasA5;
396+ using ElementDst = ElementDst_;
397+ using ElementSrc = typename TensorSrc_::Element;
398+ static constexpr auto quantPre =
399+ CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::NO_QUANT>::VALUE;
400+ static constexpr auto reluEn = ReluEnable_;
401+ 
402+ template <class TensorDst, class TensorSrc>
403+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0)
404+ {
405+ static_assert(
406+ tla::detail::isRowMajor<typename TensorDst::Layout>::value && TensorSrc::position == AscendC::TPosition::CO1
407+ && TensorDst::position == AscendC::TPosition::VECCALC,
408+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"
409+ );
410+ 
411+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;
412+ 
413+ // Fixpipe layout information
414+ intriParams.nSize = tla::get<1>(dstTensor.shape());
415+ intriParams.mSize = RoundUp(tla::get<0>(dstTensor.shape()), 2); // m must be even when spilt m
416+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
417+ intriParams.dstStride = tla::get<0>(dstTensor.stride());
418+ 
419+ // Fixpipe auxiliary arguments
420+ intriParams.quantPre = quantPre;
421+ intriParams.reluEn = reluEn;
422+ intriParams.unitFlag = unitFlag;
423+ intriParams.dualDstCtl = 1;
424+ 
425+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
426+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
427+ 
428+ // Call AscendC Fixpipe
429+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_ROW_MAJOR_UB>(
430+ dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
431+ }
432+};
433+ 
434+template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_>
435+struct CopyL0CToUBTla<
436+ NpuArch::Arch::AtlasA5,
437+ TensorSrc_,
438+ tla::Tensor<AscendC::LocalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>,
439+ CopyL0CToUBMode::SPLIT_M,
440+ ScaleGranularity::NO_QUANT,
441+ ReluEnable_,
442+ std::enable_if_t<tla::detail::iszN<ElementDst_, LayoutDst_>::value>> {
443+ using ArchTag = NpuArch::Arch::AtlasA5;
444+ using ElementDst = ElementDst_;
445+ using ElementSrc = typename TensorSrc_::Element;
446+ static constexpr auto quantPre =
447+ CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::NO_QUANT>::VALUE;
448+ static constexpr auto reluEn = ReluEnable_;
449+ 
450+ template <class TensorDst, class TensorSrc>
451+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0)
452+ {
453+ static_assert(
454+ tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
455+ && TensorSrc::position == AscendC::TPosition::CO1
456+ && TensorDst::position == AscendC::TPosition::VECCALC,
457+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and zN"
458+ );
459+ 
460+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::NZ> intriParams;
461+ 
462+ intriParams.nSize = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
463+ intriParams.mSize = RoundUp(tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape()), 2);
464+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<1, 0>(srcTensor.shape());
465+ intriParams.dstStride = intriParams.mSize * (BYTE_PER_C0 / sizeof(ElementDst));
466+ 
467+ intriParams.quantPre = quantPre;
468+ intriParams.reluEn = reluEn;
469+ intriParams.unitFlag = unitFlag;
470+ intriParams.dualDstCtl = 1;
471+ 
472+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
473+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
474+ 
475+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_NZ_UB>(
476+ dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
477+ }
478+};
479+ 
480+// NOTE: SPLIT_M + PER_CHANNEL is NOT supported by single Fixpipe with dualDstCtl=1,
481+// because per-channel dequant (scale tensor) conflicts with dual-core split mode.
482+// Workaround: manually decompose into two independent Fixpipe calls (sub0 + sub1),
483+// each with a separate sub-tile and the same full scale tensor.
484+// This keeps the dequantization correct while achieving dual-core output.
485+template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_>
486+struct CopyL0CToUBTla<
487+ NpuArch::Arch::AtlasA5,
488+ TensorSrc_,
489+ tla::Tensor<AscendC::LocalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>,
490+ CopyL0CToUBMode::SPLIT_M,
491+ ScaleGranularity::PER_CHANNEL,
492+ ReluEnable_,
493+ std::enable_if_t<tla::detail::isRowMajor<LayoutDst_>::value>> {
494+ using ArchTag = NpuArch::Arch::AtlasA5;
495+ using ElementDst = ElementDst_;
496+ using ElementSrc = typename TensorSrc_::Element;
497+ static constexpr auto quantPre =
498+ CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::PER_CHANNEL>::VALUE;
499+ static constexpr auto reluEn = ReluEnable_;
500+ 
501+ struct Params {};
502+ Params params;
503+ 
504+ __aicore__ inline
505+ CopyL0CToUBTla() = default;
506+ 
507+ __aicore__ inline
508+ CopyL0CToUBTla(Params const &params_) : params(params_) {};
509+ 
510+ template <class TensorDst, class TensorSrc>
511+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor,
512+ AscendC::LocalTensor<uint64_t> const &scale, uint8_t unitFlag = 0)
513+ {
514+ static_assert(
515+ tla::detail::isRowMajor<typename TensorDst::Layout>::value && TensorSrc::position == AscendC::TPosition::CO1
516+ && TensorDst::position == AscendC::TPosition::VECCALC,
517+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"
518+ );
519+ 
520+ uint32_t mSize = tla::get<0>(dstTensor.shape());
521+ uint32_t nSize = tla::get<1>(dstTensor.shape());
522+ uint32_t mPerSubCore = mSize / 2;
523+ 
524+ // --- prepare sub-tiles for dual sub-core fixpipe ---
525+ // CAUTION: L0C srcTensor has a nested fractal shape (tuple<tuple<C<16>, uint>, tuple<C<16>, uint>>).
526+ // Do NOT use tla::get<1>(srcTensor.shape()) as MakeShape argument here,
527+ // because it returns a nested tuple fragment, not a flat uint.
528+ // Always pass flat uint values (nSize) to MakeShape; MakeLayoutTile handles the fractal mapping internally.
529+ auto dstSub0 = tla::GetTile(dstTensor, tla::MakeCoord(0, 0), tla::MakeShape(mPerSubCore, nSize));
530+ auto srcSub0 = tla::GetTile(srcTensor, tla::MakeCoord(0, 0), tla::MakeShape(mPerSubCore, nSize));
531+ 
532+ auto dstSub1 = tla::GetTile(dstTensor, tla::MakeCoord(0, 0), tla::MakeShape(mPerSubCore, nSize));
533+ auto srcSub1 = tla::GetTile(srcTensor, tla::MakeCoord(mPerSubCore, 0), tla::MakeShape(mPerSubCore, nSize));
534+ 
535+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams0;
536+ intriParams0.nSize = tla::get<1>(dstSub0.shape());
537+ intriParams0.mSize = tla::get<0>(dstSub0.shape());
538+ intriParams0.srcStride = tla::get<1, 1>(srcSub0.stride()) / tla::get<0, 0>(srcSub0.stride());
539+ intriParams0.dstStride = tla::get<0>(dstSub0.stride());
540+ intriParams0.quantPre = quantPre;
541+ intriParams0.reluEn = reluEn;
542+ intriParams0.unitFlag = unitFlag;
543+ intriParams0.dualDstCtl = 0;
544+ intriParams0.subBlockId = 0;
545+ 
546+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams1;
547+ intriParams1.nSize = tla::get<1>(dstSub1.shape());
548+ intriParams1.mSize = tla::get<0>(dstSub1.shape());
549+ intriParams1.srcStride = tla::get<1, 1>(srcSub1.stride()) / tla::get<0, 0>(srcSub1.stride());
550+ intriParams1.dstStride = tla::get<0>(dstSub1.stride());
551+ intriParams1.quantPre = quantPre;
552+ intriParams1.reluEn = reluEn;
553+ intriParams1.unitFlag = unitFlag;
554+ intriParams1.dualDstCtl = 0;
555+ intriParams1.subBlockId = 1;
556+ 
557+ // --- execute dual fixpipe for sub-core 0 and sub-core 1 ---
558+ auto dstOffset0 = dstSub0.layout()(dstSub0.coord());
559+ auto srcOffset0 = srcSub0.layout()(srcSub0.coord());
560+ auto dstOffset1 = dstSub1.layout()(dstSub1.coord());
561+ auto srcOffset1 = srcSub1.layout()(srcSub1.coord());
562+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_ROW_MAJOR_UB>(
563+ dstSub0.data()[dstOffset0], srcSub0.data()[srcOffset0], scale, intriParams0);
564+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_ROW_MAJOR_UB>(
565+ dstSub1.data()[dstOffset1], srcSub1.data()[srcOffset1], scale, intriParams1);
566+ }
567+};
568+ 
569+template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_>
570+struct CopyL0CToUBTla<
571+ NpuArch::Arch::AtlasA5,
572+ TensorSrc_,
573+ tla::Tensor<AscendC::LocalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>,
574+ CopyL0CToUBMode::SPLIT_M,
575+ ScaleGranularity::PER_CHANNEL,
576+ ReluEnable_,
577+ std::enable_if_t<tla::detail::iszN<ElementDst_, LayoutDst_>::value>> {
578+ using ArchTag = NpuArch::Arch::AtlasA5;
579+ using ElementDst = ElementDst_;
580+ using ElementSrc = typename TensorSrc_::Element;
581+ static constexpr auto quantPre =
582+ CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::PER_CHANNEL>::VALUE;
583+ static constexpr auto reluEn = ReluEnable_;
584+ 
585+ struct Params {};
586+ Params params;
587+ 
588+ __aicore__ inline
589+ CopyL0CToUBTla() = default;
590+ 
591+ __aicore__ inline
592+ CopyL0CToUBTla(Params const &params_) : params(params_) {};
593+ 
594+ template <class TensorDst, class TensorSrc>
595+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor,
596+ AscendC::LocalTensor<uint64_t> const &scale, uint8_t unitFlag = 0)
597+ {
598+ static_assert(
599+ tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
600+ && TensorSrc::position == AscendC::TPosition::CO1
601+ && TensorDst::position == AscendC::TPosition::VECCALC,
602+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and zN"
603+ );
604+ 
605+ uint32_t mSize = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
606+ uint32_t nSize = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
607+ uint32_t mPerSubCore = mSize / 2;
608+ 
609+ auto dstSub0 = tla::GetTile(dstTensor, tla::MakeCoord(0, 0), tla::MakeShape(mPerSubCore, nSize));
610+ auto srcSub0 = tla::GetTile(srcTensor, tla::MakeCoord(0, 0), tla::MakeShape(mPerSubCore, nSize));
611+ 
612+ auto dstSub1 = tla::GetTile(dstTensor, tla::MakeCoord(0, 0), tla::MakeShape(mPerSubCore, nSize));
613+ auto srcSub1 = tla::GetTile(srcTensor, tla::MakeCoord(mPerSubCore, 0), tla::MakeShape(mPerSubCore, nSize));
614+ 
615+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::NZ> intriParams0;
616+ intriParams0.nSize = tla::get<1, 0>(dstSub0.shape()) * tla::get<1, 1>(dstSub0.shape());
617+ intriParams0.mSize = tla::get<0, 0>(dstSub0.shape()) * tla::get<0, 1>(dstSub0.shape());
618+ intriParams0.srcStride = tla::get<1, 1>(srcSub0.stride()) / tla::get<1, 0>(srcSub0.shape());
619+ intriParams0.dstStride = intriParams0.mSize * (BYTE_PER_C0 / sizeof(ElementDst));
620+ intriParams0.quantPre = quantPre;
621+ intriParams0.reluEn = reluEn;
622+ intriParams0.unitFlag = unitFlag;
623+ intriParams0.dualDstCtl = 0;
624+ intriParams0.subBlockId = 0;
625+ 
626+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::NZ> intriParams1;
627+ intriParams1.nSize = tla::get<1, 0>(dstSub1.shape()) * tla::get<1, 1>(dstSub1.shape());
628+ intriParams1.mSize = tla::get<0, 0>(dstSub1.shape()) * tla::get<0, 1>(dstSub1.shape());
629+ intriParams1.srcStride = tla::get<1, 1>(srcSub1.stride()) / tla::get<1, 0>(srcSub1.shape());
630+ intriParams1.dstStride = intriParams1.mSize * (BYTE_PER_C0 / sizeof(ElementDst));
631+ intriParams1.quantPre = quantPre;
632+ intriParams1.reluEn = reluEn;
633+ intriParams1.unitFlag = unitFlag;
634+ intriParams1.dualDstCtl = 0;
635+ intriParams1.subBlockId = 1;
636+ 
637+ auto dstOffset0 = dstSub0.layout()(dstSub0.coord());
638+ auto srcOffset0 = srcSub0.layout()(srcSub0.coord());
639+ auto dstOffset1 = dstSub1.layout()(dstSub1.coord());
640+ auto srcOffset1 = srcSub1.layout()(srcSub1.coord());
641+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_NZ_UB>(
642+ dstSub0.data()[dstOffset0], srcSub0.data()[srcOffset0], scale, intriParams0);
643+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_NZ_UB>(
644+ dstSub1.data()[dstOffset1], srcSub1.data()[srcOffset1], scale, intriParams1);
645+ }
646+};
647+ 
648+template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_>
649+struct CopyL0CToUBTla<
650+ NpuArch::Arch::AtlasA5,
651+ TensorSrc_,
652+ tla::Tensor<AscendC::LocalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>,
653+ CopyL0CToUBMode::SPLIT_N,
654+ ScaleGranularity::NO_QUANT,
655+ ReluEnable_,
656+ std::enable_if_t<tla::detail::isRowMajor<LayoutDst_>::value>> {
657+ using ArchTag = NpuArch::Arch::AtlasA5;
658+ using ElementDst = ElementDst_;
659+ using ElementSrc = typename TensorSrc_::Element;
660+ static constexpr auto quantPre =
661+ CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::NO_QUANT>::VALUE;
662+ static constexpr auto reluEn = ReluEnable_;
663+ 
664+ template <class TensorDst, class TensorSrc>
665+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint8_t unitFlag = 0)
666+ {
667+ static_assert(
668+ tla::detail::isRowMajor<typename TensorDst::Layout>::value && TensorSrc::position == AscendC::TPosition::CO1
669+ && TensorDst::position == AscendC::TPosition::VECCALC,
670+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"
671+ );
672+ 
673+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;
674+ 
675+ // Fixpipe layout information
676+ intriParams.nSize = RoundUp(tla::get<1>(dstTensor.shape()), 32);
677+ intriParams.mSize = tla::get<0>(dstTensor.shape()); // m must be even when spilt m
678+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
679+ intriParams.dstStride = tla::get<0>(dstTensor.stride());
680+ 
681+ // Fixpipe auxiliary arguments
682+ intriParams.quantPre = quantPre;
683+ intriParams.reluEn = reluEn;
684+ intriParams.unitFlag = unitFlag;
685+ intriParams.dualDstCtl = 2;
686+ 
687+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
688+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
689+ 
690+ // Call AscendC Fixpipe
691+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_ROW_MAJOR_UB>(
692+ dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], intriParams);
693+ }
694+};
695+ 
696+template <class TensorSrc_, class ElementDst_, class LayoutDst_, class CoordDst_, bool ReluEnable_>
697+struct CopyL0CToUBTla<
698+ NpuArch::Arch::AtlasA5,
699+ TensorSrc_,
700+ tla::Tensor<AscendC::LocalTensor<ElementDst_>, LayoutDst_, CoordDst_, AscendC::TPosition::VECCALC>,
701+ CopyL0CToUBMode::NO_SPLIT,
702+ ScaleGranularity::PER_CHANNEL,
703+ ReluEnable_,
704+ std::enable_if_t<tla::detail::isRowMajor<LayoutDst_>::value>> {
705+ using ArchTag = NpuArch::Arch::AtlasA5;
706+ using ElementDst = ElementDst_;
707+ using ElementSrc = typename TensorSrc_::Element;
708+ static constexpr auto quantPre =
709+ CopyL0CToDstQuantMode<ArchTag, ElementSrc, ElementDst, ScaleGranularity::PER_CHANNEL>::VALUE;
710+ static constexpr auto reluEn = ReluEnable_;
711+ 
712+ struct Params {};
713+ Params params;
714+ 
715+ __aicore__ inline
716+ CopyL0CToUBTla() = default;
717+ 
718+ __aicore__ inline
719+ CopyL0CToUBTla(Params const &params_) : params(params_) {};
720+ 
721+ template <class TensorDst, class TensorSrc>
722+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor,
723+ AscendC::LocalTensor<uint64_t> const &scale, uint8_t unitFlag = 0)
724+ {
725+ static_assert(
726+ tla::detail::isRowMajor<typename TensorDst::Layout>::value && TensorSrc::position == AscendC::TPosition::CO1
727+ && TensorDst::position == AscendC::TPosition::VECCALC,
728+ "The input parameters do not match. TensorSrc must be L0C, while TensorDst must be UB and RowMajor"
729+ );
730+ 
731+ AscendC::FixpipeParamsC310<AscendC::CO2Layout::ROW_MAJOR> intriParams;
732+ 
733+ // Fixpipe layout information
734+ intriParams.nSize = tla::get<1>(dstTensor.shape());
735+ intriParams.mSize = tla::get<0>(dstTensor.shape());
736+ intriParams.srcStride = tla::get<1, 1>(srcTensor.stride()) / tla::get<0, 0>(srcTensor.stride());
737+ intriParams.dstStride = tla::get<0>(dstTensor.stride());
738+ 
739+ // Fixpipe auxiliary arguments
740+ intriParams.quantPre = quantPre;
741+ intriParams.reluEn = reluEn;
742+ intriParams.unitFlag = unitFlag;
743+ 
744+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
745+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
746+ 
747+ // Call AscendC Fixpipe with scale tensor for per-channel quant
748+ AscendC::Fixpipe<ElementDst, ElementSrc, CFG_ROW_MAJOR_UB>(
749+ dstTensor.data()[dstOffset], srcTensor.data()[srcOffset], scale, intriParams);
750+ }
751+};
752+ 
753+/////////////////////////////////////////////////////////////////////////////////////////////////////////////
754+ 
755+} // namespace NpuArch::Gemm::Tile
756+ 
757+#endif // GEMM_TILE_COPY_L0C_TO_UB_A5_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/copy_l1_to_bt.hpp+27-0
@@ -0,0 +1,27 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_COPY_L1_TO_BT_HPP
12+#define GEMM_TILE_COPY_L1_TO_BT_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/layout/layout.hpp"
17+#include "../../../attn_infra/gemm/gemm_type.hpp"
18+ 
19+ 
20+namespace NpuArch::Gemm::Tile {
21+ 
22+ 
23+/////////////////////////////////////////////////////////////////////////////////////////////////////////////
24+ 
25+} // namespace NpuArch::Gemm::Tile
26+ 
27+#endif // GEMM_TILE_COPY_L1_TO_BT_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/copy_l1_to_l0a_a2.hpp+349-0
@@ -0,0 +1,349 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_COPY_L1_TO_L0A_HPP
12+#define GEMM_TILE_COPY_L1_TO_L0A_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/layout/layout.hpp"
17+#include "../../../attn_infra/gemm/gemm_type.hpp"
18+#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp"
19+ 
20+ 
21+namespace NpuArch::Gemm::Tile {
22+ 
23+template <
24+ class ArchTag,
25+ class L1Type,
26+ class L0Type = void
27+>
28+struct CopyL1ToL0A {
29+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy l1 to l0, can not find the specialization.");
30+};
31+ 
32+////////////////////////////////
33+/// new add gemm
34+template<class ArchTag, class Element>
35+struct CopyL1ToL0A<ArchTag, NpuArch::Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>, NpuArch::Gemm::GemmType<Element, layout::zZ, AscendC::TPosition::A2>>{
36+ using LayoutDst = layout::zZ;
37+ using LayoutSrc = layout::zN;
38+ 
39+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
40+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
41+ 
42+ __aicore__ inline
43+ CopyL1ToL0A(){}
44+ 
45+ __aicore__ inline
46+ void operator()(
47+ AscendC::LocalTensor<Element> dstTensor,
48+ AscendC::LocalTensor<Element> srcTensor,
49+ LayoutDst layoutDst, LayoutSrc layoutSrc
50+ ){
51+ AscendC::LoadData2DParams loadDataParams;
52+ loadDataParams.startIndex = 0;
53+ loadDataParams.repeatTimes = static_cast<uint16_t>(layoutDst.shape(3));
54+ loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL;
55+ loadDataParams.sid = 0;
56+ loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1;
57+ loadDataParams.ifTranspose = false;
58+ loadDataParams.addrMode = 0;
59+ 
60+ for (uint32_t i = 0; i < layoutDst.shape(1); i++) {
61+ AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams);
62+ }
63+ }
64+};
65+ 
66+template<class ArchTag, class Element>
67+struct CopyL1ToL0A<ArchTag, NpuArch::Gemm::GemmType<Element, layout::nN, AscendC::TPosition::A1>, NpuArch::Gemm::GemmType<Element, layout::zZ, AscendC::TPosition::A2>>{
68+ using LayoutDst = layout::zZ;
69+ using LayoutSrc = layout::nN;
70+ 
71+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
72+ 
73+ __aicore__ inline
74+ CopyL1ToL0A(){}
75+ 
76+ __aicore__ inline
77+ void operator()(
78+ AscendC::LocalTensor<Element> dstTensor,
79+ AscendC::LocalTensor<Element> srcTensor,
80+ LayoutDst layoutDst, LayoutSrc layoutSrc
81+ ){
82+ AscendC::LoadData2DParams loadDataParams;
83+ loadDataParams.startIndex = 0;
84+ loadDataParams.repeatTimes = static_cast<uint16_t>(CeilDiv<C0_NUM_PER_FRACTAL>(layoutDst.orgShape(1)));
85+ loadDataParams.srcStride = static_cast<uint16_t>(CeilDiv<ELE_NUM_PER_C0>(layoutSrc.orgShape(0)));;
86+ loadDataParams.sid = 0;
87+ loadDataParams.dstGap = 0;
88+ loadDataParams.ifTranspose = true;
89+ loadDataParams.addrMode = 0;
90+ for(uint32_t i = 0; i < CeilDiv<ELE_NUM_PER_C0>(layoutSrc.orgShape(0)); i++){
91+ AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams);
92+ }
93+ }
94+};
95+ 
96+template<class ArchTag>
97+struct CopyL1ToL0A<ArchTag, NpuArch::Gemm::GemmType<float, layout::nN, AscendC::TPosition::A1>, NpuArch::Gemm::GemmType<float, layout::zZ, AscendC::TPosition::A2>>{
98+ using Element = float;
99+ using LayoutDst = layout::zZ;
100+ using LayoutSrc = layout::nN;
101+ 
102+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
103+ 
104+ __aicore__ inline
105+ CopyL1ToL0A(){}
106+ 
107+ __aicore__ inline
108+ void operator()(
109+ AscendC::LocalTensor<Element> dstTensor,
110+ AscendC::LocalTensor<Element> srcTensor,
111+ LayoutDst layoutDst, LayoutSrc layoutSrc
112+ ){
113+ AscendC::LoadData2dTransposeParams loadDataParams;
114+ loadDataParams.startIndex = 0;
115+ loadDataParams.repeatTimes = static_cast<uint16_t>(CeilDiv<C0_NUM_PER_FRACTAL>(layoutDst.orgShape(1)));
116+ loadDataParams.srcStride = static_cast<uint16_t>(CeilDiv<C0_NUM_PER_FRACTAL>(layoutSrc.orgShape(0)));
117+ loadDataParams.dstGap = 1;
118+ loadDataParams.dstFracGap = 0;
119+ for(uint32_t i = 0; i < CeilDiv<C0_NUM_PER_FRACTAL>(layoutSrc.orgShape(0)); i++){
120+ AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1) * 2], loadDataParams);
121+ }
122+ }
123+};
124+ 
125+template<class ArchTag>
126+struct CopyL1ToL0A<ArchTag, NpuArch::Gemm::GemmType<int8_t, layout::nZ, AscendC::TPosition::A1>, NpuArch::Gemm::GemmType<int8_t, layout::zZ, AscendC::TPosition::A2>>{
127+ using Element = int8_t;
128+ using LayoutDst = layout::zZ;
129+ using LayoutSrc = layout::nZ;
130+ 
131+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
132+ 
133+ __aicore__ inline
134+ CopyL1ToL0A(){}
135+ 
136+ __aicore__ inline
137+ void operator()(
138+ AscendC::LocalTensor<Element> dstTensor,
139+ AscendC::LocalTensor<Element> srcTensor,
140+ LayoutDst layoutDst, LayoutSrc layoutSrc
141+ ){
142+ AscendC::LoadData2dTransposeParams loadDataParams;
143+ 
144+ loadDataParams.startIndex = 0;
145+ loadDataParams.repeatTimes = static_cast<uint16_t>(CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(1)));
146+ loadDataParams.srcStride = 1;
147+ loadDataParams.dstGap = 0;
148+ loadDataParams.dstFracGap = CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(1)) - 1;
149+ 
150+ for (uint32_t i = 0; i < CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(0)); i++) {
151+ AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2],
152+ srcTensor[i * layoutSrc.stride(1)],
153+ loadDataParams);
154+ }
155+ }
156+};
157+//////////////////////////////////////////
158+ 
159+/// Partial specialization for zN in and zZ out.
160+template <class ArchTag, class Element>
161+struct CopyL1ToL0A<ArchTag, Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>> {
162+ using LayoutDst = layout::zZ;
163+ using LayoutSrc = layout::zN;
164+ 
165+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
166+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
167+ 
168+ // Methods
169+ 
170+ __aicore__ inline
171+ CopyL1ToL0A() {};
172+ 
173+ __aicore__ inline
174+ void operator()(
175+ AscendC::LocalTensor<Element> const &dstTensor,
176+ AscendC::LocalTensor<Element> const &srcTensor,
177+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
178+ {
179+ AscendC::LoadData2DParams loadDataParams;
180+ 
181+ loadDataParams.startIndex = 0;
182+ loadDataParams.repeatTimes = static_cast<uint16_t>(layoutDst.shape(3));
183+ loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL;
184+ loadDataParams.sid = 0;
185+ loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1;
186+ loadDataParams.ifTranspose = false;
187+ loadDataParams.addrMode = 0;
188+ 
189+ for (uint32_t i = 0; i < layoutDst.shape(1); i++) {
190+ AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams);
191+ }
192+ }
193+};
194+ 
195+/// Partial specialization for float, zN in and zZ out.
196+template <class ArchTag>
197+struct CopyL1ToL0A<ArchTag, Gemm::GemmType<float, layout::zN, AscendC::TPosition::A1>> {
198+ using Element = float;
199+ using LayoutDst = layout::zZ;
200+ using LayoutSrc = layout::zN;
201+ 
202+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
203+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
204+ 
205+ // Methods
206+ 
207+ __aicore__ inline
208+ CopyL1ToL0A() {};
209+ 
210+ __aicore__ inline
211+ void operator()(
212+ AscendC::LocalTensor<Element> const &dstTensor,
213+ AscendC::LocalTensor<Element> const &srcTensor,
214+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
215+ {
216+ constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0};
217+ uint16_t l1M = layoutSrc.shape(0) * layoutSrc.shape(1);
218+ uint16_t l1K = layoutSrc.shape(2) * layoutSrc.shape(3);
219+ uint16_t l0M = layoutDst.shape(0) * layoutDst.shape(1);
220+ uint16_t l0K = layoutDst.shape(2) * layoutDst.shape(3);
221+ AscendC::SetFmatrix(1, l1M, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT);
222+ static constexpr AscendC::IsResetLoad3dConfig config = {false, false};
223+ AscendC::LoadData3DParamsV2<Element> loadDataParams;
224+ loadDataParams.kExtension = l0K;
225+ loadDataParams.mExtension = l0M;
226+ loadDataParams.channelSize = l1K;
227+ 
228+ AscendC::LoadData<Element, config>(dstTensor, srcTensor, loadDataParams);
229+ }
230+};
231+ 
232+template <class ArchTag, class Element>
233+struct CopyL1ToL0A<ArchTag, Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::A1>> {
234+ using LayoutDst = layout::zZ;
235+ using LayoutSrc = layout::nZ;
236+ 
237+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
238+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
239+ 
240+ __aicore__ inline
241+ CopyL1ToL0A() {};
242+ 
243+ __aicore__ inline
244+ void operator()(
245+ AscendC::LocalTensor<Element> const &dstTensor,
246+ AscendC::LocalTensor<Element> const &srcTensor,
247+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
248+ {
249+ AscendC::LoadData2DParams loadDataParams;
250+ 
251+ loadDataParams.startIndex = 0;
252+ loadDataParams.repeatTimes = static_cast<uint16_t>(CeilDiv<C0_NUM_PER_FRACTAL>(layoutDst.orgShape(1)));
253+ loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL;
254+ loadDataParams.sid = 0;
255+ loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1;
256+ loadDataParams.ifTranspose = true;
257+ loadDataParams.addrMode = 0;
258+ 
259+ for (uint32_t i = 0; i < CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(0)); i++) {
260+ AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams);
261+ }
262+ }
263+};
264+ 
265+/// Partial specialization for int8_t, nZ in and zZ out. (Transpose A)
266+template <class ArchTag>
267+struct CopyL1ToL0A<ArchTag, Gemm::GemmType<int8_t, layout::nZ, AscendC::TPosition::A1>> {
268+ using Element = int8_t;
269+ using LayoutDst = layout::zZ;
270+ using LayoutSrc = layout::nZ;
271+ 
272+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
273+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
274+ 
275+ // Methods
276+ 
277+ __aicore__ inline
278+ CopyL1ToL0A() {};
279+ 
280+ __aicore__ inline
281+ void operator()(
282+ AscendC::LocalTensor<Element> const &dstTensor,
283+ AscendC::LocalTensor<Element> const &srcTensor,
284+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
285+ {
286+ AscendC::LoadData2dTransposeParams loadDataParams;
287+ 
288+ loadDataParams.startIndex = 0;
289+ loadDataParams.repeatTimes = static_cast<uint16_t>(CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(1)));
290+ loadDataParams.srcStride = 1;
291+ loadDataParams.dstGap = 0;
292+ loadDataParams.dstFracGap = CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(1)) - 1;
293+ 
294+ for (uint32_t i = 0; i < CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(0)); i++) {
295+ AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2],
296+ srcTensor[i * layoutSrc.stride(1)],
297+ loadDataParams);
298+ }
299+ }
300+};
301+ 
302+/// Partial specialization for float, nZ in and zZ out. (Transpose A)
303+template <class ArchTag>
304+struct CopyL1ToL0A<ArchTag, Gemm::GemmType<float, layout::nZ, AscendC::TPosition::A1>> {
305+ using Element = float;
306+ using LayoutDst = layout::zZ;
307+ using LayoutSrc = layout::nZ;
308+ 
309+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
310+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
311+ 
312+ // Methods
313+ 
314+ __aicore__ inline
315+ CopyL1ToL0A() {};
316+ 
317+ __aicore__ inline
318+ void operator()(
319+ AscendC::LocalTensor<Element> const &dstTensor,
320+ AscendC::LocalTensor<Element> const &srcTensor,
321+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
322+ {
323+ constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0};
324+ uint16_t l1M = layoutSrc.shape(0) * layoutSrc.shape(1);
325+ uint16_t l1K = layoutSrc.shape(2) * layoutSrc.shape(3);
326+ uint16_t l0M = layoutDst.shape(0) * layoutDst.shape(1);
327+ uint16_t l0K = layoutDst.shape(2) * layoutDst.shape(3);
328+ // K, M need to be 16 aligned for f32
329+ uint16_t l1MAlign = RoundUp<C0_NUM_PER_FRACTAL>(l1M);
330+ uint16_t l1KAlign = RoundUp<C0_NUM_PER_FRACTAL>(l1K);
331+ uint16_t l0MAlign = RoundUp<C0_NUM_PER_FRACTAL>(l0M);
332+ uint16_t l0KAlign = RoundUp<C0_NUM_PER_FRACTAL>(l0K);
333+ AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_LEFT);
334+ static constexpr AscendC::IsResetLoad3dConfig config = {false, false};
335+ AscendC::LoadData3DParamsV2<Element> loadDataParams;
336+ loadDataParams.kExtension = l0MAlign;
337+ loadDataParams.mExtension = l0KAlign;
338+ loadDataParams.enTranspose = true;
339+ loadDataParams.channelSize = l1MAlign;
340+ 
341+ AscendC::LoadData<Element, config>(dstTensor, srcTensor, loadDataParams);
342+ }
343+};
344+ 
345+/////////////////////////////////////////////////////////////////////////////////////////////////////////////
346+ 
347+} // namespace NpuArch::Gemm::Tile
348+ 
349+#endif // GEMM_TILE_COPY_L1_TO_L0A_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/copy_l1_to_l0a_a5.hpp+335-0
@@ -0,0 +1,335 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_COPY_L1_TO_L0A_A5_HPP
12+#define GEMM_TILE_COPY_L1_TO_L0A_A5_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp"
17+#include "../../../tla/tensor.hpp"
18+ 
19+namespace NpuArch::Gemm::Tile {
20+ 
21+/// Partial specialization for CopyL1ToL0A, AtlasA5, zN in and zN out.
22+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
23+struct TileCopyTla<
24+ Arch::AtlasA5,
25+ tla::Tensor<AscendC::LocalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::A1>,
26+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::A2>,
27+ std::enable_if_t<tla::detail::iszN<ElementSrc, LayoutSrc>::value && tla::detail::iszN<ElementDst, LayoutDst>::value>> {
28+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
29+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(ElementSrc);
30+ 
31+ // Mehtods
32+ 
33+ __aicore__ inline
34+ TileCopyTla() {};
35+ 
36+ template <class TensorDst, class TensorSrc>
37+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor)
38+ {
39+ static_assert(
40+ tla::detail::iszN<typename TensorSrc::Element, typename TensorSrc::Layout>::value
41+ && tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
42+ && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2,
43+ "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zN"
44+ );
45+ 
46+ const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());
47+ const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());
48+ const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
49+ const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
50+ auto srcCoord = srcTensor.coord();
51+ 
52+ AscendC::LoadData2DParamsV2 loadDataParams;
53+ loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<0>(srcCoord));
54+ loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<1>(srcCoord));
55+ loadDataParams.mStep = dstOuterShapeRow;
56+ loadDataParams.kStep = dstOuterShapeCol;
57+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
58+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
59+ loadDataParams.ifTranspose = false;
60+ 
61+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
62+ AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams);
63+ }
64+ 
65+ template <class TensorDst, class TensorSrc>
66+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch)
67+ {
68+ static_assert(
69+ tla::detail::iszN<typename TensorSrc::Element, typename TensorSrc::Layout>::value
70+ && tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
71+ && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2,
72+ "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0A and zN"
73+ );
74+ 
75+ const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());
76+ const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());
77+ const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
78+ const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
79+ 
80+ AscendC::LoadData2DParamsV2 loadDataParams;
81+ loadDataParams.mStartPosition = 0;
82+ loadDataParams.kStartPosition = 0;
83+ loadDataParams.mStep = dstOuterShapeRow;
84+ loadDataParams.kStep = dstOuterShapeCol * l0Batch;
85+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
86+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
87+ loadDataParams.ifTranspose = false;
88+ 
89+ AscendC::LoadData(dstTensor.data(), srcTensor.data(), loadDataParams);
90+ }
91+};
92+ 
93+/// Partial specialization for CopyL1ToL0A, AtlasA5, not B8, nZ in and zN out. (Transpose A)
94+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
95+struct TileCopyTla<
96+ Arch::AtlasA5,
97+ tla::Tensor<AscendC::LocalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::A1>,
98+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::A2>,
99+ std::enable_if_t<
100+ !AscendC::Std::is_one_of_v<ElementSrc, int8_t, float8_e4m3_t, float8_e5m2_t> &&
101+ !AscendC::Std::is_one_of_v<ElementDst, int8_t, float8_e4m3_t, float8_e5m2_t> &&
102+ tla::detail::isnZ<ElementSrc, LayoutSrc>::value && tla::detail::iszN<ElementDst, LayoutDst>::value>> {
103+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
104+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(ElementSrc);
105+ 
106+ // Mehtods
107+ 
108+ __aicore__ inline
109+ TileCopyTla() {};
110+ 
111+ template <class TensorDst, class TensorSrc>
112+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor)
113+ {
114+ static_assert(
115+ !AscendC::Std::is_one_of_v<typename TensorSrc::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
116+ !AscendC::Std::is_one_of_v<typename TensorDst::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
117+ tla::detail::isnZ<typename TensorSrc::Element, typename TensorSrc::Layout>::value
118+ && tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
119+ && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2,
120+ "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"
121+ );
122+ 
123+ const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
124+ const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
125+ const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
126+ const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
127+ auto srcCoord = srcTensor.coord();
128+ 
129+ AscendC::LoadData2DParamsV2 loadDataParams;
130+ loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<1>(srcCoord));
131+ loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<0>(srcCoord));
132+ loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);
133+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);
134+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
135+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
136+ loadDataParams.ifTranspose = true;
137+ 
138+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
139+ AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams);
140+ }
141+ 
142+ template <class TensorDst, class TensorSrc>
143+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch)
144+ {
145+ static_assert(
146+ !AscendC::Std::is_one_of_v<typename TensorSrc::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
147+ !AscendC::Std::is_one_of_v<typename TensorDst::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
148+ tla::detail::isnZ<typename TensorSrc::Element, typename TensorSrc::Layout>::value
149+ && tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
150+ && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2,
151+ "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"
152+ );
153+ 
154+ const uint32_t L1M = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape());
155+ const uint32_t L1K = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());
156+ const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
157+ const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
158+ const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
159+ const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
160+ 
161+ AscendC::LoadData2DParamsV2 loadDataParams;
162+ loadDataParams.mStartPosition = 0;
163+ loadDataParams.kStartPosition = 0;
164+ loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);
165+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);
166+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
167+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
168+ loadDataParams.ifTranspose = true;
169+ 
170+ for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {
171+ AscendC::LoadData(
172+ dstTensor.data()[l0BatchIdx * L0M * L0K], srcTensor.data()[l0BatchIdx * L1M * L1K], loadDataParams
173+ );
174+ }
175+ }
176+};
177+ 
178+/// Partial specialization for CopyL1ToL0A, AtlasA5, B8, nZ in and zN out. (Transpose A)
179+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
180+struct TileCopyTla<
181+ Arch::AtlasA5,
182+ tla::Tensor<AscendC::LocalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::A1>,
183+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::A2>,
184+ std::enable_if_t<
185+ AscendC::Std::is_one_of_v<ElementSrc, int8_t, float8_e4m3_t, float8_e5m2_t> &&
186+ AscendC::Std::is_one_of_v<ElementDst, int8_t, float8_e4m3_t, float8_e5m2_t> &&
187+ tla::detail::isnZ<ElementSrc, LayoutSrc>::value && tla::detail::iszN<ElementDst, LayoutDst>::value>> {
188+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
189+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(ElementSrc);
190+ 
191+ // Mehtods
192+ 
193+ __aicore__ inline
194+ TileCopyTla() {};
195+ 
196+ template <class TensorDst, class TensorSrc>
197+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor)
198+ {
199+ static_assert(
200+ AscendC::Std::is_one_of_v<typename TensorSrc::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
201+ AscendC::Std::is_one_of_v<typename TensorDst::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
202+ tla::detail::isnZ<typename TensorSrc::Element, typename TensorSrc::Layout>::value
203+ && tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
204+ && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2,
205+ "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"
206+ );
207+ 
208+ const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
209+ const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
210+ const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
211+ const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
212+ auto srcCoord = srcTensor.coord();
213+ 
214+ AscendC::LoadData2DParamsV2 loadDataParams;
215+ if (L0M % ELE_NUM_PER_C0 == 0) {
216+ loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<1>(srcCoord));
217+ loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<0>(srcCoord));
218+ loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);
219+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);
220+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
221+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
222+ loadDataParams.ifTranspose = true;
223+ 
224+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
225+ AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams);
226+ } else {
227+ for (uint32_t kIdx = 0; kIdx < L0K / ELE_NUM_PER_C0; kIdx++) {
228+ loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<1>(srcCoord)) + kIdx * 2;
229+ loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<0>(srcCoord));
230+ loadDataParams.mStep = 2;
231+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);
232+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
233+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
234+ loadDataParams.ifTranspose = true;
235+ 
236+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
237+ AscendC::LoadData(
238+ dstTensor.data()[dstOffset + kIdx * L0M * ELE_NUM_PER_C0], srcTensor.data(), loadDataParams
239+ );
240+ }
241+ }
242+ }
243+ 
244+ template <class TensorDst, class TensorSrc>
245+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch)
246+ {
247+ static_assert(
248+ AscendC::Std::is_one_of_v<typename TensorSrc::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
249+ AscendC::Std::is_one_of_v<typename TensorDst::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
250+ tla::detail::isnZ<typename TensorSrc::Element, typename TensorSrc::Layout>::value
251+ && tla::detail::iszN<typename TensorDst::Element, typename TensorDst::Layout>::value
252+ && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::A2,
253+ "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0A and zN"
254+ );
255+ 
256+ const uint32_t L1M = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape());
257+ const uint32_t L1K = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());
258+ const uint32_t L0M = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
259+ const uint32_t L0K = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
260+ const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
261+ const uint32_t dstOuterStrideCol = tla::get<1, 1>(dstTensor.stride());
262+ 
263+ AscendC::LoadData2DParamsV2 loadDataParams;
264+ if (L0M % ELE_NUM_PER_C0 == 0) {
265+ loadDataParams.mStartPosition = 0;
266+ loadDataParams.kStartPosition = 0;
267+ loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);
268+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);
269+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
270+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
271+ loadDataParams.ifTranspose = true;
272+ 
273+ for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {
274+ AscendC::LoadData(
275+ dstTensor.data()[l0BatchIdx * L0M * L0K], srcTensor.data()[l0BatchIdx * L1M * L1K], loadDataParams
276+ );
277+ }
278+ } else {
279+ loadDataParams.mStartPosition = 0;
280+ loadDataParams.kStartPosition = 0;
281+ loadDataParams.mStep = 2;
282+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0M);
283+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
284+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideCol);
285+ loadDataParams.ifTranspose = true;
286+ for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {
287+ for (uint32_t kIdx = 0; kIdx < L0K / ELE_NUM_PER_C0; kIdx++) {
288+ AscendC::LoadData(
289+ dstTensor.data()[l0BatchIdx * L0M * L0K + kIdx * L0M * ELE_NUM_PER_C0],
290+ srcTensor.data()[l0BatchIdx * L1M * L1K + kIdx * ELE_NUM_PER_FRACTAL * 2], loadDataParams
291+ );
292+ }
293+ }
294+ }
295+ }
296+};
297+ 
298+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
299+struct TileCopyTla<
300+ Arch::AtlasA5,
301+ tla::Tensor<AscendC::LocalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::A1>,
302+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::A2>,
303+ std::enable_if_t<tla::detail::isVector<LayoutSrc>::value>> {
304+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
305+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(ElementSrc);
306+ 
307+ __aicore__ inline
308+ TileCopyTla() {};
309+ 
310+ template <class TensorDst, class TensorSrc>
311+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor)
312+ {
313+ uint16_t aL1M = tla::get<0, 0>(srcTensor.stride());
314+ uint16_t madM = tla::get<1, 1>(dstTensor.stride());
315+ uint16_t madK = tla::get<1, 1>(dstTensor.shape());
316+ 
317+ AscendC::LoadData2DParamsV2 loadDataParams;
318+ loadDataParams.mStartPosition = 0;
319+ loadDataParams.kStartPosition = 0;
320+ loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(madM);
321+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(madK);
322+ loadDataParams.srcStride = CeilDiv<C0_NUM_PER_FRACTAL>(aL1M);
323+ loadDataParams.dstStride = CeilDiv<C0_NUM_PER_FRACTAL>(madM);
324+ 
325+ loadDataParams.ifTranspose = false;
326+ auto srcOffset = srcTensor.layout()(srcTensor.coord());
327+ AscendC::LoadData(dstTensor.data(), srcTensor.data()[srcOffset], loadDataParams);
328+ }
329+};
330+ 
331+/////////////////////////////////////////////////////////////////////////////////////////////////////////////
332+ 
333+} // namespace NpuArch::Gemm::Tile
334+ 
335+#endif // GEMM_TILE_COPY_L1_TO_L0A_A5_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/copy_l1_to_l0b_a2.hpp+479-0
@@ -0,0 +1,479 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_COPY_L1_TO_L0B_HPP
12+#define GEMM_TILE_COPY_L1_TO_L0B_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/layout/layout.hpp"
17+#include "../../../attn_infra/gemm/gemm_type.hpp"
18+#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp"
19+ 
20+namespace NpuArch::Gemm::Tile {
21+ 
22+template <
23+ class ArchTag,
24+ class L1Type,
25+ class L0Type = void
26+>
27+struct CopyL1ToL0B {
28+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported copy l1 to l0, can not find the specialization.");
29+};
30+ 
31+////////////////////////////////////////
32+/// new add gemm
33+template<class ArchTag, class Element>
34+struct CopyL1ToL0B<ArchTag, NpuArch::Gemm::GemmType<Element, layout::zZ, AscendC::TPosition::B1>, NpuArch::Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::B2>>{
35+ using LayoutDst = layout::nZ;
36+ using LayoutSrc = layout::zZ;
37+ 
38+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
39+ 
40+ __aicore__ inline
41+ CopyL1ToL0B(){}
42+ 
43+ __aicore__ inline
44+ void operator()(
45+ AscendC::LocalTensor<Element> dstTensor,
46+ AscendC::LocalTensor<Element> srcTensor,
47+ LayoutDst layoutDst, LayoutSrc layoutSrc
48+ ){
49+ AscendC::LoadData2DParams loadDataParams;
50+ loadDataParams.startIndex = 0;
51+ loadDataParams.repeatTimes = static_cast<uint16_t>(CeilDiv<ELE_NUM_PER_C0>(layoutSrc.orgShape(1)));
52+ loadDataParams.srcStride = 1;
53+ loadDataParams.sid = 0;
54+ loadDataParams.dstGap = 0;
55+ loadDataParams.ifTranspose = true;
56+ loadDataParams.addrMode = 0;
57+ for(uint32_t i = 0; i < CeilDiv<C0_NUM_PER_FRACTAL>(layoutDst.orgShape(0)); i++){ // K N
58+ AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams);
59+ }
60+ }
61+};
62+ 
63+template<class ArchTag>
64+struct CopyL1ToL0B<ArchTag, NpuArch::Gemm::GemmType<float, layout::zZ, AscendC::TPosition::B1>, NpuArch::Gemm::GemmType<float, layout::nZ, AscendC::TPosition::B2>>{
65+ using Element = float;
66+ using LayoutDst = layout::nZ;
67+ using LayoutSrc = layout::zZ;
68+ 
69+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
70+ 
71+ __aicore__ inline
72+ CopyL1ToL0B(){}
73+ 
74+ __aicore__ inline
75+ void operator()(
76+ AscendC::LocalTensor<Element> dstTensor,
77+ AscendC::LocalTensor<Element> srcTensor,
78+ LayoutDst layoutDst, LayoutSrc layoutSrc
79+ ){
80+ AscendC::LoadData2dTransposeParams loadDataParams;
81+ loadDataParams.startIndex = 0;
82+ loadDataParams.repeatTimes = static_cast<uint16_t>(CeilDiv<C0_NUM_PER_FRACTAL>(layoutSrc.orgShape(1)));
83+ loadDataParams.srcStride = 1;
84+ loadDataParams.dstGap = 0;
85+ loadDataParams.dstFracGap = static_cast<uint16_t>(CeilDiv<C0_NUM_PER_FRACTAL>(layoutDst.orgShape(1))) - 1;
86+ for(uint32_t i = 0; i < CeilDiv<C0_NUM_PER_FRACTAL>(layoutDst.orgShape(0)); i++){ // K N
87+ AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1) * 2], srcTensor[i * layoutSrc.stride(1)], loadDataParams);
88+ }
89+ }
90+};
91+ 
92+ 
93+template<class ArchTag>
94+struct CopyL1ToL0B<ArchTag, NpuArch::Gemm::GemmType<int8_t, layout::zN, AscendC::TPosition::B1>, NpuArch::Gemm::GemmType<int8_t, layout::nZ, AscendC::TPosition::B2>>{
95+ using Element = int8_t;
96+ using LayoutDst = layout::nZ;
97+ using LayoutSrc = layout::zN;
98+ 
99+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
100+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
101+ 
102+ __aicore__ inline
103+ CopyL1ToL0B(){}
104+ 
105+ __aicore__ inline
106+ void operator()(
107+ AscendC::LocalTensor<Element> dstTensor,
108+ AscendC::LocalTensor<Element> srcTensor,
109+ LayoutDst layoutDst, LayoutSrc layoutSrc
110+ ){
111+ AscendC::LoadData2dTransposeParams loadDataParams;
112+ 
113+ loadDataParams.startIndex = 0;
114+ loadDataParams.repeatTimes = static_cast<uint16_t>(CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(1)));
115+ loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL / 2;
116+ loadDataParams.dstGap = 1;
117+ loadDataParams.dstFracGap = 0;
118+ 
119+ for (uint32_t i = 0; i < CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(0)); i++) {
120+ AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)],
121+ srcTensor[i * layoutSrc.stride(1) * 2],
122+ loadDataParams);
123+ }
124+ }
125+};
126+ 
127+template <class ArchTag, class Element>
128+struct CopyL1ToL0B<ArchTag, NpuArch::Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::B1>, NpuArch::Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::B2>> {
129+ using LayoutDst = layout::nZ;
130+ using LayoutSrc = layout::nZ;
131+ 
132+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
133+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
134+ 
135+ // Methods
136+ 
137+ __aicore__ inline
138+ CopyL1ToL0B() {};
139+ 
140+ __aicore__ inline
141+ void operator()(
142+ AscendC::LocalTensor<Element> const &dstTensor,
143+ AscendC::LocalTensor<Element> const &srcTensor,
144+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
145+ {
146+ AscendC::LoadData2DParams loadDataParams;
147+ 
148+ loadDataParams.startIndex = 0;
149+ loadDataParams.repeatTimes = static_cast<uint16_t>(layoutDst.shape(3));
150+ loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL;
151+ loadDataParams.sid = 0;
152+ loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1;
153+ loadDataParams.ifTranspose = false;
154+ loadDataParams.addrMode = 0;
155+ 
156+ for (uint32_t i = 0; i < layoutDst.shape(1); i++) {
157+ AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams);
158+ }
159+ }
160+};
161+/////////////////////////////////////////////
162+ 
163+////////////////////////////////////////////
164+/// new add gemv
165+template <class ArchTag, class Element>
166+struct CopyL1ToL0B<ArchTag, NpuArch::Gemm::GemmType<Element, layout::zN, AscendC::TPosition::B1>, NpuArch::Gemm::GemmType<Element, layout::zN, AscendC::TPosition::B2>>{
167+ using LayoutDst = layout::zN;
168+ using LayoutSrc = layout::zN;
169+ 
170+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
171+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
172+ 
173+ // Methods
174+ 
175+ __aicore__ inline
176+ CopyL1ToL0B() {};
177+ 
178+ __aicore__ inline
179+ void operator()(
180+ AscendC::LocalTensor<Element> const &dstTensor,
181+ AscendC::LocalTensor<Element> const &srcTensor,
182+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
183+ {
184+ AscendC::LoadData2DParams loadDataParams;
185+ 
186+ loadDataParams.startIndex = 0;
187+ loadDataParams.repeatTimes = static_cast<uint16_t>(layoutDst.shape(1));
188+ loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL;
189+ loadDataParams.sid = 0;
190+ loadDataParams.dstGap = layoutDst.stride(1) / ELE_NUM_PER_FRACTAL - 1;
191+ loadDataParams.ifTranspose = false;
192+ loadDataParams.addrMode = 0;
193+ 
194+ for (uint32_t i = 0; i < layoutDst.shape(3); i++)
195+ {
196+ AscendC::LoadData(dstTensor[i * layoutDst.stride(3)], srcTensor[i * layoutSrc.stride(3)], loadDataParams);
197+ }
198+ }
199+};
200+ 
201+template <class ArchTag, class Element>
202+struct CopyL1ToL0B<ArchTag, NpuArch::Gemm::GemmType<Element, layout::nN, AscendC::TPosition::B1>, NpuArch::Gemm::GemmType<Element, layout::zN, AscendC::TPosition::B2>>
203+{
204+ using LayoutDst = layout::zN;
205+ using LayoutSrc = layout::nN;
206+ 
207+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
208+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
209+ 
210+ // Methods
211+ 
212+ __aicore__ inline
213+ CopyL1ToL0B() {};
214+ 
215+ __aicore__ inline
216+ void operator()(
217+ AscendC::LocalTensor<Element> const &dstTensor,
218+ AscendC::LocalTensor<Element> const &srcTensor,
219+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
220+ {
221+ AscendC::LoadData2DParams loadDataParams;
222+ 
223+ loadDataParams.startIndex = 0;
224+ loadDataParams.repeatTimes = layoutDst.shape(1) * layoutDst.shape(3);
225+ loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL;
226+ loadDataParams.sid = 0;
227+ loadDataParams.dstGap = layoutDst.stride(1) / ELE_NUM_PER_FRACTAL - 1;
228+ loadDataParams.ifTranspose = true;
229+ loadDataParams.addrMode = 0;
230+ AscendC::LoadData(dstTensor, srcTensor, loadDataParams);
231+ };
232+};
233+ 
234+template <class ArchTag>
235+struct CopyL1ToL0B<ArchTag, NpuArch::Gemm::GemmType<float, layout::nN, AscendC::TPosition::B1>, NpuArch::Gemm::GemmType<float, layout::zN, AscendC::TPosition::B2>>{
236+ using LayoutDst = layout::zN;
237+ using LayoutSrc = layout::nN;
238+ using Element = float;
239+ 
240+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
241+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
242+ 
243+ // Methods
244+ 
245+ __aicore__ inline
246+ CopyL1ToL0B() {};
247+ 
248+ __aicore__ inline
249+ void operator()(
250+ AscendC::LocalTensor<Element> const &dstTensor,
251+ AscendC::LocalTensor<Element> const &srcTensor,
252+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
253+ {
254+ AscendC::LoadData2dTransposeParams loadDataParams;
255+ 
256+ loadDataParams.startIndex = 0;
257+ loadDataParams.repeatTimes = static_cast<uint16_t>(CeilDiv<C0_NUM_PER_FRACTAL>(layoutDst.orgShape(0)));
258+ loadDataParams.srcStride = 1;
259+ loadDataParams.dstGap = 0;
260+ loadDataParams.dstFracGap = CeilDiv<C0_NUM_PER_FRACTAL>(layoutDst.orgShape(0)) - 1;
261+ 
262+ for (uint32_t i = 0; i < CeilDiv<2 * ELE_NUM_PER_C0>(layoutDst.orgShape(1)); i++)
263+ {
264+ AscendC::LoadDataWithTranspose(
265+ dstTensor[i * layoutDst.stride(3) * 2],
266+ srcTensor[i * layoutSrc.stride(3)],
267+ loadDataParams);
268+ }
269+ };
270+};
271+ 
272+template <class ArchTag>
273+struct CopyL1ToL0B<ArchTag, NpuArch::Gemm::GemmType<int8_t, layout::nZ, AscendC::TPosition::B1>, NpuArch::Gemm::GemmType<int8_t, layout::zN, AscendC::TPosition::B2>>{
274+ using LayoutDst = layout::zN;
275+ using LayoutSrc = layout::nZ;
276+ using Element = int8_t;
277+ 
278+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
279+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
280+ 
281+ // Methods
282+ 
283+ __aicore__ inline
284+ CopyL1ToL0B() {};
285+ 
286+ __aicore__ inline
287+ void operator()(
288+ AscendC::LocalTensor<Element> const &dstTensor,
289+ AscendC::LocalTensor<Element> const &srcTensor,
290+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
291+ {
292+ AscendC::LoadData2dTransposeParams loadDataParams;
293+ 
294+ loadDataParams.startIndex = 0;
295+ loadDataParams.repeatTimes = static_cast<uint16_t>(CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(0)));
296+ loadDataParams.srcStride = layoutSrc.stride(1) / ELE_NUM_PER_FRACTAL / 2;
297+ loadDataParams.dstGap = 1;
298+ loadDataParams.dstFracGap = 0;
299+ 
300+ for (uint32_t i = 0; i < CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(1)); i++)
301+ {
302+ AscendC::LoadDataWithTranspose(
303+ dstTensor[i * layoutDst.stride(3)],
304+ srcTensor[i * layoutSrc.stride(3) * 2],
305+ loadDataParams);
306+ }
307+ }
308+};
309+////////////////////////////////////////////
310+ 
311+/// Partial specialization for int8_t, zN in and nZ out.
312+template <class ArchTag>
313+struct CopyL1ToL0B<ArchTag, Gemm::GemmType<int8_t, layout::zN, AscendC::TPosition::A1>> {
314+ using Element = int8_t;
315+ using LayoutDst = layout::nZ;
316+ using LayoutSrc = layout::zN;
317+ 
318+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
319+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
320+ 
321+ // Methods
322+ 
323+ __aicore__ inline
324+ CopyL1ToL0B() {};
325+ 
326+ __aicore__ inline
327+ void operator()(
328+ AscendC::LocalTensor<Element> const &dstTensor,
329+ AscendC::LocalTensor<Element> const &srcTensor,
330+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
331+ {
332+ AscendC::LoadData2dTransposeParams loadDataParams;
333+ 
334+ loadDataParams.startIndex = 0;
335+ loadDataParams.repeatTimes = static_cast<uint16_t>(CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(1)));
336+ loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL / 2;
337+ loadDataParams.dstGap = 1;
338+ loadDataParams.dstFracGap = 0;
339+ 
340+ for (uint32_t i = 0; i < CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(0)); i++) {
341+ AscendC::LoadDataWithTranspose(dstTensor[i * layoutDst.stride(1)],
342+ srcTensor[i * layoutSrc.stride(1) * 2],
343+ loadDataParams);
344+ }
345+ }
346+};
347+ 
348+/// Partial specialization for float, zN in and nZ out.
349+template <class ArchTag>
350+struct CopyL1ToL0B<ArchTag, Gemm::GemmType<float, layout::zN, AscendC::TPosition::A1>> {
351+ using Element = float;
352+ using LayoutDst = layout::nZ;
353+ using LayoutSrc = layout::zN;
354+ 
355+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
356+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
357+ 
358+ // Methods
359+ 
360+ __aicore__ inline
361+ CopyL1ToL0B() {};
362+ 
363+ __aicore__ inline
364+ void operator()(
365+ AscendC::LocalTensor<Element> const &dstTensor,
366+ AscendC::LocalTensor<Element> const &srcTensor,
367+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
368+ {
369+ constexpr uint8_t PAD_LIST[4] = {0, 0, 0, 0};
370+ uint16_t l1K = layoutSrc.shape(0) * layoutSrc.shape(1);
371+ uint16_t l1N = layoutSrc.shape(2) * layoutSrc.shape(3);
372+ uint16_t l0K = layoutDst.shape(0) * layoutDst.shape(1);
373+ uint16_t l0N = layoutDst.shape(2) * layoutDst.shape(3);
374+ // K, N need to be 16 aligned for f32
375+ uint16_t l1KAlign = RoundUp<C0_NUM_PER_FRACTAL>(l1K);
376+ uint16_t l1NAlign = RoundUp<C0_NUM_PER_FRACTAL>(l1N);
377+ uint16_t l0KAlign = RoundUp<C0_NUM_PER_FRACTAL>(l0K);
378+ uint16_t l0NAlign = RoundUp<C0_NUM_PER_FRACTAL>(l0N);
379+ AscendC::SetFmatrix(1, l1KAlign, PAD_LIST, AscendC::FmatrixMode::FMATRIX_RIGHT);
380+ static constexpr AscendC::IsResetLoad3dConfig config = {false, false};
381+ AscendC::LoadData3DParamsV2<Element> loadDataParams;
382+ loadDataParams.kExtension = l0NAlign;
383+ loadDataParams.mExtension = l0KAlign;
384+ loadDataParams.channelSize = l1NAlign;
385+ loadDataParams.fMatrixCtrl = true;
386+ 
387+ AscendC::LoadData<Element, config>(dstTensor, srcTensor, loadDataParams);
388+ }
389+};
390+ 
391+/// Partial specialization for zN in and nZ out.
392+template <class ArchTag, class Element>
393+struct CopyL1ToL0B<ArchTag, Gemm::GemmType<Element, layout::zN, AscendC::TPosition::A1>> {
394+ using LayoutDst = layout::nZ;
395+ using LayoutSrc = layout::zN;
396+ 
397+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
398+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
399+ 
400+ // Methods
401+ 
402+ __aicore__ inline
403+ CopyL1ToL0B() {};
404+ 
405+ __aicore__ inline
406+ void operator()(
407+ AscendC::LocalTensor<Element> const &dstTensor,
408+ AscendC::LocalTensor<Element> const &srcTensor,
409+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
410+ {
411+ AscendC::LoadData2DParams loadDataParams;
412+ 
413+ loadDataParams.startIndex = 0;
414+ loadDataParams.repeatTimes = static_cast<uint16_t>(CeilDiv<ELE_NUM_PER_C0>(layoutDst.orgShape(1)));
415+ loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL;
416+ loadDataParams.sid = 0;
417+ loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1;
418+ loadDataParams.ifTranspose = true;
419+ loadDataParams.addrMode = 0;
420+ 
421+ for (uint32_t i = 0; i < CeilDiv<C0_NUM_PER_FRACTAL>(layoutDst.orgShape(0)); i++) {
422+ AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams);
423+ }
424+ }
425+};
426+ 
427+/// Partial specialization for nZ in and nZ out. (Transpose B)
428+template <class ArchTag, class Element>
429+struct CopyL1ToL0B<ArchTag, Gemm::GemmType<Element, layout::nZ, AscendC::TPosition::A1>> {
430+ using LayoutDst = layout::nZ;
431+ using LayoutSrc = layout::nZ;
432+ 
433+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
434+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
435+ 
436+ // Methods
437+ 
438+ __aicore__ inline
439+ CopyL1ToL0B() {};
440+ 
441+ __aicore__ inline
442+ void operator()(
443+ AscendC::LocalTensor<Element> const &dstTensor,
444+ AscendC::LocalTensor<Element> const &srcTensor,
445+ LayoutDst const &layoutDst, LayoutSrc const &layoutSrc)
446+ {
447+ AscendC::LoadData2DParams loadDataParams;
448+ if (layoutSrc.shape(3) == layoutDst.shape(3)) {
449+ loadDataParams.startIndex = 0;
450+ loadDataParams.repeatTimes = static_cast<uint16_t>(layoutDst.shape(1) * layoutDst.shape(3));
451+ loadDataParams.srcStride = 1;
452+ loadDataParams.sid = 0;
453+ loadDataParams.dstGap = 0;
454+ loadDataParams.ifTranspose = false;
455+ loadDataParams.addrMode = 0;
456+ 
457+ AscendC::LoadData(dstTensor, srcTensor, loadDataParams);
458+ } else {
459+ loadDataParams.startIndex = 0;
460+ loadDataParams.repeatTimes = static_cast<uint16_t>(layoutDst.shape(3));
461+ loadDataParams.srcStride = layoutSrc.stride(3) / ELE_NUM_PER_FRACTAL;
462+ loadDataParams.sid = 0;
463+ loadDataParams.dstGap = layoutDst.stride(3) / ELE_NUM_PER_FRACTAL - 1;
464+ loadDataParams.ifTranspose = false;
465+ loadDataParams.addrMode = 0;
466+ 
467+ for (uint32_t i = 0; i < layoutDst.shape(1); i++) {
468+ AscendC::LoadData(dstTensor[i * layoutDst.stride(1)], srcTensor[i * layoutSrc.stride(1)], loadDataParams);
469+ }
470+ }
471+ 
472+ }
473+};
474+ 
475+/////////////////////////////////////////////////////////////////////////////////////////////////////////////
476+ 
477+} // namespace NpuArch::Gemm::Tile
478+ 
479+#endif // GEMM_TILE_COPY_L1_TO_L0B_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/copy_l1_to_l0b_a5.hpp+302-0
@@ -0,0 +1,302 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_COPY_L1_TO_L0B_A5_HPP
12+#define GEMM_TILE_COPY_L1_TO_L0B_A5_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp"
17+#include "../../../tla/tensor.hpp"
18+ 
19+namespace NpuArch::Gemm::Tile {
20+ 
21+/// Partial specialization for CopyL1ToL0B, AtlasA5, not B8, zN in and nZ out.
22+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
23+struct TileCopyTla<
24+ Arch::AtlasA5,
25+ tla::Tensor<AscendC::LocalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::A1>,
26+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::B2>,
27+ std::enable_if_t<
28+ !AscendC::Std::is_one_of_v<ElementSrc, int8_t, float8_e4m3_t, float8_e5m2_t> &&
29+ !AscendC::Std::is_one_of_v<ElementDst, int8_t, float8_e4m3_t, float8_e5m2_t> &&
30+ tla::detail::iszN<ElementSrc, LayoutSrc>::value && tla::detail::isnZ<ElementDst, LayoutDst>::value>> {
31+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
32+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(ElementSrc);
33+ 
34+ // Mehtods
35+ 
36+ __aicore__ inline
37+ TileCopyTla() {};
38+ 
39+ template <class TensorDst, class TensorSrc>
40+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor)
41+ {
42+ static_assert(
43+ !AscendC::Std::is_one_of_v<typename TensorSrc::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
44+ !AscendC::Std::is_one_of_v<typename TensorDst::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
45+ tla::detail::iszN<typename TensorSrc::Element, typename TensorSrc::Layout>::value
46+ && tla::detail::isnZ<typename TensorDst::Element, typename TensorDst::Layout>::value
47+ && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2,
48+ "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"
49+ );
50+ 
51+ const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
52+ const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
53+ const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
54+ const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
55+ auto srcCoord = srcTensor.coord();
56+ 
57+ AscendC::LoadData2DParamsV2 loadDataParams;
58+ loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<0>(srcCoord));
59+ loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<1>(srcCoord));
60+ loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);
61+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0N);
62+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
63+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
64+ loadDataParams.ifTranspose = true;
65+ 
66+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
67+ AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams);
68+ }
69+ 
70+ template <class TensorDst, class TensorSrc>
71+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch)
72+ {
73+ static_assert(
74+ !AscendC::Std::is_one_of_v<typename TensorSrc::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
75+ !AscendC::Std::is_one_of_v<typename TensorDst::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
76+ tla::detail::iszN<typename TensorSrc::Element, typename TensorSrc::Layout>::value
77+ && tla::detail::isnZ<typename TensorDst::Element, typename TensorDst::Layout>::value
78+ && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2,
79+ "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"
80+ );
81+ 
82+ const uint32_t L1K = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape());
83+ const uint32_t L1N = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());
84+ const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
85+ const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
86+ const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
87+ const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
88+ 
89+ AscendC::LoadData2DParamsV2 loadDataParams;
90+ loadDataParams.mStartPosition = 0;
91+ loadDataParams.kStartPosition = 0;
92+ loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);
93+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0N);
94+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
95+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
96+ loadDataParams.ifTranspose = true;
97+ 
98+ for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {
99+ AscendC::LoadData(
100+ dstTensor.data()[l0BatchIdx * L0N * L0K], srcTensor.data()[l0BatchIdx * L1N * L1K], loadDataParams
101+ );
102+ }
103+ }
104+};
105+ 
106+/// Partial specialization for CopyL1ToL0B, AtlasA5, B8, zN in and nZ out.
107+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
108+struct TileCopyTla<
109+ Arch::AtlasA5,
110+ tla::Tensor<AscendC::LocalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::A1>,
111+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::B2>,
112+ std::enable_if_t<
113+ AscendC::Std::is_one_of_v<ElementSrc, int8_t, float8_e4m3_t, float8_e5m2_t> &&
114+ AscendC::Std::is_one_of_v<ElementDst, int8_t, float8_e4m3_t, float8_e5m2_t> &&
115+ tla::detail::iszN<ElementSrc, LayoutSrc>::value && tla::detail::isnZ<ElementDst, LayoutDst>::value>> {
116+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
117+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(ElementSrc);
118+ 
119+ // Mehtods
120+ 
121+ __aicore__ inline
122+ TileCopyTla() {};
123+ 
124+ template <class TensorDst, class TensorSrc>
125+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor)
126+ {
127+ static_assert(
128+ AscendC::Std::is_one_of_v<typename TensorSrc::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
129+ AscendC::Std::is_one_of_v<typename TensorDst::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
130+ tla::detail::iszN<typename TensorSrc::Element, typename TensorSrc::Layout>::value
131+ && tla::detail::isnZ<typename TensorDst::Element, typename TensorDst::Layout>::value
132+ && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2,
133+ "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"
134+ );
135+ 
136+ const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
137+ const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
138+ const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
139+ const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
140+ auto srcCoord = srcTensor.coord();
141+ 
142+ AscendC::LoadData2DParamsV2 loadDataParams;
143+ if (L0N % ELE_NUM_PER_C0 == 0) {
144+ loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<0>(srcCoord));
145+ loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<1>(srcCoord));
146+ loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);
147+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0N);
148+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
149+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
150+ loadDataParams.ifTranspose = true;
151+ 
152+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
153+ AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams);
154+ } else {
155+ for (uint32_t kIdx = 0; kIdx < L0K / ELE_NUM_PER_C0; kIdx++) {
156+ loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<0>(srcCoord)) + kIdx * 2;
157+ loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<1>(srcCoord));
158+ loadDataParams.mStep = 2;
159+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0N);
160+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
161+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
162+ loadDataParams.ifTranspose = true;
163+ 
164+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
165+ AscendC::LoadData(
166+ dstTensor.data()[dstOffset + kIdx * L0N * ELE_NUM_PER_C0], srcTensor.data(), loadDataParams
167+ );
168+ }
169+ }
170+ }
171+ 
172+ template <class TensorDst, class TensorSrc>
173+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch)
174+ {
175+ static_assert(
176+ AscendC::Std::is_one_of_v<typename TensorSrc::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
177+ AscendC::Std::is_one_of_v<typename TensorDst::Element, int8_t, float8_e4m3_t, float8_e5m2_t> &&
178+ tla::detail::iszN<typename TensorSrc::Element, typename TensorSrc::Layout>::value
179+ && tla::detail::isnZ<typename TensorDst::Element, typename TensorDst::Layout>::value
180+ && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2,
181+ "The input parameters do not match. TensorSrc must be L1 and zN, while TensorDst must be L0B and nZ"
182+ );
183+ 
184+ const uint32_t L1K = tla::get<0, 0>(srcTensor.shape()) * tla::get<0, 1>(srcTensor.shape());
185+ const uint32_t L1N = tla::get<1, 0>(srcTensor.shape()) * tla::get<1, 1>(srcTensor.shape());
186+ const uint32_t L0K = tla::get<0, 0>(dstTensor.shape()) * tla::get<0, 1>(dstTensor.shape());
187+ const uint32_t L0N = tla::get<1, 0>(dstTensor.shape()) * tla::get<1, 1>(dstTensor.shape());
188+ const uint32_t srcOuterStrideCol = tla::get<1, 1>(srcTensor.stride());
189+ const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
190+ 
191+ AscendC::LoadData2DParamsV2 loadDataParams;
192+ if (L0N % ELE_NUM_PER_C0 == 0) {
193+ loadDataParams.mStartPosition = 0;
194+ loadDataParams.kStartPosition = 0;
195+ loadDataParams.mStep = CeilDiv<C0_NUM_PER_FRACTAL>(L0K);
196+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0N);
197+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
198+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
199+ loadDataParams.ifTranspose = true;
200+ 
201+ for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {
202+ AscendC::LoadData(
203+ dstTensor.data()[l0BatchIdx * L0N * L0K], srcTensor.data()[l0BatchIdx * L1N * L1K], loadDataParams
204+ );
205+ }
206+ } else {
207+ loadDataParams.mStartPosition = 0;
208+ loadDataParams.kStartPosition = 0;
209+ loadDataParams.mStep = 2;
210+ loadDataParams.kStep = CeilDiv<ELE_NUM_PER_C0>(L0N);
211+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideCol);
212+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
213+ loadDataParams.ifTranspose = true;
214+ for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {
215+ for (uint32_t kIdx = 0; kIdx < L0K / ELE_NUM_PER_C0; kIdx++) {
216+ AscendC::LoadData(
217+ dstTensor.data()[l0BatchIdx * L0N * L0K + kIdx * L0N * ELE_NUM_PER_C0],
218+ srcTensor.data()[l0BatchIdx * L1N * L1K + kIdx * ELE_NUM_PER_FRACTAL * 2], loadDataParams
219+ );
220+ }
221+ }
222+ }
223+ }
224+};
225+ 
226+/// Partial specialization for CopyL1ToL0B, AtlasA5, nZ in and nZ out. (Transpose B)
227+template <class ElementSrc, class ElementDst, class LayoutSrc, class LayoutDst, class CoordSrc, class CoordDst>
228+struct TileCopyTla<
229+ Arch::AtlasA5,
230+ tla::Tensor<AscendC::LocalTensor<ElementSrc>, LayoutSrc, CoordSrc, AscendC::TPosition::A1>,
231+ tla::Tensor<AscendC::LocalTensor<ElementDst>, LayoutDst, CoordDst, AscendC::TPosition::B2>,
232+ std::enable_if_t<tla::detail::isnZ<ElementSrc, LayoutSrc>::value && tla::detail::isnZ<ElementDst, LayoutDst>::value>> {
233+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(ElementSrc);
234+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(ElementSrc);
235+ 
236+ // Mehtods
237+ 
238+ __aicore__ inline
239+ TileCopyTla() {};
240+ 
241+ template <class TensorDst, class TensorSrc>
242+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor)
243+ {
244+ static_assert(
245+ tla::detail::isnZ<typename TensorSrc::Element, typename TensorSrc::Layout>::value
246+ && tla::detail::isnZ<typename TensorDst::Element, typename TensorDst::Layout>::value
247+ && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2,
248+ "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"
249+ );
250+ 
251+ const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());
252+ const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());
253+ const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
254+ const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
255+ auto srcCoord = srcTensor.coord();
256+ 
257+ AscendC::LoadData2DParamsV2 loadDataParams;
258+ loadDataParams.mStartPosition = CeilDiv<C0_NUM_PER_FRACTAL>(tla::get<1>(srcCoord));
259+ loadDataParams.kStartPosition = CeilDiv<ELE_NUM_PER_C0>(tla::get<0>(srcCoord));
260+ loadDataParams.mStep = dstOuterShapeCol;
261+ loadDataParams.kStep = dstOuterShapeRow;
262+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
263+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
264+ loadDataParams.ifTranspose = false;
265+ 
266+ auto dstOffset = dstTensor.layout()(dstTensor.coord());
267+ AscendC::LoadData(dstTensor.data()[dstOffset], srcTensor.data(), loadDataParams);
268+ }
269+ 
270+ template <class TensorDst, class TensorSrc>
271+ __aicore__ inline void operator()(TensorDst const &dstTensor, TensorSrc const &srcTensor, uint32_t l0Batch)
272+ {
273+ static_assert(
274+ tla::detail::isnZ<typename TensorSrc::Element, typename TensorSrc::Layout>::value
275+ && tla::detail::isnZ<typename TensorDst::Element, typename TensorDst::Layout>::value
276+ && TensorSrc::position == AscendC::TPosition::A1 && TensorDst::position == AscendC::TPosition::B2,
277+ "The input parameters do not match. TensorSrc must be L1 and nZ, while TensorDst must be L0B and nZ"
278+ );
279+ 
280+ const uint32_t dstOuterShapeRow = tla::get<0, 1>(dstTensor.shape());
281+ const uint32_t dstOuterShapeCol = tla::get<1, 1>(dstTensor.shape());
282+ const uint32_t srcOuterStrideRow = tla::get<0, 1>(srcTensor.stride());
283+ const uint32_t dstOuterStrideRow = tla::get<0, 1>(dstTensor.stride());
284+ 
285+ AscendC::LoadData2DParamsV2 loadDataParams;
286+ loadDataParams.mStartPosition = 0;
287+ loadDataParams.kStartPosition = 0;
288+ loadDataParams.mStep = dstOuterShapeCol;
289+ loadDataParams.kStep = dstOuterShapeRow * l0Batch;
290+ loadDataParams.srcStride = CeilDiv<ELE_NUM_PER_FRACTAL>(srcOuterStrideRow);
291+ loadDataParams.dstStride = CeilDiv<ELE_NUM_PER_FRACTAL>(dstOuterStrideRow);
292+ loadDataParams.ifTranspose = false;
293+ 
294+ AscendC::LoadData(dstTensor.data(), srcTensor.data(), loadDataParams);
295+ }
296+};
297+ 
298+/////////////////////////////////////////////////////////////////////////////////////////////////////////////
299+ 
300+} // namespace NpuArch::Gemm::Tile
301+ 
302+#endif // GEMM_TILE_COPY_L1_TO_L0B_A5_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/copy_ub_to_gm.hpp+21-0
@@ -0,0 +1,21 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_COPY_UB_TO_GM_HPP
12+#define GEMM_TILE_COPY_UB_TO_GM_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/arch/arch.hpp"
16+#include "../../../attn_infra/gemm/tile_common/tile_copy_tla.hpp"
17+namespace NpuArch::Gemm::Tile {
18+ 
19+} // NpuArch::Gemm::Tile
20+ 
21+#endif // GEMM_TILE_COPY_UB_TO_GM_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/tile_copy.hpp+203-0
@@ -0,0 +1,203 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_TILE_COPY_HPP
12+#define GEMM_TILE_TILE_COPY_HPP
13+ 
14+#include <type_traits>
15+#include "../../../attn_infra/base_defs.hpp"
16+#include "../../../attn_infra/detail/tag_to_layout.hpp"
17+#include "../../../tla/tensor.hpp"
18+#if (__CCE_AICORE__ == 310)
19+#include "../../../attn_infra/gemm/tile_common/copy_gm_to_l1_a5.hpp"
20+#include "../../../attn_infra/gemm/tile_common/copy_l0c_to_ub_a5.hpp"
21+#include "../../../attn_infra/gemm/tile_common/copy_l1_to_l0a_a5.hpp"
22+#include "../../../attn_infra/gemm/tile_common/copy_l1_to_l0b_a5.hpp"
23+#endif
24+#if (__CCE_AICORE__ == 220)
25+#include "../../../attn_infra/gemm/tile_common/copy_gm_to_l1_a2.hpp"
26+#include "../../../attn_infra/gemm/tile_common/copy_l0c_to_gm_a2.hpp"
27+#include "../../../attn_infra/gemm/tile_common/copy_l1_to_l0a_a2.hpp"
28+#include "../../../attn_infra/gemm/tile_common/copy_l1_to_l0b_a2.hpp"
29+#include "../../../attn_infra/gemm/tile_common/copy_l1_to_bt.hpp"
30+#include "../../../attn_infra/gemm/tile_common/copy_gm_to_ub.hpp"
31+#include "../../../attn_infra/gemm/tile_common/copy_ub_to_gm.hpp"
32+#endif
33+#include "../../../attn_infra/gemm/helper.hpp"
34+ 
35+ 
36+namespace NpuArch::Gemm::Tile {
37+ 
38+#if (__CCE_AICORE__ == 220)
39+template <
40+ /// Tag indicating architecture
41+ class ArchTag,
42+ /// GemmType for A matrix operand
43+ class AType,
44+ /// GemmType type for B matrix operand
45+ class BType,
46+ /// GemmType type for C matrix operand
47+ class CType,
48+ /// GemmType type for Bias operand
49+ class BiasType = void
50+>
51+struct TileCopy {
52+ using ElementA = typename AType::Element;
53+ using ElementB = typename BType::Element;
54+ using ElementAccumulator =
55+ typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
56+ 
57+ using CopyGmToL1A = Gemm::Tile::CopyGmToL1<ArchTag, AType>;
58+ using CopyGmToL1B = Gemm::Tile::CopyGmToL1<ArchTag, BType>;
59+ using CopyL1ToL0A = Gemm::Tile::CopyL1ToL0A<
60+ ArchTag, typename helper::L1ATypeSelector<AType>::L1AType>;
61+ using CopyL1ToL0B = Gemm::Tile::CopyL1ToL0B<
62+ ArchTag, typename helper::L1BTypeSelector<BType>::L1BType>;
63+ using CopyL0CToGm = Gemm::Tile::CopyL0CToGm<ArchTag, ElementAccumulator, CType>;
64+ using BiasTypeSelector = helper::L1BiasTypeSelector<BiasType, ElementAccumulator>;
65+ using CopyGmToL1Bias = std::conditional_t<std::is_same_v<BiasType, void>,
66+ void,
67+ Gemm::Tile::CopyGmToL1<ArchTag,
68+ typename BiasTypeSelector::GMBiasType,
69+ typename BiasTypeSelector::L1BiasType>>;
70+};
71+#endif
72+ 
73+template <
74+ /// Tag indicating architecture
75+ class ArchTag,
76+ class ElementA_,
77+ class LayoutTagA_,
78+ class ElementB_,
79+ class LayoutTagB_,
80+ class ElementC_,
81+ class LayoutTagC_,
82+ class ElementBias = void,
83+ bool ReluEnable_ = false,
84+ ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT,
85+ class L0CCopyMode = CopyToGM
86+>
87+struct PackedTileCopyTla {
88+ using ElementA = ElementA_;
89+ using ElementB = ElementB_;
90+ using LayoutTagA = LayoutTagA_;
91+ using LayoutTagB = LayoutTagB_;
92+ using LayoutTagC = LayoutTagC_;
93+ using ElementAccumulator =
94+ typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
95+ static constexpr bool ReluEnable = ReluEnable_;
96+ 
97+ static constexpr bool HAS_BIAS = !std::is_void_v<ElementBias>;
98+ 
99+ using LayoutTagL1A = typename helper::L1ATypeSelector<Gemm::GemmType<ElementA, LayoutTagA>>::L1AType::Layout;
100+ using LayoutTagL1B = typename helper::L1BTypeSelector<Gemm::GemmType<ElementB, LayoutTagB>>::L1BType::Layout;
101+ using LayoutTagL0A = typename helper::L0ALayoutSelector<ArchTag>::Layout;
102+ using LayoutTagL0B = layout::nZ;
103+ 
104+ using LayoutA = detail::TagToLayout_t<ElementA, LayoutTagA>;
105+ using LayoutB = detail::TagToLayout_t<ElementB, LayoutTagB>;
106+ using LayoutC = detail::TagToLayout_t<ElementC_, LayoutTagC>;
107+ 
108+ using LayoutL1A = detail::TagToLayout_t<ElementA, LayoutTagL1A>;
109+ using LayoutL1B = detail::TagToLayout_t<ElementB, LayoutTagL1B>;
110+ using LayoutL0A = detail::TagToLayout_t<ElementA, LayoutTagL0A>;
111+ using LayoutL0B = detail::TagToLayout_t<ElementB, LayoutTagL0B>;
112+ using LayoutL0C = typename detail::LayoutL0C;
113+ 
114+ using TensorL1AVectorLayout =
115+ tla::Tensor<AscendC::LocalTensor<ElementA>, LayoutL1A, tla::Coord<tla::_0>, AscendC::TPosition::A1>;
116+ using TensorL1ALayout =
117+ tla::Tensor<AscendC::LocalTensor<ElementA>, LayoutL1A, tla::Coord<tla::_0, tla::_0>, AscendC::TPosition::A1>;
118+
119+ using TensorL1A = std::conditional_t<tla::detail::isVector<LayoutTagA>::value, TensorL1AVectorLayout, TensorL1ALayout>;
120+ using TensorL1B =
121+ tla::Tensor<AscendC::LocalTensor<ElementB>, LayoutL1B, tla::Coord<tla::_0, tla::_0>, AscendC::TPosition::A1>;
122+ using TensorL0A =
123+ tla::Tensor<AscendC::LocalTensor<ElementA>, LayoutL0A, tla::Coord<tla::_0, tla::_0>, AscendC::TPosition::A2>;
124+ using TensorL0B =
125+ tla::Tensor<AscendC::LocalTensor<ElementB>, LayoutL0B, tla::Coord<tla::_0, tla::_0>, AscendC::TPosition::B2>;
126+ using TensorL0C = tla::Tensor<AscendC::LocalTensor<ElementAccumulator>, LayoutL0C, tla::Coord<tla::_0, tla::_0>,
127+ AscendC::TPosition::CO1>;
128+ using TensorL1Bias = std::conditional_t<
129+ HAS_BIAS,
130+ tla::Tensor<AscendC::LocalTensor<ElementBias>, detail::TagToLayout_t<ElementBias, layout::VectorLayout>,
131+ tla::Coord<tla::_0>, AscendC::TPosition::A1>,
132+ EmptyClass>;
133+ using TensorL0Bias = tla::Tensor<
134+ AscendC::LocalTensor<ElementAccumulator>,
135+ detail::TagToLayout_t<ElementAccumulator, layout::VectorLayout>,
136+ tla::Coord<tla::_0>,
137+ AscendC::TPosition::C2>;
138+ 
139+ using L1AAlignHelper = Gemm::helper::L1AlignHelper<ElementA, LayoutTagA>;
140+ using L1BAlignHelper = Gemm::helper::L1AlignHelper<ElementB, LayoutTagB>;
141+ 
142+ template <class TensorA>
143+ using CopyGmToL1A = Gemm::Tile::TileCopyTla<ArchTag, TensorA, TensorL1A>;
144+ 
145+ template <class TensorB>
146+ using CopyGmToL1B = Gemm::Tile::TileCopyTla<ArchTag, TensorB, TensorL1B>;
147+ 
148+ template <class TensorBias>
149+ using CopyGmToL1Bias = std::conditional_t<
150+ HAS_BIAS,
151+ Gemm::Tile::TileCopyTla<ArchTag, TensorBias, TensorL1Bias>,
152+ EmptyClass>;
153+ 
154+ using CopyL1ToL0A = Gemm::Tile::TileCopyTla<ArchTag, TensorL1A, TensorL0A>;
155+ using CopyL1ToL0B = Gemm::Tile::TileCopyTla<ArchTag, TensorL1B, TensorL0B>;
156+ using CopyL1ToBT = std::conditional_t<
157+ HAS_BIAS,
158+ Gemm::Tile::TileCopyTla<ArchTag, TensorL1Bias, TensorL0Bias>,
159+ EmptyClass>;
160+ 
161+ template <class TensorC>
162+ using CopyL0CToDst = Gemm::Tile::CopyL0CToGmTla<ArchTag, TensorL0C, TensorC, DEQUANT_GRANULARITY, ReluEnable>;
163+};
164+ 
165+#if (__CCE_AICORE__ == 310)
166+template <
167+ /// Tag indicating architecture
168+ class ArchTag,
169+ class ElementA_,
170+ class LayoutTagA,
171+ class ElementB_,
172+ class LayoutTagB,
173+ class ElementC_,
174+ class LayoutTagC,
175+ class ElementBias = void,
176+ CopyL0CToUBMode CopyMode_ = CopyL0CToUBMode::NO_SPLIT,
177+ bool ReluEnable = false,
178+ ScaleGranularity DEQUANT_GRANULARITY = ScaleGranularity::NO_QUANT
179+>
180+struct PackedTileCopyTlaToUB : public PackedTileCopyTla<ArchTag, ElementA_, LayoutTagA, ElementB_, LayoutTagB,
181+ ElementC_, LayoutTagC, ElementBias> {
182+ static constexpr CopyL0CToUBMode CopyMode = CopyMode_;
183+ // 重写 CopyL0CToDst
184+ using TensorL0C = typename PackedTileCopyTla<
185+ ArchTag,
186+ ElementA_,
187+ LayoutTagA,
188+ ElementB_,
189+ LayoutTagB,
190+ ElementC_,
191+ LayoutTagC,
192+ ElementBias>::TensorL0C;
193+ 
194+ template <class TensorC>
195+ using CopyL0CToDst =
196+ Gemm::Tile::CopyL0CToUBTla<ArchTag, TensorL0C, TensorC, CopyMode, DEQUANT_GRANULARITY, ReluEnable>;
197+};
198+#endif
199+ 
200+//////////////////////////////
201+} // namespace NpuArch::Gemm::Tile
202+ 
203+#endif // GEMM_TILE_TILE_COPY_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/tile_copy_tla.hpp+28-0
@@ -0,0 +1,28 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_TILE_COPY_TLA_HPP
12+#define GEMM_TILE_TILE_COPY_TLA_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+ 
16+namespace NpuArch::Gemm::Tile {
17+template <
18+ class ArchTag,
19+ class TensorSrc,
20+ class TensorDst,
21+ class Enable = void
22+>
23+struct TileCopyTla {
24+ static_assert(DEPENDENT_FALSE<ArchTag>, "Unsupported TileCopyTla, can not find the specialization.");
25+};
26+} // namespace NpuArch::Gemm::Tile
27+ 
28+#endif // GEMM_TILE_TILE_COPY_TLA_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm/tile_common/tile_mmad.hpp+227-0
@@ -0,0 +1,227 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef GEMM_TILE_TILE_MMAD_HPP
12+#define GEMM_TILE_TILE_MMAD_HPP
13+ 
14+#include "../../../attn_infra/base_defs.hpp"
15+#include "../../../attn_infra/gemm/helper.hpp"
16+#include "../../../tla/tensor.hpp"
17+namespace NpuArch::Gemm::Tile {
18+ 
19+///////////////////////////////////////////////////////////
20+ 
21+template <
22+ /// Tag indicating architecture
23+ class ArchTag_,
24+ /// GemmType for A matrix operand
25+ class AType_,
26+ /// GemmType type for B matrix operand
27+ class BType_,
28+ /// GemmType type for Bias operand
29+ class BiasType_
30+>
31+struct TileMmad {
32+ using ElementA = typename AType_::Element;
33+ using ElementB = typename BType_::Element;
34+ using ElementAccumulator =
35+ typename Gemm::helper::ElementAccumulatorSelector<ElementA, ElementB>::ElementAccumulator;
36+ 
37+ // Methods
38+ 
39+ __aicore__ inline
40+ TileMmad() {}
41+ 
42+ __aicore__ inline
43+ void operator()(AscendC::LocalTensor<ElementAccumulator> const &l0CTensor,
44+ AscendC::LocalTensor<ElementA> const &l0ATensor,
45+ AscendC::LocalTensor<ElementB> const &l0BTensor,
46+ uint32_t m, uint32_t n, uint32_t k,
47+ bool initC = true, uint8_t unitFlag = 0)
48+ {
49+ AscendC::MmadParams mmadParams;
50+ mmadParams.m = m;
51+ mmadParams.n = n;
52+ mmadParams.k = k;
53+ mmadParams.unitFlag = unitFlag;
54+ mmadParams.cmatrixInitVal = initC;
55+ if constexpr (std::is_same_v<ElementA, float> && std::is_same_v<typename AType_::Layout, layout::ColumnMajor>) {
56+ mmadParams.kDirectionAlign = true;
57+ }
58+ 
59+ AscendC::Mmad(l0CTensor,
60+ l0ATensor,
61+ l0BTensor,
62+ mmadParams);
63+ 
64+ const uint32_t PIPE_M_BARRIER_THRESHOLD = 10;
65+ if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) {
66+ AscendC::PipeBarrier<PIPE_M>();
67+ }
68+ }
69+ 
70+ __aicore__ inline
71+ void operator()(AscendC::LocalTensor<ElementAccumulator> const &l0CTensor,
72+ AscendC::LocalTensor<ElementA> const &l0ATensor,
73+ AscendC::LocalTensor<ElementB> const &l0BTensor,
74+ AscendC::LocalTensor<ElementAccumulator> const &l0BiasTensor,
75+ uint32_t m, uint32_t n, uint32_t k,
76+ bool initC = true, uint8_t unitFlag = 0)
77+ {
78+ AscendC::MmadParams mmadParams;
79+ mmadParams.m = m;
80+ mmadParams.n = n;
81+ mmadParams.k = k;
82+ mmadParams.unitFlag = unitFlag;
83+ mmadParams.cmatrixInitVal = false;
84+ if constexpr (std::is_same_v<ElementA, float> && std::is_same_v<typename AType_::Layout, layout::ColumnMajor>) {
85+ mmadParams.kDirectionAlign = true;
86+ }
87+ 
88+ AscendC::Mmad(l0CTensor,
89+ l0ATensor,
90+ l0BTensor,
91+ l0BiasTensor,
92+ mmadParams);
93+ 
94+ const uint32_t PIPE_M_BARRIER_THRESHOLD = 10;
95+ if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) {
96+ AscendC::PipeBarrier<PIPE_M>();
97+ }
98+ }
99+};
100+ 
101+template <
102+ /// Tag indicating architecture
103+ class ArchTag,
104+ /// Element for A matrix operand
105+ class ElementA,
106+ /// LayoutTag for A matrix operand in L1
107+ class LayoutTagL1A
108+>
109+struct TileMmadTla {
110+ // Methods
111+ 
112+ __aicore__ inline
113+ TileMmadTla() {}
114+ 
115+ template <class TensorC, class TensorA, class TensorB>
116+ __aicore__ inline
117+ void operator()(TensorC const &l0CTensor,
118+ TensorA const &l0ATensor,
119+ TensorB const &l0BTensor,
120+ uint32_t m, uint32_t n, uint32_t k,
121+ bool initC = true, uint8_t unitFlag = 0)
122+ {
123+ AscendC::MmadParams mmadParams;
124+ mmadParams.m = m;
125+ mmadParams.n = n;
126+ mmadParams.k = k;
127+ mmadParams.unitFlag = unitFlag;
128+ mmadParams.cmatrixInitVal = initC;
129+#if (__CCE_AICORE__ == 310)
130+ if constexpr(std::is_same_v<LayoutTagL1A, layout::VectorLayout>) {
131+ mmadParams.disableGemv = false;
132+ } else {
133+ mmadParams.disableGemv = true;
134+ }
135+#endif
136+#if (__CCE_AICORE__ == 220)
137+ if constexpr (std::is_same_v<ElementA, float> && std::is_same_v<LayoutTagL1A, layout::nZ>) {
138+ mmadParams.kDirectionAlign = true;
139+ }
140+#endif
141+ 
142+ AscendC::Mmad(l0CTensor.data(),
143+ l0ATensor.data(),
144+ l0BTensor.data(),
145+ mmadParams);
146+ 
147+ const uint32_t PIPE_M_BARRIER_THRESHOLD = 10;
148+ if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) {
149+ AscendC::PipeBarrier<PIPE_M>();
150+ }
151+ }
152+ 
153+ template <class TensorC, class TensorA, class TensorB, class TensorBias>
154+ __aicore__ inline
155+ void operator()(TensorC const &l0CTensor,
156+ TensorA const &l0ATensor,
157+ TensorB const &l0BTensor,
158+ TensorBias const &l0BiasTensor,
159+ uint32_t m, uint32_t n, uint32_t k,
160+ bool initC = true, uint8_t unitFlag = 0)
161+ {
162+ AscendC::MmadParams mmadParams;
163+ mmadParams.m = m;
164+ mmadParams.n = n;
165+ mmadParams.k = k;
166+ mmadParams.unitFlag = unitFlag;
167+ mmadParams.cmatrixInitVal = false;
168+#if (__CCE_AICORE__ == 310)
169+ mmadParams.disableGemv = true;
170+#endif
171+#if (__CCE_AICORE__ == 220)
172+ if constexpr (std::is_same_v<ElementA, float> && std::is_same_v<LayoutTagL1A, layout::nZ>) {
173+ mmadParams.kDirectionAlign = true;
174+ }
175+#endif
176+ 
177+ AscendC::Mmad(l0CTensor.data(),
178+ l0ATensor.data(),
179+ l0BTensor.data(),
180+ l0BiasTensor.data(),
181+ mmadParams);
182+ 
183+ const uint32_t PIPE_M_BARRIER_THRESHOLD = 10;
184+ if ((m / C0_NUM_PER_FRACTAL) * (n / C0_NUM_PER_FRACTAL) < PIPE_M_BARRIER_THRESHOLD) {
185+ AscendC::PipeBarrier<PIPE_M>();
186+ }
187+ }
188+ 
189+ template <class TensorC, class TensorA, class TensorB>
190+ __aicore__ inline
191+ void operator()(TensorC const &l0CTensor,
192+ TensorA const &l0ATensor,
193+ TensorB const &l0BTensor,
194+ uint32_t m, uint32_t n, uint32_t k,
195+ uint32_t l0Batch)
196+ {
197+ const uint32_t L0AM = tla::get<0, 0>(l0ATensor.shape()) * tla::get<0, 1>(l0ATensor.shape());
198+ const uint32_t L0AK = tla::get<1, 0>(l0ATensor.shape()) * tla::get<1, 1>(l0ATensor.shape());
199+ const uint32_t L0BK = tla::get<0, 0>(l0BTensor.shape()) * tla::get<0, 1>(l0BTensor.shape());
200+ const uint32_t L0BN = tla::get<1, 0>(l0BTensor.shape()) * tla::get<1, 1>(l0BTensor.shape());
201+ const uint32_t L0CM = tla::get<0, 0>(l0CTensor.shape()) * tla::get<0, 1>(l0CTensor.shape());
202+ const uint32_t L0CN = tla::get<1, 0>(l0CTensor.shape()) * tla::get<1, 1>(l0CTensor.shape());
203+ 
204+ AscendC::MmadParams mmadParams;
205+ mmadParams.m = m;
206+ mmadParams.n = n;
207+ mmadParams.k = k;
208+ mmadParams.unitFlag = 0;
209+ mmadParams.cmatrixInitVal = true;
210+#if (__CCE_AICORE__ == 310)
211+ mmadParams.disableGemv = true;
212+#endif
213+ 
214+ for (uint32_t l0BatchIdx = 0; l0BatchIdx < l0Batch; l0BatchIdx++) {
215+ AscendC::Mmad(l0CTensor.data()[l0BatchIdx * L0CM * L0CN],
216+ l0ATensor.data()[l0BatchIdx * L0AM * L0AK],
217+ l0BTensor.data()[l0BatchIdx * L0BK * L0BN],
218+ mmadParams);
219+ }
220+ }
221+};
222+ 
223+/////////////////////////////////////////////////////////////////////////////////////////////////////////////
224+ 
225+} // namespace NpuArch::Gemm::Tile
226+ 
227+#endif // GEMM_TILE_TILE_MMAD_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/gemm_coord.hpp+163-0
@@ -0,0 +1,163 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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 gemm_coord.hpp
13+ * \brief
14+ */
15+ 
16+#ifndef GEMM_COORD_HPP
17+#define GEMM_COORD_HPP
18+ 
19+#include "../attn_infra/coord.hpp"
20+ 
21+namespace NpuArch {
22+ 
23+/// Shape of a matrix multiply-add operation
24+template <
25+ /// Rows of matrix product
26+ uint32_t M_ = 1,
27+ /// Columns of matrix product
28+ uint32_t N_ = 1,
29+ /// Inner dimension of matrix product
30+ uint32_t K_ = 1
31+>
32+struct GemmShape {
33+ static constexpr uint32_t M = M_;
34+ static constexpr uint32_t N = N_;
35+ static constexpr uint32_t K = K_;
36+ 
37+ static constexpr int64_t MN = M * N;
38+ static constexpr int64_t MK = M * K;
39+ static constexpr int64_t KN = N * K;
40+ static constexpr int64_t MNK = M * N * K;
41+ 
42+ static constexpr int64_t COUNT = MNK;
43+ 
44+ /// Returns a Coord object
45+ HOST_DEVICE
46+ static Coord<3> ToCoord()
47+ {
48+ return MakeCoord(M, N, K);
49+ }
50+ 
51+ HOST_DEVICE
52+ static Coord<2> ToCoordMN()
53+ {
54+ return MakeCoord(M, N);
55+ }
56+ 
57+ HOST_DEVICE
58+ static Coord<2> ToCoordMK()
59+ {
60+ return MakeCoord(M, K);
61+ }
62+ 
63+ HOST_DEVICE
64+ static Coord<2> ToCoordKN()
65+ {
66+ return MakeCoord(K, N);
67+ }
68+};
69+ 
70+/// GemmCoord is a structure derived from Coord<3> that specifies a location within the
71+/// coordinate space of a Gemm problem.
72+struct GemmCoord : public Coord<3, uint32_t> {
73+ /// Integer-valued index
74+ using Index = uint32_t;
75+ 
76+ /// Base type is a Coord of rank=3
77+ using Base = Coord<3, Index>;
78+ 
79+ /// Gemm M dimension - rows of the output C matrix
80+ static constexpr int M_INDEX = 0;
81+ 
82+ /// Gemm N dimension - columns of the output C matrix
83+ static constexpr int N_INDEX = 1;
84+ 
85+ /// Gemm K dimension - inner dimension of the Gemm problem
86+ static constexpr int K_INDEX = 2;
87+ 
88+ /// Default ctor
89+ HOST_DEVICE
90+ GemmCoord() {}
91+ 
92+ /// Constructs from Coord<3> and a batch
93+ HOST_DEVICE
94+ GemmCoord(Coord<3, Index> const &coord) : Base(coord) {}
95+ 
96+ /// Helper to construct from a K, N, M, batch variables
97+ HOST_DEVICE
98+ GemmCoord(Index m, Index n, Index k) : Base(MakeCoord(m, n, k)) {}
99+ 
100+ /// Returns the Gemm M coordinate
101+ HOST_DEVICE
102+ Index const &m() const
103+ {
104+ return this->At(M_INDEX);
105+ }
106+ 
107+ /// Returns reference to the Gemm M coordinate
108+ HOST_DEVICE
109+ Index &m()
110+ {
111+ return this->At(M_INDEX);
112+ }
113+ 
114+ /// Returns the Gemm N coordinate
115+ HOST_DEVICE
116+ Index const &n() const
117+ {
118+ return this->At(N_INDEX);
119+ }
120+ 
121+ /// Returns reference to the Gemm N coordinate
122+ HOST_DEVICE
123+ Index &n()
124+ {
125+ return this->At(N_INDEX);
126+ }
127+ 
128+ /// Returns the Gemm K coordinate
129+ HOST_DEVICE
130+ Index const &k() const
131+ {
132+ return this->At(K_INDEX);
133+ }
134+ 
135+ /// Returns reference to the Gemm K coordinate
136+ HOST_DEVICE
137+ Index &k()
138+ {
139+ return this->At(K_INDEX);
140+ }
141+ 
142+ HOST_DEVICE
143+ auto GetCoordMN() const
144+ {
145+ return this->GetCoordByAxis<M_INDEX, N_INDEX>();
146+ }
147+ 
148+ HOST_DEVICE
149+ auto GetCoordMK() const
150+ {
151+ return this->GetCoordByAxis<M_INDEX, K_INDEX>();
152+ }
153+ 
154+ HOST_DEVICE
155+ auto GetCoordKN() const
156+ {
157+ return this->GetCoordByAxis<K_INDEX, N_INDEX>();
158+ }
159+};
160+ 
161+} // namespace NpuArch
162+ 
163+#endif // GEMM_COORD_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/layout/layout.hpp+18-0
@@ -0,0 +1,18 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef LAYOUT_LAYOUT_HPP
12+#define LAYOUT_LAYOUT_HPP
13+ 
14+#include "../../attn_infra/base_defs.hpp"
15+#include "../../attn_infra/layout/matrix.hpp"
16+#include "../../attn_infra/layout/vector.hpp"
17+ 
18+#endif // LAYOUT_LAYOUT_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/layout/matrix.hpp+1247-0
@@ -0,0 +1,1247 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef LAYOUT_MATRIX_HPP
12+#define LAYOUT_MATRIX_HPP
13+ 
14+#include "../../attn_infra/base_defs.hpp"
15+#include "../../attn_infra/coord.hpp"
16+#include "../../attn_infra/detail/alignment.hpp"
17+#include "../../attn_infra/matrix_coord.hpp"
18+ 
19+namespace NpuArch::layout
20+{
21+ 
22+/// Mapping function for row-major matrices
23+struct RowMajor {
24+public:
25+ /// Logical rank of tensor
26+ static constexpr int RANK = 2;
27+ 
28+ /// Index type used for coordinates
29+ using Index = uint32_t;
30+ 
31+ /// Long index type used for offsets
32+ using LongIndex = int64_t;
33+ 
34+ /// Logical coordinate
35+ using Shape = Coord<RANK, Index>;
36+ 
37+ /// Stride vector
38+ using Stride = Coord<RANK, LongIndex>;
39+ 
40+public:
41+ /// Constructor
42+ HOST_DEVICE
43+ RowMajor(Index rows = 0, Index cols = 0)
44+ : shape_(MakeCoord(rows, cols)), stride_(MakeCoord(LongIndex(cols), LongIndex(1))) {}
45+ 
46+ /// Constructor
47+ HOST_DEVICE
48+ RowMajor(Index rows, Index cols, LongIndex ldm)
49+ : shape_(MakeCoord(rows, cols)), stride_(MakeCoord(ldm, LongIndex(1))) {}
50+ 
51+ /// Ctor
52+ HOST_DEVICE
53+ RowMajor(Shape shape, Stride stride) : shape_(shape), stride_(stride) {}
54+ 
55+ template <class Element>
56+ HOST_DEVICE
57+ static RowMajor MakeLayout(Index rows, Index cols)
58+ {
59+ return RowMajor(rows, cols);
60+ }
61+ 
62+ template <class Element>
63+ HOST_DEVICE
64+ static RowMajor MakeLayoutInUb(MatrixCoord const &shape)
65+ {
66+ return RowMajor(shape.row(), shape.column(), RoundUp<BYTE_PER_C0 / sizeof(Element)>(shape.column()));
67+ }
68+ 
69+ /// Returns the offset of a coordinate in linear memory.
70+ /// Assumes coordinate has convention (row, column)
71+ HOST_DEVICE
72+ LongIndex GetOffset(MatrixCoord const &coord) const
73+ {
74+ return LongIndex(coord.row()) * stride_[0] + LongIndex(coord.column());
75+ }
76+ 
77+ /// Returns the layout of a tile_common.
78+ HOST_DEVICE
79+ RowMajor GetTileLayout(MatrixCoord const &tileShape) const
80+ {
81+ return RowMajor(tileShape, stride());
82+ }
83+ 
84+ /// Returns the shape of the layout
85+ HOST_DEVICE
86+ Shape shape() const
87+ {
88+ return shape_;
89+ }
90+ 
91+ /// Returns the shape of the layout
92+ HOST_DEVICE
93+ Shape &shape()
94+ {
95+ return shape_;
96+ }
97+ 
98+ /// Returns the shape of the layout
99+ HOST_DEVICE
100+ typename Shape::Index shape(int idx) const
101+ {
102+ return shape_[idx];
103+ }
104+ 
105+ /// Returns the shape of the layout
106+ HOST_DEVICE
107+ typename Shape::Index &shape(int idx)
108+ {
109+ return shape_[idx];
110+ }
111+ 
112+ /// Returns the stride of the layout
113+ HOST_DEVICE
114+ Stride stride() const
115+ {
116+ return stride_;
117+ }
118+ 
119+ /// Returns the stride of the layout
120+ HOST_DEVICE
121+ Stride &stride()
122+ {
123+ return stride_;
124+ }
125+ 
126+ /// Returns the stride of the layout
127+ HOST_DEVICE
128+ typename Stride::Index stride(int idx) const
129+ {
130+ return stride_[idx];
131+ }
132+ 
133+ /// Returns the stride of the layout
134+ HOST_DEVICE
135+ typename Stride::Index &stride(int idx)
136+ {
137+ return stride_[idx];
138+ }
139+ 
140+ /// Returns the length of the layout
141+ HOST_DEVICE
142+ LongIndex Capacity() const
143+ {
144+ return static_cast<LongIndex>(shape_[0]) * stride_[0];
145+ }
146+ 
147+protected:
148+ //
149+ // Data members
150+ //
151+ 
152+ /// Shape data member
153+ Shape shape_;
154+ 
155+ /// Stride data member
156+ Stride stride_;
157+};
158+ 
159+/// Mapping function for col-major matrices
160+struct ColumnMajor {
161+public:
162+ /// Logical rank of tensor
163+ static constexpr int RANK = 2;
164+ 
165+ /// Index type used for coordinates
166+ using Index = uint32_t;
167+ 
168+ /// Long index type used for offsets
169+ using LongIndex = int64_t;
170+ 
171+ /// Logical coordinate
172+ using Shape = Coord<RANK, Index>;
173+ 
174+ /// Stride vector
175+ using Stride = Coord<RANK, LongIndex>;
176+ 
177+public:
178+ // Methods
179+ 
180+ /// Constructor
181+ HOST_DEVICE
182+ ColumnMajor(Index rows = 0, Index cols = 0)
183+ : shape_(MakeCoord(rows, cols)), stride_(MakeCoord(LongIndex(1), LongIndex(rows))) {}
184+ 
185+ /// Constructor
186+ HOST_DEVICE
187+ ColumnMajor(Index rows, Index cols, LongIndex ldm)
188+ : shape_(MakeCoord(rows, cols)), stride_(MakeCoord(LongIndex(1), ldm)) {}
189+ 
190+ /// Ctor
191+ HOST_DEVICE
192+ ColumnMajor(Shape shape, Stride stride) : shape_(shape), stride_(stride) {}
193+ 
194+ template <class Element>
195+ HOST_DEVICE
196+ static ColumnMajor MakeLayout(Index rows, Index cols)
197+ {
198+ return ColumnMajor(rows, cols);
199+ }
200+ 
201+ /// Returns the offset of a coordinate in linear memory.
202+ /// Assumes coordinate has convention (row, column)
203+ HOST_DEVICE
204+ LongIndex GetOffset(MatrixCoord const &coord) const
205+ {
206+ return LongIndex(coord.row()) + LongIndex(coord.column()) * stride_[1];
207+ }
208+ 
209+ /// Returns the layout of a tile_common.
210+ HOST_DEVICE
211+ ColumnMajor GetTileLayout(MatrixCoord const &tileShape) const
212+ {
213+ return ColumnMajor(tileShape, stride());
214+ }
215+ 
216+ /// Returns the shape of the layout
217+ HOST_DEVICE
218+ Shape shape() const
219+ {
220+ return shape_;
221+ }
222+ 
223+ /// Returns the shape of the layout
224+ HOST_DEVICE
225+ Shape &shape()
226+ {
227+ return shape_;
228+ }
229+ 
230+ /// Returns the shape of the layout
231+ HOST_DEVICE
232+ typename Shape::Index shape(int idx) const
233+ {
234+ return shape_[idx];
235+ }
236+ 
237+ /// Returns the shape of the layout
238+ HOST_DEVICE
239+ typename Shape::Index &shape(int idx)
240+ {
241+ return shape_[idx];
242+ }
243+ 
244+ /// Returns the stride of the layout
245+ HOST_DEVICE
246+ Stride stride() const
247+ {
248+ return stride_;
249+ }
250+ 
251+ /// Returns the stride of the layout
252+ HOST_DEVICE
253+ Stride &stride()
254+ {
255+ return stride_;
256+ }
257+ 
258+ /// Returns the stride of the layout
259+ HOST_DEVICE
260+ typename Stride::Index stride(int idx) const
261+ {
262+ return stride_[idx];
263+ }
264+ 
265+ /// Returns the stride of the layout
266+ HOST_DEVICE
267+ typename Stride::Index &stride(int idx)
268+ {
269+ return stride_[idx];
270+ }
271+ 
272+ /// Returns the length of the layout
273+ HOST_DEVICE
274+ LongIndex Capacity() const
275+ {
276+ return static_cast<LongIndex>(shape_[1]) * stride_[1];
277+ }
278+ 
279+protected:
280+ //
281+ // Data members
282+ //
283+ 
284+ /// Shape data member
285+ Shape shape_;
286+ 
287+ /// Stride data member
288+ Stride stride_;
289+};
290+ 
291+/// Mapping function for nZ matrices which is col-major inside fractal and row-major between fractal
292+struct nZ {
293+public:
294+ /// Logical rank of tensor
295+ static constexpr int RANK = 4;
296+ 
297+ /// Index type used for coordinates
298+ using Index = uint32_t;
299+ 
300+ /// Long index type used for offsets
301+ using LongIndex = int64_t;
302+ 
303+ /// Logical rank of orgshape
304+ static constexpr int ORG_SHAPE_RANK = 2;
305+ 
306+ /// Logical coordinate
307+ using OrgShape = Coord<ORG_SHAPE_RANK, Index>;
308+ 
309+ /// Logical coordinate
310+ using Shape = Coord<RANK, Index>;
311+ 
312+ /// Stride vector
313+ using Stride = Coord<RANK, LongIndex>;
314+ 
315+public:
316+ // Methods
317+ 
318+ /// Constructor
319+ HOST_DEVICE constexpr
320+ nZ(Index orgRows = 0, /// Number of rows of origin matrices
321+ Index orgCols = 0, /// Number of cols of origin matrices
322+ Index rowsInFractal = 0, /// Number of rows inside the fractal
323+ Index rowsByFractal = 0, /// number of rows by the fractal
324+ Index colsInFractal = 0, /// number of cols inside the fractal
325+ Index colsByFractal = 0, /// number of cols by the fractal
326+ LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal
327+ LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows
328+ LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal
329+ LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols
330+ : orgShape_(MakeCoord(orgRows, orgCols)),
331+ shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)),
332+ stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) {}
333+ 
334+ /// Ctor
335+ HOST_DEVICE constexpr
336+ nZ(OrgShape orgShape, Shape shape, Stride stride) : orgShape_(orgShape), shape_(shape), stride_(stride) {}
337+ 
338+ /// Make the layout of a coordinate (row, column)
339+ template <class Element>
340+ HOST_DEVICE constexpr
341+ static nZ MakeLayout(Index orgRows, Index orgCols)
342+ {
343+ constexpr uint32_t ELE_NUM_PER_C0 = static_cast<uint32_t>(BYTE_PER_C0) / static_cast<uint32_t>(sizeof(Element));
344+ constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
345+ Index rowsRound = RoundUp<ELE_NUM_PER_C0>(orgRows);
346+ Index colsRound = RoundUp<C0_NUM_PER_FRACTAL>(orgCols);
347+ return nZ(orgRows,
348+ orgCols,
349+ ELE_NUM_PER_C0,
350+ rowsRound / ELE_NUM_PER_C0,
351+ C0_NUM_PER_FRACTAL,
352+ colsRound / C0_NUM_PER_FRACTAL,
353+ 1,
354+ colsRound * ELE_NUM_PER_C0,
355+ ELE_NUM_PER_C0,
356+ ELE_NUM_PER_FRACTAL);
357+ }
358+ 
359+ /// Returns the offset of a coordinate in linear memory.
360+ /// Assumes coordinate has convention (row, column)
361+ HOST_DEVICE
362+ LongIndex GetOffset(MatrixCoord const &coord) const
363+ {
364+ return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] +
365+ (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2];
366+ }
367+ 
368+ /// Returns the layout of a tile_common.
369+ HOST_DEVICE
370+ nZ GetTileLayout(MatrixCoord const &tileOriShape) const
371+ {
372+ auto tileShape = MakeCoord(
373+ shape(0), CeilDiv(tileOriShape.row(), shape(0)),
374+ shape(2), CeilDiv(tileOriShape.column(), shape(2))
375+ );
376+ return nZ(tileOriShape, tileShape, stride());
377+ }
378+ 
379+ /// Returns the origin shape of the layout
380+ HOST_DEVICE
381+ typename OrgShape::Index orgShape(int idx) const
382+ {
383+ return orgShape_[idx];
384+ }
385+ 
386+ /// Returns the origin shape of the layout
387+ HOST_DEVICE
388+ typename OrgShape::Index &orgShape(int idx)
389+ {
390+ return orgShape_[idx];
391+ }
392+ 
393+ /// Returns the shape of the layout
394+ HOST_DEVICE
395+ Shape shape() const
396+ {
397+ return shape_;
398+ }
399+ 
400+ /// Returns the shape of the layout
401+ HOST_DEVICE
402+ Shape &shape()
403+ {
404+ return shape_;
405+ }
406+ 
407+ /// Returns the shape of the layout
408+ HOST_DEVICE
409+ typename Shape::Index shape(int idx) const
410+ {
411+ return shape_[idx];
412+ }
413+ 
414+ /// Returns the shape of the layout
415+ HOST_DEVICE
416+ typename Shape::Index &shape(int idx)
417+ {
418+ return shape_[idx];
419+ }
420+ 
421+ /// Returns the stride of the layout
422+ HOST_DEVICE
423+ Stride stride() const
424+ {
425+ return stride_;
426+ }
427+ 
428+ /// Returns the stride of the layout
429+ HOST_DEVICE
430+ Stride &stride()
431+ {
432+ return stride_;
433+ }
434+ 
435+ /// Returns the stride of the layout
436+ HOST_DEVICE
437+ typename Stride::Index stride(int idx) const
438+ {
439+ return stride_[idx];
440+ }
441+ 
442+ /// Returns the stride of the layout
443+ HOST_DEVICE
444+ typename Stride::Index &stride(int idx)
445+ {
446+ return stride_[idx];
447+ }
448+ 
449+ /// Returns the length of the layout
450+ HOST_DEVICE
451+ LongIndex Capacity() const
452+ {
453+ return static_cast<LongIndex>(stride_[1]) * shape_[1];
454+ }
455+ 
456+private:
457+ /// Origin Shape data member
458+ OrgShape orgShape_;
459+ 
460+ /// Shape data member
461+ Shape shape_;
462+ 
463+ /// Stride data member
464+ Stride stride_;
465+};
466+ 
467+/// Mapping function for zN matrices which is row-major inside fractal and col-major between fractal
468+struct zN {
469+public:
470+ /// Logical rank of tensor
471+ static constexpr int RANK = 4;
472+ 
473+ /// Index type used for coordinates
474+ using Index = uint32_t;
475+ 
476+ /// Long index type used for offsets
477+ using LongIndex = int64_t;
478+ 
479+ /// Logical rank of orgshape
480+ static constexpr int ORG_SHAPE_RANK = 2;
481+ 
482+ /// Logical coordinate
483+ using OrgShape = Coord<ORG_SHAPE_RANK, Index>;
484+ 
485+ /// Logical coordinate
486+ using Shape = Coord<RANK, Index>;
487+ 
488+ /// Stride vector
489+ using Stride = Coord<RANK, LongIndex>;
490+ 
491+public:
492+ // Methods
493+ 
494+ /// Constructor
495+ HOST_DEVICE constexpr
496+ zN(Index orgRows = 0, /// Number of rows of origin matrices
497+ Index orgCols = 0, /// Number of cols of origin matrices
498+ Index rowsInFractal = 0, /// Number of rows inside the fractal
499+ Index rowsByFractal = 0, /// number of rows by the fractal
500+ Index colsInFractal = 0, /// number of cols inside the fractal
501+ Index colsByFractal = 0, /// number of cols by the fractal
502+ LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal
503+ LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows
504+ LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal
505+ LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols
506+ : orgShape_(MakeCoord(orgRows, orgCols)),
507+ shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)),
508+ stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) {}
509+ 
510+ /// Ctor
511+ HOST_DEVICE constexpr
512+ zN(OrgShape orgShape, Shape shape, Stride stride) : orgShape_(orgShape), shape_(shape), stride_(stride) {}
513+ 
514+ /// Make the layout of a coordinate (row, column)
515+ template <class Element>
516+ HOST_DEVICE constexpr
517+ static zN MakeLayout(Index orgRows, Index orgCols)
518+ {
519+ constexpr uint32_t ELE_NUM_PER_C0 = static_cast<uint32_t>(BYTE_PER_C0) / static_cast<uint32_t>(sizeof(Element));
520+ constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
521+ Index rowsRound = RoundUp<C0_NUM_PER_FRACTAL>(orgRows);
522+ Index colsRound = RoundUp<ELE_NUM_PER_C0>(orgCols);
523+ return zN(orgRows,
524+ orgCols,
525+ C0_NUM_PER_FRACTAL,
526+ rowsRound / C0_NUM_PER_FRACTAL,
527+ ELE_NUM_PER_C0,
528+ colsRound / ELE_NUM_PER_C0,
529+ ELE_NUM_PER_C0,
530+ ELE_NUM_PER_FRACTAL,
531+ 1,
532+ rowsRound * ELE_NUM_PER_C0);
533+ }
534+ 
535+ HOST_DEVICE
536+ static zN MakeLayoutInL0C(MatrixCoord const &shape)
537+ {
538+ return zN(shape.row(),
539+ shape.column(),
540+ C0_NUM_PER_FRACTAL,
541+ CeilDiv<C0_NUM_PER_FRACTAL>(shape.row()),
542+ C0_NUM_PER_FRACTAL,
543+ CeilDiv<C0_NUM_PER_FRACTAL>(shape.column()),
544+ C0_NUM_PER_FRACTAL,
545+ C0_NUM_PER_FRACTAL * C0_NUM_PER_FRACTAL,
546+ 1,
547+ RoundUp<C0_NUM_PER_FRACTAL>(shape.row()) * C0_NUM_PER_FRACTAL);
548+ }
549+ 
550+ /// Returns the offset of a coordinate in linear memory.
551+ /// Assumes coordinate has convention (row, column)
552+ HOST_DEVICE
553+ LongIndex GetOffset(MatrixCoord const &coord) const
554+ {
555+ return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3] +
556+ (LongIndex(coord.row()) % shape_[0]) * stride_[0] + (LongIndex(coord.column()) % shape_[2]) * stride_[2];
557+ }
558+ 
559+ /// Returns the layout of a tile_common.
560+ HOST_DEVICE
561+ zN GetTileLayout(MatrixCoord const &tileOriShape) const
562+ {
563+ auto tileShape = MakeCoord(
564+ shape(0), CeilDiv(tileOriShape.row(), shape(0)),
565+ shape(2), CeilDiv(tileOriShape.column(), shape(2))
566+ );
567+ return zN(tileOriShape, tileShape, stride());
568+ }
569+ 
570+ /// Returns the origin shape of the layout
571+ HOST_DEVICE
572+ typename OrgShape::Index orgShape(int idx) const
573+ {
574+ return orgShape_[idx];
575+ }
576+ 
577+ /// Returns the origin shape of the layout
578+ HOST_DEVICE
579+ typename OrgShape::Index &orgShape(int idx)
580+ {
581+ return orgShape_[idx];
582+ }
583+ 
584+ /// Returns the shape of the layout
585+ HOST_DEVICE
586+ Shape shape() const
587+ {
588+ return shape_;
589+ }
590+ 
591+ /// Returns the shape of the layout
592+ HOST_DEVICE
593+ Shape &shape()
594+ {
595+ return shape_;
596+ }
597+ 
598+ /// Returns the shape of the layout
599+ HOST_DEVICE
600+ typename Shape::Index shape(int idx) const
601+ {
602+ return shape_[idx];
603+ }
604+ 
605+ /// Returns the shape of the layout
606+ HOST_DEVICE
607+ typename Shape::Index &shape(int idx)
608+ {
609+ return shape_[idx];
610+ }
611+ 
612+ /// Returns the stride of the layout
613+ HOST_DEVICE
614+ Stride stride() const
615+ {
616+ return stride_;
617+ }
618+ 
619+ /// Returns the stride of the layout
620+ HOST_DEVICE
621+ Stride &stride()
622+ {
623+ return stride_;
624+ }
625+ 
626+ /// Returns the stride of the layout
627+ HOST_DEVICE
628+ typename Stride::Index stride(int idx) const
629+ {
630+ return stride_[idx];
631+ }
632+ 
633+ /// Returns the stride of the layout
634+ HOST_DEVICE
635+ typename Stride::Index &stride(int idx)
636+ {
637+ return stride_[idx];
638+ }
639+ 
640+ /// Returns the length of the layout
641+ HOST_DEVICE
642+ LongIndex Capacity() const
643+ {
644+ return static_cast<LongIndex>(stride_[3]) * shape_[3];
645+ }
646+ 
647+private:
648+ /// Origin Shape data member
649+ OrgShape orgShape_;
650+ 
651+ /// Shape data member
652+ Shape shape_;
653+ 
654+ /// Stride data member
655+ Stride stride_;
656+};
657+ 
658+/// Mapping function for zN matrices which is row-major inside fractal and row-major between fractal
659+struct zZ {
660+public:
661+ /// Logical rank of tensor
662+ static constexpr int RANK = 4;
663+ 
664+ /// Index type used for coordinates
665+ using Index = uint32_t;
666+ 
667+ /// Long index type used for offsets
668+ using LongIndex = int64_t;
669+ 
670+ /// Logical rank of orgshape
671+ static constexpr int ORG_SHAPE_RANK = 2;
672+ 
673+ /// Logical coordinate
674+ using OrgShape = Coord<ORG_SHAPE_RANK, Index>;
675+ 
676+ /// Logical coordinate
677+ using Shape = Coord<RANK, Index>;
678+ 
679+ /// Stride vector
680+ using Stride = Coord<RANK, LongIndex>;
681+ 
682+public:
683+ // Methods
684+ 
685+ /// Constructor
686+ HOST_DEVICE constexpr
687+ zZ(Index orgRows = 0, /// Number of rows of origin matrices
688+ Index orgCols = 0, /// Number of cols of origin matrices
689+ Index rowsInFractal = 0, /// Number of rows inside the fractal
690+ Index rowsByFractal = 0, /// number of rows by the fractal
691+ Index colsInFractal = 0, /// number of cols inside the fractal
692+ Index colsByFractal = 0, /// number of cols by the fractal
693+ LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal
694+ LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows
695+ LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal
696+ LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols
697+ : orgShape_(MakeCoord(orgRows, orgCols)),
698+ shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)),
699+ stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) {}
700+ 
701+ /// Ctor
702+ HOST_DEVICE constexpr
703+ zZ(OrgShape orgShape, Shape shape, Stride stride) : orgShape_(orgShape), shape_(shape), stride_(stride) {}
704+ 
705+ /// Make the layout of a coordinate (row, column)
706+ template <class Element>
707+ HOST_DEVICE constexpr
708+ static zZ MakeLayout(Index orgRows, Index orgCols)
709+ {
710+ constexpr uint32_t ELE_NUM_PER_C0 = static_cast<uint32_t>(BYTE_PER_C0) / static_cast<uint32_t>(sizeof(Element));
711+ constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
712+ Index rowsRound = RoundUp<C0_NUM_PER_FRACTAL>(orgRows);
713+ Index colsRound = RoundUp<ELE_NUM_PER_C0>(orgCols);
714+ return zZ(orgRows,
715+ orgCols,
716+ C0_NUM_PER_FRACTAL,
717+ rowsRound / C0_NUM_PER_FRACTAL,
718+ ELE_NUM_PER_C0,
719+ colsRound / ELE_NUM_PER_C0,
720+ ELE_NUM_PER_C0,
721+ colsRound * C0_NUM_PER_FRACTAL,
722+ 1,
723+ ELE_NUM_PER_FRACTAL);
724+ }
725+ 
726+ /// Returns the offset of a coordinate in linear memory.
727+ /// Assumes coordinate has convention (row, column)
728+ HOST_DEVICE
729+ LongIndex GetOffset(MatrixCoord const &coord) const
730+ {
731+ return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3];
732+ }
733+ 
734+ /// Returns the origin shape of the layout
735+ HOST_DEVICE
736+ typename OrgShape::Index orgShape(int idx) const
737+ {
738+ return orgShape_[idx];
739+ }
740+ 
741+ /// Returns the origin shape of the layout
742+ HOST_DEVICE
743+ typename OrgShape::Index &orgShape(int idx)
744+ {
745+ return orgShape_[idx];
746+ }
747+ 
748+ /// Returns the shape of the layout
749+ HOST_DEVICE
750+ Shape shape() const
751+ {
752+ return shape_;
753+ }
754+ 
755+ /// Returns the shape of the layout
756+ HOST_DEVICE
757+ Shape &shape()
758+ {
759+ return shape_;
760+ }
761+ 
762+ /// Returns the shape of the layout
763+ HOST_DEVICE
764+ typename Shape::Index shape(int idx) const
765+ {
766+ return shape_[idx];
767+ }
768+ 
769+ /// Returns the shape of the layout
770+ HOST_DEVICE
771+ typename Shape::Index &shape(int idx)
772+ {
773+ return shape_[idx];
774+ }
775+ 
776+ /// Returns the stride of the layout
777+ HOST_DEVICE
778+ Stride stride() const
779+ {
780+ return stride_;
781+ }
782+ 
783+ /// Returns the stride of the layout
784+ HOST_DEVICE
785+ Stride &stride()
786+ {
787+ return stride_;
788+ }
789+ 
790+ /// Returns the stride of the layout
791+ HOST_DEVICE
792+ typename Stride::Index stride(int idx) const
793+ {
794+ return stride_[idx];
795+ }
796+ 
797+ /// Returns the stride of the layout
798+ HOST_DEVICE
799+ typename Stride::Index &stride(int idx)
800+ {
801+ return stride_[idx];
802+ }
803+ 
804+private:
805+ /// Origin Shape data member
806+ OrgShape orgShape_;
807+ 
808+ /// Shape data member
809+ Shape shape_;
810+ 
811+ /// Stride data member
812+ Stride stride_;
813+};
814+ 
815+/// Mapping function for padding rowmajor matrices
816+/// A special data layout designed to improve the efficiency of matrix operations in non-512B aligned scenarios.
817+/// This layout is row-major within blocks and also row-major between blocks.
818+struct PaddingRowMajor {
819+public:
820+ /// Logical rank of tensor
821+ static constexpr int RANK = 4;
822+ 
823+ /// Logical rank of orgshape
824+ static constexpr int ORG_SHAPE_RANK = 2;
825+ 
826+ /// Index type used for coordinates
827+ using Index = uint32_t;
828+ 
829+ /// Long index type used for offsets
830+ using LongIndex = int64_t;
831+ 
832+ /// Logical coordinate
833+ using OrgShape = Coord<ORG_SHAPE_RANK, Index>;
834+ 
835+ /// Logical coordinate
836+ using Shape = Coord<RANK, Index>;
837+ 
838+ /// Stride vector
839+ using Stride = Coord<RANK, LongIndex>;
840+ 
841+public:
842+ /// Constructor
843+ HOST_DEVICE
844+ PaddingRowMajor(Index orgRows = 0, Index orgCols = 0, Index blockRows = 0, Index blockCols = 0) :
845+ orgShape_(MakeCoord(orgRows, orgCols)),
846+ shape_(MakeCoord(blockRows, CeilDiv(orgRows, blockRows), blockCols, CeilDiv(orgCols, blockCols))),
847+ stride_(MakeCoord((LongIndex)blockCols, (LongIndex)blockRows * (LongIndex)RoundUp(orgCols, blockCols),
848+ (LongIndex)1, (LongIndex)blockRows * (LongIndex)blockCols)) {}
849+ 
850+ /// Returns the offset of a coordinate in linear memory.
851+ /// Assumes coordinate has convention (row, column)
852+ HOST_DEVICE
853+ LongIndex GetOffset(MatrixCoord const &coord) const
854+ {
855+ LongIndex blockRows = (LongIndex)shape_[0];
856+ LongIndex blockCols = (LongIndex)shape_[2];
857+ return (LongIndex)coord.row() / blockRows * stride_[1]
858+ + (LongIndex)coord.column() / blockCols * stride_[3]
859+ + (LongIndex)coord.row() % blockRows * stride_[0]
860+ + (LongIndex)coord.column() % blockCols;
861+ }
862+ 
863+ HOST_DEVICE
864+ PaddingRowMajor GetTileLayout(MatrixCoord const &tileShape) const
865+ {
866+ return PaddingRowMajor(tileShape.row(), tileShape.column(), shape_[0], shape_[2]);
867+ }
868+ 
869+ /// Returns the origin shape of the layout
870+ HOST_DEVICE
871+ typename OrgShape::Index orgShape(int idx) const
872+ {
873+ return orgShape_[idx];
874+ }
875+ 
876+ /// Returns the origin shape of the layout
877+ HOST_DEVICE
878+ typename OrgShape::Index &orgShape(int idx)
879+ {
880+ return orgShape_[idx];
881+ }
882+ 
883+ /// Returns the shape of the layout
884+ HOST_DEVICE
885+ Shape shape() const
886+ {
887+ return shape_;
888+ }
889+ 
890+ /// Returns the shape of the layout
891+ HOST_DEVICE
892+ Shape &shape()
893+ {
894+ return shape_;
895+ }
896+ 
897+ /// Returns the shape of the layout
898+ HOST_DEVICE
899+ typename Shape::Index shape(int idx) const
900+ {
901+ return shape_[idx];
902+ }
903+ 
904+ /// Returns the shape of the layout
905+ HOST_DEVICE
906+ typename Shape::Index &shape(int idx)
907+ {
908+ return shape_[idx];
909+ }
910+ 
911+ /// Returns the stride of the layout
912+ HOST_DEVICE
913+ Stride stride() const
914+ {
915+ return stride_;
916+ }
917+ 
918+ /// Returns the stride of the layout
919+ HOST_DEVICE
920+ Stride &stride()
921+ {
922+ return stride_;
923+ }
924+ 
925+ /// Returns the stride of the layout
926+ HOST_DEVICE
927+ typename Stride::Index stride(int idx) const
928+ {
929+ return stride_[idx];
930+ }
931+ 
932+ /// Returns the stride of the layout
933+ HOST_DEVICE
934+ typename Stride::Index &stride(int idx)
935+ {
936+ return stride_[idx];
937+ }
938+ 
939+private:
940+ //
941+ // Data members
942+ //
943+ 
944+ /// Origin Shape data member
945+ OrgShape orgShape_;
946+ 
947+ /// Shape data member
948+ Shape shape_;
949+ 
950+ /// Stride data member
951+ Stride stride_;
952+};
953+ 
954+/// Mapping function for padding columnmajor matrices
955+/// A special data layout designed to improve the efficiency of matrix operations in non-512B aligned scenarios.
956+/// This layout is column-major within blocks and also column-major between blocks.
957+struct PaddingColumnMajor {
958+public:
959+ /// Logical rank of tensor
960+ static constexpr int RANK = 4;
961+ 
962+ /// Logical rank of orgshape
963+ static constexpr int ORG_SHAPE_RANK = 2;
964+ 
965+ /// Index type used for coordinates
966+ using Index = uint32_t;
967+ 
968+ /// Long index type used for offsets
969+ using LongIndex = int64_t;
970+ 
971+ /// Logical coordinate
972+ using OrgShape = Coord<ORG_SHAPE_RANK, Index>;
973+ 
974+ /// Logical coordinate
975+ using Shape = Coord<RANK, Index>;
976+ 
977+ /// Stride vector
978+ using Stride = Coord<RANK, LongIndex>;
979+ 
980+public:
981+ /// Constructor
982+ HOST_DEVICE
983+ PaddingColumnMajor(Index orgRows = 0, Index orgCols = 0, Index blockRows = 0, Index blockCols = 0) :
984+ orgShape_(MakeCoord(orgRows, orgCols)),
985+ shape_(MakeCoord(blockRows, CeilDiv(orgRows, blockRows), blockCols, CeilDiv(orgCols, blockCols))),
986+ stride_(MakeCoord((LongIndex)1, (LongIndex)blockRows * (LongIndex)blockCols, (LongIndex)blockRows,
987+ (LongIndex)RoundUp(orgRows, blockRows) * (LongIndex)blockCols)) {}
988+ 
989+ /// Returns the offset of a coordinate in linear memory.
990+ /// Assumes coordinate has convention (row, column)
991+ HOST_DEVICE
992+ LongIndex GetOffset(MatrixCoord const &coord) const
993+ {
994+ LongIndex blockRows = (LongIndex)shape_[0];
995+ LongIndex blockCols = (LongIndex)shape_[2];
996+ return (LongIndex)coord.row() / blockRows * stride_[1]
997+ + (LongIndex)coord.column() / blockCols * stride_[3]
998+ + (LongIndex)coord.row() % blockRows
999+ + (LongIndex)coord.column() % blockCols * stride_[2];
1000+ }
1001+ 
1002+ HOST_DEVICE
1003+ PaddingColumnMajor GetTileLayout(MatrixCoord const &tileShape) const
1004+ {
1005+ return PaddingColumnMajor(tileShape.row(), tileShape.column(), shape_[0], shape_[2]);
1006+ }
1007+ 
1008+ /// Returns the origin shape of the layout
1009+ HOST_DEVICE
1010+ typename OrgShape::Index orgShape(int idx) const
1011+ {
1012+ return orgShape_[idx];
1013+ }
1014+ 
1015+ /// Returns the origin shape of the layout
1016+ HOST_DEVICE
1017+ typename OrgShape::Index &orgShape(int idx)
1018+ {
1019+ return orgShape_[idx];
1020+ }
1021+ 
1022+ /// Returns the shape of the layout
1023+ HOST_DEVICE
1024+ Shape shape() const
1025+ {
1026+ return shape_;
1027+ }
1028+ 
1029+ /// Returns the shape of the layout
1030+ HOST_DEVICE
1031+ Shape &shape()
1032+ {
1033+ return shape_;
1034+ }
1035+ 
1036+ /// Returns the shape of the layout
1037+ HOST_DEVICE
1038+ typename Shape::Index shape(int idx) const
1039+ {
1040+ return shape_[idx];
1041+ }
1042+ 
1043+ /// Returns the shape of the layout
1044+ HOST_DEVICE
1045+ typename Shape::Index &shape(int idx)
1046+ {
1047+ return shape_[idx];
1048+ }
1049+ 
1050+ /// Returns the stride of the layout
1051+ HOST_DEVICE
1052+ Stride stride() const
1053+ {
1054+ return stride_;
1055+ }
1056+ 
1057+ /// Returns the stride of the layout
1058+ HOST_DEVICE
1059+ Stride &stride()
1060+ {
1061+ return stride_;
1062+ }
1063+ 
1064+ /// Returns the stride of the layout
1065+ HOST_DEVICE
1066+ typename Stride::Index stride(int idx) const
1067+ {
1068+ return stride_[idx];
1069+ }
1070+ 
1071+ /// Returns the stride of the layout
1072+ HOST_DEVICE
1073+ typename Stride::Index &stride(int idx)
1074+ {
1075+ return stride_[idx];
1076+ }
1077+ 
1078+ 
1079+private:
1080+ //
1081+ // Data members
1082+ //
1083+ 
1084+ /// Origin Shape data member
1085+ OrgShape orgShape_;
1086+ 
1087+ /// Shape data member
1088+ Shape shape_;
1089+ 
1090+ /// Stride data member
1091+ Stride stride_;
1092+};
1093+ 
1094+///////////////////////
1095+// new add layout nN
1096+// nN layout
1097+struct nN {
1098+public:
1099+ /// Logical rank of tensor
1100+ static constexpr int RANK = 4;
1101+ 
1102+ /// Index type used for coordinates
1103+ using Index = uint32_t;
1104+ 
1105+ /// Long index type used for offsets
1106+ using LongIndex = int64_t;
1107+ 
1108+ /// Logical rank of orgshape
1109+ static constexpr int ORG_SHAPE_RANK = 2;
1110+ 
1111+ /// Logical coordinate
1112+ using OrgShape = Coord<ORG_SHAPE_RANK, Index>;
1113+ 
1114+ /// Logical coordinate
1115+ using Shape = Coord<RANK, Index>;
1116+ 
1117+ /// Stride vector
1118+ using Stride = Coord<RANK, LongIndex>;
1119+ 
1120+public:
1121+ // Methods
1122+ 
1123+ /// Constructor
1124+ HOST_DEVICE
1125+ nN(Index orgRows = 0, /// Number of rows of origin matrices
1126+ Index orgCols = 0, /// Number of cols of origin matrices
1127+ 
1128+ Index rowsInFractal = 0, /// Number of rows inside the fractal
1129+ Index rowsByFractal = 0, /// number of rows by the fractal
1130+ Index colsInFractal = 0, /// number of cols inside the fractal
1131+ Index colsByFractal = 0, /// number of cols by the fractal
1132+ 
1133+ LongIndex strideRowsInFractal = 0, /// number of elements between adjacent rows inside the fractal
1134+ LongIndex strideRowsByFractal = 0, /// number of elements between adjacent fractal rows
1135+ LongIndex strideColsInFractal = 0, /// number of elements between adjacent cols inside the fractal
1136+ LongIndex strideColsByFractal = 0) /// number of elements between adjacent fractal cols
1137+ : orgShape_(MakeCoord(orgRows, orgCols)),
1138+ shape_(MakeCoord(rowsInFractal, rowsByFractal, colsInFractal, colsByFractal)),
1139+ stride_(MakeCoord(strideRowsInFractal, strideRowsByFractal, strideColsInFractal, strideColsByFractal)) {
1140+ }
1141+ 
1142+ /// Ctor
1143+ HOST_DEVICE
1144+ nN(OrgShape orgShape, Shape shape, Stride stride)
1145+ : orgShape_(orgShape), shape_(shape), stride_(stride) {}
1146+ 
1147+ /// Make the layout of a coordinate (row, column)
1148+ template <class Element>
1149+ HOST_DEVICE static nN MakeLayout(Index orgRows, Index orgCols) {
1150+ static constexpr uint32_t ELE_NUM_PER_C0 = BYTE_PER_C0 / sizeof(Element);
1151+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = BYTE_PER_FRACTAL / sizeof(Element);
1152+ Index rowsRound = RoundUp<ELE_NUM_PER_C0>(orgRows);
1153+ Index colsRound = RoundUp<C0_NUM_PER_FRACTAL>(orgCols);
1154+ return nN(orgRows,
1155+ orgCols,
1156+ 
1157+ ELE_NUM_PER_C0,
1158+ rowsRound / ELE_NUM_PER_C0,
1159+ C0_NUM_PER_FRACTAL,
1160+ colsRound / C0_NUM_PER_FRACTAL,
1161+ 
1162+ 1,
1163+ ELE_NUM_PER_FRACTAL,
1164+ ELE_NUM_PER_C0,
1165+ rowsRound * C0_NUM_PER_FRACTAL);
1166+ }
1167+ 
1168+ /// Returns the offset of a coordinate in linear memory.
1169+ /// Assumes coordinate has convention (row, column)
1170+ HOST_DEVICE
1171+ LongIndex GetOffset(MatrixCoord const& coord) const {
1172+ return LongIndex(coord.row()) / shape_[0] * stride_[1] + LongIndex(coord.column()) / shape_[2] * stride_[3];
1173+ }
1174+ 
1175+ /// Returns the origin shape of the layout
1176+ HOST_DEVICE
1177+ typename OrgShape::Index orgShape(int idx) const {
1178+ return orgShape_[idx];
1179+ }
1180+ 
1181+ /// Returns the origin shape of the layout
1182+ HOST_DEVICE
1183+ typename OrgShape::Index& orgShape(int idx) {
1184+ return orgShape_[idx];
1185+ }
1186+ 
1187+ /// Returns the shape of the layout
1188+ HOST_DEVICE
1189+ Shape shape() const {
1190+ return shape_;
1191+ }
1192+ 
1193+ /// Returns the shape of the layout
1194+ HOST_DEVICE
1195+ Shape& shape() {
1196+ return shape_;
1197+ }
1198+ 
1199+ /// Returns the shape of the layout
1200+ HOST_DEVICE
1201+ typename Shape::Index shape(int idx) const {
1202+ return shape_[idx];
1203+ }
1204+ 
1205+ /// Returns the shape of the layout
1206+ HOST_DEVICE
1207+ typename Shape::Index& shape(int idx) {
1208+ return shape_[idx];
1209+ }
1210+ 
1211+ /// Returns the stride of the layout
1212+ HOST_DEVICE
1213+ Stride stride() const {
1214+ return stride_;
1215+ }
1216+ 
1217+ /// Returns the stride of the layout
1218+ HOST_DEVICE
1219+ Stride& stride() {
1220+ return stride_;
1221+ }
1222+ 
1223+ /// Returns the stride of the layout
1224+ HOST_DEVICE
1225+ typename Stride::Index stride(int idx) const {
1226+ return stride_[idx];
1227+ }
1228+ 
1229+ /// Returns the stride of the layout
1230+ HOST_DEVICE
1231+ typename Stride::Index& stride(int idx) {
1232+ return stride_[idx];
1233+ }
1234+ 
1235+private:
1236+ /// Origin Shape data member
1237+ OrgShape orgShape_;
1238+ 
1239+ /// Shape data member
1240+ Shape shape_;
1241+ 
1242+ /// Stride data member
1243+ Stride stride_;
1244+};
1245+} // namespace NpuArch::layout
1246+ 
1247+#endif // LAYOUT_MATRIX_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/layout/vector.hpp+133-0
@@ -0,0 +1,133 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef LAYOUT_VECTOR_HPP
12+#define LAYOUT_VECTOR_HPP
13+ 
14+#include "../../attn_infra/base_defs.hpp"
15+#include "../../attn_infra/coord.hpp"
16+ 
17+namespace NpuArch::layout
18+{
19+ 
20+struct VectorLayout {
21+public:
22+ /// Logical rank of tensor
23+ static constexpr int RANK = 1;
24+ 
25+ /// Index type used for coordinates
26+ using Index = uint32_t;
27+ 
28+ /// Long index type used for offsets
29+ using LongIndex = int64_t;
30+ 
31+ /// Shape vector
32+ using Shape = Coord<RANK, Index>;
33+ 
34+ /// Stride vector
35+ using Stride = Coord<RANK, LongIndex>;
36+ 
37+ /// Logical coordinate
38+ using TensorCoord = Coord<RANK, Index>;
39+ 
40+public:
41+ // Methods
42+ 
43+ HOST_DEVICE
44+ VectorLayout(Index size = 0) : shape_(MakeCoord(size)), stride_(MakeCoord(LongIndex(1))) {}
45+ 
46+ HOST_DEVICE
47+ VectorLayout(Shape shape, Stride stride) : shape_(shape), stride_(stride) {}
48+ 
49+ template <class Element>
50+ HOST_DEVICE
51+ static VectorLayout MakeLayoutInUb(TensorCoord const &tileShape)
52+ {
53+ return VectorLayout{RoundUp<BYTE_PER_BLK / sizeof(Element)>(tileShape[0])};
54+ }
55+ 
56+ HOST_DEVICE
57+ LongIndex GetOffset(TensorCoord const &coord) const
58+ {
59+ return stride_[0] * coord[0];
60+ }
61+ 
62+ /// Returns the layout of a tile_common.
63+ HOST_DEVICE
64+ VectorLayout GetTileLayout(TensorCoord const &tileShape) const
65+ {
66+ return VectorLayout(tileShape, stride());
67+ }
68+ 
69+ /// Returns the shape of the layout
70+ HOST_DEVICE
71+ Shape shape() const
72+ {
73+ return shape_;
74+ }
75+ 
76+ /// Returns the shape of the layout
77+ HOST_DEVICE
78+ Shape &shape()
79+ {
80+ return shape_;
81+ }
82+ 
83+ /// Returns the shape of the layout
84+ HOST_DEVICE
85+ typename Shape::Index shape(int idx) const
86+ {
87+ return shape_[idx];
88+ }
89+ 
90+ /// Returns the shape of the layout
91+ HOST_DEVICE
92+ typename Shape::Index &shape(int idx)
93+ {
94+ return shape_[idx];
95+ }
96+ 
97+ /// Returns the stride of the layout
98+ HOST_DEVICE
99+ Stride stride() const
100+ {
101+ return stride_;
102+ }
103+ 
104+ /// Returns the stride of the layout
105+ HOST_DEVICE
106+ Stride &stride()
107+ {
108+ return stride_;
109+ }
110+ 
111+ /// Returns the stride of the layout
112+ HOST_DEVICE
113+ typename Stride::Index stride(int idx) const
114+ {
115+ return stride_[idx];
116+ }
117+ 
118+ /// Returns the stride of the layout
119+ HOST_DEVICE
120+ typename Stride::Index &stride(int idx)
121+ {
122+ return stride_[idx];
123+ }
124+ 
125+private:
126+ /// Stride data member
127+ Shape shape_;
128+ Stride stride_;
129+};
130+ 
131+} // namespace NpuArch::layout
132+ 
133+#endif // LAYOUT_VECTOR_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/matrix_coord.hpp+108-0
@@ -0,0 +1,108 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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 matrix_coord.hpp
13+ * \brief
14+ */
15+ 
16+#ifndef MATRIX_COORD_HPP
17+#define MATRIX_COORD_HPP
18+ 
19+#include "../attn_infra/coord.hpp"
20+ 
21+namespace NpuArch {
22+ 
23+template <
24+ uint32_t ROW_ = 1,
25+ uint32_t COLUMN_ = 1
26+>
27+struct MatrixShape {
28+ static constexpr uint32_t ROW = ROW_;
29+ static constexpr uint32_t COLUMN = COLUMN_;
30+ 
31+ static constexpr int64_t COUNT = ROW * COLUMN;
32+ 
33+ HOST_DEVICE
34+ static Coord<2> ToCoord()
35+ {
36+ return MakeCoord(ROW, COLUMN);
37+ }
38+};
39+ 
40+/// MatrixCoord wraps Coord<2, uint32_t> to provide a helper for accessing named dimensions. Classes
41+/// expecting a coordinate in the rank=2 index space of a matrix should use MatrixCoord.
42+struct MatrixCoord : public Coord<2, uint32_t> {
43+ /// Integer-valued index
44+ using Index = uint32_t;
45+ 
46+ /// Base type is a Coord of rank=2
47+ using Base = Coord<2, Index>;
48+ 
49+ /// LongIndex type
50+ using LongIndex = typename Base::LongIndex;
51+ 
52+ /// Rows dimension
53+ static constexpr uint32_t ROW_INDEX = 0;
54+ 
55+ /// Columns dimension
56+ static constexpr uint32_t COLUMN_INDEX = 1;
57+ 
58+ /// Default ctor
59+ HOST_DEVICE
60+ MatrixCoord() {}
61+ 
62+ /// Constructs from Coord<2>
63+ HOST_DEVICE
64+ MatrixCoord(Coord<2, Index> const &coord) : Base(coord) {}
65+ 
66+ /// Helper to construct from a row and column
67+ HOST_DEVICE
68+ MatrixCoord(Index row, Index column) : Base(MakeCoord(row, column)) {}
69+ 
70+ /// Helper to construct from a row and column, which are LongIndex based
71+ HOST_DEVICE
72+ MatrixCoord(LongIndex row, LongIndex column) : Base(MakeCoord(Index(row), Index(column))) {}
73+ 
74+ /// Returns the row of the coordinate
75+ HOST_DEVICE
76+ Index const &row() const { return this->At(ROW_INDEX); }
77+ 
78+ /// Returns the row of the coordinate
79+ HOST_DEVICE
80+ Index &row() { return this->At(ROW_INDEX); }
81+ 
82+ /// Returns the column of the coordinate
83+ HOST_DEVICE
84+ Index const &column() const { return this->At(COLUMN_INDEX); }
85+ 
86+ /// Returns the column of the coordinate
87+ HOST_DEVICE
88+ Index &column() { return this->At(COLUMN_INDEX); }
89+ 
90+ /// Element-wise addition
91+ HOST_DEVICE
92+ MatrixCoord operator+(Base const &b) const
93+ {
94+ return MatrixCoord(Base::operator+(b));
95+ }
96+ 
97+ /// In-place addition
98+ HOST_DEVICE
99+ MatrixCoord &operator+=(Base const &b)
100+ {
101+ Base::operator+=(b);
102+ return *this;
103+ }
104+};
105+ 
106+} // namespace NpuArch
107+ 
108+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/attn_infra/status.hpp+20-0
@@ -0,0 +1,20 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef STATUS_HPP
12+#define STATUS_HPP
13+ 
14+namespace NpuArch{
15+ 
16+enum class Status{ kSuccess, kInvalid };
17+ 
18+} // namespace NpuArch
19+ 
20+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/eagle_quant_block_sparse_attention.cpp+60-0
@@ -0,0 +1,60 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#include "kernel_operator.h"
12+#include "kernel_operator_list_tensor_intf.h"
13+#include "eagle_quant_block_sparse_attention_tilingkey.h"
14+#include "eagle_quant_block_sparse_attention_kernel_interface.cpp"
15+ 
16+extern "C" __global__ __aicore__ void eagle_quant_block_sparse_attention(__gm__ uint8_t* query, __gm__ uint8_t* key, __gm__ uint8_t* value,
17+ __gm__ uint8_t* blockSparseMask, __gm__ uint8_t* mask, __gm__ uint8_t* blockShape,
18+ __gm__ uint8_t* actualSeqLengths, __gm__ uint8_t* actualSeqLengthsKv, __gm__ uint8_t* blockTable,
19+ __gm__ uint8_t* query_scale, __gm__ uint8_t* key_scale, __gm__ uint8_t* value_scale,
20+ __gm__ uint8_t* attentionOut, __gm__ uint8_t* softmaxLse, __gm__ uint8_t* workspace, __gm__ uint8_t* tiling)
21+{
22+ if (TILING_KEY_VAR >= RFA_BASE_TILING) {
23+ __gm__ uint8_t *user = AscendC::GetUserWorkspace(workspace);
24+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
25+ // 读取tilingKey进行kernel分发
26+ __gm__ EagleQuantBlockSparseAttentionTilingData *tilingDataPtr =
27+ reinterpret_cast<__gm__ EagleQuantBlockSparseAttentionTilingData *>(tiling);
28+ uint64_t tilingKey = tilingDataPtr->tilingKey;
29+ 
30+#if (__CCE_AICORE__ == 310)
31+ TILING_KEY_IS(QKINT8_VFP8E4M3_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY);
32+ TILING_KEY_IS(QKINT8_VFP8E4M3_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY);
33+ TILING_KEY_IS(QKINT8_VFP8E4M3_QTND_KVTND_NOCACHE_SMF16_REF32_OBF16_NOMASK_KEY);
34+ TILING_KEY_IS(QKINT8_VFP8E4M3_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_OBF16_NOMASK_KEY);
35+ #if TILING_KEY_VAR == QKINT8_VFP8E4M3_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY
36+ BsaInferIntfRegular<
37+ int8_t, float8_e4m3_t, half, half, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND>(
38+ query, key, value, mask, blockTable, query_scale, key_scale, value_scale, attentionOut,
39+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, tiling);
40+ #elif TILING_KEY_VAR == QKINT8_VFP8E4M3_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY
41+ BsaInferIntfRegular<
42+ int8_t, float8_e4m3_t, half, half, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD>(
43+ query, key, value, mask, blockTable, query_scale, key_scale, value_scale, attentionOut,
44+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, tiling);
45+ #elif TILING_KEY_VAR == QKINT8_VFP8E4M3_QTND_KVTND_NOCACHE_SMF16_REF32_OBF16_NOMASK_KEY
46+ BsaInferIntfRegular<
47+ int8_t, float8_e4m3_t, bfloat16_t, half, float, BsaKernelArch35::Format::TND, BsaKernelArch35::Format::TND>(
48+ query, key, value, mask, blockTable, query_scale, key_scale, value_scale, attentionOut,
49+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, tiling);
50+ #elif TILING_KEY_VAR == QKINT8_VFP8E4M3_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_OBF16_NOMASK_KEY
51+ BsaInferIntfRegular<
52+ int8_t, float8_e4m3_t, bfloat16_t, half, float, BsaKernelArch35::Format::BNSD, BsaKernelArch35::Format::BNSD>(
53+ query, key, value, mask, blockTable, query_scale, key_scale, value_scale, attentionOut,
54+ actualSeqLengths, actualSeqLengthsKv, blockSparseMask, user, tiling);
55+
56+ #endif
57+#endif
58+ }
59+}
60+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/eagle_quant_block_sparse_attention_kernel_interface.cpp+243-0
@@ -0,0 +1,243 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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 eagle_quant_block_sparse_attention_interface.cpp
13+ * \brief Block Sparse Attention Interface
14+ */
15+#include "kernel_operator.h"
16+#if (__CCE_AICORE__ == 220)
17+#include "eagle_quant_block_sparse_attention_kernel_regular_arch32.h"
18+#endif
19+#if (__CCE_AICORE__ == 310)
20+#include "arch35/eagle_quant_block_sparse_attention_kernel_arch35_regular.h"
21+#include "arch35/eagle_quant_block_sparse_attention_kernel_arch35_qmode1.h"
22+#endif
23+ 
24+using namespace NpuArch;
25+ 
26+#if (__CCE_AICORE__ == 220)
27+namespace BlockSparse {
28+template <
29+ typename InputDtype = half,
30+ typename SoftmaxDtype = float,
31+ Epilogue::LseMode lseMode = Epilogue::LseMode::NONE,
32+ uint32_t QueryLayout = 0, // 0=TND, 1=BNSD
33+ uint32_t KvCacheLayout = 0> // 0=TND, 1=BNSD
34+__global__ __aicore__ void EagleQuantBlockSparseAttentionInfer(
35+ GM_ADDR q,
36+ GM_ADDR k,
37+ GM_ADDR v,
38+ GM_ADDR blockSparseMask,
39+ GM_ADDR mask,
40+ GM_ADDR blockTables,
41+ GM_ADDR o,
42+ GM_ADDR actualQseqlen,
43+ GM_ADDR actualKvseqlen,
44+ GM_ADDR blockShape,
45+ GM_ADDR workspace, // S matrix (QK^T result)
46+ GM_ADDR lse,
47+ GM_ADDR tiling)
48+{
49+ using ArchTag = Arch::AtlasA2;
50+ using ElementQ = InputDtype;
51+ using LayoutQ = layout::RowMajor;
52+ using ElementK = InputDtype;
53+ using LayoutK = layout::ColumnMajor;
54+ using ElementV = InputDtype;
55+ using LayoutV = layout::RowMajor;
56+ using ElementS = SoftmaxDtype;
57+ using LayoutS = layout::RowMajor;
58+ using ElementP = InputDtype;
59+ using LayoutP = layout::RowMajor;
60+ using ElementO = InputDtype;
61+ using LayoutO = layout::RowMajor;
62+ using ElementLse = float;
63+ using LayoutLse = layout::RowMajor;
64+ using ElementMask = int8_t;
65+ using LayoutMask = layout::RowMajor;
66+ using ElementOTmp = SoftmaxDtype;
67+ using LayoutOTmp = layout::RowMajor;
68+ using ElementUpdate = SoftmaxDtype;
69+ using LayoutUpdate = layout::RowMajor;
70+ 
71+ // Use sparse-specific dispatch policies
72+ using L1TileShapeQK = GemmShape<Q_TILE_CEIL, 128, 128>;
73+ using L0TileShapeQK = GemmShape<128, 128, 128>;
74+ using DispatchPolicyQK = Gemm::MmadAtlasA2SFAIQK<false, false>;
75+ using QType = Gemm::GemmType<ElementQ, LayoutQ>;
76+ using KType = Gemm::GemmType<ElementK, LayoutK>;
77+ using SType = Gemm::GemmType<ElementS, LayoutS>;
78+ using BlockMmadQK = Gemm::Block::BlockMmad<DispatchPolicyQK, L1TileShapeQK, L0TileShapeQK,
79+ QType, KType, SType>;
80+ 
81+ using L1TileShapePV = GemmShape<128, 128, 256>;
82+ using L0TileShapePV = GemmShape<128, 128, 128>;
83+ using DispatchPolicyPV = Gemm::MmadAtlasA2SFAIPV<false, false>;
84+ using PType = Gemm::GemmType<ElementP, LayoutP>;
85+ using VType = Gemm::GemmType<ElementV, LayoutV>;
86+ using OTmpType = Gemm::GemmType<ElementOTmp, LayoutOTmp>;
87+ using BlockMmadPV = Gemm::Block::BlockMmad<DispatchPolicyPV, L1TileShapePV, L0TileShapePV,
88+ PType, VType, OTmpType>;
89+ 
90+ // Epilogue policies for sparse attention
91+ using DispatchPolicyOnlineSoftmax = Epilogue::EpilogueAtlasA2OnlineSoftmax<lseMode, SoftmaxDtype>;
92+ using MaskType = Gemm::GemmType<ElementMask, LayoutMask>;
93+ using EpilogueOnlineSoftmax = Epilogue::Block::BlockEpilogue<DispatchPolicyOnlineSoftmax,
94+ PType, SType, MaskType>;
95+ using DispatchPolicyRescaleO = Epilogue::EpilogueAtlasA2RescaleO<lseMode, SoftmaxDtype>;
96+ using OType = Gemm::GemmType<ElementO, LayoutO>;
97+ using OUpdateType = Gemm::GemmType<ElementUpdate, LayoutUpdate>;
98+ using LseType = Gemm::GemmType<ElementLse, LayoutLse>;
99+ using EpilogueRescaleO = Epilogue::Block::BlockEpilogue<DispatchPolicyRescaleO,
100+ OType, OTmpType, OUpdateType, LseType>;
101+ 
102+ // Kernel instantiation
103+ using EagleQuantBlockSparseAttentionKernelType = EagleQuantBlockSparseAttentionKernel<BlockMmadQK, BlockMmadPV,
104+ EpilogueOnlineSoftmax,
105+ EpilogueRescaleO,
106+ false, // PAGED_CACHE_FLAG
107+ QueryLayout, // QUERY_LAYOUT
108+ KvCacheLayout>; // KV_CACHE_LAYOUT
109+ EagleQuantBlockSparseAttentionKernelParams params{q, k, v, blockSparseMask, mask, blockTables, actualQseqlen, actualKvseqlen,
110+ o, lse, workspace, tiling};
111+ 
112+ // Call block sparse attention kernel
113+ EagleQuantBlockSparseAttentionKernelType blockSparseAttenInfer;
114+ blockSparseAttenInfer(params);
115+}
116+}
117+#endif
118+#if (__CCE_AICORE__ == 310)
119+ 
120+// Lookup helper for ElementP type selection based on ElementQ and ElementV
121+template <typename QType, typename VType, typename DefaultType>
122+struct ElementPTypeLookup {
123+ using type = DefaultType;
124+};
125+ 
126+template <typename DefaultType>
127+struct ElementPTypeLookup<int8_t, float8_e4m3_t, DefaultType> {
128+ using type = float8_e4m3_t;
129+};
130+ 
131+template <typename QType, typename KType, typename VType, typename OUTPODType>
132+struct GetQuantMode {
133+ static constexpr int value = 0;
134+};
135+ 
136+template <>
137+struct GetQuantMode<int8_t, int8_t, float8_e4m3_t, bfloat16_t> {
138+ static constexpr int value = 1;
139+};
140+ 
141+template <>
142+struct GetQuantMode<int8_t, int8_t, float8_e4m3_t, half> {
143+ static constexpr int value = 1;
144+};
145+ 
146+ 
147+template <int quantMode>
148+struct QKScaleGranularitySelector {
149+ static constexpr NpuArch::Gemm::Tile::ScaleGranularity value = NpuArch::Gemm::Tile::ScaleGranularity::NO_QUANT;
150+};
151+ 
152+template <>
153+struct QKScaleGranularitySelector<1> {
154+ static constexpr NpuArch::Gemm::Tile::ScaleGranularity value = NpuArch::Gemm::Tile::ScaleGranularity::PER_TENSOR;
155+};
156+ 
157+template <int quantMode>
158+struct PVScaleGranularitySelector {
159+ static constexpr NpuArch::Gemm::Tile::ScaleGranularity value = NpuArch::Gemm::Tile::ScaleGranularity::NO_QUANT;
160+};
161+ 
162+template <>
163+struct PVScaleGranularitySelector<1> {
164+ static constexpr NpuArch::Gemm::Tile::ScaleGranularity value = NpuArch::Gemm::Tile::ScaleGranularity::NO_QUANT;
165+};
166+ 
167+using namespace BsaKernelArch35;
168+ 
169+template <class InQKDtype, class InVDtype, class OUTPODType, class SMDtype, class REDtype, Format qFormat, Format kvFormat>
170+__global__ __aicore__ void BsaInferIntfRegular(
171+ GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR mask, GM_ADDR blockTables,
172+ GM_ADDR query_scale, GM_ADDR key_scale, GM_ADDR value_scale,
173+ GM_ADDR o, GM_ADDR actualQseqlen, GM_ADDR actualKvseqlen,
174+ GM_ADDR blockSparseMask, GM_ADDR workspace,
175+ GM_ADDR tiling
176+) {
177+ using ArchTag = Arch::AtlasA5;
178+ using ElementSparseMask = uint8_t;
179+ using ElementSparseIdx = int32_t;
180+ using ElementSparseCount = int32_t;
181+ using ElementQ = InQKDtype;
182+ using ElementK = InQKDtype;
183+ using ElementV = InVDtype;
184+ using ElementS = SMDtype;
185+ using ElementP = typename ElementPTypeLookup<ElementQ, ElementV, OUTPODType>::type;
186+ using ElementO = OUTPODType;
187+ using ElementOTmp = REDtype;
188+ // layout tags
189+ using LayoutQ = layout::RowMajor;
190+ using LayoutK = layout::ColumnMajor;
191+ // S is rowMajor on UB(dst)
192+ using LayoutS = layout::RowMajor;
193+ // P is actually zN on UB(src), since there is no nd2nz in MTE1
194+ //tODO后续增加一个selector，支持nz和zn
195+ using LayoutPDummy = layout::nZ;
196+ using LayoutV = layout::RowMajor;
197+ using LayoutO = layout::RowMajor;
198+ // OTmp is rowMajor on UB(dst)
199+ using LayoutOTmp = layout::RowMajor;
200+ using LayoutSparseIdx = layout::RowMajor;
201+ using LayoutSparseCount = layout::RowMajor;
202+ // block mask pre-process
203+ using DispatchPolicyMask2Idx = Epilogue::EpilogueBsaMask2Idx;
204+ using EpilogueMask2Idx = Epilogue::Block::BlockEpilogue<
205+ DispatchPolicyMask2Idx, ElementSparseMask, ElementSparseIdx, ElementSparseCount>;
206+ // 处理单个tile内Q和K的matmul
207+ using L1TileShapeQK = Shape<Int<128>, Int<512>, Int<128>>;
208+ using L0TileShapeQK = Shape<Int<128>, Int<256>, Int<128>>;
209+ using DispatchPolicyQK = Gemm::MmadAtlasA5BsaQK;
210+ using TileCopyQK = Gemm::Tile::PackedTileCopyTlaToUB<
211+ ArchTag, ElementQ, LayoutQ, ElementK, LayoutK, ElementS, LayoutS,
212+ void, Gemm::Tile::CopyL0CToUBMode::NO_SPLIT, false, QKScaleGranularitySelector<GetQuantMode<ElementQ, ElementK, ElementV, OUTPODType>::value>::value>;
213+ using BlockMmadQK = Gemm::Block::BlockMmadTla<
214+ DispatchPolicyQK, L1TileShapeQK, L0TileShapeQK, ElementQ, ElementK, ElementS, void, TileCopyQK>;
215+ // online softmax
216+ using DispatchPolicyOnlineSoftmax = Epilogue::EpilogueOnlineSoftmaxBsa;
217+ using PType = Gemm::GemmType<ElementP, layout::zN>;
218+ using SType = Gemm::GemmType<ElementS, LayoutS>;
219+ using EpilogueOnlineSoftmax = Epilogue::Block::BlockEpilogue<DispatchPolicyOnlineSoftmax, PType, SType>;
220+ // 处理单个tile内P和Value的matmul
221+ using L1TileShapePV = Shape<Int<128>, Int<128>, Int<512>>;
222+ using L0TileShapePV = Shape<Int<128>, Int<128>, Int<256>>;
223+ using DispatchPolicyPV = Gemm::MmadAtlasA5BsaPV;
224+ using TileCopyPV = Gemm::Tile::PackedTileCopyTlaToUB<
225+ ArchTag, ElementP, LayoutPDummy, ElementV, LayoutV, ElementOTmp, LayoutOTmp,
226+ void, Gemm::Tile::CopyL0CToUBMode::SPLIT_M, false, PVScaleGranularitySelector<GetQuantMode<ElementQ, ElementK, ElementV, OUTPODType>::value>::value>;
227+ using BlockMmadPV = Gemm::Block::BlockMmadTla<
228+ DispatchPolicyPV, L1TileShapePV, L0TileShapePV, ElementP, ElementV, ElementOTmp, void, TileCopyPV>;
229+ // rescale O
230+ using DispatchPolicyRescaleO = Epilogue::EpilogueAtlasA5BsaRescaleO;
231+ using TileCopyRescaleO = Epilogue::Tile::TileCopyRescaleO<
232+ ArchTag, ElementO, LayoutO, LayoutOTmp>;
233+ using EpilogueRescaleO = Epilogue::Block::BlockEpilogue<
234+ DispatchPolicyRescaleO, ElementO, ElementOTmp, ElementS, TileCopyRescaleO, Arch::PositionL0C>;
235+ 
236+ using BsaRegularKernelArch35 = BsaRegularKernelArch35<
237+ EpilogueMask2Idx, BlockMmadQK, EpilogueOnlineSoftmax, BlockMmadPV, EpilogueRescaleO, qFormat, kvFormat>;
238+ BsaKernelParamsArch35 params{q, k, v, mask, blockTables, query_scale, key_scale, value_scale,
239+ actualQseqlen, actualKvseqlen, blockSparseMask, o, workspace, tiling};
240+ BsaRegularKernelArch35 bsaRegularKernelArch35;
241+ bsaRegularKernelArch35.template operator()<GetQuantMode<ElementQ, ElementK, ElementV, OUTPODType>::value>(params);
242+}
243+#endif
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/eagle_quant_block_sparse_attention_kernel_regular_arch32.h+794-0
@@ -0,0 +1,794 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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 eagle_quant_block_sparse_attention_kernel.h
13+ * \brief Block Sparse Attention Kernel Implementation
14+ */
15+ 
16+#ifndef EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_KERNEL_H
17+#define EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_KERNEL_H
18+ 
19+#include "kernel_common.hpp"
20+ 
21+using namespace NpuArch;
22+using namespace RfaKenelCommon;
23+ 
24+namespace BlockSparse {
25+ /**
26+ * @brief Block Sparse Attention Inference Kernel
27+ * This kernel implements block sparse attention where attention is computed only on
28+ * selected KV blocks specified by selectIdx. This reduces computation for long sequences
29+ * by focusing on relevant tokens.
30+ * @tparam BlockMmadQK Block-level QK matmul module
31+ * @tparam BlockMmadPV Block-level PV matmul module
32+ * @tparam EpilogueOnlineSoftmax Online softmax epilogue
33+ * @tparam EpilogueRescaleO Output rescaling epilogue
34+ * @tparam PAGED_CACHE_FLAG Whether to use paged KV cache
35+ * @tparam QUERY_LAYOUT Query tensor layout (0=TND, 1=BNSD)
36+ * @tparam KV_CACHE_LAYOUT KV cache layout (0=TND, 1=BNSD)
37+ */
38+ template <
39+ class BlockMmadQK,
40+ class BlockMmadPV,
41+ class EpilogueOnlineSoftmax,
42+ class EpilogueRescaleO,
43+ bool PAGED_CACHE_FLAG,
44+ uint32_t QUERY_LAYOUT,
45+ uint32_t KV_CACHE_LAYOUT>
46+ class EagleQuantBlockSparseAttentionKernel {
47+ public:
48+ using ArchTag = typename BlockMmadQK::ArchTag;
49+ using L1TileShape = typename BlockMmadQK::L1TileShape;
50+ using ElementQ = typename BlockMmadQK::ElementA;
51+ using LayoutQ = typename BlockMmadQK::LayoutA;
52+ using ElementK = typename BlockMmadQK::ElementB;
53+ using LayoutK = typename BlockMmadQK::LayoutB;
54+ using ElementS = typename BlockMmadQK::ElementC;
55+ using LayoutS = typename BlockMmadQK::LayoutC;
56+ 
57+ using ElementP = typename BlockMmadPV::ElementA;
58+ using LayoutP = typename BlockMmadPV::LayoutA;
59+
60+ using ElementV = typename BlockMmadPV::ElementB;
61+ using LayoutV = typename BlockMmadPV::LayoutB;
62+ 
63+ using ElementMask = typename EpilogueOnlineSoftmax::ElementMask;
64+ 
65+ using ElementO = typename EpilogueRescaleO::ElementOutput;
66+ using LayoutO = typename EpilogueRescaleO::LayoutOutput;
67+ 
68+ using ElementOTmp = typename EpilogueRescaleO::ElementInput;
69+ using LayoutOTmp = typename EpilogueRescaleO::LayoutInput;
70+ 
71+ using ElementLse = typename EpilogueRescaleO::ElementLse;
72+ using LayoutLse = typename EpilogueRescaleO::LayoutLse;
73+ 
74+ using ElementUpdate = typename EpilogueRescaleO::ElementUpdate;
75+ using LayoutUpdate = typename EpilogueRescaleO::LayoutUpdate;
76+ 
77+ static constexpr Epilogue::LseMode LSE_MODE = EpilogueRescaleO::LSE_MODE;
78+ static constexpr int32_t BASIC_BLOCK = 64;
79+ static constexpr uint32_t LS_UB_TENSOR_OFFSET = 0;
80+ static constexpr uint32_t MASK_PATTERN_HALF_OFFSET = BASIC_BLOCK * 2 + LS_UB_TENSOR_OFFSET;
81+ static constexpr uint32_t MASK_PATTERN_FLOAT_OFFSET = BASIC_BLOCK * 2 + MASK_PATTERN_HALF_OFFSET;
82+ static constexpr uint32_t MASK_BIT_OFFSET = BASIC_BLOCK * 4 + MASK_PATTERN_FLOAT_OFFSET;
83+ static constexpr uint32_t MASK_IDX_OFFSET = BASIC_BLOCK + MASK_BIT_OFFSET;
84+ static constexpr uint32_t SPARSE_IDX_OFFSET = BASIC_BLOCK * 4 + MASK_IDX_OFFSET;
85+ static constexpr uint32_t SELECT_NUM_IDX_OFFSET = BASIC_BLOCK * 4 + SPARSE_IDX_OFFSET;
86+ static constexpr uint32_t SYNC_OFFSET = BASIC_BLOCK * 4 + SELECT_NUM_IDX_OFFSET;
87+
88+ __aicore__ inline
89+ EagleQuantBlockSparseAttentionKernel() {}
90+ 
91+ __aicore__ inline void Mask2IdxAndCount(const AscendC::GlobalTensor<uint8_t> maskGM,
92+ AscendC::GlobalTensor<int32_t> selectIdxGM,
93+ AscendC::GlobalTensor<int32_t> selectNumGM)
94+ {
95+ static constexpr uint32_t PRE_ROW_TILE = 128;
96+ static constexpr uint32_t PRE_COL_TILE = 64;
97+ static constexpr uint32_t PRE_ELEM_NUM_PER_LOOP = PRE_ROW_TILE * PRE_COL_TILE;
98+ static constexpr uint32_t MASK_PAT_IN_UINT8 = 0;
99+ static constexpr uint32_t MASK_PAT_IN_FP16 = 2 * PRE_ELEM_NUM_PER_LOOP;
100+ static constexpr uint32_t MASK_PAT_IN_FP32 = 4 * PRE_ELEM_NUM_PER_LOOP;
101+ static constexpr uint32_t MASK_PAT_IN_BIT = 8 * PRE_ELEM_NUM_PER_LOOP;
102+ static constexpr uint32_t MASK_IDX = 9 * PRE_ELEM_NUM_PER_LOOP;
103+ static constexpr uint32_t RSVD_SPARSE_IDX = 13 * PRE_ELEM_NUM_PER_LOOP;
104+ static constexpr uint32_t RSVD_SPARSE_COUNT = 21 * PRE_ELEM_NUM_PER_LOOP;
105+ 
106+ AscendC::LocalTensor<uint8_t> maskPatternUbLocal[2];
107+ AscendC::LocalTensor<int32_t> selectNumIdxUbLocal[2];
108+ AscendC::LocalTensor<uint8_t> maskPatternInBitUbLocal;
109+ AscendC::LocalTensor<uint32_t> maskPatternInBitUbLocalUint32;
110+ 
111+ AscendC::LocalTensor<int32_t> maskIdxUbLocal;
112+ AscendC::LocalTensor<int32_t> sparseIdxUbLocal[2];
113+ AscendC::LocalTensor<half> maskPatternHalfLocal;
114+ AscendC::LocalTensor<float> maskPatternFloatLocal;
115+ 
116+ for (uint32_t i = 0; i < 2; i++) {
117+ maskPatternUbLocal[i] = resource.ubBuf.template GetBufferByByte<uint8_t>(
118+ MASK_PAT_IN_UINT8 + i * PRE_ELEM_NUM_PER_LOOP * sizeof(int8_t));
119+ sparseIdxUbLocal[i] = resource.ubBuf.template GetBufferByByte<int32_t>(
120+ RSVD_SPARSE_IDX + i * PRE_ELEM_NUM_PER_LOOP * sizeof(int32_t));
121+ selectNumIdxUbLocal[i] = resource.ubBuf.template GetBufferByByte<int32_t>(
122+ RSVD_SPARSE_COUNT + i * PRE_ROW_TILE * sizeof(int32_t));
123+ }
124+ maskPatternHalfLocal = resource.ubBuf.template GetBufferByByte<half>(MASK_PAT_IN_FP16);
125+ maskPatternFloatLocal = resource.ubBuf.template GetBufferByByte<float>(MASK_PAT_IN_FP32);
126+ maskPatternInBitUbLocal = resource.ubBuf.template GetBufferByByte<uint8_t>(MASK_PAT_IN_BIT);
127+ maskPatternInBitUbLocalUint32 = resource.ubBuf.template GetBufferByByte<uint32_t>(MASK_PAT_IN_BIT);
128+ maskIdxUbLocal = resource.ubBuf.template GetBufferByByte<int32_t>(MASK_IDX);
129+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(0);
130+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(1);
131+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(0);
132+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(1);
133+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(2);
134+ AscendC::SetFlag<AscendC::HardEvent::S_MTE3>(0);
135+ AscendC::SetFlag<AscendC::HardEvent::MTE3_S>(0);
136+ AscendC::SetFlag<AscendC::HardEvent::MTE3_S>(1);
137+ 
138+ uint32_t subCoreIdx = AscendC::GetBlockIdx();
139+ uint32_t curStartRowIdx = subCoreIdx * avgRowPerSubCore;
140+ uint32_t totalRowNumBlockMask = batch * qHeads * maxQBlockNum;
141+ uint32_t actDealtRow = (subCoreIdx == preActivateSubCoreNum - 1) ?
142+ (totalRowNumBlockMask - curStartRowIdx) : avgRowPerSubCore;
143+ if (subCoreIdx < preActivateSubCoreNum) {
144+ uint32_t rowLoop = CeilDiv(actDealtRow, PRE_ROW_TILE);
145+ uint32_t colLoop = CeilDiv(maxKvBlockNum, PRE_COL_TILE);
146+ uint32_t idxPingPongFlag = 0;
147+ uint32_t countPingPongFlag = 0;
148+ for (uint32_t i = 0; i < rowLoop; i++) {
149+ uint32_t curLoopRowOffset = i * PRE_ROW_TILE;
150+ int32_t actDealtRowCurLoop = (i == rowLoop - 1) ? (actDealtRow - curLoopRowOffset) : PRE_ROW_TILE;
151+ uint64_t rsvdCountPerRow[PRE_ROW_TILE] = {0};
152+ uint64_t rsvdCountPerRowCurColLoop[PRE_ROW_TILE] = {0};
153+ for (uint32_t j = 0; j < colLoop; j++) {
154+ uint32_t curLoopColOffset = j * PRE_COL_TILE;
155+ uint32_t actDealtColCurLoop =
156+ (j == colLoop - 1) ? (maxKvBlockNum - curLoopColOffset) : PRE_COL_TILE;
157+ uint32_t actDealtColCurLoop32 = CeilDiv(actDealtColCurLoop, 32) * 32;
158+ AscendC::CreateVecIndex(maskIdxUbLocal, static_cast<int32_t>(curLoopColOffset), PRE_COL_TILE);
159+ AscendC::PipeBarrier<PIPE_V>();
160+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(idxPingPongFlag);
161+ int32_t maskOffset = curStartRowIdx * maxKvBlockNum +
162+ curLoopRowOffset * maxKvBlockNum + curLoopColOffset;
163+ AscendC::DataCopyPad(
164+ maskPatternUbLocal[idxPingPongFlag],
165+ maskGM[maskOffset],
166+ AscendC::DataCopyExtParams(
167+ actDealtRowCurLoop,
168+ actDealtColCurLoop * sizeof(int8_t),
169+ (maxKvBlockNum - actDealtColCurLoop) * sizeof(int8_t),
170+ (PRE_COL_TILE - actDealtColCurLoop32) / 32, 0),
171+ AscendC::DataCopyPadExtParams<uint8_t>(
172+ true, 0, (actDealtColCurLoop32 - actDealtColCurLoop), 0));
173+
174+ AscendC::SetFlag<AscendC::HardEvent::MTE2_V>(0);
175+ AscendC::WaitFlag<AscendC::HardEvent::MTE2_V>(0);
176+ AscendC::Cast(
177+ maskPatternHalfLocal, maskPatternUbLocal[idxPingPongFlag],
178+ AscendC::RoundMode::CAST_NONE, actDealtRowCurLoop * PRE_COL_TILE);
179+ AscendC::PipeBarrier<PIPE_V>();
180+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(idxPingPongFlag);
181+ AscendC::Cast(
182+ maskPatternFloatLocal, maskPatternHalfLocal,
183+ AscendC::RoundMode::CAST_NONE, actDealtRowCurLoop * PRE_COL_TILE);
184+ AscendC::PipeBarrier<PIPE_V>();
185+ for (uint32_t k = 0; k < actDealtRowCurLoop; k++) {
186+ if (actDealtColCurLoop32 != PRE_COL_TILE) {
187+ AscendC::Duplicate(maskPatternFloatLocal[k * PRE_COL_TILE + actDealtColCurLoop32],
188+ (float)0, PRE_COL_TILE - actDealtColCurLoop32);
189+ }
190+ AscendC::PipeBarrier<PIPE_V>();
191+ AscendC::CompareScalar(
192+ maskPatternInBitUbLocal[k * PRE_COL_TILE],
193+ maskPatternFloatLocal[k * PRE_COL_TILE],
194+ (float)1.0, AscendC::CMPMODE::GE, PRE_COL_TILE);
195+ AscendC::PipeBarrier<PIPE_V>();
196+ if (k == 0) {
197+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(idxPingPongFlag);
198+ }
199+ AscendC::GatherMask(
200+ sparseIdxUbLocal[idxPingPongFlag][k * PRE_COL_TILE],
201+ maskIdxUbLocal,
202+ maskPatternInBitUbLocalUint32[k * PRE_COL_TILE / 4],
203+ false,
204+ (uint32_t)0, {1, 1, 0, 0}, rsvdCountPerRowCurColLoop[k]);
205+ AscendC::SetFlag<AscendC::HardEvent::V_MTE3>(0);
206+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE3>(0);
207+ if (k == 0) {
208+ AscendC::WaitFlag<AscendC::HardEvent::S_MTE3>(0);
209+ }
210+ AscendC::DataCopyPad(
211+ selectIdxGM[curStartRowIdx * maxKvBlockNum +
212+ curLoopRowOffset * maxKvBlockNum + k * maxKvBlockNum + rsvdCountPerRow[k]],
213+ sparseIdxUbLocal[idxPingPongFlag][k * PRE_COL_TILE],
214+ AscendC::DataCopyExtParams(
215+ 1, rsvdCountPerRowCurColLoop[k] * sizeof(int32_t), 0, 0, 0));
216+ if (k == actDealtRowCurLoop - 1) {
217+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(idxPingPongFlag);
218+ }
219+ AscendC::SetFlag<AscendC::HardEvent::MTE3_S>(0);
220+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_S>(0);
221+ rsvdCountPerRow[k] += rsvdCountPerRowCurColLoop[k];
222+ if (k == actDealtRowCurLoop - 1) {
223+ AscendC::SetFlag<AscendC::HardEvent::S_MTE3>(0);
224+ }
225+ }
226+ idxPingPongFlag = 1 - idxPingPongFlag;
227+ }
228+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_S>(countPingPongFlag);
229+ for (uint32_t k = 0; k < actDealtRowCurLoop; k++) {
230+ selectNumIdxUbLocal[countPingPongFlag].SetValue(k, static_cast<int32_t>(rsvdCountPerRow[k]));
231+ }
232+ AscendC::SetFlag<AscendC::HardEvent::S_MTE3>(countPingPongFlag + 2);
233+ AscendC::WaitFlag<AscendC::HardEvent::S_MTE3>(countPingPongFlag + 2);
234+ AscendC::DataCopyPad(
235+ selectNumGM[curStartRowIdx + curLoopRowOffset],
236+ selectNumIdxUbLocal[countPingPongFlag],
237+ AscendC::DataCopyExtParams(1, actDealtRowCurLoop * sizeof(int32_t), 0, 0, 0));
238+ AscendC::SetFlag<AscendC::HardEvent::MTE3_S>(countPingPongFlag);
239+ countPingPongFlag = 1 - countPingPongFlag;
240+ }
241+ }
242+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(0);
243+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(1);
244+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(0);
245+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(1);
246+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(2);
247+ AscendC::WaitFlag<AscendC::HardEvent::S_MTE3>(0);
248+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_S>(0);
249+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_S>(1);
250+ }
251+ 
252+ __aicore__ inline void operator()(EagleQuantBlockSparseAttentionKernelParams const &params)
253+ {
254+ __gm__ EagleQuantBlockSparseAttentionTilingData *blockSparseAttentionTilingData =
255+ reinterpret_cast<__gm__ EagleQuantBlockSparseAttentionTilingData *>(params.tiling);
256+ uint64_t mm1OutSize = blockSparseAttentionTilingData->mm1OutSize;
257+ uint64_t smOnlineOutSize = blockSparseAttentionTilingData->smOnlineOutSize;
258+ uint64_t mm2OutSize = blockSparseAttentionTilingData->mm2OutSize;
259+ uint64_t updateSize = blockSparseAttentionTilingData->updateSize;
260+ uint64_t selectNumIdxSize = blockSparseAttentionTilingData->selectNumIdxSize;
261+ uint64_t selectIdxSize = blockSparseAttentionTilingData->selectIdxSize;
262+ 
263+ batch = blockSparseAttentionTilingData->batch;
264+ qHeads = blockSparseAttentionTilingData->numHeads;
265+ uint32_t kvHeads = blockSparseAttentionTilingData->kvHeads;
266+ 
267+ uint32_t embed = blockSparseAttentionTilingData->embeddingSize;
268+ uint32_t pagedBlockSize = blockSparseAttentionTilingData->blockSize;
269+ uint32_t maxNumBlocksPerBatch = blockSparseAttentionTilingData->maxNumBlocksPerBatch;
270+ uint32_t firstBatchTaskNum = blockSparseAttentionTilingData->firstBatchTaskNum;
271+ uint32_t totalTaskNum = blockSparseAttentionTilingData->totalTaskNum;
272+ uint32_t maskType = blockSparseAttentionTilingData->maskType;
273+ ElementS scaleValue = static_cast<ElementS>(blockSparseAttentionTilingData->scaleValue);
274+ uint32_t totalQBlocks = blockSparseAttentionTilingData->totalQBlocks;
275+ maxKvBlockNum = blockSparseAttentionTilingData->maxKvBlockNum;
276+ uint32_t maxKvBlockNumPad = CeilDiv(maxKvBlockNum, 32) * 32;
277+ maxQBlockNum = blockSparseAttentionTilingData->maxQBlockNum;
278+ avgRowPerSubCore = blockSparseAttentionTilingData->avgRowNumPerSubCore;
279+ preActivateSubCoreNum = blockSparseAttentionTilingData->preActivateSubCoreNum;
280+
281+ uint32_t qBlockX = blockSparseAttentionTilingData->blockShapeX;
282+ uint32_t qBlockY = blockSparseAttentionTilingData->blockShapeY;
283+ uint32_t qBlockNum = totalQBlocks / qBlockX;
284+ uint32_t qBlockInX = (qBlockX + BASIC_BLOCK_SIZE - 1) / BASIC_BLOCK_SIZE;
285+ uint32_t firstQBlockNum = blockSparseAttentionTilingData->firstQBlockNum;
286+ uint32_t maxQSeqlen = blockSparseAttentionTilingData->maxQSeqlen;
287+ uint32_t maxKvSeqlen = blockSparseAttentionTilingData->maxKvSeqlen;
288+ uint32_t useUniformQSeqlen = blockSparseAttentionTilingData->useUniformQSeqlen;
289+ uint32_t useUniformKvSeqlen = blockSparseAttentionTilingData->useUniformKvSeqlen;
290+ 
291+ // Initialize global tensors
292+ AscendC::GlobalTensor<ElementQ> gQ;
293+ gQ.SetGlobalBuffer((__gm__ ElementQ *)params.q);
294+ AscendC::GlobalTensor<ElementK> gK;
295+ gK.SetGlobalBuffer((__gm__ ElementK *)params.k);
296+ AscendC::GlobalTensor<ElementK> gV;
297+ gV.SetGlobalBuffer((__gm__ ElementK *)params.v);
298+ AscendC::GlobalTensor<int32_t> gBlockTable;
299+ gBlockTable.SetGlobalBuffer((__gm__ int32_t *)(params.blockTables));
300+ AscendC::GlobalTensor<int64_t> gActualQseqlen;
301+ gActualQseqlen.SetGlobalBuffer((__gm__ int64_t *)params.actualQseqlen);
302+ AscendC::GlobalTensor<int64_t> gActualKvseqlen;
303+ gActualKvseqlen.SetGlobalBuffer((__gm__ int64_t *)params.actualKvseqlen);
304+ AscendC::GlobalTensor<int32_t> gSelectIdx;
305+ gSelectIdx.SetGlobalBuffer((__gm__ int32_t *)(params.workspace + mm1OutSize +
306+ smOnlineOutSize + mm2OutSize + updateSize + selectNumIdxSize));
307+ AscendC::GlobalTensor<int32_t> gSelectNumIdx;
308+ gSelectNumIdx.SetGlobalBuffer((__gm__ int32_t *)(params.workspace + mm1OutSize +
309+ smOnlineOutSize + mm2OutSize + updateSize));
310+ AscendC::GlobalTensor<uint8_t> gBlockSparseMask;
311+ gBlockSparseMask.SetGlobalBuffer((__gm__ uint8_t *)params.blockSparseMask);
312+ AscendC::GlobalTensor<ElementO> gO;
313+ gO.SetGlobalBuffer((__gm__ ElementO *)params.o);
314+ AscendC::GlobalTensor<ElementLse> gLse;
315+ gLse.SetGlobalBuffer((__gm__ ElementLse *)params.lse);
316+ AscendC::GlobalTensor<ElementS> gS;
317+ gS.SetGlobalBuffer((__gm__ ElementS *)params.workspace);
318+ AscendC::GlobalTensor<ElementP> gP;
319+ gP.SetGlobalBuffer((__gm__ ElementP *)(params.workspace + mm1OutSize));
320+ AscendC::GlobalTensor<ElementOTmp> gOTmp;
321+ gOTmp.SetGlobalBuffer((__gm__ ElementOTmp *)(params.workspace + mm1OutSize + smOnlineOutSize));
322+ AscendC::GlobalTensor<ElementOTmp> gOUpdate;
323+ gOUpdate.SetGlobalBuffer((__gm__ ElementOTmp *)(params.workspace + mm1OutSize +
324+ smOnlineOutSize + mm2OutSize));
325+
326+ uint32_t coreIdx = AscendC::GetBlockIdx();
327+ uint32_t coreNum = AscendC::GetBlockNum();
328+#ifdef __DAV_C220_VEC__
329+ Mask2IdxAndCount(gBlockSparseMask, gSelectIdx, gSelectNumIdx);
330+#endif
331+ resource.pipe.Reset();
332+ AscendC::SyncAll<false>();
333+ 
334+#ifdef __DAV_C220_CUBE__
335+ // Initialize hardware events for cube core
336+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
337+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID1);
338+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID2);
339+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID3);
340+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID4);
341+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID5);
342+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID6);
343+ AscendC::SetFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID7);
344+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
345+ AscendC::SetFlag<AscendC::HardEvent::FIX_M>(EVENT_ID1);
346+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
347+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID1);
348+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
349+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID3);
350+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID4);
351+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID5);
352+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID6);
353+ AscendC::SetFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID7);
354+
355+ static constexpr uint32_t L1_QK_SIZE =
356+ BlockMmadQK::L1TileShape::M * BlockMmadQK::L1TileShape::K * sizeof(ElementQ) +
357+ BlockMmadQK::L1TileShape::N * BlockMmadQK::L1TileShape::K * sizeof(ElementK) * 2;
358+ BlockMmadQK blockMmadQK(resource);
359+ BlockMmadPV blockMmadPV(resource, L1_QK_SIZE);
360+#endif
361+ 
362+#ifdef __DAV_C220_VEC__
363+ // Initialize hardware events for vector core
364+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
365+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID1);
366+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
367+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID2);
368+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID3);
369+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID4);
370+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID5);
371+ AscendC::SetFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID6);
372+ 
373+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
374+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID1);
375+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
376+ AscendC::SetFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID2);
377+ AscendC::SetFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
378+ 
379+ EpilogueOnlineSoftmax epilogueOnlineSoftmax(resource, scaleValue);
380+ EpilogueRescaleO epilogueRescaleO(resource);
381+ 
382+ coreIdx = AscendC::GetBlockIdx() / AscendC::GetSubBlockNum();
383+ uint32_t bn = AscendC::GetSubBlockNum();
384+#endif
385+ // Calculate strides based on layout (compile-time optimization)
386+ // For TND: [T, N, D], stride = N * D
387+ // For BNSD: [B, N, S, D], strideB = N * S * D, strideN = S * D, strideS = D
388+ uint64_t strideQO = 0;
389+ uint64_t strideKV = 0;
390+ uint64_t strideQOB = 0; // BNSD batch stride for Q
391+ uint64_t strideQON = 0; // BNSD head stride for Q
392+ uint64_t strideQOS = 0; // BNSD seq stride for Q
393+ uint64_t strideKVB = 0; // BNSD batch stride for KV
394+ uint64_t strideKVN = 0; // BNSD head stride for KV
395+ uint64_t strideKVS = 0; // BNSD seq stride for KV
396+
397+ if constexpr (QUERY_LAYOUT == 1) { // BNSD_Q
398+ // BNSD: [B, N, S, D]
399+ // strideB = N * S * D, strideN = S * D, strideS = D
400+ // maxQSeqlen is the third dimension (S) of query shape, set in tiling
401+ strideQOB = qHeads * maxQSeqlen * embed; // batch stride
402+ strideQON = maxQSeqlen * embed; // head stride
403+ strideQOS = embed; // seq stride
404+ } else {
405+ // TND: [T, N, D]
406+ strideQO = qHeads * embed;
407+ }
408+
409+ if constexpr (KV_CACHE_LAYOUT == 1) { // BNSD
410+ // BNSD: [B, N, S, D]
411+ // maxKvSeqlen is the third dimension (S) of value shape, set in tiling
412+ strideKVB = kvHeads * maxKvSeqlen * embed; // batch stride
413+ strideKVN = maxKvSeqlen * embed; // head stride
414+ strideKVS = embed; // seq stride
415+ } else {
416+ // TND: [T, N, D]
417+ strideKV = kvHeads * embed;
418+ }
419+
420+ uint32_t embedRound = AlignUp<uint32_t>(embed, BLOCK_SIZE);
421+ uint32_t groupSize = qHeads / kvHeads;
422+ 
423+ uint64_t qBOffset = 0;
424+ uint64_t kBOffset = 0;
425+ uint64_t vBOffset = 0;
426+ uint64_t oBOffset = 0;
427+ uint64_t blockBOffset = 0;
428+ uint64_t lseBOffset = 0;
429+ 
430+ uint32_t preTotalTaskNum = 0;
431+ uint32_t preTotalQBlockNum = 0;
432+ uint32_t curBatch = 0;
433+ // 根据useUniformQSeqlen标志位决定使用actualSeqLengths数组还是maxQSeqlen
434+ uint32_t qSeqlen = useUniformQSeqlen ? maxQSeqlen :
435+ static_cast<uint32_t>(static_cast<int64_t>(gActualQseqlen.GetValue(curBatch)));
436+ // 根据useUniformKvSeqlen标志位决定使用actualSeqLengthsKv数组还是maxKvSeqlen
437+ uint32_t kvSeqlen = useUniformKvSeqlen ? maxKvSeqlen :
438+ static_cast<uint32_t>(static_cast<int64_t>(gActualKvseqlen.GetValue(curBatch)));
439+ uint32_t curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize);
440+ uint32_t qNBlockNumPerGroup = curQNBlockTile == 0 ? 1 : (groupSize + curQNBlockTile - 1) / curQNBlockTile;
441+ uint32_t curQNBlockNum = qNBlockNumPerGroup * kvHeads;
442+ uint32_t curQSBlockTile = GetQSBlockTile(kvSeqlen);
443+ uint32_t curQSBlockNum = GetQBlocks(qSeqlen, qBlockX);
444+ uint32_t curTotalTaskNum = firstBatchTaskNum;
445+ uint32_t curQXBlockNum = (qSeqlen + qBlockX - 1) / qBlockX;
446+ uint32_t curTotalQBlockNum = firstQBlockNum;
447+ 
448+ // Go through each task
449+ for (uint32_t taskIdx = coreIdx; taskIdx < totalTaskNum; taskIdx += uint32_t(coreNum)) {
450+ while (taskIdx >= curTotalTaskNum) {
451+ ++curBatch;
452+ preTotalTaskNum = curTotalTaskNum;
453+ preTotalQBlockNum = curTotalQBlockNum;
454+
455+ // Update offsets based on layout (compile-time optimization)
456+ if constexpr (QUERY_LAYOUT == 1) { // BNSD_Q
457+ // BNSD: [B, N, S, D], offset = batch * strideB
458+ qBOffset = curBatch * strideQOB;
459+ oBOffset = curBatch * strideQOB;
460+ lseBOffset = curBatch * qHeads * maxQSeqlen;
461+ } else {
462+ // TND
463+ qBOffset += qSeqlen * strideQO;
464+ oBOffset += qSeqlen * strideQO;
465+ lseBOffset += qSeqlen * qHeads;
466+ }
467+
468+ if constexpr (!PAGED_CACHE_FLAG) {
469+ if constexpr (KV_CACHE_LAYOUT == 1) { // BNSD
470+ // BNSD: [B, N, S, D], offset = batch * strideB
471+ kBOffset = curBatch * strideKVB;
472+ vBOffset = curBatch * strideKVB;
473+ } else {
474+ // TND
475+ kBOffset += kvSeqlen * strideKV;
476+ vBOffset += kvSeqlen * strideKV;
477+ }
478+ } else {
479+ blockBOffset += maxNumBlocksPerBatch;
480+ }
481+ // 根据useUniformQSeqlen标志位决定使用actualSeqLengths数组还是maxQSeqlen
482+ qSeqlen = useUniformQSeqlen ? maxQSeqlen :
483+ static_cast<uint32_t>(static_cast<int64_t>(gActualQseqlen.GetValue(curBatch)));
484+ // 根据useUniformKvSeqlen标志位决定使用actualSeqLengthsKv数组还是maxKvSeqlen
485+ kvSeqlen = useUniformKvSeqlen ? maxKvSeqlen :
486+ static_cast<uint32_t>(static_cast<int64_t>(gActualKvseqlen.GetValue(curBatch)));
487+ curQNBlockTile = GetQNBlockTile(qSeqlen, groupSize);
488+ qNBlockNumPerGroup = curQNBlockTile == 0 ? 1 : (groupSize + curQNBlockTile - 1) / curQNBlockTile;
489+ curQNBlockNum = qNBlockNumPerGroup * kvHeads;
490+ curQSBlockTile = GetQSBlockTile(kvSeqlen);
491+ curQSBlockNum = GetQBlocks(qSeqlen, qBlockX);
492+ curTotalTaskNum += curQNBlockNum * curQSBlockNum;
493+ curQXBlockNum = (qSeqlen + qBlockX - 1) / qBlockX;
494+ curTotalQBlockNum += qHeads * curQXBlockNum;
495+ }
496+ 
497+ // Q task splitting按照[qNBlockNum, qHead]
498+ uint32_t taskIdxCurBatch = taskIdx - preTotalTaskNum;
499+ uint32_t qSBlockIdx = taskIdxCurBatch / curQNBlockNum;
500+ uint32_t qXIdx = qSBlockIdx / qBlockInX;
501+ uint32_t qXInnerIdx = qSBlockIdx - qXIdx * qBlockInX;
502+ uint32_t qNBlockIdx = taskIdxCurBatch - qSBlockIdx * curQNBlockNum;
503+ uint32_t qNBlockIdxCurGroup = qNBlockIdx % qNBlockNumPerGroup;
504+ uint32_t xBlockNum = qSeqlen / qBlockX;
505+ uint32_t xTailNum = qSeqlen - xBlockNum * qBlockX;
506+
507+ uint32_t kvHeadIdx = qNBlockIdx / qNBlockNumPerGroup;
508+ uint32_t qHeadIdx = kvHeadIdx * groupSize + qNBlockIdxCurGroup * curQNBlockTile;
509+ 
510+ uint32_t curSelectIdx = curBatch * qHeads * maxQBlockNum + qHeadIdx * maxQBlockNum + qXIdx;
511+ uint32_t curSelectNum = static_cast<uint32_t>(gSelectNumIdx.GetValue(curSelectIdx));
512+ if (curSelectNum == 0) {
513+ continue;
514+ }
515+ uint32_t lastSelectIdx = static_cast<int32_t>(
516+ gSelectIdx.GetValue(curSelectIdx * maxKvBlockNum + curSelectNum - 1));
517+ uint32_t kvYBlockNum = (kvSeqlen + qBlockY - 1) / qBlockY; // CeilDiv
518+ uint32_t curKvSeqLen = (lastSelectIdx == kvYBlockNum - 1 && kvSeqlen % qBlockY != 0) ?
519+ qBlockY * (curSelectNum - 1) + kvSeqlen % qBlockY : qBlockY * curSelectNum;
520+ // Calculate offsets based on layout (compile-time optimization)
521+ uint64_t gmOffsetQ = 0;
522+ uint64_t gmOffsetK = 0;
523+ uint64_t gmOffsetV = 0;
524+ uint64_t gmOffsetO = 0;
525+ uint64_t gmOffsetLse = 0;
526+ if constexpr (QUERY_LAYOUT == 1) { // BNSD_Q: [B, N, S, D]
527+ // offset = batch * strideB + head * strideN + seq * strideS
528+ uint32_t qSeqOffset = qXIdx * qBlockX + qXInnerIdx * BASIC_BLOCK_SIZE;
529+ gmOffsetQ = qBOffset + qHeadIdx * strideQON + qSeqOffset * strideQOS;
530+ gmOffsetO = oBOffset + qHeadIdx * strideQON + qSeqOffset * strideQOS;
531+ // LSE format: [B, N, S] - strideN = maxQSeqlen
532+ gmOffsetLse = lseBOffset + qHeadIdx * maxQSeqlen + qSeqOffset;
533+ } else {
534+ // TND: [T, N, D]
535+ uint32_t qSeqOffset = qXIdx * qBlockX + qXInnerIdx * BASIC_BLOCK_SIZE;
536+ gmOffsetQ = qBOffset + qSeqOffset * strideQO + qHeadIdx * embed;
537+ gmOffsetO = oBOffset + qSeqOffset * strideQO + qHeadIdx * embed;
538+ // LSE format: [T, N] - same as Q/O but without D dimension
539+ gmOffsetLse = lseBOffset + qSeqOffset * qHeads + qHeadIdx;
540+ }
541+
542+ if constexpr (KV_CACHE_LAYOUT == 1) { // BNSD: [B, N, S, D]
543+ // offset = batch * strideB + head * strideN
544+ // seq offset will be handled in blockMmadQK/blockMmadPV based on selectIdx
545+ gmOffsetK = kBOffset + kvHeadIdx * strideKVN;
546+ gmOffsetV = vBOffset + kvHeadIdx * strideKVN;
547+ } else {
548+ // TND: [T, N, D]
549+ gmOffsetK = kBOffset + kvHeadIdx * embed;
550+ gmOffsetV = vBOffset + kvHeadIdx * embed;
551+ }
552+ 
553+ uint32_t qSBlockSize = (qXIdx == xBlockNum) ?
554+ (qXInnerIdx == xTailNum / curQSBlockTile ?
555+ xTailNum - qXInnerIdx * curQSBlockTile : curQSBlockTile) :
556+ ((qXInnerIdx == qBlockInX - 1) ? qBlockX - qXInnerIdx * curQSBlockTile : curQSBlockTile);
557+ 
558+ uint32_t qNBlockSize = (qNBlockIdxCurGroup == (qNBlockNumPerGroup - 1)) ?
559+ (groupSize - qNBlockIdxCurGroup * curQNBlockTile) : curQNBlockTile;
560+ uint32_t rowNum = qSBlockSize * qNBlockSize;
561+ uint32_t rowNumRound = AlignUp<uint32_t>(rowNum, BLOCK_SIZE);
562+ 
563+ uint32_t noSkipKvS = curKvSeqLen;
564+ uint32_t kvSLoopNumTotal = (noSkipKvS + pagedBlockSize - 1) / pagedBlockSize; // CeilDiv
565+ 
566+ uint32_t blockStackNum = MAX_KV_STACK_LEN / pagedBlockSize;
567+ uint32_t stackSeqTile;
568+ uint32_t stackSeqTilePad = blockStackNum * pagedBlockSize;
569+ uint32_t preKVNum = PRE_LAUNCH * blockStackNum;
570+ int32_t stackSeqCount = 0;
571+ 
572+#ifdef __DAV_C220_CUBE__
573+ LayoutQ layoutQTemp(rowNum, embed);
574+ // For BNSD format, use strideKVS; for TND, use strideKV (compile-time)
575+ uint64_t actualStrideKV = 0;
576+ if constexpr (KV_CACHE_LAYOUT == 1) {
577+ actualStrideKV = strideKVS;
578+ } else {
579+ actualStrideKV = strideKV;
580+ }
581+ LayoutK layoutKTemp(actualStrideKV, blockStackNum * pagedBlockSize);
582+ LayoutV layoutVTemp(blockStackNum * pagedBlockSize, actualStrideKV);
583+ // Pass correct Q stride based on data format
584+ uint64_t qGmStride = 0;
585+ if constexpr (QUERY_LAYOUT == 1) { // BNSD: [B, N, S, D]
586+ qGmStride = strideQOS; // embed
587+ } else { // TND: [T, N, D]
588+ qGmStride = strideQO; // qHeads * embed
589+ }
590+ blockMmadQK.loadQGM(gQ[gmOffsetQ], layoutQTemp, rowNum, qNBlockSize, qGmStride);
591+#endif
592+ // Main computation loop: QK matmul -> Softmax -> PV matmul
593+ for (uint32_t kvSIdx = 0; kvSIdx < kvSLoopNumTotal + preKVNum; kvSIdx += blockStackNum) {
594+ // Stage 1: QK matmul (computed on CUBE core)
595+ if (kvSIdx < kvSLoopNumTotal) {
596+ stackSeqTile = noSkipKvS - kvSIdx * pagedBlockSize;
597+ if (stackSeqTile >= pagedBlockSize * blockStackNum) {
598+ stackSeqTile = pagedBlockSize * blockStackNum;
599+ }
600+ uint32_t curStackTileMod = stackSeqCount % (PRE_LAUNCH + 1);
601+ uint64_t gmOffsetS = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (PRE_LAUNCH + 1) +
602+ curStackTileMod * WORKSPACE_BLOCK_SIZE_DB;
603+ GemmCoord actualBlockShapeQK{rowNum, stackSeqTile, embed};
604+ LayoutS layOutS(rowNum, stackSeqTile, stackSeqTilePad);
605+#ifdef __DAV_C220_CUBE__
606+ // For BNSD format, pass strideKVS; for TND, pass strideKV (compile-time)
607+ uint64_t actualStrideKVForQK = 0;
608+ if constexpr (KV_CACHE_LAYOUT == 1) {
609+ actualStrideKVForQK = strideKVS;
610+ } else {
611+ actualStrideKVForQK = strideKV;
612+ }
613+ blockMmadQK(gQ[gmOffsetQ],
614+ gK[gmOffsetK],
615+ gS[gmOffsetS],
616+ gBlockTable[blockBOffset],
617+ gSelectIdx[curSelectIdx * maxKvBlockNum],
618+ layoutQTemp,
619+ layoutKTemp,
620+ layOutS,
621+ actualBlockShapeQK,
622+ kvSIdx,
623+ kvSLoopNumTotal,
624+ pagedBlockSize,
625+ actualStrideKVForQK,
626+ qBlockY,
627+ curSelectNum,
628+ kvYBlockNum,
629+ kvSeqlen);
630+ NpuArch::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(qkReady);
631+#endif
632+#ifdef __DAV_C220_VEC__
633+ // Stage 2: Online softmax (computed on VECTOR core)
634+ LayoutP layOutP(rowNum, stackSeqTile, stackSeqTilePad);
635+ uint64_t gmOffsetP = gmOffsetS;
636+ 
637+ NpuArch::Arch::CrossCoreWaitFlag(qkReady);
638+ // online softmax
639+ epilogueOnlineSoftmax(gP[gmOffsetP],
640+ gS[gmOffsetS],
641+ layOutP,
642+ layOutS,
643+ actualBlockShapeQK,
644+ (stackSeqCount == 0),
645+ 0,
646+ qSBlockSize,
647+ qNBlockSize,
648+ curStackTileMod,
649+ softmaxReady);
650+#endif
651+ }
652+ // Stage 3: PV matmul and output rescaling
653+ if (kvSIdx >= preKVNum) {
654+ uint32_t nowkvSIdx = kvSIdx - preKVNum;
655+ stackSeqTile = noSkipKvS - nowkvSIdx * pagedBlockSize;
656+ if (stackSeqTile >= pagedBlockSize * blockStackNum) {
657+ stackSeqTile = pagedBlockSize * blockStackNum;
658+ }
659+ uint32_t curStackTileMod = (stackSeqCount - PRE_LAUNCH) % (PRE_LAUNCH + 1);
660+ uint64_t gmOffsetOTmp = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (PRE_LAUNCH + 1) +
661+ curStackTileMod * WORKSPACE_BLOCK_SIZE_DB;
662+ GemmCoord actualBlockShapePV{rowNum, embed, stackSeqTile};
663+ LayoutOTmp layoutOTmp(rowNum, embed, embedRound);
664+#ifdef __DAV_C220_CUBE__
665+ LayoutP layoutPTemp(rowNum, stackSeqTile, stackSeqTilePad);
666+ uint64_t gmOffsetP = coreIdx * WORKSPACE_BLOCK_SIZE_DB * (PRE_LAUNCH + 1) +
667+ curStackTileMod * WORKSPACE_BLOCK_SIZE_DB;
668+ // For BNSD format, pass strideKVS; for TND, pass strideKV (compile-time)
669+ uint64_t actualStrideKVForPV = 0;
670+ if constexpr (KV_CACHE_LAYOUT == 1) {
671+ actualStrideKVForPV = strideKVS;
672+ } else {
673+ actualStrideKVForPV = strideKV;
674+ }
675+ blockMmadPV(gP[gmOffsetP],
676+ gV[gmOffsetV],
677+ gOTmp[gmOffsetOTmp],
678+ gBlockTable[blockBOffset],
679+ gSelectIdx[curSelectIdx * maxKvBlockNum],
680+ layoutPTemp,
681+ layoutVTemp,
682+ layoutOTmp,
683+ actualBlockShapePV,
684+ nowkvSIdx,
685+ kvSLoopNumTotal,
686+ pagedBlockSize,
687+ kvSeqlen,
688+ actualStrideKVForPV,
689+ blockStackNum,
690+ softmaxReady,
691+ qBlockY,
692+ curSelectNum,
693+ kvYBlockNum);
694+ NpuArch::Arch::CrossCoreSetFlag<0x2, PIPE_FIX>(pvReady);
695+#endif
696+#ifdef __DAV_C220_VEC__
697+ // Setup layoutO based on data format
698+ LayoutO layoutO;
699+ LayoutLse layoutLse;
700+ if constexpr (QUERY_LAYOUT == 1) { // BNSD: [B, N, S, D]
701+ // BNSD format: stride[0] = embed (strideQOS)
702+ layoutO = LayoutO(qSeqlen, embed);
703+ layoutLse = LayoutLse(qSeqlen, 1);
704+ } else { // TND: [T, N, D]
705+ // TND format: stride[0] = qHeads * embed (strideQO)
706+ layoutO = LayoutO(qSeqlen, qHeads * embed);
707+ layoutLse = LayoutLse(qSeqlen, qHeads);
708+ }
709+ LayoutUpdate layoutUpdate(rowNum, embed, embedRound);
710+ uint64_t gmOffsetUpdate = (uint64_t)(coreIdx * WORKSPACE_BLOCK_SIZE_DB);
711+ 
712+ NpuArch::Arch::CrossCoreWaitFlag(pvReady);
713+ // rescale O
714+ epilogueRescaleO(
715+ gO[gmOffsetO],
716+ gOTmp[gmOffsetOTmp],
717+ gOUpdate[gmOffsetUpdate],
718+ gLse[gmOffsetLse],
719+ layoutO,
720+ layoutOTmp,
721+ layoutUpdate,
722+ layoutLse,
723+ actualBlockShapePV,
724+ qSBlockSize,
725+ qNBlockSize,
726+ (stackSeqCount - PRE_LAUNCH == 0),
727+ nowkvSIdx + blockStackNum >= kvSLoopNumTotal,
728+ curStackTileMod);
729+#endif
730+ }
731+ stackSeqCount++;
732+ }
733+ }
734+#ifdef __DAV_C220_CUBE__
735+ // Wait for all CUBE core events
736+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID0);
737+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID1);
738+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID2);
739+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID3);
740+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID4);
741+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID5);
742+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID6);
743+ AscendC::WaitFlag<AscendC::HardEvent::M_MTE1>(EVENT_ID7);
744+ 
745+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID0);
746+ AscendC::WaitFlag<AscendC::HardEvent::FIX_M>(EVENT_ID1);
747+ 
748+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID0);
749+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID1);
750+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID2);
751+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID3);
752+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID4);
753+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID5);
754+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID6);
755+ AscendC::WaitFlag<AscendC::HardEvent::MTE1_MTE2>(EVENT_ID7);
756+#endif
757+#ifdef __DAV_C220_VEC__
758+ // Wait for all VECTOR core events
759+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID2);
760+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID3);
761+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID4);
762+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID5);
763+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_MTE2>(EVENT_ID6);
764+ 
765+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID0);
766+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID1);
767+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID2);
768+ AscendC::WaitFlag<AscendC::HardEvent::MTE3_V>(EVENT_ID4);
769+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID0);
770+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID1);
771+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID2);
772+ AscendC::WaitFlag<AscendC::HardEvent::V_MTE2>(EVENT_ID3);
773+#endif
774+ AscendC::PipeBarrier<PIPE_ALL>();
775+ }
776+ 
777+ private:
778+ NpuArch::Arch::Resource<ArchTag> resource;
779+ NpuArch::Arch::CrossCoreFlag qkReady{QK_READY_ID};
780+ NpuArch::Arch::CrossCoreFlag softmaxReady{SOFTMAX_READY_ID};
781+ NpuArch::Arch::CrossCoreFlag pvReady{PV_READY_ID};
782+ 
783+ uint32_t batch{0};
784+ uint32_t qHeads{0};
785+ uint32_t maxQBlockNum{0};
786+ uint32_t maxKvBlockNum{0};
787+ uint32_t avgRowPerSubCore{0};
788+ uint32_t preActivateSubCoreNum{0};
789+ };
790+ 
791+} // namespace BlockSparse
792+ 
793+#endif // EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_KERNEL_H
794+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/eagle_quant_block_sparse_attention_tilingkey.h+83-0
@@ -0,0 +1,83 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_TILINGKEY_H_
12+#define EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_TILINGKEY_H_
13+ 
14+/**
15+ * EagleQuantBlockSparseAttention TilingKey 定义
16+ *
17+ * TilingKey编码规则 (64-bit):
18+ * 位域: AAAABBBBCCCCDDDDEEEE
19+ * - [0-1] Q Layout: 2=TND, 3=BNSD
20+ * - [2-4] Mask Type: 0=NoMask, 3=CausalMask
21+ * - [5-7] Softmax Precision: 0=Float, 1=Half
22+ * - [8-10] PagedCache Flag: 0=NoCache, 1=WithCache
23+ * - [11-13] KV Layout: 00=TND, 20=BNSD
24+ * - [14-15] Data Type: 00=FP16, 22=BF16
25+ * - [16-18] Operator Category: 900=EagleQuantBlockSparseAttention
26+ */
27+ 
28+#define RFA_BASE_TILING 9000000000000000
29+ 
30+#if (__CCE_AICORE__ == 220)
31+// FP16, Q=TND, KV=TND, No PagedCache, Float Softmax, No Mask
32+#define QF16_KVF16_TND_TND_NOCACHE_FLOATSM_NOMASK_RFA_TILING 9000000030000002
33+ 
34+// FP16, Q=TND, KV=TND, No PagedCache, Half Softmax, No Mask
35+#define QF16_KVF16_TND_TND_NOCACHE_HALFSM_NOMASK_RFA_TILING 9000000030100002
36+ 
37+// BF16, Q=TND, KV=TND, No PagedCache, Float Softmax, No Mask
38+#define QBF16_KVBF16_TND_TND_NOCACHE_FLOATSM_NOMASK_RFA_TILING 9000000030022222
39+ 
40+// FP16, Q=BNSD, KV=BNSD, No PagedCache, Float Softmax, No Mask
41+#define QF16_KVF16_BNSD_BNSD_NOCACHE_FLOATSM_NOMASK_RFA_TILING 9000000050000003
42+ 
43+// FP16, Q=BNSD, KV=BNSD, No PagedCache, Half Softmax, No Mask
44+#define QF16_KVF16_BNSD_BNSD_NOCACHE_HALFSM_NOMASK_RFA_TILING 9000000050100003
45+ 
46+// BF16, Q=BNSD, KV=BNSD, No PagedCache, Float Softmax, No Mask
47+#define QBF16_KVBF16_BNSD_BNSD_NOCACHE_FLOATSM_NOMASK_RFA_TILING 9000000050022223
48+ 
49+ // LSE Output versions (LSE_MODE::OUT_ONLY)
50+// FP16, Q=TND, KV=TND, No PagedCache, Float Softmax, No Mask, LSE Output
51+#define QF16_KVF16_TND_TND_NOCACHE_FLOATSM_NOMASK_RFA_TILING_LSE_OUT 9000000130000002
52+ 
53+// FP16, Q=TND, KV=TND, No PagedCache, Half Softmax, No Mask, LSE Output
54+#define QF16_KVF16_TND_TND_NOCACHE_HALFSM_NOMASK_RFA_TILING_LSE_OUT 9000000130100002
55+ 
56+// BF16, Q=TND, KV=TND, No PagedCache, Float Softmax, No Mask, LSE Output
57+#define QBF16_KVBF16_TND_TND_NOCACHE_FLOATSM_NOMASK_RFA_TILING_LSE_OUT 9000000130022222
58+ 
59+// FP16, Q=BNSD, KV=BNSD, No PagedCache, Float Softmax, No Mask, LSE Output
60+#define QF16_KVF16_BNSD_BNSD_NOCACHE_FLOATSM_NOMASK_RFA_TILING_LSE_OUT 9000000150000003
61+ 
62+// FP16, Q=BNSD, KV=BNSD, No PagedCache, Half Softmax, No Mask, LSE Output
63+#define QF16_KVF16_BNSD_BNSD_NOCACHE_HALFSM_NOMASK_RFA_TILING_LSE_OUT 9000000150100003
64+ 
65+// BF16, Q=BNSD, KV=BNSD, No PagedCache, Float Softmax, No Mask, LSE Output
66+#define QBF16_KVBF16_BNSD_BNSD_NOCACHE_FLOATSM_NOMASK_RFA_TILING_LSE_OUT 9000000150022223
67+ 
68+#endif
69+ 
70+#if (__CCE_AICORE__ == 310)
71+ 
72+#define QF16_KVF16_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY 9050000030400002
73+#define QBF16_KVBF16_QTND_KVTND_NOCACHE_SMBF16_REF32_NOMASK_KEY 9050000030422222
74+#define QF16_KVF16_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY 9050000050400003
75+#define QBF16_KVBF16_QBNSD_KVBNSD_NOCACHE_SMBF16_REF32_NOMASK_KEY 9050000050422223
76+#define QKINT8_VFP8E4M3_QTND_KVTND_NOCACHE_SMF16_REF32_NOMASK_KEY 9050010030444442
77+#define QKINT8_VFP8E4M3_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_NOMASK_KEY 9050010050444443
78+#define QKINT8_VFP8E4M3_QTND_KVTND_NOCACHE_SMF16_REF32_OBF16_NOMASK_KEY 9050010030455552
79+#define QKINT8_VFP8E4M3_QBNSD_KVBNSD_NOCACHE_SMF16_REF32_OBF16_NOMASK_KEY 9050010050455553
80+ 
81+#endif
82+#endif // EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_TILINGKEY_H_
83+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/kernel_common.hpp+179-0
@@ -0,0 +1,179 @@
1+/*
2+ * Copyright (c) Huawei Technologies Co., Ltd. 2026. All rights reserved.
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+#ifndef EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_KERNEL_COMMON_HPP
12+#define EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_KERNEL_COMMON_HPP
13+ 
14+#include "attn_infra/base_defs.hpp"
15+#include "attn_infra/arch/arch.hpp"
16+#include "attn_infra/layout/layout.hpp"
17+ 
18+#include "attn_infra/gemm/block/block_mmad.hpp"
19+#include "attn_infra/gemm/dispatch_policy.hpp"
20+#include "attn_infra/gemm/gemm_type.hpp"
21+ 
22+#include "attn_infra/arch/cross_core_sync.hpp"
23+#include "attn_infra/arch/resource.hpp"
24+#include "attn_infra/epilogue/block/block_epilogue.hpp"
25+#include "attn_infra/epilogue/dispatch_policy.hpp"
26+#include "kernel_operator.h"
27+#include "lib/matmul_intf.h"
28+#include "kernel_tiling/kernel_tiling.h"
29+ 
30+using namespace AscendC;
31+using namespace matmul;
32+ 
33+namespace RfaKenelCommon {
34+ constexpr uint32_t QK_READY_ID = 1;
35+ constexpr uint32_t SOFTMAX_READY_ID = 2;
36+ constexpr uint32_t PV_READY_ID = 3;
37+ constexpr uint32_t MASKTOIDX_READY_ID = 4;
38+ 
39+ constexpr uint32_t BLOCK_SIZE = 16;
40+ constexpr uint32_t WORKSPACE_BLOCK_SIZE_DB = 131072;
41+ constexpr uint32_t TMP_SIZE_DECODER = 32768;
42+ 
43+ constexpr int32_t TILING_BATCH = 0;
44+ constexpr int32_t TILING_NUMHEADS = 1;
45+ constexpr int32_t TILING_HEADDIM = 2;
46+ constexpr int32_t TILING_NUMBLOKS = 3;
47+ constexpr int32_t TILING_BLOCKSIZE = 4;
48+ constexpr int32_t TILING_MAXBLOCKS = 5;
49+ constexpr int32_t TILING_TOR = 6;
50+ constexpr int32_t TILING_KVHEADS = 7;
51+ constexpr int32_t TILING_HEADSIZE = 8;
52+ constexpr int32_t TILING_PARASIZE = 9;
53+ constexpr int32_t TILING_HEAD_SPLIT_SIZE = 10;
54+ constexpr int32_t TILING_HEAD_SPLIT_NUM = 11;
55+ constexpr int32_t TILING_HEADDIM_ROPE = 13;
56+ constexpr int32_t TILING_MAX_KVSEQLEN = 14;
57+ constexpr int32_t TILING_KVSPLIT = 15;
58+ constexpr int32_t TILING_KVCORENUM = 16;
59+ constexpr int32_t TILING_TOTAL_QTOKENS = 18;
60+ constexpr int32_t TILING_FORMERTASKNUM = 19;
61+ constexpr int32_t TILING_TAILTASKNUM = 20;
62+ constexpr int32_t TILING_BLOCKSIZE_CALC = 25;
63+ constexpr int32_t TILING_HEADDIM_K_SPLIT = 38;
64+ constexpr int32_t TILING_HEADDIM_V_SPLIT = 39;
65+ constexpr int32_t TILING_HEADDIM_V_SPLIT_VECTOR_FORMER = 40;
66+ constexpr int32_t TILING_HEADDIM_V_SPLIT_VECTOR_TAIL = 41;
67+ 
68+ constexpr int32_t NUM1 = 1;
69+ constexpr int32_t NUM4 = 4;
70+ 
71+ constexpr int32_t NUM64 = 64;
72+ constexpr int32_t NUM512 = 512;
73+ constexpr int32_t NUM576 = 576;
74+ constexpr uint32_t BASIC_BLOCK_SIZE = 128;
75+ constexpr int32_t Q_BLK = 128;
76+ constexpr int32_t MAX_STACK_LEN = 512;
77+ constexpr int32_t PRE_LAUNCH = 2;
78+ 
79+ constexpr uint32_t FLOAT_VECTOR_SIZE = 64;
80+ 
81+ constexpr uint32_t UNIT_BLOCK_STACK_NUM = 4;
82+ 
83+ constexpr uint32_t Q_TILE_CEIL = 128;
84+ constexpr uint32_t MAX_KV_STACK_LEN = 512; //可配置1024或512
85+ 
86+ template <typename T>
87+ __aicore__ inline T AlignUp(T a, T b)
88+ {
89+ return (b == 0) ? 0 : (a + b - 1) / b * b;
90+ }
91+ 
92+ template <typename T>
93+ __aicore__ inline T Min(T a, T b)
94+ {
95+ return (a > b) ? b : a;
96+ }
97+ 
98+ enum class cvPipeLineType {
99+ FAI_COMMON_NORMAL = 0,
100+ FAI_COMMON_CHUNK_MASK = 1,
101+ FAI_SPARSE_BLOCK = 2
102+ };
103+ 
104+ enum class SparseMaskType {
105+ NO_MASK = 0,
106+ MASK_SPEC = 1,
107+ MASK_CAUSUAL = 2,
108+ SPARSE_BLOCK = 3
109+ };
110+ 
111+ __aicore__ inline
112+ uint32_t GetQNBlockTile(uint32_t qSeqlen, uint32_t groupSize)
113+ {
114+ uint32_t qNBlockTile = 1;
115+ return qNBlockTile;
116+ }
117+ 
118+ __aicore__ inline
119+ uint32_t GetQSBlockTile(uint32_t kvSeqlen)
120+ {
121+ uint32_t qSBlockTile = Q_BLK;
122+ return qSBlockTile;
123+ }
124+ 
125+ __aicore__ inline
126+ uint32_t GetQBlocks(int32_t qseqlen, int32_t x)
127+ {
128+ uint32_t qBlocksInX = (x + BASIC_BLOCK_SIZE - 1) / BASIC_BLOCK_SIZE;
129+ uint32_t completeXBlocks = x != 0 ? qseqlen / x : qseqlen / BASIC_BLOCK_SIZE;
130+ uint32_t remainingSeqlen = x != 0 ? qseqlen - completeXBlocks * x : qseqlen % BASIC_BLOCK_SIZE;
131+ uint32_t remainingBlocks = (remainingSeqlen + BASIC_BLOCK_SIZE - 1) / BASIC_BLOCK_SIZE;
132+ return qBlocksInX * completeXBlocks + remainingBlocks;
133+ }
134+ 
135+ 
136+ // EagleQuantBlockSparseAttention Kernel Parameters
137+ struct EagleQuantBlockSparseAttentionKernelParams {
138+ // 输入张量
139+ GM_ADDR q; // Query张量
140+ GM_ADDR k; // Key张量
141+ GM_ADDR v; // Value张量
142+ GM_ADDR blockSparseMask; // 适配新增的张量
143+ GM_ADDR mask; // 掩码张量
144+ GM_ADDR blockTables; // 块表
145+ GM_ADDR actualQseqlen; // 实际Q序列长度
146+ GM_ADDR actualKvseqlen; // 实际KV序列长度
147+
148+
149+ // 输出和工作空间
150+ GM_ADDR o; // 输出张量
151+ GM_ADDR lse;
152+ GM_ADDR workspace;
153+ GM_ADDR tiling; // Tiling数据
154+ 
155+ // 默认构造函数
156+ __aicore__ inline
157+ EagleQuantBlockSparseAttentionKernelParams() {}
158+ 
159+ // 带参数的构造函数(新增了参数的构造器)
160+ __aicore__ inline
161+ EagleQuantBlockSparseAttentionKernelParams(
162+ GM_ADDR q_, GM_ADDR k_, GM_ADDR v_, GM_ADDR blockSparseMask_, GM_ADDR mask_, GM_ADDR blockTables_,
163+ GM_ADDR actualQseqlen_, GM_ADDR actualKvseqlen_,
164+ GM_ADDR o_, GM_ADDR lse_, GM_ADDR workspace_,
165+ GM_ADDR tiling_)
166+ : q(q_), k(k_), v(v_), blockSparseMask(blockSparseMask_), mask(mask_), blockTables(blockTables_),
167+ actualQseqlen(actualQseqlen_), actualKvseqlen(actualKvseqlen_),
168+ o(o_),
169+ lse(lse_),
170+ workspace(workspace_),
171+ tiling(tiling_)
172+ {}
173+ };
174+ 
175+ // 为了兼容性，保留旧名称的别名（逐步废弃）
176+ using FASparseKernelParams = EagleQuantBlockSparseAttentionKernelParams;
177+}
178+#endif // EAGLE_QUANT_BLOCK_SPARSE_ATTENTION_KERNEL_COMMON_HPP
179+ 
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/tla/int_tuple.hpp+273-0
@@ -0,0 +1,273 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 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+#ifndef TLA_INT_TUPLE_HPP
12+#define TLA_INT_TUPLE_HPP
13+ 
14+#include "../tla/type_traits.hpp"
15+#include "../tla/tuple.hpp"
16+#include "../tla/numeric/integral_constant.hpp"
17+#include "../tla/numeric/integer_sequence.hpp"
18+ 
19+namespace tla {
20+ 
21+namespace detail {
22+ 
23+template <class T, class F, int... I>
24+HOST_DEVICE constexpr
25+auto apply(T&& t, F&& f, seq<I...>)
26+{
27+ return f(get<I>(static_cast<T&&>(t))...);
28+}
29+ 
30+template <class T, class F, class G, int... I>
31+HOST_DEVICE constexpr
32+auto tapply(T&& t, F&& f, G&& g, seq<I...>)
33+{
34+ return g(f(get<I>(static_cast<T&&>(t)))...);
35+}
36+ 
37+template <class T0, class T1, class F, class G, int... I>
38+HOST_DEVICE constexpr
39+auto tapply(T0&& t0, T1&& t1, F&& f, G&& g, seq<I...>)
40+{
41+ return g(f(get<I>(static_cast<T0&&>(t0)),
42+ get<I>(static_cast<T1&&>(t1)))...);
43+}
44+ 
45+} // end namespace detail
46+ 
47+template <class T, class F>
48+HOST_DEVICE constexpr
49+auto apply(T&& t, F&& f)
50+{
51+ return detail::apply(static_cast<T&&>(t), f, tuple_seq<T>{});
52+}
53+ 
54+template <class T, class F, class G>
55+HOST_DEVICE constexpr
56+auto transform_apply(T&& t, F&& f, G&& g)
57+{
58+ if constexpr (is_tuple<remove_cvref_t<T>>::value) {
59+ return detail::tapply(static_cast<T&&>(t), f, g, tuple_seq<T>{});
60+ } else {
61+ return g(f(static_cast<T&&>(t)));
62+ }
63+}
64+ 
65+template <class T0, class T1, class F, class G>
66+HOST_DEVICE constexpr
67+auto transform_apply(T0&& t0, T1&& t1, F&& f, G&& g)
68+{
69+ if constexpr (is_tuple<remove_cvref_t<T0>>::value) {
70+ return detail::tapply(static_cast<T0&&>(t0), static_cast<T1&&>(t1), f, g, tuple_seq<T0>{});
71+ } else {
72+ return g(f(static_cast<T0&&>(t0), static_cast<T1&&>(t1)));
73+ }
74+}
75+ 
76+template <class T, class F>
77+HOST_DEVICE constexpr
78+void for_each(T&& t, F&& f)
79+{
80+ if constexpr (is_tuple<remove_cvref_t<T>>::value) {
81+ return detail::apply(t, [&](auto&&... a) { (f(static_cast<decltype(a)&&>(a)), ...); }, tuple_seq<T>{});
82+ } else {
83+ return f(static_cast<T&&>(t));
84+ }
85+}
86+ 
87+struct UnpackedMakeTuple {
88+ template <class... T>
89+ HOST_DEVICE constexpr
90+ auto operator()(T const&... a) const {
91+ return tla::MakeTuple(a...);
92+ }
93+};
94+ 
95+template <class T0, class T1, class F>
96+HOST_DEVICE constexpr
97+auto transform(T0 const& t0, T1 const& t1, F&& f)
98+{
99+ if constexpr (is_tuple<T0>::value) {
100+ static_assert(tuple_size<T0>::value == tuple_size<T1>::value, "Mismatched tuple_size");
101+ return detail::tapply(t0, t1, f, UnpackedMakeTuple{}, tuple_seq<T0>{});
102+ } else {
103+ return f(t0, t1);
104+ }
105+}
106+ 
107+template <size_t I, class T,
108+ TLA_REQUIRES(tla::is_integral<tla::remove_cvref_t<T>>::value)>
109+HOST_DEVICE constexpr
110+decltype(auto) get(T&& t) noexcept
111+{
112+ static_assert(I == 0, "Index out of range");
113+ return static_cast<T&&>(t);
114+}
115+ 
116+template <size_t I0, size_t I1, size_t... Is, class T>
117+HOST_DEVICE constexpr
118+decltype(auto) get(T&& t) noexcept
119+{
120+ return get<I1, Is...>(get<I0>(static_cast<T&&>(t)));
121+}
122+ 
123+// max
124+template <class T0, class... Ts>
125+HOST_DEVICE constexpr
126+auto max(T0 const& t0, Ts const&... ts);
127+ 
128+struct UnpackedMax {
129+ template <class... T>
130+ HOST_DEVICE constexpr
131+ auto operator()(T const&... v) const {
132+ return tla::max(v...);
133+ }
134+};
135+ 
136+template <class T0, class... Ts>
137+HOST_DEVICE constexpr
138+auto max(T0 const& t0, Ts const&... ts)
139+{
140+ if constexpr (is_tuple<T0>::value) {
141+ return tla::max(tla::apply(t0, UnpackedMax{}), ts...);
142+ } else if constexpr (sizeof...(Ts) == 0) {
143+ return t0;
144+ } else {
145+ return tla::max(t0, tla::max(ts...));
146+ }
147+}
148+ 
149+// rank
150+template <int... Is, class Tuple>
151+HOST_DEVICE constexpr
152+auto rank(Tuple const& t)
153+{
154+ if constexpr (sizeof...(Is) == 0) {
155+ if constexpr (is_tuple<Tuple>::value) {
156+ return Int<tuple_size<Tuple>::value>{};
157+ } else {
158+ return Int<1>{};
159+ }
160+ } else {
161+ return rank(get<Is...>(t));
162+ }
163+}
164+ 
165+template <class Tuple>
166+using rank_t = decltype(rank(std::declval<Tuple>()));
167+ 
168+template <class Tuple>
169+constexpr auto rank_v = rank_t<Tuple>::value;
170+ 
171+// depth
172+template <int... Is, class Tuple>
173+HOST_DEVICE constexpr
174+auto depth(Tuple const& t);
175+ 
176+struct UnpackedDepth {
177+ template <class... T>
178+ HOST_DEVICE constexpr
179+ auto operator()(T const&... v) const {
180+ return tla::max(depth(v)...);
181+ }
182+};
183+ 
184+template <int... Is, class Tuple>
185+HOST_DEVICE constexpr
186+auto depth(Tuple const& t)
187+{
188+ if constexpr (sizeof...(Is) == 0) {
189+ if constexpr (is_tuple<Tuple>::value) {
190+ return Int<1>{} + tla::apply(t, UnpackedDepth{});
191+ } else {
192+ return Int<0>{};
193+ }
194+ } else {
195+ return depth(get<Is...>(t));
196+ }
197+}
198+ 
199+template <class Tuple>
200+using depth_t = decltype(depth(std::declval<Tuple>()));
201+ 
202+template <class Tuple>
203+constexpr auto depth_v = depth_t<Tuple>::value;
204+ 
205+struct MultipliesUnaryLfold {
206+ template <class... T>
207+ HOST_DEVICE constexpr
208+ auto operator()(T const&... v) const {
209+ return (... * v);
210+ }
211+};
212+ 
213+// Implementation of product as a function object
214+struct Product {
215+ template <class IntTuple>
216+ HOST_DEVICE constexpr
217+ auto operator()(IntTuple const& a) const
218+ {
219+ if constexpr (is_tuple<IntTuple>::value) {
220+ if constexpr (tuple_size<IntTuple>::value == 0) {
221+ return Int<1>{};
222+ } else {
223+ return tla::transform_apply(a, Product{}, MultipliesUnaryLfold{});
224+ }
225+ } else if constexpr (tla::is_integral<IntTuple>::value) {
226+ return a;
227+ }
228+ }
229+};
230+ 
231+namespace detail {
232+ 
233+template <size_t N, typename Sequence>
234+struct MakeZeroTupleImpl;
235+ 
236+template <size_t N, size_t... Is>
237+struct MakeZeroTupleImpl<N, tla::index_sequence<Is...>> {
238+ using type = tla::tuple<tla::Int<Is*0>...>;
239+};
240+ 
241+template <size_t N>
242+using MakeZeroTuple = typename MakeZeroTupleImpl<N, tla::make_index_sequence<N>>::type;
243+ 
244+} // end namespace detail
245+ 
246+// Add
247+template <class IntTupleA, class IntTupleB>
248+HOST_DEVICE constexpr
249+auto Add(IntTupleA const& a, IntTupleB const& b);
250+ 
251+struct UnpackedAdd {
252+ template <class IntTupleA, class IntTupleB>
253+ HOST_DEVICE constexpr
254+ auto operator()(IntTupleA const& x, IntTupleB const& y) const {
255+ return Add(x, y);
256+ }
257+};
258+ 
259+template <class IntTupleA, class IntTupleB>
260+HOST_DEVICE constexpr
261+auto Add(IntTupleA const& a, IntTupleB const& b)
262+{
263+ if constexpr (is_tuple<IntTupleA>::value && is_tuple<IntTupleB>::value) {
264+ static_assert(tuple_size<IntTupleA>::value == tuple_size<IntTupleB>::value, "Mismatched ranks");
265+ return transform(a, b, UnpackedAdd{});
266+ } else {
267+ return tla::add(a, b);
268+ }
269+}
270+ 
271+} // end namespace tla
272+ 
273+#endif // TLA_INT_TUPLE_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/tla/layout.hpp+612-0
@@ -0,0 +1,612 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 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+#ifndef TLA_LAYOUT_HPP
12+#define TLA_LAYOUT_HPP
13+ 
14+#include "../attn_infra/base_defs.hpp"
15+#include "../tla/numeric/integral_constant.hpp"
16+#include "../tla/tuple.hpp"
17+#include "../tla/int_tuple.hpp"
18+#include "../attn_infra/layout/layout.hpp"
19+ 
20+namespace tla {
21+ 
22+// Aliases
23+ 
24+template <class... Shapes>
25+using Shape = tla::tuple<Shapes...>;
26+ 
27+template <class... Strides>
28+using Stride = tla::tuple<Strides...>;
29+ 
30+template <class... Coords>
31+using Coord = tla::tuple<Coords...>;
32+ 
33+template <class... Ts>
34+HOST_DEVICE constexpr
35+Shape<Ts...> MakeShape(Ts const&... t) {
36+ return {t...};
37+}
38+template <class... Ts>
39+HOST_DEVICE constexpr
40+Stride<Ts...> MakeStride(Ts const&... t) {
41+ return {t...};
42+}
43+template <class... Ts>
44+HOST_DEVICE constexpr
45+Coord<Ts...> MakeCoord(Ts const&... t) {
46+ return {t...};
47+}
48+ 
49+//
50+// Layout
51+//
52+ 
53+template <class Shape, class Stride>
54+struct Layout : private tla::tuple<Shape, Stride> {
55+ // NOTE: This defaults static Shapes/Strides correctly, but not dynamic
56+ HOST_DEVICE constexpr
57+ Layout(Shape const& shape = {}, Stride const& stride = {})
58+ : tla::tuple<Shape, Stride>(shape, stride) {}
59+ 
60+ //
61+ // Accessors
62+ //
63+ 
64+ static constexpr int rank = rank_v<Stride>;
65+ static constexpr int depth = depth_v<Stride>;
66+ 
67+ template <int... I>
68+ HOST_DEVICE constexpr
69+ decltype(auto) shape()
70+ {
71+ return get<0, I...>(static_cast<tla::tuple<Shape, Stride>&>(*this));
72+ }
73+ 
74+ template <int... I>
75+ HOST_DEVICE constexpr
76+ decltype(auto) shape() const
77+ {
78+ return get<0, I...>(static_cast<tla::tuple<Shape, Stride> const&>(*this));
79+ }
80+ 
81+ template <int... I>
82+ HOST_DEVICE constexpr
83+ decltype(auto) stride()
84+ {
85+ return get<1, I...>(static_cast<tla::tuple<Shape, Stride>&>(*this));
86+ }
87+ 
88+ template <int... I>
89+ HOST_DEVICE constexpr
90+ decltype(auto) stride() const
91+ {
92+ return get<1, I...>(static_cast<tla::tuple<Shape, Stride> const&>(*this));
93+ }
94+ 
95+ template <class Coord>
96+ HOST_DEVICE constexpr
97+ auto operator()(Coord const& coord) const
98+ {
99+ return crd2offset(coord, shape(), stride());
100+ }
101+};
102+ 
103+// Layout construction
104+ 
105+template <class Shape, class Stride>
106+HOST_DEVICE constexpr
107+auto MakeLayout(Shape const& shape, Stride const& stride)
108+{
109+ static_assert(is_tuple<Shape>::value || is_integral<Shape>::value);
110+ static_assert(is_tuple<Stride>::value || is_integral<Stride>::value);
111+ return Layout<Shape, Stride>(shape, stride);
112+}
113+ 
114+// Convenience tags for common layouts
115+ 
116+template <class LayoutTag>
117+HOST_DEVICE constexpr
118+auto MakeLayoutFromTag(LayoutTag const& tag)
119+{
120+ static_assert(std::is_same_v<LayoutTag, NpuArch::layout::RowMajor> ||
121+ std::is_same_v<LayoutTag, NpuArch::layout::ColumnMajor> ||
122+ std::is_same_v<LayoutTag, NpuArch::layout::zN> ||
123+ std::is_same_v<LayoutTag, NpuArch::layout::nZ>,
124+ "Unsupported LayoutTag for MakeLayoutFromTag, only support NpuArch::layout::RowMajor or"
125+ "NpuArch::layout::ColumnMajor or NpuArch::layout::zN or NpuArch::layout::nZ");
126+ 
127+ if constexpr (std::is_same_v<LayoutTag, NpuArch::layout::RowMajor>) {
128+ return MakeLayout(MakeShape(tag.shape(0), tag.shape(1)), MakeStride(tag.stride(0), Int<1>{}));
129+ } else if constexpr (std::is_same_v<LayoutTag, NpuArch::layout::ColumnMajor>) {
130+ return MakeLayout(MakeShape(tag.shape(0), tag.shape(1)), MakeStride(Int<1>{}, tag.stride(1)));
131+ } else { // zN or nZ
132+ return MakeLayout(MakeShape(MakeShape(tag.shape(0), tag.shape(1)), MakeShape(tag.shape(2), tag.shape(3))),
133+ MakeStride(MakeStride(tag.stride(0), tag.stride(1)), MakeStride(tag.stride(2), tag.stride(3))));
134+ }
135+}
136+ 
137+// Return the shape of a mode
138+template <int... Is, class Shape, class Stride>
139+HOST_DEVICE constexpr
140+decltype(auto) shape(Layout<Shape, Stride>& layout)
141+{
142+ return layout.template shape<Is...>();
143+}
144+ 
145+template <int... Is, class Shape, class Stride>
146+HOST_DEVICE constexpr
147+decltype(auto) shape(Layout<Shape, Stride> const& layout)
148+{
149+ return layout.template shape<Is...>();
150+}
151+ 
152+// Return the stride of a mode
153+template <int... Is, class Shape, class Stride>
154+HOST_DEVICE constexpr
155+decltype(auto) stride(Layout<Shape, Stride>& layout)
156+{
157+ return layout.template stride<Is...>();
158+}
159+ 
160+template <int... Is, class Shape, class Stride>
161+HOST_DEVICE constexpr
162+decltype(auto) stride(Layout<Shape, Stride> const& layout)
163+{
164+ return layout.template stride<Is...>();
165+}
166+ 
167+// Return the rank of layout
168+template <int... Is, class Shape, class Stride>
169+HOST_DEVICE constexpr
170+auto rank(Layout<Shape, Stride> const& layout)
171+{
172+ return rank(shape<Is...>(layout));
173+}
174+ 
175+// Return the depth of the layout
176+template <int... Is, class Shape, class Stride>
177+HOST_DEVICE constexpr
178+auto depth(Layout<Shape, Stride> const& layout)
179+{
180+ return depth(shape<Is...>(layout));
181+}
182+ 
183+// Return the offset of coord
184+template <class Coord, class Shape, class Stride>
185+HOST_DEVICE constexpr
186+auto crd2offset(Coord const& coord, Shape const& shape, Stride const& stride);
187+ 
188+namespace detail {
189+ 
190+template <class Coord, class Shape, class Stride, int... Is>
191+HOST_DEVICE constexpr
192+auto crd2offset_ttt(Coord const& coord, Shape const& shape, Stride const& stride, seq<Is...>)
193+{
194+ return (... + crd2offset(get<Is>(coord), get<Is>(shape), get<Is>(stride)));
195+}
196+ 
197+template <class CInt, class STuple, class DTuple, int I0, int... Is>
198+HOST_DEVICE constexpr
199+auto crd2offset_itt(CInt const& coord, STuple const& shape, DTuple const& stride, seq<I0, Is...>)
200+{
201+ if constexpr (sizeof...(Is) == 0) { // Avoid recursion and mod on single/last iter
202+ return crd2offset(coord, get<I0>(shape), get<I0>(stride));
203+ } else if constexpr (is_constant<0, CInt>::value) {
204+ return crd2offset(_0{}, get<I0>(shape), get<I0>(stride)) +
205+ (_0{} + ... + crd2offset(_0{}, get<Is>(shape), get<Is>(stride)));
206+ } else { // General case
207+ return crd2offset(coord % Product{}(get<I0>(shape)), get<I0>(shape), get<I0>(stride)) +
208+ crd2offset_itt(coord / Product{}(get<I0>(shape)), shape, stride, seq<Is...>{});
209+ }
210+}
211+ 
212+} // end namespace detail
213+ 
214+template <class Coord, class Shape, class Stride>
215+HOST_DEVICE constexpr
216+auto crd2offset(Coord const& coord, Shape const& shape, Stride const& stride)
217+{
218+ if constexpr (is_tuple<Coord>::value) {
219+ if constexpr (is_tuple<Shape>::value) { // tuple tuple tuple
220+ static_assert(tuple_size<Coord>::value == tuple_size<Shape>::value, "Mismatched Ranks");
221+ static_assert(tuple_size<Coord>::value == tuple_size<Stride>::value, "Mismatched Ranks");
222+ return detail::crd2offset_ttt(coord, shape, stride, tuple_seq<Coord>{});
223+ } else { // tuple "int" "int"
224+ static_assert(sizeof(Coord) == 0, "Invalid parameters");
225+ }
226+ } else {
227+ if constexpr (is_tuple<Shape>::value) { // "int" tuple tuple
228+ static_assert(tuple_size<Shape>::value == tuple_size<Stride>::value, "Mismatched Ranks");
229+ return detail::crd2offset_itt(coord, shape, stride, tuple_seq<Shape>{});
230+ } else { // "int" "int" "int"
231+ return coord * stride;
232+ }
233+ }
234+}
235+ 
236+template <class Layout>
237+struct is_layout : false_type {};
238+template <class Shape, class Stride>
239+struct is_layout<Layout<Shape, Stride>> : true_type {};
240+ 
241+// Layout Check
242+namespace detail {
243+ 
244+template <class Layout, class Enable = void>
245+struct isVector {
246+ static bool const value = false;
247+};
248+ 
249+template <class Layout>
250+struct isVector<Layout, std::enable_if_t<Layout::depth == 1 && Layout::rank == 1>> {
251+ static bool const value = (stride<0>(Layout{}) == 1);
252+};
253+ 
254+template <class Layout, class Enable = void>
255+struct isRowMajor {
256+ static bool const value = false;
257+};
258+ 
259+template <class Layout>
260+struct isRowMajor<Layout, std::enable_if_t<Layout::depth == 1 && Layout::rank == 2>> {
261+ static bool const value = (stride<1>(Layout{}) == 1);
262+};
263+ 
264+template <class Layout, class Enable = void>
265+struct isColumnMajor {
266+ static bool const value = false;
267+};
268+ 
269+template <class Layout>
270+struct isColumnMajor<Layout, std::enable_if_t<Layout::depth == 1 && Layout::rank == 2>> {
271+ static bool const value = (stride<0>(Layout{}) == 1);
272+};
273+ 
274+template <class Element, class Layout, class Enable1 = void, class Enable2 = void>
275+struct iszN {
276+ static bool const value = false;
277+};
278+ 
279+template <class Element, class Layout>
280+struct iszN<Element, Layout,
281+ std::enable_if_t<Layout::depth == 2 && Layout::rank == 2>, std::enable_if_t<rank_v<decltype(shape<0>(Layout{}))> == 2 &&
282+ rank_v<decltype(shape<1>(Layout{}))> == 2>> {
283+ static constexpr uint32_t ELE_NUM_PER_C0 = NpuArch::BYTE_PER_C0 / sizeof(Element);
284+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = NpuArch::BYTE_PER_FRACTAL / sizeof(Element);
285+ static bool const value = (shape<0, 0>(Layout{}) == NpuArch::C0_NUM_PER_FRACTAL &&
286+ shape<1, 0>(Layout{}) == ELE_NUM_PER_C0 &&
287+ stride<1, 0>(Layout{}) == 1 &&
288+ stride<0, 1>(Layout{}) == ELE_NUM_PER_FRACTAL);
289+};
290+ 
291+template <class Element, class Layout, class Enable1 = void, class Enable2 = void>
292+struct iszZ {
293+ static bool const value = false;
294+};
295+ 
296+template <class Element, class Layout>
297+struct iszZ<Element, Layout,
298+ std::enable_if_t<Layout::depth == 2 && Layout::rank == 2>, std::enable_if_t<rank_v<decltype(shape<0>(Layout{}))> == 2 &&
299+ rank_v<decltype(shape<1>(Layout{}))> == 2>> {
300+ static constexpr uint32_t ELE_NUM_PER_C0 = NpuArch::BYTE_PER_C0 / sizeof(Element);
301+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = NpuArch::BYTE_PER_FRACTAL / sizeof(Element);
302+ static bool const value = (shape<0, 0>(Layout{}) == NpuArch::C0_NUM_PER_FRACTAL &&
303+ shape<1, 0>(Layout{}) == ELE_NUM_PER_C0 &&
304+ stride<1, 0>(Layout{}) == 1 &&
305+ stride<1, 1>(Layout{}) == ELE_NUM_PER_FRACTAL);
306+};
307+ 
308+template <class Element, class Layout, class Enable1 = void, class Enable2 = void>
309+struct isnZ {
310+ static bool const value = false;
311+};
312+ 
313+template <class Element, class Layout>
314+struct isnZ<Element, Layout,
315+ std::enable_if_t<Layout::depth == 2 && Layout::rank == 2>, std::enable_if_t<rank_v<decltype(shape<0>(Layout{}))> == 2 &&
316+ rank_v<decltype(shape<1>(Layout{}))> == 2>> {
317+ static constexpr uint32_t ELE_NUM_PER_C0 = NpuArch::BYTE_PER_C0 / sizeof(Element);
318+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = NpuArch::BYTE_PER_FRACTAL / sizeof(Element);
319+ static bool const value = (shape<0, 0>(Layout{}) == ELE_NUM_PER_C0 &&
320+ shape<1, 0>(Layout{}) == NpuArch::C0_NUM_PER_FRACTAL &&
321+ stride<0, 0>(Layout{}) == 1 &&
322+ stride<1, 1>(Layout{}) == ELE_NUM_PER_FRACTAL);
323+};
324+ 
325+#if defined(CATLASS_ARCH_A5_ENABLED)
326+template <class Element, class Layout, class Enable1 = void, class Enable2 = void>
327+struct isMxScaleANoTrans {
328+ static bool const value = false;
329+};
330+ 
331+template <class Layout>
332+struct isMxScaleANoTrans<AscendC::fp8_e8m0_t, Layout,
333+ std::enable_if_t<Layout::depth == 2 && Layout::rank == 2>, std::enable_if_t<rank_v<decltype(shape<0>(Layout{}))> == 1 &&
334+ rank_v<decltype(shape<1>(Layout{}))> == 2>> {
335+ static constexpr uint32_t ELE_NUM_PER_C0 = 2;
336+ static bool const value =
337+ (shape<1, 0>(Layout{}) == ELE_NUM_PER_C0 && stride<1, 0>(Layout{}) == 1 &&
338+ stride<1, 1>(Layout{}) == ELE_NUM_PER_C0);
339+};
340+ 
341+template <class Element, class Layout, class Enable1 = void, class Enable2 = void>
342+struct isMxScaleATrans {
343+ static bool const value = false;
344+};
345+ 
346+template <class Layout>
347+struct isMxScaleATrans<AscendC::fp8_e8m0_t, Layout,
348+ std::enable_if_t<Layout::depth == 2 && Layout::rank == 2>, std::enable_if_t<rank_v<decltype(shape<0>(Layout{}))> == 1 &&
349+ rank_v<decltype(shape<1>(Layout{}))> == 2>> {
350+ static constexpr uint32_t ELE_NUM_PER_C0 = 2;
351+ static bool const value =
352+ (shape<1, 0>(Layout{}) == ELE_NUM_PER_C0 && stride<1, 0>(Layout{}) == 1 &&
353+ stride<0>(Layout{}) == ELE_NUM_PER_C0);
354+};
355+ 
356+template <class Element, class Layout, class Enable1 = void, class Enable2 = void>
357+struct isMxScaleBNoTrans {
358+ static bool const value = false;
359+};
360+ 
361+template <class Layout>
362+struct isMxScaleBNoTrans<AscendC::fp8_e8m0_t, Layout,
363+ std::enable_if_t<Layout::depth == 2 && Layout::rank == 2>, std::enable_if_t<rank_v<decltype(shape<0>(Layout{}))> == 2 &&
364+ rank_v<decltype(shape<1>(Layout{}))> == 1>> {
365+ static constexpr uint32_t ELE_NUM_PER_C0 = 2;
366+ static bool const value =
367+ (shape<0, 0>(Layout{}) == ELE_NUM_PER_C0 && stride<0, 0>(Layout{}) == 1 &&
368+ stride<1>(Layout{}) == ELE_NUM_PER_C0);
369+};
370+ 
371+template <class Element, class Layout, class Enable1 = void, class Enable2 = void>
372+struct isMxScaleBTrans {
373+ static bool const value = false;
374+};
375+ 
376+template <class Layout>
377+struct isMxScaleBTrans<AscendC::fp8_e8m0_t, Layout,
378+ std::enable_if_t<Layout::depth == 2 && Layout::rank == 2>, std::enable_if_t<rank_v<decltype(shape<0>(Layout{}))> == 2 &&
379+ rank_v<decltype(shape<1>(Layout{}))> == 1>> {
380+ static constexpr uint32_t ELE_NUM_PER_C0 = 2;
381+ static bool const value =
382+ (shape<0, 0>(Layout{}) == ELE_NUM_PER_C0 && stride<0, 0>(Layout{}) == 1 &&
383+ stride<0, 1>(Layout{}) == ELE_NUM_PER_C0);
384+};
385+ 
386+template <class Element, class Layout, class Enable1 = void, class Enable2 = void>
387+struct isMxScalezZ {
388+ static bool const value = false;
389+};
390+ 
391+template <class Layout>
392+struct isMxScalezZ<AscendC::fp8_e8m0_t, Layout,
393+ std::enable_if_t<Layout::depth == 2 && Layout::rank == 2>, std::enable_if_t<rank_v<decltype(shape<0>(Layout{}))> == 2 &&
394+ rank_v<decltype(shape<1>(Layout{}))> == 2>> {
395+ static constexpr uint32_t ELE_NUM_PER_C0 = 2;
396+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = 32;
397+ static bool const value = (shape<0, 0>(Layout{}) == NpuArch::C0_NUM_PER_FRACTAL &&
398+ shape<1, 0>(Layout{}) == ELE_NUM_PER_C0 &&
399+ stride<1, 0>(Layout{}) == 1 &&
400+ stride<1, 1>(Layout{}) == ELE_NUM_PER_FRACTAL);
401+};
402+ 
403+template <class Element, class Layout, class Enable1 = void, class Enable2 = void>
404+struct isMxScalenN {
405+ static bool const value = false;
406+};
407+ 
408+template <class Layout>
409+struct isMxScalenN<AscendC::fp8_e8m0_t, Layout,
410+ std::enable_if_t<Layout::depth == 2 && Layout::rank == 2>, std::enable_if_t<rank_v<decltype(shape<0>(Layout{}))> == 2 &&
411+ rank_v<decltype(shape<1>(Layout{}))> == 2>> {
412+ static constexpr uint32_t ELE_NUM_PER_C0 = 2;
413+ static constexpr uint32_t ELE_NUM_PER_FRACTAL = 32;
414+ static bool const value = (shape<0, 0>(Layout{}) == ELE_NUM_PER_C0 &&
415+ shape<1, 0>(Layout{}) == NpuArch::C0_NUM_PER_FRACTAL &&
416+ stride<0, 0>(Layout{}) == 1 &&
417+ stride<0, 1>(Layout{}) == ELE_NUM_PER_FRACTAL);
418+};
419+#endif
420+ 
421+} // end namespace detail
422+ 
423+// Advanced Layout constructions
424+// Make a vector layout.
425+template <class T>
426+HOST_DEVICE constexpr
427+auto MakeLayout(T const& len)
428+{
429+ return MakeLayout(MakeShape(len), MakeStride(Int<1>{}));
430+}
431+ 
432+// Make a inner layout with Rows and Cols.
433+template <class Element, class LayoutTag, class T, class U>
434+HOST_DEVICE constexpr
435+auto MakeLayout(T const& rows, U const& cols)
436+{
437+ static_assert(std::is_same_v<LayoutTag, NpuArch::layout::RowMajor> ||
438+ std::is_same_v<LayoutTag, NpuArch::layout::ColumnMajor> ||
439+ std::is_same_v<LayoutTag, NpuArch::layout::VectorLayout> ||
440+ std::is_same_v<LayoutTag, NpuArch::layout::zN> ||
441+ std::is_same_v<LayoutTag, NpuArch::layout::nZ> ||
442+ std::is_same_v<LayoutTag, NpuArch::layout::zZ>,
443+ "Unsupported LayoutTag for MakeLayoutFromTag, only support NpuArch::layout::RowMajor or"
444+ "NpuArch::layout::ColumnMajor or NpuArch::layout::zN or NpuArch::layout::nZ or NpuArch::layout::zZ");
445+ 
446+ constexpr uint32_t ELE_NUM_PER_C0 = NpuArch::BYTE_PER_C0 / sizeof(Element);
447+ constexpr uint32_t ELE_NUM_PER_FRACTAL = NpuArch::BYTE_PER_FRACTAL / sizeof(Element);
448+ 
449+ if constexpr (std::is_same_v<LayoutTag, NpuArch::layout::VectorLayout>) {
450+ return MakeLayout(MakeShape(cols), MakeStride(Int<1>{}));
451+ } else if constexpr (std::is_same_v<LayoutTag, NpuArch::layout::RowMajor>) {
452+ return MakeLayout(MakeShape(rows, cols), MakeStride((int64_t)cols, Int<1>{}));
453+ } else if constexpr (std::is_same_v<LayoutTag, NpuArch::layout::ColumnMajor>) {
454+ return MakeLayout(MakeShape(rows, cols), MakeStride(Int<1>{}, (int64_t)rows));
455+ } else if constexpr (std::is_same_v<LayoutTag, NpuArch::layout::zN>) {
456+ return MakeLayout(
457+ MakeShape(MakeShape(Int<NpuArch::C0_NUM_PER_FRACTAL>{}, CeilDiv(rows, Int<NpuArch::C0_NUM_PER_FRACTAL>{})),
458+ MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(cols, Int<ELE_NUM_PER_C0>{}))),
459+ MakeStride(MakeStride(Int<ELE_NUM_PER_C0>{}, Int<ELE_NUM_PER_FRACTAL>{}),
460+ MakeStride(Int<1>{}, RoundUp((int64_t)rows, Int<NpuArch::C0_NUM_PER_FRACTAL>{}) * ELE_NUM_PER_C0)));
461+ } else if constexpr (std::is_same_v<LayoutTag, NpuArch::layout::zZ>) {
462+ return MakeLayout(
463+ MakeShape(MakeShape(Int<NpuArch::C0_NUM_PER_FRACTAL>{}, CeilDiv(rows, Int<NpuArch::C0_NUM_PER_FRACTAL>{})),
464+ MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(cols, Int<ELE_NUM_PER_C0>{}))),
465+ MakeStride(MakeStride(Int<ELE_NUM_PER_C0>{},
466+ RoundUp((int64_t)cols, Int<ELE_NUM_PER_C0>{}) * NpuArch::C0_NUM_PER_FRACTAL),
467+ MakeStride(Int<1>{}, Int<ELE_NUM_PER_FRACTAL>{})));
468+ } else {
469+ return MakeLayout(
470+ MakeShape(MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(rows, Int<ELE_NUM_PER_C0>{})),
471+ MakeShape(Int<NpuArch::C0_NUM_PER_FRACTAL>{}, CeilDiv(cols, Int<NpuArch::C0_NUM_PER_FRACTAL>{}))),
472+ MakeStride(
473+ MakeStride(Int<1>{}, RoundUp((int64_t)cols, Int<NpuArch::C0_NUM_PER_FRACTAL>{}) * ELE_NUM_PER_C0),
474+ MakeStride(Int<ELE_NUM_PER_C0>{}, Int<ELE_NUM_PER_FRACTAL>{})));
475+ }
476+}
477+ 
478+// Make a MxScale layout with Rows and Cols.
479+template <class Element, class LayoutTag, bool isMxScaleB, class T, class U>
480+HOST_DEVICE constexpr
481+auto MakeMxScaleLayout(T const& rows, U const& cols)
482+{
483+ static_assert(
484+ std::is_same_v<Element, AscendC::fp8_e8m0_t> &&
485+ (std::is_same_v<LayoutTag, NpuArch::layout::RowMajor> ||
486+ std::is_same_v<LayoutTag, NpuArch::layout::ColumnMajor> ||
487+ std::is_same_v<LayoutTag, NpuArch::layout::zZ> || std::is_same_v<LayoutTag, NpuArch::layout::nN>),
488+ "only support RowMajor, ColumnMajor, zZ, nN in fp8_e8m0_t dtype"
489+ );
490+ 
491+ constexpr uint32_t ELE_NUM_PER_C0 = 2;
492+ constexpr uint32_t ELE_NUM_PER_FRACTAL = 32;
493+ 
494+ if constexpr (std::is_same_v<LayoutTag, NpuArch::layout::RowMajor>) {
495+ if constexpr (!isMxScaleB) {
496+ return MakeLayout(
497+ MakeShape(rows, MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(cols, Int<ELE_NUM_PER_C0>{}))),
498+ MakeStride(RoundUp(cols, Int<ELE_NUM_PER_C0>{}), MakeStride(Int<1>{}, Int<ELE_NUM_PER_C0>{}))
499+ );
500+ } else {
501+ return MakeLayout(
502+ MakeShape(MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(rows, Int<ELE_NUM_PER_C0>{})), cols),
503+ MakeStride(MakeStride(Int<1>{}, cols * ELE_NUM_PER_C0), Int<ELE_NUM_PER_C0>{})
504+ );
505+ }
506+ } else if constexpr (std::is_same_v<LayoutTag, NpuArch::layout::ColumnMajor>) {
507+ if constexpr (!isMxScaleB) {
508+ return MakeLayout(
509+ MakeShape(rows, MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(cols, Int<ELE_NUM_PER_C0>{}))),
510+ MakeStride(Int<ELE_NUM_PER_C0>{}, MakeStride(Int<1>{}, rows * ELE_NUM_PER_C0))
511+ );
512+ } else {
513+ return MakeLayout(
514+ MakeShape(MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(rows, Int<ELE_NUM_PER_C0>{})), cols),
515+ MakeStride(MakeStride(Int<1>{}, Int<ELE_NUM_PER_C0>{}), RoundUp(rows, Int<ELE_NUM_PER_C0>{}))
516+ );
517+ }
518+ } else if constexpr (std::is_same_v<LayoutTag, NpuArch::layout::zZ>) {
519+ return MakeLayout(
520+ MakeShape(
521+ MakeShape(Int<NpuArch::C0_NUM_PER_FRACTAL>{}, CeilDiv(rows, Int<NpuArch::C0_NUM_PER_FRACTAL>{})),
522+ MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(cols, Int<ELE_NUM_PER_C0>{}))
523+ ),
524+ MakeStride(
525+ MakeStride(
526+ Int<ELE_NUM_PER_C0>{}, RoundUp((int64_t)cols, Int<ELE_NUM_PER_C0>{}) * NpuArch::C0_NUM_PER_FRACTAL
527+ ),
528+ MakeStride(Int<1>{}, Int<ELE_NUM_PER_FRACTAL>{})
529+ )
530+ );
531+ } else {
532+ return MakeLayout(
533+ MakeShape(
534+ MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(rows, Int<ELE_NUM_PER_C0>{})),
535+ MakeShape(Int<NpuArch::C0_NUM_PER_FRACTAL>{}, CeilDiv(cols, Int<NpuArch::C0_NUM_PER_FRACTAL>{}))
536+ ),
537+ MakeStride(
538+ MakeStride(Int<1>{}, Int<ELE_NUM_PER_FRACTAL>{}),
539+ MakeStride(
540+ Int<ELE_NUM_PER_C0>{}, RoundUp((int64_t)rows, Int<ELE_NUM_PER_C0>{}) * NpuArch::C0_NUM_PER_FRACTAL
541+ )
542+ )
543+ );
544+ }
545+}
546+ 
547+template <class Layout, class ShapeNew>
548+HOST_DEVICE constexpr
549+auto MakeLayoutTile(Layout const& layout, ShapeNew const& shapeNew)
550+{
551+ static_assert(
552+ is_tuple<ShapeNew>::value && depth_v<ShapeNew> == 1 && (rank_v<ShapeNew> == 1 || rank_v<ShapeNew> == 2)
553+ );
554+ 
555+ if constexpr (Layout::depth == 1 && (Layout::rank == 1 || Layout::rank == 2)) {
556+ return MakeLayout(shapeNew, layout.stride());
557+ } else if constexpr (Layout::depth == 2 && Layout::rank == 2 && rank_v<decltype(shape<0>(Layout{}))> == 1 &&
558+ rank_v<decltype(shape<1>(Layout{}))> == 2) {
559+ const uint32_t rows = get<0>(shapeNew);
560+ const uint32_t cols = get<1>(shapeNew);
561+ constexpr uint32_t ELE_NUM_PER_C0 = decltype(shape<1, 0>(layout))::value;
562+ return MakeLayout(
563+ MakeShape(rows, MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(cols, Int<ELE_NUM_PER_C0>{}))),
564+ layout.stride()
565+ );
566+ } else if constexpr (Layout::depth == 2 && Layout::rank == 2 && rank_v<decltype(shape<0>(Layout{}))> == 2 &&
567+ rank_v<decltype(shape<1>(Layout{}))> == 1) {
568+ const uint32_t rows = get<0>(shapeNew);
569+ const uint32_t cols = get<1>(shapeNew);
570+ constexpr uint32_t ELE_NUM_PER_C0 = decltype(shape<0, 0>(layout))::value;
571+ return MakeLayout(
572+ MakeShape(MakeShape(Int<ELE_NUM_PER_C0>{}, CeilDiv(rows, Int<ELE_NUM_PER_C0>{})), cols),
573+ layout.stride()
574+ );
575+ } else if constexpr (is_static<decltype(shape<0, 0>(layout))>::value &&
576+ is_static<decltype(shape<1, 0>(layout))>::value) {
577+ const uint32_t rows = get<0>(shapeNew);
578+ const uint32_t cols = get<1>(shapeNew);
579+ constexpr uint32_t dstInnerShapeRow = decltype(shape<0, 0>(layout))::value;
580+ constexpr uint32_t dstInnerShapeCol = decltype(shape<1, 0>(layout))::value;
581+ return MakeLayout(
582+ MakeShape(MakeShape(Int<dstInnerShapeRow>{}, CeilDiv<dstInnerShapeRow>(rows)),
583+ MakeShape(Int<dstInnerShapeCol>{}, CeilDiv<dstInnerShapeCol>(cols))),
584+ layout.stride());
585+ } else {
586+ const uint32_t rows = get<0>(shapeNew);
587+ const uint32_t cols = get<1>(shapeNew);
588+ const uint32_t dstInnerShapeRow = shape<0, 0>(layout);
589+ const uint32_t dstInnerShapeCol = shape<1, 0>(layout);
590+ return MakeLayout(
591+ MakeShape(MakeShape(dstInnerShapeRow, CeilDiv(rows, dstInnerShapeRow)),
592+ MakeShape(dstInnerShapeCol, CeilDiv(cols, dstInnerShapeCol))),
593+ layout.stride());
594+ }
595+}
596+ 
597+template <class T, class U>
598+HOST_DEVICE constexpr
599+auto MakeLayoutL0C(T const& rows, U const& cols)
600+{
601+ constexpr uint32_t ELE_NUM_PER_FRACTAL = 256;
602+ return MakeLayout(
603+ MakeShape(MakeShape(Int<NpuArch::C0_NUM_PER_FRACTAL>{}, CeilDiv(rows, Int<NpuArch::C0_NUM_PER_FRACTAL>{})),
604+ MakeShape(Int<NpuArch::C0_NUM_PER_FRACTAL>{}, CeilDiv(cols, Int<NpuArch::C0_NUM_PER_FRACTAL>{}))),
605+ MakeStride(MakeStride(Int<NpuArch::C0_NUM_PER_FRACTAL>{}, Int<ELE_NUM_PER_FRACTAL>{}),
606+ MakeStride(
607+ Int<1>{}, RoundUp((int64_t)rows, Int<NpuArch::C0_NUM_PER_FRACTAL>{}) * NpuArch::C0_NUM_PER_FRACTAL)));
608+}
609+ 
610+} // end namespace tla
611+ 
612+# endif // TLA_LAYOUT_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/tla/numeric/integer_sequence.hpp+68-0
@@ -0,0 +1,68 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 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+#ifndef TLA_NUMERIC_INTEGER_SEQUENCE_HPP
12+#define TLA_NUMERIC_INTEGER_SEQUENCE_HPP
13+ 
14+#include "../../tla/numeric/integral_constant.hpp"
15+#include "../../tla/type_traits.hpp"
16+ 
17+namespace tla {
18+ 
19+template <typename T, T... Ns>
20+struct IntegerSequence {
21+ using value_type = T;
22+ static constexpr size_t size() { return sizeof...(Ns); }
23+};
24+ 
25+template <typename Sequence, typename T, size_t N>
26+struct MakeIntegerSequenceImpl;
27+ 
28+template <typename T, size_t... Ns>
29+struct MakeIntegerSequenceImpl<IntegerSequence<T, Ns...>, T, 0> {
30+ typedef IntegerSequence<T, Ns...> type;
31+};
32+ 
33+template <typename T, size_t N, size_t... Ns>
34+struct MakeIntegerSequenceImpl<IntegerSequence<T, Ns...>, T, N> {
35+ typedef typename MakeIntegerSequenceImpl<IntegerSequence<T, N - 1, Ns...>, T, N - 1>::type type;
36+};
37+ 
38+template <typename T, T N>
39+using MakeIntegerSequence = typename MakeIntegerSequenceImpl<IntegerSequence<T>, T, N>::type;
40+ 
41+ 
42+// index_sequence
43+template <size_t... Ints>
44+using index_sequence = IntegerSequence<size_t, Ints...>;
45+ 
46+template <size_t N>
47+using make_index_sequence = MakeIntegerSequence<size_t, N>;
48+ 
49+// int_sequence
50+template <int... Ints>
51+using int_sequence = IntegerSequence<int, Ints...>;
52+ 
53+template <int N>
54+using make_int_sequence = MakeIntegerSequence<int, N>;
55+ 
56+// Shortcuts
57+template <int... Ints>
58+using seq = int_sequence<Ints...>;
59+ 
60+template <int N>
61+using make_seq = make_int_sequence<N>;
62+ 
63+template <class Tuple>
64+using tuple_seq = make_seq<tuple_size<tla::remove_cvref_t<Tuple>>::value>;
65+ 
66+} // end namespace tla
67+ 
68+#endif // TLA_NUMERIC_INTEGER_SEQUENCE_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/tla/numeric/integral_constant.hpp+172-0
@@ -0,0 +1,172 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 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+#ifndef TLA_NUMERIC_INTEGER_CONSTANT_HPP
12+#define TLA_NUMERIC_INTEGER_CONSTANT_HPP
13+ 
14+#include "../../attn_infra/detail/macros.hpp"
15+#include "../../tla/type_traits.hpp"
16+#include "../../tla/numeric/math.hpp"
17+ 
18+namespace tla {
19+ 
20+// A constant value: short name and type-deduction for fast compilation
21+template <auto v>
22+struct C {
23+ using type = C<v>;
24+ static constexpr auto value = v;
25+ using value_type = decltype(v);
26+ HOST_DEVICE constexpr operator value_type() const noexcept { return value; }
27+ HOST_DEVICE constexpr value_type operator()() const noexcept { return value; }
28+};
29+ 
30+// Deprecate
31+template <class T, T v>
32+using constant = C<v>;
33+ 
34+template <bool b>
35+using bool_constant = C<b>;
36+ 
37+using true_type = bool_constant<true>;
38+using false_type = bool_constant<false>;
39+ 
40+template <class T>
41+using is_std_integral = std::is_integral<T>;
42+ 
43+// A more std:: conforming integral_constant that enforces type but interops with C<v>
44+template <class T, T v>
45+struct integral_constant : C<v> {
46+ using type = integral_constant<T, v>;
47+ static constexpr T value = v;
48+ using value_type = T;
49+ HOST_DEVICE constexpr value_type operator()() const noexcept { return value; }
50+};
51+ 
52+// Use tla::is_std_integral<T> to match built-in integral types (int, int64_t, unsigned, etc)
53+// Use tla::is_integral<T> to match both built-in integral types AND static integral types.
54+ 
55+template <class T>
56+struct is_integral : bool_constant<is_std_integral<T>::value> {};
57+template <auto v>
58+struct is_integral<C<v> > : true_type {};
59+template <class T, T v>
60+struct is_integral<integral_constant<T, v>> : true_type {};
61+ 
62+// is_static detects if an (abstract) value is defined completely by its type (no members)
63+template <class T>
64+struct is_static : bool_constant<std::is_empty<remove_cvref_t<T>>::value> {};
65+ 
66+// is_constant detects if a type is a static integral type and if v is equal to a value
67+ 
68+template <auto n, class T>
69+struct is_constant : false_type {};
70+template <auto n, class T>
71+struct is_constant<n, T const > : is_constant<n, T> {};
72+template <auto n, class T>
73+struct is_constant<n, T const&> : is_constant<n, T> {};
74+template <auto n, class T>
75+struct is_constant<n, T &> : is_constant<n, T> {};
76+template <auto n, class T>
77+struct is_constant<n, T &&> : is_constant<n, T> {};
78+template <auto n, auto v>
79+struct is_constant<n, C<v> > : bool_constant<v == n> {};
80+template <auto n, class T, T v>
81+struct is_constant<n, integral_constant<T, v>> : bool_constant<v == n> {};
82+ 
83+//
84+// Specializations
85+//
86+ 
87+template <int v>
88+using Int = C<v>;
89+using _0 = Int<0>;
90+using _64 = Int<64>;
91+using _128 = Int<128>;
92+using _256 = Int<256>;
93+using _512 = Int<512>;
94+ 
95+/***************/
96+/** Operators **/
97+/***************/
98+ 
99+#define TLA_LEFT_UNARY_OP(OP) \
100+ template <auto t> \
101+ HOST_DEVICE constexpr \
102+ C<(OP t)> operator OP (C<t>) { \
103+ return {}; \
104+ }
105+#define TLA_BINARY_OP(OP) \
106+ template <auto t, auto u> \
107+ HOST_DEVICE constexpr \
108+ C<(t OP u)> operator OP (C<t>, C<u>) { \
109+ return {}; \
110+ }
111+ 
112+TLA_LEFT_UNARY_OP(+);
113+TLA_LEFT_UNARY_OP(-);
114+TLA_LEFT_UNARY_OP(~);
115+TLA_LEFT_UNARY_OP(!);
116+TLA_LEFT_UNARY_OP(*);
117+ 
118+TLA_BINARY_OP(+);
119+TLA_BINARY_OP(-);
120+TLA_BINARY_OP(*);
121+TLA_BINARY_OP(/);
122+TLA_BINARY_OP(%);
123+TLA_BINARY_OP(&);
124+TLA_BINARY_OP(|);
125+TLA_BINARY_OP(^);
126+TLA_BINARY_OP(<<);
127+TLA_BINARY_OP(>>);
128+ 
129+#undef TLA_BINARY_OP
130+#undef TLA_LEFT_UNARY_OP
131+#undef TLA_RIGHT_UNARY_OP
132+ 
133+//
134+// Named functions from math.hpp
135+//
136+ 
137+#define TLA_NAMED_UNARY_FN(OP) \
138+ template <auto t> \
139+ HOST_DEVICE constexpr \
140+ auto OP (C<t>) { \
141+ return C<OP(t)>{}; \
142+ }
143+#define TLA_NAMED_BINARY_FN(OP) \
144+ template <auto t, auto u> \
145+ HOST_DEVICE constexpr \
146+ auto OP (C<t>, C<u>) { \
147+ return C<OP(t, u)>{}; \
148+ } \
149+ template <auto t, class U, \
150+ TLA_REQUIRES(is_std_integral<U>::value)> \
151+ HOST_DEVICE constexpr \
152+ auto OP (C<t>, U u) { \
153+ return OP(t, u); \
154+ } \
155+ template <class T, auto u, \
156+ TLA_REQUIRES(is_std_integral<T>::value)> \
157+ HOST_DEVICE constexpr \
158+ auto OP (T t, C<u>) { \
159+ return OP(t, u); \
160+ }
161+ 
162+TLA_NAMED_BINARY_FN(max);
163+TLA_NAMED_BINARY_FN(min);
164+TLA_NAMED_BINARY_FN(add);
165+ 
166+#undef TLA_NAMED_UNARY_FN
167+#undef TLA_NAMED_BINARY_FN
168+ 
169+ 
170+} // end namespace tla
171+ 
172+#endif // TLA_NUMERIC_INTEGER_CONSTANT_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/tla/numeric/math.hpp+52-0
@@ -0,0 +1,52 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 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+#ifndef TLA_NUMERIC_MATH_HPP
12+#define TLA_NUMERIC_MATH_HPP
13+ 
14+#include "../../attn_infra/detail/macros.hpp"
15+#include "../../tla/type_traits.hpp"
16+ 
17+namespace tla {
18+ 
19+//
20+// Common Operations
21+//
22+ 
23+template <class T, class U,
24+ TLA_REQUIRES(std::is_arithmetic<T>::value &&
25+ std::is_arithmetic<U>::value)>
26+HOST_DEVICE constexpr
27+auto
28+max(T const& t, U const& u) {
29+ return t < u ? u : t;
30+}
31+ 
32+template <class T, class U,
33+ TLA_REQUIRES(std::is_arithmetic<T>::value &&
34+ std::is_arithmetic<U>::value)>
35+HOST_DEVICE constexpr
36+auto
37+min(T const& t, U const& u) {
38+ return t < u ? t : u;
39+}
40+ 
41+template <class T, class U,
42+ TLA_REQUIRES(std::is_arithmetic<T>::value &&
43+ std::is_arithmetic<U>::value)>
44+HOST_DEVICE constexpr
45+auto
46+add(T const& t, U const& u) {
47+ return t + u;
48+}
49+ 
50+} // namespace tla
51+ 
52+#endif // TLA_NUMERIC_MATH_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/tla/tensor.hpp+117-0
@@ -0,0 +1,117 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 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+#ifndef TLA_TENSOR_HPP
12+#define TLA_TENSOR_HPP
13+ 
14+#include "../attn_infra/arch/arch.hpp"
15+#include "../tla/layout.hpp" // tla::Shape
16+#include "../tla/numeric/integral_constant.hpp" // tla::is_integral
17+#include "../tla/int_tuple.hpp"
18+ 
19+namespace tla {
20+//
21+// Tensor
22+//
23+ 
24+template <class BuiltinTensor, class Layout_, class Coord_, AscendC::TPosition Position>
25+struct Tensor {
26+ using Element = typename BuiltinTensor::PrimType;
27+ using Layout = Layout_;
28+ using Coord = Coord_;
29+ static constexpr AscendC::TPosition position = Position;
30+ 
31+ HOST_DEVICE constexpr
32+ Tensor() {}
33+ 
34+ HOST_DEVICE constexpr
35+ Tensor(BuiltinTensor const& builtinTensor, Layout const& layout, Coord const& coord = {})
36+ : rep_(builtinTensor, layout, coord) {}
37+ 
38+ //
39+ // Accessors
40+ //
41+ 
42+ static constexpr int rank = Layout::rank;
43+ 
44+ HOST_DEVICE constexpr
45+ decltype(auto) tensor() const
46+ {
47+ return *this;
48+ }
49+ 
50+ HOST_DEVICE constexpr
51+ decltype(auto) data() const
52+ {
53+ return get<0>(rep_);
54+ }
55+ 
56+ HOST_DEVICE constexpr
57+ decltype(auto) data()
58+ {
59+ return get<0>(rep_);
60+ }
61+ 
62+ HOST_DEVICE constexpr
63+ decltype(auto) layout() const
64+ {
65+ return get<1>(rep_);
66+ }
67+ 
68+ HOST_DEVICE constexpr
69+ decltype(auto) coord() const
70+ {
71+ return get<2>(rep_);
72+ }
73+ 
74+ HOST_DEVICE constexpr
75+ decltype(auto) shape() const
76+ {
77+ return layout().shape();
78+ }
79+ 
80+ HOST_DEVICE constexpr
81+ decltype(auto) stride() const
82+ {
83+ return layout().stride();
84+ }
85+ 
86+ tla::tuple<BuiltinTensor, Layout, Coord> rep_;
87+};
88+ 
89+template <class BuiltinTensor, class Layout, class PositionType>
90+HOST_DEVICE constexpr
91+auto MakeTensor(BuiltinTensor const& builtinTensor, Layout const& layout, PositionType)
92+{
93+ using Coord = detail::MakeZeroTuple<Layout::rank>;
94+ return Tensor<BuiltinTensor, Layout, Coord, PositionType::value>(builtinTensor, layout);
95+}
96+ 
97+template <class BuiltinTensor, class Layout, class Coord, class PositionType>
98+HOST_DEVICE constexpr
99+auto MakeTensor(BuiltinTensor const& builtinTensor, Layout const& layout, Coord const& coord, PositionType)
100+{
101+ return Tensor<BuiltinTensor, Layout, Coord, PositionType::value>(builtinTensor, layout, coord);
102+}
103+ 
104+template <class Tensor, class Coord, class Shape>
105+__aicore__ inline constexpr
106+auto GetTile(Tensor const& tensor, Coord const& coord, Shape const& shape)
107+{
108+ auto layout = tensor.layout();
109+ auto builtinTensor = tensor.data();
110+ auto layoutNew = MakeLayoutTile(layout, shape);
111+ auto coordNew = Add(tensor.coord(), coord);
112+ return MakeTensor(builtinTensor, layoutNew, coordNew, NpuArch::Arch::PositionType<Tensor::position>{});
113+}
114+ 
115+} // end namespace tla
116+ 
117+#endif // TLA_TENSOR_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/tla/tuple.hpp+151-0
@@ -0,0 +1,151 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 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+#ifndef TLA_TUPLE_HPP
12+#define TLA_TUPLE_HPP
13+ 
14+#include "../tla/numeric/integral_constant.hpp"
15+#include "../tla/numeric/integer_sequence.hpp"
16+ 
17+namespace tla {
18+ 
19+namespace detail {
20+ 
21+// EBO stands for "empty base optimization."
22+template <size_t N, class T, bool IsEmpty = std::is_empty<T>::value>
23+struct EBO;
24+ 
25+// Specialization for types T that are empty;
26+template <size_t N, class T>
27+struct EBO<N, T, true> {
28+ HOST_DEVICE constexpr
29+ EBO() {}
30+ 
31+ HOST_DEVICE constexpr
32+ EBO(T const&) {}
33+};
34+ 
35+template <size_t N, class T>
36+HOST_DEVICE constexpr
37+T getv(EBO<N, T, true> const&)
38+{
39+ return {};
40+}
41+ 
42+// Specialization for types T that are not empty;
43+template <size_t N, class T>
44+struct EBO<N, T, false> {
45+ HOST_DEVICE constexpr
46+ EBO() : t_{} {}
47+ 
48+ HOST_DEVICE constexpr
49+ EBO(T const& t) : t_{t} {}
50+ 
51+ T t_;
52+};
53+ 
54+template <size_t N, class T>
55+HOST_DEVICE constexpr
56+T const& getv(EBO<N, T, false> const& x)
57+{
58+ return x.t_;
59+}
60+ 
61+template <size_t N, class T>
62+HOST_DEVICE constexpr
63+T& getv(EBO<N, T, false>& x)
64+{
65+ return x.t_;
66+}
67+ 
68+// TupleBase
69+template <class IdxSeq, class... T>
70+struct TupleBase;
71+ 
72+template <size_t... I, class... T>
73+struct TupleBase<index_sequence<I...>, T...> : EBO<I, T>... {
74+ HOST_DEVICE constexpr
75+ TupleBase() {}
76+ 
77+ HOST_DEVICE constexpr
78+ TupleBase(T const&... t) : EBO<I, T>(t)... {}
79+};
80+ 
81+} // end namespace detail
82+ 
83+// tla::tuple class.
84+template <class... T>
85+struct tuple : detail::TupleBase<make_index_sequence<sizeof...(T)>, T...> {
86+ HOST_DEVICE constexpr
87+ tuple() {}
88+ 
89+ HOST_DEVICE constexpr
90+ tuple(T const&... t) : detail::TupleBase<make_index_sequence<sizeof...(T)>, T...>(t...) {}
91+};
92+ 
93+template <>
94+struct tuple<> {};
95+ 
96+// get for tla::tuple
97+template <size_t I, class... T>
98+HOST_DEVICE constexpr
99+decltype(auto) get(tuple<T...> const& t) noexcept
100+{
101+ static_assert(I < sizeof...(T), "Index out of range");
102+ return detail::getv<I>(t);
103+}
104+ 
105+template <size_t I, class... T>
106+HOST_DEVICE constexpr
107+decltype(auto) get(tuple<T...>& t) noexcept
108+{
109+ static_assert(I < sizeof...(T), "Index out of range");
110+ return detail::getv<I>(t);
111+}
112+ 
113+template <size_t I, class... T>
114+HOST_DEVICE constexpr
115+decltype(auto) get(tuple<T...>&& t) noexcept
116+{
117+ static_assert(I < sizeof...(T), "Index out of range");
118+ return detail::getv<I>(static_cast<tuple<T...>&&>(t));
119+}
120+ 
121+namespace detail {
122+ 
123+template <class T>
124+auto has_tuple_size(T*) -> bool_constant<(0 <= tuple_size<T>::value)>;
125+auto has_tuple_size(...) -> false_type;
126+ 
127+} // end namespace detail
128+ 
129+template <class T>
130+struct is_tuple : decltype(detail::has_tuple_size((T*)0)) {};
131+ 
132+template <class... T>
133+struct tuple_size<tla::tuple<T...>>
134+ : std::integral_constant<size_t, sizeof...(T)> {};
135+ 
136+template <class... T>
137+struct tuple_size<const tla::tuple<T...>>
138+ : std::integral_constant<size_t, sizeof...(T)> {};
139+ 
140+// make_tuple
141+template <class... T>
142+HOST_DEVICE constexpr
143+tuple<T...>
144+MakeTuple(T const&... t)
145+{
146+ return {t...};
147+}
148+ 
149+} // end namespace tla
150+ 
151+#endif // TLA_TUPLE_HPP
Acsrc/ops/eagle_quant_block_sparse_attention/op_kernel/tla/type_traits.hpp+50-0
@@ -0,0 +1,50 @@
1+/**
2+ * This program is free software, you can redistribute it and/or modify.
3+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4+ * This file is a part of the CANN Open Software.
5+ * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6+ * Please refer to the License for details. You may not use this file except in compliance with the License.
7+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, 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+#ifndef TLA_UTIL_TYPE_TRAITS_HPP
12+#define TLA_UTIL_TYPE_TRAITS_HPP
13+ 
14+#ifdef ASCENDC_MODULE_OPERATOR_H
15+#undef inline
16+#endif
17+#include <tuple>
18+#ifdef ASCENDC_MODULE_OPERATOR_H
19+#define inline __inline__ __attribute__((always_inline))
20+#endif
21+ 
22+#define TLA_REQUIRES(...) typename std::enable_if<(__VA_ARGS__)>::type* = nullptr
23+ 
24+namespace tla {
25+ 
26+// using std::remove_cvref;
27+template <class T>
28+struct remove_cvref {
29+ using type = std::remove_cv_t<std::remove_reference_t<T>>;
30+};
31+ 
32+// using std::remove_cvref_t;
33+template <class T>
34+using remove_cvref_t = typename remove_cvref<T>::type;
35+ 
36+ 
37+// tuple_size, tuple_element
38+template <class T, class = void>
39+struct tuple_size;
40+ 
41+template <class T>
42+struct tuple_size<T, std::void_t<typename std::tuple_size<T>::type>>
43+ : std::integral_constant<size_t, std::tuple_size<T>::value> {};
44+ 
45+template <class T>
46+constexpr size_t tuple_size_v = tuple_size<T>::value;
47+ 
48+} // end namespace tla
49+ 
50+#endif // TLA_UTIL_TYPE_TRAITS_HPP