QKV attention programming with einsum fluency

Question: QKV attention programming with einsum fluency

Companies: Anthropic

Difficulty: Medium

Problem statement

A short programming screen around sequence models and scaled dot-product attention, plus quick ML knowledge questions. Reading equations and expressing the math with einsum-style ops is the point.

Use NumPy. Shapes: b batch, h heads, t query positions, s key positions, d head width, m model width.

import numpy as np


def softmax(x, axis=-1): ...
def attention(q, k, v, mask=None): ...            # mask: True = blocked
def causal_mask(t, s=None): ...
def multi_head_attention(x, wq, wk, wv, wo, causal=True): ...   # x: (b, t, m); w*: (m, h, d); wo: (h, d, m)
def grouped_query_attention(q, k, v, mask=None): ...            # q: (b, h, t, d); k, v: (b, g, s, d)
def decode_step(x_new, cache, wq, wk, wv, wo): ...              # one token with a KV cache dict

Approach

What the screen checks. You are given an equation (or a paper excerpt) and asked to turn it into array code quickly and correctly, then answer short ML questions. The skill is reading shapes. Before writing any line, write the shape of every tensor as a comment; einsum then almost writes itself, because its subscripts are the shapes.

The equation. For queries Q (T_q × d), keys K (T_k × d) and values V (T_k × d_v):

Attention(Q, K, V) = softmax(Q Kᵀ / √d + M) V

where M is 0 where attention is allowed and −∞ where it is masked (future tokens, padding).

Shapes for multi-head attention. Use b batch, t time (query), s time (key), m model width, h heads, d head width, with m = h × d.

Writing the projection weights as (m, h, d) instead of (m, m) removes every reshape and transpose, which is where most bugs come from.

Numerical details that interviewers check.

• Scale by 1/√d. If q and k entries have unit variance, q·k has variance d. At d = 128 the scores have a standard deviation of about 11, softmax becomes nearly one-hot, and gradients vanish. Dividing by √d brings the variance back to 1.
• Stable softmax. Subtract the row max before exp. Without it, a score of 100 overflows float16 (max ≈ 65,504, and e¹² ≈ 162,755 already exceeds it).
• Masking. Add −∞ (or a large negative number) before the softmax, never multiply after it, or the rows no longer sum to 1. A row with every position masked (all padding) gives 0/0; set it to zeros explicitly.
• Causal mask. Query position i may see key positions j ≤ i: mask j > i, the strict upper triangle.

Grouped-query attention (a common follow-up). With h query heads and g key/value heads (g divides h), each K/V head serves r = h / g query heads. Reshape q to (b, g, r, t, d) and use "bgrtd,bgsd->bgrts". KV cache memory drops by h / g; with h = 64 and g = 8 that is 8× less.

Decode with a KV cache. At generation step t only the new token's query is needed, and keys and values of earlier tokens are reused from the cache. Cost per step is O(t × d) per head instead of O(t² × d) for recomputing everything. The test below checks that step-by-step decoding gives exactly the causal full-sequence result.

Code

import numpy as np


def softmax(x, axis=-1):
    x = x - np.max(x, axis=axis, keepdims=True)
    e = np.exp(x)
    total = e.sum(axis=axis, keepdims=True)
    return np.divide(e, total, out=np.zeros_like(e), where=total > 0)


def attention(q, k, v, mask=None):
    """q: (..., t, d), k: (..., s, d), v: (..., s, dv); mask: True = blocked."""
    d = q.shape[-1]
    scores = np.einsum("...td,...sd->...ts", q, k) / np.sqrt(d)
    if mask is not None:
        scores = np.where(mask, -np.inf, scores)
        all_blocked = np.all(mask, axis=-1, keepdims=True)
        scores = np.where(all_blocked, 0.0, scores)  # avoid nan rows
    weights = softmax(scores)
    if mask is not None:
        weights = np.where(all_blocked, 0.0, weights)
    return np.einsum("...ts,...sd->...td", weights, v)


def causal_mask(t, s=None):
    s = t if s is None else s
    # query i sits at absolute position i + (s - t) when s > t (cached prefix)
    i = np.arange(t)[:, None] + (s - t)
    j = np.arange(s)[None, :]
    return j > i


def multi_head_attention(x, wq, wk, wv, wo, causal=True):
    """x: (b, t, m); wq/wk/wv: (m, h, d); wo: (h, d, m)."""
    q = np.einsum("btm,mhd->bhtd", x, wq)
    k = np.einsum("btm,mhd->bhtd", x, wk)
    v = np.einsum("btm,mhd->bhtd", x, wv)
    mask = causal_mask(x.shape[1]) if causal else None
    out = attention(q, k, v, mask)
    return np.einsum("bhtd,hdm->btm", out, wo)


def grouped_query_attention(q, k, v, mask=None):
    """q: (b, h, t, d); k, v: (b, g, s, d) with g dividing h."""
    b, h, t, d = q.shape
    g = k.shape[1]
    qg = q.reshape(b, g, h // g, t, d)
    scores = np.einsum("bgrtd,bgsd->bgrts", qg, k) / np.sqrt(d)
    if mask is not None:
        scores = np.where(mask, -np.inf, scores)
    out = np.einsum("bgrts,bgsd->bgrtd", softmax(scores), v)
    return out.reshape(b, h, t, d)


def decode_step(x_new, cache, wq, wk, wv, wo):
    """One token. x_new: (b, 1, m); cache: dict with 'k', 'v' of shape (b, h, s, d)."""
    q = np.einsum("btm,mhd->bhtd", x_new, wq)
    k_new = np.einsum("btm,mhd->bhtd", x_new, wk)
    v_new = np.einsum("btm,mhd->bhtd", x_new, wv)
    cache["k"] = k_new if cache.get("k") is None else np.concatenate([cache["k"], k_new], axis=2)
    cache["v"] = v_new if cache.get("v") is None else np.concatenate([cache["v"], v_new], axis=2)
    out = attention(q, cache["k"], cache["v"])  # the new token may see everything cached
    return np.einsum("bhtd,hdm->btm", out, wo)

Complexity

For sequence length T, model width m, h heads of width d:

• Projections: O(T × m²) FLOPs (4 matrices of m × m in total, counting W_o).
• Scores and mixing: O(h × T² × d) = O(T² × m). For T = 8,192 and m = 4,096 the two einsums together are 2 × 8,192² × 4,096 ≈ 5.5e11 multiply-adds per layer, the same as the four projections (4 × 8,192 × 4,096² ≈ 5.5e11). Attention dominates once T is longer than 2m.
• Memory: the (h, T, T) score matrix. At h = 32, T = 8,192, fp16: 32 × 8,192² × 2 bytes ≈ 4.3 GB per layer per sequence. Fused kernels (FlashAttention) never materialize it: they tile over keys and keep a running max and sum, so memory is O(T × d).
• Decoding with a cache: O(h × t × d) per token, and the cache holds 2 × layers × g × d values per token.

Test cases

import numpy as np

from solution import attention, causal_mask, softmax


def loop_attention(q, k, v):
    t, d = q.shape
    out = np.zeros((t, v.shape[1]))
    for i in range(t):
        s = np.array([q[i] @ k[j] for j in range(k.shape[0])]) / np.sqrt(d)
        w = np.exp(s - s.max())
        w /= w.sum()
        out[i] = w @ v
    return out


def test_matches_loops():
    rng = np.random.default_rng(0)
    q, k, v = rng.normal(size=(5, 8)), rng.normal(size=(7, 8)), rng.normal(size=(7, 3))
    assert np.allclose(attention(q, k, v), loop_attention(q, k, v))


def test_softmax_is_stable_and_normalized():
    p = softmax(np.array([[1000.0, 1001.0, 999.0]]))
    assert np.all(np.isfinite(p))
    assert np.allclose(p.sum(axis=-1), 1)


def test_causal_mask_shape_and_offset():
    assert causal_mask(3).tolist() == [[False, True, True], [False, False, True], [False, False, False]]
    assert causal_mask(1, 4).tolist() == [[False, False, False, False]]


def test_fully_masked_row_is_zero_not_nan():
    q, k, v = np.ones((2, 4)), np.ones((3, 4)), np.ones((3, 2))
    mask = np.array([[True, True, True], [False, True, True]])
    out = attention(q, k, v, mask)
    assert np.allclose(out[0], 0) and np.allclose(out[1], 1)

import numpy as np

from solution import attention, decode_step, grouped_query_attention, multi_head_attention


def weights(rng, m=16, h=4, d=4):
    return (rng.normal(size=(m, h, d)), rng.normal(size=(m, h, d)),
            rng.normal(size=(m, h, d)), rng.normal(size=(h, d, m)))


def test_causal_output_ignores_the_future():
    rng = np.random.default_rng(1)
    w = weights(rng)
    x = rng.normal(size=(2, 6, 16))
    y = x.copy()
    y[:, 4:] += 10.0
    a, b = multi_head_attention(x, *w), multi_head_attention(y, *w)
    assert np.allclose(a[:, :4], b[:, :4])
    assert not np.allclose(a[:, 4:], b[:, 4:])


def test_gqa_with_one_query_head_per_group_is_plain_attention():
    rng = np.random.default_rng(2)
    q, k, v = (rng.normal(size=(2, 4, 5, 8)) for _ in range(3))
    assert np.allclose(grouped_query_attention(q, k, v), attention(q, k, v))


def test_gqa_shares_kv_across_the_group():
    rng = np.random.default_rng(3)
    q = rng.normal(size=(1, 4, 3, 8))
    k, v = rng.normal(size=(1, 2, 3, 8)), rng.normal(size=(1, 2, 3, 8))
    expanded = np.repeat(k, 2, axis=1), np.repeat(v, 2, axis=1)
    assert np.allclose(grouped_query_attention(q, k, v), attention(q, *expanded))


def test_kv_cache_decoding_matches_full_causal_pass():
    rng = np.random.default_rng(4)
    w = weights(rng)
    x = rng.normal(size=(1, 5, 16))
    full = multi_head_attention(x, *w)
    cache = {}
    steps = [decode_step(x[:, i:i + 1], cache, *w) for i in range(5)]
    assert np.allclose(np.concatenate(steps, axis=1), full)

Follow-ups

These are the short knowledge questions that usually follow. One or two sentences each is the right length.

1. Why multiple heads? Each head can attend by a different criterion (syntax, position, coreference) in a d-dimensional subspace at the same total cost as one wide head. One head's softmax picks one mixture per position; many heads give many.
2. Where does position come from? Attention is permutation-equivariant without it. RoPE rotates q and k by position-dependent angles so q·k depends on the relative offset; it extrapolates better than learned absolute embeddings and works with the KV cache.
3. Pre-norm or post-norm? Pre-norm (LayerNorm before attention and MLP, residual around them) keeps an identity path through the network and trains stably at depth without careful warm-up; most modern LLMs use pre-norm, often with RMSNorm.
4. Why is decoding memory-bound? Each step does about 2 FLOPs per weight read (one multiply-add per parameter for a single token), far below the ~300 FLOPs per byte a modern GPU can sustain. Batching many sequences raises the FLOPs per byte read, which is why serving systems batch.
5. Cross-entropy and perplexity. The training loss is the mean negative log probability of the next token; perplexity is exp(loss). A loss of 2.0 nats means the model is as uncertain as a uniform choice among about 7.4 tokens.

Interviewer expectations

Mid-level (L4)

• Implements scaled dot-product attention with correct shapes, the 1/√d scale and a stable softmax.
• Adds a causal mask that blocks future positions, applied before the softmax.
• Answers the basic questions (why scale, why heads) correctly.

Senior (L5)

• Writes multi-head attention with einsum and no reshape bugs, with shapes annotated, and tests it against a loop.
• Handles fully masked rows, grouped-query attention and KV-cache decoding, and proves the cache matches the full pass.
• Gives compute and memory costs with numbers (T², score-matrix size).

Staff (L6+)

• Explains why naive attention is memory-bound and how fused kernels tile with a running max and sum.
• Connects GQA and cache size to serving cost, and decoding to arithmetic intensity.
• Moves fluently between the equation, the code and the hardware cost, and spots shape or masking mistakes in someone else's code.