20. Causal Masking for Autoregressive Attention

Problem

An autoregressive decoder generates one token at a time, but at training time it processes a whole sequence in parallel for efficiency. Position i's output must only ever depend on positions ≤ i, never on anything later — otherwise the model "cheats" by peeking at the answer during training and produces garbage at actual generation time when position i+1 doesn't exist yet.

Implement self-attention (q, k, and v all derived from the same sequence, so seq_len_q == seq_len_k) with a causal mask applied internally, so that attend(x)[i] is provably independent of x[j] for any j > i.

Source: src/20_causal_attention_mask.py

def causal_self_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor: ...

>>> q = k = v = torch.randn(1, 4, 8)
>>> out = causal_self_attention(q, k, v)
>>> out.shape
torch.Size([1, 4, 8])
>>> v2 = v.clone(); v2[:, 3, :] = 999.0
>>> torch.allclose(out[:, 0], causal_self_attention(q, k, v2)[:, 0])
True

Step-by-Step Approach

1. Compute raw scaled scores as usual: q @ k.transpose(-1, -2) / sqrt(d_model), shape (batch, seq, seq).
2. Build a lower-triangular boolean matrix with torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool)) — entry (i, j) is True exactly when j ≤ i, i.e. key position j is allowed for query position i.
3. Apply it with masked_fill, setting every disallowed (i, j) with j > i to -inf before softmax, same as any other attention mask.
4. Softmax over the key axis as usual — because every masked entry is -inf, row i's distribution only has nonzero mass on keys 0..i.
5. Multiply by v to get the output; verify the invariant directly rather than trusting the mask by eye — e.g. corrupt v at the last position and confirm every earlier output row is bit-for-bit unchanged.

The key insight is that causal masking is not a different attention mechanism, it's the same softmax(QKT/√d_k)V with one extra, sequence-length-dependent mask applied before softmax — the triangular structure is exactly what makes parallel training and strictly left-to-right generation produce identical per-position computations.

Reference solution

import math
import torch


def causal_self_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
    seq_len, d_model = q.shape[-2], q.shape[-1]

    scores = q @ k.transpose(-1, -2) / math.sqrt(d_model)  # (batch, seq, seq)

    # lower-triangular (including diagonal) = True means "position i may
    # attend to position j" only when j <= i
    causal_mask = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=q.device))
    scores = scores.masked_fill(~causal_mask, float("-inf"))

    attn_weights = torch.softmax(scores, dim=-1)
    return attn_weights @ v

Key Functions & Tricks

• torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool)) — the standard one-liner for a lower-triangular "may attend" mask; torch.triu with diagonal=1 gives the complementary "must block" set directly if preferred.
• device=q.device on the mask construction — avoids a silent CPU/GPU mismatch error when q lives on an accelerator.
• masked_fill(~causal_mask, float("-inf")) — same pre-softmax masking trick as plain attention; the boolean is inverted because tril marks the allowed set, not the blocked one.
• Corrupt-and-compare invariant check (mutate a late value, assert earlier outputs are untouched) — a much stronger correctness test than just checking output shape, since a shape-correct-but-wrong mask (e.g. off-by-one on the diagonal) would still pass a shape check.

How to Recognize This Pattern