Problem 108 · Attention & Inference

KV-Cache Incremental Decode Equivalence

HardAI engineering

01Problem

Build a one-layer, one-head toy decoder: deterministic token embeddings (seeded), projections Wq, Wk, Wv, causal attention, and an unembedding to logits. Implement `full_forward(tokens)` that recomputes every position from scratch and a `KVCache` class whose `step(token)` projects only the new token, appends its K and V, and attends over the cache. Prove they produce the same logits at every position and the same greedy continuation, and count projection calls to show the saving.

02Key insight

Causal attention means position t's output depends only on positions <= t, and their K and V never change once computed; caching them turns each decode step from O(t) projections into O(1). The equivalence test is the contract every serving optimisation (paged attention, prefix caching, speculative decoding) must keep: if cached and uncached logits drift beyond float tolerance, you have a position or masking bug. The call counter makes the asymptotics concrete: n(n+1)/2 projections versus n.

03Solution

solution.pypython71 lines

import math
import random

Vec = list[float]


class ToyLM:
    def __init__(self, vocab: int, d: int, seed: int = 0) -> None:
        rng = random.Random(seed)
        r = lambda rows, cols: [[rng.uniform(-1, 1) for _ in range(cols)] for _ in range(rows)]
        self.vocab, self.d = vocab, d
        self.E = r(vocab, d)
        self.Wq, self.Wk, self.Wv = r(d, d), r(d, d), r(d, d)
        self.U = r(d, vocab)
        self.projections = 0

    def _mv(self, x: Vec, W: list[Vec]) -> Vec:
        return [sum(x[i] * W[i][j] for i in range(len(x))) for j in range(len(W[0]))]

    def project(self, tok: int) -> tuple[Vec, Vec, Vec]:
        self.projections += 1
        x = self.E[tok]
        return self._mv(x, self.Wq), self._mv(x, self.Wk), self._mv(x, self.Wv)

    def attend_and_unembed(self, q: Vec, K: list[Vec], V: list[Vec]) -> Vec:
        s = 1 / math.sqrt(self.d)
        scores = [sum(a * b for a, b in zip(q, k)) * s for k in K]
        m = max(scores)
        w = [math.exp(x - m) for x in scores]
        z = sum(w)
        h = [sum(w[j] * V[j][c] for j in range(len(V))) / z for c in range(self.d)]
        return self._mv(h, self.U)


def full_forward(lm: ToyLM, tokens: list[int]) -> list[Vec]:
    qkv = [lm.project(t) for t in tokens]
    return [lm.attend_and_unembed(qkv[t][0], [k for _, k, _ in qkv[:t + 1]], [v for _, _, v in qkv[:t + 1]])
            for t in range(len(tokens))]


class KVCache:
    def __init__(self, lm: ToyLM) -> None:
        self.lm, self.K, self.V = lm, [], []

    def step(self, tok: int) -> Vec:
        q, k, v = self.lm.project(tok)
        self.K.append(k)
        self.V.append(v)
        return self.lm.attend_and_unembed(q, self.K, self.V)


def argmax(xs: Vec) -> int:
    return max(range(len(xs)), key=lambda i: (xs[i], -i))


def generate_full(lm: ToyLM, prompt: list[int], n: int) -> list[int]:
    toks = list(prompt)
    for _ in range(n):
        toks.append(argmax(full_forward(lm, toks)[-1]))
    return toks


def generate_cached(lm: ToyLM, prompt: list[int], n: int) -> list[int]:
    cache, toks = KVCache(lm), list(prompt)
    logits: Vec = []
    for t in prompt:
        logits = cache.step(t)
    for _ in range(n):
        toks.append(argmax(logits))
        logits = cache.step(toks[-1])
    return toks

Try it first.

TimeCached: O(t * d + d^2) per step; uncached: O(t * d^2 + t^2 * d) per step

SpaceO(T * d) cache per layer per head

tests.py · 4 passing

tests.pypython15 lines

lm = ToyLM(vocab=11, d=6, seed=3)
seq = [1, 4, 2, 9, 0, 7]
full = full_forward(lm, seq)
cache = KVCache(lm)
for t, tok in enumerate(seq):
    inc = cache.step(tok)
    assert all(abs(a - b) < 1e-9 for a, b in zip(inc, full[t]))
lm.projections = 0
g_full = generate_full(lm, [1, 2], 10)
full_calls = lm.projections
lm.projections = 0
g_cached = generate_cached(lm, [1, 2], 10)
assert g_full == g_cached and len(g_full) == 12
assert lm.projections == 12
assert full_calls == sum(range(2, 12))

04Staff-level follow-up

Production caches are paged: fixed-size blocks with a block table per sequence (vLLM's PagedAttention), so fragmentation stays low and prefixes can be shared copy-on-write across requests. Then cover what invalidates a cache: changing the system prompt, a different LoRA adapter, a different RoPE scaling, or quantising K and V (fp8 KV halves memory but needs its own equivalence tolerance).