Asked about A decode loop with and without a KV cache

Write a generation loop with a KV cache, and say what it saves.

An answer that works

Prefill runs the prompt through in one batched pass, with the causal mask, and keeps the keys and values. Each decode step then embeds the single new token, computes one query, key and value vector, appends the key and value to the cache, attends over the whole cache with no mask, and produces logits for one position. It saves recomputing every past key and value — constant work per step instead of work proportional to the sequence so far — and it computes one row of attention scores instead of the full square matrix. The output is identical; a cache is exact, not an approximation.

Then they push further

Why is it exact?

Because of the causal mask. A token's key and value depend only on that token and the ones before it, never on what follows, so they cannot change as generation continues. That is the same property as a prefix of a sequence producing unchanged outputs when the rest is truncated.

Why is there no mask in the decode step?

The mask prevents attending to the future, and during decode the future has not been generated. Everything in the cache is already in the past. The mask is needed in prefill and in training, where many positions are processed together.

What does it cost you?

Memory, and it is usually the binding constraint on concurrency. Two vectors per token per layer per KV head. For an 80-layer model with 8 KV heads and head dimension 128 in fp16, about 0.31 MB per token, so 8,000 tokens is 2.5 GB for one sequence and thirty-odd concurrent sequences fill an 80 GB card. Grouped-query attention and paged allocation are both responses to that.

What gets marked up

Anyone can say the cache stores past keys and values. The follow-ups are where it separates: that it is exact and why, that decode needs no mask, and that the cost is memory rather than compute.

Reading an answer is not being able to give one. The section that teaches this is A decode loop with and without a KV cache, 4 hours of Quantization & inference optimisation.

Learn this properly

Also asked about Quantization & inference optimisation

What is the bottleneck in LLM inference?What does quantization trade?Why is batching almost free?When would you self-host a model?Will an 8 billion parameter model serve 16 concurrent users with 8k contexts on a 24 GB GPU?Why does weight-only quantization speed up generation?Why does prompt length change time to first token but barely change time per output token?How would you work out the fastest a model could possibly generate on a given GPU?
All 14 questions on Quantization & inference optimisation