The KV Cache

The block is finished. Give it a sequence of tokens and it returns, for each position, the logits for what comes next. But look again at what attention did in Chapter 5: to score position five, it needed the keys and values of positions zero through four. Run the model one token at a time, the way generation actually works, and every step recomputes the keys and values of every token before it — the same keys and values, from the same weights, that the step before already computed. By the hundredth token the model is redoing the work of the first ninety-nine on every single step.

That work is pure waste, because those keys and values never change. Once a token is in the past, its vector is fixed, so the key and value it projects are fixed too. Compute them once, keep them, and every later step can read them instead of rebuilding them. That store is the KV cache, and it is the difference between generation that slows to a crawl as the conversation grows and generation that stays fast. This chapter is not a new piece of math — it is a piece of memory, and its correctness question is different from every chapter before it: not “does this formula match a reference,” but does the fast, cached path return the same logits as the plain, uncached one.

7.1 What the cache holds

The cache itself is almost nothing — a list with one slot per layer.

class KVCache:
    def __init__(self, n_layers: int) -> None:
        self._entries: list[KVEntry | None] = [None] * n_layers

    def get(self, layer: int) -> KVEntry | None:
        return self._entries[layer]

    def update(self, layer: int, entry: KVEntry) -> None:
        self._entries[layer] = entry

    def reset(self) -> None:
        self._entries = [None] * len(self._entries)

Listing 7.1 · forge’s KV cache. One stored entry per layer, read by get, written by update.

A KVEntry is a pair of tensors: the keys and the values for every token the layer has seen so far. The cache holds one entry per layer and does nothing clever with it — get returns a layer’s entry, update replaces it, reset empties them all. The growth happens a level up, inside attention. Recall from Chapter 5 that when a layer is given a past, it concatenates the past keys and values with the new token’s before it attends: k = torch.cat([past[0], k]). The tuple it hands back is that longer pair, and update writes it into the slot. So each step, the entry in every slot grows by exactly one token, and the slot always holds the complete history the next step will need.

This is the first structure in the engine whose size depends on the conversation, not the model. A layer’s entry is (keys, values) of shape roughly n_kv_heads × sequence_length × head_dim — and sequence_length climbs with every token generated. That is a memory cost that grows without bound as a chat goes on, and it is what Part II’s quantization exists to shrink. Here we only need it to be correct.

7.2 Prefill and decode

The cache changes nothing about the block’s math; it changes how the forward pass is called. Here is the part of Chapter 2’s forward pass that drives it.

        if cache is not None:
            pos_start = self.current_pos
            pos_end = pos_start + t
            self.current_pos = pos_end
        else:
            pos_start, pos_end = 0, t

        mask = causal_mask(t, pos_start, pos_end, in_idx.device)

        for layer, block in enumerate(self.trf_blocks):
            past = cache.get(layer) if cache is not None else None
            x, entry = block(x, mask, self.cos, self.sin, pos_start, past)
            if cache is not None:
                cache.update(layer, entry)

Listing 7.2 · the cache branch of the forward pass. Position bookkeeping, then per-layer read and write.

Everything turns on current_pos, the count of tokens the cache already holds. When a cache is passed, the forward pass reads it as the starting position, advances it by the number of new tokens t, and — crucially — hands pos_start to each block so rotary embeddings rotate the new tokens by their true positions in the sequence, not their offsets within this one call. Without a cache, current_pos is ignored and positions start at zero, exactly as they did in every chapter until now.

That single mechanism covers both phases of generation. Prefill is the first call: the whole prompt goes in at once, t is its length, the cache fills from empty, and causal_mask builds the triangular mask so each prompt token sees only the ones before it. Decode is every call after: t is 1, a single new token enters, its key and value are appended to each layer’s entry, and it attends over the entire cached history. A decode step needs no mask at all — there is only one query, at the newest position, and every cached key is in its past — which is why causal_mask returns None when t is 1. Prefill pays once for the prompt; decode pays for one token at a time, reusing everything already stored.

7.3 Same answer, less work

The cache is only worth having if it is invisible in the output: the logits from generating token by token with the cache must be the logits from one plain forward pass over the whole sequence. So the check is forge against itself — the same model, the same tokens, run both ways. The companion driver runs a six-token prompt through one uncached forward pass, then feeds the same tokens one at a time through a cache, and prints the largest disagreement between the two logit streams:

The gap is 3.58e-07 — rounding-sized, about twenty-eight times under the 1e-5 gate, and not zero for the same reason Chapter 5’s attention was not: the uncached pass scores all six positions in one batched matmul, while the cached path scores one position at a time, and floating-point addition does not reduce in exactly the same order across those two shapes. The last digits differ; the answer does not. Any real mistake in the cache — a token appended in the wrong place, a position miscounted, a write that never happens — would move the gap far past a rounding error, because the cached path would then be attending over the wrong history. The checkpoint makes exactly that happen.

The engine now generates the way Chapter 1 promised: prefill the prompt, then decode one token per step, each step cheap because the past is remembered rather than rebuilt. One black box from that first decode loop is still closed. The loop reads the logits at the final position and has to turn them into an actual token — greedily taking the most likely one, or sampling under a temperature, or restricting the draw to the top of the distribution. That choice is the sampler, and it is the last thing Part I builds.