Moonshot AI has released MoonEP, an open-source library designed to improve how expert parallelism handles communication in large Mixture-of-Experts models. The code is available under an MIT license.
The announcement arrived during Kimi K3 Open Day. Alongside the K3 model weights and a technical report, the team published three infrastructure projects: MoonEP, FlashKDA, and AgentEnv. FlashKDA had already been released previously. MoonEP and AgentEnv are new additions. This library contributes to a reported 2.5× gain in scaling efficiency for Kimi K3. That model contains 2.8 trillion parameters, supports vision, and accepts a context window of one million tokens.
The problem MoonEP targets
In expert parallelism, a router directs each token to its top-K experts. These experts live on different processing ranks. The router rarely distributes work evenly. Some experts receive far more tokens than others.
The repository measures this unevenness using maxvio. The formula is max_e (T_e / T̄) − 1. Here, T_e represents tokens routed to expert e, and T̄ is the expected count under perfect balance. A maxvio of 0 indicates perfect balance.
Imbalance causes structural costs rather than incidental delays. A collective operation’s latency depends on its slowest participant. Therefore, the busiest rank dictates the iteration time. Token counts per rank change every step. These dynamic activation shapes fragment GPU memory and force per-layer host synchronization.
The core idea: dynamic redundant experts
MoonEP enforces a hard invariant. Every rank receives exactly S × K tokens, regardless of how skewed the routing is. S is the input tokens per rank, and K is the routed top-k per token.
The system achieves this by planning a small number of redundant experts online, directly from current router outputs. Those duplicated experts are prefetched before expert computation begins. During the backward pass, their gradients reduce back to their home ranks.
The design relies on three properties:
- Perfect balance: The
S × Kguarantee holds via online-planned redundant experts. - Online planning: A near-optimal GPU planning kernel with negligible overhead. It uses the CUTLASS CuTe DSL;
setup.pypinsnvidia-cutlass-dsl==4.4.2. - Zero copy and static shapes: Fused permute/unpermute operations. Tokens write directly into their expert-grouped positions on remote ranks. Buffer views return to the computation. Only a fixed
S × Kbuffer is needed. Statically known shapes eliminate per-layer MoE host synchronization.
The memory contract
MoonEP’s agreement with a training or inference framework is specific. It requires 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. E is the 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/Rrows 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 bybuffer.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 to user-buffer copy that dominates the epilogue. Consequently, 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. DeepEP v2 — whose latency is set by the hottest rank — degrades steadily.
What it means
Developers building MoE systems gain a tool that removes the bottleneck of uneven token distribution. By enforcing a fixed token count per rank through redundant experts, MoonEP stabilises iteration time. This prevents the slowest rank from dragging down the entire batch. The static shapes also stop memory fragmentation, reducing the risk of out-of-memory errors during training.
