Multi-Head Attention from Scratch

TL;DR — The whiteboard task is always the same: given x of shape (B, T, d_model) and n_heads, produce the multi-head attention output. Four decisions account for almost all lost points: (1) you must split d_model into n_heads × d_k and move the head axis next to the batch axis, (2) you must scale by $1/ d_{k}$ before the softmax, (3) the causal mask must be applied before the softmax and must use a large finite negative number, not -inf, and (4) you must make the tensor contiguous before reshaping the heads back together. Get those four right and the rest is transcription. This chapter is the drill: the reference implementation, the four decisions explained as decisions, the follow-ups that follow each one, and the tests that prove your implementation is correct rather than merely plausible.

1. Why this exists

“Implement multi-head attention” is the single most common whiteboard question for LLM roles, for a specific reason: it is the shortest task that cannot be completed by recall. To write it you must simultaneously hold

the shapes ((B, T, d_model) → four projections → heads → back),
the semantics (why the head axis exists and what transposing it does),
the numerics (the scaling, and why -inf is a trap),
and the frameworks (what .view, .transpose, .contiguous each do).

Someone who memorised the paper can reproduce the formula $softmax (Q K^{⊤} / d_{k}) V$ and still fail the question, because the question is really “can you implement a batched tensor program”. The follow-up questions — why the transpose, what happens if you mask after the softmax — are designed to separate the two.

What an interviewer is checking. Not whether you can recall the attention equation. Whether you can (a) reason about tensor shapes out loud, (b) name a numerical hazard before it bites, and (c) say what you would test.

2. How it works

2.1 The shape journey

Work the shapes forward before writing a line. Everything else follows.

Step	Expression	Shape
input	`x`	`(B, T, d_model)`
project	`W_q(x)`, `W_k(x)`, `W_v(x)`	`(B, T, d_model)` each
split heads	`.view(B, T, h, d_k)`	`(B, T, h, d_k)`
move head axis	`.transpose(1, 2)`	`(B, h, T, d_k)`
scores	`Q @ K.transpose(-2, -1)`	`(B, h, T, T)`
scale + mask + softmax	row-wise over the last axis	`(B, h, T, T)`
mix values	`attn @ V`	`(B, h, T, d_k)`
merge heads	`.transpose(1, 2).contiguous().view(B, T, d_model)`	`(B, T, d_model)`
output projection	`W_o(...)`	`(B, T, d_model)`

where $B$ is the batch size, $T$ the sequence length, $h$ the number of heads, and $d_{k} = d_{m o d e l} / h$ the per-head dimension. The d_model must be divisible by h; assert it.

Why the head axis goes to position 1. The matrix multiplies Q @ K.transpose(-2, -1) and attn @ V are batched over every leading axis. If the head axis sits at position 2, it is treated as part of the sequence dimension and the multiply is wrong (or the shapes do not line up at all). By moving h next to B, the leading axes become (B, h) and PyTorch runs $B \times h$ independent attention operations in one call — no Python loop over heads. That is the entire point of the transpose, and it is the most common follow-up question.

2.2 The four decisions

(1) The scaling is applied to the scores. $\frac{1}{d _{k}}$ rescales the logits so their variance is roughly 1 before the softmax. The derivation is in Attention §2.3; the practical point here is where it goes. Mathematically you may fold it into $Q$ or $K$ beforehand, but writing it on the score tensor is what the derivation describes and what a reviewer expects to read.

(2) The mask is additive and pre-softmax. Masking is “make the forbidden logits so negative that they receive zero probability”. Since $softmax$ is applied along the key axis, the mask must be applied to the score matrix before the softmax. Applying it afterwards means you have to re-normalise, and zeroing weights after softmax gives rows that do not sum to 1.

(3) Use a large finite negative number, not -inf. If any row is fully masked — which happens with padding masks when a batch entry is entirely pad — then every entry of that row is $- \infty$ and the softmax evaluates $exp (- \infty) / \sum_{j} exp (- \infty_{j}) = 0/0$ , which is NaN. NaN then propagates through every subsequent operation and the loss comes out NaN with no obvious cause. torch.finfo(dtype).min (about $- 65504$ for fp16, $- 3.4 \times 1 0^{38}$ for fp32) underflows to zero in the exponent without producing NaN.

(4) .contiguous() before .view(). After transpose(1, 2) on the output the tensor’s logical shape no longer matches its memory strides. view requires a compatible layout and will raise RuntimeError: view size is not compatible with input tensor's size and stride. reshape would silently copy; view needs the explicit .contiguous() first. This is the cheapest question in the whole interview to get right and the most embarrassing to get wrong.

3. Implementation

This is the reference answer. Reproduce it from memory, on paper, in about ten minutes.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class MultiHeadAttention(nn.Module):
    """Self-attention. x: (B, T, d_model) -> (B, T, d_model)."""

    def __init__(self, d_model: int, n_heads: int, dropout: float = 0.0):
        super().__init__()
        assert d_model % n_heads == 0, "d_model must be divisible by n_heads"
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads

        # bias=False is the modern convention (LLaMA, GPT-NeoX, ...): the
        # following RMSNorm/LayerNorm already supplies a shift.
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, attn_mask=None):          # attn_mask: True = allowed
        B, T, _ = x.shape

        # (B, T, d_model) -> (B, h, T, d_k)
        def split(t):
            return t.view(B, T, self.n_heads, self.d_k).transpose(1, 2)

        Q = split(self.W_q(x))
        K = split(self.W_k(x))
        V = split(self.W_v(x))

        # (B, h, T, T): entry (i, j) = how much query i attends to key j
        scores = (Q @ K.transpose(-2, -1)) / math.sqrt(self.d_k)

        if attn_mask is not None:
            # (B, 1, T, T) or (B, T, T) broadcast against (B, h, T, T)
            scores = scores.masked_fill(~attn_mask, torch.finfo(scores.dtype).min)

        attn = self.dropout(F.softmax(scores, dim=-1))
        out = attn @ V                                  # (B, h, T, d_k)

        # merge heads: (B, h, T, d_k) -> (B, T, d_model)
        out = out.transpose(1, 2).contiguous().view(B, T, self.d_model)
        return self.W_o(out)


def causal_mask(T, device="cpu"):
    """Lower triangular: position i may attend to j <= i."""
    return torch.tril(torch.ones(T, T, dtype=torch.bool, device=device))

The compact version

If asked to write it as a bare function rather than a module, this is the same program in twelve lines — and having it in your fingers is worth more than having the class:

def mha(x, Wq, Wk, Wv, Wo, h, mask=None):
    B, T, d = x.shape
    dk = d // h
    split = lambda t: t.view(B, T, h, dk).transpose(1, 2)
    q, k, v = split(x @ Wq), split(x @ Wk), split(x @ Wv)
    s = q @ k.transpose(-2, -1) / dk ** 0.5
    if mask is not None:
        s = s.masked_fill(~mask, torch.finfo(s.dtype).min)
    o = s.softmax(-1) @ v
    return o.transpose(1, 2).reshape(B, T, d) @ Wo

Note the last line uses reshape, which handles the copy implicitly. That is correct; .view would not be.

The version you should mention, then not write

Since PyTorch 2.0 the fused kernel does all of the above in one call, without ever materialising the $T \times T$ score matrix:

out = F.scaled_dot_product_attention(Q, K, V, is_causal=True)

Say this. It demonstrates you know the production answer. Then write the explicit version, because that is what was asked and because the fused kernel hides exactly the details being probed.

If the follow-up is GQA or MQA

Grouped-query attention lets n_kv_heads < n_heads share keys and values: each KV head serves n_heads / n_kv_heads query heads. The only change to the code is that W_k and W_v project to n_kv_heads * d_k instead of d_model, and K, V are expanded (or, better, broadcast via repeat_interleave on the head axis) before the matmul:

K = K.repeat_interleave(self.n_heads // self.n_kv_heads, dim=1)  # (B, h, T, d_k)
V = V.repeat_interleave(self.n_heads // self.n_kv_heads, dim=1)

The payoff is entirely at inference: the KV cache shrinks by the ratio n_heads / n_kv_heads — for the common 32/8 configuration, 4×. Attention compute is unchanged. That trade is the subject of Modern Architecture Deltas and KV Cache.

4. Numbers and trade-offs

Parameter count, computed rather than quoted. With bias=False, each of the four projections has $d_{m o d e l}^{2}$ weights, so:

$N_{a tt n} = 4 d_{m o d e l}^{2} per layer$

For a 7B-class model with $d_{m o d e l} = 4096$ and 32 layers: $4 \times 409 6^{2} \times 32 \approx 2.1 \times 1 0^{9}$ — about 2.1B of the 7B parameters are attention projections, roughly 30%. Everything else is the feed-forward network and embeddings.

Compute, per layer, per sequence, with $B$ the batch size, $T$ the sequence length, and $d = d_{m o d e l}$ :

Term	FLOPs (forward)	Note
$Q, K, V, O$ projections	$\approx 4 \times 2 BT d^{2}$	independent of $T$ per token
$Q K^{⊤}$	$2 B h T^{2} d_{k} = 2 B T^{2} d$	quadratic in $T$
$(\cdot) V$	$2 B T^{2} d$	quadratic

The projections scale linearly with sequence length; the two attention matmuls scale quadratically. At T = 512 the projections dominate; the two are exactly equal at the crossover T = 2 d = 8192 for $d_{m o d e l} = 4096$ , so attention only dominates above that point. This crossover is why FlashAttention mattered more at long context than at training-time short context, and it is a good answer to “why does long context cost so much”.

Memory for the score matrix, per layer, before any kernel fusion:

$2 \times B \times h \times T^{2} bytes (fp16)$

For $B = 1, h = 32, T = 8192$ : $2 \times 32 \times 819 2^{2} \approx 4.3$ GB per layer, materialised and immediately discarded. This is the number that makes the fused kernel non-optional.

5. Pitfalls

Masking after the softmax. The row no longer sums to 1, so the output is a weighted sum of the wrong scale. The model still trains, which is why this survives to production. Symptom: slightly worse quality with no error, and a suspiciously good training loss if the bug is in the decoder mask — because the leak lets the model see the token it is predicting.

-inf in a fully-masked row. Symptom: NaN loss, often appearing only after many steps, only on the batch that contained a fully-padded sequence. This is the classic “it trained fine for two hours and then went to NaN” bug.

Reshaping heads in the wrong order. x.view(B, T, h, d_k) and x.view(B, T, d_k, h) both run without complaint and produce different results; so does transposing the wrong pair of axes. Nothing raises, the model just learns worse. Symptom: loss decreases but plateaus far above the expected value. Test: see §7.

Forgetting torch.no_grad() / .eval() when checking. Dropout is active by default in training mode, so two identical forward passes return different answers and it looks like a shape bug.

Assuming nn.MultiheadAttention is interchangeable. PyTorch’s built-in defaults to batch_first=False, i.e. (T, B, d_model), treats attn_mask and key_padding_mask with different conventions (a float attn_mask is added, not used as a boolean selector), and applies the scaling differently. Passing your (B, T, d) tensor to it gives a wrong answer or a shape error, not a helpful message.

Scaling with d_model instead of d_k. The variance argument is about the per-head dimension. Dividing by $d_{m o d e l}$ over-shrinks the logits by $h$ , flattening the attention distribution. Symptom: attention weights that stay suspiciously uniform.

6. Using it in a project

You will essentially never hand-write attention in production. What you will do, and what this chapter is the entry point for:

Choosing a checkpoint’s head configuration. n_heads, n_kv_heads, and d_k are properties of the checkpoint, not free parameters. When you evaluate two models for serving, the KV-cache arithmetic from §4 — not the parameter count — is what decides whether a given GPU can hold a given batch.

Deciding whether to write a custom kernel. The answer is almost always no. F.scaled_dot_product_attention already dispatches to FlashAttention or a memory-efficient fallback. Writing your own is justified only when you need a masking pattern the fused kernel does not support.

Reading attention maps for debugging. When retrieval or long-context behaviour is wrong, logging attention entropy per layer is a cheap diagnostic (see Attention §7). This is the one place where the internals are directly useful at work.

When not to care at all. For a 7B model at short context with a sensible serving stack, attention is not your bottleneck. Retrieval quality, prompt shape, and output parsing dominate. Optimising attention without measuring first is the most common wasted week in this field.

7. How you would prove it worked

This is the part candidates skip, and the part that most reliably impresses an interviewer, because it turns “I wrote some code” into “I verified some code”. Four tests, in increasing order of strength:

1. Shape test. Forward a batch with T not divisible by anything convenient and assert the output shape. Catches almost nothing, costs nothing.

2. Equivalence test against a trusted implementation. With dropout disabled and identical weights loaded, compare against F.scaled_dot_product_attention:

torch.manual_seed(0)
x = torch.randn(2, 7, 64, dtype=torch.float64)
m = MultiHeadAttention(64, 4).double().eval()
with torch.no_grad():
    mine = m(x, causal_mask(7))
    q, k, v = (m.W_q(x), m.W_k(x), m.W_v(x))
    # reshape to (B, h, T, d_k) exactly as in the module, then:
    ref = F.scaled_dot_product_attention(
        q.view(2, 7, 4, 16).transpose(1, 2),
        k.view(2, 7, 4, 16).transpose(1, 2),
        v.view(2, 7, 4, 16).transpose(1, 2),
        is_causal=True,
    ).transpose(1, 2).reshape(2, 7, 64)
    ref = m.W_o(ref)
assert torch.allclose(mine, ref, atol=1e-6)

Run it in float64. In float32 the two paths differ by ~1e-6 and the test becomes flaky, which is itself worth saying out loud — it shows you know the difference between an algorithmic mismatch and accumulated floating-point error.

3. Causality test. This is the one that catches the mask-after-softmax bug that the equivalence test cannot, because both implementations would share it if you wrote the reference the same way. Perturb a future token and assert the past outputs are unchanged:

with torch.no_grad():
    y1 = m(x, causal_mask(T))
    x2 = x.clone(); x2[:, -1] += 10.0        # change only the last token
    y2 = m(x2, causal_mask(T))
# positions 0..T-2 must be bit-identical; position T-1 may change
assert torch.equal(y1[:, :-1], y2[:, :-1])

torch.equal, not allclose. A leak that is small enough to pass allclose is still a leak, and this exact test is what the interviewer is hoping you propose.

4. Masking test with a fully-padded row. Construct a batch where one entry is entirely padding and assert the output contains no NaN. This is the test that motivates the finfo.min decision, and it is the difference between knowing the -inf warning and having internalised it.

Reporting template (example — replace with your own). “I verified the implementation three ways: bit-exact agreement with scaled_dot_product_attention in float64, invariance of past positions under perturbation of future tokens, and absence of NaN on a fully-masked batch. The first version passed the equivalence check and failed the causality check — the mask was being applied after the softmax.”

8. Interview drill

“Why the transpose?” Strong answer: the leading axes of a batched matmul are broadcasting axes. Moving the head dimension next to the batch dimension makes the operation (B, h) batched attention calls in a single kernel instead of a Python loop over heads. Common wrong answer: “to make the dimensions match up.” It is a semantic move, not a shape-declaration move, and this answer will draw a follow-up you cannot answer.

“Why divide by $d_{k}$ and not $d_{k}$ ?” Strong answer: the dot product of $d_{k}$ independent unit-variance terms has variance $d_{k}$ , so the standard deviation is $d_{k}$ ; dividing by the standard deviation restores unit variance and keeps the softmax in its non-saturated region. Common wrong answer: “to prevent overflow.” Overflow is a symptom; the reason is the gradient through the softmax.

“How do you apply the causal mask?” Strong answer: add a large negative constant to the scores above the diagonal before the softmax, so those positions get probability zero while the row still sums to 1. Follow-up: “why not -inf?” → a fully-masked row becomes 0/0 = NaN, so use finfo(dtype).min. If you can answer this follow-up you have clearly read the PyTorch docs rather than a blog post. Trap: “multiply the attention matrix by a lower-triangular matrix after the softmax” — that does not renormalise the distribution.

“What does .contiguous() do here?” Strong answer: transpose is a view operation that swaps strides without moving data; view needs a compatible layout, so the tensor must be physically reordered first. reshape would do the copy implicitly. Follow-up: “what is the cost?” → one full copy of the tensor, which is small relative to the matmuls but is why fused kernels keep everything in registers.

“How would you test it?” Strong answer: the four tests in §7 — shape, equivalence in float64 against scaled_dot_product_attention, causality by perturbing a future token with torch.equal, and no-NaN on a fully-masked batch. The causality test is the one that catches the bug the equivalence test shares with a wrong reference. Common wrong answer: “check the output shape and that it runs.” That tests that the code does not crash, not that it is correct — and a wrong head-ordering runs perfectly.

“How does this differ from nn.MultiheadAttention?” Strong answer: batch_first=False by default, a different mask convention (a float attn_mask is added rather than used as a boolean selector), a different internal scaling path, and no guarantee of weight-layout compatibility with a checkpoint trained against another implementation.

“Add GQA.” See §3. The expected answer names the two lines that change, the 4× KV-cache saving for a 32/8 configuration, and the fact that attention compute is unchanged.

Next: The Interview Playbook — how to deploy all of this in an actual interview.

Keyboard shortcuts

Henry's Notebook