Generation and the KV Cache

Dive into Deep Learning · §11.3

Generation and the KV cache
stop recomputing the past · the memory bill · GQA, MLA, windows and sinks

Generation recomputes everything

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.

First duty of an optimization: change nothing

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   9.84 ms/token, cached  9.63 ms/token
n= 1024: naive  10.44 ms/token, cached  9.81 ms/token
n= 2048: naive  21.71 ms/token, cached  9.55 ms/token
n= 4096: naive  48.61 ms/token, cached  9.62 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: 11.5s =  22.3 tokens/s
cached: 256 tokens after a 3584-token prompt:  2.9s =  87.1 tokens/s

The memory bill

\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

PyTorch trap found by measuring: caching a view of the fused QKV buffer pinned 2x the formula; contiguous() closed the gap.

Prefill is compute-bound, decode is memory-bound

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

prefill: 2048 tokens in   22 ms =  92671 tokens/s
decode:  one token in   9.92 ms =    101 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. a hundred tokens/s.
Decode speed = bytes moved per token. Every cache byte shaved is speed.

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

Cache against quality

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; the loss column moves within seed noise.
At scale: MQA costs a sliver, GQA at H/8 matches (Ainslie et al., 2023).

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 (6x smaller): loss unchanged on this passage; rank 128 degrades gently. The premise holds — MLA makes the projection trained.

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 ...

GPT-2 parks a third to a half of its attention mass on token 0 — softmax must sum to one; idle heads need a dumping ground (StreamingLLM, 2024).

A window needs a sink

Bound the cache to a 256-token window — but evicting token 0 evicts every head’s dumping ground:

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

Window alone: perplexity 36 → thousands, with the full window intact.
One retained sink token recovers almost everything; gpt-oss ships a learned sink logit per head.

The cache-relief map

Attack the factors of 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

Causality freezes the past: cache K/V — the dense work goes quadratic → linear (cache reads stay cumulatively quadratic, subdominant), logits unchanged.
Cache bytes = 2 L\, n_\textrm{kv} d_\textrm{head}\, n\, b — verified against the allocator; it moves every step.
Prefill compute-bound, decode memory-bound: cache bytes are the currency of generation speed.
GQA: 8x smaller cache, quality flat at our scale and at theirs.
MLA: compress width; windows: bound length — but keep the sink.