MoonEP routing flow — live
An interactive model of one MoonEP MoE layer. Set the shape, drag the router skew, and watch:
how the router loads each expert, how the planner's greedy balance repairs the imbalance,
which expert weights cross which link, and what an analytic cost model estimates for each rank.
Allocation and deduplicated hidden-payload counts follow a JS port of MoonEP's planning kernel
(2bd860b, 2026-08-13).
Timing and memory lifetimes are modeled assumptions, not GPU measurements.
Routing flow who sends what to whom, and what it costs
Read top to bottom, in the order the planner works. First, which experts the router loaded, and how far each rank ends up over or under capacity. Then the greedy balance that repairs the imbalance, one move at a time. Then where the replica weights fly, next to every token src→dst pair. Finally, what it all costs in time, rank by rank — and the memory it all lives in. The charts recompute from the same synthetic routing: Steps 1–4, C2 destination arithmetic, and per-token dedup determine the counts. Byte formulas and an explicitly synchronized schedule turn those counts into estimates; planning/control traffic and transport overhead are not fully modeled.
Per-expert breakdown
One bar per expert, colored by its owner rank (the same colors as the tables and the
balance log). Router skew shows up as uneven bars; a hot rank shows up as one tall colored block.
Each bar splits in two: the solid base is entries computed on their own source rank (no
wire traffic), the faded top is entries sent to another rank. Hover a bar for the
planner's full decision on it. The tallest bar is marked ▾: MaxVio
(bench_vs_deepep.py:84) is how far it sits above the dashed average line,
as a ratio − 1.
Same data, summed per rank. Add each rank's experts together and subtract its capacity (S·K real entries after balancing, before padding). Bars above zero have too many entries; bars below zero have room to spare. These are net export/import requirements; a donor may export more than its initial surplus and later import entries.
The greedy balance, step by step
The planner's two greedy loops, replayed for this exact routing
(planning.py:672–701 Step 2, :746–768 Step 3). Each move
refills the currently hungriest rank in one shot, from the currently biggest donor.
The donor can overshoot into deficit; it then becomes a receiver and gets refilled by a
later move. That is the deliberate price of one-donor-per-receiver. Each move is realized by
slicing off the donor's biggest experts. With the full prefetch budget, every foreign expert
segment gets one weight copy; smaller inference pools leave some segments to read weights remotely.
Why one donor per receiver is enough
Let C = S·K and bh = the load of home group h − C. Balances sum to zero. Each step transfers the entire largest deficit from the largest positive donor; ties select the smallest rank index. The receiver becomes zero and cannot be selected again while positive and negative balances remain. At most R−1 moves finish the plan.
A deficit is at most C. A positive donor has not previously received and retains C+bh of its own entries, so it can supply the transfer even if its balance becomes negative. That donor can then receive once from a different home group. For example, balances [+3,+2,−5] become [−2,+2,0], then [0,0,0].
Each receiver therefore gets foreign experts from only one home group, containing E/R experts. B = E/R prefetch slots suffice. At most 2·E/R nonempty local-plus-foreign segments need padding, bounding logical capacity by NvS = S·K + 2·(E/R)·(p−1). This guarantees feasibility; it does not minimize bytes or account for topology. Capacity derivation in the API.
Three counts to keep separate
One token routed to two experts on one destination and a third expert on another creates three entries, two hidden-payload rows, and three fp32 route weights. A destination expands its shared hidden row into expert slots; combine folds those slots back together. Every nonempty expert segment is rounded up to 128 rows for GEMM. Local payload rows appear on the matrix diagonal but do not cross a link.
Where the weights fly
Each orange arrow is one prefetched expert's weights, flying from its owner to the rank that will compute it. Width scales with the rows the copy serves; dashed means a hypothetical IB edge between NVLink islands. Token traffic lives in the matrix beside it — every src→dst pair.
Fewer prefetch slots in inference
The planner first balances the entries, then selects the B largest foreign
expert segments on each destination for prefetch; equal counts favor the larger expert ID.
With B = E/R every foreign segment fits. In inference, a smaller pool is valid: unselected segments
keep their original expert IDs in cu_seqlens and GEMM reads those weights remotely
through the owner's symmetric mapping. Changing B changes weight placement, not the routing or padding.
The repository suggests 3–4 slots for inference. That saves reserved memory even if a particular routing needs few copies. It does not eliminate the overflow weight traffic. The same README requires B = E/R for training, so these controls unlock only in infer MXFP4 mode.
The weight-pass setting is a sensitivity assumption. One pass charges the complete gate/up/down payload, including scales, once per overflow expert on each destination; more passes model rereads. Real traffic depends on GEMM tiling and caches, and is not determined by the planner. The timeline adds the estimated remote-read service cost to forward GEMM compute, without hiding it under compute. It is an additive approximation, not a calibrated latency bound. Orange ring arrows show explicit prefetches only; the table lists overflow reads that take place inside GEMM.
A training contract before the timeline
This model saves an independent dispatched input and BF16 intermediate
activations for backward. Dispatch and combine use boundary copies; their cost is included.
Communication-buffer views are overwritten by later calls and cannot be saved for autograd.
Hidden zero_copy and router_weights_zero_copy are independent options,
both false by default. Forward route weights are copied and retained for the expert backward path.
Save the MoonEPCommPlan and forward cu_seqlens.
Backward-of-combine dispatches output gradients using that plan; plan reuse skips planning but retains
the API's default pre-sync, and returns no new cu_seqlens. We re-prefetch weights before backward
because other layers can overwrite shared slots. Expert backward computes dX, dW and route-weight
gradients; backward-of-dispatch combines hidden gradients and gathers route-weight gradients.
reduce_grad then pulls temporary fp32 expert gradients into their owners, fences peers,
and clears consumed slots. Owners and pool users must respect completion events.
API contract.
Estimated routed-expert cost under a synchronized schedule
Every modeled phase starts together and finishes at the slowest rank. Gray means idle time imposed by this schedule. Dispatch has an exit barrier and combine an entry barrier; a simultaneous combine finish is a modeling boundary, not an extra kernel barrier. These bars include forward, backward weight prefetch, boundary copies, and gradient reduction, with no assumed overlap. Attention, router computation, shared experts, nonlinearities and their gradients, base dispatch/combine local-memory work, allocator/launch overhead, and synchronization latency are outside the estimate.
How to read these estimates
Hidden payload is always BF16. Forward dispatch sends one fp32 route weight per entry; backward combine gathers one fp32 route-weight gradient per entry. Token phases use max(egress bytes, ingress bytes) ÷ effective link bandwidth, assuming full duplex. Duplicate expansion reads one primary row per group and writes each duplicate; folding reads primary plus duplicate rows and writes the primary. Boundary copies count both reads and writes.
GEMM uses 6·Pd·H·H′ FLOPs, where Pd is that rank's padded row count. Backward matmuls use twice those FLOPs. The effective throughput is held fixed across ranks and precision modes; segment shapes can change real utilization. Planning is assumed to cost 20 µs; local HBM bandwidth is assumed to be 3.3 TB/s. Neither is calibrated here.
For scale: H100 SXM dense peaks are about 989 TFLOP/s BF16 and 1,979 FP8. B200 dense FP8 is about 4,500 TFLOP/s per GPU, approximately 2.27× H100 FP8. The default 400 is an assumption, not a promise of a particular percentage utilization. MXFP4 changes weight bytes to 3·H·H′·(½+1/32), including scales, and selects forward-only execution; it does not model a quantized GEMM.
Island size < R prices cross-island edges as hypothetical IB, independently of intra-island NVLink. MoonEP implements shared NVLink/NVSwitch memory, including multi-node NVLink fabric; it does not implement this arbitrary IB payload path. NIC sharing, switch contention, and wire overhead are omitted. The topology controls change costs, not allocation decisions.
K3's preset describes its routed expert dimensions; latent projections and its full quantized compute path are outside this model. Its training report describes shared-expert stream overlap and reconstructing expert inputs by redispatching saved token-major inputs. Those require a different memory/compute schedule. Likewise, async gradient reduction can overlap independent work only with correct events and shared-pool lifetimes; this timeline includes it in the total.
When does moving weights pay off?
The first breakeven is an overlap threshold: can independent shared-expert computation cover the time to dispatch tokens and prefetch foreign weights? Kimi K3 has two full-width shared experts per MoE layer: D = 7,168 and intermediate width 3,072 each. Its routed experts use H = 3,584, H′ = 3,072, E = 896 and K = 16. Shared weights are replicated across EP ranks, and their GEMMs run on a separate stream. Architecture · Training schedule.
For the chosen shared branch, forward GEMM time is Tshared = 6·S·Ns·D·Is / Fshared. The available window is η·Tshared. η = 100% assumes the entire branch can cover startup; lower it for late routing/latent-input readiness or work reserved for another overlap window. These are effective rates under concurrent communication, not isolated hardware peaks.
Deriving the breakeven points
For each rank, link class and direction j, let Wj be its selected foreign-weight bytes and Aj its off-rank token bytes per input token. Forward token bytes include deduplicated BF16 hidden rows and fp32 route weights. Count both donor egress and receiver ingress: many receivers can make the weight owner the bottleneck. Traffic sharing a port adds.
Tready(S) = L + maxj[(Wj + Aj·S) / BWj]. With c = η·6·Ns·D·Is/Fshared, full hiding requires Tready(S) ≤ c·S. Thus S* = ceil(maxj[(L + Wj/BWj) / (c − Aj/BWj)]). If an active port's token slope already meets or exceeds the window slope and has a positive fixed cost, no finite S hides it. Faster shared GEMMs shorten this window; that alone does not mean the layer gets slower.
The curves freeze the current route fractions, prefetch selections and effective bandwidths while scaling S. They are a local extrapolation, not new planner runs or a K3 routing trace. The model pools token and weight bytes within each direction, assumes full duplex and independent link classes, and permits scheduling them freely after planning. L = 20 µs for forward; backward reuses the plan and uses L = 0. Real ordering, SM/HBM contention, readiness, launches and synchronization can expose more time. Recheck after changing routing or B. An IB island remains hypothetical.
The backward row uses twice the shared forward matmul work to cover reverse-combine dispatch and restored weights. It is a separate window. The forward window cannot also hide the final combine, and neither startup test hides later weight-gradient reduction. Inference overflow weights are read inside GEMM and are not counted as prefetched or hidden here. The synchronized timeline and memory chart above/below retain their original schedule; this panel explores an alternative.
The second breakeven: versus a fused megakernel
Mixture-of-Kittens already covers initial dispatch with shared-expert work, pipelines token communication with routed GEMMs, and overlaps the last reverse-dispatches with weight-gradient GEMMs. It keeps routed weights on their owners. Full hiding of MoonEP's startup transfers is therefore useful, but is not itself a crossover against MoK.
For this forward-only comparison, assume shared and routed GEMMs consume the same compute resource sequentially at their stated effective rates. Both routed layouts use the simulator's 128-row segment padding. Give an ideal owner-compute fused baseline zero exposed communication, planning, copies or pipeline tails: Tfused,ideal = Tshared + Gowner,max. Our illustrative MoonEP schedule costs Tshared + max(0, Tready − η·Tshared) + Gbalanced,max + Cremaining. C includes the existing model's duplicate expansion, boundary copies, folding, final combine, and any additive overflow-GEMM service cost. Balancing has room to win when Gowner,max − Gbalanced,max exceeds exposed startup plus Cremaining. This is a comparison against an idealized reference at equal routed throughput, not a measured MoK prediction.
A measured crossover needs matching K3 inputs, routing, precision, hardware and SM allocation, including latent projections, SiTU-GLU, backward and gradient reduction. The public MoK implementation targets DSV3-style BF16/MXFP8 training; its published shapes do not establish a K3 LatentMoE or MXFP4 crossover. Balanced owners have no compute imbalance to recover; a hot expert or home group can create relief, but also concentrates weight service on its owner.
Memory: logical storage and lifetimes
One identified rank, on the timeline above. Light bands are persistent buffers or layer parameters; solid bands are independently allocated plans, API outputs, and framework tensors for the chosen saved-activation schedule. Bands show logical tensor bytes, before hardware-dependent VMM/multicast granularity and allocator rounding. They are not a measured physical allocation trace. Static capacity bounds routing storage; the API still allocates outputs and plans, and whole-model memory depends on layers, live microbatches, and the framework. Hover a band for its formula and lifetime.