Star1098
2668
代码
代码
Issues196
Pull Requests534
Actions
项目讨论
Wiki
分析
Star1098
2668
已合并
【NSAG】UT补充 #5703
xuchengyang创建于 5月21日
【NSAG】UT补充 #5703
已合并
xuchengyang创建于 5月21日
从xuchengyang/ops-transformernsag_ut合入到cann/ops-transformermaster
共 12 个文件变更+845-537
Mattention/nsa_selected_attention_grad/tests/CMakeLists.txt+11-29
@@ -8,34 +8,16 @@
8# See LICENSE in the root of the software repository for the full text of the License.8# See LICENSE in the root of the software repository for the full text of the License.
9# -----------------------------------------------------------------------------------------------------------9# -----------------------------------------------------------------------------------------------------------
10 10 
11-########################################################################################################################11+set(_nsa_selected_attention_grad_TilingSourcesExt
12-# 调用编译方法, 生成对应编译目标12+ ${OPS_TRANSFORMER_DIR}/attention/nsa_selected_attention_grad/op_host/nsa_selected_attention_grad_tiling.cpp
13-########################################################################################################################13+ ${OPS_TRANSFORMER_DIR}/attention/nsa_selected_attention_grad/op_host/nsa_selected_attention_grad_tiling_bs1.cpp
14- 14+ ${OPS_TRANSFORMER_DIR}/attention/nsa_selected_attention_grad/op_host/nsa_selected_attention_grad_tiling_bs1_basic.cpp
15-set(_nsag_TilingSourcesExt15+ ${OPS_TRANSFORMER_DIR}/attention/nsa_selected_attention_grad/op_host/nsa_selected_attention_grad_tiling_common.cpp
16- ${OPS_ADV_DIR}/attention/nsa_selected_attention_grad/op_host/nsa_selected_attention_grad_tiling_bs1.cpp
17- ${OPS_ADV_DIR}/attention/nsa_selected_attention_grad/op_host/nsa_selected_attention_grad_tiling_bs1_basic.cpp
18- ${OPS_ADV_DIR}/attention/nsa_selected_attention_grad/op_host/nsa_selected_attention_grad_tiling_common.cpp
19- ${OPS_ADV_DIR}/attention/nsa_selected_attention_grad/op_host/nsa_selected_attention_grad_tiling.cpp
20)16)
21 17 
22-set(_nsag_KernelTilingDataDefH18+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
23- ${OPS_TRANSFORMER_DIR}/common/include/op_host/data_copy_transpose_tiling_def.h19+foreach(SUB_DIR ${CURRENT_DIRS})
24- ${CMAKE_SOURCE_DIR}/attention/nsa_selected_attention_grad/op_host/nsa_selected_attention_grad_tiling.h20+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
25-)21+ add_subdirectory(${SUB_DIR})
26- 22+ endif()
27-set(_nsag_TargetPrivateLinkLibrariesExt23+endforeach()
28- ${UTest_NamePrefix}_OpTiling
29- error_manager
30-)
31-aux_source_directory(${CMAKE_SOURCE_DIR}/attention/nsa_selected_attention_grad/op_kernel _nsag_KernelSourcesExt)
32- 
33-OpsTest_Level2_AddOp(
34- SUB_SYSTEM transformer/attention
35- BRIEF nsaGrad
36- SNAKE nsa_selected_attention_grad
37- KERNEL_TILING_DATA_DEF_H ${_nsag_KernelTilingDataDefH}
38- KERNEL_SOURCES_EXT ${_nsag_KernelSourcesExt}
39- TILING_SOURCES_EXT ${_nsag_TilingSourcesExt}
40- UTEST_COMMON_PRIVATE_LINK_LIBRARIES_EXT ${_nsag_TargetPrivateLinkLibrariesExt}
41-)
Dattention/nsa_selected_attention_grad/tests/comm/nsag_case.cpp+0-228
@@ -1,228 +0,0 @@
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 nsag_case.cpp
13- * \brief NativeSelectedAttention / NativeSelectedAttentionGrad 测试用例.
14- */
15- 
16-#include "nsag_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 "op_host/tiling_base.h"
24-using namespace Ops::Transformer::OpTiling;
25-/**
26- * 以下函数声明需要保持与 CMakeList.txt 中调用 OpsTest_Level2_AddOp 函数时 KERNEL_PRIVATE_COMPILE_DEFINITIONS_EXT
27- * 参数所控制的 Kernel 入口一致.
28- */
29- 
30-#define NSAG_KERNEL_PARAM \
31- (__gm__ uint8_t * query, __gm__ uint8_t * key, __gm__ uint8_t * value, __gm__ uint8_t * attention_out, \
32- __gm__ uint8_t * attention_out_grad, __gm__ uint8_t * softmax_max, __gm__ uint8_t * softmax_sum, \
33- __gm__ uint8_t * topk_indices, __gm__ uint8_t * actual_seq_qlen, __gm__ uint8_t * actual_seq_kvlen, \
34- __gm__ uint8_t * atten_mask, __gm__ uint8_t * dq, __gm__ uint8_t * dk, __gm__ uint8_t * dv, \
35- __gm__ uint8_t * workspace, __gm__ uint8_t * tiling)
36- 
37-typedef void(*NsagKernalFunc) NSAG_KERNEL_PARAM;
38- 
39-extern "C" __global__ __aicore__ void nsa_selected_attention_grad NSAG_KERNEL_PARAM;
40- 
41-using namespace ops::adv::tests::nsaGrad;
42-using TensorIntf = ops::adv::tests::utils::TensorIntf;
43-using Case = ops::adv::tests::utils::Case;
44-using Platform = ops::adv::tests::utils::Platform;
45- 
46-bool RunNsag(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- // Kernel 运行
50- auto kernelFunc = (NsagKernalFunc)func;
51- ICPU_SET_TILING_KEY(tilingKey);
52- ICPU_RUN_KF(kernelFunc, blockDim, inputs[0]->GetDevData(), inputs[1]->GetDevData(), inputs[2]->GetDevData(),
53- inputs[3]->GetDevData(), inputs[4]->GetDevData(), inputs[5]->GetDevData(), inputs[6]->GetDevData(),
54- inputs[7]->GetDevData(), inputs[8]->GetDevData(), inputs[9]->GetDevData(), inputs[10]->GetDevData(),
55- outputs[0]->GetDevData(), // queryOut
56- outputs[1]->GetDevData(), // valueOut
57- outputs[2]->GetDevData(), // keyOut
58- workspace, tilingData);
59- return true;
60-}
61- 
62-extern "C" ge::graphStatus TilingSelectedAttentionStub(gert::TilingContext *context)
63-{
64- auto *nsaGradCase = static_cast<NsaGradCase *>(Case::GetCurrentCase());
65- if (nsaGradCase != nullptr) {
66- NsaGradCase::DoTilingParam p;
67- p.ctx = context;
68- p.ret = ge::GRAPH_SUCCESS;
69- p.actualSeqQLengthsTensor = const_cast<gert::Tensor *>(context->GetOptionalInputTensor(8));
70- p.actualSeqKVLengthsTensor = const_cast<gert::Tensor *>(context->GetOptionalInputTensor(9));
71- if (!nsaGradCase->DoOpTiling(p)) {
72- return p.ret;
73- }
74- return nsaGradCase->NsagTilingFunc(context);
75- }
76- return ge::GRAPH_FAILED;
77-}
78- 
79-bool NsaGradCase::InitParam()
80-{
81- int64_t t1 = mParam.B * mParam.S1;
82- int64_t t2 = mParam.B * mParam.S2;
83- auto queryShape = {t1, mParam.N1, mParam.D};
84- auto keyShape = {t2, mParam.N2, mParam.D};
85- auto valueShape = {t2, mParam.N2, mParam.D2};
86- auto attentionOutShape = {t1, mParam.N1, mParam.D2};
87- auto layout = "TND";
88- auto dtype = ge::DataType::DT_BF16;
89- auto format = ge::FORMAT_ND;
90- 
91- query = Tensor("query", queryShape, layout, dtype, format);
92- key = Tensor("key", keyShape, layout, dtype, format);
93- value = Tensor("value", valueShape, layout, dtype, format);
94- attention_out = Tensor("attentionOut", attentionOutShape, layout, dtype, format);
95- attention_out_grad = Tensor("attentionOutGrad", attentionOutShape, layout, dtype, format);
96- softmax_max = Tensor("softmaxMax", {t1, mParam.N1, 8}, layout, ge::DataType::DT_FLOAT, format);
97- softmax_sum = Tensor("softmaxSum", {t1, mParam.N1, 8}, layout, ge::DataType::DT_FLOAT, format);
98- topk_indices =
99- Tensor("topkIndices", {t1, mParam.N2, mParam.SelectedBlockCount}, layout, ge::DataType::DT_INT32, format);
100- 
101- // option input
102- actual_seq_qlen = Tensor("actualSeqQlen", {mParam.B}, "B", ge::DataType::DT_INT64, format);
103- actual_seq_kvlen = Tensor("actualSeqKvlen", {mParam.B}, "B", ge::DataType::DT_INT64, format);
104- atten_mask = Tensor("attenMask", {}, "None", ge::DataType::DT_BOOL, format);
105- 
106- // output
107- dq = Tensor("queryOut", queryShape, layout, dtype, format);
108- dk = Tensor("keyOut", keyShape, layout, dtype, format);
109- dv = Tensor("valueOut", valueShape, layout, dtype, format);
110- 
111- 
112- mParam.actualSeqQData = std::vector(mParam.B, static_cast<int64_t>(0));
113- mParam.actualSeqKVData = std::vector(mParam.B, static_cast<int64_t>(0));
114- std::vector<int64_t> topkData(t1 * mParam.N2 * mParam.SelectedBlockCount);
115- 
116- // init tensor data
117- for (auto i = 0; i < t1 * mParam.N2 * mParam.SelectedBlockCount; i++) {
118- topkData[i] = i;
119- }
120- for (auto i = 0; i < mParam.B; i++) {
121- mParam.actualSeqQData[i] = i + 1;
122- mParam.actualSeqKVData[i] = (i + 1) * mParam.SelectedBlockCount * mParam.SelectedBlockSize;
123- }
124- 
125- if (!InitTensor(actual_seq_qlen, mParam.actualSeqQData)) {
126- return false;
127- }
128- if (!InitTensor(actual_seq_kvlen, mParam.actualSeqKVData)) {
129- return false;
130- }
131- if (!InitTensor(topk_indices, topkData)) {
132- return false;
133- }
134- 
135- return true;
136-}
137- 
138-bool NsaGradCase::InitOpInfo()
139-{
140- auto *nsagKernalFunc = (void *)nsa_selected_attention_grad;
141- 
142- bool rst = mCtx.SetOpName("NsaSelectedAttentionGrad");
143- rst = rst && mCtx.SetDeterministic(isDeterministic);
144- rst = rst && mCtx.SetInputs({&query, &key, &value, &attention_out, &attention_out_grad, &softmax_max, &softmax_sum,
145- &topk_indices, &actual_seq_qlen, &actual_seq_kvlen, &atten_mask});
146- rst = rst && mCtx.SetOutputs({&dq, &dk, &dv});
147- rst = rst && mCtx.SetAttrs({{"scaleValue", mParam.scaleValue},
148- {"SelectedBlockCount", mParam.SelectedBlockCount},
149- {"SelectedBlockSize", mParam.SelectedBlockSize},
150- {"headNum", mParam.N1},
151- {"inputLayout", mParam.inputLayout},
152- {"sparseMode", mParam.sparseMode}});
153- rst = rst && mCtx.SetKernelRunCbf(RunNsag);
154- rst = rst && mCtx.SetTilingDataMaxSize(2280); // max tilingDataLen
155- rst = rst && mCtx.SetKernelMainFunc(nsagKernalFunc);
156- rst = rst && mOpInfo.SetContext(&mCtx);
157- 
158- auto *platform = Platform::GetGlobalPlatform();
159- if (platform == nullptr) {
160- LOG_ERR("Global Platform is null");
161- return false;
162- }
163- 
164- NsagTilingFunc =
165- (gert::OpImplRegisterV2::TilingKernelFunc)platform->LoadOpTilingSoSym("TilingNsaSelectedAttentionGrad");
166- if (NsagTilingFunc == nullptr) {
167- LOG_ERR("Can't get origin tiling func, nsag(%p)", NsagTilingFunc);
168- return false;
169- }
170- IMPL_OP(NsaSelectedAttentionGrad).Tiling(TilingSelectedAttentionStub);
171- return rst;
172-}
173- 
174-bool NsaGradCase::InitCurrentCasePtr()
175-{
176- Case::mCurrentCasePtr = this;
177- return true;
178-}
179- 
180-bool NsaGradCase::Run()
181-{
182- if (!mEnable) {
183- return true;
184- }
185- if (!mOpInfo.ProcessTiling(mName)) {
186- return false;
187- }
188- if (!mOpInfo.ProcessKernel(mName)) {
189- return false;
190- }
191- return true;
192-}
193- 
194-NsaGradCase::NsaGradCase(const char *name, bool enable, const char *dbgInfo, OpInfo incre, nsaGradParam param)
195- : Case(name, enable, dbgInfo), mOpInfo(std::move(incre)), mParam(std::move(param))
196-{
197- this->mOpInfo.mName = "NsaSelectedAttentionGrad";
198-}
199- 
200-NsaGradCase::NsaGradCase()
201-{
202-}
203- 
204-NsaGradCase::nsaGradParam::nsaGradParam()
205-{
206-}
207- 
208-NsaGradCase::nsaGradParam::nsaGradParam(int64_t pB, int64_t pS1, int64_t pS2, int64_t pN1, int64_t pD, int64_t pD2,
209- int64_t pN2, int64_t pSelectedBlockCount, int64_t pSelectedBlockSize,
210- float pScaleValue, std::string pInputLayout, int64_t pSparseMode)
211- : B(pB), S1(pS1), S2(pS2), N1(pN1), D(pD), D2(pD2), N2(pN2), SelectedBlockCount(pSelectedBlockCount),
212- SelectedBlockSize(pSelectedBlockSize), scaleValue(pScaleValue), inputLayout(pInputLayout), sparseMode(pSparseMode)
213-{
214-}
215- 
216-bool NsaGradCase::DoOpTiling(DoTilingParam &tilingParam)
217-{
218- if (tilingParam.ctx == nullptr) {
219- return false;
220- }
221- if (tilingParam.actualSeqQLengthsTensor != nullptr && mParam.actualSeqQData.size() != 0) {
222- tilingParam.actualSeqQLengthsTensor->SetData(gert::TensorData{mParam.actualSeqQData.data()});
223- }
224- if (tilingParam.actualSeqKVLengthsTensor != nullptr && mParam.actualSeqKVData.size() != 0) {
225- tilingParam.actualSeqKVLengthsTensor->SetData(gert::TensorData{mParam.actualSeqKVData.data()});
226- }
227- return true;
228-}
Dattention/nsa_selected_attention_grad/tests/comm/nsag_case.h+0-104
@@ -1,104 +0,0 @@
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 ifa_case.h
13- * \brief IncreFlashAttention 测试用例.
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::nsaGrad {
29- 
30-/**
31- * 算子 NativeSelectedAttentionGrad 参数
32- */
33- 
34-class NsaGradCase : 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 nsaGradParam {
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 D = 0;
49- int64_t D2 = 0;
50- int64_t N2 = 0;
51- int64_t SelectedBlockCount = 1;
52- int64_t SelectedBlockSize = 1;
53- float scaleValue = 1.0;
54- int64_t headNum = 1;
55- std::string inputLayout = "TND";
56- int64_t sparseMode = 0;
57- ge::DataType qDataType = ge::DataType::DT_FLOAT16;
58- ge::DataType kvDataType = ge::DataType::DT_FLOAT16;
59- ge::DataType topkIndicesDataType = ge::DataType::DT_INT32;
60- ge::DataType outDataType = ge::DataType::DT_FLOAT16;
61- std::vector<int64_t> actualSeqQData = {};
62- std::vector<int64_t> actualSeqKVData = {};
63- nsaGradParam();
64- nsaGradParam(int64_t pB, int64_t pS1, int64_t pS2, int64_t pN1, int64_t pD, int64_t pD2, int64_t pN2,
65- int64_t pSelectedBlockCount, int64_t pSelectedBlockSize, float pScaleValue,
66- std::string pInputLayout, int64_t pSparseMode);
67- };
68- class DoTilingParam {
69- public:
70- gert::TilingContext *ctx = nullptr;
71- ge::graphStatus ret = ge::GRAPH_SUCCESS;
72- gert::Tensor *actualSeqQLengthsTensor = nullptr;
73- gert::Tensor *actualSeqKVLengthsTensor = nullptr;
74- };
75- 
76- Tensor query, key, value, attention_out, attention_out_grad, softmax_max, softmax_sum, topk_indices,
77- actual_seq_qlen, actual_seq_kvlen, atten_mask, dq, dk, dv;
78- OpInfo mOpInfo;
79- Context mCtx;
80- nsaGradParam mParam;
81- gert::OpImplRegisterV2::TilingKernelFunc NsagTilingFunc = nullptr;
82- bool isDeterministic = false;
83- NsaGradCase();
84- NsaGradCase(const char *name, bool enable, const char *dbgInfo, OpInfo incre, nsaGradParam param);
85- bool Run() override;
86- bool InitParam() override;
87- bool InitOpInfo() override;
88- bool InitCurrentCasePtr() override;
89- bool DoOpTiling(DoTilingParam &tilingParam);
90- template <class T> static bool InitTensor(Tensor &tensor, std::vector<T> &hostData)
91- {
92- if (hostData.empty()) {
93- return true;
94- }
95- int64_t expMinSize = hostData.size() * sizeof(T);
96- if (tensor.AllocDevData(0, expMinSize) == nullptr) {
97- printf("Tensor(%s, %ld) AllocDevData Failed.", tensor.Name().c_str(), expMinSize);
98- return false;
99- }
100- return tensor.CopyHostToDevData(hostData);
101- }
102-};
103- 
104-} // namespace ops::adv::tests::nsaGrad
Aattention/nsa_selected_attention_grad/tests/ut/CMakeLists.txt+16-0
@@ -0,0 +1,16 @@
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+foreach(SUB_DIR ${CURRENT_DIRS})
13+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
14+ add_subdirectory(${SUB_DIR})
15+ endif()
16+endforeach()
Aattention/nsa_selected_attention_grad/tests/ut/op_api/CMakeLists.txt+13-0
@@ -0,0 +1,13 @@
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+if(UT_TEST_ALL OR OP_API_UT)
12+ add_modules_ut_sources(UT_NAME ${OP_API_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
13+endif()
Aattention/nsa_selected_attention_grad/tests/ut/op_api/test_aclnn_nsa_selected_attention_grad.cpp+245-0
@@ -0,0 +1,245 @@
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 <cstdint>
12+#include <memory>
13+#include <vector>
14+ 
15+#include "gtest/gtest.h"
16+#include "../../../op_host/op_api/aclnn_nsa_selected_attention_grad.h"
17+#include "op_api_ut_common/tensor_desc.h"
18+#include "opdev/platform.h"
19+ 
20+using namespace std;
21+using namespace op;
22+ 
23+namespace {
24+void DestroyAclTensor(aclTensor *tensor)
25+{
26+ Release(tensor);
27+}
28+ 
29+void DestroyAclIntArray(aclIntArray *arr)
30+{
31+ aclDestroyIntArray(arr);
32+}
33+ 
34+using AclTensorPtr = unique_ptr<aclTensor, decltype(&DestroyAclTensor)>;
35+using AclIntArrayPtr = unique_ptr<aclIntArray, decltype(&DestroyAclIntArray)>;
36+ 
37+AclTensorPtr MakeTensor(const vector<int64_t> &shape, aclDataType dtype, aclFormat fmt = ACL_FORMAT_ND)
38+{
39+ return AclTensorPtr(TensorDesc(shape, dtype, fmt).ToAclTypeRawPtr(), DestroyAclTensor);
40+}
41+} // namespace
42+ 
43+class NsaSelectedAttentionGradOpapiUt : public testing::Test {
44+protected:
45+ static void SetUpTestCase()
46+ {
47+ op::SetPlatformSocVersion(op::SocVersion::ASCEND910B);
48+ cout << "NsaSelectedAttentionGradOpapiUt SetUp" << endl;
49+ }
50+ 
51+ static void TearDownTestCase()
52+ {
53+ cout << "NsaSelectedAttentionGradOpapiUt TearDown" << endl;
54+ }
55+};
56+ 
57+// A1: positive - typical TND, fp16, no mask. Workspace query should succeed.
58+TEST_F(NsaSelectedAttentionGradOpapiUt, A1_tnd_fp16_no_mask)
59+{
60+ constexpr int64_t b = 1;
61+ constexpr int64_t s1 = 2;
62+ constexpr int64_t s2 = 1024;
63+ constexpr int64_t t1 = b * s1;
64+ constexpr int64_t t2 = b * s2;
65+ constexpr int64_t n1 = 4;
66+ constexpr int64_t n2 = 1;
67+ constexpr int64_t d = 192;
68+ constexpr int64_t d2 = 128;
69+ constexpr int64_t selectedBlockCount = 16;
70+ constexpr int64_t selectedBlockSize = 64;
71+ 
72+ auto query = MakeTensor({t1, n1, d}, ACL_FLOAT16);
73+ auto key = MakeTensor({t2, n2, d}, ACL_FLOAT16);
74+ auto value = MakeTensor({t2, n2, d2}, ACL_FLOAT16);
75+ auto attentionOut = MakeTensor({t1, n1, d2}, ACL_FLOAT16);
76+ auto attentionOutGrad = MakeTensor({t1, n1, d2}, ACL_FLOAT16);
77+ auto softmaxMax = MakeTensor({t1, n1, 8}, ACL_FLOAT);
78+ auto softmaxSum = MakeTensor({t1, n1, 8}, ACL_FLOAT);
79+ auto topkIndices = MakeTensor({t1, n2, selectedBlockCount}, ACL_INT32);
80+ 
81+ auto dqOut = MakeTensor({t1, n1, d}, ACL_FLOAT16);
82+ auto dkOut = MakeTensor({t2, n2, d}, ACL_FLOAT16);
83+ auto dvOut = MakeTensor({t2, n2, d2}, ACL_FLOAT16);
84+ 
85+ int64_t actQ[1] = {t1};
86+ int64_t actKv[1] = {t2};
87+ AclIntArrayPtr actSeqQ(aclCreateIntArray(actQ, 1), DestroyAclIntArray);
88+ AclIntArrayPtr actSeqKv(aclCreateIntArray(actKv, 1), DestroyAclIntArray);
89+ 
90+ char inputLayout[] = "TND";
91+ uint64_t workspaceSize = 0;
92+ aclOpExecutor *executor = nullptr;
93+ 
94+ aclnnStatus aclRet = aclnnNsaSelectedAttentionGradGetWorkspaceSize(
95+ query.get(), key.get(), value.get(), attentionOut.get(), attentionOutGrad.get(),
96+ softmaxMax.get(), softmaxSum.get(), topkIndices.get(),
97+ actSeqQ.get(), actSeqKv.get(), /*attenMaskOptional=*/nullptr,
98+ /*scaleValue=*/0.088388, selectedBlockSize, selectedBlockCount,
99+ /*headNum=*/n1, inputLayout, /*sparseMode=*/0,
100+ dqOut.get(), dkOut.get(), dvOut.get(), &workspaceSize, &executor);
101+ 
102+ EXPECT_EQ(aclRet, ACL_SUCCESS);
103+ EXPECT_NE(executor, nullptr);
104+}
105+ 
106+// A2: positive - bf16 + sparse_mode=2 with atten_mask.
107+TEST_F(NsaSelectedAttentionGradOpapiUt, A2_tnd_bf16_with_mask)
108+{
109+ constexpr int64_t b = 1;
110+ constexpr int64_t t1 = 2;
111+ constexpr int64_t t2 = 2048;
112+ constexpr int64_t n1 = 4;
113+ constexpr int64_t n2 = 1;
114+ constexpr int64_t d = 192;
115+ constexpr int64_t d2 = 128;
116+ constexpr int64_t selectedBlockCount = 16;
117+ constexpr int64_t selectedBlockSize = 64;
118+ 
119+ auto query = MakeTensor({t1, n1, d}, ACL_BF16);
120+ auto key = MakeTensor({t2, n2, d}, ACL_BF16);
121+ auto value = MakeTensor({t2, n2, d2}, ACL_BF16);
122+ auto attentionOut = MakeTensor({t1, n1, d2}, ACL_BF16);
123+ auto attentionOutGrad = MakeTensor({t1, n1, d2}, ACL_BF16);
124+ auto softmaxMax = MakeTensor({t1, n1, 8}, ACL_FLOAT);
125+ auto softmaxSum = MakeTensor({t1, n1, 8}, ACL_FLOAT);
126+ auto topkIndices = MakeTensor({t1, n2, selectedBlockCount}, ACL_INT32);
127+ auto attenMask = MakeTensor({selectedBlockSize, selectedBlockSize}, ACL_BOOL);
128+ 
129+ auto dqOut = MakeTensor({t1, n1, d}, ACL_BF16);
130+ auto dkOut = MakeTensor({t2, n2, d}, ACL_BF16);
131+ auto dvOut = MakeTensor({t2, n2, d2}, ACL_BF16);
132+ 
133+ int64_t actQ[1] = {t1};
134+ int64_t actKv[1] = {t2};
135+ AclIntArrayPtr actSeqQ(aclCreateIntArray(actQ, 1), DestroyAclIntArray);
136+ AclIntArrayPtr actSeqKv(aclCreateIntArray(actKv, 1), DestroyAclIntArray);
137+ 
138+ char inputLayout[] = "TND";
139+ uint64_t workspaceSize = 0;
140+ aclOpExecutor *executor = nullptr;
141+ 
142+ aclnnStatus aclRet = aclnnNsaSelectedAttentionGradGetWorkspaceSize(
143+ query.get(), key.get(), value.get(), attentionOut.get(), attentionOutGrad.get(),
144+ softmaxMax.get(), softmaxSum.get(), topkIndices.get(),
145+ actSeqQ.get(), actSeqKv.get(), attenMask.get(),
146+ /*scaleValue=*/0.088388, selectedBlockSize, selectedBlockCount,
147+ /*headNum=*/n1, inputLayout, /*sparseMode=*/2,
148+ dqOut.get(), dkOut.get(), dvOut.get(), &workspaceSize, &executor);
149+ 
150+ EXPECT_EQ(aclRet, ACL_SUCCESS);
151+ EXPECT_NE(executor, nullptr);
152+}
153+ 
154+// E1: negative - null query should be rejected by CheckParams.
155+TEST_F(NsaSelectedAttentionGradOpapiUt, E1_null_query)
156+{
157+ auto key = MakeTensor({1024, 1, 192}, ACL_FLOAT16);
158+ auto value = MakeTensor({1024, 1, 128}, ACL_FLOAT16);
159+ auto attentionOut = MakeTensor({2, 4, 128}, ACL_FLOAT16);
160+ auto attentionOutGrad = MakeTensor({2, 4, 128}, ACL_FLOAT16);
161+ auto softmaxMax = MakeTensor({2, 4, 8}, ACL_FLOAT);
162+ auto softmaxSum = MakeTensor({2, 4, 8}, ACL_FLOAT);
163+ auto topkIndices = MakeTensor({2, 1, 16}, ACL_INT32);
164+ auto dqOut = MakeTensor({2, 4, 192}, ACL_FLOAT16);
165+ auto dkOut = MakeTensor({1024, 1, 192}, ACL_FLOAT16);
166+ auto dvOut = MakeTensor({1024, 1, 128}, ACL_FLOAT16);
167+ 
168+ char inputLayout[] = "TND";
169+ uint64_t workspaceSize = 0;
170+ aclOpExecutor *executor = nullptr;
171+ 
172+ aclnnStatus aclRet = aclnnNsaSelectedAttentionGradGetWorkspaceSize(
173+ /*query=*/nullptr, key.get(), value.get(), attentionOut.get(), attentionOutGrad.get(),
174+ softmaxMax.get(), softmaxSum.get(), topkIndices.get(),
175+ /*actualSeqQLenOptional=*/nullptr, /*actualSeqKvLenOptional=*/nullptr,
176+ /*attenMaskOptional=*/nullptr,
177+ /*scaleValue=*/0.088388, /*selectedBlockSize=*/64, /*selectedBlockCount=*/16,
178+ /*headNum=*/4, inputLayout, /*sparseMode=*/0,
179+ dqOut.get(), dkOut.get(), dvOut.get(), &workspaceSize, &executor);
180+ 
181+ EXPECT_NE(aclRet, ACL_SUCCESS);
182+ EXPECT_EQ(executor, nullptr);
183+}
184+ 
185+// E2: negative - missing topkIndices.
186+TEST_F(NsaSelectedAttentionGradOpapiUt, E2_null_topk_indices)
187+{
188+ auto query = MakeTensor({2, 4, 192}, ACL_FLOAT16);
189+ auto key = MakeTensor({1024, 1, 192}, ACL_FLOAT16);
190+ auto value = MakeTensor({1024, 1, 128}, ACL_FLOAT16);
191+ auto attentionOut = MakeTensor({2, 4, 128}, ACL_FLOAT16);
192+ auto attentionOutGrad = MakeTensor({2, 4, 128}, ACL_FLOAT16);
193+ auto softmaxMax = MakeTensor({2, 4, 8}, ACL_FLOAT);
194+ auto softmaxSum = MakeTensor({2, 4, 8}, ACL_FLOAT);
195+ auto dqOut = MakeTensor({2, 4, 192}, ACL_FLOAT16);
196+ auto dkOut = MakeTensor({1024, 1, 192}, ACL_FLOAT16);
197+ auto dvOut = MakeTensor({1024, 1, 128}, ACL_FLOAT16);
198+ 
199+ char inputLayout[] = "TND";
200+ uint64_t workspaceSize = 0;
201+ aclOpExecutor *executor = nullptr;
202+ 
203+ aclnnStatus aclRet = aclnnNsaSelectedAttentionGradGetWorkspaceSize(
204+ query.get(), key.get(), value.get(), attentionOut.get(), attentionOutGrad.get(),
205+ softmaxMax.get(), softmaxSum.get(), /*topkIndices=*/nullptr,
206+ /*actualSeqQLenOptional=*/nullptr, /*actualSeqKvLenOptional=*/nullptr,
207+ /*attenMaskOptional=*/nullptr,
208+ /*scaleValue=*/0.088388, /*selectedBlockSize=*/64, /*selectedBlockCount=*/16,
209+ /*headNum=*/4, inputLayout, /*sparseMode=*/0,
210+ dqOut.get(), dkOut.get(), dvOut.get(), &workspaceSize, &executor);
211+ 
212+ EXPECT_NE(aclRet, ACL_SUCCESS);
213+ EXPECT_EQ(executor, nullptr);
214+}
215+ 
216+// E3: negative - missing dq output.
217+TEST_F(NsaSelectedAttentionGradOpapiUt, E3_null_dq_output)
218+{
219+ auto query = MakeTensor({2, 4, 192}, ACL_FLOAT16);
220+ auto key = MakeTensor({1024, 1, 192}, ACL_FLOAT16);
221+ auto value = MakeTensor({1024, 1, 128}, ACL_FLOAT16);
222+ auto attentionOut = MakeTensor({2, 4, 128}, ACL_FLOAT16);
223+ auto attentionOutGrad = MakeTensor({2, 4, 128}, ACL_FLOAT16);
224+ auto softmaxMax = MakeTensor({2, 4, 8}, ACL_FLOAT);
225+ auto softmaxSum = MakeTensor({2, 4, 8}, ACL_FLOAT);
226+ auto topkIndices = MakeTensor({2, 1, 16}, ACL_INT32);
227+ auto dkOut = MakeTensor({1024, 1, 192}, ACL_FLOAT16);
228+ auto dvOut = MakeTensor({1024, 1, 128}, ACL_FLOAT16);
229+ 
230+ char inputLayout[] = "TND";
231+ uint64_t workspaceSize = 0;
232+ aclOpExecutor *executor = nullptr;
233+ 
234+ aclnnStatus aclRet = aclnnNsaSelectedAttentionGradGetWorkspaceSize(
235+ query.get(), key.get(), value.get(), attentionOut.get(), attentionOutGrad.get(),
236+ softmaxMax.get(), softmaxSum.get(), topkIndices.get(),
237+ /*actualSeqQLenOptional=*/nullptr, /*actualSeqKvLenOptional=*/nullptr,
238+ /*attenMaskOptional=*/nullptr,
239+ /*scaleValue=*/0.088388, /*selectedBlockSize=*/64, /*selectedBlockCount=*/16,
240+ /*headNum=*/4, inputLayout, /*sparseMode=*/0,
241+ /*dqOut=*/nullptr, dkOut.get(), dvOut.get(), &workspaceSize, &executor);
242+ 
243+ EXPECT_NE(aclRet, ACL_SUCCESS);
244+ EXPECT_EQ(executor, nullptr);
245+}
Aattention/nsa_selected_attention_grad/tests/ut/op_host/CMakeLists.txt+20-0
@@ -0,0 +1,20 @@
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+if(UT_TEST_ALL OR OP_HOST_UT)
12+ add_modules_ut_sources(UT_NAME ${OP_INFERSHAPE_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
13+endif()
14+ 
15+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
16+foreach(SUB_DIR ${CURRENT_DIRS})
17+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
18+ add_subdirectory(${SUB_DIR})
19+ endif()
20+endforeach()
Aattention/nsa_selected_attention_grad/tests/ut/op_host/arch22/CMakeLists.txt+13-0
@@ -0,0 +1,13 @@
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+if(UT_TEST_ALL OR OP_HOST_UT)
12+ add_modules_ut_sources(UT_NAME ${OP_TILING_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
13+endif()
Aattention/nsa_selected_attention_grad/tests/ut/op_host/arch22/test_nsa_selected_attention_grad_tiling.cpp+316-0
@@ -0,0 +1,316 @@
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 <gtest/gtest.h>
12+#include <iostream>
13+#include <vector>
14+ 
15+#include "tiling_context_faker.h"
16+#include "tiling_case_executor.h"
17+#include "../../../../op_host/nsa_selected_attention_grad_tiling_common.h"
18+ 
19+namespace {
20+constexpr int64_t kBlockCountBasic = 16;
21+constexpr int64_t kBlockSizeBasic = 64;
22+constexpr int64_t kBlockCountGeneral = 32;
23+ 
24+constexpr int64_t kBatch = 1;
25+constexpr int64_t kT1 = 2;
26+constexpr int64_t kT2 = 2048;
27+constexpr int64_t kN1 = 4;
28+constexpr int64_t kN2 = 1;
29+constexpr int64_t kDhead = 192;
30+constexpr int64_t kDhead2 = 128;
31+ 
32+// Basic template (priority 1) tilingKey:
33+// sparse_mode == 0 (no mask): 10
34+// sparse_mode == 2 (with mask): 11
35+constexpr uint64_t kTilingKeyBasicNoMask = 10UL;
36+constexpr uint64_t kTilingKeyBasicWithMask = 11UL;
37+ 
38+// General template (priority 10) tilingKey:
39+// sparse_mode == 0 (no mask, deterministic=0): 0
40+// sparse_mode == 2 (with mask, deterministic=0): 1
41+constexpr uint64_t kTilingKeyGeneralNoMask = 0UL;
42+ 
43+// Prefix-sum of actual_seq_q_len / actual_seq_kv_len for batch=1 setting.
44+int64_t kActualSeqQLen[1] = {kT1};
45+int64_t kActualSeqKVLen[1] = {kT2};
46+ 
47+using TensorDesc = gert::TilingContextPara::TensorDescription;
48+using OpAttr = gert::TilingContextPara::OpAttr;
49+ 
50+inline optiling::nsa::NsaSelectedAttentionGradCompileInfo MakeCompileInfo()
51+{
52+ optiling::nsa::NsaSelectedAttentionGradCompileInfo info{};
53+ info.aivNum = 1U;
54+ info.aicNum = 1U;
55+ info.ubSize = 196608ULL;
56+ info.l1Size = 524288ULL;
57+ info.l0aSize = 65536ULL;
58+ info.l0bSize = 65536ULL;
59+ info.l0cSize = 131072ULL;
60+ info.l2CacheSize = 33554432ULL;
61+ info.coreNum = 1;
62+ return info;
63+}
64+ 
65+// 11 inputs: q, k, v, attentionOut, attentionOutGrad, softmaxMax, softmaxSum,
66+// topkIndices, actualSeqQLen(opt), actualSeqKVLen(opt), attenMsk(opt)
67+std::vector<TensorDesc> MakeBaseInputs(ge::DataType dtype, int64_t blockCount,
68+ bool withMask, int64_t d = kDhead, int64_t d2 = kDhead2,
69+ int64_t n1 = kN1, int64_t n2 = kN2,
70+ int64_t t1 = kT1, int64_t t2 = kT2)
71+{
72+ std::vector<TensorDesc> inputs = {
73+ {{{t1, n1, d}, {t1, n1, d}}, dtype, ge::FORMAT_ND}, // query
74+ {{{t2, n2, d}, {t2, n2, d}}, dtype, ge::FORMAT_ND}, // key
75+ {{{t2, n2, d2}, {t2, n2, d2}}, dtype, ge::FORMAT_ND}, // value
76+ {{{t1, n1, d2}, {t1, n1, d2}}, dtype, ge::FORMAT_ND}, // attentionOut
77+ {{{t1, n1, d2}, {t1, n1, d2}}, dtype, ge::FORMAT_ND}, // attentionOutGrad
78+ {{{t1, n1, 8}, {t1, n1, 8}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmaxMax (DT_FLOAT)
79+ {{{t1, n1, 8}, {t1, n1, 8}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmaxSum (DT_FLOAT)
80+ {{{t1, n2, blockCount}, {t1, n2, blockCount}},
81+ ge::DT_INT32, ge::FORMAT_ND}, // topkIndices
82+ {{{kBatch}, {kBatch}}, ge::DT_INT64, ge::FORMAT_ND, true,
83+ static_cast<void *>(kActualSeqQLen)}, // actualSeqQLen
84+ {{{kBatch}, {kBatch}}, ge::DT_INT64, ge::FORMAT_ND, true,
85+ static_cast<void *>(kActualSeqKVLen)}, // actualSeqKVLen
86+ };
87+ if (withMask) {
88+ inputs.push_back({{{kBlockSizeBasic, kBlockSizeBasic}, {kBlockSizeBasic, kBlockSizeBasic}},
89+ ge::DT_BOOL, ge::FORMAT_ND}); // attenMsk
90+ } else {
91+ inputs.push_back({{{}, {}}, ge::DT_BOOL, ge::FORMAT_ND}); // attenMsk empty
92+ }
93+ return inputs;
94+}
95+ 
96+std::vector<TensorDesc> MakeBaseOutputs(ge::DataType dtype, int64_t d = kDhead, int64_t d2 = kDhead2,
97+ int64_t n1 = kN1, int64_t n2 = kN2,
98+ int64_t t1 = kT1, int64_t t2 = kT2)
99+{
100+ return {
101+ {{{t1, n1, d}, {t1, n1, d}}, dtype, ge::FORMAT_ND},
102+ {{{t2, n2, d}, {t2, n2, d}}, dtype, ge::FORMAT_ND},
103+ {{{t2, n2, d2}, {t2, n2, d2}}, dtype, ge::FORMAT_ND},
104+ };
105+}
106+ 
107+std::vector<OpAttr> MakeAttrs(int64_t blockCount,
108+ int64_t blockSize = kBlockSizeBasic,
109+ int64_t sparseMode = 0,
110+ const std::string &layout = "TND",
111+ int64_t headNum = kN1,
112+ float scaleValue = 0.088388f)
113+{
114+ return {
115+ {"scaleValue", Ops::Transformer::AnyValue::CreateFrom<float>(scaleValue)},
116+ {"selectedBlockCount", Ops::Transformer::AnyValue::CreateFrom<int64_t>(blockCount)},
117+ {"selectedBlockSize", Ops::Transformer::AnyValue::CreateFrom<int64_t>(blockSize)},
118+ {"headNum", Ops::Transformer::AnyValue::CreateFrom<int64_t>(headNum)},
119+ {"inputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>(layout)},
120+ {"sparseMode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(sparseMode)},
121+ };
122+}
123+} // namespace
124+ 
125+class NsaSelectedAttentionGradTiling : public testing::Test {
126+protected:
127+ static void SetUpTestCase()
128+ {
129+ std::cout << "NsaSelectedAttentionGradTiling SetUp" << std::endl;
130+ }
131+ static void TearDownTestCase()
132+ {
133+ std::cout << "NsaSelectedAttentionGradTiling TearDown" << std::endl;
134+ }
135+};
136+ 
137+// A1: positive - basic template, fp16, sparse_mode=0 (no mask).
138+// selected_block_count=16, selected_block_size=64, d=192, d2=128 -> basic IsCapable hits.
139+TEST_F(NsaSelectedAttentionGradTiling, A1_basic_fp16_no_mask)
140+{
141+ auto compileInfo = MakeCompileInfo();
142+ gert::TilingContextPara para(
143+ "NsaSelectedAttentionGrad",
144+ MakeBaseInputs(ge::DT_FLOAT16, kBlockCountBasic, /*withMask=*/false),
145+ MakeBaseOutputs(ge::DT_FLOAT16),
146+ MakeAttrs(kBlockCountBasic, kBlockSizeBasic, /*sparseMode=*/0),
147+ &compileInfo,
148+ "Ascend910B",
149+ /*coreNum=*/1,
150+ /*ubSize=*/196608ULL,
151+ /*tilingDataSize=*/8192);
152+ 
153+ ExecuteTestCase(para, ge::GRAPH_SUCCESS, kTilingKeyBasicNoMask);
154+}
155+ 
156+// A2: positive - basic template, bf16, sparse_mode=2 (with mask).
157+TEST_F(NsaSelectedAttentionGradTiling, A2_basic_bf16_with_mask)
158+{
159+ auto compileInfo = MakeCompileInfo();
160+ gert::TilingContextPara para(
161+ "NsaSelectedAttentionGrad",
162+ MakeBaseInputs(ge::DT_BF16, kBlockCountBasic, /*withMask=*/true),
163+ MakeBaseOutputs(ge::DT_BF16),
164+ MakeAttrs(kBlockCountBasic, kBlockSizeBasic, /*sparseMode=*/2),
165+ &compileInfo,
166+ "Ascend910B",
167+ 1,
168+ 196608ULL,
169+ 8192);
170+ 
171+ ExecuteTestCase(para, ge::GRAPH_SUCCESS, kTilingKeyBasicWithMask);
172+}
173+ 
174+// A3: positive - general template, fp16, sparse_mode=0.
175+// selected_block_count=32 != 16, so basic IsCapable=false, general is used.
176+TEST_F(NsaSelectedAttentionGradTiling, A3_general_fp16_no_mask)
177+{
178+ auto compileInfo = MakeCompileInfo();
179+ gert::TilingContextPara para(
180+ "NsaSelectedAttentionGrad",
181+ MakeBaseInputs(ge::DT_FLOAT16, kBlockCountGeneral, /*withMask=*/false),
182+ MakeBaseOutputs(ge::DT_FLOAT16),
183+ MakeAttrs(kBlockCountGeneral, kBlockSizeBasic, /*sparseMode=*/0),
184+ &compileInfo,
185+ "Ascend910B",
186+ 1,
187+ 196608ULL,
188+ 8192);
189+ 
190+ ExecuteTestCase(para, ge::GRAPH_SUCCESS, kTilingKeyGeneralNoMask);
191+}
192+ 
193+// E1: negative - invalid sparse_mode=1.
194+TEST_F(NsaSelectedAttentionGradTiling, E1_invalid_sparse_mode)
195+{
196+ auto compileInfo = MakeCompileInfo();
197+ gert::TilingContextPara para(
198+ "NsaSelectedAttentionGrad",
199+ MakeBaseInputs(ge::DT_FLOAT16, kBlockCountBasic, /*withMask=*/false),
200+ MakeBaseOutputs(ge::DT_FLOAT16),
201+ MakeAttrs(kBlockCountBasic, kBlockSizeBasic, /*sparseMode=*/1),
202+ &compileInfo,
203+ "Ascend910B",
204+ 1,
205+ 196608ULL,
206+ 8192);
207+ 
208+ ExecuteTestCase(para, ge::GRAPH_FAILED);
209+}
210+ 
211+// E2: negative - invalid layout "BSH".
212+TEST_F(NsaSelectedAttentionGradTiling, E2_invalid_layout_BSH)
213+{
214+ auto compileInfo = MakeCompileInfo();
215+ gert::TilingContextPara para(
216+ "NsaSelectedAttentionGrad",
217+ MakeBaseInputs(ge::DT_FLOAT16, kBlockCountBasic, /*withMask=*/false),
218+ MakeBaseOutputs(ge::DT_FLOAT16),
219+ MakeAttrs(kBlockCountBasic, kBlockSizeBasic, /*sparseMode=*/0, "BSH"),
220+ &compileInfo,
221+ "Ascend910B",
222+ 1,
223+ 196608ULL,
224+ 8192);
225+ 
226+ ExecuteTestCase(para, ge::GRAPH_FAILED);
227+}
228+ 
229+// E3: negative - q.head_dim mismatches k.head_dim.
230+// query D=128 but key D=192 -> CheckBaseShapeInfo fails.
231+TEST_F(NsaSelectedAttentionGradTiling, E3_qkv_head_dim_mismatch)
232+{
233+ auto compileInfo = MakeCompileInfo();
234+ std::vector<TensorDesc> inputs = {
235+ {{{kT1, kN1, 128}, {kT1, kN1, 128}}, ge::DT_FLOAT16, ge::FORMAT_ND}, // query d=128
236+ {{{kT2, kN2, kDhead}, {kT2, kN2, kDhead}}, ge::DT_FLOAT16, ge::FORMAT_ND}, // key d=192
237+ {{{kT2, kN2, kDhead2}, {kT2, kN2, kDhead2}}, ge::DT_FLOAT16, ge::FORMAT_ND}, // value d2=128
238+ {{{kT1, kN1, kDhead2}, {kT1, kN1, kDhead2}}, ge::DT_FLOAT16, ge::FORMAT_ND}, // attentionOut
239+ {{{kT1, kN1, kDhead2}, {kT1, kN1, kDhead2}}, ge::DT_FLOAT16, ge::FORMAT_ND}, // attentionOutGrad
240+ {{{kT1, kN1, 8}, {kT1, kN1, 8}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmaxMax
241+ {{{kT1, kN1, 8}, {kT1, kN1, 8}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmaxSum
242+ {{{kT1, kN2, kBlockCountBasic}, {kT1, kN2, kBlockCountBasic}},
243+ ge::DT_INT32, ge::FORMAT_ND}, // topkIndices
244+ {{{kBatch}, {kBatch}}, ge::DT_INT64, ge::FORMAT_ND, true,
245+ static_cast<void *>(kActualSeqQLen)}, // actualSeqQLen
246+ {{{kBatch}, {kBatch}}, ge::DT_INT64, ge::FORMAT_ND, true,
247+ static_cast<void *>(kActualSeqKVLen)}, // actualSeqKVLen
248+ {{{}, {}}, ge::DT_BOOL, ge::FORMAT_ND}, // attenMsk empty
249+ };
250+ 
251+ gert::TilingContextPara para(
252+ "NsaSelectedAttentionGrad",
253+ inputs,
254+ MakeBaseOutputs(ge::DT_FLOAT16),
255+ MakeAttrs(kBlockCountBasic, kBlockSizeBasic, /*sparseMode=*/0),
256+ &compileInfo,
257+ "Ascend910B",
258+ 1,
259+ 196608ULL,
260+ 8192);
261+ 
262+ ExecuteTestCase(para, ge::GRAPH_FAILED);
263+}
264+ 
265+// E4: negative - missing actualSeqQLen / actualSeqKVLen optional tensors in TND.
266+// Tiling reads them mandatorily for TND, so it must FAIL.
267+TEST_F(NsaSelectedAttentionGradTiling, E4_missing_actual_seq_lens)
268+{
269+ auto compileInfo = MakeCompileInfo();
270+ std::vector<TensorDesc> inputs = {
271+ {{{kT1, kN1, kDhead}, {kT1, kN1, kDhead}}, ge::DT_FLOAT16, ge::FORMAT_ND}, // query
272+ {{{kT2, kN2, kDhead}, {kT2, kN2, kDhead}}, ge::DT_FLOAT16, ge::FORMAT_ND}, // key
273+ {{{kT2, kN2, kDhead2}, {kT2, kN2, kDhead2}}, ge::DT_FLOAT16, ge::FORMAT_ND}, // value
274+ {{{kT1, kN1, kDhead2}, {kT1, kN1, kDhead2}}, ge::DT_FLOAT16, ge::FORMAT_ND}, // attentionOut
275+ {{{kT1, kN1, kDhead2}, {kT1, kN1, kDhead2}}, ge::DT_FLOAT16, ge::FORMAT_ND}, // attentionOutGrad
276+ {{{kT1, kN1, 8}, {kT1, kN1, 8}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmaxMax
277+ {{{kT1, kN1, 8}, {kT1, kN1, 8}}, ge::DT_FLOAT, ge::FORMAT_ND}, // softmaxSum
278+ {{{kT1, kN2, kBlockCountBasic}, {kT1, kN2, kBlockCountBasic}},
279+ ge::DT_INT32, ge::FORMAT_ND}, // topkIndices
280+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND}, // actualSeqQLen empty
281+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND}, // actualSeqKVLen empty
282+ {{{}, {}}, ge::DT_BOOL, ge::FORMAT_ND}, // attenMsk empty
283+ };
284+ 
285+ gert::TilingContextPara para(
286+ "NsaSelectedAttentionGrad",
287+ inputs,
288+ MakeBaseOutputs(ge::DT_FLOAT16),
289+ MakeAttrs(kBlockCountBasic, kBlockSizeBasic, /*sparseMode=*/0),
290+ &compileInfo,
291+ "Ascend910B",
292+ 1,
293+ 196608ULL,
294+ 8192);
295+ 
296+ ExecuteTestCase(para, ge::GRAPH_FAILED);
297+}
298+ 
299+// E5: negative - invalid selected_block_count (out of [1, 128] range).
300+// basic IsCapable=false (block_count!=16) so general kicks in and rejects 200.
301+TEST_F(NsaSelectedAttentionGradTiling, E5_invalid_block_count_too_large)
302+{
303+ auto compileInfo = MakeCompileInfo();
304+ gert::TilingContextPara para(
305+ "NsaSelectedAttentionGrad",
306+ MakeBaseInputs(ge::DT_FLOAT16, /*blockCount=*/200, /*withMask=*/false),
307+ MakeBaseOutputs(ge::DT_FLOAT16),
308+ MakeAttrs(/*blockCount=*/200, kBlockSizeBasic, /*sparseMode=*/0),
309+ &compileInfo,
310+ "Ascend910B",
311+ 1,
312+ 196608ULL,
313+ 8192);
314+ 
315+ ExecuteTestCase(para, ge::GRAPH_FAILED);
316+}
atomgit-bot
atomgit-botatomgit-bot5月21日

🟡 Medium Priority

被删除的旧测试 ts_nsag_tc.cpp 包含 2 个 D2=192 的正向用例（case_nsa_selcted_attention_grad_0 / _1），新 tiling UT 中所有正向用例均使用 D2=128。若 D2=192 路径存在特殊 tiling 逻辑，将失去测试保护。

likedislike
Aattention/nsa_selected_attention_grad/tests/ut/op_host/test_nsa_selected_attention_grad_infershape.cpp+211-0
@@ -0,0 +1,211 @@
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 <gtest/gtest.h>
12+#include <iostream>
13+#include <vector>
14+ 
15+#include "infer_shape_context_faker.h"
16+#include "infer_shape_case_executor.h"
17+ 
18+namespace {
19+constexpr int64_t kBlockCount = 16;
20+constexpr int64_t kBlockSize = 64;
21+constexpr int64_t kHeadDim = 192;
22+constexpr int64_t kValueHeadDim = 128;
23+constexpr int64_t kN1 = 4;
24+constexpr int64_t kN2 = 1;
25+constexpr int64_t kT1 = 2;
26+constexpr int64_t kT2 = 2048;
27+ 
28+// 11 inputs: q, k, v, attentionOut, attentionOutGrad, softmaxMax, softmaxSum,
29+// topkIndices, actualSeqQLen(opt), actualSeqKVLen(opt), attenMsk(opt)
30+const std::vector<uint32_t> kInputIrInstanceFull{1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0};
31+const std::vector<uint32_t> kOutputIrInstance{1, 1, 1};
32+ 
33+using TensorDesc = gert::InfershapeContextPara::TensorDescription;
34+using OpAttr = gert::InfershapeContextPara::OpAttr;
35+ 
36+std::vector<TensorDesc> MakeInputs(ge::DataType qkvDtype, int64_t t1 = kT1, int64_t t2 = kT2,
37+ int64_t n1 = kN1, int64_t n2 = kN2,
38+ int64_t d = kHeadDim, int64_t d2 = kValueHeadDim)
39+{
40+ return {
41+ // query: [t1, n1, d]
42+ {{{t1, n1, d}, {t1, n1, d}}, qkvDtype, ge::FORMAT_ND},
43+ // key: [t2, n2, d]
44+ {{{t2, n2, d}, {t2, n2, d}}, qkvDtype, ge::FORMAT_ND},
45+ // value: [t2, n2, d2]
46+ {{{t2, n2, d2}, {t2, n2, d2}}, qkvDtype, ge::FORMAT_ND},
47+ // attentionOut: [t1, n1, d2]
48+ {{{t1, n1, d2}, {t1, n1, d2}}, qkvDtype, ge::FORMAT_ND},
49+ // attentionOutGrad: [t1, n1, d2]
50+ {{{t1, n1, d2}, {t1, n1, d2}}, qkvDtype, ge::FORMAT_ND},
51+ // softmaxMax: [t1, n1, 8]
52+ {{{t1, n1, 8}, {t1, n1, 8}}, ge::DT_FLOAT, ge::FORMAT_ND},
53+ // softmaxSum: [t1, n1, 8]
54+ {{{t1, n1, 8}, {t1, n1, 8}}, ge::DT_FLOAT, ge::FORMAT_ND},
55+ // topkIndices: [t1, n2, selectedBlockCount]
56+ {{{t1, n2, kBlockCount}, {t1, n2, kBlockCount}}, ge::DT_INT32, ge::FORMAT_ND},
57+ // actualSeqQLen optional - empty
58+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},
59+ // actualSeqKVLen optional - empty
60+ {{{}, {}}, ge::DT_INT64, ge::FORMAT_ND},
61+ // attenMsk optional - empty
62+ {{{}, {}}, ge::DT_BOOL, ge::FORMAT_ND},
63+ };
64+}
65+ 
66+std::vector<TensorDesc> MakeEmptyOutputs()
67+{
68+ return {
69+ {{{}, {}}, ge::DT_FLOAT16, ge::FORMAT_ND},
70+ {{{}, {}}, ge::DT_FLOAT16, ge::FORMAT_ND},
71+ {{{}, {}}, ge::DT_FLOAT16, ge::FORMAT_ND},
72+ };
73+}
74+ 
75+std::vector<OpAttr> MakeAttrs(const std::string& layout,
76+ int64_t blockCount = kBlockCount,
77+ int64_t blockSize = kBlockSize,
78+ int64_t headNum = kN1,
79+ int64_t sparseMode = 0,
80+ float scaleValue = 0.088388f)
81+{
82+ return {
83+ {"scaleValue", Ops::Transformer::AnyValue::CreateFrom<float>(scaleValue)},
84+ {"selectedBlockCount", Ops::Transformer::AnyValue::CreateFrom<int64_t>(blockCount)},
85+ {"selectedBlockSize", Ops::Transformer::AnyValue::CreateFrom<int64_t>(blockSize)},
86+ {"headNum", Ops::Transformer::AnyValue::CreateFrom<int64_t>(headNum)},
87+ {"inputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>(layout)},
88+ {"sparseMode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(sparseMode)},
89+ };
90+}
91+} // namespace
92+ 
93+class NsaSelectedAttentionGradProto : public testing::Test {
94+protected:
95+ static void SetUpTestCase()
96+ {
97+ std::cout << "NsaSelectedAttentionGradProto SetUp" << std::endl;
98+ }
99+ 
100+ static void TearDownTestCase()
101+ {
102+ std::cout << "NsaSelectedAttentionGradProto TearDown" << std::endl;
103+ }
104+};
105+ 
106+// A1: positive, TND layout, fp16, single batch
107+TEST_F(NsaSelectedAttentionGradProto, A1_TND_fp16_basic)
108+{
109+ gert::InfershapeContextPara para(
110+ "NsaSelectedAttentionGrad",
111+ MakeInputs(ge::DT_FLOAT16),
112+ MakeEmptyOutputs(),
113+ MakeAttrs("TND"),
114+ kInputIrInstanceFull,
115+ kOutputIrInstance);
116+ 
117+ std::vector<std::vector<int64_t>> expectShapes = {
118+ {kT1, kN1, kHeadDim}, // dq = query shape
119+ {kT2, kN2, kHeadDim}, // dk = key shape
120+ {kT2, kN2, kValueHeadDim}, // dv = value shape
121+ };
122+ ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectShapes);
123+}
124+ 
125+// A2: positive, TND layout, bf16
126+TEST_F(NsaSelectedAttentionGradProto, A2_TND_bf16_basic)
127+{
128+ gert::InfershapeContextPara para(
129+ "NsaSelectedAttentionGrad",
130+ MakeInputs(ge::DT_BF16),
131+ MakeEmptyOutputs(),
132+ MakeAttrs("TND", kBlockCount, kBlockSize, kN1, 2),
133+ kInputIrInstanceFull,
134+ kOutputIrInstance);
135+ 
136+ std::vector<std::vector<int64_t>> expectShapes = {
137+ {kT1, kN1, kHeadDim},
138+ {kT2, kN2, kHeadDim},
139+ {kT2, kN2, kValueHeadDim},
140+ };
141+ ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectShapes);
142+}
143+ 
144+// A3: positive, lower-case "tnd" should be accepted (case-insensitive layout in infershape)
145+TEST_F(NsaSelectedAttentionGradProto, A3_tnd_lowercase_layout)
146+{
147+ gert::InfershapeContextPara para(
148+ "NsaSelectedAttentionGrad",
149+ MakeInputs(ge::DT_FLOAT16),
150+ MakeEmptyOutputs(),
151+ MakeAttrs("tnd"),
152+ kInputIrInstanceFull,
153+ kOutputIrInstance);
154+ 
155+ std::vector<std::vector<int64_t>> expectShapes = {
156+ {kT1, kN1, kHeadDim},
157+ {kT2, kN2, kHeadDim},
158+ {kT2, kN2, kValueHeadDim},
159+ };
160+ ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectShapes);
161+}
162+ 
163+// A4: positive, multi-batch TND with different t1/t2 values
164+TEST_F(NsaSelectedAttentionGradProto, A4_TND_multi_batch)
165+{
166+ constexpr int64_t kMultiT1 = 8;
167+ constexpr int64_t kMultiT2 = 8192;
168+ 
169+ gert::InfershapeContextPara para(
170+ "NsaSelectedAttentionGrad",
171+ MakeInputs(ge::DT_FLOAT16, kMultiT1, kMultiT2),
172+ MakeEmptyOutputs(),
173+ MakeAttrs("TND"),
174+ kInputIrInstanceFull,
175+ kOutputIrInstance);
176+ 
177+ std::vector<std::vector<int64_t>> expectShapes = {
178+ {kMultiT1, kN1, kHeadDim},
179+ {kMultiT2, kN2, kHeadDim},
180+ {kMultiT2, kN2, kValueHeadDim},
181+ };
182+ ExecuteTestCase(para, ge::GRAPH_SUCCESS, expectShapes);
183+}
184+ 
185+// E1: negative, invalid layout BSH
186+TEST_F(NsaSelectedAttentionGradProto, E1_invalid_layout_BSH)
187+{
188+ gert::InfershapeContextPara para(
189+ "NsaSelectedAttentionGrad",
190+ MakeInputs(ge::DT_FLOAT16),
191+ MakeEmptyOutputs(),
192+ MakeAttrs("BSH"),
193+ kInputIrInstanceFull,
194+ kOutputIrInstance);
195+ 
196+ ExecuteTestCase(para, ge::GRAPH_FAILED);
197+}
198+ 
199+// E2: negative, invalid layout SBH
200+TEST_F(NsaSelectedAttentionGradProto, E2_invalid_layout_SBH)
201+{
202+ gert::InfershapeContextPara para(
203+ "NsaSelectedAttentionGrad",
204+ MakeInputs(ge::DT_FLOAT16),
205+ MakeEmptyOutputs(),
206+ MakeAttrs("SBH"),
207+ kInputIrInstanceFull,
208+ kOutputIrInstance);
209+ 
210+ ExecuteTestCase(para, ge::GRAPH_FAILED);
211+}
Dattention/nsa_selected_attention_grad/tests/utest/ts_nsag.h+0-22
@@ -1,22 +0,0 @@
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 "nsag_case.h"
19- 
20-using nsaGradCase = ops::adv::tests::nsaGrad::NsaGradCase;
21- 
22-class Ts_nsaGrad : public Ts<nsaGradCase> {};
Dattention/nsa_selected_attention_grad/tests/utest/ts_nsag_tc.cpp+0-154
@@ -1,154 +0,0 @@
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 nsaGrad 用例.
14- */
15- 
16-#include "ts_nsag.h"
17- 
18-using Tensor = ops::adv::tests::utils::Tensor;
19- 
20-TEST_F(Ts_nsaGrad, case_nsa_selcted_attention_grad_0)
21-{
22- nsaGradCase cs;
23- cs.mParam.B = 64;
24- cs.mParam.S1 = 1;
25- cs.mParam.S2 = 1024;
26- cs.mParam.N1 = 4;
27- cs.mParam.N2 = 1;
28- cs.mParam.D = 192;
29- cs.mParam.D2 = 192;
30- cs.mParam.SelectedBlockCount = 16;
31- cs.mParam.SelectedBlockSize = 64;
32- cs.mParam.scaleValue = 1.0;
33- cs.mParam.inputLayout = "TND";
34- cs.mParam.sparseMode = 0;
35- cs.mOpInfo.mExp.mTilingBlockDim = 24; // expected block dim
36- cs.mOpInfo.mCtr.mRunTiling = true;
37- cs.mOpInfo.mCtr.mRunKernel = false;
38- ASSERT_TRUE(cs.Init());
39- ASSERT_TRUE(cs.Run());
40-}
41- 
42- 
43-TEST_F(Ts_nsaGrad, case_nsa_selcted_attention_grad_1)
44-{
45- nsaGradCase cs;
46- cs.mParam.B = 64;
47- cs.mParam.S1 = 1;
48- cs.mParam.S2 = 1024;
49- cs.mParam.N1 = 4;
50- cs.mParam.N2 = 1;
51- cs.mParam.D = 192;
52- cs.mParam.D2 = 192;
53- cs.mParam.SelectedBlockCount = 16;
54- cs.mParam.SelectedBlockSize = 64;
55- cs.mParam.scaleValue = 1.0;
56- cs.mParam.inputLayout = "TND";
57- cs.mParam.sparseMode = 2;
58- cs.mOpInfo.mExp.mTilingBlockDim = 24; // expected block dim
59- cs.mOpInfo.mCtr.mRunTiling = true;
60- cs.mOpInfo.mCtr.mRunKernel = false;
61- ASSERT_TRUE(cs.Init());
62- ASSERT_TRUE(cs.Run());
63-}
64- 
65-TEST_F(Ts_nsaGrad, case_nsa_selcted_attention_grad_diff_headdim_0)
66-{
67- nsaGradCase cs;
68- cs.mParam.B = 64;
69- cs.mParam.S1 = 1;
70- cs.mParam.S2 = 1024;
71- cs.mParam.N1 = 4;
72- cs.mParam.N2 = 1;
73- cs.mParam.D = 192;
74- cs.mParam.D2 = 128;
75- cs.mParam.SelectedBlockCount = 16;
76- cs.mParam.SelectedBlockSize = 64;
77- cs.mParam.scaleValue = 1.0;
78- cs.mParam.inputLayout = "TND";
79- cs.mParam.sparseMode = 0;
80- cs.mOpInfo.mExp.mTilingBlockDim = 24; // expected block dim
81- cs.mOpInfo.mCtr.mRunTiling = true;
82- cs.mOpInfo.mCtr.mRunKernel = false;
83- ASSERT_TRUE(cs.Init());
84- ASSERT_TRUE(cs.Run());
85-}
86- 
87- 
88-TEST_F(Ts_nsaGrad, case_nsa_selcted_attention_grad_diff_headdim_1)
89-{
90- nsaGradCase cs;
91- cs.mParam.B = 64;
92- cs.mParam.S1 = 1;
93- cs.mParam.S2 = 1024;
94- cs.mParam.N1 = 4;
95- cs.mParam.N2 = 1;
96- cs.mParam.D = 192;
97- cs.mParam.D2 = 128;
98- cs.mParam.SelectedBlockCount = 16;
99- cs.mParam.SelectedBlockSize = 64;
100- cs.mParam.scaleValue = 1.0;
101- cs.mParam.inputLayout = "TND";
102- cs.mParam.sparseMode = 2;
103- cs.mOpInfo.mExp.mTilingBlockDim = 24; // expected block dim
104- cs.mOpInfo.mCtr.mRunTiling = true;
105- cs.mOpInfo.mCtr.mRunKernel = false;
106- ASSERT_TRUE(cs.Init());
107- ASSERT_TRUE(cs.Run());
108-}
109- 
110-TEST_F(Ts_nsaGrad, case_nsa_selcted_attention_grad_diff_headdim_1_deterministic)
111-{
112- nsaGradCase cs;
113- cs.isDeterministic = true;
114- cs.mParam.B = 64;
115- cs.mParam.S1 = 1;
116- cs.mParam.S2 = 1024;
117- cs.mParam.N1 = 4;
118- cs.mParam.N2 = 1;
119- cs.mParam.D = 192;
120- cs.mParam.D2 = 128;
121- cs.mParam.SelectedBlockCount = 16;
122- cs.mParam.SelectedBlockSize = 64;
123- cs.mParam.scaleValue = 1.0;
124- cs.mParam.inputLayout = "TND";
125- cs.mParam.sparseMode = 2;
126- cs.mOpInfo.mExp.mTilingBlockDim = 24; // expected block dim
127- cs.mOpInfo.mCtr.mRunTiling = true;
128- cs.mOpInfo.mCtr.mRunKernel = false;
129- ASSERT_TRUE(cs.Init());
130- ASSERT_TRUE(cs.Run());
131-}
132- 
133-TEST_F(Ts_nsaGrad, case_nsa_selcted_attention_grad_diff_headdim_0_deterministic)
134-{
135- nsaGradCase cs;
136- cs.isDeterministic = true;
137- cs.mParam.B = 64;
138- cs.mParam.S1 = 1;
139- cs.mParam.S2 = 1024;
140- cs.mParam.N1 = 4;
141- cs.mParam.N2 = 1;
142- cs.mParam.D = 192;
143- cs.mParam.D2 = 128;
144- cs.mParam.SelectedBlockCount = 16;
145- cs.mParam.SelectedBlockSize = 64;
146- cs.mParam.scaleValue = 1.0;
147- cs.mParam.inputLayout = "TND";
148- cs.mParam.sparseMode = 0;
149- cs.mOpInfo.mExp.mTilingBlockDim = 24; // expected block dim
150- cs.mOpInfo.mCtr.mRunTiling = true;
151- cs.mOpInfo.mCtr.mRunKernel = false;
152- ASSERT_TRUE(cs.Init());
153- ASSERT_TRUE(cs.Run());
154-}