Ascend960 moe_update_expert&barrier&combine-ops-transformer-AtomGit

Mmc2/common/op_host/op_tiling/mc2_tiling_utils.h+2-1

	@@ -254,7 +254,8 @@ inline ge::graphStatus GetEpWinSize(const gert::TilingContext* context, const ch
254	bool isLayered)	254	bool isLayered)
255	{	255	{
256	auto attrs = context->GetAttrs();	256	auto attrs = context->GetAttrs();
257	- if (mc2tiling::GetNpuArch(context) == Ops::Base::DAV_3510) {	257	+ NpuArch npuArch = mc2tiling::GetNpuArch(context);
		258	+ if (npuArch == Ops::Base::DAV_3510 \|\| npuArch == Ops::Base::DAV_9201 \|\| npuArch == Ops::Base::DAV_9202) {
		Y Y yayahello26 days ago 新增了 A6 的分支判断，是不是新增两个 inline 的函数，增加一下可读性；并且部分注释只说明了 A5，后续也需要同步修改一下 inline bool IsA5(arch); inline bool IsA6(arch); Y yezhuoxi 26 days ago Comments： 26 days ago Comments：
258	// A5 暂不支持 Hccl CommGetBufSizeCfg 接口，此处暂作规避	259	// A5 暂不支持 Hccl CommGetBufSizeCfg 接口，此处暂作规避
259	// A5 实际物理分配为 HCCL_BUFFSIZE 的 2 倍	260	// A5 实际物理分配为 HCCL_BUFFSIZE 的 2 倍
260	hcclBufferSizeEp = mc2tiling::Mc2TilingUtils::GetMaxWindowSize() * 2UL;	261	hcclBufferSizeEp = mc2tiling::Mc2TilingUtils::GetMaxWindowSize() * 2UL;

Mmc2/distribute_barrier/op_api/distribute_barrier_base.cpp+4-2

	@@ -91,7 +91,8 @@ aclnnStatus aclnnDistributeBarrierGetWorkspaceSizeBase(const aclTensor* xRef, co
91	(void)strncpy_s(groupBuf, HCCL_GROUP_NAME_MAX, group, HCCL_GROUP_NAME_MAX - 1);	91	(void)strncpy_s(groupBuf, HCCL_GROUP_NAME_MAX, group, HCCL_GROUP_NAME_MAX - 1);
92	}	92	}
93	aclTensor* xRefBuf = const_cast<aclTensor*>(xRef);	93	aclTensor* xRefBuf = const_cast<aclTensor*>(xRef);
94	- if (GetCurrentPlatformInfo().GetCurNpuArch() != Ops::Base::DAV_3510) {	94	+ auto npuArch = GetCurrentPlatformInfo().GetCurNpuArch();
		95	+ if (npuArch != Ops::Base::DAV_3510 && npuArch != Ops::Base::DAV_9201 && npuArch != Ops::Base::DAV_9202) {
95	getWorkspaceSizesRes = aclnnInnerDistributeBarrierGetWorkspaceSize(xRefBuf, timeOut, elasticInfo, groupBuf,	96	getWorkspaceSizesRes = aclnnInnerDistributeBarrierGetWorkspaceSize(xRefBuf, timeOut, elasticInfo, groupBuf,
96	worldSize, workspaceSize, executor);	97	worldSize, workspaceSize, executor);
97	} else {	98	} else {
	@@ -114,7 +115,8 @@ aclnnStatus aclnnDistributeBarrierBase(void* workspace, uint64_t workspaceSize,
114	NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE);	115	NnopbaseSetHcclServerType(executor, NNOPBASE_HCCL_SERVER_TYPE_MTE);
115	}	116	}
116	#if HCOMM_VERSION_NUM >= HCCL_CHANNEL_SUPPORT_VERSION	117	#if HCOMM_VERSION_NUM >= HCCL_CHANNEL_SUPPORT_VERSION
117	- if (GetCurrentPlatformInfo().GetCurNpuArch() == Ops::Base::DAV_3510) {	118	+ auto npuArch = GetCurrentPlatformInfo().GetCurNpuArch();
		119	+ if (npuArch == Ops::Base::DAV_3510 \|\| npuArch == Ops::Base::DAV_9201 \|\| npuArch == Ops::Base::DAV_9202) {
118	OP_LOGD("aclnn_distribute_barrier_extend inner start");	120	OP_LOGD("aclnn_distribute_barrier_extend inner start");
119	return aclnnInnerDistributeBarrierExtend(workspace, workspaceSize, executor, stream);	121	return aclnnInnerDistributeBarrierExtend(workspace, workspaceSize, executor, stream);
120	}	122	}

Mmc2/distribute_barrier/op_host/distribute_barrier_def.cpp+2-0

	@@ -76,6 +76,8 @@ public:
76		76
77	this->AICore().AddConfig("ascend910_93", aicore_config_a3);	77	this->AICore().AddConfig("ascend910_93", aicore_config_a3);
78	this->AICore().AddConfig("ascend950", aicore_config_apt);	78	this->AICore().AddConfig("ascend950", aicore_config_apt);
		79	+ this->AICore().AddConfig("ascend960dt", aicore_config_apt);
		80	+ this->AICore().AddConfig("ascend960pr", aicore_config_apt);
79	this->MC2().HcclGroup({"group"});	81	this->MC2().HcclGroup({"group"});
80	}	82	}
81	};	83	};

Mmc2/distribute_barrier/op_host/op_tiling/distribute_barrier_tiling.cpp+1-1

	@@ -120,7 +120,7 @@ static bool CheckAndSetAttrs(const gert::TilingContext* context, DistributeBarri
120	OP_TILING_CHECK(attrs == nullptr, OPS_LOG_E(A_INNER_DEBUG_BARRIER, "GetAttrs returned nullptr!"), return false);	120	OP_TILING_CHECK(attrs == nullptr, OPS_LOG_E(A_INNER_DEBUG_BARRIER, "GetAttrs returned nullptr!"), return false);
121	NpuArch npuArch = mc2tiling::GetNpuArch(context);	121	NpuArch npuArch = mc2tiling::GetNpuArch(context);
122	int64_t MAX_WORLD_SIZE = 0;	122	int64_t MAX_WORLD_SIZE = 0;
123	- if (npuArch == Ops::Base::DAV_3510) {	123	+ if (npuArch == Ops::Base::DAV_3510 \|\| npuArch == Ops::Base::DAV_9201 \|\| npuArch == Ops::Base::DAV_9202) {
124	MAX_WORLD_SIZE = MAX_WORLD_SIZE_A5;	124	MAX_WORLD_SIZE = MAX_WORLD_SIZE_A5;
125	} else {	125	} else {
126	MAX_WORLD_SIZE = MAX_WORLD_SIZE_A3;	126	MAX_WORLD_SIZE = MAX_WORLD_SIZE_A3;

Mmc2/distribute_barrier_extend/op_host/distribute_barrier_extend_def.cpp+2-0

	@@ -67,6 +67,8 @@ public:
67	.ExtendCfgInfo("opFile.value", "distribute_barrier_extend_apt");	67	.ExtendCfgInfo("opFile.value", "distribute_barrier_extend_apt");
68		68
69	this->AICore().AddConfig("ascend950", aicore_config_apt);	69	this->AICore().AddConfig("ascend950", aicore_config_apt);
		70	+ this->AICore().AddConfig("ascend960dt", aicore_config_apt);
		71	+ this->AICore().AddConfig("ascend960pr", aicore_config_apt);
70	}	72	}
71	};	73	};
72		74

Mmc2/moe_distribute_combine/op_host/moe_distribute_combine_def.cpp+2-0

	@@ -76,6 +76,8 @@ public:
76	this->AICore().AddConfig("ascend910b", aicore_config_a2);	76	this->AICore().AddConfig("ascend910b", aicore_config_a2);
77	this->AICore().AddConfig("ascend910_93", aicore_config_a3);	77	this->AICore().AddConfig("ascend910_93", aicore_config_a3);
78	this->AICore().AddConfig("ascend950", aicore_config_apt);	78	this->AICore().AddConfig("ascend950", aicore_config_apt);
		79	+ this->AICore().AddConfig("ascend960dt", aicore_config_apt);
		80	+ this->AICore().AddConfig("ascend960pr", aicore_config_apt);
79	this->MC2().HcclGroup({"group_ep", "group_tp"});	81	this->MC2().HcclGroup({"group_ep", "group_tp"});
80	}	82	}
81		83

Mmc2/moe_distribute_combine/op_host/op_tiling/arch22/moe_distribute_combine_tiling_a2a3.cpp+5-1

	@@ -429,7 +429,11 @@ ge::graphStatus MoeDistributeCombineTilingA2A3::MoeDistributeCombineA3TilingFunc
429	if (commQuantMode == INT8_COMM_QUANT) {	429	if (commQuantMode == INT8_COMM_QUANT) {
430	quantMode = TILINGKEY_INT8_QUANT;	430	quantMode = TILINGKEY_INT8_QUANT;
431	}	431	}
432	- uint32_t archTag = (mc2tiling::GetNpuArch(context) == Ops::Base::DAV_3510) ? TILINGKEY_TPL_A5 : TILINGKEY_TPL_A3;	432	+ NpuArch npuArch = mc2tiling::GetNpuArch(context);
		433	+ uint32_t archTag =
		434	+ (npuArch == Ops::Base::DAV_3510 \|\| npuArch == Ops::Base::DAV_9201 \|\| npuArch == Ops::Base::DAV_9202) ?
		435	+ TILINGKEY_TPL_A5 :
		436	+ TILINGKEY_TPL_A3;
433	const uint64_t tilingKey = GET_TPL_TILING_KEY(quantMode, layeredMode, archTag);	437	const uint64_t tilingKey = GET_TPL_TILING_KEY(quantMode, layeredMode, archTag);
434	OP_LOGD(nodeName, "tilingKey is %lu", tilingKey);	438	OP_LOGD(nodeName, "tilingKey is %lu", tilingKey);
435	context->SetTilingKey(tilingKey);	439	context->SetTilingKey(tilingKey);

Mmc2/moe_distribute_combine/op_host/op_tiling/arch35/moe_distribute_combine_tiling_a5.cpp+5-1

	@@ -107,7 +107,11 @@ ge::graphStatus MoeDistributeCombineTilingA5::MoeDistributeCombineA5TilingFuncIm
107	if (commQuantMode == INT8_COMM_QUANT) {	107	if (commQuantMode == INT8_COMM_QUANT) {
108	quantMode = TILINGKEY_INT8_QUANT;	108	quantMode = TILINGKEY_INT8_QUANT;
109	}	109	}
110	- uint32_t archTag = (mc2tiling::GetNpuArch(context) == Ops::Base::DAV_3510) ? TILINGKEY_TPL_A5 : TILINGKEY_TPL_A3;	110	+ NpuArch npuArch = mc2tiling::GetNpuArch(context);
		111	+ uint32_t archTag =
		112	+ (npuArch == Ops::Base::DAV_3510 \|\| npuArch == Ops::Base::DAV_9201 \|\| npuArch == Ops::Base::DAV_9202) ?
		113	+ TILINGKEY_TPL_A5 :
		114	+ TILINGKEY_TPL_A3;
111	const uint64_t tilingKey = GET_TPL_TILING_KEY(quantMode, layeredMode, archTag);	115	const uint64_t tilingKey = GET_TPL_TILING_KEY(quantMode, layeredMode, archTag);
112	OP_LOGD(nodeName, "tilingKey is %lu", tilingKey);	116	OP_LOGD(nodeName, "tilingKey is %lu", tilingKey);
113	context->SetTilingKey(tilingKey);	117	context->SetTilingKey(tilingKey);

Mmc2/moe_distribute_combine_add_rms_norm/op_host/moe_distribute_combine_add_rms_norm_def.cpp+2-0

	@@ -56,6 +56,8 @@ public:
56		56
57	this->AICore().AddConfig("ascend910_93", aicore_config_a3);	57	this->AICore().AddConfig("ascend910_93", aicore_config_a3);
58	this->AICore().AddConfig("ascend950", aicore_config_apt);	58	this->AICore().AddConfig("ascend950", aicore_config_apt);
		59	+ this->AICore().AddConfig("ascend960dt", aicore_config_apt);
		60	+ this->AICore().AddConfig("ascend960pr", aicore_config_apt);
59	this->MC2().HcclGroup({"group_ep", "group_tp"});	61	this->MC2().HcclGroup({"group_ep", "group_tp"});
60	}	62	}
61		63

Mmc2/moe_distribute_combine_v2/op_api/moe_distribute_combine_v2_base.cpp+15-11

	@@ -121,10 +121,10 @@ static inline void SafeCopyGroupBuf(char* dst, size_t dstSize, const char* src,
121	}	121	}
122	}	122	}
123		123
124	-static void SetCommArgs(aclOpExecutor** executor, const bool is910B, const bool isA5, const char* commAlg)	124	+static void SetCommArgs(aclOpExecutor** executor, const bool is910B, const bool isA5OrA6, const char* commAlg)
125	{	125	{
126	#if HCOMM_VERSION_NUM >= HCCL_CHANNEL_SUPPORT_VERSION	126	#if HCOMM_VERSION_NUM >= HCCL_CHANNEL_SUPPORT_VERSION
127	- if (isA5) {	127	+ if (isA5OrA6) {
128	CommType type = CommType::AIV;	128	CommType type = CommType::AIV;
129	if (commAlg != nullptr && std::strcmp(commAlg, "ccu") == 0) {	129	if (commAlg != nullptr && std::strcmp(commAlg, "ccu") == 0) {
130	type = CommType::CCU;	130	type = CommType::CCU;
	@@ -137,7 +137,7 @@ static void SetCommArgs(aclOpExecutor** executor, const bool is910B, const bool
137	if (NnopbaseSetHcclServerType) {	137	if (NnopbaseSetHcclServerType) {
138	if (is910B) {	138	if (is910B) {
139	NnopbaseSetHcclServerType(*executor, NNOPBASE_HCCL_SERVER_TYPE_AICPU);	139	NnopbaseSetHcclServerType(*executor, NNOPBASE_HCCL_SERVER_TYPE_AICPU);
140	- } else if (isA5 && commAlg != nullptr && std::strcmp(commAlg, "ccu") == 0) {	140	+ } else if (isA5OrA6 && commAlg != nullptr && std::strcmp(commAlg, "ccu") == 0) {
141	NnopbaseSetHcclServerType(*executor, NNOPBASE_HCCL_SERVER_TYPE_CCU);	141	NnopbaseSetHcclServerType(*executor, NNOPBASE_HCCL_SERVER_TYPE_CCU);
142	} else {	142	} else {
143	NnopbaseSetHcclServerType(*executor, NNOPBASE_HCCL_SERVER_TYPE_MTE);	143	NnopbaseSetHcclServerType(*executor, NNOPBASE_HCCL_SERVER_TYPE_MTE);
	@@ -181,12 +181,12 @@ static aclnnStatus ResolveInnerWorkspaceSize(
181	char* groupEpBuf, char* groupTpBuf, const char* commAlg, char* commAlgBuf, int64_t epWorldSize, int64_t epRankId,	181	char* groupEpBuf, char* groupTpBuf, const char* commAlg, char* commAlgBuf, int64_t epWorldSize, int64_t epRankId,
182	int64_t moeExpertNum, int64_t tpWorldSize, int64_t tpRankId, int64_t expertShardType, int64_t sharedExpertNum,	182	int64_t moeExpertNum, int64_t tpWorldSize, int64_t tpRankId, int64_t expertShardType, int64_t sharedExpertNum,
183	int64_t sharedExpertRankNum, int64_t globalBs, int64_t outDtype, int64_t commQuantMode, int64_t groupListType,	183	int64_t sharedExpertRankNum, int64_t globalBs, int64_t outDtype, int64_t commQuantMode, int64_t groupListType,
184	- int64_t zeroExpertNum, int64_t copyExpertNum, int64_t constExpertNum, bool is910B, bool isA5, aclTensor* xOut,	184	+ int64_t zeroExpertNum, int64_t copyExpertNum, int64_t constExpertNum, bool is910B, bool isA5OrA6, aclTensor* xOut,
185	uint64_t* workspaceSize, aclOpExecutor** executor)	185	uint64_t* workspaceSize, aclOpExecutor** executor)
186	{	186	{
187	aclnnStatus getWorkspaceSizesRes = ACLNN_ERR_INNER;	187	aclnnStatus getWorkspaceSizesRes = ACLNN_ERR_INNER;
188	// ccu暂时不支持mc2Context，回退到V2接口	188	// ccu暂时不支持mc2Context，回退到V2接口
189	- if (!isA5 \|\| (commAlg != nullptr && std::strcmp(commAlg, "ccu") == 0)) {	189	+ if (!isA5OrA6 \|\| (commAlg != nullptr && std::strcmp(commAlg, "ccu") == 0)) {
190	getWorkspaceSizesRes = aclnnInnerMoeDistributeCombineV2GetWorkspaceSize(	190	getWorkspaceSizesRes = aclnnInnerMoeDistributeCombineV2GetWorkspaceSize(
191	expandX, expertIds, assistInfoForCombine, epSendCounts, expertScales, tpSendCountsOptional,	191	expandX, expertIds, assistInfoForCombine, epSendCounts, expertScales, tpSendCountsOptional,
192	xActiveMaskOptional, activationScaleOptional, weightScaleOptional, groupListOptional, expandScalesOptional,	192	xActiveMaskOptional, activationScaleOptional, weightScaleOptional, groupListOptional, expandScalesOptional,
	@@ -212,7 +212,7 @@ static aclnnStatus ResolveInnerWorkspaceSize(
212	xOut, workspaceSize, executor);	212	xOut, workspaceSize, executor);
213	#endif	213	#endif
214	}	214	}
215	- SetCommArgs(executor, is910B, isA5, commAlg);	215	+ SetCommArgs(executor, is910B, isA5OrA6, commAlg);
216	return getWorkspaceSizesRes;	216	return getWorkspaceSizesRes;
217	}	217	}
218		218
	@@ -230,8 +230,10 @@ aclnnStatus aclnnMoeDistributeCombineBaseGetWorkspaceSize(
230	int64_t copyExpertNum, int64_t constExpertNum, aclTensor* xOut, uint64_t* workspaceSize, aclOpExecutor** executor)	230	int64_t copyExpertNum, int64_t constExpertNum, aclTensor* xOut, uint64_t* workspaceSize, aclOpExecutor** executor)
231	{	231	{
232	OP_LOGD("enter to the aclnnMoeDistributeCombineBaseGetWorkspaceSize\n");	232	OP_LOGD("enter to the aclnnMoeDistributeCombineBaseGetWorkspaceSize\n");
		233	+ auto npuArch = GetCurrentPlatformInfo().GetCurNpuArch();
233	const static bool is910B = GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910B;	234	const static bool is910B = GetCurrentPlatformInfo().GetSocVersion() == SocVersion::ASCEND910B;
234	- const static bool isA5 = GetCurrentPlatformInfo().GetCurNpuArch() == Ops::Base::DAV_3510;	235	+ const static bool isA5OrA6 =
		236	+ (npuArch == Ops::Base::DAV_3510 \|\| npuArch == Ops::Base::DAV_9201 \|\| npuArch == Ops::Base::DAV_9202);
235	auto retParam = CombineCheckParams(expandX, expertIds, assistInfoForCombine, epSendCounts, expertScales, groupEp,	237	auto retParam = CombineCheckParams(expandX, expertIds, assistInfoForCombine, epSendCounts, expertScales, groupEp,
236	groupTp, xOut);	238	groupTp, xOut);
237	CHECK_RET(retParam == ACLNN_SUCCESS, retParam);	239	CHECK_RET(retParam == ACLNN_SUCCESS, retParam);
	@@ -247,8 +249,8 @@ aclnnStatus aclnnMoeDistributeCombineBaseGetWorkspaceSize(
247	elasticInfoOptional, oriXOptional, constExpertAlpha1Optional, constExpertAlpha2Optional, constExpertVOptional,	249	elasticInfoOptional, oriXOptional, constExpertAlpha1Optional, constExpertAlpha2Optional, constExpertVOptional,
248	performanceInfoOptional, groupEp, groupEpBuf, groupTpBuf, commAlg, commAlgBuf, epWorldSize, epRankId,	250	performanceInfoOptional, groupEp, groupEpBuf, groupTpBuf, commAlg, commAlgBuf, epWorldSize, epRankId,
249	moeExpertNum, tpWorldSize, tpRankId, expertShardType, sharedExpertNum, sharedExpertRankNum, globalBs, outDtype,	251	moeExpertNum, tpWorldSize, tpRankId, expertShardType, sharedExpertNum, sharedExpertRankNum, globalBs, outDtype,
250	- commQuantMode, groupListType, zeroExpertNum, copyExpertNum, constExpertNum, is910B, isA5, xOut, workspaceSize,	252	+ commQuantMode, groupListType, zeroExpertNum, copyExpertNum, constExpertNum, is910B, isA5OrA6, xOut,
251	- executor);	253	+ workspaceSize, executor);
252	}	254	}
253		255
254	// aclnn二段式接口	256	// aclnn二段式接口
	@@ -256,8 +258,10 @@ aclnnStatus aclnnMoeDistributeCombineBase(void* workspace, uint64_t workspaceSiz
256	aclrtStream stream)	258	aclrtStream stream)
257	{	259	{
258	#if HCOMM_VERSION_NUM >= HCCL_CHANNEL_SUPPORT_VERSION	260	#if HCOMM_VERSION_NUM >= HCCL_CHANNEL_SUPPORT_VERSION
259	- const static bool isA5 = GetCurrentPlatformInfo().GetCurNpuArch() == Ops::Base::DAV_3510;	261	+ auto npuArch = GetCurrentPlatformInfo().GetCurNpuArch();
260	- if (isA5) {	262	+ const static bool isA5OrA6 =
		263	+ (npuArch == Ops::Base::DAV_3510 \|\| npuArch == Ops::Base::DAV_9201 \|\| npuArch == Ops::Base::DAV_9202);
		264	+ if (isA5OrA6) {
261	void* arg = NnopbaseGetUserHandle(executor);	265	void* arg = NnopbaseGetUserHandle(executor);
262	uintptr_t handleVal = reinterpret_cast<uintptr_t>(arg);	266	uintptr_t handleVal = reinterpret_cast<uintptr_t>(arg);
263	CommType commType = static_cast<CommType>(handleVal);	267	CommType commType = static_cast<CommType>(handleVal);

Mmc2/moe_distribute_combine_v2/op_host/moe_distribute_combine_v2_def.cpp+2-0

	@@ -78,6 +78,8 @@ public:
78	this->AICore().AddConfig("ascend910b", aicore_config_a2);	78	this->AICore().AddConfig("ascend910b", aicore_config_a2);
79	this->AICore().AddConfig("ascend910_93", aicore_config_a3);	79	this->AICore().AddConfig("ascend910_93", aicore_config_a3);
80	this->AICore().AddConfig("ascend950", aicore_config_apt);	80	this->AICore().AddConfig("ascend950", aicore_config_apt);
		81	+ this->AICore().AddConfig("ascend960dt", aicore_config_apt);
		82	+ this->AICore().AddConfig("ascend960pr", aicore_config_apt);
81	this->MC2().HcclGroup({"group_ep", "group_tp"});	83	this->MC2().HcclGroup({"group_ep", "group_tp"});
82	}	84	}
83		85

Mmc2/moe_distribute_combine_v2/op_host/op_tiling/arch35/moe_distribute_combine_tiling_arch35.cpp+1-1

	@@ -808,7 +808,7 @@ uint64_t MoeDistributeCombineV2TilingFuncA5::CalTilingKey(uint32_t commQuantMode
808		808
809	bool MoeDistributeCombineTilingA5::IsCapable()	809	bool MoeDistributeCombineTilingA5::IsCapable()
810	{	810	{
811	- if (npuArch_ == Ops::Base::DAV_3510) {	811	+ if (npuArch_ == Ops::Base::DAV_3510 \|\| npuArch_ == Ops::Base::DAV_9201 \|\| npuArch_ == Ops::Base::DAV_9202) {
812	const char* nodeName = context_->GetNodeName();	812	const char* nodeName = context_->GetNodeName();
813	OP_LOGD(nodeName, "Do MoeDistributeCombineTilingA5 tiling.");	813	OP_LOGD(nodeName, "Do MoeDistributeCombineTilingA5 tiling.");
814	return true;	814	return true;

Mmc2/moe_distribute_combine_v2/op_host/op_tiling/moe_distribute_combine_v2_tiling.cpp+22-16

Mmc2/moe_distribute_combine_v2/op_kernel/arch22/moe_distribute_combine_v2.h+6-5

	@@ -489,7 +489,7 @@ __aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::InitAttrs(
489	uint32_t hFloatSize = axisH_ * static_cast<uint32_t>(sizeof(float));	489	uint32_t hFloatSize = axisH_ * static_cast<uint32_t>(sizeof(float));
490	hFloatAlign32Size_ = Ceil(hFloatSize, UB_ALIGN) * UB_ALIGN;	490	hFloatAlign32Size_ = Ceil(hFloatSize, UB_ALIGN) * UB_ALIGN;
491	hFloatAlign256Size_ = Ceil(hFloatSize, ALIGNED_LEN) * ALIGNED_LEN;	491	hFloatAlign256Size_ = Ceil(hFloatSize, ALIGNED_LEN) * ALIGNED_LEN;
492	-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)	492	+#if defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202))
493	// A5 INT8/MXFP8 融合反量化每轮固定访问 128 floats，工作区按 512B 对齐避免尾段越界	493	// A5 INT8/MXFP8 融合反量化每轮固定访问 128 floats，工作区按 512B 对齐避免尾段越界
494	if constexpr ((QuantMode == INT8_COMM_QUANT) \|\| (QuantMode == MXFP8_E5M2_COMM_QUANT) \|\|	494	if constexpr ((QuantMode == INT8_COMM_QUANT) \|\| (QuantMode == MXFP8_E5M2_COMM_QUANT) \|\|
495	(QuantMode == MXFP8_E4M3_COMM_QUANT)) {	495	(QuantMode == MXFP8_E4M3_COMM_QUANT)) {
	@@ -575,7 +575,8 @@ template <CombineMC2TypeClass>
575	__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::BuffInit()	575	__aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::BuffInit()
576	{	576	{
577	tpipe_->Reset();	577	tpipe_->Reset();
578	-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510) // reset后 ctrl寄存器会复位为默认值	578	+#if defined(__NPU_ARCH__) && \
		579	+ ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202)) // reset后 ctrl寄存器会复位为默认值
579	// 单指令饱和模式	580	// 单指令饱和模式
580	AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);	581	AscendC::SetCtrlSpr<FLOAT_OVERFLOW_MODE_CTRL, FLOAT_OVERFLOW_MODE_CTRL>(0);
581	#endif	582	#endif
	@@ -981,7 +982,7 @@ __aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::PrefetchDispa
981	uint32_t slot = sendInputCount_ % BUFFER_NUM;	982	uint32_t slot = sendInputCount_ % BUFFER_NUM;
982	WaitFlag<HardEvent::V_MTE2>(slot);	983	WaitFlag<HardEvent::V_MTE2>(slot);
983	expandXInTensor_ = sendInputBuf_.Get<ExpandXType>()[slot * hExpandXAlignSize_ / sizeof(ExpandXType)];	984	expandXInTensor_ = sendInputBuf_.Get<ExpandXType>()[slot * hExpandXAlignSize_ / sizeof(ExpandXType)];
984	-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)	985	+#if defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202))
985	LocalTensor<uint8_t> singleByteTok = expandXInTensor_.template ReinterpretCast<uint8_t>();	986	LocalTensor<uint8_t> singleByteTok = expandXInTensor_.template ReinterpretCast<uint8_t>();
986	// 由于MX量化在计算scales时每次搬入256字节数据，所以在token搬入前需要对空间填0，避免引入脏数据	987	// 由于MX量化在计算scales时每次搬入256字节数据，所以在token搬入前需要对空间填0，避免引入脏数据
987	if constexpr ((QuantMode == MXFP8_E5M2_COMM_QUANT) \|\| (QuantMode == MXFP8_E4M3_COMM_QUANT)) {	988	if constexpr ((QuantMode == MXFP8_E5M2_COMM_QUANT) \|\| (QuantMode == MXFP8_E4M3_COMM_QUANT)) {
	@@ -1329,7 +1330,7 @@ __aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::ProcessMoeExp
1329	} else {	1330	} else {
1330	quantInst_.DeQuantProcessWithoutExpertScale(tmpUb, outLocalTensor, rowTmpFloatLocal_, sumFloatBufLocal_);	1331	quantInst_.DeQuantProcessWithoutExpertScale(tmpUb, outLocalTensor, rowTmpFloatLocal_, sumFloatBufLocal_);
1331	}	1332	}
1332	-#if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510))	1333	+#if !(defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202)))
1333	// 非 A5 在外层完成 INT8 加权求和	1334	// 非 A5 在外层完成 INT8 加权求和
1334	if constexpr (QuantMode == INT8_COMM_QUANT) {	1335	if constexpr (QuantMode == INT8_COMM_QUANT) {
1335	Cast(rowTmpFloatLocal_, outLocalTensor, AscendC::RoundMode::CAST_NONE, processLen);	1336	Cast(rowTmpFloatLocal_, outLocalTensor, AscendC::RoundMode::CAST_NONE, processLen);
	@@ -1460,7 +1461,7 @@ __aicore__ inline void MoeDistributeCombineV2<CombineMC2TypeFunc>::ProcessShared
1460	LocalTensor<XType> outLocalTensor = fp16CastTensor_.template ReinterpretCast<XType>();	1461	LocalTensor<XType> outLocalTensor = fp16CastTensor_.template ReinterpretCast<XType>();
1461	if constexpr (QuantMode > UNQUANT) {	1462	if constexpr (QuantMode > UNQUANT) {
1462	quantInst_.DeQuantProcess(tmpUb, outLocalTensor, rowTmpFloatLocal_, sumFloatBufLocal_, 1.0f);	1463	quantInst_.DeQuantProcess(tmpUb, outLocalTensor, rowTmpFloatLocal_, sumFloatBufLocal_, 1.0f);
1463	-#if !(defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510))	1464	+#if !(defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202)))
1464	// 非 A5 在外层完成 INT8 求和累加	1465	// 非 A5 在外层完成 INT8 求和累加
1465	if constexpr (QuantMode == INT8_COMM_QUANT) {	1466	if constexpr (QuantMode == INT8_COMM_QUANT) {
1466	Cast(rowTmpFloatLocal_, outLocalTensor, AscendC::RoundMode::CAST_NONE, processLen);	1467	Cast(rowTmpFloatLocal_, outLocalTensor, AscendC::RoundMode::CAST_NONE, processLen);

Mmc2/moe_distribute_combine_v2/op_kernel/moe_distribute_combine_v2_quant.h+8-8

	@@ -18,7 +18,7 @@
18		18
19	#include "../../moe_distribute_dispatch_v2/op_kernel/moe_distribute_v2_constant.h"	19	#include "../../moe_distribute_dispatch_v2/op_kernel/moe_distribute_v2_constant.h"
20	#include "../../moe_distribute_dispatch_v2/op_kernel/moe_distribute_v2_base.h"	20	#include "../../moe_distribute_dispatch_v2/op_kernel/moe_distribute_v2_base.h"
21	-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)	21	+#if defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202))
22	#include "../../moe_distribute_dispatch_v2/op_kernel/moe_distribute_dispatch_v2_common.h"	22	#include "../../moe_distribute_dispatch_v2/op_kernel/moe_distribute_dispatch_v2_common.h"
23	#include "../../common/op_kernel/quantize_functions.h"	23	#include "../../common/op_kernel/quantize_functions.h"
24	#endif	24	#endif
	@@ -82,7 +82,7 @@ public:
82	uint32_t scaleGranu =	82	uint32_t scaleGranu =
83	static_cast<uint32_t>(UB_ALIGN / sizeof(float)); // 计算每个block得到的reducemax结果数量	83	static_cast<uint32_t>(UB_ALIGN / sizeof(float)); // 计算每个block得到的reducemax结果数量
84	quantScaleNum_ = (hExpandXAlign32Size_ / sizeof(ExpandXType)) / scaleGranu; // 得到有效scale的个数	84	quantScaleNum_ = (hExpandXAlign32Size_ / sizeof(ExpandXType)) / scaleGranu; // 得到有效scale的个数
85	-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)	85	+#if defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202))
86	// A5 INT8融合反量化按512B覆盖尾段，scale数量需与Brcb的repeat覆盖范围一致	86	// A5 INT8融合反量化按512B覆盖尾段，scale数量需与Brcb的repeat覆盖范围一致
87	quantScaleNum_ = (hFloatAlign256Size_ / sizeof(float)) / scaleGranu;	87	quantScaleNum_ = (hFloatAlign256Size_ / sizeof(float)) / scaleGranu;
88	#endif	88	#endif
	@@ -93,7 +93,7 @@ public:
93	mask_ = static_cast<uint32_t>(ALIGNED_LEN / sizeof(float));	93	mask_ = static_cast<uint32_t>(ALIGNED_LEN / sizeof(float));
94	tokenScaleCnt_ = hAlign32Size_ / sizeof(ExpandXType) + quantScaleNum_; // int8_align + scale有效个数	94	tokenScaleCnt_ = hAlign32Size_ / sizeof(ExpandXType) + quantScaleNum_; // int8_align + scale有效个数
95	}	95	}
96	-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)	96	+#if defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202))
97	else if constexpr (QuantMode == MXFP8_E5M2_COMM_QUANT \|\| QuantMode == MXFP8_E4M3_COMM_QUANT) {	97	else if constexpr (QuantMode == MXFP8_E5M2_COMM_QUANT \|\| QuantMode == MXFP8_E4M3_COMM_QUANT) {
98	hExpandXAlignSize_ = Align128(axisH) * sizeof(ExpandXType);	98	hExpandXAlignSize_ = Align128(axisH) * sizeof(ExpandXType);
99	quantScaleNum_ = Align2(Ceil32(axisH));	99	quantScaleNum_ = Align2(Ceil32(axisH));
	@@ -147,7 +147,7 @@ public:
147	PipeBarrier<PIPE_V>();	147	PipeBarrier<PIPE_V>();
148	}	148	}
149		149
150	-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)	150	+#if defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202))
151	template <bool ApplyExpertScale = true>	151	template <bool ApplyExpertScale = true>
152	__aicore__ inline void Int8DequantProcessA5(LocalTensor<XType>& inLocal, LocalTensor<float>& sumFinalTensor,	152	__aicore__ inline void Int8DequantProcessA5(LocalTensor<XType>& inLocal, LocalTensor<float>& sumFinalTensor,
153	float scaleVal)	153	float scaleVal)
	@@ -361,7 +361,7 @@ public:
361	if constexpr (QuantMode == INT8_COMM_QUANT) {	361	if constexpr (QuantMode == INT8_COMM_QUANT) {
362	Int8QuantProcess(outLocal, inLocal);	362	Int8QuantProcess(outLocal, inLocal);
363	}	363	}
364	-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)	364	+#if defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202))
365	else if constexpr (QuantMode == MXFP8_E5M2_COMM_QUANT \|\| QuantMode == MXFP8_E4M3_COMM_QUANT) {	365	else if constexpr (QuantMode == MXFP8_E5M2_COMM_QUANT \|\| QuantMode == MXFP8_E4M3_COMM_QUANT) {
366	QuantMxFp8(outLocal, inLocal);	366	QuantMxFp8(outLocal, inLocal);
367	}	367	}
	@@ -372,14 +372,14 @@ public:
372	float scaleVal)	372	float scaleVal)
373	{	373	{
374	if constexpr (QuantMode == INT8_COMM_QUANT) {	374	if constexpr (QuantMode == INT8_COMM_QUANT) {
375	-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)	375	+#if defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202))
376	// A5 走反量化及加权求和融合	376	// A5 走反量化及加权求和融合
377	Int8DequantProcessA5(inLocal, sumFinalTensor, scaleVal);	377	Int8DequantProcessA5(inLocal, sumFinalTensor, scaleVal);
378	#else	378	#else
379	Int8DequantProcess(inLocal, outLocal);	379	Int8DequantProcess(inLocal, outLocal);
380	#endif	380	#endif
381	}	381	}
382	-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)	382	+#if defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202))
383	else if constexpr (QuantMode == MXFP8_E5M2_COMM_QUANT) {	383	else if constexpr (QuantMode == MXFP8_E5M2_COMM_QUANT) {
384	DeQuantMxFp8<fp8_e5m2_t>(inLocal, sumTensor, sumFinalTensor, scaleVal);	384	DeQuantMxFp8<fp8_e5m2_t>(inLocal, sumTensor, sumFinalTensor, scaleVal);
385	} else if constexpr (QuantMode == MXFP8_E4M3_COMM_QUANT) {	385	} else if constexpr (QuantMode == MXFP8_E4M3_COMM_QUANT) {
	@@ -392,7 +392,7 @@ public:
392	LocalTensor<float>& sumTensor,	392	LocalTensor<float>& sumTensor,
393	LocalTensor<float>& sumFinalTensor)	393	LocalTensor<float>& sumFinalTensor)
394	{	394	{
395	-#if defined(__NPU_ARCH__) && (__NPU_ARCH__ == 3510)	395	+#if defined(__NPU_ARCH__) && ((__NPU_ARCH__ == 3510) \|\| (__NPU_ARCH__ == 9201) \|\| (__NPU_ARCH__ == 9202))
396	if constexpr (QuantMode == INT8_COMM_QUANT) {	396	if constexpr (QuantMode == INT8_COMM_QUANT) {
397	Int8DequantProcessA5<false>(inLocal, sumFinalTensor, 0.0f);	397	Int8DequantProcessA5<false>(inLocal, sumFinalTensor, 0.0f);
398	} else if constexpr (QuantMode == MXFP8_E5M2_COMM_QUANT) {	398	} else if constexpr (QuantMode == MXFP8_E5M2_COMM_QUANT) {

Mmc2/moe_distribute_combine_v3/op_host/moe_distribute_combine_v3_def.cpp+2-0

	@@ -68,6 +68,8 @@ public:
68		68
69	this->AICore().AddConfig("ascend910_93", aicore_config_a3);	69	this->AICore().AddConfig("ascend910_93", aicore_config_a3);
70	this->AICore().AddConfig("ascend950", aicore_config_apt);	70	this->AICore().AddConfig("ascend950", aicore_config_apt);
		71	+ this->AICore().AddConfig("ascend960dt", aicore_config_apt);
		72	+ this->AICore().AddConfig("ascend960pr", aicore_config_apt);
71	}	73	}
72		74
73	private:	75	private:

Mmc2/moe_update_expert/op_host/moe_update_expert_def.cpp+2-0

	@@ -72,6 +72,8 @@ public:
72		72
73	this->AICore().AddConfig("ascend910_93", aicore_config);	73	this->AICore().AddConfig("ascend910_93", aicore_config);
74	this->AICore().AddConfig("ascend950", aicore_config);	74	this->AICore().AddConfig("ascend950", aicore_config);
		75	+ this->AICore().AddConfig("ascend960dt", aicore_config);
		76	+ this->AICore().AddConfig("ascend960pr", aicore_config);
75	}	77	}
76	};	78	};
77		79

	@@ -270,7 +270,7 @@ static ge::graphStatus CheckEpWorldSize(const gert::TilingContext* context, cons
270	{	270	{
271	NpuArch npuArch = mc2tiling::GetNpuArch(context);	271	NpuArch npuArch = mc2tiling::GetNpuArch(context);
272	int64_t maxEpworldsize = 0;	272	int64_t maxEpworldsize = 0;
273	- if (npuArch == Ops::Base::DAV_3510) {	273	+ if (npuArch == Ops::Base::DAV_3510 \|\| npuArch == Ops::Base::DAV_9201 \|\| npuArch == Ops::Base::DAV_9202) {
274	maxEpworldsize = isLayered ? MAX_EP_WORLD_SIZE_LAYERED : MAX_EP_WORLD_SIZE_A5;	274	maxEpworldsize = isLayered ? MAX_EP_WORLD_SIZE_LAYERED : MAX_EP_WORLD_SIZE_A5;
275	} else {	275	} else {
276	maxEpworldsize = isLayered ? MAX_EP_WORLD_SIZE_LAYERED : MAX_EP_WORLD_SIZE_A3;	276	maxEpworldsize = isLayered ? MAX_EP_WORLD_SIZE_LAYERED : MAX_EP_WORLD_SIZE_A3;
	@@ -352,7 +352,8 @@ static ge::graphStatus GetTpAndEpAttrAndSetTilingData(const gert::TilingContext*
352	static ge::graphStatus CheckCommunicationQuantMode(const gert::TilingContext* context, const char* nodeName,	352	static ge::graphStatus CheckCommunicationQuantMode(const gert::TilingContext* context, const char* nodeName,
353	const int64_t* commQuantModePtr)	353	const int64_t* commQuantModePtr)
354	{	354	{
355	- if (mc2tiling::GetNpuArch(context) != Ops::Base::DAV_3510) {	355	+ NpuArch npuArch = mc2tiling::GetNpuArch(context);
		356	+ if (npuArch != Ops::Base::DAV_3510 && npuArch != Ops::Base::DAV_9201 && npuArch != Ops::Base::DAV_9202) {
356	OP_TILING_CHECK((commQuantModePtr != 0) && (commQuantModePtr != INT8_COMM_QUANT),	357	OP_TILING_CHECK((commQuantModePtr != 0) && (commQuantModePtr != INT8_COMM_QUANT),
357	OP_LOGE_FOR_INVALID_VALUE(nodeName, "commQuantMode", std::to_string(*commQuantModePtr).c_str(),	358	OP_LOGE_FOR_INVALID_VALUE(nodeName, "commQuantMode", std::to_string(*commQuantModePtr).c_str(),
358	"should be 0 or 2"),	359	"should be 0 or 2"),
	@@ -1869,10 +1870,10 @@ static uint32_t GetA3PartialBufferSizeByFlag(const uint32_t axisBS, const uint32
1869	return partialBufferSize;	1870	return partialBufferSize;
1870	}	1871	}
1871		1872
1872	-static uint32_t GetFloatAlignSize(uint32_t hFloatSize, uint32_t alignSize, bool isA5FusedQuant)	1873	+static uint32_t GetFloatAlignSize(uint32_t hFloatSize, uint32_t alignSize, bool isA5OrA6FusedQuant)
1873	{	1874	{
1874	// A5 + INT8/MXFP8 场景按512B对齐，其余场景使用调用方指定的对齐大小	1875	// A5 + INT8/MXFP8 场景按512B对齐，其余场景使用调用方指定的对齐大小
1875	- if (isA5FusedQuant) {	1876	+ if (isA5OrA6FusedQuant) {
1876	alignSize = ALIGNED_LEN * 2U;	1877	alignSize = ALIGNED_LEN * 2U;
1877	}	1878	}
1878	return (hFloatSize + alignSize - 1U) / alignSize * alignSize;	1879	return (hFloatSize + alignSize - 1U) / alignSize * alignSize;
	@@ -1893,7 +1894,10 @@ static uint64_t GetComputeBufferSize(const gert::TilingContext* context,
1893	bool isInt8Quant = commQuantMode == INT8_COMM_QUANT;	1894	bool isInt8Quant = commQuantMode == INT8_COMM_QUANT;
1894	bool isMxFp8Quant = (commQuantMode == static_cast<CommQuantModeType>(CommQuantMode::MXFP8_E5M2_QUANT)) \|\|	1895	bool isMxFp8Quant = (commQuantMode == static_cast<CommQuantModeType>(CommQuantMode::MXFP8_E5M2_QUANT)) \|\|
1895	(commQuantMode == static_cast<CommQuantModeType>(CommQuantMode::MXFP8_E4M3_QUANT));	1896	(commQuantMode == static_cast<CommQuantModeType>(CommQuantMode::MXFP8_E4M3_QUANT));
1896	- bool isA5FusedQuant = mc2tiling::GetNpuArch(context) == Ops::Base::DAV_3510 && (isInt8Quant \|\| isMxFp8Quant);	1897	+ NpuArch npuArch = mc2tiling::GetNpuArch(context);
		1898	+ bool isA5OrA6FusedQuant =
		1899	+ (npuArch == Ops::Base::DAV_3510 \|\| npuArch == Ops::Base::DAV_9201 \|\| npuArch == Ops::Base::DAV_9202) &&
		1900	+ (isInt8Quant \|\| isMxFp8Quant);
1897	auto expandXDesc = context->GetInputDesc(config.expandXIndex);	1901	auto expandXDesc = context->GetInputDesc(config.expandXIndex);
1898	uint32_t expandXTypeSize = static_cast<uint32_t>(ge::GetSizeByDataType(expandXDesc->GetDataType()));	1902	uint32_t expandXTypeSize = static_cast<uint32_t>(ge::GetSizeByDataType(expandXDesc->GetDataType()));
1899	uint32_t hExpandXTypeSize = axisH * expandXTypeSize;	1903	uint32_t hExpandXTypeSize = axisH * expandXTypeSize;
	@@ -1901,8 +1905,8 @@ static uint64_t GetComputeBufferSize(const gert::TilingContext* context,
1901	uint32_t maxSizeTokenBuf = (hExpandXTypeSize + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;	1905	uint32_t maxSizeTokenBuf = (hExpandXTypeSize + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;
1902	uint32_t activeMaskAlignSize = axisBS * ((axisK * sizeof(bool) + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN);	1906	uint32_t activeMaskAlignSize = axisBS * ((axisK * sizeof(bool) + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN);
1903	uint32_t hExpandXAlign32Size = (hExpandXTypeSize + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;	1907	uint32_t hExpandXAlign32Size = (hExpandXTypeSize + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;
1904	- uint32_t hFloatAlign32Size = GetFloatAlignSize(hFloatSize, UB_ALIGN, isA5FusedQuant);	1908	+ uint32_t hFloatAlign32Size = GetFloatAlignSize(hFloatSize, UB_ALIGN, isA5OrA6FusedQuant);
1905	- uint32_t hFloatAlign256Size = GetFloatAlignSize(hFloatSize, ALIGNED_LEN, isA5FusedQuant);	1909	+ uint32_t hFloatAlign256Size = GetFloatAlignSize(hFloatSize, ALIGNED_LEN, isA5OrA6FusedQuant);
1906	uint32_t bsKFloatAlign = (axisBS * axisK * sizeof(float) + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;	1910	uint32_t bsKFloatAlign = (axisBS * axisK * sizeof(float) + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;
1907	uint32_t mulBufSize = hFloatAlign256Size > bsKFloatAlign ? hFloatAlign256Size : bsKFloatAlign;	1911	uint32_t mulBufSize = hFloatAlign256Size > bsKFloatAlign ? hFloatAlign256Size : bsKFloatAlign;
1908	uint32_t maxSizeRowTmpFloatBuf = hFloatAlign32Size;	1912	uint32_t maxSizeRowTmpFloatBuf = hFloatAlign32Size;
	@@ -2035,7 +2039,7 @@ static uint64_t GetQuantScratchSize(uint32_t axisH, uint32_t expandXTypeSize, bo
2035	static void GetCommQuantScratchAndDataSize(const gert::TilingContext* context, const CombineV2Config& config,	2039	static void GetCommQuantScratchAndDataSize(const gert::TilingContext* context, const CombineV2Config& config,
2036	uint32_t axisH, uint32_t& hExpandXTypeSize, uint32_t& hExpandXAlign32Size,	2040	uint32_t axisH, uint32_t& hExpandXTypeSize, uint32_t& hExpandXAlign32Size,
2037	uint32_t& scaleNumAlignSize, uint32_t& commDataBytes,	2041	uint32_t& scaleNumAlignSize, uint32_t& commDataBytes,
2038	- uint32_t& quantResultDataBytes, uint64_t& quantBufferSize, bool& isA5)	2042	+ uint32_t& quantResultDataBytes, uint64_t& quantBufferSize, bool& isA5OrA6)
2039	{	2043	{
2040	auto attrs = context->GetAttrs();	2044	auto attrs = context->GetAttrs();
2041	auto commQuantModePtr = attrs->GetAttrPointer<int64_t>((config.attrCommQuantModeIndex));	2045	auto commQuantModePtr = attrs->GetAttrPointer<int64_t>((config.attrCommQuantModeIndex));
	@@ -2043,27 +2047,28 @@ static void GetCommQuantScratchAndDataSize(const gert::TilingContext* context, c
2043	bool isInt8Quant = commQuantMode == INT8_COMM_QUANT;	2047	bool isInt8Quant = commQuantMode == INT8_COMM_QUANT;
2044	bool isMxFp8Quant = (commQuantMode == static_cast<CommQuantModeType>(CommQuantMode::MXFP8_E5M2_QUANT)) \|\|	2048	bool isMxFp8Quant = (commQuantMode == static_cast<CommQuantModeType>(CommQuantMode::MXFP8_E5M2_QUANT)) \|\|
2045	(commQuantMode == static_cast<CommQuantModeType>(CommQuantMode::MXFP8_E4M3_QUANT));	2049	(commQuantMode == static_cast<CommQuantModeType>(CommQuantMode::MXFP8_E4M3_QUANT));
2046	- isA5 = mc2tiling::GetNpuArch(context) == Ops::Base::DAV_3510;	2050	+ NpuArch npuArch = mc2tiling::GetNpuArch(context);
		2051	+ isA5OrA6 = (npuArch == Ops::Base::DAV_3510 \|\| npuArch == Ops::Base::DAV_9201 \|\| npuArch == Ops::Base::DAV_9202);
2047	auto expandXDesc = context->GetInputDesc(config.expandXIndex);	2052	auto expandXDesc = context->GetInputDesc(config.expandXIndex);
2048	uint32_t expandXTypeSize = static_cast<uint32_t>(ge::GetSizeByDataType(expandXDesc->GetDataType()));	2053	uint32_t expandXTypeSize = static_cast<uint32_t>(ge::GetSizeByDataType(expandXDesc->GetDataType()));
2049	hExpandXTypeSize = axisH * expandXTypeSize;	2054	hExpandXTypeSize = axisH * expandXTypeSize;
2050	hExpandXAlign32Size = (hExpandXTypeSize + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;	2055	hExpandXAlign32Size = (hExpandXTypeSize + UB_ALIGN - 1) / UB_ALIGN * UB_ALIGN;
2051	uint32_t hFloatSize = axisH * static_cast<uint32_t>(sizeof(float));	2056	uint32_t hFloatSize = axisH * static_cast<uint32_t>(sizeof(float));
2052	- uint32_t hFloatAlign32Size = GetFloatAlignSize(hFloatSize, UB_ALIGN, isA5 && (isInt8Quant \|\| isMxFp8Quant));	2057	+ uint32_t hFloatAlign32Size = GetFloatAlignSize(hFloatSize, UB_ALIGN, isA5OrA6 && (isInt8Quant \|\| isMxFp8Quant));
2053	- uint32_t hFloatAlign256Size = GetFloatAlignSize(hFloatSize, ALIGNED_LEN, isA5 && (isInt8Quant \|\| isMxFp8Quant));	2058	+ uint32_t hFloatAlign256Size = GetFloatAlignSize(hFloatSize, ALIGNED_LEN, isA5OrA6 && (isInt8Quant \|\| isMxFp8Quant));
2054	scaleNumAlignSize = 0U;	2059	scaleNumAlignSize = 0U;
2055	commDataBytes = hExpandXTypeSize;	2060	commDataBytes = hExpandXTypeSize;
2056	quantBufferSize = 0U;	2061	quantBufferSize = 0U;
2057	if (isInt8Quant \|\| isMxFp8Quant) {	2062	if (isInt8Quant \|\| isMxFp8Quant) {
2058	uint32_t scaleBaseCnt =	2063	uint32_t scaleBaseCnt =
2059	- isInt8Quant ? (isA5 ? hFloatAlign256Size / sizeof(float) : hExpandXAlign32Size / expandXTypeSize) : 0U;	2064	+ isInt8Quant ? (isA5OrA6 ? hFloatAlign256Size / sizeof(float) : hExpandXAlign32Size / expandXTypeSize) : 0U;
2060	commDataBytes = GetQuantCommDataBytes(axisH, expandXTypeSize, scaleBaseCnt, commQuantMode, scaleNumAlignSize);	2065	commDataBytes = GetQuantCommDataBytes(axisH, expandXTypeSize, scaleBaseCnt, commQuantMode, scaleNumAlignSize);
2061	quantBufferSize =	2066	quantBufferSize =
2062	GetQuantScratchSize(axisH, expandXTypeSize, isMxFp8Quant, hFloatAlign32Size, hFloatAlign256Size);	2067	GetQuantScratchSize(axisH, expandXTypeSize, isMxFp8Quant, hFloatAlign32Size, hFloatAlign256Size);
2063	}	2068	}
2064	// 量化输出UB保留原有长度，通信只计算有效scale字节	2069	// 量化输出UB保留原有长度，通信只计算有效scale字节
2065	quantResultDataBytes = commDataBytes;	2070	quantResultDataBytes = commDataBytes;
2066	- if (isA5 && isMxFp8Quant) {	2071	+ if (isA5OrA6 && isMxFp8Quant) {
2067	uint32_t scaleNum = (axisH + UB_ALIGN - 1U) / UB_ALIGN;	2072	uint32_t scaleNum = (axisH + UB_ALIGN - 1U) / UB_ALIGN;
2068	scaleNum = (scaleNum + 1U) / 2U * 2U;	2073	scaleNum = (scaleNum + 1U) / 2U * 2U;
2069	commDataBytes = (axisH + ALIGNED_LEN - 1U) / ALIGNED_LEN * ALIGNED_LEN + scaleNum * sizeof(uint8_t);	2074	commDataBytes = (axisH + ALIGNED_LEN - 1U) / ALIGNED_LEN * ALIGNED_LEN + scaleNum * sizeof(uint8_t);
	@@ -2085,9 +2090,9 @@ static uint64_t GetCommBufferSize(const gert::TilingContext* context,
2085	uint32_t commDataBytes = 0U;	2090	uint32_t commDataBytes = 0U;
2086	uint32_t quantResultDataBytes = 0U;	2091	uint32_t quantResultDataBytes = 0U;
2087	uint64_t quantBufferSize = 0U;	2092	uint64_t quantBufferSize = 0U;
2088	- bool isA5 = false;	2093	+ bool isA5OrA6 = false;
2089	GetCommQuantScratchAndDataSize(context, config, axisH, hExpandXTypeSize, hExpandXAlign32Size, scaleNumAlignSize,	2094	GetCommQuantScratchAndDataSize(context, config, axisH, hExpandXTypeSize, hExpandXAlign32Size, scaleNumAlignSize,
2090	- commDataBytes, quantResultDataBytes, quantBufferSize, isA5);	2095	+ commDataBytes, quantResultDataBytes, quantBufferSize, isA5OrA6);
2091	uint32_t blockCntPerToken = (commDataBytes + SPLIT_BLOCK_DATA_SIZE - 1) / SPLIT_BLOCK_DATA_SIZE;	2096	uint32_t blockCntPerToken = (commDataBytes + SPLIT_BLOCK_DATA_SIZE - 1) / SPLIT_BLOCK_DATA_SIZE;
2092	uint32_t packedDataBytes = blockCntPerToken * SPLIT_BLOCK_DATA_SIZE;	2097	uint32_t packedDataBytes = blockCntPerToken * SPLIT_BLOCK_DATA_SIZE;
2093	uint32_t rawPackedDataBytes =	2098	uint32_t rawPackedDataBytes =
	@@ -2138,7 +2143,8 @@ static bool UbUsedCal(const uint64_t ubSize, const gert::TilingContext* context,
2138	{	2143	{
2139	// A3 retains its raw-data/separate-flag protocol after EP16 regression checks.	2144	// A3 retains its raw-data/separate-flag protocol after EP16 regression checks.
2140	// Its prefetch queues use the original A3 UB budget; A5 keeps the packed budget.	2145	// Its prefetch queues use the original A3 UB budget; A5 keeps the packed budget.
2141	- if (mc2tiling::GetNpuArch(context) != Ops::Base::DAV_3510) {	2146	+ NpuArch npuArch = mc2tiling::GetNpuArch(context);
		2147	+ if (npuArch != Ops::Base::DAV_3510 && npuArch != Ops::Base::DAV_9201 && npuArch != Ops::Base::DAV_9202) {
2142	SetA3BufferNum(ubSize, context, tilingData, config);	2148	SetA3BufferNum(ubSize, context, tilingData, config);
2143	return true;	2149	return true;
2144	}	2150	}