Module 2 · Inference execution

Inference and the KV Cache

Generating text is two phases: a parallel prefill of the prompt, then a one-token-at-a-time decode. Naively, decode recomputes attention for the whole history every step. The KV cache stores the past keys and values so each step does only a sliver of new work — trading memory for compute.

Two phases: parallel prefill, then sequential decode

step-through

Prefill ingests the whole prompt in one parallel pass; decode then emits one token per pass, feeding each output back as the next input.

prompt / generated token compute-bound (parallel) memory-bound (sequential) KV cache

Prefill and decode run the same transformer layers — they differ only in granularity. Prefill pushes many tokens through at once (lots of parallel math), while decode pushes a single new token through but must read the entire prior sequence. Both build Q, K, V matrices with one row per token. The next two figures zoom into the attention step to see where decode wastes work — and how to stop it.

Without a KV cache: recomputing the whole history every step

step-through

Each decode step rebuilds K and V for every previous token, then multiplies Q×Kᵀ — yet the shaded block is identical to last step's result. That redundant work grows every step: quadratic total compute.

new token's Q K / V (needed) recomputed work (wasted) new attention scores
The key observation Look at the top-left corner of Q×Kᵀ: it is identical to the previous iteration. Causal masking means past tokens never attend to future ones, so their scores can't change. The only genuinely new work is the new token's Q row against the full Kᵀ, then against V. Everything else is recompute.

With the KV cache: compute one new token, append, attend

step-through

Past K and V live in the cache (amber). Each step computes only the new token's Q, K, V, appends K,V to the cache, and multiplies the one new Q row against the whole cache. Per-step work is tiny and constant-shaped — linear total instead of quadratic.

cached K / V (reused) new token's Q only-new compute

The KV cache trades memory for compute: by storing past keys and values, it turns the per-step quadratic attention into a single Q-row times the cache — linear work per token. That's why decode is memory-bound: the bottleneck shifts from doing math to reading the ever-growing cache out of HBM. Next: how big does that cache actually get?

The cost: KV cache memory grows linearly with the sequence

auto-play

Every generated token appends one K and one V vector in every head of every layer. The cache fills steadily as the sequence grows — and that is exactly what MQA / GQA / MLA shrink.

KV cache (in HBM) free HBM over the cap → eviction
Per-token cache size Bytes per token ≈ layers × heads × 2 (K and V) × d_head × bytes_per_element. It scales with sequence length and with batch size, and can grow unboundedly — so runtimes set a maximum size and evict tokens past it (carefully: dropping the earliest attention-sink tokens collapses quality). MQA / GQA / MLA from Module 1 cut the heads factor, shrinking every number above.