#Fast inference

Sampling picks a token from the model's probabilities, and that part is cheap. Producing those probabilities, one full pass through the model for every token, is the expensive part. Training a model happens once. Inference, running it to answer requests, happens billions of times, so most engineering effort in LLM serving (running models for many users at once) goes into making each token cheaper and faster.

This chapter starts with the hardware, because almost every trick in it follows from one fact about how a GPU spends its time. Then it covers the main tricks in roughly the order a request meets them.

#Why GPUs at all

A transformer is mostly large matrix products: every layer multiplies token vectors by weight matrices, and attention multiplies queries by keys. A GPU (graphics processing unit) is a chip with thousands of small cores built to do many multiply-adds at once. A CPU, the computer's main processor, has a handful of big cores built to do one thing after another quickly.

Your browser can run code on your GPU directly through WebGPU. The figure below multiplies two random matrices twice: once with a plain JavaScript loop on the CPU, and once with a tiny compute shader, a small program that runs on the GPU, which gives every output cell its own thread (one of thousands of workers running in parallel).

Figure 10.1 · Matrix multiplication on your GPU (WebGPU)

The speeds are in GFLOP/s: billions of floating-point operations (arithmetic on fractional numbers) per second. On most machines even this naive shader beats the CPU by a wide margin. A data-center GPU running tuned kernels is thousands of times faster still.

#Compute and memory

The figure measured only the arithmetic. But numbers have to reach the cores before they can be multiplied, and on a GPU that trip is often the slower part. Two parts of the chip matter here:

• Compute: the cores. An H100, NVIDIA's data-center GPU, does about 1,000 trillion operations per second (989 TFLOP/s in 16-bit arithmetic).
• Memory: where the weights and everything else live. The main memory, HBM (high-bandwidth memory), holds 80 GB on an H100. The chip also has a little SRAM, fast memory right next to the cores: only a few tens of megabytes in total, but much quicker to reach.

Memory has two different numbers, and this chapter uses both. Capacity is how much fits: 80 GB. Bandwidth is how fast it can be streamed into the cores: 3.35 TB/s (terabytes per second) for HBM.

Divide one speed by the other: 989 TFLOP/s ÷ 3.35 TB/s ≈ 300. In the time it takes one byte to arrive from HBM, the cores can do about 300 operations. So what decides the speed of a computation is its arithmetic intensity: how many operations it does for each byte it reads.

• Well above 300 operations per byte, the cores are the limit. The work is compute-bound.
• Below that, the cores sit waiting for data. The work is memory-bound: a chip with faster arithmetic would not help, only reading fewer bytes would.

Two small cases show the difference. Multiplying a matrix by another matrix reuses every number it reads many times, once for each row of the other matrix, so its intensity is high. Multiplying a matrix by a single vector uses each number once: a 16-bit weight is 2 bytes and takes part in one multiply and one add, so 2 operations for 2 bytes, an intensity of about 1. Keep these two cases in mind. They are exactly the two phases of inference.

#Two phases: prefill and decode

A request is processed in two very different phases:

1. Prefill. The whole prompt goes through the model in one pass. All prompt tokens are processed in parallel, so every weight that is read is used once for each prompt token. That is a matrix times a matrix, and the GPU is compute-bound and busy. This phase decides the time to first token.
2. Decode. New tokens come out one at a time, because each depends on the one before. Every step pushes just the newest token through all layers (the next section shows why the older tokens don't have to go through again). That is a matrix times a single vector, with an intensity around 1. This phase decides the tokens per second you see streaming.

So decode is the slow part, and the reason is memory, not math. Before measuring how slow, we need the data structure that makes "just the newest token" possible.

#The KV cache

Say the prompt is "The cat sat on" and the model has just produced "the". To produce the next token, it runs "the" through every layer. In each layer, attention compares the new token's query with the keys of all five tokens so far, and mixes their values by the resulting weights. Where do those five keys and values come from? Each is a learned matrix times that token's vector at that layer, and that vector is itself the output of all the layers below.

Without a cache, everything is recomputed. To rebuild the keys and values of "The cat sat on", the model would have to run those four tokens through all the layers again. On the next step it would rerun five tokens, then six. In effect it would redo the whole prefill for every new token, and the total work would grow with the square of the length.

Keeping them is safe. Attention is causal: the mask lets a token look only at itself and earlier tokens. So when "the" arrives, nothing that "The cat sat on" computed changes, in any layer. Their keys and values are exactly what they were one step ago. The KV cache stores every token's keys and values in every layer, and each decode step adds just one new row per layer.

Keys and values, but not queries. A query is the question a token asks while computing its own output: "which tokens matter to me?" The old tokens have already asked theirs. Their outputs are done and, thanks to the mask, will never change, so their queries are never needed again. Keys and values are the opposite: they are what later tokens read, so every future token needs them. The same goes for everything else a layer computes, such as the MLP outputs: only keys and values are ever looked at again, so only they are kept.

Figure 10.2 · Decoding with and without a KV cache

With KV cacheWithout cache

tokenKV

Theq

·catq

·satq

·onq

·the

·warm

·mat

.

computed now recomputed (wasted) read from cache

4K/V rows this step4total with cache4total without4rows kept in cache

1/5

Prefill. The whole prompt goes through the model in one parallel pass: 4 key and value rows are computed at once. Every prompt token has a query (q) here, because each one needs its own output. This is compute-heavy and fast per token.

The cache turns quadratic recomputation into linear work, but it costs memory, and that memory grows with every token of every conversation. For a model with $L$ layers, $h_{kv}$ key/value heads of $d_h$ numbers each, a context of $n$ tokens and $b$ bytes per number:

$$\text{KV cache} = 2 \times L \times h_{kv} \times d_h \times n \times b$$

The factor 2 is for keys and values. Llama 3 8B has $L = 32$ layers and $h_{kv} = 8$ key/value heads of $d_h = 128$ numbers, stored in $b = 2$ bytes: $2 \times 32 \times 8 \times 128 \times 2 = 131{,}072$ bytes, or 128 KB per token. An 8,000-token conversation holds about 1 GB of cache, on top of the 16 GB of weights. Serve a hundred conversations and the cache, not the model, fills the GPU's 80 GB.

#Sharing keys and values: GQA and MQA

One factor in that formula can be shrunk without shortening the context: the number of key/value heads. Standard multi-head attention (MHA) gives every query head its own key and value head. Multi-query attention (MQA) shares a single key/value head among all query heads. Grouped-query attention (GQA) sits in between: query heads are split into groups, and each group shares one key/value head. Llama 3, for example, has 32 or 64 query heads but only 8 key/value heads. Queries are not cached, so having many of them costs no memory.

Figure 10.3 · KV cache size: MHA vs GQA vs MQA

MHAGQAMQA

Llama 3 8BLlama 3 70B

Layers32Query heads32Context8kBatch1

fp16fp8int4

2 × 32 layers × 8 KV heads × 128 dims × 8,192 tokens × 1 × 2 B

Multi-head (MHA)4.00 GB

Grouped-query (GQA)1.00 GB

Multi-query (MQA)128 MB

128 KB of cache per token with GQA. Head size is fixed at 128.

The KV cache is valuable enough to deserve its own chapter: reusing it across requests (prompt and prefix caching), storing it in pages (PagedAttention), and compressing or evicting it are all in Chapter 11: Caching. Chapter 14 shows multi-head latent attention, which shrinks the cache further by compressing keys and values into a small latent vector.

#Why decoding is memory-bound

Back to speed. To produce one token, every weight of the model has to travel from HBM into the cores, and for one token each weight is used for just one multiply-add: an intensity of about 1. A model with 8 billion parameters stored in 16-bit floats is 16 GB, so each decode step reads 16 GB. How fast that can happen is set by the memory bandwidth:

$$\text{tokens per second} \lesssim \frac{\text{memory bandwidth}}{\text{bytes of weights}} = \frac{3.35\ \text{TB/s}}{16\ \text{GB}} \approx 200$$

That is an upper bound for an 8B model on an H100 at batch size 1 (one sequence at a time). The KV cache has to be read at every step too, so the real number is lower. Meanwhile the cores sit almost idle: they could do about 300 operations per byte, and decode gives them about one.

The roofline picture makes this visible. The x-axis is arithmetic intensity, the y-axis the speed you can reach. Left of the ridge (the bend, at about 300 operations per byte) bandwidth is the limit, so the line climbs; right of it, raw compute is, so the line is flat.

Figure 10.4 · The roofline: why one user can't keep a GPU busy

Batch size1

fp16fp8int4

Memory-bound: the math units wait for weights to arrive.

209tokens/s for each user209tokens/s for the whole GPU0.3%of the GPU’s arithmetic in use

Almost every trick below is a way to read fewer bytes per token, or to get more tokens out of every byte read:

• Batching many users together: the weights are read once and used for every sequence in the batch.
• Quantization: store weights in 8 or 4 bits instead of 16 (Chapter 13).
• A smaller KV cache: GQA above, and the tricks in Chapter 11.
• Speculative decoding: check several tokens per read of the weights (below).

#Batching: static vs continuous

Batching is the biggest lever, but requests are awkward to batch: they arrive at different times and want different numbers of tokens. The naive scheme, static batching, groups a few requests, runs them together until all are done, then starts the next group. A three-word answer has to wait in its slot for the essay next to it.

Continuous batching (also called in-flight or iteration-level batching) schedules at the level of single decode steps. After every step, finished sequences leave and waiting ones join, so a slot never idles while there is work in the queue.

Figure 10.5 · Static vs continuous batching

Static batchingContinuous batching

slot 1AE

slot 2BF

slot 3CG

slot 4DH

WaitingE·3F·8G·2H·3Finished

0steps so far0%slots doing useful work…average finish time

1/18

Step 0. A batch only ends when its longest request ends; until then, freed slots stay empty (hatched).

Every modern serving engine (vLLM, SGLang, TensorRT-LLM, TGI) batches continuously. New requests also need a prefill, which is much heavier than a decode step; engines split long prompts into chunks (chunked prefill) and mix them into decode batches so that streaming users don't stall while someone pastes a book.

#Speculative decoding

A big model reads all its weights to produce one token. Because decode is memory-bound, checking several proposed tokens costs it about the same as producing one: the weights are read once and used for every position, just like in prefill. Speculative decoding exploits that:

1. A small, fast draft model guesses the next few tokens.
2. The big model runs once over all guesses and computes what it would have produced at each position.
3. Keep the guesses up to the first disagreement, then take the big model's own token at that spot.

Figure 10.6 · Speculative decoding

Text

The·quick

Draft

0tokens produced0big-model passes0draft passes (cheap)0big passes without drafting

1/19

Two models: a small, fast draft model and the big model whose output we actually want.

When the draft is usually right, the big model emits several tokens per pass. With a careful accept/reject rule the output follows exactly the big model's probability distribution, so speculation is a pure speed-up with no change in quality. Typical gains are 2–3×. Variants drop the separate draft model and instead add extra prediction heads (small extra output layers, not attention heads) to the big model itself (Medusa, EAGLE, multi-token prediction).

#FlashAttention: attention without the big matrix

Attention compares every query with every key, producing an $n \times n$ matrix of scores. For a 32,000-token context that matrix has a billion entries per head, far too big to fit in SRAM. The straightforward implementation writes it out to HBM and reads it back, and that memory traffic, not the arithmetic, dominates the run time.

FlashAttention never builds the full matrix. It splits queries and keys into tiles that fit in SRAM, and for each block of queries streams the key/value tiles through, keeping only a running result. The difficulty is softmax: it needs the maximum and the sum over the whole row before anything can be normalized. The fix is an online softmax that keeps a running maximum $m$ and running sum $\ell$, and rescales what it has so far whenever a larger maximum appears:

$$m' = \max(m, \max_j s_j), \qquad \ell' = \ell\, e^{m - m'} + \sum_j e^{s_j - m'}$$

Figure 10.7 · FlashAttention tiling and the online softmax

HBM big, slow GPU memoryQ, K, V, output

↓ K,V block 1

SRAM tiny, ~10× fasterQ block 1 · tile (1,1)

query 2 scores:−0.5−1.5

−0.45running max m1.37running sum l0.583output so far0.583exact softmax

1/10

Query block 1 done after 1 tile. Tiles to the right are skipped by the causal mask. The full score row was never written to slow memory, yet the result equals exact softmax.

The result is exact attention, not an approximation, but with memory that grows linearly in the context length and several times less memory traffic. It matters most in prefill and training, where many queries meet many keys; in decode there is one query per step, and a variant (Flash-Decoding) instead splits its long row of keys across the cores. Every major inference and training stack uses FlashAttention or a descendant.

#When one GPU is not enough

A 70B model in 16-bit needs 140 GB for weights alone, more than one GPU's 80 GB. Serving engines split the model across several GPUs:

• Tensor parallelism cuts every weight matrix into slices, one per GPU. Each GPU computes its slice of every layer, and they exchange partial results after each layer over a fast link such as NVLink. Every GPU reads only its share of the weights, so tensor parallelism also multiplies the memory bandwidth behind each token.
• Pipeline parallelism puts different layers on different GPUs and passes activations along, like an assembly line.
• Expert parallelism places different experts of a mixture-of-experts model (Chapter 13) on different GPUs.

Large deployments combine all three, and often run prefill and decode on separate pools of GPUs (disaggregated serving), because one phase is compute-bound and the other memory-bound, so they want different hardware trade-offs.

#Summary

• A GPU can do about 300 operations in the time one byte arrives from its main memory. Work that does fewer operations per byte is memory-bound.
• Prefill processes the prompt in parallel and is compute-bound; decode emits one token per pass and is limited by memory bandwidth.
• The KV cache keeps every token's keys and values, the only things later tokens read, so each decode step processes just the new token. Its memory grows with context length.
• GQA and MQA shrink that cache by sharing key/value heads between query heads.
• Batching reuses every byte of weights for many users; continuous batching keeps the batch full step by step.
• Speculative decoding lets a cheap draft model propose tokens that the big model verifies in one pass.
• FlashAttention computes exact attention tile by tile in fast on-chip memory, using an online softmax.
• Models too big for one GPU are split with tensor, pipeline and expert parallelism.

The KV cache turned out to be the central data structure of serving. The next chapter is about making the most of it: caching.