Attention variants, and the cache that eats your memory

Why there is a cache at all

Generating token 500 requires attending to the keys and values of tokens 1 to 499. Recomputing them every step would be quadratic and absurd, so they are computed once and kept. That store is the KV cache, and on any serious deployment it is the thing that runs out, not the weights.

The size is exact and worth memorising:

bytes per token = 2 * num_layers * num_kv_heads * head_dim * bytes_per_element

The leading 2 is for keys and values. For Llama-3.1-8B (32 layers, 8 key-value heads, head dimension 128, fp16) that is 2 × 32 × 8 × 128 × 2 = 131,072 bytes per token, or 128 KB.

• 8,000 tokens of context: about 1.0 GB
• 32,000 tokens: about 4.2 GB
• 128,000 tokens: about 16.8 GB, which is as much as the weights

And that is for one conversation. Ten concurrent users at 8k context each need 10 GB of cache on top of the model. This is why a server that runs beautifully in a demo falls over with twelve people on it.

Three attention layouts

Multi-head attention gives every query head its own key and value heads. Maximum quality, maximum cache. Llama 2’s 7B has 32 query heads and 32 key-value heads, so its cache is 512 KB per token: four times the Llama 3 model above, at a similar parameter count.

Multi-query attention gives all query heads a single shared key-value head. Smallest possible cache, some quality cost.

Grouped-query attention is the compromise everything now uses: several query heads share one key-value head. The ratio num_attention_heads / num_key_value_heads is the saving. At 32 and 8, the cache is a quarter of what full multi-head would need, with a quality difference small enough that the whole field adopted it within a year.

This is why num_key_value_heads is the field to look at in a config when you care about serving. Two models of identical parameter count can differ by four times in how many users one GPU will hold.

Sliding window attention attacks the problem differently: each token attends only to the last N tokens. Mistral’s early models used a 4,096-token window; Gemma alternates local sliding-window layers with occasional global layers, so most layers keep a small fixed cache while a few see everything. The cache stops growing past the window, which turns a linear memory cost into a constant one. The trade is that information from far back reaches the present only indirectly, through several layers of local attention.

What you can actually do about it

Quantise the cache. Most serving stacks can store keys and values in 8-bit rather than 16-bit, halving the memory for a small quality cost that is usually invisible on ordinary tasks:

llama-server -m model.gguf -c 32768 --cache-type-k q8_0 --cache-type-v q8_0
vllm serve Qwen/Qwen3-8B --kv-cache-dtype fp8 --max-model-len 32768

Test it before you trust it. Cache quantisation hurts more on long-context retrieval than on short chat, because subtle key distinctions are what let the model find a specific fact far back.

Cap the context you actually allow. Serving a 128k model with --max-model-len 8192 is not cowardice; it is choosing to serve sixteen times as many users. Most production traffic never needs more.

Let the serving layer manage it. vLLM’s paged attention allocates the cache in fixed blocks like operating-system pages instead of reserving a contiguous maximum per request. That removes most of the waste from over-allocating for requests that turn out short, and is a large part of why vLLM serves several times more concurrent users than a naive implementation on the same hardware.

Reuse the shared prefix. If every request starts with the same 2,000-token system prompt, prefix caching computes it once and shares it across requests. On a chat product this is often the single biggest saving available, and it is a flag rather than a project.

The arithmetic to do before you buy

Take your GPU memory. Subtract the weights at your chosen precision. Subtract about 1 to 2 GB of runtime overhead. Divide what remains by (bytes per token × your context limit). That is your concurrency.

For a 24 GB card running the 8B at Q4 with 8k context: 24 − 5 − 2 = 17 GB, divided by 1.0 GB per conversation, is about 16 concurrent conversations. If that number is smaller than your traffic, the fix is a shorter context limit or cache quantisation long before it is a bigger model.