44  KV-Cache: Autoregressive Inference Efficiency

← Back to: BERT & GPT

TipTL;DR

A decoder generates one token at a time, and each new token attends to every previous one. Recomputing all past keys and values at every step is quadratic waste; the KV-cache stores them once so each step is linear in context length. Its memory cost — batch × layers × heads × seq × \(d_k\) — then becomes the inference bottleneck, which MQA, GQA, and MLA attack by shrinking the K/V footprint.

44.1 Why decoding recomputes

In autoregressive decoding the model appends one token per step and re-runs attention. Naively, generating token \(t\) recomputes \(K\) and \(V\) for all positions \(1..t\) — but those projections never change for tokens already produced. Only the new token’s query needs the fresh past. Caching \(K,V\) turns each step from \(O(t)\) recompute into \(O(1)\) projection plus an \(O(t)\) attention read.

\[ \text{step } t:\quad q_t = x_t W_Q,\quad K_{1:t} = [\,K_{1:t-1}\,;\, x_t W_K\,],\quad V_{1:t} = [\,V_{1:t-1}\,;\, x_t W_V\,]. \]

The bracketed history is read from cache; only the last row is computed.

44.2 The memory bottleneck

The cache size grows with everything at once:

\[ \text{KV memory} = 2 \times b \times L \times h \times n \times d_k \times \text{(bytes)}, \]

for batch \(b\), layers \(L\), heads \(h\), sequence \(n\). For a long-context LLM this dwarfs the parameter memory and caps batch size — inference becomes memory-bandwidth bound, not compute bound. Reducing the K/V term is therefore the highest-leverage inference optimization.

44.3 Attention-head reduction lineage

All three methods keep full query heads but share or compress K/V:

• MQA (Shazeer 2019 (Shazeer, 2019)) — all query heads share a single K/V head. Shrinks the cache by a factor of \(h\), at some quality cost.
• GQA (Ainslie 2023 (Ainslie et al., 2023)) — query heads share \(G\) K/V groups, interpolating between MHA (\(G{=}h\)) and MQA (\(G{=}1\)). The pragmatic default in LLaMA-2/3 and most open models — most of the savings, little quality loss.
• MLA (DeepSeek-V2 2024 (DeepSeek-AI, 2024)) — low-rank joint compression of K and V into a small latent that is cached and re-expanded per head, cutting cache further while retaining multi-head expressivity.

44.4 Related inference trick: speculative decoding

Orthogonal to cache size: speculative decoding (Chen et al., 2023) uses a small draft model to propose several tokens, which the large model then verifies in one parallel forward pass. Accepted tokens are kept, so multiple tokens can be produced per expensive step without changing the output distribution.

NoteKey takeaway

The KV-cache makes autoregressive decoding linear instead of quadratic by storing past keys/values; its memory footprint then becomes the dominant inference constraint. MQA → GQA → MLA progressively shrink that footprint, and speculative decoding attacks the per-step latency from a different angle.

← Back to: BERT & GPT · Output Layer & Decoding

Shazeer, N. (2019). Fast transformer decoding: One write-head is all you need. arXiv Preprint arXiv:1911.02150.

Ainslie, J., Lee-Thorp, J., Jong, M. de, Zemlyanskiy, Y., Lebrón, F., & Sanghai, S. (2023). GQA: Training generalized multi-query transformer models from multi-head checkpoints. arXiv Preprint arXiv:2305.13245.

Chen, C., Borgeaud, S., Irving, G., Lespiau, J.-B., Sifre, L., & Jumper, J. (2023). Accelerating large language model decoding with speculative sampling. arXiv Preprint arXiv:2302.01318.

DeepSeek-AI. (2024). DeepSeek-V2: A strong, economical, and efficient mixture-of-experts language model. arXiv Preprint arXiv:2405.04434. https://arxiv.org/abs/2405.04434