The part that never changes: the KV cache

Last post ended with an accusation. A token’s poster and synopsis are made from that token alone. They do not change when the sentence gets longer. And yet the loop from three posts ago rebuilds every one of them, from scratch, on every step.

Time to see what that costs.

The next token comes out of one row

Start from the end, because the end is narrow.

The model produces a vector for every position, and the LM head turns each of those into scores over the vocabulary. To pick the next token we take the scores at the last position. Every other position is scored and dropped, which was already true in the second post and is worth saying again now that we know what is underneath it.

Follow that backwards through a layer. The last position’s output came from its own query, scored against the keys of everything before it, pulling a blend out of the values. One query. One row of attention. One output.

Everything the model must produce, on any given step, is a single row of the square.

The loop computes the whole square

Here is the loop again.

for _ in range(max_new_tokens):
    logits = model(ids, use_cache=False).logits
    next_id = logits[:, -1, :].argmax(dim=-1, keepdim=True)
    ids = torch.cat([ids, next_id], dim=-1)

ids is the whole sequence and it gets longer every time round. On the tenth step the model runs a complete forward pass over ten positions: ten embedding lookups, ten sets of queries, keys and values in every layer, a ten by ten attention matrix, ten trips through the MLP, ten rows of logits.

Then the next line takes one row.

That is not a small overhead, and you can price it without a benchmark. Count the token positions the model runs over. With a prompt of p tokens and n tokens to generate, step k processes p + k positions, so the whole generation costs

p + (p+1) + (p+2) + ... + (p+n-1)

Take a twenty token prompt and two hundred generated tokens. That sum is 23,900 token positions. The same output with the keys and values kept around costs 20 for the prompt plus 199 single tokens: 219. About a hundred and ten times the work, and the ratio gets worse the longer you generate, because the sum is quadratic and the alternative is not.

Nothing was measured to get those numbers. They are counting, not timing, which is why they can be stated in a post that has no benchmark in it. If you want the felt version, time each step and watch the number climb. Step one is cheap. Step two hundred is not.

The other rows were already correct

The interesting part is not that the rest of the square is unused. It is that it is identical to last step.

Row 2 of that picture, the third token deciding how to divide its attention, was computed on the previous step too, from the same query and the same keys. Nothing that happened afterwards could have changed it.

That is not luck. It is causality, from last post, doing a second job.

A token only ever attends backwards, so nothing that arrives later can reach it, so nothing later can change its residual stream, so nothing later can change the key and the value it produces. The mask that was introduced to stop the model reading ahead during training turns out to be the thing that makes generation cheap.

A token’s key and value depend only on that token and the place it sits. They are computed once and they are correct forever.

Put two consecutive steps next to each other and the invariant is the picture.

The horizontal edges did not move. One cell was appended to each of them, and one row appeared in the square. Everything else in the second picture is the same tensor as in the first, and the loop we have been running rebuilds all of it anyway.

The keys are different in every layer

One complication, and it is the reason the store is bigger than you might first guess.

Layer 12’s key for token 5 is not built from token 5’s embedding. It is built from token 5’s residual stream as it stands after eleven layers, and that stream holds everything token 5’s own queries pulled in on the way up. The poster a token presents at layer 12 is one it wrote using what it learned at layers 1 through 11. Every layer has its own keys and its own values, and they are not transformations of each other in any way you could shortcut.

Which invites the obvious objection. If a deep key depends on shallow attention, is it really fixed?

It is, and for exactly the same reason. The attention at layer 3 that helped build token 5’s layer 12 key also only ever looked backwards. Token 5’s stream, at every height in the stack, is a function of tokens 0 through 5 and nothing else. Add a sixth token and nothing anywhere in that column moves.

So the thing to keep is not a key and a value per token. It is a key and a value per token per layer: a whole column through the stack, frozen the moment that token exists.

Keep the horizontal edges

Look at the square again, with the edges named.

The top and the bottom belong to the sentence. Keys and values are indexed by position, they sit on the axis the columns run along, and they are what the invariant is about. The left and the right belong to this step. Queries and outputs are indexed by the position doing the asking, they sit on the axis the rows run along, and they are different every time.

So keep the horizontal pair, and compute the vertical pair. That store is the KV cache. It is not a clever data structure. It is a pair of tensors per layer that get one column longer each step.

The picture is the whole idea. The horizontal edges stay long. The vertical edges shrink to a point. The square between them, the thing that was growing quadratically, becomes one row.

In code

@torch.no_grad()
def generate(model, ids, max_new_tokens):
    out = model(ids, use_cache=True)          # the prompt, once
    cache = out.past_key_values
    token = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
    produced = [token]

    for _ in range(max_new_tokens - 1):
        out = model(token, past_key_values=cache, use_cache=True)   # one token in
        cache = out.past_key_values
        token = out.logits[:, -1, :].argmax(dim=-1, keepdim=True)
        produced.append(token)

    return torch.cat([ids] + produced, dim=-1)

The line that matters is model(token, ...). Not ids. One token goes in, however long the sentence has become. It gets embedded, produces one query, one key and one value per layer, appends its key and value to the store, scores its query against every key already there, and pulls one blended value back out. The sentence is present in the computation only as the cache the new token reads from.

The output is the same. Not approximately, and not usually.

n = reference.shape[-1]
assert torch.equal(mine[:, :n], reference)

Every arithmetic operation that produced the bottom row is still being performed. The operations that produced the rows above it are the only thing that went away, and they were producing values we already had.

What it costs instead

Nothing is free, and the price here is memory.

The cache holds a key and a value for every position, in every layer, for every head. It starts when the request starts, it grows by one position per generated token, and it cannot be released until the request finishes. Two requests need two caches. A long conversation needs a long one.

That is a different kind of cost from the one just removed. Recomputation costs time, and time is something you can spend more of. The cache costs space on a device that has a fixed amount of it, and when that runs out it does not get slower. It fails.

Almost everything ahead in this series is a consequence of that sentence.

What this does not do yet

• No numbers. The counting above is arithmetic. Nothing here has been timed, and the memory has not been measured either. Both are next.
• No limits. This cache grows until the loop stops. Nothing here says what happens when it does not fit, because nothing here knows how big it is.
• One request. Still a single sequence, still one cache, still nobody else waiting.
• Nothing about position. The invariant quietly says “and the place it sits”, and that phrase has not been paid for yet.

There is also something odd sitting in plain sight in the code above.

The first call and the loop calls are not the same kind of work. The first takes the whole prompt and fills the square in one go. Every call after it takes a single token and fills one row. Same model, same weights, same code path, and two completely different shapes of computation.

Nothing so far has asked whether a GPU feels the same way about both.