Compressing the KV cache

The size equation#

For a transformer with $L$ layers, $H_{kv}$ key–value heads, head dimension $d_h$, batch $B$, sequence length $S$ and $b$ bytes per element:

$$ M_{\text{KV}} = 2 \cdot L \cdot H_{kv} \cdot d_h \cdot B \cdot S \cdot b $$

Linear in sequence length and batch, and independent of the parameter count. Since weights are fixed while $B \cdot S$ grows with load, there is a crossover past which the cache dominates — and for any serious serving deployment it is already past.

Architectural reductions to $H_{kv} \cdot d_h$:

• Multi-query attention (Shazeer, 2019): $H_{kv} = 1$.
• Grouped-query attention (Ainslie et al., 2023): $H_{kv} = H_q / g$, typically $g \in {4, 8}$. The standard choice, because MQA loses quality and GQA does not.
• Multi-head latent attention (DeepSeek-V2, 2024): cache a low-rank latent $c_t^{KV}$ and reconstruct K and V by projection. Reported at roughly 1/10 the KV size of MHA with quality above GQA.

Where the outliers are#

The empirical structure that determines quantisation design (Liu et al., 2024, KIVI):

• Key cache: a small number of channels have magnitudes one to two orders larger than the rest, consistently across tokens. This is the same fixed-channel outlier phenomenon Dettmers et al., 2022 documented in activations for LLM.int8().
• Value cache: no such channel structure; magnitudes are comparatively uniform.

Under symmetric quantisation, error scales with the group's dynamic range. Grouping along the axis that contains the outliers puts them in their own group; grouping across it forces every ordinary value to share a scale set by an outlier. Hence KIVI's asymmetric design:

• K: per-channel, in chunks along the token axis (needed because tokens arrive incrementally, so a whole-sequence per-channel scale would have to be recomputed).
• V: per-token, grouped along channels.

They report 2-bit KIVI reaching near-fp16 quality with 2.6× less peak memory and up to 3.5× throughput. The demonstration above reproduces the qualitative ordering at 4 bits.

Hooper et al., 2024 (KVQuant) push to 3-bit and below with four components: per-channel key quantisation applied before RoPE (rotary mixes channels and destroys the fixed-channel structure), non-uniform (sensitivity-weighted) quantisation levels, per-vector dense-and-sparse decomposition isolating outliers, and normalisation of the attention-sink token. The pre-RoPE detail is the one most likely to be missed in a reimplementation.

fp8, and why hardware prefers it#

Integer quantisation needs scales, dequantisation, and awkward kernels. fp8 needs none of that on Hopper and later, where the tensor cores consume it natively. e4m3 has range ±448 and is the standard choice for the cache; e5m2 trades mantissa for range and is used for gradients.

vLLM's --kv-cache-dtype fp8 is close to free on supported hardware, which is why it, rather than int4, is the default recommendation in production.

Attention sinks#

Xiao et al., 2024 (StreamingLLM) observed that the first few tokens of a sequence receive disproportionate attention regardless of content — attention "sinks" absorbing probability mass the softmax must place somewhere. Two consequences:

1. Evicting the first tokens from a sliding-window cache collapses quality. Keeping four of them restores it.
2. Sink tokens have extreme activation magnitudes and must be excluded from, or handled separately by, any quantisation scheme.

Beyond quantisation#

PagedAttention is the one with no downside and is the reason vLLM became the default serving engine. Kwon et al., 2023 report 2–4× throughput over the then-standard systems purely from eliminating memory waste, with the gain largest at long sequences and large batches.

Papers#

• Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, 2019 — arxiv.org/abs/1911.02150
• Dettmers et al., LLM.int8(), NeurIPS 2022 — arxiv.org/abs/2208.07339
• Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models, EMNLP 2023 — arxiv.org/abs/2305.13245
• Kwon et al., Efficient Memory Management for LLM Serving with PagedAttention, SOSP 2023 — arxiv.org/abs/2309.06180
• Zhang et al., H2O: Heavy-Hitter Oracle for Efficient Generative Inference, NeurIPS 2023 — arxiv.org/abs/2306.14048
• Xiao et al., Efficient Streaming Language Models with Attention Sinks, ICLR 2024 — arxiv.org/abs/2309.17453
• Liu et al., KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache, ICML 2024 — arxiv.org/abs/2402.02750
• Hooper et al., KVQuant, NeurIPS 2024 — arxiv.org/abs/2401.18079
• DeepSeek-AI, DeepSeek-V2, 2024 — arxiv.org/abs/2405.04434 (multi-head latent attention)

What to learn next#

• fp32, bf16, fp8 and int4 — the formats this lesson has been assuming.
• Measuring what quantisation costs you — how to know whether 17% attention error matters.
• GPU memory anatomy — where all the memory goes.