KV Cache

Why Does This Exist?

You generate text one token at a time. To pick token 1,000, the model must weigh it against the 999 tokens before it, which means self-attention needs every one of those tokens' key and value vectors close at hand. If those terms are new, how LLMs work sets up next-token prediction first.

The naive way rebuilds all of them. At step 1 you project 1 token, at step 2 you re-project 2, at step 3 you re-project 3, and so on. Nothing about token 500's key changes when token 999 arrives, yet the naive loop recomputes it hundreds of times. Across a full generation of $n$ tokens that is $1 + 2 + \ldots + n$ projections, which grows like $n$ squared. For 3 tokens that is 6 projections. For 1,000 it is 500,500.

The KV cache deletes that waste. Compute each token's keys and values once, store them, and append only the newcomer's pair at each step. Projection work across the whole generation drops from about $n$ squared to $n$.

This page covers the per-sequence cache itself. Its neighbours answer the follow-ups: paged attention is how servers store the cache without wasting memory, prefix caching shares one cache across requests, and grouped-query attention shrinks it.

Think of It Like This

How It Actually Works

Where keys and values come from

At each attention layer, token $i$ carries an embedding $x_i$, a list of numbers describing it in context. Two fixed weight matrices, $W_K$ and $W_V$, project it:

In plain words: $k_i$ is the key vector, what token $i$ offers up for matching. $v_i$ is the value vector, what token $i$ contributes to the output if it gets picked. Fixed means frozen after training; neither matrix changes while you generate.

To extend the sequence, the newest token forms a query $q$, what it is looking for. It scores $q$ against every key, normalises the scores into weights, and blends the values:

Here $s_i$ is the score saying how much the new token cares about position $i$, $d$ is the head dimension (one number shared by the whole head), and $\mathrm{softmax}$ exponentiates each score and divides by the total, so the weights $w_i$ are positive and sum to 1. The output is a weighted blend of stored values. Dividing by $\sqrt{d}$ keeps the scores in a range where softmax can still tell them apart.

Why old rows never change

Two facts make caching exact rather than approximate. First, causality: in a decoder, token $i$ may only look at positions up to $i$ (causal masking enforces this), so later tokens never retroactively alter an earlier token's embedding. Second, the matrices $W_K$ and $W_V$ are frozen. Same input through the same matrices gives the same output, so the $k_1$ you computed at step 1 is identical to the $k_1$ you would recompute at step 50. Storing it loses nothing.

That rests on one assumption: causal, decoder-side generation with a fixed prefix. Encoders that read the whole sequence at once, and any edit to an earlier token, break it, which is what the first pitfall below is about.

Worked example

One head, head dimension $d = 2$, two tokens cached:

The new query is $q = [1, 1]$. The scale is $\sqrt{2} \approx 1.41$.

Scores, showing every step:

The scores tie, so softmax splits evenly: $w = [0.5, 0.5]$. Quick check: $e^{0.71} \approx 2.03$ twice, total $\approx 4.06$, and $2.03 / 4.06 = 0.5$ per share. Then:

The new token now computes its own pair, call them $K_3$ and $V_3$, and the cache grows from 2 rows to 3. Next step scores its query against all three stored rows. Count the projections: step 3 did exactly 1, the newcomer's. Across 3 steps the cache did 3 total, while the naive loop did $1 + 2 + 3 = 6$. Same answers, half the projection work, and the gap widens every step.

What it costs in memory

The cache trades memory for speed. The number of stored values is:

where $L$ is the layer count, $H$ the key-value heads per layer, $d$ the numbers per head, and $n$ the tokens stored so far. The factor of 2 counts keys and values separately. Multiply by $b$, the bytes per stored number, for bytes.

Toy check against the fixture above: 1 layer, 1 head, $d = 2$, 3 tokens, 2 bytes each. That is $2 \times 1 \times 1 \times 2 \times 3 = 12$ numbers, or 24 bytes. Count them directly: 3 rows times a key pair plus a value pair is $3 \times 4 = 12$. It agrees.

One honesty note before you celebrate: the cache removes re-projection, not re-scoring. Each step still scores the new query against $n$ stored keys, so total scoring across a generation still grows like $n$ squared. That half is what flash attention attacks. And when there is nothing to reuse, skip the cache entirely: scoring a full prompt in one parallel pass or generating a single label gains nothing from it. At long contexts the cache itself turns into the bottleneck, which is why sliding-window attention caps how many rows you keep.

Code

This snippet replays the worked example: two cached rows, one new query, and the append that grows the cache. No ML library needed.

1import math2
3
4def softmax(xs: list[float]) -> list[float]:5    exps = [math.exp(x) for x in xs]6    total = sum(exps)7    return [e / total for e in exps]8
9
10def dot(a: list[float], b: list[float]) -> float:11    return sum(x * y for x, y in zip(a, b))12
13
14# Rows stored at steps 1-2; step 3 only reads them.15keys: list[list[float]] = [[1.0, 0.0], [0.0, 1.0]]16values: list[list[float]] = [[4.0, 0.0], [0.0, 2.0]]17
18query: list[float] = [1.0, 1.0]19scale: float = math.sqrt(2.0)  # 1 / sqrt(d) with head dim d = 220scores = [dot(query, k) / scale for k in keys]21weights = softmax(scores)22out = [sum(w * v[i] for w, v in zip(weights, values)) for i in range(2)]23
24print([round(s, 2) for s in scores])    # -> [0.71, 0.71]25print([round(w, 2) for w in weights])   # -> [0.5, 0.5]26print([round(x, 2) for x in out])       # -> [2.0, 1.0]27
28# Step 3 appends its own pair; the old rows are never touched again.29keys.append([1.0, 1.0])30values.append([1.0, -1.0])31print(len(keys))                        # -> 3

The prints match the hand computation: tied scores split 50/50 and blend the stored values into $[2.0, 1.0]$. The append is the whole mechanism. One projection per step, forever.

Watch Out For

The Quick Version

• A KV cache stores each past token's key and value vectors because causal generation never changes them once written.
• Each decoding step then computes one new key-value pair and scores one new query against the stored rows, with weights that sum to 1.
• Across a full generation this cuts projection work from about $n$ squared to $n$, while per-step scoring against stored keys still costs proportional to $n$.
• Its memory is 2 times layers times key-value heads times head dim times tokens times bytes per value, and it grows with every token: the toy above holds 12 numbers in 24 bytes.
• Never reuse a cache after its prefix changes, and budget cache memory per concurrent request before blaming compute.

What to Read Next

• Paged Attention
• Prefix Caching
• Speculative Decoding
• Grouped-Query Attention