Moonshot AI 开源 MoonEP：面向 MoE 训练的完美平衡专家并行库

Moonshot AI 在 Kimi K3 开放日上正式开源了 MoonEP，这是一个专为分布式混合专家（MoE）工作负载设计的专家并行（EP）通信库。该库旨在使大规模专家并行通信更加高效，并以 MIT 许可证发布。

MoonEP 是 Kimi K3 开放日的一部分。除了 K3 模型权重和技术报告外，Moonshot 还发布了三个基础设施代码库：MoonEP、FlashKDA 和 AgentEnv。其中 FlashKDA 此前已开源，而 MoonEP 和 AgentEnv 则是本次新发布的。MoonEP 是声称将 Kimi K3 的扩展效率提升 2.5 倍的关键创新之一。Kimi K3 是一个拥有 2.8 万亿参数、原生视觉能力和 1M 令牌上下文窗口的 MoE 模型。

MoonEP 针对的问题：在专家并行中，路由器将每个令牌发送给其 top-K 个专家，这些专家位于不同的排名上。路由器很少是平衡的，一些专家接收的令牌远多于其他专家。该代码库使用 maxvio 来量化偏斜，定义为 max_e (T_e / T̄) − 1，其中 T_e 是路由到专家 e 的令牌数，T̄ 是完美平衡下的预期数量。maxvio 为 0 表示完全平衡。不平衡的成本是结构性的，而非偶然。集合的延迟由其最慢的参与者决定，因此最热的排名决定迭代时间。更糟的是，每个步骤的令牌计数都会变化，这些动态激活形状会碎片化 GPU 内存，并强制每层主机同步。

核心思想：动态冗余专家。MoonEP 的主要创新是一个硬性不变量：无论路由多么偏斜，每个排名接收恰好 S × K 个令牌——其中 S 是每个排名的输入令牌数，K 是每个令牌的 routed top-k。它通过直接从当前路由器的输出在线规划少量冗余专家来实现这一点。这些重复的专家在专家计算之前被预取。在反向传播中，它们的梯度被归约回其主排名。

该设计具有三个属性：完美平衡（通过在线规划的冗余专家实现 S × K 保证）、在线规划（一个近最优的 GPU 规划内核，开销可忽略，使用 CUTLASS CuTe DSL 实现，setup.py 锁定 nvidia-cutlass-dsl==4.4.2）、零复制和静态形状（融合的排列/反排列，令牌直接写入远程排名上的专家分组位置，缓冲区视图返回给计算，只需固定的 S × K 缓冲区，静态已知形状消除每层 MoE 主机同步）。

内存契约：MoonEP 与训练或推理框架的契约是具体的：每个专家投影一个连续的对称内存权重张量，加上一个规划器产生的 cu_seqlens。VM 组 GEMM 消耗一个单一的 [E+B, H, H'] 权重张量，其中 E 是总路由专家数，B 是每个排名的预取槽数，H 是隐藏大小，H' 是专家 FFN 中间大小。dispatch 返回的 cu_seqlens[E+B] 选择哪些专家行是活动的。连续性是一个硬性要求，因为组 GEMM 仅通过行索引寻址专家。布局清晰划分：行 [0, E) 保存所有排名的本地专家，每个排名 E/R 行，每个块物理上是主排名的参数内存，通过对称内存映射到所有地方；行 [E, E+B) 是本地预取槽，由 buffer.prefetch_weight 填充。预取槽来自一个所有层共享的进程全局池。这意味着额外内存成本是总投影的 B 个专家权重，而不是每层。

如何设置 B 取决于工作负载。训练必须使用 B = E/R，因为规划器每个排名最多从一个远程主组复制专家，该边界保证组 GEMM 接触的每个专家都是本地的。推理允许 B < E/R，README 推荐 B = 3–4。如果一个排名需要多于 B 的不同远程专家，组 GEMM 会通过对称映射直接从主排名读取溢出权重——稍慢，但不影响正确性。

训练中，权重布局在 fp32 中镜像，每个投影有一个 [E+B, H, H'] 梯度缓冲区。关键的是，行 [E, E+B) 由一个单独的归约缓冲区支持，而非参数梯度。重复专家的梯度是临时的，必须对框架自身的梯度归约不可见。每个排名将所有 R 个归约缓冲区映射为一个 [R, B, H, H'] 视图，然后 reduce_grad 从每个排名通过 NVLink 读取其自己专家的槽，累积到本地参数梯度，并清零已消耗的槽。

与 DeepEP v2 的基准测试：两个已发布的基准测试在 H20 上运行，EP=8，扫描路由器不平衡。比较脚本 benchmarks/bench_vs_deepep.py 默认使用 S=8192, E=384, H=7168, K=8, H'=2048, 32 个 SM，maxvio 目标为 0.2, 1, 10, 20。两个库接收来自共享种子的相同路由矩阵。报告了三个发现：零复制使原始通信更快，消除了困扰 DeepEP v2 后段的复制开销，因此 MoonEP 的通信时间在所有不平衡水平上都低于 DeepEP v2；完美平衡使 MoonEP 几乎不受偏斜影响，通信时间随 maxvio 增长几乎不变，而 DeepEP v2 的延迟由最热排名决定，稳定下降。

关键要点：MoonEP 声称每个 EP 排名接收恰好 S × K 个令牌，无论路由器偏斜如何；平衡来自在线在 GPU 上规划的冗余专家，然后在专家计算前预取；静态形状消除每层 MoE 主机同步，并阻止导致 DeepEP 内存溢出的碎片；训练需要 B = E/R 预取槽，推理可降至 B = 3–4 而无正确性成本；以 MIT 许可证在 Kimi K3 开放日发布，与 FlashKDA 和 AgentEnv 一同。