The KV cache | Latent

Generating with and without a cache

Each row is one generation step. Squares mark the positions whose keys and values the step needs.

Keys and values computed in this step
Read from the cache

Without a cache

Without a cache: which positions have their keys and values computed or read from the cache at each step
Step	The	·cat	·sat	·on	·the	·warm	·mat	.	·It	Predicts
1										·the
2
3
4
5
6

Token passes so far: 4
Attention scores so far: 10
This step: 4 passes, 10 scores

With a KV cache

With a KV cache: which positions have their keys and values computed or read from the cache at each step
Step	The	·cat	·sat	·on	·the	·warm	·mat	.	·It	Predicts
1										·the
2
3
4
5
6

Token passes so far: 4
Attention scores so far: 10
This step: 4 passes, 10 scores

Step1 of 6

Token passes, 100-token prompt then 1,000 new tokens599,500 without, 1,099 with

Challenges

0 of 1 complete

Step until generating without a cache has cost more than four times the token passes of generating with one. (not yet)

A token pass pushes one position through every layer's projections and MLP. Attention scores count query-key dot products per layer and head. The text is an example; the counts depend only on the lengths.

Try this

Step to the end and compare the token passes in the two grids.
Notice that attention scores still grow each step even with the cache.

Attention with and without a cachePython with NumPy

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

How big is the cache?

Keys and values for every layer, head, and token, compared with the model's own weights.

Weights12.6 GiB

KV cache2.00 GiB

Llama 2 7B: 32 layers, 32 query heads, one key/value head per query head, head size 128. Trained context: 4,096 tokens.

Cache per token512 KiB

Whole cache2.00 GiB

Cache compared with weights0.16x

Challenges

0 of 3 complete

With one sequence, find the context length where the cache just overtakes the weights: between 1 and 1.1 times their size. (not yet)
Make the cache bigger than the weights without going past the model's trained context length. (not yet)
Find a model whose cache would be 8 times larger without grouped-query attention. (not yet)

Try this

Find the context length at which the cache outgrows the weights for the smaller models. (For Llama 3 70B at batch size 1 it would take about 430,000 tokens, beyond the slider.)
For a grouped-query model, turn on the toggle to see the cache it avoids.

Size the cache for two modelsPython with NumPy

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