Attention and the KV Cache
Mental model
For hidden states (X), learned projections produce (Q), (K), and (V). Each query compares with eligible keys; scaling controls logit magnitude, a causal mask removes future positions, and softmax turns the remaining scores into weights over values.
During generation, the prefix is unchanged when one new token arrives. Its keys and values therefore remain valid. Storing them lets the next step compute only the new token’s projections and its attention against cached history. The trade is direct: less repeated projection work and prefix reading, but persistent device memory that grows with layers, batch, sequence length, KV heads, head dimension, and bytes per element.
Grouped-query and multi-query attention reduce the number of KV heads without necessarily reducing query heads. Cache paging and block allocation address fragmentation and variable request lengths; they do not change the model’s mathematical attention rule.
Practical estimate
A first-order cache budget is:
2 × layers × batch × cached_tokens × kv_heads × head_dim × bytes_per_element
The factor two is for keys and values. Add allocator overhead and any replicated or speculative state before treating the estimate as capacity.
Connections
- Implement the core operation in Tiny Attention Lab.
- Apply the memory model in Inference and Serving.
- Revisit data movement in CUDA Execution Concepts.
Limitations
This model omits implementation-specific layouts, quantized caches, prefix sharing, sliding windows, offloading, and recomputation policies. It also does not predict kernel efficiency or output quality; those require measurements under the chosen model and serving engine.
Self-test
- For batch 2, sequence 1024, and 8 KV heads, which terms double if sequence grows to 2048?
- Why is caching useful for decoding but less decisive for processing a prefix once?
- What semantic error occurs if a causal mask exposes a future key?
- Which assumptions must be added before converting the estimate into a concurrency limit?