You benchmark a fine-grained MoE model — DeepSeek-V3 style, 256 routed experts, top-8 — against a dense model of the same active parameter count. Prefill throughput is great: 3-4x the tokens/sec you expected from parameter count alone. Then you measure decode at batch 32 and it's slower per token than the dense model, on more GPUs. A profiler shows the MoE layers spending most of their time in all-to-all collectives, and one expert-parallel rank finishing its grouped GEMM well after the others while everyone waits.
That gap is MoE expert load imbalance, and it is not a router bug. It's a batch-size problem with a hard arithmetic floor.
TL;DR
- MoE expert load imbalance is a decode-time problem, not a prefill problem. Prefill pushes thousands of tokens through the router per step, so expert loads average out. Decode pushes one token per sequence, so a batch of 64 sends 512 token-expert pairs across 256 experts — 2 tokens each, wildly uneven.
- Every MoE layer ends in a synchronizing all-to-all. Step latency is set by the slowest rank, so the hottest expert's queue is your decode latency.
-
You need roughly
E/ktokens per step just to reuse each expert's weights once, and aroundE·TILE/ktokens (thousands, for fine-grained MoE) to fill grouped-GEMM tiles. Below that, MoE saves FLOPs but not memory bandwidth — you read nearly all expert weights to serve a handful of tokens. - Load-balancing losses don't help at decode. They balance over training batches of hundreds of thousands of tokens. Balance is a large-number effect.
- Fixes that work: raise tokens-per-step (DP attention, chunked prefill mixing, prefill/decode disaggregation), replicate hot experts (EPLB-style redundant experts), or drop expert parallelism for tensor parallelism at low batch.
Why does MoE expert load imbalance hurt decode but not prefill?
Because the number of tokens hitting the router per forward pass differs by two or three orders of magnitude.
In prefill, a single 4k-token request produces 4096 router decisions in one step. With 256 experts and top-8, that's 32,768 token-expert assignments — an average of 128 tokens per expert. Balls-in-bins variance at that scale is small; max/mean rank load lands near 1.02. Every expert GEMM is fat and compute-bound. MoE looks exactly as good as the FLOP math promises.
In decode, each sequence contributes exactly one token. A batch of 64 sequences produces 64 router decisions, 512 assignments, 2 tokens per expert on average — and many experts get 0 while some get 6. The grouped GEMM for the hot expert is 3x the rows of the mean, and since MoE dispatch/combine is an all-to-all pair per layer, ranks synchronize twice per layer. Over ~60 layers, the tail dominates.
The structural point: expert parallelism converts statistical variance into wall-clock latency. Tensor parallelism doesn't have this property — every rank does an identical slice of identical work. EP shards work by routing outcome, which is data-dependent.
How many tokens per step do you actually need?
Two thresholds, both easy to derive.
Weight reuse. With E experts and top-k routing, a batch of B tokens touches an expected E·(1 - (1 - k/E)^B) distinct experts. For E=256, k=8:
| B (tokens/step) | fraction of experts touched | expert weight bytes read per token |
|---|---|---|
| 1 | 3.1% | 3.1% of all expert weights |
| 8 | 22% | 2.8% |
| 64 | 87% | 1.4% |
| 256 | ~100% | 0.39% |
| 2048 | 100% | 0.049% |
At B=64 you are already reading 87% of the expert weights to produce 64 tokens. The FLOPs are 8/256 of dense, but the bytes are nearly all of it — and decode is bandwidth-bound. Your 671B-parameter model behaves like a 671B-parameter dense model on the memory bus while doing 37B parameters' worth of math. That is the worst of both worlds. Reuse only starts paying above roughly B ≫ E/k = 32 tokens per step, and you want an order of magnitude past it.
Tile quantization. Grouped GEMM kernels tile the M dimension in blocks of 64 or 128 rows. An expert with 3 assigned tokens costs a 64-row tile anyway. Averaging 64 tokens per expert requires B = E·TILE/k = 256·64/8 = 2048 tokens per decode step. That is why large-scale MoE serving pushes toward hundreds of concurrent sequences per EP group and disaggregates prefill from decode — decode needs an enormous running batch to stop wasting the tensor cores.
Here's the simulation. It assumes a perfectly uniform router, so treat every number as an optimistic bound:
import numpy as np
rng = np.random.default_rng(0)
def sim(B, E=256, k=8, ranks=8, tile=64, trials=600):
"""Return (mean rank imbalance, p99 imbalance, tile-inflated work ratio)."""
per_rank, imb, tile_imb = E // ranks, [], []
for _ in range(trials):
counts = np.zeros(E, dtype=np.int64)
for _t in range(B): # top-k without replacement
counts[rng.choice(E, size=k, replace=False)] += 1
rank_load = counts.reshape(ranks, per_rank).sum(axis=1)
imb.append(rank_load.max() / rank_load.mean())
padded = np.ceil(counts / tile) * tile # grouped-GEMM M padding
tile_imb.append(padded.reshape(ranks, per_rank).sum(axis=1).max()
/ rank_load.mean())
return np.mean(imb), np.percentile(imb, 99), np.mean(tile_imb)
for B in [8, 64, 256, 1024, 4096]:
m, p99, t = sim(B)
print(f"B={B:5d} mean imb={m:.2f} p99={p99:.2f} tile-inflated work={t:.1f}x")
Output:
B= 8 mean imb=1.51 p99=2.00 tile-inflated work=84.9x
B= 64 mean imb=1.18 p99=1.33 tile-inflated work=30.3x
B= 256 mean imb=1.09 p99=1.18 tile-inflated work=8.0x
B= 1024 mean imb=1.04 p99=1.09 tile-inflated work=2.0x
B= 4096 mean imb=1.02 p99=1.04 tile-inflated work=1.3x
Rank-level imbalance at B=64 is only 18% on average — but the p99 step is 33% over, and every step pays it because latency is a max over ranks. The tile-inflation column is the bigger scandal: at B=64 the grouped GEMMs do ~30x the useful row-work. Coarse MoE (Mixtral-style E=8, k=2) has the opposite profile — worse rank imbalance at small batch (4x at B=1, since a single token can only light up 2 of 8 ranks) but far less tile waste, because each active expert gets many more tokens.
Why doesn't the router's load-balancing loss fix this?
Because auxiliary load-balancing losses and DeepSeek-V3-style bias-based balancing operate on training batches of hundreds of thousands of tokens. They make the router balanced in expectation. Expectation is not what your 64-token decode step samples.
Three effects make production worse than the uniform simulation above:
- Domain skew. Routers specialize. A batch of Python code requests and a batch of Korean customer-support requests light up measurably different expert subsets. A single-tenant workload can sit permanently on a skewed slice of the router's distribution — a persistent hot expert, not a random one.
- Intra-request correlation. Tokens from the same sequence route similarly. Fifty concurrent sequences from one tenant are not fifty independent draws.
- Layer-wise variance. Balance is per-layer. One badly balanced layer out of 60 still adds its tail to every step, because the all-to-all barrier is per-layer.
Sensible expectation: real max/mean rank load at moderate decode batch runs meaningfully above the uniform bound, and the persistent component is what redundant-expert placement is designed to absorb.
Expert parallelism or tensor parallelism for MoE?
Pick by tokens-per-step, not by model size.
| Expert parallelism (EP) | Tensor parallelism (TP) | |
|---|---|---|
| Work split | By routing outcome (data-dependent) | By hidden dim (fixed, equal) |
| Imbalance | Yes — latency = max rank | None by construction |
| Comms | 2 all-to-all per MoE layer | 1 all-reduce per MoE layer |
| Per-rank weight traffic | Only local experts | A slice of every activated expert |
| Best at | Large decode batch, high throughput | Small batch, latency-sensitive |
TP for MoE looks wasteful on paper — every rank reads a slice of every activated expert — but at B=64 you're touching 87% of experts anyway, so EP's traffic advantage barely exists. What EP does give up at that batch is real: data-dependent tail latency plus two latency-bound collectives per layer with tiny messages. That's why low-batch, low-latency MoE deployments often win with TP, and why high-throughput deployments need EP plus a very large running batch to amortize it.
Rough serving config for the throughput regime (verify flag names against your engine version — these move):
# vLLM: EP for MoE + DP attention, aiming for a fat decode batch
vllm serve deepseek-ai/DeepSeek-V3 \
--tensor-parallel-size 1 \
--data-parallel-size 8 \
--enable-expert-parallel \
--max-num-seqs 512 \ # tokens per decode step ~= running seqs
--enable-chunked-prefill \
--max-num-batched-tokens 8192 # prefill chunks ride with decode tokens
Two things matter here. --max-num-seqs is effectively your decode tokens-per-step ceiling; if it's 32, no amount of EP tuning will save you. And chunked prefill is unusually valuable for MoE: mixing prefill chunks into decode steps fattens every expert GEMM, so decode tokens ride along in tiles that would otherwise be mostly padding. The usual chunked-prefill trade-off (slight ITL cost) is offset by a much larger MoE efficiency gain.
How do I fix MoE expert load imbalance in production?
In priority order:
- Raise tokens per step. Everything else is secondary. Data-parallel attention with expert parallelism across the DP group aggregates all ranks' decode tokens into one shared expert pool — that's the difference between 64 tokens/step and 512. Prefill/decode disaggregation exists largely so decode can run at a batch that makes MoE arithmetic work.
- Replicate hot experts. EPLB-style redundant experts: measure per-expert token counts over a window, place a second copy of the top offenders on under-loaded ranks, and split their traffic. This targets the persistent skew (domain effects), not the per-step noise. Rebalance on a slow cadence — minutes, not steps — since moving expert weights isn't free.
- Use low-latency all-to-all kernels. At decode, dispatch/combine messages are small and latency-bound; DeepEP-class kernels (NVLink intra-node, RDMA inter-node, communication overlapped with compute) matter more than raw bandwidth.
- Check your grouped-GEMM tile size. If your kernel uses a 128-row M tile and you're serving at 2 tokens/expert, a smaller tile can cut padded work substantially at the cost of some peak efficiency.
- Shrink the EP group before you shrink the batch. EP=8 over 256 experts means 32 experts per rank and decent aggregation. EP=64 means 4 experts per rank — variance per rank rises sharply and the all-to-all gets wider. Large EP only pays with a matching large batch.
Instrument first. Log per-expert token counts per layer over a few hundred decode steps and compute max_rank_load / mean_rank_load and the fraction of expert-GEMM rows that are padding. If mean imbalance is under ~1.1 and padding is under ~2x, your MoE layers are fine and the bottleneck is elsewhere. If padding is 20x, no kernel tuning will help — you need more tokens per step.
Direct answer
MoE expert load imbalance stalls decode on one GPU because expert parallelism splits work by routing outcome, and each MoE layer ends in a synchronizing all-to-all, so step latency equals the slowest rank's grouped GEMM. At decode you have one token per sequence, so a small batch spreads a handful of token-expert assignments across hundreds of experts: loads are uneven, most GEMM tiles are padding, and you read nearly all expert weights to produce very few tokens — MoE's FLOP savings without its bandwidth savings. Prefill hides all of this because thousands of tokens per step average the router out. The fix is more tokens per decode step (DP attention, chunked prefill, prefill/decode disaggregation), redundant copies of persistently hot experts, and choosing tensor parallelism over expert parallelism when your batch is genuinely small.
Top comments (0)