企业官网建设流程全解析

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

基于 DeepGEMMsm90_fp8_mega_moe.cuh（Hopper SM90）与 SGLang Mega-MoE 路径整理。
目标：把教科书式的 MoE All-to-All 流程，映射到真实 fused kernel 里的 warp 分工、同步原语与 GEMM pipeline。

一、Mega-MoE 是什么

传统 MoE 推理往往拆成多个 kernel + 多次dist.all_to_all：

1. Router 算 topk
2. 统计 send_count
3. All-to-All 交换 count
4. 分配 recv buffer
5. All-to-All 交换 token 数据
6. Expert FFN（gate_up + SwiGLU + down）
7. All-to-All 把结果发回
8. topk 加权合并

Mega-MoE把这些阶段熔进一个 kernel（SM100 上还有 FP4 等扩展），在单卡内用NVLink 对称内存 + TMA + WGMMA完成跨 rank 通信与计算 overlap。

核心文件：

• Kernel：deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh（~1936 行）
• Scheduler：deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh
• SGLang 入口：sglang/srt/layers/moe/mega_moe.py

文件头注释里的五段流水线：

Dispatch pull → TMA load → WGMMA → Epilogue → Combine

二、Rank、SM、Expert 是什么关系

全局 expert 编号与 rank 的映射：

dst_rank_idx = expert_idx / kNumExpertsPerRank local_expert = expert_idx % kNumExpertsPerRank

例：kNumExpertsPerRank = 4

rank 0: global expert 0..3 → local 0..3 rank 1: global expert 4..7 → local 0..3

Token 在哪、Expert 在哪要分开：token 可在任意 rank；gating 后若 expert 在别的 rank，就 dispatch 过去；目标 rank 用本地 expert 编号索引权重和 recv 表。

三、Warp 分工：一张表看懂整个 CTA

Kernel 入口：

constuint32_twarp_idx=cutlass::canonical_warp_idx_sync();

设N = kNumDispatchWarps = kNumDispatchThreads / 32，kNumNonEpilogueThreads == 128时 TMA 占 4 warp（2 干活 + 2 占位）：

Init 阶段（359–386 行）另有一套一次性分工：warp0 清 expert 计数，warp1 init dispatch barrier，warp2 init full/empty/combine barrier。

四、Dispatch 阶段详解

4.1read_topk_idx：并行遍历 token × topk

Dispatch 里的局部 lambda，不是全局函数：

constautoread_topk_idx=[&](constauto&process){for(uint32_ti=(sm_idx*kNumDispatchWarps+warp_idx)*kNumTokensPerWarp;i<num_tokens;i+=kNumSMs*kNumDispatchWarps*kNumTokensPerWarp){if(i+(lane_idx/kNumTopk)<num_tokensandlane_idx<kNumActivateLanes){expert_idx=__ldg(input_topk_idx_buffer+i*kNumTopk+lane_idx);if(expert_idx>=0)process(i*kNumTopk+lane_idx,expert_idx);}__syncwarp();}};

关键常量：kNumTokensPerWarp = 32 / kNumTopk（例：topk=8 → 每 warp 一次处理 4 个 token）。

Lane 映射（固定外层 token 起点i）：

调用两次：第一次计数，第二次写索引。

4.2 Step1：统计 send_count + 抢起始 slot

// 第一次 read_topk_idx：计数atomicAdd_block(smem_expert_count+expert_idx,1);// 全局 atomic：打包 (1<<32) | 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 个 → 全局(2<<32)|5；SM1 拿到的起始 slot = 3。

起始 slot= 本 SM 在接收队列里从第几个位置开始写（电影院抢座：SM0 用 0,1,2；SM1 用 3,4）。

4.3 Step2：写 token 索引到远端

dst_rank_idx=expert_idx/kNumExpertsPerRank;dst_slot_idx=atomicAdd_block(smem_expert_count+expert_idx,1);dst_ptr=get_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_idx
Pull 时反查：去哪个 rank、拉哪个 token、哪个 topk 槽。

4.4 grid_sync、SM0 广播、NVLink barrier

• grid_sync：本 GPU 所有 SM 写完索引表后再继续
• SM0：把 send_count 写到远端expert_recv_count/recv_count_sum
• nvlink_barrier：跨 rank 握手

对应 naive MoE 的Step2 统计 count和Step3 all_to_all(count)，只是用对称内存点对点替代all_to_allAPI。

4.5 Pull token（Step5 数据交换）

533 行起：按 recv_count 和索引表，TMA 从源 rankpullFP8 hidden + SF + topk_weight，写入本地l1_token_buffer。

五、与 Naive MoE 流程对照

六、L1 Linear1 与 L2 Linear2

MoE expert FFN（SwiGLU）：

hidden ──[L1: gate_up]──► gate|up ── SwiGLU ──► intermediate intermediate ──[L2: down]──► hidden

Scheduler 状态机：先一批 expert 跑完 L1 blocks，再跑 L2 blocks，wave 交替。

七、同步原语速查

7.1__syncwarpvs__syncthreads

• __syncwarp：32 lane 对齐；用于 TMAelect_one、pull、topk 遍历
• __syncthreads：整个 block；init 后仅用一次
• bar.sync（sync_aligned）：block 内前 N 个 dispatch thread对齐
• mbarrier：TMA ↔ Math 的 async pipeline
• warpgroup_arrive：4 warp 集体发 WGMMA 前的 fence

Warp 绑在一起执行，但分支分化、单 lane 发 TMA、smem 共享仍需要显式 sync。

7.2ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx)

asmvolatile("bar.sync %0, %1;"::"r"(barrier_idx),"r"(num_threads));

只同步前 kNumDispatchThreads 个 thread（Dispatch 角色），不必让 TMA/Math 空等。

7.3full_barriers/empty_barriers

经典生产者–消费者pipeline（每个 stage 一块 smem）：

• 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. 坐标 setup（valid_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 scatter

8.1 典型 WGMMA 序列（L1）

full_barriers[stage_idx]->wait(phase);// 等 TMA 填好 smemwarpgroup_fence_operand(accum[i]);// 寄存器对 WGMMA 可见warpgroup_arrive();// warpgroup 集体 fencefor(k=0;k<BLOCK_K/WGMMA::K;++k)WGMMA::wgmma(desc_a,desc_b,accum,k);// 发 WGMMAwarpgroup_commit_batch();// 提交一批warpgroup_wait<0>();// 等算完if(lane_idx==0)empty_barriers[stage_idx]->arrive();// 释放 stagefinal_accum+=scale_a*weight_sf*accum;// FP8 scale 融合

8.2warpgroup_arrive是什么

asmvolatile("wgmma.fence.sync.aligned;\n":::"memory");

• Warpgroup = 4 warp = 128 thread，WGMMA 的集体执行单位
• 必须在 smem 读取之后、WGMMA 指令之前
• 否则下一 pipeline stage 可能覆盖 smem，或 warpgroup 内步调不一致

配套原语：fence_operand→arrive→wgmma→commit_batch→wait。

L2 会把BLOCK_K=128拆成两个 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 → smem（full/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

十、学习建议

1. 先对照 naive MoE 四步（count → 换 count → 换 data → ffn → 换回来 → sum），再在 kernel 里找对应行号。
2. 分清三层并行：rank（跨 GPU）、SM/block（同 GPU）、warp/lane（block 内）。
3. 同步要分层记：warp 内__syncwarp，角色间bar.sync，TMA/Math 间mbarrier，WGMMA 前warpgroup_arrive。
4. L1/L2 不是抽象概念，就是 gate_up 和 down 两层 GEMM，epilogue 不同。

十一、关键代码索引

结语

Mega-MoE 的难点不在于「又写了一个 GEMM」，而在于把 MoE 的通信模式（count、索引、pull、scatter、combine）和 Hopper 的异步拷贝（TMA）、集体矩阵乘（WGMMA）、多级 barrier 绑在同一条流水线里。理解时建议始终带着两个问题读代码：数据在哪个 rank 的哪块 buffer？当前 warp 在生产还是消费 smem 的哪个 stage？答清楚这两个问题，大部分行号都能自己归类。