21. KV-Cache for Incremental Decoding

Problem

A real-time voice or text decoder generates output one token at a time. If every new token recomputed attention over the full history from scratch, generating a T-token response would cost O(T²) attention work with massive redundancy, since the keys and values for every already-seen token never change once computed — only the newest token contributes a new key/value pair. Caching the previously computed keys and values, and appending to that cache at each step instead of recomputing it, turns incremental generation into O(T) total work.

Implement one incremental decoding step: given the new token's query, key, and value, and the cache accumulated from all previous steps, compute that token's attention output (attending causally over everything cached so far, plus itself) and return the updated cache for the next step.

Source: src/21_kv_cache_incremental_decoding.py

def kv_cache_step(
    q_new: torch.Tensor, k_new: torch.Tensor, v_new: torch.Tensor,
    cache: tuple[torch.Tensor, torch.Tensor] | None,
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: ...

>>> q0 = k0 = v0 = torch.randn(1, 1, 4)
>>> out0, cache = kv_cache_step(q0, k0, v0, cache=None)
>>> out0.shape, cache[0].shape
(torch.Size([1, 1, 4]), torch.Size([1, 1, 4]))
>>> q1 = k1 = v1 = torch.randn(1, 1, 4)
>>> out1, cache = kv_cache_step(q1, k1, v1, cache=cache)
>>> cache[0].shape
torch.Size([1, 2, 4])

Step-by-Step Approach

1. Handle the first-step case: if cache is None, there's no history to append to, so the cache after this step is simply (k_new, v_new).
2. On later steps, concatenate the new key/value onto the existing cache along the sequence axis: torch.cat([k_prev, k_new], dim=1) (and the same for v) — this is the entire "cache update," no recomputation of past keys/values.
3. Compute attention scores between the single new query and every cached key: q_new @ k_cache.transpose(-1, -2) / sqrt(d_model), shape (batch, 1, t_so_far).
4. Notice no explicit causal mask is needed here — the cache by construction only ever contains positions up to and including the current one, so "attend over everything in the cache" already is the causal constraint.
5. Softmax over the cached-key axis, then multiply by v_cache to get the (batch, 1, d_model) output for this token.
6. Verify the core correctness property directly: running this step function token-by-token over a whole sequence and concatenating the per-step outputs must exactly equal ordinary full-sequence causal self-attention computed in one shot — the cache is purely an optimization, not a different computation.

The key insight is that a KV-cache trades recomputation for storage: every past key/value is computed exactly once and never touched again, so the total attention work across a full generation drops from O(T²) (recompute the whole prefix every step) to O(T) cache growth plus O(t) attention cost at step t, and the growing-cache-attend-with-no-mask formulation is mathematically identical to full causal attention because the cache never contains a future position.

Reference solution

import math
import torch


def kv_cache_step(
    q_new: torch.Tensor, k_new: torch.Tensor, v_new: torch.Tensor,
    cache: tuple[torch.Tensor, torch.Tensor] | None,
) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]:
    if cache is None:
        # first token: nothing to prepend, cache starts as just this step
        k_cache, v_cache = k_new, v_new
    else:
        k_prev, v_prev = cache
        # append along the sequence axis (dim=1); the new token's own k/v
        # must be included so the new token can attend to itself
        k_cache = torch.cat([k_prev, k_new], dim=1)
        v_cache = torch.cat([v_prev, v_new], dim=1)

    d_model = q_new.shape[-1]
    # q_new is a single query position; k_cache/v_cache hold every position
    # up to and including this one -- no explicit causal mask is needed
    # because the cache by construction never contains a future position
    scores = q_new @ k_cache.transpose(-1, -2) / math.sqrt(d_model)  # (batch, 1, t_so_far)
    attn_weights = torch.softmax(scores, dim=-1)
    output = attn_weights @ v_cache  # (batch, 1, d_model)

    return output, (k_cache, v_cache)

Key Functions & Tricks

• torch.cat([k_prev, k_new], dim=1) — O(1) append along the sequence axis; the operation that replaces recomputing keys/values for every previously seen token.
• Cache represented as a plain (k_cache, v_cache) tuple — in a real serving stack this is usually a pre-allocated fixed-size buffer per layer per request instead of a tuple that reallocates on every cat, but the logical behavior is the same.
• No explicit mask against the cache — recognizing that "attend over the cache" already encodes causality once you trust the cache only ever holds past-and-current positions is the crux of why this is O(t) per step instead of O(t²).
• torch.testing.assert_close(incremental_out, full_out, ...) — the equivalence-to-full-attention check is the real correctness bar for a cache implementation, not just shape checking each step.

How to Recognize This Pattern