Mixture-of-experts router

Router. Each row is a token and each column an expert. Every token’s router logits go through a softmax (or, as in DeepSeek-V3, a sigmoid per expert). The token keeps its top-k experts, and their scores are renormalised into gates that sum to 1. The layer’s output is the gated sum \$y = \sum g_i E_i(x)\$. Capacity factor sets each expert’s capacity at ⌊cf × T × k ÷ E⌋ token slots. Slots beyond it are dropped (red ✕). As in GShard, every token’s first choice is placed before any second choice. A token that loses every slot passes through on the residual connection only. Popularity skew makes some experts more attractive, which is how an unbalanced router behaves. Switch Routing to Expert choice and each expert picks its own top tokens instead: every expert is exactly full, but some tokens get nothing. The step-by-step panel works one token through logits → scores → top-k → gates. The table below it shows the same token at six layers: routing is decided again at every layer.

Load balancing. The same skewed router is trained for a number of steps, balancing only, four ways. None is the baseline. Auxiliary loss is Switch Transformer’s α·E·Σ fᵢ·Pᵢ, minimised by gradient descent on a per-expert logit offset. Bias balancing is DeepSeek-V3’s method: a per-expert bias is added to the score only when choosing the top-k, and nudged by ±γ after each step. Expert Choice is the fourth. The charts show the busiest expert’s load ÷ the mean, and the share of token slots dropped at the capacity factor you set. The table shows what each method did to the gate values. The auxiliary loss balances by bending the router’s probabilities, so it competes with the model’s real objective. The bias never changes a kept expert’s score.

Active vs total. Published MoE models with their experts computed from the configuration (hidden size d, expert width f, three d×f matrices per SwiGLU expert): routed experts, always-on shared experts, the rest (attention, embeddings, dense layers), and the parameters one token actually uses. Compute follows the active count, but memory must hold the total.

Experts touched. With uniform routing, B tokens use E × (1 − (1 − k/E)^B) of a layer’s E experts. Mixtral’s 8 experts are almost all busy by batch 16. DeepSeek-V3 needs about 22 tokens to touch half of its 256. Add draft tokens to see why speculative decoding helps an MoE less (Section 33 of the book). Verifying γ drafts per sequence sends B × (γ + 1) tokens through the layer, and at small batch those extra tokens pull in experts that plain decoding would never read. The speed-up counts weight bytes only, which is what bounds a decode step.

All-to-all. Under expert parallelism every MoE layer runs two all-to-alls. Dispatch sends each token’s hidden vector to the GPUs holding its k experts. Combine brings the outputs back to the token’s home GPU. The tab gives bytes per token, per GPU per layer, and per forward step, split between NVLink and the network, and a bandwidth-only time. Network copies: once per node models DeepSeek-V3’s scheme. A token crosses the network once per node it needs and is fanned out over NVLink on arrival. Max nodes per token caps those nodes (DeepSeek-V3 uses 4).

The router is a toy: logits are random normals plus a fixed per-expert popularity, not a trained model’s hidden states. Real routing is skewed and correlated, so tokens next to each other share experts more often than here. Experts touched and all-to-all assume uniform, independent routing, which is the most scattered case.

The balancing simulation trains only the balancing, with no language-modelling loss to trade against. Its learning rates and γ are in the toy’s units. DeepSeek-V3’s γ = 0.001 acts on its own sigmoid scores.

Bandwidth-only all-to-all times ignore latency, congestion, and the overlap with computation that production systems (DeepSeek’s DualPipe, for example) rely on.

Active counts computed from the configuration differ slightly from the published ones. The remainder ("other") absorbs attention, embeddings, norms and biases. OpenAI’s gpt-oss active counts leave out the embedding table.

Lepikhin et al., GShard (2020): top-2 routing, expert capacity, and first choices placed before second choices.

Zhou et al., Mixture-of-Experts with Expert Choice Routing (2022).

Leviathan, Kalman & Matias, Fast Inference from Transformers via Speculative Decoding (2023): expected tokens per verification (1 − α^γ+1) ÷ (1 − α).

Link speeds: NVIDIA H100 (NVLink 900 GB/s bidirectional, 450 GB/s each way) and GB200 NVL72 (1.8 TB/s per GPU), ConnectX-7 (400 Gb/s InfiniBand NDR) and ConnectX-8 (800 Gb/s XDR); the DeepSeek-V3 report’s measured 160 GB/s NVLink and 50 GB/s InfiniBand on H800. Every one is listed with its source on the All-to-all tab.