Codetail

Article 12 of 15

Context Windows and the KV Cache

Why longer conversations cost more.

20 min read

Attention's quadratic cost: every token against every token

It's tempting to think a bigger context window is just a bigger number in a config file. Double it and you double the cost, the same way doubling a batch size roughly doubles training time. That mental model is wrong, and it's wrong in a way that costs real money the moment you put a model into production.

The Attention article established the mechanism: every token computes a query, and that query gets compared against the key of every other token in the sequence, including itself. For a sequence of length n, that's n tokens each comparing against n tokens, n × n pairwise comparisons, per attention head, per layer. Double nand you don't double the comparisons, you roughly quadruple them.

Put actual numbers on it. A prompt of 1,024 tokens needs 1,024 × 1,024, about 1.05 million pairwise comparisons for one head in one layer. Grow that to 4,096 tokens, a 4x increase in length, and it's 4,096 × 4,096, about 16.8 million comparisons, 16x the work for 4x the tokens. Push to a 32,768-token prompt, 32x the original length, and it's roughly 1.07 billion comparisons, 1,024x the original work. That 1,024x isn't a coincidence, it's 32 squared, exactly what “quadratic” predicts.

And that's per head, per layer. A real model runs this same arithmetic across every attention head in every transformer block, dozens of times over, so the constant multiplier in front of n² is large before the quadratic term even starts dominating. This is the direct, mechanical reason a 128K-token context window isn't just “the same thing but bigger,” it's a different cost regime entirely, and it's the first of two costs this article covers. The second shows up not when the model reads a long prompt, but when it has to keep generating one token at a time afterward.

What the KV cache actually caches

Here's the naive picture of generating a long reply: the model writes token 500, then to write token 501 it reruns the full forward pass over all 500 tokens again, from scratch, just to produce one more word. Token 502 reruns the whole thing over 501 tokens. If that were actually how it worked, a thousand-token reply would cost roughly as much as generating a thousand separate prompts of increasing length, and the Sampling and Generation article's one-token-at-a-time loop would be brutally expensive in practice, not just in theory.

That's not what happens, because most of that rerun is wasted work. When generating token 501, the new query only needs to be computed once, for the new token. But attending to the sequence still means comparing that new query against the key of every earlier token, and taking a weighted sum using the value of every earlier token, exactly the mechanism from the Attention article. Those earlier keys and values are a fixed function of tokens that already exist and never change again. Recomputing them at every single step recomputes the exact same numbers, over and over, forever, for the life of the conversation.

The KV cacheis the fix, and it's almost embarrassingly simple: the first time a token's key and value vectors get computed, store them. On every later step, reuse the stored versions instead of recomputing them, and only do fresh work for the one new token: compute its query, key, and value, append its key and value to the cache, and attend the new query against the full cache, old entries plus the one just added.

Generation loop, reusing cached keys and values

Python
1def generate(prompt_tokens, model, max_new_tokens):
2 # first pass: process the whole prompt at once, cache builds up here
3 logits, kv_cache = model.forward(prompt_tokens, kv_cache=None)
4 generated = []
5
6 for _ in range(max_new_tokens):
7 next_token = sample(logits)
8 generated.append(next_token)
9
10 # only the new token runs through the model, cache supplies the rest
11 logits, kv_cache = model.forward([next_token], kv_cache=kv_cache)
12
13 return generated

Rule: the KV cache doesn't shrink the attention computation itself, the new token still compares against every previous key. What it eliminates is redundant re-projection, recomputing the same old tokens' key and value vectors again and again. Without it, generating a 1,000-token reply would redo token 1's key and value projection a thousand times over, for no new information.

This is why generation is described as having two distinct phases with two different cost shapes: prefill, processing the whole prompt at once and building the initial cache, and decode, generating one token at a time afterward, each step doing new work proportional to one token, not the whole sequence. The cache is what makes decode cheap per step. It is not, however, free to keep around.

The memory cost of the cache itself

It's easy to read the previous section and think the KV cache is a clean win, more speed, no downside. It's a trade, not a discount: it buys back compute by spending memory, and that memory bill is a real, growing number that sits on a GPU for as long as the conversation stays open.

The cache has to store one key vector and one value vector, per token, per layer, per attention head. That's four multipliers stacked on top of sequence length: two vectors (key and value), times the number of layers, times the number of heads, times however many bytes each number takes up. Sequence length is the only one of those that grows as the conversation goes on, but it's multiplied by three constants that are already large in any real model.

Take a 7B-parameter-class model as a concrete example: 32 layers, 32 attention heads, a 128-dimension head size (4,096 hidden size total), weights stored in fp16, 2 bytes per number. Per token, the cache needs 2 (key and value) × 32 layers × 32 heads × 128 dimensions × 2 bytes, which comes out to 512 KB, per token, before the conversation has produced a single reply. At 1,024 tokens of conversation, that's about 512 MB. At 4,096 tokens, about 2 GB. At 32,768 tokens, about 16 GB, and that's one conversation, held by one user, for as long as that chat stays active.

That 16 GB number is worth sitting with. The model's own weights, for this same 7B-class model in fp16, are only around 14 GB, loaded once and shared across every request being served. A handful of long, active conversations can carry more memory in KV caches than the model itself takes up. This is frequently the actual ceiling on how many simultaneous conversations a server can serve, not the size of the model's weights. Providers don't under-provision GPUs by accident, the KV cache is where a server's memory budget quietly disappears.

Why doubling context length isn't a free upgrade

Put the last three sections together and a “we now support 128K context” announcement stops reading like a config change and starts reading like a capacity planning decision. Every long conversation pays two separate costs, on two separate curves, at the same time.

The first cost is prefill: processing the initial prompt, all at once, pays the full n² attention cost from the first section, up front, before the model writes a single word of reply. Doubling the prompt length roughly quadruples that one-time bill. The second cost is the cache: once generation starts, each new token is cheap thanks to the KV cache, but the cache itself grows linearly with total sequence length, and it sits in GPU memory for the entire time that conversation stays open, multiplied by however many conversations are running at once. A longer context window is a quadratic compute commitment at the start and a linear, but large and long-lived, memory commitment for the whole conversation after that.

Slide the sequence length below and watch the two costs pull apart. At small n they look similar. By 32K tokens they don't, and neither curve is exaggerated, they're the same n² and n from the first three sections, just plotted side by side.

Attention compute vs. KV cache memory4K tokens
1281K4K16K32K
Attention compute

16.78M

pairwise comparisons, per head per layer

KV cache memoryn

2.00 GB

for this conversation, this model config

relative to a 1,024 token baseline: attention compute is 16.0×, KV cache memory is 4.0×

Memory assumes a 7B-class model: 32 layers, 32 attention heads, 128-dim head size, fp16 weights. Bar heights plot growth shape on a shared log scale, not the raw comparisons-vs-bytes units, which is why the compute bar is always exactly twice the memory bar's height: n² grows twice as fast as n on a log scale. The numbers above each bar are the real, unnormalized values.

In practice:this is exactly why context-window increases in real products come with real infrastructure cost, rationed rollout, or a higher price per token at long context, not just a flag flipped in a config file. It's also why techniques like sliding-window attention, capping how far back a token can attend, and grouped-query attention, sharing key and value projections across multiple query heads to shrink the cache's head-count multiplier, exist. Both are active research areas aimed squarely at bending these two curves back down, not covered in depth here.

None of this changes what a language model fundamentally is, a next-word guesser, or how it decides which word comes next, covered in the Sampling and Generation article. It changes what that guesser costs to run at scale, and it's the reason context length shows up as a pricing tier, not just a feature flag. Next, the series turns from cost to correctness: how do you actually know if a model is good, covered in Evaluating LLMs.