How LLMs Generate Text, and What the KV Cache Actually Fixes

Every chat model writes the same way: one token at a time, each new token fed back in to produce the next. The simplest version of that loop is quadratic. The standard fix, the KV cache, should make it hundreds of times faster.

I built GPT-2 inference from scratch in plain PyTorch (repo) and measured it on a MacBook CPU. The cache made generation 6x faster, not 260x. A GPU changes the absolute numbers. The reasons behind them carry over.

The argument runs in order: why generation is sequential, why the naive loop fails, and what the KV cache fixes and what it costs. Why the cached step is still slow is the subject of part 2.

Generation Is a Loop You Can’t Parallelize

A language model is a function. Token ids go in; for every position, a score (a logit) comes out for each of the 50,257 words in GPT-2’s vocabulary. To generate text, we pick a token from the last position’s scores, append it, and run the function again.

Two properties of this loop drive everything that follows.

It is sequential in time. Token n+1 cannot be computed until token n exists. Training sees a whole sentence at once and parallelizes across its positions. Generation never gets that luxury, because the thing it is producing is also its input.

Every step needs the whole model. All 12 layers and all 124M weights, 498 MB in fp32, take part in every single step, however short the sequence is.

Sequential steps, each touching every weight. Keep those two facts in mind; the rest of this series is their consequences.

(Turning scores into a token is a separate, cheap step: temperature, top-k, top-p, then a random draw. It sits outside the model and does not change the performance story, so I leave it out here. The repo has it.)

The Naive Loop, and a Number That Doesn’t Add Up

The obvious implementation feeds the entire sequence back into the model at every step, and reads the scores at the last position.

The problem is visible in the arithmetic. With a prompt of P tokens and N new tokens, step k processes P + k tokens:

naive total tokens processed = P·N + N(N+1)/2
with a cache                 = P + N

For a 10-token prompt and 512 new tokens, that is about 136,000 tokens of work against 522. Nearly all of it is redundant: the model recomputes what it computed one step earlier.

So the cache should make generation roughly 260x faster. Here is what I measured, in tokens per second:

New tokens	16	32	64	128	256	512
Naive	30.7	29.2	25.0	19.4	13.5	8.0
Cached	52.2	54.9	48.8	50.9	51.8	48.2
Speedup	1.7x	1.9x	1.9x	2.6x	3.9x	6.1x

The shape is right: cached throughput stays flat, naive collapses as sequences grow. But the size is wrong. The cache removed 99.6% of the work and only 84% of the time. At 512 tokens the speedup is 6x, not 260x.

The explanation is that the work the cache removed was the efficient work. Re-feeding hundreds of tokens at once is one big matrix-matrix multiply, which a CPU runs well. A cached step processes a single token, and the chip spends most of that step idle. The naive loop’s last step takes about 240 ms to process roughly 520 tokens. The cached step takes about 20 ms to process one.

Hold on to that. Doing 500x less work and getting 12x less time means the cheap step is expensive per unit of work. Part 2 is about why. First, what the cache actually stores.

The KV Cache: What to Remember and Why

In attention, every token produces a query (Q), a key (K) and a value (V). The output for token i is:

out_i = softmax(q_i · K[0..i]ᵀ / √d) · V[0..i]

Look at what the formula needs from the past: the keys and values of earlier tokens. Not their queries. And because of the causal mask, K_j and V_j depend only on tokens up to j, which never change once generated. Compute them once, keep them forever. That stored set is the KV cache.

Each decode step then does three things: compute q, k, v for the one new token, append k and v to the cache, and attend over the whole cache. The first pass over the prompt, which fills the cache in one parallel sweep, is called prefill. Every step after it is decode.

Why K and V, and not something else? Old queries were only ever used to compute old outputs, which are finished, so there is nothing to reuse. You could instead cache each layer’s input vectors, which is half the memory. But then every step has to redo the K and V projections for the whole history, which is a smaller version of the naive loop. The KV cache is the extreme end of a trade: spend memory to avoid recomputation.

Where it breaks. Two bugs are easy to write and silent. The new token’s position id must start at the cache length; start it at 0 and the model believes the token sits at the beginning of the text. And the causal mask is no longer square: with T new tokens and S total, it is tril(ones(T, S), diagonal=S-T). A test checks that cached greedy generation matches the naive loop token for token.

The bill. The cache costs memory, and the formula is worth memorizing:

bytes per token = 2 (K and V) × layers × kv_heads × head_dim × bytes per value

For GPT-2 in fp32 that is 2 × 12 × 12 × 64 × 4 = 72 KB per token, which is nothing (0.3 GB at 4k tokens). For a Llama-2-7B-class model in fp16 it is 512 KB per token, or 2.15 GB for a single 4k-token request. Part 2 shows how that number ends up capping how many users one GPU can serve.

Limits of These Measurements

One CPU, fp32, GPT-2 small with full multi-head attention. Runs vary by a few percent: the 512-token speedup came out at 6.6x in one run and 6.1x in another. A GPU will shift every absolute number.

What the Cache Leaves Open

The naive loop was quadratic: every step redid the work of every step before it. The KV cache fixed that. Generation now costs P + N tokens of work instead of P·N + N(N+1)/2, and throughput stays flat however long the output grows.

But it exposed a stranger problem. A cached step processes one token, needs about a third of a millisecond of math, and still takes 20 ms. Removing the redundant work didn’t make each token cheap. It made it obvious that something other than math sets the cost.

Part 2 follows that 20 ms: why a faster chip doesn’t help, what does, and what that fix costs in turn.

Code, tests and benchmarks for every step: github.com/sanjay23singh/llm-inference-from-scratch.