Self-Attention and Multi-Head Attention

1. Quick anchor

Self-attention is a learned, content-addressable lookup that runs once per token, in parallel, across the whole sequence. Each token emits a query ("what am I looking for?"), every token exposes a key ("what do I offer?") and a value ("here's my actual content"). A token's new representation is a weighted average of all values, where the weights come from how well its query matches each key. That's the entire mechanism: softmax(QKᵀ/√d)V. Multi-head attention just runs several of these lookups in parallel subspaces so the model can simultaneously track grammar, coreference, and long-range structure. Everything hard about attention in production — the quadratic cost, FlashAttention, GQA, the KV cache — is downstream of this one equation.

2. Why interviewers probe this

• IC4 — Can you write attention from the formula without a library, explain the shapes, and say why the √d scaling exists (not just that it does)? Do you understand the causal mask well enough to not leak the future?
• IC5 — Do you know that attention is memory-bandwidth bound, not FLOP bound, on real GPUs? Can you explain what FlashAttention actually changes (tiling + online softmax, never materializing the N×N matrix) and why the output is bit-for-bit equivalent? Can you reason about the O(n²) wall for long context?
• IC6 — Can you connect the mechanism to a serving cost model? KV-cache sizing, GQA/MQA tradeoffs, prefill vs. decode being two different bottlenecks, FP8 attention on Hopper, and when you'd reach for a different attention variant entirely. You should be opinionated about what the current architecture consensus is and where it breaks.

3. Concept build-up

3.1 The three projections: why Q, K, and V are different

Start with a sequence of n token vectors stacked into a matrix X ∈ ℝ^{n×d_model}. A naive idea: let token i attend to token j by x_i · x_j. That fails for a subtle reason — it's symmetric and self-obsessed. The dot product x_i · x_j equals x_j · x_i, so "how much should i read from j" would be forced equal to "how much j reads from i", and every token would attend maximally to itself. Language isn't symmetric: in "the cat that she owned", cat and owned relate asymmetrically.

The fix is three separate learned projections:

where $W^Q, W^K \in \mathbb{R}^{d_{model}\times d_k}$ and $W^V \in \mathbb{R}^{d_{model}\times d_v}$. Now the score q_i · k_j is asymmetric (different weight matrices), and V is decoupled from the matching: the thing we read out (value) is separate from the thing we match on (key). A token can advertise itself as a verb (its key) while carrying tense information (its value). That decoupling is the whole reason attention is expressive.

3.2 Scaled dot-product attention and the √d question

The full operation, from Vaswani et al. (2017):

Symbols: Q ∈ ℝ^{n×d_k} (queries), K ∈ ℝ^{n×d_k} (keys), V ∈ ℝ^{n×d_v} (values), d_k the per-head key dimension. QKᵀ ∈ ℝ^{n×n} is the score matrix — entry (i,j) is q_i · k_j. Softmax is applied row-wise, so each row is a probability distribution over the n keys.

Why divide by √d_k? This is the canonical IC4 question and there's a clean first-principles answer. Suppose the entries of q and k are independent with mean 0 and variance 1. Their dot product is a sum of d_k independent products: $q\cdot k = \sum_{m=1}^{d_k} q_m k_m$. Each product term has mean 0 and variance 1, so the sum has variance d_k and standard deviation √d_k. As d_k grows (64, 128…), raw scores have magnitude ~√d_k, which can be ±10 or more. Feed numbers that large into softmax and it saturates: one entry goes to ≈1, the rest to ≈0. Softmax's gradient is p_i(δ_ij − p_j); when p is one-hot, that gradient is ≈0 everywhere. Training stalls. Dividing by √d_k rescales the dot product back to unit variance, keeping softmax in its responsive, well-conditioned region. It's variance normalization, full stop.

3.3 The causal mask

For a decoder generating left-to-right, token i must not see tokens j > i, or it would cheat during training (the next token is right there in the input). We enforce this by adding −∞ to the upper triangle of the score matrix before softmax:

exp(−∞) = 0, so future positions get exactly zero weight after softmax, and each row still normalizes correctly over the visible past. The elegance: this masking is what lets us train in parallel. We feed the whole sequence at once, compute all n positions' losses simultaneously, and the mask guarantees position i's prediction only used positions ≤ i. No recurrence, no sequential bottleneck during training — that parallelism is the reason transformers scaled when RNNs couldn't.

3.4 Multi-head attention

One attention function is a single "lookup pattern". Real language needs many at once — subject-verb agreement, coreference, positional locality. Multi-head attention runs h independent attention functions in lower-dimensional subspaces:

Crucially d_k = d_model / h, so the heads partition the dimensionality rather than multiplying total compute — h heads of size d_model/h cost about the same as one head of size d_model, but give the model h distinct attention patterns. The final W^O ∈ ℝ^{d_model×d_model} mixes the concatenated heads back into the residual stream. Typical configs: 8–32 heads, d_head of 64–128. Empirically different heads specialize (induction heads, positional heads, "previous-token" heads), though specialization is messy and heads are partly redundant — which is exactly what GQA later exploits.

3.5 Complexity: the O(n²) wall

The score matrix QKᵀ is n×n. Both compute and memory are quadratic in sequence length: O(n²·d) FLOPs and O(n²) memory for the scores. At n = 100k tokens, the score matrix alone is 10 billion entries per head per layer. This is the central scaling pain of transformers and the reason for everything in §4–§5: FlashAttention (avoid materializing the matrix), positional interpolation / RoPE scaling (extend context cheaply), and KV caching (avoid recomputing keys/values at decode). Naive attention doesn't change the asymptotic FLOPs; it changes whether you survive the memory traffic.

4. Minimal implementation

The entire mechanism is about six load-bearing lines. Full runnable file: examples/transformers/attention.py.

import math, torch
import torch.nn.functional as F
 
def scaled_dot_product_attention(q, k, v, causal=False):
    # q: (..., n_q, d_k)  k: (..., n_kv, d_k)  v: (..., n_kv, d_v)
    d_k = q.size(-1)
    scores = q @ k.transpose(-2, -1) / math.sqrt(d_k)        # (..., n_q, n_kv)
    if causal:
        n_q, n_kv = scores.shape[-2:]
        mask = torch.ones(n_q, n_kv, dtype=torch.bool,
                          device=q.device).triu(1)           # True above diagonal
        scores = scores.masked_fill(mask, float("-inf"))     # future -> -inf
    weights = scores.softmax(dim=-1)                         # rows sum to 1
    return weights @ v                                       # (..., n_q, d_v)

Read it against §3.2: q @ k.transpose is QKᵀ, the / math.sqrt(d_k) is the variance fix, .triu(1) builds the strictly-upper-triangular causal mask, softmax(dim=-1) normalizes each query's row over keys, and the final weights @ v is the weighted value-average. The (...) leading dims let the same function handle (batch, heads, seq, d_head) tensors unchanged — multi-head is just batching this over a head axis.

Multi-head adds projections around it. The file's MultiHeadAttention splits d_model into (n_heads, d_head), runs the function above, merges, and applies o_proj. It also implements GQA: project fewer KV heads than Q heads and repeat_interleave them, which is the single change that shrinks the KV cache (§5). The file's _check() asserts our output matches PyTorch's fused F.scaled_dot_product_attention to 1e-5, confirming the hand-written version is the same math the optimized kernel computes:

ours = scaled_dot_product_attention(q, k, v, causal=True)
ref  = F.scaled_dot_product_attention(q, k, v, is_causal=True)
assert torch.allclose(ours, ref, atol=1e-5)   # FlashAttention == naive, numerically

In production you never call the naive version — F.scaled_dot_product_attention dispatches to a FlashAttention kernel. You write the naive one in interviews and to understand what the kernel is approximating (it isn't approximating — it's exact, just IO-aware).

5. Production tradeoffs

The core production fact: attention is memory-bandwidth bound, not FLOP bound. Standard attention reads and writes the N×N score matrix to GPU HBM (high-bandwidth memory), and HBM is ~10–20× slower than on-chip SRAM. FlashAttention (Dao et al., 2022) recognizes this and tiles Q, K, V into blocks that fit in SRAM, computing softmax incrementally ("online softmax") so the full N×N matrix is never materialized. Same output, far less HBM traffic: HBM reads/writes drop from O(n² + n·d) to roughly O(n·d). Reported: 7.6× speedup on attention, 15% on BERT-large, enabling 64K-token contexts. FlashAttention-2 fixed work partitioning (Q in the outer loop) for 2× more; FlashAttention-3 (2024) added Hopper async + FP8.

At decode time the bottleneck flips to the KV cache. During generation you cache every past token's keys and values so each new token does O(n) work instead of O(n²). Cache size = 2 × n_layers × n_kv_heads × d_head × seq_len × bytes per sequence. For a 70B model at 128K context this runs into tens of GB per request — it dominates memory and caps your batch size (throughput). GQA is the lever: use, say, 8 KV heads for 64 Q heads, cutting the cache 8× with minimal quality loss because heads were partly redundant anyway. Llama-3 and most 2024–2026 production decoders ship GQA. MQA (1 KV head) shrinks it further but usually costs measurable quality unless trained or distilled in.

What changes at scale: prefill and decode are two different machines. Prefill (processing the prompt) is compute/FLOP-heavy and parallel — FlashAttention wins here. Decode (one token at a time) is memory-bound on the KV cache — GQA and quantized KV win here. Mature inference stacks (vLLM, TensorRT-LLM) optimize them separately, with paged KV caches and continuous batching. See /inference/kv-cache and /inference for the serving side.

6. How it's asked

7. Pitfalls & flashcards

• Masking after softmax instead of before. The −∞ must be added to the scores, pre-softmax, so exp zeroes them and the row renormalizes. Masking the weights afterward breaks normalization.
• Forgetting the scale at large d_k. Skipping /√d_k works at toy sizes (d=2) and silently saturates at d=128. It's a "passes the smoke test, fails at scale" bug.
• Confusing FLOP-bound and bandwidth-bound. Attention is bandwidth-bound; the FlashAttention win is IO, not arithmetic. Saying "it does fewer FLOPs" is wrong and a red flag in interviews.
• Treating prefill and decode the same. Prefill is compute-bound and parallel; decode is memory-bound on the KV cache. Different optimizations apply.
• Assuming more heads = strictly better. Heads are partly redundant; that redundancy is why GQA works. Head count is a capacity/compute tradeoff, not a free quality dial.
• Numerical softmax. Always subtract the row max before exp (PyTorch does this internally; your hand-rolled numpy version must too) or you'll overflow.

Flashcard. Attention = softmax(QKᵀ/√d)V: query·key gives relevance, √d keeps softmax unsaturated, the mask hides the future, the value is what you read out. Multi-head runs it in parallel subspaces; FlashAttention makes it IO-cheap; GQA makes its KV cache serveable.

8. Further reading

• Vaswani et al., 2017 — Attention Is All You Need — the original formulation of scaled dot-product and multi-head attention. https://arxiv.org/abs/1706.03762 (NeurIPS: https://papers.nips.cc/paper/2017/file/3f5ee243-7b1b-4501-ba6c-0ff5015a3da0-Paper.pdf)
• Dao et al., 2022 — FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — the tiling + online-softmax algorithm and the bandwidth-bound insight. https://arxiv.org/abs/2205.14135
• Dao et al., 2024 — FlashAttention-3 — Hopper async Tensor Cores, TMA, and FP8 attention. https://arxiv.org/abs/2407.08608
• Su et al., 2021 — RoFormer (RoPE) — how relative position rides on top of QK, prerequisite for long-context scaling. https://arxiv.org/abs/2104.09864
• Holtzman et al., 2020 — The Curious Case of Neural Text Degeneration — what happens after attention, when you decode from the logits. https://arxiv.org/abs/1904.09751

Next: /transformers/positional-encodings — attention is permutation-invariant by itself; positional encodings are how the model learns where tokens are.