Implement causal attention and a KV cache from scratch

The question

Implement causal multi-head self-attention in NumPy. Then add a KV cache so generation doesn’t recompute attention over the whole sequence for every new token. Show that it’s correct.

Candidates report being asked to code attention, the KV cache and grouped-query attention from scratch in ML and AI engineer loops. Even when it isn’t asked as code, every serving question assumes you understand it: the KV cache is why long context is expensive.

What to say before writing code

• Attention is softmax(Q·Kᵀ / √d) · V, computed per head.
• Causal means token i can only attend to tokens up to and including i. A mask enforces it.
• During generation, the keys and values of past tokens never change. Recomputing them for every new token is wasted work, so cache K and V, and compute only the new token’s Q, K and V.
• Grouped-query attention (GQA): several query heads share one KV head. It shrinks the cache by a factor of n_heads / n_kv_heads, which is why Llama-family models use 8 KV heads.
• How to prove it’s correct: prefill, then decode token by token with the cache, and you must get exactly the same outputs as one full forward pass. That’s the test.

The implementation

import numpy as np


def softmax(x: np.ndarray, axis: int = -1) -> np.ndarray:
    x = x - x.max(axis=axis, keepdims=True)  # numerical stability
    e = np.exp(x)
    return e / e.sum(axis=axis, keepdims=True)


class MultiHeadAttention:
    """Causal multi-head self-attention with grouped-query attention (GQA) and a KV cache."""

    def __init__(self, d_model: int, n_heads: int, n_kv_heads: int, seed: int = 0):
        assert d_model % n_heads == 0 and n_heads % n_kv_heads == 0
        self.d_model, self.n_heads, self.n_kv_heads = d_model, n_heads, n_kv_heads
        self.head_dim = d_model // n_heads
        rng = np.random.default_rng(seed)
        s = 1 / np.sqrt(d_model)
        self.w_q = rng.normal(0, s, (d_model, n_heads * self.head_dim))
        self.w_k = rng.normal(0, s, (d_model, n_kv_heads * self.head_dim))
        self.w_v = rng.normal(0, s, (d_model, n_kv_heads * self.head_dim))
        self.w_o = rng.normal(0, s, (n_heads * self.head_dim, d_model))
        self.k_cache = None  # (n_kv_heads, tokens_so_far, head_dim)
        self.v_cache = None

    def _split(self, x: np.ndarray, heads: int) -> np.ndarray:
        # (T, heads * head_dim) -> (heads, T, head_dim)
        return x.reshape(x.shape[0], heads, self.head_dim).transpose(1, 0, 2)

    def _attend(self, q, k, v, q_offset: int) -> np.ndarray:
        # q: (H, Tq, D); k, v: (Hkv, Tk, D). Each KV head serves H / Hkv query heads.
        group = self.n_heads // self.n_kv_heads
        k = np.repeat(k, group, axis=0)
        v = np.repeat(v, group, axis=0)
        scores = q @ k.transpose(0, 2, 1) / np.sqrt(self.head_dim)  # (H, Tq, Tk)
        tq, tk = q.shape[1], k.shape[1]
        # Query i sits at absolute position q_offset + i and may see keys 0..q_offset + i.
        q_pos = q_offset + np.arange(tq)[:, None]
        mask = np.arange(tk)[None, :] > q_pos
        scores = np.where(mask, -np.inf, scores)
        out = softmax(scores) @ v  # (H, Tq, D)
        return out.transpose(1, 0, 2).reshape(tq, -1) @ self.w_o

    def forward(self, x: np.ndarray) -> np.ndarray:
        """Full causal attention over a whole sequence x of shape (T, d_model). No cache."""
        q = self._split(x @ self.w_q, self.n_heads)
        k = self._split(x @ self.w_k, self.n_kv_heads)
        v = self._split(x @ self.w_v, self.n_kv_heads)
        return self._attend(q, k, v, q_offset=0)

    def reset_cache(self) -> None:
        self.k_cache = self.v_cache = None

    def step(self, x_new: np.ndarray) -> np.ndarray:
        """Process new tokens (prefill if T > 1, decode if T == 1), reusing cached keys/values."""
        offset = 0 if self.k_cache is None else self.k_cache.shape[1]
        q = self._split(x_new @ self.w_q, self.n_heads)
        k = self._split(x_new @ self.w_k, self.n_kv_heads)
        v = self._split(x_new @ self.w_v, self.n_kv_heads)
        if self.k_cache is None:
            self.k_cache, self.v_cache = k, v
        else:
            self.k_cache = np.concatenate([self.k_cache, k], axis=1)
            self.v_cache = np.concatenate([self.v_cache, v], axis=1)
        return self._attend(q, self.k_cache, self.v_cache, q_offset=offset)

    def kv_cache_bytes(self, bytes_per_value: int = 2) -> int:
        if self.k_cache is None:
            return 0
        return 2 * self.k_cache.size * bytes_per_value  # keys + values

The one subtle line is the mask. With a cache, the new queries don’t start at position 0. They start at q_offset, the number of tokens already cached. Get that wrong and decoding attends to the wrong positions without raising any error.

Tests

import numpy as np

from attention_kv_cache import MultiHeadAttention, softmax


def test_softmax_rows_sum_to_one_and_handle_large_values():
    p = softmax(np.array([[1000.0, 1000.0], [0.0, -np.inf]]))
    assert np.allclose(p, [[0.5, 0.5], [1.0, 0.0]])


def test_output_shape():
    attn = MultiHeadAttention(d_model=64, n_heads=8, n_kv_heads=2)
    x = np.random.default_rng(1).normal(size=(10, 64))
    assert attn.forward(x).shape == (10, 64)


def test_causality_future_tokens_do_not_change_past_outputs():
    attn = MultiHeadAttention(d_model=32, n_heads=4, n_kv_heads=4)
    x = np.random.default_rng(2).normal(size=(6, 32))
    full = attn.forward(x)
    prefix = attn.forward(x[:4])
    assert np.allclose(full[:4], prefix)


def test_kv_cache_decode_matches_full_forward():
    attn = MultiHeadAttention(d_model=64, n_heads=8, n_kv_heads=2)
    x = np.random.default_rng(3).normal(size=(12, 64))
    expected = attn.forward(x)

    attn.reset_cache()
    out = [attn.step(x[:5])]  # prefill the first 5 tokens
    for t in range(5, 12):  # then decode one token at a time
        out.append(attn.step(x[t : t + 1]))
    assert np.allclose(np.concatenate(out), expected)


def test_gqa_shrinks_cache_by_head_ratio():
    x = np.random.default_rng(4).normal(size=(16, 64))
    mha = MultiHeadAttention(d_model=64, n_heads=8, n_kv_heads=8)
    gqa = MultiHeadAttention(d_model=64, n_heads=8, n_kv_heads=2)
    mha.step(x)
    gqa.step(x)
    assert mha.kv_cache_bytes() == 4 * gqa.kv_cache_bytes()
    # 2 (K and V) x 2 KV heads x 16 tokens x 8 head_dim x 2 bytes
    assert gqa.kv_cache_bytes() == 2 * 2 * 16 * 8 * 2

test_kv_cache_decode_matches_full_forward is the test that matters. It prefills 5 tokens, decodes 7 more one at a time, and checks that every output equals the full forward pass.

Complexity

• Without a cache, step t recomputes K and V for all t tokens of the prefix. Generating n tokens costs O(n²) projections in total.
• With a cache, each step projects only the new token and attends over the cached keys: O(1) projections and O(t) attention work per step.
• The price is memory: 2 × layers × kv_heads × head_dim × bytes per token. For Llama 3.1 70B in BF16 that’s 320 KB per token. See the numbers and the KV cache calculator.

Follow-ups the interviewer will ask

1. Why is decode memory-bound? Each step does very little compute per byte read: every weight and the whole cache are read to produce one token per sequence. Batching many sequences shares the weight reads.
2. What is paged attention? Instead of one contiguous, maximum-length buffer per sequence, the cache lives in fixed-size blocks tracked by a block table, like virtual memory pages. That removes fragmentation and lets sequences share prefix blocks. It’s the core idea of vLLM.
3. What’s prefix caching? If many requests share a system prompt, compute its K and V once and reuse the blocks. Prefill cost drops and time to first token improves.
4. How would you shrink the cache further? Quantise it to FP8 (half the size), use fewer KV heads (GQA or multi-query attention), or use sliding-window attention in some layers.
5. Add RoPE. Rotary position embeddings are applied to Q and K before K is cached, using each token’s absolute position, so the cached keys are already rotated.

Common mistakes

• Forgetting the √d scale, which saturates the softmax.
• A softmax that overflows on large scores because the maximum isn’t subtracted first.
• Masking with the wrong offset after the first decode step.
• Testing only the output shapes, not equivalence with the full forward pass.