Attention variants

1. Intuition

In self-attention each query scores every key, and a mask sets some scores to $-\infty$ before the softmax so those keys get zero weight. A decoder uses the causal mask: token $i$ reads keys $0 \dots i$, never the future. An encoder like BERT uses no mask: every token reads every token.

Every variant in this lesson is just a different mask:

1. Sliding window (local). Token $i$ reads only the last $w$ tokens, itself included. Cost per query is $w$, not $i$. A sliding-window layer only needs to keep the last $w$ keys and values in its KV cache.
2. Local + strided (sparse). Add a few long-range links, such as every 4th earlier token, so distant information has a shortcut (Child et al., 2019).
3. Global tokens. Pick a handful of positions (often a [CLS] token) that read everything and that everyone reads. Longformer and BigBird combined a local window with global tokens (Beltagy et al., 2020; Zaheer et al., 2020).
4. Interleaving. Most current long-context models keep some layers fully global. Gemma 2 alternated local and global layers 1:1 with a 4,096-token window; Gemma 3 uses five local layers (window 1,024) per global layer (Gemma Team, 2025).

A local mask doesn't make distant tokens invisible forever. Layer 1 lets token 15 read tokens 12 to 15. But token 12 already read 9 to 12 in the layer before, so after two layers token 15 carries information from 9 onward. Each layer extends the reach by $w - 1$ positions. The catch: information that travels many hops gets blurred, so pure local models are weak at retrieving one exact fact from far back. That is why global layers stay in the mix.

Note

These are changes to which pairs get scored. FlashAttention is different: it computes exactly the same full attention, only faster and with less memory traffic. Multi-query and grouped-query attention are different again: they share keys and values across heads to shrink the KV cache (see LLM inference).

2. Visualization

Try it (predict first, then check):

1. Choose causal, then sliding window with $w = 4$ on 16 tokens. Predict the attended-cell count for each before reading it. (136 and 58.)
2. Keep sliding window and raise the window to 16. What mask do you get, and why?
3. Raise the layers slider for sliding window $w = 4$. After how many layers can the last of 16 tokens see the first?
4. Switch to global tokens and pick query 0. Why does it see every key, while query 12 sees only 6?

3. The math

Symbols

Masked attention

Each pattern is a rule for $M_{ij}$:

Strided attention with $s \approx \sqrt{n}$ costs $O(n\sqrt{n})$; sliding window costs $O(nw)$, linear in $n$ for fixed $w$.

Receptive field. One sliding-window layer lets position $i$ read back $w - 1$ positions. Stack $L$ layers and the reach adds:

Worked example: 16 tokens

Take $n = 16$, $w = 4$.

• Causal: $16 \cdot 17 / 2 = 272/2 = 136$ cells, 53% of $16^2 = 256$.
• Sliding window: rows 0, 1, 2 have only 1, 2, 3 earlier keys, so $1 + 2 + 3 = 6$. Rows 3 to 15 are 13 rows of 4 keys: $13 \cdot 4 = 52$. Total $6 + 52 = 58$. The formula agrees: $16 \cdot 4 - 4 \cdot 3/2 = 64 - 6 = 58$.
• Local $w = 2$ + stride 4: long-range keys are $j \in \{3, 7, 11, 15\}$. Query 12 reads its window $\{11, 12\}$ plus $\{3, 7, 11\}$, which is $\{3, 7, 11, 12\}$: 4 keys. Summing over all 16 rows gives 52 cells.
• Local radius 2 + global token 0: query 12 reads $\{10, \dots, 14\}$ (5 keys) plus key 0: 6 keys. Query 0 is global and reads all 16. Total 100 cells.

Receptive field of the last token (position 15), $w = 4$:

Five layers of a 4-token window connect the whole 16-token sequence. A Gemma-3-style window of 1,024 over a few dozen layers reaches tens of thousands of tokens in principle, but each hop dilutes the signal, which is why exact long-range recall still needs global layers.

Worked example: KV cache for a long context

A synthetic 80-layer model with 8 KV heads, $d_{\text{head}} = 128$, BF16 (2 bytes), one 131,072-token sequence. KV bytes $= 2 \cdot \text{layers} \cdot n_{kv} \cdot d_{\text{head}} \cdot \text{tokens} \cdot 2$.

• All 80 layers global: $2 \cdot 80 \cdot 8 \cdot 128 \cdot 131{,}072 \cdot 2 = 42{,}949{,}672{,}960$ bytes, 42.95 GB.
• About 5:1 local:global, so 13 global and 67 local layers with $w = 1{,}024$:
  • Global part: $42.95 \cdot 13/80 = 6.98$ GB.
  • Local part: $2 \cdot 67 \cdot 8 \cdot 128 \cdot 1{,}024 \cdot 2 = 281{,}018{,}368$ bytes, 0.28 GB.
  • Total 7.26 GB, which is $42.95 / 7.26 = 5.9\times$ smaller.

The global layers carry almost all of the cost. The ratio of global to local layers is the first number to check on a long-context model.

4. Implementation

Build it from scratch

The full script is code/14-transformers/attention_variants.py (NumPy only). Every mask is a boolean matrix built from two index grids: i holds each cell's query position and j its key position. A rule over i and j is the whole mask.

def causal(n):
    i, j = np.indices((n, n))
    return j <= i                                   # [n, n], lower triangle

def sliding_window(n, w):
    i, j = np.indices((n, n))
    return (j <= i) & (i - j < w)                   # the last w tokens, self included

Predict: what does sliding_window(n, n) return, and what about sliding_window(n, 1)? (The script asserts the first equals causal(n).)

Sparse patterns are unions of simple rules. The strided mask adds every stride-th earlier key to the local window. Global tokens set whole rows and columns to True.

def local_strided(n, w, stride):
    i, j = np.indices((n, n))
    return sliding_window(n, w) | ((j <= i) & (j % stride == stride - 1))

def with_global(mask, global_ids):
    m = mask.copy()
    m[:, global_ids] = True                         # everyone reads the global tokens
    m[global_ids, :] = True                         # global tokens read everyone
    return m

Predict: with with_global applied to a causal mask, token 0 now reads token 15. Is that still a valid decoder for generation? (No: it leaks the future. Global tokens belong in encoders, or must stay causal in a decoder.)

The receptive field is a reachability question. Treat the mask as a graph: one layer is one hop. Multiplying the mask by the current reach matrix adds one hop.

def reach_after(mask, layers):
    """Which keys can influence each query after stacking `layers` copies of the mask."""
    R = np.eye(len(mask), dtype=bool)
    for _ in range(layers):
        R = (mask.astype(int) @ R.astype(int)) > 0  # one more attention hop
    return R

Predict: for a causal mask, how many hops until the last token reaches everything? (One. Causal attention is already fully connected to the past.)

The KV cache counts keys and values per layer. A sliding-window layer stores w tokens instead of the whole context.

def kv_bytes(layers, n_kv, d_head, tokens, bytes_per=2):
    return 2 * layers * n_kv * d_head * tokens * bytes_per   # K and V

Predict: if all 80 layers used $w = 1{,}024$ at a 131,072-token context, how much smaller would the cache be than all-global? ($131{,}072 / 1{,}024 = 128\times$; the script prints 0.34 GB.)

Use the library

PyTorch's scaled_dot_product_attention takes a boolean mask where True means "may attend", the same convention as above. For long sequences, materialising an $n \times n$ mask wastes memory, so FlexAttention (PyTorch 2.5+) takes the rule as a function and compiles a kernel that skips fully masked blocks:

import torch.nn.functional as F
from torch.nn.attention.flex_attention import flex_attention, create_block_mask

out = F.scaled_dot_product_attention(q, k, v, attn_mask=torch.from_numpy(sliding_window(n, 4)))

def swa(b, h, q_idx, kv_idx):                       # same rule, as a function
    return (kv_idx <= q_idx) & (q_idx - kv_idx < 1024)

block_mask = create_block_mask(swa, B=None, H=None, Q_LEN=n, KV_LEN=n)
out = flex_attention(q, k, v, block_mask=block_mask)   # q, k, v: [batch, heads, n, d_head]

What the library adds: block-level skipping (the speed comes from never loading zero blocks, not from the mask itself), fused softmax, and GPU kernels. In Hugging Face, a model's config tells you its pattern: look for sliding_window and, in recent models, layer_types (a list such as "sliding_attention" and "full_attention" per layer).

Run it

From guide/:

uv run code/14-transformers/attention_variants.py

=== attended cells on a 16-token sequence ===
  full (bidirectional)                             256 of 256
  causal                                           136 of 256
  sliding window w=4                                58 of 256
  local w=2 + stride 4                              52 of 256
  local radius 2 + global token 0 (bidirectional)  100 of 256
...
  3 layer(s): 10 tokens   (formula L*(w-1)+1 = 10, capped at 16 -> 10)
  5 layer(s): 16 tokens   (formula L*(w-1)+1 = 16, capped at 16 -> 16)
...
  every layer full          :  42.95 GB
  every layer window 1024   :   0.34 GB
  5:1 local:global, W=1024  :   7.26 GB   (13 global + 67 local layers, 5.9x smaller)

The reachability matrix agrees with $L(w-1)+1$ at every depth (the script asserts it), and the KV numbers match the worked example.

5. Engineering

Where each pattern fits.

Interleave, don't replace. A model made only of local layers is cheap but weak at "find the one sentence 50,000 tokens back." Keeping a fraction of global layers restores direct long-range paths. The global layers then dominate the KV cache, as the worked example shows.

The kernel matters more than the mask. A boolean mask applied to a dense score matrix still computes all $n^2$ scores, so nothing gets faster. Real savings need a kernel that skips masked blocks (FlexAttention, FlashAttention's sliding-window support, or a serving engine that implements the model's pattern).

Learned sparsity. Newer work selects which blocks to attend per query instead of fixing the pattern in advance, for example Native Sparse Attention (Yuan et al., 2025). Differential attention (Ye et al., 2024) is a different idea: it subtracts two softmax maps to cancel weight that piles onto uninformative tokens such as the first one (the "attention sink"). Both are worth knowing; neither is a default yet.

Serving. Inference engines allocate KV memory per layer type. Check that your engine supports the model's sliding window; one that silently treats every layer as global still gives correct outputs but uses far more memory.

Common mistakes

• Off-by-one windows: deciding whether $w$ counts the query itself. Here it does ($i - j < w$). Check against the model's reference code.
• Applying an encoder pattern (global tokens reading everything) inside a decoder, which leaks future tokens during training.
• Expecting speed from a mask alone, without a block-sparse kernel.
• Assuming a large theoretical receptive field means good long-range recall.
• Sizing KV memory for a hybrid model as if every layer were global, or as if every layer were local.

6. Knowledge check

Practice exercise. A 32-token sequence uses causal sliding-window attention with $w = 8$. How many cells are attended? How many layers until the last token can be influenced by the first?

Implementation challenge. Add a dilated(n, w, d) mask to the script: each query reads $w$ earlier keys spaced $d$ apart ($i, i-d, i-2d, \dots$). Print its cell count for $n = 16, w = 4, d = 2$ and its receptive field per layer. Why does stacking a dilated layer on a plain window layer grow the reach faster than two window layers?

def dilated(n, w, d):
    i, j = np.indices((n, n))
    gap = i - j
    return (gap >= 0) & (gap % d == 0) & (gap < w * d)

Summary

• Every attention variant here is a mask: a rule for which query-key pairs get scored.
• Sliding window costs $O(nw)$ and keeps only $w$ tokens of KV cache per layer; causal costs $O(n^2)$.
• Stacked windows reach $L(w-1)+1$ tokens, but many-hop paths blur, so models keep some global layers.
• In a local/global hybrid the global layers dominate KV memory; the ratio is the key number.
• Masks save compute only with kernels that skip masked blocks.

Next: MDPs, states, actions, and rewards

Related: Self-attention · LLM inference · Open models: architecture walkthroughs

Go deeper, code-first (AI Engineering from Scratch): Differential Attention (V2) · Native Sparse Attention (DeepSeek NSA)

Parts of this lesson are adapted from AI Engineering from Scratch by Rohit Ghumare (MIT license, notice): Attention Variants — Sliding Window, Sparse, Differential.

Interview angle