Thomas Piacentini
← Atelier

Fiche

Attention Variants Compared

Texte original · anglais

Full, sliding window, strided, BigBird, differential and native sparse attention. Cost, KV cache, code, fixes.

Variants side by side

VariantComputeKV cacheQuality vs full
Full attentionO(N^2)O(N) per layerbaseline
SWA, window 1024N times WO(W) per layer-0.1 ppl, good with global layers
Local + strided sparseN times sqrt Nmixedsimilar to SWA
BigBird, local + global + randomO(N) approxmixedmatches full at 2× context
Native Sparse (DeepSeek-V3.2)N times active fractionO(N)within 0.05 ppl
Differential2 times N squaredO(2N)-5 to -10% ppl
Sources

Use X when

  • 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.
Sources

Mask builders

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.00.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.0every 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)")
Sources

SWA with FlexAttention

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
Sources

NSA: three branches

attendtop_k_blocksattendload blocksout_selout_cmpout_winQueryone per positionclientCompressed branchblock means, N/l keysprocessSelected branchtop-k blocks, k*l keysprocessSliding windowlast W tokensprocessKV cachefine-grained tokensstoreGatesigmoid, 3 weightsmodel
Nodes and connections
NodeRoleKind
Queryone per positionclient
Compressed branchblock means, N/l keysprocess
Selected branchtop-k blocks, k*l keysprocess
Sliding windowlast W tokensprocess
KV cachefine-grained tokensstore
Gatesigmoid, 3 weightsmodel
FromToConnection
QueryCompressed branchattend
Compressed branchSelected branchtop_k_blocks
QuerySliding windowattend
KV cacheSelected branchload blocks
Selected branchGateout_sel
Compressed branchGateout_cmp
Sliding windowGateout_win
Sources

Flash Attention tile loop

  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).
Sources

Tasks in code

TaskCode
SWA causal maskswa_mask(n, window)
Local plus strided maskstrided_mask(n, window, stride)
Count attended cellscount_nonmasked(M)
KV cache byteskv_cache_bytes(N, n_layers, n_heads_kv, d_head, dtype=2)
Flash style softmaxtiled_softmax_dot(q, Ks, Vs, tile=4)
Signal to noise of a rowsnr(weights_row, signal_pos)
DIFF V1 and V2 paramscompute_param_diff(hidden=4096, n_heads=32, kv_heads=8)
NSA keys per querycount_nsa(N, l, k, W)
Sources

Numbers to know

SettingDefaultNote
Llama 3 70B KV per token per layer4096 bytes80 layers, 8 KV heads, fp16
Llama 3 70B KV, 32K context10 GiBGQA with 8 KV heads
Gemma 3 layer mix5:1 SWA:globalwindow=1024
SWA KV at 128K context128×window 1024, last W tokens
Gemma-3 mix KV~6xvs pure full attention
DIFF lambda_init0.8learned per head
NSA window W512l = 64, k = 16 at 64k, 25x
NSA at 128k3536 keys36x reduction
Sources

Troubleshooting

SymptomCauseFix
SWA model fails long-range retrievalNo global layersInterleave full layers, 5:1
Weight bleeds to token 0Attention sink from softmaxDifferential attention
Lost in the middle past 32kSoftmax noise floor grows with NDIFF from the start, or LoRA on Q
DIFF V1 decode slower than baselineValue cache loaded twiceDIFF V2
Late training unstable at 70B-plusV1 per-head RMSNormV2 removes it
NSA on a dense pre-trained modelNo retrofitContinued training
No NSA gain under 16kThree-branch overheadFull attention
Inference-time sparsity gives small speedupModel never routed through itNatively trained sparsity
Sources

Q&A

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.
Sources

Glossary

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.
Sources
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.
Code et livrables 1
Ce que vous en retenez

Enregistré dans votre carnet sur cet appareil.

Sources et limites

Les repères de lecture sont déclaratifs. Les résultats de tests sont enregistrés séparément ; ils valent pour les cas exécutés.

AI Engineering from Scratch · 3f9715c · MIT

Télécharger ce fichier ↓