Sliding Window Attention - Transformer Architecture

Sliding Window Attention Mask

What you'll build

A causal mask that also limits how far back each token can look: position i may attend only to the window_size most recent positions, itself included. Mistral and Gemma 2/3 (in alternating layers) use this pattern.

Background

Full causal attention costs O(n^2) time and the KV cache grows with the whole context. With a window of W tokens, each query reads at most W keys, so attention costs O(n · W) and the cache can be a fixed-size ring buffer of W entries. Information still travels further than W: stacking layers lets a token see about L · (W - 1) positions back after L layers, like the receptive field of stacked convolutions.

Formula

allowed(i, j) = (j <= i) and (i - j < window_size)
mask[i, j]    = 0 if allowed(i, j) else -inf

For seq_len = 5, window_size = 3:

[[   0, -inf, -inf, -inf, -inf],
 [   0,    0, -inf, -inf, -inf],
 [   0,    0,    0, -inf, -inf],
 [-inf,    0,    0,    0, -inf],
 [-inf, -inf,    0,    0,    0]]

Inputs

• seq_len: positive int
• window_size: int >= 1, number of visible positions including the token itself

Expected Return Format

Float array (seq_len, seq_len) with 0.0 for allowed and -inf for masked entries (rows = queries, columns = keys).

Edge cases and conventions

• window_size = 1: only the diagonal is allowed.
• window_size >= seq_len: identical to the plain Causal Attention Mask.
• The diagonal is always allowed, so no row is fully masked.

Constraints

NumPy only, no loops (broadcast row and column index arrays).

Follow-up

• How would you implement the KV cache for a sliding window model as a ring buffer?
• Off-by-one check: with a window of 4096, can token 5000 see token 904? (5000 - 904 = 4096, which is not < 4096, so no.)

Input: sliding_window_mask(4, 2)
Output: [[0, -inf, -inf, -inf], [0, 0, -inf, -inf], [-inf, 0, 0, -inf], [-inf, -inf, 0, 0]]

Input: sliding_window_mask(3, 10)
Output: [[0, -inf, -inf], [0, 0, -inf], [0, 0, 0]]

import numpy as np

def sliding_window_mask(seq_len, window_size):
    """
    Create a sliding window causal attention mask.

    Args:
        seq_len: Sequence length
        window_size: Number of positions each token can attend to,
                     including itself

    Returns:
        mask: (seq_len, seq_len), 0.0 where j <= i and i - j < window_size,
              -inf elsewhere
    """
    # Your code here
    pass