The KV Cache: Why It's Allowed, and What It Costs

Part 1 of 4 on how a transformer remembers. → KV cache · Linear attention · Kernelization · Recurrent memory

A language model writes one token at a time. Taken literally, that should be catastrophic: to produce the 500th token you’d re-read the previous 499, and to produce the 501st you’d re-read 500. Generation would be quadratic in its own output length.

It isn’t, because of one table. This post is about why that table is allowed to exist, exactly what goes into it, and the bill it quietly runs up. Every diagram below is animated — press play.

Only two of the three

Every attention head projects every token three ways — into a query, a key, and a value. Only two of the three ever get cached, and the name tells you which. Worth asking why, because the answer turns out to be the cleanest door into everything else.

Here is one head at position $t$:

\[o_t \;=\; \mathrm{softmax}\!\left(\frac{q_t K^{\top}}{\sqrt{d}}\right) V\]

Ignore the operation for a second and look at the shapes. $q_t$ is a single row — this token’s query, and nothing else’s. $K$ and $V$ are stacks: the keys and values of every token so far, all of them, needed in full, at every step.

So a decoding step consumes one query and the entire history of keys and values. Which means $q_t$ gets read exactly once — at the step that built it — while $k_i$ gets read at step $i$, and $i+1$, and $i+2$, and every step until the sequence ends.

A cache only pays for itself on something you read more than once. Watch the counters.

So the cache holds what the new row has to look at, not the row itself. Queries are scaffolding — built, fired, discarded. Keys and values are the record.

What that one row actually does

“q is one row” is easy to say and easy to skip past, so here is the row doing its job, number by number. There is nothing in it beyond a dot product, a softmax, and a weighted average — but watching where each value comes from is what makes the rest of this post concrete.

What actually goes into the two tables

So much for what the rows are for. Where do they come from? Watch a sentence turn into them.

Why you’re allowed to keep it

Here’s the part that usually gets skipped. Caching a computation is only valid if the cached value can’t change. So why can’t $k_3$ change?

Because of the direction of the arrows. In a causal model, the hidden state $h_i$ is built from tokens $1\ldots i$ and nothing else — the attention mask forbids position $i$ from looking right. And $k_i = h_i W_K$ is a per-position map. So the moment token 3 has been seen, $k_3$ and $v_3$ are final. No future token can reach back and perturb them.

The same causality gives you the second half: appending a token doesn’t change any earlier output either, because query $t$ never attended to anything past $t$ in the first place. Old rows of the attention matrix are untouched, bit for bit.

Toggle the mask off in the panel below and watch both properties die at once.

This is also the cleanest way to see why BERT-style bidirectional models can’t stream. There, every token sees every other token, so adding one word at the end rewrites every representation in the sentence. Nothing is final until the whole input is final, and a cache of “final” values has nothing to hold. Autoregressive generation and the KV cache are the same design decision, seen from two sides.

Prefill and decode are different machines

Because the prompt is known all at once, its rows don’t have to be built one at a time. That splits inference into two phases with completely different characters.

Prefill runs the whole prompt in one pass. Every prompt token’s $k$ and $v$ is independent of the others, so the GPU does it as one large matmul. Lots of arithmetic per byte of weights moved — the hardware is busy.

Decode then goes one token at a time, and each step has to attend over everything already in the cache. The arithmetic is tiny — one query row — but the entire cache gets pulled out of memory again. The hardware mostly waits.

That asymmetry is why serving stacks talk about the two phases separately, why time-to-first-token and tokens-per-second have different bottlenecks, and why people bother to run prefill and decode on separate machines. Prefill is a compute problem. Decode is a memory problem. And the thing decode is bottlenecked on is precisely the cache.

What it buys

Strip the cache out and the cost of generation is easy to write down. To produce token $t$ you must rebuild the $k,v$ of all $t$ preceding tokens, so over $N$ generated tokens you pay

\[1 + 2 + \cdots + N \;=\; \frac{N(N+1)}{2} \quad\text{token-passes, instead of}\quad N\]

Drag $N$ and watch the top bar bend.

What it costs

Nothing is free. The cache is a list, and a list grows:

\[\text{bytes} \;=\; 2 \times L \times H_{kv} \times d_{\text{head}} \times N \times \text{sizeof(dtype)}\]

The 2 is K and V. $L$ is layers, $H_{kv}$ is key/value heads, $N$ is context length. Everything except $N$ is fixed by the architecture, so the cache is exactly linear in how much you’ve said.

Put real numbers in: 32 layers, 8 KV heads, head dim 128, fp16 — a shape typical of an 8B-class open model — comes out at 128 KB per token. A 32K conversation is 4 GB. Of cache. For one user.

Two consequences, and the second is the one people underrate.

It eats your batch size. Concurrency is what makes serving economical, and every concurrent user needs their own cache. The memory left after the weights divides by 4 GB, and that quotient is how many conversations your GPU can hold. Grouped-query attention (GQA) exists mostly to shrink that number’s denominator — fewer KV heads, proportionally smaller cache. It works, and it’s a constant factor applied to a term that still grows linearly.

It eats your bandwidth. Look again at the last stat in the panel: the bytes re-read per generated token. Decode has to stream the whole cache through the memory bus for every single token it emits. Your tokens-per-second is set by that number far more than by your GPU’s FLOPs. Long context doesn’t just cost storage — it makes every subsequent token slower.

Where this goes

So the KV cache is not a hack. It’s the direct consequence of causal masking, and it converts generation from quadratic to linear. It is also, structurally, a lossless list that grows without bound, and that is the thing everything after this post is trying to get rid of.

The next move is not to make the list smaller. It’s to notice that the list is only ever consumed through one operation — $\mathrm{softmax}(q K^{\top})V$ — and to ask whether that operation could be rearranged so the list never has to be stored at all. Take the softmax out and it can: the whole history collapses into a single fixed-size matrix. That’s the next post.

---

One table, one row per token, and every arrow pointing forward.

References

• Attention / the transformer — Vaswani et al., Attention Is All You Need (arXiv:1706.03762)
• Where the cache is first spelled out for inference — Shazeer, Fast Transformer Decoding: One Write-Head is All You Need (arXiv:1911.02150) — also the origin of multi-query attention
• GQA — Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (arXiv:2305.13245)
• Paging the cache instead of over-allocating it — Kwon et al., Efficient Memory Management for LLM Serving with PagedAttention (vLLM, arXiv:2309.06180)
• Why decode is memory-bound — Pope et al., Efficiently Scaling Transformer Inference (arXiv:2211.05102)