The KV cache

When a model writes text, each new token needs the keys and values of every token before it. Those never change, so they are computed once and stored. This KV cache turns generation from work that grows with the square of the length into work that grows linearly, at the price of a lot of memory.

Key idea: Prefill and decode, caching keys and values, the cache size formula, grouped-query attention, paged attention
Real world: Reading a novel with a notebook. After each chapter you jot down who is who and what happened. To read the next chapter you check your notes, instead of rereading the whole book from page one.
In an LLM: Llama 3 8B stores 131 kB of keys and values for every token (across all 32 layers). A 128k-token conversation needs 17 GB of cache, more than the 16 GB of weights.

Two phases

Generating a reply has two phases with very different shapes:

. The whole prompt goes through the model at once. All of its tokens are processed in parallel, as big matrix multiplications, and the last position predicts the first new token.
Decode. Every new token after that is produced one at a time: feed in the newest token, get the next one, repeat (the generation loop).

At every decode step, attention needs the newest token’s query and the keys and values of every earlier token, in every layer. Here is the important fact: because of the causal mask, an earlier token’s key and value never change when later tokens arrive. Token 3 cannot see token 50, so its key is the same whether the text is 10 tokens long or 1,000.

So there is no need to recompute them. Store each token’s keys and values the first time they are computed, and read them back afterwards. That store is the .

Watch the difference

Step through generation with and without a cache. Highlighted rows are the ones computed in that step; faded rows come from memory.

prefill

newest token’s attention:

K = X·W_K5 × 3

The0.20.02.0cat2.00.20.0sat0.42.00.0on0.00.60.0the0.20.00.0matandsleptallday

V = X·W_V5 × 3

0.10.00.51.00.10.00.21.00.00.00.30.40.10.00.5

newnewnewnewnew·····

rows computed this step5

total, with cache5

would be, no cache5

rows kept in memory5

Prefill: the whole prompt goes through the model at once, so all five key and value rows are computed in one big matrix multiplication. The last row’s output predicts the first new token, “mat”.

Writing “The cat sat on the mat and slept all day” with the toy attention head. Prefill reads the prompt; each decode step feeds in one new token. Switch between a cache and no cache and compare how many key and value rows get computed.

Without a cache, step $t$ recomputes all $t$ rows, so producing $n$ tokens costs about $1 + 2 + \dots + n \approx n^2/2$ row computations per layer. With a cache it is one row per step: $n$ in total. Only the new token goes through the matrix multiplications with the weights; the cache supplies the rest.

example · the reading notebook

You are reading a long mystery novel, one chapter a night. Without notes, each night you would have to reread the book from page one to remember who the butler is. With a notebook, each chapter costs you one chapter of reading plus a glance at your notes. The notebook grows by a page a chapter, so by the end it is thick: you have swapped rereading time for shelf space. That is exactly the trade the KV cache makes.

What the cache costs

Every token leaves one key vector and one value vector in each layer, for each key/value head. So the cache needs

\text{KV bytes} = 2 \times \text{layers} \times \text{KV heads} \times \text{head size} \times \text{bytes per number} \times \text{tokens}

The 2 counts keys and values. Multiply by the number of conversations the server holds at once.

context8k tokensconversations1no GQA (32 KV heads)

per token

2 × 32 layers × 8 KV heads × 128 × 2 bytes = 131,072 bytes

KV cache: 8k tokens × 11.1 GB

weights (16-bit)16.1 GB

per token131 kB

per conversation1.1 GB

total cache1.1 GB

cache ÷ weights0.07×

The cache is 1.1 GB, 7% of the size of the weights (16.1 GB). It grows in a straight line with the context and with the number of conversations: try the sliders. Llama 3 8B shares each key/value head between 4 query heads (grouped-query attention), which makes the cache 4× smaller than it would otherwise be.

How much memory the KV cache needs. Pick a model, then drag the context length and the number of conversations served at once.

Some numbers worth remembering, in 16-bit:

Model	Layers	KV heads × head size	Per token	8k tokens	128k tokens
GPT-2 small	12	12 × 64	36.9 kB	(max 1,024 tokens: 38 MB)
Llama 3 8B	32	8 × 128	131 kB	1.07 GB	17.2 GB
Llama 3 70B	80	8 × 128	328 kB	2.7 GB	42.9 GB

The weights of Llama 3 8B are 16 GB in 16-bit. One long conversation, or a few dozen short ones, needs as much memory again for the cache.

example · serving many users

A chat service runs Llama 3 8B on one 80 GB GPU. The weights take 16 GB, leaving about 64 GB. At 131 kB per token, that holds about 490,000 tokens of cache: about 60 conversations of 8,000 tokens each, or 3 of 128,000. How many users one GPU can serve at once is decided by the cache, not by the model.

Making the cache smaller

The formula shows the levers:

Fewer key/value heads. In (GQA), several query heads share one key/value head. Llama 3 8B has 32 query heads but only 8 key/value heads, so its cache is 4 times smaller than it would be with one per query head; Llama 3 70B shares each among 8. Taken to the extreme, multi-query attention keeps a single key/value head. See efficient attention.
Fewer bytes per number. Storing the cache in 8 or 4 bits halves or quarters it, usually with little loss: see quantizing the KV cache.
Fewer tokens. Sliding-window attention keeps only the last few thousand tokens in some layers; other methods compress or evict old entries.

Paged attention

A server does not know in advance how long each conversation will grow. Reserving room for the maximum length for everyone wastes most of the memory. (introduced by vLLM) stores the cache in small fixed-size blocks, say 16 tokens each, allocated as a conversation grows, like pages of virtual memory in an operating system. A table per conversation lists where its blocks are. Little memory is wasted, so many more conversations fit, and conversations that start with the same prompt can share its blocks.

Caching also works across requests: if many requests begin with the same long system prompt, its keys and values can be computed once and reused (prefix caching), so prefill only runs on the new part.

going deeper

Why decode with a cache is exact

For the newest position $t$ in a layer, attention needs $q_t$ , and $k_j, v_j$ for $j \le t$ . With the causal mask, $k_j$ and $v_j$ depend only on tokens $1..j$ , so they are identical at every later step. Reusing them gives exactly the same output as recomputing, up to floating-point rounding (different kernel shapes can change the order of additions).

Cost per decode step

With a cache, each decode step does about $2N$ FLOPs for the weights (one token through every matrix), plus attention over the cache: $q_t$ against $t$ keys and a weighted sum of $t$ values, about $4\,t\,d_{\text{model}}$ FLOPs per layer. It also reads the whole cache from memory every step. See what a forward pass costs.

In code

# Hugging Face transformers keeps the cache for you:
out = model(input_ids=prompt_ids, use_cache=True)            # prefill
past = out.past_key_values
next_id = out.logits[:, -1].argmax(-1, keepdim=True)
out = model(input_ids=next_id, past_key_values=past, use_cache=True)   # one decode step

Numbers

Per token, 16-bit: GPT-2 small $2 \times 12 \times 12 \times 64 \times 2 = 36{,}864$ bytes; Llama 3 8B $2 \times 32 \times 8 \times 128 \times 2 = 131{,}072$ ; Llama 3 70B $2 \times 80 \times 8 \times 128 \times 2 = 327{,}680$ .

In papers

Multi-query attention: Shazeer, 2019.
Grouped-query attention: Ainslie et al., 2023.
Paged attention and vLLM: Kwon et al., 2023.
Llama 3 architecture details: Llama Team, 2024.