The KV cache

This is where GPU fundamentals stop being an abstract exercise and start explaining why inference engines look the way they do. Autoregressive sampling generates one token at a time, and at every step, self-attention needs the key and value vectors for every token already in the sequence, not just the new one. Recomputing them from scratch each step would mean redoing the same matrix multiplies over and over as the sequence grows. Instead, engines store those vectors in a KV cache the first time they are computed, and read them back on every later step.

The size of that cache follows from the shape of the thing being stored. Per token the engine keeps two vectors, one key and one value, for every layer, and each is an n_heads by d_head matrix, so at 16-bit precision the storage per token is 2 * 2 * n_layers * n_heads * d_head bytes: two for the pair of vectors, two for the bytes per number. Nothing in that expression mentions the prompt, which is the point. It is a fixed cost per token of context, paid again for every sequence in the batch, and it grows without bound as generation continues.

What the cache buys is a change in the shape of the work. Computing the key and value vectors for one token costs exactly one sixth of passing that token through the whole model, so the saving at each sampling step is a sixth of the step multiplied by the number of tokens already in the sequence, and that multiplier grows with every token emitted. Without the cache, sampling would be quadratic in time complexity as the sequence gets longer. With it, each step re-reads state that has already been computed instead of rebuilding it.

The trade is not unconditional, and the boundary from the previous section is what decides it. A decode step reads every weight in the model to produce a single token, so it sits far on the memory-bound side, and time spent moving bytes is the budget that matters. When batches are small enough that the step is memory bound rather than flop bound, kipply’s arithmetic points the other way: an engine may prefer to recompute the vectors and spend the flops, since it is already paying the memory cost of sampling. Caching wins when the arithmetic it removes is worth more than the bytes it adds, which is why the rest of this book keeps returning to the size of this structure rather than treating it as free.