Moonshot AI Open-Sources MoonEP: A Perfectly Balanced Expert Parallelism Library for MoE Training

Moonshot AI has open-sourced MoonEP, an Expert Parallelism (EP) communication library for distributed Mixture-of-Experts (MoE) workloads. The team announced the release as a library built to make expert-parallel communication more efficient at scale. It ships under an MIT license.

MoonEP arrived as part of Kimi K3 Open Day. Alongside the K3 model weights and technical report, Moonshot released three infrastructure codebases: MoonEP, FlashKDA, and AgentEnv. FlashKDA had already been open-sourced; MoonEP and AgentEnv were published with this release. MoonEP is one of the innovations behind a claimed 2.5× improvement in scaling efficiency for Kimi K3, a 2.8-trillion-parameter MoE model with native vision and a 1M-token context window.

The problem MoonEP targets

In expert parallelism, a router sends each token to its top-K experts, which live on different ranks. Routers are rarely balanced. Some experts get far more tokens than others.

The repository quantifies skew with maxvio, defined as max_e (T_e / T̄) − 1, where T_e is tokens routed to expert e and T̄ is the expected count under perfect balance. A maxvio of 0 means perfectly balanced.

Imbalance costs are structural, not incidental. A collective’s latency is set by its slowest participant, so the hottest rank determines iteration time. Worse, token counts per rank change every step. Those dynamic activation shapes fragment GPU memory and force per-layer host synchronization.

The core idea: dynamic redundant experts

MoonEP’s main mention is a hard invariant. Every rank receives exactly S × K tokens, no matter how skewed the routing is — where S is input tokens per rank and K is routed top-k per token.

It achieves this by planning a small number of redundant experts online, directly from the current router outputs. Those duplicated experts are prefetched before expert computation. In the backward pass, their gradients are reduced back to their home ranks.

The design is set up into three properties:

Perfect balance: the S × K guarantee above, via online-planned redundant experts.

Online planning: a near-optimal GPU planning kernel with negligible overhead. It is implemented in the CUTLASS CuTe DSL; setup.py pins nvidia-cutlass-dsl==4.4.2.

Zero copy and static shapes: fused permute/unpermute. Tokens are written directly into their expert-grouped positions on remote ranks, and buffer views are returned to the computation. Only a fixed S × K buffer is needed, and statically known shapes eliminate per-layer MoE host synchronization.

The interactive explainer below computes the resulting buffer and prefetch-pool sizes live from a config you control.

The memory contract

MoonEP’s contract with a training or inference framework is specific: one contiguous symmetric-memory weight tensor per expert projection, plus a planner-produced cu_seqlens. The VM group GEMM consumes a single [E+B, H, H’] weight tensor, where E is total routed experts, B is prefetch slots per rank, H is hidden size, and H’ is the expert FFN intermediate size. The cu_seqlens[E+B] returned by dispatch selects which expert rows are active.

Contiguity is a hard requirement, because the group GEMM addresses experts purely by row index. The layout splits cleanly:

Rows [0, E) hold all ranks’ local experts, E/R rows per rank. Each chunk physically is the home rank’s parameter memory, mapped everywhere via symmetric memory.

Rows [E, E+B) are local prefetch slots, filled by buffer.prefetch_weight.

The prefetch slots draw from a process-global pool shared by all layers. That detail matters: the extra memory cost is B expert weights per projection in total, not per layer.

How you set B depends on the workload. Training must use B = E/R, because the planner duplicates experts from at most one remote home group per rank. That bound guarantees every expert the group GEMM touches is local. Inference allows B < E/R, and the README recommends B = 3–4. If a rank needs more distinct remote experts than B, the group GEMM reads overflow weights straight from the home rank through the symmetric mapping — slightly slower, with no impact on correctness.

Training mirrors the weight layout in fp32 with a [E+B, H, H’] grad buffer per projection. Critically, rows [E, E+B) are backed by a separate reduce buffer, not by the parameter grads. Duplicated experts’ gradients are temporary and must stay invisible to the framework’s own grad reduce. Each rank maps all R reduce buffers as one [R, B, H, H’] view, then reduce_grad reads its own experts’ slots from every rank over NVLink, accumulates into the local parameter grad, and zeroes consumed slots.

Benchmarks against DeepEP v2

Both published benchmarks run on H20 with EP=8, sweeping router imbalance. The comparison script benchmarks/bench_vs_deepep.py defaults to S=8192, E=384, H=7168, K=8, H’=2048, and 32 SMs, with maxvio targets of 0.2, 1, 10, and 20. Both libraries receive an identical routing matrix from a shared seed.

Three findings were reported on their github page. Zero copy makes raw communication faster by eliminating the comm-buffer → user-buffer copy that dominates the epilogue, so MoonEP’s comm time sits consistently below DeepEP v2 at every imbalance level. Perfect balance makes MoonEP nearly immune to skew: its comm time stays almost flat as maxvio grows, while DeepEP v2 — whose latency is set by the hottest rank — degrades steadily.

Key Takeaways

MoonEP claims every EP rank receives exactly S × K tokens, regardless of router skew.

Balance comes from redundant experts planned online on-GPU, then prefetched before expert compute.

Static shapes remove per-layer MoE host sync and stop the memory fragmentation that OOMs DeepEP.

Training requires B = E/R prefetch slots; inference can drop to B = 3–4 with no correctness cost.

Released MIT under Kimi K3 Open Day, alongside FlashKDA and AgentEnv.

Sources: MoonshotAI/MoonEP on GitHub and @Kimi_Moonshot announcement