fix(mega_moe):35-92分核判断分离-ops-transformer-AtomGit

Mmc2/mega_moe/op_host/CMakeLists.txt+4-0

	@@ -46,6 +46,10 @@ else() # 回黄host
46	set(SUB_MC2_COMPILE TRUE PARENT_SCOPE)	46	set(SUB_MC2_COMPILE TRUE PARENT_SCOPE)
47	endif()	47	endif()
48		48
		49	+# arch35_92 为共用 kernel 源码目录，不能按硬件架构目录过滤。
		50	+install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/../op_kernel/arch35_92
		51	+ DESTINATION ${IMPL_INSTALL_DIR}/ascendc/mega_moe/op_kernel)
		52	+
49	# Torch 复用算子包中的纯布局头，不再单独维护 CCL 容量公式。	53	# Torch 复用算子包中的纯布局头，不再单独维护 CCL 容量公式。
50	if (BUILD_OPEN_PROJECT OR ENABLE_BUILT_IN)	54	if (BUILD_OPEN_PROJECT OR ENABLE_BUILT_IN)
51	set(MEGA_MOE_LAYOUT_INC_INSTALL_DIR ${ACLNN_INC_INSTALL_DIR})	55	set(MEGA_MOE_LAYOUT_INC_INSTALL_DIR ${ACLNN_INC_INSTALL_DIR})

Mmc2/mega_moe/op_kernel/arch35/mega_moe_arch35.h+23-3

	@@ -31,10 +31,11 @@
31	#include "../arch35_92/stage/mega_moe_send_mask.h"	31	#include "../arch35_92/stage/mega_moe_send_mask.h"
32	#include "../arch35_92/stage/mega_moe_workspace_reset.h"	32	#include "../arch35_92/stage/mega_moe_workspace_reset.h"
33	#include "../arch35_92/stage/mega_moe_token_dispatch_dedup.h"	33	#include "../arch35_92/stage/mega_moe_token_dispatch_dedup.h"
34	-#include "../arch35_92/stage/mega_moe_gmm1_activation.h"
35	#include "../arch35_92/stage/mega_moe_gmm2_combine_dedup.h"	34	#include "../arch35_92/stage/mega_moe_gmm2_combine_dedup.h"
36	#include "../arch35_92/stage/mega_moe_unpermute.h"	35	#include "../arch35_92/stage/mega_moe_unpermute.h"
37	#include "../../../common/op_kernel/quantize_functions.h"	36	#include "../../../common/op_kernel/quantize_functions.h"
		37	+#include "stage/mega_moe_gmm1_activation_arch35.h"
		38	+#include "stage/mega_moe_gmm2_combine_arch35.h"
38		39
39	namespace MegaMoeImpl {	40	namespace MegaMoeImpl {
40		41
	@@ -94,6 +95,7 @@ protected:
94	__aicore__ inline void DispatchBuffInit();	95	__aicore__ inline void DispatchBuffInit();
95	template <bool NeedCumsumReload = false>	96	template <bool NeedCumsumReload = false>
96	__aicore__ inline void PrepareMoeExpertTokenCountTable();	97	__aicore__ inline void PrepareMoeExpertTokenCountTable();
		98	+ __aicore__ inline void BuildMoeDedupRowDescTable();
97	__aicore__ inline void ReloadDispatchCumsumRange(uint32_t firstExpertIdx, uint32_t lastExpertIdx);	99	__aicore__ inline void ReloadDispatchCumsumRange(uint32_t firstExpertIdx, uint32_t lastExpertIdx);
98	__aicore__ inline void ExportExpertTokenCounts();	100	__aicore__ inline void ExportExpertTokenCounts();
99	__aicore__ inline void EnterSteadyDispatch();	101	__aicore__ inline void EnterSteadyDispatch();
	@@ -466,6 +468,24 @@ __aicore__ inline void MegaMoe<TemplateMegaMoeTypeFunc>::PrepareMoeExpertTokenCo
466	GmWrite(countTableReady, static_cast<int32_t>(1));	468	GmWrite(countTableReady, static_cast<int32_t>(1));
467	}	469	}
468		470
		471	+template <TemplateMegaMoeTypeClass>
		472	+__aicore__ inline void MegaMoe<TemplateMegaMoeTypeFunc>::BuildMoeDedupRowDescTable()
		W W wang-minbo1 day ago 【一般】这次分核判断拆得挺彻底，dedup 建表的角色编排、GMM1/GMM2 入口、combine buffer 初始化都挪了位置，PR 描述声明是保持原有计算和同步流程的重构，但测试一节是空的：UT 通过数、真机冒烟/泛化结果一个没列。这种千行级搬移建议把 A5、A6 各自的回归结论补到描述里（至少 UT 用例数加二级冒烟通过情况），方便检视人确认重构前后行为一致。
		473	+{
		474	+ const auto& blockJob = gmmExecutionConfig_.blockJob;
		475	+ const bool aiv0Joins = params_.tilingData->compute.hiddenDim <= DEDUP_AIV0_PREPASS_MAX_HIDDEN_DIM;
		476	+ const bool isAiv1 = MegaMoeHardwarePolicy::GetSubBlockIdx() == 1U;
		477	+ if (!isAiv1 && !aiv0Joins) {
		478	+ return;
		479	+ }
		480	+ BlockJobContext workerJob{isAiv1 ? blockJob.jobIndex : blockJob.jobIndex + blockJob.totalJobs,
		481	+ aiv0Joins ? blockJob.totalJobs * 2U : blockJob.totalJobs};
		482	+ // AIV1 计低 16 位，AIV0 计高 16 位；保持 dispatch/combine 的既有等待协议。
		483	+ const int32_t doneIncrement =
		484	+ isAiv1 ? static_cast<int32_t>(1) : static_cast<int32_t>(1) << DEDUP_PREPASS_DONE_AIV0_SHIFT;
		485	+ BuildDedupRowDescTable<TopkIndexType>(tokenDispatchConfig_.dedup, commonConfig_, workerJob, doneIncrement,
		486	+ syncWorkspaceLayout_, params_, tokenDispatchScratch_);
		487	+}
		488	+
469	// Activation 可能覆盖前缀表，Dispatch 前恢复所需专家范围。	489	// Activation 可能覆盖前缀表，Dispatch 前恢复所需专家范围。
470	template <TemplateMegaMoeTypeClass>	490	template <TemplateMegaMoeTypeClass>
471	__aicore__ inline void MegaMoe<TemplateMegaMoeTypeFunc>::ReloadDispatchCumsumRange(uint32_t firstExpertIdx,	491	__aicore__ inline void MegaMoe<TemplateMegaMoeTypeFunc>::ReloadDispatchCumsumRange(uint32_t firstExpertIdx,
	@@ -714,7 +734,7 @@ __aicore__ inline void MegaMoe<TemplateMegaMoeTypeFunc>::RunSharedExpertGmm1Acti
714	true, true>(sharedEpilogueOp_, params_, problemShape, gmmAddrInfo,	734	true, true>(sharedEpilogueOp_, params_, problemShape, gmmAddrInfo,
715	runtimeState.startBlockIdx, runtimeState.vecSetSyncCom, gmmConfig.blockJob,	735	runtimeState.startBlockIdx, runtimeState.vecSetSyncCom, gmmConfig.blockJob,
716	expertBeforeCnt, sharedExpertIdx, runtimeState.pingpongIdx,	736	expertBeforeCnt, sharedExpertIdx, runtimeState.pingpongIdx,
717	- commonConfig_.activationOutputDivisor, nullptr,	737	+ commonConfig_.activationOutputDivisor,
718	SharedQuantConfig::AXW_MODE == AxWMode::A8W8, layoutArgs...);	738	SharedQuantConfig::AXW_MODE == AxWMode::A8W8, layoutArgs...);
719	}	739	}
720	};	740	};
	@@ -862,7 +882,7 @@ __aicore__ inline void MegaMoe<TemplateMegaMoeTypeFunc>::RunSharedExpertGmm2(con
862	RunGmm2Generic<COMBINE_NO_QUANT, typename SharedQuantConfig::QuantOutType,	882	RunGmm2Generic<COMBINE_NO_QUANT, typename SharedQuantConfig::QuantOutType,
863	typename SharedQuantConfig::QuantOutType, bfloat16_t, SharedQuantScaleType, SharedQuantScaleType,	883	typename SharedQuantConfig::QuantOutType, bfloat16_t, SharedQuantScaleType, SharedQuantScaleType,
864	SharedWeight2Format != FORMAT_ND, false, GMM1_TILE_M, false, true, true>(	884	SharedWeight2Format != FORMAT_ND, false, GMM1_TILE_M, false, true, true>(
865	- problemShape, gmmAddrInfo, startBlockIdx_, gmmExecutionConfig_.blockJob, nullptr, true, 0U, nullptr,	885	+ problemShape, gmmAddrInfo, startBlockIdx_, gmmExecutionConfig_.blockJob, true, 0U, nullptr,
866	commonConfig_.activationOutputDivisor);	886	commonConfig_.activationOutputDivisor);
867	}	887	}
868	}	888	}

Mmc2/mega_moe/op_kernel/arch35/mega_moe_layered.h+4-4

	@@ -47,9 +47,9 @@
47	#include "../arch35_92/stage/mega_moe_token_quant.h"	47	#include "../arch35_92/stage/mega_moe_token_quant.h"
48	#include "../arch35_92/stage/mega_moe_expert_count.h"	48	#include "../arch35_92/stage/mega_moe_expert_count.h"
49	#include "../arch35_92/mega_moe_vf/mega_moe_layered_vf.h"	49	#include "../arch35_92/mega_moe_vf/mega_moe_layered_vf.h"
50	-#include "../arch35_92/stage/mega_moe_gmm1_activation.h"
51	-#include "../arch35_92/stage/mega_moe_gmm2_combine.h"
52	#include "../arch35_92/stage/mega_moe_layered_combine_send.h"	50	#include "../arch35_92/stage/mega_moe_layered_combine_send.h"
		51	+#include "stage/mega_moe_gmm1_activation_arch35.h"
		Y Y yangzeheng3 days ago 也和上面的一样改一下
		52	+#include "stage/mega_moe_gmm2_combine_arch35.h"
53		53
54	namespace MegaMoeImpl {	54	namespace MegaMoeImpl {
55		55
	@@ -933,7 +933,7 @@ __aicore__ inline void MegaMoeLayered<TemplateMegaMoeLayeredTypeFunc>::RunGmm1Wi
933	QuantScaleOutType, Weight1Format != FORMAT_ND, GMM1_TILE_M, epilogueTileM, prefetchWeights,	933	QuantScaleOutType, Weight1Format != FORMAT_ND, GMM1_TILE_M, epilogueTileM, prefetchWeights,
934	IsShared, !IsShared>(epilogue, params_, state.problemShape, gmmAddrInfo, startBlockIdx_,	934	IsShared, !IsShared>(epilogue, params_, state.problemShape, gmmAddrInfo, startBlockIdx_,
935	vecSetSyncCom, blockJob, expertRowStart, expertIdx, gmm1PingPongIdx_,	935	vecSetSyncCom, blockJob, expertRowStart, expertIdx, gmm1PingPongIdx_,
936	- activationOutputDivisor_, nullptr, false, layoutArgs...);	936	+ activationOutputDivisor_, false, layoutArgs...);
937	}	937	}
938	}	938	}
939		939
	@@ -1006,7 +1006,7 @@ __aicore__ inline void MegaMoeLayered<TemplateMegaMoeLayeredTypeFunc>::GroupMatm
1006	} else {	1006	} else {
1007	RunGmm2Generic<CombineQuantMode, QuantOutType, QuantOutType, bfloat16_t, QuantScaleOutType, QuantScaleOutType,	1007	RunGmm2Generic<CombineQuantMode, QuantOutType, QuantOutType, bfloat16_t, QuantScaleOutType, QuantScaleOutType,
1008	Weight2Format != FORMAT_ND, true, L1_TILE_M_256, TopkWeightsPrefetch, IsShared, !IsShared>(	1008	Weight2Format != FORMAT_ND, true, L1_TILE_M_256, TopkWeightsPrefetch, IsShared, !IsShared>(
1009	- state.problemShape, gmmAddrInfo, startBlockIdx_, nullptr, false, rowOffsetInExpert, nullptr,	1009	+ state.problemShape, gmmAddrInfo, startBlockIdx_, false, rowOffsetInExpert, nullptr,
1010	activationOutputDivisor_, expertTokenCount);	1010	activationOutputDivisor_, expertTokenCount);
1011	}	1011	}
1012	}	1012	}

Mmc2/mega_moe/op_kernel/arch35/mega_moe_wave_a4w4.h+9-3

Mmc2/mega_moe/op_kernel/arch35/mega_moe_wave_a8w4.h+9-3

Mmc2/mega_moe/op_kernel/arch35/mega_moe_wave_a8w8.h+21-10

	@@ -58,6 +58,7 @@ private:
58		58
59	using MegaMoeBase::DispatchBuffInit;	59	using MegaMoeBase::DispatchBuffInit;
60	using MegaMoeBase::PrepareMoeExpertTokenCountTable;	60	using MegaMoeBase::PrepareMoeExpertTokenCountTable;
		61	+ using MegaMoeBase::BuildMoeDedupRowDescTable;
61	using MegaMoeBase::EnterSteadyDispatch;	62	using MegaMoeBase::EnterSteadyDispatch;
62	using MegaMoeBase::commonConfig_;	63	using MegaMoeBase::commonConfig_;
63	using MegaMoeBase::countWorkspace_;	64	using MegaMoeBase::countWorkspace_;
	@@ -163,8 +164,13 @@ template <TemplateMegaMoeA8W8WaveTypeClass>
163	__aicore__ inline typename MegaMoeA8W8Wave<TemplateMegaMoeA8W8WaveTypeFunc>::CombineBufferConfig	164	__aicore__ inline typename MegaMoeA8W8Wave<TemplateMegaMoeA8W8WaveTypeFunc>::CombineBufferConfig
164	MegaMoeA8W8Wave<TemplateMegaMoeA8W8WaveTypeFunc>::InitCombineBuffers()	165	MegaMoeA8W8Wave<TemplateMegaMoeA8W8WaveTypeFunc>::InitCombineBuffers()
165	{	166	{
166	- return InitWaveCombineBuffers<CombineQuantMode, false, WAVE_COMBINE_STEADY_ROW_BUFFER_COUNT>(commonConfig_,	167	+ if constexpr (g_coreType == AIV) {
167	- waveCombineScratch_);	168	+ if (MegaMoeHardwarePolicy::GetSubBlockIdx() == 1U) {
		169	+ return InitWaveCombineBuffers<CombineQuantMode, WAVE_COMBINE_STEADY_ROW_BUFFER_COUNT>(commonConfig_,
		170	+ waveCombineScratch_);
		171	+ }
		172	+ }
		173	+ return {};
168	}	174	}
169		175
170	template <TemplateMegaMoeA8W8WaveTypeClass>	176	template <TemplateMegaMoeA8W8WaveTypeClass>
	@@ -192,9 +198,13 @@ __aicore__ inline ExpertTokenPosition MegaMoeA8W8Wave<TemplateMegaMoeA8W8WaveTyp
192	plannedDispatchPosition = nextDispatchRange.end;	198	plannedDispatchPosition = nextDispatchRange.end;
193	dispatchRange.end = nextDispatchRange.end;	199	dispatchRange.end = nextDispatchRange.end;
194	}	200	}
195	- DispatchTokenRange<TopkIndexType, ActivationType, QuantScaleOutType, GMM1_TILE_M, TopkWeightsPrefetch>(	201	+ if constexpr (g_coreType == AIV) {
196	- tokenDispatchConfig_, commonConfig_, gmmExecutionConfig_.blockJob, syncWorkspaceLayout_, params_,	202	+ if (MegaMoeHardwarePolicy::GetSubBlockIdx() == 1U) {
197	- g_winRankAddr_, tokenDispatchScratch_, dispatchRange);	203	+ DispatchTokenRange<TopkIndexType, ActivationType, QuantScaleOutType, GMM1_TILE_M, TopkWeightsPrefetch>(
		204	+ tokenDispatchConfig_, commonConfig_, gmmExecutionConfig_.blockJob, syncWorkspaceLayout_, params_,
		205	+ g_winRankAddr_, tokenDispatchScratch_, dispatchRange);
		206	+ }
		207	+ }
198	// 整 WAVE 的数据和 ready flag 发布完成后，再提交 Dispatch 进度。	208	// 整 WAVE 的数据和 ready flag 发布完成后，再提交 Dispatch 进度。
199	dispatchPosition = plannedDispatchPosition;	209	dispatchPosition = plannedDispatchPosition;
200	return dispatchPosition;	210	return dispatchPosition;
	@@ -305,7 +315,7 @@ __aicore__ inline ExpertTokenPosition MegaMoeA8W8Wave<TemplateMegaMoeA8W8WaveTyp
305	MoeWeight1Format != FORMAT_ND, GMM1_TILE_M, EPILOGUE_TILE_M, TopkWeightsPrefetch, false, true>(	315	MoeWeight1Format != FORMAT_ND, GMM1_TILE_M, EPILOGUE_TILE_M, TopkWeightsPrefetch, false, true>(
306	epilogueOp_, params_, gmm1WaveProblemShape, gmm1AddrInfo, runtimeState.startBlockIdx,	316	epilogueOp_, params_, gmm1WaveProblemShape, gmm1AddrInfo, runtimeState.startBlockIdx,
307	runtimeState.vecSetSyncCom, gmmExecutionConfig_.blockJob, waveTokenStartIndex, gmm1Position.expertIdx,	317	runtimeState.vecSetSyncCom, gmmExecutionConfig_.blockJob, waveTokenStartIndex, gmm1Position.expertIdx,
308	- runtimeState.pingpongIdx, commonConfig_.activationOutputDivisor, nullptr, isWholeExpert);	318	+ runtimeState.pingpongIdx, commonConfig_.activationOutputDivisor, isWholeExpert);
309		319
310	processedMGroupCount += problemMGroupCount;	320	processedMGroupCount += problemMGroupCount;
311	gmm1Position.tokenIndexInExpert = waveEndTokenIndexInExpert;	321	gmm1Position.tokenIndexInExpert = waveEndTokenIndexInExpert;
	@@ -351,13 +361,13 @@ __aicore__ inline void MegaMoeA8W8Wave<TemplateMegaMoeA8W8WaveTypeFunc>::Process
351	if (IsCombineDedupOn(params_.tilingData->compute.dedupMode)) {	361	if (IsCombineDedupOn(params_.tilingData->compute.dedupMode)) {
352	RunGmm2Generic<COMBINE_NO_QUANT, QuantOutType, QuantOutType, bfloat16_t, QuantScaleOutType, QuantScaleOutType,	362	RunGmm2Generic<COMBINE_NO_QUANT, QuantOutType, QuantOutType, bfloat16_t, QuantScaleOutType, QuantScaleOutType,
353	MoeWeight2Format != FORMAT_ND, false, GMM1_TILE_M, TopkWeightsPrefetch, false, true, false>(	363	MoeWeight2Format != FORMAT_ND, false, GMM1_TILE_M, TopkWeightsPrefetch, false, true, false>(
354	- gmm2WaveProblemShape, gmm2AddrInfo, startBlockIdx, gmmExecutionConfig_.blockJob, nullptr, isWholeExpert,	364	+ gmm2WaveProblemShape, gmm2AddrInfo, startBlockIdx, gmmExecutionConfig_.blockJob, isWholeExpert,
355	gmm2Position.tokenIndexInExpert, &params_, commonConfig_.activationOutputDivisor);	365	gmm2Position.tokenIndexInExpert, &params_, commonConfig_.activationOutputDivisor);
356	} else {	366	} else {
357	RunGmm2Generic<COMBINE_NO_QUANT, QuantOutType, QuantOutType, bfloat16_t, QuantScaleOutType, QuantScaleOutType,	367	RunGmm2Generic<COMBINE_NO_QUANT, QuantOutType, QuantOutType, bfloat16_t, QuantScaleOutType, QuantScaleOutType,
358	MoeWeight2Format != FORMAT_ND, false, GMM1_TILE_M, TopkWeightsPrefetch, false, true,	368	MoeWeight2Format != FORMAT_ND, false, GMM1_TILE_M, TopkWeightsPrefetch, false, true,
359	CombineQuantMode == COMBINE_NO_QUANT>(	369	CombineQuantMode == COMBINE_NO_QUANT>(
360	- gmm2WaveProblemShape, gmm2AddrInfo, startBlockIdx, gmmExecutionConfig_.blockJob, nullptr, isWholeExpert,	370	+ gmm2WaveProblemShape, gmm2AddrInfo, startBlockIdx, gmmExecutionConfig_.blockJob, isWholeExpert,
361	gmm2Position.tokenIndexInExpert, &params_, commonConfig_.activationOutputDivisor);	371	gmm2Position.tokenIndexInExpert, &params_, commonConfig_.activationOutputDivisor);
362	}	372	}
363	}	373	}
	@@ -631,8 +641,9 @@ __aicore__ inline void MegaMoeA8W8Wave<TemplateMegaMoeA8W8WaveTypeFunc>::Prepare
631	if (tokenDispatchConfig_.dedup.dedupMode != 0) {	641	if (tokenDispatchConfig_.dedup.dedupMode != 0) {
632	// 去重建表：route 槽到达保证与 count epoch 校验同源；函数尾部含 AIV1 间栅栏。	642	// 去重建表：route 槽到达保证与 count epoch 校验同源；函数尾部含 AIV1 间栅栏。
633	PrepareDedupPrepassAiv0Local(tokenDispatchConfig_.dedup, commonConfig_, params_, tokenDispatchScratch_);	643	PrepareDedupPrepassAiv0Local(tokenDispatchConfig_.dedup, commonConfig_, params_, tokenDispatchScratch_);
634	- BuildDedupRowDescTable<TopkIndexType>(tokenDispatchConfig_.dedup, commonConfig_, gmmExecutionConfig_.blockJob,	644	+ if constexpr (g_coreType == AIV) {
635	- syncWorkspaceLayout_, params_, tokenDispatchScratch_);	645	+ BuildMoeDedupRowDescTable();
		646	+ }
636	}	647	}
637	}	648	}
638		649

Amc2/mega_moe/op_kernel/arch35/stage/mega_moe_gmm1_activation_arch35.h+344-0

@@ -0,0 +1,344 @@
	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	+#ifndef MEGA_MOE_GMM1_ACTIVATION_ARCH35_H
	12	+#define MEGA_MOE_GMM1_ACTIVATION_ARCH35_H
	13	+
	14	+#include "../../arch35_92/stage/mega_moe_gmm1_activation.h"
	15	+
	16	+namespace MegaMoeImpl {
	17	+namespace GmmKernel {
	18	+
	19	+template <typename BlockMmad, typename ElementC, typename MakeLayoutC, bool TopkWeightsPrefetch, bool IsShared = false,
	20	+ bool IsWaveFlagGrained = false, typename Scheduler, typename Config, typename ActivationQuantOp,
	21	+ typename LayoutBundle>
	22	+__aicore__ inline void Gmm1ExecGeneric(Scheduler& scheduler, const Params& params, const GMMAddrInfo& gmmAddrInfo,
	23	+ const Config& config, uint32_t startLoopIdx, uint32_t tileNum,
	24	+ ActivationQuantOp& activationQuantOp, int32_t& vecSetSyncCom,
	25	+ uint32_t expertBeforeCnt, uint32_t expertIdx, uint16_t& pingpongIdx,
	26	+ bool allowWeightL2Bypass, const LayoutBundle& layouts)
	27	+{
	28	+ using KernelConfig = typename Config::KernelConfig;
	29	+ using ElementA = typename KernelConfig::ElementAType;
	30	+ using ElementB = typename KernelConfig::ElementBType;
	31	+ using ElementMxScaleA = typename KernelConfig::ElementMxScaleAType;
	32	+ using ElementMxScaleB = typename KernelConfig::ElementMxScaleBType;
	33	+ auto gmA = asc::te::make_tensor(
	34	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementA*>(gmmAddrInfo.aGlobal)),
	35	+ layouts.a);
	36	+ auto gmB = asc::te::make_tensor(
	37	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementB*>(gmmAddrInfo.bGlobal)),
	38	+ layouts.b);
	39	+ auto gmScaleA = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
	40	+ reinterpret_cast<__gm__ ElementMxScaleA*>(gmmAddrInfo.aScaleGlobal)),
	41	+ layouts.scaleA);
	42	+ auto gmScaleB = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
	43	+ reinterpret_cast<__gm__ ElementMxScaleB*>(gmmAddrInfo.bScaleGlobal)),
	44	+ layouts.scaleB);
	45	+ if constexpr (IsWaveFlagGrained && g_coreType == AscendC::AIC) {
	46	+ SetWaveWeightL2CacheHint<KernelConfig::IS_WEIGHT_NZ, KernelConfig>(config, allowWeightL2Bypass, gmB, gmScaleB);
	47	+ }
	48	+ GM_ADDR cGlobal = 0UL;
	49	+ if constexpr (TopkWeightsPrefetch) {
	50	+ cGlobal = gmmAddrInfo.gmm1OutGlobal;
	51	+ }
	52	+ auto gmC = asc::te::make_tensor(
	53	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(cGlobal)), layouts.c);
	54	+ __gm__ float* metaInfoGm = nullptr;
	55	+ if constexpr (TopkWeightsPrefetch && g_coreType == AscendC::AIV) {
	56	+ metaInfoGm = reinterpret_cast<__gm__ float*>(gmmAddrInfo.metaInfoGlobal);
	57	+ }
	58	+
	59	+ using WorkSetType = Gmm1WorkSet<Scheduler, decltype(gmA), decltype(gmB), decltype(gmScaleA), decltype(gmScaleB),
	60	+ decltype(gmC), decltype(metaInfoGm)>;
	61	+ WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmC, metaInfoGm};
	62	+ if constexpr (g_coreType == AscendC::AIC) {
	63	+ BlockMmadContext<BlockMmad> localBlockMmadContext;
	64	+ if constexpr (TopkWeightsPrefetch) {
	65	+ Gmm1AicMmadToGmGeneric<BlockMmadContext<BlockMmad>, IsShared, TopkWeightsPrefetch>(
	66	+ workSet, params, gmmAddrInfo, config, startLoopIdx, tileNum, expertIdx, &localBlockMmadContext);
	67	+ } else {
	68	+ Gmm1AicMmadGeneric<BlockMmadContext<BlockMmad>, IsShared>(workSet, gmmAddrInfo, config, startLoopIdx,
	69	+ tileNum, vecSetSyncCom, pingpongIdx,
	70	+ &localBlockMmadContext);
	71	+ }
	72	+ } else {
	73	+ if constexpr (TopkWeightsPrefetch) {
	74	+ Gmm1Aiv0PrefetchEpilogueGeneric<MakeLayoutC, ElementC, IsWaveFlagGrained>(
	75	+ workSet, params, gmmAddrInfo, config, startLoopIdx, tileNum, activationQuantOp, expertBeforeCnt,
	76	+ expertIdx);
	77	+ } else {
	78	+ Gmm1Aiv0EpilogueGeneric<IsWaveFlagGrained>(workSet, gmmAddrInfo, config, startLoopIdx, tileNum,
	79	+ activationQuantOp, pingpongIdx);
	80	+ }
	81	+ }
	82	+}
	83	+
	84	+template <typename BlockMmad, typename BlockPrologue, typename ElementC, typename MakeLayoutC, bool IsShared,
	85	+ typename Scheduler, typename Config, bool TopkWeightsPrefetch, bool IsWaveFlagGrained,
	86	+ typename ActivationQuantOp, typename LayoutBundle>
	87	+__aicore__ inline void Gmm1ExecA8W4(Scheduler& scheduler, const Params& params, const GMMAddrInfo& gmmAddrInfo,
	88	+ const Config& config, uint32_t startLoopIdx, uint32_t tileNum,
	89	+ ActivationQuantOp& activationQuantOp, uint32_t expertBeforeCnt, uint32_t expertIdx,
	90	+ const LayoutBundle& layouts,
	91	+ const A8W4BlockContext<BlockMmad, BlockPrologue>& pipeline)
	92	+{
	93	+ using KernelConfig = typename Config::KernelConfig;
	94	+ using ElementA = typename KernelConfig::ElementAType;
	95	+ using ElementB = typename KernelConfig::ElementBType;
	96	+ using ElementMxScaleA = typename KernelConfig::ElementMxScaleAType;
	97	+ using ElementMxScaleB = typename KernelConfig::ElementMxScaleBType;
	98	+ auto gmC = asc::te::make_tensor(
	99	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(gmmAddrInfo.gmm1OutGlobal)),
	100	+ layouts.c);
	101	+ auto gmA = asc::te::make_tensor(
	102	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementA*>(gmmAddrInfo.aGlobal)),
	103	+ layouts.a);
	104	+ auto gmB = asc::te::make_tensor(
	105	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementB*>(gmmAddrInfo.bGlobal)),
	106	+ layouts.b);
	107	+ auto gmScaleA = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
	108	+ reinterpret_cast<__gm__ ElementMxScaleA*>(gmmAddrInfo.aScaleGlobal)),
	109	+ layouts.scaleA);
	110	+ auto gmScaleB = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
	111	+ reinterpret_cast<__gm__ ElementMxScaleB*>(gmmAddrInfo.bScaleGlobal)),
	112	+ layouts.scaleB);
	113	+ __gm__ float* metaInfoGm = nullptr;
	114	+ if constexpr (TopkWeightsPrefetch && g_coreType == AscendC::AIV) {
	115	+ metaInfoGm = reinterpret_cast<__gm__ float*>(gmmAddrInfo.metaInfoGlobal);
	116	+ }
	117	+
	118	+ using WorkSetType = Gmm1WorkSet<Scheduler, decltype(gmA), decltype(gmB), decltype(gmScaleA), decltype(gmScaleB),
	119	+ decltype(gmC), decltype(metaInfoGm)>;
	120	+ WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmC, metaInfoGm};
	121	+ if constexpr (g_coreType == AscendC::AIC) {
	122	+ Gmm1AicMmadA8W4<BlockMmad, IsShared, TopkWeightsPrefetch>(*pipeline.block, workSet, params, gmmAddrInfo, config,
	123	+ startLoopIdx, tileNum, expertIdx);
	124	+ } else if (MegaMoeHardwarePolicy::GetSubBlockIdx() == 0U) {
	125	+ Gmm1Aiv0PrologueA8W4<BlockPrologue>(*pipeline.block, workSet.scheduler, workSet.gmB, config, startLoopIdx,
	126	+ tileNum);
	127	+ } else {
	128	+ if constexpr (TopkWeightsPrefetch) {
	129	+ Gmm1Aiv1PrefetchEpilogueA8W4<ElementC, MakeLayoutC, IsWaveFlagGrained>(
	130	+ workSet, params, gmmAddrInfo, config, startLoopIdx, tileNum, activationQuantOp, expertBeforeCnt,
	131	+ expertIdx);
	132	+ } else {
	133	+ Gmm1Aiv1EpilogueA8W4<ElementC, MakeLayoutC, IsWaveFlagGrained>(workSet, gmmAddrInfo, config, startLoopIdx,
	134	+ tileNum, activationQuantOp);
	135	+ }
	136	+ }
	137	+}
	138	+
	139	+} // namespace GmmKernel
	140	+
	141	+template <typename ElementA, typename EpilogueElementA, typename ElementB, typename ElementC, typename ElementMxScaleA,
	142	+ typename ElementMxScaleB, bool IsWeightNZ = false, uint32_t Gmm1TileM = L1_TILE_M_256,
	143	+ uint32_t EpilogueTileM = Gmm1TileM, bool TopkWeightsPrefetch = false, bool IsShared = false,
	144	+ bool IsWaveFlagGrained = false, typename BlockEpilogue, typename... LayoutArgs>
	145	+__aicore__ inline void RunGmm1Generic(BlockEpilogue& epilogueOp, const Params& params,
	146	+ const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
	147	+ const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx, int32_t& vecSetSyncCom,
	148	+ const BlockJobContext& blockJob, uint32_t expertBeforeCnt, uint32_t expertIdx,
	149	+ uint16_t& pingpongIdx, uint32_t activationOutputDivisor = ACTIVATION_N_HALF,
	150	+ bool allowWeightL2Bypass = false, const LayoutArgs&... layoutArgs)
	151	+{
	152	+ using GmmConfig =
	153	+ GmmKernel::Config<false, COMBINE_NO_QUANT, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB,
	154	+ IsWeightNZ, TopkWeightsPrefetch, IsShared, false, IsWaveFlagGrained>;
	155	+ auto config = GmmConfig::BuildGmm1ProblemConfig(problemShape, blockJob, Gmm1TileM, activationOutputDivisor);
	156	+
	157	+ GmmKernel::BlockScheduler scheduler({config.m, config.schedulerN, config.k},
	158	+ GmmKernel::BlockScheduler::Params{asc::te::make_coord(
	159	+ static_cast<int64_t>(config.tileM), static_cast<int64_t>(L1_TILE_N))});
	160	+ uint32_t tileNum = scheduler.GetTileNum();
	161	+
	162	+ if (MegaMoeHardwarePolicy::GetSubBlockIdx() != 0U) {
	163	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
	164	+ return;
	165	+ }
	166	+ if constexpr (g_coreType == AIV) {
	167	+ epilogueOp.UpdateNextProblem({config.m, config.outputN, config.k, 0});
	168	+ }
	169	+
	170	+ uint32_t startLoopIdx =
	171	+ (config.blockIdx < startBlockIdx ? config.blockIdx + config.blockNum : config.blockIdx) - startBlockIdx;
	172	+ using BlockMmad = typename GmmConfig::BlockMmad;
	173	+ using MakeLayoutC = typename GmmConfig::MakeLayoutC;
	174	+ auto layouts = GmmConfig::BuildLayouts(config, layoutArgs...);
	175	+ GmmKernel::Gmm1ExecGeneric<BlockMmad, ElementC, MakeLayoutC, TopkWeightsPrefetch, IsShared, IsWaveFlagGrained>(
	176	+ scheduler, params, gmmAddrInfo, config, startLoopIdx, tileNum, epilogueOp, vecSetSyncCom, expertBeforeCnt,
	177	+ expertIdx, pingpongIdx, allowWeightL2Bypass, layouts);
	178	+
	179	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
	180	+}
	181	+
	182	+template <typename ElementA, typename EpilogueElementA, typename ElementB, typename ElementC, typename ElementMxScaleA,
	183	+ typename ElementMxScaleB, bool IsWeightNZ = false, uint32_t Gmm1TileM = L1_TILE_M_256,
	184	+ uint32_t EpilogueTileM = Gmm1TileM, bool TopkWeightsPrefetch = false, bool IsShared = false,
	185	+ bool IsWaveFlagGrained = false, typename BlockEpilogue>
	186	+__aicore__ inline void RunGmm1Generic(BlockEpilogue& epilogueOp, const Params& params,
	187	+ const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
	188	+ const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx, int32_t& vecSetSyncCom,
	189	+ uint32_t expertBeforeCnt, uint32_t expertIdx, uint16_t& pingpongIdx,
	190	+ uint32_t activationOutputDivisor = ACTIVATION_N_HALF,
	191	+ bool allowWeightL2Bypass = false)
	192	+{
	193	+ BlockJobContext blockJob{MegaMoeHardwarePolicy::GetAicBlockIdx(),
	194	+ static_cast<uint32_t>(MegaMoeHardwarePolicy::GetBlockNum())};
	195	+ RunGmm1Generic<ElementA, EpilogueElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, IsWeightNZ,
	196	+ Gmm1TileM, EpilogueTileM, TopkWeightsPrefetch, IsShared, IsWaveFlagGrained>(
	197	+ epilogueOp, params, problemShape, gmmAddrInfo, startBlockIdx, vecSetSyncCom, blockJob, expertBeforeCnt,
	198	+ expertIdx, pingpongIdx, activationOutputDivisor, allowWeightL2Bypass);
	199	+}
	200	+
	201	+template <typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA, typename ElementMxScaleB,
	202	+ uint32_t Gmm1TileM = L1_TILE_M_256, uint32_t EpilogueTileM = Gmm1TileM, bool TopkWeightsPrefetch = false,
	203	+ bool IsShared = false, bool IsWaveFlagGrained = false, typename BlockContext, typename BlockEpilogue,
	204	+ typename... LayoutArgs>
	205	+__aicore__ inline void RunGmm1A8W4(const BlockContext& pipeline, BlockEpilogue& activationQuantOp, const Params& params,
	206	+ const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
	207	+ const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,
	208	+ const BlockJobContext& blockJob, uint32_t expertBeforeCnt, uint32_t expertIdx,
	209	+ uint32_t activationOutputDivisor = ACTIVATION_N_HALF,
	210	+ const LayoutArgs&... layoutArgs)
	211	+{
	212	+ static_assert(std::is_same_v<ElementA, __fp8e4m3>, "Activation must be __fp8e4m3");
	213	+ static_assert(std::is_same_v<ElementB, __fp4e2m1x2>, "Weight must be __fp4e2m1x2");
	214	+
	215	+ using GmmConfig = GmmKernel::Config<true, 0, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, false,
	216	+ TopkWeightsPrefetch, IsShared, false, IsWaveFlagGrained>;
	217	+ auto config = GmmConfig::BuildGmm1ProblemConfig(problemShape, blockJob, Gmm1TileM, activationOutputDivisor);
	218	+
	219	+ if constexpr (g_coreType == AIV) {
	220	+ activationQuantOp.UpdateNextProblem({config.m, config.outputN, config.k, 0});
	221	+ }
	222	+
	223	+ using BlockMmad = typename GmmConfig::BlockMmad;
	224	+ using BlockPrologue = typename GmmConfig::BlockPrologue;
	225	+ using MakeLayoutC = typename GmmConfig::MakeLayoutC;
	226	+
	227	+ GmmKernel::BlockScheduler scheduler({config.m, config.schedulerN, config.k},
	228	+ GmmKernel::BlockScheduler::Params{asc::te::make_coord(
	229	+ static_cast<int64_t>(config.tileM), static_cast<int64_t>(L1_TILE_N))});
	230	+ uint32_t tileNum = scheduler.GetTileNum();
	231	+ uint32_t startLoopIdx =
	232	+ (config.blockIdx < startBlockIdx ? config.blockIdx + config.blockNum : config.blockIdx) - startBlockIdx;
	233	+ if (startLoopIdx >= tileNum) {
	234	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
	235	+ return;
	236	+ }
	237	+
	238	+ auto layouts = GmmConfig::BuildLayouts(config, layoutArgs...);
	239	+ GmmKernel::Gmm1ExecA8W4<BlockMmad, BlockPrologue, ElementC, MakeLayoutC, IsShared, GmmKernel::BlockScheduler,
	240	+ decltype(config), TopkWeightsPrefetch, IsWaveFlagGrained>(
	241	+ scheduler, params, gmmAddrInfo, config, startLoopIdx, tileNum, activationQuantOp, expertBeforeCnt, expertIdx,
	242	+ layouts, pipeline);
	243	+
	244	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
	245	+}
	246	+
	247	+template <typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA, typename ElementMxScaleB,
	248	+ uint32_t Gmm1TileM = L1_TILE_M_256, uint32_t EpilogueTileM = Gmm1TileM, bool TopkWeightsPrefetch = false,
	249	+ bool IsShared = false, bool IsWaveFlagGrained = false, typename BlockContext, typename BlockEpilogue>
	250	+__aicore__ inline void RunGmm1A8W4(const BlockContext& pipeline, BlockEpilogue& activationQuantOp, const Params& params,
	251	+ const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
	252	+ const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx, uint32_t expertBeforeCnt,
	253	+ uint32_t expertIdx = 0, uint32_t activationOutputDivisor = ACTIVATION_N_HALF)
	254	+{
	255	+ BlockJobContext blockJob{MegaMoeHardwarePolicy::GetAicBlockIdx(),
	256	+ static_cast<uint32_t>(MegaMoeHardwarePolicy::GetBlockNum())};
	257	+ RunGmm1A8W4<ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, Gmm1TileM, EpilogueTileM,
	258	+ TopkWeightsPrefetch, IsShared, IsWaveFlagGrained>(pipeline, activationQuantOp, params, problemShape,
	259	+ gmmAddrInfo, startBlockIdx, blockJob, expertBeforeCnt,
	260	+ expertIdx, activationOutputDivisor);
	261	+}
	262	+
	263	+// A5 的地址准备与 Activation 绑定按本架构 AIV 分工执行。
	264	+template <typename ActivationType, typename WeightType, typename ActivationOutType, typename QuantScaleType,
	265	+ uint32_t ActivationElementsPerByte, bool UseA8W4Gmm1, bool TopkWeightsPrefetch, typename BlockEpilogue>
	266	+__aicore__ inline void UpdateMoeExpertGmm1GlobalBuffer(const GmmExecutionConfig& gmmConfig,
	267	+ const MoeSyncWorkspaceLayout& syncLayout,
	268	+ const WorkspaceLayout& workspace, GM_ADDR workspaceBase,
	269	+ const ExpertWeightTensorListAddrs& weights,
	270	+ BlockEpilogue& epilogueOp, GMMAddrInfo& gmmAddrInfo,
	271	+ const ExpertLoopState& state, uint32_t activationOutputWidth,
	272	+ uint32_t rowOffsetInExpert = 0U)
	273	+{
	274	+ if constexpr (g_coreType == AIV && !UseA8W4Gmm1) {
	275	+ if (MegaMoeHardwarePolicy::GetSubBlockIdx() != 0U) {
	276	+ return;
	277	+ }
	278	+ }
	279	+
	280	+ constexpr uint32_t weightElementsPerByte = PackedElementTraits<WeightType>::ELEMENTS_PER_BYTE;
	281	+ constexpr uint32_t outputElementsPerByte = PackedElementTraits<ActivationOutType>::ELEMENTS_PER_BYTE;
	282	+ int64_t n = Get<N_VALUE>(state.problemShape);
	283	+ int64_t k = Get<K_VALUE>(state.problemShape);
	284	+ int64_t expertOffset = state.expertIdx;
	285	+ int64_t globalTokenStartIndex = state.globalTokenStartIndex + rowOffsetInExpert;
	286	+ uint64_t physicalTokenStartIndex =
	287	+ gmmConfig.ring.gmm1 ?
	288	+ CalcMteRingPhysicalRow(state.paddedTokenStartIndex + rowOffsetInExpert, gmmConfig.ring.rows) :
	289	+ static_cast<uint64_t>(globalTokenStartIndex);
	290	+ // Dispatch 输入按自身循环开关选择行号；GMM1 输出和激活仍按各自的循环物理行寻址。
	291	+ const uint64_t dispatchTokenStartIndex =
	292	+ gmmConfig.ring.dispatch ? physicalTokenStartIndex : static_cast<uint64_t>(globalTokenStartIndex);
	293	+ uint32_t expertMGroupOffset = rowOffsetInExpert / L1_TILE_M_256;
	294	+ int64_t scaleK = Ops::Base::CeilDiv(k, static_cast<int64_t>(MXFP_DIVISOR_SIZE)) * MXFP_MULTI_BASE_SIZE;
	295	+
	296	+ if constexpr (UseA8W4Gmm1 \|\| TopkWeightsPrefetch) {
	297	+ gmmAddrInfo.gmm1OutGlobal =
	298	+ workspace.routed.gmm1Output.get_ptr(workspaceBase) + physicalTokenStartIndex * n * sizeof(bfloat16_t);
	299	+ }
	300	+ gmmAddrInfo.metaInfoGlobal = workspace.routing.receivedTokenMetadata.get_ptr(workspaceBase);
	301	+ gmmAddrInfo.activationToGmm2Flag =
	302	+ reinterpret_cast<__gm__ int32_t*>(workspace.sync.activationReady.get_ptr(workspaceBase)) +
	303	+ expertOffset * syncLayout.activationFlagSlotCountPerExpert +
	304	+ static_cast<uint64_t>(expertMGroupOffset) * INT_CACHELINE;
	305	+ gmmAddrInfo.aGlobal = workspace.routed.gmm1Input.data.get_ptr(workspaceBase) +
	306	+ dispatchTokenStartIndex * k / ActivationElementsPerByte * sizeof(ActivationType);
	307	+ gmmAddrInfo.aScaleGlobal = workspace.routed.gmm1Input.scale.get_ptr(workspaceBase) +
	308	+ dispatchTokenStartIndex * scaleK * sizeof(QuantScaleType);
	309	+ gmmAddrInfo.bGlobal =
	310	+ GetExpertWeightAddr<WeightType>(weights.weight1, gmmConfig.isPerExpertWeightTensor, state.expertIdx,
	311	+ static_cast<uint64_t>(state.expertIdx) * n * k / weightElementsPerByte);
	312	+ gmmAddrInfo.bScaleGlobal =
	313	+ GetExpertWeightAddr<QuantScaleType>(weights.weightScales1, gmmConfig.isPerExpertWeightTensor, state.expertIdx,
	314	+ static_cast<uint64_t>(state.expertIdx) * n * scaleK);
	315	+ if constexpr (g_coreType == AIV) {
	316	+ bool runsActivation = true;
	317	+ if constexpr (UseA8W4Gmm1) {
	318	+ runsActivation = MegaMoeHardwarePolicy::GetSubBlockIdx() == 1U;
	319	+ }
	320	+ if (runsActivation) {
	321	+ int64_t scaleN = Ops::Base::CeilDiv(static_cast<int64_t>(activationOutputWidth),
	322	+ static_cast<int64_t>(MXFP_DIVISOR_SIZE)) *
	323	+ MXFP_MULTI_BASE_SIZE;
	324	+ AscendC::Coord<int64_t, int64_t, int64_t, int64_t, int64_t, int64_t> vecBaseOffset{
	325	+ static_cast<int64_t>(physicalTokenStartIndex) * activationOutputWidth / outputElementsPerByte,
	326	+ static_cast<int64_t>(physicalTokenStartIndex) * scaleN,
	327	+ expertOffset * syncLayout.activationFlagSlotCountPerExpert / INT_CACHELINE + expertMGroupOffset,
	328	+ 0L,
	329	+ 0L,
	330	+ 0L};
	331	+ epilogueOp.UpdateGlobalAddr(vecBaseOffset);
	332	+ }
	333	+ }
	334	+ if constexpr (g_coreType == AIC) {
	335	+ gmmAddrInfo.dispatchToGmm1Flag =
	336	+ reinterpret_cast<__gm__ int32_t*>(workspace.sync.dispatchReady.get_ptr(workspaceBase)) +
	337	+ expertOffset * syncLayout.dispatchFlagSlotCountPerExpert +
	338	+ static_cast<uint64_t>(expertMGroupOffset) * INT_CACHELINE;
	339	+ }
	340	+}
	341	+
	342	+} // namespace MegaMoeImpl
	343	+
	344	+#endif // MEGA_MOE_GMM1_ACTIVATION_ARCH35_H

Amc2/mega_moe/op_kernel/arch35/stage/mega_moe_gmm2_combine_arch35.h+256-0

@@ -0,0 +1,256 @@
	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	+#ifndef MEGA_MOE_GMM2_COMBINE_ARCH35_H
	12	+#define MEGA_MOE_GMM2_COMBINE_ARCH35_H
	13	+
	14	+#include "../../arch35_92/stage/mega_moe_gmm2_combine.h"
	15	+
	16	+namespace MegaMoeImpl {
	17	+namespace GmmKernel {
	18	+
	19	+template <uint8_t CombineQuantMode, typename BlockMmad, typename ElementC, bool IsLayered = false,
	20	+ bool IsShared = false, bool NotifyCombineTileReady = false, typename Scheduler, typename Config>
	21	+__aicore__ inline void Gmm2ExecGeneric(Scheduler& scheduler, const GMMAddrInfo& gmmAddrInfo, const Config& config,
	22	+ uint32_t startLoopIdx, uint32_t tileNum, bool allowWeightL2Bypass,
	23	+ uint32_t rowOffsetInExpert = 0U, const Params* params = nullptr,
	24	+ uint32_t expertTokenCount = 0U)
	25	+{
	26	+ using KernelConfig = typename Config::KernelConfig;
	27	+ using ElementA = typename KernelConfig::ElementAType;
	28	+ using ElementB = typename KernelConfig::ElementBType;
	29	+ using ElementMxScaleA = typename KernelConfig::ElementMxScaleAType;
	30	+ using ElementMxScaleB = typename KernelConfig::ElementMxScaleBType;
	31	+ using BiasType = typename KernelConfig::BiasType;
	32	+
	33	+ auto layouts = KernelConfig::BuildLayouts(config);
	34	+ auto gmA = asc::te::make_tensor(
	35	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementA*>(gmmAddrInfo.aGlobal)),
	36	+ layouts.a);
	37	+ auto gmB = asc::te::make_tensor(
	38	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementB*>(gmmAddrInfo.bGlobal)),
	39	+ layouts.b);
	40	+ auto gmScaleA = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
	41	+ reinterpret_cast<__gm__ ElementMxScaleA*>(gmmAddrInfo.aScaleGlobal)),
	42	+ layouts.scaleA);
	43	+ auto gmScaleB = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
	44	+ reinterpret_cast<__gm__ ElementMxScaleB*>(gmmAddrInfo.bScaleGlobal)),
	45	+ layouts.scaleB);
	46	+ if constexpr (Config::IS_WAVE_FLAG_GRAINED && g_coreType == AscendC::AIC) {
	47	+ SetWaveWeightL2CacheHint<KernelConfig::IS_WEIGHT_NZ, KernelConfig>(config, allowWeightL2Bypass, gmB, gmScaleB);
	48	+ }
	49	+ auto gmBias = asc::te::make_tensor(
	50	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ BiasType*>(0UL)), layouts.bias);
	51	+ auto gmC = asc::te::make_tensor(
	52	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(gmmAddrInfo.gmm2OutGlobal)),
	53	+ layouts.c);
	54	+
	55	+ using WorkSetType = GroupMatmulWorkSet<Scheduler, decltype(gmA), decltype(gmB), decltype(gmScaleA),
	56	+ decltype(gmScaleB), decltype(gmBias), decltype(gmC)>;
	57	+ WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmBias, gmC};
	58	+
	59	+ if constexpr (g_coreType == AscendC::AIC) {
	60	+ BlockMmadContext<BlockMmad> localBlockMmadContext;
	61	+ localBlockMmadContext.template Init<GmmStage::GMM2>(config);
	62	+ Gmm2AicMmadGeneric<CombineQuantMode, BlockMmad, IsShared, IsLayered, NotifyCombineTileReady>(
	63	+ localBlockMmadContext.blockMmad, workSet, gmmAddrInfo, config, startLoopIdx, tileNum, rowOffsetInExpert,
	64	+ expertTokenCount);
	65	+ } else if constexpr (NotifyCombineTileReady) {
	66	+ // NotifyCombineTileReady 为真时 AIV1 逐 tile 锁步消费。combine 去重的独立消费路径
	67	+ // （RunWaveCombineStageDedup）由调用方把本参数实例化为 false，整个分支编译期裁掉。
	68	+ if (MegaMoeHardwarePolicy::GetSubBlockIdx() == 1U) {
	69	+ using MakeLayoutC = typename KernelConfig::MakeLayoutC;
	70	+ CombineTokenRange<ElementC, MakeLayoutC>(workSet.scheduler, workSet.gmC, *params, gmmAddrInfo, config,
	71	+ startLoopIdx, tileNum);
	72	+ }
	73	+ }
	74	+}
	75	+
	76	+template <typename BlockMmad, typename BlockPrologue, typename ElementC, bool IsShared, bool IsLayered,
	77	+ bool NotifyCombineTileReady, typename Scheduler, typename Config>
	78	+__aicore__ inline void Gmm2ExecA8W4(Scheduler& scheduler, const GMMAddrInfo& gmmAddrInfo, const Config& config,
	79	+ uint32_t startLoopIdx, uint32_t tileNum, uint32_t expertTokenCount,
	80	+ uint32_t rowOffsetInExpert, const Params* params,
	81	+ const A8W4BlockContext<BlockMmad, BlockPrologue>& pipeline)
	82	+{
	83	+ using KernelConfig = typename Config::KernelConfig;
	84	+ using ElementA = typename KernelConfig::ElementAType;
	85	+ using ElementB = typename KernelConfig::ElementBType;
	86	+ using ElementMxScaleA = typename KernelConfig::ElementMxScaleAType;
	87	+ using ElementMxScaleB = typename KernelConfig::ElementMxScaleBType;
	88	+ using BiasType = typename KernelConfig::BiasType;
	89	+
	90	+ auto layouts = KernelConfig::BuildLayouts(config);
	91	+ auto gmC = asc::te::make_tensor(
	92	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(gmmAddrInfo.gmm2OutGlobal)),
	93	+ layouts.c);
	94	+ auto gmA = asc::te::make_tensor(
	95	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementA*>(gmmAddrInfo.aGlobal)),
	96	+ layouts.a);
	97	+ auto gmB = asc::te::make_tensor(
	98	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementB*>(gmmAddrInfo.bGlobal)),
	99	+ layouts.b);
	100	+ auto gmScaleA = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
	101	+ reinterpret_cast<__gm__ ElementMxScaleA*>(gmmAddrInfo.aScaleGlobal)),
	102	+ layouts.scaleA);
	103	+ auto gmScaleB = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
	104	+ reinterpret_cast<__gm__ ElementMxScaleB*>(gmmAddrInfo.bScaleGlobal)),
	105	+ layouts.scaleB);
	106	+ auto gmBias = asc::te::make_tensor(
	107	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ BiasType*>(0UL)), layouts.bias);
	108	+
	109	+ using WorkSetType = GroupMatmulWorkSet<Scheduler, decltype(gmA), decltype(gmB), decltype(gmScaleA),
	110	+ decltype(gmScaleB), decltype(gmBias), decltype(gmC)>;
	111	+ WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmBias, gmC};
	112	+
	113	+ if constexpr (g_coreType == AscendC::AIC) {
	114	+ auto& blockMmad = *pipeline.block;
	115	+ typename BlockMmad::BlockShape l0TileShape{config.blockMmadTiling.tileM, config.blockMmadTiling.tileN,
	116	+ L0_TILE_K, 0};
	117	+ typename BlockMmad::ProblemShape matmulShape{config.m, config.outputN, config.k, 0};
	118	+ // GMM2 has no gate/up concatenation; preserve the persistent buffer indices.
	119	+ blockMmad.template Init<false>(matmulShape, l0TileShape, false);
	120	+ Gmm2AicMmadA8W4<IsShared, IsLayered, NotifyCombineTileReady, BlockMmad, decltype(workSet.scheduler),
	121	+ decltype(workSet.gmA), decltype(workSet.gmScaleA), decltype(workSet.gmScaleB),
	122	+ std::remove_reference_t<decltype(workSet.gmC)>, Config>(
	123	+ blockMmad, workSet.scheduler, workSet.gmA, workSet.gmScaleA, workSet.gmScaleB, workSet.gmC, gmmAddrInfo,
	124	+ config, startLoopIdx, tileNum, expertTokenCount, rowOffsetInExpert);
	125	+ } else {
	126	+ if (MegaMoeHardwarePolicy::GetSubBlockIdx() == 0U) {
	127	+ Gmm2Aiv0PrologueA8W4(*pipeline.block, workSet.scheduler, workSet.gmB, config, startLoopIdx, tileNum);
	128	+ } else if constexpr (NotifyCombineTileReady) {
	129	+ // 同上：W4 路径去重开关走运行时判断（嵌套模板参数触发编译器代码生成问题）。
	130	+ if (!IsCombineDedupOn(params->tilingData->compute.dedupMode)) {
	131	+ using MakeLayoutC = typename KernelConfig::MakeLayoutC;
	132	+ CombineTokenRange<ElementC, MakeLayoutC>(workSet.scheduler, workSet.gmC, *params, gmmAddrInfo, config,
	133	+ startLoopIdx, tileNum);
	134	+ }
	135	+ }
	136	+ }
	137	+}
	138	+
	139	+} // namespace GmmKernel
	140	+
	141	+template <uint8_t CombineQuantMode, typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA,
	142	+ typename ElementMxScaleB, bool IsWeightNZ = false, bool IsLayered = false, uint32_t Gmm1TileM = L1_TILE_M_256,
	143	+ bool TopkWeightsPrefetch = false, bool IsShared = false, bool IsWaveFlagGrained = false,
	144	+ bool NotifyCombineTileReady = false>
	145	+__aicore__ inline void RunGmm2Generic(const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
	146	+ const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,
	147	+ const BlockJobContext& blockJob, bool allowWeightL2Bypass = false,
	148	+ uint32_t rowOffsetInExpert = 0U, const Params* params = nullptr,
	149	+ uint32_t activationOutputDivisor = ACTIVATION_N_HALF,
	150	+ uint32_t expertTokenCount = 0U)
	151	+{
	152	+ static_assert(!IsShared \|\| !NotifyCombineTileReady, "Shared expert GMM2 has its own completion counter");
	153	+ static_assert(!IsLayered \|\| !NotifyCombineTileReady, "Layered GMM2 selects its synchronization through IsLayered");
	154	+ using GmmConfig =
	155	+ GmmKernel::Config<false, CombineQuantMode, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB,
	156	+ IsWeightNZ, TopkWeightsPrefetch, IsShared, IsLayered, IsWaveFlagGrained>;
	157	+ auto config = GmmConfig::BuildGmm2ProblemConfig(problemShape, blockJob, Gmm1TileM, activationOutputDivisor);
	158	+
	159	+ GmmKernel::BlockScheduler scheduler({config.m, config.schedulerN, config.k},
	160	+ GmmKernel::BlockScheduler::Params{asc::te::make_coord(
	161	+ static_cast<int64_t>(config.tileM), static_cast<int64_t>(L1_TILE_N))});
	162	+ uint32_t tileNum = scheduler.GetTileNum();
	163	+
	164	+ if constexpr (CombineQuantMode == COMBINE_NO_QUANT && g_coreType == AscendC::AIV && !NotifyCombineTileReady) {
	165	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
	166	+ return;
	167	+ }
	168	+
	169	+ uint32_t startLoopIdx =
	170	+ (config.blockIdx < startBlockIdx ? config.blockIdx + config.blockNum : config.blockIdx) - startBlockIdx;
	171	+ using BlockMmad = typename GmmConfig::BlockMmad;
	172	+ GmmKernel::Gmm2ExecGeneric<CombineQuantMode, BlockMmad, ElementC, IsLayered, IsShared, NotifyCombineTileReady>(
	173	+ scheduler, gmmAddrInfo, config, startLoopIdx, tileNum, allowWeightL2Bypass, rowOffsetInExpert, params,
	174	+ expertTokenCount);
	175	+
	176	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
	177	+}
	178	+
	179	+template <uint8_t CombineQuantMode, typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA,
	180	+ typename ElementMxScaleB, bool IsWeightNZ = false, bool IsLayered = false, uint32_t Gmm1TileM = L1_TILE_M_256,
	181	+ bool TopkWeightsPrefetch = false, bool IsShared = false, bool IsWaveFlagGrained = false,
	182	+ bool NotifyCombineTileReady = false>
	183	+__aicore__ inline void RunGmm2Generic(const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
	184	+ const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,
	185	+ bool allowWeightL2Bypass = false, uint32_t rowOffsetInExpert = 0U,
	186	+ const Params* params = nullptr,
	187	+ uint32_t activationOutputDivisor = ACTIVATION_N_HALF,
	188	+ uint32_t expertTokenCount = 0U)
	189	+{
	190	+ BlockJobContext blockJob{MegaMoeHardwarePolicy::GetAicBlockIdx(),
	191	+ static_cast<uint32_t>(MegaMoeHardwarePolicy::GetBlockNum())};
	192	+ RunGmm2Generic<CombineQuantMode, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, IsWeightNZ,
	193	+ IsLayered, Gmm1TileM, TopkWeightsPrefetch, IsShared, IsWaveFlagGrained, NotifyCombineTileReady>(
	194	+ problemShape, gmmAddrInfo, startBlockIdx, blockJob, allowWeightL2Bypass, rowOffsetInExpert, params,
	195	+ activationOutputDivisor, expertTokenCount);
	196	+}
	197	+
	198	+template <typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA, typename ElementMxScaleB,
	199	+ uint32_t Gmm1TileM = L1_TILE_M_256, bool TopkWeightsPrefetch = false, bool IsShared = false,
	200	+ bool IsLayered = false, bool IsWaveFlagGrained = false, bool NotifyCombineTileReady = false,
	201	+ typename BlockContext>
	202	+__aicore__ inline void RunGmm2A8W4(const BlockContext& pipeline,
	203	+ const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
	204	+ const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,
	205	+ const BlockJobContext& blockJob, uint32_t expertTokenCount,
	206	+ uint32_t rowOffsetInExpert, const Params* params = nullptr,
	207	+ uint32_t activationOutputDivisor = ACTIVATION_N_HALF)
	208	+{
	209	+ static_assert(!IsShared \|\| !NotifyCombineTileReady, "Shared expert GMM2 has its own completion counter");
	210	+ static_assert(!IsLayered \|\| !NotifyCombineTileReady, "Layered GMM2 selects its synchronization through IsLayered");
	211	+ static_assert(std::is_same_v<ElementA, __fp8e4m3>, "Activation must be __fp8e4m3");
	212	+ static_assert(std::is_same_v<ElementB, __fp4e2m1x2>, "Weight must be __fp4e2m1x2");
	213	+
	214	+ using GmmConfig = GmmKernel::Config<true, 0, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, false,
	215	+ TopkWeightsPrefetch, IsShared, IsLayered, IsWaveFlagGrained>;
	216	+ auto config = GmmConfig::BuildGmm2ProblemConfig(problemShape, blockJob, Gmm1TileM, activationOutputDivisor);
	217	+
	218	+ using BlockMmad = typename GmmConfig::BlockMmad;
	219	+ using BlockPrologue = typename GmmConfig::BlockPrologue;
	220	+ GmmKernel::BlockScheduler scheduler({config.m, config.outputN, config.k},
	221	+ GmmKernel::BlockScheduler::Params{asc::te::make_coord(
	222	+ static_cast<int64_t>(config.tileM), static_cast<int64_t>(L1_TILE_N))});
	223	+ // tileNum 与 mega_moe_utils.h 的 GetGmmProblemTileCount 同源，预判无活的口径不得各算一份。
	224	+ uint32_t tileNum = GetGmmProblemTileCount(config.m, config.outputN);
	225	+ uint32_t startLoopIdx =
	226	+ (config.blockIdx < startBlockIdx ? config.blockIdx + config.blockNum : config.blockIdx) - startBlockIdx;
	227	+ if (startLoopIdx >= tileNum) {
	228	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
	229	+ return;
	230	+ }
	231	+
	232	+ GmmKernel::Gmm2ExecA8W4<BlockMmad, BlockPrologue, ElementC, IsShared, IsLayered, NotifyCombineTileReady,
	233	+ GmmKernel::BlockScheduler, decltype(config)>(
	234	+ scheduler, gmmAddrInfo, config, startLoopIdx, tileNum, expertTokenCount, rowOffsetInExpert, params, pipeline);
	235	+
	236	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
	237	+}
	238	+
	239	+// 默认由 AIV1 复用既有配置，仅为 AIV0 绑定最终 Combine 的 UB 视图。steady Wave 使用较小行环，
	240	+// 调用方排空旧行环后通过 ReinitializeAllAiv 让两个 AIV 都恢复最终 Wave 的最大容量。
	241	+template <uint8_t CombineMode, bool ReinitializeAllAiv = false>
	242	+__aicore__ inline WaveCombineBufferConfig PrepareFinalWaveCombineBuffers(
	243	+ const MoeStageCommonConfig& common, const WaveCombineBufferConfig& currentBufferConfig, WaveCombineScratch& scratch)
	244	+{
	245	+ if constexpr (g_coreType == AIC) {
	246	+ return currentBufferConfig;
	247	+ }
	248	+ if (ReinitializeAllAiv \|\| MegaMoeHardwarePolicy::GetSubBlockIdx() == 0U) {
	249	+ return InitWaveCombineBuffers<CombineMode>(common, scratch);
	250	+ }
	251	+ return currentBufferConfig;
	252	+}
	253	+
	254	+} // namespace MegaMoeImpl
	255	+
	256	+#endif // MEGA_MOE_GMM2_COMBINE_ARCH35_H

Mmc2/mega_moe/op_kernel/arch35_92/common/mega_moe_arch_policy.h+0-11

	@@ -114,17 +114,6 @@ using MegaMoeHardwarePolicy = MegaMoeArch92HardwarePolicy;
114	using MegaMoeHardwarePolicy = MegaMoeArch35HardwarePolicy;	114	using MegaMoeHardwarePolicy = MegaMoeArch35HardwarePolicy;
115	#endif	115	#endif
116		116
117	-// 单 AIV 执行原先按子核号分配的分支；双 AIV 按真实子核号选择。
118	-template <uint32_t SubBlockIdx>
119	-__aicore__ inline bool IsSubBlockIdx()
120	-{
121	- if constexpr (MegaMoeHardwarePolicy::VECTOR_CORE_COUNT_PER_CUBE == 1U) {
122	- return true;
123	- } else {
124	- return MegaMoeHardwarePolicy::GetSubBlockIdx() == SubBlockIdx;
125	- }
126	-}
127	-
128	} // namespace MegaMoeImpl	117	} // namespace MegaMoeImpl
		W W wang-minbo1 day ago 【建议】IsSubBlockIdx 删干净了，不过 PR 描述的类型标签一个都没勾，这次明显是重构，把「♻️ 重构」勾上方便后续按类型检索；「关联的Issue」也是空的，有对应的需求单/问题单建议补个链接。
129		118
130	#endif // MEGA_MOE_ARCH_POLICY_H	119	#endif // MEGA_MOE_ARCH_POLICY_H

Mmc2/mega_moe/op_kernel/arch35_92/stage/mega_moe_gmm1_activation.h+0-341

	@@ -452,85 +452,6 @@ __aicore__ inline void Gmm1Aiv0EpilogueGeneric(WorkSet& workSet, const GMMAddrIn
452	}	452	}
453	}	453	}
454		454
455	-// 根据 GM 地址建立执行资源，并执行通用 GMM1/Activation 阶段。
456	-template <typename BlockMmad, typename ElementC, typename MakeLayoutC, bool TopkWeightsPrefetch, bool IsShared = false,
457	- bool IsWaveFlagGrained = false, typename Scheduler, typename Config, typename ActivationQuantOp,
458	- typename LayoutBundle>
459	-__aicore__ inline void Gmm1ExecGeneric(Scheduler& scheduler, const Params& params, const GMMAddrInfo& gmmAddrInfo,
460	- const Config& config, uint32_t startLoopIdx, uint32_t tileNum,
461	- ActivationQuantOp& activationQuantOp, int32_t& vecSetSyncCom,
462	- uint32_t expertBeforeCnt, uint32_t expertIdx, uint16_t& pingpongIdx,
463	- BlockMmadContext<BlockMmad>* blockMmadContext, bool allowWeightL2Bypass,
464	- const LayoutBundle& layouts)
465	-{
466	- using KernelConfig = typename Config::KernelConfig;
467	- using ElementA = typename KernelConfig::ElementAType;
468	- using ElementB = typename KernelConfig::ElementBType;
469	- using ElementMxScaleA = typename KernelConfig::ElementMxScaleAType;
470	- using ElementMxScaleB = typename KernelConfig::ElementMxScaleBType;
471	- auto gmA = asc::te::make_tensor(
472	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementA*>(gmmAddrInfo.aGlobal)),
473	- layouts.a);
474	- auto gmB = asc::te::make_tensor(
475	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementB*>(gmmAddrInfo.bGlobal)),
476	- layouts.b);
477	- auto gmScaleA = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
478	- reinterpret_cast<__gm__ ElementMxScaleA*>(gmmAddrInfo.aScaleGlobal)),
479	- layouts.scaleA);
480	- auto gmScaleB = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
481	- reinterpret_cast<__gm__ ElementMxScaleB*>(gmmAddrInfo.bScaleGlobal)),
482	- layouts.scaleB);
483	- if constexpr (IsWaveFlagGrained && g_coreType == AscendC::AIC) {
484	- SetWaveWeightL2CacheHint<KernelConfig::IS_WEIGHT_NZ, KernelConfig>(config, allowWeightL2Bypass, gmB, gmScaleB);
485	- }
486	- GM_ADDR cGlobal = 0UL;
487	- if constexpr (TopkWeightsPrefetch) {
488	- cGlobal = gmmAddrInfo.gmm1OutGlobal;
489	- }
490	- auto gmC = asc::te::make_tensor(
491	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(cGlobal)), layouts.c);
492	- __gm__ float* metaInfoGm = nullptr;
493	- if constexpr (TopkWeightsPrefetch && g_coreType == AscendC::AIV) {
494	- metaInfoGm = reinterpret_cast<__gm__ float*>(gmmAddrInfo.metaInfoGlobal);
495	- }
496	-
497	- using WorkSetType = Gmm1WorkSet<Scheduler, decltype(gmA), decltype(gmB), decltype(gmScaleA), decltype(gmScaleB),
498	- decltype(gmC), decltype(metaInfoGm)>;
499	- WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmC, metaInfoGm};
500	- if constexpr (g_coreType == AscendC::AIC) {
501	- if (blockMmadContext != nullptr) {
502	- if constexpr (TopkWeightsPrefetch) {
503	- Gmm1AicMmadToGmGeneric<BlockMmadContext<BlockMmad>, IsShared, TopkWeightsPrefetch>(
504	- workSet, params, gmmAddrInfo, config, startLoopIdx, tileNum, expertIdx, blockMmadContext);
505	- } else {
506	- Gmm1AicMmadGeneric<BlockMmadContext<BlockMmad>, IsShared>(
507	- workSet, gmmAddrInfo, config, startLoopIdx, tileNum, vecSetSyncCom, pingpongIdx, blockMmadContext);
508	- }
509	- } else {
510	- BlockMmadContext<BlockMmad> localBlockMmadContext;
511	- if constexpr (TopkWeightsPrefetch) {
512	- Gmm1AicMmadToGmGeneric<BlockMmadContext<BlockMmad>, IsShared, TopkWeightsPrefetch>(
513	- workSet, params, gmmAddrInfo, config, startLoopIdx, tileNum, expertIdx, &localBlockMmadContext);
514	- } else {
515	- Gmm1AicMmadGeneric<BlockMmadContext<BlockMmad>, IsShared>(workSet, gmmAddrInfo, config, startLoopIdx,
516	- tileNum, vecSetSyncCom, pingpongIdx,
517	- &localBlockMmadContext);
518	- }
519	- }
520	- } else {
521	- if constexpr (TopkWeightsPrefetch) {
522	- if (IsSubBlockIdx<0U>()) {
523	- Gmm1Aiv0PrefetchEpilogueGeneric<MakeLayoutC, ElementC, IsWaveFlagGrained>(
524	- workSet, params, gmmAddrInfo, config, startLoopIdx, tileNum, activationQuantOp, expertBeforeCnt,
525	- expertIdx);
526	- }
527	- } else if (IsSubBlockIdx<0U>()) {
528	- Gmm1Aiv0EpilogueGeneric<IsWaveFlagGrained>(workSet, gmmAddrInfo, config, startLoopIdx, tileNum,
529	- activationQuantOp, pingpongIdx);
530	- }
531	- }
532	-}
533	-
534	// Prefetch 路径的 AIV1 通过硬同步等待 tile 就绪，再执行 A8W4 Activation/量化。	455	// Prefetch 路径的 AIV1 通过硬同步等待 tile 就绪，再执行 A8W4 Activation/量化。
535	template <typename ElementC, typename MakeLayoutC, bool IsWaveFlagGrained, typename WorkSet, typename Config,	456	template <typename ElementC, typename MakeLayoutC, bool IsWaveFlagGrained, typename WorkSet, typename Config,
536	typename ActivationQuantOp>	457	typename ActivationQuantOp>
	@@ -574,269 +495,7 @@ __aicore__ inline void Gmm1Aiv1EpilogueA8W4(WorkSet& workSet, const GMMAddrInfo&
574	}	495	}
575	}	496	}
576		497
577	-// 根据 GM 地址建立执行资源，并执行 A8W4 GMM1/Activation 阶段。
578	-template <typename BlockMmad, typename BlockPrologue, typename ElementC, typename MakeLayoutC, bool IsShared,
579	- typename Scheduler, typename Config, bool TopkWeightsPrefetch, bool IsWaveFlagGrained,
580	- typename ActivationQuantOp, typename LayoutBundle>
581	-__aicore__ inline void Gmm1ExecA8W4(Scheduler& scheduler, const Params& params, const GMMAddrInfo& gmmAddrInfo,
582	- const Config& config, uint32_t startLoopIdx, uint32_t tileNum,
583	- ActivationQuantOp& activationQuantOp, uint32_t expertBeforeCnt, uint32_t expertIdx,
584	- const LayoutBundle& layouts,
585	- const A8W4BlockContext<BlockMmad, BlockPrologue>& pipeline)
586	-{
587	- using KernelConfig = typename Config::KernelConfig;
588	- using ElementA = typename KernelConfig::ElementAType;
589	- using ElementB = typename KernelConfig::ElementBType;
590	- using ElementMxScaleA = typename KernelConfig::ElementMxScaleAType;
591	- using ElementMxScaleB = typename KernelConfig::ElementMxScaleBType;
592	- auto gmC = asc::te::make_tensor(
593	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(gmmAddrInfo.gmm1OutGlobal)),
594	- layouts.c);
595	- auto gmA = asc::te::make_tensor(
596	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementA*>(gmmAddrInfo.aGlobal)),
597	- layouts.a);
598	- auto gmB = asc::te::make_tensor(
599	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementB*>(gmmAddrInfo.bGlobal)),
600	- layouts.b);
601	- auto gmScaleA = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
602	- reinterpret_cast<__gm__ ElementMxScaleA*>(gmmAddrInfo.aScaleGlobal)),
603	- layouts.scaleA);
604	- auto gmScaleB = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
605	- reinterpret_cast<__gm__ ElementMxScaleB*>(gmmAddrInfo.bScaleGlobal)),
606	- layouts.scaleB);
607	- __gm__ float* metaInfoGm = nullptr;
608	- if constexpr (TopkWeightsPrefetch && g_coreType == AscendC::AIV) {
609	- metaInfoGm = reinterpret_cast<__gm__ float*>(gmmAddrInfo.metaInfoGlobal);
610	- }
611	-
612	- using WorkSetType = Gmm1WorkSet<Scheduler, decltype(gmA), decltype(gmB), decltype(gmScaleA), decltype(gmScaleB),
613	- decltype(gmC), decltype(metaInfoGm)>;
614	- WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmC, metaInfoGm};
615	- if constexpr (g_coreType == AscendC::AIC) {
616	- Gmm1AicMmadA8W4<BlockMmad, IsShared, TopkWeightsPrefetch>(*pipeline.block, workSet, params, gmmAddrInfo, config,
617	- startLoopIdx, tileNum, expertIdx);
618	- } else if (IsSubBlockIdx<0U>()) {
619	- Gmm1Aiv0PrologueA8W4<BlockPrologue>(*pipeline.block, workSet.scheduler, workSet.gmB, config, startLoopIdx,
620	- tileNum);
621	- } else {
622	- if constexpr (TopkWeightsPrefetch) {
623	- Gmm1Aiv1PrefetchEpilogueA8W4<ElementC, MakeLayoutC, IsWaveFlagGrained>(
624	- workSet, params, gmmAddrInfo, config, startLoopIdx, tileNum, activationQuantOp, expertBeforeCnt,
625	- expertIdx);
626	- } else {
627	- Gmm1Aiv1EpilogueA8W4<ElementC, MakeLayoutC, IsWaveFlagGrained>(workSet, gmmAddrInfo, config, startLoopIdx,
628	- tileNum, activationQuantOp);
629	- }
630	- }
631	-}
632	-
633	} // namespace GmmKernel	498	} // namespace GmmKernel
634	-
635	-template <typename ElementA, typename EpilogueElementA, typename ElementB, typename ElementC, typename ElementMxScaleA,
636	- typename ElementMxScaleB, bool IsWeightNZ = false, uint32_t Gmm1TileM = L1_TILE_M_256,
637	- uint32_t EpilogueTileM = Gmm1TileM, bool TopkWeightsPrefetch = false, bool IsShared = false,
638	- bool IsWaveFlagGrained = false, typename BlockEpilogue, typename... LayoutArgs>
639	-__aicore__ inline void RunGmm1Generic(BlockEpilogue& epilogueOp, const Params& params,
640	- const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
641	- const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx, int32_t& vecSetSyncCom,
642	- const BlockJobContext& blockJob, uint32_t expertBeforeCnt, uint32_t expertIdx,
643	- uint16_t& pingpongIdx, uint32_t activationOutputDivisor = ACTIVATION_N_HALF,
644	- void* blockMmadContext = nullptr, bool allowWeightL2Bypass = false,
645	- const LayoutArgs&... layoutArgs)
646	-{
647	- using GmmConfig =
648	- GmmKernel::Config<false, COMBINE_NO_QUANT, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB,
649	- IsWeightNZ, TopkWeightsPrefetch, IsShared, false, IsWaveFlagGrained>;
650	- auto config = GmmConfig::BuildGmm1ProblemConfig(problemShape, blockJob, Gmm1TileM, activationOutputDivisor);
651	-
652	- GmmKernel::BlockScheduler scheduler({config.m, config.schedulerN, config.k},
653	- GmmKernel::BlockScheduler::Params{asc::te::make_coord(
654	- static_cast<int64_t>(config.tileM), static_cast<int64_t>(L1_TILE_N))});
655	- uint32_t tileNum = scheduler.GetTileNum();
656	-
657	- if (!IsSubBlockIdx<0U>()) {
658	- startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
659	- return;
660	- }
661	- if constexpr (g_coreType == AIV) {
662	- epilogueOp.UpdateNextProblem({config.m, config.outputN, config.k, 0});
663	- }
664	-
665	- uint32_t startLoopIdx =
666	- (config.blockIdx < startBlockIdx ? config.blockIdx + config.blockNum : config.blockIdx) - startBlockIdx;
667	- using BlockMmad = typename GmmConfig::BlockMmad;
668	- using MmadContext = GmmKernel::BlockMmadContext<BlockMmad>;
669	- auto* typedBlockMmadContext = reinterpret_cast<MmadContext*>(blockMmadContext);
670	- using MakeLayoutC = typename GmmConfig::MakeLayoutC;
671	- auto layouts = GmmConfig::BuildLayouts(config, layoutArgs...);
672	- GmmKernel::Gmm1ExecGeneric<BlockMmad, ElementC, MakeLayoutC, TopkWeightsPrefetch, IsShared, IsWaveFlagGrained>(
673	- scheduler, params, gmmAddrInfo, config, startLoopIdx, tileNum, epilogueOp, vecSetSyncCom, expertBeforeCnt,
674	- expertIdx, pingpongIdx, typedBlockMmadContext, allowWeightL2Bypass, layouts);
675	-
676	- startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
677	-}
678	-
679	-template <typename ElementA, typename EpilogueElementA, typename ElementB, typename ElementC, typename ElementMxScaleA,
680	- typename ElementMxScaleB, bool IsWeightNZ = false, uint32_t Gmm1TileM = L1_TILE_M_256,
681	- uint32_t EpilogueTileM = Gmm1TileM, bool TopkWeightsPrefetch = false, bool IsShared = false,
682	- bool IsWaveFlagGrained = false, typename BlockEpilogue>
683	-__aicore__ inline void RunGmm1Generic(BlockEpilogue& epilogueOp, const Params& params,
684	- const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
685	- const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx, int32_t& vecSetSyncCom,
686	- uint32_t expertBeforeCnt, uint32_t expertIdx, uint16_t& pingpongIdx,
687	- uint32_t activationOutputDivisor = ACTIVATION_N_HALF,
688	- void* blockMmadContext = nullptr, bool allowWeightL2Bypass = false)
689	-{
690	- BlockJobContext blockJob{MegaMoeHardwarePolicy::GetAicBlockIdx(),
691	- static_cast<uint32_t>(MegaMoeHardwarePolicy::GetBlockNum())};
692	- RunGmm1Generic<ElementA, EpilogueElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, IsWeightNZ,
693	- Gmm1TileM, EpilogueTileM, TopkWeightsPrefetch, IsShared, IsWaveFlagGrained>(
694	- epilogueOp, params, problemShape, gmmAddrInfo, startBlockIdx, vecSetSyncCom, blockJob, expertBeforeCnt,
695	- expertIdx, pingpongIdx, activationOutputDivisor, blockMmadContext, allowWeightL2Bypass);
696	-}
697	-
698	-// RunGmm1A8W4：执行 A8W4 prologue（W4→W8）、GMM1、Activation 和量化。
699	-template <typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA, typename ElementMxScaleB,
700	- uint32_t Gmm1TileM = L1_TILE_M_256, uint32_t EpilogueTileM = Gmm1TileM, bool TopkWeightsPrefetch = false,
701	- bool IsShared = false, bool IsWaveFlagGrained = false, typename BlockContext, typename BlockEpilogue,
702	- typename... LayoutArgs>
703	-__aicore__ inline void RunGmm1A8W4(const BlockContext& pipeline, BlockEpilogue& activationQuantOp, const Params& params,
704	- const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
705	- const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,
706	- const BlockJobContext& blockJob, uint32_t expertBeforeCnt, uint32_t expertIdx,
707	- uint32_t activationOutputDivisor = ACTIVATION_N_HALF,
708	- const LayoutArgs&... layoutArgs)
709	-{
710	- static_assert(std::is_same_v<ElementA, __fp8e4m3>, "Activation must be __fp8e4m3");
711	- static_assert(std::is_same_v<ElementB, __fp4e2m1x2>, "Weight must be __fp4e2m1x2");
712	-
713	- using GmmConfig = GmmKernel::Config<true, 0, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, false,
714	- TopkWeightsPrefetch, IsShared, false, IsWaveFlagGrained>;
715	- auto config = GmmConfig::BuildGmm1ProblemConfig(problemShape, blockJob, Gmm1TileM, activationOutputDivisor);
716	-
717	- if constexpr (g_coreType == AIV) {
718	- activationQuantOp.UpdateNextProblem({config.m, config.outputN, config.k, 0});
719	- }
720	-
721	- using BlockMmad = typename GmmConfig::BlockMmad;
722	- using BlockPrologue = typename GmmConfig::BlockPrologue;
723	- using MakeLayoutC = typename GmmConfig::MakeLayoutC;
724	-
725	- GmmKernel::BlockScheduler scheduler({config.m, config.schedulerN, config.k},
726	- GmmKernel::BlockScheduler::Params{asc::te::make_coord(
727	- static_cast<int64_t>(config.tileM), static_cast<int64_t>(L1_TILE_N))});
728	- uint32_t tileNum = scheduler.GetTileNum();
729	- uint32_t startLoopIdx =
730	- (config.blockIdx < startBlockIdx ? config.blockIdx + config.blockNum : config.blockIdx) - startBlockIdx;
731	- if (startLoopIdx >= tileNum) {
732	- startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
733	- return;
734	- }
735	-
736	- auto layouts = GmmConfig::BuildLayouts(config, layoutArgs...);
737	- GmmKernel::Gmm1ExecA8W4<BlockMmad, BlockPrologue, ElementC, MakeLayoutC, IsShared, GmmKernel::BlockScheduler,
738	- decltype(config), TopkWeightsPrefetch, IsWaveFlagGrained>(
739	- scheduler, params, gmmAddrInfo, config, startLoopIdx, tileNum, activationQuantOp, expertBeforeCnt, expertIdx,
740	- layouts, pipeline);
741	-
742	- startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
743	-}
744	-
745	-// 使用当前核的任务配置执行 A8W4 GMM1。
746	-template <typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA, typename ElementMxScaleB,
747	- uint32_t Gmm1TileM = L1_TILE_M_256, uint32_t EpilogueTileM = Gmm1TileM, bool TopkWeightsPrefetch = false,
748	- bool IsShared = false, bool IsWaveFlagGrained = false, typename BlockContext, typename BlockEpilogue>
749	-__aicore__ inline void RunGmm1A8W4(const BlockContext& pipeline, BlockEpilogue& activationQuantOp, const Params& params,
750	- const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
751	- const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx, uint32_t expertBeforeCnt,
752	- uint32_t expertIdx = 0, uint32_t activationOutputDivisor = ACTIVATION_N_HALF)
753	-{
754	- BlockJobContext blockJob{MegaMoeHardwarePolicy::GetAicBlockIdx(),
755	- static_cast<uint32_t>(MegaMoeHardwarePolicy::GetBlockNum())};
756	- RunGmm1A8W4<ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, Gmm1TileM, EpilogueTileM,
757	- TopkWeightsPrefetch, IsShared, IsWaveFlagGrained>(pipeline, activationQuantOp, params, problemShape,
758	- gmmAddrInfo, startBlockIdx, blockJob, expertBeforeCnt,
759	- expertIdx, activationOutputDivisor);
760	-}
761	-
762	-template <typename ActivationType, typename WeightType, typename ActivationOutType, typename QuantScaleType,
763	- uint32_t ActivationElementsPerByte, bool UseA8W4Gmm1, bool TopkWeightsPrefetch, typename BlockEpilogue>
764	-__aicore__ inline void UpdateMoeExpertGmm1GlobalBuffer(const GmmExecutionConfig& gmmConfig,
765	- const MoeSyncWorkspaceLayout& syncLayout,
766	- const WorkspaceLayout& workspace, GM_ADDR workspaceBase,
767	- const ExpertWeightTensorListAddrs& weights,
768	- BlockEpilogue& epilogueOp, GMMAddrInfo& gmmAddrInfo,
769	- const ExpertLoopState& state, uint32_t activationOutputWidth,
770	- uint32_t rowOffsetInExpert = 0U)
771	-{
772	- if constexpr (g_coreType == AIV && !UseA8W4Gmm1) {
773	- if (!IsSubBlockIdx<0U>()) {
774	- return;
775	- }
776	- }
777	-
778	- constexpr uint32_t weightElementsPerByte = PackedElementTraits<WeightType>::ELEMENTS_PER_BYTE;
779	- constexpr uint32_t outputElementsPerByte = PackedElementTraits<ActivationOutType>::ELEMENTS_PER_BYTE;
780	- int64_t n = Get<N_VALUE>(state.problemShape);
781	- int64_t k = Get<K_VALUE>(state.problemShape);
782	- int64_t expertOffset = state.expertIdx;
783	- int64_t globalTokenStartIndex = state.globalTokenStartIndex + rowOffsetInExpert;
784	- uint64_t physicalTokenStartIndex =
785	- gmmConfig.ring.gmm1 ?
786	- CalcMteRingPhysicalRow(state.paddedTokenStartIndex + rowOffsetInExpert, gmmConfig.ring.rows) :
787	- static_cast<uint64_t>(globalTokenStartIndex);
788	- // Dispatch 输入按自身循环开关选择行号；GMM1 输出和激活仍按各自的循环物理行寻址。
789	- const uint64_t dispatchTokenStartIndex =
790	- gmmConfig.ring.dispatch ? physicalTokenStartIndex : static_cast<uint64_t>(globalTokenStartIndex);
791	- uint32_t expertMGroupOffset = rowOffsetInExpert / L1_TILE_M_256;
792	- int64_t scaleK = Ops::Base::CeilDiv(k, static_cast<int64_t>(MXFP_DIVISOR_SIZE)) * MXFP_MULTI_BASE_SIZE;
793	-
794	- if constexpr (UseA8W4Gmm1 \|\| TopkWeightsPrefetch) {
795	- gmmAddrInfo.gmm1OutGlobal =
796	- workspace.routed.gmm1Output.get_ptr(workspaceBase) + physicalTokenStartIndex * n * sizeof(bfloat16_t);
797	- }
798	- gmmAddrInfo.metaInfoGlobal = workspace.routing.receivedTokenMetadata.get_ptr(workspaceBase);
799	- gmmAddrInfo.activationToGmm2Flag =
800	- reinterpret_cast<__gm__ int32_t*>(workspace.sync.activationReady.get_ptr(workspaceBase)) +
801	- expertOffset * syncLayout.activationFlagSlotCountPerExpert +
802	- static_cast<uint64_t>(expertMGroupOffset) * INT_CACHELINE;
803	- gmmAddrInfo.aGlobal = workspace.routed.gmm1Input.data.get_ptr(workspaceBase) +
804	- dispatchTokenStartIndex * k / ActivationElementsPerByte * sizeof(ActivationType);
805	- gmmAddrInfo.aScaleGlobal = workspace.routed.gmm1Input.scale.get_ptr(workspaceBase) +
806	- dispatchTokenStartIndex * scaleK * sizeof(QuantScaleType);
807	- gmmAddrInfo.bGlobal =
808	- GetExpertWeightAddr<WeightType>(weights.weight1, gmmConfig.isPerExpertWeightTensor, state.expertIdx,
809	- static_cast<uint64_t>(state.expertIdx) * n * k / weightElementsPerByte);
810	- gmmAddrInfo.bScaleGlobal =
811	- GetExpertWeightAddr<QuantScaleType>(weights.weightScales1, gmmConfig.isPerExpertWeightTensor, state.expertIdx,
812	- static_cast<uint64_t>(state.expertIdx) * n * scaleK);
813	- if constexpr (g_coreType == AIV) {
814	- bool runsActivation = true;
815	- if constexpr (UseA8W4Gmm1) {
816	- runsActivation = IsSubBlockIdx<1U>();
817	- }
818	- if (runsActivation) {
819	- int64_t scaleN = Ops::Base::CeilDiv(static_cast<int64_t>(activationOutputWidth),
820	- static_cast<int64_t>(MXFP_DIVISOR_SIZE)) *
821	- MXFP_MULTI_BASE_SIZE;
822	- AscendC::Coord<int64_t, int64_t, int64_t, int64_t, int64_t, int64_t> vecBaseOffset{
823	- static_cast<int64_t>(physicalTokenStartIndex) * activationOutputWidth / outputElementsPerByte,
824	- static_cast<int64_t>(physicalTokenStartIndex) * scaleN,
825	- expertOffset * syncLayout.activationFlagSlotCountPerExpert / INT_CACHELINE + expertMGroupOffset,
826	- 0L,
827	- 0L,
828	- 0L};
829	- epilogueOp.UpdateGlobalAddr(vecBaseOffset);
830	- }
831	- }
832	- if constexpr (g_coreType == AIC) {
833	- gmmAddrInfo.dispatchToGmm1Flag =
834	- reinterpret_cast<__gm__ int32_t*>(workspace.sync.dispatchReady.get_ptr(workspaceBase)) +
835	- expertOffset * syncLayout.dispatchFlagSlotCountPerExpert +
836	- static_cast<uint64_t>(expertMGroupOffset) * INT_CACHELINE;
837	- }
838	-}
839	-
840	// 构造一个共享专家的 GMM1 输入、权重、输出与 Activation 量化输出 GM 视图。	499	// 构造一个共享专家的 GMM1 输入、权重、输出与 Activation 量化输出 GM 视图。
841	template <typename WeightType, typename ActivationOutType, typename QuantScaleType, bool UseGmOutput,	500	template <typename WeightType, typename ActivationOutType, typename QuantScaleType, bool UseGmOutput,
842	typename BlockEpilogue>	501	typename BlockEpilogue>

Mmc2/mega_moe/op_kernel/arch35_92/stage/mega_moe_gmm2_combine.h+8-273

	@@ -149,21 +149,14 @@ __aicore__ inline void CombineQuantizedTokens(uint32_t batchStart, uint32_t curR
149		149
150	} // namespace CombineImpl	150	} // namespace CombineImpl
151		151
152	-// 非 layered 路径统一使用的 Combine 行环。常规 Wave 只有 AIV1 建立并使用这些 UB 视图；	152	+// 调用方决定参与核与阶段；这里只计算 Combine 行环布局并绑定当前核的 UB 视图。
153	-// 最后一轮不再与 GMM1/Activation 并行时，调用方才为 AIV0 补充初始化。	153	+template <uint8_t CombineMode, uint32_t MaxRowBufferCount = WAVE_COMBINE_MAX_ROW_BUFFER_COUNT>
154	-template <uint8_t CombineMode, bool IncludeAiv0 = false, uint32_t MaxRowBufferCount = WAVE_COMBINE_MAX_ROW_BUFFER_COUNT>
155	__aicore__ inline WaveCombineBufferConfig InitWaveCombineBuffers(const MoeStageCommonConfig& common,	154	__aicore__ inline WaveCombineBufferConfig InitWaveCombineBuffers(const MoeStageCommonConfig& common,
156	WaveCombineScratch& scratch)	155	WaveCombineScratch& scratch)
157	{	156	{
158	static_assert(MaxRowBufferCount > 0U && MaxRowBufferCount <= WAVE_COMBINE_MAX_ROW_BUFFER_COUNT,	157	static_assert(MaxRowBufferCount > 0U && MaxRowBufferCount <= WAVE_COMBINE_MAX_ROW_BUFFER_COUNT,
159	"invalid Combine row buffer count");	158	"invalid Combine row buffer count");
160	WaveCombineBufferConfig bufferConfig{};	159	WaveCombineBufferConfig bufferConfig{};
161	- if constexpr (g_coreType == AIC) {
162	- return bufferConfig;
163	- }
164	- if (!IncludeAiv0 && !IsSubBlockIdx<1U>()) {
165	- return bufferConfig;
166	- }
167	bufferConfig.rowBytes = common.tokenHiddenDim * sizeof(bfloat16_t);	160	bufferConfig.rowBytes = common.tokenHiddenDim * sizeof(bfloat16_t);
168	bufferConfig.rowStrideBytes = static_cast<uint32_t>(	161	bufferConfig.rowStrideBytes = static_cast<uint32_t>(
169	Ops::Base::CeilAlign(static_cast<uint64_t>(bufferConfig.rowBytes), static_cast<uint64_t>(ALIGN_32)));	162	Ops::Base::CeilAlign(static_cast<uint64_t>(bufferConfig.rowBytes), static_cast<uint64_t>(ALIGN_32)));
	@@ -201,21 +194,6 @@ __aicore__ inline WaveCombineBufferConfig InitWaveCombineBuffers(const MoeStageC
201	return bufferConfig;	194	return bufferConfig;
202	}	195	}
203		196
204	-// 默认由 AIV1 复用既有配置，仅为 AIV0 绑定最终 Combine 的 UB 视图。steady Wave 使用较小行环，
205	-// 调用方排空旧行环后通过 ReinitializeAllAiv 让两个 AIV 都恢复最终 Wave 的最大容量。
206	-template <uint8_t CombineMode, bool ReinitializeAllAiv = false>
207	-__aicore__ inline WaveCombineBufferConfig PrepareFinalWaveCombineBuffers(
208	- const MoeStageCommonConfig& common, const WaveCombineBufferConfig& currentBufferConfig, WaveCombineScratch& scratch)
209	-{
210	- if constexpr (g_coreType == AIC) {
211	- return currentBufferConfig;
212	- }
213	- if (ReinitializeAllAiv \|\| IsSubBlockIdx<0U>()) {
214	- return InitWaveCombineBuffers<CombineMode, true>(common, scratch);
215	- }
216	- return currentBufferConfig;
217	-}
218	-
219	// 构造量化 Unpermute 输入记录的单 token 存储布局。	197	// 构造量化 Unpermute 输入记录的单 token 存储布局。
220	__aicore__ inline QuantTokenBufferConfig CreateQuantTokenBufferConfig(uint32_t tokenHiddenDim)	198	__aicore__ inline QuantTokenBufferConfig CreateQuantTokenBufferConfig(uint32_t tokenHiddenDim)
221	{	199	{
	@@ -228,7 +206,7 @@ __aicore__ inline QuantTokenBufferConfig CreateQuantTokenBufferConfig(uint32_t t
228	}	206	}
229		207
230	// 按专家累计 token 行数滚动物理 AIV 起点，每个 AIV 在当前专家内仍负责一个连续均衡区间。	208	// 按专家累计 token 行数滚动物理 AIV 起点，每个 AIV 在当前专家内仍负责一个连续均衡区间。
231	-// 常规 Wave 只使用每个 AIC 配对的 AIV1；最终专家可把同组 AIV0/AIV1 展开成两个独立任务。	209	+// arch35 常规 Wave 使用 AIV1，最终专家可展开两组 AIV；arch92 始终使用唯一 AIV 的原始任务编号。
232	template <bool IncludeAiv0 = false>	210	template <bool IncludeAiv0 = false>
233	__aicore__ inline WorkRange GetWaveCombineOwnedRange(const AivJobContext& job, uint32_t tokenCount,	211	__aicore__ inline WorkRange GetWaveCombineOwnedRange(const AivJobContext& job, uint32_t tokenCount,
234	uint64_t expertRowPrefix)	212	uint64_t expertRowPrefix)
	@@ -236,18 +214,22 @@ __aicore__ inline WorkRange GetWaveCombineOwnedRange(const AivJobContext& job, u
236	if constexpr (g_coreType == AIC) {	214	if constexpr (g_coreType == AIC) {
237	return {};	215	return {};
238	}	216	}
		217	+#ifdef MEGA_MOE_ARCH92_DEVICE
		218	+ return GetRotatedBalancedWorkRange(tokenCount, job.jobIndex, job.totalJobs, expertRowPrefix);
		219	+#else
239	AivJobContext combineJob = job;	220	AivJobContext combineJob = job;
240	if constexpr (IncludeAiv0) {	221	if constexpr (IncludeAiv0) {
241	constexpr uint32_t AIV_CORE_COUNT_PER_AIC = MegaMoeHardwarePolicy::VECTOR_CORE_COUNT_PER_CUBE;	222	constexpr uint32_t AIV_CORE_COUNT_PER_AIC = MegaMoeHardwarePolicy::VECTOR_CORE_COUNT_PER_CUBE;
242	uint32_t vectorRoleIdx = MegaMoeHardwarePolicy::GetSubBlockIdx();	223	uint32_t vectorRoleIdx = MegaMoeHardwarePolicy::GetSubBlockIdx();
243	combineJob = {job.jobIndex * AIV_CORE_COUNT_PER_AIC + vectorRoleIdx, job.totalJobs * AIV_CORE_COUNT_PER_AIC};	224	combineJob = {job.jobIndex * AIV_CORE_COUNT_PER_AIC + vectorRoleIdx, job.totalJobs * AIV_CORE_COUNT_PER_AIC};
244	- } else if (!IsSubBlockIdx<1U>()) {	225	+ } else if (MegaMoeHardwarePolicy::GetSubBlockIdx() != 1U) {
245	return {};	226	return {};
246	}	227	}
247	if (combineJob.totalJobs == 0U \|\| combineJob.jobIndex >= combineJob.totalJobs) {	228	if (combineJob.totalJobs == 0U \|\| combineJob.jobIndex >= combineJob.totalJobs) {
248	return {};	229	return {};
249	}	230	}
250	return GetRotatedBalancedWorkRange(tokenCount, combineJob.jobIndex, combineJob.totalJobs, expertRowPrefix);	231	return GetRotatedBalancedWorkRange(tokenCount, combineJob.jobIndex, combineJob.totalJobs, expertRowPrefix);
		232	+#endif
251	}	233	}
252		234
253	// 等待并消费当前 Wave 的 Combine 发送在行环形 buffer 上产生的全部完成事件，使该环不再有在途 slot。	235	// 等待并消费当前 Wave 的 Combine 发送在行环形 buffer 上产生的全部完成事件，使该环不再有在途 slot。
	@@ -609,258 +591,11 @@ __aicore__ inline void Gmm2Aiv0PrologueA8W4(BlockPrologue& blockPrologue, Schedu
609	}	591	}
610		592
611	// 根据 GM 地址建立执行资源，并执行通用 GMM2 阶段。	593	// 根据 GM 地址建立执行资源，并执行通用 GMM2 阶段。
612	-template <uint8_t CombineQuantMode, typename BlockMmad, typename ElementC, bool IsLayered = false,
613	- bool IsShared = false, bool NotifyCombineTileReady = false, typename Scheduler, typename Config>
614	-__aicore__ inline void Gmm2ExecGeneric(Scheduler& scheduler, const GMMAddrInfo& gmmAddrInfo, const Config& config,
615	- uint32_t startLoopIdx, uint32_t tileNum,
616	- BlockMmadContext<BlockMmad>* blockMmadContext, bool allowWeightL2Bypass,
617	- uint32_t rowOffsetInExpert = 0U, const Params* params = nullptr,
618	- uint32_t expertTokenCount = 0U)
619	-{
620	- using KernelConfig = typename Config::KernelConfig;
621	- using ElementA = typename KernelConfig::ElementAType;
622	- using ElementB = typename KernelConfig::ElementBType;
623	- using ElementMxScaleA = typename KernelConfig::ElementMxScaleAType;
624	- using ElementMxScaleB = typename KernelConfig::ElementMxScaleBType;
625	- using BiasType = typename KernelConfig::BiasType;
626	-
627	- auto layouts = KernelConfig::BuildLayouts(config);
628	- auto gmA = asc::te::make_tensor(
629	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementA*>(gmmAddrInfo.aGlobal)),
630	- layouts.a);
631	- auto gmB = asc::te::make_tensor(
632	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementB*>(gmmAddrInfo.bGlobal)),
633	- layouts.b);
634	- auto gmScaleA = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
635	- reinterpret_cast<__gm__ ElementMxScaleA*>(gmmAddrInfo.aScaleGlobal)),
636	- layouts.scaleA);
637	- auto gmScaleB = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
638	- reinterpret_cast<__gm__ ElementMxScaleB*>(gmmAddrInfo.bScaleGlobal)),
639	- layouts.scaleB);
640	- if constexpr (Config::IS_WAVE_FLAG_GRAINED && g_coreType == AscendC::AIC) {
641	- SetWaveWeightL2CacheHint<KernelConfig::IS_WEIGHT_NZ, KernelConfig>(config, allowWeightL2Bypass, gmB, gmScaleB);
642	- }
643	- auto gmBias = asc::te::make_tensor(
644	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ BiasType*>(0UL)), layouts.bias);
645	- auto gmC = asc::te::make_tensor(
646	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(gmmAddrInfo.gmm2OutGlobal)),
647	- layouts.c);
648	-
649	- using WorkSetType = GroupMatmulWorkSet<Scheduler, decltype(gmA), decltype(gmB), decltype(gmScaleA),
650	- decltype(gmScaleB), decltype(gmBias), decltype(gmC)>;
651	- WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmBias, gmC};
652	-
653	- if constexpr (g_coreType == AscendC::AIC) {
654	- if (blockMmadContext != nullptr) {
655	- blockMmadContext->template Init<GmmStage::GMM2>(config);
656	- Gmm2AicMmadGeneric<CombineQuantMode, BlockMmad, IsShared, IsLayered, NotifyCombineTileReady>(
657	- blockMmadContext->blockMmad, workSet, gmmAddrInfo, config, startLoopIdx, tileNum, rowOffsetInExpert,
658	- expertTokenCount);
659	- } else {
660	- BlockMmadContext<BlockMmad> localBlockMmadContext;
661	- localBlockMmadContext.template Init<GmmStage::GMM2>(config);
662	- Gmm2AicMmadGeneric<CombineQuantMode, BlockMmad, IsShared, IsLayered, NotifyCombineTileReady>(
663	- localBlockMmadContext.blockMmad, workSet, gmmAddrInfo, config, startLoopIdx, tileNum, rowOffsetInExpert,
664	- expertTokenCount);
665	- }
666	- } else if constexpr (NotifyCombineTileReady) {
667	- // NotifyCombineTileReady 为真时 AIV1 逐 tile 锁步消费。combine 去重的独立消费路径
668	- // （RunWaveCombineStageDedup）由调用方把本参数实例化为 false，整个分支编译期裁掉。
669	- if (IsSubBlockIdx<1U>()) {
670	- using MakeLayoutC = typename KernelConfig::MakeLayoutC;
671	- CombineTokenRange<ElementC, MakeLayoutC>(workSet.scheduler, workSet.gmC, *params, gmmAddrInfo, config,
672	- startLoopIdx, tileNum);
673	- }
674	- }
675	-}
676		594
677	// 根据 GM 地址建立执行资源，并执行 A8W4 GMM2 阶段。	595	// 根据 GM 地址建立执行资源，并执行 A8W4 GMM2 阶段。
678	-template <typename BlockMmad, typename BlockPrologue, typename ElementC, bool IsShared, bool IsLayered,
679	- bool NotifyCombineTileReady, typename Scheduler, typename Config>
680	-__aicore__ inline void Gmm2ExecA8W4(Scheduler& scheduler, const GMMAddrInfo& gmmAddrInfo, const Config& config,
681	- uint32_t startLoopIdx, uint32_t tileNum, uint32_t expertTokenCount,
682	- uint32_t rowOffsetInExpert, const Params* params,
683	- const A8W4BlockContext<BlockMmad, BlockPrologue>& pipeline)
684	-{
685	- using KernelConfig = typename Config::KernelConfig;
686	- using ElementA = typename KernelConfig::ElementAType;
687	- using ElementB = typename KernelConfig::ElementBType;
688	- using ElementMxScaleA = typename KernelConfig::ElementMxScaleAType;
689	- using ElementMxScaleB = typename KernelConfig::ElementMxScaleBType;
690	- using BiasType = typename KernelConfig::BiasType;
691	-
692	- auto layouts = KernelConfig::BuildLayouts(config);
693	- auto gmC = asc::te::make_tensor(
694	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(gmmAddrInfo.gmm2OutGlobal)),
695	- layouts.c);
696	- auto gmA = asc::te::make_tensor(
697	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementA*>(gmmAddrInfo.aGlobal)),
698	- layouts.a);
699	- auto gmB = asc::te::make_tensor(
700	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementB*>(gmmAddrInfo.bGlobal)),
701	- layouts.b);
702	- auto gmScaleA = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
703	- reinterpret_cast<__gm__ ElementMxScaleA*>(gmmAddrInfo.aScaleGlobal)),
704	- layouts.scaleA);
705	- auto gmScaleB = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
706	- reinterpret_cast<__gm__ ElementMxScaleB*>(gmmAddrInfo.bScaleGlobal)),
707	- layouts.scaleB);
708	- auto gmBias = asc::te::make_tensor(
709	- asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ BiasType*>(0UL)), layouts.bias);
710	-
711	- using WorkSetType = GroupMatmulWorkSet<Scheduler, decltype(gmA), decltype(gmB), decltype(gmScaleA),
712	- decltype(gmScaleB), decltype(gmBias), decltype(gmC)>;
713	- WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmBias, gmC};
714	-
715	- if constexpr (g_coreType == AscendC::AIC) {
716	- auto& blockMmad = *pipeline.block;
717	- typename BlockMmad::BlockShape l0TileShape{config.blockMmadTiling.tileM, config.blockMmadTiling.tileN,
718	- L0_TILE_K, 0};
719	- typename BlockMmad::ProblemShape matmulShape{config.m, config.outputN, config.k, 0};
720	- // GMM2 has no gate/up concatenation; preserve the persistent buffer indices.
721	- blockMmad.template Init<false>(matmulShape, l0TileShape, false);
722	- Gmm2AicMmadA8W4<IsShared, IsLayered, NotifyCombineTileReady, BlockMmad, decltype(workSet.scheduler),
723	- decltype(workSet.gmA), decltype(workSet.gmScaleA), decltype(workSet.gmScaleB),
724	- std::remove_reference_t<decltype(workSet.gmC)>, Config>(
725	- blockMmad, workSet.scheduler, workSet.gmA, workSet.gmScaleA, workSet.gmScaleB, workSet.gmC, gmmAddrInfo,
726	- config, startLoopIdx, tileNum, expertTokenCount, rowOffsetInExpert);
727	- } else {
728	- if (IsSubBlockIdx<0U>()) {
729	- Gmm2Aiv0PrologueA8W4(*pipeline.block, workSet.scheduler, workSet.gmB, config, startLoopIdx, tileNum);
730	- } else if constexpr (NotifyCombineTileReady) {
731	- // 同上：W4 路径去重开关走运行时判断（嵌套模板参数触发编译器代码生成问题）。
732	- if (!IsCombineDedupOn(params->tilingData->compute.dedupMode)) {
733	- using MakeLayoutC = typename KernelConfig::MakeLayoutC;
734	- CombineTokenRange<ElementC, MakeLayoutC>(workSet.scheduler, workSet.gmC, *params, gmmAddrInfo, config,
735	- startLoopIdx, tileNum);
736	- }
737	- }
738	- }
739	-}
740		596
741	} // namespace GmmKernel	597	} // namespace GmmKernel
742		598
743	-// RunGmm2Generic：执行 Generic GMM2，支持量化和非量化 Combine 模式。
744	-// =================================================================================================
745	-template <uint8_t CombineQuantMode, typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA,
746	- typename ElementMxScaleB, bool IsWeightNZ = false, bool IsLayered = false, uint32_t Gmm1TileM = L1_TILE_M_256,
747	- bool TopkWeightsPrefetch = false, bool IsShared = false, bool IsWaveFlagGrained = false,
748	- bool NotifyCombineTileReady = false>
749	-__aicore__ inline void RunGmm2Generic(const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
750	- const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,
751	- const BlockJobContext& blockJob, void* blockMmadContext = nullptr,
752	- bool allowWeightL2Bypass = false, uint32_t rowOffsetInExpert = 0U,
753	- const Params* params = nullptr,
754	- uint32_t activationOutputDivisor = ACTIVATION_N_HALF,
755	- uint32_t expertTokenCount = 0U)
756	-{
757	- static_assert(!IsShared \|\| !NotifyCombineTileReady, "Shared expert GMM2 has its own completion counter");
758	- static_assert(!IsLayered \|\| !NotifyCombineTileReady, "Layered GMM2 selects its synchronization through IsLayered");
759	- using GmmConfig =
760	- GmmKernel::Config<false, CombineQuantMode, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB,
761	- IsWeightNZ, TopkWeightsPrefetch, IsShared, IsLayered, IsWaveFlagGrained>;
762	- auto config = GmmConfig::BuildGmm2ProblemConfig(problemShape, blockJob, Gmm1TileM, activationOutputDivisor);
763	-
764	- GmmKernel::BlockScheduler scheduler({config.m, config.schedulerN, config.k},
765	- GmmKernel::BlockScheduler::Params{asc::te::make_coord(
766	- static_cast<int64_t>(config.tileM), static_cast<int64_t>(L1_TILE_N))});
767	- uint32_t tileNum = scheduler.GetTileNum();
768	-
769	- if constexpr (CombineQuantMode == COMBINE_NO_QUANT && g_coreType == AscendC::AIV && !NotifyCombineTileReady) {
770	- startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
771	- return;
772	- }
773	-
774	- uint32_t startLoopIdx =
775	- (config.blockIdx < startBlockIdx ? config.blockIdx + config.blockNum : config.blockIdx) - startBlockIdx;
776	- using BlockMmad = typename GmmConfig::BlockMmad;
777	- using MmadContext = GmmKernel::BlockMmadContext<BlockMmad>;
778	- auto* typedBlockMmadContext = reinterpret_cast<MmadContext*>(blockMmadContext);
779	- GmmKernel::Gmm2ExecGeneric<CombineQuantMode, BlockMmad, ElementC, IsLayered, IsShared, NotifyCombineTileReady>(
780	- scheduler, gmmAddrInfo, config, startLoopIdx, tileNum, typedBlockMmadContext, allowWeightL2Bypass,
781	- rowOffsetInExpert, params, expertTokenCount);
782	-
783	- startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
784	-}
785	-
786	-template <uint8_t CombineQuantMode, typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA,
787	- typename ElementMxScaleB, bool IsWeightNZ = false, bool IsLayered = false, uint32_t Gmm1TileM = L1_TILE_M_256,
788	- bool TopkWeightsPrefetch = false, bool IsShared = false, bool IsWaveFlagGrained = false,
789	- bool NotifyCombineTileReady = false>
790	-__aicore__ inline void RunGmm2Generic(const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
791	- const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,
792	- void* blockMmadContext = nullptr, bool allowWeightL2Bypass = false,
793	- uint32_t rowOffsetInExpert = 0U, const Params* params = nullptr,
794	- uint32_t activationOutputDivisor = ACTIVATION_N_HALF,
795	- uint32_t expertTokenCount = 0U)
796	-{
797	- BlockJobContext blockJob{MegaMoeHardwarePolicy::GetAicBlockIdx(),
798	- static_cast<uint32_t>(MegaMoeHardwarePolicy::GetBlockNum())};
799	- RunGmm2Generic<CombineQuantMode, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, IsWeightNZ,
800	- IsLayered, Gmm1TileM, TopkWeightsPrefetch, IsShared, IsWaveFlagGrained, NotifyCombineTileReady>(
801	- problemShape, gmmAddrInfo, startBlockIdx, blockJob, blockMmadContext, allowWeightL2Bypass, rowOffsetInExpert,
802	- params, activationOutputDivisor, expertTokenCount);
803	-}
804	-
805	-// RunGmm2A8W4：AIV0执行W4→W8 prologue，AIC执行GMM2；可选由配对AIV1逐tile Combine。
806	-template <typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA, typename ElementMxScaleB,
807	- uint32_t Gmm1TileM = L1_TILE_M_256, bool TopkWeightsPrefetch = false, bool IsShared = false,
808	- bool IsLayered = false, bool IsWaveFlagGrained = false, bool NotifyCombineTileReady = false,
809	- typename BlockContext>
810	-__aicore__ inline void RunGmm2A8W4(const BlockContext& pipeline,
811	- const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
812	- const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,
813	- const BlockJobContext& blockJob, uint32_t expertTokenCount,
814	- uint32_t rowOffsetInExpert, const Params* params = nullptr,
815	- uint32_t activationOutputDivisor = ACTIVATION_N_HALF)
816	-{
817	- static_assert(!IsShared \|\| !NotifyCombineTileReady, "Shared expert GMM2 has its own completion counter");
818	- static_assert(!IsLayered \|\| !NotifyCombineTileReady, "Layered GMM2 selects its synchronization through IsLayered");
819	- static_assert(std::is_same_v<ElementA, __fp8e4m3>, "Activation must be __fp8e4m3");
820	- static_assert(std::is_same_v<ElementB, __fp4e2m1x2>, "Weight must be __fp4e2m1x2");
821	-
822	- using GmmConfig = GmmKernel::Config<true, 0, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, false,
823	- TopkWeightsPrefetch, IsShared, IsLayered, IsWaveFlagGrained>;
824	- auto config = GmmConfig::BuildGmm2ProblemConfig(problemShape, blockJob, Gmm1TileM, activationOutputDivisor);
825	-
826	- using BlockMmad = typename GmmConfig::BlockMmad;
827	- using BlockPrologue = typename GmmConfig::BlockPrologue;
828	- GmmKernel::BlockScheduler scheduler({config.m, config.outputN, config.k},
829	- GmmKernel::BlockScheduler::Params{asc::te::make_coord(
830	- static_cast<int64_t>(config.tileM), static_cast<int64_t>(L1_TILE_N))});
831	- // tileNum 与 mega_moe_utils.h 的 GetGmmProblemTileCount 同源，预判无活的口径不得各算一份。
832	- uint32_t tileNum = GetGmmProblemTileCount(config.m, config.outputN);
833	- uint32_t startLoopIdx =
834	- (config.blockIdx < startBlockIdx ? config.blockIdx + config.blockNum : config.blockIdx) - startBlockIdx;
835	- if (startLoopIdx >= tileNum) {
836	- startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
837	- return;
838	- }
839	-
840	- GmmKernel::Gmm2ExecA8W4<BlockMmad, BlockPrologue, ElementC, IsShared, IsLayered, NotifyCombineTileReady,
841	- GmmKernel::BlockScheduler, decltype(config)>(
842	- scheduler, gmmAddrInfo, config, startLoopIdx, tileNum, expertTokenCount, rowOffsetInExpert, params, pipeline);
843	-
844	- startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
845	-}
846	-
847	-template <typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA, typename ElementMxScaleB,
848	- uint32_t Gmm1TileM = L1_TILE_M_256, bool TopkWeightsPrefetch = false, bool IsShared = false,
849	- bool IsLayered = false, bool IsWaveFlagGrained = false, bool NotifyCombineTileReady = false,
850	- typename BlockContext>
851	-__aicore__ inline void RunGmm2A8W4(const BlockContext& pipeline,
852	- const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
853	- const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,
854	- const Params* params = nullptr, uint32_t activationOutputDivisor = ACTIVATION_N_HALF)
855	-{
856	- BlockJobContext blockJob{MegaMoeHardwarePolicy::GetAicBlockIdx(),
857	- static_cast<uint32_t>(MegaMoeHardwarePolicy::GetBlockNum())};
858	- RunGmm2A8W4<ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, Gmm1TileM, TopkWeightsPrefetch,
859	- IsShared, IsLayered, IsWaveFlagGrained, NotifyCombineTileReady>(
860	- pipeline, problemShape, gmmAddrInfo, startBlockIdx, blockJob, static_cast<uint32_t>(Get<M_VALUE>(problemShape)),
861	- 0U, params, activationOutputDivisor);
862	-}
863	-
864	// 紧凑 token 资源按累计行偏移寻址，专家固定资源按 expertIdx 寻址。	599	// 紧凑 token 资源按累计行偏移寻址，专家固定资源按 expertIdx 寻址。
865	template <typename WeightType, typename ActivationOutType, typename QuantScaleType>	600	template <typename WeightType, typename ActivationOutType, typename QuantScaleType>
866	__aicore__ inline void UpdateMoeExpertGmm2GlobalBuffer(const GmmExecutionConfig& gmmConfig,	601	__aicore__ inline void UpdateMoeExpertGmm2GlobalBuffer(const GmmExecutionConfig& gmmConfig,

Mmc2/mega_moe/op_kernel/arch35_92/stage/mega_moe_token_dedup.h+10-48

	@@ -421,34 +421,6 @@ __aicore__ inline void DedupStoreRowDescChunk(DedupPrepassWorker<TopkIndexType>&
421	}	421	}
422	}	422	}
423		423
424	-#ifndef MEGA_MOE_ARCH92_DEVICE
425	-// 建表工人集按档切换（判据见 DEDUP_AIV0_PREPASS_MAX_HIDDEN_DIM）：AIV0 参与档按
426	-// AIV1=jobIndex / AIV0=jobIndex+totalJobs 平铺；不参与档保持 AIV1 单方建表。
427	-// 返回 false = 本子核不参与建表。
428	-__aicore__ inline bool GetDedupPrepassWorkerSlot(const BlockJobContext& blockJob, const Params& params,
429	- uint32_t& workerIndex, uint32_t& totalWorkers)
430	-{
431	- const bool aiv0Joins = params.tilingData->compute.hiddenDim <= DEDUP_AIV0_PREPASS_MAX_HIDDEN_DIM;
432	- if (!IsSubBlockIdx<1U>() && !aiv0Joins) {
433	- return false;
434	- }
435	- workerIndex = IsSubBlockIdx<1U>() ? blockJob.jobIndex : blockJob.jobIndex + blockJob.totalJobs;
436	- totalWorkers = aiv0Joins ? blockJob.totalJobs * 2U : blockJob.totalJobs;
437	- return true;
438	-}
439	-#endif
440	-
441	-#ifdef MEGA_MOE_ARCH92_DEVICE
442	-// Arch92 单 AIV 配比，直接使用 block 的任务编号与总数。
443	-__aicore__ inline bool GetDedupPrepassWorkerSlot(const BlockJobContext& blockJob, const Params&, uint32_t& workerIndex,
444	- uint32_t& totalWorkers)
445	-{
446	- workerIndex = blockJob.jobIndex;
447	- totalWorkers = blockJob.totalJobs;
448	- return true;
449	-}
450	-#endif
451	-
452	// 计算工作项 chunkIdx 的 token 窗口（末段按 numTokens 截断）。	424	// 计算工作项 chunkIdx 的 token 窗口（末段按 numTokens 截断）。
453	__aicore__ inline DedupChunkWindow MakeDedupChunkWindow(uint32_t chunkIdx, uint32_t chunkTokens, uint32_t numTokens,	425	__aicore__ inline DedupChunkWindow MakeDedupChunkWindow(uint32_t chunkIdx, uint32_t chunkTokens, uint32_t numTokens,
454	uint32_t topK, int32_t& tokenHi)	426	uint32_t topK, int32_t& tokenHi)
	@@ -504,27 +476,20 @@ __aicore__ inline void RunDedupPrepassWorkItem(DedupPrepassWorker<TopkIndexType>
504		476
505	/*	477	/*
506	* 去重 prepass 主入口：扫描 route 槽，为每个接收行写一条 rowDesc（布局见 mega_moe_constants.h）。	478	* 去重 prepass 主入口：扫描 route 槽，为每个接收行写一条 rowDesc（布局见 mega_moe_constants.h）。
507	- * 分工：工作项 = (srcRank, token 段)，全体 AIV1+AIV0 平铺认领（AIV0 的本地前缀表与 UB 布局由	479	+ * 仅由参与建表的 AIV 调用；架构编排提供 worker 编号、总数和完成计数增量，并先准备本地前缀表与 UB。
508	- * PrepareDedupPrepassAiv0Local 先行就位）。同一 (srcRank, token) 的所有出现恒落在同一工作项内，	480	+ * 工作项 = (srcRank, token 段)，按 workerJob 平铺认领。同一 (srcRank, token) 的所有出现恒落在同一
509	- * 工作项之间零交叠，无需原子。	481	+ * 工作项内，工作项之间零交叠，无需原子。
510	* 依赖：count 表已就绪（cumsumInfoTensor；count 到达即 route 槽到达，是 dispatch 沿用的既有契约）。	482	* 依赖：count 表已就绪（cumsumInfoTensor；count 到达即 route 槽到达，是 dispatch 沿用的既有契约）。
511	* 不得把 prepass 整体挪给 AIV0 提前跑：大 bs 下 AIV0 早期密集访问 route 窗口会触发 aicore trap	483	* 不得把 prepass 整体挪给 AIV0 提前跑：大 bs 下 AIV0 早期密集访问 route 窗口会触发 aicore trap
512	* （驱动侧机制未明，实测复现），当前平铺形态是稳定边界。	484	* （驱动侧机制未明，实测复现），当前平铺形态是稳定边界。
513	*/	485	*/
514	template <typename TopkIndexType, typename DispatchScratch>	486	template <typename TopkIndexType, typename DispatchScratch>
515	__aicore__ inline void BuildDedupRowDescTable(const DedupConfig& dedup, const MoeStageCommonConfig& common,	487	__aicore__ inline void BuildDedupRowDescTable(const DedupConfig& dedup, const MoeStageCommonConfig& common,
516	- const BlockJobContext& blockJob, const MoeSyncWorkspaceLayout& syncLayout,	488	+ const BlockJobContext& workerJob, int32_t doneIncrement,
517	- const Params& params, DispatchScratch& scratch)	489	+ const MoeSyncWorkspaceLayout& syncLayout, const Params& params,
		490	+ DispatchScratch& scratch)
518	{	491	{
519	- if constexpr (g_coreType == AIC) {	492	+ if (workerJob.totalJobs == 0U) {
520	- return;
521	- }
522	- if (blockJob.totalJobs == 0U) {
523	- return;
524	- }
525	- uint32_t workerIndex = 0U;
526	- uint32_t totalWorkers = 0U;
527	- if (!GetDedupPrepassWorkerSlot(blockJob, params, workerIndex, totalWorkers)) {
528	return;	493	return;
529	}	494	}
530		495
	@@ -538,18 +503,15 @@ __aicore__ inline void BuildDedupRowDescTable(const DedupConfig& dedup, const Mo
538	PipeNotify<PIPE_MTE3, PIPE_S>(static_cast<TEventID>(SYNC_EVENT_ID3));	503	PipeNotify<PIPE_MTE3, PIPE_S>(static_cast<TEventID>(SYNC_EVENT_ID3));
539	PipeNotify<PIPE_MTE3, PIPE_S>(static_cast<TEventID>(SYNC_EVENT_ID5));	504	PipeNotify<PIPE_MTE3, PIPE_S>(static_cast<TEventID>(SYNC_EVENT_ID5));
540		505
541	- for (uint32_t workItem = workerIndex; workItem < totalWorkItems; workItem += totalWorkers) {	506	+ for (uint32_t workItem = workerJob.jobIndex; workItem < totalWorkItems; workItem += workerJob.totalJobs) {
542	RunDedupPrepassWorkItem(worker, common, params, scratch, dedup, workItem);	507	RunDedupPrepassWorkItem(worker, common, params, scratch, dedup, workItem);
543	}	508	}
544	// 消费掉最后一次武装，不留悬挂 flag 毒化后续同 id 同步。	509	// 消费掉最后一次武装，不留悬挂 flag 毒化后续同 id 同步。
545	PipeWait<PIPE_MTE3, PIPE_S>(static_cast<TEventID>(SYNC_EVENT_ID3));	510	PipeWait<PIPE_MTE3, PIPE_S>(static_cast<TEventID>(SYNC_EVENT_ID3));
546	PipeWait<PIPE_MTE3, PIPE_S>(static_cast<TEventID>(SYNC_EVENT_ID5));	511	PipeWait<PIPE_MTE3, PIPE_S>(static_cast<TEventID>(SYNC_EVENT_ID5));
547		512
548	- // 槽 0 = done 计数，AIV1 计低 16 位、AIV0 计高 16 位。dispatch 发布 gate 只等高位（全体 AIV0）：	513	+ // 架构编排决定完成计数编码；参与核即使未分到工作项，也必须上报一次。
549	- // 每个 AIV1 发布前自己必然已建完，AIV1 互等会打断 prepass 与 dispatch 的流水交叠（实测大形状	514	+ AtomicAdd(readyBase, doneIncrement);
550	- // 明显劣化）；combine 兜底等全值。本核不等任何人。
551	- AtomicAdd(readyBase,
552	- IsSubBlockIdx<1U>() ? static_cast<int32_t>(1) : static_cast<int32_t>(1) << DEDUP_PREPASS_DONE_AIV0_SHIFT);
553	}	515	}
554		516
555	} // namespace MegaMoeImpl	517	} // namespace MegaMoeImpl

Mmc2/mega_moe/op_kernel/arch35_92/stage/mega_moe_token_dispatch.h+0-73

	@@ -480,79 +480,6 @@ __aicore__ inline void LoadAndComputeExpertCountTable(const MoeStageCommonConfig
480	PipeFence<PIPE_V, PIPE_S, SYNC_EVENT_ID2>();	480	PipeFence<PIPE_V, PIPE_S, SYNC_EVENT_ID2>();
481	}	481	}
482		482
483	-// Wave 进入逐专家流水前一次准备完整 count 表；每个物理 block 只发布一次 ready。
484	-// W4 Wave-ahead 路径同时将 prefix 持久化到 GM，避免后续 Activation 覆盖 UB 后丢失 Dispatch 状态。
485	-template <bool NeedCumsumReload = false, typename ActivationType, typename TopkIndexType>
486	-__aicore__ inline void PrepareMoeExpertTokenCountTable(const MoeStageCommonConfig& common,
487	- const BlockWorkspaceContext& countWorkspace,
488	- const Params& params,
489	- TokenDispatchScratch<ActivationType, TopkIndexType>& scratch)
490	-{
491	- const auto& workspace = params.workspace.layout;
492	- const auto& buffer = params.workspace.base;
493	- if constexpr (g_coreType == AIC) {
494	- return;
495	- }
496	- if (!IsSubBlockIdx<1U>()) {
497	- // AIV0 不读取 count，但 mode 0 需要全部 AIV 每轮各上报一次；参与去重建表的 AIV0
498	- // 经 PrepareDedupPrepassAiv0Local -> LoadAndComputeExpertCountTable 已上报，此处不得重复。
499	- // 条件必须与 PrepareDedupPrepassAiv0Local 的启用门严格互补（含 hiddenDim 档位），
500	- // 否则计算主导档的 AIV0 无人上报，ExportAndResetExpertCounts 等 read-done 挂死。
501	- if (params.tilingData->compute.dedupMode == 0 \|\|
502	- params.tilingData->compute.hiddenDim > DEDUP_AIV0_PREPASS_MAX_HIDDEN_DIM) {
503	- SyncInterArrive<PIPE_V>(COUNT_TABLE_READ_DONE_FLAG);
504	- }
505	- return;
506	- }
507	- LoadAndComputeExpertCountTable(common, params, scratch);
508	- uint32_t rawCountElementCount = common.worldSize * common.moeExpertPerRank;
509	- uint64_t countOffset = GetExpertCountWorkspaceOffset(countWorkspace, common.moeExpertPerRank, 0U, true);
510	- PipeFence<PIPE_V, PIPE_MTE3, SYNC_EVENT_ID2>();
511	- CopyUbToGm(scratch.expertRecvCountsGm + countOffset, scratch.expertTokenNumsOutTensor,
512	- common.moeExpertPerRank * static_cast<uint32_t>(sizeof(int32_t)));
513	- if constexpr (NeedCumsumReload) {
514	- CopyUbToGm(scratch.cumsumGm, scratch.cumsumInfoTensor,
515	- rawCountElementCount * static_cast<uint32_t>(sizeof(int32_t)));
516	- }
517	- PipeFence<PIPE_MTE3, PIPE_S, SYNC_EVENT_ID2>();
518	-
519	- __gm__ int32_t* countTableReady = workspace.sync.countsReady.get_ptr<int32_t>(buffer) +
520	- static_cast<uint64_t>(countWorkspace.blockIdx) * INT_CACHELINE;
521	- GmWrite(countTableReady, static_cast<int32_t>(1));
522	-}
523	-
524	-/*
525	- * W4 的 Activation/SwiGLU 会覆盖 AIV1 UB 0 开始的 prefix 表。调用方在执行 DispatchTokenRange 前，
526	- * 从当前物理 block 的 GM 备份恢复“首专家前缀 + 本次范围覆盖的专家前缀”。
527	- */
528	-template <typename ActivationType, typename TopkIndexType>
529	-__aicore__ inline void ReloadDispatchCumsumRange(const MoeStageCommonConfig& common,
530	- TokenDispatchScratch<ActivationType, TopkIndexType>& scratch,
531	- uint32_t firstExpertIdx, uint32_t lastExpertIdx)
532	-{
533	- if constexpr (g_coreType == AIC) {
534	- return;
535	- }
536	- if (!IsSubBlockIdx<1U>()) {
537	- return;
538	- }
539	-
540	- if (firstExpertIdx >= common.moeExpertPerRank \|\| firstExpertIdx > lastExpertIdx) {
541	- return;
542	- }
543	- if (lastExpertIdx >= common.moeExpertPerRank) {
544	- lastExpertIdx = common.moeExpertPerRank - 1U;
545	- }
546	- constexpr uint32_t INT32_PER_32B = ALIGN_32 / sizeof(int32_t);
547	- uint32_t requiredBeginIndex = firstExpertIdx == 0U ? 0U : firstExpertIdx * common.worldSize - 1U;
548	- uint32_t alignedBeginIndex = requiredBeginIndex / INT32_PER_32B * INT32_PER_32B;
549	- uint32_t requiredEndIndex = (lastExpertIdx + 1U) * common.worldSize;
550	- uint32_t reloadElementCount = requiredEndIndex - alignedBeginIndex;
551	- CopyGmToUb((scratch.cumsumInfoTensor + (alignedBeginIndex)), scratch.cumsumGm + alignedBeginIndex,
552	- reloadElementCount * static_cast<uint32_t>(sizeof(int32_t)), GmLoadCacheMode::NORMAL_FIRST_VICTIM, true);
553	- PipeFence<PIPE_MTE2, PIPE_S, SYNC_EVENT_ID2>();
554	-}
555	-
556	} // namespace MegaMoeImpl	483	} // namespace MegaMoeImpl
557		484
558	#endif // MEGA_MOE_TOKEN_DISPATCH_H	485	#endif // MEGA_MOE_TOKEN_DISPATCH_H

Mmc2/mega_moe/op_kernel/arch35_92/stage/mega_moe_token_dispatch_dedup.h+1-1

	@@ -487,7 +487,7 @@ __aicore__ inline void DispatchTokenRange(const TokenDispatchConfig& context, co
487	if constexpr (g_coreType == AIC) {	487	if constexpr (g_coreType == AIC) {
488	return;	488	return;
489	}	489	}
490	- if (!IsSubBlockIdx<1U>() \|\| blockJob.totalJobs == 0U) {	490	+ if (blockJob.totalJobs == 0U) {
491	return;	491	return;
492	}	492	}
493	WaitDedupPrepassPublishGate(context, params);	493	WaitDedupPrepassPublishGate(context, params);

Mmc2/mega_moe/op_kernel/arch92/mega_moe_arch92.h+7-111

	@@ -30,7 +30,8 @@
30	#include "../arch35_92/stage/mega_moe_send_mask.h"	30	#include "../arch35_92/stage/mega_moe_send_mask.h"
31	#include "../arch35_92/stage/mega_moe_workspace_reset.h"	31	#include "../arch35_92/stage/mega_moe_workspace_reset.h"
32	#include "../arch35_92/stage/mega_moe_token_dispatch_dedup.h"	32	#include "../arch35_92/stage/mega_moe_token_dispatch_dedup.h"
33	-#include "../arch35_92/stage/mega_moe_gmm1_activation.h"	33	+#include "stage/mega_moe_gmm1_activation_arch92.h"
		34	+#include "stage/mega_moe_gmm2_combine_arch92.h"
34	#include "../arch35_92/stage/mega_moe_gmm2_combine_dedup.h"	35	#include "../arch35_92/stage/mega_moe_gmm2_combine_dedup.h"
35	#include "../arch35_92/stage/mega_moe_unpermute.h"	36	#include "../arch35_92/stage/mega_moe_unpermute.h"
36	#include "../../../common/op_kernel/quantize_functions.h"	37	#include "../../../common/op_kernel/quantize_functions.h"
	@@ -108,33 +109,25 @@ protected:
108	template <typename Derived>	109	template <typename Derived>
109	__aicore__ inline void SyncBeforeMoeExpertStages(Derived& derived, SharedGmmStartCondition sharedGmmStartCondition);	110	__aicore__ inline void SyncBeforeMoeExpertStages(Derived& derived, SharedGmmStartCondition sharedGmmStartCondition);
110	__aicore__ inline void ExportAndResetExpertCounts();	111	__aicore__ inline void ExportAndResetExpertCounts();
111	- __aicore__ inline void RunGmm2CombineForExpert(ExpertLoopState& state, GMMAddrInfo& gmmAddrInfo,
112	- uint32_t& startBlockIdx, uint32_t tokenStartIndexInExpert,
113	- uint32_t sliceTokenCount,
114	- WaveCombineBufferConfig& combineBufferConfig,
115	- uint32_t& combineRowSequence, bool isFinalCombine,
116	- const A8W4BlockContext& pipeline);
117	- template <bool WaitForTokenCountReady>
118	- __aicore__ inline void PrepareGmmExpertState(ExpertLoopState& state, uint32_t expertIdx);
119	__aicore__ inline void ConfigureSharedGmm1Input(GMMAddrInfo& gmmAddrInfo, GmmExecutionConfig& gmmConfig,	112	__aicore__ inline void ConfigureSharedGmm1Input(GMMAddrInfo& gmmAddrInfo, GmmExecutionConfig& gmmConfig,
120	Gmm1ActivationSync& sharedGmm1ActivationSync) const;	113	Gmm1ActivationSync& sharedGmm1ActivationSync) const;
121	__aicore__ inline void RunSharedExpertGmm1Activation(const GMMAddrInfo& gmmAddrInfo,	114	__aicore__ inline void RunSharedExpertGmm1Activation(const GMMAddrInfo& gmmAddrInfo,
122	const ProblemShape& problemShape,	115	const ProblemShape& problemShape,
123	const GmmExecutionConfig& gmmConfig,	116	const GmmExecutionConfig& gmmConfig,
124	GmmRuntimeState& runtimeState, uint32_t sharedExpertIdx,	117	GmmRuntimeState& runtimeState, uint32_t sharedExpertIdx,
125	- void* blockMmadContext = nullptr);	118	+ void* blockMmadContext);
126	__aicore__ inline void RunSharedExpertGmm2(const GMMAddrInfo& gmmAddrInfo, const ProblemShape& problemShape,	119	__aicore__ inline void RunSharedExpertGmm2(const GMMAddrInfo& gmmAddrInfo, const ProblemShape& problemShape,
127	- void* blockMmadContext = nullptr);	120	+ void* blockMmadContext);
128	__aicore__ inline void ProcessSharedExpertGmm1Loop(GMMAddrInfo& gmmAddrInfo, const ProblemShape& problemShape,	121	__aicore__ inline void ProcessSharedExpertGmm1Loop(GMMAddrInfo& gmmAddrInfo, const ProblemShape& problemShape,
129	const GmmExecutionConfig& sharedGmmConfig,	122	const GmmExecutionConfig& sharedGmmConfig,
130	GmmRuntimeState& runtimeState, void* blockMmadContext);	123	GmmRuntimeState& runtimeState, void* blockMmadContext);
131	__aicore__ inline void ProcessSharedExpertGmm1(Gmm1ActivationSync& sharedGmm1ActivationSync,	124	__aicore__ inline void ProcessSharedExpertGmm1(Gmm1ActivationSync& sharedGmm1ActivationSync,
132	- void* blockMmadContext = nullptr);	125	+ void* blockMmadContext);
133	__aicore__ inline void ProcessSharedExpertGmm2Loop(GMMAddrInfo& gmmAddrInfo, const ProblemShape& problemShape,	126	__aicore__ inline void ProcessSharedExpertGmm2Loop(GMMAddrInfo& gmmAddrInfo, const ProblemShape& problemShape,
134	void* blockMmadContext);	127	void* blockMmadContext);
135	- __aicore__ inline void ProcessSharedExpertGmm2(void* blockMmadContext = nullptr);	128	+ __aicore__ inline void ProcessSharedExpertGmm2(void* blockMmadContext);
136	template <typename Derived>	129	template <typename Derived>
137	- __aicore__ inline void ProcessWave(Derived& derived, void* blockMmadContext = nullptr);	130	+ __aicore__ inline void ProcessWave(Derived& derived, void* blockMmadContext);
138		131
139	__gm__ Mc2MoeContext* mc2Context_{nullptr};	132	__gm__ Mc2MoeContext* mc2Context_{nullptr};
140	Params params_{};	133	Params params_{};
	@@ -810,103 +803,6 @@ __aicore__ inline void MegaMoe<TemplateMegaMoeTypeFunc>::ProcessSharedExpertGmm2
810	}	803	}
811	}	804	}
812		805
813	-// W4 基类统一执行原 A8W4/A4W4 的专家 GMM2/Combine；派生模板只提供当前专家 slice。
814	-// 非量化路径由 AIV1 在 GMM2 scheduler 内逐 tile 调用统一的 CombineTokenRange。
815	-template <TemplateMegaMoeTypeClass>
816	-__aicore__ inline void MegaMoe<TemplateMegaMoeTypeFunc>::RunGmm2CombineForExpert(
817	- ExpertLoopState& state, GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx, uint32_t tokenStartIndexInExpert,
818	- uint32_t sliceTokenCount, WaveCombineBufferConfig& combineBufferConfig, uint32_t& combineRowSequence,
819	- bool isFinalCombine, const A8W4BlockContext& pipeline)
820	-{
821	- const auto& workspace = params_.workspace.layout;
822	- const auto& buffer = params_.workspace.base;
823	- uint32_t expertTokenCount = static_cast<uint32_t>(Get<M_VALUE>(state.problemShape));
824	- uint32_t gmm2NTileCount = Ops::Base::CeilDiv(commonConfig_.tokenHiddenDim, static_cast<uint32_t>(L1_TILE_N));
825	- uint32_t problemTileCount = GetMGroupCountForRows(sliceTokenCount, GMM1_TILE_M) * gmm2NTileCount;
826	- if (!HandleWaveProblemWithoutWork(problemTileCount, gmmExecutionConfig_.blockJob, startBlockIdx)) {
827	- UpdateMoeExpertGmm2GlobalBuffer<MoeWeightType, ActivationQuantOutType, QuantScaleOutType>(
828	- gmmExecutionConfig_, syncWorkspaceLayout_, workspace, buffer, moeWeightTensorListAddrs_, gmmAddrInfo, state,
829	- commonConfig_.gmm1OutputDim / commonConfig_.activationOutputDivisor, tokenStartIndexInExpert);
830	- ProblemShape sliceProblemShape = state.problemShape;
831	- Get<M_VALUE>(sliceProblemShape) = sliceTokenCount;
832	- RunGmm2A8W4<ActivationQuantOutType, MoeWeightType, bfloat16_t, QuantScaleOutType, QuantScaleOutType,
833	- GMM1_TILE_M, TopkWeightsPrefetch, false, false, true, CombineQuantMode == COMBINE_NO_QUANT>(
834	- pipeline, sliceProblemShape, gmmAddrInfo, startBlockIdx, gmmExecutionConfig_.blockJob, expertTokenCount,
835	- tokenStartIndexInExpert, &params_, commonConfig_.activationOutputDivisor);
836	- }
837	-
838	- if constexpr (CombineQuantMode == COMBINE_NO_QUANT) {
839	- if (IsCombineDedupOn(static_cast<int32_t>(params_.tilingData->compute.dedupMode))) {
840	- // combine 去重：跨 WAVE 的专家等最后一个 slice 完成后按专家粒度独立消费（与量化路径
841	- // 同构；AIC 被跳过 problem 的情形也会到达本分支，per-expert 标记不漏发）。
842	- if (tokenStartIndexInExpert + sliceTokenCount >= expertTokenCount) {
843	- NotifyWaveGmm2Ready(waveCombineJob_, params_, state.expertIdx);
844	- if (isFinalCombine) {
845	- if constexpr (g_coreType == AIV) {
846	- if (GetSubBlockIdx() == 0U) {
847	- // W4 prologue 完成后，末轮 AIV0 才能通过 MTE2 复用同一 UB 区域。
848	- PipeFence<PIPE_MTE3, PIPE_MTE2, SYNC_EVENT_ID0>();
849	- }
850	- }
851	- combineBufferConfig =
852	- InitWaveCombineBuffersDedup(commonConfig_, waveCombineScratch_, waveCombineDedupScratch_);
853	- RunWaveCombineStageDedup<true>(commonConfig_, waveCombineJob_, combineBufferConfig,
854	- waveCombineScratch_, waveCombineDedupScratch_, params_, state,
855	- state.expertIdx, combineRowSequence);
856	- } else {
857	- RunWaveCombineStageDedup<false>(commonConfig_, waveCombineJob_, combineBufferConfig,
858	- waveCombineScratch_, waveCombineDedupScratch_, params_, state,
859	- state.expertIdx, combineRowSequence);
860	- }
861	- }
862	- return;
863	- }
864	- // 基线路径：AIV1 已在上述 GMM2 调用中逐 tile 消费当前 slice；无需在外部重建 scheduler。
865	- return;
866	- }
867	-
868	- // 量化 Combine 保持原 W4 语义：跨 WAVE 的专家必须等最后一个 slice 完成后再消费。
869	- if (tokenStartIndexInExpert + sliceTokenCount < expertTokenCount) {
870	- return;
871	- }
872	-
873	- NotifyWaveGmm2Ready(waveCombineJob_, params_, state.expertIdx);
874	- UpdateMoeExpertCombineGlobalBuffer(workspace, buffer, gmmAddrInfo, state);
875	- if (isFinalCombine) {
876	- if constexpr (g_coreType == AIV) {
877	- // W4 prologue 完成后，末轮 AIV0 才能通过 MTE2 复用同一 UB 区域。
878	- PipeFence<PIPE_MTE3, PIPE_MTE2, SYNC_EVENT_ID0>();
879	- }
880	- combineBufferConfig =
881	- PrepareFinalWaveCombineBuffers<CombineQuantMode>(commonConfig_, combineBufferConfig, waveCombineScratch_);
882	- RunWaveCombineStage<CombineQuantMode, true>(commonConfig_, waveCombineJob_, combineBufferConfig,
883	- waveCombineScratch_, params_, gmmAddrInfo, state, state.expertIdx,
884	- combineRowSequence);
885	- return;
886	- }
887	- RunWaveCombineStage<CombineQuantMode>(commonConfig_, waveCombineJob_, combineBufferConfig, waveCombineScratch_,
888	- params_, gmmAddrInfo, state, state.expertIdx, combineRowSequence);
889	-}
890	-
891	-// 从 workspace 读取 token 数并准备 GMM 专家状态；GMM1 可按需先等待 token count ready。
892	-template <TemplateMegaMoeTypeClass>
893	-template <bool WaitForTokenCountReady>
894	-__aicore__ inline void MegaMoe<TemplateMegaMoeTypeFunc>::PrepareGmmExpertState(ExpertLoopState& state,
895	- uint32_t expertIdx)
896	-{
897	- const auto& workspace = params_.workspace.layout;
898	- const auto& buffer = params_.workspace.base;
899	- if constexpr (WaitForTokenCountReady) {
900	- if (!state.expertCountTableReady) {
901	- WaitForMoeExpertTokenCountReady(workspace.sync.countsReady.get_ptr(buffer), countWorkspace_, 0U);
902	- state.expertCountTableReady = true;
903	- }
904	- }
905	- uint32_t expertTokenCount = GetExpertTokenCountFromWorkspace(
906	- workspace.routing.recvCounts.get_ptr(buffer), countWorkspace_, commonConfig_.moeExpertPerRank, expertIdx);
907	- UpdateExpertLoopState(state, expertIdx, expertTokenCount);
908	-}
909	-
910	// 只有量化类型不同且启用了共享专家，才传入独立共享量化参数。	806	// 只有量化类型不同且启用了共享专家，才传入独立共享量化参数。
911	template <TemplateMegaMoeTypeClass>	807	template <TemplateMegaMoeTypeClass>
912	__aicore__ inline void MegaMoe<TemplateMegaMoeTypeFunc>::QuantizeInputTokens(const WorkRange& tokenRange)	808	__aicore__ inline void MegaMoe<TemplateMegaMoeTypeFunc>::QuantizeInputTokens(const WorkRange& tokenRange)

Mmc2/mega_moe/op_kernel/arch92/mega_moe_wave.h+23-15

	@@ -87,8 +87,12 @@ private:
87	std::conditional_t<g_coreType == AIC, GmmKernel::BlockMmadContext<BlockMmad>, EmptyMmadContext>;	87	std::conditional_t<g_coreType == AIC, GmmKernel::BlockMmadContext<BlockMmad>, EmptyMmadContext>;
88	static constexpr bool ROUTED_AND_SHARED_USE_SAME_MMAD = std::is_same_v<Gmm1BlockMmad, SharedGmm1BlockMmad>;	88	static constexpr bool ROUTED_AND_SHARED_USE_SAME_MMAD = std::is_same_v<Gmm1BlockMmad, SharedGmm1BlockMmad>;
89	static constexpr bool SHARED_GMM_STAGES_USE_SAME_MMAD = std::is_same_v<SharedGmm1BlockMmad, SharedGmm2BlockMmad>;	89	static constexpr bool SHARED_GMM_STAGES_USE_SAME_MMAD = std::is_same_v<SharedGmm1BlockMmad, SharedGmm2BlockMmad>;
		90	+ static_assert(std::is_same_v<Gmm1BlockMmad, Gmm2BlockMmad>,
		91	+ "Arch92 routed GMM1/GMM2 must share the persistent MMAD context type");
		92	+ static_assert(SHARED_GMM_STAGES_USE_SAME_MMAD,
		93	+ "Arch92 shared GMM1/GMM2 must share the persistent MMAD context type");
90	using SharedMmadContextStorage =	94	using SharedMmadContextStorage =
91	- std::conditional_t<g_coreType == AIC && SHARED_GMM_STAGES_USE_SAME_MMAD && !ROUTED_AND_SHARED_USE_SAME_MMAD,	95	+ std::conditional_t<g_coreType == AIC && !ROUTED_AND_SHARED_USE_SAME_MMAD,
92	GmmKernel::BlockMmadContext<SharedGmm1BlockMmad>, EmptyMmadContext>;	96	GmmKernel::BlockMmadContext<SharedGmm1BlockMmad>, EmptyMmadContext>;
93		97
94	__aicore__ inline void* GetRoutedMmadContext()	98	__aicore__ inline void* GetRoutedMmadContext()
	@@ -101,15 +105,12 @@ private:
101		105
102	__aicore__ inline void* GetRoutedGmm2MmadContext()	106	__aicore__ inline void* GetRoutedGmm2MmadContext()
103	{	107	{
104	- if constexpr (std::is_same_v<Gmm1BlockMmad, Gmm2BlockMmad>) {	108	+ return GetRoutedMmadContext();
105	- return GetRoutedMmadContext();
106	- }
107	- return nullptr;
108	}	109	}
109		110
110	__aicore__ inline void* GetSharedMmadContext()	111	__aicore__ inline void* GetSharedMmadContext()
111	{	112	{
112	- if constexpr (g_coreType == AIC && SHARED_GMM_STAGES_USE_SAME_MMAD) {	113	+ if constexpr (g_coreType == AIC) {
113	if constexpr (ROUTED_AND_SHARED_USE_SAME_MMAD) {	114	if constexpr (ROUTED_AND_SHARED_USE_SAME_MMAD) {
114	return static_cast<void*>(&routedMmadContext_);	115	return static_cast<void*>(&routedMmadContext_);
115	}	116	}
	@@ -201,8 +202,11 @@ template <TemplateMegaMoeArch92WaveTypeClass>
201	__aicore__ inline typename MegaMoeArch92Wave<TemplateMegaMoeArch92WaveTypeFunc>::CombineBufferConfig	202	__aicore__ inline typename MegaMoeArch92Wave<TemplateMegaMoeArch92WaveTypeFunc>::CombineBufferConfig
202	MegaMoeArch92Wave<TemplateMegaMoeArch92WaveTypeFunc>::InitCombineBuffers()	203	MegaMoeArch92Wave<TemplateMegaMoeArch92WaveTypeFunc>::InitCombineBuffers()
203	{	204	{
204	- return InitWaveCombineBuffers<CombineQuantMode, false, WAVE_COMBINE_STEADY_ROW_BUFFER_COUNT>(commonConfig_,	205	+ if constexpr (g_coreType == AIV) {
205	- waveCombineScratch_);	206	+ return InitWaveCombineBuffers<CombineQuantMode, WAVE_COMBINE_STEADY_ROW_BUFFER_COUNT>(commonConfig_,
		207	+ waveCombineScratch_);
		208	+ }
		209	+ return {};
206	}	210	}
207		211
208	template <TemplateMegaMoeArch92WaveTypeClass>	212	template <TemplateMegaMoeArch92WaveTypeClass>
	@@ -358,12 +362,13 @@ MegaMoeArch92Wave<TemplateMegaMoeArch92WaveTypeFunc>::ProcessNextDispatchAndCurr
358	gmm1Position.tokenIndexInExpert == 0U && static_cast<uint64_t>(waveEndTokenIndexInExpert) == expertRowCount;	362	gmm1Position.tokenIndexInExpert == 0U && static_cast<uint64_t>(waveEndTokenIndexInExpert) == expertRowCount;
359	uint32_t waveTokenStartIndex =	363	uint32_t waveTokenStartIndex =
360	static_cast<uint32_t>(gmm1ExpertState.globalTokenStartIndex) + gmm1Position.tokenIndexInExpert;	364	static_cast<uint32_t>(gmm1ExpertState.globalTokenStartIndex) + gmm1Position.tokenIndexInExpert;
		365	+ int32_t vecSetSyncCom = 0;
361	void* blockMmadContext = GetRoutedMmadContext();	366	void* blockMmadContext = GetRoutedMmadContext();
362	- RunGmm1ToGm<QuantOutType, ActivationType, MoeWeightType, bfloat16_t, QuantScaleOutType, QuantScaleOutType,	367	+ RunGmm1Generic<QuantOutType, ActivationType, MoeWeightType, bfloat16_t, QuantScaleOutType, QuantScaleOutType,
363	- MoeWeight1Format != FORMAT_ND, GMM1_TILE_M, EPILOGUE_TILE_M, TopkWeightsPrefetch, false, true>(	368	+ MoeWeight1Format != FORMAT_ND, GMM1_TILE_M, EPILOGUE_TILE_M, TopkWeightsPrefetch, false, true>(
364	- epilogueOp_, params_, gmm1WaveProblemShape, gmm1AddrInfo, startBlockIdx_, gmmExecutionConfig_.blockJob,	369	+ epilogueOp_, params_, gmm1WaveProblemShape, gmm1AddrInfo, startBlockIdx_, vecSetSyncCom,
365	- waveTokenStartIndex, gmm1Position.expertIdx, gmm1PingPongIdx_, commonConfig_.activationOutputDivisor,	370	+ gmmExecutionConfig_.blockJob, waveTokenStartIndex, gmm1Position.expertIdx, gmm1PingPongIdx_,
366	- blockMmadContext, isWholeExpert);	371	+ commonConfig_.activationOutputDivisor, blockMmadContext, isWholeExpert);
367		372
368	processedMGroupCount += problemMGroupCount;	373	processedMGroupCount += problemMGroupCount;
369	gmm1Position.tokenIndexInExpert = waveEndTokenIndexInExpert;	374	gmm1Position.tokenIndexInExpert = waveEndTokenIndexInExpert;
	@@ -570,8 +575,11 @@ __aicore__ inline void MegaMoeArch92Wave<TemplateMegaMoeArch92WaveTypeFunc>::Pro
570	PrepareMoeExpertTokenCountTable();	575	PrepareMoeExpertTokenCountTable();
571	if (tokenDispatchConfig_.dedup.dedupMode != 0) {	576	if (tokenDispatchConfig_.dedup.dedupMode != 0) {
572	// Arch92 没有独立的 AIV0 工作核，也不使用完成计数的高 16 位。	577	// Arch92 没有独立的 AIV0 工作核，也不使用完成计数的高 16 位。
573	- BuildDedupRowDescTable<TopkIndexType>(tokenDispatchConfig_.dedup, commonConfig_, gmmExecutionConfig_.blockJob,	578	+ if constexpr (g_coreType == AIV) {
574	- syncWorkspaceLayout_, params_, tokenDispatchScratch_);	579	+ BuildDedupRowDescTable<TopkIndexType>(tokenDispatchConfig_.dedup, commonConfig_,
		580	+ gmmExecutionConfig_.blockJob, static_cast<int32_t>(1),
		581	+ syncWorkspaceLayout_, params_, tokenDispatchScratch_);
		582	+ }
575	}	583	}
576		584
577	GMMAddrInfo gmm1AddrInfo{};	585	GMMAddrInfo gmm1AddrInfo{};

Mmc2/mega_moe/op_kernel/arch92/stage/mega_moe_gmm1_activation_arch92.h+120-22

	@@ -16,15 +16,15 @@
16	namespace MegaMoeImpl {	16	namespace MegaMoeImpl {
17	namespace GmmKernel {	17	namespace GmmKernel {
18		18
19	-// Arch92 routed GMM1 固定写 GM，Activation 使用既有预取或顺序消费协议。
20	template <typename BlockMmad, typename ElementC, typename MakeLayoutC, bool TopkWeightsPrefetch, bool IsShared = false,	19	template <typename BlockMmad, typename ElementC, typename MakeLayoutC, bool TopkWeightsPrefetch, bool IsShared = false,
21	bool IsWaveFlagGrained = false, typename Scheduler, typename Config, typename ActivationQuantOp,	20	bool IsWaveFlagGrained = false, typename Scheduler, typename Config, typename ActivationQuantOp,
22	typename LayoutBundle>	21	typename LayoutBundle>
23	-__aicore__ inline void Gmm1ExecToGm(Scheduler& scheduler, const Params& params, const GMMAddrInfo& gmmAddrInfo,	22	+__aicore__ inline void Gmm1ExecGeneric(Scheduler& scheduler, const Params& params, const GMMAddrInfo& gmmAddrInfo,
24	- const Config& config, uint32_t startLoopIdx, uint32_t tileNum,	23	+ const Config& config, uint32_t startLoopIdx, uint32_t tileNum,
25	- ActivationQuantOp& activationQuantOp, uint32_t expertBeforeCnt, uint32_t expertIdx,	24	+ ActivationQuantOp& activationQuantOp, int32_t& vecSetSyncCom,
26	- uint16_t& pingpongIdx, BlockMmadContext<BlockMmad>* blockMmadContext,	25	+ uint32_t expertBeforeCnt, uint32_t expertIdx, uint16_t& pingpongIdx,
27	- bool allowWeightL2Bypass, const LayoutBundle& layouts)	26	+ BlockMmadContext<BlockMmad>* blockMmadContext, bool allowWeightL2Bypass,
		27	+ const LayoutBundle& layouts)
28	{	28	{
29	using KernelConfig = typename Config::KernelConfig;	29	using KernelConfig = typename Config::KernelConfig;
30	using ElementA = typename KernelConfig::ElementAType;	30	using ElementA = typename KernelConfig::ElementAType;
	@@ -46,7 +46,12 @@ __aicore__ inline void Gmm1ExecToGm(Scheduler& scheduler, const Params& params,
46	if constexpr (IsWaveFlagGrained && g_coreType == AscendC::AIC) {	46	if constexpr (IsWaveFlagGrained && g_coreType == AscendC::AIC) {
47	SetWaveWeightL2CacheHint<KernelConfig::IS_WEIGHT_NZ, KernelConfig>(config, allowWeightL2Bypass, gmB, gmScaleB);	47	SetWaveWeightL2CacheHint<KernelConfig::IS_WEIGHT_NZ, KernelConfig>(config, allowWeightL2Bypass, gmB, gmScaleB);
48	}	48	}
49	- GM_ADDR cGlobal = gmmAddrInfo.gmm1OutGlobal;	49	+ // Shared 保留 UB 交接；routed 即使关闭 prefetch 也必须写 GM。
		50	+ constexpr bool useGmOutput = !IsShared \|\| TopkWeightsPrefetch;
		51	+ GM_ADDR cGlobal = 0UL;
		52	+ if constexpr (useGmOutput) {
		53	+ cGlobal = gmmAddrInfo.gmm1OutGlobal;
		54	+ }
50	auto gmC = asc::te::make_tensor(	55	auto gmC = asc::te::make_tensor(
51	asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(cGlobal)), layouts.c);	56	asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(cGlobal)), layouts.c);
52	__gm__ float* metaInfoGm = nullptr;	57	__gm__ float* metaInfoGm = nullptr;
	@@ -58,19 +63,21 @@ __aicore__ inline void Gmm1ExecToGm(Scheduler& scheduler, const Params& params,
58	decltype(gmC), decltype(metaInfoGm)>;	63	decltype(gmC), decltype(metaInfoGm)>;
59	WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmC, metaInfoGm};	64	WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmC, metaInfoGm};
60	if constexpr (g_coreType == AscendC::AIC) {	65	if constexpr (g_coreType == AscendC::AIC) {
61	- if (blockMmadContext != nullptr) {	66	+ if constexpr (useGmOutput) {
62	Gmm1AicMmadToGmGeneric<BlockMmadContext<BlockMmad>, IsShared, TopkWeightsPrefetch>(	67	Gmm1AicMmadToGmGeneric<BlockMmadContext<BlockMmad>, IsShared, TopkWeightsPrefetch>(
63	workSet, params, gmmAddrInfo, config, startLoopIdx, tileNum, expertIdx, blockMmadContext);	68	workSet, params, gmmAddrInfo, config, startLoopIdx, tileNum, expertIdx, blockMmadContext);
64	} else {	69	} else {
65	- BlockMmadContext<BlockMmad> localBlockMmadContext;	70	+ Gmm1AicMmadGeneric<BlockMmadContext<BlockMmad>, IsShared>(
66	- Gmm1AicMmadToGmGeneric<BlockMmadContext<BlockMmad>, IsShared, TopkWeightsPrefetch>(	71	+ workSet, gmmAddrInfo, config, startLoopIdx, tileNum, vecSetSyncCom, pingpongIdx, blockMmadContext);
67	- workSet, params, gmmAddrInfo, config, startLoopIdx, tileNum, expertIdx, &localBlockMmadContext);
68	}	72	}
69	} else {	73	} else {
70	if constexpr (TopkWeightsPrefetch) {	74	if constexpr (TopkWeightsPrefetch) {
71	Gmm1Aiv0PrefetchEpilogueGeneric<MakeLayoutC, ElementC, IsWaveFlagGrained>(	75	Gmm1Aiv0PrefetchEpilogueGeneric<MakeLayoutC, ElementC, IsWaveFlagGrained>(
72	workSet, params, gmmAddrInfo, config, startLoopIdx, tileNum, activationQuantOp, expertBeforeCnt,	76	workSet, params, gmmAddrInfo, config, startLoopIdx, tileNum, activationQuantOp, expertBeforeCnt,
73	expertIdx);	77	expertIdx);
		78	+ } else if constexpr (IsShared) {
		79	+ Gmm1Aiv0EpilogueGeneric<IsWaveFlagGrained>(workSet, gmmAddrInfo, config, startLoopIdx, tileNum,
		80	+ activationQuantOp, pingpongIdx);
74	} else {	81	} else {
75	auto& sync = *gmmAddrInfo.gmm1ActivationSync;	82	auto& sync = *gmmAddrInfo.gmm1ActivationSync;
76	for (uint32_t loopIdx = startLoopIdx; loopIdx < tileNum; loopIdx += config.blockNum) {	83	for (uint32_t loopIdx = startLoopIdx; loopIdx < tileNum; loopIdx += config.blockNum) {
	@@ -87,20 +94,65 @@ __aicore__ inline void Gmm1ExecToGm(Scheduler& scheduler, const Params& params,
87	}	94	}
88	}	95	}
89		96
		97	+// 保留既有未开放分支的模板引用，不扩展 Arch92 能力范围。
		98	+template <typename BlockMmad, typename BlockPrologue, typename ElementC, typename MakeLayoutC, bool IsShared,
		99	+ typename Scheduler, typename Config, bool TopkWeightsPrefetch, bool IsWaveFlagGrained,
		100	+ typename ActivationQuantOp, typename LayoutBundle>
		101	+__aicore__ inline void Gmm1ExecA8W4(Scheduler& scheduler, const Params& params, const GMMAddrInfo& gmmAddrInfo,
		102	+ const Config& config, uint32_t startLoopIdx, uint32_t tileNum,
		103	+ ActivationQuantOp& activationQuantOp, uint32_t expertBeforeCnt, uint32_t expertIdx,
		104	+ const LayoutBundle& layouts,
		105	+ const A8W4BlockContext<BlockMmad, BlockPrologue>& pipeline)
		106	+{
		107	+ using KernelConfig = typename Config::KernelConfig;
		108	+ using ElementA = typename KernelConfig::ElementAType;
		109	+ using ElementB = typename KernelConfig::ElementBType;
		110	+ using ElementMxScaleA = typename KernelConfig::ElementMxScaleAType;
		111	+ using ElementMxScaleB = typename KernelConfig::ElementMxScaleBType;
		112	+ auto gmC = asc::te::make_tensor(
		113	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(gmmAddrInfo.gmm1OutGlobal)),
		114	+ layouts.c);
		115	+ auto gmA = asc::te::make_tensor(
		116	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementA*>(gmmAddrInfo.aGlobal)),
		117	+ layouts.a);
		118	+ auto gmB = asc::te::make_tensor(
		119	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementB*>(gmmAddrInfo.bGlobal)),
		120	+ layouts.b);
		121	+ auto gmScaleA = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
		122	+ reinterpret_cast<__gm__ ElementMxScaleA*>(gmmAddrInfo.aScaleGlobal)),
		123	+ layouts.scaleA);
		124	+ auto gmScaleB = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
		125	+ reinterpret_cast<__gm__ ElementMxScaleB*>(gmmAddrInfo.bScaleGlobal)),
		126	+ layouts.scaleB);
		127	+ __gm__ float* metaInfoGm = nullptr;
		128	+ if constexpr (TopkWeightsPrefetch && g_coreType == AscendC::AIV) {
		129	+ metaInfoGm = reinterpret_cast<__gm__ float*>(gmmAddrInfo.metaInfoGlobal);
		130	+ }
		131	+
		132	+ using WorkSetType = Gmm1WorkSet<Scheduler, decltype(gmA), decltype(gmB), decltype(gmScaleA), decltype(gmScaleB),
		133	+ decltype(gmC), decltype(metaInfoGm)>;
		134	+ WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmC, metaInfoGm};
		135	+ if constexpr (g_coreType == AscendC::AIC) {
		136	+ Gmm1AicMmadA8W4<BlockMmad, IsShared, TopkWeightsPrefetch>(*pipeline.block, workSet, params, gmmAddrInfo, config,
		137	+ startLoopIdx, tileNum, expertIdx);
		138	+ } else {
		139	+ Gmm1Aiv0PrologueA8W4<BlockPrologue>(*pipeline.block, workSet.scheduler, workSet.gmB, config, startLoopIdx,
		140	+ tileNum);
		141	+ }
		142	+}
		143	+
90	} // namespace GmmKernel	144	} // namespace GmmKernel
91		145
92	-// 构造 routed GMM1 的矩阵配置和分核状态；共享专家仍使用公共入口。
93	template <typename ElementA, typename EpilogueElementA, typename ElementB, typename ElementC, typename ElementMxScaleA,	146	template <typename ElementA, typename EpilogueElementA, typename ElementB, typename ElementC, typename ElementMxScaleA,
94	typename ElementMxScaleB, bool IsWeightNZ = false, uint32_t Gmm1TileM = L1_TILE_M_256,	147	typename ElementMxScaleB, bool IsWeightNZ = false, uint32_t Gmm1TileM = L1_TILE_M_256,
95	uint32_t EpilogueTileM = Gmm1TileM, bool TopkWeightsPrefetch = false, bool IsShared = false,	148	uint32_t EpilogueTileM = Gmm1TileM, bool TopkWeightsPrefetch = false, bool IsShared = false,
96	bool IsWaveFlagGrained = false, typename BlockEpilogue, typename... LayoutArgs>	149	bool IsWaveFlagGrained = false, typename BlockEpilogue, typename... LayoutArgs>
97	-__aicore__ inline void RunGmm1ToGm(BlockEpilogue& epilogueOp, const Params& params,	150	+__aicore__ inline void RunGmm1Generic(BlockEpilogue& epilogueOp, const Params& params,
98	- const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,	151	+ const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
99	- const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,	152	+ const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx, int32_t& vecSetSyncCom,
100	- const BlockJobContext& blockJob, uint32_t expertBeforeCnt, uint32_t expertIdx,	153	+ const BlockJobContext& blockJob, uint32_t expertBeforeCnt, uint32_t expertIdx,
101	- uint16_t& pingpongIdx, uint32_t activationOutputDivisor = ACTIVATION_N_HALF,	154	+ uint16_t& pingpongIdx, uint32_t activationOutputDivisor, void* blockMmadContext,
102	- void* blockMmadContext = nullptr, bool allowWeightL2Bypass = false,	155	+ bool allowWeightL2Bypass = false, const LayoutArgs&... layoutArgs)
103	- const LayoutArgs&... layoutArgs)
104	{	156	{
105	using GmmConfig =	157	using GmmConfig =
106	GmmKernel::Config<false, COMBINE_NO_QUANT, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB,	158	GmmKernel::Config<false, COMBINE_NO_QUANT, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB,
	@@ -123,9 +175,55 @@ __aicore__ inline void RunGmm1ToGm(BlockEpilogue& epilogueOp, const Params& para
123	auto* typedBlockMmadContext = reinterpret_cast<MmadContext*>(blockMmadContext);	175	auto* typedBlockMmadContext = reinterpret_cast<MmadContext*>(blockMmadContext);
124	using MakeLayoutC = typename GmmConfig::MakeLayoutC;	176	using MakeLayoutC = typename GmmConfig::MakeLayoutC;
125	auto layouts = GmmConfig::BuildLayouts(config, layoutArgs...);	177	auto layouts = GmmConfig::BuildLayouts(config, layoutArgs...);
126	- GmmKernel::Gmm1ExecToGm<BlockMmad, ElementC, MakeLayoutC, TopkWeightsPrefetch, IsShared, IsWaveFlagGrained>(	178	+ GmmKernel::Gmm1ExecGeneric<BlockMmad, ElementC, MakeLayoutC, TopkWeightsPrefetch, IsShared, IsWaveFlagGrained>(
127	- scheduler, params, gmmAddrInfo, config, startLoopIdx, tileNum, epilogueOp, expertBeforeCnt, expertIdx,	179	+ scheduler, params, gmmAddrInfo, config, startLoopIdx, tileNum, epilogueOp, vecSetSyncCom, expertBeforeCnt,
128	- pingpongIdx, typedBlockMmadContext, allowWeightL2Bypass, layouts);	180	+ expertIdx, pingpongIdx, typedBlockMmadContext, allowWeightL2Bypass, layouts);
		181	+
		182	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
		183	+}
		184	+
		185	+template <typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA, typename ElementMxScaleB,
		186	+ uint32_t Gmm1TileM = L1_TILE_M_256, uint32_t EpilogueTileM = Gmm1TileM, bool TopkWeightsPrefetch = false,
		187	+ bool IsShared = false, bool IsWaveFlagGrained = false, typename BlockContext, typename BlockEpilogue,
		188	+ typename... LayoutArgs>
		189	+__aicore__ inline void RunGmm1A8W4(const BlockContext& pipeline, BlockEpilogue& activationQuantOp, const Params& params,
		190	+ const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
		191	+ const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,
		192	+ const BlockJobContext& blockJob, uint32_t expertBeforeCnt, uint32_t expertIdx,
		193	+ uint32_t activationOutputDivisor = ACTIVATION_N_HALF,
		194	+ const LayoutArgs&... layoutArgs)
		195	+{
		196	+ static_assert(std::is_same_v<ElementA, __fp8e4m3>, "Activation must be __fp8e4m3");
		197	+ static_assert(std::is_same_v<ElementB, __fp4e2m1x2>, "Weight must be __fp4e2m1x2");
		198	+
		199	+ using GmmConfig = GmmKernel::Config<true, 0, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, false,
		200	+ TopkWeightsPrefetch, IsShared, false, IsWaveFlagGrained>;
		201	+ auto config = GmmConfig::BuildGmm1ProblemConfig(problemShape, blockJob, Gmm1TileM, activationOutputDivisor);
		202	+
		203	+ if constexpr (g_coreType == AIV) {
		204	+ activationQuantOp.UpdateNextProblem({config.m, config.outputN, config.k, 0});
		205	+ }
		206	+
		207	+ using BlockMmad = typename GmmConfig::BlockMmad;
		208	+ using BlockPrologue = typename GmmConfig::BlockPrologue;
		209	+ using MakeLayoutC = typename GmmConfig::MakeLayoutC;
		210	+
		211	+ GmmKernel::BlockScheduler scheduler({config.m, config.schedulerN, config.k},
		212	+ GmmKernel::BlockScheduler::Params{asc::te::make_coord(
		213	+ static_cast<int64_t>(config.tileM), static_cast<int64_t>(L1_TILE_N))});
		214	+ uint32_t tileNum = scheduler.GetTileNum();
		215	+ uint32_t startLoopIdx =
		216	+ (config.blockIdx < startBlockIdx ? config.blockIdx + config.blockNum : config.blockIdx) - startBlockIdx;
		217	+ if (startLoopIdx >= tileNum) {
		218	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
		219	+ return;
		220	+ }
		221	+
		222	+ auto layouts = GmmConfig::BuildLayouts(config, layoutArgs...);
		223	+ GmmKernel::Gmm1ExecA8W4<BlockMmad, BlockPrologue, ElementC, MakeLayoutC, IsShared, GmmKernel::BlockScheduler,
		224	+ decltype(config), TopkWeightsPrefetch, IsWaveFlagGrained>(
		225	+ scheduler, params, gmmAddrInfo, config, startLoopIdx, tileNum, activationQuantOp, expertBeforeCnt, expertIdx,
		226	+ layouts, pipeline);
129		227
130	startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;	228	startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
131	}	229	}

Amc2/mega_moe/op_kernel/arch92/stage/mega_moe_gmm2_combine_arch92.h+215-0

@@ -0,0 +1,215 @@
	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	+#ifndef MEGA_MOE_GMM2_COMBINE_ARCH92_H
	12	+#define MEGA_MOE_GMM2_COMBINE_ARCH92_H
	13	+
	14	+#include "../../arch35_92/stage/mega_moe_gmm2_combine.h"
	15	+
	16	+namespace MegaMoeImpl {
	17	+namespace GmmKernel {
	18	+
	19	+template <uint8_t CombineQuantMode, typename BlockMmad, typename ElementC, bool IsLayered = false,
	20	+ bool IsShared = false, bool NotifyCombineTileReady = false, typename Scheduler, typename Config>
	21	+__aicore__ inline void Gmm2ExecGeneric(Scheduler& scheduler, const GMMAddrInfo& gmmAddrInfo, const Config& config,
	22	+ uint32_t startLoopIdx, uint32_t tileNum,
	23	+ BlockMmadContext<BlockMmad>* blockMmadContext, bool allowWeightL2Bypass,
	24	+ uint32_t rowOffsetInExpert = 0U, const Params* params = nullptr,
	25	+ uint32_t expertTokenCount = 0U)
	26	+{
	27	+ using KernelConfig = typename Config::KernelConfig;
	28	+ using ElementA = typename KernelConfig::ElementAType;
	29	+ using ElementB = typename KernelConfig::ElementBType;
	30	+ using ElementMxScaleA = typename KernelConfig::ElementMxScaleAType;
	31	+ using ElementMxScaleB = typename KernelConfig::ElementMxScaleBType;
	32	+ using BiasType = typename KernelConfig::BiasType;
	33	+
	34	+ auto layouts = KernelConfig::BuildLayouts(config);
	35	+ auto gmA = asc::te::make_tensor(
	36	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementA*>(gmmAddrInfo.aGlobal)),
	37	+ layouts.a);
	38	+ auto gmB = asc::te::make_tensor(
	39	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementB*>(gmmAddrInfo.bGlobal)),
	40	+ layouts.b);
	41	+ auto gmScaleA = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
	42	+ reinterpret_cast<__gm__ ElementMxScaleA*>(gmmAddrInfo.aScaleGlobal)),
	43	+ layouts.scaleA);
	44	+ auto gmScaleB = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
	45	+ reinterpret_cast<__gm__ ElementMxScaleB*>(gmmAddrInfo.bScaleGlobal)),
	46	+ layouts.scaleB);
	47	+ if constexpr (Config::IS_WAVE_FLAG_GRAINED && g_coreType == AscendC::AIC) {
	48	+ SetWaveWeightL2CacheHint<KernelConfig::IS_WEIGHT_NZ, KernelConfig>(config, allowWeightL2Bypass, gmB, gmScaleB);
	49	+ }
	50	+ auto gmBias = asc::te::make_tensor(
	51	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ BiasType*>(0UL)), layouts.bias);
	52	+ auto gmC = asc::te::make_tensor(
	53	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(gmmAddrInfo.gmm2OutGlobal)),
	54	+ layouts.c);
	55	+
	56	+ using WorkSetType = GroupMatmulWorkSet<Scheduler, decltype(gmA), decltype(gmB), decltype(gmScaleA),
	57	+ decltype(gmScaleB), decltype(gmBias), decltype(gmC)>;
	58	+ WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmBias, gmC};
	59	+
	60	+ if constexpr (g_coreType == AscendC::AIC) {
	61	+ blockMmadContext->template Init<GmmStage::GMM2>(config);
	62	+ Gmm2AicMmadGeneric<CombineQuantMode, BlockMmad, IsShared, IsLayered, NotifyCombineTileReady>(
	63	+ blockMmadContext->blockMmad, workSet, gmmAddrInfo, config, startLoopIdx, tileNum, rowOffsetInExpert,
	64	+ expertTokenCount);
	65	+ } else if constexpr (NotifyCombineTileReady) {
	66	+ // NotifyCombineTileReady 为真时唯一 AIV 逐 tile 锁步消费。combine 去重的独立消费路径
	67	+ // （RunWaveCombineStageDedup）由调用方把本参数实例化为 false，整个分支编译期裁掉。
	68	+ using MakeLayoutC = typename KernelConfig::MakeLayoutC;
	69	+ CombineTokenRange<ElementC, MakeLayoutC>(workSet.scheduler, workSet.gmC, *params, gmmAddrInfo, config,
	70	+ startLoopIdx, tileNum);
	71	+ }
	72	+}
	73	+
	74	+// 保留既有未开放分支的模板引用，不扩展 Arch92 能力范围。
	75	+template <typename BlockMmad, typename BlockPrologue, typename ElementC, bool IsShared, bool IsLayered,
	76	+ bool NotifyCombineTileReady, typename Scheduler, typename Config>
	77	+__aicore__ inline void Gmm2ExecA8W4(Scheduler& scheduler, const GMMAddrInfo& gmmAddrInfo, const Config& config,
	78	+ uint32_t startLoopIdx, uint32_t tileNum, uint32_t expertTokenCount,
	79	+ uint32_t rowOffsetInExpert, const Params* params,
	80	+ const A8W4BlockContext<BlockMmad, BlockPrologue>& pipeline)
	81	+{
	82	+ using KernelConfig = typename Config::KernelConfig;
	83	+ using ElementA = typename KernelConfig::ElementAType;
	84	+ using ElementB = typename KernelConfig::ElementBType;
	85	+ using ElementMxScaleA = typename KernelConfig::ElementMxScaleAType;
	86	+ using ElementMxScaleB = typename KernelConfig::ElementMxScaleBType;
	87	+ using BiasType = typename KernelConfig::BiasType;
	88	+
	89	+ auto layouts = KernelConfig::BuildLayouts(config);
	90	+ auto gmC = asc::te::make_tensor(
	91	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementC*>(gmmAddrInfo.gmm2OutGlobal)),
	92	+ layouts.c);
	93	+ auto gmA = asc::te::make_tensor(
	94	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementA*>(gmmAddrInfo.aGlobal)),
	95	+ layouts.a);
	96	+ auto gmB = asc::te::make_tensor(
	97	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ ElementB*>(gmmAddrInfo.bGlobal)),
	98	+ layouts.b);
	99	+ auto gmScaleA = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
	100	+ reinterpret_cast<__gm__ ElementMxScaleA*>(gmmAddrInfo.aScaleGlobal)),
	101	+ layouts.scaleA);
	102	+ auto gmScaleB = asc::te::make_tensor(asc::te::make_mem_ptr<asc::te::location::gm>(
	103	+ reinterpret_cast<__gm__ ElementMxScaleB*>(gmmAddrInfo.bScaleGlobal)),
	104	+ layouts.scaleB);
	105	+ auto gmBias = asc::te::make_tensor(
	106	+ asc::te::make_mem_ptr<asc::te::location::gm>(reinterpret_cast<__gm__ BiasType*>(0UL)), layouts.bias);
	107	+
	108	+ using WorkSetType = GroupMatmulWorkSet<Scheduler, decltype(gmA), decltype(gmB), decltype(gmScaleA),
	109	+ decltype(gmScaleB), decltype(gmBias), decltype(gmC)>;
	110	+ WorkSetType workSet{scheduler, gmA, gmB, gmScaleA, gmScaleB, gmBias, gmC};
	111	+
	112	+ if constexpr (g_coreType == AscendC::AIC) {
	113	+ auto& blockMmad = *pipeline.block;
	114	+ typename BlockMmad::BlockShape l0TileShape{config.blockMmadTiling.tileM, config.blockMmadTiling.tileN,
	115	+ L0_TILE_K, 0};
	116	+ typename BlockMmad::ProblemShape matmulShape{config.m, config.outputN, config.k, 0};
	117	+ // GMM2 has no gate/up concatenation; preserve the persistent buffer indices.
	118	+ blockMmad.template Init<false>(matmulShape, l0TileShape, false);
	119	+ Gmm2AicMmadA8W4<IsShared, IsLayered, NotifyCombineTileReady, BlockMmad, decltype(workSet.scheduler),
	120	+ decltype(workSet.gmA), decltype(workSet.gmScaleA), decltype(workSet.gmScaleB),
	121	+ std::remove_reference_t<decltype(workSet.gmC)>, Config>(
	122	+ blockMmad, workSet.scheduler, workSet.gmA, workSet.gmScaleA, workSet.gmScaleB, workSet.gmC, gmmAddrInfo,
	123	+ config, startLoopIdx, tileNum, expertTokenCount, rowOffsetInExpert);
	124	+ } else {
	125	+ Gmm2Aiv0PrologueA8W4(*pipeline.block, workSet.scheduler, workSet.gmB, config, startLoopIdx, tileNum);
	126	+ }
	127	+}
	128	+
	129	+} // namespace GmmKernel
	130	+
	131	+template <uint8_t CombineQuantMode, typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA,
	132	+ typename ElementMxScaleB, bool IsWeightNZ = false, bool IsLayered = false, uint32_t Gmm1TileM = L1_TILE_M_256,
	133	+ bool TopkWeightsPrefetch = false, bool IsShared = false, bool IsWaveFlagGrained = false,
	134	+ bool NotifyCombineTileReady = false>
	135	+__aicore__ inline void RunGmm2Generic(const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
	136	+ const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,
	137	+ const BlockJobContext& blockJob, void* blockMmadContext,
	138	+ bool allowWeightL2Bypass = false, uint32_t rowOffsetInExpert = 0U,
	139	+ const Params* params = nullptr,
	140	+ uint32_t activationOutputDivisor = ACTIVATION_N_HALF,
	141	+ uint32_t expertTokenCount = 0U)
	142	+{
	143	+ static_assert(!IsShared \|\| !NotifyCombineTileReady, "Shared expert GMM2 has its own completion counter");
	144	+ static_assert(!IsLayered \|\| !NotifyCombineTileReady, "Layered GMM2 selects its synchronization through IsLayered");
	145	+ using GmmConfig =
	146	+ GmmKernel::Config<false, CombineQuantMode, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB,
	147	+ IsWeightNZ, TopkWeightsPrefetch, IsShared, IsLayered, IsWaveFlagGrained>;
	148	+ auto config = GmmConfig::BuildGmm2ProblemConfig(problemShape, blockJob, Gmm1TileM, activationOutputDivisor);
	149	+
	150	+ GmmKernel::BlockScheduler scheduler({config.m, config.schedulerN, config.k},
	151	+ GmmKernel::BlockScheduler::Params{asc::te::make_coord(
	152	+ static_cast<int64_t>(config.tileM), static_cast<int64_t>(L1_TILE_N))});
	153	+ uint32_t tileNum = scheduler.GetTileNum();
	154	+
	155	+ if constexpr (CombineQuantMode == COMBINE_NO_QUANT && g_coreType == AscendC::AIV && !NotifyCombineTileReady) {
	156	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
	157	+ return;
	158	+ }
	159	+
	160	+ uint32_t startLoopIdx =
	161	+ (config.blockIdx < startBlockIdx ? config.blockIdx + config.blockNum : config.blockIdx) - startBlockIdx;
	162	+ using BlockMmad = typename GmmConfig::BlockMmad;
	163	+ using MmadContext = GmmKernel::BlockMmadContext<BlockMmad>;
	164	+ auto* typedBlockMmadContext = reinterpret_cast<MmadContext*>(blockMmadContext);
	165	+ GmmKernel::Gmm2ExecGeneric<CombineQuantMode, BlockMmad, ElementC, IsLayered, IsShared, NotifyCombineTileReady>(
	166	+ scheduler, gmmAddrInfo, config, startLoopIdx, tileNum, typedBlockMmadContext, allowWeightL2Bypass,
	167	+ rowOffsetInExpert, params, expertTokenCount);
	168	+
	169	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
	170	+}
	171	+
	172	+template <typename ElementA, typename ElementB, typename ElementC, typename ElementMxScaleA, typename ElementMxScaleB,
	173	+ uint32_t Gmm1TileM = L1_TILE_M_256, bool TopkWeightsPrefetch = false, bool IsShared = false,
	174	+ bool IsLayered = false, bool IsWaveFlagGrained = false, bool NotifyCombineTileReady = false,
	175	+ typename BlockContext>
	176	+__aicore__ inline void RunGmm2A8W4(const BlockContext& pipeline,
	177	+ const AscendC::Shape<int64_t, int64_t, int64_t, int64_t>& problemShape,
	178	+ const GMMAddrInfo& gmmAddrInfo, uint32_t& startBlockIdx,
	179	+ const BlockJobContext& blockJob, uint32_t expertTokenCount,
	180	+ uint32_t rowOffsetInExpert, const Params* params = nullptr,
	181	+ uint32_t activationOutputDivisor = ACTIVATION_N_HALF)
	182	+{
	183	+ static_assert(!IsShared \|\| !NotifyCombineTileReady, "Shared expert GMM2 has its own completion counter");
	184	+ static_assert(!IsLayered \|\| !NotifyCombineTileReady, "Layered GMM2 selects its synchronization through IsLayered");
	185	+ static_assert(std::is_same_v<ElementA, __fp8e4m3>, "Activation must be __fp8e4m3");
	186	+ static_assert(std::is_same_v<ElementB, __fp4e2m1x2>, "Weight must be __fp4e2m1x2");
	187	+
	188	+ using GmmConfig = GmmKernel::Config<true, 0, ElementA, ElementB, ElementC, ElementMxScaleA, ElementMxScaleB, false,
	189	+ TopkWeightsPrefetch, IsShared, IsLayered, IsWaveFlagGrained>;
	190	+ auto config = GmmConfig::BuildGmm2ProblemConfig(problemShape, blockJob, Gmm1TileM, activationOutputDivisor);
	191	+
	192	+ using BlockMmad = typename GmmConfig::BlockMmad;
	193	+ using BlockPrologue = typename GmmConfig::BlockPrologue;
	194	+ GmmKernel::BlockScheduler scheduler({config.m, config.outputN, config.k},
	195	+ GmmKernel::BlockScheduler::Params{asc::te::make_coord(
	196	+ static_cast<int64_t>(config.tileM), static_cast<int64_t>(L1_TILE_N))});
	197	+ // tileNum 与 mega_moe_utils.h 的 GetGmmProblemTileCount 同源，预判无活的口径不得各算一份。
	198	+ uint32_t tileNum = GetGmmProblemTileCount(config.m, config.outputN);
	199	+ uint32_t startLoopIdx =
	200	+ (config.blockIdx < startBlockIdx ? config.blockIdx + config.blockNum : config.blockIdx) - startBlockIdx;
	201	+ if (startLoopIdx >= tileNum) {
	202	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
	203	+ return;
	204	+ }
	205	+
	206	+ GmmKernel::Gmm2ExecA8W4<BlockMmad, BlockPrologue, ElementC, IsShared, IsLayered, NotifyCombineTileReady,
	207	+ GmmKernel::BlockScheduler, decltype(config)>(
	208	+ scheduler, gmmAddrInfo, config, startLoopIdx, tileNum, expertTokenCount, rowOffsetInExpert, params, pipeline);
	209	+
	210	+ startBlockIdx = (startBlockIdx + tileNum) % config.blockNum;
	211	+}
	212	+
	213	+} // namespace MegaMoeImpl
	214	+
	215	+#endif // MEGA_MOE_GMM2_COMBINE_ARCH92_H

	@@ -51,6 +51,7 @@ private:
51		51
52	using MegaMoeBase::DispatchBuffInit;	52	using MegaMoeBase::DispatchBuffInit;
53	using MegaMoeBase::PrepareMoeExpertTokenCountTable;	53	using MegaMoeBase::PrepareMoeExpertTokenCountTable;
		54	+ using MegaMoeBase::BuildMoeDedupRowDescTable;
54	using MegaMoeBase::ReloadDispatchCumsumRange;	55	using MegaMoeBase::ReloadDispatchCumsumRange;
55	using MegaMoeBase::EnterSteadyDispatch;	56	using MegaMoeBase::EnterSteadyDispatch;
56	using MegaMoeBase::RunGmm2CombineForExpert;	57	using MegaMoeBase::RunGmm2CombineForExpert;
	@@ -317,12 +318,17 @@ __aicore__ inline void MegaMoeA4W4Wave<TemplateMegaMoeA4W4WaveTypeFunc>::Process
317	if (tokenDispatchConfig_.dedup.dedupMode != 0) {	318	if (tokenDispatchConfig_.dedup.dedupMode != 0) {
318	// 去重建表：route 槽到达保证与 count epoch 校验同源；函数尾部含 AIV1 间栅栏。	319	// 去重建表：route 槽到达保证与 count epoch 校验同源；函数尾部含 AIV1 间栅栏。
319	PrepareDedupPrepassAiv0Local(tokenDispatchConfig_.dedup, commonConfig_, params_, tokenDispatchScratch_);	320	PrepareDedupPrepassAiv0Local(tokenDispatchConfig_.dedup, commonConfig_, params_, tokenDispatchScratch_);
320	- BuildDedupRowDescTable<TopkIndexType>(tokenDispatchConfig_.dedup, commonConfig_, gmmExecutionConfig_.blockJob,	321	+ if constexpr (g_coreType == AIV) {
321	- syncWorkspaceLayout_, params_, tokenDispatchScratch_);	322	+ BuildMoeDedupRowDescTable();
		323	+ }
322	}	324	}
323	WaveCombineBufferConfig combineBufferConfig{};	325	WaveCombineBufferConfig combineBufferConfig{};
324	if constexpr (CombineQuantMode != COMBINE_NO_QUANT) {	326	if constexpr (CombineQuantMode != COMBINE_NO_QUANT) {
325	- combineBufferConfig = InitWaveCombineBuffers<CombineQuantMode>(commonConfig_, waveCombineScratch_);	327	+ if constexpr (g_coreType == AIV) {
		328	+ if (MegaMoeHardwarePolicy::GetSubBlockIdx() == 1U) {
		329	+ combineBufferConfig = InitWaveCombineBuffers<CombineQuantMode>(commonConfig_, waveCombineScratch_);
		330	+ }
		331	+ }
326	} else if (IsCombineDedupOn(params_.tilingData->compute.dedupMode)) {	332	} else if (IsCombineDedupOn(params_.tilingData->compute.dedupMode)) {
327	// 非量化 combine 去重：专家粒度独立消费（基类 RunGmm2CombineForExpert 内接入）。	333	// 非量化 combine 去重：专家粒度独立消费（基类 RunGmm2CombineForExpert 内接入）。
328	if constexpr (g_coreType == AIV) {	334	if constexpr (g_coreType == AIV) {

	@@ -52,6 +52,7 @@ private:
52		52
53	using MegaMoeBase::DispatchBuffInit;	53	using MegaMoeBase::DispatchBuffInit;
54	using MegaMoeBase::PrepareMoeExpertTokenCountTable;	54	using MegaMoeBase::PrepareMoeExpertTokenCountTable;
		55	+ using MegaMoeBase::BuildMoeDedupRowDescTable;
55	using MegaMoeBase::ReloadDispatchCumsumRange;	56	using MegaMoeBase::ReloadDispatchCumsumRange;
56	using MegaMoeBase::EnterSteadyDispatch;	57	using MegaMoeBase::EnterSteadyDispatch;
57	using MegaMoeBase::RunGmm2CombineForExpert;	58	using MegaMoeBase::RunGmm2CombineForExpert;
	@@ -291,12 +292,17 @@ __aicore__ inline void MegaMoeA8W4Wave<TemplateMegaMoeA8W4WaveTypeFunc>::Process
291	if (tokenDispatchConfig_.dedup.dedupMode != 0) {	292	if (tokenDispatchConfig_.dedup.dedupMode != 0) {
292	// 去重建表：route 槽到达保证与 count epoch 校验同源；函数尾部含 AIV1 间栅栏。	293	// 去重建表：route 槽到达保证与 count epoch 校验同源；函数尾部含 AIV1 间栅栏。
293	PrepareDedupPrepassAiv0Local(tokenDispatchConfig_.dedup, commonConfig_, params_, tokenDispatchScratch_);	294	PrepareDedupPrepassAiv0Local(tokenDispatchConfig_.dedup, commonConfig_, params_, tokenDispatchScratch_);
294	- BuildDedupRowDescTable<TopkIndexType>(tokenDispatchConfig_.dedup, commonConfig_, gmmExecutionConfig_.blockJob,	295	+ if constexpr (g_coreType == AIV) {
295	- syncWorkspaceLayout_, params_, tokenDispatchScratch_);	296	+ BuildMoeDedupRowDescTable();
		297	+ }
296	}	298	}
297	WaveCombineBufferConfig combineBufferConfig{};	299	WaveCombineBufferConfig combineBufferConfig{};
298	if constexpr (CombineQuantMode != COMBINE_NO_QUANT) {	300	if constexpr (CombineQuantMode != COMBINE_NO_QUANT) {
299	- combineBufferConfig = InitWaveCombineBuffers<CombineQuantMode>(commonConfig_, waveCombineScratch_);	301	+ if constexpr (g_coreType == AIV) {
		302	+ if (MegaMoeHardwarePolicy::GetSubBlockIdx() == 1U) {
		303	+ combineBufferConfig = InitWaveCombineBuffers<CombineQuantMode>(commonConfig_, waveCombineScratch_);
		304	+ }
		305	+ }
300	} else if (IsCombineDedupOn(params_.tilingData->compute.dedupMode)) {	306	} else if (IsCombineDedupOn(params_.tilingData->compute.dedupMode)) {
301	// 非量化 combine 去重：专家粒度独立消费（基类 RunGmm2CombineForExpert 内接入）。	307	// 非量化 combine 去重：专家粒度独立消费（基类 RunGmm2CombineForExpert 内接入）。
302	if constexpr (g_coreType == AIV) {	308	if constexpr (g_coreType == AIV) {