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. window.addEventListener("message",function(e){ if(e.data&&e.data.mepFrame==="mep1"){ var f=document.getElementById("mep1-frame"); if(f) f.style.height=e.data.height+"px"; } }); 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, becau