05 — Expert Parallelism (MoE)

Reference reading: GShard (Lepikhin et al. 2020), Switch Transformer (Fedus et al. 2021), DeepSeekMoE (Dai et al. 2024), DeepSeek-V3 tech report (arxiv:<phone_number_or_numberic_id_or_random_id_82>), Mixtral-8x7B report (Jiang et al. 2024). Kernel awareness: DeepEP (DeepSeek’s open-source MoE all-to-all kernels, Feb 2025).

Why MoE exists at all

Dense scaling is expensive: doubling parameters roughly doubles training and inference FLOPs. Mixture of Experts changes that ratio: you scale total parameters (memory) without scaling active parameters (compute) per token.

A MoE layer replaces a single MLP with N MLP “experts” plus a router. For each token, the router picks the top-k experts (typically k=2 or k=8); only those experts execute. N = 64, top_k = 2: token sees 2 experts out of 64, but the model holds 64 experts’ worth of parameters.

The economic result:

  • Compute cost per token ≈ 2/64 of dense equivalent (for k=2, N=64).

  • Memory cost per token = same as dense equivalent — all 64 experts must be resident.

  • Best-case quality: close to a dense model with parameters ≈ N × expert_size, at the compute of a dense model with k × expert_size.

DeepSeek-V3: 671B total, 37B active. That’s 18× compute reduction vs a dense 671B — the reason MoE has won at frontier scale.

The serving asymmetry (this is the exam question)

Dense LLM serving: memory bandwidth is the enemy at decode. Every token streams all weights.

MoE serving: memory capacity is the enemy. All experts must be resident somewhere, but only k execute per token. So:

  • If you serve a MoE on one GPU (all experts local), you waste ~(N-k)/N of your HBM capacity to hold experts that don’t run for this token.

  • If you shard experts across GPUs (expert parallelism), you fit the model economically — but now each token must be dispatched to the right GPU (all-to-all on the way in, all-to-all on the way out).

The crossover: EP wins as the number of GPUs grows and as batch size grows (all-to-all latency amortizes over more tokens).

The all-to-all pattern

For a MoE layer:

tokens (b*s, hidden) 
  → router assigns each token to top_k experts
  → dispatch: all-to-all sends each token to the GPU that owns its expert
  → local expert MLP (bunch of GEMMs on token subset)
  → combine: all-to-all sends results back to origin GPU
  → weighted sum by router probability

Bytes per token per MoE layer: 2 × top_k × hidden × bytes across the all-to-all fabric.

Example (Mixtral-8x7B): hidden=4096, top_k=2, bf16. 2 × 2 × 4096 × 2 = 32 KB per token per MoE layer for dispatch+combine.

All-to-all is bandwidth-heavy and topology-sensitive. Ring topologies are terrible for it (worst-case pairs). NVSwitch (any-to-any full bisection) is designed for it. On InfiniBand it’s tolerable but slower.

Load balancing — the hard part

If the router directs 90% of tokens to expert 3, expert 3’s GPU is a bottleneck and the others idle. Two families of fixes:

(a) Auxiliary loss (Switch Transformer, GShard)

Add a loss term that penalizes imbalance:

L_aux = N * Σ_i (f_i * P_i)

where f_i is the fraction of tokens routed to expert i and P_i is the average router probability for expert i. This nudges the router toward uniform distribution. Downside: adds noise to the primary loss, needs a tuned coefficient, and imbalances still occur.

(b) Auxiliary-loss-free (DeepSeek-V3)

The insight: instead of penalizing imbalance in the loss, directly adjust router logits with a per-expert bias. After each step, count tokens per expert. If expert i is under-loaded, add a small positive bias to its logits. If over-loaded, subtract.

chosen = TopK(logits + bias)   # bias is not part of gradient

# after step:
bias[i] += lr_bias * (avg_load - load[i])

The bias is not part of the gradient — pure control-loop feedback. This eliminates the aux-loss noise and gets better final quality. DeepSeek-V3 uses this and reports better performance than aux-loss variants at same compute. This is the modern default.

(c) Top-k with capacity + token drop / restore

Set a maximum tokens-per-expert-per-batch (capacity factor). If an expert overflows, drop the excess tokens (they skip this layer) or route to next-best expert. Both DeepSeek-V3 and Mixtral use capacity limits during training but avoid dropping at inference.

Fine-grained vs coarse experts (DeepSeekMoE insight)

Most early MoE (Switch, GShard) had ~8-64 large experts and k=1 or k=2. DeepSeekMoE argued for many small experts (256+) with top-k=8+ and shared experts always active:

  • Small experts specialize better (empirical).

  • Shared experts absorb common patterns, freeing routed experts for niche knowledge.

  • Higher k gives smoother combinations, reducing router variance.

DeepSeek-V3: 1 shared expert + 256 routed experts, top-k=8. Mixtral: 8 experts, top-k=2. Fine-grained has won the argument for frontier scale.

Wide-EP serving economics

For DeepSeek-R1-class models (671B, ~250 GB in bf16 or ~125 GB in fp8), you need EP ≥ 32 for realistic economics. The math:

  • 32× H100 SXM (80GB) = 2.56 TB total HBM.

  • Weights fit at bf16 with room for KV cache and activations.

  • But: at TP=8, PP=1, you need EP across the DP dimension = 32 to shard experts.

Actual DeepSeek reference serving stack (as reported in their inference cookbook):

  • EP=32 or EP=64 for the MoE layers.

  • TP=8 for dense parts (attention, shared expert).

  • Prefill and decode disaggregated (Mooncake pattern).

  • DeepEP kernels for the all-to-all step (see below).

DeepEP — the kernel you need to know exists

February 2025, DeepSeek open-sourced DeepEP (github.com/deepseek-ai/DeepEP): a communication library for MoE dispatch/combine, hand-tuned CUDA + NVSHMEM kernels for intra-node NVLink and inter-node RDMA. It provides:

  • Low-latency all-to-all for decode (small batch, latency-sensitive).

  • High-throughput all-to-all for training/prefill (large batch, bandwidth-sensitive).

  • Overlap of comm with compute via CUDA streams.

  • FP8 dispatch (halves bytes moved).

This library is what makes EP=64 serving of R1 fast enough to be profitable. vLLM and SGLang integrate DeepEP for MoE serving on H100+.

Awareness level for you (2026): know DeepEP exists, know that it’s why wide-EP is fast, know that upstream MoE-serving in vLLM/SGLang uses it. You don’t need to implement all-to-all kernels yourself unless you go deep on kernels.

MoE + serving — the failure modes

  1. Uneven expert selection at inference — no aux loss active, router imbalance can degrade latency. Solution: warm up the bias table with representative traffic, or use DeepSeek’s bias-adjustment loop even during serving.

  2. All-to-all latency dominates at small batch — decode with batch 1 pays full all-to-all cost for one token. Increases per-token latency to ~50-200 μs (DeepEP is very good but not free). This is why MoE serving loves larger batches.

  3. KV cache is dense — experts are only in MLP layers; attention is dense. So KV cache size doesn’t shrink with MoE. A 671B MoE has the same KV footprint as a 37B dense at the same architecture depth. This is often a pleasant surprise for capacity planning.

  4. Speculative decoding is weird for MoE — the drafter and target may route tokens to different experts, complicating verification. MTP (multi-token prediction), used by DeepSeek-V3, is the MoE-native form of speculative decoding.

Practical code entrypoints

  • vLLM MoE: --enable-expert-parallel, --data-parallel-size N. Reads model’s MoE config from config.json.

  • SGLang MoE: --tp N --enable-ep-moe and related flags.

  • Megatron MoE: --num-experts, --moe-router-topk, --expert-model-parallel-size.

  • DeepSpeed-MoE: legacy but still used; moe.experts.deepspeed.

Exercises

  1. For Mixtral-8x7B on 8xH100 with EP=8, compute all-to-all bytes per token per MoE layer at batch 32, seq 512 decode.

  2. Argue: for a 400B MoE with 128 experts, top-k=4, on 32 GPUs — what’s the min EP degree that fits weights in fp8 with 8k context KV?

  3. Sketch on paper the aux-loss-free bias update: given per-expert token counts [100, 50, 200, 50] after a step (target avg = 100), how do the biases update?

Papers in reading order

  1. Switch Transformer — the modern MoE template.

  2. GShard — all-to-all serving of MoE.

  3. DeepSeekMoE — fine-grained + shared experts argument.

  4. DeepSeek-V3 (arxiv:<phone_number_or_numberic_id_or_random_id_82>) — aux-loss-free routing, MTP, full frontier system.

  5. Mixtral report — practical prod deployment reference.