50. Causal Masking for Autoregressive Attention

Problem

An autoregressive (decoder-only) transformer must never let a token attend to a token that comes after it -- otherwise the model could "see the future" at training time and would be useless for generation, where later tokens don't exist yet. This is enforced with a causal mask: before the softmax, every attention score for a (query position i, key position j) pair with j > i is set to -inf, so it contributes exactly 0 attention weight after softmax. Real batches also mix in padding: shorter sequences are padded to the batch's max length, and no position -- past or future -- should ever attend to a padding key. Both masks apply at once: a key is masked out if it's in the future OR if it's padding.

Implement scaled dot-product attention with a causal mask combined with an optional key-padding mask.

Source: src/50_causal_attention_mask.py

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

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

Step-by-Step Approach

1. Compute raw attention scores: q @ k.transpose(-2, -1), giving shape (batch, heads, seq, seq).
2. Scale by 1/sqrt(head_dim) so the softmax input's variance doesn't grow with head_dim -- without this, softmax saturates for large head dimensions and gradients vanish.
3. Build the causal mask with torch.triu(ones(seq, seq, dtype=bool), diagonal=1) -- True strictly above the diagonal, exactly the j > i "future" positions.
4. Broadcast the causal mask from (seq, seq) up to (batch, heads, seq, seq) so it can be combined with the per-batch padding mask.
5. If a padding_mask is given, invert it (True for real tokens becomes True for padding) and reshape it to (batch, 1, 1, seq) so it broadcasts over every query position and head, then OR it into the combined mask -- a key must never be attended to if it's a padding position, regardless of which query is asking.
6. Apply the combined mask with scores.masked_fill(mask, float("-inf")), softmax over the last dimension, then multiply by v.

The key insight is that the two masks combine with a logical OR, not an AND -- a key position is excluded if it fails either constraint, and using -inf rather than a large negative number guarantees exactly zero attention weight regardless of the other unmasked scores in that row.

Reference solution

def causal_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    padding_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    batch, num_heads, seq_len, head_dim = q.shape

    # (batch, heads, seq, seq) raw scores, scaled by 1/sqrt(head_dim) so
    # softmax input variance doesn't grow with head_dim.
    scores = (q @ k.transpose(-2, -1)) / math.sqrt(head_dim)

    # triu(diagonal=1): True strictly above the main diagonal, i.e. exactly
    # the (query i, key j) pairs with j > i -- the "future" positions.
    causal = torch.triu(
        torch.ones(seq_len, seq_len, dtype=torch.bool, device=q.device), diagonal=1
    )
    mask = causal.expand(batch, num_heads, seq_len, seq_len)

    if padding_mask is not None:
        # padding_mask is True for real tokens; a KEY at a padding position
        # must never be attended to by ANY query, so invert it and
        # broadcast: (batch, seq) -> (batch, 1, 1, seq).
        key_is_padding = ~padding_mask[:, None, None, :]
        mask = mask | key_is_padding

    # -inf guarantees exactly 0 weight after softmax, independent of the
    # other scores in that row.
    scores = scores.masked_fill(mask, float("-inf"))
    weights = torch.softmax(scores, dim=-1)

    return weights @ v

Key Functions & Tricks

• torch.triu(ones, diagonal=1) — the standard one-liner for a strictly-upper-triangular boolean causal mask; diagonal=1 excludes the main diagonal, so a position can always attend to itself.
• tensor.expand(...) — broadcasts a smaller mask up to the full batch/head shape without copying memory, since expand creates a view, not a new allocation.
• padding_mask[:, None, None, :] — inserts two new size-1 dimensions so a (batch, seq) mask broadcasts correctly against a (batch, heads, seq, seq) score tensor.
• scores.masked_fill(mask, float("-inf")) — sets masked positions to -inf so softmax assigns them exactly zero weight.
• mask_a | mask_b — boolean OR combines two masks so a position is excluded if it violates either constraint.

How to Recognize This Pattern