KV caching: why generation is one token at a time

The mental model: a reader who can't skim

A transformer decoder writes the way you would fill in a sentence you're not allowed to plan: pick one word, commit to it, look at everything so far, pick the next. That is what autoregressive means. The model outputs a probability distribution over the next token, conditioned on every token before it. You sample one, append it, and run the whole thing again. There is no batch of "the next ten tokens." Token N+1 depends on the token you just sampled at position N. (This is the exact loop Lesson 6 was choosing tokens inside — decoding is just how you pick each sample; here we look at what that loop costs to run.)

Done naively this is brutal. To produce token 1000 you re-run attention over all 999 tokens before it; token 1001 re-runs over 1000. Generating a sequence of length n costs on the order of n² work if you recompute from scratch each step. That n² is exactly the attention-compute cost Lesson 4 pinned on the all-pairs lookup, and it is real. But watch where the bottleneck goes once you stop recomputing: the KV cache is the trick that turns most of that quadratic wall back into something you can serve — and in doing so it moves the true ceiling off compute and onto memory. The cache grows only linearly with tokens, yet it is what saturates VRAM and bandwidth first, long before attention compute becomes the limit again. That relocation of the bottleneck — quadratic compute avoided, linear memory now dominant — is the whole story of this lesson.

What attention recomputes, and what it doesn't

Inside each attention layer, every token is projected into three vectors: a query (Q), a key (K), and a value (V). A token attends by taking its own query and dotting it against the keys of all earlier tokens, then using those scores to average their values. Here is the asymmetry everything hinges on: when you generate a new token, only the new token's query is new. The keys and values of every earlier token are identical to what they were last step. They do not depend on the new token at all.

So don't recompute them. Cache them.

Decode step for a new token (with KV cache):

  new token ──> [Q_new] [K_new] [V_new]
                   │        │       │
                   │        └───────┴──> append to cache
                   ▼
              ┌──────────────────────────────┐
   attention: Q_new · [K_0 K_1 ... K_new]     │  <- keys/values for 0..N-1
              softmax · [V_0 V_1 ... V_new]    │     read straight from cache
              └──────────────────────────────┘
                   │
                   ▼
              next-token logits

Each decode step you compute exactly one new K and one new V, append them to the cache, and run a single query against the whole stored history. Per-step work drops from "attend over N tokens and rebuild all of them" to "compute one token, read N from memory." Be careful how you state the win, though — it's three regimes, not one. Prefill is still quadratic: it runs full dense attention over all N prompt tokens in one pass, and no cache removes that. A single cached decode step is what goes linear — one new query dotted against N stored positions, roughly O(N) work instead of rebuilding the whole history. Generating a full response then sums those growing steps, so total decode work still climbs across the sequence; what the cache eliminates is the recomputation of past keys and values, not the fact that each step reads a longer history. The price for even that is holding the entire history in fast memory.

The memory math (this is the part that bites)

The cache size per token is simple and worth memorizing:

bytes_per_token = 2  ·  n_layers  ·  n_kv_heads  ·  head_dim  ·  bytes_per_elem
                  ▲
                  K and V

Put real numbers in. Llama 3 70B has 80 layers, head dimension 128, and thanks to grouped-query attention only 8 key/value heads instead of 64. In fp16:

2 × 80 × 8 × 128 × 2 = 327,680 bytes ≈ 320 KB per token.

Fill its 8192-token context and one sequence's cache is 8192 × 320 KB ≈ 2.6 GB. Batch ten users and that is 26 GB of VRAM spent on cache alone, before a single weight. The vLLM paper reports the same shape for OPT-13B: about 800 KB per token, 1.6 GB for one 2048-token sequence.

Now look at what grouped-query attention bought. Full multi-head attention would use all 64 KV heads, making the per-token cost 8× larger: roughly 2.5 MB per token, over 20 GB for one filled 8192-token sequence. GQA (Ainslie et al., 2023), which generalizes Shazeer's multi-query attention (2019), shrinks the cache by sharing K and V across groups of query heads. Modern open-weight models ship 8 KV heads for exactly this reason. The cache, not the compute, is what stops scaling first.

Researcher: most "long context" and "KV compression" work is attacking this one equation. Quantizing the cache to int8 or int4, evicting low-attention tokens, sharing K/V across layers: each trades a term in bytes_per_token for some quality risk.

Prefill vs decode: two different machines

Generation runs in two phases with opposite performance profiles, and conflating them is the classic serving mistake.

Prefill processes the whole prompt at once. Every prompt token passes through the model in parallel and produces the initial KV cache. You are multiplying big matrices against weights you have already loaded, so arithmetic intensity is high (hundreds of FLOPs per byte moved) and the GPU runs near saturation. Prefill is compute-bound.

Decode emits one token per forward pass. Each step you stream the entire set of model weights out of HBM to produce a single token's worth of arithmetic. For a 70B model in fp16 that is about 140 GB read per token. On a GPU with roughly 3 TB/s of memory bandwidth, that sets a floor near 45 ms per token at batch size 1, in the ballpark of 20 tokens per second, with the compute units mostly idle waiting on memory. Decode is memory-bandwidth-bound.

PREFILL                          DECODE
prompt: [t0 t1 t2 ... t500]      one token per pass
all tokens at once               ├─ read ALL weights from HBM
compute-bound, ~95% util         ├─ read growing KV cache
builds the initial KV cache      └─ memory-bound, ~20-40% util

That split explains most serving behavior. Time-to-first-token is dominated by prefill and grows with prompt length. Inter-token latency is dominated by decode bandwidth. Batching is the lever. Decode is already stalling on the weight read, so routing many sequences through the same weight load spreads that read across all of them: throughput climbs steeply with batch size while per-token latency barely moves, right up until the KV caches stop fitting in VRAM.

What goes wrong in practice

The failure mode is almost never "out of compute." It is "out of KV cache." Under naive allocation each sequence reserves cache for its maximum possible length up front, and most of that sits empty. The vLLM authors measured 60 to 80 percent of KV memory lost to internal fragmentation and over-reservation. PagedAttention (Kwon et al., 2023) borrowed a trick from operating systems: chop the cache into fixed pages of 16 tokens, allocate them on demand, and keep a per-sequence page table. Fragmentation drops to near zero, sequences that share a prompt prefix can share physical pages, and effective batch size (and therefore throughput) roughly doubles.

Builder: when a serving stack reports it can hold N concurrent requests, that number is KV_memory_budget / (avg_context_len × bytes_per_token). Cut context, quantize the cache, or pick a GQA model and N rises in proportion. Set max_model_len too high and you pre-reserve cache you will never touch.

Defender: prefix sharing is a side channel. When prefix pages are cached and reused across requests, a difference in response timing can reveal whether another user recently sent the same prompt prefix, a prompt-cache timing oracle. Treat cross-tenant prefix sharing as a data-isolation decision, not a free speedup.

One token at a time is also why streaming exists, why speculative decoding (draft cheap tokens, verify them in one batched pass) earns its complexity, and why the KV cache is the object every serious inference optimization is really fighting over. This memory-bound framing is the one Lesson 8 picks up when it separates the several distinct reasons long context degrades — the VRAM ceiling here is only one of them. We come back to the cache itself in the module on continuous batching and paged attention, where it stops being a data structure and becomes a scheduler.

Sources

• Kwon et al. (2023), Efficient Memory Management for Large Language Model Serving with PagedAttention (the vLLM paper), arXiv:2309.06180 — OPT-13B ~800 KB/token figure, 60–80% fragmentation, and paging results.
• Ainslie et al. (2023), GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, arXiv:2305.13245.
• Shazeer (2019), Fast Transformer Decoding: One Write-Head is All You Need (multi-query attention), arXiv:1911.02150.
• Vaswani et al. (2017), Attention Is All You Need, arXiv:1706.03762 — the Q/K/V attention mechanism.
• Williams, Waterman & Patterson (2009), Roofline: An Insightful Visual Performance Model for Multicore Architectures, Communications of the ACM — the compute-bound vs memory-bound (arithmetic-intensity) framing behind prefill/decode. — https://www2.eecs.berkeley.edu/Pubs/TechRpts/2008/EECS-2008-134.html