Mega-MoE 内核代码导读：从 Dispatch 到 WGMMA 的融合流水线

Mega-MoE 内核代码导读从 Dispatch 到 WGMMA 的融合流水线基于 DeepGEMMsm90_fp8_mega_moe.cuhHopper SM90与 SGLang Mega-MoE 路径整理。目标把教科书式的 MoE All-to-All 流程映射到真实 fused kernel 里的 warp 分工、同步原语与 GEMM pipeline。一、Mega-MoE 是什么传统 MoE 推理往往拆成多个 kernel 多次dist.all_to_allRouter 算 topk统计 send_countAll-to-All 交换 count分配 recv bufferAll-to-All 交换 token 数据Expert FFNgate_up SwiGLU downAll-to-All 把结果发回topk 加权合并Mega-MoE把这些阶段熔进一个 kernelSM100 上还有 FP4 等扩展在单卡内用NVLink 对称内存 TMA WGMMA完成跨 rank 通信与计算 overlap。核心文件Kerneldeep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh~1936 行Schedulerdeep_gemm/include/deep_gemm/scheduler/mega_moe.cuhSGLang 入口sglang/srt/layers/moe/mega_moe.py文件头注释里的五段流水线Dispatch pull → TMA load → WGMMA → Epilogue → Combine二、Rank、SM、Expert 是什么关系概念含义rank参与 Expert Parallelism 的一块 GPU不是整机一机多卡 多 rankSMGPU 上的一个 Streaming MultiprocessorMega kernel 通常每 SM 一个 blockkNumExpertsPerRank总 expert 数 / rank 数每个 rank 存一部分 expert 权重全局 expert 编号与 rank 的映射dst_rank_idx expert_idx / kNumExpertsPerRank local_expert expert_idx % kNumExpertsPerRank例kNumExpertsPerRank 4rank 0: global expert 0..3 → local 0..3 rank 1: global expert 4..7 → local 0..3Token 在哪、Expert 在哪要分开token 可在任意 rankgating 后若 expert 在别的 rank就 dispatch 过去目标 rank 用本地 expert 编号索引权重和 recv 表。三、Warp 分工一张表看懂整个 CTAKernel 入口constuint32_twarp_idxcutlass::canonical_warp_idx_sync();设N kNumDispatchWarps kNumDispatchThreads / 32kNumNonEpilogueThreads 128时 TMA 占 4 warp2 干活 2 占位Warp角色代码分支0 … N-1Dispatch计数、写索引、NVLink pullif (warp_idx kNumDispatchWarps)NTMA Load A SFAelse if (warp_idx kNumDispatchWarps)N1TMA Load B不搬 weight SFelse if (warp_idx kNumDispatchWarps 1)N2, N3空闲参与 warpgroup reg decelse if (warp_idx N kNumMMANonEpilogueWarps)N4 … endMath WarpgroupsWGMMA epilogue combineelseInit 阶段359–386 行另有一套一次性分工warp0 清 expert 计数warp1 init dispatch barrierwarp2 init full/empty/combine barrier。四、Dispatch 阶段详解4.1read_topk_idx并行遍历 token × topkDispatch 里的局部 lambda不是全局函数constautoread_topk_idx[](constautoprocess){for(uint32_ti(sm_idx*kNumDispatchWarpswarp_idx)*kNumTokensPerWarp;inum_tokens;ikNumSMs*kNumDispatchWarps*kNumTokensPerWarp){if(i(lane_idx/kNumTopk)num_tokensandlane_idxkNumActivateLanes){expert_idx__ldg(input_topk_idx_bufferi*kNumTopklane_idx);if(expert_idx0)process(i*kNumTopklane_idx,expert_idx);}__syncwarp();}};关键常量kNumTokensPerWarp 32 / kNumTopk例topk8 → 每 warp 一次处理 4 个 token。Lane 映射固定外层 token 起点i表达式含义lane_idx / kNumTopk本轮第几个 tokenlane_idx % kNumTopk该 token 的第几个 topk 槽i * kNumTopk lane_idx扁平下标token_topk_idx调用两次第一次计数第二次写索引。4.2 Step1统计 send_count 抢起始 slot// 第一次 read_topk_idx计数atomicAdd_block(smem_expert_countexpert_idx,1);// 全局 atomic打包 (132) | local_countsmem_expert_count[i]atomic_add(expert_send_count_ptr(i),send_value)的低32位;64 位send_value含义低 32 位本 SM 发给 expert i 的 token 数高 32 位固定 1又一个 SM 来报到例SM0 发 3 个SM1 发 2 个 → 全局(232)|5SM1 拿到的起始 slot 3。起始 slot 本 SM 在接收队列里从第几个位置开始写电影院抢座SM0 用 0,1,2SM1 用 3,4。4.3 Step2写 token 索引到远端dst_rank_idxexpert_idx/kNumExpertsPerRank;dst_slot_idxatomicAdd_block(smem_expert_countexpert_idx,1);dst_ptrget_src_token_topk_idx_ptr(local_expert,sym_buffer.rank_idx,dst_slot_idx);*sym_buffer.map(dst_ptr,dst_rank_idx)token_topk_idx;接收表布局[local_expert][source_rank][slot] → token_topk_idxPull 时反查去哪个 rank、拉哪个 token、哪个 topk 槽。4.4 grid_sync、SM0 广播、NVLink barriergrid_sync本 GPU 所有 SM 写完索引表后再继续SM0把 send_count 写到远端expert_recv_count/recv_count_sumnvlink_barrier跨 rank 握手对应 naive MoE 的Step2 统计 count和Step3 all_to_all(count)只是用对称内存点对点替代all_to_allAPI。4.5 Pull tokenStep5 数据交换533 行起按 recv_count 和索引表TMA 从源 rankpullFP8 hidden SF topk_weight写入本地l1_token_buffer。五、与 Naive MoE 流程对照Naive 步骤Mega-MoEStep2 count479–492 per-expert send_countStep3 all_to_all(count)508–528 SM0 广播 nvlink_barrierStep4 分配 recv预分配 L1 poolfetch_expert_recv_count决定实际 pull 数Step5 数据 all_to_all533 NVLink pullStep6 expert_ffn877–1796 L1/L2 WGMMA epilogueStep7 结果 all_to_all1755 L2 NVLink scatterStep8 weighted_sumL1 已 × topk_weight1852 COMBINE 累加六、L1 Linear1 与 L2 Linear2MoE expert FFNSwiGLUhidden ──[L1: gate_up]──► gate|up ── SwiGLU ──► intermediate intermediate ──[L2: down]──► hidden阶段权重GEMM 形状默认EpilogueLinear1l1_weightsKhidden, N2×intermediateSwiGLU × weight → FP8 → l2_actsLinear2l2_weightsKintermediate, NhiddenBF16 → NVLink scatter → combine bufferScheduler 状态机先一批 expert 跑完 L1 blocks再跑 L2 blockswave 交替。七、同步原语速查7.1__syncwarpvs__syncthreads__syncwarp32 lane 对齐用于 TMAelect_one、pull、topk 遍历__syncthreads整个 blockinit 后仅用一次bar.syncsync_alignedblock 内前 N 个 dispatch thread对齐mbarrierTMA ↔ Math 的 async pipelinewarpgroup_arrive4 warp 集体发 WGMMA 前的 fenceWarp 绑在一起执行但分支分化、单 lane 发 TMA、smem 共享仍需要显式 sync。7.2ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx)asmvolatile(bar.sync %0, %1;::r(barrier_idx),r(num_threads));只同步前 kNumDispatchThreads 个 threadDispatch 角色不必让 TMA/Math 空等。7.3full_barriers/empty_barriers经典生产者–消费者pipeline每个 stage 一块 smem屏障wait 谁arrive 谁语义fullMathTMAarrive_and_expect_tx(N)smem写满了可以读emptyTMAMathsmem用完了可以覆盖arrive 报到计数 1 声明 TMA 字节数wait 阻塞直到条件满足phase在 stage 绕回 0 时翻转配合 parity wait 区分多轮复用。八、process_math_block一个 GEMM tile 的一生Scheduler 对每个(expert, m_block, n_block, L1/L2)调用process_math_block(block_phase,local_expert_idx,num_k_blocks,m_block_idx,n_block_idx);结构A. 坐标 setupvalid_m、warpgroup 在 smem 里的 offset B. K 维 pipeline 循环 full.wait → 读 scale → WGMMA → empty.arrive → 融合 scale 到 final_accum C. Epilogue Linear1: SwiGLU FP8 quant TMA store Linear2: BF16 NVLink scatter8.1 典型 WGMMA 序列L1full_barriers[stage_idx]-wait(phase);// 等 TMA 填好 smemwarpgroup_fence_operand(accum[i]);// 寄存器对 WGMMA 可见warpgroup_arrive();// warpgroup 集体 fencefor(k0;kBLOCK_K/WGMMA::K;k)WGMMA::wgmma(desc_a,desc_b,accum,k);// 发 WGMMAwarpgroup_commit_batch();// 提交一批warpgroup_wait0();// 等算完if(lane_idx0)empty_barriers[stage_idx]-arrive();// 释放 stagefinal_accumscale_a*weight_sf*accum;// FP8 scale 融合8.2warpgroup_arrive是什么asmvolatile(wgmma.fence.sync.aligned;\n:::memory);Warpgroup 4 warp 128 threadWGMMA 的集体执行单位必须在 smem 读取之后、WGMMA 指令之前否则下一 pipeline stage 可能覆盖 smem或 warpgroup 内步调不一致配套原语fence_operand→arrive→wgmma→commit_batch→wait。L2 会把BLOCK_K128拆成两个 64因 L2 activation SF 是 per-64 粒度。九、端到端数据流一图流[Kernel 外] topk → input_topk_idx_buffer ↓ Dispatch: 计数 → 抢 slot → 写索引表 → grid_sync → SM0 广播 count → nvlink_barrier ↓ Dispatch: NVLink pull token/SF/weight → l1_token_buffer ↓ TMA warp: 异步 load A/B/SFA → smemfull/empty pipeline ↓ Math WG: WGMMA L1 → SwiGLU×weight → FP8 l2_acts ↓ Math WG: WGMMA L2 → BF16 scatter → combine_token_buffer远端 ↓ nvlink_barrier → COMBINE: 各 topk slot 累加 → 写 y十、学习建议先对照 naive MoE 四步count → 换 count → 换 data → ffn → 换回来 → sum再在 kernel 里找对应行号。分清三层并行rank跨 GPU、SM/block同 GPU、warp/laneblock 内。同步要分层记warp 内__syncwarp角色间bar.syncTMA/Math 间mbarrierWGMMA 前warpgroup_arrive。L1/L2 不是抽象概念就是 gate_up 和 down 两层 GEMMepilogue 不同。十一、关键代码索引主题约略行号sm90_fp8_mega_moe.cuhPipeline 注释108–123Init barriers359–386read_topk_idx463–477send_count slot479–501grid_sync503–506SM0 广播 recv508–521NVLink pull533–673TMA A/B load725–867process_math_block903–1770L1 epilogue SwiGLU1428L2 scatter1687–1765COMBINE1798–1925结语Mega-MoE 的难点不在于「又写了一个 GEMM」而在于把 MoE 的通信模式count、索引、pull、scatter、combine和 Hopper 的异步拷贝TMA、集体矩阵乘WGMMA、多级 barrier 绑在同一条流水线里。理解时建议始终带着两个问题读代码数据在哪个 rank 的哪块 buffer当前 warp 在生产还是消费 smem 的哪个 stage答清楚这两个问题大部分行号都能自己归类。