Mixture-of-Experts: routing and capacity

Trace router scores, top-k, shared/routed experts, capacity, and communication.

Applied AI · advanced · Session 19

Mechanism map

           h_t (representation of token t)
                      │
                      ▼
              ┌───────────────┐  E = 4 logits per token;
              │    ROUTER     │  worked example over 3:
              └───────┬───────┘  [2.1, 1.8, 0.2]
                      ▼  softmax → [0.529, 0.392, 0.079]
              ┌───────────────┐
              │  TOP-k (k=2)  │  keeps E1, E2
              └───────┬───────┘
      ┌──────────┬────┴─────┬──────────┬──────────┐
      ▼          ▼          ▼          ▼          ▼
 ┌────────┐ ┌────────┐ ┌────────┐ ┌────────┐ ┌──────────┐
 │ E1 C=5 │ │ E2 C=5 │ │ E3 C=5 │ │ E4 C=5 │ │  SHARED  │
 │ 6 req. │ │ 4 req. │ │ 3 req. │ │ 3 req. │ │  always  │
 │ ❌ 1   │ │   ok   │ │   ok   │ │   ok   │ │   on     │
 │dropped │ │        │ │        │ │        │ │          │
 └───┬────┘ └───┬────┘ └───┬────┘ └───┬────┘ └────┬─────┘
     └──────────┴──────────┴──────────┴───────────┘
                      ▼  weighted sum + residual
                   h_t′  (all-to-all return if multi-GPU)

The problem — Why experts

To gain capacity, a dense block must grow — and every token pays for the whole block: ×20 parameters, ×20 FLOPs per token, even to predict “the”. The compute bill tracks capacity instead of tracking need.

The idea — Why experts

MoE decouples the two: E subnetworks (experts) exist, only k activate per token. In the trace: 64 routed experts + 1 shared = 26B parameters, but (2+1) × 0.4 = 1.2B active per token — a ratio of ≈ 21×.

Why / at what price — Why experts

Total capacity without the equivalent dense cost — that is the pitch. The immediate price: the 26B must reside in GPU memory at load time. The 21× ratio speaks of FLOPs — never of memory nor, as we will see, of latency.

Check: Here 26B total parameters but 1.2B active per token. Which quantity determines GPU memory at load time, and which determines FLOPs per token?

Visual support — Why experts

                    dense 26B          MoE 64 + 1 experts
parameters          26B                26B
in GPU memory       26B                26B      ← identical!
FLOPs per token     ∝ 26B              ∝ 1.2B  (top-2 + shared)

                    ratio ≈ 21× — on FLOPs, and only there

The problem — Router

Who decides which experts see which token? A hand-written switchboard — by language? by topic? — would be rigid and wrong; you need a per-token decision, learned with the rest of the model.

The idea — Router

The router is a small projection: h_t → one logit per expert. Worked example over 3 experts: [2.1, 1.8, 0.2] → softmax [0.529, 0.392, 0.079]. These scores are decisions learned through the global loss — not human categories.

Why / at what price — Router

The learned switchboard adapts without intervention. The price: illegibility — nothing guarantees an expert “is” math or code; observed specialization is statistical and can shift with the next rebalancing.

Check: The router gives [0.529, 0.392, 0.079] for t1. What if the third logit went from 0.2 to 1.9 — which top-2 decision changes, and has the router thereby “understood” the token?

The problem — Top-k and mixture

Running every expert “a little” would ruin the economics; keeping only one makes the choice brutal and the gradient fragile. How many experts to activate, and how to recombine their outputs?

The idea — Top-k and mixture

Top-k decides: k = 2 experts per token, weights renormalized — 0.529 and 0.392 become 0.574 and 0.426 (dividing by 0.921). An always-on shared expert completes it: output = Σ weight × expert + shared + residual.

Why / at what price — Top-k and mixture

Small k preserves the economics; the mixture keeps a gradient through two paths. The price: every token now hangs on a discrete decision — one flipped logit changes the whole compute path, less continuous behavior than a dense block.

Check: The renormalized weights for t1 are 0.574 and 0.426. Recompute them from 0.529 and 0.392, then say what the shared expert adds to that sum and why it is not part of the top-2.

The problem — Capacity

8 tokens, top-2, 4 experts: 16 assignments to seat. But a GPU tensor has fixed dimensions — an expert cannot accept “as many as needed”. What happens when demand exceeds room?

The idea — Capacity

Capacity is a tensor dimension: C = ceil(T×k/E × f) = ceil(4 × 1.25) = 5 here. The trace shows it: E1 receives 6 requests for 5 seats — the 6th arrival is dropped, its output carrying only 0.426 of the intended mixture (logits assumed equal to t1’s).

Why / at what price — Capacity

C bounds per-expert memory and compute — the tensor requires it. The price: the overflow policy (drop, reroute, queue) becomes an architecture choice that touches quality; and a generous f buys fewer drops at the cost of padding — f = 2 leaves 50% of slots empty.

Check: C = ceil(8×2/4 × 1.25) = 5 and E1 requests 6. Recompute C for f = 1.0 then f = 2.0: how many tokens are dropped in each case, and how many slots stay empty?

Visual support — Capacity

dispatch buffer [E, C, d]  with C = 5

       slot →   1     2     3     4     5     requests
E1             t1    t2    t3    t4    t5     6 → ❌ t6 dropped
E2             t1    t3    t6    t7     ·     4
E3             t2    t5    t8     ·     ·     3
E4             t4    t7    t8     ·     ·     3

16 assignments, 20 slots: 15 filled, 5 empty, 1 dropped

The problem — Load balancing

The trace gives 6/4/3/3: E1 saturates while E3 and E4 run at 60%. Taken to the extreme, a “lazy” router sends everything to one expert — and the others never learn anything again.

The idea — Load balancing

An auxiliary loss pushes the distribution toward uniform. It is a slider, not a switch: too weak, collapse onto one expert; too strong, it overrules the router even when E1 genuinely is the right choice.

Why / at what price — Load balancing

The right setting keeps every expert alive without crushing useful specialization. The price: one more loss term to watch, whose effect reads in load distribution — a systems metric — as much as in quality.

Check: The observed split is 6/4/3/3 instead of 4/4/4/4. A very strong auxiliary loss would force it uniform: what would be lost if E1 were the only useful expert for those tokens?

The problem — Communication

The 4 experts live on 4 GPUs. Every top-2 token must travel to its experts and back — and the 20 reserved slots, 25% padding included, cross the network in both directions. Where does the latency go?

The idea — Communication

The all-to-all paces the step: two round trips per top-2 token, a rhythm set by the busiest expert (E1 at 100%). The 1.2B active parameters measure compute; the network bills the slots — used or not.

Why / at what price — Communication

Well placed — co-located experts, grouped batches — MoE keeps its promises. The generic price: latency is often communication-dominated; “active parameters” is an excellent FLOPs indicator and a very poor speed indicator.

Check: 15 used slots out of 20 reserved still travel over the network. If the 4 experts live on 4 GPUs, how many all-to-all hops does a single top-2 token trigger, and why do 1.2B active parameters not guarantee speed?

Visual support — Communication

the journey of ONE top-2 token (experts on separate GPUs)

        t1 (top-2: E1, E2)
origin GPU ──▶ GPU(E1) ──▶ return      ┐ two round trips
origin GPU ──▶ GPU(E2) ──▶ return      ┘ per top-2 token

the [4 experts × 5 slots] buffer makes the same trip —
15 useful slots and 5 empty ones pay the same all-to-all

Worked case — full trace

Routing example reduced to three experts for the calculation: scores [2.1,1.8,0.2] give softmax about [0.529,0.392,0.079]. With top-2, experts 1 and 2 are active. If expert 1 has reached capacity, routing must apply the overflow policy.

BATCH: T = 8 tokens, E = 4 experts, k = 2, capacity factor f = 1.25
CAPACITY C = ceil(T×k/E × f) = ceil(8×2/4 × 1.25) = ceil(5) = 5

ROUTER on t1 (worked example reduced to 3 experts; the batch uses E = 4):
  logits [2.1, 1.8, 0.2]
  exp = [8.166, 6.050, 1.221]   sum = 15.437
  softmax = [0.529, 0.392, 0.079]
  top-2 → E1, E2; renormalized weights 0.529/0.921 = 0.574 and 0.426

ASSIGNMENTS (arrival order within the batch):
  t1 → E1,E2   t2 → E1,E3   t3 → E1,E2   t4 → E1,E4
  t5 → E1,E3   t6 → E1,E2   t7 → E2,E4   t8 → E3,E4

COUNTS        E1 = 6   E2 = 4   E3 = 3   E4 = 3   (total 16 = 8×2 ✅)

E2, E3, E4 ≤ 5 → accepted                              ✅
E1 = 6 > C = 5 → the 6th arrival (t6) is dropped       ❌
  (illustration assumption: t6 shares t1’s logits)
  t6 only gets E2: its output carries just 0.426 of the
  intended mixture; 57.4% of the expert signal is lost,
  and only the residual + shared expert rescue the token.

UTILIZATION: E1 5/5=100%  E2 4/5=80%  E3 3/5=60%  E4 3/5=60%
  slots reserved 4×5 = 20, used 15 → 25% padding paid over all-to-all

SHAPE CHECK: logits [T=8, E=4] → dispatch [E=4, C=5, d_model].
A [4, 4, d] buffer would refuse E1 at its 5th token: capacity is a
tensor dimension, not an optional policy.

PARAMETERS: 64 routed experts × 0.4B + 1 shared × 0.4B = 26B total
  active per token = (2 + 1) × 0.4B = 1.2B → ratio ≈ 21×

Overflow policies when E1 receives 6 tokens for C = 5

Policy What happens to t6 Real cost
Drop loses E1, keeps 0.426 of the mixture (t1’s logits assumed) degraded quality, stable latency
Reroute to E2/E3 handled by the 3rd-best score router decision overridden
Capacity f = 2.0 (C = 8) accepted, nothing dropped 32 slots for 16 tokens: 50% padding
Queue to next batch processed in the next micro-batch extra all-to-all, sawtooth latency

Causal lab

Predict → change one variable → run → explain the delta

/interactives/curriculum/moe-attnres.html?lang=en

Common errors

“Each expert specializes in a domain: math, code, French.”

The router emits learned logits ([2.1, 1.8, 0.2]), not labels. In the trace t1 and t3 both go to E1+E2 without sharing a topic, and t6 ends on E2 only because E1 was full. Observed specialization is statistical, often unreadable, and reshaped by the balancing loss alone.

“1.2B active out of 26B, so it is 21× faster than an equivalent dense model.”

The 21× ratio concerns FLOPs only. All 26B must sit in memory, the 20 reserved slots (25% of them padding) cross the all-to-all in both directions, and the step is paced by the busiest expert — E1 at 100% occupancy. In practice latency is often dominated by communication, not active compute.

Boundary, evidence, and sources

Activation names and exact expert counts attributed to a named model remain source-reported until primary evidence is attached.

Evidence status: Mixed: established mechanisms + source-reported Kimi K3-style choices.

  • Owner-supplied bilingual course packet, Chapter 14.
  • Shazeer et al., “Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer”, ICLR (2017).
  • Fedus, Zoph & Shazeer, “Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity”, JMLR (2022).
  • Dai et al., “DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models”, ACL (2024).
  • Course source packet supplied by the owner; named-product details remain source-reported until primary verification.

Transfer challenge

Rerun the trace’s batch (T = 8, E = 4, k = 2, f = 1.25).

  1. Pre-register: what happens to the drop count if k goes to 1 (f unchanged)? And if f goes to 1.5 (k = 2)?
  2. Compute C and the drops in both cases, reusing the trace’s assignments (first expert of each pair for k = 1).
  3. Conclude: which of the two levers buys quality, and in what currency?

Synthesis and exit ticket

  • Why experts
  • Router
  • Top-k and mixture
  • Capacity
  • Load balancing
  • Communication

Ticket: mechanism · trace · observation · boundary · evidence · next experiment

Instructor notes: Frame the problem before naming the mechanism. Collect an initial prediction and retain it for the exit ticket.

Instructor notes: Connect every step to the next with a causal verb. Flag any merely decorative arrow.

Instructor notes: Opening vote: “26B or 1.2B — which sizes your GPU order?”. Keep the tally posted: the beat shows both numbers matter, but on different budget lines.

Instructor notes: Open with a vote: “26B or 1.2B — which number do you quote your GPU vendor?” Both numbers have a use, and confusing them is what wrecks real budgets.

Instructor notes: Answer: load-time memory follows the 26B — every expert resides; FLOPs per token follow the 1.2B active. Expected wrong answer, worth harvesting explicitly: “the 26B determines FLOPs”. Both numbers matter, on two different budget lines.

Instructor notes: Have the “in GPU memory” row circled: it is the one that breaks the “MoE = small model” intuition. Return to this row in the communication beat, when latency joins the picture.

Instructor notes: Ask first: “how would YOU route the tokens?” — by language? by topic? Every proposed rule is rigid: exactly the intended contrast with the beat’s learned switchboard.

Instructor notes: Have the softmax computed by hand exactly once, then erase the scores and ask them to guess the token’s domain. The impossibility of answering is the lesson, not a failed exercise.

Instructor notes: Answer: logits [2.1, 1.8, 1.9] → softmax ≈ [0.391, 0.289, 0.320]: the top-2 becomes E1 and E3 — E3 replaces E2. And no, the router “understood” nothing: a learned score moved, the decision flipped. Expected wrong answer: hunting for a semantic explanation of the flip.

Instructor notes: Flash question: “why not k = 64? why not k = 1?”. Each extreme refutes itself in one sentence — dense cost on one side, fragile gradient on the other. The beat justifies the middle.

Instructor notes: Physically hand the 8 tokens to 4 learner-“experts” in the room, 5 chips each. Overflow is better lived than explained.

Instructor notes: Answer: 0.529/0.921 = 0.574 and 0.392/0.921 = 0.426, with 0.921 = 0.529 + 0.392. The shared expert adds outside the renormalization — always on, it competes with nobody: it processes every token. Expected wrong answer: renormalizing over all three scores.

Instructor notes: Stage it before the formula: 4 learner-experts, 5 chips each, 16 assignments dealt out per the trace. t6’s rejection must physically HAPPEN before being explained.

Instructor notes: Do not announce C = 5: derive the formula, then let the room discover that E1 wants 6. The drop must be a surprise they computed themselves.

Instructor notes: Answer: f = 1.0 → C = 4: E1 drops 2 tokens (6 − 4), 16 slots with 14 filled (2 empty); f = 2.0 → C = 8: zero drops, 32 slots with 16 used (16 empty). Expected misses: forgetting the ceil, or counting drops at E2-E4 — only E1 overflows.

Instructor notes: Have the counts checked row by row against the trace’s assignment list, then ask: “who decides t6 is the one to go, and not t1?” — arrival order, an implementation decision, not a quality one.

Instructor notes: Write 6/4/3/3 on the board and ask: “should we force 4/4/4/4?”. Let the for/against debate run two minutes — that debate is precisely the slider the beat formalizes.

Instructor notes: Frame the dilemma as a slider, not a setting: “push balancing to 4/4/4/4 — what did you break?” Have them name the loss before you supply the word “specialization”.

Instructor notes: Answer: if E1 genuinely is best for those tokens, forcing 4/4/4/4 sends them to worse experts — quality pays for uniformity. The balancing loss then fights useful specialization: it is a slider, not a goal. Expected wrong answer: “uniform is always the target”.

Instructor notes: Have the 4 GPUs drawn and t7’s full path traced (to E2, to E4, two returns). Count the crossings aloud BEFORE introducing the word “all-to-all”.

Instructor notes: Have the 4 GPUs drawn on the board and trace t7’s path (E2 then E4) there and back. Counting arrows aloud reframes the “MoE is free” debate in thirty seconds.

Instructor notes: Answer: two remote experts = two out + two back = four network crossings for ONE token; and the [E, C, d] buffer travels whole, padding included. The 1.2B active parameters measure compute, not the network. Expected wrong answer: “fewer FLOPs = faster”.

Instructor notes: Have them multiply: 8 tokens × 2 experts × 2 directions = 32 crossings for a single MoE layer step. Then ask what changes if all 4 experts share one GPU — the whole beat sits in that answer.

Instructor notes: Walk line by line. Locate an inconsistency at the first faulty step, not only on the final line.

Instructor notes: Have learners fill the final row before revealing it: that trade-off is what decides in production.

Instructor notes: Retain initial and final values. Do not allow simultaneous changes that make the delta impossible to attribute.

Instructor notes: For each claim, have the group produce the smallest counterexample before giving the correction.

Instructor notes: Separate verifiable mechanism, reported implementation choice, and experimental result. Evidence precision must match claim precision.

Instructor notes: Answers: k = 1 → C = ceil(8×1/4 × 1.25) = 3; the top-1 load becomes 6/1/1/0 and E1 drops 3 of its 6 — less compute, more breakage. f = 1.5 → C = 6: zero drops, but 24 slots for 16 assignments = 33% padding paid over the all-to-all. Expected conclusion: f buys quality (fewer drops) in the currency of padding and latency; k buys compute in the currency of mixture robustness. Check pre-registrations before any computation. Ten minutes.

Instructor notes: Rebuild the chain without looking at the slides, then fill the ticket in at most six lines. Compare with the opening prediction and name what actually changed.