Section IX.05 covered expert parallelism as a distributed systems problem. This file covers the serving engineering: kernels, scheduling, and what actually makes MoE fast.
1. The critical kernel: grouped GEMM#
The single thing that determines whether MoE serving is fast.
THE PROBLEM
After routing, you have tokens assigned to different experts:
expert 0: tokens [3, 17, 42, 88]
expert 1: tokens [1, 5, 91]
expert 2: tokens [0, 2, 4, 7, 12, 30]
...
Each expert has DIFFERENT weights and a DIFFERENT number of tokens.
NAIVE: loop over experts, one GEMM each
✗ E small GEMMs → E kernel launches, poor tile utilization
✗ at E=256 that's 256 launches per layer per step
GROUPED GEMM: one kernel launch, many different-shaped GEMMs
✓ one launch
✓ each thread block works on one expert's tile
✓ the token→expert assignment is passed as metadataTHE PREREQUISITE: SORT TOKENS BY EXPERT
before: tokens in arbitrary order, expert ids scattered
after: [expert 0's tokens][expert 1's tokens][expert 2's tokens]...
plus offsets: [0, 4, 7, 13, ...]
→ each expert's tokens are CONTIGUOUS
→ the grouped GEMM reads them coalesced
→ no warp divergence (all threads in a block handle one expert)Sorting is not optional. Without it you have maximal divergence (Section VI.10) and uncoalesced access (Section VI.09). The sort itself is cheap (a counting sort over E buckets).
def moe_forward(x, router_logits, expert_weights, k):
# 1. route
topk_probs, topk_idx = router_logits.softmax(-1).topk(k, -1)
# 2. SORT tokens by expert (the essential step)
flat_expert = topk_idx.flatten() # (T*k,)
flat_token = torch.arange(len(x)).repeat_interleave(k)
order = flat_expert.argsort()
sorted_expert = flat_expert[order]
sorted_token = flat_token[order]
# offsets: where each expert's segment starts
offsets = torch.searchsorted(sorted_expert,
torch.arange(num_experts + 1))
# 3. gather the tokens in sorted order
x_sorted = x[sorted_token] # (T*k, d)
# 4. GROUPED GEMM: one launch, per-expert weights and offsets
y_sorted = grouped_gemm(x_sorted, expert_weights, offsets)
# 5. scatter back and combine with routing weights
out = torch.zeros_like(x)
out.index_add_(0, sorted_token,
y_sorted * topk_probs.flatten()[order].unsqueeze(-1))
return out
2. Where the time goes#
MoE LAYER TIME BREAKDOWN (measured, DeepSeek-scale, EP across a node)
router (small GEMM + topk) 3%
sort + gather 8%
ALL-TO-ALL dispatch 18% ← communication
grouped GEMM (up) 22%
activation 4%
grouped GEMM (down) 22%
ALL-TO-ALL combine 17% ← communication
scatter + weighted sum 6%
─────────────────────────────────────
communication total: 35%
compute total: 48%
data movement (sort/gather/scatter):17%35% in All-to-All is typical and is the dominant optimization target. The two big levers:
1. OVERLAP communication with computation
While the All-to-All for layer L is in flight, compute layer L's
attention or layer L+1's router.
→ can hide 50-80% of the communication cost
→ requires careful stream management (Section VI.06)
2. KEEP THE ALL-TO-ALL WITHIN A NODE
NVLink: 450 GB/s. InfiniBand: 50 GB/s. 9x difference.
→ if E experts fit on one node, don't spread them across nodes3. Expert placement#
STRATEGY DESCRIPTION WHEN
Pure EP each GPU holds E/N experts memory-constrained
Expert TP (ETP) each expert sharded across GPUs small E, model fits
Hybrid EP+TP groups of GPUs hold groups of large deployments
experts, each sharded within
Replicated hot experts duplicate frequently-used measured imbalance
experts on multiple GPUs
DeepSeek-V3's approach (per their report):
fine-grained experts + node-limited routing: a token's experts are
constrained to a limited number of nodes, bounding cross-node
All-to-All.“Node-limited routing” is an elegant idea worth noting: constrain the router (during training) so that a token’s k experts live on at most M nodes. This bounds the cross-node communication by construction rather than hoping the routing cooperates.
4. Load balancing at serving time#
THE PROBLEM (Section IX.05): routing is uneven, the All-to-All is a barrier,
so the hottest expert sets the pace.
SERVING-TIME MITIGATIONS
1. LARGER BATCHES
imbalance averages out. Section XIII.02's formula shows why:
at large B, every expert gets tokens.
→ the strongest and simplest mitigation
2. CAPACITY FACTOR
cap tokens per expert at C × (tokens/E). Excess tokens SKIP the
expert (pass through the residual).
C = 1.25 typical.
✓ bounds the imbalance hard
✗ dropped tokens = quality loss
→ measure the drop rate; if it's > 2%, raise C
3. EXPERT REPLICATION
measure which experts are hot (offline, over real traffic).
Replicate them on multiple GPUs; route round-robin among copies.
✓ directly addresses the imbalance
✗ extra memory; requires the measurement
4. DYNAMIC REBALANCING [EMERGING]
periodically reassign experts to GPUs based on measured load.
✓ adapts to distribution shift
✗ moving an expert costs a weight transferMeasure the expert load distribution before choosing. If your imbalance is mild (hottest expert at 1.2x fair share), larger batches suffice. If it’s severe (2x+), replication is worth the memory.
5. Batching implications#
MoE inverts a normal serving intuition:
DENSE MODEL
larger batch → better (amortizes weight reads)
MoE MODEL
larger batch → better for compute AND for load balance
but WORSE for the memory-traffic advantage
(Section XIII.02: more experts touched)
→ MoE has an OPTIMAL batch size, not a monotonic preference
small batch: few experts touched (good bandwidth) but severe
imbalance and poor All-to-All efficiency
large batch: all experts touched (bandwidth like a dense model
of the TOTAL size) but good balance and amortizationMEASURED (illustrative, Mixtral 8x7B on 4 GPUs):
batch tok/s experts touched imbalance ratio
1 48 2.0 n/a
8 340 6.7 2.4x
32 1,180 8.0 1.6x
128 2,900 8.0 1.2x
512 4,100 8.0 1.05x
→ throughput keeps improving, but the marginal gain shrinks and
the bandwidth advantage is gone by batch 32.Practical guidance: run MoE at high batch. The imbalance and communication overhead only amortize there, and the memory-traffic advantage is already gone.
6. Engine support#
CHECK, FOR YOUR ENGINE AND MODEL:
□ grouped GEMM (not a per-expert loop)
□ fused sort/gather/scatter
□ All-to-All overlapped with compute
□ EP, ETP, or hybrid — which is supported?
□ capacity factor configurable
□ expert load metrics exposed
□ FP8 support for expert weights
RED FLAGS IN A PROFILE
✗ E separate GEMM kernels per MoE layer → no grouped GEMM
✗ All-to-All not overlapped (visible gaps around it)
✗ large sort/gather time → not fusedSupport quality varies substantially between engines and model architectures. A model that runs correctly may run at half the achievable speed. Profile it.
7. Production implications#
- Run MoE at high batch. Low-batch MoE serving is inefficient on every axis except raw memory traffic.
- Keep the All-to-All within a node if the experts fit.
- Monitor expert load distribution. It’s a first-class metric for MoE.
- Verify grouped GEMM in a profile.
- Set the capacity factor deliberately and monitor the token drop rate.
- Quantize aggressively — memory is the constraint (Section XIII.02).
- Budget for the All-to-All in your latency model; it’s 30-40% of MoE layer time unoverlapped.
- MoE is a throughput architecture. If you need low-latency low-batch serving, a dense model is usually the better choice.
8. Common mistakes#
Per-expert loops instead of grouped GEMM. Several-fold slowdown.
Not sorting tokens by expert. Divergence and uncoalesced access.
Cross-node All-to-All when experts would fit in a node.
Not overlapping communication. 35% of layer time exposed.
Running MoE at low batch. All the overhead, none of the amortization.
Capacity factor too low. Silent token drops and quality loss.
Not monitoring expert load. Invisible throughput loss.
Sizing hardware from active parameters (Section XIII.02).
9. Hands-on exercise#
A. Implement grouped GEMM MoE. Write the forward pass from section 1, including the sort. Compare against a naive per-expert loop. Measure both at batch 32 and 256.
B. Measure the breakdown. Profile an MoE model and reproduce the time breakdown from section 2. What fraction is communication? Is it overlapped?
C. Load distribution. Instrument the router to record tokens per expert. Plot the distribution at several batch sizes. Compute the imbalance ratio (hottest/mean) at each.
D. Capacity factor sweep. Vary C from 1.0 to 2.0. Measure throughput and the token drop rate. Where’s the knee? What drop rate is acceptable for your task?
E. Batch sweep. Measure MoE throughput, experts touched, and imbalance at batch 1 through 512. Reproduce the table in section 5 for your model.
F. Expert replication. Identify the top-3 hottest experts from your measurement. Simulate replicating them. What’s the throughput improvement and the memory cost?
10. Interview questions#
- Why must tokens be sorted by expert before the GEMM?
- What is grouped GEMM and why does MoE need it?
- Where does time go in an MoE layer? What’s the dominant optimization target?
- Why does MoE have an optimal batch size rather than “bigger is better”?
- How would you address expert load imbalance at serving time?
- What is node-limited routing and what problem does it solve?
- What would you check in a profile to verify good MoE support?
11. Further reading#
- [ESTABLISHED] Rajbhandari et al., “DeepSpeed-MoE” (2022)
- [ESTABLISHED] DeepSeek-V3 technical report — the most detailed public account of large-scale MoE inference engineering
- [REFERENCE] CUTLASS grouped GEMM; Megablocks
- [ESTABLISHED] Gale et al., “MegaBlocks: Efficient Sparse Training with Mixture-of-Experts”
- Next: 04 — Long-context inference