Tutorials › Understanding Transformers › KV Caching

Understanding Transformers · Part 14 of 16

KV Caching

Why generation stays fast even as the text grows — by not recomputing the past.

The intuition: the past doesn't change, so don't recompute it

During training, the model processes a whole sequence in parallel. During generation it produces one token at a time, feeding each new token back in to produce the next. Done naively, every new step would recompute the Keys and Values for every earlier token, in every layer — repeating almost all of the previous step's work.

But under a causal mask (Part 8), an earlier token's representation never depends on a token that comes after it. Once the model has processed “The cat sat on the,” adding “mat” doesn't change anything about “cat.” So “cat”'s Key and Value, computed once, stay valid for the rest of the generation. Storing them and reusing them is the KV cache.

KV cache during generation Keys and values for past tokens are stored; only the new token's projections are computed each step. cached from earlier steps K,VThe K,Vcat K,Vsat K,Von/the new: q,k,v “mat” computed now the new Query compares against every cached Key, then blends the cached Values the new K,V are appended to the cache for the next step
Each step computes only the newest token's projections and reuses everything before it.

Cache the Keys and Values, not the old Queries

A Query is only used while its own token is the current position being processed. Once that step is done, no future token needs that old Query again. Keys and Values are different: every future token will compare its new Query against the old Keys and pull from the old Values. So the cache holds K and V, and discards each Query after it's used.

At generation step t, for each layer: compute the new token's qt, kt, vt; append kt, vt to that layer's stored tensors; score qt against all cached Keys; blend the cached Values.

scorest = qt [K1 … Kt]T / √dk
outputt = softmax(scorest) [V1 … Vt]
There's a separate cache at every layer. A token's Key and Value in layer 1 are not its Key and Value in layer 12 — each layer sees a different hidden representation and uses different projection matrices (Part 12). The cache therefore stores K and V for every past position, in every layer, for every head.

What it saves, and what it costs

BenefitCost
Skips recomputing every past token's K/V at each step, so token-by-token generation stays fast.Memory grows with context length — every layer must retain all past Keys and Values.
Lets each new token attend to the full history.Reading a large cache eats memory bandwidth and eventually becomes the bottleneck.

Longer contexts need proportionally more cache. Grouped-query attention trims the cost by using a smaller number of Key–Value heads, each shared by a group of Query heads (Ainslie et al., 2023).

Previous ← The Complete Transformer Block
Next How GPT Trains and Generates →