OLMo 3's Sliding-Window Attention Pattern

The problem

Full attention lets every token look at all the tokens before it, and its cost grows with the square of the text length. OLMo 3 cuts the cost with sliding-window attention in three layers out of four: a token attends only to the most recent window tokens, itself included. Every fourth layer uses full causal attention, and the last layer is always full. In the released configurations the window is 4,096 tokens.

You will build the masks that implement this. A mask is a table of True (allowed) and False (forbidden), with one row per query position $t$ and one column per key position $s$ (section 7).

Which layers are full

Layer $\ell$, counting from 0, is "full" when $\ell \bmod 4 = 3$ ($\ell$ divided by 4 leaves remainder 3), and the last layer is "full" whatever its index; all the others are "sliding".

Example. 8 layers: sliding, sliding, sliding, full, sliding, sliding, sliding, full. For 6 layers: sliding, sliding, sliding, full, sliding, full: layer 5 is the last, so it is full although $5 \bmod 4 = 1$.

The sliding-window mask

Query $t$ may attend to key $s$ when both conditions hold:

• $s \le t$: no looking at the future (causal);
• $t - s < \text{window}$: the key is among the last window tokens.

Example. $T = 5$ and window = 2. Row $t$ allows the keys $t - 1$ and $t$ (when they exist):

	$s=0$	$s=1$	$s=2$	$s=3$	$s=4$
$t=0$	1	0	0	0	0
$t=1$	1	1	0	0	0
$t=2$	0	1	1	0	0
$t=3$	0	0	1	1	0
$t=4$	0	0	0	1	1

(1 is True.) A full layer would also allow the ones at the bottom-left.

With window >= T every key $s \le t$ is within the window, so the mask is the ordinary causal mask: a lower triangle.

What to implement

• layer_types(n_layers): the list of "sliding" and "full".
• sliding_window_mask(T, window): the boolean $(T, T)$ mask.
• attention_masks(T, n_layers, window): one mask per layer, following layer_types. A full layer gets the plain causal mask.

Requirements

• sliding_window_mask(T, window) with window >= T is the ordinary causal mask.
• No loops over positions.

Hint

Compare a column of query positions np.arange(T)[:, None] with a row of key positions np.arange(T)[None, :]: broadcasting builds the whole table at once, and the two conditions are two comparisons joined with &. In row $t$ the number of allowed positions is min(t + 1, window).