The KV Cache: Speed at a Price

The base transformer is complete. From here, every module answers a practical pressure that only appears when you make a model big and actually run it. The first pressure shows up the moment a model generates text, and it is about memory at inference time. This chapter builds the fix that every modern model relies on — the KV cache — and then exposes the problem it creates, which the rest of the module is devoted to solving.

A tax on every generated token

Recall how generation works: the model predicts one token, appends it to the input, and runs the whole thing again to predict the next. Watch what that means for the computation. To produce the 100th token, the model processes a 99-token sequence; to produce the 101st, it reprocesses all 100; and so on. Each step redoes almost everything the previous step did. Generation is drowning in repeated work.

Look inside attention to see exactly what repeats. At each step the model computes a query, key and value for every token, scores every query against every key, and blends the values. But the keys and values of the earlier tokens do not change when a new token is appended — the token at position 5 has the same key it had three steps ago. We are recomputing identical keys and values again and again.

The one insight that unlocks the fix

Here is the observation that makes everything efficient. To predict the next token, the model only needs the output for the last position. The final layer produces a score vector for every position, but we read off only the last one — the prediction of what follows the most recent token. Everything upstream exists only to produce that single last-position output, and that output depends on the last token's query attending over all the keys and values.

So for each new token we genuinely need: its own fresh query, key and value — and the keys and values of every earlier token. The earlier queries we never need again. And the earlier keys and values we already computed. So why recompute them? Store them.

The cache

The fix is exactly that. During generation, keep a running store — the KV cache — of the keys and values of every token processed so far, in every layer. When a new token arrives:

1. Compute only its query, key and value.
2. Append its key and value to the cache.
3. Score its query against all cached keys, blend all cached values.
4. Produce the one output needed to predict the next token.

No past key or value is ever recomputed. Queries are not cached, because a past token's query is never needed again — hence the name: we cache K and V only.

The payoff is a change in how cost scales. Without the cache, producing a sequence of length $n$ costs on the order of $n^2$ work, because step $t$ redoes the work of all earlier steps. With the cache, each step adds a fixed amount, so the total is linear in $n$. In practice this is a large, easily-measured speed-up — generation with caching runs several times faster than without, and the gap widens the longer the text.

The dark side: the cache grows

Caching trades computation for memory, and memory is not free. The size of the KV cache is the product of everything that makes a model big:

Two of those factors are the trouble. The cache stores a key and a value for every head ($n\cdot h$), in every layer, for every token in the context ($S$). So it grows linearly with the conversation length — and modern context windows are enormous.

The numbers are sobering. For a small model the cache is tens of megabytes. For a GPT-3-scale model it is a few gigabytes. For a frontier-scale model — dozens of layers, 128 heads, a context window of 100,000 tokens — the KV cache balloons to the order of hundreds of gigabytes, larger than the model's own weights. At that point the cache, not the parameters, is what fills the accelerator's memory, and holding it there slows everything else down.

So we have a genuine dilemma. The KV cache is essential — without it, generation is hopelessly slow. But in its plain form it consumes more memory than we can afford at long context. The question that drives the rest of this module is sharp: can we keep the cache's speed while drastically shrinking its size? The next three chapters are three increasingly clever answers — multi-query attention, grouped-query attention, and finally latent attention.