LESSON 6.2 · Derivation · 100 min
KV cache: trade memory for repeated compute
When generating token t, the Keys and Values for earlier tokens do not change. A KV cache stores them at every layer, so the model computes new Q, K, and V only for the new token a
DIRECT ANSWER · VERIFIED SOURCES ·
Why does a KV cache speed up autoregressive decoding, and why can it consume so much memory?
When generating token t, the Keys and Values for earlier tokens do not change. A KV cache stores them at every layer, so the model computes new Q, K, and V only for the new token and lets its Query attend to cached history instead of recomputing the full prefix. For a conventional cache, memory is roughly proportional to 2 × layers × batch × sequence length × KV heads × head_dim × bytes per element, where 2 accounts for K and V.
Lecture alignment
Lecture alignment: diagram prefill separately from token-by-token decode, label the new Q, appended K/V, and historical K/V at each layer, then plug the actual model configuration into the cache formula.
Key takeaways
- KV caching mainly removes repeated historical K/V computation during decode; it does not eliminate full prompt prefill.
- Cache usage grows with active sequence length and concurrent requests, so long-context serving can become KV-memory-bound before weight-memory-bound.
- GQA and MQA reduce cache size by using fewer KV heads while retaining more Query heads.
Boundary and caveat
The formula assumes full, same-precision K/V at every layer. Sliding windows, quantization, offloading, cross-layer sharing, and sparse attention change it. A changed prefix, position scheme, or attention mask may also invalidate reuse.
Primary sources
Learning objectives
- Separate prefill, which writes prompt K/V in parallel, from decode, which appends one position per step.
- Derive the KV-cache shape and byte formula, including the effect of GQA and MQA.
- Prove that cached attention matches recomputing every prefix.
- Distinguish the KV cache itself from paged allocation and prefix caching.
Core concepts
Prefill and decode
Prefill processes the prompt in parallel and writes each layer’s K/V. Decode computes and appends one new K/V pair, then lets the current query read the prefix cache.
Cache shape
A useful logical shape is [layers, 2, batch, tokens, kv_heads, head_dim]. Implementations may transpose or page it, but these dimensions still determine the element count.
Capacity formula
Bytes = layers × 2(K,V) × batch × tokens × kv_heads × head_dim × bytes per element. Doubling context doubles capacity when all other dimensions stay fixed.
MHA, GQA, and MQA
MHA commonly uses one KV head per query head; GQA shares fewer KV heads across query groups; MQA keeps one K/V head. The reduction directly lowers cache size and reads.
Output equivalence
A correct cache reuses previously computed K/V without changing the attention equation. Fixed inputs and weights must match a reference that recomputes every prefix.
Build and verify
Implement cached attention and prove equivalence
- Use three 2-D tokens and fixed WQ/WK/WV; hand-calculate the first K/V and output.
- Implement decode_step so each new token computes K/V once and appends them before attention.
- Implement a reference that recomputes every prefix, then assert elementwise equality at every step.
- Build a capacity table for layers, tokens, kv_heads, head_dim, and dtype; compare doubled context and GQA.