transformers·llm-inference·kv-cache·mixture-of-experts

The forward pass: what actually happens to a token, and what it costs

Gaurav Nigam·July 28, 2026·22 min read

An open card-catalog drawer packed with blank index cards, connected by a heavy sagging copper cable to a rail of eight brass cylinders of which only two are lit, on a graphite table in raking window light.

Attention is a lookup whose dictionary gets re-read once per generated token — and after the first token, that re-reading, not the quadratic FLOP count everybody quotes, is what you are paying for.

Understanding LLMs From the Inside — part 2 of 5.

← Previous: Tokens: the alphabet an LLM actually reads · Next: Pretraining: the law everyone quotes governs a number nobody optimizes anymore →

29 milliseconds per token, and almost none of it is multiplication

In November 2022, ten authors at Google published Efficiently Scaling Transformer Inference and reported two numbers about the same system. On PaLM 540B, on TPU v4, with a 2048-token context: 76% model FLOPs utilisation while processing input tokens at large batch size, and 29 milliseconds per token at low batch size during generation.

Read those as measurements of one machine and they don't reconcile. Read them as measurements of two regimes and everything about how LLMs are priced falls out.

At 76% MFU the arithmetic units are close to saturated. That is a hardware-efficiency number most people writing CUDA never see, and it says the machine is doing what it was built for: dense matrix multiplication, at scale, with enough work in flight to keep the pipelines full. Then generation starts, and the same silicon spends 29ms producing a single token. Nothing about the model changed. What changed is that the work per byte moved collapsed.

Treat those specific figures as artefacts of a 2022 TPU v4 run on a 540B dense model, because that is what they are. The quantity that transfers to 2026 hardware is the ratio between the two regimes, not the millisecond count: as long as a single decode step has to read a parameter or a cached key roughly once per multiply it performs, arithmetic throughput cannot be the binding constraint, and no generation of accelerator has yet closed that gap. Faster chips move both numbers; they do not reconcile them.

The diagnosis is older than the measurement. Noam Shazeer wrote it down in November 2019, in the abstract of Fast Transformer Decoding: training these layers is fast because it parallelises across sequence length, but incremental inference "is often slow, due to the memory-bandwidth cost of repeatedly loading the large 'keys' and 'values' tensors." Memory bandwidth. Not FLOPs. Stated plainly, in a four-page paper, six years before "the KV cache is the bottleneck" became a thing people say at meetups.

So here is the claim this piece defends. After the first token, a transformer's forward pass is a bandwidth problem rather than an arithmetic one, and that single fact explains three things that otherwise look like unrelated engineering trivia: why attention is best read as a lookup, why parameter count turned into two numbers, and why long context costs what it costs.

A two-panel figure separated by a dashed vertical rule labelled "different quantity, different unit, heights not comparable". The left panel, headed PREFILL (reading the 2048-token input, large batch), is a vertical bar on a 0–100% axis of peak TPU v4 FLOPS; the bar reaches 76%, just under a dashed line at 100% marked "peak arithmetic throughput". The right panel, headed DECODE (emitting one token at a time, low batch), is a horizontal gold bar on a 0–60 millisecond time axis; the bar spans 0 to 29 ms and is labelled "29 ms — achieved low-batch generation latency", followed by a dashed gold outline of equal length labelled "and again for the next", indicating the cost repeats per token. Footer text notes that a ratio with a ceiling cannot be converted into a duration without one, and gives the source as arXiv:2211.05102, PaLM 540B on TPU v4 slices at 2048-token context with int8 weight quantisation. — **Two regimes, two units: prefill utilisation vs. decode latency in PaLM 540B inference.** The same system, measured two incompatible ways — 76% of peak arithmetic throughput while chewing through the 2048-token input, and 29 ms of wall clock for each token it then emits one at a time (Pope et al., "Efficiently Scaling Transformer Inference", arXiv:2211.05102, 9 Nov 2022, abstract; PaLM 540B, TPU v4, 2048-token context, int8 weights).

Attention is a lookup, and the dictionary is re-read once per token

The 2017 paper, Attention Is All You Need, describes an architecture "based solely on attention mechanisms, dispensing with recurrence and convolutions entirely." That sentence is usually read as a claim about elegance. It is better read as a constraint: attention is the only mechanism in the model that moves information between positions. Every MLP block, every normalisation, every nonlinearity operates on one position at a time. If token 400 is going to know anything about token 12, an attention head has to fetch it.

The fetch works like an addressable table. Each token's residual vector gets projected three ways: into a query, a key, and a value. Keys are addresses. Values are contents. The query is compared against every key by dot product, the scores are normalised, and the values are summed in proportion. That is the whole operation. It is a soft dictionary lookup where "soft" means you retrieve a weighted blend of all entries instead of exactly one. Multi-head attention, in the 2017 formulation, runs h of these lookups in parallel with separate projections and concatenates the results, so a single layer can perform several unrelated retrievals at once. Nothing here is archaeology: the models being served in production right now do exactly this, and the well-known variants are refinements of the same lookup rather than replacements for it. Multi-query attention shares one key/value head across all the query heads (Shazeer, 2019); DeepSeek-V3 uses multi-head latent attention. Both change what gets stored and re-read, not what the query does with it.

The residual stream is the bus this all hangs off. Each token carries one vector through the depth of the network. Every block reads that vector, computes something, and adds its output back. Nothing is overwritten. Attention blocks write retrieved information into the stream; MLP blocks write position-local transformations into it. By the final layer the stream has accumulated everything the model is going to use to produce logits over the vocabulary. The reason this matters for cost is that the stream is small — one vector per token per layer — while the machinery reading it is enormous.

Now run it one step at a time, which is what generation is. At decode step t, exactly one new token enters. It produces one query. That query is matched against t keys and blended with t values, all of which were computed on earlier steps and are sitting in memory. Recomputing them would be insane; Shazeer's 2019 framing is that we store them and re-read them, which is where the memory-bandwidth cost comes from.

Count the work. One query against t keys is a matrix-vector product. Small. Now count the bytes: every key and value for every layer and every head, for all t positions, has to travel from high-bandwidth memory to the compute units. Plus every weight in the model, because the MLP blocks and the projections all have to run too.

That ratio — a few FLOPs per byte loaded — is the whole story. The dictionary is large and resident. The lookup against it is cheap. And the dictionary gets re-read once per generated token, forever, for the length of the response.

A horizontal band across the top represents the residual stream, a single running vector for one token, flowing left to right with the note that every block adds to it and none replaces it. Three vertical taps drop from the band into three panels below, and three arrows return upward into plus-in-a-circle symbols sitting on the band. The first panel, ATTENTION, is labelled "a lookup, not a metaphor": step 1 reads a copy of the stream; step 2 builds a QUERY box — what this token is looking for, built from this token alone; step 3 compares that query against a table drawn to the right. The table, outlined in gold and titled "KV cache · one row per token seen so far", has columns token, K and V, and six rows: 1 "The", 2 "cat", 3 "sat", 4 "on", 5 "the", and a gold-highlighted row 6 "mat · new row". A note says rows 1–5 were stored on earlier steps and row 6 was appended on this step. Step 4 adds attention's answer back to the stream. The second panel, MLP, says it reads the same copy, transforms it using this token's vector alone, adds the result back, and sees no other token. The third panel, drawn with a dashed outline and titled "A block with nothing to add", says it still reads, it writes back almost nothing, and the vector arrives at the next block unchanged; its write arrow is dashed and labelled "adds ≈ nothing". A gold-bordered callout reads: why this is the expensive part — nothing in the table is recomputed, it is re-read, every stored row, once per token generated, and that re-reading is what long context costs. — **One transformer block as read–modify–write on the residual stream.** A single block never replaces the token's vector — attention reads a copy, builds one query for the current token, looks it up against a stored key/value row for every earlier token, and adds its answer back; the MLP does the same without any lookup, and a block with nothing to add leaves the vector untouched. Mechanism from Vaswani et al., arXiv:1706.03762 (2017); the cost framing — that incremental decoding is bound by repeatedly loading the keys and values tensors — from Shazeer, arXiv:1911.02150 (2019).

Every generated token re-reads the whole model, which is why the arithmetic units idle

Take a dense model at batch size 1. To produce token t+1, the hardware must load every weight matrix from HBM into the compute units, use each one exactly once against a single vector, and discard it. A matrix-vector product does roughly two floating-point operations per parameter. In 8-bit, a parameter is one byte. So decode runs at something close to two FLOPs per byte moved.

Modern accelerators are built for ratios two orders of magnitude higher than that. Feed them two FLOPs per byte and the multipliers wait on the memory system. This is why the 29ms figure and the 76% figure coexist without contradiction: prefill has a whole sequence of tokens to process at once, so each loaded weight does work proportional to sequence length, and the arithmetic intensity is high enough to saturate the units. Decode has one token. Same weights, one-token's worth of work.

Look at what the Google authors actually did to reach 29ms and the diagnosis confirms itself. The paper is explicit that the figure is obtained using int8 weight quantization. Quantising weights from 16-bit to 8-bit does not remove a single multiplication. It halves the bytes. If decode were arithmetic-bound, that change would buy nothing.

The same logic drove multi-query attention. Shazeer's 2019 proposal was to share one key/value head across all query heads, "greatly reducing the size of these tensors and hence the memory bandwidth requirements of incremental decoding." The reported result was that models "can indeed be much faster to decode, and incur only minor quality degradation from the baseline." Note the shape of that trade: it is a real quality cost, accepted in exchange for bytes. Nobody would take that deal to save FLOPs, because MQA saves almost none.

And the payoff compounds. The 2022 paper reports that with appropriate partitioning, the lower memory requirements of multi-query attention "enables scaling up to 32x larger context lengths." Thirty-two times the context, purchased entirely by shrinking a tensor. That is not a FLOP argument in any form.

The practical upshot for anyone sizing a deployment: at batch size 1, your tokens per second has a ceiling you can compute on the back of an envelope, and it is your device's memory bandwidth divided by the bytes of weights the model must traverse per token. Batching is what breaks that ceiling, because one weight load can serve many sequences. Everything else in serving is a variation on that theme.

Mixture-of-experts is a bandwidth trick, which is why parameter count became two numbers

Once you accept that per-token cost is dominated by parameters moved, an obvious exploit appears. Store far more parameters than you move.

That is the architecture. Switch Transformers, published January 2021 by Fedus, Zoph and Shazeer, opens by noting that models "typically reuse the same parameters for all inputs," and that mixture-of-experts "defies this and instead selects different parameters for each incoming example." The result, in their phrasing, is "a sparsely-activated model — with outrageous numbers of parameters — but a constant computational cost." They report up to 7× pre-training speedups at matched compute against T5-Base and T5-Large, and a 4× speedup over T5-XXL at trillion-parameter scale.

Which is why every MoE model card now carries two figures. Sebastian Raschka's architecture gallery MoE entry, added 14 March 2026, makes the mechanism explicit: as expert count grows, total parameters rise much faster than the parameters active per token. The DeepSeek-V3 technical report gives the canonical instance in its first sentence. 671B total parameters, 37B activated for each token, in a model built on Multi-head Latent Attention and DeepSeekMoE. That 18-to-1 ratio is the whole design, stated as a spec.

Here is how to read the pair, and it is not the way most coverage reads it. Total parameters tell you your memory footprint. Active parameters tell you your per-token arithmetic. All 671B must be resident, because the router can send the next token to any expert and you cannot page the unused ones out and hope. An MoE model is therefore cheap in FLOPs and expensive in VRAM, which is close to the exact opposite of the intuition "37B active" produces in someone used to dense models.

And neither number tells you the third bill. Switch Transformer's own abstract is candid that MoE adoption "has been hindered by complexity, communication costs and training instability." Communication cost is the one that survives into inference: experts live on different devices, and routing tokens to them means moving activations across an interconnect. That cost appears in no parameter count on any model card.

A two-panel diagram. Left panel, "Dense MLP block": a token vector arrow enters one solid gold FFN box labelled d_model → d_ff → d_model and exits to the residual stream; notes read "Loaded into memory: every weight in the block" and "Read for this token: every weight in the block." Right panel, "Mixture-of-experts block": the token enters a router box that scores N experts and selects k; below it a grid of eight expert boxes E1 to E8, of which only E3 is filled solid gold and marked READ, while the other seven are drawn as dashed ghost boxes; notes read "Loaded into memory: all N experts" and "Read for this token: the k the router picked. The other N − k stay resident and untouched." A bottom ledger compares the two: for dense, parameters stored = W and parameters read per token = W, with a fully gold bar meaning the whole layer's bytes move each token; for mixture-of-experts, parameters stored = N × W but parameters read per token = k × W, shown as a bar of eight segments with only one segment gold. Captions state that in the dense case stored and read are the same number so one count suffices, and in the sparse case cards must quote two counts, total (stored) and active (read). Footnotes mark the figure as schematic with N = 8 and k = 1 drawn for legibility, and quote Switch Transformers describing "a sparsely-activated model — with outrageous numbers of parameters — but a constant computational cost." — **Parameter count became two numbers.** A dense MLP block stores and reads the same weights, so one count describes it; a Switch-style mixture-of-experts layer stores all N experts but reads only the k the router picks for each token, which is why model cards must quote "total" and "active" separately — schematic after Fedus, Zoph & Shazeer, *Switch Transformers*, arXiv:2101.03961.

The KV cache is the bill for long context, and it is a per-user bill

The cache is the stored dictionary from section two: every key and value, for every layer and head, for every token in the sequence. Its size grows linearly in context length, linearly in layer count, and linearly in the number of KV heads. Multiply by batch size, because every concurrent request has its own.

That last property is the one that shapes pricing. Weights are shared — one copy in HBM serves every request on the device. The KV cache is not shareable across users, and it competes with the weights for the same memory. Adding a concurrent user with a long context does not just add compute; it takes HBM away from the thing that makes batching work.

This is why the entire optimisation literature points at those tensors rather than at attention arithmetic. Shazeer's 2019 target was shrinking "these tensors and hence the memory bandwidth requirements." The 2022 Google paper's 32× context gain came from MQA's lower memory footprint under appropriate partitioning, in a setting they describe as the hard one: 500B+ parameter models with tight latency targets and long sequence lengths. Grouped-query attention is the same lever with a gentler quality trade.

The 2026 direction of travel is precision. NVIDIA's NVFP4 KV cache work, published 5 December 2025, is worth reading in the original rather than through summaries: it positions FP8 as the outgoing standard and reports halving KV memory against it, which converts directly into roughly double the context budget at fixed memory. Latent-compression schemes in the MLA family attack the same quantity by projecting keys and values into a smaller shared space before caching. Secondary sources circulating in April 2026 quote further FP4-era kernel throughput figures that I could not confirm against a primary source; treat them as leads.

Which brings me to the open question, and it is not a small one. Every KV-compression win I can point to in these sources is a memory-footprint number. Gigabytes saved. Context multiplied. What the evidence base surfaced here does not contain is a held-out quality eval at the context lengths that motivate the compression. The missing ablation is specific and cheap: hold the model fixed, sweep KV precision and KV head count, and report task accuracy at 128k tokens and beyond — not just cache size. Shazeer at least reported his cost honestly in 2019 ("minor quality degradation from the baseline"). A 2026 press release that reports only bytes has not told you whether the trade was good.

I also cannot tell you how much of long-context pricing reflects KV bandwidth and how much reflects margin. No vendor primary sources were reached for this piece. The mechanism explains why long context must be expensive to serve; it does not audit anyone's price list.

A stacked bar chart with six bars for context lengths 4K, 16K, 32K, 64K, 128K and 256K tokens, plotting gigabytes of HBM occupied on one inference node at a fixed batch of 8 concurrent sequences. Each bar has an identical bottom segment of 140 GB labelled "140" for model weights, constant across all six bars. Stacked on top, in gold, is the KV cache: 11 GB at 4K, 43 GB at 16K, 86 GB at 32K, 172 GB at 64K, 344 GB at 128K and 687 GB at 256K, so totals run 151, 183, 226, 312, 484 and 827 GB. A dashed horizontal line at 640 GB marks an illustrative eight-by-80 GB device budget; the final bar pierces it, with the portion above the line outlined in dashes and annotated as 187 GB over budget. An annotation notes the KV cache overtakes the weights between the 32K and 64K bars, at about 53,000 tokens, and that the node fills at about 191,000 tokens. A note in the empty upper-left states that weights load once per node while every concurrent sequence carries its own KV cache. A footer gives the formula: KV bytes equals 2 times layers times KV-heads times head-dim times bytes per element times batch times context length, substituted as 2 x 80 x 8 x 128 x 2 bytes x 8 sequences, giving 2.62 MB per token of context, and flags the whole figure as schematic because no source in the evidence set supplies verified per-token KV bytes for a named deployment. — **The KV cache is the bill for long context — and it is charged per sequence.** A schematic stacked-budget chart: at fixed batch, model weights occupy a constant slice of a node's HBM while the KV cache grows linearly with context length, overtaking the weights at roughly 53K tokens and exhausting the node near 191K — computed from the KV-cache scaling relation with stated illustrative parameters (80 layers, 8 GQA KV-heads, head-dim 128, FP16, batch 8), not measured; relation and its bandwidth consequences per Shazeer, arXiv:1911.02150 and Pope et al., arXiv:2211.05102.

The case against "bytes, not FLOPs" — and why it's partly right

The strongest objection is in the same paper I opened with. That 76% MFU figure is real, and it means prefill genuinely is compute-bound. The arithmetic units saturate. For a workload dominated by input tokens — RAG over long documents, classification of large inputs, anything where you read far more than you write — the FLOP-shaped mental model is the correct one, and the quadratic term in attention is a real arithmetic cost that bites at long sequence length.

You can see the whole field agreeing with that objection. The survey Efficient Attention Mechanisms for Large Language Models, revised to v3 on 7 February 2026, still describes the quadratic time and memory complexity of self-attention as the fundamental obstacle to efficient long-context modelling, and organises the literature into linear attention and sparse attention. The linear-attention branch exists precisely because attention arithmetic becomes binding at extreme context. If bytes were the only story, that branch would be a solution to a non-problem.

Two more concessions. High batch sizes change the regime: one weight load amortises across many sequences, so the weight-bandwidth term shrinks per token while KV traffic and attention work scale with batch × length. Production serving is a mixed regime, and the 2022 paper's own method — an analytical model used to select partitioning strategies per application — is an admission that there is no single answer. And a separate warning worth keeping: not every efficiency technique makes the same trade. FlashAttention-style tiling changes the memory access pattern and returns mathematically identical output. GQA and cache quantization change the answer. A piece that files them under one heading has confused an exact optimisation with an approximation.

The deepest objection is that arithmetic-to-bandwidth ratio is a property of 2026 accelerators, not of transformers. "Bytes, not FLOPs" is a claim about hardware. If memory bandwidth grows faster than peak FLOPs in the next generation, the slogan inverts, and anyone who memorised it will have overfit to a hardware moment rather than learned a mechanism. The July 2026 ACL Findings work on attention internals in masked diffusion models is a reminder that even the autoregressive framing has a shelf life. Learn the ratio, not the slogan.

How to check this on your own hardware this week

Two measurements, plotted separately, on any endpoint or local runtime. Sweep input length and record time-to-first-token and inter-token latency as distinct series. TTFT should climb faster than linearly, because prefill is arithmetic and attention has a quadratic term. Inter-token latency should climb roughly linearly with cached tokens, because that is KV bytes being re-read. Then turn on prefix caching: TTFT should collapse while inter-token latency barely moves, because the cache still has to be read even when it didn't have to be built.

If inter-token latency is flat in context length on your stack, the thesis is wrong there. Say so. That is the falsification condition, and it is worth publishing when you find it.

Two ablations sharpen it. Swap MHA for GQA at fixed weights: decode speed should move considerably more than prefill speed. Swap a dense model for an MoE of matched active-parameter count: if bytes bind, decode latency tracks active parameters while memory footprint tracks total.

And a back-of-envelope to anchor all of it: at batch 1, divide your device's HBM bandwidth by the bytes of weights traversed per token. If measured tokens/sec lands near that, you are bandwidth-bound. Raise batch size until throughput stops tracking bandwidth; that crossover is where your serving economics actually live. Quantize the weights but not the cache, and watch it move.

Which is the question to put to a vendor, or to your own capacity plan: tokens per second at my batch size and my context length. Not parameter count.

Dated 26 July 2026. The mechanism here is background — 2017, 2019, 2021, 2022. What has shifted in the last six months is the accounting. Part 3 takes up where the weights came from in the first place.

The series

Tokens: the alphabet an LLM actually reads
The forward pass: what actually happens to a token, and what it costs (you are here)
Pretraining: the law everyone quotes governs a number nobody optimizes anymore
Post-training: where behavior comes from
Making it yours: distillation and fine-tuning a domain model

— Gaurav

ShareX LinkedIn

The forward pass: what actually happens to a token, and what it costs

29 milliseconds per token, and almost none of it is multiplication

Attention is a lookup, and the dictionary is re-read once per token

Every generated token re-reads the whole model, which is why the arithmetic units idle

Mixture-of-experts is a bandwidth trick, which is why parameter count became two numbers

The KV cache is the bill for long context, and it is a per-user bill

The case against "bytes, not FLOPs" — and why it's partly right

How to check this on your own hardware this week

The series

New essays by email

Responses

Making it yours: distillation and fine-tuning a domain model

Post-training: where behavior comes from