Implement multi-head self-attention from scratch in NumPy, with a causal mask.

TL;DR: Project the input into queries, keys, and values, reshape into H heads, compute softmax(QKᵀ/√d_k + mask)·V per head in parallel, concatenate the heads, and apply an output projection. The causal mask sets future positions to -inf before the softmax so a token never attends ahead. The details that matter: the √d_k scale, numerically stable softmax, and getting the head reshape and mask broadcasting right.

How to approach it. State the pipeline (project to Q/K/V, split heads, scaled-dot-product attention per head with the causal mask, concat, output project) and the tensor shapes before coding. Call out the two correctness traps up front: the √d_k scale and applying the causal mask as -inf pre-softmax, not zeroing post-softmax.

A strong answer.

import numpy as np

def softmax(x):                                  # stable softmax over last axis
    x = x - x.max(axis=-1, keepdims=True)
    e = np.exp(x)
    return e / e.sum(axis=-1, keepdims=True)

def mha(X, Wq, Wk, Wv, Wo, n_heads, causal=True):
    # X: (T, d_model); W*: (d_model, d_model)
    T, d_model = X.shape
    d_k = d_model // n_heads
    Q, K, V = X @ Wq, X @ Wk, X @ Wv             # each (T, d_model)

    def split(M):                                # (T, d_model) -> (H, T, d_k)
        return M.reshape(T, n_heads, d_k).transpose(1, 0, 2)
    Qh, Kh, Vh = split(Q), split(K), split(V)

    scores = Qh @ Kh.transpose(0, 2, 1) / np.sqrt(d_k)   # (H, T, T)
    if causal:
        mask = np.triu(np.full((T, T), -np.inf), k=1)    # block attending ahead
        scores = scores + mask                           # broadcasts over heads
    out = softmax(scores) @ Vh                           # (H, T, d_k)

    out = out.transpose(1, 0, 2).reshape(T, d_model)     # concat heads
    return out @ Wo                                      # output projection

What a strong candidate narrates while writing it:

• Shapes are the whole game. (T, d_model) splits into (H, T, d_k) with d_k = d_model / n_heads; heads are an independent batch dimension, which is why multi-head attention parallelizes cleanly.
• The √d_k scale keeps the logits from saturating the softmax (the variance argument).
• Causal mask as -inf before softmax, not zeroing after: exp(-inf)=0, so masked positions get exactly zero weight; zeroing post-softmax would leave the distribution unnormalized.
• Numerically stable softmax (subtract the max), or large logits overflow.
• Concat then output-project: reshape heads back together and mix them with Wo; the output projection is not optional, it lets heads' outputs interact.

The shape journey, which is the part interviewers watch your hands on:

Key takeaways

• Heads are just an extra batch axis: split d_model into H × d_k, run scaled-dot-product attention independently, concat, project.
• The √d_k scale controls logit variance so softmax does not saturate; dividing by d_k instead is a common silent bug.
• Causal masking is -inf added before softmax (so exp zeroes it), never zeroing after, which would leave the row unnormalized.
• The output projection Wo is load-bearing: it lets the heads' subspaces mix.

What interviewers probe next.

• "Why split into heads at all?" Each head can attend to a different relationship (syntax, coreference, position) in its own subspace; concatenating gives a richer representation than one big head.
• "Complexity?" O(T²·d) time and an O(T²) attention matrix per head in memory, which is why long sequences are expensive and motivate FlashAttention.
• "Add the KV cache for inference." At decode you append the new token's K/V to cached K/V and attend over the whole cache, so you do not recompute past keys/values each step.
• "MQA/GQA change here?" Use fewer K/V heads than Q heads (share them), shrinking the KV cache; the reshape for K/V uses the smaller head count.

Common mistakes.

• Forgetting the √d_k scale, or dividing by d_k.
• Applying the causal mask after the softmax (must be -inf before) or getting the triangular direction wrong (k=1 to exclude the diagonal-future, keep self).
• Botching the head reshape/transpose so heads and the time axis get mixed.
• A non-stable softmax that overflows on large logits.