The KV Cache

Chapter 1 left us with a number: single-stream decoding runs at 0.34% of an H100's peak, because generating one token requires reading all 16 GB of weights to do 16 GFLOPs of work. Arithmetic intensity exactly 1, against a machine that wants 295.

The fix is obvious and I've already named it: batch. Serve 64 requests at once and the same weight read does 64× the work. Intensity 64, still short of the ridge but two orders of magnitude better.

This chapter is about why that doesn't work as well as it should, and what the field did about it. The short version: the fix for the weight bottleneck creates a new bottleneck, and the new one is worse, because unlike the weights it grows without bound.

Why the cache exists#

Autoregressive generation has an obvious redundancy. To produce token $t+1$ you run attention over tokens $1 \ldots t$. To produce token $t+2$ you run attention over $1 \ldots t+1$ — recomputing the keys and values for all $t$ tokens you already processed, whose values haven't changed and can't, because causal masking means position $i$'s representation never depends on anything after it.

So cache them. Keep $K$ and $V$ for every position, and each new token computes only its own $K$ and $V$, appends them, and attends against the accumulated cache.

The saving is real: without a cache, generating $n$ tokens costs $O(n^2)$ forward passes' worth of work; with one, $O(n)$. Nobody has ever seriously proposed not doing this.

But note what it is. You have converted a compute problem into a memory problem, and taken on a state that grows linearly with every token generated, for every concurrent user, and never shrinks until the request ends.

The arithmetic#

Per token, across all layers:

The 2 is $K$ and $V$. Note $n_{kv}$ — the number of key/value heads, not query heads. That distinction is the entire subject of the next section.

For Llama 3 8B ($L = 32$, $n_{kv} = 8$, $d_h = 128$, bf16):

A clean number, worth memorizing. And its consequences:

An 8B model's weights are 15 GiB, leaving 65 GiB on an H100. At 8k context with GQA that's 65 concurrent sequences. With plain MHA it would be 16. And a single 128k-context request under MHA would need 64 GiB of cache — four times the model itself.

That's the whole motivation. Now the responses, in order of increasing cleverness.

Response 1: share the KV heads#

The observation behind MQA is almost embarrassingly simple. The cache scales with $n_{kv}$. So make $n_{kv}$ smaller.

Multi-Query Attention takes it to the limit: all query heads share a single key head and a single value head.Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, 2019 — three years before anyone needed it badly enough to adopt. Each head still computes its own queries, so heads still attend differently; they just consult a shared key/value space. For Llama 3 8B that's a 32× cache reduction, 512 KiB down to 16 KiB per token.

It also degrades quality, and it destabilizes training. The empirical finding is that going all the way to one head is too far.

Grouped-Query Attention is the interpolation: partition the $n_h$ query heads into $g$ groups, one KV head per group.Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, EMNLP 2023. The paper's other contribution is "uptraining" — converting an existing MHA checkpoint to GQA by mean-pooling the KV heads and fine-tuning on ~5% of the original compute, which is why adoption was so fast. $g = n_h$ is MHA; $g = 1$ is MQA; Llama 3 uses $g = 8$ against 32 query heads, for a 4× reduction at essentially no quality cost.

The reason this works is worth stating, since it's the only part that isn't mechanical: attention heads within a layer are substantially redundant in what they attend to, even when they differ in what they do with it. Queries carry most of the head-specific behaviour; keys and values carry more shared structure. GQA exploits that asymmetry, and the fact that it costs so little quality is genuine evidence about how transformers use their heads.

Every open model that matters now ships GQA. It is the default, and it is one of the few free lunches in this book.

Response 2: compress instead of share#

GQA reduces the cache by storing fewer heads. Multi-head Latent Attention reduces it by storing something else entirely.

Instead of caching $K$ and $V$, project the hidden state down into a low-rank latent $c_t = W^{DKV} h_t$ of dimension $d_c \ll n_h d_h$, and cache that. Reconstruct keys and values on the fly via up-projections $W^{UK}, W^{UV}$.DeepSeek-AI, DeepSeek-V2, 2024, §2.1. Refined in DeepSeek-V3, 2024.

Stated that way it sounds like it just moves the cost — you've traded cache bytes for reconstruction FLOPs, and decode is bandwidth-bound, so maybe that's fine. But the actual trick is better than that. Since attention scores are $q^\top k = (W^{UQ}c^Q)^\top(W^{UK}c^{KV})$, the up-projection matrix can be absorbed into the query projection at inference time: $(W^{UK})^\top W^{UQ}$ is a fixed matrix you fold in once. The keys are never reconstructed at all. The same absorption works for $W^{UV}$ into the output projection.

DeepSeek-V3 uses $d_c = 512$ plus a 64-wide decoupled RoPE key — decoupled because rotary embeddings have to be applied at full head resolution, so that component can't be compressed. That's 576 elements per token per layer, versus $2 \times n_h \times d_h$ for MHA.

And the surprise: MLA outperforms MHA on quality in DeepSeek's ablations. The low-rank bottleneck appears to act as a useful regularizer rather than a lossy compression. That's the rare case in this book of a technique that improves both budgets at once.

The cost is architectural complexity — MLA is not a drop-in change, it interacts awkwardly with RoPE, and it requires custom kernels. Which is why, despite being better, it hasn't displaced GQA outside of DeepSeek's own models.

Compare all four at your own configuration:

Worth doing: switch to 128k context and watch MHA's per-sequence cache exceed the entire GPU. Then switch the cache precision to fp8 and note that it's a straight 2× on every variant — quantizing the cache composes with all of them, and is often the easiest win available.

Response 3: stop wasting the memory you have#

Everything above shrinks the cache. This one stops squandering it, and the gains turned out to be comparable.

The problem: a naive server allocates a contiguous buffer per request, sized for the maximum possible length, because you can't know in advance how long a generation runs. A request that could generate 2048 tokens gets 2048 tokens of cache reserved, and if it stops after 60 you've wasted 97% of it. Add internal fragmentation from over-allocation and external fragmentation from varying request sizes, and measured utilization in pre-2023 serving systems was 20–40%.

PagedAttention applies the oldest idea in operating systems.Kwon et al., Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023 — the paper behind vLLM. Split the cache into fixed-size blocks (typically 16 tokens), maintain a per-sequence block table mapping logical positions to physical blocks, and allocate blocks on demand. Non-contiguous physical storage, contiguous logical view, and the attention kernel is modified to gather through the block table.

Waste drops to under 4% — at most one partially-filled block per sequence. That's a 2–4× increase in serving throughput from nothing but allocation strategy, no change to the model at all.

It also enables copy-on-write sharing. Parallel samples from one prompt, or beam search branches, share the prompt's physical blocks with a reference count and only diverge on write.

Response 4: don't recompute the prefix#

The last one is the technique the source notes call "stateful caching," and it's the highest-leverage optimization in production serving that nobody outside serving teams talks about.

In a chat application, turn $n$'s prompt is turn $n-1$'s prompt plus the exchange since. In an agent loop, every step re-sends a system prompt and tool definitions that haven't changed. In a RAG system, thousands of requests share an identical instruction preamble. The KV cache for a shared prefix is identical across all of them, because causal attention means a prefix's keys and values don't depend on anything that follows.

So: hash prefixes, keep the computed blocks, and on a new request find the longest cached prefix and prefill only the suffix.

The clean implementation is a radix tree over token sequences, where each node owns the cache blocks for its path from the root.Zheng et al., SGLang: Efficient Execution of Structured Language Model Programs, NeurIPS 2024 — RadixAttention. vLLM ships equivalent automatic prefix caching. Shared prefixes are shared paths, which is exactly the structure a conversation tree already has. Eviction is LRU over leaves. PagedAttention's block table is what makes this cheap — sharing is just pointing two block tables at the same physical blocks.

The payoff scales with prefix reuse, and in agent workloads that's most of the prompt. Cutting time-to-first-token by 5–10× on a long-system-prompt workload is routine. If you run an inference service and haven't turned this on, it is almost certainly the largest single win available to you.

What decode looks like once you've done all this#

Return to the roofline and re-derive the number that started the chapter. With batch $B$ and context $S$, decode's intensity is:

Batching amortizes the weight term — that was the point. But the cache term scales with $B$ too, exactly like the numerator. So as $B$ grows, intensity asymptotes:

a constant. For Llama 3 8B at 8k context in bf16, that ceiling is around 64 FLOP/byte — still below the ridge point of 295. No batch size crosses it.

That is the real reason this chapter has five sections instead of one. Once the KV cache dominates HBM traffic, batching stops helping, and the only remaining moves are on the cache itself: fewer bytes per token (GQA, MLA), fewer bits per byte (fp8/int4 cache), or fewer tokens (sliding window, prefix sharing). Try it in the chapter 1 explorer — push the batch slider to 512 at 8k context and watch the decode point stall short of the ridge.