Implementing attention from scratch

Why it matters

"Implement attention" is the research-engineering analogue of "reverse a linked list": short to state and easy to get subtly wrong. Bugs hide in reshapes that mix heads with positions, masks with inverted conventions, softmax overflow in half precision, and caches whose positions drift by one. Because a buggy output is still a plausible tensor, the only real defence is testing against a trusted reference. Getting this right also makes FlashAttention, PagedAttention and serving metrics readable instead of magical.

Core ideas

Shapes first. With input x of shape (B, T, C) and H heads, project to Q, K and V of shape (B, T, C), reshape to (B, T, H, D) with D = C/H, and transpose to (B, H, T, D). Scores QKᵀ have shape (B, H, T, S), where S is the number of key positions (S = T without a cache). Weights times V gives (B, H, T, D); transpose back and merge heads into (B, T, C), using reshape or contiguous() because the transposed tensor is not contiguous. Writing the shape after every line is the most effective habit you can build.

Why scale by √D. If query and key components are roughly independent with unit variance, their dot product has variance D. Unscaled scores grow with head size and push the softmax into saturated regions where gradients vanish.

Masks. A causal mask lets position i attend only to positions j ≤ i; a padding mask removes keys that are padding. Conventions differ: in PyTorch's scaled_dot_product_attention a boolean True means "may attend", while nn.MultiheadAttention's key_padding_mask uses True for "ignore". Additive masks use 0 and −inf. If every key in a row is masked, softmax over all −inf is NaN, so padded query rows need explicit handling or must be excluded from the loss.

Stable softmax. Subtract the row maximum before exponentiating; the result is mathematically identical and exp can no longer overflow. The online variant processes keys in blocks while keeping a running maximum and a running denominator, rescaling earlier partial results by exp(m_old − m_new) whenever the maximum grows. Tiled kernels such as FlashAttention are built on exactly this step.

KV caching. During generation, keys and values of past tokens never change, so store them per layer and compute projections only for new tokens. A decode step then attends one new query over S cached positions instead of recomputing the whole prefix. Cache memory per token is 2 × layers × KV heads × D × bytes per element; multi-query and grouped-query attention shrink it by sharing KV heads across query heads. Positions (absolute indices or rotary angles) must use the true offset: cache length plus position within the new chunk. When queries and keys differ in length, a causal mask must be aligned to the bottom-right; PyTorch's is_causal flag aligns it to the top-left, so a single new query would see only the first cached key.

Testing against a reference. Compare with scaled_dot_product_attention (its documentation includes a plain reference formulation) or with a naive float64 version. Use torch.testing.assert_close with dtype-appropriate tolerances, random shapes including length 1 and odd lengths, ragged padding, and gradient comparisons. For caching, assert that logits from incremental decoding equal those of a full forward pass at every step.

float64 random (B,H,T,D)  -> match reference to ~1e-10
float32 / bf16            -> looser tolerances; watch error growth with T
edge cases                -> T = 1, odd T, fully padded row, causal + padding
cache test                -> incremental logits == full-forward logits per step

Cost model. Scores and weights are T × S per head, so naive training-time attention needs memory quadratic in sequence length, and the two matrix products cost about 4·B·H·T·S·D FLOPs. Decoding is memory-bandwidth bound: every step reads the whole cache to produce one token.

Common misconceptions

• "view() after transpose is fine." The transposed tensor is not contiguous; use reshape or contiguous(), and check you are not interleaving heads with positions.
• "Masking with −inf is always safe." A fully masked row yields NaN; handle padding rows explicitly.
• "is_causal=True works for cached decoding." Its implicit mask is top-left aligned when query and key lengths differ, which is wrong after a cache.
• "Matching on one random input proves correctness." Test lengths, padding, dtypes and gradients, not one happy path.
• "A KV cache changes outputs slightly; that is expected." In exact arithmetic it changes nothing, and float32 differences should be tiny; larger gaps mean position or mask bugs.
• "Dropout is off in eval mode automatically." The functional API applies whatever dropout probability you pass; set it to zero outside training.

How to approach it

Write a naive float64 version first, then a batched multi-head version, then masks, then caching, each tested against the previous stage and the reference. Use nanoGPT and the Annotated Transformer as second opinions, not as code to copy. In interviews, state shapes aloud, mention the stable softmax unprompted and propose the cache-equivalence test; senior candidates are expected to discuss cache memory and why decoding is bandwidth bound.