Here is the generation loop from §6.4, stripped to its core. Run the model on everything so far, take the last position's logits, pick a token, append it, repeat. It's correct, and it is wildly wasteful. To produce token 1,000 it runs the model on 999 tokens, and to produce token 1,001 it runs it on 1,000, recomputing 999 rows it computed one step earlier.

Here's the puzzle. Of everything the model computes for the 1,000-token prefix, all but one row is identical to what it computed for the 999-token prefix. Which row is new, and what is the least we need to remember from earlier steps to compute it? The answer is the KV cache, the single most important optimization in LLM inference. Every serving system you'll ever use is organized around it, and its size is what limits how many users one GPU can serve at once.

Why naive generation is quadratic

Let $N$ be the number of parameters that every token is multiplied through (the blocks' matrices plus the output head, but not the embedding lookup, which is free), $L$ the number of layers, and $d$ the width. A forward pass over $t$ tokens costs about $2Nt$ floating-point operations for the matrix multiplies (one multiply and one add per parameter per token, as in §5.6), plus attention, where each of the $t$ queries is compared with up to $t$ keys and mixes up to $t$ values. With a causal mask that's about $2Ldt^2$ more.

The naive loop runs that forward pass at every step. Generating $n$ tokens from a short prompt therefore costs

In plain words, doubling the length of the output quadruples the matmul work and multiplies the attention work by eight. Generating a 1,000-token answer redoes, in total, about five hundred times the work of reading a 1,000-token prompt once.

What changes when you append a token

Here's the observation that fixes it. In a causal model, the hidden state at position $j$ depends only on tokens $0, \dots, j$. That's true at the input, where $\mathbf{x}^0_j = \mathbf{E}_{w_j,:} + \mathbf{P}_{j,:}$. And if it's true at layer $\ell$, it stays true at layer $\ell + 1$. The MLP and LayerNorm act on each position separately, and the attention at position $j$ reads only keys and values from positions $\le j$, because the mask forbids anything later. By induction over the layers, appending token $t$ changes nothing at positions $0, \dots, t-1$, at any layer.

So at step $t$ the only new row is position $t$'s, and to compute it, each layer needs exactly three things.

• the new token's own vector, which comes from the layer below,
• its query $\mathbf{q}_t$, computed from that vector,
• the keys and values of all positions $0, \dots, t$ at that layer.

Old queries are never used again, since past positions are finished. Old MLP outputs are never used again either. Only the keys and values of past positions are read by future tokens. Store them, one $\mathbf{K}$ and one $\mathbf{V}$ per layer, and each step needs only the new token's computation,

where the subscript $i$ selects the $d_k = d/h$ columns that belong to head $i$, one of the $h$ heads. No mask is needed in decoding, since every cached position is in the past. In plain words, each layer appends one row to its key table and one to its value table, and the new query reads both tables in full.

The cost per generated token drops to about $2N + 4Ldt$ (every weight used once, plus one query against $t$ keys and $t$ values in every layer), so

The matmul work is now linear in the length. Attention is still quadratic in total, because the $t$-th token really does have to look at $t$ earlier tokens, but its per-token cost is linear and small. Let's put numbers on it for GPT-2 small.

The ratio grows roughly like $n/2$, about 500 times less work at the full 1,024-token context.

Building it

We need the §6.2 model again, written for one sequence at a time (no batch axis) so that the shapes stay readable. gpt_forward is the reference with no cache. forward_cached runs any number of new tokens. It looks up their positions starting where the cache ends, appends their keys and values, and lets each new token attend to the whole cache plus the new tokens up to itself.

The mask line deserves a second look. With $t_0$ tokens cached and $n$ new ones, the score matrix is $n \times (t_0 + n)$, not square. New token $i$ sits at absolute position $t_0 + i$, so it may see columns $0$ to $t_0 + i$. That's a lower triangle shifted right by $t_0$, which is np.tril(..., k=t0). With $n = 1$ the mask is all True, and with $t_0 = 0$ it's the ordinary causal mask.

Generation now has two phases. Prefill runs the whole prompt through forward_cached in one pass, filling the cache and producing the first new token's logits. Decode then feeds one token at a time. Let's check that the cached generator produces exactly the same tokens as the naive one, and time both.

The tokens are identical, and the logits agree to rounding error, about $10^{-14}$. (They aren't always bit-identical, because a matrix-vector product and a matrix-matrix product may add the same numbers in a different order.) Even for this tiny model the cache is many times faster, and the gap grows with length, as the widget below shows by timing the real tiny GPT in your browser.

import numpy as np

def decode_step(P, token, cache, h):
    """one token in, (V,) logits out; appends this token's K and V rows to every layer of cache"""
    raise NotImplementedError

def generate_cached(P, prompt, n_new, h):
    raise NotImplementedError

import numpy as np

def prefill_chunk(P, ids, cache, h):
    """(n,) new ids -> (n, V) logits; extends every layer of cache by n rows"""
    raise NotImplementedError

How big is the cache?

The cache holds one key vector and one value vector per token, per layer, per key-value head. With $n_{kv}$ key-value heads of size $d_k$ (for GPT-2, $n_{kv} = h$ and $n_{kv} d_k = d$), a context of $T$ tokens, a batch of $B$ sequences, and $s$ bytes per number,

In plain words, two tables (keys and values), times layers, times the width of a key, times tokens, times sequences, times bytes per number. The leading 2 and the use of $n_{kv}$ rather than $h$ are the two places people slip.

Work three examples in 16-bit precision ($s = 2$).

• GPT-2 small has $L = 12$ and $h d_k = 768$, so each token costs $2 \cdot 12 \cdot 768 \cdot 2 = 36{,}864$ bytes, 36 KiB, and a full 1,024-token sequence 36 MiB.
• LLaMA-2 7B has $L = 32$, $n_{kv} = 32$, $d_k = 128$, so each token costs $2 \cdot 32 \cdot 32 \cdot 128 \cdot 2 = 524{,}288$ bytes, exactly 0.5 MiB. Its 4,096-token context is 2 GiB per sequence. A batch of 16 such sequences needs 32 GiB of cache, more than twice the 13.5 GB of weights.
• LLaMA-2 70B uses grouped-query attention with only $n_{kv} = 8$ key-value heads shared by its 64 query heads (Ainslie et al., 2023). With $L = 80$, a token costs $2 \cdot 80 \cdot 8 \cdot 128 \cdot 2 = 327{,}680$ bytes, 320 KiB, less than the 7B model despite ten times the parameters. Without GQA it would be 2.5 MiB.

When does the cache outgrow the weights? Take multi-head attention ($n_{kv} d_k = d$), with weights and cache stored at the same precision. The weights are about $12Ld^2$ numbers (§5.6), and the cache for $B$ sequences of $T$ tokens is $2LdTB$, so

In plain words, the cache matches the weights once the batch holds about $6d$ tokens in total, whatever the depth. For LLaMA-2 7B that is $6 \times 4096 = 24{,}576$ tokens, six sequences at its 4,096-token context, which agrees with the 2 GiB per sequence against 13.5 GB of weights above. Grouped-query attention multiplies the ratio by $n_{kv}/h$, so a model with 8 key-value heads for 64 query heads holds eight times as many tokens before its cache catches up with its weights.

So at long context the cache, not the weights, is what limits how many users a GPU can serve. That single fact drives a whole family of techniques. Sharing keys and values across heads (multi-query and grouped-query attention) and compressing them into a latent vector (multi-head latent attention) are derived in §12.5. Storing the cache in 8 or 4 bits and splitting it into fixed-size pages that are allocated on demand, so that sequences of different lengths don't waste memory, are the domain of serving systems built around paged attention and quantized KV caches.

LLaMA-2 70B's weights alone exceed one 80 GB GPU, so it shows a batch of zero. It has to be split across GPUs (§10.4), and the cache is split with it. Explore other configurations below.

def kv_cache_bytes(n_layers, n_kv_heads, head_dim, seq_len, batch, bytes_per_elem=2):
    raise NotImplementedError

def max_batch(gpu_bytes, weight_bytes, n_layers, n_kv_heads, head_dim, seq_len, bytes_per_elem=2):
    raise NotImplementedError

Prefill versus decode

The two phases of generation stress the hardware in opposite ways, and understanding why is the starting point for everything in inference engineering.

Prefill processes all $p$ prompt tokens at once. Every weight matrix is multiplied by a $p \times d$ block of activations, so each weight read from memory is used $p$ times. The work is $2Np$ FLOPs for about $2N$ bytes of weights in 16-bit, an arithmetic intensity of roughly $p$ FLOPs per byte. For a prompt of hundreds of tokens that is far above what the GPU needs to stay busy (§11.2), so prefill is compute-bound. Its duration sets the time to first token.

Decode processes one token per sequence. Every weight is read from memory to be used once per sequence in the batch, so at batch size 1 the intensity is about 1 FLOP per byte, and the step also reads the entire KV cache. Decode is memory-bound. The time per token is about the bytes moved divided by the memory bandwidth, whatever the GPU's FLOPS.

Put numbers on a 7B model in bfloat16, about 13.5 GB of weights, on a GPU with roughly 2 TB/s of memory bandwidth and 312 TFLOPS of bfloat16 compute (an A100 80GB). One decode step at batch 1 must read the weights at least once, $13.5\,\text{GB} / 2\,\text{TB/s} \approx 7$ ms, so at most about 150 tokens per second, while the arithmetic would take $2 \times 6.7 \times 10^9 / 3.12 \times 10^{14} \approx 0.04$ ms. Prefilling a 1,000-token prompt needs $1.3 \times 10^{13}$ FLOPs, about 43 ms at full utilization, the price of six or seven decode steps for a thousand tokens of work.

The consequence is that decode is almost free to batch. Serving 32 sequences at once reads the same weights once per step and does 32 times the arithmetic, so throughput rises nearly 32-fold until the cache reads or the compute catch up. With the cache counted, one decode step takes about $(M_{\text{weights}} + M_{\text{KV}})$ divided by the memory bandwidth. For the 7B model at batch 16 with full 4,096-token contexts that is 13.5 GB of weights plus 34 GB of cache, about 24 ms per step, so at large batch and long context the cache, not the weights, sets the speed, a second reason to shrink it. Inference servers therefore batch decode steps from many users together, and treat the prefill of a long prompt as a separate, compute-heavy job, which they often split into chunks (exactly Exercise 6.5b) so it doesn't stall everyone else's next token. Serving systems track latency as two numbers, time to first token (set by prefill) and time per output token (set by decode), and schedule the two phases separately to keep both low, which is the core of fast inference engineering.

In PyTorch

Production caches are preallocated, one tensor of shape $(L, B, n_{kv}, T_{\max}, d_k)$ for keys and one for values, written in place. Concatenating as our NumPy code does copies the whole cache at every step, which turns linear memory traffic quadratic and fragments GPU memory. A static shape also lets the decode step be captured once as a CUDA graph (§11.7).

One line in that block hides a famous bug. For the cached path we build the mask by hand instead of passing is_causal=True. When the query length differs from the key length, PyTorch's is_causal aligns the triangle to the top-left corner, so a single decode query would be allowed to see only the first key. The correct mask for a cache is aligned to the bottom-right, which is what tril(diagonal=t0) builds.

Batching adds one more wrinkle. Prompts in a batch have different lengths, but a preallocated cache wants every row to write at the same index. The simple fix is left padding. Pad the short prompts at the start, mask the pad positions out of every attention row, and give each real token its position counted from its own first real token, so that all sequences then decode in lockstep. Serving systems avoid the waste altogether with paged caches, where each sequence owns its own list of fixed-size blocks and the attention kernel gathers them (Kwon et al., 2023).

Read more

• Pope et al. (2022). Efficiently Scaling Transformer Inference. Prefill versus decode, memory bandwidth, and KV cache costs at scale.
• Shazeer (2019). Fast Transformer Decoding: One Write-Head is All You Need. Multi-query attention, motivated by the cache.
• Ainslie et al. (2023). GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. Grouped-query attention, the middle ground used by LLaMA-2 70B and most models since.
• Kwon et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. The vLLM paper, paging the cache like virtual memory.
• Touvron et al. (2023). Llama 2: Open Foundation and Fine-Tuned Chat Models. The configurations used for the cache sizes above.
• Karpathy. nanoGPT model.py. Its generate is the naive loop, a good exercise to retrofit with this section's cache.

Watch next

▶

▶