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.
image/svg+xml Matplotlib v3.10.8, https://matplotlib.org/

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 + jnp.arange(x.shape[1])
    inv_freq = 10000.0 ** (-jnp.arange(0, d, 2) / d)
    theta = pos[:, None] * inv_freq[None, :]
    cos = jnp.cos(theta)[:, None, :]  # broadcast over heads
    sin = jnp.sin(theta)[:, None, :]
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return jnp.stack([x1 * cos - x2 * sin,
                      x1 * sin + x2 * cos], -1).reshape(x.shape)

def init_cache(model, batch_size, max_len):
    """Preallocated K and V buffers, one slab per layer."""
    attn = model.blks[0].attention
    head_dim = attn.W_o.in_features // attn.num_heads
    shape = (len(model.blks), batch_size, max_len, attn.num_heads, head_dim)
    return jnp.zeros(shape), jnp.zeros(shape)

Prefill once, then decode one token at a time

@nnx.jit
def cached_forward(model, ks, vs, X, t):
    """Logits for the last of the new tokens X at positions t, t+1, ...;
    writes their keys and values into the cache slots at those positions."""
    B, T = X.shape
    H = model.token_emb(X)
    if model.pos == 'learned':
        H = H + model.pos_emb(t + jnp.arange(T))
    mask = (jnp.arange(ks.shape[2])[None, :]
            <= (t + jnp.arange(T))[:, None])[None, None]
    for i, blk in enumerate(model.blks):  # pre-norm arrangement
        attn, Y = blk.attention, blk.norm1(H)
        q, k, v = jnp.split(attn.W_qkv(Y), 3, axis=-1)
        q, k, v = (u.reshape(B, T, attn.num_heads, -1) for u in (q, k, v))
        if attn.rope:
            q, k = rope(q, t), rope(k, t)
        # Callers must keep t + T <= max_len: an out-of-range write does
        # not fail here -- dynamic_update_slice clamps its start index
        # silently. (t is traced under jit, so generate_cached asserts.)
        ks = jax.lax.dynamic_update_slice(ks, k[None], (i, 0, t, 0, 0))
        vs = jax.lax.dynamic_update_slice(vs, v[None], (i, 0, t, 0, 0))
        Y = jax.nn.dot_product_attention(q, ks[i], vs[i], mask=mask)
        H = H + attn.W_o(Y.reshape(B, T, -1))
        H = H + blk.ffn(blk.norm2(H))
    return model.token_emb.attend(model.norm(H))[:, -1], ks, vs

def generate_cached(self, prefix, num_tokens, seed=0, temperature=1.0,
                    top_k=None):
    """Sample as generate does, but never recompute the prefix."""
    assert len(prefix) + num_tokens <= self.max_len, 'cache overflow'
    ks, vs = init_cache(self, 1, self.max_len)
    logits, ks, vs = cached_forward(self, ks, vs,
                                    jnp.asarray(prefix)[None], jnp.array(0))
    ids, key = list(prefix), jax.random.key(seed)
    for t in range(len(prefix), len(prefix) + num_tokens):
        l = logits[0] / temperature
        if top_k is not None:
            l = jnp.where(l < jnp.sort(l)[-top_k], -jnp.inf, l)
        key, sub = jax.random.split(key)
        ids.append(int(jax.random.categorical(sub, l)))
        logits, ks, vs = cached_forward(self, ks, vs,
                                        jnp.asarray([[ids[-1]]]),
                                        jnp.array(t))
    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| = 9.39e-07
learned: max |full - cached| = 8.34e-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  31.27 ms/token, cached 11.53 ms/token
n= 1024: naive  14.83 ms/token, cached 11.30 ms/token
n= 2048: naive  42.23 ms/token, cached 12.44 ms/token
n= 4096: naive  76.04 ms/token, cached 12.12 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: 24.5s =  10.5 tokens/s
cached: 256 tokens after a 3584-token prompt:  8.8s =  29.2 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.

cache buffers: 288 MiB, formula says 288 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   34 ms =  60455 tokens/s
decode:  one token in  14.44 ms =     69 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.

Sharing keys and values: MQA and GQA

Queries ask; keys and values are the library. Nothing forces n_\textrm{kv} = H.

image/svg+xml Matplotlib v3.10.8, https://matplotlib.org/
  • MQA (Shazeer, 2019): one KV head — cache \div H.
  • GQA (Ainslie et al., 2023): H_{kv} heads, one per query group — Llama 3 and Mistral run H/4; uptraining made adoption cheap.

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.55, cache 12.0 KiB per token
H_kv=4: best validation loss 1.51, cache  6.0 KiB per token
H_kv=2: best validation loss 1.51, cache  3.0 KiB per token
H_kv=1: best validation loss 1.54, 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 3842
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.