AI Engineer Exhaustive
EasyPlusgpt-oss

Sliding Window

gpt-oss

Easy

Causal mask that restricts each query to the previous W tokens, used on alternating attention layers.

Independent study note. Written from the public paper and official code. This is not TensorTonic Plus and does not reproduce their exercises, starter code, or tests. For the official version, subscribe on TensorTonic.

Overview

Following GPT-3, gpt-oss alternates attention layers between a banded local pattern and a fully dense causal pattern. Section 2.2 of the model card sets the bandwidth at 128 tokens. Local layers keep the KV cache and the score matrix cheap at the 131,072-token YaRN context; dense layers every other block still let information travel the full sequence.

The paper says “banded window and fully dense” and “every other” in spirit; it does not specify which parity is local. The reference model does: AttentionBlock sets self.sliding_window = config.sliding_window if layer_idx % 2 == 0 else 0. Even layers ((0,2,4,\ldots)) use (W=128). Odd layers use (W=0), which sdpa treats as “no extra band,” i.e. ordinary causal attention over the whole prompt.

The mask is applied to logits, together with the usual future-token triangle. Attention sinks (a learned extra logit) are not masked; a local head can dump mass into its sink when the 128-token band is uninformative.

How it works

For a sequence of length (T), build a (T\times T) score mask (M). Index rows as queries (i) and columns as keys (j), both 0-based. A banded causal layer with window (W) allows key (j) only when it is in the past-or-present band of width (W):

M_{ij} = \begin{cases} 0 & \text{if } i-W < j \le i,\\ -\infty & \text{otherwise.} \end{cases}

The allowed set has exactly (W) positions once (i\ge W-1), including the current token. Query (i=200), (W=128) sees keys (73,\ldots,200). Near the start of the sequence the band clips at (j=0), so query (i=10) sees (11) keys, not 128.

The official sdpa builds this as the sum of two additive masks:

  1. torch.triu(..., diagonal=1) — future keys (j>i) get (-\infty).
  2. If (W>0), torch.tril(..., diagonal=-W) — keys that are at least (W) steps behind ((j\le i-W)) get (-\infty).

When (W=0) the second term is skipped, leaving a standard causal triangle. Scores are scaled by (1/\sqrt{d}) before the mask is added, then the sink logit is concatenated, then softmax runs on the last axis.

The same (M) is broadcast over KV groups and the (q_{\mathrm{mult}}=8) query heads that share a group. Windowing does not depend on head index. RoPE is applied to (Q) and (K) using each token’s absolute position, not a position-within-window index. A token at index 10,000 in a local layer still rotates by 10,000, even though it may only attend 128 keys.

Layers: 36 (120b) or 24 (20b). That is 18 or 12 local layers and the same number of dense layers, interleaved. Dense layers are the ones that actually use the YaRN-extended 131k context; local layers always see at most 128 keys per query, independent of how far YaRN stretches the rotary basis.

Official code

gpt_oss/torch/model.py: ModelConfig.sliding_window = 128. AttentionBlock.__init__ gates that value on layer_idx % 2. sdpa applies the triu/tril pair when sliding_window > 0. There is no separate mask module.

Watch-outs

Sources