Star1443
3089
代码
代码
Issues297
Pull Requests644
Actions
项目讨论
Wiki
分析
Star1443
3089
已合并
【NSA算子性能优化】importance score计算性能优化 #2615
dingxu创建于 3月12日
【NSA算子性能优化】importance score计算性能优化 #2615
已合并
dingxu创建于 3月12日
从xdnjust/ops-transformeropt_nsa合入到cann/ops-transformermaster
共 21 个文件变更+5202-0
Aexperimental/attention/nsa_compress_attention_infer/CMakeLists.txt+19-0
@@ -0,0 +1,19 @@
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+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
12+if(NOT ENABLE_TEST AND NOT BENCHMARK)
13+ list(REMOVE_ITEM CURRENT_DIRS tests)
14+endif()
15+foreach(SUB_DIR ${CURRENT_DIRS})
16+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
17+ add_subdirectory(${SUB_DIR})
18+ endif()
19+endforeach()
Aexperimental/attention/nsa_compress_attention_infer/examples/test_aclnn_nsa_compress_attention_infer.cpp+218-0
@@ -0,0 +1,218 @@
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+/*!
12+ * \file test_nsa_compress_attention_infer_2.cpp
13+ * \brief
14+ */
15+ 
16+ 
17+#include <iostream>
18+#include <vector>
19+#include <cstring>
20+#include "acl/acl.h"
21+#include "aclnn/opdev/fp16_t.h"
22+#include "aclnnop/aclnn_nsa_compress_attention_infer.h"
23+ 
24+using namespace std;
25+ 
26+#define CHECK_RET(cond, return_expr) \
27+ do { \
28+ if (!(cond)) { \
29+ return_expr; \
30+ } \
31+ } while (0)
32+ 
33+#define LOG_PRINT(message, ...) \
34+ do { \
35+ printf(message, ##__VA_ARGS__); \
36+ } while (0)
37+ 
38+int64_t GetShapeSize(const std::vector<int64_t>& shape) {
39+ int64_t shapeSize = 1;
40+ for (auto i : shape) {
41+ shapeSize *= i;
42+ }
43+ return shapeSize;
44+}
45+ 
46+int Init(int32_t deviceId, aclrtStream* stream) {
47+ // 固定写法，资源初始化
48+ auto ret = aclInit(nullptr);
49+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
50+ ret = aclrtSetDevice(deviceId);
51+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
52+ ret = aclrtCreateStream(stream);
53+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
54+ return 0;
55+}
56+ 
57+template <typename T>
58+int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
59+ aclDataType dataType, aclTensor** tensor) {
60+ auto size = GetShapeSize(shape) * sizeof(T);
61+ // 调用aclrtMalloc申请device侧内存
62+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
63+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
64+ // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
65+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
66+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
67+
68+ // 计算连续tensor的strides
69+ std::vector<int64_t> strides(shape.size(), 1);
70+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
71+ strides[i] = shape[i + 1] * strides[i + 1];
72+ }
73+ 
74+ // 调用aclCreateTensor接口创建aclTensor
75+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
76+ shape.data(), shape.size(), *deviceAddr);
77+ return 0;
78+}
79+ 
80+int main() {
81+ // 1. （固定写法）device/stream初始化，参考AscendCL对外接口列表
82+ // 根据自己的实际device填写deviceId
83+ int32_t deviceId = 0;
84+ aclrtStream stream;
85+ auto ret = Init(deviceId, &stream);
86+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
87+ 
88+ // 2. 构造输入与输出，需要根据API的接口自定义构造
89+ int32_t batchSize = 20;
90+ int32_t headDimsQK = 192;
91+ int32_t blockNum = 640;
92+ int32_t headDimsV = 128;
93+ int32_t sequenceLengthK = 4096;
94+ int32_t maxNumBlocksPerSeq = 32;
95+ // attr
96+ int64_t numHeads = 64;
97+ int64_t numKeyValueHeads = 4;
98+ int64_t selectBlockSize = 64;
99+ int64_t selectBlockCount = 16;
100+ int64_t compressBlockSize = 32;
101+ int64_t compressStride = 16;
102+ double scaleValue = 0.088388;
103+ string sLayerOut = "TND";
104+ char layOut[sLayerOut.length()];
105+ strcpy(layOut, sLayerOut.c_str());
106+ int64_t pageBlockSize = 128;
107+ int64_t sparseMod = 0;
108+ std::vector<int64_t> queryShape = {batchSize, numHeads, headDimsQK};
109+ std::vector<int64_t> keyShape = {blockNum, pageBlockSize, numKeyValueHeads * headDimsQK};
110+ std::vector<int64_t> valueShape = {blockNum, pageBlockSize, numKeyValueHeads * headDimsV};
111+ std::vector<int64_t> blockTableOptionalShape = {batchSize, maxNumBlocksPerSeq};
112+ std::vector<int64_t> outputShape = {batchSize, numHeads, headDimsV};
113+ std::vector<int64_t> topkIndicesShape = {batchSize, numKeyValueHeads, selectBlockCount};
114+ void *queryDeviceAddr = nullptr;
115+ void *keyDeviceAddr = nullptr;
116+ void *valueDeviceAddr = nullptr;
117+ void *blockTableOptionalDeviceAddr = nullptr;
118+ void *outputDeviceAddr = nullptr;
119+ void *topkIndicesDeviceAddr = nullptr;
120+ aclTensor *queryTensor = nullptr;
121+ aclTensor *keyTensor = nullptr;
122+ aclTensor *valueTensor = nullptr;
123+ aclTensor *blockTableOptionalTensor = nullptr;
124+ aclTensor *outputTensor = nullptr;
125+ aclTensor *topkIndicesTensor = nullptr;
126+ std::vector<op::fp16_t> queryHostData(batchSize * numHeads * headDimsQK, 1.0);
127+ std::vector<op::fp16_t> keyHostData(blockNum * pageBlockSize * numKeyValueHeads * headDimsQK, 1.0);
128+ std::vector<op::fp16_t> valueHostData(blockNum * pageBlockSize * numKeyValueHeads * headDimsV, 1.0);
129+ std::vector<int32_t> blockTableOptionalHostData(batchSize * maxNumBlocksPerSeq, 1);
130+ std::vector<op::fp16_t> outputHostData(batchSize * numHeads * headDimsV, 1.0);
131+ std::vector<int32_t> topkIndicesHostData(batchSize * numKeyValueHeads * selectBlockCount, 1);
132+ 
133+ // 创建query aclTensor
134+ ret = CreateAclTensor(queryHostData, queryShape, &queryDeviceAddr, aclDataType::ACL_FLOAT16, &queryTensor);
135+ CHECK_RET(ret == ACL_SUCCESS, return ret);
136+ // 创建key aclTensor
137+ ret = CreateAclTensor(keyHostData, keyShape, &keyDeviceAddr, aclDataType::ACL_FLOAT16, &keyTensor);
138+ CHECK_RET(ret == ACL_SUCCESS, return ret);
139+ // 创建v aclTensor
140+ ret = CreateAclTensor(valueHostData, valueShape, &valueDeviceAddr, aclDataType::ACL_FLOAT16, &valueTensor);
141+ CHECK_RET(ret == ACL_SUCCESS, return ret);
142+ // 创建blockTableOptional aclTensor
143+ ret = CreateAclTensor(blockTableOptionalHostData, blockTableOptionalShape, &blockTableOptionalDeviceAddr, aclDataType::ACL_INT32, &blockTableOptionalTensor);
144+ CHECK_RET(ret == ACL_SUCCESS, return ret);
145+ // 创建output aclTensor
146+ ret = CreateAclTensor(outputHostData, outputShape, &outputDeviceAddr, aclDataType::ACL_FLOAT16, &outputTensor);
147+ CHECK_RET(ret == ACL_SUCCESS, return ret);
148+ // 创建topkIndices aclTensor
149+ ret = CreateAclTensor(topkIndicesHostData, topkIndicesShape, &topkIndicesDeviceAddr, aclDataType::ACL_INT32, &topkIndicesTensor);
150+ CHECK_RET(ret == ACL_SUCCESS, return ret);
151+ 
152+ std::vector<int64_t> actualCmpKvSeqLenVector(batchSize, sequenceLengthK);
153+ auto actualCmpKvSeqLen = aclCreateIntArray(actualCmpKvSeqLenVector.data(), actualCmpKvSeqLenVector.size());
154+
155+ // 3. 调用CANN算子库API
156+ uint64_t workspaceSize = 0;
157+ aclOpExecutor* executor;
158+ // 调用第一段接口
159+ ret = aclnnNsaCompressAttentionInferGetWorkspaceSize(queryTensor, keyTensor, valueTensor, nullptr, blockTableOptionalTensor, nullptr, actualCmpKvSeqLen,
160+ nullptr, nullptr,
161+ numHeads, numKeyValueHeads, selectBlockSize, selectBlockCount, compressBlockSize, compressStride,
162+ scaleValue, layOut, pageBlockSize, sparseMod, outputTensor, topkIndicesTensor, &workspaceSize, &executor);
163+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNsaCompressAttentionInferGetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
164+ // 根据第一段接口计算出的workspaceSize申请device内存
165+ void* workspaceAddr = nullptr;
166+ if (workspaceSize > 0) {
167+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
168+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
169+ }
170+ // 调用第二段接口
171+ ret = aclnnNsaCompressAttentionInfer(workspaceAddr, workspaceSize, executor, stream);
172+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnNsaCompressAttentionInfer failed. ERROR: %d\n", ret); return ret);
173+ 
174+ // 4. （固定写法）同步等待任务执行结束
175+ ret = aclrtSynchronizeStream(stream);
176+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
177+ 
178+ // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
179+ auto size = GetShapeSize(outputShape);
180+ std::vector<op::fp16_t> resultData(size, 0);
181+ ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), outputDeviceAddr,
182+ size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
183+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy [attn] result from device to host failed. ERROR: %d\n", ret); return ret);
184+ uint64_t printNum = 10;
185+ for (int64_t i = 0; i < printNum; i++) {
186+ std::cout << "index: " << i << ": " << static_cast<float>(resultData[i]) << std::endl;
187+ }
188+ auto topksize = GetShapeSize(topkIndicesShape);
189+ std::vector<op::fp16_t> topkresultData(topksize, 0);
190+ ret = aclrtMemcpy(topkresultData.data(), topkresultData.size() * sizeof(topkresultData[0]), topkIndicesDeviceAddr,
191+ topksize * sizeof(topkresultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
192+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy [top k] result from device to host failed. ERROR: %d\n", ret); return ret);
193+ for (int64_t i = 0; i < printNum; i++) {
194+ std::cout << "topk index: " << i << ": " << static_cast<int32_t>(topkresultData[i]) << std::endl;
195+ }
196+ 
197+ // 6. 释放资源
198+ aclDestroyTensor(queryTensor);
199+ aclDestroyTensor(keyTensor);
200+ aclDestroyTensor(valueTensor);
201+ aclDestroyTensor(blockTableOptionalTensor);
202+ aclDestroyIntArray(actualCmpKvSeqLen);
203+ aclDestroyTensor(outputTensor);
204+ aclDestroyTensor(topkIndicesTensor);
205+ aclrtFree(queryDeviceAddr);
206+ aclrtFree(keyDeviceAddr);
207+ aclrtFree(valueDeviceAddr);
208+ aclrtFree(blockTableOptionalDeviceAddr);
209+ aclrtFree(outputDeviceAddr);
210+ aclrtFree(topkIndicesDeviceAddr);
211+ if (workspaceSize > 0) {
212+ aclrtFree(workspaceAddr);
213+ }
214+ aclrtDestroyStream(stream);
215+ aclrtResetDevice(deviceId);
216+ aclFinalize();
217+ return 0;
218+}
Aexperimental/attention/nsa_compress_attention_infer/op_host/CMakeLists.txt+29-0
@@ -0,0 +1,29 @@
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+add_op_to_compiled_list()
11+ 
12+if (BUILD_OPEN_PROJECT)
13+ target_sources(op_host_aclnnExc PRIVATE
14+ nsa_compress_attention_infer_def.cpp
15+ )
16+endif()
17+ 
18+if (BUILD_OPS_RTY_KERNEL) # 回黄kernel
19+ add_ops_compile_options(
20+ OP_NAME NsaCompressAttentionInfer
21+ OPTIONS --cce-auto-sync=off
22+ -Wno-deprecated-declarations
23+ -Werror
24+ -fpermissive
25+ )
26+else()
27+ add_modules_sources(OPTYPE nsa_compress_attention_infer ACLNNTYPE aclnn_exclude)
28+endif()
29+ 
Aexperimental/attention/nsa_compress_attention_infer/op_host/nsa_compress_attention_infer_def.cpp+97-0
@@ -0,0 +1,97 @@
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+/*!
12+ * \file nsa_compress_attention_infer_def.cpp
13+ * \brief
14+ */
15+ 
16+#include "register/op_def_registry.h"
17+ 
18+ 
19+namespace ops {
20+class NsaCompressAttentionInfer : public OpDef {
21+public:
22+ explicit NsaCompressAttentionInfer(const char* name) : OpDef(name)
23+ {
24+ this->Input("query")
25+ .ParamType(REQUIRED)
26+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
27+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
28+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
29+ this->Input("key")
30+ .ParamType(REQUIRED)
31+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
32+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
33+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
34+ this->Input("value")
35+ .ParamType(REQUIRED)
36+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
37+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
38+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
39+ this->Input("attentionMaskOptional")
40+ .ParamType(OPTIONAL)
41+ .DataType({ge::DT_BOOL, ge::DT_BOOL})
42+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
43+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
44+ this->Input("blockTableOptional")
45+ .ParamType(OPTIONAL)
46+ .DataType({ge::DT_INT32, ge::DT_INT32})
47+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
48+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
49+ this->Input("actualQSeqLenOptional")
50+ .ParamType(OPTIONAL)
51+ .DataType({ge::DT_INT64, ge::DT_INT64})
52+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
53+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
54+ this->Input("actualCmpKvSeqLenOptional")
55+ .ParamType(OPTIONAL)
56+ .DataType({ge::DT_INT64, ge::DT_INT64})
57+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
58+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
59+ this->Input("actualSelKvSeqLenOptional")
60+ .ParamType(OPTIONAL)
61+ .DataType({ge::DT_INT64, ge::DT_INT64})
62+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
63+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
64+ this->Input("topkMaskOptional")
65+ .ParamType(OPTIONAL)
66+ .DataType({ge::DT_BOOL, ge::DT_BOOL})
67+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
68+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
69+ this->Output("output")
70+ .ParamType(REQUIRED)
71+ .DataType({ge::DT_FLOAT16, ge::DT_BF16})
72+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
73+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
74+ this->Output("topkIndicesOut")
75+ .ParamType(REQUIRED)
76+ .DataType({ge::DT_INT32, ge::DT_INT32})
77+ .Format({ge::FORMAT_ND, ge::FORMAT_ND})
78+ .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
79+
80+ this->Attr("numHeads").AttrType(OPTIONAL).Int(1);
81+ this->Attr("numKeyValueHeads").AttrType(OPTIONAL).Int(0);
82+ this->Attr("selectBlockSize").AttrType(OPTIONAL).Int(0);
83+ this->Attr("selectBlockCount").AttrType(OPTIONAL).Int(0);
84+ this->Attr("compressBlockSize").AttrType(OPTIONAL).Int(0);
85+ this->Attr("compressStride").AttrType(OPTIONAL).Int(0);
86+ this->Attr("scaleValue").AttrType(OPTIONAL).Float(1.0);
87+ this->Attr("layoutOptional").AttrType(OPTIONAL).String();
88+ this->Attr("pageBlockSize").AttrType(OPTIONAL).Int(0);
89+ this->Attr("sparseMode").AttrType(OPTIONAL).Int(0);
90+ 
91+ this->AICore().AddConfig("ascend910b");
92+ this->AICore().AddConfig("ascend910_93");
93+ }
94+};
95+ 
96+OP_ADD(NsaCompressAttentionInfer);
97+}
Aexperimental/attention/nsa_compress_attention_infer/op_host/nsa_compress_attention_infer_infershape.cpp+94-0
@@ -0,0 +1,94 @@
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+/*!
CANN-robot
CANN-robotCANN-robot3月12日
代码结构与可维护性: 文件头注释中的文件名与实际文件名不一致。注释中写的是'nsa_compress_attention_infer_proto.cpp'，但实际文件名是'nsa_compress_attention_infer_infershape.cpp'。这会导致维护困惑。
问题类型: 代码结构与可维护性
文件路径: experimental/attention/nsa_compress_attention_infer/op_host/nsa_compress_attention_infer_infershape.cpp
行号: 11
问题代码:
/*!
 * \file nsa_compress_attention_infer_proto.cpp
 * \brief
 */
修改建议:
将文件头注释中的文件名修改为与实际文件名一致：'nsa_compress_attention_infer_infershape.cpp'
---
此评论由代码审查工具自动生成
likedislike
D
dingxu
3月12日 评论：
12+ * \file nsa_compress_attention_infer_infershape.cpp
13+ * \brief
14+ */
15+ 
16+#include <register/op_impl_registry.h>
17+#include "log/log.h"
18+#include "platform/platform_info.h"
19+ 
20+using namespace ge;
21+ 
22+namespace ops {
23+ 
24+constexpr uint32_t QUERY_INPUT_INDEX = 0;
25+constexpr uint32_t VALUE_INPUT_INDEX = 2;
26+constexpr uint32_t ATTEN_OUTPUT_INDEX = 0;
27+constexpr uint32_t TOPK_OUTPUT_INDEX = 1;
28+constexpr uint32_t KV_NUM_HEADS_ATTR_INDEX = 1;
29+constexpr uint32_t SELECT_BLOCK_ATTR_INDEX = 3;
30+constexpr uint32_t LAYOUT_ATTR_INDEX = 7;
31+static const uint64_t DIM_NUM_1 = 1;
32+static const uint64_t DIM_NUM_2 = 2;
33+static const uint64_t DIM_NUM_3 = 3;
34+ 
35+static ge::graphStatus InferShapeNsaCompressAttentionInfer(gert::InferShapeContext *context)
36+{
37+ if (context == nullptr) {
38+ return ge::GRAPH_FAILED;
39+ }
40+ const gert::Shape *queryShape = context->GetInputShape(QUERY_INPUT_INDEX);
41+ OP_CHECK_NULL_WITH_CONTEXT(context, queryShape);
42+ gert::Shape *outShape = context->GetOutputShape(ATTEN_OUTPUT_INDEX);
43+ OP_CHECK_NULL_WITH_CONTEXT(context, outShape);
44+ const gert::Shape *valueShape = context->GetInputShape(VALUE_INPUT_INDEX);
45+ OP_CHECK_NULL_WITH_CONTEXT(context, valueShape);
46+
47+ auto attrs = context->GetAttrs();
48+ OP_CHECK_NULL_WITH_CONTEXT(context, attrs);
49+ const uint32_t *kvHeadNumsPtr = attrs->GetAttrPointer<uint32_t>(KV_NUM_HEADS_ATTR_INDEX);
50+ OP_CHECK_NULL_WITH_CONTEXT(context, kvHeadNumsPtr);
51+ const int *selectBlockCountPtr = attrs->GetAttrPointer<int32_t>(SELECT_BLOCK_ATTR_INDEX);
52+ OP_CHECK_NULL_WITH_CONTEXT(context, selectBlockCountPtr);
53+ const char *layOutPtr = attrs->GetAttrPointer<char>(LAYOUT_ATTR_INDEX);
54+ OP_CHECK_NULL_WITH_CONTEXT(context, layOutPtr);
55+ 
56+ auto kvHeadNums = *kvHeadNumsPtr;
57+ auto selectBlockCount = *selectBlockCountPtr;
58+ std::string layOutStr = std::string(layOutPtr);
59+ if (layOutStr != "TND" && layOutStr != "BSND") {
60+ OP_LOGE(context, "The layout should be TND or BSND, but got %s.", layOutStr.c_str());
61+ return GRAPH_FAILED;
62+ }
63+ 
64+ auto shapeD2 = valueShape->GetDim(DIM_NUM_2);
65+ *outShape = *queryShape;
M
Mmonologue8153月12日

40行获取queryShape指针之后没有对指针进行校验，此处若为空指针则可能发生段错误。需要增加校验

likedislike
D
dingxu
3月12日 评论：
66+ 
67+ gert::Shape *topkIndicesOutShape = context->GetOutputShape(TOPK_OUTPUT_INDEX);
68+ OP_CHECK_NULL_WITH_CONTEXT(context, topkIndicesOutShape);
69+ *topkIndicesOutShape = *queryShape;
70+ 
71+ if (layOutStr == "TND") {
72+ outShape->SetDim(DIM_NUM_2, shapeD2 / kvHeadNums);
73+ topkIndicesOutShape->SetDim(DIM_NUM_1, kvHeadNums);
74+ topkIndicesOutShape->SetDim(DIM_NUM_2, selectBlockCount);
75+ } else if (layOutStr == "BSND") {
76+ outShape->SetDim(DIM_NUM_3, shapeD2 / kvHeadNums);
77+ topkIndicesOutShape->SetDim(DIM_NUM_2, kvHeadNums);
78+ topkIndicesOutShape->SetDim(DIM_NUM_3, selectBlockCount);
79+ }
80+ return GRAPH_SUCCESS;
81+}
82+ 
83+static ge::graphStatus InferDataTypeNsaCompressAttentionInfer(gert::InferDataTypeContext *context)
84+{
85+ const auto inputDataType = context->GetInputDataType(0);
86+ context->SetOutputDataType(0, inputDataType);
87+ context->SetOutputDataType(1, ge::DataType::DT_INT32);
88+ return GRAPH_SUCCESS;
89+}
90+ 
91+IMPL_OP_INFERSHAPE(NsaCompressAttentionInfer)
92+ .InferShape(InferShapeNsaCompressAttentionInfer)
93+ .InferDataType(InferDataTypeNsaCompressAttentionInfer);
94+} // namespace ops
Aexperimental/attention/nsa_compress_attention_infer/op_host/nsa_compress_attention_infer_tiling.h+269-0
@@ -0,0 +1,269 @@
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+/*!
12+ * \file nsa_compress_attention_infer_tiling.h
13+ * \brief
14+ */
15+ 
16+#ifndef NSA_COMPRESS_ATTENTION_INFER_TILING_H
17+#define NSA_COMPRESS_ATTENTION_INFER_TILING_H
18+#include <cstdint>
19+#include "register/tilingdata_base.h"
20+#include "tiling/platform/platform_ascendc.h"
21+#include "tiling/tiling_api.h"
22+#include "register/op_def_registry.h"
23+ 
24+namespace optiling {
25+ 
26+BEGIN_TILING_DATA_DEF(NsaCompressAttentionInferBaseParams)
27+TILING_DATA_FIELD_DEF(uint32_t, batchSize)
28+TILING_DATA_FIELD_DEF(uint32_t, qSeqSize)
29+TILING_DATA_FIELD_DEF(uint32_t, seqSize)
30+TILING_DATA_FIELD_DEF(uint32_t, qHeadNum)
31+TILING_DATA_FIELD_DEF(uint32_t, kvHeadNum)
32+TILING_DATA_FIELD_DEF(uint32_t, headSizeQk)
33+TILING_DATA_FIELD_DEF(uint32_t, headSizeVo)
34+TILING_DATA_FIELD_DEF(uint32_t, blockSize)
35+TILING_DATA_FIELD_DEF(uint32_t, groupSize)
36+TILING_DATA_FIELD_DEF(uint32_t, maxBlockNumPerBatch)
37+TILING_DATA_FIELD_DEF(float, scaleValue)
38+TILING_DATA_FIELD_DEF(uint32_t, attenMaskFlag)
39+TILING_DATA_FIELD_DEF(uint32_t, selectSize)
40+TILING_DATA_FIELD_DEF(uint32_t, selectNum)
41+TILING_DATA_FIELD_DEF(uint32_t, compSizeL)
42+TILING_DATA_FIELD_DEF(uint32_t, compStrideD)
43+TILING_DATA_FIELD_DEF(uint32_t, ubSize)
44+TILING_DATA_FIELD_DEF(uint32_t, topKInWorkSpaceSize);
45+ 
46+TILING_DATA_FIELD_DEF(uint32_t, tilingKey);
47+// 新增
48+TILING_DATA_FIELD_DEF(uint32_t, workSpaceElemNum);
49+TILING_DATA_FIELD_DEF(uint32_t, mm1ResWorkSpaceSize);
50+TILING_DATA_FIELD_DEF(uint32_t, mm2InWorkSpaceSize);
51+TILING_DATA_FIELD_DEF(uint32_t, scoreInWorkSpaceSize);
52+TILING_DATA_FIELD_DEF(uint32_t, impScoreResultEleNum);
53+TILING_DATA_FIELD_DEF(uint32_t, impScoreResultCol);
54+TILING_DATA_FIELD_DEF(uint32_t, impScoreResultColPad);
55+TILING_DATA_FIELD_DEF(uint32_t, impScoreReduceRound);
56+TILING_DATA_FIELD_DEF(uint32_t, impScoreReduceCol);
57+TILING_DATA_FIELD_DEF(uint32_t, impScoreReduceColPad);
58+TILING_DATA_FIELD_DEF(uint32_t, impScoreReduceTailCol);
59+TILING_DATA_FIELD_DEF(uint32_t, impScoreReduceTailColPad);
60+TILING_DATA_FIELD_DEF(uint32_t, impScoreWorkSpaceSize);
61+TILING_DATA_FIELD_DEF(uint32_t, impScoreParamValidCol);
62+TILING_DATA_FIELD_DEF(uint32_t, impScoreParamHasPrefix);
63+END_TILING_DATA_DEF
64+REGISTER_TILING_DATA_CLASS(NsaCompressAttentionInferBaseParamsOp, NsaCompressAttentionInferBaseParams)
65+ 
66+BEGIN_TILING_DATA_DEF(NsaCompressAttentionInferSplitBNParams)
67+TILING_DATA_FIELD_DEF(uint32_t, coreNumUsed)
68+TILING_DATA_FIELD_DEF(uint32_t, kvHeadSplitSize)
69+TILING_DATA_FIELD_DEF(uint32_t, kvHeadSplitNum)
70+TILING_DATA_FIELD_DEF(uint32_t, qSeqLenSplitSize)
71+TILING_DATA_FIELD_DEF(uint32_t, processPerBatch)
72+TILING_DATA_FIELD_DEF(uint32_t, processNum)
73+ 
74+TILING_DATA_FIELD_DEF(uint32_t, bnTotalSize)
75+TILING_DATA_FIELD_DEF(uint32_t, headCoreNum)
76+TILING_DATA_FIELD_DEF(uint32_t, tailCoreNum)
77+TILING_DATA_FIELD_DEF(uint32_t, bnPerHeadCore)
78+TILING_DATA_FIELD_DEF(uint32_t, bnPerTailCore)
79+TILING_DATA_FIELD_DEF(uint32_t, rowLenPerHeadCore)
80+TILING_DATA_FIELD_DEF(uint32_t, rowLenPerTailCore)
81+TILING_DATA_FIELD_DEF(uint32_t, baseRowLenHeadCore)
82+TILING_DATA_FIELD_DEF(uint32_t, baseRowLenTailCore)
83+END_TILING_DATA_DEF
84+REGISTER_TILING_DATA_CLASS(NsaCompressAttentionInferSplitBNParamsOp, NsaCompressAttentionInferSplitBNParams)
85+ 
86+BEGIN_TILING_DATA_DEF(NsaCompressAttentionInferTilingData)
87+TILING_DATA_FIELD_DEF_STRUCT(NsaCompressAttentionInferBaseParams, baseParams);
88+TILING_DATA_FIELD_DEF_STRUCT(NsaCompressAttentionInferSplitBNParams, splitBNParams);
89+TILING_DATA_FIELD_DEF_STRUCT(SoftMaxTiling, softmaxTilingData);
90+TILING_DATA_FIELD_DEF_STRUCT(TopkTiling, topkTilingData);
91+END_TILING_DATA_DEF;
92+REGISTER_TILING_DATA_CLASS(NsaCompressAttentionInfer, NsaCompressAttentionInferTilingData)
93+ 
94+struct NsaCompressAttentionInferCompileInfo {
95+ uint32_t inputDataByte = 2;
96+ ge::DataType inputDataType;
97+ 
98+ uint32_t dataNumSingleUb = 1; // UB空间可处理的最大数据量
99+ uint32_t blockNum = 1; // 32B对齐使用 //
100+ uint32_t cacheLineLen = 1; // 512B对齐使用
101+ 
102+ uint32_t coreNum = 1;
103+ uint32_t aivNum = 0;
104+ uint32_t aicNum = 0;
105+ uint64_t ubSize = 0;
106+ uint64_t l1Size = 0;
107+ uint64_t l0ASize = 0;
108+ uint64_t l0BSize = 0;
109+ uint64_t l0CSize = 0;
110+ uint64_t sysWorkspaceSize = 0;
111+ platform_ascendc::SocVersion socVersion;
112+};
113+ 
114+ 
115+struct RequiredParaInfo {
116+ const gert::CompileTimeTensorDesc *desc;
117+ const gert::StorageShape *shape;
118+};
119+ 
120+struct OptionalParaInfo {
121+ const gert::CompileTimeTensorDesc *desc;
122+ const gert::Tensor *tensor;
123+};
124+ 
125+struct NsaCompressAttentionInferContext {
126+ const char *opName;
127+ fe::PlatFormInfos *platformInfo;
128+ RequiredParaInfo query;
129+ RequiredParaInfo key;
130+ RequiredParaInfo value;
131+ OptionalParaInfo attenMask;
132+ OptionalParaInfo blockTable;
133+ OptionalParaInfo actualQSeqLengths;
134+ OptionalParaInfo actualCmpKvSeqLengths;
135+ OptionalParaInfo actualSelKvSeqLengths;
136+ OptionalParaInfo topkMask;
137+ 
138+ RequiredParaInfo attenOut;
139+ RequiredParaInfo selectOut;
140+ 
141+ const uint32_t *numHeads;
142+ const uint32_t *kvHeadNums;
143+ const uint32_t *selectSize;
144+ const uint32_t *selectNum;
145+ const uint32_t *compSizeL;
146+ const uint32_t *compStrideD;
147+ const float *scaleValue;
148+ 
149+ const char *layOut;
150+ const uint32_t *blockSize;
151+ const uint32_t *sparseMode;
152+ 
153+ size_t *workSpaces;
154+ uint64_t tilingKey;
155+ uint32_t blockDim;
156+};
157+ 
158+enum NSICALayout : uint32_t {
159+ TND = 0,
160+};
161+ 
162+class NCAITiling {
163+public:
164+ NCAITiling() = default;
165+ ~NCAITiling() = default;
166+ 
167+ ge::graphStatus GetNCAITiling(NsaCompressAttentionInferContext &ncaiContext,
168+ NsaCompressAttentionInferTilingData &tilingData, bool isWorkspace = false);
169+ ge::graphStatus NSICASetTilingData(gert::TilingContext &context,
170+ NsaCompressAttentionInferTilingData &tilingData);
171+ static ge::graphStatus ConvertContext(gert::TilingContext &context,
172+ NsaCompressAttentionInferContext &ncaiContext);
173+private:
174+ ge::graphStatus GetNpuInfo();
175+ ge::graphStatus Split();
176+ ge::graphStatus SplitBN();
177+ ge::graphStatus ProcessQkv();
178+ ge::graphStatus CheckQkvShape();
179+ ge::graphStatus CheckAttr();
180+ ge::graphStatus CheckTopK();
181+ ge::graphStatus ProcessEpilogue();
182+ ge::graphStatus ProcessInput();
183+ ge::graphStatus CheckQkv();
184+ ge::graphStatus ParamsPostCheck();
185+ ge::graphStatus SoftmaxTiling();
186+ ge::graphStatus TopKTiling();
187+ void FillImpScoreTilingData();
188+ ge::graphStatus FillTilingData();
189+ ge::graphStatus CalcWorkSpace();
190+ ge::graphStatus CalcTilingKey();
191+ ge::graphStatus GetMaxMinSeqlen();
192+ ge::graphStatus CheckSelKvSeqlen();
193+ ge::graphStatus GetMaxQSeqlen();
194+ ge::graphStatus CheckAttenMask();
195+ ge::graphStatus CheckImpScoreParam();
196+ 
197+private:
198+ uint32_t batchSize_ = 0;
199+ uint32_t qSeqSize_ = 0;
200+ uint32_t sMax_ = 0;
201+ uint32_t sMin_ = 0;
202+ uint32_t qHeadNum_ = 0;
203+ uint32_t kvHeadNum_ = 0;
204+ uint32_t groupSize_ = 0;
205+ uint32_t headSizeQk_ = 0;
206+ uint32_t headSizeVo_ = 0;
207+ uint32_t blockSize_ = 0;
208+ uint32_t maxBlockNumPerBatch_ = 0;
209+ float scaleValue_ = 0;
210+ uint32_t attenMaskFlag_ = 0;
211+ uint32_t selectSize_ = 0;
212+ uint32_t selectNum_ = 0;
213+ uint32_t compSizeL_ = 0;
214+ uint32_t compStrideD_ = 0;
215+ uint32_t kvHeadSplitSize_ = 0;
216+ uint32_t kvHeadSplitNum_ = 0;
217+ uint32_t qSeqLenSplitSize_ = 1;
218+ uint32_t qSeqLenCumSum_ = 0;
219+ uint32_t processPerBatch_ = 0;
220+ uint32_t processNum_ = 0;
221+ bool pagedAttentionFlag_ = true;
222+ 
223+ bool isWorkspace_ = false;
224+ uint32_t blockDim_ = 0;
225+ uint32_t aivNum_ = 0;
226+ uint32_t aicNum_ = 0;
227+ uint32_t libapiSize_ = 0;
228+ uint32_t workSpaceSize_ = 0;
229+ uint64_t ubSize_ = 0;
230+ 
231+ uint32_t bnTotalSize_ = 0;
232+ uint32_t coreNumUsed_ = 0;
233+ uint32_t headCoreNum_ = 0;
234+ uint32_t tailCoreNum_ = 0;
235+ uint32_t bnPerHeadCore_ = 0;
236+ uint32_t bnPerTailCore_ = 0;
237+ 
238+ uint32_t sInnerLoopTimes_ = 0;
239+ 
240+ uint32_t workSpaceElemNum_ = 0;
241+ uint32_t mm1ResWorkSpaceSize_ = 0;
242+ uint32_t mm2InWorkSpaceSize_ = 0;
243+ uint32_t scoreInWorkSpaceSize_ = 0;
244+ uint32_t topKInWorkSpaceSize_ = 0;
245+ 
246+ bool useImpScoreOpt_ = false;
247+ uint32_t impScoreParamValidCol_ = 0;
248+ bool impScoreParamHasPrefix_ = true;
249+ uint32_t impScoreReduceRound_ = 0;
250+ uint32_t impScoreReduceCol_ = 0;
251+ uint32_t impScoreReduceColPad_ = 0;
252+ uint32_t impScoreReduceTailCol_ = 0;
253+ uint32_t impScoreReduceTailColPad_ = 0;
254+ uint32_t impScoreResultEleNum_ = 0;
255+ uint32_t impScoreResultCol_ = 0;
256+ uint32_t impScoreResultColPad_ = 0;
257+ uint32_t impScoreWorkSpaceSize_ = 0;
258+ 
259+ uint32_t rowLenPerHeadCore_ = 0;
260+ uint32_t rowLenPerTailCore_ = 0;
261+ uint32_t baseRowLenHeadCore_ = 0;
262+ uint32_t baseRowLenTailCore_ = 0;
263+ 
264+ NsaCompressAttentionInferContext *ncaiContext_ = nullptr;
265+ NsaCompressAttentionInferTilingData *tilingData_ = nullptr;
266+};
267+ 
268+} // namespace optiling
269+#endif // NSA_COMPRESS_ATTENTION_INFER_TILING_H
Aexperimental/attention/nsa_compress_attention_infer/op_host/nsa_compress_attention_infer_tiling_utils.h+75-0
@@ -0,0 +1,75 @@
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+/*!
12+ * \file nsa_compress_attention_infer_tiling_utils.h
13+ * \brief
14+ */
15+ 
16+#ifndef NSA_COMPRESS_ATTENTION_INFER_TILING_UTILS_H
17+#define NSA_COMPRESS_ATTENTION_INFER_TILING_UTILS_H
18+ 
19+#include "nsa_compress_attention_infer_tiling.h"
20+ 
21+namespace optiling {
22+ 
23+constexpr uint32_t ZERO = 0;
24+constexpr uint32_t ONE = 1;
25+ 
26+constexpr size_t DIM_0 = 0;
27+constexpr size_t DIM_1 = 1;
28+constexpr size_t DIM_2 = 2;
29+constexpr size_t DIM_3 = 3;
30+ 
31+constexpr uint32_t QUERY_INPUT_INDEX = 0;
32+constexpr uint32_t KEY_INPUT_INDEX = 1;
33+constexpr uint32_t VALUE_INPUT_INDEX = 2;
34+constexpr uint32_t ATTEN_MASK_INPUT_INDEX = 3;
35+constexpr uint32_t BLOCK_TABLE_INPUT_INDEX = 4;
36+constexpr uint32_t ACT_Q_SEQ_LEN_INPUT_INDEX = 5;
37+constexpr uint32_t ACT_CMP_KV_SEQ_LEN_INPUT_INDEX = 6;
38+constexpr uint32_t ACT_SEL_KV_SEQ_LEN_INPUT_INDEX = 7;
39+constexpr uint32_t TOPK_MASK_INPUT_INDEX = 8;
40+ 
41+constexpr uint32_t ATTEN_OUTPUT_INDEX = 0;
42+constexpr uint32_t SELECT_OUTPUT_INDEX = 1;
43+ 
44+constexpr uint32_t NUM_HEADS_ATTR_INDEX = 0;
45+constexpr uint32_t KV_NUM_HEADS_ATTR_INDEX = 1;
46+constexpr uint32_t SELECT_SIZE_ATTR_INDEX = 2;
47+constexpr uint32_t SELECT_NUM_ATTR_INDEX = 3;
48+constexpr uint32_t COMP_SIZE_ATTR_INDEX = 4;
49+constexpr uint32_t COMP_STRIDE_ATTR_INDEX = 5;
50+constexpr uint32_t SCALE_VALUE_ATTR_INDEX = 6;
51+ 
52+constexpr uint32_t LAYOUT_ATTR_INDEX = 7;
53+constexpr uint32_t BLOCK_SIZE_ATTR_INDEX = 8;
54+constexpr uint32_t SPARSE_MODE_ATTR_INDEX = 9;
55+ 
56+constexpr uint32_t ROW_LIMIT_QP = 128;
57+constexpr uint32_t SEQLEN_LIMIT = 8192;
58+ 
59+constexpr uint32_t SINGLE_UB_SIZE_BFLOAT16 = 4; // 分成4份
60+constexpr uint32_t SOFTMAX_INPUT_DATA_BYTE = 4;
61+constexpr uint32_t COMPARE_INT = 64;
62+ 
63+constexpr uint32_t BLOCK_SIZE = 16;
64+constexpr uint32_t COLLEN_ALIGNED_REQUIRED = 16;
65+constexpr uint32_t NUM_VECTOR_PER_CUBE = 2;
66+constexpr uint32_t ALIGNED_32 = 32;
67+constexpr uint32_t ALIGNED_8 = 8;
68+constexpr uint32_t IMP_SCORE_PARAM_ROW = 128;
69+ 
70+constexpr int64_t MAX_ACTUALQSEQLEN = 4;
71+constexpr int64_t MIN_ACTUALQSEQLEN = 1;
72+ 
73+} // namespace optiling
74+ 
75+#endif // NSA_COMPRESS_ATTENTION_INFER_TILING_UTILS_H
Aexperimental/attention/nsa_compress_attention_infer/op_host/op_api/aclnn_nsa_compress_attention_infer.cpp+342-0
@@ -0,0 +1,342 @@
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+#include <iostream>
12+#include <algorithm>
13+#include <unordered_map>
14+#include "nsa_compress_attention_infer.h"
15+#include "aclnn_nsa_compress_attention_infer.h"
16+#include "aclnn_kernels/contiguous.h"
17+#include "aclnn_kernels/pad.h"
18+#include "aclnn_kernels/reshape.h"
19+#include "aclnn_kernels/slice.h"
20+#include "aclnn_kernels/transpose.h"
21+#include "opdev/common_types.h"
22+#include "opdev/fast_vector.h"
23+#include "opdev/op_errno.h"
24+#include "opdev/op_executor.h"
25+ 
26+using namespace op;
27+ 
28+#ifdef __cplusplus
29+extern "C" {
30+#endif
31+ 
32+namespace {
33+ 
34+__attribute__((visibility("default"))) aclnnStatus aclnnNsaCompressAttentionInferGetMaxWorkspaceSize(
35+ const aclTensor *query,
36+ const aclTensor *key,
37+ const aclTensor *value,
38+ const aclTensor *attentionMaskOptional,
39+ const aclTensor *blockTableOptional,
40+ const aclIntArray *actualQSeqLenOptional,
41+ const aclIntArray *actualCmpKvSeqLenOptional,
42+ const aclIntArray *actualSelKvSeqLenOptional,
43+ const aclTensor *topKMaskOptional,
44+ int64_t numHeads,
45+ int64_t numKeyValueHeads,
46+ int64_t selectBlockSize,
47+ int64_t selectBlockCount,
48+ int64_t compressBlockSize,
49+ int64_t compressBlockStride,
50+ double scaleValue,
51+ char *layoutOptional,
52+ int64_t pageBlockSize,
53+ int64_t sparseMode,
54+ const aclTensor *output,
55+ const aclTensor *topKOutput,
56+ uint64_t *workspaceSize,
57+ aclOpExecutor **executor);
58+ 
59+aclnnStatus nsaCompressAttentionInferContiguous(const aclTensor *&query, const aclTensor *&key, const aclTensor *&value,
60+ const aclTensor *attentionMaskOptional, const aclTensor *&blockTableOptional,
61+ const aclTensor *&topKMaskOptional,
62+ aclOpExecutor *executor)
63+{
64+ query = l0op::Contiguous(query, executor);
65+ CHECK_RET(query != nullptr, ACLNN_ERR_PARAM_NULLPTR);
66+ key = l0op::Contiguous(key, executor);
67+ CHECK_RET(key != nullptr, ACLNN_ERR_PARAM_NULLPTR);
68+ value = l0op::Contiguous(value, executor);
69+ CHECK_RET(value != nullptr, ACLNN_ERR_PARAM_NULLPTR);
70+ if (attentionMaskOptional) {
71+ attentionMaskOptional = l0op::Contiguous(attentionMaskOptional, executor);
72+ CHECK_RET(attentionMaskOptional != nullptr, ACLNN_ERR_INNER_NULLPTR);
73+ }
74+ if (blockTableOptional) {
75+ blockTableOptional = l0op::Contiguous(blockTableOptional, executor);
76+ CHECK_RET(blockTableOptional != nullptr, ACLNN_ERR_INNER_NULLPTR);
77+ }
78+ if (topKMaskOptional) {
79+ topKMaskOptional = l0op::Contiguous(topKMaskOptional, executor);
80+ CHECK_RET(topKMaskOptional != nullptr, ACLNN_ERR_INNER_NULLPTR);
81+ }
82+ return ACLNN_SUCCESS;
83+}
84+ 
85+static bool nsaCompressAttentionInferCheckDataType(
86+ const aclTensor *query,
87+ const aclTensor *key,
88+ const aclTensor *value,
89+ const aclTensor *blockTableOptional,
90+ const aclTensor *output,
91+ const aclTensor *topKOutput)
92+{
93+ auto qDtype = query->GetDataType();
94+ auto kDtype = key->GetDataType();
95+ auto vDtype = value->GetDataType();
96+ auto blockTableDtype = blockTableOptional->GetDataType();
97+ auto attnOutDtype = output->GetDataType();
98+ auto topKOutDtype = topKOutput->GetDataType();
99+ 
100+ static const std::unordered_map<DataType, std::vector<DataType>> dTypeMappingQtoKv = {
101+ {DataType::DT_FLOAT16, {DataType::DT_FLOAT16}},
102+ {DataType::DT_BF16, {DataType::DT_BF16}},
103+ };
104+ if (dTypeMappingQtoKv.find(qDtype) == dTypeMappingQtoKv.end()) {
105+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Input Q data type is invalid, please check.");
106+ return false;
107+ } else {
108+ auto validKvDtypeList = dTypeMappingQtoKv.at(qDtype);
109+ if (std::find(validKvDtypeList.begin(), validKvDtypeList.end(), kDtype) == validKvDtypeList.end() ||
110+ std::find(validKvDtypeList.begin(), validKvDtypeList.end(), vDtype) == validKvDtypeList.end()) {
111+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Input K or V data type is invalid, please check.");
112+ return false;
113+ }
114+ if (attnOutDtype != qDtype) {
115+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Output attnOut data type is invalid, please check.");
116+ return false;
117+ }
118+ }
119+ if (blockTableDtype != DataType::DT_INT32) {
120+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "blockTable data type must be int32, please check.");
121+ return false;
122+ }
123+ if (topKOutDtype != DataType::DT_INT32) {
124+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Output topkOut data type must be int32, please check.");
125+ return false;
126+ }
127+ return true;
128+}
129+ 
130+aclnnStatus nsaCompressAttentionInferValidateParam(
131+ const aclTensor *query,
132+ const aclTensor *key,
133+ const aclTensor *value,
134+ const aclTensor *blockTableOptional,
135+ const aclTensor *output,
136+ const aclTensor *topKOutput)
137+{
138+ Shape qShape = query->GetViewShape();
139+ auto batchSize = qShape[0];
140+ if (batchSize > 10000U || batchSize < 1U) {
141+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "batchSize %ld not in [1, 10000], please check", batchSize);
142+ return ACLNN_ERR_PARAM_INVALID;
143+ }
144+ CHECK_RET(nsaCompressAttentionInferCheckDataType(query, key, value, blockTableOptional, output, topKOutput),
145+ ACLNN_ERR_PARAM_INVALID);
146+ return ACLNN_SUCCESS;
147+}
148+ 
149+aclnnStatus nsaCompressAttentionInferCheckTensorNull(
150+ const aclTensor *query,
151+ const aclTensor *key,
152+ const aclTensor *value,
153+ const aclTensor *blockTableOptional,
154+ const aclIntArray *actualCmpKvSeqLenOptional,
155+ const char *layoutOptional,
156+ const aclTensor *output,
157+ const aclTensor *topKOutput,
158+ const uint64_t *workspaceSize)
159+{
160+ // 必须的参数指针判空
161+ CHECK_RET(query != nullptr, ACLNN_ERR_PARAM_NULLPTR);
162+ CHECK_RET(key != nullptr, ACLNN_ERR_PARAM_NULLPTR);
163+ CHECK_RET(value != nullptr, ACLNN_ERR_PARAM_NULLPTR);
164+ CHECK_RET(blockTableOptional != nullptr, ACLNN_ERR_PARAM_NULLPTR);
165+ CHECK_RET(actualCmpKvSeqLenOptional != nullptr, ACLNN_ERR_PARAM_NULLPTR);
166+ CHECK_RET(layoutOptional != nullptr, ACLNN_ERR_PARAM_NULLPTR);
167+ CHECK_RET(output != nullptr, ACLNN_ERR_PARAM_NULLPTR);
168+ CHECK_RET(topKOutput != nullptr, ACLNN_ERR_PARAM_NULLPTR);
169+ CHECK_RET(workspaceSize != nullptr, ACLNN_ERR_PARAM_NULLPTR);
170+ return ACLNN_SUCCESS;
171+}
172+ 
173+aclnnStatus FakeArray_cai(const aclIntArray *inArray, aclIntArray *&outArray)
174+{
175+ OP_LOGD("start fake array");
176+ if (inArray != nullptr) {
177+ OP_LOGD("input array is not nullptr");
178+ uint64_t size = inArray->Size();
179+ // tiling侧认为有tensor但没有data就是计算最大workspace
180+ std::vector<int64_t> fake_array(size, -1);
181+ outArray = aclCreateIntArray(fake_array.data(), size);
182+ if (outArray == nullptr) {
183+ OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "Try alloc array failed");
184+ return ACLNN_ERR_INNER_NULLPTR;
185+ }
186+ }
187+ OP_LOGD("end fake array");
188+ return ACLNN_SUCCESS;
189+}
190+ 
191+aclnnStatus aclnnNsaCompressAttentionInferGetMaxWorkspaceSize(
192+ const aclTensor *query,
193+ const aclTensor *key,
194+ const aclTensor *value,
195+ const aclTensor *attentionMaskOptional,
196+ const aclTensor *blockTableOptional,
197+ const aclIntArray *actualQSeqLenOptional,
198+ const aclIntArray *actualCmpKvSeqLenOptional,
199+ const aclIntArray *actualSelKvSeqLenOptional,
200+ const aclTensor *topKMaskOptional,
201+ int64_t numHeads,
202+ int64_t numKeyValueHeads,
203+ int64_t selectBlockSize,
204+ int64_t selectBlockCount,
205+ int64_t compressBlockSize,
206+ int64_t compressBlockStride,
207+ double scaleValue,
208+ char *layoutOptional,
209+ int64_t pageBlockSize,
210+ int64_t sparseMode,
211+ const aclTensor *output,
212+ const aclTensor *topKOutput,
213+ uint64_t *workspaceSize,
214+ aclOpExecutor **executor)
215+{
216+ aclIntArray *fakeActualCmpKvSeqLenOptional{nullptr};
217+ // nullptr不处理， nullptr是空指针，这样不会影响原来就不传入actual seq length为空的逻辑
218+ aclnnStatus ret = FakeArray_cai(actualCmpKvSeqLenOptional, fakeActualCmpKvSeqLenOptional);
219+ CHECK_RET_CODE(ret, "Try alloc fake actualSeqLengthsKvOptional failed");
220+ ret = aclnnNsaCompressAttentionInferGetWorkspaceSize(query,
221+ key,
222+ value,
223+ attentionMaskOptional,
224+ blockTableOptional,
225+ actualQSeqLenOptional,
226+ fakeActualCmpKvSeqLenOptional,
227+ actualSelKvSeqLenOptional,
228+ topKMaskOptional,
229+ numHeads,
230+ numKeyValueHeads,
231+ selectBlockSize,
232+ selectBlockCount,
233+ compressBlockSize,
234+ compressBlockStride,
235+ scaleValue,
236+ layoutOptional,
237+ pageBlockSize,
238+ sparseMode,
239+ output,
240+ topKOutput,
241+ workspaceSize,
242+ executor);
243+ aclDestroyIntArray(fakeActualCmpKvSeqLenOptional);
244+ return ret;
245+ }
246+ 
247+aclnnStatus aclnnNsaCompressAttentionInferGetWorkspaceSize(
248+ const aclTensor *query,
249+ const aclTensor *key,
250+ const aclTensor *value,
251+ const aclTensor *attentionMaskOptional,
252+ const aclTensor *blockTableOptional,
253+ const aclIntArray *actualQSeqLenOptional,
254+ const aclIntArray *actualCmpKvSeqLenOptional,
255+ const aclIntArray *actualSelKvSeqLenOptional,
256+ const aclTensor *topKMaskOptional,
257+ int64_t numHeads,
258+ int64_t numKeyValueHeads,
259+ int64_t selectBlockSize,
260+ int64_t selectBlockCount,
261+ int64_t compressBlockSize,
262+ int64_t compressBlockStride,
263+ double scaleValue,
264+ char *layoutOptional,
265+ int64_t pageBlockSize,
266+ int64_t sparseMode,
267+ const aclTensor *output,
268+ const aclTensor *topKOutput,
269+ uint64_t *workspaceSize,
270+ aclOpExecutor **executor)
271+{
272+ CHECK_RET(nsaCompressAttentionInferCheckTensorNull(query, key, value, blockTableOptional,
273+ actualCmpKvSeqLenOptional, layoutOptional, output, topKOutput,
274+ workspaceSize) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_NULLPTR);
275+ CHECK_RET(nsaCompressAttentionInferValidateParam(query, key, value, blockTableOptional,
276+ output, topKOutput) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID);
277+ L2_DFX_PHASE_1(aclnnNsaCompressAttentionInfer,
278+ DFX_IN(query, key, value, attentionMaskOptional, blockTableOptional, actualQSeqLenOptional,
279+ actualCmpKvSeqLenOptional, actualSelKvSeqLenOptional, topKMaskOptional, numHeads, numKeyValueHeads,
280+ selectBlockSize, selectBlockCount, compressBlockSize, compressBlockStride, scaleValue,
281+ layoutOptional, pageBlockSize, sparseMode),
282+ DFX_OUT(output, topKOutput));
283+ 
284+ auto uniqueExecutor = CREATE_EXECUTOR();
285+ CHECK_RET(uniqueExecutor.get() != nullptr, ACLNN_ERR_INNER_CREATE_EXECUTOR);
286+ if (output->IsEmpty() && topKOutput->IsEmpty()) {
287+ *workspaceSize = 0;
288+ uniqueExecutor.ReleaseTo(executor);
289+ return ACLNN_SUCCESS;
290+ }
291+ if (strcmp(layoutOptional, "TND") != 0 && strcmp(layoutOptional, "BSND") != 0) {
292+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "Layout %s is not TND and BSND, invalid shape, please check", layoutOptional);
293+ *workspaceSize = 0;
294+ uniqueExecutor.ReleaseTo(executor);
295+ return ACLNN_ERR_PARAM_INVALID;
296+ }
297+ 
298+ aclOpExecutor *l0Executor = uniqueExecutor.get();
299+ 
300+ CHECK_RET(nsaCompressAttentionInferContiguous(query, key, value, attentionMaskOptional, blockTableOptional,
301+ topKMaskOptional, l0Executor) == ACLNN_SUCCESS,
302+ ACLNN_ERR_INNER_NULLPTR);
303+ string inputLayoutStr = op::ToString(layoutOptional).GetString();
304+ auto l0NsaCompressAttentionInferOuts = l0op::NsaCompressAttentionInfer(
305+ query, key, value, attentionMaskOptional, blockTableOptional,
306+ actualQSeqLenOptional, actualCmpKvSeqLenOptional, actualSelKvSeqLenOptional, topKMaskOptional, numHeads,
307+ numKeyValueHeads, selectBlockSize, selectBlockCount, compressBlockSize, compressBlockStride, scaleValue,
308+ inputLayoutStr.c_str(), pageBlockSize, sparseMode, l0Executor);
309+ 
310+ CHECK_RET(l0NsaCompressAttentionInferOuts[0] != nullptr, ACLNN_ERR_INNER_NULLPTR);
311+ CHECK_RET(l0NsaCompressAttentionInferOuts[1] != nullptr, ACLNN_ERR_INNER_NULLPTR);
312+ auto l0output = l0NsaCompressAttentionInferOuts[0];
313+ auto l0topKOutput = l0NsaCompressAttentionInferOuts[1];
314+ 
315+ if (l0output == nullptr || l0topKOutput == nullptr) {
316+ OP_LOGE(ACLNN_ERR_INNER_NULLPTR, "l0output or l0topKOutput is null");
317+ *workspaceSize = 0;
318+ uniqueExecutor.ReleaseTo(executor);
319+ return ACLNN_ERR_INNER_NULLPTR;
320+ }
321+ auto viewCopyResult0 = l0op::ViewCopy(l0output, output, l0Executor);
322+ CHECK_RET(viewCopyResult0 != nullptr, ACLNN_ERR_INNER_NULLPTR);
323+ auto viewCopyResult1 = l0op::ViewCopy(l0topKOutput, topKOutput, l0Executor);
324+ CHECK_RET(viewCopyResult1 != nullptr, ACLNN_ERR_INNER_NULLPTR);
325+ 
326+ *workspaceSize = uniqueExecutor->GetWorkspaceSize();
327+ uniqueExecutor.ReleaseTo(executor);
328+ return ACLNN_SUCCESS;
329+}
330+ 
331+aclnnStatus aclnnNsaCompressAttentionInfer(void *workspace, uint64_t workspaceSize, aclOpExecutor *executor,
332+ aclrtStream stream)
333+{
334+ L2_DFX_PHASE_2(aclnnNsaCompressAttentionInfer);
335+ // 固定写法，调用框架能力，完成计算
336+ return CommonOpExecutorRun(workspace, workspaceSize, executor, stream);
337+}
338+} // namespace
339+ 
340+#ifdef __cplusplus
341+}
342+#endif
Aexperimental/attention/nsa_compress_attention_infer/op_host/op_api/aclnn_nsa_compress_attention_infer.h+94-0
@@ -0,0 +1,94 @@
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 OP_API_INC_LEVEL2_ACLNN_NSA_COMPRESS_ATTENTION_INFER_H_
12+#define OP_API_INC_LEVEL2_ACLNN_NSA_COMPRESS_ATTENTION_INFER_H_
13+ 
14+#include "aclnn/aclnn_base.h"
15+ 
16+#ifdef __cplusplus
17+extern "C" {
18+#endif
19+ 
20+/**
21+ * @brief aclnnNsaCompressAttentionInfer的第一段接口，根据具体的计算流程，计算workspace大小。
22+ * @domain aclnn_ops_infer
23+ * funtion: aclnnNsaCompressAttentionInferGetWorkspaceSize
24+ * param [in] query : required
25+ * param [in] key : required
26+ * param [in] value : required
27+ * param [in] attentionMaskOptional : optional
28+ * param [in] blockTableOptional : optional
29+ * param [in] actualQSeqLenOptional : optional
30+ * param [in] actualCmpKvSeqLenOptional : optional
31+ * param [in] actualSelKvSeqLenOptional : optional
32+ * param [in] topkMaskOptional : optional
33+ * param [in] numHeads : optional
34+ * param [in] numKeyValueHeads : optional
35+ * param [in] selectBlockSize : optional
36+ * param [in] selectBlockCount : optional
37+ * param [in] compressBlockSize : optional
38+ * param [in] compressStride : optional 【aclnn中是compressBlockStride】
39+ * param [in] scaleValue : optional
40+ * param [in] layoutOptional : optional
41+ * param [in] pageBlockSize : optional
42+ * param [in] sparseMode : optional
43+ * @param [out] output : required
44+ * @param [out] topkIndicesOut : required 【aclnn中是topKOutput】
45+ * @param [out] workspaceSize : size of workspace(output).
46+ * @param [out] executor : executor context(output).
47+ * @return aclnnStatus: 返回状态码
48+ */
49+ 
50+aclnnStatus aclnnNsaCompressAttentionInferGetWorkspaceSize(
51+ const aclTensor *query,
52+ const aclTensor *key,
53+ const aclTensor *value,
54+ const aclTensor *attentionMaskOptional,
55+ const aclTensor *blockTableOptional,
56+ const aclIntArray *actualQSeqLenOptional,
57+ const aclIntArray *actualCmpKvSeqLenOptional,
58+ const aclIntArray *actualSelKvSeqLenOptional,
59+ const aclTensor *topKMaskOptional,
60+ int64_t numHeads,
61+ int64_t numKeyValueHeads,
62+ int64_t selectBlockSize,
63+ int64_t selectBlockCount,
64+ int64_t compressBlockSize,
65+ int64_t compressBlockStride,
66+ double scaleValue,
67+ char *layoutOptional,
68+ int64_t pageBlockSize,
69+ int64_t sparseMode,
70+ const aclTensor *output,
71+ const aclTensor *topKOutput,
72+ uint64_t *workspaceSize,
73+ aclOpExecutor **executor);
74+ 
75+/**
76+ * funtion: aclnnNsaCompressAttentionInfer
77+ * param [in] workspace : workspace memory addr(input).
78+ * param [in] workspaceSize : size of workspace(input).
79+ * param [in] executor : executor context(input).
80+ * param [in] stream : acl stream.
81+ * @return aclnnStatus: 返回状态码
82+ */
83+ 
84+aclnnStatus aclnnNsaCompressAttentionInfer(
85+ void *workspace,
86+ uint64_t workspaceSize,
87+ aclOpExecutor *executor,
88+ aclrtStream stream);
89+ 
90+#ifdef __cplusplus
91+}
92+#endif
93+ 
94+#endif
Aexperimental/attention/nsa_compress_attention_infer/op_host/op_api/nsa_compress_attention_infer.cpp+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+#include "opdev/make_op_executor.h"
12+#include "opdev/op_dfx.h"
13+ 
14+using namespace op;
15+ 
16+namespace l0op {
17+ 
18+OP_TYPE_REGISTER(NsaCompressAttentionInfer);
19+ 
20+const std::array<const aclTensor *, 2>
21+NsaCompressAttentionInfer(const aclTensor *query,
22+ const aclTensor *key,
23+ const aclTensor *value,
24+ const aclTensor *attentionMaskOptional,
25+ const aclTensor *blockTableOptional,
26+ const aclIntArray *actualQSeqLenOptional,
27+ const aclIntArray *actualCmpKvSeqLenOptional,
28+ const aclIntArray *actualSelKvSeqLenOptional,
29+ const aclTensor *topKMaskOptional,
30+ int64_t numHeads,
31+ int64_t numKeyValueHeads,
32+ int64_t selectBlockSize,
33+ int64_t selectBlockCount,
34+ int64_t compressBlockSize,
35+ int64_t compressBlockStride,
36+ double scaleValue,
37+ const char *layoutOptional,
38+ int64_t pageBlockSize,
39+ int64_t sparseMode,
40+ aclOpExecutor *executor)
41+{
42+ L0_DFX(NsaCompressAttentionInfer, query, key, value, attentionMaskOptional, blockTableOptional,
43+ actualQSeqLenOptional, actualCmpKvSeqLenOptional, actualSelKvSeqLenOptional, topKMaskOptional, numHeads,
44+ numKeyValueHeads, selectBlockSize, selectBlockCount, compressBlockSize, compressBlockStride, scaleValue,
45+ layoutOptional, pageBlockSize, sparseMode);
46+ 
47+ if (attentionMaskOptional == nullptr) {
48+ attentionMaskOptional = executor->AllocTensor(DataType::DT_BOOL, Format::FORMAT_ND, Format::FORMAT_ND);
49+ }
50+ 
51+ const aclTensor *actualSeqQLen = nullptr;
52+ if (actualQSeqLenOptional) {
53+ actualSeqQLen = executor->ConvertToTensor(actualQSeqLenOptional, DataType::DT_INT64);
54+ const_cast<aclTensor *>(actualSeqQLen)->SetStorageFormat(Format::FORMAT_ND);
55+ const_cast<aclTensor *>(actualSeqQLen)->SetViewFormat(Format::FORMAT_ND);
56+ const_cast<aclTensor *>(actualSeqQLen)->SetOriginalFormat(Format::FORMAT_ND);
57+ } else {
58+ actualSeqQLen = executor->AllocTensor(DataType::DT_INT64, Format::FORMAT_ND, Format::FORMAT_ND);
59+ }
60+ 
61+ const aclTensor *actualCmpSeqKvLen = nullptr;
62+ if (actualCmpKvSeqLenOptional) {
63+ actualCmpSeqKvLen = executor->ConvertToTensor(actualCmpKvSeqLenOptional, DataType::DT_INT64);
64+ const_cast<aclTensor *>(actualCmpSeqKvLen)->SetStorageFormat(Format::FORMAT_ND);
65+ const_cast<aclTensor *>(actualCmpSeqKvLen)->SetViewFormat(Format::FORMAT_ND);
66+ const_cast<aclTensor *>(actualCmpSeqKvLen)->SetOriginalFormat(Format::FORMAT_ND);
67+ } else {
68+ actualCmpSeqKvLen = executor->AllocTensor(DataType::DT_INT64, Format::FORMAT_ND, Format::FORMAT_ND);
69+ }
70+ 
71+ const aclTensor *actualSelKvLen = nullptr;
72+ if (actualSelKvSeqLenOptional) {
73+ actualSelKvLen = executor->ConvertToTensor(actualSelKvSeqLenOptional, DataType::DT_INT64);
74+ const_cast<aclTensor *>(actualSelKvLen)->SetStorageFormat(Format::FORMAT_ND);
75+ const_cast<aclTensor *>(actualSelKvLen)->SetViewFormat(Format::FORMAT_ND);
76+ const_cast<aclTensor *>(actualSelKvLen)->SetOriginalFormat(Format::FORMAT_ND);
77+ } else {
78+ actualSelKvLen = executor->AllocTensor(DataType::DT_INT64, Format::FORMAT_ND, Format::FORMAT_ND);
79+ }
80+ 
81+ auto output = executor->AllocTensor(query->GetDataType(), Format::FORMAT_ND, Format::FORMAT_ND);
82+ auto topKOutput = executor->AllocTensor(DataType::DT_INT32, Format::FORMAT_ND, Format::FORMAT_ND);
83+ 
84+ auto ret = INFER_SHAPE(NsaCompressAttentionInfer,
85+ OP_INPUT(query, key, value, attentionMaskOptional, blockTableOptional, actualSeqQLen,
86+ actualCmpSeqKvLen, actualSelKvLen, topKMaskOptional),
87+ OP_OUTPUT(output, topKOutput),
88+ OP_ATTR(numHeads, numKeyValueHeads, selectBlockSize, selectBlockCount, compressBlockSize,
89+ compressBlockStride, static_cast<float>(scaleValue), layoutOptional, pageBlockSize,
90+ sparseMode));
91+ if (ret != ACLNN_SUCCESS) {
92+ OP_LOGE(ACLNN_ERR_PARAM_INVALID, "NsaCompressAttentionInfer InferShape failed.");
93+ return {nullptr, nullptr};
94+ }
95+ 
96+ ADD_TO_LAUNCHER_LIST_AICORE(NsaCompressAttentionInfer,
97+ OP_INPUT(query, key, value, attentionMaskOptional, blockTableOptional, actualSeqQLen,
98+ actualCmpSeqKvLen, actualSelKvLen, topKMaskOptional),
99+ OP_OUTPUT(output, topKOutput),
100+ OP_ATTR(numHeads, numKeyValueHeads, selectBlockSize, selectBlockCount,
101+ compressBlockSize, compressBlockStride, static_cast<float>(scaleValue),
102+ layoutOptional, pageBlockSize, sparseMode));
103+ return {output, topKOutput};
104+}
105+ 
106+} // namespace l0op
Aexperimental/attention/nsa_compress_attention_infer/op_host/op_api/nsa_compress_attention_infer.h+41-0
@@ -0,0 +1,41 @@
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 OP_API_INC_LEVEL0_OP_NSA_COMPRESS_ATTENTION_INFER_OP_H_
12+#define OP_API_INC_LEVEL0_OP_NSA_COMPRESS_ATTENTION_INFER_OP_H_
13+ 
14+#include "opdev/op_executor.h"
15+ 
16+namespace l0op {
17+ 
18+const std::array<const aclTensor *, 2>
19+NsaCompressAttentionInfer(const aclTensor *query,
20+ const aclTensor *key,
21+ const aclTensor *value,
22+ const aclTensor *attentionMaskOptional,
23+ const aclTensor *blockTableOptional,
24+ const aclIntArray *actualQSeqLenOptional,
25+ const aclIntArray *actualCmpKvSeqLenOptional,
26+ const aclIntArray *actualSelKvSeqLenOptional,
27+ const aclTensor *topKMaskOptional,
28+ int64_t numHeads,
29+ int64_t numKeyValueHeads,
30+ int64_t selectBlockSize,
31+ int64_t selectBlockCount,
32+ int64_t compressBlockSize,
33+ int64_t compressBlockStride,
34+ double scaleValue,
35+ const char *layoutOptional,
36+ int64_t pageBlockSize,
37+ int64_t sparseMode,
38+ aclOpExecutor *executor);
39+}
40+ 
41+#endif // OP_API_INC_LEVEL0_OP_NSA_NSA_COMPRESS_ATTENTION_INFER_OP_H_
Aexperimental/attention/nsa_compress_attention_infer/op_kernel/ncai_public_define.h+67-0
@@ -0,0 +1,67 @@
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+/*!
12+ * \file ncai_public_define.h
13+ * \brief
14+ */
15+ 
16+#ifndef NCAI_PUBLIC_DEFINE_H
17+#define NCAI_PUBLIC_DEFINE_H
18+ 
19+#include "kernel_operator.h"
20+#include "lib/matmul_intf.h"
21+#include "lib/matrix/matmul/tiling.h"
22+ 
23+using namespace AscendC;
24+ 
25+namespace NSA_COMPRESS_ATTENTION_INFER {
26+ 
27+constexpr uint32_t BASE_L1_BLOCK_ELEM_NUM_QKV = 128 * 192;
28+constexpr uint32_t BASE_L0C_BLOCK_ELEM_NUM = 128 * 128;
29+constexpr uint32_t BASE_L0AB_BLOCK_ELEM_NUM = 128 * 128;
30+constexpr uint32_t BASE_L1_BUF_ADDR_OFFSET = 128 * 192 * 2 * 2;
31+constexpr uint32_t DOUBLE_BUFFER = 2;
32+constexpr uint32_t BLOCK_SIZE = 16;
33+constexpr uint32_t SOFTMAX_BLOCK_SIZE = 32;
34+constexpr int32_t MM1_READY = 0;
35+constexpr int32_t SOFTMAX_READY = 1;
36+constexpr int32_t IMPORTANCE_SCORE_READY = 2;
37+constexpr int32_t MM2_READY = 3;
38+constexpr uint32_t IMPORTANCE_SCORE_PARAM_ROW = 128;
39+constexpr uint32_t IMPORTANCE_SCORE_EVENT_ID = 4;
40+constexpr uint32_t FP32_CUBE_BLOCK_SIZE = 8 * 16;
41+constexpr uint32_t CUBE_BLOCK_SIZE = 16 * 16;
42+constexpr uint32_t ALIGNED_8 = 8;
43+constexpr uint32_t L1_SIZE = 512 * 1024 / sizeof(float);
44+constexpr uint32_t L0A_PING_PONG_SIZE = 64 * 1024 / 2 / sizeof(float);
45+constexpr uint32_t L0C_PING_PONG_SIZE = 128 * 1024 / 2 / sizeof(float);
46+constexpr uint32_t BUFFER_NUM = 1;
47+constexpr uint64_t SYNC_MODE2 = 2;
48+constexpr uint32_t BASE_TOPK_ELEM_NUM_OFFSET = 8192;
49+constexpr uint32_t NUM64 = 64;
50+ 
51+enum class LAYOUT : uint8_t {
52+ TND = 0,
53+ BSND
54+};
55+ 
56+template <typename Q_T, typename KV_T, typename OUT_T, LAYOUT LAYOUT_T = LAYOUT::TND, bool IMP_SCORE_OPT_ = false,
57+ typename... Args>
58+struct NCAIType {
59+ using queryType = Q_T;
60+ using kvType = KV_T;
61+ using outputType = OUT_T;
62+ static constexpr LAYOUT layout = LAYOUT_T;
63+ static constexpr bool IMP_SCORE_OPT = IMP_SCORE_OPT_;
64+};
65+ 
66+} // namespace NSA_COMPRESS_ATTENTION_INFER
67+#endif // NCAI_PUBLIC_DEFINE_H
Aexperimental/attention/nsa_compress_attention_infer/op_kernel/nsa_compress_attention_infer.cpp+72-0
@@ -0,0 +1,72 @@
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+/*!
12+ * \file nsa_compress_attention_infer.cpp
13+ * \brief
14+ */
15+ 
16+#include "kernel_operator.h"
17+#include "nsa_compress_attention_infer.h"
18+ 
19+using namespace AscendC;
20+using namespace NSA_COMPRESS_ATTENTION_INFER;
21+ 
22+extern "C" __global__ __aicore__ void
23+nsa_compress_attention_infer(
24+ GM_ADDR query, GM_ADDR key, GM_ADDR value, GM_ADDR attentionMaskOptional, GM_ADDR blockTableOptional,
25+ GM_ADDR actualQSeqLenOptional, GM_ADDR actualCmpKvSeqLenOptional, GM_ADDR actualSelKvSeqLenOptional,
26+ GM_ADDR topkMaskOptional, GM_ADDR output, GM_ADDR topkIndicesOut, GM_ADDR workspace, GM_ADDR tiling)
27+{
28+ GM_ADDR user = GetUserWorkspace(workspace);
29+ if (user == nullptr) {
30+ return;
31+ }
32+ KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
33+ GET_TILING_DATA(tilingData, tiling);
34+#if (ORIG_DTYPE_QUERY == DT_FLOAT16)
35+ if (TILING_KEY_IS(0)) {
36+ NsaCompressAttentionInfer<NCAIType<half, half, half>> op;
37+ op.Run(query, key, value, blockTableOptional, actualQSeqLenOptional, actualCmpKvSeqLenOptional,
38+ actualSelKvSeqLenOptional, output, topkIndicesOut, user, &tilingData);
39+ } else if (TILING_KEY_IS(1)) {
40+ NsaCompressAttentionInfer<NCAIType<half, half, half, LAYOUT::BSND>> op;
41+ op.Run(query, key, value, blockTableOptional, actualQSeqLenOptional, actualCmpKvSeqLenOptional,
42+ actualSelKvSeqLenOptional, output, topkIndicesOut, user, &tilingData);
43+ } else if (TILING_KEY_IS(2)) {
44+ NsaCompressAttentionInfer<NCAIType<half, half, half, LAYOUT::TND, true>> op;
45+ op.Run(query, key, value, blockTableOptional, actualQSeqLenOptional, actualCmpKvSeqLenOptional,
46+ actualSelKvSeqLenOptional, output, topkIndicesOut, user, &tilingData);
47+ } else if (TILING_KEY_IS(3)) {
48+ NsaCompressAttentionInfer<NCAIType<half, half, half, LAYOUT::BSND, true>> op;
49+ op.Run(query, key, value, blockTableOptional, actualQSeqLenOptional, actualCmpKvSeqLenOptional,
50+ actualSelKvSeqLenOptional, output, topkIndicesOut, user, &tilingData);
51+ }
52+#endif
53+#if (ORIG_DTYPE_QUERY == DT_BF16)
54+ if (TILING_KEY_IS(0)) {
55+ NsaCompressAttentionInfer<NCAIType<bfloat16_t, bfloat16_t, bfloat16_t>> op;
56+ op.Run(query, key, value, blockTableOptional, actualQSeqLenOptional, actualCmpKvSeqLenOptional,
57+ actualSelKvSeqLenOptional, output, topkIndicesOut, user, &tilingData);
58+ } else if (TILING_KEY_IS(1)) {
59+ NsaCompressAttentionInfer<NCAIType<bfloat16_t, bfloat16_t, bfloat16_t, LAYOUT::BSND>> op;
60+ op.Run(query, key, value, blockTableOptional, actualQSeqLenOptional, actualCmpKvSeqLenOptional,
61+ actualSelKvSeqLenOptional, output, topkIndicesOut, user, &tilingData);
62+ } else if (TILING_KEY_IS(2)) {
63+ NsaCompressAttentionInfer<NCAIType<bfloat16_t, bfloat16_t, bfloat16_t, LAYOUT::TND, true>> op;
64+ op.Run(query, key, value, blockTableOptional, actualQSeqLenOptional, actualCmpKvSeqLenOptional,
65+ actualSelKvSeqLenOptional, output, topkIndicesOut, user, &tilingData);
66+ } else if (TILING_KEY_IS(3)) {
67+ NsaCompressAttentionInfer<NCAIType<bfloat16_t, bfloat16_t, bfloat16_t, LAYOUT::BSND, true>> op;
68+ op.Run(query, key, value, blockTableOptional, actualQSeqLenOptional, actualCmpKvSeqLenOptional,
69+ actualSelKvSeqLenOptional, output, topkIndicesOut, user, &tilingData);
70+ }
71+#endif
72+}
Aexperimental/attention/nsa_compress_attention_infer/tests/CMakeLists.txt+80-0
@@ -0,0 +1,80 @@
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+########################################################################################################################
12+# 调用编译方法, 生成对应编译目标
13+########################################################################################################################
14+ 
15+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wno-error=deprecated-declarations -Wno-deprecated-declarations")
16+ 
17+set(_ncai_TilingSourcesExt
18+ ${OPS_TRANSFORMER_DIR}/attention/nsa_compress_attention_infer/op_host/nsa_compress_attention_infer_tiling.cpp
19+)
20+ 
21+set(_ncai_KernelTilingDataDefH
22+ ${OPS_TRANSFORMER_DIR}/common/include/tiling_base/data_copy_transpose_tiling_def.h
23+ ${CMAKE_SOURCE_DIR}/attention/nsa_compress_attention_infer/op_host/nsa_compress_attention_infer_tiling.h
24+)
25+ 
26+set(_ncai_CompileDefintions
27+ ORIG_DTYPE_QUERY=DT_FLOAT16
28+)
29+ 
30+set(_ncai_KernelTilingExtInclude
31+ ${CMAKE_SOURCE_DIR}/attention/nsa_compress_attention_infer/op_host/
32+ ${CMAKE_SOURCE_DIR}/attention/nsa_compress_attention_infer/op_host/op_api/
33+)
34+ 
35+set(_ncai_UTestCommonPrivateIncludeExt
36+ ${OPBASE_INC_DIRS}
37+ ${OPS_TRANSFORMER_DIR}/common/include
38+)
39+ 
40+set(_ncai_TargetPrivateLinkLibrariesExt
41+ ${UTest_NamePrefix}_OpTiling
42+ error_manager
43+)
44+ 
45+set(_ncai_OpKernelPrivateIncludesExt
46+ ${OPS_TRANSFORMER_DIR}/common/include/kernel
47+)
48+ 
49+set(_NCAI_UTestSourcesExtForce)
50+ 
51+if (UT_TEST_ALL)
52+ file(GLOB_RECURSE _Src1 "${CMAKE_CURRENT_SOURCE_DIR}/utest/*.cc")
53+ file(GLOB_RECURSE _Src2 "${CMAKE_CURRENT_SOURCE_DIR}/utest/*.cpp")
54+elseif(OP_HOST_UT)
55+ file(GLOB_RECURSE _Src1 "${CMAKE_CURRENT_SOURCE_DIR}/utest/op_host/*tiling.cc")
56+ file(GLOB_RECURSE _Src2 "${CMAKE_CURRENT_SOURCE_DIR}/utest/op_host/*tiling.cpp")
57+elseif(OP_KERNEL_UT)
58+ file(GLOB_RECURSE _Src1 "${CMAKE_CURRENT_SOURCE_DIR}/utest/op_kernel/*kernel.cc")
59+ file(GLOB_RECURSE _Src2 "${CMAKE_CURRENT_SOURCE_DIR}/utest/op_kernel/*kernel.cpp")
60+endif()
61+ 
62+list(APPEND _NCAI_UTestSourcesExtForce ${_Src1} ${_Src2})
63+ 
64+aux_source_directory(${CMAKE_SOURCE_DIR}/attention/nsa_compress_attention_infer/op_kernel _ncai_KernelSourcesExt)
65+ 
66+OpsTest_Level2_AddOp(
67+ SUB_SYSTEM transformer
68+ BRIEF ncaInfer
69+ SNAKE nsa_compress_attention_infer
70+ KERNEL_TILING_DATA_DEF_H ${_ncai_KernelTilingDataDefH}
71+ KERNEL_PRIVATE_INCLUDES_EXT ${_ncai_OpKernelPrivateIncludesExt}
72+ KERNEL_SOURCES_EXT ${_ncai_KernelSourcesExt}
73+ TILING_SOURCES_EXT ${_ncai_TilingSourcesExt}
74+ TILING_PRIVATE_INCLUDES_EXT ${_ncai_KernelTilingExtInclude}
75+ KERNEL_PRIVATE_COMPILE_DEFINITIONS_EXT ${_ncai_CompileDefintions}
76+ UTEST_SOURCES_EXT ${_NCAI_UTestSourcesExtForce}
77+ UTEST_SOURCES_EXT_FORCE ON
78+ UTEST_COMMON_PRIVATE_LINK_LIBRARIES_EXT ${_ncai_TargetPrivateLinkLibrariesExt}
79+ UTEST_COMMON_PRIVATE_INCLUDES_EXT ${_ncai_UTestCommonPrivateIncludeExt}
80+)
Aexperimental/attention/nsa_compress_attention_infer/tests/comm/ncai_case.cpp+258-0
@@ -0,0 +1,258 @@
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+/*!
12+ * \file ncai_case.cpp
13+ * \brief NativeSelectedAttention / NativeSelectedAttentionGrad 测试用例.
14+ */
15+ 
16+#include "ncai_case.h"
17+#include <utility>
18+#include <tikicpulib.h>
19+#include <graph/utils/type_utils.h>
20+#include <register/op_impl_registry.h>
21+#include "tests/utils/log.h"
22+#include "tests/utils/platform.h"
23+#include "tiling_base/tiling_base.h"
24+ 
25+/**
26+ * 以下函数声明需要保持与 CMakeList.txt 中调用 OpsTest_Level2_AddOp 函数时 KERNEL_PRIVATE_COMPILE_DEFINITIONS_EXT
27+ * 参数所控制的 Kernel 入口一致.
28+ */
29+ 
30+#define NCAI_KERNEL_PARAM \
31+ (__gm__ uint8_t * query, __gm__ uint8_t * key, __gm__ uint8_t * value, __gm__ uint8_t * attention_mask, \
32+ __gm__ uint8_t * block_table, __gm__ uint8_t * actualQSeqLen, __gm__ uint8_t * actualCmpKvSeqLen, \
33+ __gm__ uint8_t * actualSelKvSeqLen, __gm__ uint8_t * topkMask, __gm__ uint8_t * output, \
34+ __gm__ uint8_t * topkIndicesOut, __gm__ uint8_t * workspace, __gm__ uint8_t * tiling)
35+ 
36+typedef void(*NcaiKernalFunc) NCAI_KERNEL_PARAM;
37+ 
38+extern "C" __global__ __aicore__ void nsa_compress_attention_infer NCAI_KERNEL_PARAM;
39+ 
40+using namespace ops::adv::tests::ncaInfer;
41+using TensorIntf = ops::adv::tests::utils::TensorIntf;
42+using Case = ops::adv::tests::utils::Case;
43+using Platform = ops::adv::tests::utils::Platform;
44+ 
45+ 
46+bool RunNcai(void *func, uint64_t tilingKey, int64_t blockDim, std::vector<TensorIntf *> &inputs,
47+ std::vector<TensorIntf *> &outputs, uint8_t *workspace, uint8_t *tilingData)
48+{
49+ constexpr uint32_t ATTEN_OUTPUT_INDEX = 0;
50+ constexpr uint32_t SELECT_OUTPUT_INDEX = 1;
51+ constexpr uint32_t QUERY_INPUT_INDEX = 0;
52+ constexpr uint32_t KEY_INPUT_INDEX = 1;
53+ constexpr uint32_t VALUE_INPUT_INDEX = 2;
54+ constexpr uint32_t ATTEN_MASK_INPUT_INDEX = 3;
55+ constexpr uint32_t BLOCK_TABLE_INPUT_INDEX = 4;
56+ constexpr uint32_t ACT_Q_SEQ_LEN_INPUT_INDEX = 5;
57+ constexpr uint32_t ACT_CMP_KV_SEQ_LEN_INPUT_INDEX = 6;
58+ constexpr uint32_t ACT_SEL_KV_SEQ_LEN_INPUT_INDEX = 7;
59+ constexpr uint32_t TOPK_MASK_INPUT_INDEX = 8;
60+ 
61+ // Kernel 运行
62+ auto kernelFunc = (NcaiKernalFunc)func;
63+ ICPU_SET_TILING_KEY(tilingKey);
64+ ICPU_RUN_KF(kernelFunc, blockDim,
65+ inputs[QUERY_INPUT_INDEX]->GetDevData(), // query
66+ inputs[KEY_INPUT_INDEX]->GetDevData(), // key
67+ inputs[VALUE_INPUT_INDEX]->GetDevData(), // value
68+ inputs[ATTEN_MASK_INPUT_INDEX]->GetDevData(), // attentionMaskOptional
69+ inputs[BLOCK_TABLE_INPUT_INDEX]->GetDevData(), // blockTableOptional
70+ inputs[ACT_Q_SEQ_LEN_INPUT_INDEX]->GetDevData(), // actualQSeqLenOptional
71+ inputs[ACT_CMP_KV_SEQ_LEN_INPUT_INDEX]->GetDevData(), // actualCmpKvSeqLenOptional
72+ inputs[ACT_SEL_KV_SEQ_LEN_INPUT_INDEX]->GetDevData(), // actualSelKvSeqLenOptional
73+ inputs[TOPK_MASK_INPUT_INDEX]->GetDevData(), // topkMaskOptional
74+ outputs[ATTEN_OUTPUT_INDEX]->GetDevData(), // output
75+ outputs[SELECT_OUTPUT_INDEX]->GetDevData(), // topkIndicesOut
76+ workspace, tilingData);
77+ return true;
78+}
79+ 
80+extern "C" ge::graphStatus TilingNsaCompressAttentionInferStub(gert::TilingContext *context)
81+{
82+ constexpr uint32_t ACT_Q_SEQ_LEN_INPUT_INDEX = 5;
83+ constexpr uint32_t ACT_CMP_KV_SEQ_LEN_INPUT_INDEX = 6;
84+ constexpr uint32_t ACT_SEL_KV_SEQ_LEN_INPUT_INDEX = 7;
85+ auto *ncaInferCase = static_cast<NcaInferCase *>(Case::GetCurrentCase());
86+ if (ncaInferCase != nullptr) {
87+ NcaInferCase::DoTilingParam p;
88+ p.ctx = context;
89+ p.ret = ge::GRAPH_SUCCESS;
90+ p.actQSeqLenTensor = const_cast<gert::Tensor *>(context->GetOptionalInputTensor(ACT_Q_SEQ_LEN_INPUT_INDEX));
91+ p.actSeqCmpKVLenTensor = const_cast<gert::Tensor *>(context->GetOptionalInputTensor(ACT_CMP_KV_SEQ_LEN_INPUT_INDEX));
92+ p.actSeqSelKVLenTensor = const_cast<gert::Tensor *>(context->GetOptionalInputTensor(ACT_SEL_KV_SEQ_LEN_INPUT_INDEX));
93+ if (!ncaInferCase->DoOpTiling(p)) {
94+ return p.ret;
95+ }
96+ return ncaInferCase->NcaiTilingFunc(context);
97+ }
98+ return ge::GRAPH_FAILED;
99+}
100+ 
101+bool NcaInferCase::InitParam()
102+{
103+ std::vector<int64_t> qShape;
104+ std::vector<int64_t> outShape;
105+ std::vector<int64_t> topkOutShape;
106+ int64_t headSizeVO = 128;
107+ int64_t maxBlockPerQuery = 32;
108+ int64_t defaultBlockNum = 1000;
109+ mParam.actualQSeqLenOptionalLenList = std::vector(mParam.B, static_cast<int64_t>(mParam.S1));
110+ mParam.actualCmpKvSeqLenOptionalLenList = std::vector(mParam.B, static_cast<int64_t>(mParam.S2));
111+ mParam.actualSelKvSeqLenOptionalLenList = std::vector(mParam.B, static_cast<int64_t>((mParam.S2 - 1) * mParam.compressStride + mParam.compressBlockSize));
112+ 
113+ int64_t qSeqLenSum = 0;
114+ int64_t kvSeqLenSum = 0;
115+ for (long &it : mParam.actualQSeqLenOptionalLenList) {
116+ qSeqLenSum += it;
117+ }
118+ for (long &it : mParam.actualCmpKvSeqLenOptionalLenList) {
119+ kvSeqLenSum += it;
120+ }
121+ int64_t blockNum = (kvSeqLenSum + mParam.pageBlockSize - 1) / mParam.pageBlockSize;
122+ if (mParam.S2 == -1) {
123+ blockNum = defaultBlockNum;
124+ maxBlockPerQuery = (mParam.S2 + mParam.pageBlockSize - 1) / mParam.pageBlockSize;
125+ }
126+ 
127+ if (mParam.layoutOptional == "TND") {
128+ qShape = {qSeqLenSum, mParam.N1, mParam.D};
129+ outShape = {mParam.B, mParam.N1, headSizeVO};
130+ topkOutShape = {mParam.B, mParam.N2, mParam.selectBlockCount};
131+ } else {
132+ qShape = {mParam.B, mParam.S1, mParam.N1, mParam.D};
133+ outShape = {mParam.B, mParam.S1, mParam.N1, headSizeVO};
134+ topkOutShape = {mParam.B, mParam.S1, mParam.N2, mParam.selectBlockCount};
135+ }
136+ query = Tensor("query", qShape, mParam.layoutOptional.c_str(), ge::DataType::DT_FLOAT16, ge::FORMAT_ND);
137+ key = Tensor("key", {blockNum, mParam.pageBlockSize, mParam.N2 * mParam.D}, mParam.layoutOptional.c_str(), ge::DataType::DT_FLOAT16, ge::FORMAT_ND);
138+ value = Tensor("value", {blockNum, mParam.pageBlockSize, mParam.N2 * headSizeVO}, mParam.layoutOptional.c_str(), ge::DataType::DT_FLOAT16, ge::FORMAT_ND);
139+ attentionMaskOptional = Tensor("attentionMaskOptional", {128, 128}, mParam.layoutOptional.c_str(), ge::DataType::DT_BOOL, ge::FORMAT_ND);
140+ blockTableOptional = Tensor("blockTableOptional", {mParam.B, maxBlockPerQuery}, mParam.layoutOptional.c_str(), ge::DataType::DT_INT32, ge::FORMAT_ND);
141+
142+ actualQSeqLenOptional = Tensor("actualQSeqLenOptional", {mParam.B}, mParam.layoutOptional.c_str(), ge::DataType::DT_INT64, ge::FORMAT_ND);
143+ actualCmpKvSeqLenOptional = Tensor("actualCmpKvSeqLenOptional", {mParam.B}, mParam.layoutOptional.c_str(), ge::DataType::DT_INT64, ge::FORMAT_ND);
144+ actualSelKvSeqLenOptional = Tensor("actualSelKvSeqLenOptional", {mParam.B}, mParam.layoutOptional.c_str(), ge::DataType::DT_INT64, ge::FORMAT_ND);
145+
146+ topkMaskOptional = Tensor("topkMaskOptional", {20, 4096}, mParam.layoutOptional.c_str(), ge::DataType::DT_BOOL, ge::FORMAT_ND);
147+ 
148+ output = Tensor("output", outShape, mParam.layoutOptional.c_str(), ge::DataType::DT_FLOAT16, ge::FORMAT_ND);
149+ topkIndicesOut = Tensor("topkIndicesOut", topkOutShape, mParam.layoutOptional.c_str(), ge::DataType::DT_INT32, ge::FORMAT_ND);
150+
151+ if (!InitTensor(actualQSeqLenOptional, mParam.actualQSeqLenOptionalLenList) ||
152+ !InitTensor(actualCmpKvSeqLenOptional, mParam.actualCmpKvSeqLenOptionalLenList) ||
153+ !InitTensor(actualSelKvSeqLenOptional, mParam.actualSelKvSeqLenOptionalLenList)) {
154+ return false;
155+ }
156+ return true;
157+}
158+ 
159+bool NcaInferCase::InitOpInfo()
160+{
161+ constexpr uint32_t TILING_DATA_MAX_SIZE = 2280;
162+ auto *ncaiKernalFunc = (void *)nsa_compress_attention_infer;
163+ 
164+ bool rst = mCtx.SetOpName("NsaCompressAttentionInfer");
165+ rst = rst && mCtx.SetDeterministic(false);
166+ rst = rst && mCtx.SetInputs({&query, &key, &value, &attentionMaskOptional, &blockTableOptional, &actualQSeqLenOptional, &actualCmpKvSeqLenOptional, &actualSelKvSeqLenOptional, &topkMaskOptional});
167+ rst = rst && mCtx.SetOutputs({&output, &topkIndicesOut});
168+ rst = rst && mCtx.SetAttrs({{"numHeads", mParam.numHeads},
169+ {"numKeyValueHeads", mParam.numKeyValueHeads},
170+ {"selectBlockSize", mParam.selectBlockSize},
171+ {"selectBlockCount", mParam.selectBlockCount},
172+ {"compressBlockSize", mParam.compressBlockSize},
173+ {"compressStride", mParam.compressStride},
174+ {"scaleValue", mParam.scaleValue},
175+ {"layoutOptional", mParam.layoutOptional},
176+ {"pageBlockSize", mParam.pageBlockSize},
177+ {"sparseMode", mParam.sparseMode}});
178+ rst = rst && mCtx.SetKernelRunCbf(RunNcai);
179+ rst = rst && mCtx.SetTilingDataMaxSize(TILING_DATA_MAX_SIZE); // max tilingDataLen
180+ rst = rst && mCtx.SetKernelMainFunc(ncaiKernalFunc);
181+ rst = rst && mOpInfo.SetContext(&mCtx);
182+ 
183+ auto *platform = Platform::GetGlobalPlatform();
184+ if (platform == nullptr) {
185+ LOG_ERR("Global Platform is null");
186+ return false;
187+ }
188+ 
189+ NcaiTilingFunc = (gert::OpImplRegisterV2::TilingKernelFunc)platform->LoadOpTilingSoSym("TilingNsaCompressAttentionInfer");
190+ if (NcaiTilingFunc == nullptr) {
191+ LOG_ERR("Can't get origin tiling func, ncai(%p)", NcaiTilingFunc);
192+ return false;
193+ }
194+ IMPL_OP(NsaCompressAttentionInfer).Tiling(TilingNsaCompressAttentionInferStub);
195+ 
196+ return rst;
197+}
198+ 
199+bool NcaInferCase::InitCurrentCasePtr()
200+{
201+ Case::mCurrentCasePtr = this;
202+ return true;
203+}
204+ 
205+bool NcaInferCase::Run()
206+{
207+ if (!mEnable) {
208+ return true;
209+ }
210+ if (!mOpInfo.ProcessTiling(mName)) {
211+ return false;
212+ }
213+ if (!mOpInfo.ProcessKernel(mName)) {
214+ return false;
215+ }
216+ return true;
217+}
218+ 
219+NcaInferCase::NcaInferCase(const char *name, bool enable, const char *dbgInfo, OpInfo incre, ncaInferParam param)
220+ : Case(name, enable, dbgInfo), mOpInfo(std::move(incre)), mParam(std::move(param))
221+{
222+ this->mOpInfo.mName = "NsaCompressAttentionInfer";
223+}
224+ 
225+NcaInferCase::NcaInferCase()
226+{
227+}
228+ 
229+NcaInferCase::ncaInferParam::ncaInferParam()
230+{
231+}
232+ 
233+NcaInferCase::ncaInferParam::ncaInferParam(int64_t pB, int64_t pS1, int64_t pS2, int64_t pN1, int64_t pG, int64_t pD, int64_t pN2,
234+ int64_t pnumHeads, int64_t pnumKeyValueHeads, int64_t pselectBlockSize,
235+ int64_t pselectBlockCount,int64_t pcompressBlockSize,int64_t pcompressStride, float pscaleValue,
236+ std::string playoutOptional, int64_t ppageBlockSize, int64_t pSparseMode)
237+ : B(pB), S1(pS1), S2(pS2), N1(pN1), G(pG), D(pD), N2(pN2), numHeads(pnumHeads),
238+ numKeyValueHeads(pnumKeyValueHeads), selectBlockSize(pselectBlockSize), selectBlockCount(pselectBlockCount),
239+ compressBlockSize(pcompressBlockSize), compressStride(pcompressStride), scaleValue(pscaleValue),
240+ layoutOptional(playoutOptional), pageBlockSize(ppageBlockSize), sparseMode(pSparseMode)
241+{
242+}
243+ 
244+bool NcaInferCase::DoOpTiling(DoTilingParam& tilingParam) {
245+ if (tilingParam.ctx == nullptr) {
246+ return false;
247+ }
248+ if (tilingParam.actQSeqLenTensor != nullptr && mParam.actualQSeqLenOptionalLenList.size() != 0) {
249+ tilingParam.actQSeqLenTensor->SetData(gert::TensorData{mParam.actualQSeqLenOptionalLenList.data()});
250+ }
251+ if (tilingParam.actSeqCmpKVLenTensor != nullptr && mParam.actualCmpKvSeqLenOptionalLenList.size() != 0) {
252+ tilingParam.actSeqCmpKVLenTensor->SetData(gert::TensorData{mParam.actualCmpKvSeqLenOptionalLenList.data()});
253+ }
254+ if (tilingParam.actSeqSelKVLenTensor != nullptr && mParam.actualSelKvSeqLenOptionalLenList.size() != 0) {
255+ tilingParam.actSeqSelKVLenTensor->SetData(gert::TensorData{mParam.actualSelKvSeqLenOptionalLenList.data()});
256+ }
257+ return true;
258+}
Aexperimental/attention/nsa_compress_attention_infer/tests/comm/ncai_case.h+108-0
@@ -0,0 +1,108 @@
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+/*!
12+ * \file ncai_case.h
13+ * \brief NsaCompressAttentionInfer 测试用例.
14+ */
15+ 
16+#pragma once
17+#include <vector>
18+#include <cstdint>
19+#include "graph/types.h"
20+#include "tests/utils/case.h"
21+#include "tests/utils/op_info.h"
22+#include "tests/utils/context.h"
23+#include "tests/utils/tensor.h"
24+#include "tests/utils/tensor_list.h"
25+#include <exe_graph/runtime/tiling_context.h>
26+#include <register/op_impl_registry.h>
27+ 
28+namespace ops::adv::tests::ncaInfer {
29+ 
30+/**
31+ * 算子 NsaCompressAttentionInfer 参数
32+ */
33+ 
34+class NcaInferCase : public ops::adv::tests::utils::Case {
35+using OpInfo = ops::adv::tests::utils::OpInfo;
36+using Context = ops::adv::tests::utils::Context;
37+using Tensor = ops::adv::tests::utils::Tensor;
38+using TensorList = ops::adv::tests::utils::TensorList;
39+ 
40+public:
41+ class ncaInferParam {
42+ public:
43+ /* 设置参数 */
44+ int64_t B = 0;
45+ int64_t S1 = 0;
46+ int64_t S2 = 0;
47+ int64_t N1 = 0;
48+ int64_t G = 0;
49+ int64_t D = 0;
50+ int64_t N2 = 0;
51+
52+ // attr
53+ int64_t numHeads = 1;
54+ int64_t numKeyValueHeads = 1;
55+ int64_t selectBlockSize = 1;
56+ int64_t selectBlockCount = 1;
57+ int64_t compressBlockSize = 1;
58+ int64_t compressStride = 1;
59+ float scaleValue = 1.0;
60+ std::string layoutOptional = "TND";
61+ int64_t pageBlockSize = 1;
62+ int64_t sparseMode = 0;
63+ 
64+ std::vector<int64_t> actualQSeqLenOptionalLenList = {};
65+ std::vector<int64_t> actualCmpKvSeqLenOptionalLenList = {};
66+ std::vector<int64_t> actualSelKvSeqLenOptionalLenList = {};
67+ ncaInferParam();
68+ ncaInferParam(int64_t pB, int64_t pS1, int64_t pS2, int64_t pN1, int64_t pG, int64_t pD, int64_t pN2,
69+ int64_t pnumHeads, int64_t pnumKeyValueHeads, int64_t pselectBlockSize,
70+ int64_t pselectBlockCount,int64_t pcompressBlockSize,int64_t pcompressStride, float pscaleValue,
71+ std::string pInputLayout, int64_t ppageBlockSize, int64_t pSparseMode);
72+ };
73+ class DoTilingParam {
74+ public:
75+ gert::TilingContext *ctx = nullptr;
76+ ge::graphStatus ret = ge::GRAPH_SUCCESS;
77+ gert::Tensor *actQSeqLenTensor = nullptr;
78+ gert::Tensor *actSeqCmpKVLenTensor = nullptr;
79+ gert::Tensor *actSeqSelKVLenTensor = nullptr;
80+ };
81+ 
82+ Tensor query, key, value, attentionMaskOptional, blockTableOptional, actualQSeqLenOptional, actualCmpKvSeqLenOptional, actualSelKvSeqLenOptional, topkMaskOptional, output, topkIndicesOut;
83+ OpInfo mOpInfo;
84+ Context mCtx;
85+ ncaInferParam mParam;
86+ gert::OpImplRegisterV2::TilingKernelFunc NcaiTilingFunc = nullptr;
87+ NcaInferCase();
88+ NcaInferCase(const char *name, bool enable, const char *dbgInfo, OpInfo incre, ncaInferParam param);
89+ bool Run() override;
90+ bool InitParam() override;
91+ bool InitOpInfo() override;
92+ bool InitCurrentCasePtr() override;
93+ bool DoOpTiling(DoTilingParam& tilingParam);
94+ template <class T> static bool InitTensor(Tensor &tensor, std::vector<T> &hostData)
95+ {
96+ if (hostData.empty()) {
97+ return true;
98+ }
99+ int64_t expMinSize = hostData.size() * sizeof(T);
100+ if (tensor.AllocDevData(0, expMinSize) == nullptr) {
101+ printf("Tensor(%s, %ld) AllocDevData Failed.", tensor.Name().c_str(), expMinSize);
102+ return false;
103+ }
104+ return tensor.CopyHostToDevData(hostData);
105+ }
106+};
107+ 
108+}
Aexperimental/attention/nsa_compress_attention_infer/tests/comm/ts_ncai.h+22-0
@@ -0,0 +1,22 @@
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+/*!
12+ * \file ts_nsag_param.h
13+ * \brief nsaGrad UTest 相关基类定义.
14+ */
15+ 
16+ 
17+#include "tests/utest/ts.h"
18+#include "ncai_case.h"
19+ 
20+using ncaInferCase = ops::adv::tests::ncaInfer::NcaInferCase;
21+ 
22+class Ts_ncaInfer : public Ts<ncaInferCase> {};
Aexperimental/attention/nsa_compress_attention_infer/tests/utest/op_host/test_nsa_compress_attention_infer_tiling.cpp+192-0
@@ -0,0 +1,192 @@
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+/*!
12+ * \file ts_nsag_param_tc.cpp
13+ * \brief nsa compress attention infer 用例.
14+ */
15+ 
16+#include "ts_ncai.h"
17+ 
18+using Tensor = ops::adv::tests::utils::Tensor;
19+ 
20+TEST_F(Ts_ncaInfer, case_nsa_compress_attention_infer_0_tiling)
21+{
22+ constexpr uint32_t CASE1_B = 20;
23+ constexpr uint32_t CASE1_S2 = 4096;
24+ constexpr uint32_t CASE1_NUMHEADS = 64;
25+ constexpr uint32_t CASE1_NUMKVHEADS = 4;
26+ constexpr uint32_t CASE1_SELECTBLOCKSIZE = 64;
27+ constexpr uint32_t CASE1_SELECTBLOCKCOUNT = 16;
28+ constexpr uint32_t CASE1_COMPRESSBLOCKSIZE = 32;
29+ constexpr uint32_t CASE1_COMPRESSSTRIDE = 16;
30+ constexpr uint32_t CASE1_PAGEBLOCKSIZE = 128;
31+ constexpr uint32_t CASE1_MTILINGBLOCKDIM = 24;
32+ constexpr uint32_t CASE1_HEADSIZE_QK = 192;
33+ 
34+ ncaInferCase cs;
35+ /* 设置参数 */
36+ cs.mParam.B = CASE1_B;
37+ cs.mParam.S1 = 1;
38+ cs.mParam.S2 = CASE1_S2;
39+ cs.mParam.N1 = CASE1_NUMHEADS;
40+ cs.mParam.G = CASE1_NUMHEADS / CASE1_NUMKVHEADS;
41+ cs.mParam.D = CASE1_HEADSIZE_QK;
42+ cs.mParam.N2 = CASE1_NUMKVHEADS;
43+
44+ // attr
45+ cs.mParam.numHeads = CASE1_NUMHEADS;
46+ cs.mParam.numKeyValueHeads = CASE1_NUMKVHEADS;
47+ cs.mParam.selectBlockSize = CASE1_SELECTBLOCKSIZE;
48+ cs.mParam.selectBlockCount = CASE1_SELECTBLOCKCOUNT;
49+ cs.mParam.compressBlockSize = CASE1_COMPRESSBLOCKSIZE;
50+ cs.mParam.compressStride = CASE1_COMPRESSSTRIDE;
51+ cs.mParam.scaleValue = 1.0;
52+ cs.mParam.layoutOptional = "TND";
53+ cs.mParam.pageBlockSize = CASE1_PAGEBLOCKSIZE;
54+ cs.mParam.sparseMode = 0;
55+ 
56+ cs.mOpInfo.mExp.mTilingBlockDim = CASE1_MTILINGBLOCKDIM; // expected block dim
57+ cs.mOpInfo.mCtr.mRunTiling = true;
58+ cs.mOpInfo.mCtr.mRunKernel = false;
59+ ASSERT_TRUE(cs.Init());
60+ ASSERT_TRUE(cs.Run());
61+}
62+ 
63+TEST_F(Ts_ncaInfer, case_nsa_compress_attention_infer_1_tiling)
64+{
65+ constexpr uint32_t CASE2_B = 20;
66+ constexpr uint32_t CASE2_S2 = 2048;
67+ constexpr uint32_t CASE2_NUMHEADS = 64;
68+ constexpr uint32_t CASE2_NUMKVHEADS = 4;
69+ constexpr uint32_t CASE2_SELECTBLOCKSIZE = 64;
70+ constexpr uint32_t CASE2_SELECTBLOCKCOUNT = 16;
71+ constexpr uint32_t CASE2_COMPRESSBLOCKSIZE = 32;
72+ constexpr uint32_t CASE2_COMPRESSSTRIDE = 16;
73+ constexpr uint32_t CASE2_PAGEBLOCKSIZE = 128;
74+ constexpr uint32_t CASE2_MTILINGBLOCKDIM = 24;
75+ constexpr uint32_t CASE2_HEADSIZE_QK = 192;
76+ 
77+ ncaInferCase cs;
78+ /* 设置参数 */
79+ cs.mParam.B = CASE2_B;
80+ cs.mParam.S1 = 1;
81+ cs.mParam.S2 = CASE2_S2;
82+ cs.mParam.N1 = CASE2_NUMHEADS;
83+ cs.mParam.G = CASE2_NUMHEADS / CASE2_NUMKVHEADS;
84+ cs.mParam.D = CASE2_HEADSIZE_QK;
85+ cs.mParam.N2 = CASE2_NUMKVHEADS;
86+
87+ // attr
88+ cs.mParam.numHeads = CASE2_NUMHEADS;
89+ cs.mParam.numKeyValueHeads = CASE2_NUMKVHEADS;
90+ cs.mParam.selectBlockSize = CASE2_SELECTBLOCKSIZE;
91+ cs.mParam.selectBlockCount = CASE2_SELECTBLOCKCOUNT;
92+ cs.mParam.compressBlockSize = CASE2_COMPRESSBLOCKSIZE;
93+ cs.mParam.compressStride = CASE2_COMPRESSSTRIDE;
94+ cs.mParam.scaleValue = 1.0;
95+ cs.mParam.layoutOptional = "BSND";
96+ cs.mParam.pageBlockSize = CASE2_PAGEBLOCKSIZE;
97+ cs.mParam.sparseMode = 0;
98+ 
99+ cs.mOpInfo.mExp.mTilingBlockDim = CASE2_MTILINGBLOCKDIM;
100+ cs.mOpInfo.mCtr.mRunTiling = true;
101+ cs.mOpInfo.mCtr.mRunKernel = false;
102+ ASSERT_TRUE(cs.Init());
103+ ASSERT_TRUE(cs.Run());
104+}
105+ 
106+// getmaxworkspace test
107+TEST_F(Ts_ncaInfer, case_nsa_compress_attention_infer_aclgraph_tiling)
108+{
109+ constexpr uint32_t CASE3_B = 20;
110+ constexpr int32_t CASE3_S2 = -1;
111+ constexpr uint32_t CASE3_NUMHEADS = 64;
112+ constexpr uint32_t CASE3_NUMKVHEADS = 4;
113+ constexpr uint32_t CASE3_SELECTBLOCKSIZE = 64;
114+ constexpr uint32_t CASE3_SELECTBLOCKCOUNT = 16;
115+ constexpr uint32_t CASE3_COMPRESSBLOCKSIZE = 32;
116+ constexpr uint32_t CASE3_COMPRESSSTRIDE = 16;
117+ constexpr uint32_t CASE3_PAGEBLOCKSIZE = 128;
118+ constexpr uint32_t CASE3_MTILINGBLOCKDIM = 24;
119+ constexpr uint32_t CASE3_HEADSIZE_QK = 192;
120+ 
121+ ncaInferCase cs;
122+ /* 设置参数 */
123+ cs.mParam.B = CASE3_B;
124+ cs.mParam.S1 = 1;
125+ cs.mParam.S2 = CASE3_S2;
126+ cs.mParam.N1 = CASE3_NUMHEADS;
127+ cs.mParam.G = CASE3_NUMHEADS / CASE3_NUMKVHEADS;
128+ cs.mParam.D = CASE3_HEADSIZE_QK;
129+ cs.mParam.N2 = CASE3_NUMKVHEADS;
130+
131+ // attr
132+ cs.mParam.numHeads = CASE3_NUMHEADS;
133+ cs.mParam.numKeyValueHeads = CASE3_NUMKVHEADS;
134+ cs.mParam.selectBlockSize = CASE3_SELECTBLOCKSIZE;
135+ cs.mParam.selectBlockCount = CASE3_SELECTBLOCKCOUNT;
136+ cs.mParam.compressBlockSize = CASE3_COMPRESSBLOCKSIZE;
137+ cs.mParam.compressStride = CASE3_COMPRESSSTRIDE;
138+ cs.mParam.scaleValue = 1.0;
139+ cs.mParam.layoutOptional = "TND";
140+ cs.mParam.pageBlockSize = CASE3_PAGEBLOCKSIZE;
141+ cs.mParam.sparseMode = 0;
142+ 
143+ cs.mOpInfo.mExp.mTilingBlockDim = CASE3_MTILINGBLOCKDIM;
144+ cs.mOpInfo.mCtr.mRunTiling = true;
145+ cs.mOpInfo.mCtr.mRunKernel = false;
146+ ASSERT_TRUE(cs.Init());
147+ ASSERT_TRUE(cs.Run());
148+}
149+ 
150+TEST_F(Ts_ncaInfer, case_nsa_compress_attention_infer_mtp_tiling)
151+{
152+ constexpr uint32_t CASE4_B = 20;
153+ constexpr uint32_t CASE4_S1 = 2;
154+ constexpr uint32_t CASE4_S2 = 4096;
155+ constexpr uint32_t CASE4_NUMHEADS = 64;
156+ constexpr uint32_t CASE4_NUMKVHEADS = 4;
157+ constexpr uint32_t CASE4_SELECTBLOCKSIZE = 64;
158+ constexpr uint32_t CASE4_SELECTBLOCKCOUNT = 16;
159+ constexpr uint32_t CASE4_COMPRESSBLOCKSIZE = 32;
160+ constexpr uint32_t CASE4_COMPRESSSTRIDE = 16;
161+ constexpr uint32_t CASE4_PAGEBLOCKSIZE = 128;
162+ constexpr uint32_t CASE4_MTILINGBLOCKDIM = 24;
163+ constexpr uint32_t CASE4_HEADSIZE_QK = 192;
164+ 
165+ ncaInferCase cs;
166+ /* 设置参数 */
167+ cs.mParam.B = CASE4_B;
168+ cs.mParam.S1 = CASE4_S1;
169+ cs.mParam.S2 = CASE4_S2;
170+ cs.mParam.N1 = CASE4_NUMHEADS;
171+ cs.mParam.G = CASE4_NUMHEADS / CASE4_NUMKVHEADS;
172+ cs.mParam.D = CASE4_HEADSIZE_QK;
173+ cs.mParam.N2 = CASE4_NUMKVHEADS;
174+
175+ // attr
176+ cs.mParam.numHeads = CASE4_NUMHEADS;
177+ cs.mParam.numKeyValueHeads = CASE4_NUMKVHEADS;
178+ cs.mParam.selectBlockSize = CASE4_SELECTBLOCKSIZE;
179+ cs.mParam.selectBlockCount = CASE4_SELECTBLOCKCOUNT;
180+ cs.mParam.compressBlockSize = CASE4_COMPRESSBLOCKSIZE;
181+ cs.mParam.compressStride = CASE4_COMPRESSSTRIDE;
182+ cs.mParam.scaleValue = 1.0;
183+ cs.mParam.layoutOptional = "TND";
184+ cs.mParam.pageBlockSize = CASE4_PAGEBLOCKSIZE;
185+ cs.mParam.sparseMode = 0;
186+ 
187+ cs.mOpInfo.mExp.mTilingBlockDim = CASE4_MTILINGBLOCKDIM;
188+ cs.mOpInfo.mCtr.mRunTiling = true;
189+ cs.mOpInfo.mCtr.mRunKernel = false;
190+ ASSERT_TRUE(cs.Init());
191+ ASSERT_TRUE(cs.Run());
192+}
Aexperimental/attention/nsa_compress_attention_infer/tests/utest/op_kernel/test_nsa_compress_attention_infer_kernel.cpp+192-0
@@ -0,0 +1,192 @@
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+/*!
12+ * \file ts_nsag_param_tc.cpp
13+ * \brief nsa compress attention infer 用例.
14+ */
15+ 
16+#include "ts_ncai.h"
17+ 
18+using Tensor = ops::adv::tests::utils::Tensor;
19+ 
20+TEST_F(Ts_ncaInfer, case_nsa_compress_attention_infer_0_kernel)
21+{
22+ constexpr uint32_t CASE1_B = 20;
23+ constexpr uint32_t CASE1_S2 = 4096;
24+ constexpr uint32_t CASE1_NUMHEADS = 64;
25+ constexpr uint32_t CASE1_NUMKVHEADS = 4;
26+ constexpr uint32_t CASE1_SELECTBLOCKSIZE = 64;
27+ constexpr uint32_t CASE1_SELECTBLOCKCOUNT = 16;
28+ constexpr uint32_t CASE1_COMPRESSBLOCKSIZE = 32;
29+ constexpr uint32_t CASE1_COMPRESSSTRIDE = 16;
30+ constexpr uint32_t CASE1_PAGEBLOCKSIZE = 128;
31+ constexpr uint32_t CASE1_MTILINGBLOCKDIM = 24;
32+ constexpr uint32_t CASE1_HEADSIZE_QK = 192;
33+ 
34+ ncaInferCase cs;
35+ /* 设置参数 */
36+ cs.mParam.B = CASE1_B;
37+ cs.mParam.S1 = 1;
38+ cs.mParam.S2 = CASE1_S2;
39+ cs.mParam.N1 = CASE1_NUMHEADS;
40+ cs.mParam.G = CASE1_NUMHEADS / CASE1_NUMKVHEADS;
41+ cs.mParam.D = CASE1_HEADSIZE_QK;
42+ cs.mParam.N2 = CASE1_NUMKVHEADS;
43+
44+ // attr
45+ cs.mParam.numHeads = CASE1_NUMHEADS;
46+ cs.mParam.numKeyValueHeads = CASE1_NUMKVHEADS;
47+ cs.mParam.selectBlockSize = CASE1_SELECTBLOCKSIZE;
48+ cs.mParam.selectBlockCount = CASE1_SELECTBLOCKCOUNT;
49+ cs.mParam.compressBlockSize = CASE1_COMPRESSBLOCKSIZE;
50+ cs.mParam.compressStride = CASE1_COMPRESSSTRIDE;
51+ cs.mParam.scaleValue = 1.0;
52+ cs.mParam.layoutOptional = "TND";
53+ cs.mParam.pageBlockSize = CASE1_PAGEBLOCKSIZE;
54+ cs.mParam.sparseMode = 0;
55+ 
56+ cs.mOpInfo.mExp.mTilingBlockDim = CASE1_MTILINGBLOCKDIM; // expected block dim
57+ cs.mOpInfo.mCtr.mRunTiling = true;
58+ cs.mOpInfo.mCtr.mRunKernel = true;
59+ ASSERT_TRUE(cs.Init());
60+ ASSERT_TRUE(cs.Run());
61+}
62+ 
63+TEST_F(Ts_ncaInfer, case_nsa_compress_attention_infer_1_kernel)
64+{
65+ constexpr uint32_t CASE2_B = 20;
66+ constexpr uint32_t CASE2_S2 = 2048;
67+ constexpr uint32_t CASE2_NUMHEADS = 64;
68+ constexpr uint32_t CASE2_NUMKVHEADS = 4;
69+ constexpr uint32_t CASE2_SELECTBLOCKSIZE = 64;
70+ constexpr uint32_t CASE2_SELECTBLOCKCOUNT = 16;
71+ constexpr uint32_t CASE2_COMPRESSBLOCKSIZE = 32;
72+ constexpr uint32_t CASE2_COMPRESSSTRIDE = 16;
73+ constexpr uint32_t CASE2_PAGEBLOCKSIZE = 128;
74+ constexpr uint32_t CASE2_MTILINGBLOCKDIM = 24;
75+ constexpr uint32_t CASE2_HEADSIZE_QK = 192;
76+ 
77+ ncaInferCase cs;
78+ /* 设置参数 */
79+ cs.mParam.B = CASE2_B;
80+ cs.mParam.S1 = 1;
81+ cs.mParam.S2 = CASE2_S2;
82+ cs.mParam.N1 = CASE2_NUMHEADS;
83+ cs.mParam.G = CASE2_NUMHEADS / CASE2_NUMKVHEADS;
84+ cs.mParam.D = CASE2_HEADSIZE_QK;
85+ cs.mParam.N2 = CASE2_NUMKVHEADS;
86+
87+ // attr
88+ cs.mParam.numHeads = CASE2_NUMHEADS;
89+ cs.mParam.numKeyValueHeads = CASE2_NUMKVHEADS;
90+ cs.mParam.selectBlockSize = CASE2_SELECTBLOCKSIZE;
91+ cs.mParam.selectBlockCount = CASE2_SELECTBLOCKCOUNT;
92+ cs.mParam.compressBlockSize = CASE2_COMPRESSBLOCKSIZE;
93+ cs.mParam.compressStride = CASE2_COMPRESSSTRIDE;
94+ cs.mParam.scaleValue = 1.0;
95+ cs.mParam.layoutOptional = "BSND";
96+ cs.mParam.pageBlockSize = CASE2_PAGEBLOCKSIZE;
97+ cs.mParam.sparseMode = 0;
98+ 
99+ cs.mOpInfo.mExp.mTilingBlockDim = CASE2_MTILINGBLOCKDIM;
100+ cs.mOpInfo.mCtr.mRunTiling = true;
101+ cs.mOpInfo.mCtr.mRunKernel = true;
102+ ASSERT_TRUE(cs.Init());
103+ ASSERT_TRUE(cs.Run());
104+}
105+ 
106+// getmaxworkspace test
107+TEST_F(Ts_ncaInfer, case_nsa_compress_attention_infer_aclgraph_kernel)
108+{
109+ constexpr uint32_t CASE3_B = 20;
110+ constexpr int32_t CASE3_S2 = -1;
111+ constexpr uint32_t CASE3_NUMHEADS = 64;
112+ constexpr uint32_t CASE3_NUMKVHEADS = 4;
113+ constexpr uint32_t CASE3_SELECTBLOCKSIZE = 64;
114+ constexpr uint32_t CASE3_SELECTBLOCKCOUNT = 16;
115+ constexpr uint32_t CASE3_COMPRESSBLOCKSIZE = 32;
116+ constexpr uint32_t CASE3_COMPRESSSTRIDE = 16;
117+ constexpr uint32_t CASE3_PAGEBLOCKSIZE = 128;
118+ constexpr uint32_t CASE3_MTILINGBLOCKDIM = 24;
119+ constexpr uint32_t CASE3_HEADSIZE_QK = 192;
120+ 
121+ ncaInferCase cs;
122+ /* 设置参数 */
123+ cs.mParam.B = CASE3_B;
124+ cs.mParam.S1 = 1;
125+ cs.mParam.S2 = CASE3_S2;
126+ cs.mParam.N1 = CASE3_NUMHEADS;
127+ cs.mParam.G = CASE3_NUMHEADS / CASE3_NUMKVHEADS;
128+ cs.mParam.D = CASE3_HEADSIZE_QK;
129+ cs.mParam.N2 = CASE3_NUMKVHEADS;
130+
131+ // attr
132+ cs.mParam.numHeads = CASE3_NUMHEADS;
133+ cs.mParam.numKeyValueHeads = CASE3_NUMKVHEADS;
134+ cs.mParam.selectBlockSize = CASE3_SELECTBLOCKSIZE;
135+ cs.mParam.selectBlockCount = CASE3_SELECTBLOCKCOUNT;
136+ cs.mParam.compressBlockSize = CASE3_COMPRESSBLOCKSIZE;
137+ cs.mParam.compressStride = CASE3_COMPRESSSTRIDE;
138+ cs.mParam.scaleValue = 1.0;
139+ cs.mParam.layoutOptional = "TND";
140+ cs.mParam.pageBlockSize = CASE3_PAGEBLOCKSIZE;
141+ cs.mParam.sparseMode = 0;
142+ 
143+ cs.mOpInfo.mExp.mTilingBlockDim = CASE3_MTILINGBLOCKDIM;
144+ cs.mOpInfo.mCtr.mRunTiling = true;
145+ cs.mOpInfo.mCtr.mRunKernel = true;
146+ ASSERT_TRUE(cs.Init());
147+ ASSERT_TRUE(cs.Run());
148+}
149+ 
150+TEST_F(Ts_ncaInfer, case_nsa_compress_attention_infer_mtp_kernel)
151+{
152+ constexpr uint32_t CASE4_B = 20;
153+ constexpr uint32_t CASE4_S1 = 2;
154+ constexpr uint32_t CASE4_S2 = 4096;
155+ constexpr uint32_t CASE4_NUMHEADS = 64;
156+ constexpr uint32_t CASE4_NUMKVHEADS = 4;
157+ constexpr uint32_t CASE4_SELECTBLOCKSIZE = 64;
158+ constexpr uint32_t CASE4_SELECTBLOCKCOUNT = 16;
159+ constexpr uint32_t CASE4_COMPRESSBLOCKSIZE = 32;
160+ constexpr uint32_t CASE4_COMPRESSSTRIDE = 16;
161+ constexpr uint32_t CASE4_PAGEBLOCKSIZE = 128;
162+ constexpr uint32_t CASE4_MTILINGBLOCKDIM = 24;
163+ constexpr uint32_t CASE4_HEADSIZE_QK = 192;
164+ 
165+ ncaInferCase cs;
166+ /* 设置参数 */
167+ cs.mParam.B = CASE4_B;
168+ cs.mParam.S1 = CASE4_S1;
169+ cs.mParam.S2 = CASE4_S2;
170+ cs.mParam.N1 = CASE4_NUMHEADS;
171+ cs.mParam.G = CASE4_NUMHEADS / CASE4_NUMKVHEADS;
172+ cs.mParam.D = CASE4_HEADSIZE_QK;
173+ cs.mParam.N2 = CASE4_NUMKVHEADS;
174+
175+ // attr
176+ cs.mParam.numHeads = CASE4_NUMHEADS;
177+ cs.mParam.numKeyValueHeads = CASE4_NUMKVHEADS;
178+ cs.mParam.selectBlockSize = CASE4_SELECTBLOCKSIZE;
179+ cs.mParam.selectBlockCount = CASE4_SELECTBLOCKCOUNT;
180+ cs.mParam.compressBlockSize = CASE4_COMPRESSBLOCKSIZE;
181+ cs.mParam.compressStride = CASE4_COMPRESSSTRIDE;
182+ cs.mParam.scaleValue = 1.0;
183+ cs.mParam.layoutOptional = "TND";
184+ cs.mParam.pageBlockSize = CASE4_PAGEBLOCKSIZE;
185+ cs.mParam.sparseMode = 0;
186+ 
187+ cs.mOpInfo.mExp.mTilingBlockDim = CASE4_MTILINGBLOCKDIM;
188+ cs.mOpInfo.mCtr.mRunTiling = true;
189+ cs.mOpInfo.mCtr.mRunKernel = true;
190+ ASSERT_TRUE(cs.Init());
191+ ASSERT_TRUE(cs.Run());
192+}