viewof layers = Inputs.range([8, 128], {value: 32, step: 1, label: "Layers"})
viewof kvheads = Inputs.range([1, 64], {value: 8, step: 1, label: "KV heads (8 = GQA, 32 = full MHA)"})
viewof headdim = Inputs.range([32, 256], {value: 128, step: 32, label: "Head dimension"})
viewof ctx = Inputs.range([1000, 128000], {value: 8000, step: 1000, label: "Context length (tokens)"})
viewof bits = Inputs.radio([16, 8], {value: 16, label: "Cache precision (bits)"})The model stops redoing work it has already done. Nothing is approximated — it just keeps notes in the margin instead of re-reading the whole chart every time.
Doing it all again, every time
Generation is one word at a time, and each word attends to everything before it.
To write word 500, the model considers all 499 words before it. To write word 501, it does all 500 again. Then 501 again for the next one. The same arithmetic, over and over.
A doctor writing a long note, who re-reads the patient’s entire chart before adding each new line. Twenty lines in, she’s read the chart twenty times. Or — she could jot the key facts in the margin as she goes, and only read the new line.
What the trick is worth
I measured both approaches on the same model, generating the same text:
| Approach | Time taken |
|---|---|
| Re-reading everything each time (no cache) | 2.97 seconds |
| Keeping notes (KV cache) | 0.18 seconds |
Sixteen times faster — and the output was word-for-word identical. Nothing was approximated or skipped. It simply stopped doing the same work twice.
That’s the rare kind of optimisation: free. The bill arrives as memory.
Why K and V, but never Q
From the attention chapter, every token produces a query, a key and a value. When the model generates a new token:
- The new token’s query is compared against the keys of every token so far.
- Those keys and values have not changed since they were computed, and never will.
- The new query, once used, is thrown away. It’s never needed again.
So you cache K and V, not Q. That’s the entire idea, and it’s why it’s called the KV cache.
Without it, generating \(n\) tokens costs \(O(n^3)\) work overall. With it, \(O(n^2)\) — each step computes one new K and V and attends over the stored ones.
It grows with every token
\[ \text{cache bytes} \approx 2 \times L \times H_{kv} \times d_{head} \times n \times \text{bytes} \]
where the 2 is K and V, \(L\) is layers, \(H_{kv}\) is key/value heads, \(d_{head}\) is head dimension, and \(n\) is tokens so far. It grows linearly with every token generated, for every concurrent user.
Play with it — these defaults are roughly an 8B-class model in fp16:
Push the context to 128k and watch how few users fit. This is why long context is expensive to serve, not just to compute.
Two phases, two bottlenecks
| Phase | What happens | Limited by |
|---|---|---|
| Prefill | The whole prompt is processed at once, filling the cache | Compute — it’s one big matrix multiply |
| Decode | One token at a time, each step dragging the entire cache through memory | Memory bandwidth |
This is the single most useful mental model for inference performance. Time-to-first-token is a prefill problem; tokens-per-second after that is a bandwidth problem. Adding compute rarely fixes slow decoding — during decode a GPU spends most of its time waiting on memory, not calculating.
How the cache is made smaller
- GQA / MQA. Several query heads share one set of keys and values. Multi-query attention is the extreme (one KV head). Grouped-query attention is the common middle ground — it’s why the slider above defaults to 8 KV heads rather than 32, and it shrinks the cache by that same 4×.
- Sliding-window attention. Only keep the last \(w\) tokens’ keys and values, capping the cache no matter how long the conversation runs. You trade away exact long-range recall.
- Quantised cache. Store K and V in 8 bits instead of 16. Halves the memory for a small quality cost — try the toggle above.
- PagedAttention (vLLM). Don’t shrink the cache; stop wasting it. Allocate in small pages like virtual memory instead of one contiguous block per request, and fragmentation mostly disappears.
- Prefix sharing. Many requests start with the same system prompt. Compute that prefix’s cache once and share it across requests.
The same trick, one level up
Prompt caching in a hosted API is this idea at the product layer: a long, unchanging system prompt gets its cache computed once and reused, so you’re billed a fraction for those tokens on later calls. Same principle — stop redoing old work — different altitude.
Where it breaks
- Memory, not compute, is usually your serving limit. Batch size is capped by cache size. People buy faster GPUs when they needed more memory or a smaller cache.
- The cache is per-conversation. A long chat history is re-billed and re-carried on every turn. Trim aggressively.
- Change anything early in the prompt and the cache dies. Cached prefixes are only valid while the prefix is byte-identical. A timestamp at the top of a system prompt quietly destroys every prefix hit.
- Sliding windows forget. If a fact left the window, no amount of attention brings it back. It’s not “compressed” — it’s gone.
Whiteboard check
Marker in hand, someone watching. Could you get through these without notes?
NoteWhat exactly is stored in a KV cache, and why not the queries?
The key and value vectors for every past token, at every layer and every KV head. Past keys and values are reused at every future step and never change. A query is used once, at the step that created it, and is then irrelevant.
NoteEstimate the KV cache for a 32-layer model, 8 KV heads, head dim 128, fp16, at 8k context.
\(2 \times 32 \times 8 \times 128 \times 8000 \times 2\) bytes ≈ 1.05 GB per request. On an 80GB GPU with ~16GB of weights, that’s around 60 concurrent users before memory runs out — which is why batch size, not FLOPs, usually sets your throughput.
NoteYour time-to-first-token is fine but tokens-per-second is poor. What do you investigate?
That’s decode, so suspect memory bandwidth, not compute: cache size per request, batch size, GQA vs MHA, cache precision, and whether paged attention is in use. Bigger GPUs with the same bandwidth won’t help; larger batches (amortising weight reads) usually will.
NoteDoes a KV cache change the model’s output?
No. It’s exact — identical tokens, identical probabilities, just without recomputation. That’s what separates it from sliding-window or quantised variants, which do trade quality for memory.
TL;DR
- Cache the keys and values of past tokens; queries are single-use.
- 16× faster, identical output. Pure win on compute, paid for in memory.
- Cache size grows linearly with tokens and users: \(2 \times L \times H_{kv} \times d_{head} \times n\).
- Prefill is compute-bound; decode is bandwidth-bound. Diagnose them separately.
- GQA, sliding windows, quantisation and paged attention all exist to fight the same bill.