12. KV-Cache for Autoregressive Decoding

Problem

Recomputing attention over the entire generated sequence at every decoding step is wasted work: the key/value projections of already-generated tokens are fixed once computed. A KV-cache stores those past keys and values and, at each new step, only computes attention between the new query and the concatenation of the new key/value with everything cached so far — turning each step's cost into an O(1) amortized append plus a single attention read over the growing cache, and is the single biggest lever for autoregressive inference throughput.

Given a new token's q, k_new, v_new (each shape (batch, heads, 1, head_dim)) and either None or a previous cache dict of (batch, heads, seq_so_far, head_dim) tensors, concatenate, attend, and return the updated cache.

Source: src/12_kv_cache_decoding.py

def attend_with_kv_cache(q, k_new, v_new, cache) -> tuple[torch.Tensor, dict[str, torch.Tensor]]

>>> import torch
>>> torch.manual_seed(0)
>>> q = torch.randn(1, 2, 1, 4)
>>> k = torch.randn(1, 2, 1, 4)
>>> v = torch.randn(1, 2, 1, 4)
>>> out, cache = attend_with_kv_cache(q, k, v, None)
>>> out.shape
torch.Size([1, 2, 1, 4])
>>> cache["k"].shape
torch.Size([1, 2, 1, 4])

Step-by-Step Approach

1. Handle the first step: if cache is None, there's nothing to concatenate onto yet, so k_full = k_new and v_full = v_new directly.
2. Otherwise, append the new key/value onto the cached ones with torch.cat([cache["k"], k_new], dim=2) along the sequence dimension (dim 2, since shapes are (batch, heads, seq, head_dim)) — and the same for values.
3. Compute attention scores between the single new query and the full key tensor: q @ k_full.transpose(-2, -1) / sqrt(head_dim), giving shape (batch, heads, 1, seq_so_far+1).
4. Softmax over the last (key/sequence) dimension and matrix-multiply against v_full to get the attention output, shape (batch, heads, 1, head_dim).
5. Return the output alongside a new cache dict holding the concatenated k_full/v_full, ready to be passed into the next call.
6. Sanity-check correctness by running several steps in sequence and comparing the cached output at the final step against a naive full recompute from scratch over all keys/values — they must match exactly, since caching changes only how attention is computed, never what it computes.

The key insight is that no explicit causal mask is needed here: because the cache only ever contains past and current tokens (nothing from the future has been appended yet), attending over "everything in the cache" is automatically causal by construction — the mask that full-sequence training-time attention needs becomes unnecessary once generation is done one token at a time with an append-only cache.

Reference solution

import torch


def attend_with_kv_cache(q, k_new, v_new, cache):
    if cache is None:
        # first token: nothing to concatenate onto yet
        k_full, v_full = k_new, v_new  # (batch, heads, 1, head_dim)
    else:
        # append along the sequence dim (2) -- this is the entire "cache" idea
        k_full = torch.cat([cache["k"], k_new], dim=2)  # (batch, heads, seq_so_far+1, head_dim)
        v_full = torch.cat([cache["v"], v_new], dim=2)  # (batch, heads, seq_so_far+1, head_dim)

    head_dim = q.shape[-1]
    # q is a single new token; it attends over ALL keys seen so far (no mask
    # needed -- causality is automatic because only past/current keys exist)
    scores = torch.matmul(q, k_full.transpose(-2, -1)) / (head_dim ** 0.5)  # (batch, heads, 1, seq_so_far+1)
    attn = torch.softmax(scores, dim=-1)
    out = torch.matmul(attn, v_full)  # (batch, heads, 1, head_dim)

    new_cache = {"k": k_full, "v": v_full}
    return out, new_cache

Key Functions & Tricks

• torch.cat([cache["k"], k_new], dim=2) — appends along the sequence axis; getting the dim right (2, not 0 or -1) is the most common bug in a from-scratch cache.
• torch.matmul(q, k_full.transpose(-2, -1)) — batched matrix multiply that broadcasts over the leading (batch, heads) dims automatically.
• / (head_dim ** 0.5) — the standard scaled dot-product attention scaling that keeps softmax inputs from growing with head_dim.
• torch.softmax(scores, dim=-1) — normalizes attention weights over the key/sequence axis specifically, not any other dimension.
• torch.testing.assert_close against a from-scratch recompute — the real correctness bar for any cache: the cached and uncached paths must be numerically identical, not just similarly shaped.

How to Recognize This Pattern