A forward pass enters multi-head attention with input shape [batch=4, seq=512, d_model=1024] and num_heads=16. After the Q and K projections are reshaped to per head form and Q @ K^T runs, a debugger breakpoint fires right before softmax. Predict the full shape of the attention score tensor at that breakpoint.
Q and K reshape to [4, 16, 512, 64]; Q @ K^T contracts the head_dim axis to give the pre-softmax score tensor [4, 16, 512, 512], one square matrix per head per batch.
Think of attention as filling out a square seating chart for every guest. There are sixteen separate planners (the heads) each filling their own chart, and you are running four parties in parallel (the batch). Each planner asks every guest to rate every other guest, producing a 512 by 512 table of scores. Stack the tables: four parties, sixteen planners per party, each with a 512 by 512 table, and you get a four by sixteen by 512 by 512 box of numbers. That box is what enters softmax. Each planner uses 64 features (the head dim) to come up with each score, but those features collapse during the matrix multiply.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Concept explanation~2 min read
Everything you need to truly understand this topic: intuition, mechanics, step by step explanation, code, formulas, and worked example. Click to expand.
Tracing the shape of a tensor through a transformer block is the single best diagnostic skill for a transformer engineer. The attention score tensor specifically is where most candidates trip, because the per-head reshape and the K transpose move axes in ways that are easy to get wrong by one.
This deep dive walks through the [4, 16, 512, 512] answer in detail: where each axis comes from, what each axis means, and what changes in modern variants (GQA, FlashAttention, sliding windows). By the end the trace should be muscle memory.
From [B, T, d_model] to per-head [B, H, T, d_h]
The block input is shaped [B, T, d_model] = [4, 512, 1024]. The Q, K, and V projections are three linear layers, each mapping from d_model to d_model. After each projection, the tensor still has shape [4, 512, 1024]. No reshape yet.
Multi-head attention computes attention independently in H parallel subspaces. The shape transformation is a pure reshape and transpose: split the d_model axis into [H, d_h] where d_h = d_model / H. For our case, d_h = 1024 / 16 = 64.
In PyTorch this is two lines:
q = q.view(B, T, H, d_h) # [4, 512, 16, 64]
q = q.transpose(1, 2) # [4, 16, 512, 64]
The transpose puts the head axis before the seq axis. This makes the per-head matmul a batched matmul over [B, H], which maps cleanly to the GPU's batched-GEMM kernels. Q, K, V all end up shape [B, H, T, d_h] = [4, 16, 512, 64].
A common mistake is to forget the transpose and run the matmul on the wrong axes. The math still type-checks but the heads bleed into each other, and the model is silently broken. Always trace through the transpose explicitly.
Situations where this technique stops working.
2–4 min · Everything important, quickly.
Real products, models, and research that use this idea.
- PyTorch's `torch.nn.functional.scaled_dot_product_attention` operates on inputs shaped [B, H, T, d_h] and returns the same, internally dispatching to FlashAttention when available
- Llama 3 8B at d_model 4096 with 32 heads has head_dim 128, so a 4 x 32 x 8192 x 8192 score tensor in fp16 is 32 GB per layer, the reason FlashAttention is mandatory at production context lengths
What an interviewer would ask next. Try answering before peeking at the approach.
QHow does the score shape change for GQA with 32 query heads and 8 KV heads?
Q is [B, 32, T, d_h]. K and V are [B, 8, T, d_h]. Before the matmul, K is broadcast or repeated to [B, 32, T, d_h], so the score still ends up [B, 32, T, T]. The shape after broadcast is unchanged; the memory savings live in the KV cache, not in the score tensor.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Red flags & common mistakes
The phrases that signal junior thinking. Click to expand.
Leaving the head_dim axis in the output, or forgetting that the K transpose only swaps the last two axes, not all axes, so the contraction is on head_dim and not on seq.
60 second bullets to scan on the way to the call.
How d_model splits into heads x head_dim and what that means for the reshape
Why the K transpose is on the last two axes only, not a full axis permutation
Primary sources. Browse if you want the original framing.
Same topic, related formats. Practice these next.