Caching keys and values

Inference and Serving, lesson 2

Keep each token's keys and values so decode never recomputes the past.

Each decode step in the last lesson took only the newest token as input. That works because the model keeps every earlier token's keys and values instead of working them out again. This lesson shows why keeping them is safe, what it saves and what it still costs. You need queries, keys and values ("Attention as a soft lookup") and the causal mask ("Masking future tokens").

By the end, you can:

Show why earlier tokens' keys and values stay the same as the text grows.
Count the work a decode step saves by reading from the cache.
Work out how a decode step's memory reads grow with the cache.

Part 1. The past doesn't change

Show why earlier tokens' keys and values stay the same as the text grows.

At every layer, the newest token's query is compared with the keys of the tokens before it, and their values are mixed. Those keys and values come from each earlier token's own vector at that layer. So do those vectors move when a new token arrives? Try it on a layer small enough to work by hand.

Example 1.1

A one-number attention layer

Tokens are single numbers: 1.0, then −1.0, then 0.5. A token's query is 0.5 × itself, its key is 1 × itself and its value is 2 × itself. What does attention give token 2, before and after token 3 arrives?

Make token 2's query0.5 × (−1.0) = −0.5
It may compare with token 1's key, 1.0, and its own, −1.0. The causal mask hides anything later.
Score the two keys−0.5 and 0.5
−0.5 × 1.0 = −0.5 for token 1, and −0.5 × (−1.0) = 0.5 for itself.
Turn scores into weights0.61 ÷ 2.26 = 0.27
Softmax: e to the −0.5 is 0.61 and e to the 0.5 is 1.65, each divided by their sum, 2.26. The other weight is 0.73.
Mix the values0.27 × 2 + 0.73 × (−2) = −0.92
The values are 2 × 1.0 and 2 × (−1.0).
Add token 3still −0.92
Token 3's key isn't among the ones token 2 may see, so nothing in the sum changes.

Answer−0.92, both before and after token 3 arrives.

The same holds at every layer. Token 2's output at layer 1 becomes its input at layer 2, where it again sees only tokens 1 and 2, and so on up. A token added later never enters any of those sums, so every earlier key and value, at every layer, stays as it was.

Figure 1Each row is one token's attention. A new token adds a row at the bottom. The rows above never gain a square, so nothing they worked out can change.

KV cache

The stored keys and values of every earlier token, for every layer and every head. Prefill fills it for the prompt, and each decode step adds the newest token's keys and values.

The guarantee only runs forward. Change a token in the middle, as an app does when it edits its chat history, and every token after it saw the old version. Their stored keys and values are wrong from there on. Only the part before the change can be kept.

Your turn

In the same layer, what does attention give token 1 when it's alone, and after tokens 2 and 3 arrive?

Token 1 sees only itself, so its one weight is 1 and its output is its own value, 2 × 1.0 = 2.0. The mask hides every later token from it, so it stays 2.0.

Show the answer

C. Token 1 sees only itself, so its one weight is 1 and its output is its own value, 2 × 1.0 = 2.0. The mask hides every later token from it, so it stays 2.0.

Part 2. Only the newest token goes through the weights

Count the work a decode step saves by reading from the cache.

Without a cache, each decode step runs the whole text through the model again, just to get one new token. With a cache, the step runs only the newest token: its query, key and value at each layer, then attention over the keys and values already stored.

Example 2.1

Llama 3 8B, 1,000 tokens in

Pushing one token through Llama 3 8B's 8.03 billion weights takes about 2 operations per weight. The text so far is 1,000 tokens. How much weight work does the next decode step need without a cache, and with one?

One token's weight work2 × 8.03 B = 16.06 G
About 16 billion operations to take one token through every layer.
Without a cache1,000 × 16.06 G = 16.06 T
All 1,000 tokens go through again, to produce one new token.
With a cache1 × 16.06 G = 16.06 G
Only the newest token. The other 999 have their keys and values stored.
Compare16.06 T ÷ 16.06 G = 1,000
The saving is the length of the text so far.

Answer16 trillion operations without a cache, 16 billion with one: 1,000 times less.

The outputs are exactly the same, because the stored numbers are the ones recomputing would give. Over a whole reply the gap adds up. Without a cache, a 500-token reply to a 1,000-token prompt pushes 624,750 tokens through the weights. With one, it pushes 1,499: the prompt once, then one token per step.

Figure 2Five decode steps after a 4-token prompt. Without a cache, every step runs every token so far. With one, each step runs only the newest, in orange, and reads the rest from the cache.

Incremental decoding

Decoding that runs only the newest token through the model at each step and reads everything earlier from the KV cache. Serving engines decode this way, and so does Hugging Face transformers unless you pass use_cache=False.

weight work per step = 1 token, not n

Each token makes three vectors at each layer, and the cache holds only two of them. The third is left out on purpose.

Your turn

A KV cache keeps every earlier token's keys and values. Why doesn't it keep their queries too?

A token's query is used once, in the step where that token is the newest, to decide what to look at. Later tokens need its key, to find it, and its value, to read it, but never its query. Storing queries would only waste memory.

Show the answer

D. A token's query is used once, in the step where that token is the newest, to decide what to look at. Later tokens need its key, to find it, and its value, to read it, but never its query. Storing queries would only waste memory.

Your turn

The same model now has 4,000 tokens so far. About how much less weight work does the next step need with a cache?

Without a cache, 4,000 tokens go through the weights: 4,000 × 16.06 G ≈ 64 trillion operations. With one, a single token: 16.06 billion. The saving is the length so far. 16 million is 4,000 squared, but one step's weight work grows only in step with the length.

Show the answer

D. Without a cache, 4,000 tokens go through the weights: 4,000 × 16.06 G ≈ 64 trillion operations. With one, a single token: 16.06 billion. The saving is the length so far. 16 million is 4,000 squared, but one step's weight work grows only in step with the length.

Part 3. What a step still costs

Work out how a decode step's memory reads grow with the cache.

The cache skips the weight work for old tokens, but not attention. The newest query still meets every stored key, and its output still mixes every stored value. Decode is memory-bound, so what matters is that every one of those keys and values is read from memory, on every step.

Example 3.1

Llama 3.1 8B at 1,000 and 32,000 tokens

Llama 3.1 8B has Llama 3 8B's shape but reads up to 131,072 tokens. It keeps 128 KiB, or 131,072 bytes, per cached token (the next lesson works this out), and each step also reads all 16.06 GB of weights. At 3.35 TB/s, how long is a step with 1,000 tokens cached, and with 32,000?

The cache at 1,000 tokens1,000 × 131,072 B = 0.13 GB
Under 1% of what the weights take.
A step at 1,000 tokens16.19 GB ÷ 3.35 TB/s = 4.8 ms
16.06 + 0.13 = 16.19 GB read per step.
The cache at 32,000 tokens32,000 × 131,072 B = 4.19 GB
Now a fifth of what each step reads.
A step at 32,000 tokens20.25 GB ÷ 3.35 TB/s = 6.0 ms
16.06 + 4.19 = 20.25 GB read per step.

Answer4.8 ms a step at 1,000 tokens and 6.0 ms at 32,000, so each token takes 25% longer.

The weights are a fixed cost per step. The cache adds 131,072 bytes for every token in the text, and the attention arithmetic grows the same way. On short texts the fixed part dominates and steps look flat. On long ones the cache's share keeps growing, and on a long enough text it outweighs the model.

Figure 3What each decode step reads. The weights stay at 16.06 GB. The cache, in orange, grows with every token.

O(n)

Read it as "order n": grows in step with n. With a cache, a decode step does fixed weight work plus attention and reads that are O(n) in the n cached tokens. Without one, its weight work is O(n) too, and its attention, redone for every token, is O(n²).

bytes per step = weight bytes + n × cache bytes per token

So "the KV cache makes each step O(1)", a line you'll see in blog posts, is half right. The part that skips recomputing is fixed. The part that reads the cache is not, and on long texts it's the part that grows. "The price of a long context" picks this up.

Your turn

Same model and GPU, now with 16,000 tokens cached. About how long does a decode step take?

The cache is 16,000 × 131,072 bytes ≈ 2.10 GB, so a step reads 16.06 + 2.10 = 18.16 GB. At 3.35 TB/s that's about 5.4 ms. 9.6 ms would double the whole step, but only the cache part grows with length.

Show the answer

D. The cache is 16,000 × 131,072 bytes ≈ 2.10 GB, so a step reads 16.06 + 2.10 = 18.16 GB. At 3.35 TB/s that's about 5.4 ms. 9.6 ms would double the whole step, but only the cache part grows with length.

Where you'll meet it

In interviews

A classic question: what does the KV cache store, and why is it safe? Answer: the keys and values of every earlier token, at every layer and head. It's safe because the causal mask stops a new token from changing anything before it. Then name the cost: memory that grows with every token of every request, read in full on every decode step.

On the job

Every inference engine you're likely to use, including llama.cpp, vLLM and Hugging Face transformers, keeps a KV cache by default. In transformers you can see it: a forward pass returns past_key_values, and the generation loop passes it into the next step along with the one new token. Change an early token and everything stored after it has to be worked out again.

Attention with and without a KV cache, in Python

import math

W_q, W_k, W_v = 0.5, 1.0, 2.0   # one-number projections
xs = [1.0, -1.0, 0.5, 2.0]      # the tokens, already as numbers
work = {"no cache": 0, "cache": 0}

def attend(q, ks, vs):
    w = [math.exp(q * k) for k in ks]
    return round(sum(a * v for a, v in zip(w, vs)) / sum(w), 3)

no_cache = []
for t in range(1, len(xs) + 1):
    ks = [W_k * x for x in xs[:t]]   # every key again
    vs = [W_v * x for x in xs[:t]]   # and every value
    work["no cache"] += 2 * t
    no_cache.append(attend(W_q * xs[t - 1], ks, vs))

K, V, cached = [], [], []
for x in xs:
    K.append(W_k * x)                # one new key
    V.append(W_v * x)                # one new value
    work["cache"] += 2
    cached.append(attend(W_q * x, K, V))

print(no_cache)
print(cached)
print(work)
# [2.0, -0.924, 0.671, 2.962]
# [2.0, -0.924, 0.671, 2.962]
# {'no cache': 20, 'cache': 8}

The first two outputs are part 1's 2.0 and −0.92. Without the cache the key and value projections add up as 2 + 4 + 6 + 8 = 20. With it, each token is projected once: 8.

Check yourself

Four questions from across the lesson. Try each one before you look back.

Two decode steps of the same model with a KV cache: one with 2,000 tokens cached, one with 20,000. Which part of the step's work is the same in both?

One token goes through the weights either way, about 2 operations per weight. Reading the cache and comparing the query with every stored key both grow with the tokens cached, so the 20,000-token step does ten times as much of each.

Show the answer

B. One token goes through the weights either way, about 2 operations per weight. Reading the cache and comparing the query with every stored key both grow with the tokens cached, so the 20,000-token step does ten times as much of each.

Token 7 is added to a text. What happens to the keys and values of tokens 1 to 6, at every layer?

With the causal mask, each token's work at every layer depends only on itself and the tokens before it. Token 7 comes after all of them, so none of their numbers can depend on it. Recomputing them on every step gives the same results every time.

Show the answer

A. With the causal mask, each token's work at every layer depends only on itself and the tokens before it. Token 7 comes after all of them, so none of their numbers can depend on it. Recomputing them on every step gives the same results every time.

A blog says "the KV cache makes decoding O(1) per token". What's the accurate version?

Pushing one token through the weights is a fixed cost, which is where the O(1) idea comes from. But reading and attending over the cache grows with every token added. At short lengths the fixed part dominates and steps look flat. At long lengths the cache part takes over.

Show the answer

A. Pushing one token through the weights is a fixed cost, which is where the O(1) idea comes from. But reading and attending over the cache grows with every token added. At short lengths the fixed part dominates and steps look flat. At long lengths the cache part takes over.

A chat app edits a word in the middle of the conversation history, then asks for the next reply. What happens to the cache?

Tokens before the edit never saw it, so their keys and values still hold. Every token after it attended to the old word, at every layer, so all of those entries are stale. That's why servers that reuse caches across requests only reuse a matching start of the prompt, which "Prefix caching" covers later in this module.

Show the answer

D. Tokens before the edit never saw it, so their keys and values still hold. Every token after it attended to the old word, at every layer, so all of those entries are stale. That's why servers that reuse caches across requests only reuse a matching start of the prompt, which "Prefix caching" covers later in this module.

What to remember

1The past doesn't change
Under the causal mask, a new token can't change anything before it, so earlier keys and values can be stored: token 2's output stays −0.92 when token 3 arrives.
2Only the newest token goes through the weights
With a KV cache, each decode step pushes one token through the weights instead of all n: 16 billion operations instead of 16 trillion at 1,000 tokens.
3What a step still costs
Each step still reads the whole cache: Llama 3.1 8B's steps grow from 4.8 ms at 1,000 tokens to 6.0 ms at 32,000.

Recap and going further

The lesson on one page

Why it's safeUnder the causal mask, past keys and values never change
Recomputing them gives the same numbers
What's keptKeys and values, for every layer and head
Not queries: a past query is never used again
What it savesOnly the newest token goes through the weights
1,000 times less weight work at 1,000 tokens

Keep the past, compute only the new token

What remains

Every step still reads the whole cache
Llama 3.1 8B: 4.8 ms a step at 1,000 tokens, 6.0 ms at 32,000

The cost

Memory that grows with every token
Edit something early and everything after it is stale

Try this

Generate with and without the cache

In Python with Hugging Face transformers, load GPT-2 and generate 300 new tokens greedily twice: once with use_cache=True and once with use_cache=False. Time both. Then run one forward pass on a short prompt and look at the past_key_values it returns: find the first layer's keys and print their shape.

You need: Python with torch and transformers installed. A laptop CPU is enough.

Look for

Whether the two generated texts are exactly the same.
How many times faster the cached run is, and how that changes at 50 new tokens instead of 300.
The numbers in the key tensor's shape, and which one grows as the text gets longer.

The words, in plain language

KV cache: The stored keys and values of every earlier token, for every layer and every head. Prefill fills it for the prompt, and each decode step adds the newest token's keys and values.
Incremental decoding: Decoding that runs only the newest token through the model at each step and reads everything earlier from the KV cache. Serving engines decode this way, and so does Hugging Face transformers unless you pass use_cache=False.
O(n): Read it as "order n": grows in step with n. With a cache, a decode step does fixed weight work plus attention and reads that are O(n) in the n cached tokens. Without one, its weight work is O(n) too, and its attention, redone for every token, is O(n²).

Watch

LLaMA explained: KV-Cache, Rotary Positional Embedding, RMS Norm, Grouped Query Attention, SwiGLUUmar Jamil, on YouTubeThe KV cache section draws exactly which numbers are kept and reused.

02 · Caching keys and values

Where you'll meet it

Check yourself

What to remember

Recap and going further

Generate with and without the cache