mini-infer系统实战-20-MoE 推理主线：关键不在 Router，而在 Dispatch 与 All-to-All

做完 Phase 16 之后，mini-infer 在 dense 模型推理这条线上已经比较完整了：调度、KV cache、attention kernel、TP、MLA、PD、量化都已经有真实实现和 benchmark。接下来如果还想继续往“大厂推理 / AI infra 核心组”那条能力线靠，Sparse MoE 和 Expert Parallelism 基本绕不过去。

但真开始做的时候，很容易掉进一个误区：把 MoE 理解成“加一个 router，然后让每个 token 选 top-k expert”。这件事只说对了一半。对推理系统来说，router 只是入口，真正决定系统复杂度的是后面的三件事：

1. token 副本怎么按 expert 和 rank 重排
2. rank 之间怎么做 all-to-all
3. benchmark 到底在测 ideal EP，还是在测你当前 prototype 的真实通信实现

Phase 17 的目标，就是把这几件事做成一个能跑、能测、能解释的闭环，而不是停留在“我实现了一个 MoE demo”。

问题定义

这一阶段我想回答的其实是两个非常具体的问题：

1. 如果先写一个单卡 dense MoELayer 作为 oracle，能不能把 2 卡 EPMoELayer 的 dispatch / gather / combine 路径做成数值严格对齐？
2. 在一个最小但真实的 synthetic workload 上，2 卡 EP 能不能明确快过 1 卡 dense？

这背后还有一个更重要的约束：benchmark 口径必须干净。

如果 dense 和 EP 各自随机初始化一套权重，最后拿 throughput 去比，那个对比没有意义。
如果 EP 的吞吐里包含 mp.spawn 和 init_process_group，而 dense 只测 steady-state forward，这个对比也没有意义。
如果实现里用的是 padded all_to_all_single，但报告里写的是 ideal EP bytes，这个通信量结论同样站不住。

所以这篇文章不会把重点放在 router 数学本身，而是放在推理系统里真正容易做错的那部分。

方案设计

这一版 Phase 17 的方案分成三层。

1. 先用 dense MoELayer 建一个数值 oracle

[moe_layer.py](https://github.com/psmarter/mini-infer/blob/main/mini_infer/moe_layer.py) 里的 MoELayer 故意没有做任何复杂优化。它的逻辑非常直接：

• router 先对每个 token 做 top-k
• 把 top-k 分数重新归一化
• 对每个 expert 显式挑出属于它的 token
• 跑 expert FFN
• 再按 router weight 把多个 expert 输出加权合并回原 token

这条路径看起来“笨”，但它的价值恰恰在于简单。后面所有 EP 版本的正确性，都是拿它来对。

对应的 router 输出和统计结构也一起固定下来：

• RouterOutput
• RoutingStats
• DispatchLayout
• EPForwardAux

这一步很重要，因为后面 benchmark 里要解释 send_counts、expert_loads、expert_score_sums，如果一开始没有把这些结构收成稳定接口，后面很容易一边实现一边漂。

2. 把“选了哪个 expert”变成真正可通信的 dispatch layout

真正的系统问题不在 top-k，而在 top-k 之后。

一个 token 选了两个 expert，不等于系统已经知道：

• 这个 token 的两个副本该发到哪两个 rank
• 这些副本在发送 buffer 里的顺序应该是什么
• 返回结果以后怎么按原 token 顺序 gather 回去

我在 [moe_layer.py](https://github.com/psmarter/mini-infer/blob/main/mini_infer/moe_layer.py) 里把这一步显式做成了 build_dispatch_layout()。它生成的不是抽象“路由结果”，而是一份真正能用于通信的数据布局：

layout = DispatchLayout(
    token_indices=token_indices,
    expert_ids=expert_ids,
    local_expert_ids=local_expert_ids,
    weights=weights,
    dest_ranks=dest_ranks,
    send_counts=send_counts,
)

这几个字段里，最关键的是三项：

• token_indices：后面 gather 时知道结果该加回哪个原 token
• dest_ranks：知道每个 token 副本要发往哪个 rank
• send_counts：知道每个 rank 实际要收多少 token 副本

也正是从这一步开始，MoE 不再只是一个“模型结构问题”，而变成了标准的推理系统通信问题。

3. EPMoELayer 先做功能正确，再明确承认它还是 prototype

[moe_layer.py](https://github.com/psmarter/mini-infer/blob/main/mini_infer/moe_layer.py) 里的 EPMoELayer 是这一阶段的核心。

它做的事情是：

• src_rank 持有输入 hidden states
• 用 router 算出 top-k 路由
• 按 DispatchLayout 重排 token 副本
• 用 all_to_all_single 把 token 副本发到对应 rank
• 每个 rank 只执行自己负责的 expert 子集
• 再把 expert 输出通过一次 return-path all_to_all_single 发回 src_rank
• src_rank 用原来的 token_indices + weights 做 combine

不过这版实现有一个非常明确的工程取舍：我没有继续强行追“理想变长 split-size all-to-all”，而是收回到了 padded all_to_all_single + valid mask。

原因很现实。前一版如果按真实 send_counts 去做变长 split，为了构造 split sizes，就必须在 per-layer forward 里把 CUDA 上的 count tensor 拉回 CPU，代码里会出现 .item() / .tolist() 这种 host sync。对功能来说这没问题，但对 steady-state EP benchmark 来说，这是污染计时窗口的硬伤。

所以最后我选的是：

• 牺牲一部分通信效率
• 保住 per-layer forward 没有 host sync
• 然后在 benchmark 里明确区分：
  • ideal EP bytes
  • prototype padded bytes

这不是“偷懒”，而是把 prototype 的真实边界说清楚。

实现细节

Dense 路径：先把 oracle 写死

MoELayer.forward() 的核心是对每个 expert 做显式 token 选择：

for expert_id, expert in enumerate(self.experts):
    token_idx, slot_idx = torch.where(route.expert_indices == expert_id)
    if token_idx.numel() == 0:
        continue
    expert_in = flat_x.index_select(0, token_idx)
    expert_out = expert(expert_in)
    expert_weight = route.expert_weights[token_idx, slot_idx].to(expert_out.dtype).unsqueeze(-1)
    combined.index_add_(0, token_idx, expert_out * expert_weight)

这段代码一点也不“快”，但它把数学路径钉得很死。后面 review 里只要看到 EP 输出和它不一致，问题范围就会一下子缩小很多。

EP 路径：真正困难的是通信生命周期

EPMoELayer._forward_distributed() 的结构可以概括成四段：

1. src_rank 构造 send_hidden / send_expert_ids / send_valid
2. forward all_to_all_single
3. 每个 rank 只对本地 expert 子集执行 _apply_experts()
4. return-path all_to_all_single，然后 src_rank 做 combine

代码上最重要的设计约束有两个。

第一，每个 rank 当前仍持有完整 expert 权重。
这是刻意保守的第一版原型。这样做没有真实 expert-sharding 的内存收益，但可以先把路由、dispatch、通信和 combine 路径做对。

第二，通信层优先避免 per-layer host sync。
这也是为什么最终版本没有继续保留变长 split-size 的原因。对于推理系统 benchmark 来说，“prototype 的真实通信量更大”是可以接受的；但“benchmark 被 host sync 污染，还假装自己测的是纯 GPU steady-state”不可以。

EPEngine：把计时放到 worker 内部

[ep_engine.py](https://github.com/psmarter/mini-infer/blob/main/mini_infer/ep_engine.py) 复用了 Phase 13 的 mp.spawn + file:// rendezvous + NCCL 约定，但这里有一个很容易被忽略的点：吞吐计时不能包住 mp.spawn 和进程组初始化。

如果把这些开销都算进来，EP benchmark 基本只会测出“多进程启动很慢”，这不是真正想回答的问题。

所以这版 EPEngine 做了两件事：

• worker 启动后先 warmup
• 真正的 elapsed_s 在 worker 内部、steady-state forward 周围测

最终 benchmark 输出里也明确写成：

ep_note=2-GPU EPEngine steady-state worker timing; spawn/init excluded

这让 dense 和 EP 的吞吐至少处在同一统计口径上。

Benchmark：必须同时打印理想公式和 prototype 公式

[benchmark_moe.py](https://github.com/psmarter/mini-infer/blob/main/benchmarks/benchmark_moe.py) 最后稳定下来的一个关键点，是通信量摘要函数 build_comm_summary()。

它不会只打印一个“EP bytes”，而是同时输出三套量：

• tp_bytes_per_layer
• ep_ideal_bytes_per_layer
• ep_prototype_bytes_per_layer

对应公式也会一起打印：

• 2 * num_tokens * hidden_size * dtype_bytes
• 2 * num_tokens * top_k * hidden_size * dtype_bytes
• 2 * ep_size * num_tokens * top_k * hidden_size * dtype_bytes

这样写的意义很大。它强迫我在写 benchmark 结论时明确承认：

• 理想 EP 通信量应该是多少
• 当前 prototype 因为 padded all-to-all，多付出了多少 hidden-state bytes

这会让结果更“丑”一点，但工程上更可信。

实验结果

下文使用的数据来自该阶段的正式 benchmark 记录。

环境：

• Ubuntu 24.04
• torch 2.1.2+cu121
• RTX 4090 × 2
• conda run -n ai-infra

正式 workload：

• batch_size=4
• seq_len=16
• hidden_size=512
• intermediate_size=1024
• num_experts=8
• top_k=2
• dtype=float16
• warmup=2
• runs=5
• src_rank=1

命令是：

conda run -n ai-infra python benchmarks/benchmark_moe.py \
  --compare \
  --batch-size 4 \
  --seq-len 16 \
  --hidden-size 512 \
  --intermediate-size 1024 \
  --num-experts 8 \
  --top-k 2 \
  --dtype float16 \
  --warmup 2 \
  --runs 5 \
  --src-rank 1

两次正式复跑的结果是：

指标	Run 1	Run 2	平均
dense throughput	22730.35 tok/s	22513.92 tok/s	22622.14 tok/s
EP throughput	42552.69 tok/s	43004.13 tok/s	42778.41 tok/s
EP / dense	1.872x	1.910x	1.891x
max abs diff	0.000000	0.000000	0.000000

辅助统计也比较稳定：

• send_counts = [69, 59]
• expert_loads = [20, 12, 23, 14, 19, 15, 14, 11]

通信量口径是：

• TP bytes / layer = 131072
• EP ideal bytes / layer = 262144
• EP prototype bytes / layer = 524288

这组结果说明了两件事。

第一，当前 EPMoELayer 和 dense oracle 的数学路径是干净对齐的，max_abs_diff = 0.000000 已经足够说明问题。
第二，即使当前 prototype 仍然使用 padded all-to-all，它在这个 synthetic workload 上仍然拿到了 1.891x 的吞吐提升，说明 EP 主链路确实已经形成了可测量的并行收益。

这阶段真正踩到的坑

坑 1：变长 split-size all-to-all 看起来更优雅，但会把 host sync 带进热路径

这是整个 Phase 17 里我认为最有代表性的工程取舍。

一开始我其实做过按真实 send_counts 组织 split sizes 的版本。它表面上更接近理想 EP，因为不会 padding 那些空位。但只要 split sizes 依赖 CUDA count tensor，最终就很容易在 per-layer forward 里出现 .item() / .tolist()，把 GPU 计数同步回 CPU。

对 correctness 来说这没问题。
对 benchmark 来说，这是不可接受的。

最后我选择退回 padded all-to-all，并在 benchmark 里显式区分 ideal bytes 和 prototype bytes。这个决定的代价是 prototype 通信量更大，收益是 benchmark 的计时窗口更干净。

坑 2：如果 dense 和 EP 不是同权重、同输入，对照结果没有意义

这个问题一开始看起来很低级，但实际很常见。

如果 dense benchmark 和 EP benchmark 各自随机初始化一套 MoELayer，或者输入 hidden states 不是同一组，那么最后 throughput、expert_loads 甚至 max_abs_diff 都没有真正可比性。

所以 benchmark_moe.py 最后专门收成了 --compare 路径，用同一个 shared layer 和同一组 hidden states 同时喂给 dense 和 EP。只有这样，max_abs_diff 才真的是在验证同一件事。

坑 3：dense 的计时窗口也会被你自己写坏

后面 review 里还暴露出一个很典型的问题：dense benchmark 的计时窗口里一度混进了 output.cpu()、expert_loads.cpu()、expert_score_sums.cpu()，而 EP 的 steady-state 计时是在 worker 内部完成的，不包含这些 D2H copy。

如果不把这些 copy 挪出计时窗口，dense baseline 会被额外惩罚，最后得到的 “EP 更快” 结论其实是脏的。

这个 bug 修完之后，Phase 17 的 benchmark 才真正具备说服力。

总结

Phase 17 真正交付的，不是“一个有 router 的 MoE 层”，而是一版 synthetic MoE + 2-GPU Expert Parallelism 的完整闭环：

• dense MoELayer 作为数值 oracle
• DispatchLayout 把 top-k 路由变成真正可通信的数据布局
• EPMoELayer 打通 dispatch / all-to-all / gather / combine
• EPEngine 提供 2 卡最小可运行原型
• benchmark_moe.py 明确区分 ideal EP 与 prototype padded bytes
• 正式 benchmark 下，2 卡 EP 达到 1.891x dense，且 max_abs_diff = 0.000000

当然，这还不是最终生产实现。当前版本仍有明确边界：

• 每个 rank 仍持有完整 expert 权重
• 通信仍是 padded all_to_all_single
• benchmark 还是 synthetic layer-level workload
• 只验证了 ep_size=2

但从学习项目的角度看，这一阶段已经足够回答一个关键问题：MoE 推理最难的部分，不是 router，而是把 router 之后的 dispatch、通信、benchmark 和工程口径都做对。

下一步如果继续往前走，最值得做的三件事会是：

1. 真正做 expert weight shard，验证内存收益
2. 去掉 padded all-to-all，逼近 ideal EP 通信口径
3. 把 MoE / EP 接到更完整的生成或服务链路里

---

延伸阅读

• AI-Infra学习之旅-RoPE、RMSNorm、GQA与MQA、SwiGLU与MoE

---

系列导航

• 导读：mini-infer系统实战-00-导读：从最小推理链路到 MoE Expert Parallel 的项目路线
• 上一篇：mini-infer系统实战-19-量化推理：把 Linear 改成 INT8 之后，问题才刚开始
• 下一篇：mini-infer系统实战-21-Expert Sharding：真正的收益不只是“把权重切一半”