LLM Infra · Series · 02 of 2

What Actually Lives in the KV Cache

Phase 2TypeDeep DiveTime~15 min readPrereqThe Journey of a Request

// TL;DR

Deep dive on the question everyone gets slightly wrong: from a token ID you look up an embedding — but that is not what the KV cache stores.
The tokenizer only makes integer IDs. The embedding is the model’s input. What gets cached is K and V — projections computed inside every layer — so a 32-layer model caches 32 (K, V) pairs per token, not one vector.
K and V come from multiplying the input by learned matrices W_K and W_V — each output is a dot product with one column. Q (query) is computed the same way, used once to score attention, then thrown away and never cached.
Caching is only valid because attention is causal: a token’s K and V never change when later tokens arrive, so they can be frozen and reused.
Worked on a tiny 4-dim model where every number is real: watch prefill fill the cache token by token, then watch two full decode turns reuse it — Q·K → softmax → weighted V.

In the last post we said the KV cache is what makes decode affordable — and that at real batch sizes it can eat more memory than the model weights. That raises the obvious question: what exactly is sitting in it? You have a token ID and a lookup table. Is it the vector from the lookup? Something the model computes later? This post answers that precisely, on a model small enough that every number fits on screen.

The confusion, settled first

Three different vectors get conflated all the time, and only one of them is cached. Here is who produces each, and what happens to it.

Embedding

Token ID → a row looked up in the model’s embedding matrix (learned weights — not the tokenizer). It’s the input to layer 0.

→ not cached

Q — Query

Computed inside each layer from that layer’s input. Used once, immediately, to score against the keys — then discarded.

→ not cached

K, V — Key & Value

Also computed inside each layer, from the same input, via learned W_K/W_V. Every future token needs them.

✓ this is the KV cache

A model tiny enough to watch

Real models hide behind big numbers — d_model = 4096, 32 layers. So we shrink one until the arithmetic is visible. Everything below is this model, computed honestly.

d_model = 4heads = 1head_dim = 4layers = 2prompt = “the cat sat on …”

The tokenizer already ran. These are the IDs, and the emb vector each ID looks up (a row of the embedding matrix — the model’s weights):

the 10emb2331

cat 12emb2232

sat 13emb1233

on 16emb2133

One token, all the way through

Inside a layer, the input vector is multiplied by three learned matrices to produce Q, K, and V. This is just a projection of one token — no attention yet, no token looking at any other; that comes later. Multiplying a vector by a matrix means each output entry is the input dotted with one column — multiply pairwise, then add. Here it is on “the”, with the model’s actual matrices:

the 10enters layer 0 withemb2331

W_K → K

→ K₀

→ K₁

→ K₂

→ K₃

e₁

e₂

e₃

e₄

K₀ = 2·3 + 3·3 + 3·1 + 1·2 = 6 + 9 + 3 + 2 = 20

K₁ = 2·2 + 3·1 + 3·3 + 1·2 = 4 + 3 + 9 + 2 = 18

K₂ = 2·2 + 3·2 + 3·1 + 1·2 = 4 + 6 + 3 + 2 = 15

K₃ = 2·2 + 3·3 + 3·2 + 1·3 = 4 + 9 + 6 + 3 = 22

K =20181522

W_V → V

→ V₀

→ V₁

→ V₂

→ V₃

e₁

e₂

e₃

e₄

V₀ = 2·1 + 3·2 + 3·1 + 1·2 = 2 + 6 + 3 + 2 = 13

V₁ = 2·2 + 3·1 + 3·2 + 1·1 = 4 + 3 + 6 + 1 = 14

V₂ = 2·2 + 3·1 + 3·1 + 1·2 = 4 + 3 + 3 + 2 = 12

V₃ = 2·1 + 3·2 + 3·2 + 1·1 = 2 + 6 + 6 + 1 = 15

V =13141215

Q =18221525computed the same way (through W_Q) — used now and discarded, never stored

Prefill: fill the cache, token by token

Prefill runs that same projection for every prompt token — in one parallel pass — and drops each token’s K and V into the cache. Step through it and watch the rows fill:

0 / 4 tokens cached

the 10cat 12sat 13on 16

Press Step ▶ to send the first token through layer 0.

Layer 0 cache

stored K & V per token

row 0 · the

K20181522V13141215

row 1 · cat

K19191522V13141314

row 2 · sat

K18191523V14131314

row 3 · on

K18201522V13141413

The cache grows by exactly one row per token. Nothing here is ever recomputed.

// question

For each token, what does the model actually write into the KV cache?

The storage stack — it’s per layer

The grid above was one layer. The real cache is that grid repeated for every layer, because each layer computes its own K and V from its own input. The token flows up; each floor deposits a pair.

▲ hidden state flows up ▲

Layer 1deposits K¹ + V¹ per token — its own W_K/W_V

Layer 0deposits K⁰ + V⁰ per token — the grid above

▲ embedding (input) ▲

So the cache is a 4-D block, and its shape is the memory formula from the last post, read left to right:

text

KV cache = [ layers · 2 (K,V) · heads · seq_len · head_dim ] × bytes

Only seq_len grows at runtime. For our toy model one token costs 2 × 1 × 4 × 2 = 16 numbers. For Llama-2-13B it’s ~0.8 MB per token — and at a batch of 32 and 4k context, that block is larger than the weights themselves. Same formula, just bigger dimensions.

Decode: what the cache is for

Now generate the token after “on”. We compute a fresh Q for the current token, then reach into the cache for every earlier token’s K and V. This is the moment the storage pays off.

on 16→Q20181921the only new projection this step

The one subtlety worth saying out loud

We kept it out of the numbers to stay legible, but it’s a favourite interview follow-up: the same word at two positions produces different K and V. Modern models inject position with RoPE (rotary position embedding), which rotates Q and K by the token’s position before the dot product. So position is already baked into K at the moment it’s written to the cache — the cache stores position-aware keys, which is exactly why “cat” at position 1 and “cat” at position 5 don’t collide.

The loop, two full turns

The step above was turn one. Generation just repeats it. Below are turns 1 and 2 with every number shown — watch three things each turn: only the current token’s Q is new, the cache is read (never recomputed), and each turn appends exactly one row before handing off to the next.

TURN 1currenton 16→Q20181921the only new projection

1. Score Q against the 4 cached keys, scale by 1/√4, softmax:

cached token	K (from cache)	Q·K	÷√4	softmax
the	[20, 18, 15, 22]	1471	735.5	0.47
cat	[19, 19, 15, 22]	1469	734.5	0.17
sat	[18, 19, 15, 23]	1470	735.0	0.29
on	[18, 20, 15, 22]	1467	733.5	0.06

2. Blend the cached V by those weights → the attention output:

out13.2913.7112.5914.41

3. out → rest of the network (MLP, layer 1, unembedding) → next token mat 14. Append its K/V → cache now 5 rows, and it becomes the current token for the next turn.

TURN 2currentmat 14→Q1111913the only new projection

1. Score Q against the 5 cached keys, scale by 1/√4, softmax:

cached token	K (from cache)	Q·K	÷√4	softmax
the	[20, 18, 15, 22]	839	419.5	0.17
cat	[19, 19, 15, 22]	839	419.5	0.17
sat	[18, 19, 15, 23]	841	420.5	0.48
on	[18, 20, 15, 22]	839	419.5	0.17
mat	[10, 11, 8, 12]	459	229.5	0.00

2. Blend the cached V by those weights → the attention output:

out13.4813.5213.0014.00

3. out → rest of the network (MLP, layer 1, unembedding) → next token . 3. Append its K/V → cache now 6 rows, and it becomes the current token for the next turn.

Check yourself

The KV cache

0/3 answered

// question

A model has 32 layers. How many (K, V) pairs does it cache per token?

// question

Why is it correct to cache a token’s K and V and never recompute them?

// question

What does the KV cache do to the cost of each decode step?

What’s next

You now know precisely what’s in the cache and why. The natural next question is the one that made this whole topic matter for serving: if the KV cache is this big and grows per request, how does a server pack dozens of them into one GPU without wasting memory? That’s PagedAttention — treating the KV cache like virtual memory, in pages — and it’s the next deep dive.

The cache stores K and V — per-layer projections — not the embedding and not Q.
One (K, V) pair per token per layer; only sequence length grows at runtime.
Caching is valid because causal attention freezes past K/V.
Decode reuses the cache to turn O(n²) into O(n) per step — the reason it’s fast.

series progress0%