Generation and the KV Cache

Dive into Deep Learning · §11.3

Generation and the KV cache
reusing past states · cache memory · GQA, MLA, windows, and sinks

Naive Autoregressive Generation Recomputes the Prefix

The naive generate: one full forward pass per token.

Step t costs \approx 2Nt FLOPs; T tokens cost \approx NT^2 — quadratic, and all but the last logit row discarded.
But causality freezes the past: \mathbf{k}_{1..t-1}, \mathbf{v}_{1..t-1} are identical at every future step. Store them.

The cached attention step

Two details: RoPE needs the absolute position offset, and the step receives only the new tokens plus the cache.

def rope(x, offset=0):
    """Rotary rotation of x at absolute positions offset, offset+1, ..."""
    d = x.shape[-1]
    pos = offset + torch.arange(x.shape[-2], dtype=torch.float32,
                                device=x.device)
    inv_freq = 10000.0 ** (-torch.arange(0, d, 2, device=x.device) / d)
    theta = pos[:, None] * inv_freq[None, :]
    cos, sin = torch.cos(theta), torch.sin(theta)
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack([x1 * cos - x2 * sin,
                        x1 * sin + x2 * cos], -1).flatten(-2)

def forward_step(self, X, cache):
    """Attention for the new tokens X against cached keys and values."""
    B, T, D = X.shape
    assert not (cache and T > 1), \
        'chunked append to a non-empty cache is unsupported ' \
        '(is_causal would misalign the mask)'
    q, k, v = self.W_qkv(X).chunk(3, -1)
    q, k, v = (u.reshape(B, T, self.num_heads, -1).transpose(1, 2)
               for u in (q, k, v))
    offset = cache['k'].shape[2] if cache else 0
    if self.rope:
        q, k = rope(q, offset), rope(k, offset)
    if cache:
        k = torch.cat([cache['k'], k], dim=2)
        v = torch.cat([cache['v'], v], dim=2)
    # contiguous(): at prefill k and v are still views into the fused QKV
    # projection, and caching a view would pin the whole 3x-wide buffer
    cache['k'], cache['v'] = k.contiguous(), v.contiguous()
    # Prefill (T > 1, empty cache) needs the causal mask; a single decoded
    # token (T = 1) attends to the entire cache
    Y = F.scaled_dot_product_attention(q, k, v, is_causal=(T > 1))
    return self.W_o(Y.transpose(1, 2).reshape(B, T, D))

Prefill once, then decode one token at a time

@torch.no_grad()
def forward_cached(self, X, caches):
    """Forward for new tokens only, extending one cache per block."""
    offset = caches[0]['k'].shape[2] if caches[0] else 0
    H = self.token_emb(X)
    if self.pos == 'learned':
        H = H + self.pos_emb(torch.arange(offset, offset + X.shape[1],
                                          device=X.device))
    for blk, cache in zip(self.blks, caches):  # pre-norm arrangement
        H = H + blk.attention.forward_step(blk.norm1(H), cache)
        H = H + blk.ffn(blk.norm2(H))
    return F.linear(self.norm(H), self.token_emb.weight)

@torch.no_grad()
def generate_cached(self, prefix, num_tokens, temperature=1.0, top_k=None):
    """Sample as generate does, but never recompute the prefix."""
    device = next(self.parameters()).device
    caches = [{} for _ in self.blks]
    ids = list(prefix)
    X = torch.tensor(prefix, device=device)[None]
    for _ in range(num_tokens):
        logits = self.forward_cached(X, caches)[0, -1] / temperature
        if top_k is not None:
            cutoff = torch.topk(logits, top_k).values[-1]
            logits[logits < cutoff] = -torch.inf
        ids.append(int(torch.multinomial(F.softmax(logits, -1), 1)))
        X = torch.tensor([ids[-1]], device=device)[None]
    return ids

JAX: growing shapes would recompile every step — preallocate (layers, batch, max_len, heads, head_dim) and write slots with dynamic_update_slice. Static shapes are how real serving works.

Cached and Uncached Logits Agree

rope: max |full - cached| = 3.65e-07
learned: max |full - cached| = 3.32e-07

Prefill 16, decode 32, stack the logits against one full forward pass: agreement to floating-point rounding, both positional schemes.

Per-step latency vs. context

124M-parameter model, untrained (arithmetic does not care):

123.6M parameters
n=  512: naive  10.12 ms/token, cached  9.77 ms/token
n= 1024: naive  15.22 ms/token, cached  9.83 ms/token
n= 2048: naive  33.74 ms/token, cached  9.79 ms/token
n= 4096: naive  77.45 ms/token, cached 11.91 ms/token

Reading the curve

Naive grows with context (2Nn); cached stays flat.
About 5\times at a 4k context in our runs, doubling with every further doubling of context.
Curves merge at short context: a 12-block step is launch-bound there — the cache removes a term that grows, not a constant factor.

naive : 256 tokens after a 3584-token prompt: 19.1s =  13.4 tokens/s
cached: 256 tokens after a 3584-token prompt:  4.3s =  59.5 tokens/s

KV-cache memory

\textrm{cache bytes} = 2 \cdot n_\textrm{layers} \cdot n_\textrm{kv} \cdot d_\textrm{head} \cdot n \cdot b \cdot (\textrm{bytes/elem})

72 KiB per token here: at 4k context, 288 MiB — more than half the model.

n=1024: allocator grew   72.0 MiB, formula says   72.0 MiB
n=2048: allocator grew  144.0 MiB, formula says  144.0 MiB
n=4096: allocator grew  288.0 MiB, formula says  288.0 MiB

In PyTorch, caching a view of the fused QKV buffer retained twice the storage predicted by the formula; contiguous() removed the excess allocation.

Prefill is compute-bound, decode is memory-bound

Every decoded token reads all weights + the whole cache for 2N FLOPs:

prefill: 2048 tokens in   37 ms =  54992 tokens/s
decode:  one token in  13.93 ms =     72 tokens/s
decode intensity: 247 MFLOP / 645 MB = 0.4 FLOP/byte (GPU ridge is near 80)
prefill intensity: about 784 FLOP/byte
bandwidth ceiling for decode: about 1550 tokens/s

Same layers, same GPU: tens of thousands vs. under a hundred tokens/s.
Decode speed depends on bytes moved per token, so reducing the cache can increase throughput.

A pluggable GQAAttention

W_k, W_v produce H_{kv} heads; the fused kernel broadcasts per group. Drops into TransformerBlock’s attn_factory and into d2l.GPT; H_{kv} = H recovers multi-head attention exactly.

H_kv=8: 262144 parameters, cache 2048 bytes per token per layer
H_kv=4: 196608 parameters, cache 1024 bytes per token per layer
H_kv=2: 163840 parameters, cache  512 bytes per token per layer
H_kv=1: 147456 parameters, cache  256 bytes per token per layer

GQA Cache Size and Validation Loss

Same GPT config, H_{kv} \in \{8, 4, 2, 1\}, best validation loss over a fixed budget:

H_kv=8: best validation loss 1.53, cache 12.0 KiB per token
H_kv=4: best validation loss 1.53, cache  6.0 KiB per token
H_kv=2: best validation loss 1.50, cache  3.0 KiB per token
H_kv=1: best validation loss 1.51, cache  1.5 KiB per token

Cache shrinks 8\times; with one seed and values rounded to 0.01 nat, the experiment supplies no uncertainty estimate or equivalence test.
Ainslie et al. (2023) provide separate large-model ablations of MQA and GQA.

Low-rank keys and values — the MLA idea

Keep all heads, store a low-rank latent instead (DeepSeek-V2/V3). Oracle check on GPT-2: truncate every layer’s realized K/V to rank r of 1536.

rank 256 of 1536 (6x smaller cache): loss 3.49, perplexity 33
rank 128 of 1536 (12x smaller cache): loss 3.66, perplexity 39

Rank 256 changes loss by less than displayed precision on this passage; rank 128 increases it. This oracle, passage-specific SVD is not MLA; MLA learns a shared projection during training.

Where trained attention actually goes

mean attention weight on token 0 (uniform would be ~0.003):
layer 0: 0.00  layer 1: 0.01  layer 2: 0.11  layer 3: 0.30  layer 4: 0.32
    layer 5: 0.51  layer 6: 0.46  layer 7: 0.57  layer 8: 0.48  layer 9:
    0.48  layer 10: 0.45  layer 11: 0.34

On the displayed passage, several GPT-2 layers assign a third to a half of their attention mass to token 0. StreamingLLM (2024) calls this pattern an attention sink and tests its role more broadly.

A window needs a sink

Compare a 256-token window with variants that retain initial sink positions:

window 256 + 0 sink tokens: loss 8.25, perplexity 3836
window 256 + 1 sink tokens: loss 3.86, perplexity 47
window 256 + 4 sink tokens: loss 3.76, perplexity 43
full attention:            loss 3.58, perplexity 36

On this passage, the window-only mask raises perplexity from 36 to thousands.
One retained initial token removes most of the increase; gpt-oss ships a learned sink logit per head.

Cache-reduction methods

The methods address different factors in the cache formula:

Width: GQA (fewer KV heads), MLA (low-rank latent).
Length: sliding window + sinks.
Remove it: linear attention → fixed recurrent state (§10.5); hybrids interleave a few full-attention layers (ch. 12).
Decode bandwidth arithmetic also powers speculative decoding — draft cheap, verify at prefill price (→ Computational Performance).

Recap

Because causality leaves past states unchanged, caching their keys and values reduces cumulative dense-layer work from quadratic to linear in generated length without changing logits. Cache reads remain cumulatively quadratic but are subdominant here.
Cache bytes = 2 L\, n_\textrm{kv} d_\textrm{head}\, n\, b — verified against the allocator; it moves every step.
Prefill is compute-bound, while decode speed is often limited by the number of cache bytes transferred per token.
GQA: 8x smaller cache in the fixed-seed experiment; quality equivalence is not established.
MLA compresses width; windows bound length while retaining designated sink positions.