35. Scaled Dot-Product Attention From Scratch

Problem

Every transformer-based model an interviewer might ask about next — causal decoders, cross-attention, ALiBi/relative-position variants — is built on one primitive: scaled dot-product attention, attention(Q, K, V) = softmax(QKT / √d_k + bias) V. Before touching multi-head splits or KV-caching, you need to be able to write the two-matmul-plus-softmax core cold, with the scaling and masking made explicit rather than hidden behind a single fused call.

This version uses the additive masking convention rather than a boolean keep/block mask: a bias tensor is added directly to the raw scores before softmax. That's the more general convention — a hard block is just -inf in one entry, but the same mechanism also expresses soft relative-position penalties (ALiBi-style) with finite values, which a boolean mask cannot.

The 1/√d_k scale factor matters for more than convention: if each component of q and k is roughly unit-variance and independent, their dot product has variance that grows linearly with d_k. Without rescaling, scores blow up in magnitude for large d_k, softmax saturates onto one logit, and gradients through the softmax vanish almost everywhere.

Source: src/35_scaled_dot_product_attention.py

def scaled_dot_product_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, attn_bias: torch.Tensor | None = None
) -> torch.Tensor: ...

>>> q = torch.randn(1, 2, 4)
>>> k = torch.randn(1, 3, 4)
>>> v = torch.randn(1, 3, 3)
>>> scaled_dot_product_attention(q, k, v).shape
torch.Size([1, 2, 3])

Step-by-Step Approach

1. Compute raw similarity scores between every query and every key: q @ k.transpose(-1, -2), giving shape (..., seq_q, seq_k).
2. Scale by 1 / sqrt(d_k), where d_k is the last dimension of q (and k).
3. If attn_bias is given, add it directly to the scores before softmax — never zero the post-softmax weights instead, that breaks the sum-to-1 invariant rather than cleanly blocking a position.
4. Apply softmax along the last axis (the key axis) so every query's attention weights form a probability distribution over keys.
5. Multiply the attention weights by v: attn_weights @ v, giving shape (..., seq_q, d_v) — each output row is a weighted average of value vectors.

The key insight is that attention is two matrix multiplies with a normalization step sandwiched between them: the first matmul turns similarity into a score, softmax turns scores into a distribution, and the second matmul turns that distribution into a weighted blend of values — everything else (masks, multiple heads, caching) is decoration around this core.

Reference solution

import math
import torch


def scaled_dot_product_attention(
    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, attn_bias: torch.Tensor | None = None
) -> torch.Tensor:
    d_k = q.shape[-1]
    # (..., seq_q, d_k) @ (..., d_k, seq_k) -> (..., seq_q, seq_k)
    scores = q @ k.transpose(-1, -2) / math.sqrt(d_k)
    if attn_bias is not None:
        # additive convention: -inf entries hard-block (exp(-inf) == 0 after
        # softmax), finite entries just tilt the distribution
        scores = scores + attn_bias
    attn_weights = torch.softmax(scores, dim=-1)
    # (..., seq_q, seq_k) @ (..., seq_k, d_v) -> (..., seq_q, d_v)
    return attn_weights @ v

Key Functions & Tricks

• tensor.transpose(-1, -2) — swaps the last two dims so the key matmul contracts over d_k instead of the sequence axis; works regardless of how many leading batch/head dims are present.
• math.sqrt(d_k) as a Python float — cheaper than a tensor op and broadcasts fine against the score tensor.
• Additive bias (scores + attn_bias) — the same mechanism handles hard -inf masks and soft finite biases (ALiBi, relative position) with one code path.
• torch.softmax(scores, dim=-1) — internally uses the max-subtraction trick, so it stays numerically stable even with -inf entries in the row.
• Batched @ (torch.matmul) — broadcasts over any number of leading batch/head dimensions automatically, as long as the last two dims align for matrix multiply.
• torch.nn.functional.scaled_dot_product_attention — PyTorch's own fused kernel; a solid oracle to sanity-check a from-scratch implementation against.

How to Recognize This Pattern