Attention With a KV Cache, From Scratch

TL;DR: Cache each layer's K and V for every token already processed, project only the new tokens, and attend the new queries over the whole cache with a causal mask whose diagonal is offset by S minus T. That makes each decode step one token of projection work instead of a full re-run: generating 1,000 tokens after a 1,000-token prompt costs about 1,500x fewer FLOPs than recomputing from scratch, in exchange for 128 KiB of KV per token on Llama 3.1 8B.

What the cache stores and why

Attention for token $t$ needs the keys and values of tokens $0..t$. Those depend only on earlier tokens (causality), so they never change once computed. Without a cache, step $t$ re-runs the whole forward pass over $t$ tokens; with one, it runs the forward pass for one token and reads the stored K and V.

no cache : step t costs about 2N x t FLOPs ; with cache : about 2N per token
1,000 prompt + 1,000 generated:
  recompute = 2N x (1,001 + ... + 2,000) = 2N x 1,500,500
  cached    = 2N x 1,000 (after prefill)
  ratio     = 1,500x  (attention score FLOPs ignored; they favor the cache further)

price, Llama 3.1 8B (32 layers, 8 KV heads, head_dim 128, BF16):
  KV/token  = 2 x 32 x 8 x 128 x 2 B = 131,072 B = 128 KiB
  at 8,192 tokens = 1 GiB per sequence
  decode step, batch 1, H100 SXM, 4,096 context:
  (16.06 GB weights + 0.54 GB KV) / 3.35 TB/s = 4.95 ms floor, bandwidth-bound

The cache converts a compute problem into a memory problem. Decode becomes a stream over weights and KV, which is why every later optimization (paged KV, GQA, KV quantization) is about bytes.

The offset causal mask

During prefill, queries and keys are the same tokens, so the mask is the usual lower triangle. During decode or chunked prefill there are T new queries and S total keys, with T smaller than S. Query $i$ sits at absolute position $S - T + i$ and may see key $j$ exactly when:

np.tril(np.ones((T, S))) gets this wrong. It anchors the diagonal at the top left, so with T = 3 and S = 7 the first new query sees one key instead of five. The output is plausible numbers and the test against full recompute fails, which is why that test exists.

The implementation

import numpy as np

def attention(q, k, v):
    """q: [H, T, d], k, v: [Hkv, S, d] with S >= T. Query i sits at S - T + i."""
    H, T, d = q.shape
    Hkv, S, _ = k.shape
    if T == 0:
        return np.zeros_like(q)
    k = np.repeat(k, H // Hkv, axis=0)                  # GQA: share KV heads
    v = np.repeat(v, H // Hkv, axis=0)
    s = q @ k.transpose(0, 2, 1) / np.sqrt(d)           # [H, T, S]
    pos_q = np.arange(T)[:, None] + (S - T)             # offset diagonal
    s = np.where(np.arange(S)[None, :] <= pos_q, s, -np.inf)
    s = s - s.max(axis=-1, keepdims=True)
    p = np.exp(s)
    return (p / p.sum(axis=-1, keepdims=True)) @ v      # [H, T, d]

class KVCache:
    def __init__(self, Hkv, d):
        self.k = np.zeros((Hkv, 0, d))
        self.v = np.zeros((Hkv, 0, d))

    def append(self, k_new, v_new):                     # [Hkv, t, d]
        self.k = np.concatenate([self.k, k_new], axis=1)
        self.v = np.concatenate([self.v, v_new], axis=1)
        return self.k, self.v

def step(x, Wq, Wk, Wv, cache, H, Hkv):
    """x: [t, D] new tokens only. Projects them once, attends over the cache."""
    t, D = x.shape
    d = Wq.shape[1] // H
    q = (x @ Wq).reshape(t, H, d).transpose(1, 0, 2)
    k = (x @ Wk).reshape(t, Hkv, d).transpose(1, 0, 2)
    v = (x @ Wv).reshape(t, Hkv, d).transpose(1, 0, 2)
    K, V = cache.append(k, v)
    return attention(q, K, V).transpose(1, 0, 2).reshape(t, H * d)

Three details carry the grade. The max subtraction before exp keeps the softmax finite; every row has at least one unmasked key (itself), so no row is all -inf. GQA is a repeat of KV heads to query heads, and the cache stores only the 2 KV heads, which is the whole memory saving. And step takes any number of new tokens, so the same function does prefill, chunked prefill and decode.

The tests

rng = np.random.default_rng(0)
H, Hkv, d, D, N = 8, 2, 16, 64, 13
Wq = rng.standard_normal((D, H * d)) / 8
Wk = rng.standard_normal((D, Hkv * d)) / 8
Wv = rng.standard_normal((D, Hkv * d)) / 8
x = rng.standard_normal((N, D))

full = step(x, Wq, Wk, Wv, KVCache(Hkv, d), H, Hkv)          # one prefill
c = KVCache(Hkv, d)
parts = [step(x[:5], Wq, Wk, Wv, c, H, Hkv)]                 # prefill 5
parts += [step(x[i:i + 1], Wq, Wk, Wv, c, H, Hkv) for i in range(5, N)]
print(np.abs(full - np.concatenate(parts)).max() < 1e-12)    # True
print(step(x[:0], Wq, Wk, Wv, KVCache(Hkv, d), H, Hkv).shape)  # (0, 128)

Prefill-then-decode must match one full pass to float64 rounding (the measured gap is about 1e-15). The prefill of 5 exercises T = 5 with S = 5, each decode exercises T = 1 with S up to 13, and an empty input returns an empty result instead of a softmax over nothing.

Where it breaks

• Concatenation is quadratic in copies. Growing an array per step copies the cache every token. Real engines preallocate pages and write in place, which is what paged KV memory does.
• Batching ragged sequences. Different lengths need per-sequence S, so the mask and the cache become per-sequence offsets plus a block table.
• Positional encoding. RoPE must rotate each new key by its absolute position, S minus T plus i, before it is cached; rotating by the local index is the decode-only version of the mask bug.
• Sliding-window layers. gpt-oss-120b alternates full attention with a 128-token window. Those layers add a lower bound, key $j$ must also satisfy $j > S - T + i - 128$, and their cache can drop every key older than the window, so the KV for half the layers stops growing.
• Precision. BF16 scores over long contexts lose accuracy; production kernels accumulate in FP32 and use the online softmax so the full score row is never materialized.

What interviewers listen for

Shapes stated up front, the offset mask derived rather than guessed, and a test against full recompute written without being asked. They watch for whether the cache stores KV heads or query heads, whether T = 0 and T < S work, and whether the candidate can price the cache in bytes per token and explain why it makes decode bandwidth-bound.

Say it in 30 seconds

"Keys and values of past tokens never change, so I cache them per layer and project only new tokens. New queries attend over the whole cache with a causal mask offset by S minus T, so query i sees keys up to S minus T plus i. GQA stores only the KV heads and repeats them at attention time. I test that prefill plus incremental decode matches one full pass. The cache trades quadratic recompute for 128 KiB per token on an 8B model, which is why decode is bound by memory bandwidth."