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.

=1073741824) return (b/1073741824).toFixed(b/1073741824<10?3:2)+’ GiB’; if(b>=1048576) return (b/1048576).toFixed(b/1048576<10?2:1)+’ MiB’; if(b>=1024) return (b/1024).toFixed(1)+’ KiB’; return b+’ B’; } function stat(l,v,u){return ”} var CMD={invariant:’moonep –explain invariant’,memory:’moonep –plan-memory’, maxvio:’moonep –sweep maxvio’,pillars:’moonep –arch’}; function render(){ var c=cfg(); warn.className = (c.E && c.R && c.E % c.R !== 0) ? ‘warn on’ : ‘warn’; var B = c.M===’train’ ? (c.R? Math.floor(c.E/c.R):0) : 4; line.textContent = CMD[cur] + (cur===’memory’ ? ‘ –mode ‘+c.M : ”); var h=””; if(cur===’invariant’){ var tok=c.S*c.K, buf=tok*c.H*2; h+=’// every rank receives exactly S × K tokens, regardless of routing skew‘; h+=’tokens_per_rank = S × K = ‘+num(c.S)+’ × ‘+c.K+’ = ‘+num(tok)+’‘; h+=’dispatch_buffer = S × K × H × 2B (bf16) = ‘+bytes(buf)+’ // fixed, per rank‘; h+=’‘+stat(‘tokens / rank’,num(tok),’constant at any maxvio’) +stat(‘dispatch buffer’,bytes(buf),’bf16, statically shaped’) +stat(‘host syncs / MoE layer’,’0′,’shapes known ahead of time’)+’‘; h+=’Because the shape is static, MoonEP needs only a fixed S × K buffer and removes the per-layer MoE host synchronization that dynamic shapes force. Activation shapes never change, so memory does not fragment across iterations.‘; } if(cur===’memory’){ var per=B*c.H*c.P*2, three=per*3, red=B*c.H*c.P*4; h+=’// prefetch slots live in rows [E, E+B) of one contiguous [E+B, H, H′] weight tensor‘; h+=’B = ‘+(c.M===’train’ ? ‘E / R = ‘+num(c.E)+’ / ‘+c.R+’ = ‘+B+’ // training REQUIRES B = E/R‘ : ‘4 // inference: B < E/R allowed, 3–4 recommended‘); h+=’‘; h+=’‘+stat(‘prefetch slots B’,num(B),c.M===’train’?’forced by gradient locality’:’tunable, 3–4 recommended’) +stat(‘per projection’,bytes(per),’bf16, process-global pool’) +stat(‘gate + up + down’,bytes(three),’total, NOT per layer’)+’‘; if(c.M===’train’){ h+=’reduce_buffer (local, fp32) = B × H × H′ × 4B = ‘+bytes(red)+’, mapped across ranks as one [R, B, H, H′] view.‘; h+=’Duplicated experts’ grads land in the reduce buffer, not in the parameter grads, so they stay invisible to the framework’s own grad reduce. reduce_grad then pulls each rank’s own slots back over NVLink and zeroes them for the next microbatch.‘; } else { h+=’If a rank ever needs more distinct remote experts than B, the group GEMM reads the overflow weights straight from the home rank through the symmetric mapping. Slightly slower — correctness is unaffected.‘; } h+=’// the pool is shared by all layers, so the extra cost above is paid once per process, not once per MoE layer‘; } if(cur===’maxvio’){ var total=c.S*c.K*c.R, mean=c.E? total/c.E : 0, tgts=[0.2,1,10,20]; h+=’// maxvio = max_e ( T_e / T̄ ) − 1 · 0 means perfectly balanced‘; h+=’total routed tokens = S × K × R = ‘+num(total)+’ · T̄ = ‘+mean.toFixed(1)+’ tokens/expert‘; h+=’<table><tr><th>maxvio</th><th>hottest expert</th><th>MoonEP tokens/rank</th><th>MoonEP behaviour</th></tr>‘; for(var i=0;i<tgts.length h=""><td class=""f"">‘+tgts[i]+’</td><td class=""n"">‘+num(Math.round((1+tgts[i])*mean))+’</td><td class=""n"">‘+num(c.S*c.K)+’</td><td>flat</td>‘; } h+=’</tgts.length></table>‘; h+=’These four points are the sweep run in the repo’s H20, EP=8 benchmark. Per the README, MoonEP’s comm time stays almost flat as maxvio grows, while DeepEP v2 — whose latency is set by the hottest rank — degrades steadily and eventually OOMs end-to-end at high imbalance.‘; h+=’// hottest-expert counts above are derived from the maxvio definition, not measured timings‘; } if(cur===’pillars’){ h+=’[1] perfect balance — a small number of redundant experts is planned online from the current router outputs and prefetched before expert compute; their gradients are reduced back to their home ranks in the backward pass.‘; h+=’[2] online planning — a near-optimal GPU planning kernel with negligible overhead, written in the CUTLASS CuTe DSL. Planning and prefetch kernels are already counted inside MoonEP’s published bars.‘; h+=’[3] zero copy + static shapes — fused permute/unpermute. Tokens are written directly into their expert-grouped positions on remote ranks, and views of the comm buffer are handed to the computation, eliminating the comm-buffer → user-buffer copy that dominates the epilogue.‘; h+=’‘+stat(‘devices’,’NVIDIA GPU’,’Zhenwu PPU under review’) +stat(‘license’,’MIT’,’commercial use permitted’) +stat(‘interconnect’,’NVLink’,’symmetric memory / VMM’)+’‘; h+=’// inspired by DeepEP, Echo and UltraEP — see the repo acknowledgments‘; } out.innerHTML=h; resize(); } var btns=document.querySelectorAll(‘.cmd’); for(var i=0;i<btns.length btns="" for="" j="0;j<btns.length;j++)" this.classname="cmd on" cur="this.getAttribute('data-c');" render="" var="" ins="document.querySelectorAll('.f" input="" select="" k="0;k<ins.length;k++)" function="" resize="" el="document.getElementById('mep1');" if="" window.parent.postmessage="" window.addeventlistener=""> </btns.length>

Source link