AI Engineer Exhaustive
MediumPlusGemma 3

Sliding Window Attention

Gemma 3

Medium

Local attention with fixed window size

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

Local layers in Gemma 3 do not attend over the full context. Each query may only see keys inside a fixed-width sliding window, in the style of Longformer (Beltagy et al., 2020). The report assigns that window a span of 1024 tokens and uses it as the main lever, together with the 5:1 local:global ratio, to keep the KV cache from growing linearly with a 128K context.

Only global layers pay the full-sequence cache cost. A local layer’s cache can be bounded by the window. Figure 5 in the report compares a global-only 2B stack (about 60% extra memory from the KV cache at 32K) to a 1:3 local:global stack with sw=1024 (under 15%). Figure 4 shows that shrinking the window further barely moves perplexity in their text-only ablation.

The 1024 figure is the paper’s default for the long-context models. The official configs split by size: 4B, 12B, and 27B use sliding_window_size=1024; 270M and 1B use 512. The 1B model is also the one limited to 32K context.

How it works

create_sliding_mask in gemma/gm/nn/_modules.py builds a boolean mask from absolute positions. Let (t) be a query position and (s) a key/cache position, with window width (W):

M_{t,s} = \mathbf{1}\!\left[s > t-W\right]\,\mathbf{1}\!\left[s < t+W\right].

Shapes: positions is [B, L], cache_positions is [B, S] (or a copy of positions when there is no cache). Broadcasting yields (M) of shape [B, L, S]. That mask is multiplied onto the already-built causal / padding mask attn_mask. Blocked logits are then replaced by the large negative constant K_MASK = -2.3819763e38 before softmax.

The (s < t+W) side looks bidirectional. In the decoder it is almost redundant because the incoming attn_mask is causal ((s \le t)). The conjunction that actually fires at train time is

t-W < s \le t,

which, for integer positions, is a closed window of (W) tokens including the query. Example: (t=2000), (W=1024) keeps keys (s \in {977,\ldots,2000}).

During decoding the cache is left-aligned and may hold stale or padded slots. The library therefore stores positions inside each layer cache and passes cache_positions into create_sliding_mask, so the window is computed from true token indices, not from buffer offsets.

Local layers also use the short RoPE base (\theta=10^4). The report keeps that base on local attention and only raises the global base to (10^6). A short window plus a short RoPE period is consistent: those layers are not asked to encode 128K of relative phase.

The paper specifies the 1024-token span, the Longformer citation, and the KV-cache motivation. The exact inequalities, the K_MASK fill, the cache-position trick, and the 512-token window on the smallest two models are from the official code.

Official code

Watch-outs

Sources