MoonshotAI/MoonEP

MoonshotAI/MoonEP

Description: MoonEP: A Perfectly Balanced Expert Parallelism Library via Dynamic Redundant Experts

Language: Python

License: MIT

Stars: 626

Forks: 62

Open issues: 5

Created: 2026-07-24T09:07:21Z

Pushed: 2026-07-28T01:37:11Z

Default branch: master

Fork: no

Archived: no

README:

MoonEP

MoonEP is an Expert Parallelism communication library that keeps token loads perfectly balanced across ranks via dynamic redundant experts.

Notation: S = input tokens per rank, K = routed top-k per token.

1. Perfect balance: every rank receives exactly S × K tokens, no matter how skewed the routing is. A small number of redundant experts is planned online from the current router outputs and prefetched before expert computation; their gradients are reduced back to their home ranks in the backward pass. 2. Online planning: a near-optimal GPU planning kernel with negligible overhead 3. Zero copy and static shapes: fused permute/unpermute — tokens are sent directly to 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.

Performance

Both benchmarks run on H20 with EP=8, sweeping the router imbalance:

$$\text{maxvio} = \max_e \left( \frac{T_e}{\bar{T}} \right) - 1$$

where $T_e$ is the number of tokens routed to expert $e$, and $\bar{T}$ is the expected tokens per expert under perfect balance (maxvio = 0 means perfectly balanced).

Communication vs DeepEP v2 ([benchmarks/bench_vs_deepep.py](benchmarks/bench_vs_deepep.py)):

• Zero copy makes raw communication faster: tokens are written directly to their final expert-grouped positions on remote ranks — no permute in, no permute out — and views of the communication buffer are handed straight to the computation, eliminating the comm-buffer → user-buffer copy that dominates the epilogue. MoonEP's comm time is consistently below DeepEP v2 at every imbalance level.
• Perfect balance makes it immune to imbalance: MoonEP's comm time stays almost flat as maxvio grows, while DeepEP v2 — whose latency is set by the hottest rank — degrades steadily.
• The comparison counts MoonEP's extra kernels: MoonEP adds planning and weight-prefetch kernels that DeepEP does not need, and they are already stacked in the bars above. Even with the whole critical path included, total dispatch time is on par with DeepEP v2's dispatch alone and pulls ahead under imbalance, while combine is significantly faster at every level.

End-to-end training:

• DeepEP degrades with imbalance: the hottest ranks receive more tokens, so iteration time climbs steadily as maxvio grows; meanwhile the ever-changing activation shapes fragment GPU memory, until training OOMs at high imbalance.
• MoonEP is unaffected: every rank always computes exactly S × K tokens per layer, so iteration time stays flat at every imbalance level; fully static memory shapes mean no fragmentation, and training never OOMs.

Supported Devices

• NVIDIA GPU
• Zhenwu PPU (under review, coming soon)

Usage

Integration

Notation: S = input tokens per rank, K = routed top-k per token, E = total routed experts in the EP group, R = number of EP ranks (EP comm size), B = weight prefetch slots per rank, NvS = dispatched token slots per rank (S × K real tokens plus per-VM-group padding), H = hidden size, H' = expert FFN intermediate size.

MoonEP's contract with a training or inference framework is 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; cu_seqlens[E+B] (returned by dispatch) selects which expert rows are active for the current step.

Weight buffer

For each expert projection (gate/up/down), every layer holds one contiguous VMM range [E+B, H, H'], identically laid out on every rank. Contiguity is a hard requirement: the group GEMM addresses experts purely by row index.

• Rows `[0, E)`: 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)`: local prefetch slots, filled by buffer.prefetch_weight; the planner points duplicated experts' token segments at these slots via cu_seqlens. Their physical memory comes from a process-global pool shared by all layers, so the extra cost is B expert weights per projection in total, not per layer.

How to set B.

• Training: must use `B = E/R` — the planner duplicates experts from at most one remote home group per rank (≤ E/R experts), so every expert the group GEMM touches is local.
• Inference (prefetch only, no gradients): `B

Training mirrors the weight layout in fp32: one contiguous [E+B, H, H'] grad buffer per projection.

• Rows `[0, E)`: the owner ranks' parameter grads.
• Rows `[E, E+B)`: prefetch-slot grads, backed by a separate reduce buffer, not by the parameter grads — duplicated experts' grads are temporary and must stay invisible to the framework's own grad reduce. The physical memory comes from a process-global pool shared across layers, like the prefetch pool.
• Reduce buffer: every rank maps all R reduce buffers as one [R, B, H, H'] view. reduce_grad lets each rank read the slots holding its own experts' grads from every rank's reduce buffer (remote reads over NVLink), accumulate them into its local parameter grad, then zero its own consumed slots for the next microbatch.

API walkthrough

from moonep import Buffer

buffer = Buffer(S=4096, H=7168, K=8, E=256, num_ep_ranks=8,
num_sms=32, token_padding=128)

• num_sms=None defaults to 32. B defaults to E // num_ep_ranks; an explicit value like B=4 may also be passed.
• dispatch / combine / prefetch_weight / reduce_grad all accept async_finish=True to run on the comm stream and return a CUDA event.

dispatch fwd

hidden_nvsh, route_weights_nvs, cu_seqlens, plan = buffer.dispatch(
hidden_sh, # [S, H] bf16
route_weights_sk, # [S, K] fp32
topk_experts_sk, # [S, K] int32
tokens_per_expert, # [E] int32, local count
)
# hidden_nvsh: [NvS, H] bf16 — dispatched tokens in physical VM group order
# route_weights_nvs: [NvS] fp32
# cu_seqlens: [E+B] int32 — padded token end offset per VM group row...