KV Cache

A key-value cache stores attention keys and values from earlier tokens in an autoregressive decoder. When the next token arrives, the model can reuse that state instead of projecting the same prefix again.

The saving is in avoiding repeated work on past tokens. The current query still reads the relevant cached keys and values, so attention work and cache memory continue to grow with the context.

In a causal decoder, token $t$ attends to tokens $1..t$. When you later generate token $t+1$, tokens $1..t$ are unchanged, so their keys and values are identical to what you already computed. Recomputing them is pure waste. The KV cache stores them; generation appends the new token's K and V and reuses everything else.

The cache trades repeated computation for stored state. It grows with context length and must be read during attention. Decoding can become memory-bandwidth-bound, depending on batch size, hardware, and the attention implementation.

Without a cache, a naive decoder reprojects the full prefix. For dense $d \times d$ key/value projections, step $t$ takes $O(t d^2)$ projection work. Summing over $n$ generated positions gives $O(n^2 d^2)$ for those projections alone.

With a cache, each step projects only the new token ($O(d^2)$ under the same assumption). Its query still attends to the cached keys and values, requiring $O(t d)$ attention work. Caching removes repeated prefix projection; it does not make attention or memory use constant in context length.

Memory cost of the cache for one sequence:

The factor of 2 is for K and V. This grows linearly with sequence length and batch size and quickly dominates memory for long contexts. Two standard mitigations shrink it by sharing keys/values across query heads: Multi-Query Attention (MQA) uses a single K/V head for all query heads; Grouped-Query Attention (GQA) uses a few: cutting cache size (and bandwidth) by the head-sharing factor with minimal quality loss.

Key takeaway: Caching is safe because causal attention makes past keys/values immutable. Decoding can become memory-bandwidth-bound: you read the whole cache every step, which is why cache size, not FLOPs, limits long-context serving.

1. Conceptual: What does the KV cache store and why is caching it correct? (Past tokens' keys and values: they don't change as new tokens are generated in a causal decoder, so they can be reused.)
2. Implementation: How does the cache change the per-step cost of decoding? (From recomputing K/V over the whole prefix each step to computing only the new token's projections plus attention over the cache.)
3. Applied: Write the memory cost of the cache and name what makes it grow. (2 · layers · heads · d_head · seq_len · bytes: grows linearly with sequence length and batch size.)
4. Systems-level: Why can autoregressive decoding become memory-bandwidth-bound? (Each step does little arithmetic but must read the entire growing cache from memory: bandwidth, not FLOPs, is the bottleneck.)
5. Failure modes: How do MQA and GQA reduce the KV cache? (They share K/V heads across query heads (MQA uses one K/V head, GQA a few) shrinking cache size and bandwidth with little quality loss.)

Without looking: explain what the KV cache stores and why reuse is valid, give the memory-cost formula, and say why decoding is memory-bound. Then name two techniques that shrink the cache. Check against Stages 1–3.