Mixture-of-Kittens: NVL72 上的 MoE 训练 Megakernel
一句话结论
面向 Nvidia NVL72 单跳高带宽 scale-up 域，把 token dispatch / shared+routed expert FFN / combine 融成确定性 megakernel；相对最强公开基线吞吐最高 2.37×（MXFP8 forward），生产栈 512 GPU / 多 GB300 NVL72 端到端训练吞吐 1.41×。
动机
- 加速器向 scale-up（NVL72→144/576/1152）收敛：域内像大内存，不再像 scale-out 网络；为 RoCE/IB 优化的 MoE 系统迁到 NVL72 时常 慢于 朴素 PyTorch+NCCL。
- MoE 在大训中占一半以上时间，但既有重叠/调度假设在单跳胖互联上失效（文中约半数配置下朴素基线反超）。
对照 LLM Distributed Training Collectives、NVSwitch、Weave、FlashMoE。
方案
- 按算子选 push/pull：scale-up 上 pull 与 push 同级；pull dispatch + push combine 等组合把信令开销压到 <1%，调度开销最高 2.2× 更小，并共享一张调度表。
- 可调重叠粒度：吞吐最优 minibatch 约 512–32,768 tokens；相对固定劣粒度最高 3.53×。
- 消除 CPU–GPU 同步：NVL72 集成 CPU 上常见 PyTorch 主机路径最高 2.97× 更慢；设备侧定长 ring buffer + 反向 replay，相对过度分配平均仅 +1.7% 延迟。
- 生产特性：确定性、FSDP RDMA 重叠、MXFP8 融合量化、router 梯度融合、可调 SM 分区。评测形状：Kimi K2.7 / GLM 5.2 / Qwen3.5-397B-A17B / DeepSeek V4 Pro。
效果（仅论文数字）
| 指标 | 数字 |
|---|---|
| vs 最强公开基线（MXFP8 fwd / bwd） | 最高 2.37× / 1.78× |
| vs 最强公开基线（BF16 fwd / bwd） | 最高 1.92× / 1.58× |
| 生产 e2e（512 GPU，多 NVL72，tok/s/GPU） | 1.41×（相对 DeepEP 前代） |
| 调度开销 vs Comet 类 | 1.4–2.2× 更小；占 MoE runtime 2.1–7.7% |
与 wiki 的关系
- NVSwitch Scale-Up Fabric — NVL72 域内 EP 通信代价重写
- LLM Distributed Training Collectives — MoE All-to-All / megakernel 重叠轴
- Weave / FlashMoE — 同类 SM 分区与单核融合；本文专攻 NVL72 训练