PTO MegaMoEdispatch_mega_combine七阶段流水解析从 FrontReorder 到 Unpermute 的数据流、伪码与同步机制【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa本文以 CANN / pto-isa 仓库中kernels/manual/a2a3/dispatch_mega_combine/pseudocode.md为主体系统讲解 PTO Manual kernel 实现的 MegaMoE Dispatch FFN Combine 融合算子传统 MoE 中多次重排、AlltoAllV 通信与 grouped FFN 计算被合并进单个大 kernel并以 local expert 为粒度做 AIC/AIV 重叠。读完本文你将掌握 FrontReorder、Dispatch、GMM1、SwiGLU、GMM2、Combine、Unpermute 七个阶段各自的数据流、核心分工、producer/consumer epoch 同步模型以及如何把伪码对应到仓库中op_kernel/下的真实实现。一、总体数据流七阶段主路径伪码文档按当前 A2/A3dispatch_mega_combine生产主路径描述七个阶段的数据流、核心分工和同步关系。整体执行顺序在MegaMoe::Process()中串起见 dispatch_mega_combine.h先执行FrontReorderProcess再按固定分组推进ProcessFixedGroups其中 Dispatch/GMM1/SwiGLU/GMM2/Combine 按 expert 级流水重叠执行Unpermute 最后分两阶段还原。FrontReorder: x[M, K] expertId[M, topK] - source rank remoteWindow.offsetA[expert-major rows, K 32] - expandedRowIdx / tokenPerExpert / preSumBeforeRank / cumsumMM / expertTokenNums Dispatch: destination rank pulls its local-expert rows from every source rank offsetA - gmA[expert-major rows, K] perTokenScale1[rows] GMM1: gmA[int8] weight1[int8] scale1 - gmC[rows, N] half SwiGLU: gmC * perTokenScale1 - silu(up) * gate - dynamic quant - gmPermutedToken[rows, N/2] int8 perTokenScale2[rows] GMM2: gmPermutedToken[int8] weight2[int8] scale2 - gmm2Output[rows, K] half Combine: gmm2Output * perTokenScale2 - source rank remoteWindow.offsetD[expanded row, K] - per-source-rank expert progress / DataReady Unpermute: offsetD probs expandedRowIdx - phase 1: 32 AIV process tokens whose routes are ready - phase 2: 48 AIV process the remaining tokens - out[M, K]伪码中省略了边界检查、UB 地址计算和具体 event ID分组参数及 shape 配置见 overview_v1.md。一张卡包含 24 个 AIC 和 48 个 AIV按物理 AI Core 划分固定分组GMM1 用物理核 0..15 的 16 个 AIC开头短暂扩到 24GMM2 用物理核 16..23 的 8 个 AIC尾部扩到 24Dispatch 用物理核 0..15 的 AIV0SwiGLU 用同批核的 AIV1Combine 用物理核 16..23 的 AIV08 个Unpermute 最终复用全部 48 个 AIV。二、阶段同步模型producer/consumer epoch slot组内阶段同步使用 producer/consumer slot 表达数组下标是逻辑 core IDepoch是 slot 中保存的 expert 进度值。对 experteproducer 发布2 * e 1coordinator 汇总后向 consumer 发布2 * e 2。这一约定在源码中有精确对应。在 mega_expert_sync.hpp 中PublishGroupArrival计算epoch notifyCall * 2 1由每个 producer 写入自己的 arrival slotscalar GM store dcci保证可见性CoordinateGroupConsumersMte中arriveEpoch notifyCall * 2 1、releaseEpoch arriveEpoch 1coordinator 用WaitArrivalMinMte批量读 N 个 arrival 计数并取共同进度MTE 批量TLOAD到 UB 快照区逐 peer 求 min再通过PublishEpochRangeMte批量写 M 个 ready slotFixedSyncLayout依次为每个阶段预留dispatchArrival/Ready、gmm1Arrival、swigluReady/Arrival、gmm2Ready/Arrival、combineReady、gmm1DoneSlot、gmm2JoinSlot等 slot 区并在头尾放置 canary 值用于调试校验。epoch是单次 kernel launch 内按 expert 单调推进的进度值而非一次性布尔 flagcoordinator 读到的共同进度若已超过目标 expert一次检查即可覆盖此前连续完成的多个 expertconsumer 后续检查直接命中不需要重新等待。跨 launch 的通知使用独立的 DataReady epoch 区分轮次避免复用 HCCL window 时命中上一轮的状态见 overview_v1.md 的“阶段同步”一节。扩缩组只在 expert 边界发生producer 和 consumer 使用同一个切换条件因此支持 GMM124→16、GMM28→24的动态分组扩缩容。三、阶段 1FrontReorder——源 rank 内 route 排序、量化、发布 count关键流程AIV-only 阶段将M * topK条 route 按 global expert 排序生成 expert-major row 布局。host 根据 UB 工作集选择FullLoad、OneCore或MultiCoreFullLoad排序、反排、count 和量化工作集留在 UB活跃 AIV 各自持有完整排序结果并分担量化OneCore单 AIV 在 UB 内完成排序随后全部 AIV 进入共享后处理MultiCore多 AIV 生成有序 run再通过多轮 4 路归并得到最终顺序。排序结果生成expandedRowIdx[srcRoute] dstRow同一 token 的 topK route 复用一次动态量化结果scatter 到offsetA[dstRow]。localTokenPerExpert统计本 rank 发往每个 global expert 的 route 数。count row 加 marker 后写到各 peerpeer 使用TWAIT等待并恢复真实 count随后生成preSumBeforeRank、cumsumMM和expertTokenNums。关键 PTO 接口排序和抽取TSORT32、TMRGSORT、TGATHERTile 绑定和搬运TASSIGN、TLOAD、TSTORE动态量化TCVT、TABS、TROWMAX、TMAX、TDIV跨 rank count readyTWAITAIV 阶段边界SYNCALLAIVOnly对应实现分散在 front_reorder.h、front_fullload_sort.h、front_vms_sort.h 中dispatch 主流程入口见 dispatch_mega_combine.h。count-as-flag机制的具体做法是整行 count 加上0x800000marker 写到 peer 的tokenPerExpert[srcRank, :]peer 用TWAIT ! 0等待读回后减 marker 恢复真实 count——count row 同时充当数据与到达 flag从而消除 AlltoAll count 后的整机全同步。伪码已还原为可读形式if frontCase FullLoad: each active AIV: expertUb PtoLoadVector(expertId[0 : M * topK]) srcRouteUb 0 .. M * topK - 1 packedUb TSORT32(expertUb, srcRouteUb) packedUb TMRGSORT(packedUb) sortedExpert, sortedSrcRoute TGATHER(packedUb) expandedRowIdxUb inverse_sort(sortedSrcRoute) # srcRoute - expert-major dstRow owner AIV: localTokenPerExpert count_runs(sortedExpert) AIV0: PtoStoreVector(expandedRowIdx, expandedRowIdxUb) each active AIV for assigned source-token rows: qInt8, scale dynamic_quant_once(x[token, 0:K]) for topkSlot in 0 .. topK - 1: dstRow expandedRowIdxUb[token * topK topkSlot] PtoStoreVector(offsetA[dstRow], packed(qInt8, scale)) else: if frontCase OneCore: AIV0: packedRuns TSORT32(expertId, srcRoute) packedRuns TMRGSORT(packedRuns) else: # MultiCore each sort AIV builds sorted runs with TSORT32 TMRGSORT merge sorted runs in multiple 4-way TMRGSORT rounds AIV0: sortedExpert, sortedSrcRoute TGATHER(packedRuns) store sortedExpert / sortedSrcRoute to front workspace SYNCALLAIVOnly() all AIVs: localTokenPerExpert count(sortedExpert) expandedRowIdx[srcRoute] dstRow qInt8, scale dynamic_quant_once(x[token, 0:K]) scatter packed(qInt8, scale) - offsetA[expandedRowIdx[route]] SYNCALLAIVOnly() for dstRank assigned to this AIV: countRow PtoLoadVector(localTokenPerExpert) countRow countMarker # zero count can also represent arrival PtoStoreVector(peer[dstRank].tokenPerExpert[myRank], countRow) for srcRank assigned to this AIV: TWAIT(tokenPerExpert[srcRank].markers ! 0) countRow tokenPerExpert[srcRank] - countMarker preSumBeforeRank[srcRank] prefix_before_my_local_experts(countRow) SYNCALLAIVOnly() AIV0: cumsumMM inclusive_prefix_across_source_ranks(tokenPerExpert) expertTokenNums cumsumMM[lastSourceRank]四、阶段 2Dispatch——目的 rank 拉取 source rank 的 packed A关键流程物理核0..15的 AIV0 组成 Dispatch 组其中前rankSize个 AIV 分别负责一个 source rank。按 local expert 顺序处理每个 worker 根据preSumBeforeRank从 peeroffsetA找到读起点根据cumsumMM找到本地gmA写起点。每次最多拉取 2 行 packed row两个 96 KiB UB buffer 做 ping-pongpayload 和 per-token scale 分别写入gmA和perTokenScale1。每个 expert 的所有 source-rank worker 完成后由 coordinator 发布该 expert 的 GMM1 ready。GMM1 前部放行 24 个 AIC后续放行 16 个 AIC。offsetA每行是 packed rowint8[K] fp32 scale padding行跨度K 32。Dispatch 实现见 dispatch.h其在MegaMoe::ProcessFixedGroups中由stageNum 9U且groupLocalId rankSize的 Dispatch 角色 AIV 执行。关键 PTO 接口远端 packed row 读取TLOADpayload / scale 拆包写回TSTORETile 绑定TASSIGN组间通知GM arrival/ready epoch MTE 批量汇总/发布伪码for expert in 0 .. expertPerRank - 1: if dispatchLocalId rankSize: srcRank dispatchLocalId rows tokenPerExpert[srcRank, globalExpert(myRank, expert)] srcRowBase preSumBeforeRank[srcRank, expert] dstRowBase groupBase cumsum_before_source(srcRank, expert) for rowOffset in 0 .. rows step 2: bufferId next_pingpong_buffer() packedTile TLOAD(peer[srcRank].offsetA[srcRowBase rowOffset], maxRows2) payloadTile packedTile[:, 0:K] scaleTile packedTile[:, K] TSTORE(gmA[dstRowBase rowOffset], payloadTile) TSTORE(perTokenScale1[dstRowBase rowOffset], scaleTile) publish arrival[dispatchLocalId] 2 * expert 1 coordinator: wait min(arrival[0 : rankSize]) 2 * expert 1 gmm1Consumers (expert fullAicGmm1ExpertCount) ? 24 : 16 publish ready[0 : gmm1Consumers] 2 * expert 2 groupBase cumsumMM[lastSourceRank, expert]五、阶段 3GMM1——按 expert 分组做第一个 int8 GEMM关键流程按 local expert 顺序计算。前fullAicGmm1ExpertCount个 expert 使用全部 24 个 AIC之后缩为物理核0..15的 16 个 AIC。每个 expert 等待自己的 Dispatch ready slot按128 x 256output tile 在参与 AIC 间轮转分配。小 M 且默认 tile 数不足以覆盖参与 AIC 时沿 N 维按 32 列粒度重新均衡使更多 AIC 参与。GMM1/GMM2 复用同一套 GMM pipelineL1 A/B、L0A/L0B 双缓冲L0C 单缓冲并使用 N 方向 9 列蛇形 swizzle。每个 expert 的参与 AIC 全部完成后发布该 expert 的 SwiGLU readyGMM1 全部完成后发布 done供 GMM2 决定扩组时机。GMM 公共逻辑在 gmm_common.h 中GmmCommonPipeline PtoGmmPreloadPreloadAsyncFixpipe..., 128, 256, 512, 128, 256, 128, int8_t, int8_t, half即 L1(M,N,K)(128,256,512)、L0(128,256,128)GmmCommonGetBlockCoordMN实现“N 方向每 9 个 tile 为一组、奇数组反转 blockM”的蛇形 swizzlekGmmCommonSwizzleOffset 9U提升 L1 B 侧权重复用GmmCommonNextStartCoreIdx让下一个 expert 从不同 AIC 起步避免总是 core0 吃第一个 tile。GMM1 的具体实现见 gmm1.h。关键 PTO 接口Tile 绑定和搬运TASSIGN、TLOAD、TEXTRACTCube 计算TMATMUL、TMATMUL_ACCFixPipe scale 和输出TMOV、TSTORE_FP阶段同步GM ready/arrival epoch伪码groupBase 0 startCore 0 for expert in 0 .. expertPerRank - 1: activeAic (expert fullAicGmm1ExpertCount) ? 24 : 16 if physicalAicId activeAic: leave GMM1 and enter the primary GMM2 group wait dispatchReady[physicalAicId] 2 * expert 2 currentM clip(cumsumMM[lastSourceRank, expert], groupBase, maxOutputSize) tiles build_output_tiles(currentM, N, tileM128, tileN256) tiles balance_small_m_along_n_if_needed(tiles, activeAic) for tile assigned to this AIC with rotating startCore: blockM, blockN swizzle_9_columns_snake_m(tile) for kTile in 0 .. K step 512: TLOAD(A_l1[pingpong], gmA[groupBase blockM, kTile]) TLOAD(B_l1[pingpong], weight1[expert, kTile, blockN]) for l0k in 0 .. 512 step 128: TEXTRACT(A_l0, A_l1, l0k) TEXTRACT(B_l0, B_l1, l0k) firstK ? TMATMUL(acc, A_l0, B_l0) : TMATMUL_ACC(acc, A_l0, B_l0) TMOV(fixpipeScale, scale1[expert, blockN]) TSTORE_FP(gmC[groupBase blockM, blockN], acc, fixpipeScale) synchronize participating AICs publish SwiGLU ready for expert groupBase currentM startCore rotate_start_core(startCore, tileCount, activeAic) GMM1 coordinator publishes gmm1Done六、阶段 4SwiGLU——GMM1 输出反量化、激活、再量化关键流程物理核0..15的 AIV1 组成 SwiGLU 组M16 使用 8 个 AIV其余目标 shape 使用 16 个。按 expert 逐个等待 GMM1 ready并将当前 expert 的 row 均分给活跃 AIV。每个 worker 使用双 UB stage 做 full-row load/compute/store 流水。gmC转 fp32 后乘perTokenScale1计算silu(up) * gate再按 row 动态量化为 int8。perTokenScale2在 UB 中按最多 128 行聚合后写回。当前 expert 的全部活跃 AIV 完成后coordinator 向 8 个 GMM2 AIC 发布 ready。其中silu(x) x / (1 exp(-x))用TEXP TADDS TDIV完成动态量化按 row 求max(abs(y))scale2 maxAbs / 127qInt8 CINT(y / scale2)。perTokenScale2不逐行写 GM而是缓存在 UB 的 scale chunk buffer每 128 行批量PtoStoreVector一次减少小粒度 GM 写。实现见 swiglu.h。关键 PTO 接口读写和类型转换TLOAD、TSTORE、TCVT反量化和激活TMULS、TEXP、TADDS、TDIV、TMUL动态量化归约TABS、TROWMAX、TMAXscale 批量写回PtoStoreVector阶段同步GM arrival/ready epoch伪码groupBase 0 for expert in 0 .. expertPerRank - 1: coordinator waits until all GMM1 producers finish expert wait swigluReady[swigluLocalId] 2 * expert 2 currentM clip(cumsumMM[lastSourceRank, expert], groupBase, maxOutputSize) localRowStart, localRows split_rows(currentM, swigluActiveAiv) prefetch first gmC row for row in assigned rows: bufferId row % 2 if hasNextRow: prefetch next gmC row into the other UB stage cFp32 TCVT(gmC[row], CAST_NONE) dequant TMULS(cFp32, perTokenScale1[row]) expNegUp TEXP(-dequant[0 : N/2]) silu TDIV(dequant[0 : N/2], TADDS(expNegUp, 1.0)) y TMUL(silu, dequant[N/2 : N]) maxAbs reduce_max(TABS(y)) scale2 max(maxAbs, eps) / 127 qInt8 TCVT(y / scale2, CAST_RINT) TSTORE(gmPermutedToken[row], qInt8) append scale2 to current 128-row scale buffer flush perTokenScale2 scale buffers publish arrival[swigluLocalId] 2 * expert 1 coordinator waits all active SwiGLU arrivals coordinator publishes gmm2Ready[0 : 8] 2 * expert 2 groupBase currentM七、阶段 5GMM2——动态扩组的第二个 int8 GEMM关键流程物理核16..23的 8 个 AIC 从 expert 0 开始 GMM2完成 GMM1 的 16 个 AIC 等待加入时机。原 GMM2 组从配置的gmm2JoinCheckStartExpert开始在每个 expert 边界检查 GMM1 done。coordinator 将统一决策写入 join slot确保 producer 和 consumer 对同一个 join expert 达成一致。join 前由 8 个 AIC 计算join expert 及之后由全部 24 个 AIC 计算。每个 expert 等待 SwiGLU ready随后复用 GMM1 的 tile、swizzle 和多级双缓冲 pipeline。expert 完成后实际参与的 8 或 24 个 AIC 分别发布 arrival供 Combine coordinator 汇总。在 gmm2.h 中可以看到扩组决策的完整机制PrimaryJoinDecision由 primary 组 localId 0 检查gmm1DoneSlotkMegaMoeFixedGmm1DoneMarker把“是否已 join join expert”编码写入gmm2JoinSlotHelperJoinExpert等待带 join 标志位的决策并解码出 join experthelper 组从该 expert 开始重建groupBase和轮转startCoreIdx未成功 join 时 primary 组在循环末尾发布哨兵值收尾。这也是伪码“动态扩组和阶段同步GM join/ready/arrival epoch”的源码落地。关键 PTO 接口Tile 绑定和搬运TASSIGN、TLOAD、TEXTRACTCube 计算TMATMUL、TMATMUL_ACCFixPipe scale 和输出TMOV、TSTORE_FP动态扩组和阶段同步GM join/ready/arrival epoch伪码primaryGroup physical AIC 16 .. 23 helperGroup physical AIC 0 .. 15 joined false helperGroup: joinExpert wait for join decision rebuild groupBase and rotating startCore at joinExpert for expert from (helper ? joinExpert : 0) to expertPerRank - 1: wait gmm2Ready[logicalAicId % 8] 2 * expert 2 if primary and not joined and expert gmm2JoinCheckStartExpert: coordinator checks gmm1Done and publishes a unified decision for this expert if decision joins at this expert: joined true activeAic 24 currentM clip(cumsumMM[lastSourceRank, expert], groupBase, maxOutputSize) for tile assigned to this AIC with rotating startCore: blockM, blockN swizzle_9_columns_snake_m(tile) for kTile in 0 .. N/2 step 512: TLOAD(A_l1[pingpong], gmPermutedToken[groupBase blockM, kTile]) TLOAD(B_l1[pingpong], weight2[expert, kTile, blockN]) for l0k in 0 .. 512 step 128: TEXTRACT(A_l0, A_l1, l0k) TEXTRACT(B_l0, B_l1, l0k) firstK ? TMATMUL(acc, A_l0, B_l0) : TMATMUL_ACC(acc, A_l0, B_l0) TMOV(fixpipeScale, scale2[expert, blockN]) TSTORE_FP(gmm2Output[groupBase blockM, blockN], acc, fixpipeScale) synchronize participating AICs each active AIC publishes gmm2Arrival[logicalAicId] 2 * expert 1 groupBase currentM注意 GMM2 的 reduction 维是N/2SwiGLU 输出维度输出维是Kscale2[expert, K]是 per-channel fixpipe scale随TSTORE_FP写 C而perTokenScale2留给 Combine 做逐 row 反量化。八、阶段 6Combine——8 个 AIV 按 source rank 写回并发布进度关键流程物理核16..23的 AIV0 组成 8-AIV Combine 组任务按 source rank 分配EP8 时每个 AIV 负责一个 source rankEP16 时每个 AIV 轮转负责两个。Combine 先等待combineStartAfterGmm2Expert指定的 GMM2 expert ready再从 expert 0 开始追赶配置值 0 表示等待 expert 0而不是跳过等待。每行gmm2Output转 fp32乘perTokenScale2反量化再转换为输出类型并写入 source rank 的offsetD。达到unpermutePhase1ReadyExpertCount时向各 source rank 发布已完成的 expert 数全部 expert 完成后发布最终 expert progress 和 DataReady。本卡 8 条 Combine lane 全部完成后清理 count window并放行后 16 个 Unpermute worker。在 combine.h 中ProcessDirectLargeSegmentRows逐 row 完成TLOAD(gmm2Output) - TCVT(fp32) - TMULS(perTokenScale2) - TCVT(CAST_RINT) - TSTORE(remote offsetD)的完整反量化写回链FinalizeRankStreamingLane在全部 Combine lane 完成后重置tokenPerExpert为下次 launch 复用 count window并通过PublishUnpermuteStartRangeMte(initialWorkerCount, helperCount, epoch)放行物理核 16..23 上的后 16 个 AIV。关键 PTO 接口GMM2 结果读取和远端写回TLOAD、TSTORE类型转换和反量化TCVT、TMULS跨 rank 进度通知TNOTIFY组间同步GM GMM2 arrival/Combine ready epoch伪码initialReadyExpert combineStartAfterGmm2Expert coordinator waits until all GMM2 producers finish initialReadyExpert all Combine lanes wait combineReady 2 * initialReadyExpert 2 groupBase 0 for expert in 0 .. expertPerRank - 1: if expert initialReadyExpert: coordinator waits until all 8-or-24 GMM2 producers finish expert wait combineReady[combineLocalId] 2 * expert 2 currentM cumsumMM[lastSourceRank, expert] for srcRank assigned to this Combine AIV: rows tokenPerExpert[srcRank, globalExpert(myRank, expert)] srcRowBase groupBase cumsum_before_source(srcRank, expert) dstRowBase preSumBeforeRank[srcRank, expert] for row in 0 .. rows - 1: c TLOAD(gmm2Output[srcRowBase row, 0:K]) fp32 TCVT(c, CAST_NONE) fp32 TMULS(fp32, perTokenScale2[srcRowBase row]) d TCVT(fp32, CAST_RINT) TSTORE(peer[srcRank].offsetD[dstRowBase row, 0:K], d) if expert 1 unpermutePhase1ReadyExpertCount: drain remote stores for srcRank assigned to this AIV: TNOTIFY(peer[srcRank].ExpertProgress, readyExpertCountexpert 1, opSet) groupBase currentM drain remote stores for srcRank assigned to this AIV: TNOTIFY(peer[srcRank].ExpertProgress, readyExpertCountexpertPerRank, opSet) TNOTIFY(peer[srcRank].DataReady, currentLaunchEpoch, opSet) publish localCombineDone[combineLocalId] Combine coordinator: wait all 8 Combine lanes done clear tokenPerExpert for next launch publish UnpermuteStart[32 : 48]九、阶段 7Unpermute——3216 AIV 两阶段恢复原 token 顺序关键流程Unpermute 不单独占核Dispatch/SwiGLU 所在物理核的 32 个 AIV 先加入Combine 所在物理核的 16 个 AIV 在本卡 Combine 完成后加入。第一阶段 coordinator 等待 Dispatch/SwiGLU 释放以及所有 producer rank 的 expert progress 达到unpermutePhase1ReadyExpertCount随后向前 32 个 worker 发布 start。第一阶段按 32 worker 切 token仅当一个 token 的所有有效 topK route 都已由对应 producer rank 写回时才处理该 token。前 32 个 worker 完成第一阶段后coordinator 等待所有 producer rank 全部完成并等待后 16 个 worker 可用然后向全部 48 个 worker 发布第二阶段 ready。第二阶段按 48 worker 重新切 token只处理第一阶段未满足 ready 条件的 token。每个 token 沿 K 维分块读取 topK 对应的offsetDrow乘路由权重后用 fp32 累加最后转换并写回。在 unpermute.h 中TokenReadyForExpertProgress精确实现“一个 token 的所有有效 topK route 都已由对应 producer rank 写回”的判定对每个 topK slot 解析globalExpert - producerRank / localExpert只要存在任何一条 route 的 localExpert 大于该 producer rank 的已就绪 expert 数则该 token 不在第一阶段处理ProcessRankStreamingRange(..., processPhase1)分别执行第一阶段32 worker 的 ready 子集与第二阶段48 worker 的剩余 tokenmetadataexpandedRowIdx与probs采用双 buffer 预取。关键 PTO 接口metadata 预取PtoLoadVectoroffsetD 读取和输出写回TLOAD、TSTORE类型转换、权重缩放和累加TCVT、TMULS、TADD组内进度读取/发布MTE 批量TLOAD/TSTORE GM epoch伪码phase1Coordinator: wait all Dispatch AIVs and SwiGLU AIVs released wait every producerRank.ExpertProgress unpermutePhase1ReadyExpertCount snapshot phase1ReadyExpertCount[producerRank] publish UnpermuteStart[0 : 32] first 32 workers: tokenRange split_tokens(M, workerCount32, workerId) for token in tokenRange: phase1Ready true for topkSlot in 0 .. topK - 1: globalExpert expertId[token, topkSlot] producerRank, localExpert split_global_expert(globalExpert) phase1Ready localExpert phase1ReadyExpertCount[producerRank] if phase1Ready: process_token(token) publish phase1Done[workerId] phase1Coordinator: wait phase1Done[0 : 32] wait every producerRank.ExpertProgress expertPerRank wait UnpermuteStart[32] # 后 16 个 AIV 已完成本卡 Combine publish phase2Ready[0 : 48] all 48 workers: wait phase2Ready[workerId] tokenRange split_tokens(M, workerCount48, workerId) for token in tokenRange: if token was not ready in phase 1: process_token(token) process_token(token): expandedRows PtoLoadVector(expandedRowIdx[token, 0:topK]) routeProbs PtoLoadVector(probs[token, 0:topK]) for col in 0 .. K step unpermuteTileCols: acc 0.0f for topkSlot in 0 .. topK - 1: row expandedRows[topkSlot] if row is valid: d TLOAD(offsetD[row, col : col unpermuteTileCols]) fp32 TCVT(d, CAST_NONE) weighted TMULS(fp32, routeProbs[topkSlot]) acc TADD(acc, weighted) outTile TCVT(acc, CAST_RINT) TSTORE(out[token, col : col unpermuteTileCols], outTile)十、关键 tiling 参数与 shape 配置以下配置对应K7168、N4096、topK8、expertPerRank16、worldSize8表中 expert ID 从 0 开始完整字段说明见 overview_v1.md 的“关键 tiling 参数”一节MGMM1 全 AIC expert 数SwiGLU AIV 数GMM2 扩组最早检查 expertCombine 启动前等待的 GMM2 expert162812332/64/5122161301282161331024/2048216132字段对应关系fullAicGmm1ExpertCountGMM1 开头使用全部 24 个 AIC 的 expert 数默认 2。swigluActiveGroupSize实际参与 SwiGLU 的 AIV 数M16 为 8其余为 16。gmm2JoinCheckStartExpert由(EP, M)配置 GMM2 最早从哪个 expert 开始检查 GMM1 是否结束成功扩组后后续 expert 不再重复决策。combineStartAfterGmm2ExpertCombine 启动前等待的 GMM2 expert ID。例如值为 2 时先等 GMM2 expert 2 ready再从 Combine expert 0 开始处理值为 0 表示等待 expert 0不跳过等待。Combine 数据量约为 Dispatch 的 2 倍延后启动可避免与 Dispatch 抢占同一段 HBM 带宽。unpermutePhase1ReadyExpertCountUnpermute 第一阶段启动所需的各 source rank 最小 expert 进度。第一阶段使用 32 个 AIV第二阶段扩展为全部 48 个 AIV。十一、源码结构、运行方式与性能参考目录结构算子主体位于kernels/manual/a2a3/dispatch_mega_combine/其中pseudocode.md本文主体七阶段数据流与伪码overview.md 与 overview_v1.md总体设计、overlap 优化原则、固定分组调度与性能对比README_zh.md算子说明、支持芯片Ascend910B1 / Ascend910B / Ascend910C 等、构建运行与 FAQop_kernel/dispatch_mega_combine.h主流程、front_reorder.h/front_fullload_sort.h/front_vms_sort.h、dispatch.h、gmm_common.h/gmm1.h/gmm2.h、swiglu.h、combine.h、unpermute.h以及utils/下的mega_expert_sync.hppepoch slot 同步、hccl_window.hppHCCL RDMA window、pto_gmm_preload_async_fixpipe_quant.hppGMM 多级 pipeline等 helpermain.cpp/kernel_launch.cpp/runtime_context.*/tiling_builder.*/data_utils.*/comm_mpi.hhost 侧入口、launch 包装、运行时与 HCCL window 管理、tiling 构造、数据生成与校验。构建与运行先配置 CANN 环境source cann-install/set_env.sh在kernels/manual/a2a3/dispatch_mega_combine/下执行./run.sh # 默认 worldSize8, M2048, K7168, N4096, topK8, expertPerRank16, maxOutputSize81940切换典型 M 档位例如bash run.sh --world-size 8 --m 512 --k 7168 --n 4096 --topk 8 --experts 16 --max-output-size 81940 --reuse-data常用约束详见 README_zh.mdK需满足 packed row、GMM1/GMM2 tile 与量化路径对齐N是 GMM1 输出维度SwiGLU 后进入 GMM2 的维度为N/2maxOutputSize必须覆盖单 rank 接收的 routed rows 上限。环境变量方面ASCEND_HOME_PATH由 CANNset_env.sh设置MPI_ENV_BIN/MPI_ENV_LIB/MPI_LIB_PATH需指向同一 MPI 环境HCCL_BUFFSIZE由run.sh按 case 自动抬高以容纳 remote window。内存布局要点HCCL remote window 承载跨 rank 可见的数据offsetAFrontReorder 写入 packed int8 token rowDispatch 拉取、offsetDCombine 写回、Unpermute 消费、tokenPerExpertcount-as-flag 的跨 rank count row、ExpertProgress / DataReady各 rank Combine 进度与 launch 完成通知workspace GM 则承载gmA、gmC、gmPermutedToken、gmm2Output、expandedRowIdx、cumsumMM / preSumBeforeRank等中间结果。性能参考overview_v1.md 记录了 A3 EP8 / EP16 下优化版本与基线的实测对比如 EP8 下 M2048 约4285.13us - 3756.88us、EP16 下 M2048 约4380.82us - 3714.16us大 case 优化幅度约 10%~15%小 case 基本持平或略有优化overview.md 同时给出 2048 case 的实际 overlap 实况——AIC 基本满载AIV 在 Dispatch/Combine 阶段受 HBM 带宽限制未用满进一步优化方向是降低 AIC/AIV 对 HBM 的竞争。十二、小结dispatch_mega_combine的核心价值在于把传统 MoE 中“重排 AlltoAllV grouped FFN”的多段串行流程折叠进单 kernel 的 expert 级流水AIV 侧的 FrontReorder / Dispatch / SwiGLU / Combine 与 AIC 侧的 GMM1 / GMM2 通过2*e1 / 2*e2的 epoch slot 单向通知衔接AIC 分组按24→16GMM1与8→24GMM2动态扩缩Unpermute 以 3216 两阶段追赶式启动。掌握七阶段的数据流、同步模型与gmm_common.h/mega_expert_sync.hpp/gmm2.h/combine.h/unpermute.h中的实现细节即可据此理解甚至复刻类似的 PTO 多阶段融合算子流水设计。【免费下载链接】pto-isaParallel Tile Operation (PTO) is a virtual instruction set architecture designed by Ascend CANN, focusing on tile-level operations. This repository offers high-performance, cross-platform tile operations across Ascend platforms.项目地址: https://gitcode.com/cann/pto-isa创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考