Mega-MoE 代码导读（下篇）：Pull、对称内存与 WGMMA Pipel

接续上篇warp 分工、Dispatch metadata、L1/L2、与 naive All-to-All 对照。本篇聚焦索引表与 Pull 的对应关系、对称内存 / rendezvous、L1 pool 握手、process_math_block与mbarrier / WGMMA协议。一、索引表写 vs Pull 读三维坐标是同一套1.1 写入494–501dst_rank_idxexpert_idx/kNumExpertsPerRank;dst_slot_idxatomicAdd_block(smem_expert_countexpert_idx,1);dst_ptrget_src_token_topk_idx_ptr(expert_idx%kNumExpertsPerRank,// local_expert在目标 rank 上sym_buffer.rank_idx,// source_rank发送方 本 rankdst_slot_idx);// slot*sym_buffer.map(dst_ptr,dst_rank_idx)token_topk_idx;维度变量含义[local_expert]expert % kNumExpertsPerRank目标 rank 上的 expert 编号[source_rank]sym_buffer.rank_idx谁发的[slot]dst_slot_idx该源 rank 发给此 expert 的第几个slot 按 (source_rank, expert) 独立编号rank0 有 0,1,2…rank1 也有自己的 0,1,2…。1.2 目标 rank 在哪不在表的第 4 维而在map的第二参数dst_rank_idxexpert_idx/kNumExpertsPerRank;*sym_buffer.map(dst_ptr,dst_rank_idx)...;get_src_token_topk_idx_ptr(...)按接收方 workspace layout算逻辑偏移map(ptr, dst_rank_idx)映到那台 GPU 的物理地址再写表物理上在 expert 所在 rank 的 workspace里所以不需要再存 dst_rank。1.3 Pull 读614–615src_token_topk_idx*get_src_token_topk_idx_ptr(current_expert_idx,// local_expertcurrent_rank_in_expert_idx,// source_ranktoken_idx_in_rank);// slot与写入一一对应。Round-robin 只负责pool 第 k 个位置 → 该查[local][哪个 source_rank][哪个 slot]。二、Pull Token 详解533 行起对应 naive MoEStep5 数据 All-to-All实现是接收方主动 pull不是dist.all_to_all。2.1 前置条件索引表、recv_count 已写好nvlink_barrier完成531 行sync_unaligned(Dispatch Math)两边对齐后再填/用 L1 pool2.2 资源绑定pull_buffersmem_send_buffers[warp_idx];// smem 中转pull_mbarrierdispatch_barriers[warp_idx];scheduler.fetch_expert_recv_count();// 每个 local expert 收几个 token2.3 工作分配for(token_idxsm_idx*kNumDispatchWarpswarp_idx;;token_idxkNumSMs*kNumDispatchWarps)token_idx 本 rank所有 received token 的全局序号跨 expert 排成长队。2.4 按 expert 分段548–567例expert0 收 100 个expert1 收 10 个expert2 收 30 个token_idx: 0..99 → expert0 100..109 → expert1 110..139 → expert2变量含义token_idx_in_expert该 expert 内第几个expert_pool_block_offset前面 expert 占几个 M-blockceil_div(tokens, BLOCK_M)pool_token_idxL1 pool 线性下标 offset * BLOCK_M token_idx_in_expert换 expert 时加载stored_rank_count[j] rank j 给当前 local expert发几个 token。2.5 Round-robin 选源 rank569–612目的pool 里 token 在源 rank 间交替排列避免 M 维上连续 token 都来自同一 rank。算法每轮取length min(remaining)每个还有货的 rank 各贡献length个 slot重复直到分完。例rank0→3rank1→5rank2→2共 10 个pool 序号source_rankslot0–50,1,2 各 2 个各 0,16–7rank0, rank1各 28–9rank13,4offset含义跳过完整轮次后offset length表示各 rank下一 token 从 slotoffset开始。token_idx_in_expert 7walkthrough第 1 轮num_round_tokens67≥6 →slot_idx1offset022第 2 轮num_round_tokens212 → rank1token_idx_in_rank 2(1/2)2查表[local_expert][rank1][slot2]2.6 三次搬运 通知 Math① TMA load远端 input_token_buffer[src_token_idx] → smem pull_buffer ② 并行 copy SF → l1_sf_buffer ③ 读 topk_weight → l1_topk_weights_buffer ④ TMA store → l1_token_buffer[pool_token_idx] ⑤ 写 TokenSrcMetadata {rank, token, topk}L2 scatter 用 ⑥ l1_arrival_countMath 侧 spin 等到够数再 TMA load A三、Dispatch 与 Math 对齐L1 pool 何时可填531 ↔ 901同一 barrierkDispatchWithEpilogueBarrierIdx侧时机Dispatch 531metadata 做完pull 前Math 901进scheduler.for_each_block前含义共享l1_token_buffer、l1_arrival_count等对齐后Dispatch开始填 poolMath开始等/读 pool按 M-block handshake不必等全部 pull 完676 ↔ 1809结束配对Math 算完后 Dispatch 才能清 workspace四、对称内存、rendezvous、map4.1 对称内存是什么各 rank同 layout、不同 base。buffersymm_mem.empty(num_bytes,...)handlesymm_mem.rendezvous(buffer,groupgroup)# kernel: handle.buffer_ptrs → 各 rank 基址表offsets[i]peer_base[i]-local_base;map(ptr,dst)ptroffsets[dst];≈ NVSHMEM 的「对称堆 相同 offset」不是表里的某一维。4.2 Rendezvous 是什么集体握手、交换各 rank buffer 的 GPU 物理地址建立 P2P/NVLink 寻址。之前只知道本机 base之后map(local_ptr, peer)能访问 peer 对称堆这是setup 阶段不是传 token 数据。之后还有group.barrier()保证全员就绪。4.3 为何不用nvshmem_float_pMega-MoE 用PyTorchsymm_mem 手写map操作做法小 metadata*map(ptr, dst) val、atomic_add_sys大块 tokentma_load_1d/tma_store_1d从 mapped 远端地址同步自研nvlink_barrier、mbarrier非 NVSHMEM quiet/barriermap≈nvshmem_ptrrendezvous 后直接 * / TMA 合法不要求 NVSHMEM 库 API。五、process_math_block与 WGMMA 协议5.1 结构setup 坐标valid_m、warpgroup 在 smem 的 offset ↓ K 维 pipelinefull.wait → 读 scale → WGMMA → empty.arrive → 融合 scale ↓ EpilogueL1 SwiGLU×weightFP8 / L2 BF16 scatter5.2 典型 WGMMA 序列L1full_barriers[stage]-wait(phase);warpgroup_fence_operand(accum[i]);warpgroup_arrive();// wgmma.fence.sync.alignedfor(k)WGMMA::wgmma(...);warpgroup_commit_batch();warpgroup_wait0();empty_barriers[stage]-arrive();// lane 0final_accumscale_a*weight_sf*accum;5.3 五步协议含义步骤作用fence_operand寄存器操作数对 WGMMA 可见arrivewarpgroup4 warp对齐smem 读必须在此之前wgmma异步集体 MMAcommit_batch把本批 WGMMA 划为一组无则wait无目标 → 死锁或脏读wait0等该批算完5.4full/emptymbarrier屏障waitarrive语义fullMathTMAarrive_and_expect_tx(N)smem满了可读emptyTMAMathsmem空了可覆盖arrive 报到 full 侧声明 TMA 字节数wait 阻塞直到条件满足phasestage 绕回时翻转parity 区分多轮复用与 WGMMA 的commit/wait分工full/empty管 TMA↔smemcommit/wait管 WGMMA↔寄存器。六、下篇知识串联单 token 生命周期[写索引] rank0 上 token 选中 global expert 5在 rank1 → map(..., dst_rank1) 写 rank1 表[local1][src0][slotk] [Pull] rank1 上 round-robin 得到 (src0, slotk) → 查表 → TMA 从 rank0 pull hidden/SF/weight → 写 l1_token_buffer[pool_idx]metadatal1_arrival_count [Math] full/empty pipeline WGMMA L1 → SwiGLU×weight → L2 WGMMA → map scatter 回源 rank combine buffer [Combine] 各 topk slot 累加 → y七、易混点速查问题答案表里有 dst_rank 吗没有表在接收方dst 由map(..., dst_rank_idx)指定token_idxvstoken_idx_in_expert前者跨 expert 全局序号后者 expert 内序号offset在 round-robin 里已跳过轮次中每个 rank 用掉的 slot 数累加lengthrendezvous vs nvlink_barrier前者地址 setup启动时后者kernel 内跨 rank 握手为何 531 要对齐 Math同时启动 pool 生产者–消费者避免竞态commit_batch 能省吗不能WGMMA 异步批次必须 commit 才能 wait八、代码索引下篇主题约略行号写索引 map494–501Pull 前 sync531Pull 主循环533–673expert 分段548–567round-robin569–612Math 侧 sync901process_math_block903–1770L2 scatter metadata1755–1765full/empty init374–379COMBINE1798–1925结语上篇回答「MoE 通信阶段在 kernel 里哪一段」下篇回答「表、对称内存、pool、GEMM pipeline 如何咬合」。读 Pull 时始终带三个问题查哪一格表map 到哪台 rankpool 第几行读 Math 时带两个问题full/empty 到哪一 stageWGMMA 批 commit 了吗把这两个链条理清Mega-MoE 后半段基本就通了。