Inference Engines

It is September 2023. ChatGPT has been public for ten months, every major lab is racing to serve transformer models to millions of concurrent users, and a team at UC Berkeley publishes a paper that treats a language model’s memory the way an operating system treats RAM — in pages. The idea it introduces, together with a quieter fix from four years earlier that most of the industry had ignored until inference bills forced the issue, turns serving a large language model from a batch-of-one guessing game into an engineering discipline with its own arithmetic.

Where We Were

By the early 2020s, the transformer’s training-time cost had a well-worked-out playbook: mixed precision, sharded optimizer state, tensor and pipeline parallelism, and IO-aware kernels all attacked one side of the ledger — the one-time cost of producing a trained model. Serving that model, generating text one token at a time for a paying user, is a different problem, and by the time GPT-3-class models were answering millions of requests a day, it had become the more expensive one at the margin.

A transformer trained with causal masking can compute an entire training sequence in one parallel pass — the property Chapter 5 of the language-models book credits for the architecture displacing recurrent networks. Generation does not get that parallelism for free. Producing token $t+1$ requires the model’s output at token $t$, so autoregressive decoding is inherently sequential: one forward pass per output token, with no way to compute token 50 before token 49 exists. Recomputing every previous token’s keys and values from scratch on each step would make that sequential cost scale quadratically with output length, so every practical implementation caches them instead: after each forward pass, the per-layer key and value vectors for the tokens processed so far are stored in GPU memory, and the next step reads that cache rather than recomputing it. This key-value cache, or KV cache, is what makes incremental decoding tractable — and what makes it a memory problem rather than a compute problem. Generating one new token touches nearly the model’s entire set of weights and its entire growing cache for a tiny amount of arithmetic per token, which is the textbook memory-bound workload Chapter 4 already flags as the reason a chip’s peak FLOPs figure barely predicts single-request decoding speed.

Three separate problems compounded on top of that bandwidth ceiling as serving moved from research demos to production traffic. First, the cache itself is large and grows with every token generated and with every concurrent request a server tries to batch together, so a server juggling many users at once can run out of memory long before it runs out of compute. Second, early serving systems allocated each request’s cache as one large contiguous block sized for the longest sequence the request might ever reach, which — because requests finish at different, unpredictable lengths — left most of that reserved memory sitting empty for the request’s entire lifetime, the same kind of internal fragmentation operating systems solved for general-purpose RAM decades earlier. Third, the one-token-per-forward-pass constraint meant a large model’s full latency and cost was paid on every single step, even though many of the tokens a model outputs are, in hindsight, easy to predict — sequential decoding spent a full expensive pass on tokens a much cheaper model could often have guessed correctly.

Each of the three problems above got its own answer, on its own timeline, and none of the answers required retraining a model from scratch or changing what it computes — which is precisely why they could be adopted by serving teams who had no say over how a model was trained in the first place.

The Breakthrough

The earliest fix targeted the cache’s size directly, and arrived years before the problem it solved was widely felt. Shazeer observed that standard multi-head attention keeps a separate set of key and value projections for every attention head, so the cache a decoder must store scales with the number of heads even though the actual bottleneck during incremental decoding is not arithmetic but the bandwidth needed to repeatedly load those keys and values (Shazeer, 2019). His proposed variant, multi-query attention (MQA), keeps the usual multiple query heads but collapses the key and value projections down to a single shared set used by every head. Queries still differ per head, so the model retains most of multi-head attention’s representational flexibility on the query side, but the tensors that have to be cached and repeatedly streamed from memory shrink by a factor equal to the number of heads. On WMT14 English-to-German translation, Shazeer reports per-token decoder inference dropping from 46 microseconds with standard multi-head attention to 3.8 microseconds with MQA at a sequence length of 128 — roughly a 12x reduction — with beam-4 search dropping from 203 to 32 microseconds, and translation quality essentially unchanged (26.7 versus 26.5 BLEU on the validation set, with the MQA model actually scoring slightly higher on the beam-4 test set, 28.5 versus 28.4). Because it changes only how keys and values are projected, MQA is a drop-in architectural choice, not a serving trick — but through the mid-2020s it saw limited production adoption outside models built with it from the start, in part because collapsing all the way to one shared key-value set is a large enough change to a model’s inductive structure that teams were reluctant to apply it to an already-trained checkpoint.

Ainslie et al. closed that gap in 2023 with two contributions bundled into one paper. First, an uptraining recipe: rather than requiring a model to be pretrained with MQA from the start, an existing multi-head checkpoint can be converted by averaging its existing key and value heads down into a shared set and then continuing training briefly, at roughly 5% of the model’s original pretraining compute — for a model comparable to a T5-XXL checkpoint, about 600 TPUv3 chip-days (Ainslie et al., 2023). Second, and more consequentially for later adoption, grouped-query attention (GQA): rather than collapsing all key-value heads down to exactly one, GQA partitions the query heads into a small number of groups, each sharing its own set of keys and values — a tunable middle point between full multi-head attention (as many groups as heads) and MQA (exactly one group). The uptraining recipe applies to any point on that spectrum, letting a team pick how aggressively to trade cache size against quality rather than being locked into MQA’s single-group extreme. Because GQA is a drop-in attention variant reachable from an already-trained model with a modest compute budget, it became the default choice for new model releases far faster than MQA had: it is one of the reasons Mistral 7B shipped competitive inference speed in an open 7-billion-parameter model, and it is the attention pattern behind the larger models in the Llama 2 family, which use grouped-query attention specifically in the 34-billion- and 70-billion-parameter configurations while the smaller 7B and 13B models keep standard multi-head attention.

Shrinking the cache and managing the memory that holds it are different problems, and Kwon et al. addressed the second directly. They observed that existing serving systems allocated each request’s KV cache as one contiguous slab sized for a worst-case sequence length, and measured that in practice only 20.4% to 38.2% of that reserved memory ever held real token data — the rest sat wasted to fragmentation and over-reservation, exactly the failure mode virtual memory paging was invented to solve for operating systems (Kwon et al., 2023). Their answer, PagedAttention, stores each sequence’s KV cache in fixed-size blocks — 16 tokens by default — that need not be contiguous in physical GPU memory, with a per-request block table mapping logical cache positions to wherever their physical blocks actually live, mirroring how an OS maps virtual pages to physical ones. On top of that algorithm they built vLLM, a serving engine that combines paged allocation with continuous batching — admitting and evicting requests from a running batch at the granularity of individual decoding steps rather than waiting for an entire batch to finish together, a scheduling idea earlier explored by the Orca serving system — so that a GPU’s compute is rarely left idle waiting for the slowest sequence in a fixed batch to finish.

The third problem, sequential latency, got an answer that owes nothing to how attention is computed at all. Leviathan, Kalman, and Matias proposed speculative decoding: run a small, cheap draft model forward several steps to guess a short run of upcoming tokens, then run the large target model once, in parallel, over that entire drafted sequence, and use a specific accept/reject rule to decide which of the draft’s guesses to keep (Leviathan, Kalman, and Matias, 2022). The rule is built to guarantee the output distribution is identical to what the large model would have produced by itself, token for token — this is lossless acceleration, not an approximation. Where the draft and target models agree, several tokens are accepted from a single expensive forward pass instead of one; where they disagree, the target model’s own corrected sample is substituted and drafting resumes from there. On T5-XXL, the paper reports a 2x to 3x wall-clock speedup over the standard T5X implementation with byte-for-byte identical outputs.

None of these four ideas depends on the others, which is why production serving stacks by the mid-2020s stacked all of them at once: a model trained or uptrained with GQA, served through a paged, continuously batched engine, with a small draft model speculating ahead where latency allows.

The Mathematics

The KV cache’s size is the number every technique in this chapter is ultimately negotiating with. For a decoder with $L$ layers, storing keys and values for $n_{kv}$ key-value heads of dimension $d_h$, across a sequence of length $s$, for a batch of $b$ concurrent sequences, at $p$ bytes per stored value:

$$M_{\mathrm{kv}} = 2 \cdot L \cdot n_{kv} \cdot d_h \cdot s \cdot b \cdot p$$

Symbol	Meaning
$M_{\mathrm{kv}}$	KV cache memory, in bytes
$L$	number of decoder layers
$n_{kv}$	number of key-value heads (groups)
$d_h$	dimension of each attention head
$s$	sequence length so far, tokens
$b$	batch size, concurrent sequences
$p$	bytes stored per cached value (precision)

Every term in that formula scales linearly except $n_{kv}$ — and $n_{kv}$ is exactly the term MQA and GQA target. In standard multi-head attention $n_{kv}$ equals the number of query heads, $n_h$; MQA and GQA decouple the two, letting $n_{kv}$ be smaller:

$$\frac{M_{\mathrm{kv}}^{\text{MHA}}}{M_{\mathrm{kv}}^{\text{GQA}}} = \frac{n_h}{n_{kv}}, \qquad n_{kv} = 1 \ \text{is the MQA special case}$$

Speculative decoding’s payoff is a different kind of arithmetic: not how many bytes are stored, but how many expensive forward passes are needed per output token. Leviathan, Kalman, and Matias’s accept/reject rule guarantees that, for a draft model proposing $\gamma$ tokens ahead with per-token acceptance probability $\alpha$, the expected number of tokens produced per invocation of the large target model is:

$$\mathbb{E}[\text{tokens per step}] = \frac{1 - \alpha^{\gamma+1}}{1-\alpha}$$

That token count alone overstates the real speedup, because the draft model itself is not free. Accounting for its cost relative to the target model’s, with $c$ the ratio of a single draft-model run’s time to a single target-model run’s time, the paper’s walltime improvement factor is:

$$\text{Speedup} = \frac{1 - \alpha^{\gamma+1}}{(1-\alpha)(\gamma c + 1)}$$

Symbol	Meaning
$\alpha$	per-token acceptance probability of the draft model’s guesses
$\gamma$	number of tokens the draft model speculates ahead per round
$c$	cost ratio, one draft-model step’s time divided by one target-model step’s time
$\mathbb{E}[\cdot]$	expected tokens accepted per invocation of the target model

What It Enabled

The cache-shrinking and memory-management threads compose, and the numbers from the papers that introduced them show why serving engines adopted both rather than choosing one.

Attention pattern	KV-cache size vs. multi-head	Inference time per sample, T5-XXL scale (avg. across 7 eval tasks)	Average score across eval tasks
Multi-head (baseline)	1x	1.51 s	47.2
Multi-query	$1/n_h$	0.24 s	46.6
Grouped-query, 8 groups	$8/n_h$	0.28 s	47.1

Source: Ainslie et al., 2023, Table 1. GQA-8 recovers almost all of multi-head attention’s quality (47.1 vs. 47.2 average score) at nearly MQA’s speed (0.28 s vs. 0.24 s per sample) — the middle point the paper’s tunable group count was built to find.

Serving system compared	vLLM throughput advantage	Reported by
FasterTransformer	up to 22x	Kwon et al., 2023
Orca, oracle scheduling	1.7x–2.7x	Kwon et al., 2023
Orca, max-batch scheduling	2.7x–8x	Kwon et al., 2023
Shared-prefix workload, LLaMA-13B, five-shot	3.58x	Kwon et al., 2023

The paper’s own headline figure — vLLM improving throughput 2x to 4x over prior state-of-the-art systems at matched latency — is the aggregate of gains like these, concentrated in exactly the requests earlier systems handled worst: long sequences, large models, and complex decoding strategies like beam search or parallel sampling, where fixed contiguous allocation wasted the most memory.

These techniques close a loop this book opens as early as Chapter 4: training-time optimizations buy a model’s existence, but a model’s ongoing cost is paid per token, per user, indefinitely, and that cost is disproportionately a memory-bandwidth cost rather than a compute one. Chapter 18 covers architectures that sidestep the KV cache’s growth altogether by replacing attention with a fixed-size recurrent state, which is the more radical answer to the same accounting problem this chapter treats with attention kept intact. Chapter 19 picks up the economics directly: a reasoning model that emits a long chain of thought before answering is, from a serving engine’s point of view, a request whose sequence length — and therefore whose cache size and decode cost — has grown by exactly the length of that reasoning trace, and the long-context and reasoning chapters of the language-models book both point back to this chapter’s cache arithmetic as the bill that capability comes with. Chapter 20 treats the whole arc — cache-shrinking attention variants, paged memory management, and speculative execution — as one instance of a recurring pattern: a fixed hardware generation’s bandwidth ceiling provokes software techniques that buy a serving fleet several more years of useful life before the next accelerator generation is needed.

Key Papers

Paper	Authors	Year	Link	Why it matters
Fast Transformer Decoding: One Write-Head is All You Need	Shazeer	2019	arXiv:1911.02150	Introduced multi-query attention; shared keys and values across heads to shrink the KV cache and its bandwidth cost
Fast Inference from Transformers via Speculative Decoding	Leviathan, Kalman, and Matias	2022	arXiv:2211.17192	Draft-then-verify decoding that accelerates generation losslessly by computing several tokens per expensive forward pass
GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints	Ainslie et al.	2023	arXiv:2305.13245	Generalized MQA into a tunable group count and showed existing checkpoints could be uptrained into it cheaply
Efficient Memory Management for Large Language Model Serving with PagedAttention	Kwon et al.	2023	arXiv:2309.06180	Paged, non-contiguous KV-cache allocation plus continuous batching, built into the vLLM serving engine

PM Lens