Context windows and the KV cache

The big idea

The context windowThe maximum number of tokens a model can take into account at once: the prompt plus everything it has generated so far. is how much text a model can look at in one go. It looks like a simple number on a spec sheet, but two hard costs sit behind it.

First, attention compares every token with every other token. Double the context and you have four times as many comparisons. Second, while generating, the model keeps the keys and values of every earlier token in memory so it does not recompute them. That store, the KV cacheThe stored keys and values of every earlier token, per layer, kept during generation so each new token only computes its own key and value and attends to the stored ones., grows by one entry per token per layer, and for long contexts it can be bigger than the model itself.

Think of a court stenographer reading back testimony. Without notes, answering each new question means replaying the whole recording from the start. With a running transcript, they only add the newest line and glance over the page. The transcript saves huge amounts of time, but the binder gets thicker with every sentence, and eventually the binder, not the stenographer, is the problem.

How it works

Attention is quadratic in context length

With $n$ tokens, self-attention computes an $n\times n$ matrix of scores in every head of every layer. At 4,096 tokens that is $4,096^2=16,777,216$ scores per head. At 131,072 tokens, 32 times longer, it is about 17.2 billion: $32^2=1,024$ times more. The rest of the model (the feed-forward layers and projections) grows only linearly with $n$, so at short contexts attention is a small share of the cost and at very long contexts it dominates.

Clever kernels such as FlashAttention compute exact attention without ever storing the full $n\times n$ matrix, which cuts memory traffic a lot. They do not change the number of comparisons: the arithmetic is still quadratic.

Prefill and decode

Generation has two phases with very different shapes.

• PrefillThe first phase of generation: the whole prompt is processed in one parallel pass, filling the KV cache. Usually limited by compute.: the whole prompt goes through the model in one parallel pass. Every prompt token computes its key and value in every layer, and these fill the cache. Lots of arithmetic, done in big matrix multiplies, so the accelerator is busy computing.
• DecodeThe second phase of generation: one token at a time, each step reading all weights and the whole KV cache to produce a single token. Usually limited by memory bandwidth.: one new token per step. The new token computes its own query, key and value, appends its key and value to the cache, and attends over the whole cache.

Why the cache exists

A token's key and value in a decoder-only model depend only on that token and the tokens before it, thanks to the causal mask. So once computed, they never change. Without a cache, generating token 1,000 would recompute keys and values for all 999 earlier tokens, then token 1,001 would do it all again. With a cache, each step computes one new key and one new value per layer and reuses the rest. The outputs are identical; only the work changes. The Python below checks both claims.

Why decode is memory-bound

In decode, each step produces one token per sequence. To do it, the accelerator must read every weight of the model and the whole KV cache from memory, yet it does only a little arithmetic with each number it reads. Modern accelerators can compute far faster than they can move data, so decode spends most of its time waiting on memory bandwidth. Prefill is the opposite: many tokens share each weight read, so it keeps the arithmetic units busy.

This is why serving systems batch many users together (each weight read is shared across the batch) and why the size of the KV cache, not the number of FLOPs, often sets how many users fit on one GPU ( GPU memory math , Batching and PagedAttention ).

How big is the cache?

For each token, each layer stores one key and one value vector per key/value head:

Take a model with 32 layers, 32 heads of dimension 128, stored in 16-bit (2 bytes). Per token: $2\times 32\times 32\times 128\times 2=524,288$ bytes, which is 512 KiB. At 131,072 tokens that is 64 GiB for a single sequence.

That is why most recent models use grouped-query attentionAn attention variant where several query heads share one key/value head, shrinking the KV cache by the group size with little quality loss. Multi-query attention is the extreme case of a single shared key/value head. (GQA) from Multi-head attention : query heads share fewer key/value heads. With 8 key/value heads instead of 32, the same model needs 128 KiB per token and 16 GiB at 131,072 tokens. With a single shared key/value head (multi-query attention), 2 GiB.

Try it

Try this:

1. At the defaults (32 layers, 8 key/value heads, 128 dims, 131,072 tokens, an 8-billion-parameter model) the cache is 16.0 GiB and the weights are 14.9 GiB. One long conversation needs more memory for its cache than for the model.
2. Set the context to 4,096 and the batch to 32. The cache is again 16.0 GiB: a long context for one user costs the same as short contexts for 32 users. Serving systems trade these against each other constantly.
3. Set key/value heads to 32 (plain multi-head attention), then to 1. The cache scales directly with the number of key/value heads. Switch the precision to 8-bit and watch it halve.

The same mechanics in Python, with a toy single-head model:

# Generation with and without a KV cache: same outputs, very different work.
import numpy as np

rng = np.random.default_rng(0)
d, steps = 16, 64
Wq, Wk, Wv = (rng.normal(size=(d, d)) / 4 for _ in range(3))
xs = rng.normal(size=(steps, d))        # stand-ins for token vectors

def attend(q, K, V):
    s = K @ q / np.sqrt(d)
    w = np.exp(s - s.max()); w /= w.sum()
    return w @ V

# Without a cache: at every step recompute K and V for the whole prefix.
outs_a, work_a = [], 0
for t in range(steps):
    K, V = xs[: t + 1] @ Wk, xs[: t + 1] @ Wv
    work_a += 2 * (t + 1)               # rows of K and V computed this step
    outs_a.append(attend(xs[t] @ Wq, K, V))

# With a cache: compute K and V for the new token only and append.
K_cache, V_cache, outs_b, work_b = [], [], [], 0
for t in range(steps):
    K_cache.append(xs[t] @ Wk); V_cache.append(xs[t] @ Wv)
    work_b += 2
    outs_b.append(attend(xs[t] @ Wq, np.array(K_cache), np.array(V_cache)))

print("same outputs:", np.allclose(outs_a, outs_b))
print(f"K/V rows computed over {steps} steps: no cache {work_a:,}, cache {work_b:,}")

def kv_bytes(layers, kv_heads, head_dim, tokens, batch=1, nbytes=2):
    return 2 * layers * kv_heads * head_dim * tokens * batch * nbytes

GiB = 1024 ** 3
for name, kvh in [("32 kv heads (MHA)", 32), ("8 kv heads (GQA)", 8),
                  ("1 kv head (MQA)", 1)]:
    per_tok = kv_bytes(32, kvh, 128, 1)
    full = kv_bytes(32, kvh, 128, 131_072)
    print(f"{name:<18} {per_tok / 1024:6.0f} KiB per token, "
          f"{full / GiB:5.1f} GiB at 131,072 tokens")

same outputs: True
K/V rows computed over 64 steps: no cache 4,160, cache 128
32 kv heads (MHA)     512 KiB per token,  64.0 GiB at 131,072 tokens
8 kv heads (GQA)      128 KiB per token,  16.0 GiB at 131,072 tokens
1 kv head (MQA)        16 KiB per token,   2.0 GiB at 131,072 tokens

Over 64 steps, recomputing the prefix every time computes 4,160 key and value rows; the cache computes 128, and the outputs match exactly. The second half reproduces the memory numbers from the formula above.

Going deeper

Sliding-window attention

In sliding-window attentionAttention where each token attends only to the most recent W tokens instead of the whole prefix, so cost per token and KV cache size stay bounded by W. each token attends only to the last $W$ tokens. Cost per token is bounded by $W$, and the cache can drop anything older than $W$, so memory stops growing. Information from further back is not lost entirely: layer 2 sees layer 1 outputs that already summarised a window, so stacking $L$ layers lets information travel roughly $L\times W$ tokens. Some models alternate sliding-window layers with full-attention layers to get most of the savings while keeping some exact long-range access.

Sparse attention

Sparse attentionAttention patterns where each token attends to a chosen subset of positions (local neighbours, strided positions, a few global tokens) rather than all of them, cutting the quadratic cost. generalises the window: each token attends to a fixed pattern such as its neighbours, every $k$-th earlier token, and a few global tokens that everyone can see. The number of comparisons grows close to linearly with length. The cost is that whatever the pattern skips, the model can only reach indirectly.

Stretching positions: RoPE scaling

A model trained on 8,192-token sequences has never seen rotary positions ( Embeddings and positional encoding ) beyond 8,192, and quality usually collapses past that point. RoPE scalingTechniques that extend a RoPE model's context by rescaling positions or rotation frequencies so longer sequences map into ranges the model has seen, usually followed by a short fine-tune on long data. fixes this cheaply. Position interpolation divides positions by a factor, so 32,768 tokens are squeezed into the trained range of 0 to 8,192. Later variants rescale the rotation frequencies unevenly, keeping fast-rotating dimensions (which encode nearby order) intact while stretching slow ones. Either way, a short fine-tune on long documents lets the model adapt, which is far cheaper than pretraining at full length.

A long window is not the same as long understanding

Fitting a document in the window does not guarantee the model uses all of it. Models tend to use information near the start and the end of the context more reliably than information buried in the middle, the lost-in-the-middleThe tendency of language models to use information at the beginning or end of a long context more reliably than information in the middle. effect. Practical consequences: put the instructions and the most important material at the start or end, retrieve a few relevant passages instead of pasting everything ( Why RAG? ), and test long-context behaviour on your own data instead of trusting the advertised window.

Reusing the cache across requests

If many requests share the same prefix, such as a long system prompt, the server can keep that prefix's KV cache and skip its prefill next time. Providers expose this as prompt caching, usually at a discount on cached input tokens ( Prompt caching and cost control ).

Trade-offs

Common mistakes

• Sizing GPU memory for the weights only. The KV cache for long contexts or large batches can exceed the weights.
• Assuming a model "reads" everything in a 100,000-token prompt equally well. Test retrieval from the middle of your own documents.
• Thinking FlashAttention makes attention linear. It removes the stored $n\times n$ matrix and much memory traffic, not the quadratic number of comparisons.
• Putting the key instruction in the middle of a long prompt, where it is most likely to be underused.
• Confusing the context window with memory across conversations. The cache lives for one sequence; the next request starts empty unless the prefix is explicitly cached.

In the wild

Check yourself

Recap