Star1703
3413
Code
Code
Issues469
Pull Requests731
Actions
Discussion
Wiki
Analysis
Star1703
3413
Merged
feat: vector算子适配ascend960pr(DV122)平台并补齐仓库级960pr支持 #14537
biabuAdd on 22 days ago
feat: vector算子适配ascend960pr(DV122)平台并补齐仓库级960pr支持 #14537
Merged
biabuAdd on 22 days ago
Merged frombiabu111/ops-transformer_dv122feat/dv122-ascend960pr-adaptationintocann/ops-transformermaster
Total: 89 files changed+13879-194
M.gitignore+1-0
@@ -55,3 +55,4 @@ codespell_reports/
55 55 
56# OAT 远程 hook 报告目录56# OAT 远程 hook 报告目录
57oat_reports/57oat_reports/
58+pre-commit_reports/
Mattention/attention_update/op_host/CMakeLists.txt+10-0
@@ -42,6 +42,16 @@ if (BUILD_OPEN_PROJECT)
42 -Wno-deprecated-declarations42 -Wno-deprecated-declarations
43 -Werror43 -Werror
44 )44 )
45+ elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend960pr")
46+ add_ops_compile_options(
47+ OP_NAME AttentionUpdate
48+ COMPUTE_UNIT Ascend960PR_969A1
49+ OPTIONS --cce-auto-sync=off
50+ -mllvm
51+ -cce-aicore-dcci-before-kernel-end=false
52+ -Wno-deprecated-declarations
53+ -Werror
54+ )
45 else()55 else()
46 add_ops_compile_options(56 add_ops_compile_options(
47 OP_NAME AttentionUpdate57 OP_NAME AttentionUpdate
Mattention/attention_update/op_host/attention_update_def.cpp+1-0
@@ -60,6 +60,7 @@ public:
60 60 
61 this->AICore().AddConfig("ascend950", config_95);61 this->AICore().AddConfig("ascend950", config_95);
62 this->AICore().AddConfig("ascend960dt", config_95);62 this->AICore().AddConfig("ascend960dt", config_95);
63+ this->AICore().AddConfig("ascend960pr", config_95);
63 this->AICore().AddConfig("ascend350", config_95);64 this->AICore().AddConfig("ascend350", config_95);
64 }65 }
65};66};
Aattention/attention_update/op_host/config/ascend960pr/attention_update_binary.json+194-0
@@ -0,0 +1,194 @@
1+{
2+ "op_type": "AttentionUpdate",
3+ "op_list": [
4+ {
5+ "bin_filename": "AttentionUpdate_ba64f6ac73b9d4f14511676cbb9f0632",
6+ "inputs": [
7+ [
8+ {
9+ "name": "lse",
10+ "index": 0,
11+ "dtype": "float32",
12+ "format": "ND",
13+ "paramType": "dynamic",
14+ "shape": [
15+ -2
16+ ]
17+ }
18+ ],
19+ [
20+ {
21+ "name": "go",
22+ "index": 1,
23+ "dtype": "float32",
24+ "format": "ND",
25+ "paramType": "dynamic",
26+ "shape": [
27+ -2
28+ ]
29+ }
30+ ]
31+ ],
32+ "outputs": [
33+ {
34+ "name": "output",
35+ "index": 0,
36+ "dtype": "float32",
37+ "format": "ND",
38+ "paramType": "required",
39+ "shape": [
40+ -2
41+ ]
42+ },
43+ {
44+ "name": "lse_m",
45+ "index": 1,
46+ "dtype": "float32",
47+ "format": "ND",
48+ "paramType": "required",
49+ "shape": [
50+ -2
51+ ]
52+ }
53+ ],
54+ "attrs": [
55+ {
56+ "name": "update_type",
57+ "dtype": "int",
58+ "value": 0
59+ },
60+ {
61+ "name": "sp",
62+ "dtype": "int",
63+ "value": null
64+ }
65+ ]
66+ },
67+ {
68+ "bin_filename": "AttentionUpdate_ba64e3ac73b9d4f14511676cbb9f0632",
69+ "inputs": [
70+ [
71+ {
72+ "name": "lse",
73+ "index": 0,
74+ "dtype": "float32",
75+ "format": "ND",
76+ "paramType": "dynamic",
77+ "shape": [
78+ -2
79+ ]
80+ }
81+ ],
82+ [
83+ {
84+ "name": "go",
85+ "index": 1,
86+ "dtype": "float16",
87+ "format": "ND",
88+ "paramType": "dynamic",
89+ "shape": [
90+ -2
91+ ]
92+ }
93+ ]
94+ ],
95+ "outputs": [
96+ {
97+ "name": "output",
98+ "index": 0,
99+ "dtype": "float16",
100+ "format": "ND",
101+ "paramType": "required",
102+ "shape": [
103+ -2
104+ ]
105+ },
106+ {
107+ "name": "lse_m",
108+ "index": 1,
109+ "dtype": "float32",
110+ "format": "ND",
111+ "paramType": "required",
112+ "shape": [
113+ -2
114+ ]
115+ }
116+ ],
117+ "attrs": [
118+ {
119+ "name": "update_type",
120+ "dtype": "int",
121+ "value": 0
122+ },
123+ {
124+ "name": "sp",
125+ "dtype": "int",
126+ "value": null
127+ }
128+ ]
129+ },
130+ {
131+ "bin_filename": "AttentionUpdate_ba69m3ac73b9d4f14511676cbb9f0632",
132+ "inputs": [
133+ [
134+ {
135+ "name": "lse",
136+ "index": 0,
137+ "dtype": "float32",
138+ "format": "ND",
139+ "paramType": "dynamic",
140+ "shape": [
141+ -2
142+ ]
143+ }
144+ ],
145+ [
146+ {
147+ "name": "go",
148+ "index": 1,
149+ "dtype": "bfloat16",
150+ "format": "ND",
151+ "paramType": "dynamic",
152+ "shape": [
153+ -2
154+ ]
155+ }
156+ ]
157+ ],
158+ "outputs": [
159+ {
160+ "name": "output",
161+ "index": 0,
162+ "dtype": "bfloat16",
163+ "format": "ND",
164+ "paramType": "required",
165+ "shape": [
166+ -2
167+ ]
168+ },
169+ {
170+ "name": "lse_m",
171+ "index": 1,
172+ "dtype": "float32",
173+ "format": "ND",
174+ "paramType": "required",
175+ "shape": [
176+ -2
177+ ]
178+ }
179+ ],
180+ "attrs": [
181+ {
182+ "name": "update_type",
183+ "dtype": "int",
184+ "value": 0
185+ },
186+ {
187+ "name": "sp",
188+ "dtype": "int",
189+ "value": null
190+ }
191+ ]
192+ }
193+ ]
194+}
Mattention/fused_causal_conv1d/op_host/CMakeLists.txt+10-0
@@ -42,6 +42,16 @@ if (BUILD_OPEN_PROJECT)
42 -Wno-deprecated-declarations42 -Wno-deprecated-declarations
43 -Werror43 -Werror
44 )44 )
45+ elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend960pr")
46+ add_ops_compile_options(
47+ OP_NAME FusedCausalConv1d
48+ COMPUTE_UNIT Ascend960PR_969A1
49+ OPTIONS --cce-auto-sync=off
50+ -mllvm
51+ -cce-aicore-dcci-before-kernel-end=false
52+ -Wno-deprecated-declarations
53+ -Werror
54+ )
45 else()55 else()
46 add_ops_compile_options(56 add_ops_compile_options(
47 OP_NAME FusedCausalConv1d57 OP_NAME FusedCausalConv1d
Mattention/fused_causal_conv1d/op_host/fused_causal_conv1d_def.cpp+1-0
@@ -66,6 +66,7 @@ public:
66 this->AICore().AddConfig("ascend950", regbaseCfg);66 this->AICore().AddConfig("ascend950", regbaseCfg);
67 this->AICore().AddConfig("ascend350", regbaseCfg);67 this->AICore().AddConfig("ascend350", regbaseCfg);
68 this->AICore().AddConfig("ascend960dt", regbaseCfg);68 this->AICore().AddConfig("ascend960dt", regbaseCfg);
69+ this->AICore().AddConfig("ascend960pr", regbaseCfg);
69 }70 }
70};71};
71 72 
Mattention/gather_pa_kv_cache/op_host/CMakeLists.txt+6-0
@@ -31,6 +31,12 @@ if (BUILD_OPEN_PROJECT)
31 COMPUTE_UNIT ascend960dt_968B731 COMPUTE_UNIT ascend960dt_968B7
32 OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false32 OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
33 )33 )
34+ elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend960pr")
35+ add_ops_compile_options(
36+ OP_NAME GatherPaKvCache
37+ COMPUTE_UNIT Ascend960PR_969A1
38+ OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
39+ )
34 else()40 else()
35 add_ops_compile_options(41 add_ops_compile_options(
36 OP_NAME GatherPaKvCache42 OP_NAME GatherPaKvCache
Mattention/gather_pa_kv_cache/op_host/gather_pa_kv_cache_def.cpp+1-0
@@ -290,6 +290,7 @@ private:
290 .NeedCheckSupportFlag(false)290 .NeedCheckSupportFlag(false)
291 .ExtendCfgInfo("opFile.value", "gather_pa_kv_cache_apt");291 .ExtendCfgInfo("opFile.value", "gather_pa_kv_cache_apt");
292 this->AICore().AddConfig("ascend960dt", config_96);292 this->AICore().AddConfig("ascend960dt", config_96);
293+ this->AICore().AddConfig("ascend960pr", config_96);
293 }294 }
294};295};
295 296 
Mattention/gather_pa_kv_cache/op_host/gather_pa_kv_cache_tiling_arch35.cpp+2-1
@@ -594,7 +594,8 @@ ge::graphStatus GatherPaKvCacheTiling::GetShapeAttrsInfo()
594 auto platformInfo = context_->GetPlatformInfo();594 auto platformInfo = context_->GetPlatformInfo();
595 OP_CHECK_IF(platformInfo == nullptr, OP_LOGE(context_, "platformInfo is nullptr."), return ge::GRAPH_FAILED);595 OP_CHECK_IF(platformInfo == nullptr, OP_LOGE(context_, "platformInfo is nullptr."), return ge::GRAPH_FAILED);
596 auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);596 auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
597- isAscend96_ = ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_9201;597+ isAscend96_ =
598+ ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_9201 || ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_9202;
598 OP_CHECK_IF(context_ == nullptr, OP_LOGE("GatherPaKvCacheTiling", "context is null."), return ge::GRAPH_FAILED);599 OP_CHECK_IF(context_ == nullptr, OP_LOGE("GatherPaKvCacheTiling", "context is null."), return ge::GRAPH_FAILED);
599 OP_CHECK_IF(GetAttrs() != ge::GRAPH_SUCCESS, OP_LOGE(context_, "get attrs failed."), return ge::GRAPH_FAILED);600 OP_CHECK_IF(GetAttrs() != ge::GRAPH_SUCCESS, OP_LOGE(context_, "get attrs failed."), return ge::GRAPH_FAILED);
600 OP_CHECK_IF(GetInputKeyCache() != ge::GRAPH_SUCCESS, OP_LOGE(context_, "get input key_cache failed."),601 OP_CHECK_IF(GetInputKeyCache() != ge::GRAPH_SUCCESS, OP_LOGE(context_, "get input key_cache failed."),
Mattention/gather_pa_kv_cache/op_host/op_api/aclnn_gather_pa_kv_cache.cpp+1-1
@@ -80,7 +80,7 @@ static aclnnStatus CheckDtypeValid(const aclTensor* keyCache, const aclTensor* v
80 auto& platformInfo = op::GetCurrentPlatformInfo();80 auto& platformInfo = op::GetCurrentPlatformInfo();
81 auto socVersion = platformInfo.GetCurNpuArch();81 auto socVersion = platformInfo.GetCurNpuArch();
82 bool isArch35 = (socVersion == NpuArch::DAV_3510);82 bool isArch35 = (socVersion == NpuArch::DAV_3510);
83- bool isAscend96 = (socVersion == NpuArch::DAV_9201);83+ bool isAscend96 = (socVersion == NpuArch::DAV_9201 || socVersion == NpuArch::DAV_9202);
84 const auto& kvDtypeList =84 const auto& kvDtypeList =
85 isArch35 ? DTYPE_SUPPORT_LIST_A5 : (isAscend96 ? DTYPE_SUPPORT_LIST_A6 : DTYPE_SUPPORT_LIST_910B);85 isArch35 ? DTYPE_SUPPORT_LIST_A5 : (isAscend96 ? DTYPE_SUPPORT_LIST_A6 : DTYPE_SUPPORT_LIST_910B);
86 86 
Mattention/indexer_quant_cache/op_host/CMakeLists.txt+6-0
@@ -25,6 +25,12 @@ if (BUILD_OPEN_PROJECT)
25 COMPUTE_UNIT ascend960dt_968B725 COMPUTE_UNIT ascend960dt_968B7
26 OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false26 OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
27 )27 )
28+ elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend960pr")
29+ add_ops_compile_options(
30+ OP_NAME IndexerQuantCache
31+ COMPUTE_UNIT Ascend960PR_969A1
32+ OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
33+ )
28 else()34 else()
29 add_ops_compile_options(35 add_ops_compile_options(
30 OP_NAME IndexerQuantCache36 OP_NAME IndexerQuantCache
Rmoe/moe_gating_top_k/op_host/config/ascend960pr/moe_gating_top_k_simplified_key.ini→attention/indexer_quant_cache/op_host/config/ascend960pr/indexer_quant_cache_simplified_key.ini+1-1
@@ -9,5 +9,5 @@
9; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。9; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式10; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定11; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12-[MoeGatingTopK]12+[IndexerQuantCache]
13default=013default=0
Mattention/indexer_quant_cache/op_host/indexer_quant_cache_def.cpp+1-0
@@ -171,6 +171,7 @@ public:
171 ge::FORMAT_ND});171 ge::FORMAT_ND});
172 172 
173 this->AICore().AddConfig("ascend960dt", config960dt);173 this->AICore().AddConfig("ascend960dt", config960dt);
174+ this->AICore().AddConfig("ascend960pr", config960dt);
174 }175 }
175};176};
176 177 
Mattention/indexer_quant_cache/op_host/indexer_quant_cache_tiling_arch35.cpp+2-1
@@ -222,7 +222,8 @@ ge::graphStatus IndexerQuantCacheTiling::GetShapeAttrsInfoInner()
222 OP_LOGE(context_->GetNodeName(), "quant_mode should be in [0,4], got %ld", quantMode_),222 OP_LOGE(context_->GetNodeName(), "quant_mode should be in [0,4], got %ld", quantMode_),
223 return ge::GRAPH_FAILED);223 return ge::GRAPH_FAILED);
224 224 
225- OP_CHECK_IF(quantMode_ == HIFP4_QUANT_MODE && (npuArch_ != Ops::Base::DAV_9201 || d_ % HIFP4_GROUP_SIZE != 0),225+ OP_CHECK_IF(quantMode_ == HIFP4_QUANT_MODE &&
226+ (npuArch_ != Ops::Base::DAV_9201 && npuArch_ != Ops::Base::DAV_9202 || d_ % HIFP4_GROUP_SIZE != 0),
226 OP_LOGE(context_->GetNodeName(), "HiFloat4 requires ASCEND960DT and d aligned to 64."),227 OP_LOGE(context_->GetNodeName(), "HiFloat4 requires ASCEND960DT and d aligned to 64."),
227 return ge::GRAPH_FAILED);228 return ge::GRAPH_FAILED);
228 229 
Mattention/indexer_quant_cache/op_kernel/indexer_quant_cache.cpp+3-3
@@ -20,7 +20,7 @@
20#include "arch35/indexer_quant_cache_multi_row_mx_fp4.h"20#include "arch35/indexer_quant_cache_multi_row_mx_fp4.h"
21#include "arch35/indexer_quant_cache_single_row_mx_fp4.h"21#include "arch35/indexer_quant_cache_single_row_mx_fp4.h"
22 22 
23-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 9201)23+#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 9201 || __NPU_ARCH__ == 9202)
24#include "arch35/indexer_quant_cache_multi_row_hifp4.h"24#include "arch35/indexer_quant_cache_multi_row_hifp4.h"
25#include "arch35/indexer_quant_cache_single_row_hifp4.h"25#include "arch35/indexer_quant_cache_single_row_hifp4.h"
26#endif26#endif
@@ -49,8 +49,8 @@ extern "C" __global__ __aicore__ void indexer_quant_cache(GM_ADDR cache, GM_ADDR
49 GET_TILING_DATA(tilingData, tiling);49 GET_TILING_DATA(tilingData, tiling);
50 TPipe pipe;50 TPipe pipe;
51 int64_t oriOverflowMode = AscendC::GetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>();51 int64_t oriOverflowMode = AscendC::GetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>();
52-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 9201) && defined(ORIG_DTYPE_CACHE) && defined(DT_HIFLOAT4) && \52+#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 9201 || __NPU_ARCH__ == 9202) && defined(ORIG_DTYPE_CACHE) && \
53- (ORIG_DTYPE_CACHE == DT_HIFLOAT4)53+ defined(DT_HIFLOAT4) && (ORIG_DTYPE_CACHE == DT_HIFLOAT4)
54 // HiFloat4 cache is byte packed; one 64-element group produces one uint32 scale.54 // HiFloat4 cache is byte packed; one 64-element group produces one uint32 scale.
55 if (TILING_KEY_IS(10013)) {55 if (TILING_KEY_IS(10013)) {
56 IndexerQuantCache::IndexerQuantCacheMultiRowHiFp4<DTYPE_X, int8_t, uint32_t> op;56 IndexerQuantCache::IndexerQuantCacheMultiRowHiFp4<DTYPE_X, int8_t, uint32_t> op;
Mattention/kv_compress_epilog/op_host/CMakeLists.txt+8-0
@@ -37,6 +37,14 @@ if (BUILD_OPEN_PROJECT)
37 -mllvm -cce-aicore-hoist-movemask=false37 -mllvm -cce-aicore-hoist-movemask=false
38 --op_relocatable_kernel_binary=true38 --op_relocatable_kernel_binary=true
39 )39 )
40+ elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend960pr")
41+ add_ops_compile_options(
42+ OP_NAME KvCompressEpilog
43+ COMPUTE_UNIT Ascend960PR_969A1
44+ OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
45+ -mllvm -cce-aicore-hoist-movemask=false
46+ --op_relocatable_kernel_binary=true
47+ )
40 else()48 else()
41 add_ops_compile_options(49 add_ops_compile_options(
42 OP_NAME KvCompressEpilog50 OP_NAME KvCompressEpilog
Aattention/kv_compress_epilog/op_host/config/ascend960pr/kv_compress_epilog_binary.json+143-0
@@ -0,0 +1,143 @@
1+{
2+ "op_type": "KvCompressEpilog",
3+ "op_list": [
4+ {
5+ "bin_filename": "KvCompressEpilog_uint8_bf16",
6+ "inputs": [
7+ {
8+ "name": "cache",
9+ "index": 0,
10+ "dtype": "uint8",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [
14+ -2
15+ ]
16+ },
17+ {
18+ "name": "x",
19+ "index": 1,
20+ "dtype": "bfloat16",
21+ "format": "ND",
22+ "paramType": "required",
23+ "shape": [
24+ -2
25+ ]
26+ },
27+ {
28+ "name": "slot_mapping",
29+ "index": 2,
30+ "dtype": "int32",
31+ "format": "ND",
32+ "paramType": "required",
33+ "shape": [
34+ -2
35+ ]
36+ }
37+ ],
38+ "outputs": [
39+ {
40+ "name": "cache",
41+ "index": 0,
42+ "dtype": "uint8",
43+ "format": "ND",
44+ "paramType": "required",
45+ "shape": [
46+ -2
47+ ]
48+ }
49+ ],
50+ "attrs": [
51+ {
52+ "name": "quant_group_size",
53+ "dtype": "int",
54+ "value": 128
55+ },
56+ {
57+ "name": "quant_mode",
58+ "dtype": "int",
59+ "value": 1
60+ },
61+ {
62+ "name": "round_scale",
63+ "dtype": "bool",
64+ "value": true
65+ },
66+ {
67+ "name": "x_scale",
68+ "dtype": "float",
69+ "value": 1.0
70+ }
71+ ]
72+ },
73+ {
74+ "bin_filename": "KvCompressEpilog_uint8_bf16_int64",
75+ "inputs": [
76+ {
77+ "name": "cache",
78+ "index": 0,
79+ "dtype": "uint8",
80+ "format": "ND",
81+ "paramType": "required",
82+ "shape": [
83+ -2
84+ ]
85+ },
86+ {
87+ "name": "x",
88+ "index": 1,
89+ "dtype": "bfloat16",
90+ "format": "ND",
91+ "paramType": "required",
92+ "shape": [
93+ -2
94+ ]
95+ },
96+ {
97+ "name": "slot_mapping",
98+ "index": 2,
99+ "dtype": "int64",
100+ "format": "ND",
101+ "paramType": "required",
102+ "shape": [
103+ -2
104+ ]
105+ }
106+ ],
107+ "outputs": [
108+ {
109+ "name": "cache",
110+ "index": 0,
111+ "dtype": "uint8",
112+ "format": "ND",
113+ "paramType": "required",
114+ "shape": [
115+ -2
116+ ]
117+ }
118+ ],
119+ "attrs": [
120+ {
121+ "name": "quant_group_size",
122+ "dtype": "int",
123+ "value": 128
124+ },
125+ {
126+ "name": "quant_mode",
127+ "dtype": "int",
128+ "value": 1
129+ },
130+ {
131+ "name": "round_scale",
132+ "dtype": "bool",
133+ "value": true
134+ },
135+ {
136+ "name": "x_scale",
137+ "dtype": "float",
138+ "value": 1.0
139+ }
140+ ]
141+ }
142+ ]
143+}
Mattention/kv_compress_epilog/op_host/kv_compress_epilog_def.cpp+1-0
@@ -50,6 +50,7 @@ public:
50 this->AICore().AddConfig("ascend950");50 this->AICore().AddConfig("ascend950");
51 this->AICore().AddConfig("ascend350");51 this->AICore().AddConfig("ascend350");
52 this->AICore().AddConfig("ascend960dt");52 this->AICore().AddConfig("ascend960dt");
53+ this->AICore().AddConfig("ascend960pr");
53 }54 }
54};55};
55 56 
Mattention/scatter_pa_kv_cache/op_host/CMakeLists.txt+6-0
@@ -32,6 +32,12 @@ if (BUILD_OPEN_PROJECT)
32 COMPUTE_UNIT ascend960dt_968B732 COMPUTE_UNIT ascend960dt_968B7
33 OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false33 OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
34 )34 )
35+ elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend960pr")
36+ add_ops_compile_options(
37+ OP_NAME ScatterPaKvCache
38+ COMPUTE_UNIT Ascend960PR_969A1
39+ OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
40+ )
35 else()41 else()
36 add_ops_compile_options(42 add_ops_compile_options(
37 OP_NAME ScatterPaKvCache43 OP_NAME ScatterPaKvCache
Mattention/scatter_pa_kv_cache/op_host/op_api/aclnn_scatter_pa_kv_cache.cpp+1-1
@@ -84,7 +84,7 @@ static aclnnStatus CheckDtypeValid(const aclTensor* key, const aclTensor* keyCac
84 auto& platformInfo = op::GetCurrentPlatformInfo();84 auto& platformInfo = op::GetCurrentPlatformInfo();
85 auto socVersion = platformInfo.GetCurNpuArch();85 auto socVersion = platformInfo.GetCurNpuArch();
86 bool isArch35 = (socVersion == NpuArch::DAV_3510);86 bool isArch35 = (socVersion == NpuArch::DAV_3510);
87- bool isAscend96 = (socVersion == NpuArch::DAV_9201);87+ bool isAscend96 = (socVersion == NpuArch::DAV_9201 || socVersion == NpuArch::DAV_9202);
88 auto dtypeSupportList = isArch35 ?88 auto dtypeSupportList = isArch35 ?
89 KEY_VALUE_DTYPE_SUPPORT_LIST :89 KEY_VALUE_DTYPE_SUPPORT_LIST :
90 (isAscend96 ? KEY_VALUE_DTYPE_SUPPORT_LIST_960 : KEY_VALUE_DTYPE_SUPPORT_LIST_910);90 (isAscend96 ? KEY_VALUE_DTYPE_SUPPORT_LIST_960 : KEY_VALUE_DTYPE_SUPPORT_LIST_910);
Mattention/scatter_pa_kv_cache/op_host/scatter_pa_kv_cache_def.cpp+1-0
@@ -724,6 +724,7 @@ public:
724 .PrecisionReduceFlag(true)724 .PrecisionReduceFlag(true)
725 .ExtendCfgInfo("opFile.value", "scatter_pa_kv_cache_apt");725 .ExtendCfgInfo("opFile.value", "scatter_pa_kv_cache_apt");
726 this->AICore().AddConfig("ascend960dt", config_960);726 this->AICore().AddConfig("ascend960dt", config_960);
727+ this->AICore().AddConfig("ascend960pr", config_960);
727 }728 }
728};729};
729OP_ADD(ScatterPaKvCache);730OP_ADD(ScatterPaKvCache);
Mattention/scatter_pa_kv_cache/op_host/scatter_pa_kv_cache_tiling_arch35.cpp+2-1
@@ -1237,7 +1237,8 @@ ge::graphStatus ScatterPaKvCacheTiling::GetShapeAttrsInfo()
1237 auto platformInfo = context_->GetPlatformInfo();1237 auto platformInfo = context_->GetPlatformInfo();
1238 OP_CHECK_IF(platformInfo == nullptr, OP_LOGE(context_, "platformInfo is nullptr."), return ge::GRAPH_FAILED);1238 OP_CHECK_IF(platformInfo == nullptr, OP_LOGE(context_, "platformInfo is nullptr."), return ge::GRAPH_FAILED);
1239 auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);1239 auto ascendcPlatform = platform_ascendc::PlatformAscendC(platformInfo);
1240- isAscend96_ = ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_9201;1240+ isAscend96_ =
1241+ ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_9201 || ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_9202;
1241 if (GetTensorInfo(inputKeyShape_, keyStride_, inputKey_) != ge::GRAPH_SUCCESS) {1242 if (GetTensorInfo(inputKeyShape_, keyStride_, inputKey_) != ge::GRAPH_SUCCESS) {
1242 return ge::GRAPH_FAILED;1243 return ge::GRAPH_FAILED;
1243 }1244 }
Mattention/scatter_pa_kv_cache/op_kernel/arch35/scatter_pa_kv_cache_normal_fully_load.h+5-5
@@ -219,7 +219,7 @@ __aicore__ inline void ScatterPaKvCacheNormalFullyLoad<T, IndexDtype, OffsetDtyp
219 GlobalTensor<T>& outputGm, LocalTensor<T>& inputLocal, LocalTensor<OffsetDtype>& offsetLocal,219 GlobalTensor<T>& outputGm, LocalTensor<T>& inputLocal, LocalTensor<OffsetDtype>& offsetLocal,
220 LocalTensor<IndexDtype>& slotMapLocal, int64_t stride0, int64_t stride1, int64_t handleNum, int64_t curFactor)220 LocalTensor<IndexDtype>& slotMapLocal, int64_t stride0, int64_t stride1, int64_t handleNum, int64_t curFactor)
221{221{
222-#if defined(__NPU_ARCH__) && __NPU_ARCH__ == 9201222+#if defined(__NPU_ARCH__) && __NPU_ARCH__ == 9201 || __NPU_ARCH__ == 9202
223 ScatterParams scatterDataCopyParams;223 ScatterParams scatterDataCopyParams;
224 scatterDataCopyParams.offsetNum = curFactor;224 scatterDataCopyParams.offsetNum = curFactor;
225 scatterDataCopyParams.burstLen = handleNum * sizeof(T);225 scatterDataCopyParams.burstLen = handleNum * sizeof(T);
@@ -263,7 +263,7 @@ __aicore__ inline void ScatterPaKvCacheNormalFullyLoad<T, IndexDtype, OffsetDtyp
263 kSlotMappingQue_.EnQue(kSlotMappingLocal);263 kSlotMappingQue_.EnQue(kSlotMappingLocal);
264 kSlotMappingLocal = kSlotMappingQue_.DeQue<IndexDtype>();264 kSlotMappingLocal = kSlotMappingQue_.DeQue<IndexDtype>();
265 LocalTensor<OffsetDtype> kOffsetLocal;265 LocalTensor<OffsetDtype> kOffsetLocal;
266-#if defined(__NPU_ARCH__) && __NPU_ARCH__ == 9201266+#if defined(__NPU_ARCH__) && __NPU_ARCH__ == 9201 || __NPU_ARCH__ == 9202
267 kOffsetLocal = kOffsetQue_.AllocTensor<OffsetDtype>();267 kOffsetLocal = kOffsetQue_.AllocTensor<OffsetDtype>();
268 ComputeOutOffset(kOffsetLocal, kSlotMappingLocal, curUbFactor, tilingData_->keyCacheStride0,268 ComputeOutOffset(kOffsetLocal, kSlotMappingLocal, curUbFactor, tilingData_->keyCacheStride0,
269 tilingData_->keyCacheStride1);269 tilingData_->keyCacheStride1);
@@ -276,7 +276,7 @@ __aicore__ inline void ScatterPaKvCacheNormalFullyLoad<T, IndexDtype, OffsetDtyp
276 tilingData_->keyCacheStride1, tilingData_->kHandleNumPerCore, curUbFactor);276 tilingData_->keyCacheStride1, tilingData_->kHandleNumPerCore, curUbFactor);
277 inputKeyQueue_.FreeTensor(inputKeyLocal);277 inputKeyQueue_.FreeTensor(inputKeyLocal);
278 kSlotMappingQue_.FreeTensor(kSlotMappingLocal);278 kSlotMappingQue_.FreeTensor(kSlotMappingLocal);
279-#if defined(__NPU_ARCH__) && __NPU_ARCH__ == 9201279+#if defined(__NPU_ARCH__) && __NPU_ARCH__ == 9201 || __NPU_ARCH__ == 9202
280 kOffsetQue_.FreeTensor(kOffsetLocal);280 kOffsetQue_.FreeTensor(kOffsetLocal);
281#endif281#endif
282 if constexpr (InOutMode == DUAL_IN_OUT) {282 if constexpr (InOutMode == DUAL_IN_OUT) {
@@ -288,7 +288,7 @@ __aicore__ inline void ScatterPaKvCacheNormalFullyLoad<T, IndexDtype, OffsetDtyp
288 vSlotMappingQue_.EnQue(vSlotMappingLocal);288 vSlotMappingQue_.EnQue(vSlotMappingLocal);
289 vSlotMappingLocal = vSlotMappingQue_.DeQue<IndexDtype>();289 vSlotMappingLocal = vSlotMappingQue_.DeQue<IndexDtype>();
290 LocalTensor<OffsetDtype> vOffsetLocal;290 LocalTensor<OffsetDtype> vOffsetLocal;
291-#if defined(__NPU_ARCH__) && __NPU_ARCH__ == 9201291+#if defined(__NPU_ARCH__) && __NPU_ARCH__ == 9201 || __NPU_ARCH__ == 9202
292 vOffsetLocal = vOffsetQue_.AllocTensor<OffsetDtype>();292 vOffsetLocal = vOffsetQue_.AllocTensor<OffsetDtype>();
293 ComputeOutOffset(vOffsetLocal, vSlotMappingLocal, curUbFactor, tilingData_->valueCacheStride0,293 ComputeOutOffset(vOffsetLocal, vSlotMappingLocal, curUbFactor, tilingData_->valueCacheStride0,
294 tilingData_->valueCacheStride1);294 tilingData_->valueCacheStride1);
@@ -302,7 +302,7 @@ __aicore__ inline void ScatterPaKvCacheNormalFullyLoad<T, IndexDtype, OffsetDtyp
302 curUbFactor);302 curUbFactor);
303 inputValueQueue_.FreeTensor(inputValueLocal);303 inputValueQueue_.FreeTensor(inputValueLocal);
304 vSlotMappingQue_.FreeTensor(vSlotMappingLocal);304 vSlotMappingQue_.FreeTensor(vSlotMappingLocal);
305-#if defined(__NPU_ARCH__) && __NPU_ARCH__ == 9201305+#if defined(__NPU_ARCH__) && __NPU_ARCH__ == 9201 || __NPU_ARCH__ == 9202
306 vOffsetQue_.FreeTensor(vOffsetLocal);306 vOffsetQue_.FreeTensor(vOffsetLocal);
307#endif307#endif
308 }308 }
Mattention/stem_oam_prep_varlen_q/op_host/stem_oam_prep_varlen_q_def.cpp+1-0
@@ -57,6 +57,7 @@ public:
57 this->AICore().AddConfig("ascend950", aicore_config);57 this->AICore().AddConfig("ascend950", aicore_config);
58 this->AICore().AddConfig("ascend350", aicore_config);58 this->AICore().AddConfig("ascend350", aicore_config);
59 this->AICore().AddConfig("ascend960dt", aicore_config);59 this->AICore().AddConfig("ascend960dt", aicore_config);
60+ this->AICore().AddConfig("ascend960pr", aicore_config);
60 }61 }
61};62};
62OP_ADD(StemOamPrepVarlenQ);63OP_ADD(StemOamPrepVarlenQ);
Mgmm/grouped_matmul/op_host/op_tiling/grouped_matmul_tiling.cpp+2-2
@@ -1359,8 +1359,8 @@ ge::graphStatus GMMTiling::GMMGetAttrs(const gert::TilingContext* context)
1359 OP_CHECK_IF(1359 OP_CHECK_IF(
1360 !((wFormat0 == ge::FORMAT_FRACTAL_NZ) || (wFormat0 != FORMAT_FRACTAL_NZ && transposeWeight_ == 0)),1360 !((wFormat0 == ge::FORMAT_FRACTAL_NZ) || (wFormat0 != FORMAT_FRACTAL_NZ && transposeWeight_ == 0)),
1361 OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(),1361 OPS_REPORT_VECTOR_INNER_ERR(context->GetNodeName(),
1362- "A4W4 GMM currently supports only weight tensor nz transpose/untranspose input "1362+ "A4W4 GMM currently supports only weight tensor nz transpose/untranspose input \
1363- "or nd format untranspose input."),1363+or nd format untranspose input."),
1364 return ge::GRAPH_FAILED);1364 return ge::GRAPH_FAILED);
1365 }1365 }
1366 return ge::GRAPH_SUCCESS;1366 return ge::GRAPH_SUCCESS;
Mmhc/mhc_post_backward/op_host/CMakeLists.txt+6-0
@@ -31,6 +31,12 @@ if (BUILD_OPEN_PROJECT)
31 COMPUTE_UNIT ascend960dt_968B731 COMPUTE_UNIT ascend960dt_968B7
32 OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false32 OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
33 )33 )
34+ elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend960pr")
35+ add_ops_compile_options(
36+ OP_NAME MhcPostBackward
37+ COMPUTE_UNIT Ascend960PR_969A1
38+ OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
39+ )
34 else()40 else()
35 add_ops_compile_options(41 add_ops_compile_options(
36 OP_NAME MhcPostBackward42 OP_NAME MhcPostBackward
Amhc/mhc_post_backward/op_host/config/ascend960pr/mhc_post_backward_binary.json+199-0
@@ -0,0 +1,199 @@
1+{
2+ "op_type": "MhcPostBackward",
3+ "op_list": [
4+ {
5+ "bin_filename": "MhcPostBackward_FP16_FP16_FP32_FP16_FP32",
6+ "inputs": [
7+ {
8+ "name": "grad_y",
9+ "index": 0,
10+ "dtype": "float16",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [
14+ -2
15+ ]
16+ },
17+ {
18+ "name": "x",
19+ "index": 0,
20+ "dtype": "float16",
21+ "format": "ND",
22+ "paramType": "required",
23+ "shape": [
24+ -2
25+ ]
26+ },
27+ {
28+ "name": "h_res",
29+ "index": 0,
30+ "dtype": "float32",
31+ "format": "ND",
32+ "paramType": "required",
33+ "shape": [
34+ -2
35+ ]
36+ },
37+ {
38+ "name": "h_out",
39+ "index": 0,
40+ "dtype": "float16",
41+ "format": "ND",
42+ "paramType": "required",
43+ "shape": [
44+ -2
45+ ]
46+ },
47+ {
48+ "name": "h_post",
49+ "index": 0,
50+ "dtype": "float32",
51+ "format": "ND",
52+ "paramType": "required",
53+ "shape": [
54+ -2
55+ ]
56+ }
57+ ],
58+ "outputs": [
59+ {
60+ "name": "grad_x",
61+ "index": 0,
62+ "dtype": "float16",
63+ "format": "ND",
64+ "paramType": "required",
65+ "shape": [
66+ -2
67+ ]
68+ },
69+ {
70+ "name": "grad_h_res",
71+ "index": 1,
72+ "dtype": "float32",
73+ "format": "ND",
74+ "paramType": "required",
75+ "shape": [
76+ -2
77+ ]
78+ },
79+ {
80+ "name": "grad_h_out",
81+ "index": 2,
82+ "dtype": "float16",
83+ "format": "ND",
84+ "paramType": "required",
85+ "shape": [
86+ -2
87+ ]
88+ },
89+ {
90+ "name": "grad_h_post",
91+ "index": 3,
92+ "dtype": "float32",
93+ "format": "ND",
94+ "paramType": "required",
95+ "shape": [
96+ -2
97+ ]
98+ }
99+ ]
100+ },
101+ {
102+ "bin_filename": "MhcPostBackward_BF16_BF16_FP32_BF16_FP32",
103+ "inputs": [
104+ {
105+ "name": "grad_y",
106+ "index": 0,
107+ "dtype": "bfloat16",
108+ "format": "ND",
109+ "paramType": "required",
110+ "shape": [
111+ -2
112+ ]
113+ },
114+ {
115+ "name": "x",
116+ "index": 0,
117+ "dtype": "bfloat16",
118+ "format": "ND",
119+ "paramType": "required",
120+ "shape": [
121+ -2
122+ ]
123+ },
124+ {
125+ "name": "h_res",
126+ "index": 0,
127+ "dtype": "float32",
128+ "format": "ND",
129+ "paramType": "required",
130+ "shape": [
131+ -2
132+ ]
133+ },
134+ {
135+ "name": "h_out",
136+ "index": 0,
137+ "dtype": "bfloat16",
138+ "format": "ND",
139+ "paramType": "required",
140+ "shape": [
141+ -2
142+ ]
143+ },
144+ {
145+ "name": "h_post",
146+ "index": 0,
147+ "dtype": "float32",
148+ "format": "ND",
149+ "paramType": "required",
150+ "shape": [
151+ -2
152+ ]
153+ }
154+ ],
155+ "outputs": [
156+ {
157+ "name": "grad_x",
158+ "index": 0,
159+ "dtype": "bfloat16",
160+ "format": "ND",
161+ "paramType": "required",
162+ "shape": [
163+ -2
164+ ]
165+ },
166+ {
167+ "name": "grad_h_res",
168+ "index": 1,
169+ "dtype": "float32",
170+ "format": "ND",
171+ "paramType": "required",
172+ "shape": [
173+ -2
174+ ]
175+ },
176+ {
177+ "name": "grad_h_out",
178+ "index": 2,
179+ "dtype": "bfloat16",
180+ "format": "ND",
181+ "paramType": "required",
182+ "shape": [
183+ -2
184+ ]
185+ },
186+ {
187+ "name": "grad_h_post",
188+ "index": 3,
189+ "dtype": "float32",
190+ "format": "ND",
191+ "paramType": "required",
192+ "shape": [
193+ -2
194+ ]
195+ }
196+ ]
197+ }
198+ ]
199+ }
Mmhc/mhc_post_backward/op_host/mhc_post_backward_def.cpp+1-0
@@ -106,6 +106,7 @@ public:
106 this->AICore().AddConfig("ascend950", aicore_Config_950);106 this->AICore().AddConfig("ascend950", aicore_Config_950);
107 this->AICore().AddConfig("ascend350", aicore_Config_950);107 this->AICore().AddConfig("ascend350", aicore_Config_950);
108 this->AICore().AddConfig("ascend960dt", aicore_Config_950);108 this->AICore().AddConfig("ascend960dt", aicore_Config_950);
109+ this->AICore().AddConfig("ascend960pr", aicore_Config_950);
109 }110 }
110};111};
111 112 
Mmhc/mhc_post_backward/op_host/op_api/aclnn_mhc_post_backward.cpp+4-4
@@ -67,9 +67,9 @@ static aclnnStatus CheckNotNull(const aclTensor* gradOutput, const aclTensor* x,
67{67{
68 CHECK_COND(gradOutput != nullptr, ACLNN_ERR_PARAM_NULLPTR, "gradOutput must not be nullptr.");68 CHECK_COND(gradOutput != nullptr, ACLNN_ERR_PARAM_NULLPTR, "gradOutput must not be nullptr.");
69 CHECK_COND(x != nullptr, ACLNN_ERR_PARAM_NULLPTR, "x must not be nullptr.");69 CHECK_COND(x != nullptr, ACLNN_ERR_PARAM_NULLPTR, "x must not be nullptr.");
70- // hRes is optional on DAV_3510 / DAV_9201, required on A2/A370+ // hRes is optional on DAV_3510 / DAV_9201 / DAV_9202, required on A2/A3
71 auto npuArch = GetCurrentPlatformInfo().GetCurNpuArch();71 auto npuArch = GetCurrentPlatformInfo().GetCurNpuArch();
72- if (npuArch != NpuArch::DAV_3510 && npuArch != NpuArch::DAV_9201) {72+ if (npuArch != NpuArch::DAV_3510 && npuArch != NpuArch::DAV_9201 && npuArch != NpuArch::DAV_9202) {
73 CHECK_COND(hRes != nullptr, ACLNN_ERR_PARAM_NULLPTR, "hRes must not be nullptr.");73 CHECK_COND(hRes != nullptr, ACLNN_ERR_PARAM_NULLPTR, "hRes must not be nullptr.");
74 }74 }
75 CHECK_COND(hOut != nullptr, ACLNN_ERR_PARAM_NULLPTR, "hOut must not be nullptr.");75 CHECK_COND(hOut != nullptr, ACLNN_ERR_PARAM_NULLPTR, "hOut must not be nullptr.");
@@ -386,7 +386,7 @@ static aclnnStatus CheckShape(const MhcPostBackwardParams& params)
386 386 
387 if (IsTNDFormat(xDimNum)) {387 if (IsTNDFormat(xDimNum)) {
388 auto xDim1 = params.x->GetViewShape().GetDim(DIM_IDX_1);388 auto xDim1 = params.x->GetViewShape().GetDim(DIM_IDX_1);
389- if (npuArch == NpuArch::DAV_3510 || npuArch == NpuArch::DAV_9201) {389+ if (npuArch == NpuArch::DAV_3510 || npuArch == NpuArch::DAV_9201 || npuArch == NpuArch::DAV_9202) {
390 CHECK_COND((xDim1 == 4) || (xDim1 == 6) || (xDim1 == 8), ACLNN_ERR_PARAM_INVALID,390 CHECK_COND((xDim1 == 4) || (xDim1 == 6) || (xDim1 == 8), ACLNN_ERR_PARAM_INVALID,
391 "the dim n should be 4 / 6 / 8 for 3D format");391 "the dim n should be 4 / 6 / 8 for 3D format");
392 } else {392 } else {
@@ -416,7 +416,7 @@ static aclnnStatus CheckShape(const MhcPostBackwardParams& params)
416 CHECK_COND(CheckShape3D(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, "invalid shape for 3D format");416 CHECK_COND(CheckShape3D(params) == ACLNN_SUCCESS, ACLNN_ERR_PARAM_INVALID, "invalid shape for 3D format");
417 } else {417 } else {
418 auto xDim2 = params.x->GetViewShape().GetDim(DIM_IDX_2);418 auto xDim2 = params.x->GetViewShape().GetDim(DIM_IDX_2);
419- if (npuArch == NpuArch::DAV_3510 || npuArch == NpuArch::DAV_9201) {419+ if (npuArch == NpuArch::DAV_3510 || npuArch == NpuArch::DAV_9201 || npuArch == NpuArch::DAV_9202) {
420 CHECK_COND((xDim2 == 4) || (xDim2 == 6) || (xDim2 == 8), ACLNN_ERR_PARAM_INVALID,420 CHECK_COND((xDim2 == 4) || (xDim2 == 6) || (xDim2 == 8), ACLNN_ERR_PARAM_INVALID,
421 "the dim n should be 4 / 6 / 8 for 4D format");421 "the dim n should be 4 / 6 / 8 for 4D format");
422 } else {422 } else {
Mmhc/mhc_post_backward/op_host/op_tiling/arch35/mhc_post_backward_tiling_base_arch35.cpp+2-1
@@ -876,6 +876,7 @@ void MhcPostBackwardTilingBaseArch35::Reset()
876 876 
877REGISTER_TILING_TEMPLATE_WITH_ARCH(MhcPostBackward, MhcPostBackwardTilingBaseArch35,877REGISTER_TILING_TEMPLATE_WITH_ARCH(MhcPostBackward, MhcPostBackwardTilingBaseArch35,
878 std::vector<int32_t>({static_cast<int32_t>(NpuArch::DAV_3510),878 std::vector<int32_t>({static_cast<int32_t>(NpuArch::DAV_3510),
879- static_cast<int32_t>(NpuArch::DAV_9201)}),879+ static_cast<int32_t>(NpuArch::DAV_9201),
880+ static_cast<int32_t>(NpuArch::DAV_9202)}),
880 0);881 0);
881} // namespace optiling882} // namespace optiling
Mmhc/mhc_post_backward/op_host/op_tiling/mhc_post_backward_tiling.cpp+1-1
@@ -29,7 +29,7 @@ static ge::graphStatus Tiling4MhcPostBackward(gert::TilingContext* context)
29 auto platformInfo = context->GetPlatformInfo();29 auto platformInfo = context->GetPlatformInfo();
30 platform_ascendc::PlatformAscendC ascendcPlatform(platformInfo);30 platform_ascendc::PlatformAscendC ascendcPlatform(platformInfo);
31 NpuArch npuArch = ascendcPlatform.GetCurNpuArch();31 NpuArch npuArch = ascendcPlatform.GetCurNpuArch();
32- if (npuArch == NpuArch::DAV_3510 || npuArch == NpuArch::DAV_9201) {32+ if (npuArch == NpuArch::DAV_3510 || npuArch == NpuArch::DAV_9201 || npuArch == NpuArch::DAV_9202) {
33 return TilingRegistryArch::GetInstance().DoTilingImpl(context);33 return TilingRegistryArch::GetInstance().DoTilingImpl(context);
34 }34 }
35 return TilingRegistryNew::GetInstance().DoTilingImpl(context);35 return TilingRegistryNew::GetInstance().DoTilingImpl(context);
Mmhc/mhc_pre_sinkhorn/op_host/CMakeLists.txt+11-0
@@ -33,6 +33,17 @@ if (BUILD_OPEN_PROJECT)
33 -mllvm -cce-aicore-dcci-before-kernel-end=false33 -mllvm -cce-aicore-dcci-before-kernel-end=false
34 --op_relocatable_kernel_binary=true34 --op_relocatable_kernel_binary=true
35 )35 )
36+ elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend960pr")
37+ add_ops_compile_options(
38+ OP_NAME MhcPreSinkhorn
39+ COMPUTE_UNIT Ascend960PR_969A1
40+ OPTIONS --cce-auto-sync=off
41+ -Wno-deprecated-declarations
42+ -Werror
43+ -mllvm -cce-aicore-hoist-movemask=false
44+ -mllvm -cce-aicore-dcci-before-kernel-end=false
45+ --op_relocatable_kernel_binary=true
46+ )
36 else()47 else()
37 add_ops_compile_options(48 add_ops_compile_options(
38 OP_NAME MhcPreSinkhorn49 OP_NAME MhcPreSinkhorn
Amhc/mhc_pre_sinkhorn/op_host/config/ascend960pr/mhc_pre_sinkhorn_binary.json+313-0
@@ -0,0 +1,313 @@
1+{
2+ "op_type": "MhcPreSinkhorn",
3+ "op_list": [
4+ {
5+ "bin_filename": "MhcPreSinkhorn_bd8e37c67f2294e1437f03bb2bace96c",
6+ "inputs": [
7+ {
8+ "name": "x",
9+ "index": 0,
10+ "dtype": "bfloat16",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [
14+ -2
15+ ]
16+ },
17+ {
18+ "name": "phi",
19+ "index": 1,
20+ "dtype": "float32",
21+ "format": "ND",
22+ "paramType": "required",
23+ "shape": [
24+ -2
25+ ]
26+ },
27+ {
28+ "name": "alpha",
29+ "index": 2,
30+ "dtype": "float32",
31+ "format": "ND",
32+ "paramType": "required",
33+ "shape": [
34+ -2
35+ ]
36+ },
37+ {
38+ "name": "bias",
39+ "index": 3,
40+ "dtype": "float32",
41+ "format": "ND",
42+ "paramType": "required",
43+ "shape": [
44+ -2
45+ ]
46+ }
47+ ],
48+ "outputs": [
49+ {
50+ "name": "hin",
51+ "index": 0,
52+ "dtype": "bfloat16",
53+ "format": "ND",
54+ "paramType": "required",
55+ "shape": [
56+ -2
57+ ]
58+ },
59+ {
60+ "name": "hPost",
61+ "index": 1,
62+ "dtype": "float32",
63+ "format": "ND",
64+ "paramType": "required",
65+ "shape": [
66+ -2
67+ ]
68+ },
69+ {
70+ "name": "hRes",
71+ "index": 2,
72+ "dtype": "float32",
73+ "format": "ND",
74+ "paramType": "required",
75+ "shape": [
76+ -2
77+ ]
78+ },
79+ {
80+ "name": "hPre",
81+ "index": 3,
82+ "dtype": "float32",
83+ "format": "ND",
84+ "paramType": "required",
85+ "shape": [
86+ -2
87+ ]
88+ },
89+ {
90+ "name": "hcBeforeNorm",
91+ "index": 4,
92+ "dtype": "float32",
93+ "format": "ND",
94+ "paramType": "required",
95+ "shape": [
96+ -2
97+ ]
98+ },
99+ {
100+ "name": "invRms",
101+ "index": 5,
102+ "dtype": "float32",
103+ "format": "ND",
104+ "paramType": "required",
105+ "shape": [
106+ -2
107+ ]
108+ },
109+ {
110+ "name": "sumOut",
111+ "index": 6,
112+ "dtype": "float32",
113+ "format": "ND",
114+ "paramType": "required",
115+ "shape": [
116+ -2
117+ ]
118+ },
119+ {
120+ "name": "normOut",
121+ "index": 7,
122+ "dtype": "float32",
123+ "format": "ND",
124+ "paramType": "required",
125+ "shape": [
126+ -2
127+ ]
128+ }
129+ ],
130+ "attrs": [
131+ {
132+ "name": "hc_mult",
133+ "dtype": "int",
134+ "value": 4
135+ },
136+ {
137+ "name": "num_iters",
138+ "dtype": "int",
139+ "value": 20
140+ },
141+ {
142+ "name": "hc_eps",
143+ "dtype": "float32",
144+ "value": 1e-6
145+ },
146+ {
147+ "name": "norm_eps",
148+ "dtype": "float32",
149+ "value": 1e-6
150+ },
151+ {
152+ "name": "need_backward",
153+ "dtype": "bool",
154+ "value": true
155+ }
156+ ]
157+ },
158+ {
159+ "bin_filename": "MhcPreSinkhorn_5d39354b089359785ccbee74e3ea8182",
160+ "inputs": [
161+ {
162+ "name": "x",
163+ "index": 0,
164+ "dtype": "float16",
165+ "format": "ND",
166+ "paramType": "required",
167+ "shape": [
168+ -2
169+ ]
170+ },
171+ {
172+ "name": "phi",
173+ "index": 1,
174+ "dtype": "float32",
175+ "format": "ND",
176+ "paramType": "required",
177+ "shape": [
178+ -2
179+ ]
180+ },
181+ {
182+ "name": "alpha",
183+ "index": 2,
184+ "dtype": "float32",
185+ "format": "ND",
186+ "paramType": "required",
187+ "shape": [
188+ -2
189+ ]
190+ },
191+ {
192+ "name": "bias",
193+ "index": 3,
194+ "dtype": "float32",
195+ "format": "ND",
196+ "paramType": "required",
197+ "shape": [
198+ -2
199+ ]
200+ }
201+ ],
202+ "outputs": [
203+ {
204+ "name": "hin",
205+ "index": 0,
206+ "dtype": "float16",
207+ "format": "ND",
208+ "paramType": "required",
209+ "shape": [
210+ -2
211+ ]
212+ },
213+ {
214+ "name": "hPost",
215+ "index": 1,
216+ "dtype": "float32",
217+ "format": "ND",
218+ "paramType": "required",
219+ "shape": [
220+ -2
221+ ]
222+ },
223+ {
224+ "name": "hRes",
225+ "index": 2,
226+ "dtype": "float32",
227+ "format": "ND",
228+ "paramType": "required",
229+ "shape": [
230+ -2
231+ ]
232+ },
233+ {
234+ "name": "hPre",
235+ "index": 3,
236+ "dtype": "float32",
237+ "format": "ND",
238+ "paramType": "required",
239+ "shape": [
240+ -2
241+ ]
242+ },
243+ {
244+ "name": "hcBeforeNorm",
245+ "index": 4,
246+ "dtype": "float32",
247+ "format": "ND",
248+ "paramType": "required",
249+ "shape": [
250+ -2
251+ ]
252+ },
253+ {
254+ "name": "invRms",
255+ "index": 5,
256+ "dtype": "float32",
257+ "format": "ND",
258+ "paramType": "required",
259+ "shape": [
260+ -2
261+ ]
262+ },
263+ {
264+ "name": "sumOut",
265+ "index": 6,
266+ "dtype": "float32",
267+ "format": "ND",
268+ "paramType": "required",
269+ "shape": [
270+ -2
271+ ]
272+ },
273+ {
274+ "name": "normOut",
275+ "index": 7,
276+ "dtype": "float32",
277+ "format": "ND",
278+ "paramType": "required",
279+ "shape": [
280+ -2
281+ ]
282+ }
283+ ],
284+ "attrs": [
285+ {
286+ "name": "hc_mult",
287+ "dtype": "int",
288+ "value": 4
289+ },
290+ {
291+ "name": "num_iters",
292+ "dtype": "int",
293+ "value": 20
294+ },
295+ {
296+ "name": "hc_eps",
297+ "dtype": "float32",
298+ "value": 1e-6
299+ },
300+ {
301+ "name": "norm_eps",
302+ "dtype": "float32",
303+ "value": 1e-6
304+ },
305+ {
306+ "name": "need_backward",
307+ "dtype": "bool",
308+ "value": true
309+ }
310+ ]
311+ }
312+ ]
313+}
Mmhc/mhc_pre_sinkhorn/op_host/mhc_pre_sinkhorn_def.cpp+1-0
@@ -97,6 +97,7 @@ public:
97 .ExtendCfgInfo("opFile.value", "mhc_pre_sinkhorn_apt");97 .ExtendCfgInfo("opFile.value", "mhc_pre_sinkhorn_apt");
98 this->AICore().AddConfig("ascend950", regbaseCfg);98 this->AICore().AddConfig("ascend950", regbaseCfg);
99 this->AICore().AddConfig("ascend960dt", regbaseCfg);99 this->AICore().AddConfig("ascend960dt", regbaseCfg);
100+ this->AICore().AddConfig("ascend960pr", regbaseCfg);
100 }101 }
101};102};
102 103 
Mmhc/mhc_pre_sinkhorn/op_host/op_api/aclnn_mhc_pre_sinkhorn.cpp+1-1
@@ -45,7 +45,7 @@ static constexpr int64_t NUM_ITERS_DEFAULT_VALUE = 20;
45static bool IsPlatformSupportTnd()45static bool IsPlatformSupportTnd()
46{46{
47 auto npuArch = GetCurrentPlatformInfo().GetCurNpuArch();47 auto npuArch = GetCurrentPlatformInfo().GetCurNpuArch();
48- return (npuArch == NpuArch::DAV_3510) || (npuArch == NpuArch::DAV_9201);48+ return (npuArch == NpuArch::DAV_3510) || (npuArch == NpuArch::DAV_9201 || npuArch == NpuArch::DAV_9202);
49}49}
50 50 
51static bool CheckNotNull(const aclTensor* x, const aclTensor* phi, const aclTensor* alpha, const aclTensor* bias,51static bool CheckNotNull(const aclTensor* x, const aclTensor* phi, const aclTensor* alpha, const aclTensor* bias,
Mmhc/mhc_pre_sinkhorn_backward/op_host/CMakeLists.txt+11-0
@@ -30,6 +30,17 @@ if (BUILD_OPEN_PROJECT)
30 -mllvm -cce-aicore-dcci-before-kernel-end=false30 -mllvm -cce-aicore-dcci-before-kernel-end=false
31 --op_relocatable_kernel_binary=true31 --op_relocatable_kernel_binary=true
32 )32 )
33+ elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend960pr")
34+ add_ops_compile_options(
35+ OP_NAME MhcPreSinkhornBackward
36+ COMPUTE_UNIT Ascend960PR_969A1
37+ OPTIONS --cce-auto-sync=off
38+ -Wno-deprecated-declarations
39+ -Werror
40+ -mllvm -cce-aicore-hoist-movemask=false
41+ -mllvm -cce-aicore-dcci-before-kernel-end=false
42+ --op_relocatable_kernel_binary=true
43+ )
33 else()44 else()
34 add_ops_compile_options(45 add_ops_compile_options(
35 OP_NAME MhcPreSinkhornBackward46 OP_NAME MhcPreSinkhornBackward
Amhc/mhc_pre_sinkhorn_backward/op_host/config/ascend960pr/mhc_pre_sinkhorn_backward_binary.json+353-0
@@ -0,0 +1,353 @@
1+{
2+ "op_type": "MhcPreSinkhornBackward",
3+ "op_list": [
4+ {
5+ "bin_filename": "MhcPreSinkhornBackward_bf16",
6+ "inputs": [
7+ {
8+ "name": "grad_hin",
9+ "index": 0,
10+ "dtype": "bfloat16",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [
14+ -2
15+ ]
16+ },
17+ {
18+ "name": "grad_h_post",
19+ "index": 1,
20+ "dtype": "float32",
21+ "format": "ND",
22+ "paramType": "required",
23+ "shape": [
24+ -2
25+ ]
26+ },
27+ {
28+ "name": "grad_h_res",
29+ "index": 2,
30+ "dtype": "float32",
31+ "format": "ND",
32+ "paramType": "required",
33+ "shape": [
34+ -2
35+ ]
36+ },
37+ {
38+ "name": "x",
39+ "index": 3,
40+ "dtype": "bfloat16",
41+ "format": "ND",
42+ "paramType": "required",
43+ "shape": [
44+ -2
45+ ]
46+ },
47+ {
48+ "name": "phi",
49+ "index": 4,
50+ "dtype": "float32",
51+ "format": "ND",
52+ "paramType": "required",
53+ "shape": [
54+ -2
55+ ]
56+ },
57+ {
58+ "name": "alpha",
59+ "index": 5,
60+ "dtype": "float32",
61+ "format": "ND",
62+ "paramType": "required",
63+ "shape": [
64+ -2
65+ ]
66+ },
67+ {
68+ "name": "bias",
69+ "index": 6,
70+ "dtype": "float32",
71+ "format": "ND",
72+ "paramType": "required",
73+ "shape": [
74+ -2
75+ ]
76+ },
77+ {
78+ "name": "h_pre",
79+ "index": 7,
80+ "dtype": "float32",
81+ "format": "ND",
82+ "paramType": "required",
83+ "shape": [
84+ -2
85+ ]
86+ },
87+ {
88+ "name": "hc_before_norm",
89+ "index": 8,
90+ "dtype": "float32",
91+ "format": "ND",
92+ "paramType": "required",
93+ "shape": [
94+ -2
95+ ]
96+ },
97+ {
98+ "name": "inv_rms",
99+ "index": 9,
100+ "dtype": "float32",
101+ "format": "ND",
102+ "paramType": "required",
103+ "shape": [
104+ -2
105+ ]
106+ },
107+ {
108+ "name": "sum_out",
109+ "index": 10,
110+ "dtype": "float32",
111+ "format": "ND",
112+ "paramType": "required",
113+ "shape": [
114+ -2
115+ ]
116+ },
117+ {
118+ "name": "norm_out",
119+ "index": 11,
120+ "dtype": "float32",
121+ "format": "ND",
122+ "paramType": "required",
123+ "shape": [
124+ -2
125+ ]
126+ }
127+ ],
128+ "outputs": [
129+ {
130+ "name": "grad_x",
131+ "index": 0,
132+ "dtype": "bfloat16",
133+ "format": "ND",
134+ "paramType": "required",
135+ "shape": [
136+ -2
137+ ]
138+ },
139+ {
140+ "name": "grad_phi",
141+ "index": 1,
142+ "dtype": "float32",
143+ "format": "ND",
144+ "paramType": "required",
145+ "shape": [
146+ -2
147+ ]
148+ },
149+ {
150+ "name": "grad_alpha",
151+ "index": 2,
152+ "dtype": "float32",
153+ "format": "ND",
154+ "paramType": "required",
155+ "shape": [
156+ -2
157+ ]
158+ },
159+ {
160+ "name": "grad_bias",
161+ "index": 3,
162+ "dtype": "float32",
163+ "format": "ND",
164+ "paramType": "required",
165+ "shape": [
166+ -2
167+ ]
168+ }
169+ ],
170+ "attrs": [
171+ {
172+ "name": "hc_eps",
173+ "dtype": "float32",
174+ "value": 1e-6
175+ }
176+ ]
177+ },
178+ {
179+ "bin_filename": "MhcPreSinkhornBackward_fp16",
180+ "inputs": [
181+ {
182+ "name": "grad_hin",
183+ "index": 0,
184+ "dtype": "float16",
185+ "format": "ND",
186+ "paramType": "required",
187+ "shape": [
188+ -2
189+ ]
190+ },
191+ {
192+ "name": "grad_h_post",
193+ "index": 1,
194+ "dtype": "float32",
195+ "format": "ND",
196+ "paramType": "required",
197+ "shape": [
198+ -2
199+ ]
200+ },
201+ {
202+ "name": "grad_h_res",
203+ "index": 2,
204+ "dtype": "float32",
205+ "format": "ND",
206+ "paramType": "required",
207+ "shape": [
208+ -2
209+ ]
210+ },
211+ {
212+ "name": "x",
213+ "index": 3,
214+ "dtype": "float16",
215+ "format": "ND",
216+ "paramType": "required",
217+ "shape": [
218+ -2
219+ ]
220+ },
221+ {
222+ "name": "phi",
223+ "index": 4,
224+ "dtype": "float32",
225+ "format": "ND",
226+ "paramType": "required",
227+ "shape": [
228+ -2
229+ ]
230+ },
231+ {
232+ "name": "alpha",
233+ "index": 5,
234+ "dtype": "float32",
235+ "format": "ND",
236+ "paramType": "required",
237+ "shape": [
238+ -2
239+ ]
240+ },
241+ {
242+ "name": "bias",
243+ "index": 6,
244+ "dtype": "float32",
245+ "format": "ND",
246+ "paramType": "required",
247+ "shape": [
248+ -2
249+ ]
250+ },
251+ {
252+ "name": "h_pre",
253+ "index": 7,
254+ "dtype": "float32",
255+ "format": "ND",
256+ "paramType": "required",
257+ "shape": [
258+ -2
259+ ]
260+ },
261+ {
262+ "name": "hc_before_norm",
263+ "index": 8,
264+ "dtype": "float32",
265+ "format": "ND",
266+ "paramType": "required",
267+ "shape": [
268+ -2
269+ ]
270+ },
271+ {
272+ "name": "inv_rms",
273+ "index": 9,
274+ "dtype": "float32",
275+ "format": "ND",
276+ "paramType": "required",
277+ "shape": [
278+ -2
279+ ]
280+ },
281+ {
282+ "name": "sum_out",
283+ "index": 10,
284+ "dtype": "float32",
285+ "format": "ND",
286+ "paramType": "required",
287+ "shape": [
288+ -2
289+ ]
290+ },
291+ {
292+ "name": "norm_out",
293+ "index": 11,
294+ "dtype": "float32",
295+ "format": "ND",
296+ "paramType": "required",
297+ "shape": [
298+ -2
299+ ]
300+ }
301+ ],
302+ "outputs": [
303+ {
304+ "name": "grad_x",
305+ "index": 0,
306+ "dtype": "float16",
307+ "format": "ND",
308+ "paramType": "required",
309+ "shape": [
310+ -2
311+ ]
312+ },
313+ {
314+ "name": "grad_phi",
315+ "index": 1,
316+ "dtype": "float32",
317+ "format": "ND",
318+ "paramType": "required",
319+ "shape": [
320+ -2
321+ ]
322+ },
323+ {
324+ "name": "grad_alpha",
325+ "index": 2,
326+ "dtype": "float32",
327+ "format": "ND",
328+ "paramType": "required",
329+ "shape": [
330+ -2
331+ ]
332+ },
333+ {
334+ "name": "grad_bias",
335+ "index": 3,
336+ "dtype": "float32",
337+ "format": "ND",
338+ "paramType": "required",
339+ "shape": [
340+ -2
341+ ]
342+ }
343+ ],
344+ "attrs": [
345+ {
346+ "name": "hc_eps",
347+ "dtype": "float32",
348+ "value": 1e-6
349+ }
350+ ]
351+ }
352+ ]
353+}
Mmhc/mhc_pre_sinkhorn_backward/op_host/mhc_pre_sinkhorn_backward_def.cpp+1-0
@@ -123,6 +123,7 @@ public:
123 .ExtendCfgInfo("opFile.value", "mhc_pre_sinkhorn_backward_apt");123 .ExtendCfgInfo("opFile.value", "mhc_pre_sinkhorn_backward_apt");
124 this->AICore().AddConfig("ascend950", aicore_Config_950);124 this->AICore().AddConfig("ascend950", aicore_Config_950);
125 this->AICore().AddConfig("ascend960dt", aicore_Config_950);125 this->AICore().AddConfig("ascend960dt", aicore_Config_950);
126+ this->AICore().AddConfig("ascend960pr", aicore_Config_950);
126 }127 }
127};128};
128OP_ADD(MhcPreSinkhornBackward);129OP_ADD(MhcPreSinkhornBackward);
Mmhc/mhc_pre_sinkhorn_backward/op_host/op_tiling/arch92/mhc_pre_sinkhorn_backward_arch92_tiling_base.cpp+1-1
@@ -137,7 +137,7 @@ ge::graphStatus MhcPreSinkhornBackward960Tiling::GetPlatformInfo()
137 OP_CHECK_IF((coreNumAic <= 0), OP_LOGE(opName, "MhcPreSinkhornBackwardTiling fails to get coreNumAic."),137 OP_CHECK_IF((coreNumAic <= 0), OP_LOGE(opName, "MhcPreSinkhornBackwardTiling fails to get coreNumAic."),
138 return ge::GRAPH_FAILED);138 return ge::GRAPH_FAILED);
139 coreNumAic_ = coreNumAic;139 coreNumAic_ = coreNumAic;
140- // Ascend960DT 的 AIC:AIV 为 1:1（KERNEL_TYPE_MIX_AIC_1_1），AIV 核数按 AIC 核数对齐140+ // Ascend960系列(960DT/960PR) 的 AIC:AIV 为 1:1（KERNEL_TYPE_MIX_AIC_1_1），AIV 核数按 AIC 核数对齐
141 if (ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_9201 || ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_9202) {141 if (ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_9201 || ascendcPlatform.GetCurNpuArch() == NpuArch::DAV_9202) {
142 coreNumAiv_ = coreNumAic_;142 coreNumAiv_ = coreNumAic_;
143 }143 }
Mmhc/mhc_pre_sinkhorn_backward/op_kernel/arch92/mhc_pre_sinkhorn_backward_entry_arch92.h+2-1
@@ -37,7 +37,8 @@ extern "C" __global__ __aicore__ void mhc_pre_sinkhorn_backward(
37 }37 }
38 38 
39 REGISTER_TILING_DEFAULT(MhcPreSinkhornBackward960TilingData);39 REGISTER_TILING_DEFAULT(MhcPreSinkhornBackward960TilingData);
40-#if defined(__DAV_920R1__) || (defined(__NPU_ARCH__) && (__NPU_ARCH__ == 9201))40+#if defined(__DAV_920R1__) || defined(__DAV_920R2__) || \
41+ (defined(__NPU_ARCH__) && (__NPU_ARCH__ == 9201 || __NPU_ARCH__ == 9202))
41 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_1);42 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_1);
42#else43#else
43 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);44 KERNEL_TASK_TYPE_DEFAULT(KERNEL_TYPE_MIX_AIC_1_2);
Mmhc/mhc_pre_sinkhorn_backward/op_kernel/arch92/mhc_pre_sinkhorn_backward_one_stage.h+4-2
@@ -126,7 +126,8 @@ private:
126 batchSize_ = tilingData->batchSize;126 batchSize_ = tilingData->batchSize;
127 seqLength_ = tilingData->seqLength;127 seqLength_ = tilingData->seqLength;
128 aivNum_ = tilingData->aivNum;128 aivNum_ = tilingData->aivNum;
129-#if defined(__DAV_920R1__) || (defined(__NPU_ARCH__) && (__NPU_ARCH__ == 9201))129+#if defined(__DAV_920R1__) || defined(__DAV_920R2__) || \
130+ (defined(__NPU_ARCH__) && (__NPU_ARCH__ == 9201 || __NPU_ARCH__ == 9202))
130 aicNum_ = tilingData->aivNum;131 aicNum_ = tilingData->aivNum;
131#else132#else
132 aicNum_ = tilingData->aivNum / 2;133 aicNum_ = tilingData->aivNum / 2;
@@ -493,7 +494,8 @@ __aicore__ inline void MhcPreSinkhornBackwardOneStage<T, U>::Process()
493 494 
494 if ASCEND_IS_AIC {495 if ASCEND_IS_AIC {
495 int8_t ping = 0;496 int8_t ping = 0;
496-#if defined(__DAV_920R1__) || (defined(__NPU_ARCH__) && (__NPU_ARCH__ == 9201))497+#if defined(__DAV_920R1__) || defined(__DAV_920R2__) || \
498+ (defined(__NPU_ARCH__) && (__NPU_ARCH__ == 9201 || __NPU_ARCH__ == 9202))
497 // Ascend960 (MIX_AIC_1_1, AIC:AIV = 1:1)：AIC i 与 AIV i 一一配对，499 // Ascend960 (MIX_AIC_1_1, AIC:AIV = 1:1)：AIC i 与 AIV i 一一配对，
498 // AIV i 第 k 轮产出段 [i+kA, i+kA+1)*tile_，AIC i 第 k 轮消费同段，500 // AIV i 第 k 轮产出段 [i+kA, i+kA+1)*tile_，AIC i 第 k 轮消费同段，
499 // 逐轮锁步、单 tile 消费，避免 AIC_1_2 的"每迭代聚 2 tile"几何501 // 逐轮锁步、单 tile 消费，避免 AIC_1_2 的"每迭代聚 2 tile"几何
Mmhc/mhc_sinkhorn/op_host/CMakeLists.txt+6-0
@@ -25,6 +25,12 @@ if (BUILD_OPEN_PROJECT)
25 COMPUTE_UNIT Ascend960DT_968B725 COMPUTE_UNIT Ascend960DT_968B7
26 OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false26 OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
27 )27 )
28+ elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend960pr")
29+ add_ops_compile_options(
30+ OP_NAME MhcSinkhorn
31+ COMPUTE_UNIT Ascend960PR_969A1
32+ OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
33+ )
28 elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend350")34 elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend350")
29 add_ops_compile_options(35 add_ops_compile_options(
30 OP_NAME MhcSinkhorn36 OP_NAME MhcSinkhorn
Amhc/mhc_sinkhorn/op_host/config/ascend960pr/mhc_sinkhorn_binary.json+69-0
@@ -0,0 +1,69 @@
1+{
2+ "op_type": "MhcSinkhorn",
3+ "op_list": [
4+ {
5+ "bin_filename": "MhcSinkhorn_8460c4b727ab5cc93a9595333a3fd571",
6+ "inputs": [
7+ {
8+ "name": "h_res",
9+ "index": 0,
10+ "dtype": "float32",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [
14+ -2
15+ ]
16+ }
17+ ],
18+ "outputs": [
19+ {
20+ "name": "y",
21+ "index": 0,
22+ "dtype": "float32",
23+ "format": "ND",
24+ "paramType": "required",
25+ "shape": [
26+ -2
27+ ]
28+ },
29+ {
30+ "name": "norm_out",
31+ "index": 1,
32+ "dtype": "float32",
33+ "format": "ND",
34+ "paramType": "required",
35+ "shape": [
36+ -2
37+ ]
38+ },
39+ {
40+ "name": "sum_out",
41+ "index": 2,
42+ "dtype": "float32",
43+ "format": "ND",
44+ "paramType": "required",
45+ "shape": [
46+ -2
47+ ]
48+ }
49+ ],
50+ "attrs": [
51+ {
52+ "name": "eps",
53+ "dtype": "float32",
54+ "value": 1E-6
55+ },
56+ {
57+ "name": "num_iters",
58+ "dtype": "int",
59+ "value": 20
60+ },
61+ {
62+ "name": "out_flag",
63+ "dtype": "int",
64+ "value": 0
65+ }
66+ ]
67+ }
68+ ]
69+ }
Mmhc/mhc_sinkhorn/op_host/mhc_sinkhorn_def.cpp+1-0
@@ -58,6 +58,7 @@ public:
58 .ExtendCfgInfo("opFile.value", "mhc_sinkhorn_apt");58 .ExtendCfgInfo("opFile.value", "mhc_sinkhorn_apt");
59 this->AICore().AddConfig("ascend950", aicore_config);59 this->AICore().AddConfig("ascend950", aicore_config);
60 this->AICore().AddConfig("ascend960dt", aicore_config);60 this->AICore().AddConfig("ascend960dt", aicore_config);
61+ this->AICore().AddConfig("ascend960pr", aicore_config);
61 this->AICore().AddConfig("ascend350", aicore_config);62 this->AICore().AddConfig("ascend350", aicore_config);
62 }63 }
63};64};
Rmhc/mhc_sinkhorn/test/CMakeLists.txt→mhc/mhc_sinkhorn/tests/CMakeLists.txt+0-0
File renamed without changes.
Rmhc/mhc_sinkhorn/test/ut/CMakeLists.txt→mhc/mhc_sinkhorn/tests/ut/CMakeLists.txt+0-0
File renamed without changes.
Rmhc/mhc_sinkhorn/test/ut/op_host/CMakeLists.txt→mhc/mhc_sinkhorn/tests/ut/op_host/CMakeLists.txt+0-0
File renamed without changes.
Rmhc/mhc_sinkhorn/test/ut/op_host/op_api/CMakeLists.txt→mhc/mhc_sinkhorn/tests/ut/op_host/op_api/CMakeLists.txt+0-0
File renamed without changes.
Rmhc/mhc_sinkhorn/test/ut/op_host/op_api/test_aclnn_mhc_sinkhorn.cpp→mhc/mhc_sinkhorn/tests/ut/op_host/op_api/test_aclnn_mhc_sinkhorn.cpp+0-0
File renamed without changes.
Rmhc/mhc_sinkhorn/test/ut/op_host/test_mhc_sinkhorn_infershape.cpp→mhc/mhc_sinkhorn/tests/ut/op_host/test_mhc_sinkhorn_infershape.cpp+0-0
File renamed without changes.
Rmhc/mhc_sinkhorn/test/ut/op_host/test_mhc_sinkhorn_tiling.cpp→mhc/mhc_sinkhorn/tests/ut/op_host/test_mhc_sinkhorn_tiling.cpp+163-30
@@ -34,32 +34,31 @@ struct MhcSinkhornCompileInfo {
34 int64_t ubSize;34 int64_t ubSize;
35};35};
36 36 
37-TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_4d_fp16_success)37+// ==================== 950 用例 ====================
38+ 
39+TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_4d_fp16_950_success)
38{40{
39 MhcSinkhornCompileInfo compileInfo = {64, 64};41 MhcSinkhornCompileInfo compileInfo = {64, 64};
40 gert::TilingContextPara tilingContextPara("MhcSinkhorn",42 gert::TilingContextPara tilingContextPara("MhcSinkhorn",
41 {43 {
42- // input info: h_res
43 {{{1, 128, 4, 4}, {1, 128, 4, 4}}, ge::DT_FLOAT, ge::FORMAT_ND},44 {{{1, 128, 4, 4}, {1, 128, 4, 4}}, ge::DT_FLOAT, ge::FORMAT_ND},
44 },45 },
45 {46 {
46- // output info
47 {{{1, 128, 4, 4}, {1024, 4, 5120}}, ge::DT_FLOAT, ge::FORMAT_ND},47 {{{1, 128, 4, 4}, {1024, 4, 5120}}, ge::DT_FLOAT, ge::FORMAT_ND},
48 },48 },
49 {49 {
50- // attr
51 {"eps", Ops::Transformer::AnyValue::CreateFrom<float>(1e-06)},50 {"eps", Ops::Transformer::AnyValue::CreateFrom<float>(1e-06)},
52 {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},51 {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},
53 {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},52 {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
54 },53 },
55- &compileInfo);54+ &compileInfo, "Ascend950");
56 uint64_t expectTilingKey = 1;55 uint64_t expectTilingKey = 1;
57 string expectTilingDataStr = "4 5120 64 16 16 1 1024 1 5120 1 5120 5120 ";56 string expectTilingDataStr = "4 5120 64 16 16 1 1024 1 5120 1 5120 5120 ";
58 std::vector<size_t> expectWorkspaces = {16777216};57 std::vector<size_t> expectWorkspaces = {16777216};
59 ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingDataStr, expectWorkspaces);58 ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingDataStr, expectWorkspaces);
60}59}
61 60 
62-TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_3d_fp16_success)61+TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_3d_fp16_950_success)
63{62{
64 MhcSinkhornCompileInfo compileInfo = {64, 64};63 MhcSinkhornCompileInfo compileInfo = {64, 64};
65 gert::TilingContextPara tilingContextPara("MhcSinkhorn",64 gert::TilingContextPara tilingContextPara("MhcSinkhorn",
@@ -74,54 +73,188 @@ TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_3d_fp16_success)
74 {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},73 {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},
75 {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},74 {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
76 },75 },
77- &compileInfo);76+ &compileInfo, "Ascend950");
78 uint64_t expectTilingKey = 1;77 uint64_t expectTilingKey = 1;
79 string expectTilingDataStr = "4 5120 64 16 16 1 1024 1 5120 1 5120 5120 ";78 string expectTilingDataStr = "4 5120 64 16 16 1 1024 1 5120 1 5120 5120 ";
80 std::vector<size_t> expectWorkspaces = {16777216};79 std::vector<size_t> expectWorkspaces = {16777216};
81 ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingDataStr, expectWorkspaces);80 ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingDataStr, expectWorkspaces);
82}81}
83 82 
84-TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_invalid_n)83+TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_4d_fp32_950_success)
85{84{
86 MhcSinkhornCompileInfo compileInfo = {64, 64};85 MhcSinkhornCompileInfo compileInfo = {64, 64};
87 gert::TilingContextPara tilingContextPara("MhcSinkhorn",86 gert::TilingContextPara tilingContextPara("MhcSinkhorn",
88 {87 {
89- {{{1024, 5, 5}, {1024, 5, 5}}, ge::DT_FLOAT, ge::FORMAT_ND},88+ {{{1, 128, 4, 4}, {1, 128, 4, 4}}, ge::DT_FLOAT, ge::FORMAT_ND},
90 },89 },
91 {90 {
92- {{{1024, 5, 5}, {1024, 5, 5}}, ge::DT_FLOAT, ge::FORMAT_ND},91+ {{{1, 128, 4, 4}, {1024, 4, 5120}}, ge::DT_FLOAT, ge::FORMAT_ND},
93 },92 },
94 {93 {
95 {"eps", Ops::Transformer::AnyValue::CreateFrom<float>(1e-06)},94 {"eps", Ops::Transformer::AnyValue::CreateFrom<float>(1e-06)},
96 {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},95 {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},
97 {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},96 {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
98 },97 },
99- &compileInfo);98+ &compileInfo, "Ascend950");
100 uint64_t expectTilingKey = 1;99 uint64_t expectTilingKey = 1;
101- string expectTilingDataStr = "7 2048 64 8 8 1 512 1 2048 1 2048 2048 ";100+ string expectTilingDataStr = "4 5120 64 16 16 1 1024 1 5120 1 5120 5120 ";
102 std::vector<size_t> expectWorkspaces = {16777216};101 std::vector<size_t> expectWorkspaces = {16777216};
103- ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED, expectTilingKey, expectTilingDataStr, expectWorkspaces);102+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingDataStr, expectWorkspaces);
104}103}
105 104 
106-TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_invalid_shape)105+// ==================== 960PR 用例 ====================
106+ 
107+TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_4d_fp16_960pr_success)
107{108{
108 MhcSinkhornCompileInfo compileInfo = {64, 64};109 MhcSinkhornCompileInfo compileInfo = {64, 64};
109- gert::TilingContextPara tilingContextPara(110+ gert::TilingContextPara tilingContextPara("MhcSinkhorn",
110- "MhcSinkhorn",111+ {
111- {112+ {{{1, 128, 4, 4}, {1, 128, 4, 4}}, ge::DT_FLOAT, ge::FORMAT_ND},
112- {{{1, 2, 1024, 5, 5}, {1, 2, 1024, 5, 5}}, ge::DT_FLOAT, ge::FORMAT_ND},113+ },
113- },114+ {
114- {115+ {{{1, 128, 4, 4}, {1024, 4, 5120}}, ge::DT_FLOAT, ge::FORMAT_ND},
115- {{{1, 2, 1024, 5, 5}, {1, 2, 1024, 5, 5}}, ge::DT_FLOAT, ge::FORMAT_ND},116+ },
116- },117+ {
117- {118+ {"eps", Ops::Transformer::AnyValue::CreateFrom<float>(1e-06)},
118- {"eps", Ops::Transformer::AnyValue::CreateFrom<float>(1e-06)},119+ {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},
119- {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},120+ {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
120- {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},121+ },
121- },122+ &compileInfo, "Ascend960PR");
122- &compileInfo);
123 uint64_t expectTilingKey = 1;123 uint64_t expectTilingKey = 1;
124- string expectTilingDataStr = "7 2048 64 8 8 1 512 1 2048 1 2048 2048 ";124+ string expectTilingDataStr = "4 5120 64 16 16 1 1024 1 5120 1 5120 5120 ";
125 std::vector<size_t> expectWorkspaces = {16777216};125 std::vector<size_t> expectWorkspaces = {16777216};
126- ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED, expectTilingKey, expectTilingDataStr, expectWorkspaces);126+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingDataStr, expectWorkspaces);
127+}
128+ 
129+TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_3d_fp16_960pr_success)
130+{
131+ MhcSinkhornCompileInfo compileInfo = {64, 64};
132+ gert::TilingContextPara tilingContextPara("MhcSinkhorn",
133+ {
134+ {{{1024, 6, 6}, {1024, 6, 6}}, ge::DT_FLOAT, ge::FORMAT_ND},
135+ },
136+ {
137+ {{{1024, 6, 6}, {1024, 6, 6}}, ge::DT_FLOAT, ge::FORMAT_ND},
138+ },
139+ {
140+ {"eps", Ops::Transformer::AnyValue::CreateFrom<float>(1e-06)},
141+ {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},
142+ {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
143+ },
144+ &compileInfo, "Ascend960PR");
145+ uint64_t expectTilingKey = 1;
146+ string expectTilingDataStr = "4 5120 64 16 16 1 1024 1 5120 1 5120 5120 ";
147+ std::vector<size_t> expectWorkspaces = {16777216};
148+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingDataStr, expectWorkspaces);
149+}
150+ 
151+TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_4d_fp32_960pr_success)
152+{
153+ MhcSinkhornCompileInfo compileInfo = {64, 64};
154+ gert::TilingContextPara tilingContextPara("MhcSinkhorn",
155+ {
156+ {{{1, 128, 4, 4}, {1, 128, 4, 4}}, ge::DT_FLOAT, ge::FORMAT_ND},
157+ },
158+ {
159+ {{{1, 128, 4, 4}, {1024, 4, 5120}}, ge::DT_FLOAT, ge::FORMAT_ND},
160+ },
161+ {
162+ {"eps", Ops::Transformer::AnyValue::CreateFrom<float>(1e-06)},
163+ {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},
164+ {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
165+ },
166+ &compileInfo, "Ascend960PR");
167+ uint64_t expectTilingKey = 1;
168+ string expectTilingDataStr = "4 5120 64 16 16 1 1024 1 5120 1 5120 5120 ";
169+ std::vector<size_t> expectWorkspaces = {16777216};
170+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingDataStr, expectWorkspaces);
171+}
172+ 
173+// ==================== 950 异常/边界用例 ====================
174+ 
175+TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_null_input_950_failed)
176+{
177+ MhcSinkhornCompileInfo compileInfo = {64, 64};
178+ gert::TilingContextPara tilingContextPara("MhcSinkhorn",
179+ {
180+ // 空输入
181+ },
182+ {
183+ {{{1, 128, 4, 4}, {1024, 4, 5120}}, ge::DT_FLOAT, ge::FORMAT_ND},
184+ },
185+ {
186+ {"eps", Ops::Transformer::AnyValue::CreateFrom<float>(1e-06)},
187+ {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},
188+ {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
189+ },
190+ &compileInfo, "Ascend950");
191+ ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);
192+}
193+ 
194+// ==================== 960DT 用例 ====================
195+ 
196+TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_4d_fp16_960dt_success)
197+{
198+ MhcSinkhornCompileInfo compileInfo = {64, 64};
199+ gert::TilingContextPara tilingContextPara("MhcSinkhorn",
200+ {
201+ {{{1, 128, 4, 4}, {1, 128, 4, 4}}, ge::DT_FLOAT, ge::FORMAT_ND},
202+ },
203+ {
204+ {{{1, 128, 4, 4}, {1024, 4, 5120}}, ge::DT_FLOAT, ge::FORMAT_ND},
205+ },
206+ {
207+ {"eps", Ops::Transformer::AnyValue::CreateFrom<float>(1e-06)},
208+ {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},
209+ {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
210+ },
211+ &compileInfo, "Ascend960DT");
212+ uint64_t expectTilingKey = 1;
213+ string expectTilingDataStr = "4 5120 64 16 16 1 1024 1 5120 1 5120 5120 ";
214+ std::vector<size_t> expectWorkspaces = {16777216};
215+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingDataStr, expectWorkspaces);
216+}
217+ 
218+TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_3d_fp16_960dt_success)
219+{
220+ MhcSinkhornCompileInfo compileInfo = {64, 64};
221+ gert::TilingContextPara tilingContextPara("MhcSinkhorn",
222+ {
223+ {{{1024, 6, 6}, {1024, 6, 6}}, ge::DT_FLOAT, ge::FORMAT_ND},
224+ },
225+ {
226+ {{{1024, 6, 6}, {1024, 6, 6}}, ge::DT_FLOAT, ge::FORMAT_ND},
227+ },
228+ {
229+ {"eps", Ops::Transformer::AnyValue::CreateFrom<float>(1e-06)},
230+ {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},
231+ {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
232+ },
233+ &compileInfo, "Ascend960DT");
234+ uint64_t expectTilingKey = 1;
235+ string expectTilingDataStr = "4 5120 64 16 16 1 1024 1 5120 1 5120 5120 ";
236+ std::vector<size_t> expectWorkspaces = {16777216};
237+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingDataStr, expectWorkspaces);
238+}
239+ 
240+TEST_F(MhcSinkhornTiling, test_mhc_sinkhorn_4d_fp32_960dt_success)
241+{
242+ MhcSinkhornCompileInfo compileInfo = {64, 64};
243+ gert::TilingContextPara tilingContextPara("MhcSinkhorn",
244+ {
245+ {{{1, 128, 4, 4}, {1, 128, 4, 4}}, ge::DT_FLOAT, ge::FORMAT_ND},
246+ },
247+ {
248+ {{{1, 128, 4, 4}, {1024, 4, 5120}}, ge::DT_FLOAT, ge::FORMAT_ND},
249+ },
250+ {
251+ {"eps", Ops::Transformer::AnyValue::CreateFrom<float>(1e-06)},
252+ {"num_iters", Ops::Transformer::AnyValue::CreateFrom<int64_t>(20)},
253+ {"out_flag", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
254+ },
255+ &compileInfo, "Ascend960DT");
256+ uint64_t expectTilingKey = 1;
257+ string expectTilingDataStr = "4 5120 64 16 16 1 1024 1 5120 1 5120 5120 ";
258+ std::vector<size_t> expectWorkspaces = {16777216};
259+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS, expectTilingKey, expectTilingDataStr, expectWorkspaces);
127}260}
Mmhc/mhc_sinkhorn_backward/op_host/CMakeLists.txt+6-0
@@ -25,6 +25,12 @@ if (BUILD_OPEN_PROJECT)
25 COMPUTE_UNIT Ascend960DT_968B725 COMPUTE_UNIT Ascend960DT_968B7
26 OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false26 OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
27 )27 )
28+ elseif("${ASCEND_COMPUTE_UNIT}" STREQUAL "ascend960pr")
29+ add_ops_compile_options(
30+ OP_NAME MhcSinkhornBackward
31+ COMPUTE_UNIT Ascend960PR_969A1
32+ OPTIONS -mllvm -cce-aicore-dcci-before-kernel-end=false
33+ )
28 else()34 else()
29 add_ops_compile_options(35 add_ops_compile_options(
30 OP_NAME MhcSinkhornBackward36 OP_NAME MhcSinkhornBackward
Amhc/mhc_sinkhorn_backward/op_host/config/ascend960pr/mhc_sinkhorn_backward.json+60-0
@@ -0,0 +1,60 @@
1+{
2+ "op_type": "MhcSinkhornBackward",
3+ "op_list": [
4+ {
5+ "bin_filename": "MhcSinkhornBackward_cd345a94623931ac8fbf9b085a5f740d",
6+ "inputs": [
7+ {
8+ "name": "grad_y",
9+ "index": 0,
10+ "dtype": "float32",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [
14+ -2
15+ ],
16+ "format_match_mode": "FormatAgnostic",
17+ "dtype_match_mode": "DtypeByte"
18+ },
19+ {
20+ "name": "norm",
21+ "index": 1,
22+ "dtype": "float32",
23+ "format": "ND",
24+ "paramType": "required",
25+ "shape": [
26+ -2
27+ ],
28+ "format_match_mode": "FormatAgnostic",
29+ "dtype_match_mode": "DtypeByte"
30+ },
31+ {
32+ "name": "sum",
33+ "index": 2,
34+ "dtype": "float32",
35+ "format": "ND",
36+ "paramType": "required",
37+ "shape": [
38+ -2
39+ ],
40+ "format_match_mode": "FormatAgnostic",
41+ "dtype_match_mode": "DtypeByte"
42+ }
43+ ],
44+ "outputs": [
45+ {
46+ "name": "grad_input",
47+ "index": 0,
48+ "dtype": "float32",
49+ "format": "ND",
50+ "paramType": "required",
51+ "shape": [
52+ -2
53+ ],
54+ "format_match_mode": "FormatAgnostic",
55+ "dtype_match_mode": "DtypeByte"
56+ }
57+ ]
58+ }
59+ ]
60+ }
Mmhc/mhc_sinkhorn_backward/op_host/mhc_sinkhorn_backward_def.cpp+1-0
@@ -54,6 +54,7 @@ public:
54 .ExtendCfgInfo("opFile.value", "mhc_sinkhorn_backward_apt");54 .ExtendCfgInfo("opFile.value", "mhc_sinkhorn_backward_apt");
55 this->AICore().AddConfig("ascend950", aicore_config);55 this->AICore().AddConfig("ascend950", aicore_config);
56 this->AICore().AddConfig("ascend960dt", aicore_config);56 this->AICore().AddConfig("ascend960dt", aicore_config);
57+ this->AICore().AddConfig("ascend960pr", aicore_config);
57 }58 }
58};59};
59 60 
Dmoe/moe_compute_expert_tokens/op_host/config/ascend960pr/moe_compute_expert_tokens_simplified_key.ini+0-13
@@ -1,13 +0,0 @@
1-; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
2-; [某算子]
3-; default=xx
4-; ascendxx=xx
5-; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6-; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7-; 2)如果仅有default配置：各个版本按default配置
8-; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9-; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10-; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11-; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12-[MoeComputeExpertTokens]
13-default=0
Dmoe/moe_finalize_routing_v2/op_host/config/ascend960pr/moe_finalize_routing_v2_simplified_key.ini+0-13
@@ -1,13 +0,0 @@
1-; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
2-; [某算子]
3-; default=xx
4-; ascendxx=xx
5-; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6-; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7-; 2)如果仅有default配置：各个版本按default配置
8-; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9-; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10-; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11-; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12-[MoeFinalizeRoutingV2]
13-default=0
Dmoe/moe_finalize_routing_v2_grad/op_host/config/ascend960pr/moe_finalize_routing_v2_grad_simplified_key.ini+0-13
@@ -1,13 +0,0 @@
1-; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
2-; [某算子]
3-; default=xx
4-; ascendxx=xx
5-; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6-; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7-; 2)如果仅有default配置：各个版本按default配置
8-; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9-; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10-; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11-; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12-[MoeFinalizeRoutingV2Grad]
13-default=0
Dmoe/moe_gating_top_k_softmax/op_host/config/ascend960pr/moe_gating_top_k_softmax_simplified_key.ini+0-13
@@ -1,13 +0,0 @@
1-; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
2-; [某算子]
3-; default=xx
4-; ascendxx=xx
5-; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6-; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7-; 2)如果仅有default配置：各个版本按default配置
8-; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9-; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10-; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11-; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12-[MoeGatingTopKSoftmax]
13-default=0
Dmoe/moe_gating_top_k_softmax_v2/op_host/config/ascend960pr/moe_gating_top_k_softmax_v2_simplified_key.ini+0-13
@@ -1,13 +0,0 @@
1-; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
2-; [某算子]
3-; default=xx
4-; ascendxx=xx
5-; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6-; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7-; 2)如果仅有default配置：各个版本按default配置
8-; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9-; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10-; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11-; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12-[MoeGatingTopKSoftmaxV2]
13-default=0
Dmoe/moe_init_routing/op_host/config/ascend960pr/moe_init_routing_simplified_key.ini+0-13
@@ -1,13 +0,0 @@
1-; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
2-; [某算子]
3-; default=xx
4-; ascendxx=xx
5-; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6-; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7-; 2)如果仅有default配置：各个版本按default配置
8-; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9-; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10-; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11-; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12-[MoeInitRouting]
13-default=0
Dmoe/moe_init_routing_quant_v2/op_host/config/ascend960pr/moe_init_routing_quant_v2_simplified_key.ini+0-13
@@ -1,13 +0,0 @@
1-; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
2-; [某算子]
3-; default=xx
4-; ascendxx=xx
5-; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6-; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7-; 2)如果仅有default配置：各个版本按default配置
8-; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9-; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10-; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11-; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12-[MoeInitRoutingQuantV2]
13-default=0
Dmoe/moe_init_routing_v2_grad/op_host/config/ascend960pr/moe_init_routing_v2_grad_simplified_key.ini+0-13
@@ -1,13 +0,0 @@
1-; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
2-; [某算子]
3-; default=xx
4-; ascendxx=xx
5-; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6-; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7-; 2)如果仅有default配置：各个版本按default配置
8-; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9-; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10-; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11-; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12-[MoeInitRoutingV2Grad]
13-default=0
Dmoe/moe_init_routing_v3/op_host/config/ascend960pr/moe_init_routing_v3_simplified_key.ini+0-13
@@ -1,13 +0,0 @@
1-; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
2-; [某算子]
3-; default=xx
4-; ascendxx=xx
5-; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6-; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7-; 2)如果仅有default配置：各个版本按default配置
8-; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9-; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10-; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11-; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12-[MoeInitRoutingV3]
13-default=0
Dmoe/moe_re_routing/op_host/config/ascend960pr/moe_re_routing_simplified_key.ini+0-13
@@ -1,13 +0,0 @@
1-; 该文件主要影响 opc 工具 编译二进制kernel时， --simplified_key_mode 选项中填写的值，格式如下所示：
2-; [某算子]
3-; default=xx
4-; ascendxx=xx
5-; 其中，default为默认mode，ascnedxx为可选mode，如果不同芯片有差异化要求时，需要配置；
6-; 1)如果没有配置：非ascendC算子继续按空处理，即opc编译命令中不添加 --simplified_key_mode 选项，AscendC算子按照 simplified_key_mode=0 处理
7-; 2)如果仅有default配置：各个版本按default配置
8-; 3)如果仅有某些平台的配置，没有default配置：对应平台的按照配置的值传递，非对应平台的：非AscendC算子继续按空处理，AscendC算子按照 simplified_key_mode=0 处理
9-; 4)如果default配置和平台配置都有：对应平台的使用平台的配置，非对应的平台的以default值配置。
10-; 5)对于自定义simplified key的情况，需要在binary_simplified_key_mode.ini 文件中显式配置为None，不传入 --simplified_key_mode 选项，由opc工具和FE框架自行判断使用何种模式
11-; 6)是否是AscendC算子，由 ops/build-in/tbe/op_info_cfg/parser/ascendc_config.json 中配置的算子名字和对于的平台决定
12-[MoeReRouting]
13-default=0
Mposembedding/apply_rotary_pos_emb/op_host/apply_rotary_pos_emb_def.cpp+1-0
@@ -68,6 +68,7 @@ public:
68 .ExtendCfgInfo("opFile.value", "apply_rotary_pos_emb_apt");68 .ExtendCfgInfo("opFile.value", "apply_rotary_pos_emb_apt");
69 this->AICore().AddConfig("ascend950", regbaseCfg);69 this->AICore().AddConfig("ascend950", regbaseCfg);
70 this->AICore().AddConfig("ascend960dt", regbaseCfg);70 this->AICore().AddConfig("ascend960dt", regbaseCfg);
71+ this->AICore().AddConfig("ascend960pr", regbaseCfg);
71 this->AICore().AddConfig("ascend350", regbaseCfg);72 this->AICore().AddConfig("ascend350", regbaseCfg);
72 this->AICore().AddConfig("mc62", regbaseCfg);73 this->AICore().AddConfig("mc62", regbaseCfg);
73 OpAICoreConfig aicore_config;74 OpAICoreConfig aicore_config;
Aposembedding/apply_rotary_pos_emb/op_host/config/ascend960pr/apply_rotary_pos_emb_binary.json+260-0
@@ -0,0 +1,260 @@
1+{
2+ "op_type": "ApplyRotaryPosEmb",
3+ "op_list": [
4+ {
5+ "bin_filename": "ApplyRotaryPosEmb_f3486ed63528a62ec8ee094f509b8879",
6+ "inputs": [
7+ {
8+ "name": "query",
9+ "index": 0,
10+ "dtype": "float16",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [
14+ -2
15+ ],
16+ "format_match_mode": "FormatAgnostic"
17+ },
18+ {
19+ "name": "key",
20+ "index": 1,
21+ "dtype": "float16",
22+ "format": "ND",
23+ "paramType": "required",
24+ "shape": [
25+ -2
26+ ],
27+ "format_match_mode": "FormatAgnostic"
28+ },
29+ {
30+ "name": "cos",
31+ "index": 2,
32+ "dtype": "float16",
33+ "format": "ND",
34+ "paramType": "required",
35+ "shape": [
36+ -2
37+ ],
38+ "format_match_mode": "FormatAgnostic"
39+ },
40+ {
41+ "name": "sin",
42+ "index": 3,
43+ "dtype": "float16",
44+ "format": "ND",
45+ "paramType": "required",
46+ "shape": [
47+ -2
48+ ],
49+ "format_match_mode": "FormatAgnostic"
50+ }
51+ ],
52+ "outputs": [
53+ {
54+ "name": "query",
55+ "index": 0,
56+ "dtype": "float16",
57+ "format": "ND",
58+ "paramType": "required",
59+ "shape": [
60+ -2
61+ ],
62+ "format_match_mode": "FormatAgnostic"
63+ },
64+ {
65+ "name": "key",
66+ "index": 1,
67+ "dtype": "float16",
68+ "format": "ND",
69+ "paramType": "required",
70+ "shape": [
71+ -2
72+ ],
73+ "format_match_mode": "FormatAgnostic"
74+ }
75+ ],
76+ "attrs": [
77+ {
78+ "name": "layout",
79+ "dtype": "int",
80+ "value": 1
81+ },
82+ {
83+ "name": "rotary_mode",
84+ "dtype": "string",
85+ "value": "half"
86+ }
87+ ]
88+ },
89+ {
90+ "bin_filename": "ApplyRotaryPosEmb_31fab5b1b1c1a3eec396cca8eb1570b2",
91+ "inputs": [
92+ {
93+ "name": "query",
94+ "index": 0,
95+ "dtype": "bfloat16",
96+ "format": "ND",
97+ "paramType": "required",
98+ "shape": [
99+ -2
100+ ],
101+ "format_match_mode": "FormatAgnostic"
102+ },
103+ {
104+ "name": "key",
105+ "index": 1,
106+ "dtype": "bfloat16",
107+ "format": "ND",
108+ "paramType": "required",
109+ "shape": [
110+ -2
111+ ],
112+ "format_match_mode": "FormatAgnostic"
113+ },
114+ {
115+ "name": "cos",
116+ "index": 2,
117+ "dtype": "bfloat16",
118+ "format": "ND",
119+ "paramType": "required",
120+ "shape": [
121+ -2
122+ ],
123+ "format_match_mode": "FormatAgnostic"
124+ },
125+ {
126+ "name": "sin",
127+ "index": 3,
128+ "dtype": "bfloat16",
129+ "format": "ND",
130+ "paramType": "required",
131+ "shape": [
132+ -2
133+ ],
134+ "format_match_mode": "FormatAgnostic"
135+ }
136+ ],
137+ "outputs": [
138+ {
139+ "name": "query",
140+ "index": 0,
141+ "dtype": "bfloat16",
142+ "format": "ND",
143+ "paramType": "required",
144+ "shape": [
145+ -2
146+ ],
147+ "format_match_mode": "FormatAgnostic"
148+ },
149+ {
150+ "name": "key",
151+ "index": 1,
152+ "dtype": "bfloat16",
153+ "format": "ND",
154+ "paramType": "required",
155+ "shape": [
156+ -2
157+ ],
158+ "format_match_mode": "FormatAgnostic"
159+ }
160+ ],
161+ "attrs": [
162+ {
163+ "name": "layout",
164+ "dtype": "int",
165+ "value": 1
166+ },
167+ {
168+ "name": "rotary_mode",
169+ "dtype": "string",
170+ "value": "half"
171+ }
172+ ]
173+ },
174+ {
175+ "bin_filename": "ApplyRotaryPosEmb_16182dad2cad5c4ace1e467e65142759",
176+ "inputs": [
177+ {
178+ "name": "query",
179+ "index": 0,
180+ "dtype": "float32",
181+ "format": "ND",
182+ "paramType": "required",
183+ "shape": [
184+ -2
185+ ],
186+ "format_match_mode": "FormatAgnostic"
187+ },
188+ {
189+ "name": "key",
190+ "index": 1,
191+ "dtype": "float32",
192+ "format": "ND",
193+ "paramType": "required",
194+ "shape": [
195+ -2
196+ ],
197+ "format_match_mode": "FormatAgnostic"
198+ },
199+ {
200+ "name": "cos",
201+ "index": 2,
202+ "dtype": "float32",
203+ "format": "ND",
204+ "paramType": "required",
205+ "shape": [
206+ -2
207+ ],
208+ "format_match_mode": "FormatAgnostic"
209+ },
210+ {
211+ "name": "sin",
212+ "index": 3,
213+ "dtype": "float32",
214+ "format": "ND",
215+ "paramType": "required",
216+ "shape": [
217+ -2
218+ ],
219+ "format_match_mode": "FormatAgnostic"
220+ }
221+ ],
222+ "outputs": [
223+ {
224+ "name": "query",
225+ "index": 0,
226+ "dtype": "float32",
227+ "format": "ND",
228+ "paramType": "required",
229+ "shape": [
230+ -2
231+ ],
232+ "format_match_mode": "FormatAgnostic"
233+ },
234+ {
235+ "name": "key",
236+ "index": 1,
237+ "dtype": "float32",
238+ "format": "ND",
239+ "paramType": "required",
240+ "shape": [
241+ -2
242+ ],
243+ "format_match_mode": "FormatAgnostic"
244+ }
245+ ],
246+ "attrs": [
247+ {
248+ "name": "layout",
249+ "dtype": "int",
250+ "value": 1
251+ },
252+ {
253+ "name": "rotary_mode",
254+ "dtype": "string",
255+ "value": "half"
256+ }
257+ ]
258+ }
259+ ]
260+}
Mposembedding/inplace_partial_rotary_mul/op_host/CMakeLists.txt+6-0
@@ -34,6 +34,12 @@ elseif("${CONDITION_UNIT}" STREQUAL "ascend960dt")
34 COMPUTE_UNIT ascend960dt_968B734 COMPUTE_UNIT ascend960dt_968B7
35 OPTIONS -DENABLE_CV_COMM_VIA_SSBUF=true35 OPTIONS -DENABLE_CV_COMM_VIA_SSBUF=true
36 )36 )
37+elseif("${CONDITION_UNIT}" STREQUAL "ascend960pr")
38+ add_ops_compile_options(
39+ OP_NAME InplacePartialRotaryMul
40+ COMPUTE_UNIT Ascend960PR_969A1
41+ OPTIONS -DENABLE_CV_COMM_VIA_SSBUF=true
42+ )
37endif()43endif()
38 44 
39if (BUILD_OPS_RTY_KERNEL)45if (BUILD_OPS_RTY_KERNEL)
Aposembedding/inplace_partial_rotary_mul/op_host/config/ascend960pr/inplace_partial_rotary_mul_binary.json+315-0
@@ -0,0 +1,315 @@
1+{
2+ "op_type": "InplacePartialRotaryMul",
3+ "op_list": [
4+ {
5+ "bin_filename": "InplacePartialRotaryMul_0",
6+ "inputs": [
7+ {
8+ "name": "x",
9+ "index": 0,
10+ "dtype": "float16",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [
14+ -2
15+ ]
16+ },
17+ {
18+ "name": "cos",
19+ "index": 1,
20+ "dtype": "float16",
21+ "format": "ND",
22+ "paramType": "required",
23+ "shape": [
24+ -2
25+ ]
26+ },
27+ {
28+ "name": "sin",
29+ "index": 2,
30+ "dtype": "float16",
31+ "format": "ND",
32+ "paramType": "required",
33+ "shape": [
34+ -2
35+ ]
36+ }
37+ ],
38+ "outputs": [
39+ {
40+ "name": "x",
41+ "index": 0,
42+ "dtype": "float16",
43+ "format": "ND",
44+ "paramType": "required",
45+ "shape": [
46+ -2
47+ ]
48+ }
49+ ],
50+ "attrs": [
51+ {
52+ "name": "rotary_mode",
53+ "dtype": "int",
54+ "value": 1
55+ },
56+ {
57+ "name": "partial_slice",
58+ "dtype": "list_int",
59+ "value": [
60+ 0,
61+ 0
62+ ]
63+ }
64+ ]
65+ },
66+ {
67+ "bin_filename": "InplacePartialRotaryMul_1",
68+ "inputs": [
69+ {
70+ "name": "x",
71+ "index": 0,
72+ "dtype": "float32",
73+ "format": "ND",
74+ "paramType": "required",
75+ "shape": [
76+ -2
77+ ]
78+ },
79+ {
80+ "name": "cos",
81+ "index": 1,
82+ "dtype": "float32",
83+ "format": "ND",
84+ "paramType": "required",
85+ "shape": [
86+ -2
87+ ]
88+ },
89+ {
90+ "name": "sin",
91+ "index": 2,
92+ "dtype": "float32",
93+ "format": "ND",
94+ "paramType": "required",
95+ "shape": [
96+ -2
97+ ]
98+ }
99+ ],
100+ "outputs": [
101+ {
102+ "name": "x",
103+ "index": 0,
104+ "dtype": "float32",
105+ "format": "ND",
106+ "paramType": "required",
107+ "shape": [
108+ -2
109+ ]
110+ }
111+ ],
112+ "attrs": [
113+ {
114+ "name": "rotary_mode",
115+ "dtype": "int",
116+ "value": 1
117+ },
118+ {
119+ "name": "partial_slice",
120+ "dtype": "list_int",
121+ "value": [
122+ 0,
123+ 0
124+ ]
125+ }
126+ ]
127+ },
128+ {
129+ "bin_filename": "InplacePartialRotaryMul_2",
130+ "inputs": [
131+ {
132+ "name": "x",
133+ "index": 0,
134+ "dtype": "bfloat16",
135+ "format": "ND",
136+ "paramType": "required",
137+ "shape": [
138+ -2
139+ ]
140+ },
141+ {
142+ "name": "cos",
143+ "index": 1,
144+ "dtype": "bfloat16",
145+ "format": "ND",
146+ "paramType": "required",
147+ "shape": [
148+ -2
149+ ]
150+ },
151+ {
152+ "name": "sin",
153+ "index": 2,
154+ "dtype": "bfloat16",
155+ "format": "ND",
156+ "paramType": "required",
157+ "shape": [
158+ -2
159+ ]
160+ }
161+ ],
162+ "outputs": [
163+ {
164+ "name": "x",
165+ "index": 0,
166+ "dtype": "bfloat16",
167+ "format": "ND",
168+ "paramType": "required",
169+ "shape": [
170+ -2
171+ ]
172+ }
173+ ],
174+ "attrs": [
175+ {
176+ "name": "rotary_mode",
177+ "dtype": "int",
178+ "value": 1
179+ },
180+ {
181+ "name": "partial_slice",
182+ "dtype": "list_int",
183+ "value": [
184+ 0,
185+ 0
186+ ]
187+ }
188+ ]
189+ },
190+ {
191+ "bin_filename": "InplacePartialRotaryMul_3",
192+ "inputs": [
193+ {
194+ "name": "x",
195+ "index": 0,
196+ "dtype": "float16",
197+ "format": "ND",
198+ "paramType": "required",
199+ "shape": [
200+ -2
201+ ]
202+ },
203+ {
204+ "name": "cos",
205+ "index": 1,
206+ "dtype": "float32",
207+ "format": "ND",
208+ "paramType": "required",
209+ "shape": [
210+ -2
211+ ]
212+ },
213+ {
214+ "name": "sin",
215+ "index": 2,
216+ "dtype": "float32",
217+ "format": "ND",
218+ "paramType": "required",
219+ "shape": [
220+ -2
221+ ]
222+ }
223+ ],
224+ "outputs": [
225+ {
226+ "name": "x",
227+ "index": 0,
228+ "dtype": "float16",
229+ "format": "ND",
230+ "paramType": "required",
231+ "shape": [
232+ -2
233+ ]
234+ }
235+ ],
236+ "attrs": [
237+ {
238+ "name": "rotary_mode",
239+ "dtype": "int",
240+ "value": 1
241+ },
242+ {
243+ "name": "partial_slice",
244+ "dtype": "list_int",
245+ "value": [
246+ 0,
247+ 0
248+ ]
249+ }
250+ ]
251+ },
252+ {
253+ "bin_filename": "InplacePartialRotaryMul_4",
254+ "inputs": [
255+ {
256+ "name": "x",
257+ "index": 0,
258+ "dtype": "bfloat16",
259+ "format": "ND",
260+ "paramType": "required",
261+ "shape": [
262+ -2
263+ ]
264+ },
265+ {
266+ "name": "cos",
267+ "index": 1,
268+ "dtype": "float32",
269+ "format": "ND",
270+ "paramType": "required",
271+ "shape": [
272+ -2
273+ ]
274+ },
275+ {
276+ "name": "sin",
277+ "index": 2,
278+ "dtype": "float32",
279+ "format": "ND",
280+ "paramType": "required",
281+ "shape": [
282+ -2
283+ ]
284+ }
285+ ],
286+ "outputs": [
287+ {
288+ "name": "x",
289+ "index": 0,
290+ "dtype": "bfloat16",
291+ "format": "ND",
292+ "paramType": "required",
293+ "shape": [
294+ -2
295+ ]
296+ }
297+ ],
298+ "attrs": [
299+ {
300+ "name": "rotary_mode",
301+ "dtype": "int",
302+ "value": 1
303+ },
304+ {
305+ "name": "partial_slice",
306+ "dtype": "list_int",
307+ "value": [
308+ 0,
309+ 0
310+ ]
311+ }
312+ ]
313+ }
314+ ]
315+}
Mposembedding/inplace_partial_rotary_mul/op_host/inplace_partial_rotary_mul_def.cpp+1-0
@@ -83,6 +83,7 @@ public:
83 .ExtendCfgInfo("opFile.value", "inplace_partial_rotary_mul_apt");83 .ExtendCfgInfo("opFile.value", "inplace_partial_rotary_mul_apt");
84 this->AICore().AddConfig("ascend950", config950);84 this->AICore().AddConfig("ascend950", config950);
85 this->AICore().AddConfig("ascend960dt", config950);85 this->AICore().AddConfig("ascend960dt", config950);
86+ this->AICore().AddConfig("ascend960pr", config950);
86 this->AICore().AddConfig("ascend350", config950);87 this->AICore().AddConfig("ascend350", config950);
87 }88 }
88};89};
Mposembedding/partial_rotary_mul_quant/op_graph/partial_rotary_mul_quant_proto.h+1-1
@@ -27,7 +27,7 @@ namespace ge {
27 * @li sin: A tensor for rotary position embedding sin. Must be one of the following types: float16, float32, bfloat16.27 * @li sin: A tensor for rotary position embedding sin. Must be one of the following types: float16, float32, bfloat16.
28 * @li scale: Optional. HiFloat8 multiplier: single-element float32 tensor; absent means 1.0.28 * @li scale: Optional. HiFloat8 multiplier: single-element float32 tensor; absent means 1.0.
29 * @par Attributes:29 * @par Attributes:
30- * @li y_dtype: Optional. 0: HiFloat8 (Ascend950). Defaults to 0.30+ * @li y_dtype: Optional. 0: HiFloat8 (Ascend950), 1: HiFloat4 (Ascend960DT/Ascend960PR). Defaults to 0.
31 * @li quant_mode: Optional. Quantization mode. Defaults to 0.31 * @li quant_mode: Optional. Quantization mode. Defaults to 0.
32 * @li round_mode: Optional. HiFloat8 supports "round" and "hybrid" (CAST_HYBRID). Defaults to "round".32 * @li round_mode: Optional. HiFloat8 supports "round" and "hybrid" (CAST_HYBRID). Defaults to "round".
33 * @li axis: Optional. Axis for quantization. Ignored for HiFloat8. Defaults to -1.33 * @li axis: Optional. Axis for quantization. Ignored for HiFloat8. Defaults to -1.
Mposembedding/partial_rotary_mul_quant/op_host/CMakeLists.txt+2-2
@@ -10,8 +10,8 @@
10add_op_to_compiled_list()10add_op_to_compiled_list()
11 11 
12if (BUILD_OPEN_PROJECT)12if (BUILD_OPEN_PROJECT)
13- if (ASCEND_COMPUTE_UNIT STREQUAL "ascend960dt")13+ if (ASCEND_COMPUTE_UNIT STREQUAL "ascend960dt" OR ASCEND_COMPUTE_UNIT STREQUAL "ascend960pr")
14- # Register HiFloat4 outputs only when building for Ascend960DT.14+ # Register HiFloat4 outputs only when building for Ascend960DT/Ascend960PR.
15 set_property(SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/partial_rotary_mul_quant_def.cpp15 set_property(SOURCE ${CMAKE_CURRENT_SOURCE_DIR}/partial_rotary_mul_quant_def.cpp
16 DIRECTORY ${PROJECT_SOURCE_DIR} APPEND PROPERTY COMPILE_OPTIONS16 DIRECTORY ${PROJECT_SOURCE_DIR} APPEND PROPERTY COMPILE_OPTIONS
17 -DPARTIAL_ROTARY_MUL_QUANT_ASCEND960_ONLY)17 -DPARTIAL_ROTARY_MUL_QUANT_ASCEND960_ONLY)
Aposembedding/partial_rotary_mul_quant/op_host/config/ascend960pr/partial_rotary_mul_quant_binary.json+284-0
@@ -0,0 +1,284 @@
1+{
2+ "op_type": "PartialRotaryMulQuant",
3+ "op_list": [
4+ {
5+ "bin_filename": "PartialRotaryMulQuant_fp16",
6+ "inputs": [
7+ {
8+ "name": "x",
9+ "index": 0,
10+ "dtype": "float16",
11+ "format": "ND",
12+ "paramType": "required",
13+ "shape": [-2],
14+ "format_match_mode": "FormatAgnostic"
15+ },
16+ {
17+ "name": "cos",
18+ "index": 1,
19+ "dtype": "float16",
20+ "format": "ND",
21+ "paramType": "required",
22+ "shape": [-2],
23+ "format_match_mode": "FormatAgnostic"
24+ },
25+ {
26+ "name": "sin",
27+ "index": 2,
28+ "dtype": "float16",
29+ "format": "ND",
30+ "paramType": "required",
31+ "shape": [-2],
32+ "format_match_mode": "FormatAgnostic"
33+ },
34+ {
35+ "name": "scale",
36+ "index": 3,
37+ "dtype": "float32",
38+ "format": "ND",
39+ "paramType": "optional",
40+ "shape": [-2],
41+ "format_match_mode": "FormatAgnostic"
42+ }
43+ ],
44+ "outputs": [
45+ {
46+ "name": "y",
47+ "index": 0,
48+ "dtype": "hifloat4",
49+ "format": "ND",
50+ "paramType": "required",
51+ "shape": [-2],
52+ "format_match_mode": "FormatAgnostic"
53+ },
54+ {
55+ "name": "y_scale",
56+ "index": 1,
57+ "dtype": "hifloat4_scale",
58+ "format": "ND",
59+ "paramType": "required",
60+ "shape": [-2],
61+ "format_match_mode": "FormatAgnostic"
62+ }
63+ ],
64+ "attrs": [
65+ {
66+ "name": "y_dtype",
67+ "dtype": "int",
68+ "value": 0
69+ },
70+ {
71+ "name": "quant_mode",
72+ "dtype": "int",
73+ "value": 0
74+ },
75+ {
76+ "name": "round_mode",
77+ "dtype": "string",
78+ "value": "round"
79+ },
80+ {
81+ "name": "axis",
82+ "dtype": "int",
83+ "value": -1
84+ },
85+ {
86+ "name": "rotary_mode",
87+ "dtype": "int",
88+ "value": 0
89+ },
90+ {
91+ "name": "partial_slice",
92+ "dtype": "list_int",
93+ "value": [0, 0]
94+ }
95+ ]
96+ },
97+ {
98+ "bin_filename": "PartialRotaryMulQuant_bf16",
99+ "inputs": [
100+ {
101+ "name": "x",
102+ "index": 0,
103+ "dtype": "bfloat16",
104+ "format": "ND",
105+ "paramType": "required",
106+ "shape": [-2],
107+ "format_match_mode": "FormatAgnostic"
108+ },
109+ {
110+ "name": "cos",
111+ "index": 1,
112+ "dtype": "bfloat16",
113+ "format": "ND",
114+ "paramType": "required",
115+ "shape": [-2],
116+ "format_match_mode": "FormatAgnostic"
117+ },
118+ {
119+ "name": "sin",
120+ "index": 2,
121+ "dtype": "bfloat16",
122+ "format": "ND",
123+ "paramType": "required",
124+ "shape": [-2],
125+ "format_match_mode": "FormatAgnostic"
126+ },
127+ {
128+ "name": "scale",
129+ "index": 3,
130+ "dtype": "float32",
131+ "format": "ND",
132+ "paramType": "optional",
133+ "shape": [-2],
134+ "format_match_mode": "FormatAgnostic"
135+ }
136+ ],
137+ "outputs": [
138+ {
139+ "name": "y",
140+ "index": 0,
141+ "dtype": "hifloat4",
142+ "format": "ND",
143+ "paramType": "required",
144+ "shape": [-2],
145+ "format_match_mode": "FormatAgnostic"
146+ },
147+ {
148+ "name": "y_scale",
149+ "index": 1,
150+ "dtype": "hifloat4_scale",
151+ "format": "ND",
152+ "paramType": "required",
153+ "shape": [-2],
154+ "format_match_mode": "FormatAgnostic"
155+ }
156+ ],
157+ "attrs": [
158+ {
159+ "name": "y_dtype",
160+ "dtype": "int",
161+ "value": 0
162+ },
163+ {
164+ "name": "quant_mode",
165+ "dtype": "int",
166+ "value": 0
167+ },
168+ {
169+ "name": "round_mode",
170+ "dtype": "string",
171+ "value": "round"
172+ },
173+ {
174+ "name": "axis",
175+ "dtype": "int",
176+ "value": -1
177+ },
178+ {
179+ "name": "rotary_mode",
180+ "dtype": "int",
181+ "value": 0
182+ },
183+ {
184+ "name": "partial_slice",
185+ "dtype": "list_int",
186+ "value": [0, 0]
187+ }
188+ ]
189+ },
190+ {
191+ "bin_filename": "PartialRotaryMulQuant_fp32",
192+ "inputs": [
193+ {
194+ "name": "x",
195+ "index": 0,
196+ "dtype": "float32",
197+ "format": "ND",
198+ "paramType": "required",
199+ "shape": [-2],
200+ "format_match_mode": "FormatAgnostic"
201+ },
202+ {
203+ "name": "cos",
204+ "index": 1,
205+ "dtype": "float32",
206+ "format": "ND",
207+ "paramType": "required",
208+ "shape": [-2],
209+ "format_match_mode": "FormatAgnostic"
210+ },
211+ {
212+ "name": "sin",
213+ "index": 2,
214+ "dtype": "float32",
215+ "format": "ND",
216+ "paramType": "required",
217+ "shape": [-2],
218+ "format_match_mode": "FormatAgnostic"
219+ },
220+ {
221+ "name": "scale",
222+ "index": 3,
223+ "dtype": "float32",
224+ "format": "ND",
225+ "paramType": "optional",
226+ "shape": [-2],
227+ "format_match_mode": "FormatAgnostic"
228+ }
229+ ],
230+ "outputs": [
231+ {
232+ "name": "y",
233+ "index": 0,
234+ "dtype": "hifloat4",
235+ "format": "ND",
236+ "paramType": "required",
237+ "shape": [-2],
238+ "format_match_mode": "FormatAgnostic"
239+ },
240+ {
241+ "name": "y_scale",
242+ "index": 1,
243+ "dtype": "hifloat4_scale",
244+ "format": "ND",
245+ "paramType": "required",
246+ "shape": [-2],
247+ "format_match_mode": "FormatAgnostic"
248+ }
249+ ],
250+ "attrs": [
251+ {
252+ "name": "y_dtype",
253+ "dtype": "int",
254+ "value": 0
255+ },
256+ {
257+ "name": "quant_mode",
258+ "dtype": "int",
259+ "value": 0
260+ },
261+ {
262+ "name": "round_mode",
263+ "dtype": "string",
264+ "value": "round"
265+ },
266+ {
267+ "name": "axis",
268+ "dtype": "int",
269+ "value": -1
270+ },
271+ {
272+ "name": "rotary_mode",
273+ "dtype": "int",
274+ "value": 0
275+ },
276+ {
277+ "name": "partial_slice",
278+ "dtype": "list_int",
279+ "value": [0, 0]
280+ }
281+ ]
282+ }
283+ ]
284+}
Mposembedding/partial_rotary_mul_quant/op_host/partial_rotary_mul_quant_def.cpp+1-0
@@ -78,6 +78,7 @@ public:
78 .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})78 .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
79 .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});79 .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
80 this->AICore().AddConfig("ascend960dt", ascend960dtCfg);80 this->AICore().AddConfig("ascend960dt", ascend960dtCfg);
81+ this->AICore().AddConfig("ascend960pr", ascend960dtCfg);
81#endif82#endif
82 83 
83 OpAICoreConfig ascend950Cfg;84 OpAICoreConfig ascend950Cfg;
Aposembedding/partial_rotary_mul_quant/tests/ut/CMakeLists.txt+16-0
@@ -0,0 +1,16 @@
1+# -----------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# -----------------------------------------------------------------------------------------------------------
10+ 
11+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
12+foreach(SUB_DIR ${CURRENT_DIRS})
13+ if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
14+ add_subdirectory(${SUB_DIR})
15+ endif()
16+endforeach()
Aposembedding/partial_rotary_mul_quant/tests/ut/op_host/CMakeLists.txt+14-0
@@ -0,0 +1,14 @@
1+# -----------------------------------------------------------------------------------------------------------
2+# Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+# This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+# CANN Open Software License Agreement Version 2.0 (the "License").
5+# Please refer to the License for details. You may not use this file except in compliance with the License.
6+# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+# See LICENSE in the root of the software repository for the full text of the License.
9+# -----------------------------------------------------------------------------------------------------------
10+ 
11+message(STATUS "=== Debug: add test_partial_rotary_mul_quant")
12+if(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})
14+endif()
Aposembedding/partial_rotary_mul_quant/tests/ut/op_host/test_partial_rotary_mul_quant_tiling.cpp+227-0
@@ -0,0 +1,227 @@
1+/**
2+ * Copyright (c) 2026 Huawei Technologies Co., Ltd.
3+ * This program is free software, you can redistribute it and/or modify it under the terms and conditions of
4+ * CANN Open Software License Agreement Version 2.0 (the "License").
5+ * Please refer to the License for details. You may not use this file except in compliance with the License.
6+ * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
7+ * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
8+ * See LICENSE in the root of the software repository for the full text of the License.
9+ */
10+ 
11+#include <iostream>
12+#include <gtest/gtest.h>
13+#include "../../../op_host/partial_rotary_mul_quant_tiling.h"
14+#include "tiling_context_faker.h"
15+#include "tiling_case_executor.h"
16+ 
17+using namespace std;
18+ 
19+class PartialRotaryMulQuantTiling : public testing::Test {
20+protected:
21+ static void SetUpTestCase()
22+ {
23+ std::cout << "PartialRotaryMulQuantTiling SetUp" << std::endl;
24+ }
25+ 
26+ static void TearDownTestCase()
27+ {
28+ std::cout << "PartialRotaryMulQuantTiling TearDown" << std::endl;
29+ }
30+};
31+ 
32+struct PartialRotaryMulQuantCompileInfo {
33+ int64_t coreNum;
34+ int64_t ubSize;
35+ int64_t l2Size;
36+};
37+ 
38+// ==================== 950 HiFloat8 用例 ====================
39+ 
40+TEST_F(PartialRotaryMulQuantTiling, test_partial_rotary_mul_quant_fp16_hif8_950_success)
41+{
42+ PartialRotaryMulQuantCompileInfo compileInfo = {40, 196608, 33554432};
43+ gert::TilingContextPara tilingContextPara("PartialRotaryMulQuant",
44+ {
45+ {{{1, 128, 4, 128}, {1, 128, 4, 128}}, ge::DT_FLOAT16, ge::FORMAT_ND},
46+ },
47+ {
48+ {{{1, 128, 4, 64}, {1, 128, 4, 64}}, ge::DT_FLOAT16, ge::FORMAT_ND},
49+ {{{1}, {1}}, ge::DT_FLOAT, ge::FORMAT_ND},
50+ },
51+ {
52+ {"y_dtype", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
53+ },
54+ &compileInfo, "Ascend950");
55+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS);
56+}
57+ 
58+TEST_F(PartialRotaryMulQuantTiling, test_partial_rotary_mul_quant_fp32_hif8_950_success)
59+{
60+ PartialRotaryMulQuantCompileInfo compileInfo = {40, 196608, 33554432};
61+ gert::TilingContextPara tilingContextPara("PartialRotaryMulQuant",
62+ {
63+ {{{1, 64, 2, 64}, {1, 64, 2, 64}}, ge::DT_FLOAT, ge::FORMAT_ND},
64+ },
65+ {
66+ {{{1, 64, 2, 32}, {1, 64, 2, 32}}, ge::DT_FLOAT, ge::FORMAT_ND},
67+ {{{1}, {1}}, ge::DT_FLOAT, ge::FORMAT_ND},
68+ },
69+ {
70+ {"y_dtype", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
71+ },
72+ &compileInfo, "Ascend950");
73+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS);
74+}
75+ 
76+TEST_F(PartialRotaryMulQuantTiling, test_partial_rotary_mul_quant_bf16_hif8_950_success)
77+{
78+ PartialRotaryMulQuantCompileInfo compileInfo = {40, 196608, 33554432};
79+ gert::TilingContextPara tilingContextPara("PartialRotaryMulQuant",
80+ {
81+ {{{2, 256, 8, 64}, {2, 256, 8, 64}}, ge::DT_BF16, ge::FORMAT_ND},
82+ },
83+ {
84+ {{{2, 256, 8, 32}, {2, 256, 8, 32}}, ge::DT_BF16, ge::FORMAT_ND},
85+ {{{1}, {1}}, ge::DT_FLOAT, ge::FORMAT_ND},
86+ },
87+ {
88+ {"y_dtype", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
89+ },
90+ &compileInfo, "Ascend950");
91+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS);
92+}
93+ 
94+// ==================== 960PR HiFloat8 + HiFloat4 用例 ====================
95+ 
96+TEST_F(PartialRotaryMulQuantTiling, test_partial_rotary_mul_quant_fp16_hif8_960pr_success)
97+{
98+ PartialRotaryMulQuantCompileInfo compileInfo = {40, 196608, 33554432};
99+ gert::TilingContextPara tilingContextPara("PartialRotaryMulQuant",
100+ {
101+ {{{1, 128, 4, 128}, {1, 128, 4, 128}}, ge::DT_FLOAT16, ge::FORMAT_ND},
102+ },
103+ {
104+ {{{1, 128, 4, 64}, {1, 128, 4, 64}}, ge::DT_FLOAT16, ge::FORMAT_ND},
105+ {{{1}, {1}}, ge::DT_FLOAT, ge::FORMAT_ND},
106+ },
107+ {
108+ {"y_dtype", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
109+ },
110+ &compileInfo, "Ascend960PR");
111+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS);
112+}
113+ 
114+TEST_F(PartialRotaryMulQuantTiling, test_partial_rotary_mul_quant_fp16_hif4_960pr_success)
115+{
116+ PartialRotaryMulQuantCompileInfo compileInfo = {40, 196608, 33554432};
117+ gert::TilingContextPara tilingContextPara("PartialRotaryMulQuant",
118+ {
119+ {{{1, 128, 4, 128}, {1, 128, 4, 128}}, ge::DT_FLOAT16, ge::FORMAT_ND},
120+ },
121+ {
122+ {{{1, 128, 4, 64}, {1, 128, 4, 64}}, ge::DT_FLOAT16, ge::FORMAT_ND},
123+ {{{1}, {1}}, ge::DT_FLOAT, ge::FORMAT_ND},
124+ },
125+ {
126+ // y_dtype=1 仅 960 系列（含 960PR）支持 HiFloat4
127+ {"y_dtype", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
128+ },
129+ &compileInfo, "Ascend960PR");
130+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS);
131+}
132+ 
133+TEST_F(PartialRotaryMulQuantTiling, test_partial_rotary_mul_quant_fp32_hif8_960pr_success)
134+{
135+ PartialRotaryMulQuantCompileInfo compileInfo = {40, 196608, 33554432};
136+ gert::TilingContextPara tilingContextPara("PartialRotaryMulQuant",
137+ {
138+ {{{1, 64, 2, 64}, {1, 64, 2, 64}}, ge::DT_FLOAT, ge::FORMAT_ND},
139+ },
140+ {
141+ {{{1, 64, 2, 32}, {1, 64, 2, 32}}, ge::DT_FLOAT, ge::FORMAT_ND},
142+ {{{1}, {1}}, ge::DT_FLOAT, ge::FORMAT_ND},
143+ },
144+ {
145+ {"y_dtype", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
146+ },
147+ &compileInfo, "Ascend960PR");
148+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS);
149+}
150+ 
151+// ==================== 950 边界/异常用例 ====================
152+ 
153+TEST_F(PartialRotaryMulQuantTiling, test_partial_rotary_mul_quant_invalid_dtype_950_failed)
154+{
155+ PartialRotaryMulQuantCompileInfo compileInfo = {40, 196608, 33554432};
156+ gert::TilingContextPara tilingContextPara("PartialRotaryMulQuant",
157+ {
158+ // int32 不支持
159+ {{{1, 128, 4, 128}, {1, 128, 4, 128}}, ge::DT_INT32, ge::FORMAT_ND},
160+ },
161+ {
162+ {{{1, 128, 4, 64}, {1, 128, 4, 64}}, ge::DT_FLOAT16, ge::FORMAT_ND},
163+ {{{1}, {1}}, ge::DT_FLOAT, ge::FORMAT_ND},
164+ },
165+ {
166+ {"y_dtype", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
167+ },
168+ &compileInfo, "Ascend950");
169+ ExecuteTestCase(tilingContextPara, ge::GRAPH_FAILED);
170+}
171+ 
172+// ==================== 960DT HiFloat8 + HiFloat4 用例 ====================
173+ 
174+TEST_F(PartialRotaryMulQuantTiling, test_partial_rotary_mul_quant_fp16_hif8_960dt_success)
175+{
176+ PartialRotaryMulQuantCompileInfo compileInfo = {40, 196608, 33554432};
177+ gert::TilingContextPara tilingContextPara("PartialRotaryMulQuant",
178+ {
179+ {{{1, 128, 4, 128}, {1, 128, 4, 128}}, ge::DT_FLOAT16, ge::FORMAT_ND},
180+ },
181+ {
182+ {{{1, 128, 4, 64}, {1, 128, 4, 64}}, ge::DT_FLOAT16, ge::FORMAT_ND},
183+ {{{1}, {1}}, ge::DT_FLOAT, ge::FORMAT_ND},
184+ },
185+ {
186+ {"y_dtype", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
187+ },
188+ &compileInfo, "Ascend960DT");
189+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS);
190+}
191+ 
192+TEST_F(PartialRotaryMulQuantTiling, test_partial_rotary_mul_quant_fp16_hif4_960dt_success)
193+{
194+ PartialRotaryMulQuantCompileInfo compileInfo = {40, 196608, 33554432};
195+ gert::TilingContextPara tilingContextPara("PartialRotaryMulQuant",
196+ {
197+ {{{1, 128, 4, 128}, {1, 128, 4, 128}}, ge::DT_FLOAT16, ge::FORMAT_ND},
198+ },
199+ {
200+ {{{1, 128, 4, 64}, {1, 128, 4, 64}}, ge::DT_FLOAT16, ge::FORMAT_ND},
201+ {{{1}, {1}}, ge::DT_FLOAT, ge::FORMAT_ND},
202+ },
203+ {
204+ // y_dtype=1 仅 960 系列（含 960DT）支持 HiFloat4
205+ {"y_dtype", Ops::Transformer::AnyValue::CreateFrom<int64_t>(1)},
206+ },
207+ &compileInfo, "Ascend960DT");
208+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS);
209+}
210+ 
211+TEST_F(PartialRotaryMulQuantTiling, test_partial_rotary_mul_quant_fp32_hif8_960dt_success)
212+{
213+ PartialRotaryMulQuantCompileInfo compileInfo = {40, 196608, 33554432};
214+ gert::TilingContextPara tilingContextPara("PartialRotaryMulQuant",
215+ {
216+ {{{1, 64, 2, 64}, {1, 64, 2, 64}}, ge::DT_FLOAT, ge::FORMAT_ND},
217+ },
218+ {
219+ {{{1, 64, 2, 32}, {1, 64, 2, 32}}, ge::DT_FLOAT, ge::FORMAT_ND},
220+ {{{1}, {1}}, ge::DT_FLOAT, ge::FORMAT_ND},
221+ },
222+ {
223+ {"y_dtype", Ops::Transformer::AnyValue::CreateFrom<int64_t>(0)},
224+ },
225+ &compileInfo, "Ascend960DT");
226+ ExecuteTestCase(tilingContextPara, ge::GRAPH_SUCCESS);
227+}
Mposembedding/rotary_position_embedding/op_graph/fusion_pass/rotary_position_embedding_fusion_pass.cpp+3-2
@@ -40,8 +40,8 @@ constexpr size_t kSinIdx = 2;
40 40 
41// Same gating as the built-in pass: fuse only on regbase platforms (NpuArch DAV_3510/DAV_9201/DAV_5102).41// Same gating as the built-in pass: fuse only on regbase platforms (NpuArch DAV_3510/DAV_9201/DAV_5102).
42// Gating on NpuArch instead of a short_soc_version string list: one arch covers multiple SoCs42// Gating on NpuArch instead of a short_soc_version string list: one arch covers multiple SoCs
43-// (DAV_3510: Ascend950/Ascend350, DAV_9201: Ascend960DT, DAV_5102: MC62CM12A), so SoCs sharing the regbase43+// (DAV_3510: Ascend950/Ascend350, DAV_9201/DAV_9202: Ascend960DT/Ascend960PR, DAV_5102: MC62CM12A), so SoCs sharing the
44-// arch are covered without maintaining the soc list.44+// regbase arch are covered without maintaining the soc list.
45bool IsRegBasePlatform()45bool IsRegBasePlatform()
46{46{
47 char npuArchVal[16] = {0};47 char npuArchVal[16] = {0};
@@ -53,6 +53,7 @@ bool IsRegBasePlatform()
53 const std::string npuArch(npuArchVal);53 const std::string npuArch(npuArchVal);
54 return npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_3510)) ||54 return npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_3510)) ||
55 npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_9201)) ||55 npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_9201)) ||
56+ npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_9202)) ||
56 npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_5102));57 npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_5102));
57}58}
58 59 
Mposembedding/rotary_position_embedding/op_host/CMakeLists.txt+7-0
@@ -37,6 +37,13 @@ elseif("${CONDITION_UNIT}" STREQUAL "ascend960dt")
37 COMPUTE_UNIT ascend960dt_968B737 COMPUTE_UNIT ascend960dt_968B7
38 OPTIONS -DENABLE_CV_COMM_VIA_SSBUF=true38 OPTIONS -DENABLE_CV_COMM_VIA_SSBUF=true
39 )39 )
40+elseif("${CONDITION_UNIT}" STREQUAL "ascend960pr")
41+ set(rotary_position_embedding_apt_depends posembedding/apply_rotary_pos_emb CACHE INTERNAL "Dependencies for rotary_position_embedding")
42+ add_ops_compile_options(
43+ OP_NAME RotaryPositionEmbedding
44+ COMPUTE_UNIT Ascend960PR_969A1
45+ OPTIONS -DENABLE_CV_COMM_VIA_SSBUF=true
46+ )
40endif()47endif()
41 48 
42if (BUILD_OPS_RTY_KERNEL)49if (BUILD_OPS_RTY_KERNEL)
Mposembedding/rotary_position_embedding/op_host/rotary_position_embedding_def.cpp+1-0
@@ -64,6 +64,7 @@ public:
64 .ExtendCfgInfo("opFile.value", "rotary_position_embedding_apt");64 .ExtendCfgInfo("opFile.value", "rotary_position_embedding_apt");
65 this->AICore().AddConfig("ascend950", regbaseCfg);65 this->AICore().AddConfig("ascend950", regbaseCfg);
66 this->AICore().AddConfig("ascend960dt", regbaseCfg);66 this->AICore().AddConfig("ascend960dt", regbaseCfg);
67+ this->AICore().AddConfig("ascend960pr", regbaseCfg);
67 this->AICore().AddConfig("ascend350", regbaseCfg);68 this->AICore().AddConfig("ascend350", regbaseCfg);
68 this->AICore().AddConfig("ascend5162", regbaseCfg);69 this->AICore().AddConfig("ascend5162", regbaseCfg);
69 70 
Mposembedding/rotary_position_embedding_grad/op_graph/fusion_pass/rotary_position_embedding_grad_fusion_pass.cpp+3-2
@@ -42,8 +42,8 @@ constexpr size_t kDyIdx = 3;
42 42 
43// Same gating as the built-in pass: fuse only on regbase platforms (NpuArch DAV_3510/DAV_9201/DAV_5102).43// Same gating as the built-in pass: fuse only on regbase platforms (NpuArch DAV_3510/DAV_9201/DAV_5102).
44// Gating on NpuArch instead of a short_soc_version string list: one arch covers multiple SoCs44// Gating on NpuArch instead of a short_soc_version string list: one arch covers multiple SoCs
45-// (DAV_3510: Ascend950/Ascend350, DAV_9201: Ascend960DT, DAV_5102: MC62CM12A), so SoCs sharing the regbase45+// (DAV_3510: Ascend950/Ascend350, DAV_9201/DAV_9202: Ascend960DT/Ascend960PR, DAV_5102: MC62CM12A), so SoCs sharing the
46-// arch are covered without maintaining the soc list.46+// regbase arch are covered without maintaining the soc list.
47bool IsRegBasePlatform()47bool IsRegBasePlatform()
48{48{
49 char npuArchVal[16] = {0};49 char npuArchVal[16] = {0};
@@ -55,6 +55,7 @@ bool IsRegBasePlatform()
55 const std::string npuArch(npuArchVal);55 const std::string npuArch(npuArchVal);
56 return npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_3510)) ||56 return npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_3510)) ||
57 npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_9201)) ||57 npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_9201)) ||
58+ npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_9202)) ||
58 npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_5102));59 npuArch == std::to_string(static_cast<uint32_t>(Ops::Base::DAV_5102));
59}60}
60 61 
Mposembedding/rotary_position_embedding_grad/op_host/CMakeLists.txt+2-0
@@ -22,6 +22,8 @@ elseif("${CONDITION_UNIT}" STREQUAL "ascend350")
22 set(rotary_position_embedding_grad_apt_depends posembedding/apply_rotary_pos_emb CACHE INTERNAL "Dependencies for rotary_position_embedding_grad")22 set(rotary_position_embedding_grad_apt_depends posembedding/apply_rotary_pos_emb CACHE INTERNAL "Dependencies for rotary_position_embedding_grad")
23elseif("${CONDITION_UNIT}" STREQUAL "ascend960dt")23elseif("${CONDITION_UNIT}" STREQUAL "ascend960dt")
24 set(rotary_position_embedding_grad_apt_depends posembedding/apply_rotary_pos_emb CACHE INTERNAL "Dependencies for rotary_position_embedding_grad")24 set(rotary_position_embedding_grad_apt_depends posembedding/apply_rotary_pos_emb CACHE INTERNAL "Dependencies for rotary_position_embedding_grad")
25+elseif("${CONDITION_UNIT}" STREQUAL "ascend960pr")
26+ set(rotary_position_embedding_grad_apt_depends posembedding/apply_rotary_pos_emb CACHE INTERNAL "Dependencies for rotary_position_embedding_grad")
25endif()27endif()
26 28 
27if (BUILD_OPS_RTY_KERNEL)29if (BUILD_OPS_RTY_KERNEL)
Mposembedding/rotary_position_embedding_grad/op_host/rotary_position_embedding_grad_def.cpp+1-0
@@ -77,6 +77,7 @@ public:
77 .ExtendCfgInfo("opFile.value", "rotary_position_embedding_grad_apt");77 .ExtendCfgInfo("opFile.value", "rotary_position_embedding_grad_apt");
78 this->AICore().AddConfig("ascend950", regbaseCfg);78 this->AICore().AddConfig("ascend950", regbaseCfg);
79 this->AICore().AddConfig("ascend960dt", regbaseCfg);79 this->AICore().AddConfig("ascend960dt", regbaseCfg);
80+ this->AICore().AddConfig("ascend960pr", regbaseCfg);
80 this->AICore().AddConfig("ascend350", regbaseCfg);81 this->AICore().AddConfig("ascend350", regbaseCfg);
81 }82 }
82};83};