Chapter 03 · The cost of stored context
Generation happens one token at a time. A naive implementation would recompute projections (K,V) for the entire past for every new token. The obvious solution is to cache keys and values that were previously computed. But there is a cost to it and that is memory. As the context expands so do the memory demands. In this chapter we will look at how to calculate that cost and how to start optimizing the memory requirements.
At decode step t, only one new token needs a query. But that query must be scored against the keys of all previous tokens, and their keys and values haven't changed since we computed them. So we keep them:
k_all = concat(k_cached, k_new) # keys: T-1 old + 1 new
v_all = concat(v_cached, v_new) # values: T-1 old + 1 new
output = softmax(q_new @ k_all.T / sqrt(d)) @ v_all
cache = (k_all, v_all) # grows by one slot, every step, foreverAdapted from the source article by @waterloo_intern.
Caching saves computation by using memory instead. This tradeoffcontext Computing systems often store a result to avoid calculating it again. Here the stored results can fill an accelerator's fast memory, so cache size becomes an architectural constraint. also means each new query must read the stored keys and values.
Inference has two phasescontext Prefill processes the prompt before the first output token appears. Decode generates the following tokens one by one. Systems measure time to first token and time per output token separately because the phases stress different hardware resources.:
FlashAttention (Dao et al., 2022) reorganizes attentioncontextThe method uses tiling and an online softmax. It processes score blocks in fast on-chip memory while carrying a running maximum and sum. The result remains exact attention up to floating-point rounding, with fewer trips to main memory. to avoid storing the full score matrix. It reduces memory traffic, but exact attention still compares each query with every key and the KV cache still grows with the context.
Multiply the dimensions of the cache to get its size:
For GPT-2 small, the calculation is 2 × 12 layers × 12 KV heads × 64 dimensions × 2 bytes = 36,864 bytes, or about 36 KiB per token. Try the calculator with the published model configurations. Cache size rises in direct proportion to context length.
Model designs reduce KV-cache cost in three broad ways:
| Route | Idea | Where covered |
|---|---|---|
| Share the slots | Fewer KV heads serve many query heads (GQA), or compress each token's K/V into a small latent (MLA) | Ch. 07 |
| Abandon the slots | Replace per-token storage with a fixed-size recurrent state, linear attention and its descendants | Ch. 04–06 |
| Mix both | Cheap recurrent memory most layers, exact (compressed) attention periodically | Ch. 07, 10 |
An H100 SXM moves ~3.35 TB/s from its memory while sustaining on the order of a petaFLOP of half-precision math, roughly hundreds of arithmetic operations per byte moved. To keep such a chip busy you must do lots of math per bytecontextThe ratio has a name, arithmetic intensity, and a famous chart, the roofline model: below the machine's ridge point you are bandwidth-bound and extra FLOPs are literally free; above it, compute-bound. Big matrix multiplies sit far above the ridge; cached decode attention scrapes along the floor at roughly one operation per byte.. Decode does the opposite: almost no math per cached byte it rereads, so the arithmetic units can wait for memory. For the designs in the following chapters, bytes read per token can matter more than FLOPs.
Many trained attention heads behave similarly. Instead of giving each query head its own keys and values, several query heads can share one K/V pair. Reducing the number of KV heads reduces the cache by the same factor. Models trained or adapted for this arrangement can retain much of their quality:
| Scheme | K/V heads (per 32 query heads) | Relative cache | Trade |
|---|---|---|---|
| MHA: one K/V per query head | 32 | 1× (full bill) | the original; most expensive |
| GQA: groups share | 8 | ¼ | usually small quality loss; common in recent open models |
| MQA: all share one | 1 | 1⁄32 | cheapest; a bit more quality loss |
These ratios cut KV-cache storage only: the model still attends over every prior token position; it just keeps fewer distinct K/V projections to do it.
Multi-Query Attention (all query heads share a single K/V; Shazeer, 2019) was the aggressive version; Grouped-Query Attention (Ainslie et al., 2023) found the sweet spot by letting groups of query heads share, and is now common across recent open models. Llama-3-8B uses 8 KV heads for 32 query heads, giving it one quarter as many KV heads as standard multi-head attention. Chapter 07 covers MLA, which compresses each token's stored representation.
MQA and GQA reduce the "KV heads" factor in the cache formula. Other designs reduce different factors: chapter 10 examines bytes per number, while chapters 04–07 reconsider whether every token needs its own stored slot.
A server handles many requests at once, and each request's cache grows as it generates. Reserving one large, contiguous region for each request wastes memory when the answer is shorter than the reservation. The vLLM authors measured 60–80% KV-memory waste from fragmentation and over-reservation in earlier serving systems.
vLLM applies paging, a technique used by operating systems. PagedAttention (Kwon et al., the vLLM system, SOSP 2023) divides each cache into fixed-size blocks that can live anywhere in memory. An index maps logical token positions to physical blocks. Requests receive blocks as needed, leaving only the final block partly unused; the authors report waste falling to a few percent. Requests with an identical prefix can share its blocks, then use copy-on-write when their continuations diverge. Cache allocation is therefore part of the memory cost of serving a model.
What this chapter established
Go deeper
Contents · Glossary · Cache formula and phase framing adapted from the source article; hardware figures are vendor spec-sheet values.