The KV Cache
The KV cache stores every past token's key and value vectors so attention never recomputes them. It is the thing decode streams in full on every token, which makes it the platform's most consequential data structure.
CUDA equivalent: the same structure everywhere (vLLM's paged KV made it famous); the platform inflection is that on bandwidth-bound hardware its size converts directly into decode latency.
The arithmetic that runs everything: per token, per layer, the cache holds one
K and one V vector per KV head. For a typical dense 7-8B model
(n_layers ≈ 32, n_kv_heads × head_dim ≈ 1024, fp16) that is ~128 KB per
token, so a 32K-token context carries a ~4 GB cache, and every decoded token
streams all of it through the attention kernels. Three
consequences the glossary keeps meeting:
- Long context turns decode into a cache-streaming benchmark. The sparse-V war story's +22.8% at 32K is precisely a KV-cache traffic deletion.
- GQA is a KV-cache compression scheme as much as an attention variant:
fewer KV heads = a smaller cache = faster decode, which is why
steel attention implements it as a stride
trick (
kv_head_idx = tid.y / gqa_factor). - Quantizing the cache is the next lever after quantizing weights, and it is asymmetric in a way weight quantization is not.
The K/V asymmetry. The
turboquant-mlx project↗
measured that quantizing K to 4 bits or below breaks greedy decode (softmax is
sensitive to small score perturbations, and K sits upstream of the softmax),
while V tolerates 3-bit quantization fine (V is consumed by a weighted
average, which absorbs noise). Its mixed K8/V4 configuration cut a 6.21 GB
fp16 cache to 5.08 GB with greedy output verified identical to baseline. The
general rule: treat K's precision as accuracy-critical and V's as a bandwidth
knob. Practitioner reports of kv_bits=4 corrupting long prefilled contexts
are this asymmetry being ignored, and note the flip side: quantized KV does
not speed up prefill at all, since prefill's cost is
the matmuls, not cache streaming.
Layout matters at the kernel level. The simple layout is one contiguous
ring per layer; serving stacks use paged layouts (fixed-size blocks plus an
indirection table) so many sequences share memory without fragmentation. That
world exists on Metal: the official vLLM plugin ships a paged varlen
attention kernel, and MTPLX carries a
paged-attention Metal kernel family↗
(pagedattention.metal, reshape_and_cache.metal, gather_kv_cache.metal).
Serving-grade batching is its own topic; for single-user decode, the
contiguous cache plus the split-KV decode
kernels is the whole story.