What Happens at Generation Time: KV Cache, RoPE and Sampling

Training a language model is one big batched forward-and-backward pass. Generating from it is a loop that produces one token per iteration, and almost everything that makes inference engineering its own discipline follows from that one fact. This post walks through the loop, does the KV-cache memory arithmetic for a 7B-class model with explicit assumptions, verifies the key property of rotary position embeddings numerically, and implements the sampling knobs that sit at the end of every step.

Prefill processes the whole prompt in parallel and fills the cache; each decode step then processes exactly one new token, reading the entire cache and appending one entry per layer.

The decode loop

A decoder-only transformer defines $p (x_{t + 1} ∣ x_{\leq t})$ . Generation samples from it repeatedly:

Prefill. Run the prompt tokens $x_{1}, \dots, x_{P}$ through the model in one batched pass. For every layer, keep the key and value vectors of every position; this is the KV cache. Take the logits at the last position.
Sample a token from those logits (temperature, top-k, top-p, or greedy).
Decode step. Feed only the new token. In each layer, compute its query, key and value; append the key and value to the cache; attend the single query over all cached keys and values; continue to the next layer.
Repeat from step 2 until an end-of-sequence token or a length limit.

Without the cache, step 3 would have to recompute keys and values for all previous tokens, making the cost of generating $n$ tokens $O (n^{2})$ in model forward passes instead of $O (n)$ . With it, each step is one token's worth of matrix-vector products plus one attention over $t$ cached positions.

The two phases have opposite performance characters. Prefill multiplies a $P \times d$ matrix by the weights: lots of arithmetic per byte of weights loaded, so it is compute-bound. A decode step multiplies a $1 \times d$ vector by the same weights: every weight byte is read from memory to do two floating-point operations, so it is bound by memory bandwidth, and so is reading the cache.

TIP

A quick lower bound on single-sequence decode latency is bytes that must be read per step divided by memory bandwidth. A 7B model in 16-bit weights is about 13.5 GB; on an accelerator with 2 TB/s of memory bandwidth that is at least 6.7 ms per token, or at most about 150 tokens per second at batch size 1, before the cache is even counted. Batching more sequences reuses each weight read and is how serving systems recover throughput.

KV-cache memory arithmetic

Per token, the cache stores one key and one value vector per layer per key/value head:

bytes per token = 2 \times L \times n_{k v} \times d_{head} \times bytes per element,

where the leading 2 counts keys and values. Total cache memory is that number times context length times batch size. The table uses these assumptions: 32 layers, head dimension 128, and either 32 KV heads (multi-head attention, as in Llama-2-7B) or 8 KV heads (grouped-query attention, as in Mistral-7B and Llama-3-8B), with 16-bit or 8-bit cache elements.

Configuration (assumed)	n_kv	Bytes / token	4,096 tokens	32,768 tokens	Batch 16 × 4,096
L = 32, d_head = 128, MHA, fp16	32	524,288 B = 0.5 MiB	2.0 GiB	16 GiB	32 GiB
L = 32, d_head = 128, MHA, int8 cache	32	262,144 B = 0.25 MiB	1.0 GiB	8 GiB	16 GiB
L = 32, d_head = 128, GQA (8), fp16	8	131,072 B = 128 KiB	0.5 GiB	4 GiB	8 GiB
L = 32, d_head = 128, GQA (8), int8	8	65,536 B = 64 KiB	0.25 GiB	2 GiB	4 GiB

The worked number for the first row: $2 \times 32 \times 32 \times 128 \times 2 = 524,288$ bytes, exactly half a mebibyte per token. At a 32k context that is 16 GiB for a single sequence, more than the 13.5 GB of weights. This is why grouped-query attention, which shares each key/value head among several query heads, became standard in the 7B class: it divides the cache by the head ratio (here 4×) at a small cost in quality.

KV cache size versus context length for MHA and GQA configurations — KV-cache size for a single sequence as a function of context length, with reference lines for a 7B model's fp16 weights and two common accelerator memory sizes. The cache crosses the weight size around 25.7k tokens for the MHA configuration.

Memory is only half of the problem: the cache is allocated per sequence and grows unpredictably, which fragments accelerator memory. Paged attention, which stores the cache in fixed-size blocks indexed through a page table, is the systems answer and is the core idea in vLLM.

Rotary position embeddings

Attention has no notion of order, so position must be injected. Rotary position embeddings (RoPE) do it inside the attention scores. Split the $d$ -dimensional query into $d / 2$ pairs and rotate pair $i$ by the angle $m θ_{i}$ , where $m$ is the token's position and $θ_{i} = 10000^{- 2 i / d}$ is a per-pair frequency:

R_{m} q = (q^{(i)} \cos m θ_{i} - q^{(i + d / 2)} \sin m θ_{i}, q^{(i)} \sin m θ_{i} + q^{(i + d / 2)} \cos m θ_{i})_{i < d / 2} .

The property that makes it work: rotating the query by $m θ$ and the key by $n θ$ and then taking the dot product depends only on $m - n$ ,

⟨ R_{m} q, R_{n} k ⟩ = ⟨ R_{m - n} q, k ⟩,

because a product of two rotations is a rotation by the difference of the angles. Relative position therefore enters every attention score, absolute position does not, and nothing is added to the residual stream.

def rope(x, pos, base=10000.0):
    """Rotate the pairs (x[i], x[i + d/2]) of each row by pos * theta_i."""
    d = x.shape[-1]
    half = d // 2
    theta = base ** (-np.arange(half) * 2.0 / d)              # theta_i = base^(-2i/d)
    ang = np.asarray(pos)[:, None] * theta[None, :]            # (T, d/2) angles
    c, s = np.cos(ang), np.sin(ang)
    x1, x2 = x[..., :half], x[..., half:]
    return np.concatenate([x1 * c - x2 * s, x1 * s + x2 * c], axis=-1)

q, k = rng.normal(size=(1, 64)), rng.normal(size=(1, 64))
score = lambda m, n: float(rope(q, [m])[0] @ rope(k, [n])[0])
print(max(abs(score(10 + s, 6 + s) - score(10, 6)) for s in range(0, 512, 8)))

The printed maximum deviation over shifts of up to 512 positions is $3.9 \times 10^{- 14}$ , floating-point noise. At offset zero the score equals the plain dot product $q \cdot k$ , as it should, because both rotations cancel.

RoPE score invariance to absolute position and the per-pair wavelengths — Left: for a fixed random query and key, the attention score is flat as the absolute position slides, for four relative offsets. Right: the wavelength of each rotating pair for d = 64 and base 10,000 spans from about 6 tokens to about 47,000, which is what lets the same mechanism encode both local order and long-range distance.

Why the dot product depends only on m − n

Each pair $(q^{(i)}, q^{(i + d / 2)})$ is a 2-D vector, and $R_{m}$ rotates it by $m θ_{i}$ . For 2-D rotations, $R_{a}^{⊤} R_{b} = R_{b - a}$ , so for every pair $⟨ R_{m} q, R_{n} k ⟩_{pair} = q_{pair}^{⊤} R_{n - m} k_{pair}$ , which depends only on $n - m$ . Summing over the $d / 2$ pairs gives the full dot product. The same argument shows the rotation does not change the norm of $q$ or $k$ , so it does not interfere with the $1 / \sqrt{d_{k}}$ scaling.

At decode time this matters for the cache: the cached keys are stored already rotated by their own positions, so a new query only needs its own rotation. Extending context beyond the training length is then a question of what to do with angles larger than any seen in training, which is what position interpolation and frequency-scaling schemes address.

Sampling

The last step of each iteration turns logits $z \in R^{V}$ into one token. Greedy decoding takes the argmax and is deterministic but prone to repetition. Sampling from $softmax (z)$ is faithful to the model but puts real probability mass on the long tail of unlikely tokens, and a single bad draw derails the continuation. The practical knobs reshape or truncate the distribution before sampling:

Divides the logits by $T$ before the softmax, $p_{i} \propto \exp (z_{i} / T)$ . $T < 1$ sharpens (the top token gained from 0.27 to 0.63 of the mass at $T = 0.5$ in the example below), $T > 1$ flattens, $T \to 0$ is greedy. It changes every probability but never removes a token.
```
def sample_next(logits, rng, temperature=1.0):
    z = logits / max(temperature, 1e-8)
    p = np.exp(z - z.max())
    p /= p.sum()
    return int(rng.choice(len(p), p=p))
```

Keeps the $k$ highest logits and renormalizes. Simple, but $k$ is a poor fit for both peaked distributions (where $k = 50$ includes junk) and flat ones (where it cuts off plausible tokens).

def sample_next(logits, rng, top_k=50):
    z = np.where(logits >= np.sort(logits)[-top_k], logits, -np.inf)
    p = np.exp(z - z.max())
    p /= p.sum()
    return int(rng.choice(len(p), p=p))

Keeps the smallest set of tokens whose cumulative probability reaches $p$ . The set adapts to the shape of the distribution: in the example, $p = 0.9$ kept 13 of 20 tokens.

def sample_next(logits, rng, top_p=0.9):
    p = np.exp(logits - logits.max())
    p /= p.sum()
    order = np.argsort(-p)
    cum = np.cumsum(p[order])
    keep = order[: int(np.searchsorted(cum, top_p)) + 1]   # smallest prefix reaching top_p
    mask = np.zeros_like(p, dtype=bool)
    mask[keep] = True
    p = np.where(mask, p, 0.0)
    p /= p.sum()
    return int(rng.choice(len(p), p=p))

A toy next-token distribution under temperature 0.5, temperature 1 and top-p 0.9 — A synthetic 20-token next-token distribution. Temperature 0.5 concentrates mass on the top token (entropy 1.50 nats versus 2.56 at T = 1); top-p 0.9 removes the seven least likely tokens and renormalizes (entropy 2.28 nats).

The knobs compose: temperature is applied first, then truncation, then the draw. Repetition penalties and the various "min-p" schemes are further edits to the same vector. All of them are cheap compared with the forward pass, so the choice is about output quality, not speed.

Where the time goes, in one sentence

Per generated token, a 7B-class model reads all of its weights and the whole KV cache once; everything else about inference engineering (batching, quantized weights and caches, grouped-query attention, paged memory, speculative decoding) is a way to read fewer bytes per useful token.

What Happens at Generation Time: KV Cache, RoPE and Sampling

The decode loop

TIP

KV-cache memory arithmetic

Rotary position embeddings

Sampling

Where the time goes, in one sentence

Further reading

Cite this post

Keep reading