MoonEP: balanced expert parallelism with dynamic redundant experts

What MoonEP fixes about expert parallelism

In a standard expert-parallel MoE layer, the router decides where tokens go and the communication layer obeys. If one expert is hot, the rank hosting it receives more tokens than its peers, and every other rank waits for the slowest one. The README frames this with a metric it calls maxvio, the maximum ratio of tokens routed to an expert over the expected count under perfect balance, minus one. DeepEP degrades as maxvio grows because latency is set by the hottest rank. MoonEP's claim is different: every rank receives exactly S x K tokens regardless of routing skew, where S is input tokens per rank and K is the routed top-k per token.

The mechanism behind that claim is redundancy. A small number of redundant experts is planned online from the current router outputs and prefetched before expert computation, and their gradients are reduced back to their home ranks in the backward pass. The audience is narrow and specific: engineers integrating a communication layer into an existing MoE training or inference stack, who can satisfy the library's buffer contract and who care more about flat iteration time under imbalance than about a packaged end-user experience.

The weight buffer contract and why B matters

MoonEP's contract with a 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'] tensor, where E is total routed experts in the EP group, B is weight prefetch slots per rank, H is hidden size and H' is the expert FFN intermediate size. cu_seqlens[E+B] selects which expert rows are active for the step. Contiguity is stated as a hard requirement because the group GEMM addresses experts purely by row index.

Rows [0, E) hold all ranks' local experts, E/R rows per rank, and each chunk physically is the home rank's parameter memory mapped everywhere through symmetric memory. Rows [E, E+B) are local prefetch slots filled by buffer.prefetch_weight, and the planner points duplicated experts' token segments at those slots through 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 rather than per layer. That pooling detail is the part worth reading twice: it is what keeps redundancy affordable as layer count grows.

How B is set is not a free tuning knob. The README states that training must use B = E/R, because the planner duplicates experts from at most one remote home group per rank, so every expert the group GEMM touches is local. Inference, where prefetching happens without gradients, allows B < E/R with B = 3-4 recommended. If a rank needs more distinct remote experts than B, the group GEMM reads the overflow weights straight from the home rank through the symmetric mapping, addressed by cu_seqlens. That path is described as slightly slower with no impact on correctness.

Installing MoonEP and running a first dispatch

There are no published releases. setup.py builds the CUDA extension moonep._C from csrc/bindings.cu into the moonep package, and its docstring names the editable install as the preferred route, with a quick in-place build as the alternative. The only declared runtime dependency is nvidia-cutlass-dsl==4.4.2, and the extension links against cuda.

pip install -e .

After the build, the shared object lands at moonep/_C.<abi-tag>.so and the package can be imported as from moonep import _C. For a build that skips installation entirely, the docstring gives this command instead:

python setup.py build_ext --inplace

Once the extension imports, the first real use is constructing a Buffer, then dispatching one layer's tokens. The README's example passes S=4096, H=7168, K=8, E=256, num_ep_ranks=8, num_sms=32 and token_padding=128, which is a realistic MoE shape rather than a toy:

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 defaults to 32 when passed as None, and B defaults to E // num_ep_ranks, so the constructor above implicitly sets B = 32. An explicit value such as B=4 may also be passed. The dispatch call takes the hidden states, route weights, top-k expert indices and the local per-expert token counts, and returns the dispatched tokens in physical VM group order, the route weights, the cu_seqlens offsets and a MoonEPCommPlan object that must be saved for the prefetch step:

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
)

dispatch, combine, prefetch_weight and reduce_grad all accept async_finish=True to run on the comm stream and return a CUDA event.

Gradient buffers, reduce_grad and the fp32 mirror

Training mirrors the weight layout in fp32: one contiguous [E+B, H, H'] grad buffer per projection. Rows [0, E) are the owner ranks' parameter grads. Rows [E, E+B) are prefetch-slot grads, and the README is explicit that these are backed by a separate reduce buffer rather than the parameter grads, because duplicated experts' grads are temporary and must stay invisible to the framework's own grad reduce. That separation is the design decision most likely to trip an integrator who assumes all grad rows land in the same place.

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 over NVLink, accumulate them into its local parameter grad, then zero its own consumed slots for the next microbatch. The zeroing is part of the API's expected usage, not an optimization, so a caller that skips it will accumulate across microbatches. Nothing in the README documents a rollback or recovery path if a reduce is interrupted mid-step.

Where the design costs you, and when to pick DeepEP instead

MoonEP adds planning and weight-prefetch kernels that DeepEP does not need. The README says the comparison charts already stack those kernels into MoonEP's bars, and that total dispatch time is on par with DeepEP v2's dispatch alone while pulling ahead under imbalance. That is a fair way to present it, but it also means the extra work is real and only pays off once routing skew is large enough. At near-zero maxvio the two approaches are closer than the headline suggests.

DeepEP remains the better choice if you want a communication layer without a weight-layout contract. DeepEP does not require one contiguous symmetric-memory tensor per projection, does not require a planner-produced cu_seqlens, and does not require you to reason about prefetch slots. MoonEP's balance guarantee is bought with those obligations. There is also a hardware boundary: supported devices are listed as NVIDIA GPU, with Zhenwu PPU under review and coming soon, so on non-NVIDIA accelerators MoonEP is not yet an option.

Licence and the cost of tracking master

MoonEP is MIT licensed, which permits commercial use, modification and redistribution provided the copyright notice and permission notice are included. That is a permissive arrangement, but the licence text governs, so read LICENSE rather than this summary.

The upgrade cost is the more practical concern. setup.py pins version 0.0.1 and the only dependency is pinned exactly at nvidia-cutlass-dsl==4.4.2, so a CUTLASS DSL bump is a code change rather than a version-range resolution. The extension compiles with -std=c++20, --expt-extended-lambda and --expt-relaxed-constexpr, which ties builds to toolchains new enough to support them. With no releases retrieved, there is no versioned artifact to pin against: integrators track the master branch, and the last push was on 2026-08-13. Budget for rebuilding the extension whenever you move your CUDA or PyTorch stack.

Editorial conclusion

Adopt MoonEP if you run MoE training on NVIDIA GPUs with an expert-parallel group and you already control the weight layout your framework hands to the GEMM, because the library requires one contiguous symmetric-memory weight tensor per projection plus a planner-produced cu_seqlens. Do not adopt it if you need a supported release, an inference-only path with gradients disabled but no reduce_grad wiring, or a non-NVIDIA backend today, since Zhenwu PPU is listed as under review. Before committing, verify three things: that nvidia-cutlass-dsl==4.4.2 resolves for your toolchain, that your projection weights can be laid out as a single [E+B, H, H'] range on every rank, and that training uses B = E/R, because a smaller B is documented only for inference.

Frequently asked questions

Official sources

1. Issues
2. License: MIT
3. MoonshotAI/MoonEP on GitHub
4. README