Star1662
3346
Code
Code
Issues459
Pull Requests680
Actions
Discussion
Wiki
Analysis
Star1662
3346
Merged
fix(pre-commit): 全仓OAT合规整改——许可证/版权头存量修复与豁免策略对齐 #13363
景明Add on Oct 7
fix(pre-commit): 全仓OAT合规整改——许可证/版权头存量修复与豁免策略对齐 #13363
Merged
景明Add on Oct 7
Merged fromgitcode_lijd/ops-transformerfix_3intocann/ops-transformermaster
Total: 168 files changed+4476-3901
MOAT.xml+6-0
@@ -33,6 +33,8 @@
33 <filteritem type="filename" name="*.yaml" desc="YAML configuration files, no license header required"/>33 <filteritem type="filename" name="*.yaml" desc="YAML configuration files, no license header required"/>
34 <filteritem type="filename" name="*.yml" desc="YAML configuration files, no license header required"/>34 <filteritem type="filename" name="*.yml" desc="YAML configuration files, no license header required"/>
35 <filteritem type="filename" name="*.csv" desc="CSV test case files, no license header required"/>35 <filteritem type="filename" name="*.csv" desc="CSV test case files, no license header required"/>
36+ <filteritem type="filename" name="LICENSE" desc="Root LICENSE file is the license text itself, no license header applicable"/>
37+ <filteritem type="filepath" name="experimental/.*" desc="Experimental incubator code, not subject to release compliance scanning"/>
36 </filefilter>38 </filefilter>
37 <filefilter name="copyrightPolicyFilter" desc="Filters for copyright header policies">39 <filefilter name="copyrightPolicyFilter" desc="Filters for copyright header policies">
38 <filteritem type="filename" name="requirements*.txt" desc="Dependency manifest files, copyright header not applicable"/>40 <filteritem type="filename" name="requirements*.txt" desc="Dependency manifest files, copyright header not applicable"/>
@@ -41,6 +43,8 @@
41 <!--filteritem type="filepath" name="abcdefg/.*.uvwxyz" desc="Describe the reason for filtering scan results"/-->43 <!--filteritem type="filepath" name="abcdefg/.*.uvwxyz" desc="Describe the reason for filtering scan results"/-->
42 <!--filteritem type="filepath" name="projectroot/[a-zA-Z0-9]{20,}.sh" desc="Temp files"/-->44 <!--filteritem type="filepath" name="projectroot/[a-zA-Z0-9]{20,}.sh" desc="Temp files"/-->
43 <filteritem type="filename" name="*.csv" desc="CSV test case files cannot contain a copyright header without breaking their schema"/>45 <filteritem type="filename" name="*.csv" desc="CSV test case files cannot contain a copyright header without breaking their schema"/>
46+ <filteritem type="filename" name="Third_Party_Open_Source_Software_Notice" desc="Third-party open source software notice file, contains third-party copyrights by design"/>
47+ <filteritem type="filepath" name="experimental/.*" desc="Experimental incubator code, not subject to release compliance scanning"/>
44 </filefilter>48 </filefilter>
45 <filefilter name="licenseFileNamePolicyFilter" desc="Filters for LICENSE file policies">49 <filefilter name="licenseFileNamePolicyFilter" desc="Filters for LICENSE file policies">
46 <!--filteritem type="filename" name="*.uvwxyz" desc="Describe the reason for filtering scan results"/-->50 <!--filteritem type="filename" name="*.uvwxyz" desc="Describe the reason for filtering scan results"/-->
@@ -62,6 +66,8 @@
62 <!--filteritem type="filepath" name="abcdefg/.*.uvwxyz" desc="Describe the reason for filtering scan results"/-->66 <!--filteritem type="filepath" name="abcdefg/.*.uvwxyz" desc="Describe the reason for filtering scan results"/-->
63 <!--filteritem type="filepath" name="projectroot/[a-zA-Z0-9]{20,}.sh" desc="Temp files"/-->67 <!--filteritem type="filepath" name="projectroot/[a-zA-Z0-9]{20,}.sh" desc="Temp files"/-->
64 <filteritem type="filename" name="*.png" desc="PNG image files for documentation"/>68 <filteritem type="filename" name="*.png" desc="PNG image files for documentation"/>
69+ <filteritem type="filename" name="*.jpg" desc="JPG image files for documentation"/>
70+ <filteritem type="filepath" name="experimental/.*" desc="Experimental incubator code, not subject to release compliance scanning"/>
65 </filefilter>71 </filefilter>
66 72 
67 </filefilterlist>73 </filefilterlist>
Mattention/block_sparse_attention_grad/tests/ut/op_host/test_block_sparse_attention_grad_infershape.cpp+62-67
@@ -1,6 +1,11 @@
1/**1/**
2 * Copyright (c) 2026 Huawei Technologies Co., Ltd.2 * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3- * This file is a part of the CANN Open Software.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.
4 */9 */
5 10 
6#include <gtest/gtest.h>11#include <gtest/gtest.h>
@@ -13,11 +18,13 @@
13// 测试类：专门测试 BlockSparseAttentionGrad 的 InferShape 逻辑18// 测试类：专门测试 BlockSparseAttentionGrad 的 InferShape 逻辑
14class BlockSparseAttentionGradInferShapeTest : public testing::Test {19class BlockSparseAttentionGradInferShapeTest : public testing::Test {
15protected:20protected:
16- static void SetUpTestCase() {21+ static void SetUpTestCase()
22+ {
17 std::cout << "--- BlockSparseAttentionGrad InferShape UT SetUp ---" << std::endl;23 std::cout << "--- BlockSparseAttentionGrad InferShape UT SetUp ---" << std::endl;
18 }24 }
19 25 
20- static void TearDownTestCase() {26+ static void TearDownTestCase()
27+ {
21 std::cout << "--- BlockSparseAttentionGrad InferShape UT TearDown ---" << std::endl;28 std::cout << "--- BlockSparseAttentionGrad InferShape UT TearDown ---" << std::endl;
22 }29 }
23};30};
@@ -31,53 +38,46 @@ TEST_F(BlockSparseAttentionGradInferShapeTest, infershape_bnsd_layout)
31 38 
32 gert::InfershapeContextPara infershapeContextPara(39 gert::InfershapeContextPara infershapeContextPara(
33 "BlockSparseAttentionGrad",40 "BlockSparseAttentionGrad",
34- {41+ {// 0: dout [B, N, S, D]
35- // 0: dout [B, N, S, D]42+ {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
36- {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},43+ // 1: query [B, N, S, D]
37- // 1: query [B, N, S, D] 44+ {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
38- {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},45+ // 2: key [B, N_kv, S_kv, D]
39- // 2: key [B, N_kv, S_kv, D] 46+ {{{b, n_kv, s_kv, d}, {b, n_kv, s_kv, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
40- {{{b, n_kv, s_kv, d}, {b, n_kv, s_kv, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},47+ // 3: value [B, N_kv, S_kv, D]
41- // 3: value [B, N_kv, S_kv, D] 48+ {{{b, n_kv, s_kv, d}, {b, n_kv, s_kv, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
42- {{{b, n_kv, s_kv, d}, {b, n_kv, s_kv, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},49+ // 4: attentionOut (REQUIRED)
43- // 4: attentionOut (REQUIRED)50+ {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
44- {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},51+ // 5: softmaxLse (REQUIRED, Float32)
45- // 5: softmaxLse (REQUIRED, Float32)52+ {{{b, n, s}, {b, n, s}}, ge::DT_FLOAT, ge::FORMAT_ND},
46- {{{b, n, s}, {b, n, s}}, ge::DT_FLOAT, ge::FORMAT_ND}, 53+ // 6: blockSparseMaskOptional
47- // 6: blockSparseMaskOptional54+ {{{-1}, {-1}}, ge::DT_UINT8, ge::FORMAT_ND},
48- {{{-1}, {-1}}, ge::DT_UINT8, ge::FORMAT_ND}, 55+ // 7: attenMaskOptional
49- // 7: attenMaskOptional56+ {{{-1}, {-1}}, ge::DT_UNDEFINED, ge::FORMAT_ND},
50- {{{-1}, {-1}}, ge::DT_UNDEFINED, ge::FORMAT_ND}, 57+ // 8: blockShapeOptional
51- // 8: blockShapeOptional58+ {{{2}, {2}}, ge::DT_INT64, ge::FORMAT_ND},
52- {{{2}, {2}}, ge::DT_INT64, ge::FORMAT_ND}, 59+ // 9: actualSeqLengthsOptional
53- // 9: actualSeqLengthsOptional60+ {{{b}, {b}}, ge::DT_INT64, ge::FORMAT_ND},
54- {{{b}, {b}}, ge::DT_INT64, ge::FORMAT_ND}, 61+ // 10: actualSeqLengthsKvOptional
55- // 10: actualSeqLengthsKvOptional62+ {{{b}, {b}}, ge::DT_INT64, ge::FORMAT_ND}},
56- {{{b}, {b}}, ge::DT_INT64, ge::FORMAT_ND} 63+ {// 预期的输出列表占位 (dq, dk, dv)
57- },64+ {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND},
58- {65+ {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND},
59- // 预期的输出列表占位 (dq, dk, dv)66+ {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND}},
60- {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND}, 67+ {// 属性占位 (对齐 OpDef 里的 7 个属性顺序)
61- {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND},68+ {"qInputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>("BNSD")},
62- {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND}69+ {"kvInputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>("BNSD")},
63- },70+ {"numKeyValueHeads", Ops::Transformer::AnyValue::CreateFrom<int64_t>(n_kv)},
64- {71+ {"maskType", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
65- // 属性占位 (对齐 OpDef 里的 7 个属性顺序)72+ {"scaleValue", Ops::Transformer::AnyValue::CreateFrom<float>(0.088388f)},
66- {"qInputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>("BNSD")},73+ {"preTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)},
67- {"kvInputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>("BNSD")},74+ {"nextTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)}});
68- {"numKeyValueHeads", Ops::Transformer::AnyValue::CreateFrom<int64_t>(n_kv)},
69- {"maskType", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
70- {"scaleValue", Ops::Transformer::AnyValue::CreateFrom<float>(0.088388f)},
71- {"preTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)},
72- {"nextTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)}
73- }
74- );
75 75 
76 // 断言期待结果：dQ=Q, dK=K, dV=V76 // 断言期待结果：dQ=Q, dK=K, dV=V
77 std::vector<std::vector<int64_t>> expectOutputShape = {77 std::vector<std::vector<int64_t>> expectOutputShape = {
78- {b, n, s, d}, // 0: dQ 输出78+ {b, n, s, d}, // 0: dQ 输出
79- {b, n_kv, s_kv, d}, // 1: dK 输出79+ {b, n_kv, s_kv, d}, // 1: dK 输出
80- {b, n_kv, s_kv, d} // 2: dV 输出80+ {b, n_kv, s_kv, d} // 2: dV 输出
81 };81 };
82 82 
83 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);83 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);
@@ -105,29 +105,24 @@ TEST_F(BlockSparseAttentionGradInferShapeTest, infershape_tnd_layout)
105 {{{-1}, {-1}}, ge::DT_UNDEFINED, ge::FORMAT_ND}, // 7: attenMaskOptional105 {{{-1}, {-1}}, ge::DT_UNDEFINED, ge::FORMAT_ND}, // 7: attenMaskOptional
106 {{{2}, {2}}, ge::DT_INT64, ge::FORMAT_ND}, // 8: blockShapeOptional106 {{{2}, {2}}, ge::DT_INT64, ge::FORMAT_ND}, // 8: blockShapeOptional
107 {{{batch}, {batch}}, ge::DT_INT64, ge::FORMAT_ND}, // 9: actualSeqLengthsOptional107 {{{batch}, {batch}}, ge::DT_INT64, ge::FORMAT_ND}, // 9: actualSeqLengthsOptional
108- {{{batch}, {batch}}, ge::DT_INT64, ge::FORMAT_ND} // 10: actualSeqLengthsKvOptional108+ {{{batch}, {batch}}, ge::DT_INT64, ge::FORMAT_ND} // 10: actualSeqLengthsKvOptional
109 },109 },
110- {110+ {{{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND},
111- {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND}, 111+ {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND},
112- {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND},112+ {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND}},
113- {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND}113+ {{"qInputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
114- },114+ {"kvInputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
115- {115+ {"numKeyValueHeads", Ops::Transformer::AnyValue::CreateFrom<int64_t>(n_kv)},
116- {"qInputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},116+ {"maskType", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
117- {"kvInputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},117+ {"scaleValue", Ops::Transformer::AnyValue::CreateFrom<float>(0.125f)},
118- {"numKeyValueHeads", Ops::Transformer::AnyValue::CreateFrom<int64_t>(n_kv)},118+ {"preTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)},
119- {"maskType", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},119+ {"nextTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)}});
120- {"scaleValue", Ops::Transformer::AnyValue::CreateFrom<float>(0.125f)},
121- {"preTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)},
122- {"nextTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)}
123- }
124- );
125 120 
126 std::vector<std::vector<int64_t>> expectOutputShape = {121 std::vector<std::vector<int64_t>> expectOutputShape = {
127- {total_q, n, d}, // 0: dQ 输出122+ {total_q, n, d}, // 0: dQ 输出
128- {total_kv, n_kv, d}, // 1: dK 输出123+ {total_kv, n_kv, d}, // 1: dK 输出
129- {total_kv, n_kv, d} // 2: dV 输出124+ {total_kv, n_kv, d} // 2: dV 输出
130 };125 };
131 126 
132 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);127 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);
133-}128+}
Mattention/block_sparse_attention_grad/tests/ut/op_host/test_block_sparse_attention_grad_tiling.cpp+68-72
@@ -2,6 +2,10 @@
2 * Copyright (c) 2026 Huawei Technologies Co., Ltd.2 * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * 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").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.
5 */9 */
6 10 
7#include <iostream>11#include <iostream>
@@ -9,7 +13,7 @@
9#include <string>13#include <string>
10 14 
11// 引入自动生成的 Tiling 头文件15// 引入自动生成的 Tiling 头文件
12-#include "../../../op_host/block_sparse_attention_grad_tiling.h" 16+#include "../../../op_host/block_sparse_attention_grad_tiling.h"
13#include "tiling_context_faker.h"17#include "tiling_context_faker.h"
14#include "tiling_case_executor.h"18#include "tiling_case_executor.h"
15 19 
@@ -18,22 +22,23 @@ using namespace ge;
18 22 
19class BlockSparseAttentionGradTilingTest : public testing::Test {23class BlockSparseAttentionGradTilingTest : public testing::Test {
20protected:24protected:
21- static void SetUpTestCase() {25+ static void SetUpTestCase()
26+ {
22 std::cout << "--- BlockSparseAttentionGradTiling UT SetUp ---" << std::endl;27 std::cout << "--- BlockSparseAttentionGradTiling UT SetUp ---" << std::endl;
23 }28 }
24- static void TearDownTestCase() {29+ static void TearDownTestCase()
30+ {
25 std::cout << "--- BlockSparseAttentionGradTiling UT TearDown ---" << std::endl;31 std::cout << "--- BlockSparseAttentionGradTiling UT TearDown ---" << std::endl;
26 }32 }
27};33};
28 34 
29-namespace {35+namespace {
30-void RunArch35BnsdTilingCase(int64_t s, int64_t s_kv, int64_t blockX, int64_t blockY,36+void RunArch35BnsdTilingCase(int64_t s, int64_t s_kv, int64_t blockX, int64_t blockY, ge::graphStatus expectedResult)
31- ge::graphStatus expectedResult)37+{
32-{38+ optiling::BlockSparseAttentionGradCompileInfo compileInfo;
33- optiling::BlockSparseAttentionGradCompileInfo compileInfo;39+ compileInfo.socVersion = platform_ascendc::SocVersion::ASCEND950;
34- compileInfo.socVersion = platform_ascendc::SocVersion::ASCEND950;40+ 
35- 41+ int64_t b = 1, n = 4, n_kv = 2, d = 128;
36- int64_t b = 1, n = 4, n_kv = 2, d = 128;
37 int64_t ceilQ = (s + blockX - 1) / blockX;42 int64_t ceilQ = (s + blockX - 1) / blockX;
38 int64_t ceilKv = (s_kv + blockY - 1) / blockY;43 int64_t ceilKv = (s_kv + blockY - 1) / blockY;
39 44 
@@ -43,65 +48,56 @@ void RunArch35BnsdTilingCase(int64_t s, int64_t s_kv, int64_t blockX, int64_t bl
43 48 
44 gert::TilingContextPara tilingContextPara(49 gert::TilingContextPara tilingContextPara(
45 "BlockSparseAttentionGrad",50 "BlockSparseAttentionGrad",
46- {51+ {// --- Input Info ---
47- // --- Input Info ---52+ {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
48- {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},53+ {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
49- {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},54+ {{{b, n_kv, s_kv, d}, {b, n_kv, s_kv, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
50- {{{b, n_kv, s_kv, d}, {b, n_kv, s_kv, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},55+ {{{b, n_kv, s_kv, d}, {b, n_kv, s_kv, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
51- {{{b, n_kv, s_kv, d}, {b, n_kv, s_kv, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},56+ {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
52- {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},57+ {{{b, n, s}, {b, n, s}}, ge::DT_FLOAT, ge::FORMAT_ND},
53- {{{b, n, s}, {b, n, s}}, ge::DT_FLOAT, ge::FORMAT_ND},58+ {{{b, n, ceilQ, ceilKv}, {b, n, ceilQ, ceilKv}}, ge::DT_UINT8, ge::FORMAT_ND},
54- {{{b, n, ceilQ, ceilKv}, {b, n, ceilQ, ceilKv}}, ge::DT_UINT8, ge::FORMAT_ND},59+ {{{1}, {1}}, ge::DT_UINT8, ge::FORMAT_ND},
55- {{{1}, {1}}, ge::DT_UINT8, ge::FORMAT_ND},
56-
57- // 将真实的内存地址传给框架，供 Tiling 内部解析
58- {{{2}, {2}}, ge::DT_INT64, ge::FORMAT_ND, true, (void*)blockShapeData},
59- {{{b}, {b}}, ge::DT_INT64, ge::FORMAT_ND, true, (void*)actualSeqData},
60- {{{b}, {b}}, ge::DT_INT64, ge::FORMAT_ND, true, (void*)actualSeqKvData}
61- },
62- {
63- // --- Output Info ---
64- {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
65- {{{b, n_kv, s_kv, d}, {b, n_kv, s_kv, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
66- {{{b, n_kv, s_kv, d}, {b, n_kv, s_kv, d}}, ge::DT_FLOAT16, ge::FORMAT_ND}
67- },
68- {
69- // --- Attr Info ---
70 60 
71- {"qInputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>("BNSD")},61+ // 将真实的内存地址传给框架，供 Tiling 内部解析
72- {"kvInputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>("BNSD")},62+ {{{2}, {2}}, ge::DT_INT64, ge::FORMAT_ND, true, (void*)blockShapeData},
73- {"numKeyValueHeads", Ops::Transformer::AnyValue::CreateFrom<int64_t>(n_kv)},63+ {{{b}, {b}}, ge::DT_INT64, ge::FORMAT_ND, true, (void*)actualSeqData},
74- {"maskType", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},64+ {{{b}, {b}}, ge::DT_INT64, ge::FORMAT_ND, true, (void*)actualSeqKvData}},
75- {"scaleValue", Ops::Transformer::AnyValue::CreateFrom<float>(0.088388f)}, 65+ {// --- Output Info ---
76- {"preTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)},66+ {{{b, n, s, d}, {b, n, s, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
77- {"nextTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)}67+ {{{b, n_kv, s_kv, d}, {b, n_kv, s_kv, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
78- },68+ {{{b, n_kv, s_kv, d}, {b, n_kv, s_kv, d}}, ge::DT_FLOAT16, ge::FORMAT_ND}},
79- &compileInfo,69+ {// --- Attr Info ---
80- "3510",70+ 
81- 32,71+ {"qInputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>("BNSD")},
82- 262144);72+ {"kvInputLayout", Ops::Transformer::AnyValue::CreateFrom<std::string>("BNSD")},
83- 73+ {"numKeyValueHeads", Ops::Transformer::AnyValue::CreateFrom<int64_t>(n_kv)},
84- constexpr uint64_t expectTilingKey = 1003UL; // FP16 BNSD on DAV_3510.74+ {"maskType", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
85- ExecuteTestCase(tilingContextPara, expectedResult, expectTilingKey);75+ {"scaleValue", Ops::Transformer::AnyValue::CreateFrom<float>(0.088388f)},
86-}76+ {"preTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)},
87-} // namespace77+ {"nextTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)}},
88- 78+ &compileInfo, "3510", 32, 262144);
89-TEST_F(BlockSparseAttentionGradTilingTest, tiling_arch35_block64_regression)79+ 
90-{80+ constexpr uint64_t expectTilingKey = 1003UL; // FP16 BNSD on DAV_3510.
91- RunArch35BnsdTilingCase(128, 128, 64, 64, ge::GRAPH_SUCCESS);81+ ExecuteTestCase(tilingContextPara, expectedResult, expectTilingKey);
92-}82+}
93- 83+} // namespace
94-TEST_F(BlockSparseAttentionGradTilingTest, tiling_arch35_block16_with_sequence_tail)84+ 
95-{85+TEST_F(BlockSparseAttentionGradTilingTest, tiling_arch35_block64_regression)
96- RunArch35BnsdTilingCase(70, 70, 16, 16, ge::GRAPH_SUCCESS);86+{
97-}87+ RunArch35BnsdTilingCase(128, 128, 64, 64, ge::GRAPH_SUCCESS);
98- 88+}
99-TEST_F(BlockSparseAttentionGradTilingTest, tiling_arch35_asymmetric_block16x32)89+ 
100-{90+TEST_F(BlockSparseAttentionGradTilingTest, tiling_arch35_block16_with_sequence_tail)
101- RunArch35BnsdTilingCase(70, 70, 16, 32, ge::GRAPH_SUCCESS);91+{
102-}92+ RunArch35BnsdTilingCase(70, 70, 16, 16, ge::GRAPH_SUCCESS);
103- 93+}
104-TEST_F(BlockSparseAttentionGradTilingTest, tiling_arch35_rejects_non_multiple_of_16)94+ 
105-{95+TEST_F(BlockSparseAttentionGradTilingTest, tiling_arch35_asymmetric_block16x32)
106- RunArch35BnsdTilingCase(96, 96, 24, 16, ge::GRAPH_FAILED);96+{
107-}97+ RunArch35BnsdTilingCase(70, 70, 16, 32, ge::GRAPH_SUCCESS);
98+}
99+ 
100+TEST_F(BlockSparseAttentionGradTilingTest, tiling_arch35_rejects_non_multiple_of_16)
101+{
102+ RunArch35BnsdTilingCase(96, 96, 24, 16, ge::GRAPH_FAILED);
103+}
Mattention/chunk_gated_delta_rule/op_host/chunk_gated_delta_rule_def.cpp+4-5
@@ -1,11 +1,10 @@
1/**1/**
2- * This program is free software, you can redistribute it and/or modify it.
3 * Copyright (c) 2026 Huawei Technologies Co., Ltd.2 * Copyright (c) 2026 Huawei Technologies Co., Ltd.
4- * This file is a part of the CANN Open Software.3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5- * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").4+ * CANN Open Software License Agreement Version 2.0 (the "License").
6 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
7- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8-BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9 * 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.
10 */9 */
11 10 
Mattention/chunk_kda_fwd/CMakeLists.txt+2-0
@@ -5,7 +5,9 @@
5# Please refer to the License for details. You may not use this file except in compliance with 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,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.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.
8# -----------------------------------------------------------------------------------------------------------9# -----------------------------------------------------------------------------------------------------------
10+ 
9file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)11file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
10if(NOT ENABLE_TEST AND NOT BENCHMARK)12if(NOT ENABLE_TEST AND NOT BENCHMARK)
11 list(REMOVE_ITEM CURRENT_DIRS tests)13 list(REMOVE_ITEM CURRENT_DIRS tests)
Mattention/chunk_kda_fwd/tests/CMakeLists.txt+1-0
@@ -5,4 +5,5 @@
5# Please refer to the License for details. You may not use this file except in compliance with 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,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.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.
8# -----------------------------------------------------------------------------------------------------------9# -----------------------------------------------------------------------------------------------------------
Mattention/compressor/tests/assets/impl/inputs.py+3-4
@@ -1,15 +1,14 @@
1#!/usr/bin/python1#!/usr/bin/python
2# -*- coding: utf-8 -*-2# -*- coding: utf-8 -*-
3-# -----------------------------------------------------------------------------------------------------------3+# ----------------------------------------------------------------------------
4-# Copyend (c) 2026 Huawei Technologies Co., Ltd.4+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of5# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
6# CANN Open Software License Agreement Version 2.0 (the "License").6# CANN Open Software License Agreement Version 2.0 (the "License").
7# Please refer to the License for details. You may not use this file except in compliance with the License.7# Please refer to the License for details. You may not use this file except in compliance with the License.
8# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,8# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
9# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.9# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
10# See LICENSE in the root of the software repository for the full text of the License.10# See LICENSE in the root of the software repository for the full text of the License.
11-# -----------------------------------------------------------------------------------------------------------11+# ----------------------------------------------------------------------------
12- 
13 12 
14import torch13import torch
15import random14import random
Mattention/compressor_grad/op_host/CMakeLists.txt+7-6
@@ -1,11 +1,13 @@
1-# This program is free software, you can redistribute it and/or modify.1+# -----------------------------------------------------------------------------------------------------------
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3-# This file is a part of the CANN Open Software.3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4-# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").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.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, WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.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.
7# 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.
8-# ======================================================================================================================9+# -----------------------------------------------------------------------------------------------------------
10+ 
9require_pypto_pro(compressor_grad)11require_pypto_pro(compressor_grad)
10add_op_to_compiled_list()12add_op_to_compiled_list()
11 13 
@@ -49,4 +51,3 @@ endif()
49if(NOT BUILD_OPS_RTY_KERNEL)51if(NOT BUILD_OPS_RTY_KERNEL)
50 add_tiling_modules()52 add_tiling_modules()
51endif()53endif()
52- 
Mattention/dense_lightning_indexer_grad_kl_loss/tests/ut/op_host/test_dense_lightning_indexer_grad_kl_loss_infershape.cpp+57-74
@@ -1,6 +1,11 @@
1/**1/**
2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * This file is a part of the CANN Open Software.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.
4 */9 */
5 10 
6#include <gtest/gtest.h>11#include <gtest/gtest.h>
@@ -12,11 +17,13 @@
12 17 
13class DenseLightningIndexerGradKLLossInferShapeTest : public testing::Test {18class DenseLightningIndexerGradKLLossInferShapeTest : public testing::Test {
14protected:19protected:
15- static void SetUpTestCase() {20+ static void SetUpTestCase()
21+ {
16 std::cout << "--- DenseLightningIndexerGradKLLoss InferShape UT SetUp ---" << std::endl;22 std::cout << "--- DenseLightningIndexerGradKLLoss InferShape UT SetUp ---" << std::endl;
17 }23 }
18 24 
19- static void TearDownTestCase() {25+ static void TearDownTestCase()
26+ {
20 std::cout << "--- DenseLightningIndexerGradKLLoss InferShape UT TearDown ---" << std::endl;27 std::cout << "--- DenseLightningIndexerGradKLLoss InferShape UT TearDown ---" << std::endl;
21 }28 }
22};29};
@@ -28,42 +35,30 @@ TEST_F(DenseLightningIndexerGradKLLossInferShapeTest, infershape_bsnd_layout)
28 35 
29 gert::InfershapeContextPara infershapeContextPara(36 gert::InfershapeContextPara infershapeContextPara(
30 "DenseLightningIndexerGradKLLoss",37 "DenseLightningIndexerGradKLLoss",
31- {38+ {{{{b, s1, n1, d}, {b, s1, n1, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
32- {{ {b, s1, n1, d}, {b, s1, n1, d} }, ge::DT_FLOAT16, ge::FORMAT_ND},39+ {{{b, s2, n2, d}, {b, s2, n2, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
33- {{ {b, s2, n2, d}, {b, s2, n2, d} }, ge::DT_FLOAT16, ge::FORMAT_ND},40+ {{{b, s1, nidx1, d}, {b, s1, nidx1, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
34- {{ {b, s1, nidx1, d}, {b, s1, nidx1, d} }, ge::DT_FLOAT16, ge::FORMAT_ND},41+ {{{b, s2, nidx2, d}, {b, s2, nidx2, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
35- {{ {b, s2, nidx2, d}, {b, s2, nidx2, d} }, ge::DT_FLOAT16, ge::FORMAT_ND},42+ {{{b, s1, nidx1}, {b, s1, nidx1}}, ge::DT_FLOAT16, ge::FORMAT_ND},
36- {{ {b, s1, nidx1}, {b, s1, nidx1} }, ge::DT_FLOAT16, ge::FORMAT_ND},43+ {{{b, n2, s1, g}, {b, n2, s1, g}}, ge::DT_FLOAT, ge::FORMAT_ND},
37- {{ {b, n2, s1, g}, {b, n2, s1, g} }, ge::DT_FLOAT, ge::FORMAT_ND},44+ {{{b, n2, s1, g}, {b, n2, s1, g}}, ge::DT_FLOAT, ge::FORMAT_ND},
38- {{ {b, n2, s1, g}, {b, n2, s1, g} }, ge::DT_FLOAT, ge::FORMAT_ND},45+ {{{b, nidx2, s1}, {b, nidx2, s1}}, ge::DT_FLOAT, ge::FORMAT_ND},
39- {{ {b, nidx2, s1}, {b, nidx2, s1} }, ge::DT_FLOAT, ge::FORMAT_ND},46+ {{{b, nidx2, s1}, {b, nidx2, s1}}, ge::DT_FLOAT, ge::FORMAT_ND},
40- {{ {b, nidx2, s1}, {b, nidx2, s1} }, ge::DT_FLOAT, ge::FORMAT_ND},47+ {{{b, s1, n1, dr}, {b, s1, n1, dr}}, ge::DT_FLOAT16, ge::FORMAT_ND},
41- {{ {b, s1, n1, dr}, {b, s1, n1, dr} }, ge::DT_FLOAT16, ge::FORMAT_ND},48+ {{{b, s2, n2, dr}, {b, s2, n2, dr}}, ge::DT_FLOAT16, ge::FORMAT_ND},
42- {{ {b, s2, n2, dr}, {b, s2, n2, dr} }, ge::DT_FLOAT16, ge::FORMAT_ND},49+ {{{b}, {b}}, ge::DT_INT64, ge::FORMAT_ND},
43- {{ {b}, {b} }, ge::DT_INT64, ge::FORMAT_ND},50+ {{{b}, {b}}, ge::DT_INT64, ge::FORMAT_ND}},
44- {{ {b}, {b} }, ge::DT_INT64, ge::FORMAT_ND}51+ {{{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND},
45- },52+ {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND},
46- {53+ {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND},
47- {{ {-1}, {-1} }, ge::DT_FLOAT16, ge::FORMAT_ND},54+ {{{-1}, {-1}}, ge::DT_FLOAT, ge::FORMAT_ND}},
48- {{ {-1}, {-1} }, ge::DT_FLOAT16, ge::FORMAT_ND},55+ {{"scaleValue", Ops::Transformer::AnyValue::CreateFrom<float>(0.088388)},
49- {{ {-1}, {-1} }, ge::DT_FLOAT16, ge::FORMAT_ND},56+ {"layout", Ops::Transformer::AnyValue::CreateFrom<std::string>("BSND")},
50- {{ {-1}, {-1} }, ge::DT_FLOAT, ge::FORMAT_ND}57+ {"sparseMode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(3)},
51- },58+ {"preTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)},
52- {59+ {"nextTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)}});
53- {"scaleValue", Ops::Transformer::AnyValue::CreateFrom<float>(0.088388)},
54- {"layout", Ops::Transformer::AnyValue::CreateFrom<std::string>("BSND")},
55- {"sparseMode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(3)},
56- {"preTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)},
57- {"nextTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)}
58- }
59- );
60 60 
61- std::vector<std::vector<int64_t>> expectOutputShape = {61+ std::vector<std::vector<int64_t>> expectOutputShape = {{b, s1, nidx1, d}, {b, s2, nidx2, d}, {b, s1, nidx1}, {1}};
62- {b, s1, nidx1, d},
63- {b, s2, nidx2, d},
64- {b, s1, nidx1},
65- {1}
66- };
67 62 
68 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);63 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);
69}64}
@@ -76,42 +71,30 @@ TEST_F(DenseLightningIndexerGradKLLossInferShapeTest, infershape_tnd_layout)
76 71 
77 gert::InfershapeContextPara infershapeContextPara(72 gert::InfershapeContextPara infershapeContextPara(
78 "DenseLightningIndexerGradKLLoss",73 "DenseLightningIndexerGradKLLoss",
79- {74+ {{{{t1, n1, d}, {t1, n1, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
80- {{ {t1, n1, d}, {t1, n1, d} }, ge::DT_FLOAT16, ge::FORMAT_ND},75+ {{{t2, n2, d}, {t2, n2, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
81- {{ {t2, n2, d}, {t2, n2, d} }, ge::DT_FLOAT16, ge::FORMAT_ND},76+ {{{t1, nidx1, d}, {t1, nidx1, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
82- {{ {t1, nidx1, d}, {t1, nidx1, d} }, ge::DT_FLOAT16, ge::FORMAT_ND},77+ {{{t2, nidx2, d}, {t2, nidx2, d}}, ge::DT_FLOAT16, ge::FORMAT_ND},
83- {{ {t2, nidx2, d}, {t2, nidx2, d} }, ge::DT_FLOAT16, ge::FORMAT_ND},78+ {{{t1, nidx1}, {t1, nidx1}}, ge::DT_FLOAT16, ge::FORMAT_ND},
84- {{ {t1, nidx1}, {t1, nidx1} }, ge::DT_FLOAT16, ge::FORMAT_ND},79+ {{{n2, t1, g}, {n2, t1, g}}, ge::DT_FLOAT, ge::FORMAT_ND},
85- {{ {n2, t1, g}, {n2, t1, g} }, ge::DT_FLOAT, ge::FORMAT_ND},80+ {{{n2, t1, g}, {n2, t1, g}}, ge::DT_FLOAT, ge::FORMAT_ND},
86- {{ {n2, t1, g}, {n2, t1, g} }, ge::DT_FLOAT, ge::FORMAT_ND},81+ {{{nidx2, t1}, {nidx2, t1}}, ge::DT_FLOAT, ge::FORMAT_ND},
87- {{ {nidx2, t1}, {nidx2, t1} }, ge::DT_FLOAT, ge::FORMAT_ND},82+ {{{nidx2, t1}, {nidx2, t1}}, ge::DT_FLOAT, ge::FORMAT_ND},
88- {{ {nidx2, t1}, {nidx2, t1} }, ge::DT_FLOAT, ge::FORMAT_ND},83+ {{{t1, n1, dr}, {t1, n1, dr}}, ge::DT_FLOAT16, ge::FORMAT_ND},
89- {{ {t1, n1, dr}, {t1, n1, dr} }, ge::DT_FLOAT16, ge::FORMAT_ND},84+ {{{t2, n2, dr}, {t2, n2, dr}}, ge::DT_FLOAT16, ge::FORMAT_ND},
90- {{ {t2, n2, dr}, {t2, n2, dr} }, ge::DT_FLOAT16, ge::FORMAT_ND},85+ {{{batch}, {batch}}, ge::DT_INT64, ge::FORMAT_ND},
91- {{ {batch}, {batch} }, ge::DT_INT64, ge::FORMAT_ND},86+ {{{batch}, {batch}}, ge::DT_INT64, ge::FORMAT_ND}},
92- {{ {batch}, {batch} }, ge::DT_INT64, ge::FORMAT_ND}87+ {{{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND},
93- },88+ {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND},
94- {89+ {{{-1}, {-1}}, ge::DT_FLOAT16, ge::FORMAT_ND},
95- {{ {-1}, {-1} }, ge::DT_FLOAT16, ge::FORMAT_ND},90+ {{{-1}, {-1}}, ge::DT_FLOAT, ge::FORMAT_ND}},
96- {{ {-1}, {-1} }, ge::DT_FLOAT16, ge::FORMAT_ND},91+ {{"scaleValue", Ops::Transformer::AnyValue::CreateFrom<float>(0.088388)},
97- {{ {-1}, {-1} }, ge::DT_FLOAT16, ge::FORMAT_ND},92+ {"layout", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
98- {{ {-1}, {-1} }, ge::DT_FLOAT, ge::FORMAT_ND}93+ {"sparseMode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(3)},
99- },94+ {"preTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)},
100- {95+ {"nextTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)}});
101- {"scaleValue", Ops::Transformer::AnyValue::CreateFrom<float>(0.088388)},
102- {"layout", Ops::Transformer::AnyValue::CreateFrom<std::string>("TND")},
103- {"sparseMode", Ops::Transformer::AnyValue::CreateFrom<int64_t>(3)},
104- {"preTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)},
105- {"nextTokens", Ops::Transformer::AnyValue::CreateFrom<int64_t>(2147483647)}
106- }
107- );
108 96 
109- std::vector<std::vector<int64_t>> expectOutputShape = {97+ std::vector<std::vector<int64_t>> expectOutputShape = {{t1, nidx1, d}, {t2, nidx2, d}, {t1, nidx1}, {1}};
110- {t1, nidx1, d},
111- {t2, nidx2, d},
112- {t1, nidx1},
113- {1}
114- };
115 98 
116 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);99 ExecuteTestCase(infershapeContextPara, ge::GRAPH_SUCCESS, expectOutputShape);
117-}100+}
Mattention/dense_lightning_indexer_softmax_lse/CMakeLists.txt+7-8
@@ -1,13 +1,12 @@
1-# ----------------------------------------------------------------------------1+# -----------------------------------------------------------------------------------------------------------
2-# This program is free software, you can redistribute it and/or modify it.
3# Copyright (c) 2025 Huawei Technologies Co., Ltd.2# Copyright (c) 2025 Huawei Technologies Co., Ltd.
4-# This file is a part of the CANN Open Software.3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5-# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").4+# CANN Open Software License Agreement Version 2.0 (the "License").
6# Please refer to the License for details. You may not use this file except in compliance with the License.5# Please refer to the License for details. You may not use this file except in compliance with the License.
7-# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
8-# BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
9# 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.
10-# ----------------------------------------------------------------------------9+# -----------------------------------------------------------------------------------------------------------
11 10 
12file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)11file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
13if(NOT ENABLE_TEST AND NOT BENCHMARK)12if(NOT ENABLE_TEST AND NOT BENCHMARK)
@@ -17,4 +16,4 @@ foreach(SUB_DIR ${CURRENT_DIRS})
17 if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")16 if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
18 add_subdirectory(${SUB_DIR})17 add_subdirectory(${SUB_DIR})
19 endif()18 endif()
20-endforeach()19+endforeach()
Mattention/dense_lightning_indexer_softmax_lse/op_host/CMakeLists.txt+8-6
@@ -1,11 +1,13 @@
1-# This program is free software, you can redistribute it and/or modify.1+# -----------------------------------------------------------------------------------------------------------
2# Copyright (c) 2025 Huawei Technologies Co., Ltd.2# Copyright (c) 2025 Huawei Technologies Co., Ltd.
3-# This file is a part of the CANN Open Software.3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4-# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").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.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, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.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.
7# 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.
8-# ======================================================================================================================9+# -----------------------------------------------------------------------------------------------------------
10+ 
9add_op_to_compiled_list()11add_op_to_compiled_list()
10 12 
11if (BUILD_OPEN_PROJECT)13if (BUILD_OPEN_PROJECT)
@@ -26,4 +28,4 @@ if (BUILD_OPS_RTY_KERNEL)
26 )28 )
27else()29else()
28 add_modules_sources(OPTYPE dense_lightning_indexer_softmax_lse ACLNNTYPE aclnn)30 add_modules_sources(OPTYPE dense_lightning_indexer_softmax_lse ACLNNTYPE aclnn)
29-endif()31+endif()
Mattention/dense_lightning_indexer_softmax_lse_v2/CMakeLists.txt+6-6
@@ -1,12 +1,12 @@
1-# ----------------------------------------------------------------------------1+# -----------------------------------------------------------------------------------------------------------
2-# This program is free software, you can redistribute it and/or modify it.
3# Copyright (c) 2025 Huawei Technologies Co., Ltd.2# Copyright (c) 2025 Huawei Technologies Co., Ltd.
4-# This file is a part of the CANN Open Software.3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5-# Licensed under CANN Open Software License Agreement Version 2.0 (the "License").4+# CANN Open Software License Agreement Version 2.0 (the "License").
6# Please refer to the License for details. You may not use this file except in compliance with the License.5# Please refer to the License for details. You may not use this file except in compliance with the License.
7-# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.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.8# See LICENSE in the root of the software repository for the full text of the License.
9-# ----------------------------------------------------------------------------9+# -----------------------------------------------------------------------------------------------------------
10 10 
11file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)11file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
12if(NOT ENABLE_TEST AND NOT BENCHMARK)12if(NOT ENABLE_TEST AND NOT BENCHMARK)
Mattention/dense_lightning_indexer_softmax_lse_v2/op_host/CMakeLists.txt+2-1
@@ -3,10 +3,11 @@
3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of3# 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").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.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, WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,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.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.8# See LICENSE in the root of the software repository for the full text of the License.
9# -----------------------------------------------------------------------------------------------------------9# -----------------------------------------------------------------------------------------------------------
10+ 
10cmake_minimum_required(VERSION 3.18.4)11cmake_minimum_required(VERSION 3.18.4)
11require_pypto_pro(dense_lightning_indexer_softmax_lse_v2)12require_pypto_pro(dense_lightning_indexer_softmax_lse_v2)
12add_op_to_compiled_list()13add_op_to_compiled_list()
Mattention/flash_attention_score/examples/test_aclnn_flash_attention_score.cpp+204-195
@@ -1,7 +1,7 @@
1/**1/**
2- * Copyright (c) Huawei Technologies Co., Ltd. 2025. All rights reserved.2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * This file is a part of the CANN Open Software.3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").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.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,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.7 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
@@ -21,218 +21,227 @@
21#include "aclnnop/aclnn_flash_attention_score.h"21#include "aclnnop/aclnn_flash_attention_score.h"
22 22 
23#define CHECK_RET(cond, return_expr) \23#define CHECK_RET(cond, return_expr) \
24- do { \24+ do { \
25- if (!(cond)) { \25+ if (!(cond)) { \
26- return_expr; \26+ return_expr; \
27- } \27+ } \
28- } while (0)28+ } while (0)
29 29 
30-#define LOG_PRINT(message, ...) \30+#define LOG_PRINT(message, ...) \
31- do { \31+ do { \
32- printf(message, ##__VA_ARGS__); \32+ printf(message, ##__VA_ARGS__); \
33- } while (0)33+ } while (0)
34 34 
35-int64_t GetShapeSize(const std::vector<int64_t>& shape) {35+int64_t GetShapeSize(const std::vector<int64_t>& shape)
36- int64_t shapeSize = 1;36+{
37- for (auto i : shape) {37+ int64_t shapeSize = 1;
38- shapeSize *= i;38+ for (auto i : shape) {
39- }39+ shapeSize *= i;
40- return shapeSize;40+ }
41+ return shapeSize;
41}42}
42 43 
43-void PrintOutResult(std::vector<int64_t> &shape, void** deviceAddr) {44+void PrintOutResult(std::vector<int64_t>& shape, void** deviceAddr)
44- auto size = GetShapeSize(shape);45+{
45- std::vector<float> resultData(size, 0);46+ auto size = GetShapeSize(shape);
46- auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]),47+ std::vector<float> resultData(size, 0);
47- *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);48+ auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr,
48- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return);49+ size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
49- for (int64_t i = 0; i < size; i++) {50+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return);
50- LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);51+ for (int64_t i = 0; i < size; i++) {
51- }52+ LOG_PRINT("result[%ld] is: %f\n", i, resultData[i]);
53+ }
52}54}
53 55 
54-int Init(int32_t deviceId, aclrtStream* stream) {56+int Init(int32_t deviceId, aclrtStream* stream)
55- // 固定写法，AscendCL初始化57+{
56- auto ret = aclInit(nullptr);58+ // 固定写法，AscendCL初始化
57- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);59+ auto ret = aclInit(nullptr);
58- ret = aclrtSetDevice(deviceId);60+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
59- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);61+ ret = aclrtSetDevice(deviceId);
60- ret = aclrtCreateStream(stream);62+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
61- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);63+ ret = aclrtCreateStream(stream);
62- return 0;64+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
65+ return 0;
63}66}
64 67 
65template <typename T>68template <typename T>
66int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,69int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
67- aclDataType dataType, aclTensor** tensor) {70+ aclDataType dataType, aclTensor** tensor)
68- auto size = GetShapeSize(shape) * sizeof(T);71+{
69- // 调用aclrtMalloc申请device侧内存72+ auto size = GetShapeSize(shape) * sizeof(T);
70- auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);73+ // 调用aclrtMalloc申请device侧内存
71- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);74+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
72- // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上75+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
73- ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);76+ // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
74- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);77+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
78+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
75 79 
76- // 计算连续tensor的strides80+ // 计算连续tensor的strides
77- std::vector<int64_t> strides(shape.size(), 1);81+ std::vector<int64_t> strides(shape.size(), 1);
78- for (int64_t i = shape.size() - 2; i >= 0; i--) {82+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
79- strides[i] = shape[i + 1] * strides[i + 1];83+ strides[i] = shape[i + 1] * strides[i + 1];
80- }84+ }
81 85 
82- // 调用aclCreateTensor接口创建aclTensor86+ // 调用aclCreateTensor接口创建aclTensor
83- *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,87+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
84- shape.data(), shape.size(), *deviceAddr);88+ shape.data(), shape.size(), *deviceAddr);
85- return 0;89+ return 0;
86}90}
87 91 
88-int main() {92+int main()
89- // 1. （固定写法）device/stream初始化，参考AscendCL对外接口列表93+{
90- // 根据自己的实际device填写deviceId94+ // 1. （固定写法）device/stream初始化，参考AscendCL对外接口列表
91- int32_t deviceId = 0;95+ // 根据自己的实际device填写deviceId
92- aclrtStream stream;96+ int32_t deviceId = 0;
93- auto ret = Init(deviceId, &stream);97+ aclrtStream stream;
94- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);98+ auto ret = Init(deviceId, &stream);
99+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
95 100 
96- // 2. 构造输入与输出，需要根据API的接口自定义构造101+ // 2. 构造输入与输出，需要根据API的接口自定义构造
97- int64_t B = 1;102+ int64_t B = 1;
98- int64_t N1 = 1;103+ int64_t N1 = 1;
99- int64_t N2 = 1;104+ int64_t N2 = 1;
100- int64_t S1 = 256;105+ int64_t S1 = 256;
101- int64_t S2 = 256;106+ int64_t S2 = 256;
102- int64_t D = 128;107+ int64_t D = 128;
103- int64_t H1 = N1 * D;108+ int64_t H1 = N1 * D;
104- int64_t H2 = N2 * D;109+ int64_t H2 = N2 * D;
105 110 
106- int64_t q_size = S1 * B * H1;111+ int64_t q_size = S1 * B * H1;
107- int64_t kv_size = S2 * B * H2;112+ int64_t kv_size = S2 * B * H2;
108- int64_t atten_mask_size = S1 * S2;113+ int64_t atten_mask_size = S1 * S2;
109- int64_t softmax_size = B * N1 * S1 * 8;114+ int64_t softmax_size = B * N1 * S1 * 8;
110 115 
111- std::vector<int64_t> qShape = {S1, B, H1};116+ std::vector<int64_t> qShape = {S1, B, H1};
112- std::vector<int64_t> kShape = {S2, B, H2};117+ std::vector<int64_t> kShape = {S2, B, H2};
113- std::vector<int64_t> vShape = {S2, B, H2};118+ std::vector<int64_t> vShape = {S2, B, H2};
114- std::vector<int64_t> attenmaskShape = {S1, S2};119+ std::vector<int64_t> attenmaskShape = {S1, S2};
115 120 
116- std::vector<int64_t> attentionOutShape = {S1, B, H1};121+ std::vector<int64_t> attentionOutShape = {S1, B, H1};
117- std::vector<int64_t> softmaxMaxShape = {B, N1, S1, 8};122+ std::vector<int64_t> softmaxMaxShape = {B, N1, S1, 8};
118- std::vector<int64_t> softmaxSumShape = {B, N1, S1, 8};123+ std::vector<int64_t> softmaxSumShape = {B, N1, S1, 8};
119 124 
120- void* qDeviceAddr = nullptr;125+ void* qDeviceAddr = nullptr;
121- void* kDeviceAddr = nullptr;126+ void* kDeviceAddr = nullptr;
122- void* vDeviceAddr = nullptr;127+ void* vDeviceAddr = nullptr;
123- void* attenmaskDeviceAddr = nullptr;128+ void* attenmaskDeviceAddr = nullptr;
124- void* attentionOutDeviceAddr = nullptr;129+ void* attentionOutDeviceAddr = nullptr;
125- void* softmaxMaxDeviceAddr = nullptr;130+ void* softmaxMaxDeviceAddr = nullptr;
126- void* softmaxSumDeviceAddr = nullptr;131+ void* softmaxSumDeviceAddr = nullptr;
127 132 
128- aclTensor* q = nullptr;133+ aclTensor* q = nullptr;
129- aclTensor* k = nullptr;134+ aclTensor* k = nullptr;
130- aclTensor* v = nullptr;135+ aclTensor* v = nullptr;
131- aclTensor* pse = nullptr;136+ aclTensor* pse = nullptr;
132- aclTensor* dropMask = nullptr;137+ aclTensor* dropMask = nullptr;
133- aclTensor* padding = nullptr;138+ aclTensor* padding = nullptr;
134- aclTensor* attenmask = nullptr;139+ aclTensor* attenmask = nullptr;
135- aclTensor* attentionOut = nullptr;140+ aclTensor* attentionOut = nullptr;
136- aclTensor* softmaxMax = nullptr;141+ aclTensor* softmaxMax = nullptr;
137- aclTensor* softmaxSum = nullptr;142+ aclTensor* softmaxSum = nullptr;
138- aclTensor* softmaxOut = nullptr;143+ aclTensor* softmaxOut = nullptr;
139 144 
140- std::vector<float> qHostData(q_size, 1.0);145+ std::vector<float> qHostData(q_size, 1.0);
141- std::vector<float> kHostData(kv_size, 1.0);146+ std::vector<float> kHostData(kv_size, 1.0);
142- std::vector<float> vHostData(kv_size, 1.0);147+ std::vector<float> vHostData(kv_size, 1.0);
143- std::vector<uint8_t> attenmaskHostData(atten_mask_size, 0);148+ std::vector<uint8_t> attenmaskHostData(atten_mask_size, 0);
144- std::vector<float> attentionOutHostData(q_size, 0);149+ std::vector<float> attentionOutHostData(q_size, 0);
145- std::vector<float> softmaxMaxHostData(softmax_size, 3.0); 150+ std::vector<float> softmaxMaxHostData(softmax_size, 3.0);
146- std::vector<float> softmaxSumHostData(softmax_size, 3.0);151+ std::vector<float> softmaxSumHostData(softmax_size, 3.0);
147 152 
148- ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_FLOAT, &q);153+ ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_FLOAT, &q);
149- CHECK_RET(ret == ACL_SUCCESS, return ret);154+ CHECK_RET(ret == ACL_SUCCESS, return ret);
150- ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_FLOAT, &k);155+ ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_FLOAT, &k);
151- CHECK_RET(ret == ACL_SUCCESS, return ret);156+ CHECK_RET(ret == ACL_SUCCESS, return ret);
152- ret = CreateAclTensor(vHostData, vShape, &vDeviceAddr, aclDataType::ACL_FLOAT, &v);157+ ret = CreateAclTensor(vHostData, vShape, &vDeviceAddr, aclDataType::ACL_FLOAT, &v);
153- CHECK_RET(ret == ACL_SUCCESS, return ret);158+ CHECK_RET(ret == ACL_SUCCESS, return ret);
154- ret = CreateAclTensor(attenmaskHostData, attenmaskShape, &attenmaskDeviceAddr, aclDataType::ACL_UINT8, &attenmask);159+ ret = CreateAclTensor(attenmaskHostData, attenmaskShape, &attenmaskDeviceAddr, aclDataType::ACL_UINT8, &attenmask);
155- CHECK_RET(ret == ACL_SUCCESS, return ret);160+ CHECK_RET(ret == ACL_SUCCESS, return ret);
156- ret = CreateAclTensor(attentionOutHostData, attentionOutShape , &attentionOutDeviceAddr, aclDataType::ACL_FLOAT, &attentionOut);161+ ret = CreateAclTensor(attentionOutHostData, attentionOutShape, &attentionOutDeviceAddr, aclDataType::ACL_FLOAT,
157- CHECK_RET(ret == ACL_SUCCESS, return ret);162+ &attentionOut);
158- ret = CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, &softmaxMaxDeviceAddr, aclDataType::ACL_FLOAT, &softmaxMax);163+ CHECK_RET(ret == ACL_SUCCESS, return ret);
159- CHECK_RET(ret == ACL_SUCCESS, return ret);164+ ret = CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, &softmaxMaxDeviceAddr, aclDataType::ACL_FLOAT,
160- ret = CreateAclTensor(softmaxSumHostData, softmaxSumShape, &softmaxSumDeviceAddr, aclDataType::ACL_FLOAT, &softmaxSum);165+ &softmaxMax);
161- CHECK_RET(ret == ACL_SUCCESS, return ret);166+ CHECK_RET(ret == ACL_SUCCESS, return ret);
162- 167+ ret = CreateAclTensor(softmaxSumHostData, softmaxSumShape, &softmaxSumDeviceAddr, aclDataType::ACL_FLOAT,
163- std::vector<int64_t> prefixOp = {0};168+ &softmaxSum);
164- std::vector<int64_t> qStartIdxOp = {0};169+ CHECK_RET(ret == ACL_SUCCESS, return ret);
165- std::vector<int64_t> kvStartIdxOp = {0};170+ 
166- aclIntArray *prefix = aclCreateIntArray(prefixOp.data(), 1);171+ std::vector<int64_t> prefixOp = {0};
167- aclIntArray *qStartIdx = aclCreateIntArray(qStartIdxOp.data(), 1);172+ std::vector<int64_t> qStartIdxOp = {0};
168- aclIntArray *kvStartIdx = aclCreateIntArray(kvStartIdxOp.data(), 1);173+ std::vector<int64_t> kvStartIdxOp = {0};
169- double scaleValue = 1.0/sqrt(128);174+ aclIntArray* prefix = aclCreateIntArray(prefixOp.data(), 1);
170- double keepProb = 1.0;175+ aclIntArray* qStartIdx = aclCreateIntArray(qStartIdxOp.data(), 1);
171- int64_t preTokens = INT32_MAX;176+ aclIntArray* kvStartIdx = aclCreateIntArray(kvStartIdxOp.data(), 1);
172- int64_t nextTokens = INT32_MAX;177+ double scaleValue = 1.0 / sqrt(128);
173- int64_t headNum = 1;178+ double keepProb = 1.0;
174- int64_t innerPrecise = 0;179+ int64_t preTokens = INT32_MAX;
175- int64_t sparseMode = 0;180+ int64_t nextTokens = INT32_MAX;
176- int64_t pseType = 1;181+ int64_t headNum = 1;
177- char layOut[5] = {'S', 'B', 'H', 0};182+ int64_t innerPrecise = 0;
178- 183+ int64_t sparseMode = 0;
179- // 3. 调用CANN算子库API，需要修改为具体的Api名称184+ int64_t pseType = 1;
180- uint64_t workspaceSize = 0;185+ char layOut[5] = {'S', 'B', 'H', 0};
181- aclOpExecutor* executor;186+ 
182- 187+ // 3. 调用CANN算子库API，需要修改为具体的Api名称
183- // 调用aclnnFlashAttentionScoreV2第一段接口188+ uint64_t workspaceSize = 0;
184- ret = aclnnFlashAttentionScoreV2GetWorkspaceSize(189+ aclOpExecutor* executor;
185- q, k, v, pse, dropMask, padding, attenmask, prefix, qStartIdx, kvStartIdx, scaleValue,190+ 
186- keepProb, preTokens, nextTokens, headNum, layOut, innerPrecise,191+ // 调用aclnnFlashAttentionScoreV2第一段接口
187- sparseMode, pseType, softmaxMax, softmaxSum, softmaxOut, attentionOut, &workspaceSize, &executor);192+ ret = aclnnFlashAttentionScoreV2GetWorkspaceSize(q, k, v, pse, dropMask, padding, attenmask, prefix, qStartIdx,
188- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreV2GetWorkspaceSize failed. ERROR: %d\n", ret); return ret);193+ kvStartIdx, scaleValue, keepProb, preTokens, nextTokens, headNum,
189- 194+ layOut, innerPrecise, sparseMode, pseType, softmaxMax, softmaxSum,
190- // 根据第一段接口计算出的workspaceSize申请device内存195+ softmaxOut, attentionOut, &workspaceSize, &executor);
191- void* workspaceAddr = nullptr;196+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreV2GetWorkspaceSize failed. ERROR: %d\n", ret);
192- if (workspaceSize > 0) {197+ return ret);
193- ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);198+ 
194- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);199+ // 根据第一段接口计算出的workspaceSize申请device内存
195- }200+ void* workspaceAddr = nullptr;
196- 201+ if (workspaceSize > 0) {
197- // 调用aclnnFlashAttentionScoreV2第二段接口202+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
198- ret = aclnnFlashAttentionScoreV2(workspaceAddr, workspaceSize, executor, stream);203+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
199- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreV2 failed. ERROR: %d\n", ret); return ret);204+ }
200- 205+ 
201- // 4. （固定写法）同步等待任务执行结束206+ // 调用aclnnFlashAttentionScoreV2第二段接口
202- ret = aclrtSynchronizeStream(stream);207+ ret = aclnnFlashAttentionScoreV2(workspaceAddr, workspaceSize, executor, stream);
203- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);208+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreV2 failed. ERROR: %d\n", ret); return ret);
204- 209+ 
205- // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改210+ // 4. （固定写法）同步等待任务执行结束
206- PrintOutResult(attentionOutShape, &attentionOutDeviceAddr);211+ ret = aclrtSynchronizeStream(stream);
207- PrintOutResult(softmaxMaxShape, &softmaxMaxDeviceAddr);212+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
208- PrintOutResult(softmaxSumShape, &softmaxSumDeviceAddr);213+ 
209- 214+ // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
210- // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改215+ PrintOutResult(attentionOutShape, &attentionOutDeviceAddr);
211- aclDestroyTensor(q);216+ PrintOutResult(softmaxMaxShape, &softmaxMaxDeviceAddr);
212- aclDestroyTensor(k);217+ PrintOutResult(softmaxSumShape, &softmaxSumDeviceAddr);
213- aclDestroyTensor(v);218+ 
214- aclDestroyTensor(attenmask);219+ // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
215- aclDestroyTensor(attentionOut);220+ aclDestroyTensor(q);
216- aclDestroyTensor(softmaxMax);221+ aclDestroyTensor(k);
217- aclDestroyTensor(softmaxSum);222+ aclDestroyTensor(v);
218- aclDestroyIntArray(prefix);223+ aclDestroyTensor(attenmask);
219- aclDestroyIntArray(qStartIdx);224+ aclDestroyTensor(attentionOut);
220- aclDestroyIntArray(kvStartIdx);225+ aclDestroyTensor(softmaxMax);
221- 226+ aclDestroyTensor(softmaxSum);
222- // 7. 释放device资源227+ aclDestroyIntArray(prefix);
223- aclrtFree(qDeviceAddr);228+ aclDestroyIntArray(qStartIdx);
224- aclrtFree(kDeviceAddr);229+ aclDestroyIntArray(kvStartIdx);
225- aclrtFree(vDeviceAddr);230+ 
226- aclrtFree(attenmaskDeviceAddr);231+ // 7. 释放device资源
227- aclrtFree(attentionOutDeviceAddr);232+ aclrtFree(qDeviceAddr);
228- aclrtFree(softmaxMaxDeviceAddr);233+ aclrtFree(kDeviceAddr);
229- aclrtFree(softmaxSumDeviceAddr);234+ aclrtFree(vDeviceAddr);
230- if (workspaceSize > 0) {235+ aclrtFree(attenmaskDeviceAddr);
231- aclrtFree(workspaceAddr);236+ aclrtFree(attentionOutDeviceAddr);
232- }237+ aclrtFree(softmaxMaxDeviceAddr);
233- aclrtDestroyStream(stream);238+ aclrtFree(softmaxSumDeviceAddr);
234- aclrtResetDevice(deviceId);239+ if (workspaceSize > 0) {
235- aclFinalize();240+ aclrtFree(workspaceAddr);
236- 241+ }
237- return 0;242+ aclrtDestroyStream(stream);
243+ aclrtResetDevice(deviceId);
244+ aclFinalize();
245+ 
246+ return 0;
238}247}
Mattention/flash_attention_score_grad/examples/test_aclnn_flash_attention_score_grad_v4_fp16.cpp+233-224
@@ -1,10 +1,10 @@
1/**1/**
2- * This program is free software, you can redistribute it and/or modify.
3 * Copyright (c) 2025 Huawei Technologies Co., Ltd.2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4- * This file is a part of the CANN Open Software.3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5- * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").4+ * CANN Open Software License Agreement Version 2.0 (the "License").
6 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
7- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.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.8 * See LICENSE in the root of the software repository for the full text of the License.
9 */9 */
10 10 
@@ -21,255 +21,264 @@
21#include "aclnnop/aclnn_flash_attention_score_grad.h"21#include "aclnnop/aclnn_flash_attention_score_grad.h"
22 22 
23#define CHECK_RET(cond, return_expr) \23#define CHECK_RET(cond, return_expr) \
24- do { \24+ do { \
25- if (!(cond)) { \25+ if (!(cond)) { \
26- return_expr; \26+ return_expr; \
27- } \27+ } \
28- } while (0)28+ } while (0)
29 29 
30-#define LOG_PRINT(message, ...) \30+#define LOG_PRINT(message, ...) \
31- do { \31+ do { \
32- printf(message, ##__VA_ARGS__); \32+ printf(message, ##__VA_ARGS__); \
33- } while (0)33+ } while (0)
34 34 
35-int64_t GetShapeSize(const std::vector<int64_t>& shape) {35+int64_t GetShapeSize(const std::vector<int64_t>& shape)
36- int64_t shapeSize = 1;36+{
37- for (auto i : shape) {37+ int64_t shapeSize = 1;
38- shapeSize *= i;38+ for (auto i : shape) {
39- }39+ shapeSize *= i;
40- return shapeSize;40+ }
41+ return shapeSize;
41}42}
42 43 
43-void PrintOutResult(std::vector<int64_t> &shape, void** deviceAddr) {44+void PrintOutResult(std::vector<int64_t>& shape, void** deviceAddr)
44- auto size = GetShapeSize(shape);45+{
45- // 输入类型为float1646+ auto size = GetShapeSize(shape);
46- std::vector<uint16_t> resultData(size, 0);47+ // 输入类型为float16
47- auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]),48+ std::vector<uint16_t> resultData(size, 0);
48- *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);49+ auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr,
49- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return);50+ size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
50- for (int64_t i = 0; i < size; i++) {51+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return);
51- LOG_PRINT("mean result[%ld] is: %u\n", i, resultData[i]);52+ for (int64_t i = 0; i < size; i++) {
52- }53+ LOG_PRINT("mean result[%ld] is: %u\n", i, resultData[i]);
54+ }
53}55}
54 56 
55-int Init(int32_t deviceId, aclrtStream* stream) {57+int Init(int32_t deviceId, aclrtStream* stream)
56- // 固定写法，AscendCL初始化58+{
57- auto ret = aclInit(nullptr);59+ // 固定写法，AscendCL初始化
58- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);60+ auto ret = aclInit(nullptr);
59- ret = aclrtSetDevice(deviceId);61+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
60- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);62+ ret = aclrtSetDevice(deviceId);
61- ret = aclrtCreateStream(stream);63+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
62- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);64+ ret = aclrtCreateStream(stream);
63- return 0;65+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
66+ return 0;
64}67}
65 68 
66template <typename T>69template <typename T>
67int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,70int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
68- aclDataType dataType, aclTensor** tensor) {71+ aclDataType dataType, aclTensor** tensor)
69- auto size = GetShapeSize(shape) * sizeof(T);72+{
70- // 调用aclrtMalloc申请device侧内存73+ auto size = GetShapeSize(shape) * sizeof(T);
71- auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);74+ // 调用aclrtMalloc申请device侧内存
72- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);75+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
73- // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上76+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
74- ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);77+ // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
75- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);78+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
79+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
76 80 
77- // 计算连续tensor的strides81+ // 计算连续tensor的strides
78- std::vector<int64_t> strides(shape.size(), 1);82+ std::vector<int64_t> strides(shape.size(), 1);
79- for (int64_t i = shape.size() - 2; i >= 0; i--) {83+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
80- strides[i] = shape[i + 1] * strides[i + 1];84+ strides[i] = shape[i + 1] * strides[i + 1];
81- }85+ }
82 86 
83- // 调用aclCreateTensor接口创建aclTensor87+ // 调用aclCreateTensor接口创建aclTensor
84- *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,88+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
85- shape.data(), shape.size(), *deviceAddr);89+ shape.data(), shape.size(), *deviceAddr);
86- return 0;90+ return 0;
87}91}
88 92 
89-int main() {93+int main()
90- // 1. （固定写法）device/stream初始化，参考AscendCL对外接口列表94+{
91- // 根据自己的实际device填写deviceId95+ // 1. （固定写法）device/stream初始化，参考AscendCL对外接口列表
92- int32_t deviceId = 0;96+ // 根据自己的实际device填写deviceId
93- aclrtStream stream;97+ int32_t deviceId = 0;
94- auto ret = Init(deviceId, &stream);98+ aclrtStream stream;
95- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);99+ auto ret = Init(deviceId, &stream);
100+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
96 101 
97- // 2. 构造输入与输出，需要根据API的接口自定义构造102+ // 2. 构造输入与输出，需要根据API的接口自定义构造
98- int64_t B = 1;103+ int64_t B = 1;
99- int64_t N1 = 1;104+ int64_t N1 = 1;
100- int64_t N2 = 1;105+ int64_t N2 = 1;
101- int64_t S1 = 256;106+ int64_t S1 = 256;
102- int64_t S2 = 256;107+ int64_t S2 = 256;
103- int64_t D = 192;108+ int64_t D = 192;
104 109 
105- int64_t H1 = N1 * D;110+ int64_t H1 = N1 * D;
106- int64_t H2 = N2 * D;111+ int64_t H2 = N2 * D;
107 112 
108- int64_t q_size = S1 * B * H1;113+ int64_t q_size = S1 * B * H1;
109- int64_t kv_size = S2 * B * H2;114+ int64_t kv_size = S2 * B * H2;
110- int64_t atten_mask_size = S1 * S2;115+ int64_t atten_mask_size = S1 * S2;
111- int64_t softmax_size = B * N1 * S1 * 8;116+ int64_t softmax_size = B * N1 * S1 * 8;
112 117 
113- std::vector<int64_t> qShape = {S1, B, H1};118+ std::vector<int64_t> qShape = {S1, B, H1};
114- std::vector<int64_t> kShape = {S2, B, H2};119+ std::vector<int64_t> kShape = {S2, B, H2};
115- std::vector<int64_t> vShape = {S2, B, H2};120+ std::vector<int64_t> vShape = {S2, B, H2};
116- std::vector<int64_t> dxShape = {S1, B, H1};121+ std::vector<int64_t> dxShape = {S1, B, H1};
117- std::vector<int64_t> attenmaskShape = {S1, S2};122+ std::vector<int64_t> attenmaskShape = {S1, S2};
118- std::vector<int64_t> softmaxMaxShape = {B, N1, S1, 8};123+ std::vector<int64_t> softmaxMaxShape = {B, N1, S1, 8};
119- std::vector<int64_t> softmaxSumShape = {B, N1, S1, 8};124+ std::vector<int64_t> softmaxSumShape = {B, N1, S1, 8};
120- std::vector<int64_t> attentionInShape = {S1, B, H1};125+ std::vector<int64_t> attentionInShape = {S1, B, H1};
121- std::vector<int64_t> dqShape = {S1, B, H1};126+ std::vector<int64_t> dqShape = {S1, B, H1};
122- std::vector<int64_t> dkShape = {S2, B, H2};127+ std::vector<int64_t> dkShape = {S2, B, H2};
123- std::vector<int64_t> dvShape = {S2, B, H2};128+ std::vector<int64_t> dvShape = {S2, B, H2};
124 129 
125- void* qDeviceAddr = nullptr;130+ void* qDeviceAddr = nullptr;
126- void* kDeviceAddr = nullptr;131+ void* kDeviceAddr = nullptr;
127- void* vDeviceAddr = nullptr;132+ void* vDeviceAddr = nullptr;
128- void* dxDeviceAddr = nullptr;133+ void* dxDeviceAddr = nullptr;
129- void* attenmaskDeviceAddr = nullptr;134+ void* attenmaskDeviceAddr = nullptr;
130- void* softmaxMaxDeviceAddr = nullptr;135+ void* softmaxMaxDeviceAddr = nullptr;
131- void* softmaxSumDeviceAddr = nullptr;136+ void* softmaxSumDeviceAddr = nullptr;
132- void* attentionInDeviceAddr = nullptr;137+ void* attentionInDeviceAddr = nullptr;
133- void* dqDeviceAddr = nullptr;138+ void* dqDeviceAddr = nullptr;
134- void* dkDeviceAddr = nullptr;139+ void* dkDeviceAddr = nullptr;
135- void* dvDeviceAddr = nullptr;140+ void* dvDeviceAddr = nullptr;
136 141 
137- aclTensor* q = nullptr;142+ aclTensor* q = nullptr;
138- aclTensor* k = nullptr;143+ aclTensor* k = nullptr;
139- aclTensor* v = nullptr;144+ aclTensor* v = nullptr;
140- aclTensor* dx = nullptr;145+ aclTensor* dx = nullptr;
141- aclTensor* pse = nullptr;146+ aclTensor* pse = nullptr;
142- aclTensor* dropMask = nullptr;147+ aclTensor* dropMask = nullptr;
143- aclTensor* padding = nullptr;148+ aclTensor* padding = nullptr;
144- aclTensor* attenmask = nullptr;149+ aclTensor* attenmask = nullptr;
145- aclTensor* queryRope = nullptr;150+ aclTensor* queryRope = nullptr;
146- aclTensor* keyRope = nullptr;151+ aclTensor* keyRope = nullptr;
147- aclTensor* dScaleQ = nullptr;152+ aclTensor* dScaleQ = nullptr;
148- aclTensor* dScaleK = nullptr;153+ aclTensor* dScaleK = nullptr;
149- aclTensor* dScaleV = nullptr;154+ aclTensor* dScaleV = nullptr;
150- aclTensor* dScaleDy = nullptr;155+ aclTensor* dScaleDy = nullptr;
151- aclTensor* dScaleO = nullptr;156+ aclTensor* dScaleO = nullptr;
152- aclTensor* softmaxMax = nullptr;157+ aclTensor* softmaxMax = nullptr;
153- aclTensor* softmaxSum = nullptr;158+ aclTensor* softmaxSum = nullptr;
154- aclTensor* softmaxIn = nullptr;159+ aclTensor* softmaxIn = nullptr;
155- aclTensor* attentionIn = nullptr;160+ aclTensor* attentionIn = nullptr;
156- aclTensor* dq = nullptr;161+ aclTensor* dq = nullptr;
157- aclTensor* dk = nullptr;162+ aclTensor* dk = nullptr;
158- aclTensor* dv = nullptr;163+ aclTensor* dv = nullptr;
159- aclTensor* dpse = nullptr;164+ aclTensor* dpse = nullptr;
160- aclTensor* dqRope = nullptr;165+ aclTensor* dqRope = nullptr;
161- aclTensor* dkRope = nullptr;166+ aclTensor* dkRope = nullptr;
162-
163- std::vector<uint16_t> qHostData(q_size, 0x3C00);
164- std::vector<uint16_t> kHostData(kv_size, 0x3800);
165- std::vector<uint16_t> vHostData(kv_size, 0x3800);
166- std::vector<uint16_t> dxHostData(q_size, 0x3C00);
167- std::vector<uint8_t> attenmaskHostData(atten_mask_size, 0);
168- std::vector<float> softmaxMaxHostData(softmax_size, 3.0);
169- std::vector<float> softmaxSumHostData(softmax_size, 3.0);
170- std::vector<uint16_t> attentionInHostData(q_size, 0x3C00);
171- std::vector<uint16_t> dqHostData(q_size, 0);
172- std::vector<uint16_t> dkHostData(kv_size, 0);
173- std::vector<uint16_t> dvHostData(kv_size, 0);
174 167 
175- ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_FLOAT16, &q);168+ std::vector<uint16_t> qHostData(q_size, 0x3C00);
176- CHECK_RET(ret == ACL_SUCCESS, return ret);169+ std::vector<uint16_t> kHostData(kv_size, 0x3800);
177- ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_FLOAT16, &k);170+ std::vector<uint16_t> vHostData(kv_size, 0x3800);
178- CHECK_RET(ret == ACL_SUCCESS, return ret);171+ std::vector<uint16_t> dxHostData(q_size, 0x3C00);
179- ret = CreateAclTensor(vHostData, vShape, &vDeviceAddr, aclDataType::ACL_FLOAT16, &v);172+ std::vector<uint8_t> attenmaskHostData(atten_mask_size, 0);
180- CHECK_RET(ret == ACL_SUCCESS, return ret);173+ std::vector<float> softmaxMaxHostData(softmax_size, 3.0);
181- ret = CreateAclTensor(dxHostData, dxShape, &dxDeviceAddr, aclDataType::ACL_FLOAT16, &dx);174+ std::vector<float> softmaxSumHostData(softmax_size, 3.0);
182- CHECK_RET(ret == ACL_SUCCESS, return ret);175+ std::vector<uint16_t> attentionInHostData(q_size, 0x3C00);
183- ret = CreateAclTensor(attenmaskHostData, attenmaskShape, &attenmaskDeviceAddr, aclDataType::ACL_UINT8, &attenmask);176+ std::vector<uint16_t> dqHostData(q_size, 0);
184- CHECK_RET(ret == ACL_SUCCESS, return ret);177+ std::vector<uint16_t> dkHostData(kv_size, 0);
185- ret = CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, &softmaxMaxDeviceAddr, aclDataType::ACL_FLOAT, &softmaxMax);178+ std::vector<uint16_t> dvHostData(kv_size, 0);
186- CHECK_RET(ret == ACL_SUCCESS, return ret);
187- ret = CreateAclTensor(softmaxSumHostData, softmaxSumShape, &softmaxSumDeviceAddr, aclDataType::ACL_FLOAT, &softmaxSum);
188- CHECK_RET(ret == ACL_SUCCESS, return ret);
189- ret = CreateAclTensor(attentionInHostData, attentionInShape, &attentionInDeviceAddr, aclDataType::ACL_FLOAT16, &attentionIn);
190- CHECK_RET(ret == ACL_SUCCESS, return ret);
191- ret = CreateAclTensor(dqHostData, dqShape, &dqDeviceAddr, aclDataType::ACL_FLOAT16, &dq);
192- CHECK_RET(ret == ACL_SUCCESS, return ret);
193- ret = CreateAclTensor(dkHostData, dkShape, &dkDeviceAddr, aclDataType::ACL_FLOAT16, &dk);
194- CHECK_RET(ret == ACL_SUCCESS, return ret);
195- ret = CreateAclTensor(dvHostData, dvShape, &dvDeviceAddr, aclDataType::ACL_FLOAT16, &dv);
196- CHECK_RET(ret == ACL_SUCCESS, return ret);
197 179 
198- double scaleValue = 1.0/sqrt(128);180+ ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_FLOAT16, &q);
199- double keepProb = 1.0;181+ CHECK_RET(ret == ACL_SUCCESS, return ret);
200- int64_t preTokens = 65536;182+ ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_FLOAT16, &k);
201- int64_t nextTokens = 65536;183+ CHECK_RET(ret == ACL_SUCCESS, return ret);
202- int64_t headNum = 1;184+ ret = CreateAclTensor(vHostData, vShape, &vDeviceAddr, aclDataType::ACL_FLOAT16, &v);
203- int64_t innerPrecise = 0;185+ CHECK_RET(ret == ACL_SUCCESS, return ret);
204- int64_t sparseMode = 0;186+ ret = CreateAclTensor(dxHostData, dxShape, &dxDeviceAddr, aclDataType::ACL_FLOAT16, &dx);
205- int64_t pseType = 1;187+ CHECK_RET(ret == ACL_SUCCESS, return ret);
206- int64_t outDtype = 1;188+ ret = CreateAclTensor(attenmaskHostData, attenmaskShape, &attenmaskDeviceAddr, aclDataType::ACL_UINT8, &attenmask);
207- int64_t seed = 0;189+ CHECK_RET(ret == ACL_SUCCESS, return ret);
208- int64_t offset = 0;190+ ret = CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, &softmaxMaxDeviceAddr, aclDataType::ACL_FLOAT,
209- char inputlayOut[5] = {'S', 'B', 'H', 0};191+ &softmaxMax);
192+ CHECK_RET(ret == ACL_SUCCESS, return ret);
193+ ret = CreateAclTensor(softmaxSumHostData, softmaxSumShape, &softmaxSumDeviceAddr, aclDataType::ACL_FLOAT,
194+ &softmaxSum);
195+ CHECK_RET(ret == ACL_SUCCESS, return ret);
196+ ret = CreateAclTensor(attentionInHostData, attentionInShape, &attentionInDeviceAddr, aclDataType::ACL_FLOAT16,
197+ &attentionIn);
198+ CHECK_RET(ret == ACL_SUCCESS, return ret);
199+ ret = CreateAclTensor(dqHostData, dqShape, &dqDeviceAddr, aclDataType::ACL_FLOAT16, &dq);
200+ CHECK_RET(ret == ACL_SUCCESS, return ret);
201+ ret = CreateAclTensor(dkHostData, dkShape, &dkDeviceAddr, aclDataType::ACL_FLOAT16, &dk);
202+ CHECK_RET(ret == ACL_SUCCESS, return ret);
203+ ret = CreateAclTensor(dvHostData, dvShape, &dvDeviceAddr, aclDataType::ACL_FLOAT16, &dv);
204+ CHECK_RET(ret == ACL_SUCCESS, return ret);
210 205 
211- // 3. 调用CANN算子库API，需要修改为具体的Api名称206+ double scaleValue = 1.0 / sqrt(128);
212- uint64_t workspaceSize = 0;207+ double keepProb = 1.0;
213- aclOpExecutor* executor;208+ int64_t preTokens = 65536;
209+ int64_t nextTokens = 65536;
210+ int64_t headNum = 1;
211+ int64_t innerPrecise = 0;
212+ int64_t sparseMode = 0;
213+ int64_t pseType = 1;
214+ int64_t outDtype = 1;
215+ int64_t seed = 0;
216+ int64_t offset = 0;
217+ char inputlayOut[5] = {'S', 'B', 'H', 0};
214 218 
215- // 调用aclnnFlashAttentionScoreGradV4第一段接口219+ // 3. 调用CANN算子库API，需要修改为具体的Api名称
216- ret = aclnnFlashAttentionScoreGradV4GetWorkspaceSize(q, k, v, dx, pse, dropMask, padding,220+ uint64_t workspaceSize = 0;
217- attenmask, softmaxMax, softmaxSum, softmaxIn, attentionIn, nullptr, queryRope, keyRope, dScaleQ, dScaleK, dScaleV, 221+ aclOpExecutor* executor;
218- dScaleDy, dScaleO, nullptr, nullptr, nullptr, nullptr, nullptr, scaleValue, keepProb,
219- preTokens, nextTokens, headNum, inputlayOut, nullptr, innerPrecise, sparseMode,outDtype, pseType, seed, offset,
220- dq,dk,dv,dqRope,dkRope,dpse, nullptr, &workspaceSize, &executor);
221- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreGradV4GetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
222 222 
223- // 根据第一段接口计算出的workspaceSize申请device内存223+ // 调用aclnnFlashAttentionScoreGradV4第一段接口
224- void* workspaceAddr = nullptr;224+ ret = aclnnFlashAttentionScoreGradV4GetWorkspaceSize(
225- if (workspaceSize > 0) {225+ q, k, v, dx, pse, dropMask, padding, attenmask, softmaxMax, softmaxSum, softmaxIn, attentionIn, nullptr,
226- ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);226+ queryRope, keyRope, dScaleQ, dScaleK, dScaleV, dScaleDy, dScaleO, nullptr, nullptr, nullptr, nullptr, nullptr,
227- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);227+ scaleValue, keepProb, preTokens, nextTokens, headNum, inputlayOut, nullptr, innerPrecise, sparseMode, outDtype,
228- }228+ pseType, seed, offset, dq, dk, dv, dqRope, dkRope, dpse, nullptr, &workspaceSize, &executor);
229+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreGradV4GetWorkspaceSize failed. ERROR: %d\n", ret);
230+ return ret);
229 231 
230- // 调用aclnnFlashAttentionScoreGradV4第二段接口232+ // 根据第一段接口计算出的workspaceSize申请device内存
231- ret = aclnnFlashAttentionScoreGradV4(workspaceAddr, workspaceSize, executor, stream);233+ void* workspaceAddr = nullptr;
232- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreGradV4 failed. ERROR: %d\n", ret); return ret);234+ if (workspaceSize > 0) {
235+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
236+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
237+ }
233 238 
234- // 4. （固定写法）同步等待任务执行结束239+ // 调用aclnnFlashAttentionScoreGradV4第二段接口
235- ret = aclrtSynchronizeStream(stream);240+ ret = aclnnFlashAttentionScoreGradV4(workspaceAddr, workspaceSize, executor, stream);
236- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);241+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreGradV4 failed. ERROR: %d\n", ret); return ret);
237 242 
238- // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改243+ // 4. （固定写法）同步等待任务执行结束
239- // PrintOutResult(dqShape, &dqDeviceAddr);244+ ret = aclrtSynchronizeStream(stream);
240- // PrintOutResult(dkShape, &dkDeviceAddr);245+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
241- // PrintOutResult(dvShape, &dvDeviceAddr);
242 246 
243- // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改247+ // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
244- aclDestroyTensor(q);248+ // PrintOutResult(dqShape, &dqDeviceAddr);
245- aclDestroyTensor(k);249+ // PrintOutResult(dkShape, &dkDeviceAddr);
246- aclDestroyTensor(v);250+ // PrintOutResult(dvShape, &dvDeviceAddr);
247- aclDestroyTensor(dx);
248- aclDestroyTensor(attenmask);
249- aclDestroyTensor(softmaxMax);
250- aclDestroyTensor(softmaxSum);
251- aclDestroyTensor(attentionIn);
252- aclDestroyTensor(dq);
253- aclDestroyTensor(dk);
254 251 
255- // 7. 释放device资源252+ // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
256- aclrtFree(qDeviceAddr);253+ aclDestroyTensor(q);
257- aclrtFree(kDeviceAddr);254+ aclDestroyTensor(k);
258- aclrtFree(vDeviceAddr);255+ aclDestroyTensor(v);
259- aclrtFree(dxDeviceAddr);256+ aclDestroyTensor(dx);
260- aclrtFree(attenmaskDeviceAddr);257+ aclDestroyTensor(attenmask);
261- aclrtFree(softmaxMaxDeviceAddr);258+ aclDestroyTensor(softmaxMax);
262- aclrtFree(softmaxSumDeviceAddr);259+ aclDestroyTensor(softmaxSum);
263- aclrtFree(attentionInDeviceAddr);260+ aclDestroyTensor(attentionIn);
264- aclrtFree(dqDeviceAddr);261+ aclDestroyTensor(dq);
265- aclrtFree(dkDeviceAddr);262+ aclDestroyTensor(dk);
266- aclrtFree(dvDeviceAddr);
267- if (workspaceSize > 0) {
268- aclrtFree(workspaceAddr);
269- }
270- aclrtDestroyStream(stream);
271- aclrtResetDevice(deviceId);
272- aclFinalize();
273 263 
274- return 0;264+ // 7. 释放device资源
265+ aclrtFree(qDeviceAddr);
266+ aclrtFree(kDeviceAddr);
267+ aclrtFree(vDeviceAddr);
268+ aclrtFree(dxDeviceAddr);
269+ aclrtFree(attenmaskDeviceAddr);
270+ aclrtFree(softmaxMaxDeviceAddr);
271+ aclrtFree(softmaxSumDeviceAddr);
272+ aclrtFree(attentionInDeviceAddr);
273+ aclrtFree(dqDeviceAddr);
274+ aclrtFree(dkDeviceAddr);
275+ aclrtFree(dvDeviceAddr);
276+ if (workspaceSize > 0) {
277+ aclrtFree(workspaceAddr);
278+ }
279+ aclrtDestroyStream(stream);
280+ aclrtResetDevice(deviceId);
281+ aclFinalize();
282+ 
283+ return 0;
275}284}
Mattention/flash_attention_score_grad/examples/test_aclnn_flash_attention_score_grad_v4_fp32.cpp+246-237
@@ -1,10 +1,10 @@
1/**1/**
2- * This program is free software, you can redistribute it and/or modify.
3 * Copyright (c) 2025 Huawei Technologies Co., Ltd.2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4- * This file is a part of the CANN Open Software.3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5- * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").4+ * CANN Open Software License Agreement Version 2.0 (the "License").
6 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
7- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.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.8 * See LICENSE in the root of the software repository for the full text of the License.
9 */9 */
10 10 
Mattention/flash_attention_score_grad/examples/test_aclnn_flash_attention_score_grad_v4_sprase_mode.cpp+233-224
@@ -1,10 +1,10 @@
1/**1/**
2- * This program is free software, you can redistribute it and/or modify.
3 * Copyright (c) 2025 Huawei Technologies Co., Ltd.2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4- * This file is a part of the CANN Open Software.3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5- * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").4+ * CANN Open Software License Agreement Version 2.0 (the "License").
6 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
7- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.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.8 * See LICENSE in the root of the software repository for the full text of the License.
9 */9 */
10 10 
@@ -21,255 +21,264 @@
21#include "aclnnop/aclnn_flash_attention_score_grad.h"21#include "aclnnop/aclnn_flash_attention_score_grad.h"
22 22 
23#define CHECK_RET(cond, return_expr) \23#define CHECK_RET(cond, return_expr) \
24- do { \24+ do { \
25- if (!(cond)) { \25+ if (!(cond)) { \
26- return_expr; \26+ return_expr; \
27- } \27+ } \
28- } while (0)28+ } while (0)
29 29 
30-#define LOG_PRINT(message, ...) \30+#define LOG_PRINT(message, ...) \
31- do { \31+ do { \
32- printf(message, ##__VA_ARGS__); \32+ printf(message, ##__VA_ARGS__); \
33- } while (0)33+ } while (0)
34 34 
35-int64_t GetShapeSize(const std::vector<int64_t>& shape) {35+int64_t GetShapeSize(const std::vector<int64_t>& shape)
36- int64_t shapeSize = 1;36+{
37- for (auto i : shape) {37+ int64_t shapeSize = 1;
38- shapeSize *= i;38+ for (auto i : shape) {
39- }39+ shapeSize *= i;
40- return shapeSize;40+ }
41+ return shapeSize;
41}42}
42 43 
43-void PrintOutResult(std::vector<int64_t> &shape, void** deviceAddr) {44+void PrintOutResult(std::vector<int64_t>& shape, void** deviceAddr)
44- auto size = GetShapeSize(shape);45+{
45- // 输入类型为float1646+ auto size = GetShapeSize(shape);
46- std::vector<uint16_t> resultData(size, 0);47+ // 输入类型为float16
47- auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]),48+ std::vector<uint16_t> resultData(size, 0);
48- *deviceAddr, size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);49+ auto ret = aclrtMemcpy(resultData.data(), resultData.size() * sizeof(resultData[0]), *deviceAddr,
49- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return);50+ size * sizeof(resultData[0]), ACL_MEMCPY_DEVICE_TO_HOST);
50- for (int64_t i = 0; i < size; i++) {51+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("copy result from device to host failed. ERROR: %d\n", ret); return);
51- LOG_PRINT("mean result[%ld] is: %u\n", i, resultData[i]);52+ for (int64_t i = 0; i < size; i++) {
52- }53+ LOG_PRINT("mean result[%ld] is: %u\n", i, resultData[i]);
54+ }
53}55}
54 56 
55-int Init(int32_t deviceId, aclrtStream* stream) {57+int Init(int32_t deviceId, aclrtStream* stream)
56- // 固定写法，AscendCL初始化58+{
57- auto ret = aclInit(nullptr);59+ // 固定写法，AscendCL初始化
58- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);60+ auto ret = aclInit(nullptr);
59- ret = aclrtSetDevice(deviceId);61+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclInit failed. ERROR: %d\n", ret); return ret);
60- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);62+ ret = aclrtSetDevice(deviceId);
61- ret = aclrtCreateStream(stream);63+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSetDevice failed. ERROR: %d\n", ret); return ret);
62- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);64+ ret = aclrtCreateStream(stream);
63- return 0;65+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtCreateStream failed. ERROR: %d\n", ret); return ret);
66+ return 0;
64}67}
65 68 
66template <typename T>69template <typename T>
67int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,70int CreateAclTensor(const std::vector<T>& hostData, const std::vector<int64_t>& shape, void** deviceAddr,
68- aclDataType dataType, aclTensor** tensor) {71+ aclDataType dataType, aclTensor** tensor)
69- auto size = GetShapeSize(shape) * sizeof(T);72+{
70- // 调用aclrtMalloc申请device侧内存73+ auto size = GetShapeSize(shape) * sizeof(T);
71- auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);74+ // 调用aclrtMalloc申请device侧内存
72- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);75+ auto ret = aclrtMalloc(deviceAddr, size, ACL_MEM_MALLOC_HUGE_FIRST);
73- // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上76+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMalloc failed. ERROR: %d\n", ret); return ret);
74- ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);77+ // 调用aclrtMemcpy将host侧数据拷贝到device侧内存上
75- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);78+ ret = aclrtMemcpy(*deviceAddr, size, hostData.data(), size, ACL_MEMCPY_HOST_TO_DEVICE);
79+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtMemcpy failed. ERROR: %d\n", ret); return ret);
76 80 
77- // 计算连续tensor的strides81+ // 计算连续tensor的strides
78- std::vector<int64_t> strides(shape.size(), 1);82+ std::vector<int64_t> strides(shape.size(), 1);
79- for (int64_t i = shape.size() - 2; i >= 0; i--) {83+ for (int64_t i = shape.size() - 2; i >= 0; i--) {
80- strides[i] = shape[i + 1] * strides[i + 1];84+ strides[i] = shape[i + 1] * strides[i + 1];
81- }85+ }
82 86 
83- // 调用aclCreateTensor接口创建aclTensor87+ // 调用aclCreateTensor接口创建aclTensor
84- *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,88+ *tensor = aclCreateTensor(shape.data(), shape.size(), dataType, strides.data(), 0, aclFormat::ACL_FORMAT_ND,
85- shape.data(), shape.size(), *deviceAddr);89+ shape.data(), shape.size(), *deviceAddr);
86- return 0;90+ return 0;
87}91}
88 92 
89-int main() {93+int main()
90- // 1. （固定写法）device/stream初始化，参考AscendCL对外接口列表94+{
91- // 根据自己的实际device填写deviceId95+ // 1. （固定写法）device/stream初始化，参考AscendCL对外接口列表
92- int32_t deviceId = 0;96+ // 根据自己的实际device填写deviceId
93- aclrtStream stream;97+ int32_t deviceId = 0;
94- auto ret = Init(deviceId, &stream);98+ aclrtStream stream;
95- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);99+ auto ret = Init(deviceId, &stream);
100+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("Init acl failed. ERROR: %d\n", ret); return ret);
96 101 
97- // 2. 构造输入与输出，需要根据API的接口自定义构造102+ // 2. 构造输入与输出，需要根据API的接口自定义构造
98- int64_t B = 1;103+ int64_t B = 1;
99- int64_t N1 = 1;104+ int64_t N1 = 1;
100- int64_t N2 = 1;105+ int64_t N2 = 1;
101- int64_t S1 = 256;106+ int64_t S1 = 256;
102- int64_t S2 = 256;107+ int64_t S2 = 256;
103- int64_t D = 192;108+ int64_t D = 192;
104 109 
105- int64_t H1 = N1 * D;110+ int64_t H1 = N1 * D;
106- int64_t H2 = N2 * D;111+ int64_t H2 = N2 * D;
107 112 
108- int64_t q_size = S1 * B * H1;113+ int64_t q_size = S1 * B * H1;
109- int64_t kv_size = S2 * B * H2;114+ int64_t kv_size = S2 * B * H2;
110- int64_t atten_mask_size = 2048 * 2048;115+ int64_t atten_mask_size = 2048 * 2048;
111- int64_t softmax_size = B * N1 * S1 * 8;116+ int64_t softmax_size = B * N1 * S1 * 8;
112 117 
113- std::vector<int64_t> qShape = {S1, B, H1};118+ std::vector<int64_t> qShape = {S1, B, H1};
114- std::vector<int64_t> kShape = {S2, B, H2};119+ std::vector<int64_t> kShape = {S2, B, H2};
115- std::vector<int64_t> vShape = {S2, B, H2};120+ std::vector<int64_t> vShape = {S2, B, H2};
116- std::vector<int64_t> dxShape = {S1, B, H1};121+ std::vector<int64_t> dxShape = {S1, B, H1};
117- std::vector<int64_t> attenmaskShape = {2048, 2048};122+ std::vector<int64_t> attenmaskShape = {2048, 2048};
118- std::vector<int64_t> softmaxMaxShape = {B, N1, S1, 8};123+ std::vector<int64_t> softmaxMaxShape = {B, N1, S1, 8};
119- std::vector<int64_t> softmaxSumShape = {B, N1, S1, 8};124+ std::vector<int64_t> softmaxSumShape = {B, N1, S1, 8};
120- std::vector<int64_t> attentionInShape = {S1, B, H1};125+ std::vector<int64_t> attentionInShape = {S1, B, H1};
121- std::vector<int64_t> dqShape = {S1, B, H1};126+ std::vector<int64_t> dqShape = {S1, B, H1};
122- std::vector<int64_t> dkShape = {S2, B, H2};127+ std::vector<int64_t> dkShape = {S2, B, H2};
123- std::vector<int64_t> dvShape = {S2, B, H2};128+ std::vector<int64_t> dvShape = {S2, B, H2};
124 129 
125- void* qDeviceAddr = nullptr;130+ void* qDeviceAddr = nullptr;
126- void* kDeviceAddr = nullptr;131+ void* kDeviceAddr = nullptr;
127- void* vDeviceAddr = nullptr;132+ void* vDeviceAddr = nullptr;
128- void* dxDeviceAddr = nullptr;133+ void* dxDeviceAddr = nullptr;
129- void* attenmaskDeviceAddr = nullptr;134+ void* attenmaskDeviceAddr = nullptr;
130- void* softmaxMaxDeviceAddr = nullptr;135+ void* softmaxMaxDeviceAddr = nullptr;
131- void* softmaxSumDeviceAddr = nullptr;136+ void* softmaxSumDeviceAddr = nullptr;
132- void* attentionInDeviceAddr = nullptr;137+ void* attentionInDeviceAddr = nullptr;
133- void* dqDeviceAddr = nullptr;138+ void* dqDeviceAddr = nullptr;
134- void* dkDeviceAddr = nullptr;139+ void* dkDeviceAddr = nullptr;
135- void* dvDeviceAddr = nullptr;140+ void* dvDeviceAddr = nullptr;
136 141 
137- aclTensor* q = nullptr;142+ aclTensor* q = nullptr;
138- aclTensor* k = nullptr;143+ aclTensor* k = nullptr;
139- aclTensor* v = nullptr;144+ aclTensor* v = nullptr;
140- aclTensor* dx = nullptr;145+ aclTensor* dx = nullptr;
141- aclTensor* pse = nullptr;146+ aclTensor* pse = nullptr;
142- aclTensor* dropMask = nullptr;147+ aclTensor* dropMask = nullptr;
143- aclTensor* padding = nullptr;148+ aclTensor* padding = nullptr;
144- aclTensor* attenmask = nullptr;149+ aclTensor* attenmask = nullptr;
145- aclTensor* queryRope = nullptr;150+ aclTensor* queryRope = nullptr;
146- aclTensor* keyRope = nullptr;151+ aclTensor* keyRope = nullptr;
147- aclTensor* dScaleQ = nullptr;152+ aclTensor* dScaleQ = nullptr;
148- aclTensor* dScaleK = nullptr;153+ aclTensor* dScaleK = nullptr;
149- aclTensor* dScaleV = nullptr;154+ aclTensor* dScaleV = nullptr;
150- aclTensor* dScaleDy = nullptr;155+ aclTensor* dScaleDy = nullptr;
151- aclTensor* dScaleO = nullptr;156+ aclTensor* dScaleO = nullptr;
152- aclTensor* softmaxMax = nullptr;157+ aclTensor* softmaxMax = nullptr;
153- aclTensor* softmaxSum = nullptr;158+ aclTensor* softmaxSum = nullptr;
154- aclTensor* softmaxIn = nullptr;159+ aclTensor* softmaxIn = nullptr;
155- aclTensor* attentionIn = nullptr;160+ aclTensor* attentionIn = nullptr;
156- aclTensor* dq = nullptr;161+ aclTensor* dq = nullptr;
157- aclTensor* dk = nullptr;162+ aclTensor* dk = nullptr;
158- aclTensor* dv = nullptr;163+ aclTensor* dv = nullptr;
159- aclTensor* dpse = nullptr;164+ aclTensor* dpse = nullptr;
160- aclTensor* dqRope = nullptr;165+ aclTensor* dqRope = nullptr;
161- aclTensor* dkRope = nullptr;166+ aclTensor* dkRope = nullptr;
162-
163- std::vector<uint16_t> qHostData(q_size, 0x3C00);
164- std::vector<uint16_t> kHostData(kv_size, 0x3800);
165- std::vector<uint16_t> vHostData(kv_size, 0x3800);
166- std::vector<uint16_t> dxHostData(q_size, 0x3C00);
167- std::vector<uint8_t> attenmaskHostData(atten_mask_size, 0);
168- std::vector<float> softmaxMaxHostData(softmax_size, 3.0);
169- std::vector<float> softmaxSumHostData(softmax_size, 3.0);
170- std::vector<uint16_t> attentionInHostData(q_size, 0x3C00);
171- std::vector<uint16_t> dqHostData(q_size, 0);
172- std::vector<uint16_t> dkHostData(kv_size, 0);
173- std::vector<uint16_t> dvHostData(kv_size, 0);
174 167 
175- ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_FLOAT16, &q);168+ std::vector<uint16_t> qHostData(q_size, 0x3C00);
176- CHECK_RET(ret == ACL_SUCCESS, return ret);169+ std::vector<uint16_t> kHostData(kv_size, 0x3800);
177- ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_FLOAT16, &k);170+ std::vector<uint16_t> vHostData(kv_size, 0x3800);
178- CHECK_RET(ret == ACL_SUCCESS, return ret);171+ std::vector<uint16_t> dxHostData(q_size, 0x3C00);
179- ret = CreateAclTensor(vHostData, vShape, &vDeviceAddr, aclDataType::ACL_FLOAT16, &v);172+ std::vector<uint8_t> attenmaskHostData(atten_mask_size, 0);
180- CHECK_RET(ret == ACL_SUCCESS, return ret);173+ std::vector<float> softmaxMaxHostData(softmax_size, 3.0);
181- ret = CreateAclTensor(dxHostData, dxShape, &dxDeviceAddr, aclDataType::ACL_FLOAT16, &dx);174+ std::vector<float> softmaxSumHostData(softmax_size, 3.0);
182- CHECK_RET(ret == ACL_SUCCESS, return ret);175+ std::vector<uint16_t> attentionInHostData(q_size, 0x3C00);
183- ret = CreateAclTensor(attenmaskHostData, attenmaskShape, &attenmaskDeviceAddr, aclDataType::ACL_UINT8, &attenmask);176+ std::vector<uint16_t> dqHostData(q_size, 0);
184- CHECK_RET(ret == ACL_SUCCESS, return ret);177+ std::vector<uint16_t> dkHostData(kv_size, 0);
185- ret = CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, &softmaxMaxDeviceAddr, aclDataType::ACL_FLOAT, &softmaxMax);178+ std::vector<uint16_t> dvHostData(kv_size, 0);
186- CHECK_RET(ret == ACL_SUCCESS, return ret);
187- ret = CreateAclTensor(softmaxSumHostData, softmaxSumShape, &softmaxSumDeviceAddr, aclDataType::ACL_FLOAT, &softmaxSum);
188- CHECK_RET(ret == ACL_SUCCESS, return ret);
189- ret = CreateAclTensor(attentionInHostData, attentionInShape, &attentionInDeviceAddr, aclDataType::ACL_FLOAT16, &attentionIn);
190- CHECK_RET(ret == ACL_SUCCESS, return ret);
191- ret = CreateAclTensor(dqHostData, dqShape, &dqDeviceAddr, aclDataType::ACL_FLOAT16, &dq);
192- CHECK_RET(ret == ACL_SUCCESS, return ret);
193- ret = CreateAclTensor(dkHostData, dkShape, &dkDeviceAddr, aclDataType::ACL_FLOAT16, &dk);
194- CHECK_RET(ret == ACL_SUCCESS, return ret);
195- ret = CreateAclTensor(dvHostData, dvShape, &dvDeviceAddr, aclDataType::ACL_FLOAT16, &dv);
196- CHECK_RET(ret == ACL_SUCCESS, return ret);
197 179 
198- double scaleValue = 1.0/sqrt(128);180+ ret = CreateAclTensor(qHostData, qShape, &qDeviceAddr, aclDataType::ACL_FLOAT16, &q);
199- double keepProb = 1.0;181+ CHECK_RET(ret == ACL_SUCCESS, return ret);
200- int64_t preTokens = 65536;182+ ret = CreateAclTensor(kHostData, kShape, &kDeviceAddr, aclDataType::ACL_FLOAT16, &k);
201- int64_t nextTokens = 0;183+ CHECK_RET(ret == ACL_SUCCESS, return ret);
202- int64_t headNum = 1;184+ ret = CreateAclTensor(vHostData, vShape, &vDeviceAddr, aclDataType::ACL_FLOAT16, &v);
203- int64_t innerPrecise = 0;185+ CHECK_RET(ret == ACL_SUCCESS, return ret);
204- int64_t sparseMode = 2;186+ ret = CreateAclTensor(dxHostData, dxShape, &dxDeviceAddr, aclDataType::ACL_FLOAT16, &dx);
205- int64_t pseType = 1;187+ CHECK_RET(ret == ACL_SUCCESS, return ret);
206- int64_t outDtype = 1;188+ ret = CreateAclTensor(attenmaskHostData, attenmaskShape, &attenmaskDeviceAddr, aclDataType::ACL_UINT8, &attenmask);
207- int64_t seed = 0;189+ CHECK_RET(ret == ACL_SUCCESS, return ret);
208- int64_t offset = 0;190+ ret = CreateAclTensor(softmaxMaxHostData, softmaxMaxShape, &softmaxMaxDeviceAddr, aclDataType::ACL_FLOAT,
209- char inputlayOut[5] = {'S', 'B', 'H', 0};191+ &softmaxMax);
192+ CHECK_RET(ret == ACL_SUCCESS, return ret);
193+ ret = CreateAclTensor(softmaxSumHostData, softmaxSumShape, &softmaxSumDeviceAddr, aclDataType::ACL_FLOAT,
194+ &softmaxSum);
195+ CHECK_RET(ret == ACL_SUCCESS, return ret);
196+ ret = CreateAclTensor(attentionInHostData, attentionInShape, &attentionInDeviceAddr, aclDataType::ACL_FLOAT16,
197+ &attentionIn);
198+ CHECK_RET(ret == ACL_SUCCESS, return ret);
199+ ret = CreateAclTensor(dqHostData, dqShape, &dqDeviceAddr, aclDataType::ACL_FLOAT16, &dq);
200+ CHECK_RET(ret == ACL_SUCCESS, return ret);
201+ ret = CreateAclTensor(dkHostData, dkShape, &dkDeviceAddr, aclDataType::ACL_FLOAT16, &dk);
202+ CHECK_RET(ret == ACL_SUCCESS, return ret);
203+ ret = CreateAclTensor(dvHostData, dvShape, &dvDeviceAddr, aclDataType::ACL_FLOAT16, &dv);
204+ CHECK_RET(ret == ACL_SUCCESS, return ret);
210 205 
211- // 3. 调用CANN算子库API，需要修改为具体的Api名称206+ double scaleValue = 1.0 / sqrt(128);
212- uint64_t workspaceSize = 0;207+ double keepProb = 1.0;
213- aclOpExecutor* executor;208+ int64_t preTokens = 65536;
209+ int64_t nextTokens = 0;
210+ int64_t headNum = 1;
211+ int64_t innerPrecise = 0;
212+ int64_t sparseMode = 2;
213+ int64_t pseType = 1;
214+ int64_t outDtype = 1;
215+ int64_t seed = 0;
216+ int64_t offset = 0;
217+ char inputlayOut[5] = {'S', 'B', 'H', 0};
214 218 
215- // 调用aclnnFlashAttentionScoreGradV4第一段接口219+ // 3. 调用CANN算子库API，需要修改为具体的Api名称
216- ret = aclnnFlashAttentionScoreGradV4GetWorkspaceSize(q, k, v, dx, pse, dropMask, padding,220+ uint64_t workspaceSize = 0;
217- attenmask, softmaxMax, softmaxSum, softmaxIn, attentionIn, nullptr, queryRope, keyRope, dScaleQ, dScaleK, dScaleV, 221+ aclOpExecutor* executor;
218- dScaleDy, dScaleO, nullptr, nullptr, nullptr, nullptr, nullptr, scaleValue, keepProb,
219- preTokens, nextTokens, headNum, inputlayOut, nullptr, innerPrecise, sparseMode,outDtype, pseType, seed, offset,
220- dq,dk,dv,dqRope,dkRope,dpse, nullptr, &workspaceSize, &executor);
221- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreGradV4GetWorkspaceSize failed. ERROR: %d\n", ret); return ret);
222 222 
223- // 根据第一段接口计算出的workspaceSize申请device内存223+ // 调用aclnnFlashAttentionScoreGradV4第一段接口
224- void* workspaceAddr = nullptr;224+ ret = aclnnFlashAttentionScoreGradV4GetWorkspaceSize(
225- if (workspaceSize > 0) {225+ q, k, v, dx, pse, dropMask, padding, attenmask, softmaxMax, softmaxSum, softmaxIn, attentionIn, nullptr,
226- ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);226+ queryRope, keyRope, dScaleQ, dScaleK, dScaleV, dScaleDy, dScaleO, nullptr, nullptr, nullptr, nullptr, nullptr,
227- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);227+ scaleValue, keepProb, preTokens, nextTokens, headNum, inputlayOut, nullptr, innerPrecise, sparseMode, outDtype,
228- }228+ pseType, seed, offset, dq, dk, dv, dqRope, dkRope, dpse, nullptr, &workspaceSize, &executor);
229+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreGradV4GetWorkspaceSize failed. ERROR: %d\n", ret);
230+ return ret);
229 231 
230- // 调用aclnnFlashAttentionScoreGradV4第二段接口232+ // 根据第一段接口计算出的workspaceSize申请device内存
231- ret = aclnnFlashAttentionScoreGradV4(workspaceAddr, workspaceSize, executor, stream);233+ void* workspaceAddr = nullptr;
232- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreGradV4 failed. ERROR: %d\n", ret); return ret);234+ if (workspaceSize > 0) {
235+ ret = aclrtMalloc(&workspaceAddr, workspaceSize, ACL_MEM_MALLOC_HUGE_FIRST);
236+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("allocate workspace failed. ERROR: %d\n", ret); return ret);
237+ }
233 238 
234- // 4. （固定写法）同步等待任务执行结束239+ // 调用aclnnFlashAttentionScoreGradV4第二段接口
235- ret = aclrtSynchronizeStream(stream);240+ ret = aclnnFlashAttentionScoreGradV4(workspaceAddr, workspaceSize, executor, stream);
236- CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);241+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclnnFlashAttentionScoreGradV4 failed. ERROR: %d\n", ret); return ret);
237 242 
238- // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改243+ // 4. （固定写法）同步等待任务执行结束
239- // PrintOutResult(dqShape, &dqDeviceAddr);244+ ret = aclrtSynchronizeStream(stream);
240- // PrintOutResult(dkShape, &dkDeviceAddr);245+ CHECK_RET(ret == ACL_SUCCESS, LOG_PRINT("aclrtSynchronizeStream failed. ERROR: %d\n", ret); return ret);
241- // PrintOutResult(dvShape, &dvDeviceAddr);
242 246 
243- // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改247+ // 5. 获取输出的值，将device侧内存上的结果拷贝至host侧，需要根据具体API的接口定义修改
244- aclDestroyTensor(q);248+ // PrintOutResult(dqShape, &dqDeviceAddr);
245- aclDestroyTensor(k);249+ // PrintOutResult(dkShape, &dkDeviceAddr);
246- aclDestroyTensor(v);250+ // PrintOutResult(dvShape, &dvDeviceAddr);
247- aclDestroyTensor(dx);
248- aclDestroyTensor(attenmask);
249- aclDestroyTensor(softmaxMax);
250- aclDestroyTensor(softmaxSum);
251- aclDestroyTensor(attentionIn);
252- aclDestroyTensor(dq);
253- aclDestroyTensor(dk);
254 251 
255- // 7. 释放device资源252+ // 6. 释放aclTensor和aclScalar，需要根据具体API的接口定义修改
256- aclrtFree(qDeviceAddr);253+ aclDestroyTensor(q);
257- aclrtFree(kDeviceAddr);254+ aclDestroyTensor(k);
258- aclrtFree(vDeviceAddr);255+ aclDestroyTensor(v);
259- aclrtFree(dxDeviceAddr);256+ aclDestroyTensor(dx);
260- aclrtFree(attenmaskDeviceAddr);257+ aclDestroyTensor(attenmask);
261- aclrtFree(softmaxMaxDeviceAddr);258+ aclDestroyTensor(softmaxMax);
262- aclrtFree(softmaxSumDeviceAddr);259+ aclDestroyTensor(softmaxSum);
263- aclrtFree(attentionInDeviceAddr);260+ aclDestroyTensor(attentionIn);
264- aclrtFree(dqDeviceAddr);261+ aclDestroyTensor(dq);
265- aclrtFree(dkDeviceAddr);262+ aclDestroyTensor(dk);
266- aclrtFree(dvDeviceAddr);
267- if (workspaceSize > 0) {
268- aclrtFree(workspaceAddr);
269- }
270- aclrtDestroyStream(stream);
271- aclrtResetDevice(deviceId);
272- aclFinalize();
273 263 
274- return 0;264+ // 7. 释放device资源
265+ aclrtFree(qDeviceAddr);
266+ aclrtFree(kDeviceAddr);
267+ aclrtFree(vDeviceAddr);
268+ aclrtFree(dxDeviceAddr);
269+ aclrtFree(attenmaskDeviceAddr);
270+ aclrtFree(softmaxMaxDeviceAddr);
271+ aclrtFree(softmaxSumDeviceAddr);
272+ aclrtFree(attentionInDeviceAddr);
273+ aclrtFree(dqDeviceAddr);
274+ aclrtFree(dkDeviceAddr);
275+ aclrtFree(dvDeviceAddr);
276+ if (workspaceSize > 0) {
277+ aclrtFree(workspaceAddr);
278+ }
279+ aclrtDestroyStream(stream);
280+ aclrtResetDevice(deviceId);
281+ aclFinalize();
282+ 
283+ return 0;
275}284}
Mattention/flash_attention_score_grad/examples/test_aclnn_flash_attention_score_grad_v4_varlen.cpp+258-248
@@ -1,10 +1,10 @@
1/**1/**
2- * This program is free software, you can redistribute it and/or modify.
3 * Copyright (c) 2025 Huawei Technologies Co., Ltd.2 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
4- * This file is a part of the CANN Open Software.3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
5- * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").4+ * CANN Open Software License Agreement Version 2.0 (the "License").
6 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
7- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.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.8 * See LICENSE in the root of the software repository for the full text of the License.
9 */9 */
10 10 
Mattention/flash_attention_score_grad/tests/ut/op_api/test_aclnn_flash_attention_score_grad_v4.cpp+38-69
@@ -1,10 +1,10 @@
1/**1/**
2- * This program is free software, you can redistribute it and/or modify.2+ * Copyright (c) 2025 Huawei Technologies Co., Ltd.
3- * 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- * This file is a part of the CANN Open Software.4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5- * Licensed under CANN Open Software License Agreement Version 2.0 (the "License").
6 * Please refer to the License for details. You may not use this file except in compliance with the License.5 * Please refer to the License for details. You may not use this file except in compliance with the License.
7- * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.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.8 * See LICENSE in the root of the software repository for the full text of the License.
9 */9 */
10#include <vector>10#include <vector>
@@ -19,8 +19,7 @@
19using namespace std;19using namespace std;
20using namespace op;20using namespace op;
21 21 
22-class flash_attention_score_grad_v4_opapi_ut : public testing::Test22+class flash_attention_score_grad_v4_opapi_ut : public testing::Test {
23-{
24protected:23protected:
25 static void SetUpTestCase()24 static void SetUpTestCase()
26 {25 {
@@ -33,17 +32,13 @@ protected:
33 }32 }
34};33};
35 34 
36-TEST_F(flash_attention_score_grad_v4_opapi_ut, flash_attention_score_grad_aclnn_0) {35+TEST_F(flash_attention_score_grad_v4_opapi_ut, flash_attention_score_grad_aclnn_0)
37- auto tensorQ = TensorDesc({256, 1, 192}, ACL_FLOAT, ACL_FORMAT_ND)36+{
38- .ValueRange(-1, 1);37+ auto tensorQ = TensorDesc({256, 1, 192}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-1, 1);
39- auto tensorK = TensorDesc({256, 1, 192}, ACL_FLOAT, ACL_FORMAT_ND)38+ auto tensorK = TensorDesc({256, 1, 192}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-1, 1);
40- .ValueRange(-1, 1);39+ auto tensorV = TensorDesc({256, 1, 192}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-1, 1);
41- auto tensorV = TensorDesc({256, 1, 192}, ACL_FLOAT, ACL_FORMAT_ND)40+ auto tensorDx = TensorDesc({256, 1, 192}, ACL_FLOAT, ACL_FORMAT_ND).ValueRange(-1, 1);
42- .ValueRange(-1, 1);41+ auto tensorAttenMask = TensorDesc({256, 256}, ACL_UINT8, ACL_FORMAT_ND).Value(vector<uint8_t>{0});
43- auto tensorDx = TensorDesc({256, 1, 192}, ACL_FLOAT, ACL_FORMAT_ND)
44- .ValueRange(-1, 1);
45- auto tensorAttenMask = TensorDesc({256, 256}, ACL_UINT8, ACL_FORMAT_ND)
46- .Value(vector<uint8_t>{0});
47 42 
48 auto tensorSoftmaxMax =43 auto tensorSoftmaxMax =
49 TensorDesc({1, 1, 256, 8}, ACL_FLOAT, ACL_FORMAT_ND).Value(vector<float>(1 * 1 * 256 * 8, 3.0f));44 TensorDesc({1, 1, 256, 8}, ACL_FLOAT, ACL_FORMAT_ND).Value(vector<float>(1 * 1 * 256 * 8, 3.0f));
@@ -74,57 +69,31 @@ TEST_F(flash_attention_score_grad_v4_opapi_ut, flash_attention_score_grad_aclnn_
74 69 
75 auto ut = OP_API_UT(70 auto ut = OP_API_UT(
76 aclnnFlashAttentionScoreGradV4,71 aclnnFlashAttentionScoreGradV4,
77- INPUT(72+ INPUT(tensorQ, tensorK, tensorV, tensorDx,
78- tensorQ,73+ nullptr, // pse
79- tensorK,74+ nullptr, // dropmask
80- tensorV,75+ nullptr, // psepaddingmask
81- tensorDx,76+ tensorAttenMask, tensorSoftmaxMax, tensorSoftmaxSum,
82- nullptr, // pse77+ nullptr, // softmaxin
83- nullptr, // dropmask78+ tensorAttentionIn,
84- nullptr, // psepaddingmask79+ nullptr, // sink
85- tensorAttenMask,80+ nullptr, // queryrope
86- tensorSoftmaxMax,81+ nullptr, // keyrope
87- tensorSoftmaxSum,82+ nullptr, // dscaleq
88- nullptr, // softmaxin83+ nullptr, // dscalek
89- tensorAttentionIn,84+ nullptr, // dscalev
90- nullptr, // sink85+ nullptr, // dscaledy
91- nullptr, // queryrope86+ nullptr, // dscaleo
92- nullptr, // keyrope87+ nullptr, // Prefix
93- nullptr, // dscaleq88+ nullptr, // ActualSeqQLen
94- nullptr, // dscalek89+ nullptr, // ActualSeqKVLen
95- nullptr, // dscalev90+ nullptr, // QStartIdx
96- nullptr, // dscaledy91+ nullptr, // KVStartIdx
97- nullptr, // dscaleo92+ scaleValue, keepProb, preTokens, nextTokens, headNum, inputlayout,
98- nullptr, // Prefix93+ nullptr, // softmaxInLayout
99- nullptr, // ActualSeqQLen94+ innerPrecise, sparseMod, pseType, seed, offset, outDtype),
100- nullptr, // ActualSeqKVLen95+ OUTPUT(tensorDQDesc, tensorDKDesc, tensorDVDesc, tensorDQROPEDesc, tensorDKROPEDesc, tensorDPseDesc, nullptr));
101- nullptr, // QStartIdx
102- nullptr, // KVStartIdx
103- scaleValue,
104- keepProb,
105- preTokens,
106- nextTokens,
107- headNum,
108- inputlayout,
109- nullptr, // softmaxInLayout
110- innerPrecise,
111- sparseMod,
112- pseType,
113- seed,
114- offset,
115- outDtype
116- ),
117- OUTPUT(
118- tensorDQDesc,
119- tensorDKDesc,
120- tensorDVDesc,
121- tensorDQROPEDesc,
122- tensorDKROPEDesc,
123- tensorDPseDesc,
124- nullptr
125- )
126- );
127 uint64_t workspaceSize = 0;96 uint64_t workspaceSize = 0;
128 aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspaceSize);97 aclnnStatus aclRet = ut.TestGetWorkspaceSize(&workspaceSize);
129 EXPECT_EQ(aclRet, ACL_SUCCESS);98 EXPECT_EQ(aclRet, ACL_SUCCESS);
130-}99+}
Mattention/flash_attn/tests/CMakeLists.txt+2-2
@@ -2,7 +2,7 @@
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of3# 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").4# CANN Open Software License Agreement Version 2.0 (the "License").
5-# Please refer to the License for details. You can not use this file except in compliance with 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,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.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.8# See LICENSE in the root of the software repository for the full text of the License.
@@ -15,4 +15,4 @@ foreach(SUB_DIR ${CURRENT_DIRS})
15 if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")15 if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
16 add_subdirectory(${SUB_DIR})16 add_subdirectory(${SUB_DIR})
17 endif()17 endif()
18-endforeach()18+endforeach()
Mattention/flash_attn/tests/ut/CMakeLists.txt+2-2
@@ -2,7 +2,7 @@
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of3# 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").4# CANN Open Software License Agreement Version 2.0 (the "License").
5-# Please refer to the License for details. You can not use this file except in compliance with 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,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.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.8# See LICENSE in the root of the software repository for the full text of the License.
@@ -13,4 +13,4 @@ foreach(SUB_DIR ${CURRENT_DIRS})
13 if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")13 if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
14 add_subdirectory(${SUB_DIR})14 add_subdirectory(${SUB_DIR})
15 endif()15 endif()
16-endforeach()16+endforeach()
Mattention/flash_attn/tests/ut/op_host/CMakeLists.txt+2-3
@@ -2,13 +2,12 @@
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of3# 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").4# CANN Open Software License Agreement Version 2.0 (the "License").
5-# Please refer to the License for details. You can not use this file except in compliance with 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,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.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.8# See LICENSE in the root of the software repository for the full text of the License.
9# -----------------------------------------------------------------------------------------------------------9# -----------------------------------------------------------------------------------------------------------
10 10 
11- 
12if(UT_TEST_ALL OR OP_HOST_UT)11if(UT_TEST_ALL OR OP_HOST_UT)
13 add_modules_ut_sources(UT_NAME ${OP_TILING_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})12 add_modules_ut_sources(UT_NAME ${OP_TILING_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
14 add_modules_ut_sources(UT_NAME ${OP_INFERSHAPE_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})13 add_modules_ut_sources(UT_NAME ${OP_INFERSHAPE_MODULE_NAME} MODE PRIVATE DIR ${CMAKE_CURRENT_SOURCE_DIR})
@@ -19,4 +18,4 @@ foreach(SUB_DIR ${CURRENT_DIRS})
19 if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")18 if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
20 add_subdirectory(${SUB_DIR})19 add_subdirectory(${SUB_DIR})
21 endif()20 endif()
22-endforeach()21+endforeach()
Mattention/generic_block_sparse_attention/CMakeLists.txt+1-0
@@ -5,6 +5,7 @@
5# Please refer to the License for details. You may not use this file except in compliance with 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,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.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.
8# -----------------------------------------------------------------------------------------------------------9# -----------------------------------------------------------------------------------------------------------
9 10 
10file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)11file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
Mattention/generic_block_sparse_attention/op_host/CMakeLists.txt+3-1
@@ -4,8 +4,10 @@
4# CANN Open Software License Agreement Version 2.0 (the "License").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.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,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.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.
8# -----------------------------------------------------------------------------------------------------------9# -----------------------------------------------------------------------------------------------------------
10+ 
9add_op_to_compiled_list()11add_op_to_compiled_list()
10 12 
11if (BUILD_OPEN_PROJECT)13if (BUILD_OPEN_PROJECT)
Mattention/generic_block_sparse_attention_metadata/op_host/CMakeLists.txt+8-0
@@ -1,4 +1,12 @@
1+# -----------------------------------------------------------------------------------------------------------
1# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# -----------------------------------------------------------------------------------------------------------
2 10 
3# Keep the public Generic ACLNN header in the configured install source list.11# Keep the public Generic ACLNN header in the configured install source list.
4set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS12set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS
Mattention/generic_block_sparse_attention_metadata/op_kernel_aicpu/CMakeLists.txt+8-0
@@ -1,4 +1,12 @@
1+# -----------------------------------------------------------------------------------------------------------
1# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# -----------------------------------------------------------------------------------------------------------
2 10 
3if (BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG)11if (BUILD_WITH_INSTALLED_DEPENDENCY_CANN_PKG)
4 if (NOT (UT_TEST_ALL OR OP_KERNEL_AICPU_UT))12 if (NOT (UT_TEST_ALL OR OP_KERNEL_AICPU_UT))
Mattention/lightning_indexer_grad/tests/CMakeLists.txt+1-1
@@ -1,7 +1,7 @@
1# -----------------------------------------------------------------------------------------------------------1# -----------------------------------------------------------------------------------------------------------
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4-# the CANN Open Software License Agreement Version 2.0 (the "License").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.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,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.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
Mattention/lightning_indexer_grad/tests/ut/CMakeLists.txt+1-1
@@ -1,7 +1,7 @@
1# -----------------------------------------------------------------------------------------------------------1# -----------------------------------------------------------------------------------------------------------
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4-# the CANN Open Software License Agreement Version 2.0 (the "License").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.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,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.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
Mattention/lightning_indexer_grad/tests/ut/op_api/CMakeLists.txt+1-1
@@ -1,7 +1,7 @@
1# -----------------------------------------------------------------------------------------------------------1# -----------------------------------------------------------------------------------------------------------
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4-# the CANN Open Software License Agreement Version 2.0 (the "License").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.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,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.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
Mattention/lightning_indexer_grad/tests/ut/op_api/test_aclnn_lightning_indexer_grad.cpp+18-23
@@ -1,7 +1,7 @@
1/**1/**
2 * Copyright (c) 2026 Huawei Technologies Co., Ltd.2 * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of3 * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4- * the CANN Open Software License Agreement Version 2.0 (the "License").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.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,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.7 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
@@ -23,14 +23,14 @@ using namespace op;
23using namespace std;23using namespace std;
24 24 
25namespace {25namespace {
26-void DestroyAclTensor(aclTensor *tensor)26+void DestroyAclTensor(aclTensor* tensor)
27{27{
28 Release(tensor);28 Release(tensor);
29}29}
30 30 
31using AclTensorPtr = unique_ptr<aclTensor, decltype(&DestroyAclTensor)>;31using AclTensorPtr = unique_ptr<aclTensor, decltype(&DestroyAclTensor)>;
32 32 
33-AclTensorPtr MakeTensor(const vector<int64_t> &shape, aclDataType dtype, aclFormat format = ACL_FORMAT_ND)33+AclTensorPtr MakeTensor(const vector<int64_t>& shape, aclDataType dtype, aclFormat format = ACL_FORMAT_ND)
34{34{
35 return AclTensorPtr(TensorDesc(shape, dtype, format).ToAclTypeRawPtr(), DestroyAclTensor);35 return AclTensorPtr(TensorDesc(shape, dtype, format).ToAclTypeRawPtr(), DestroyAclTensor);
36}36}
@@ -63,12 +63,11 @@ TEST_F(LightningIndexerGradOpapiUt, A1_bsnd_fp16_success)
63 63 
64 char inputLayout[] = "BSND";64 char inputLayout[] = "BSND";
65 uint64_t workspaceSize = 0;65 uint64_t workspaceSize = 0;
66- aclOpExecutor *executor = nullptr;66+ aclOpExecutor* executor = nullptr;
67 67 
68 aclnnStatus aclRet = aclnnLightningIndexerGradGetWorkspaceSize(68 aclnnStatus aclRet = aclnnLightningIndexerGradGetWorkspaceSize(
69- query.get(), key.get(), dy.get(), sparseIndices.get(), weights.get(),69+ query.get(), key.get(), dy.get(), sparseIndices.get(), weights.get(), nullptr, nullptr, 64, inputLayout, 3,
70- nullptr, nullptr, 64, inputLayout, 3, 65536, 65536, false,70+ 65536, 65536, false, dqOut.get(), dkOut.get(), dweightsOut.get(), &workspaceSize, &executor);
71- dqOut.get(), dkOut.get(), dweightsOut.get(), &workspaceSize, &executor);
72 71 
73 EXPECT_EQ(aclRet, ACLNN_SUCCESS);72 EXPECT_EQ(aclRet, ACLNN_SUCCESS);
74 EXPECT_NE(executor, nullptr);73 EXPECT_NE(executor, nullptr);
@@ -89,12 +88,11 @@ TEST_F(LightningIndexerGradOpapiUt, A2_tnd_bf16_success)
89 88 
90 char inputLayout[] = "TND";89 char inputLayout[] = "TND";
91 uint64_t workspaceSize = 0;90 uint64_t workspaceSize = 0;
92- aclOpExecutor *executor = nullptr;91+ aclOpExecutor* executor = nullptr;
93 92 
94 aclnnStatus aclRet = aclnnLightningIndexerGradGetWorkspaceSize(93 aclnnStatus aclRet = aclnnLightningIndexerGradGetWorkspaceSize(
95- query.get(), key.get(), dy.get(), sparseIndices.get(), weights.get(),94+ query.get(), key.get(), dy.get(), sparseIndices.get(), weights.get(), actualSeqQ.get(), actualSeqK.get(), 64,
96- actualSeqQ.get(), actualSeqK.get(), 64, inputLayout, 3, 65536, 65536, false,95+ inputLayout, 3, 65536, 65536, false, dqOut.get(), dkOut.get(), dweightsOut.get(), &workspaceSize, &executor);
97- dqOut.get(), dkOut.get(), dweightsOut.get(), &workspaceSize, &executor);
98 96 
99 EXPECT_EQ(aclRet, ACLNN_SUCCESS);97 EXPECT_EQ(aclRet, ACLNN_SUCCESS);
100 EXPECT_NE(executor, nullptr);98 EXPECT_NE(executor, nullptr);
@@ -112,12 +110,11 @@ TEST_F(LightningIndexerGradOpapiUt, E1_null_query_failed)
112 110 
113 char inputLayout[] = "BSND";111 char inputLayout[] = "BSND";
114 uint64_t workspaceSize = 0;112 uint64_t workspaceSize = 0;
115- aclOpExecutor *executor = nullptr;113+ aclOpExecutor* executor = nullptr;
116 114 
117 aclnnStatus aclRet = aclnnLightningIndexerGradGetWorkspaceSize(115 aclnnStatus aclRet = aclnnLightningIndexerGradGetWorkspaceSize(
118- nullptr, key.get(), dy.get(), sparseIndices.get(), weights.get(),116+ nullptr, key.get(), dy.get(), sparseIndices.get(), weights.get(), nullptr, nullptr, 64, inputLayout, 3, 65536,
119- nullptr, nullptr, 64, inputLayout, 3, 65536, 65536, false,117+ 65536, false, dqOut.get(), dkOut.get(), dweightsOut.get(), &workspaceSize, &executor);
120- dqOut.get(), dkOut.get(), dweightsOut.get(), &workspaceSize, &executor);
121 118 
122 EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_NULLPTR);119 EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_NULLPTR);
123 EXPECT_EQ(executor, nullptr);120 EXPECT_EQ(executor, nullptr);
@@ -136,12 +133,11 @@ TEST_F(LightningIndexerGradOpapiUt, E2_dtype_mismatch_failed)
136 133 
137 char inputLayout[] = "BSND";134 char inputLayout[] = "BSND";
138 uint64_t workspaceSize = 0;135 uint64_t workspaceSize = 0;
139- aclOpExecutor *executor = nullptr;136+ aclOpExecutor* executor = nullptr;
140 137 
141 aclnnStatus aclRet = aclnnLightningIndexerGradGetWorkspaceSize(138 aclnnStatus aclRet = aclnnLightningIndexerGradGetWorkspaceSize(
142- query.get(), key.get(), dy.get(), sparseIndices.get(), weights.get(),139+ query.get(), key.get(), dy.get(), sparseIndices.get(), weights.get(), nullptr, nullptr, 64, inputLayout, 3,
143- nullptr, nullptr, 64, inputLayout, 3, 65536, 65536, false,140+ 65536, 65536, false, dqOut.get(), dkOut.get(), dweightsOut.get(), &workspaceSize, &executor);
144- dqOut.get(), dkOut.get(), dweightsOut.get(), &workspaceSize, &executor);
145 141 
146 EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);142 EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_INVALID);
147 EXPECT_EQ(executor, nullptr);143 EXPECT_EQ(executor, nullptr);
@@ -159,12 +155,11 @@ TEST_F(LightningIndexerGradOpapiUt, E3_null_dq_output_failed)
159 155 
160 char inputLayout[] = "BSND";156 char inputLayout[] = "BSND";
161 uint64_t workspaceSize = 0;157 uint64_t workspaceSize = 0;
162- aclOpExecutor *executor = nullptr;158+ aclOpExecutor* executor = nullptr;
163 159 
164 aclnnStatus aclRet = aclnnLightningIndexerGradGetWorkspaceSize(160 aclnnStatus aclRet = aclnnLightningIndexerGradGetWorkspaceSize(
165- query.get(), key.get(), dy.get(), sparseIndices.get(), weights.get(),161+ query.get(), key.get(), dy.get(), sparseIndices.get(), weights.get(), nullptr, nullptr, 64, inputLayout, 3,
166- nullptr, nullptr, 64, inputLayout, 3, 65536, 65536, false,162+ 65536, 65536, false, nullptr, dkOut.get(), dweightsOut.get(), &workspaceSize, &executor);
167- nullptr, dkOut.get(), dweightsOut.get(), &workspaceSize, &executor);
168 163 
169 EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_NULLPTR);164 EXPECT_EQ(aclRet, ACLNN_ERR_PARAM_NULLPTR);
170 EXPECT_EQ(executor, nullptr);165 EXPECT_EQ(executor, nullptr);
Mattention/lightning_indexer_grad/tests/ut/op_host/CMakeLists.txt+1-1
@@ -1,7 +1,7 @@
1# -----------------------------------------------------------------------------------------------------------1# -----------------------------------------------------------------------------------------------------------
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4-# the CANN Open Software License Agreement Version 2.0 (the "License").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.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,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.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
Mattention/lightning_indexer_grad/tests/ut/op_host/arch22/CMakeLists.txt+1-1
@@ -1,7 +1,7 @@
1# -----------------------------------------------------------------------------------------------------------1# -----------------------------------------------------------------------------------------------------------
2# Copyright (c) 2026 Huawei Technologies Co., Ltd.2# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of3# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4-# the CANN Open Software License Agreement Version 2.0 (the "License").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.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,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.7# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
Mmhc/mhc_post_backward/tests/ut/op_host/test_mhc_post_backward_tiling.cpp+78-43