Generating fast: the KV cache

A language model writes one token at a time. Each new token is appended to the input and the model runs again to predict the next one. Done naively, the 500th token of a reply would reprocess all 499 tokens before it, and the 501st would reprocess 500. Almost all of that work repeats work already done. The KV cacheKV cacheStored keys and values for every token already processed, kept between generation steps so each new token only needs its own keys and values computed. It grows with every token in the context.Open in glossary is the reason it does not have to.

What each step actually needs

Go back to attention. To produce the output for position $t$, a head compares the query $q_t$ with the keys of every position up to $t$ and takes a weighted sum of their values:

So the new position needs $k_i$ and $v_i$ for every earlier position $i$, in every layer and every head. Here is the important fact: those keys and values do not change when later tokens arrive. The causal maskCausal maskA mask that stops each token from attending to tokens after it, by setting those attention scores to negative infinity before the softmax. Used by models that generate text left to right.Taught in Many heads and the causal maskOpen in glossary means position $i$ only ever attends to positions at or before it, so everything computed at position $i$, in every layer, depends only on tokens $1$ through $i$. Appending a new token cannot change it.

That makes the cache possible. Compute each position’s keys and values once, store them, and on every later step compute only the new token’s query, key, and value.

import numpy as np

rng = np.random.default_rng(0)
d = 4                                    # head size
Wq, Wk, Wv = (rng.normal(size=(d, d)) for _ in range(3))
# Vectors for 4 prompt tokens and the first 5 generated ones (random here;
# in a real model each comes from the token just generated).
x = rng.normal(size=(9, d))

def attend(q, K, V):
    scores = K @ q / np.sqrt(d)
    w = np.exp(scores - scores.max())
    return (w / w.sum()) @ V

# Without a cache: every step recomputes keys and values for the whole context.
no_cache, passes = [], 0
for t in range(4, 10):                   # 6 steps, context of 4 to 9 tokens
    K, V = x[:t] @ Wk, x[:t] @ Wv
    no_cache.append(attend(x[t - 1] @ Wq, K, V))
    passes += t
print("token passes without a cache:", passes)

# With a cache: the first step fills it, then each step adds one key and value.
K, V = list(x[:4] @ Wk), list(x[:4] @ Wv)
cached, passes = [attend(x[3] @ Wq, np.array(K), np.array(V))], 4
for t in range(5, 10):
    K.append(x[t - 1] @ Wk)
    V.append(x[t - 1] @ Wv)
    cached.append(attend(x[t - 1] @ Wq, np.array(K), np.array(V)))
    passes += 1
print("token passes with a cache:", passes)
print("same outputs:", np.allclose(no_cache, cached))

token passes without a cache: 39
token passes with a cache: 9
same outputs: True

What the cache saves, and what it does not

Without a cache, generating $n$ tokens after a $p$-token prompt processes $p + (p+1) + \dots + (p+n-1)$ token positions, which grows with the square of the length. With a cache it processes the prompt once and then one position per step: $p + n - 1$ in total.

The first step is special. The whole prompt is processed at once, in parallel, to fill the cache. This phase is called prefill, and it keeps the hardware busy with large matrix multiplications. Every step after it, called decode, processes a single token.

Attention itself does not go away. The new token’s query still has to be compared with all $t$ cached keys, so each decode step costs more than the last. And decode has a different bottleneck. To produce one token, the hardware must read every weight of the model and every cached key and value from memory, while doing only about one multiply and one add with each number. With small batches, moving the data takes longer than the arithmetic, so decoding speed is set by memory bandwidth (Shazeer, 2019). Anything that makes the cache smaller makes decoding faster.

How big the cache gets

Every layer stores a key and a value vector for every key/value head and every token:

where $L$ is the number of layers, $h_{\text{kv}}$ the number of key/value heads, $d_{\text{head}}$ the head size, $b$ the bytes per stored number, $T$ the tokens in the context, and $B$ the number of sequences processed together. Llama 2 7B, with 32 layers, 32 key/value heads, head size 128, and 16-bit numbers, needs 512 KiB per token. A full 4,096-token context is 2 GiB, for a single sequence.

import numpy as np

# Cache per token: 2 (keys and values) x layers x KV heads x head size x bytes.
def per_token(layers, kv_heads, head_dim, bytes_per=2):
    return 2 * layers * kv_heads * head_dim * bytes_per

tokens = np.array([1024, 4096, 32768])
models = [("Llama 2 7B", 32, 32, 6_738_415_616),   # layers, KV heads, parameters
          ("Llama 3 8B", 32, 8, 8_030_261_248)]
for name, layers, kv_heads, params in models:
    b = per_token(layers, kv_heads, head_dim=128)
    gib = b * tokens / 1024**3           # one sequence, 16-bit numbers
    print(f"{name}: {b // 1024} KiB per token, GiB at {tokens.tolist()}: {gib.tolist()}")
    # The 16-bit weights take 2 x params bytes.
    print(f"  cache outgrows the weights past {2 * params / b:,.0f} tokens")

Llama 2 7B: 512 KiB per token, GiB at [1024, 4096, 32768]: [0.5, 2.0, 16.0]
  cache outgrows the weights past 25,705 tokens
Llama 3 8B: 128 KiB per token, GiB at [1024, 4096, 32768]: [0.125, 0.5, 4.0]
  cache outgrows the weights past 122,532 tokens

Sharing keys and values

The cache scales with the number of key/value heads, and nothing says that has to equal the number of query heads. Multi-query attention (Shazeer, 2019) lets all query heads share a single key head and a single value head, which shrinks the cache by the number of heads, at some cost in quality. Grouped-query attentionGrouped-query attentionAn attention variant in which groups of query heads share one key head and one value head. It shrinks the KV cache with little loss in quality. Multi-query attention is the extreme case of a single shared key/value head.Open in glossary (Ainslie et al., 2023) sits in between: query heads are split into groups, and each group shares one key/value head. Ainslie et al. found that it recovers quality close to full multi-head attention with speed close to multi-query attention.

It is now the common choice. Llama 2 70B and both Llama 3 8B and 70B use 8 key/value heads; Llama 3 8B has 32 query heads, so its cache is a quarter of what full multi-head attention would need. SmolLM2-135M, which you ran in your browser in Inside a real language model, shares 3 key/value heads among 9 query heads.