Original source
---
id: attention-variants
title: Attention Variants Compared
code: AV
subtitle: Full, sliding window, strided, BigBird, differential and native sparse attention. Cost, KV cache, code, fixes.
hue: 265
sources: phases/07-transformers-deep-dive/12-kv-cache-flash-attention, phases/07-transformers-deep-dive/15-attention-variants, phases/10-llms-from-scratch/16-differential-attention-v2, phases/10-llms-from-scratch/17-native-sparse-attention
source_hash: 2d7b00775e872f21
---
## Variants side by side
type: comparison
from: phases/07-transformers-deep-dive/15-attention-variants/docs/en.md
| Variant | Compute | KV cache | Quality vs full |
|---|---|---|---|
| Full attention | `O(N^2)` | `O(N)` per layer | baseline |
| SWA, window 1024 | N times W | `O(W)` per layer | -0.1 ppl, good with global layers |
| Local + strided sparse | N times sqrt N | mixed | similar to SWA |
| BigBird, local + global + random | `O(N)` approx | mixed | matches full at 2× context |
| Native Sparse (DeepSeek-V3.2) | N times active fraction | `O(N)` | within 0.05 ppl |
| Differential | 2 times N squared | `O(2N)` | -5 to -10% ppl |
## Use X when
type: tips
from: phases/07-transformers-deep-dive/15-attention-variants/docs/en.md, phases/10-llms-from-scratch/16-differential-attention-v2/docs/en.md, phases/10-llms-from-scratch/17-native-sparse-attention/docs/en.md
- **Full attention**: every layer up to ~16K context, or when retrieval quality is paramount.
- **SWA + global mix**: long context above 32K. The 2026 default above 32K.
- **Sparse block attention**: custom kernel, custom pattern. For retrieval and audio.
- **Differential attention**: attention sinks hurt, as in long-context RAG and needle-in-haystack.
- **DIFF from the start**: a new model that targets 64k-plus effective context.
- **NSA**: pre-training or continued training that targets 64k-plus context.
- **Not DIFF or NSA**: a pre-trained dense model, or context under 16k.
## Mask builders
type: annotated-code
lang: python
from: phases/07-transformers-deep-dive/15-attention-variants/code/main.py, phases/07-transformers-deep-dive/15-attention-variants/docs/en.md
```python
def swa_mask(n, window):
M = [[NEG_INF] * n for _ in range(n)]
for i in range(n):
lo = max(0, i - window + 1) # window >= n gives full causal
for j in range(lo, i + 1):
M[i][j] = 0.0 # 0.0 attends, NEG_INF masks
return M
def strided_mask(n, window, stride): # local window, then:
for j in range(0, i + 1, stride):
M[i][j] = 0.0 # every stride-th token back to 0
render(swa_mask(8, window=4), "sliding window (W=4)")
render(strided_mask(8, window=2, stride=3), "local (W=2) + strided (stride=3)")
```
## SWA with FlexAttention
type: annotated-code
lang: python
from: phases/07-transformers-deep-dive/15-attention-variants/docs/en.md
```python
from torch.nn.attention.flex_attention import flex_attention, create_block_mask
def swa_pattern(b, h, q_idx, kv_idx):
return (q_idx - kv_idx < 1024) & (q_idx >= kv_idx) # window 1024, causal
mask = create_block_mask(swa_pattern, B=batch, H=heads, Q_LEN=n, KV_LEN=n)
out = flex_attention(q, k, v, block_mask=mask) # compiles to a Triton kernel
```
## NSA: three branches
type: architecture
from: phases/10-llms-from-scratch/17-native-sparse-attention/docs/en.md, phases/10-llms-from-scratch/17-native-sparse-attention/code/main.py
```arch
node | q | 1,2 | Query | one per position | client
node | cmp | 2,1 | Compressed branch | block means, N/l keys | process
node | sel | 2,2 | Selected branch | top-k blocks, k*l keys | process
node | win | 2,3 | Sliding window | last W tokens | process
node | kv | 3,2 | KV cache | fine-grained tokens | store
node | gate | 4,2 | Gate | sigmoid, 3 weights | model
edge | q | cmp | attend
edge | cmp | sel | top_k_blocks
edge | q | win | attend
edge | kv | sel | load blocks
edge | sel | gate | out_sel
edge | cmp | gate | out_cmp
edge | win | gate | out_win
```
## Flash Attention tile loop
type: flow
from: phases/07-transformers-deep-dive/12-kv-cache-flash-attention/docs/en.md, phases/07-transformers-deep-dive/12-kv-cache-flash-attention/code/main.py
1. **Outer loop**: load one Q tile, about 128 × 128, into SRAM.
2. **Inner loop**: load each K tile and V tile into SRAM.
3. **SRAM**: compute `S_tile = Q_tile @ K_tile^T`.
4. **SRAM**: keep a running `(max, sum)` across tiles.
5. **SRAM**: rescale the old output by `exp_old` and add the new block.
6. **HBM**: write `O_tile` to HBM. One HBM trip per tile.
7. **Result**: bit-identical to standard attention, memory O(N) not O(N^2).
## Tasks in code
type: task-table
from: phases/07-transformers-deep-dive/15-attention-variants/code/main.py, phases/07-transformers-deep-dive/12-kv-cache-flash-attention/code/main.py, phases/10-llms-from-scratch/16-differential-attention-v2/code/main.py, phases/10-llms-from-scratch/17-native-sparse-attention/code/main.py
| Task | Code |
|---|---|
| SWA causal mask | `swa_mask(n, window)` |
| Local plus strided mask | `strided_mask(n, window, stride)` |
| Count attended cells | `count_nonmasked(M)` |
| KV cache bytes | `kv_cache_bytes(N, n_layers, n_heads_kv, d_head, dtype=2)` |
| Flash style softmax | `tiled_softmax_dot(q, Ks, Vs, tile=4)` |
| Signal to noise of a row | `snr(weights_row, signal_pos)` |
| DIFF V1 and V2 params | `compute_param_diff(hidden=4096, n_heads=32, kv_heads=8)` |
| NSA keys per query | `count_nsa(N, l, k, W)` |
## Numbers to know
type: defaults
from: phases/07-transformers-deep-dive/12-kv-cache-flash-attention/docs/en.md, phases/07-transformers-deep-dive/15-attention-variants/docs/en.md, phases/07-transformers-deep-dive/15-attention-variants/code/main.py, phases/10-llms-from-scratch/16-differential-attention-v2/docs/en.md, phases/10-llms-from-scratch/17-native-sparse-attention/docs/en.md
| Setting | Default | Note |
|---|---|---|
| Llama 3 70B KV per token per layer | 4096 bytes | 80 layers, 8 KV heads, fp16 |
| Llama 3 70B KV, 32K context | 10 GiB | GQA with 8 KV heads |
| Gemma 3 layer mix | 5:1 SWA:global | window=1024 |
| SWA KV at 128K context | 128× | window 1024, last W tokens |
| Gemma-3 mix KV | ~6x | vs pure full attention |
| DIFF `lambda_init` | 0.8 | learned per head |
| NSA window `W` | 512 | `l = 64, k = 16` at 64k, 25x |
| NSA at 128k | 3536 keys | 36x reduction |
## Troubleshooting
type: troubleshooting
from: phases/07-transformers-deep-dive/15-attention-variants/docs/en.md, phases/10-llms-from-scratch/16-differential-attention-v2/docs/en.md, phases/10-llms-from-scratch/17-native-sparse-attention/docs/en.md, phases/07-transformers-deep-dive/12-kv-cache-flash-attention/docs/en.md
| Symptom | Cause | Fix |
|---|---|---|
| SWA model fails long-range retrieval | No global layers | Interleave full layers, 5:1 |
| Weight bleeds to token 0 | Attention sink from softmax | Differential attention |
| Lost in the middle past 32k | Softmax noise floor grows with N | DIFF from the start, or LoRA on Q |
| DIFF V1 decode slower than baseline | Value cache loaded twice | DIFF V2 |
| Late training unstable at 70B-plus | V1 per-head RMSNorm | V2 removes it |
| NSA on a dense pre-trained model | No retrofit | Continued training |
| No NSA gain under 16k | Three-branch overhead | Full attention |
| Inference-time sparsity gives small speedup | Model never routed through it | Natively trained sparsity |
## Q&A
type: qa
from: phases/07-transformers-deep-dive/15-attention-variants/docs/en.md, phases/10-llms-from-scratch/16-differential-attention-v2/docs/en.md, phases/10-llms-from-scratch/17-native-sparse-attention/docs/en.md, phases/07-transformers-deep-dive/12-kv-cache-flash-attention/docs/en.md
- **Does Flash Attention cut compute?** No. It hides the activation memory. Every token still attends to every token.
- **Why is NSA trainable?** Gradients flow through the compressed-branch scores. `top_k` only picks blocks to load.
- **Why does DIFF help more at 128k?** The noise floor grows with context. At 4k it is small.
- **Is DIFF V2 GQA compatible?** Yes. V2 adds Q heads, not KV heads.
- **NSA keys at 1024 tokens, l 32, k 4, w 128?** 32 + 128 + 128 = 288 keys, 3.5x fewer.
## Glossary
type: glossary
span: 2
from: phases/07-transformers-deep-dive/15-attention-variants/docs/en.md, phases/07-transformers-deep-dive/12-kv-cache-flash-attention/docs/en.md, phases/10-llms-from-scratch/17-native-sparse-attention/docs/en.md
- **SWA**: sliding window attention. Each query attends to its last W tokens.
- **Receptive field**: an L-layer SWA stack with window W sees up to L × W tokens.
- **Attention sink**: rows sum to 1, so uninformative queries put weight on position 0.
- **NSA**: Native Sparse Attention. Learned block sparsity, zero blocks skipped in the kernel.
- **DSA**: Deepseek Sparse Attention, the successor of NSA.
- **HBM**: High Bandwidth Memory. About 3 TB/s on H100.
- **SRAM**: per-SM on-chip memory. About 30 TB/s.