A KV cache from scratch, and what it costs

What I built

This is the fourth module of llm-lab. It adds a KV cache to the GPT from module 3.

To write the next token, the plain way is to run the model over everything written so far. But the causal mask from the attention module means earlier positions never change when a token is appended. So their keys and values can be kept, and each new token only needs one position's work.

Generation doesn't need gradients, so this module skips the autograd engine and runs the model in plain NumPy. The core is a forward pass that takes only the new tokens, writes their keys and values into the cache, and attends over everything stored so far:

cache['k'][i, :, :, n:n + t] = heads(h @ g('Wk'), H)
cache['v'][i, :, :, n:n + t] = heads(h @ g('Wv'), H)
k, v = cache['k'][i, :, :, :n + t], cache['v'][i, :, :, :n + t]
a = softmax(q @ k.transpose(0, 1, 3, 2) / math.sqrt(cfg['C'] // H) + mask)

The same function handles a whole prompt at once (prefill) and one token at a time.

Proving it's right

A cache is only allowed if it changes nothing, so most of the 11 tests check exactly that. The cached model gives the same logits as module 3's engine to 2e-14 at every step up to its 128-character limit, and writes the same text, greedy and sampled. Prefill in one go, prefill in chunks and one token at a time all match the full forward pass. Another test checks that the cache refuses to overflow.

How much faster

The trained model only has 128 positions, so for longer contexts I used the same architecture with random weights. The output is nonsense, but the work per token doesn't depend on the weights. Without the cache, trained and random took the same time within a few percent. The cached steps are so short (~1 ms) that their timings jumped around by up to 40% in both directions from run to run, which is noise, not the weights.

There are two baselines without a cache. My first uncached forward built the attention scores for all pairs of positions at once. The fairer one does attention 128 query rows at a time, which gives the same logits (to 1e-15) with much less memory, and is faster at long contexts.

Each point is the fastest of several tries. The first and last columns are ranges over eight runs on the same laptop, the middle one over two. Single timings moved by up to 2x between runs, so the shape matters more than the exact numbers.

At 2,048 tokens the cache is ~430–510x faster than the fairer baseline, and ~530–800x faster than building all the scores at once.

Without the cache, every new token redoes every position, plus attention over every pair, so each doubling of the context costs about 3–3.5x more by 2,048 tokens. With the cache, a step stays at about 1 ms up to roughly 512 tokens and then grows linearly. Counting multiply-adds, attention over the cache catches up with the rest of a step at 6C = 768 tokens for this model. In measured time the bend comes later: at 768 tokens attention only adds ~0.3 ms to a step that's ~0.7 ms of fixed overhead.

The measured speedups are always smaller than counting multiply-adds predicts: 3–4x at 16 tokens, where the count says 16x. At this size a cached step is mostly that fixed overhead from Python and NumPy, not arithmetic. Writing 512 tokens from scratch took ~27–35 s without the cache and ~0.45–0.7 s with it.

What it costs in memory

The cache holds a key and a value for every layer and every position. That's 8 KB per token here, 16.8 MB at 2,048 tokens. It passes the size of the model itself (6.55 MB) at 800 tokens.

Without a cache nothing is kept between steps, but each step needs memory for attention. My first uncached forward peaked at 581 MB at 2,048 tokens, and I first wrote that up as 35x the cache. That was how I'd written the forward, not a cost of having no cache. Doing attention 128 rows at a time peaks at ~61 MB, about 3.6x the cache. FlashAttention takes the same idea further on GPUs. So for one sequence, the cache is still smaller than what a step without it needs, but not by much.

The real cost shows up when serving many sequences at once, because every sequence needs its own cache. Scaled up, calculated rather than measured: Llama-2-7B needs 0.5 MB of cache per token in fp16, so 2.1 GB for one 4,096-token sequence. Eight such sequences need more memory than the model's weights. That's why Llama-2-70B uses grouped-query attention, where several query heads share one key and value. It has 10x the parameters of the 7B model but a smaller cache per token, 0.33 MB.

The limit

With learned position embeddings, the cache can't slide past the model's 128 positions. Every cached key was computed with its position baked in, and after dropping the oldest token they're all off by one. I tried a naive sliding cache anyway: over 50 steps past the limit, it picked the same next character as the real model only 10% of the time.

So past 128 characters, generation goes back to recomputing the last 128 every step, like module 3 did. For short outputs the cache is much faster: 100 characters take ~0.1 s. For 300 characters it only goes from ~4.2 s to ~3.5 s, because most of them are past the limit.

Two practical notes

The first version prefilled a 16k-token context in one call. That builds a 16k × 16k score matrix per head, about 8.6 GB in float64, and the run took 7.7 minutes instead of 2. Prefilling 256 tokens at a time fills the cache identically. Real inference servers do chunked prefill for the same reason.

Skipping the autograd graph is a speedup on its own. Without any cache, 300 characters take ~4.2 s with the plain NumPy forward, against ~10 s through module 3's engine.

This was one laptop, float64 NumPy on a CPU, and random weights for the long contexts. At this size the cached step is overhead rather than arithmetic, so a GPU or a bigger model would give a different speedup.