Mixture-of-Experts Inference

Mixtral 8x7B computes each token with 13 billion parameters and stores 47 billion. That sounds like a model that is three and a half times cheaper to run than its size suggests, and for one token at a time it nearly is. On a layer built to Mixtral’s shape, one token chose two of eight experts and the layer took 3.95 milliseconds, reading only the experts it needed. Sixteen tokens chose between them all eight, and the same layer took 9.7 milliseconds, while a dense layer with the same active width took 1.75. The saving that sparsity promises is a saving per token. A batch pays for the union of what its tokens chose.

This chapter builds a mixture-of-experts layer, runs it four ways, and measures what the sparsity buys at each batch size and what the machinery for exploiting it costs. It also shows a consequence that a dense model does not have: once the layer has a fixed capacity per expert, a token’s output can depend on which other requests happened to be in the batch.

One limit has to be stated at the start. No trained mixture-of-experts model is cached on the book’s machine, and I did not download one. Everything below is therefore about the machinery: what must be stored, what must be read, how a layer can be executed, and how uneven routing behaves. None of it says anything about what a trained router learns, which experts specialise in what, or how balanced a real model’s routing is. Where a number depends on that, it is labelled synthetic.

Where this sits

Chapter 11’s three ceilings (launch, bandwidth, compute) decide which cost dominates, and Chapter 18 showed that at small batch the launch ceiling dominates and that a fixed-shape step can be recorded and replayed. Chapters 14 to 16 spent their effort on the memory a sequence occupies; this chapter is about the memory a model occupies, and what a batch makes the hardware read of it. Chapter 17’s speculative decoding verifies several tokens in one pass; in a mixture model each extra token can wake further experts, an interaction I have not measured and flag as an open question rather than a result.

The Inference Lab piece added here is the mixture layer with four dispatch strategies, an exact reference to test them against, a sparse-upcycling harness that converts GPT-2 into a mixture model, and the experiments. Chapter 20 treats routing and load balancing in depth; this chapter takes only as much of it as the cost measurements need.

For the Pi series, the idea to carry forward is about identity. In a dense model the same weights act on every token. In a mixture model the weights that act on a token depend on the token, so which experts ran is part of what was computed. A request’s output is a function of the weights, the input and the routing, and the routing is a runtime decision (and, with finite capacity, a decision that can depend on other requests).

What a mixture layer is

In a transformer block the feed-forward layer is replaced by \(E\) experts, each a feed-forward network of the usual shape, and a router that picks \(k\) of them per token. For a token representation \(x\) of size \(d\):

Total and active

Each expert holds \(2 d f\) parameters for a two-matrix MLP (\(3 d f\) with a gate: Mixtral’s experts are SwiGLU blocks). The layer stores \(E\) of them plus the router and uses \(k\) per token:

A and B hold the same number of parameters; B has more, smaller experts (the fine-grained design) and activates half as many parameters per token. C has Mixtral’s expert shape, with two-matrix experts, so 1.88 GB in fp16.

For whole models I rely on the published abstracts (this is arithmetic on their figures, not a measurement). Mixtral reports 47 billion parameters accessible and 13 billion active per token; DeepSeek-V2 reports 236 billion total and 21 billion activated. At two bytes a parameter:

Neither fits on the 12 GB card of this book. That is the first practical fact about mixture models: the memory bill is the stored size, and only the arithmetic bill is the active size. Everything that must be resident for any token to be servable has to be somewhere the server can reach, and the cost of making it reachable is what expert placement and parallelism, in the later chapters on distributed execution, are about.

What a batch makes the layer read

For one token, the layer needs only its \(k\) experts’ weights. For a batch, it needs the union. If each of \(T\) tokens picks \(k\) of \(E\) experts uniformly, an expert is missed by one token with probability \(1 - k/E\), so the expected number of experts a batch touches is

Real routers are not uniform, and I could not measure a trained one. What I could measure is the effect of the geometry of real hidden states on an untrained router: I took the inputs to block 6’s MLP from GPT-2 on 1,024 tokens of this book’s text, applied a random linear router, and drew random batches. Two things came out.

First, the untrained router was strongly skewed on real states: the busiest of 8 experts took 31.3% of all assignments (uniform would be 12.5%), and of 64 experts 7.7% (uniform, 1.6%). The residual stream of a real model is far from isotropic, and a random projection of it is not uniform. A trained router is a different object, and I do not infer anything about it from this, but the result is a warning against assuming uniformity.

Second, the skew lowers the number of experts a batch touches:

Skew delays full coverage, but only a little. A batch of 256 still touches 62 of 64 experts. This sets the regime in which sparsity pays: batch sizes small enough that the union is a fraction of the experts. For 64 fine-grained experts with 8 active that is up to a few dozen tokens; for 8 experts with 2 active, a handful.

Four ways to run the layer

All four compute the same function, and the reference is the definition: for each token, run its \(k\) experts and add the weighted outputs. It is a Python loop over tokens, slow and the oracle for tests.

• Run everything (dense_all): apply every expert to every token with batched matrix multiplies, then mask with the gate. Fixed shapes, \(E/k\) times the arithmetic, and it reads every expert whatever the routing.
• Per-expert loop (expert_loop): for each expert, gather the tokens that chose it, run them, scatter-add the results. Only the needed arithmetic and only the touched experts’ weights, but each step needs the number of tokens an expert received, which means a host-device synchronisation per expert, and the shapes depend on the data.
• Fixed capacity (capacity): give each expert a fixed number of slots, sort the (token, expert) assignments by expert, place each in its expert’s next free slot, run one batched matmul over the \([E, C, d]\) buffer, and scatter the weighted results back. All shapes are fixed, so it needs no synchronisation and can be recorded in a CUDA graph (Chapter 18); assignments beyond an expert’s capacity are dropped. The capacity is \(\lceil c \cdot Tk/E \rceil\) for a capacity factor \(c\), which is how the Switch Transformer paper defines it for \(k = 1\) (tokens per batch over experts, times the factor). With \(C = T\) nothing can be dropped, at \(E\) times the arithmetic.
• The reference.

The core of the fixed-capacity dispatch:

order = torch.sort(flat_e, stable=True).indices            # assignments grouped by expert
pos = torch.arange(T * k) - starts[flat_e[order]]          # rank within the expert
keep = pos < capacity
slot = torch.where(keep, flat_e[order] * capacity + pos, E * capacity)   # dropped ones go to a spare row
buf.index_copy_(0, slot, x[token_of])
y = torch.bmm(act(torch.bmm(buf[:-1].view(E, C, d), W1) + b1), W2) + b2
out.index_add_(0, token_of, y_flat[slot] * (gate * keep))

One detail came from Chapter 18’s discipline. My first version returned the number of dropped assignments as a Python integer, which reads a value back to the host inside the dispatch; CUDA graph capture refused it (“operation not permitted when stream is capturing”). It now returns a tensor.

Checking them

The unit tests compare all four against the reference on random layers, with and without shared experts, for \(k = 1\), \(k = E\), a single token, and experts that receive nothing. They also check that the capacity dispatch drops exactly the overflow when every token wants the same two experts, and that tokens within capacity are unaffected.

A stronger check uses a real model. Sparse upcycling (the idea, due to Komatsuzaki et al., is to start a mixture model from a dense checkpoint) copies the dense MLP into every expert. If the experts are identical, then whatever the router chooses and however it weights, the output of the layer is the dense MLP’s, so the whole model must reproduce the dense one. I did this to all 12 blocks of GPT-2 with 4 and with 8 experts and a random router, on a 17-token prompt. Every dispatch matched the dense logits to 1.1e-4 or better, with the same arg-max at every position. That verifies the plumbing (routing, gather, scatter, weighting, residuals) on real weights. It cannot verify that the router matters, because with identical experts it does not.

To exercise the router, I perturbed the 8 copies with 5% Gaussian noise so they differ. The mixture then departed from the dense model by up to 2.67 in logit and agreed with it on the arg-max at 94% of positions, so the routing decisions changed the answer, and the three optimised dispatches agreed with the reference to 9.2e-5. That is the test the layer has to pass: the same answer under every execution, on a model where the routing is not vacuous.

What it costs

The timing experiment runs each dispatch at several token counts on the GPU in float16, with the fixed-capacity one also recorded as a graph. Because this GPU is shared with a desktop and, during this chapter’s work, with another program’s model server, each point was measured in three separate passes, and I report the best pass’s median. For the fixed-shape graph the passes agreed to within 1.21 times (usually within 1.1 times), and each began with the GPU at or below 40% utilisation. An earlier run, made while that other server held 5.4 GB of the card’s memory and used it in bursts, showed slowdowns of an order of magnitude on some methods of the large layer and is discarded.

Small layers: the launch ceiling

Layers A and B, \(d = 768\), 75 MB of experts each in fp16. Milliseconds, the dense layer being an MLP with the same active width (graph):

(The dense columns are for A / B: their active widths differ.) The bandwidth floor for reading all the experts of either layer is 0.25 ms, and the floor for reading the touched ones at one token is 0.06 ms (A) or 0.03 ms (B). Nothing here is bandwidth-bound. The fixed-capacity dispatch takes about half a millisecond at every batch up to 64, regardless of how many experts are touched, because it is several dozen small kernels (router, top-k, sort, scan, scatter, two batched matmuls, an unfused activation, gather; I did not count them) whose individual costs set the time even when replayed from a graph. That is two and a half to nearly four times the dense layer with the same active compute. The per-expert loop is worse and gets worse with more experts: layer B’s 64 experts cost 64 host-device synchronisations, 11 to 49 milliseconds. The tiny layers show what the dispatch machinery costs when there is no weight traffic to save.

A Mixtral-shaped layer: the bandwidth ceiling

Layer C has \(d = 4096\), \(f = 14336\), eight experts: 1.88 GB of experts, against 470 MB for a dense MLP of the active width. Weight reading is now the cost: the floor at the machine’s measured 304 GB/s is 1.55 ms for two experts and 6.18 ms for all eight.

(Milliseconds; fixed capacity at factor 1.25.) Four things stand out.

1. At one token only the loop benefits from sparsity. It reads two experts and takes 3.95 ms (2.5 times its floor, the rest being launches and synchronisations). The fixed-shape dispatch and run-everything both read all eight experts, take about 6.7 ms, and sit within a tenth of the all-experts floor. The fixed-shape code has no way to skip an expert, because its shapes cannot depend on the data.
2. The advantage reverses by four tokens. With five of eight experts touched, the loop is already no faster than the dispatch that reads all of them, and beyond that it is slower: the loop pays for synchronisations, the fixed-shape dispatch only for arithmetic.
3. Against a dense layer with the same active compute, the mixture layer costs about four times as much once the batch has touched every expert (6.6 against 1.75 ms at 16 tokens). Same arithmetic per token, four times the weight traffic.
4. At 256 tokens the layer has become compute-bound, and the bookkeeping matters again: the fixed-capacity dispatch at factor 1.25 does 1.25 times the active arithmetic on smaller matrices, plus the sort and scatter, and takes 9.85 ms against the dense layer’s 5.12. Run-everything does four times the arithmetic and takes 22.3.

The consequence for serving is that a mixture model’s decode throughput does not come from its active parameters at any practical batch size. At batch one it can; at the batch sizes that make a GPU efficient it reads nearly everything, and what it gains over a dense model of the same stored size is arithmetic, not bandwidth. Where the mixture wins is quality per unit of training compute, which is outside this book; where it costs is memory capacity and, at serving batch sizes, memory traffic.

When routing is uneven

Real routers are not uniform, and a layer with fixed capacity has to decide what to do with the excess. Chapter 20 will go into routing and balancing; this section measures the three effects that the layer’s cost depends on.

How many assignments are dropped. An expert’s load across \(T\) tokens is binomial (each token includes it with probability \(k/E\), independently), so the expected overflow beyond a capacity is an exact sum. I computed it and simulated routing 400 times for each of 32 settings (8 and 64 experts, 16 to 1,024 tokens, capacity factors 1.0 to 2.0); the two agreed to within 0.0013 in every case. For 64 experts, top-8 and 64 tokens, the expected dropped fraction is 13.1% at capacity factor 1.0, 4.5% at 1.25, 1.2% at 1.5 and 0.03% at 2.0. Small batches are the hard case, since the load per expert is small and random: at 1,024 tokens, factor 1.0 drops 3.3% and factors of 1.25 and above drop essentially nothing.

What skew does. With synthetic Zipf-shaped popularity over 64 experts (exponent \(s\)), top-8 and capacity factor 1.25, the dropped fraction was 4.5% for uniform routing at 64 tokens, then 13 to 15% at \(s = 0.5\), 33% at \(s = 1.0\) and 48% at \(s = 1.5\), almost independent of batch size for \(s \geq 0.5\). The busiest expert’s load was up to 8 times the mean (the bound, \(E/k\), since one token cannot choose an expert twice).

Stragglers. If the 64 experts are spread over 8 devices, a layer waits for the busiest device. With contiguous placement the busiest device carried 1.04 to 1.17 times the mean load for uniform routing, but 2.2, 3.7 and 4.9 times for \(s = 0.5, 1.0, 1.5\); with round-robin placement, 1.0 to 1.2, 1.2, 1.5 and 1.5 times. The gap is an artefact of how the synthetic popularity was drawn: it falls with the expert’s index, so contiguous placement puts all the popular experts on one device and round-robin spreads them. Where a trained router’s popular experts sit is unknown here, so the lesson is only that placement matters and that it should follow the measured load, not the index.

A token’s answer depends on its neighbours

Dropping has a consequence specific to batched serving. When an expert is full, which assignments are dropped depends on which tokens arrived and in what order; so a token can lose an expert because other requests filled it. I measured it with 64 tokens per batch and compared each token’s output in the batch with its output when processed alone (where nothing is dropped):

(Fraction of tokens whose output differed; random router and inputs, 20 batches. Routing here is close to uniform, so this is the favourable case.) At a capacity factor of 1.25, a modest headroom over the average load, more than a quarter of tokens in the fine-grained layer got a different answer from the one they would have had alone. A request therefore does not have an output independent of its co-batched traffic. That qualifies the assumption, made since Chapter 6, that a request’s result is a function of the request, the artefact and the seed, and it makes testing and reproducibility depend on the batch. It would also undermine prefix caching as Chapter 16 built it, which assumes the keys and values of a prefix are a function of its tokens: in a capacity-limited mixture model the deeper layers’ keys and values for those tokens could depend on who shared their batch. Dispatch without dropping restores determinism of this kind, at the price of arithmetic proportional to the number of experts for the fixed-shape form, or of synchronisations for the data-dependent form. This is a property of capacity-limited dispatch; whether a given production system drops tokens at inference, and at what factor, I cannot say.

What the numbers do not show

• No trained model. All routing is untrained or synthetic. Specialisation, learned balance and the actual popularity of experts are unmeasured, and a trained model could be better or worse balanced than any case here.
• My dispatches, not an optimised library’s. Production systems use fused grouped matrix multiplies that avoid the padding and the synchronisations my versions have. The fixed-shape dispatch in particular pays for a sort and scatter that a fused kernel would not. The conclusions about shape (fixed shapes cannot skip experts; data-dependent shapes need synchronisation) are structural; the millisecond figures are for this code.
• One layer, one GPU. No expert parallelism, no interconnect, no all-to-all. The device-straggler numbers are a simulation of loads, not a measurement of communication.
• Two-matrix experts in float16. Gated experts have three matrices, 1.5 times the weights of layer C’s experts.
• Shared experts were tested for correctness but not timed.
• A shared, drifting machine. Pass-to-pass spread was at most 1.21 times, and the discarded early runs show what contention does.
• Quality. Nothing here measures what dropping tokens does to model quality.
• GPT-2 for upcycling. Its experts are copies; real experts differ, so the upcycling check proves the plumbing but not the routing quality.

Inference Lab increment

The Lab now has:

• route: top-\(k\) routing with a softmax over the selected logits.
• moe_reference, moe_dense_all, moe_expert_loop, moe_capacity: the four dispatches, with the last returning the dropped count as a tensor so that it can be captured in a graph.
• MoEParams and upcycle_layer, forward_moe: stacked expert weights, optional shared branch, and GPT-2 with its MLPs replaced by mixtures.
• expected_distinct_experts, expert_capacity, expected_drop_fraction, moe_param_counts, bytes_read_by_experts: the arithmetic.
• sample_routing, device_loads, imbalance: synthetic routing and placement.
• capture: records any function in a CUDA graph.
• layer, layer_large, coverage, upcycle, balance and batch_dependence commands.

The rule this adds: a mixture layer stores every expert and reads the union of what its batch chose; sparsity saves bandwidth only while that union is small, and the dispatch that exploits it needs data-dependent shapes. Every choice in the layer sits on that tension: fixed shapes give graphs and determinism but read everything; data-dependent shapes read less but synchronise; a finite capacity makes shapes fixed but makes a token’s answer depend on the batch.

What you should now be able to do

• Compute the stored and active size of a mixture model from its configuration, and say which determines the memory bill and which the arithmetic.
• Estimate how many experts a batch will touch and at what batch size a mixture layer stops saving bandwidth.
• Implement top-\(k\) routing and at least two dispatch strategies, and test them against a reference and by sparse upcycling.
• Explain why a graph-capturable dispatch cannot skip experts, and what a data-dependent one costs.
• Calculate the expected fraction of dropped assignments for a capacity factor, and say when dropping makes outputs depend on the batch.
• Say what a measurement on an untrained router can and cannot show.

Key terms

• Expert: one of the parallel feed-forward networks in a mixture layer.
• Router: the linear map from a token’s representation to a score per expert.
• Top-\(k\): the number of experts used per token.
• Total and active parameters: all the parameters stored, and those used for one token.
• Shared expert: an expert every token uses.
• Capacity (factor): slots per expert, set as a multiple of the average load; assignments beyond it are dropped.
• Dispatch: moving tokens to their experts and the results back.
• Sparse upcycling: initialising a mixture model by copying a dense model’s MLP into every expert.
• Straggler: the busiest device, which sets the time of an expert-parallel layer.

Reader exercises

1. Add a gated (SwiGLU) expert. How do the stored and active counts and the weight-read floors change, and where is the crossover with the loop?
2. Implement a grouped-matrix-multiply dispatch that sorts tokens by expert and runs one matmul per expert without padding or synchronisation (a segment-wise batched multiply). How close does it get to the loop’s reads and the fixed-shape dispatch’s synchronisation-free timing?
3. Time the mixture layer with a batch of 1 token verified with 4 speculative tokens (Chapter 17). How many extra experts does the verification wake?
4. Find, for the 64-expert layer, the capacity factor at which fewer than 1% of tokens depend on their neighbours, at 64 tokens per batch.
5. Use a different tie-break from the stable sort so that dropped assignments are chosen by lowest gate weight. How do the dropped fraction and the batch dependence change?
6. With a trained mixture model (not available here), measure the coverage table. How does it compare with the untrained router’s?
7. Place the 64 experts across 8 devices by the measured popularity from exercise 6 instead of by index. How much does the straggler factor fall?
8. Compute, for Mixtral’s published shape, how many tokens per batch make the layer compute-bound on this card, using the 304 GB/s and the roughly 24 TFLOP/s implied by this chapter’s dense layer at 256 tokens (0.12 TFLOP in 5.1 ms).

References

• Switch Transformers — Fedus, Zoph and Shazeer, 2021 (JMLR); top-1 routing, expert capacity as tokens per expert times a capacity factor, dropped tokens passed through the residual connection; read on 21 October 2026.
• DeepSeekMoE — Dai et al., 2024; fine-grained experts and shared experts; read on 21 October 2026.
• DeepSeek-V2 — DeepSeek-AI, 2024; 236B total and 21B activated parameters (abstract); read on 21 October 2026.
• Mixtral of Experts — Jiang et al., 2024; 8 experts, 2 per token, 47B total and 13B active parameters (abstract); read on 21 October 2026.
• Komatsuzaki et al., “Sparse Upcycling: Training Mixture-of-Experts from Dense Checkpoints,” 2022 — the upcycling idea, cited from memory and not re-read.

Evidence in this chapter

Implementation status

• Manuscript: content/books/inference/19-chapter.md
• Runnable examples: examples/ch19/moe.py (original toy layer), examples/ch19/moe_lab.py (dispatches, upcycling and experiments)
• Tests: tests/test_ch19_moe.py, tests/test_ch19_moe_lab.py (counting and capacity formulas against simulation, routing identity, every dispatch against the reference including shared experts, k = 1 and k = E, empty experts, exact overflow dropping, batch dependence, upcycling on a small GPT-2; a CUDA graph test runs only when a GPU is present)
• Evidence: research/ch19/sources.md, research/ch19/results.json, research/evidence-register.md
• Verification command: python -m unittest discover -s tests -v

The chapter is a draft with draft = false. It makes no claim about trained mixture-of-experts models, expert parallelism or any production dispatch kernel. Chapter 20 takes up routing algorithms and load balancing, the part of the picture this chapter had to treat as given.