Appearance
12.7.6 — Inference, the KV Cache and Batching
Cartwise decided to run an open 7-billion-parameter model on its own GPU for the support assistant, to keep customer conversations in-house. The first load test produced three puzzling numbers:
- With one customer at a time, answers streamed at about 100 tokens a second, yet the GPU's arithmetic units showed as mostly idle.
- With 32 customers at once, each customer still got nearly the same speed, so total output was about thirty times higher, on the same GPU.
- At 64 customers with long conversations, the server ran out of memory, although the model's weights took only 14 gigabytes of the GPU's 80.
All three come from how a trained model is actually run to produce text, which is called inference. It has two phases that behave like two different machines, and one data structure, the KV cache, that is both what makes generation fast and what fills the memory. Understanding them explains what you pay for when you call an LLM API, what you wait for, and how to serve a model yourself without wasting most of the hardware.
1. Two phases: prefill and decode
12.7.1 described the loop: process the prompt, choose a token, append it, repeat. Those steps fall into two phases with very different characters.
Prefill runs the whole prompt through the model in one go. All the prompt's tokens are known, so, exactly as in training, every position is computed in parallel: one enormous set of matrix multiplications. The GPU's thousands of arithmetic units are kept busy, so prefill is limited by how fast the chip can calculate. A forward pass costs about 2 operations per parameter per token (the forward-only part of the "6 per token" rule in 12.7.5), so a 2,000-token prompt through a 7-billion-parameter model is about 2 \times 7 \times 10^9 \times 2{,}000 = 2.8 \times 10^{13} operations. At an achieved 1.25 \times 10^{14} operations per second, that is about 0.22 seconds.
Decode produces the answer, one token per step, because each new token depends on the one before. Each step processes a single new token. Here is the surprise: to process even one token, the GPU must read every weight of the model from its memory, all 14 gigabytes of them in 16-bit format, because every layer's matrices are needed. The arithmetic for one token is tiny; reading the weights is not. A data-centre GPU can read memory at about 2 terabytes a second, so one decode step takes at least 14 / 2{,}000 \approx 0.007 seconds, about 7 milliseconds, a ceiling of roughly 150 tokens a second for a single conversation. During those 7 milliseconds the arithmetic units are mostly waiting for data. Decode is limited by memory bandwidth, not by calculation.
That one fact explains the first load-test puzzle, the idle GPU at 100 tokens a second, and it shapes the rest of this page.
Two numbers you will measure everywhere
- Time to first token (TTFT): how long until the first word appears. It is mostly prefill, so it grows with prompt length.
- Time per output token (TPOT): how long each further token takes. It is a decode step, roughly constant per token.
Total time for an answer is roughly TTFT plus the number of output tokens times TPOT. A 300-token answer at 15 milliseconds per token is 4.5 seconds of decode, usually far more than the prefill. So long answers, not long prompts, usually dominate how long a request takes, and asking the model to be concise is one of the cheapest latency improvements there is. Streaming, from 12.8.5, hides most of the decode time from the user because they start reading after TTFT.
This also explains API pricing. Output tokens are typically priced several times higher than input tokens, because input tokens are processed cheaply in parallel during prefill, while every output token costs a full trip through the model in the expensive decode phase.
2. The KV cache: compute each token's keys and values once
Look again at attention from 12.7.3. To produce the 101st output token, the new token's query must be compared with the keys of every earlier token, and their values mixed. Without any saving, every decode step would recompute the keys and values for the entire prompt and all output so far, in every layer: step 100 would redo the work of steps 1 to 99. Generating n tokens would cost work proportional to n^2.
But because of the causal mask, a token's key and value never change once computed: later tokens cannot affect earlier ones. So they can be computed once and stored. That store is the KV cache: for every layer and every head, the key and value vector of every token processed so far. Each decode step then computes the key, value and query for just the one new token, appends the key and value to the cache, and attends over the cached ones.
How big the cache gets
The cache holds two vectors (a key and a value) per token, per layer, per key-value head. Its size is
\text{KV bytes per token} = 2 \times \text{layers} \times \text{KV heads} \times \text{head size} \times \text{bytes per number}
For Llama 2 7B, with 32 layers, 32 heads of 128 numbers each, in 16-bit (2-byte) format: 2 \times 32 \times 32 \times 128 \times 2 = 524{,}288 bytes, half a megabyte per token. A 4,096-token conversation needs about 2.1 gigabytes of cache. Thirty-two such conversations at once need about 69 gigabytes: five times the model's own weights. That is the third load-test puzzle: the model fit easily; the caches of 64 long conversations did not.
This is exactly why the grouped-query attention of 12.7.4 became standard. With 8 key-value heads instead of 32, the same model's cache drops to 128 kilobytes per token, a quarter of the size, and four times as many conversations fit. Caches can also be quantised to 8 bits or fewer, halving them again (12.7.8). In LLM serving, memory for the KV cache, not the weights, usually decides how many users a GPU can serve at once.
3. Batching: share each read of the weights
Now the second puzzle: why 32 customers at once cost each of them almost nothing in speed. A decode step is slow because it reads all 14 gigabytes of weights, and the arithmetic for one token barely uses the chip. But once the weights are being read, they can be applied to the next token of many conversations at the same time: instead of multiplying the weight matrices by one token's vector, multiply them by a batch of 32 vectors, one per conversation. The weights are read once; the arithmetic, which was idle anyway, does 32 times the work. Each conversation's step takes nearly the same 7 milliseconds, and the GPU produces 32 tokens in that time instead of one.
This is the single most important fact about the economics of LLMs. The ratio of arithmetic to memory reading, called arithmetic intensity, is very low for one decode at a time and rises with the batch size, until eventually arithmetic, not memory, becomes the limit. Serving many requests together is dramatically cheaper per token than serving them one at a time, which is why API providers, who batch thousands of requests, can charge less than it would cost most companies to run the same model for a handful of users.
Continuous batching
Simple batching has a flaw. Conversations finish at different times: one answer is 20 tokens, another is 800. If the server waits for every request in a batch to finish before starting the next batch, the GPU spends long stretches working for a single slow request while finished slots sit empty, and new requests queue up.
Continuous batching, introduced in a 2022 system called Orca, fixes this by changing the batch at every decode step: as soon as one conversation finishes, its slot is given to a waiting request, which runs its prefill and joins the decoding. The batch stays full all the time. This raises throughput several-fold over simple batching and is standard in every modern serving system.
Managing cache memory: PagedAttention
Continuous batching makes memory management hard. Each conversation's KV cache grows by one entry per token, and nobody knows in advance how long a conversation will get. Reserving the maximum possible space for every conversation wastes most of the memory; allocating as you go leaves scattered gaps. In 2023, Woosuk Kwon and colleagues at Berkeley published PagedAttention, which borrows the idea of virtual memory pages from operating systems, covered in Chapter 2.5: store each conversation's cache in small fixed-size blocks that need not sit next to each other in memory, with a table recording where each block is. Almost no memory is wasted, and conversations that share a beginning can even share blocks. Their open-source serving engine, vLLM, built around it, became one of the most widely used ways to serve open models.
4. Prefix caching: never compute the same prompt twice
Look at what Cartwise sends with every support message: the same 1,500-token block of system instructions and policies, then the conversation, then the new message. Every request recomputes the KV cache for that identical opening block during prefill. Prefix caching keeps the computed KV cache for common prompt beginnings and reuses it whenever a new request starts with exactly the same tokens, skipping their prefill entirely.
Hosted APIs expose this as prompt caching, and charge cached input tokens at a large discount, often around a tenth of the normal input price, with cached requests also reaching their first token faster. The cache only matches identical leading tokens, so it rewards one simple habit: put the parts of the prompt that never change first, and the parts that change last. System instructions and reference documents first; the conversation history next; the new question at the very end. Putting today's date or the customer's name at the top of the system prompt breaks the match for every request, and costs real money. Chapter 12.12.2 covers caching across the whole system.
5. Generation in code
Libraries manage all of this for you, but it helps to see the cache in an explicit loop. Using Hugging Face transformers and GPT-2 again:
python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
tok = AutoTokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2").eval()
ids = tok("Your order has been", return_tensors="pt").input_ids
with torch.no_grad():
out = model(ids, use_cache=True) # (1) prefill: whole prompt at once
cache = out.past_key_values
next_id = out.logits[0, -1].argmax().view(1, 1) # (2) greedy choice of token 1
generated = [next_id.item()]
for _ in range(10):
out = model(next_id, past_key_values=cache, use_cache=True) # (3) decode: ONE token in
cache = out.past_key_values # (4) cache grew by one position
next_id = out.logits[0, -1].argmax().view(1, 1)
generated.append(next_id.item())
print(tok.decode(generated)) # e.g. " shipped to you…" (model-dependent)Line (1) is the prefill: the whole prompt goes in, and the model returns both logits and the KV cache for every layer. Line (2) picks the most likely next token. Line (3) is a decode step: notice that only one token, next_id, is passed in, together with the cache, so the model computes keys and values for that one token and attends over the stored rest. Line (4) keeps the updated cache, now one position longer. In practice you would call model.generate(...), which runs this same loop, or use a serving engine like vLLM, which also handles continuous batching and paged memory across many users.
What the interviewer will push on
"Why is generating text slow, and what limits it?" They want the prefill and decode split, and the key insight that decode is memory-bandwidth-bound because each step reads all the weights to produce one token. The tell of a strong answer is the arithmetic: weight bytes divided by memory bandwidth gives the time per token.
"What is the KV cache, and what does it cost?" Stored keys and values for every earlier token, so each step does one token's work instead of recomputing the prefix; its size is 2 × layers × KV heads × head size × bytes per token, and at long contexts and large batches it exceeds the weights. Mentioning grouped-query attention as the fix is expected.
"How would you increase throughput on one GPU?" Batching, because decode leaves the arithmetic idle; continuous batching to keep the batch full; paged memory to fit more caches; prefix caching for shared prompts; quantisation to read fewer bytes. A strong answer ranks them by impact.
"Why are output tokens more expensive than input tokens?" Input is processed in parallel in prefill; each output token is a full sequential decode step.
Volunteer this: order the prompt for the cache. Stable instructions and documents first, variable content last; putting something that changes per request at the top silently disables prompt caching.
Recall
- Inference has two phases. Prefill processes the whole prompt in parallel and is limited by arithmetic. Decode makes one token per step, reads every weight each time, and is limited by memory bandwidth: 14 GB ÷ 2 TB/s ≈ 7 ms per token for a 7-billion-parameter model.
- Time to first token grows with prompt length; time per output token is roughly constant; long answers dominate latency, and output tokens cost more.
- The KV cache stores every earlier token's keys and values so each step does one token's work. It costs 2 × layers × KV heads × head size × bytes per token: 0.5 MB per token for Llama 2 7B, about 69 GB for 32 conversations of 4,096 tokens. Grouped-query attention cuts it.
- Batching shares each read of the weights across many conversations, making serving far cheaper per token. Continuous batching refills finished slots every step; PagedAttention (vLLM) stores caches in pages to avoid waste.
- Prefix caching reuses the KV cache of identical prompt beginnings; put stable content first and variable content last.
Self-test: Estimate time per token for a model with 140 GB of 16-bit weights on GPUs with 4 TB/s of combined bandwidth. How much KV cache does a model with 80 layers, 8 KV heads of 128 numbers, in 16-bit, need per token? Why does moving the customer's name from the top of the system prompt to the end reduce cost?
Next: 12.7.7 Decoding and Sampling — each decode step ends with a probability for every token; how one token is chosen from those probabilities decides whether the output is precise, creative, repetitive or nonsense.