The Problem
Mixture-of-Experts (MoE) training suffers from load imbalance: when routing skews, some ranks receive far more tokens than others, making iteration time track the hottest rank. DeepEP v2 degrades as imbalance grows and OOMs at high skew due to memory fragmentation from dynamic shapes.
What This Does
MoonEP is an Expert Parallelism communication library that keeps token loads perfectly balanced across ranks via dynamic redundant experts. Every rank receives exactly S × K tokens regardless of routing skew. A planning kernel (moonep/planning.py) computes which redundant experts to prefetch, and the dispatch/combine kernels (moonep/dispatch.py, moonep/combine.py) move tokens directly to their final expert-grouped positions—no permute-in/permute-out, no comm-buffer to user-buffer copy.
The repo is 30 Python files plus two CUDA sources (csrc/bindings.cu, csrc/nvl_shared_buffer.cuh). Core logic lives in moonep/; benchmarks under benchmarks/; tests under tests/.
How It Is Wired
Execution starts at dispatch in moonep/api.py:685, which reaches 33 functions and is called from 4 places. The API layer (moonep/api.py, 23 functions) orchestrates planning, dispatch, combine, prefetch, and gradient reduction. moonep/planning.py is the hub: 41 functions, imported by 11 other files, carrying the widest blast radius. The internal call graph shows 429 resolved edges; init_case (called from 13 places) and _require_ctx (11 places) are the most-coupled helpers.
The only external effect is filesystem I/O: benchmarks/bench_comm.py:497 (main) reaches 68 functions and touches the filesystem via create_nvl_single_owner_tensor → _exchange_ipc_fds, using shutil.rmtree. No database or network effects are mapped.
File-by-file: moonep/buffer.py handles NVLink distributed tensor creation and IPC fd exchange; moonep/planning.py owns the online planning kernel; moonep/combine.py and moonep/dispatch.py handle communication; moonep/prefetch.py and moonep/grad_reduce.py manage weight prefetch and backward gradient reduction.
How To Use It
Setup: Install via pip install . from the repo root (uses setup.py). Requires an NVIDIA GPU with CUDA support.
Configuration: No environment variables documented. The library is integrated by providing a contiguous weight tensor [E+B, H, H'] per expert projection and using the planner-produced cu_seqlens returned by dispatch.
Running it: No CLI entry point. Integration is programmatic—call dispatch from moonep/api.py in your training loop. Benchmarks run directly: python benchmarks/bench_vs_deepep.py. The README documents the integration contract but no standalone runnable example.
Real-World Use
In an MoE training loop per layer: router outputs → planning computes redundant experts → prefetch fetches weights → dispatch sends tokens to expert-grouped positions on remote ranks → expert GEMM consumes the single contiguous weight tensor → combine gathers outputs → backward reduces gradients via grad_reduce. Static shapes mean no per-layer host synchronization.
Code Health & Issues
Static analysis found 27 findings (13 high, 14 medium) across 6 kinds:
- High – Deep nesting (max depth 10) in
moonep/planning.py,moonep/buffer.py,moonep/dispatch.py; duplicated code blocks (211 repeated 6-line blocks across 15 files, includingbenchmarks/bench_grad_reduce.py,tests/test_grad_reduce.py,moonep/combine.py). - Medium – Oversized files (
moonep/planning.pyat 1094 lines); high branching density intests/kernel_test_utils.py(119 branches over 365 lines); unclosed file handle and broad exception handling inbenchmarks/bench_comm.py.
SDLC observations: no CI/CD pipeline exists—every change merges without a build or test run. Tests are present (11 files) but not gated. License is present; no Dockerfile or lockfile.
The Bottom Line
MoonEP's perfect-balance approach is a sound answer to MoE imbalance, and the zero-copy design is genuinely faster than DeepEP v2 per the benchmarks. The codebase needs refactoring (nesting, duplication) and CI before production use. Suitable for teams running MoE training on NVIDIA GPUs who accept the integration contract.