33. Batched Sliding-Window Causal Mask Construction

Problem

Real batches are padded to a common sequence length, so a production sliding-window attention layer needs a mask that combines three separate conditions at once: causality (never attend to the future), the local window (never attend further than window_size back), and padding (never attend to a key position that isn't a real token). Getting this mask right — and building it with pure tensor broadcasting rather than a Python loop over every (batch, query, key) triple — is the piece that most often trips people up when they otherwise understand sliding-window attention's math fine, because there are now three conditions to AND together instead of one.

Given seq_len, window_size, and an optional per-batch padding mask, build the combined boolean attention mask: True means "this query may attend to this key."

Source: src/33_sliding_window_causal_mask.py

def sliding_window_causal_mask(
    seq_len: int, window_size: int, pad_mask: torch.Tensor | None = None,
) -> torch.Tensor: ...

>>> sliding_window_causal_mask(5, window_size=3).shape
torch.Size([1, 5, 5])

Step-by-Step Approach

1. Build the pairwise position grid: i = arange(seq_len).view(seq_len, 1), j = arange(seq_len).view(1, seq_len) — broadcasting these against each other gives every (query, key) index pair at once.
2. Combine causal and window conditions: base = (j <= i) & ((i - j) < window_size), an (seq_len, seq_len) boolean grid.
3. If no pad_mask is given, return base.unsqueeze(0) — shape (1, seq_len, seq_len), which broadcasts cleanly against any batch size downstream.
4. If a pad_mask (shape (batch, seq_len), True = real token) is given, reshape it to (batch, 1, seq_len) so it broadcasts across the query axis — a padded position must be unattendable as a key regardless of which query is asking.
5. AND the padding condition into the base mask: base.unsqueeze(0) & key_ok, giving the final (batch, seq_len, seq_len) mask.
6. Sanity-check with window_size >= seq_len and no padding: the result must equal torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool)) exactly, i.e. plain causal masking.

The key insight is that every one of these conditions is expressible as a broadcastable boolean comparison over index tensors, so the entire mask — for any batch size, any window, with or without padding — comes from a handful of vectorized comparisons ANDed together, never a loop.

Reference solution

import torch


def sliding_window_causal_mask(
    seq_len: int, window_size: int, pad_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    idx = torch.arange(seq_len)
    i, j = idx.view(seq_len, 1), idx.view(1, seq_len)
    # causal AND local-window, purely via broadcasting -- an (seq, seq) grid,
    # no python loop over query/key pairs
    base = (j <= i) & ((i - j) < window_size)  # (seq_len, seq_len)

    if pad_mask is None:
        return base.unsqueeze(0)  # (1, seq_len, seq_len), broadcastable over any batch

    # a padded key position must never be attended to, regardless of the
    # causal/window conditions -- AND it in as a third broadcast condition
    batch = pad_mask.shape[0]
    key_ok = pad_mask.view(batch, 1, seq_len)  # True where the key is a real (non-pad) token
    return base.unsqueeze(0) & key_ok          # (batch, seq_len, seq_len)

Key Functions & Tricks

• torch.arange(seq_len).view(seq_len, 1) / .view(1, seq_len) — the standard "outer comparison" trick, forms an (N, N) grid from two 1D ranges via broadcasting.
• pad_mask.view(batch, 1, seq_len) — the inserted middle dimension of size 1 is what lets a per-batch, per-key mask broadcast identically across every query position.
• Boolean & across broadcastable shapes — composes independent conditions (causal, window, padding) without any control flow.
• tensor.unsqueeze(0) — adds the batch axis so the same mask-construction code path works whether or not a padding mask is supplied.
• torch.tril(torch.ones(...)) — used in the test suite as the known-correct reference for "plain causal masking," the degenerate case when the window covers the whole sequence.

How to Recognize This Pattern