From Attention to the KV Cache

We opened this series with a puzzle from the beginner track: if a model re-reads the whole sequence for every token, why don't long generations grind to a halt? Now, with the attention equations and the causal mask in hand, we can answer it precisely — not with an analogy but with a derivation. The KV cache is the inevitable consequence of one fact: under causal masking, a past token's key and value never change. This closing article derives the cache from the math, walks the decode loop line by line, and quantifies the complexity it saves. It stops short of serving-level cache optimizations, which live in the inference-optimization and speculative-decoding articles.

The invariant hiding in causal attention

Consider decoding at step t: the model has generated tokens 1 through t-1 and is producing token t. The attention output for the current position is softmax(q_t K^T / sqrt(d_k)) V, where q_t is the query for position t, and K and V stack the keys and values of positions 1 through t. The only genuinely new object at this step is q_t. The rows of K and V for positions 1 through t-1 are byte-for-byte identical to what they were at step t-1.

This is not an approximation — it's forced by the causal mask. Because position j's key and value are computed from token j alone (through the fixed projections W_K and W_V) and never attend forward, they cannot depend on anything generated later. So recomputing them every step, as a naive implementation does, reproduces the exact same numbers. That redundancy is the whole opportunity.

The cache: store K and V, append each step

The optimization writes itself. Keep two running buffers per layer — the key cache and the value cache. At each decode step, compute only the new token's q, k, and v; append its k and v to the caches; then attend the single new query over the full cached K and V. No past key or value is ever recomputed. The generated token feeds back in as the next input, and the loop repeats.

That is the entire KV cache. Everything else — how the buffers are laid out in memory, how they're paged or shared across requests — is engineering on top of this core recurrence, and belongs to the optimization article. The mechanism itself is just: compute one column, append, attend.

Prefill fills, decode appends

The two inference phases map cleanly onto the cache. Prefill takes the entire prompt of length n and computes all its keys and values at once — a batched, matrix-matrix pass that populates the cache with n rows per layer in one shot. This is the compute-bound burst from the arithmetic article, and it's why the first token after a long prompt takes a moment.

Decode then extends the cache one row at a time. Each step is a matrix-vector pass: one new query, one appended key and value, one attention scan over the now-(n+1)-row cache. Prefill and decode run identical attention math; they differ only in how many query rows they process at once — many for prefill, exactly one for decode.

The complexity shift

Quantify what the cache buys. Without it, generating token t re-runs attention over the whole prefix from scratch, an order t^2 x d cost per step; summed over a full sequence of length n that's order n^3 x d of wasted work across the generation. With the cache, step t does only the fresh dense projections plus a single-query scan over t cached rows — order t x d per step, or order n^2 x d over the whole sequence, with no recomputation of past keys and values at all.

So the cache collapses per-step attention work from reprocessing the entire prefix to touching each cached row once: from an order n^2 x d recompute at every step down to an order n x d scan. That is the difference between a generation that slows quadratically per token and one that stays flat enough to feel instant.

What the cache costs

The savings are paid for in memory. The cache holds two vectors per token per layer, so its size is roughly 2 x layers x kv-heads x head-dim x bytes, multiplied by sequence length — linear growth that, as we saw, often hits the practical ceiling before compute does. And because decode is memory-bound, every step must stream both the model weights and the growing cache out of memory, so a larger cache directly slows token generation.

This is exactly why architectures reduce the number of key/value heads. Multi-query attention shares a single K/V head across all query heads; grouped-query attention uses a small number of K/V groups. Both shrink the cache and its bandwidth cost dramatically while keeping most of multi-head attention's expressiveness — a direct lever on the invariant we derived.

Why this matters

You've now derived, rather than merely accepted, the central optimization of transformer inference. The causal mask makes past keys and values immutable; immutability makes caching correct; caching turns a quadratic-per-step recompute into a linear scan; and the memory that costs is why grouped-query attention and context limits exist. Every performance property from the beginner series now has a mechanical reason underneath it.

From here the road forks into serving: how to lay the cache out in memory, share it across requests, quantize it, and speculate past the memory wall. Those are the subjects of the inference-optimization and speculative-decoding articles. But the foundation is complete — you can read the attention equations and see exactly why a long chat stays fast, and what it costs to keep it that way.

Glossary

Decode step

One iteration of autoregressive generation producing a single new token.

Key/value invariant

The fact that a past token's key and value are fixed once computed under causal masking.

KV cache

Per-layer buffers of keys and values, appended each step so the past is never recomputed.

Prefill

The batched pass that computes and caches keys and values for the whole prompt at once.

Multi-query attention (MQA)

Sharing a single key/value head across all query heads to shrink the cache.

Grouped-query attention (GQA)

Using a small number of key/value groups as a middle ground between multi-head and multi-query.

Memory-bound

Limited by how fast weights and cache can be read, as in single-token decode.