Module 09: Efficient Attention
Introduction
You have built a classic GPT: multi-head attention — and a first look at its Grouped-Query variant — in m05, a decoder stack (m06), training (m07), and a generation loop with a from-scratch KV-cache (m08). It works. It is also, by frontier standards, slow to serve — and the reason is not the arithmetic.
Efficient attention is the set of changes that let a model serve long contexts to many users at once. Two walls stand in the way, and this module knocks down both from first principles:
- The memory wall. During generation the bottleneck is the KV-cache: it holds a key and value for every head, every layer, every position, and it is re-read on every decode step. It — not the FLOPs — is what caps context length and batch size. Multi-Query and Grouped-Query Attention shrink it by sharing K/V heads across query heads; Multi-head Latent Attention goes further and caches a single low-rank latent instead.
- The compute wall. The score matrix is N \times N; writing it to memory costs more than the multiply. FlashAttention never materializes it — it streams the softmax over blocks while keeping a running total.
Why it matters for LLMs:
- Every open model since 2023 (Llama-2/3, Mistral, Qwen, Gemma, DeepSeek) uses GQA, not textbook MHA. A from-scratch GPT that only knows MHA is a generation behind.
- Long context (100k+ tokens) is impossible without shrinking the cache and the memory traffic — the topics of the next modules build directly on this one.
What You’ll Learn
After this module, you can:
- Explain why the KV-cache, not compute, limits inference — and calculate its size.
- Take the Grouped-Query Attention you met in m05 and build its cache-aware, causal form, with MHA and MQA as the two endpoints of a single
num_kv_headsdial. - Prove that GQA decoded through a cache matches a full forward pass exactly.
- Build Multi-head Latent Attention — cache a low-rank latent, reconstruct K/V on the fly, and understand why RoPE must be decoupled from it.
- Implement FlashAttention’s online softmax and show it equals ordinary attention without ever forming the N \times N matrix.
- Build PagedAttention — store the KV cache in fixed-size blocks addressed through a block table — and see why it cuts serving memory waste to near zero.
- Build continuous batching — an iteration-level scheduler that admits and evicts requests between decode steps — and measure why it beats static batches on latency and slot occupancy without changing a single output.
- Build chunked prefill — cap the tokens per iteration with a token budget, slice a long prompt’s prefill into chunks and piggyback decodes — and see why it bounds the inter-token-latency stall, dialing throughput against latency.
- Read the common attention mask shapes: full, causal, sliding-window, dilated.
- Build Forgetting Attention (FoX) — add a learned, data-dependent forget-gate decay D_{ij}=\log\prod f_l to the softmax scores; prove f\equiv1 is plain attention and a constant f=e^{-m} is exactly ALiBi, and see a gate reset the context.
- Build Compressed Attention (DeepSeek-V4’s CSA/HCA) — pool the far KV cache into a 1/m summary, keep a recent window uncompressed, and attend over selected compressed chunks plus that window; see why it reduces to dense at m=1 and trades cache for blur.
Prerequisites
This module requires familiarity with:
- Module 05: Attention — scaled dot-product and multi-head attention, where Grouped-Query Attention was introduced as a multi-head variation. Here we build its cache-aware, causal form.
- Module 08: Generation — the autoregressive loop and the
KVCachewe are about to shrink.
Intuition: The Memory Wall
When you decode one token, the model does very little arithmetic: one query attends over the cached keys and values. What it does a lot of is memory movement — it must read the entire KV-cache back from memory to compute that one step. So the size of the cache, and how many times you re-read it, is the real cost of serving.
How big is the cache? It stores K and V for every layer and every head:
\text{cache bytes} = 2 \cdot L \cdot n_{kv} \cdot n \cdot d_{\text{head}} \cdot b
where L is layers, n_{kv} the number of key/value heads, n the sequence length, d_{\text{head}} the head dimension, and b the bytes per element (2 for fp16/bf16). The only free lever is n_{kv}: use fewer K/V heads and the cache shrinks in exact proportion, with everything else untouched.
Drive the numbers yourself — stretch the context and watch the cache explode, then cut the K/V heads and watch it collapse:
NoteKey Insight
The cache scales with the number of key/value heads, never the query heads. That is the entire idea: keep all the query heads (they do the expressive work), but let them share a smaller set of K/V heads. MQA takes it to the limit — one K/V head for the whole layer.
The Spectrum: MHA → GQA → MQA
Multi-head, grouped-query, and multi-query attention are not three mechanisms. They are one mechanism with one dial: how many key/value heads do the query heads share?
- MHA (n_{kv} = n_\text{heads}): every query head has its own K/V. Full quality, full cache — the m05 baseline.
- GQA (1 < n_{kv} < n_\text{heads}): query heads split into n_{kv} groups; each group shares one K/V head. Near-MHA quality at a fraction of the cache.
- MQA (n_{kv} = 1): all query heads share a single K/V head. Smallest cache, a small quality cost.
Step through the dial and watch the query heads (top) rewire onto fewer key/value heads (bottom):
TipTry This
- MHA → MQA: slide from step 0 to step 2 and watch 8 K/V boxes collapse to 1. Each collapse divides the cache — and the memory traffic per token — by that factor.
- Notice the query heads never disappear. The model keeps all 8 ways of asking; it just shares the answers (K/V).
The Math: Sharing Keys and Values
Nothing about the attention formula changes. GQA computes
\text{Attention}(Q, K, V) = \text{softmax}\!\left(\frac{Q K^\top}{\sqrt{d_k}}\right) V
exactly as before — the only difference is where the K/V heads come from. With n_\text{heads} query heads and n_{kv} key/value heads, define the group size
n_{\text{rep}} = \frac{n_\text{heads}}{n_{kv}}.
Query head h uses key/value head \lfloor h / n_{\text{rep}} \rfloor. In code this is a single repeat_interleave: each cached K/V head is duplicated n_{\text{rep}} times to line back up with its group of query heads, after which the attention math is identical to MHA. The duplication happens on the way into the multiply; the cache still stores only n_{kv} heads — that is where the saving lives.
Because n_{kv} = n_\text{heads} makes n_{\text{rep}} = 1 (no duplication at all), MHA is just GQA with the dial turned all the way up. One implementation covers all three.
Code: Cache-Aware Grouped-Query Attention
m05 built a GroupedQueryAttention layer as a general multi-head variation; here attention.py builds its cache-aware, causal counterpart — the whole family still as one layer. The query projection keeps all num_heads heads; the key and value projections are narrower — they produce only num_kv_heads heads. That narrower projection is the concrete reason the cache is smaller.
import torch
from attention import GroupedQueryAttention
# 8 query heads sharing 2 K/V heads → Grouped-Query Attention
gqa = GroupedQueryAttention(embed_dim=32, num_heads=8, num_kv_heads=2).eval()
x = torch.randn(1, 6, 32) # (batch, seq, embed)
out = gqa(x)
print(f"Input shape: {tuple(x.shape)}")
print(f"Output shape: {tuple(out.shape)}")
print(f"n_rep (query heads per K/V head): {gqa.n_rep}")
print(f"Q projection out features: {gqa.q_proj.out_features}")
print(f"K/V projection out features: {gqa.k_proj.out_features} <- smaller")Input shape: (1, 6, 32)
Output shape: (1, 6, 32)
n_rep (query heads per K/V head): 4
Q projection out features: 32
K/V projection out features: 8 <- smallerThe K/V projection has one quarter the output width of the query projection — that ratio (n_{\text{rep}} = 4) is exactly the cache saving. Switch the dial to recover the endpoints:
# num_kv_heads == num_heads → plain Multi-Head Attention
mha = GroupedQueryAttention(embed_dim=32, num_heads=8, num_kv_heads=8)
# num_kv_heads == 1 → Multi-Query Attention
mqa = GroupedQueryAttention(embed_dim=32, num_heads=8, num_kv_heads=1)
for name, layer in [("MHA", mha), ("GQA", gqa), ("MQA", mqa)]:
print(f"{name}: num_kv_heads={layer.num_kv_heads:>2}, n_rep={layer.n_rep}, "
f"K/V params={layer.k_proj.out_features}")MHA: num_kv_heads= 8, n_rep=1, K/V params=32
GQA: num_kv_heads= 2, n_rep=4, K/V params=8
MQA: num_kv_heads= 1, n_rep=8, K/V params=4Just like the m08 layer, GroupedQueryAttention is cache-aware: pass a KVCache and it appends only the new token’s K/V (the small ones) and attends over the full history, using the same offset-aware causal mask that unifies prefill and decode.
from attention import KVCache
layer = GroupedQueryAttention(embed_dim=32, num_heads=8, num_kv_heads=2).eval()
cache = KVCache()
# Decode 4 tokens one at a time; the cache grows by one each step.
for t in range(4):
step = torch.randn(1, 1, 32)
layer(step, cache=cache)
print(f"after token {t}: cache holds {len(cache)} tokens, "
f"K shape {tuple(cache.keys.shape)}")after token 0: cache holds 1 tokens, K shape (1, 2, 1, 4)
after token 1: cache holds 2 tokens, K shape (1, 2, 2, 4)
after token 2: cache holds 3 tokens, K shape (1, 2, 3, 4)
after token 3: cache holds 4 tokens, K shape (1, 2, 4, 4)Notice the cached K has shape (1, 2, seq, head_dim) — 2 heads, not 8. An MHA cache at the same width would store four times as much.
Proving GQA Generalizes MHA and Caches Exactly
Two properties make this trustworthy rather than a hopeful approximation, and both are checked in tests/test_attention.py:
- Decoding through the cache is exact. Running a layer token-by-token with a
KVCachereturns the identical result to one full forward pass. - The cache shrinks in proportion to n_{kv}, following
kv_cache_bytes.
demonstrate_gqa shows both at once:
from attention import demonstrate_gqa
max_diff = demonstrate_gqa(embed_dim=64, num_heads=8, num_kv_heads=2, seq_len=6)
print(f"\nReturned max |diff| = {max_diff:.2e} (≈ 0 ⇒ cached == full recompute)")============================================================
GROUPED-QUERY ATTENTION
============================================================
embed_dim=64, num_heads=8, num_kv_heads=2, head_dim=8
Incremental (cached) vs full forward: max |diff| = 1.19e-07
(cache stores 6 tokens x 2 KV heads)
KV-cache at seq_len=6, 32 layers (fp16):
MHA (kv=8): 49,152 bytes
GQA (kv=2): 12,288 bytes (4x smaller)
MQA (kv=1): 6,144 bytes (8x smaller)
Returned max |diff| = 1.19e-07 (≈ 0 ⇒ cached == full recompute)The max |diff| is at the level of float32 rounding (~1e-7): incremental, cached decoding is the same computation as a full pass, just reorganized to avoid recomputing the past — exactly the guarantee m08 established for MHA, now holding across the whole GQA family.
Multi-head Latent Attention: A Different Lever
GQA shrinks the cache by keeping fewer key/value heads. But there is a second way to attack the same bottleneck, and it is what DeepSeek-V2 and V3 actually ship: keep all the heads, and instead shrink what you store per head to zero.
Multi-head Latent Attention (MLA) compresses each token into a single small latent vector and caches that — not per-head keys and values. At attention time it reconstructs the full per-head K and V from the latent on the fly. The cache holds one low-rank latent (dimension d_c) plus one small decoupled key that carries position, and nothing scales with the head count.
The trade is different from GQA’s. GQA loses a little quality by collapsing heads; MLA keeps every head’s own key and value — they are just derived from a shared latent rather than stored independently. DeepSeek reports it matches full-MHA quality while caching less than typical GQA.
The Math: Compress, Reconstruct, Decouple
Start from the token embedding h_t. MLA splits attention into a content path (low-rank, cached as a latent) and a position path (a small decoupled RoPE key).
Compress the KV into a latent — the one thing the cache keeps:
c^{KV}_t = W^{DKV} h_t \in \mathbb{R}^{d_c}, \qquad d_c \ll n_\text{heads}\, d_\text{head}.
Reconstruct per-head content keys and values from it at attention time:
k^{C}_t = W^{UK} c^{KV}_t, \qquad v_t = W^{UV} c^{KV}_t.
Decouple RoPE. Rotary position embeddings apply a position-dependent rotation between query and key, so they cannot be absorbed into a static up-projection (more on that below). MLA therefore routes position through a separate, shared key that is rotated and is cached — but is only d_h^{R} wide:
k^{R}_t = \text{RoPE}\!\left(W^{KR} h_t\right) \in \mathbb{R}^{d_h^{R}}.
Each head’s key is the concatenation of its reconstructed content and the one shared positional key; queries are built the same way (with per-head RoPE):
k_{t,i} = \big[\,k^{C}_{t,i}\;;\;k^{R}_t\,\big], \qquad q_{t,i} = \big[\,q^{C}_{t,i}\;;\;q^{R}_{t,i}\,\big],
\text{score}_{t,j,i} = \frac{q_{t,i} \cdot k_{j,i}}{\sqrt{d_\text{head} + d_h^{R}}}.
(Queries are also compressed to a latent c^{Q}_t = W^{DQ} h_t before the up-projections. That saves activation memory during training but, unlike the KV latent, is never cached — it is not where the inference win comes from.)
The cache now stores, per token per layer, exactly d_c + d_h^{R} numbers — no factor of two for K and V, no scaling with the head count:
\text{MLA cache bytes} = L \cdot n \cdot (d_c + d_h^{R}) \cdot b.
NoteKey Insight: A latent is worth a fraction of a head
DeepSeek-V2 uses d_c = 512, d_h^{R} = 64, with n_\text{heads} = 128 heads of dimension d_\text{head} = 128. Its cache is 512 + 64 = 576 numbers per token — the same as GQA with \frac{576}{2 \cdot 128} = 2.25 groups, yet it keeps all 128 heads’ worth of expressiveness. The paper reports a 93.3% smaller KV-cache than the dense MHA model it replaced, and 5.76× higher generation throughput.
Code: Latent Attention from Scratch
attention.py builds MultiHeadLatentAttention and its MLACache from these equations. The forward pass reconstructs K and V from the cached latent, rotates the decoupled key, and attends — the same offset-aware causal mask as the GQA layer.
import torch
from attention import MultiHeadLatentAttention
# 8 heads, but K and V are funneled through a 24-dim latent + 16-dim rope key.
mla = MultiHeadLatentAttention(
embed_dim=64, num_heads=8,
kv_latent_dim=24, # d_c — the cached latent
q_latent_dim=32, # d_c' — query latent (activation memory only)
rope_dim=16, # d_h^R — decoupled RoPE key, shared across heads
).eval()
x = torch.randn(1, 6, 64)
out = mla(x)
print(f"Output shape: {tuple(out.shape)}")
print(f"KV latent d_c = {mla.kv_latent_dim}, decoupled rope d_h^R = {mla.rope_dim}")
print(f"Cached per token: {mla.kv_latent_dim + mla.rope_dim} numbers "
f"(vs {2 * mla.num_heads * mla.head_dim} for MHA)")Output shape: (1, 6, 64)
KV latent d_c = 24, decoupled rope d_h^R = 16
Cached per token: 40 numbers (vs 128 for MHA)The MLACache stores only the latent and the decoupled key — look at the shapes, and notice there is no head axis at all:
from attention import MLACache
cache = MLACache()
for t in range(4):
mla(x[:, t:t+1], cache=cache) # decode one token, growing the cache
print(f"latent cache shape: {tuple(cache.latent.shape)} # (batch, seq, d_c)")
print(f"rope cache shape: {tuple(cache.k_rope.shape)} # (batch, seq, d_h^R)")latent cache shape: (1, 4, 24) # (batch, seq, d_c)
rope cache shape: (1, 4, 16) # (batch, seq, d_h^R)Just like GQA, decoding through the cache is exact — the reconstructed K/V and the decoupled key reproduce a full forward pass to floating-point rounding. demonstrate_mla checks that and prints the cache comparison:
from attention import demonstrate_mla
max_diff = demonstrate_mla(seq_len=6, verbose=True)
print(f"\nmax |full - cached| = {max_diff:.2e} (≈ 0 ⇒ exact)")============================================================
MULTI-HEAD LATENT ATTENTION
============================================================
embed_dim=64, num_heads=8, head_dim=8
d_c (kv latent)=24, d_h^R (rope)=16
Incremental (cached) vs full forward: max |diff| = 8.94e-08
(cache stores 6 tokens x (24+16) numbers)
KV-cache at seq_len=6, 60 layers (fp16):
MHA (kv=8): 92,160 bytes
GQA (kv=2): 23,040 bytes
MLA: 28,800 bytes (3.2x smaller than MHA)
MLA cache ≈ 2.5-group GQA (per DeepSeek-V2's comparison)
max |full - cached| = 8.94e-08 (≈ 0 ⇒ exact)Why RoPE Has To Be Decoupled
The decoupled key looks like an odd wart until you try to remove it. The content score for head i is
q^{C}_{t,i} \cdot k^{C}_{j,i} = \big(W^{UQ}_i c^{Q}_t\big) \cdot \big(W^{UK}_i c^{KV}_j\big) = (c^{Q}_t)^\top \underbrace{\big(W^{UQ\,\top}_i W^{UK}_i\big)}_{\text{one fixed matrix}} c^{KV}_j.
The two up-projections collapse into a single matrix you can precompute — so you never actually build k^{C}; the query attends the cached latent directly. That absorption is what makes the tiny latent enough. Verify the identity:
# Content score via explicit K equals the score via the absorbed matrix W_UQ^T W_UK.
c_kv = mla.w_dkv(x) # (1, 6, d_c)
c_q = mla.w_dq(x) # (1, 6, d_c')
k_c = mla.w_uk(c_kv).view(1, 6, mla.num_heads, mla.head_dim).transpose(1, 2)
q_c = mla.w_uq(c_q ).view(1, 6, mla.num_heads, mla.head_dim).transpose(1, 2)
explicit = torch.matmul(q_c, k_c.transpose(-2, -1)) # reconstruct K, then score
absorbed_w = mla.absorbed_qk_weight() # (heads, d_c', d_c)
proj = torch.matmul(c_q.unsqueeze(1), absorbed_w.unsqueeze(0))
absorbed = torch.matmul(proj, c_kv.unsqueeze(1).transpose(-2, -1))
print(f"max |explicit - absorbed| = {(explicit - absorbed).abs().max():.2e}")max |explicit - absorbed| = 1.79e-07Now try the same with RoPE. Rotary embeddings insert a rotation R_{t-j} that depends on the relative position t-j between the query and key: q_{t,i}^\top R_{t-j}\, k_{j,i}. That rotation sits between W^{UQ} and W^{UK} and changes every step, so the two matrices no longer collapse into one fixed matrix — the absorption breaks. MLA’s fix is to give position its own small key that is not reconstructed from the latent: a single shared k^{R}, rotated and cached directly. Content compresses; position rides separately.
Interactive: The Cache Ladder
Set a realistic model and climb the ladder from MHA down to MLA. Drag the MLA latent width and watch how many GQA groups it is worth — below about 2–3 groups GQA starts to hurt quality, but MLA gets there while keeping every head:
TipTry This
- Set heads to 128, head dim 128 (DeepSeek-V2’s shape). MHA caches 32,768 numbers per token; MLA at d_c=512, d_h^R=64 caches 576 — worth 2.25 GQA groups but with all 128 heads intact.
- Push the latent width up: the cache grows and the “groups worth” climbs. MLA is a continuous dial on the cache, where GQA can only step by whole heads.
- Drop the decoupled rope to 0: the cache shrinks, but the model loses its ability to place tokens — position had nowhere to ride.
Paged Attention: The KV Cache as Virtual Memory
GQA and MLA shrink what you cache per token. There is a second, orthogonal lever: how you store the cache in memory. It turns out most serving systems throw away the majority of their KV memory before a single clever attention trick is applied.
The reason is that the cache grows one token at a time and every request is a different length — but a GPU wants contiguous memory. So the naive server does the obvious thing: reserve one contiguous buffer of max_seq_len per request, up front. A request that will only ever produce 40 tokens still holds a 4096-slot reservation. Three kinds of waste follow:
- Internal fragmentation — the reserved-but-unused tail of every buffer (the 4056 slots that never fill).
- Reservation waste — memory pinned for tokens that may never be generated.
- External fragmentation — free gaps between buffers, too small to seat a new request even when their total would fit.
Drive it yourself. Hold a fixed max context and set the actual lengths a few requests reach — the contiguous scheme reserves the full context for each, while the tokens that arrive fill only a sliver:
NoteKey Insight
The contiguous scheme’s waste grows with the gap between max_seq_len and the length a request actually reaches — often 90%+ of the reservation. Paging caps the waste at one partially-filled block per request (block_size − 1 tokens), independent of the context you allow. Shrink the block size and the waste shrinks with it.
The Block Table
The fix is the oldest idea in operating systems: paging. An OS does not give a process one contiguous slab of physical RAM; it hands out fixed-size pages from anywhere in memory and keeps a page table mapping the process’s logical addresses to physical ones. PagedAttention (the vLLM paper, Kwon et al., 2023) does exactly this for the KV cache:
- Carve GPU memory into fixed-size physical blocks, each holding the K/V for
block_sizetokens (vLLM’s default is 16). They live in one shared pool. - Give every sequence a block table — a list mapping its logical block index
0, 1, 2, …to a physical block number, which can sit anywhere in the pool. - Allocate blocks on demand: a sequence of
ntokens claims exactly \lceil n / \text{block\_size} \rceil blocks, growing by one whenever its last block fills. Nothing is reserved ahead of time.
Attention then gathers K and V by walking the block table — a scatter across physical blocks — instead of reading one contiguous run. The numbers are identical; only the memory layout changed.
Step through a real trace: append tokens to a sequence and watch its logical blocks fill, each mapping to a scattered physical block claimed from the pool on demand.
TipTry This
- Step to token 16, then 17: the first block fills, and token 17 claims a second physical block — allocation is lazy, one block at a time.
- Notice the physical block numbers are not
0, 1, 2— they are whatever the free-list handed back. The block table is what makes a scattered layout read as one contiguous sequence. - Shrink
block_sizein the fragmentation chart above to 1: waste vanishes, but real systems keep it at 16 so the attention kernel still reads a useful run at once. Paging trades a little bandwidth for near-zero waste.
Code: A Paged KV Cache from Scratch
Two pieces do all the work. A BlockAllocator is a free-list over the physical pool — allocate() pops a block, free() returns it, and num_free + num_used never changes. A PagedKVCache owns the physical K/V pool of shape (num_blocks, block_size, num_kv_heads, head_dim) plus one block table per sequence; appending a token writes into (block_table[t // block_size], t % block_size) and claims a fresh block whenever the offset wraps to zero. Both live in paged_attention.py.
import torch
from paged_attention import PagedKVCache, blocks_needed
# A tiny pool: 8 physical blocks of 4 tokens each, 2 KV heads, head_dim 3.
cache = PagedKVCache(num_blocks=8, block_size=4, num_kv_heads=2, head_dim=3)
seq = cache.add_sequence()
# Append 5 tokens one at a time (the decode loop).
torch.manual_seed(0)
tokens = [(torch.randn(2, 3), torch.randn(2, 3)) for _ in range(5)]
for k, v in tokens:
cache.append(seq, k, v)
print(f"Tokens cached: {cache.seq_len(seq)}")
print(f"Blocks used: {len(cache.block_table(seq))} (ceil(5/4) = {blocks_needed(5, 4)})")
print(f"Block table: {cache.block_table(seq)} # logical -> physical")Tokens cached: 5
Blocks used: 2 (ceil(5/4) = 2)
Block table: [7, 6] # logical -> physicalThe block numbers come off the free-list, so they are scattered — but gathering the sequence back through its table reproduces the contiguous cache exactly, the same bit-for-bit guarantee GQA and MLA gave through their caches:
# Reconstruct the full K/V by walking the block table (a scatter-gather).
gathered_k, gathered_v = cache.gather(seq)
# The contiguous reference: just the tokens we appended, stacked.
ref_k = torch.stack([k for k, v in tokens])
print(f"Gathered shape: {tuple(gathered_k.shape)} (len, kv_heads, head_dim)")
print(f"Paged == contiguous: {torch.equal(gathered_k, ref_k)}")Gathered shape: (5, 2, 3) (len, kv_heads, head_dim)
Paged == contiguous: TrueNow the payoff, on a realistic workload. Serve four variable-length requests from one pool and compare against the naive contiguous reservation:
from paged_attention import demonstrate_paging
result = demonstrate_paging(max_seq_len=512, block_size=16)Serving 4 sequences, lengths [40, 17, 8, 31], block_size 16:
Contiguous reserve (512/seq): 2,048 slots
Paged actual use: 128 slots (8 blocks)
Contiguous reserves 21x the tokens; it wastes 93.8% of its reservation.
Paged waste: 32 tokens total (<= 4x15, one partial block per sequence).
gather() == contiguous reference for every sequence: TrueThe contiguous scheme reserves 512 slots per request — 2048 for four short requests, 94% of it empty. Paging uses 128 slots (8 blocks) with 32 wasted tokens total, at most one partial block per request. The saved memory is what a real server spends on a bigger batch, which is where the throughput comes from.
Continuous Batching
Paging fixed space; the other half of a serving system is time. A GPU runs requests in a batch — one forward pass advances every sequence in the batch by a token. The question is which requests share each pass.
The obvious answer, static batching, groups requests and runs the whole group to completion together. Because generation is autoregressive, a batch of mixed-length requests runs for as many steps as its longest member. A batch is a bus: it does not leave until the last passenger is done, and no new passenger boards until it returns to the depot. One long generation holds the batch hostage while short requests finish early and sit idle in their slots.
Continuous batching — also called iteration-level scheduling (Orca, Yu et al., OSDI 2022) — reschedules at the granularity of a single decode iteration instead of a whole request. Between every step it (1) evicts any request that just emitted its last token, returning its blocks to the pool, and (2) admits a waiting request into the freed slot. A batch becomes a taxi rank: the instant a seat opens, the next rider takes it. Paging is what makes that cheap — a request joins or leaves by editing a block table, never by copying a contiguous cache.
Code: An Iteration-Level Scheduler
continuous_batching.py builds both schedulers over the very PagedKVCache from the last section. A Request prefills its prompt on admission, then emits one decode token per step, occupying one slot until it finishes. static_batch runs fixed groups to completion; continuous_batch runs the admit → decode → evict loop. We serve the classic hostage workload: one 16-token generation among five 2-token ones.
from continuous_batching import (
Request, static_batch, continuous_batch, demonstrate_continuous_batching,
occupancy, mean_latency, outputs_agree,
)
# One long request (16 output tokens) + five short ones (2 each), 3 slots.
workload = [Request(0, 0, 4, 16)] + [Request(i, 0, 4, 2) for i in range(1, 6)]
num_blocks = sum(r.footprint(4) for r in workload)
static = static_batch(workload, num_blocks, block_size=4, max_running=3)
cont = continuous_batch(workload, num_blocks, block_size=4, max_running=3)
print(f"static : makespan {static.makespan:>2} steps, "
f"occupancy {occupancy(static, workload):.0%}, "
f"mean latency {mean_latency(static, workload):.1f}")
print(f"continuous : makespan {cont.makespan:>2} steps, "
f"occupancy {occupancy(cont, workload):.0%}, "
f"mean latency {mean_latency(cont, workload):.1f}")static : makespan 18 steps, occupancy 48%, mean latency 17.0
continuous : makespan 16 steps, occupancy 54%, mean latency 5.7The short requests are the tell. Under static batching every request in the first group returns only when the 16-token generation finishes — so a 2-token request waits 16 steps. Under continuous batching it leaves after 2 and its slot is immediately refilled:
print("finish step (continuous):", dict(sorted(cont.finish_step.items())))
print("finish step (static): ", dict(sorted(static.finish_step.items())))finish step (continuous): {0: 16, 1: 2, 2: 2, 3: 4, 4: 4, 5: 6}
finish step (static): {0: 16, 1: 16, 2: 16, 3: 18, 4: 18, 5: 18}The Math Is Untouched
Continuous batching is pure scheduling — it never changes what a request computes. Each sequence owns its own block table, so its K/V is identical no matter who it shared a step with. We prove it: every request’s gathered cache is bit-for-bit equal under both schedulers.
print("static and continuous produce identical K/V:", outputs_agree(static, cont))static and continuous produce identical K/V: TrueWatch the two policies run on the same workload. Static (top) drains one group before starting the next — empty lanes are wasted GPU. Continuous (bottom) refills a slot the moment it frees, keeping the batch full.
TipTry This
- Watch the idle lanes. Step through the static batch: after step 2, two of its three lanes go dashed (idle) while the 16-token request runs on alone — that emptiness is wasted GPU. The continuous lanes stay full.
- Find the freed slot. In the continuous row, the moment a short request hits
2/2it vanishes and a waiting request appears in its lane on the next step. - Change the workload. Call
demonstrate_continuous_batchingwith your ownRequestlist — try many short requests and one very long one, or spreadarrivaltimes out — and re-read the latency reduction.
NoteKey Insight
GQA, MLA, and FlashAttention change the math of attention to move less data. Paging and continuous batching change nothing about the math — they are pure memory management and scheduling — yet they often buy a larger throughput win than any attention variant, because on a real server the binding constraint is how many requests fit and how full the batch stays.
Chunked Prefill: Bounding the Stall
Continuous batching keeps the batch full, but it hides a cost. It treats a request’s prefill — the one forward pass that ingests the whole prompt — as a free event on admission. On a real GPU it is the opposite of free, and it is the last stall left to remove.
Generation has two phases that cost nothing alike:
- Prefill runs all
prompt_lenprompt tokens through the model in parallel, in a single pass. It keeps the matrix units busy — it is compute-bound. A loaded truck. - Decode emits one token per request per step. The pass is dominated by streaming the weights and KV cache in from memory to do a single token’s worth of math, leaving the compute units mostly idle — it is memory-bound. A scooter.
Now picture the taxi rank from the last section with one long prompt in it. The instant that request is admitted, one iteration balloons from “a handful of decode tokens” to “hundreds of prefill tokens.” Every other request in the batch — all of them mid-generation — waits for its next token behind that prefill. That is a spike in inter-token latency (TBT, “time between tokens”): the smooth stream of tokens a user is reading suddenly hitches. And a batch with no prefill in it wastes the compute a prefill would have used. You cannot get low tail latency and high throughput by scheduling whole prefills.
Chunked prefill (Sarathi, Agrawal et al. 2023; Sarathi-Serve, OSDI 2024 — now the default scheduler in vLLM, TensorRT-LLM, and SGLang) removes both problems with one lever: a token budget B on the work any single iteration may do. Build each iteration stall-free:
- Serve every running decode one token first. Decodes are cheap and latency-critical, so generation never pauses.
- Fill the leftover budget
B − (#decodes)with one chunk of a prefill. AP-token prompt is spread over⌈P / chunk⌉iterations, and the decodes piggyback in the same batch — filling the compute slack of those memory-bound steps.
Peak work per iteration is now bounded by B, so the TBT tail is bounded; the price is a longer TTFT (“time to first token”) for the chunked prompt. That is the whole throughput ⇆ latency dial, and B is the knob.
The Math: A Budget Bounds the Iteration
Let an iteration serve d decode tokens and one prefill chunk of c tokens. Its work is
w = d + c, \qquad c \le \max(0,\; B - d),
so with B \ge the batch size (every decode fits), w \le B — the peak iteration cost, and hence the worst-case time between two of a user’s tokens, is capped by a number you choose. A prompt of P tokens finishes its prefill after
\left\lceil \frac{P}{\,B - d\,} \right\rceil \text{ iterations,}
which is exactly what stretches its TTFT as B shrinks. Smaller B: flatter TBT, later first token. Setting B = \infty recovers the un-chunked full prefill — the whole prompt in one iteration — so both policies are one function with one argument.
Code: A Stall-Free Scheduler
chunked_prefill.py builds the scheduler over the very same PagedKVCache and Request from the last two sections — nothing about the cache changes; this is pure scheduling. token_budget=None is the full-prefill baseline; any integer is the chunked, budgeted version.
from chunked_prefill import (
Request, chunked_prefill_batch, demonstrate_chunked_prefill,
peak_iteration_work, mean_ttft, outputs_agree,
)
# Three short requests already decoding; one 32-token prompt arrives at step 2.
workload = [Request(i, 0, 4, 8) for i in range(3)] + [Request(3, 2, 32, 4)]
num_blocks = sum(r.footprint(4) for r in workload)
full = chunked_prefill_batch(workload, num_blocks, 4, max_running=4, token_budget=None)
chunked = chunked_prefill_batch(workload, num_blocks, 4, max_running=4, token_budget=8)
print(f"full prefill : peak work {peak_iteration_work(full):>3} tok/iter, "
f"mean TTFT {mean_ttft(full, workload):.1f}")
print(f"chunked B=8 : peak work {peak_iteration_work(chunked):>3} tok/iter, "
f"mean TTFT {mean_ttft(chunked, workload):.1f}")full prefill : peak work 35 tok/iter, mean TTFT 1.0
chunked B=8 : peak work 8 tok/iter, mean TTFT 2.8The 32-token prompt is the tell. Under full prefill its whole prompt lands in one iteration — that step does 35 tokens of work while the neighbours’ decode steps do 3, and every decoding request’s next token stalls behind it. Chunked prefill holds every iteration at or under the budget of 8, at the cost of stretching that prompt’s own first token from step 1 to step 7:
from chunked_prefill import ttfts
print("TTFT per request (full) :", ttfts(full, workload))
print("TTFT per request (chunked):", ttfts(chunked, workload))TTFT per request (full) : {0: 1, 1: 1, 2: 1, 3: 1}
TTFT per request (chunked): {0: 1, 1: 1, 2: 2, 3: 7}The Math Is Untouched
Slicing a prefill into chunks is scheduling, not arithmetic. A request appends its prompt positions 0..P-1 and then its decode positions P..P+O-1 in that order no matter how the prefill was cut, so its gathered K/V is bit-for-bit identical — to full prefill, to continuous batching, and to a standalone single-sequence reference.
print("full and chunked produce identical K/V:", outputs_agree(full, chunked))full and chunked produce identical K/V: TrueWatch the two policies iteration by iteration. Each bar is one iteration, split into decode tokens (the steady base) and the prefill chunk (stacked on top). Full prefill (top) spikes far past the budget line the instant the long prompt arrives; chunked (bottom) never crosses it.
TipTry This
- Find the stall. Step to the iteration where the long prompt arrives. In the full-prefill lane one bar shoots far past the dashed budget line and flashes stall! — that single tall bar is every other request’s frozen next token. The chunked lane holds a flat, budget-height bar instead.
- Watch the chunk crawl. In the chunked lane the prefill fills a little of the budget each iteration for several steps — that is the
⌈P/chunk⌉slices — while the decode base keeps ticking underneath. Nobody stalls. - Turn the dial. Call
chunked_prefill_batchwithtoken_budgetset to 32, 16, 8, then 4 and readpeak_iteration_workandmean_ttfteach time: peak work falls toward the budget while TTFT climbs. That is the throughput ⇆ latency trade in one number each way.
The dial in full: sweep the budget on the same workload and watch peak iteration work (the TBT tail) trade off against makespan and TTFT.
NoteKey Insight
Continuous batching bounds who is in the batch; chunked prefill bounds how much work the batch does in one step. Both are pure scheduling over the paged cache — neither touches the attention math — yet together they are what let one server hold thousands of streams smooth: decodes never stall, and the memory-bound decode steps are topped up with compute-bound prefill tokens. The single knob B slides the whole system between “answer this one prompt fastest” (large B) and “keep everyone’s tokens flowing evenly” (small B).
FlashAttention: The Compute Wall
Shrinking the cache fixes decode-time memory. But there is a second wall, and it bites hardest during prefill (and training), when a long prompt attends to itself: the score matrix S = QK^\top is N \times N. Materializing it in memory — writing N^2 numbers out and reading them back for the softmax — is the dominant cost, far more than the multiply.
FlashAttention never writes that matrix. It walks the keys in blocks and maintains, per query row, a running maximum m, a running denominator \ell, and a running output. When a new block raises the max, the prior total is rescaled by e^{m_{\text{old}} - m_{\text{new}}} and the new block folded in. The final output is identical to a full softmax — this is the online softmax identity — but only one block is ever in memory.
online_softmax_attention in attention.py is a faithful, runnable version of that recurrence:
import torch
from attention import online_softmax_attention
torch.manual_seed(0)
q = torch.randn(2, 4, 12, 16) # (batch, heads, seq, head_dim)
k = torch.randn(2, 4, 12, 16)
v = torch.randn(2, 4, 12, 16)
# The ordinary way: form the full (12 x 12) matrix, softmax, multiply.
reference = torch.softmax(q @ k.transpose(-2, -1) / 16 ** 0.5, dim=-1) @ v
# FlashAttention's way: stream over blocks of 4 keys, never forming 12 x 12.
streamed = online_softmax_attention(q, k, v, block_size=4)
print(f"max |streamed - reference| = {(streamed - reference).abs().max():.2e}")
print(f"Match: {torch.allclose(streamed, reference, atol=1e-5)}")max |streamed - reference| = 2.38e-07
Match: TrueSame answer, to rounding. The block size changes only how much memory the computation touches at once, never the result. Watch the running statistics absorb one key block at a time — the mechanism that lets FlashAttention keep the full softmax correct while seeing only a slice of it:
NoteKey Insight
The online-softmax trick is what makes FlashAttention exact, not approximate: rescaling the running total by e^{m_{\text{old}} - m_{\text{new}}} whenever the max grows means the final normalization is identical to a one-shot softmax. You trade a bit of recomputation for never storing the N \times N matrix — and on real hardware that memory traffic was the bottleneck.
Attention Mask Gallery
Efficiency also comes from attending to fewer positions. A mask is just a boolean grid — “may query i attend to key j?” — and different shapes trade coverage for cost. Pick a shape and read its structure:
The sliding-window shape is available as sliding_window_mask in attention.py, so you can feed real local attention to the layer:
from attention import sliding_window_mask
mask = sliding_window_mask(6, window=3)
print("Sliding-window mask (window=3), 1 = attend:")
print(mask.int())Sliding-window mask (window=3), 1 = attend:
tensor([[1, 0, 0, 0, 0, 0],
[1, 1, 0, 0, 0, 0],
[1, 1, 1, 0, 0, 0],
[0, 1, 1, 1, 0, 0],
[0, 0, 1, 1, 1, 0],
[0, 0, 0, 1, 1, 1]], dtype=torch.int32)
TipTry This
- Switch to Sliding window — every row keeps a fixed width, so total cost grows linearly, not quadratically, with sequence length.
- Switch to Dilated and notice how position 11 still reaches position 1 in one hop despite skipping every other key: far context, fewer connections.
Native Sparse Attention
The masks above are fixed: a sliding window always keeps the last w keys, whatever the query. But a query looking for a name defined 4,000 tokens ago and a query continuing a local phrase want to read completely different parts of the past. Native Sparse Attention (NSA) makes the sparse pattern content- dependent and learned — each query decides which slice of the past to read — and, crucially, it is trained sparse from scratch rather than approximating a dense model after the fact. It is the attention in DeepSeek’s 2025 line of work.
NSA gives every query three ways to look back, and lets it learn how much to trust each one:
- Compression — a cheap, coarse glance at the whole past: squash each block of keys into one summary token and attend over the summaries.
- Selection — a sharp look at a few whole blocks, at full resolution. Which blocks? The ones the compression glance already scored highest — so choosing them costs almost nothing.
- Sliding window — the recent tokens, at full resolution, for local detail.
A per-query gate blends the three outputs. That is the entire idea: coarse everywhere, sharp in a few chosen places, plus the local neighborhood.
NoteKey Insight
The three branches divide the labor by resolution. Compression never misses anything but sees it blurrily; selection sees a handful of places perfectly; the window covers the tokens too recent to have been compressed yet. The selection branch is what makes NSA more than a fixed pattern — and it is nearly free because its “which blocks matter?” scores are the compression attention already computed.
The Math: Three Branches, One Gate
For a query q_t at position t, NSA computes (Eq. 5 in the paper):
o^*_t = \sum_{c \in \{\text{cmp},\,\text{slc},\,\text{win}\}} g^c_t \cdot \operatorname{Attn}\!\left(q_t,\, \tilde K^c_t,\, \tilde V^c_t\right), \qquad g^c_t = \sigma\!\big(\mathrm{MLP}(q_t)\big) \in [0, 1].
Each branch is ordinary attention over a different set of keys/values:
Compression (Eq. 7). Slice the past keys into blocks of length l (stride d) and map each block to one token with a small learnable network \varphi (an intra-block position weighting):
\tilde K^{\text{cmp}}_t = \Big\{\, \varphi\big(k_{id+1:\,id+l}\big) \;\Big|\; 0 \le i \le \big\lfloor (t-l)/d \big\rfloor \Big\}.
The \lfloor (t-l)/d \rfloor bound means only completed blocks are compressed — the tokens too recent to fill a block are the window branch’s job.
Selection (Eq. 8–9). The query already attended the compressed tokens, giving scores p^{\text{cmp}}_t = \operatorname{softmax}(q_t^\top \tilde K^{\text{cmp}}_t). Reuse them as block importance — no new query·key products:
p^{\text{slc}}_t[j] = \sum_{m=0}^{l'/d - 1}\ \sum_{n=0}^{l/d - 1} p^{\text{cmp}}_t\big[(l'/d)\,j + m + n\big].
Take the top-n blocks by importance (Eq. 11–12), always keeping the first block and the local block, and attend their original keys/values at full resolution.
Sliding window. \tilde K^{\text{win}}_t = k_{t-w:t} — the last w tokens.
The paper’s models use l=32, d=16, selection block l'=64, n=16 selected blocks (including 1 initial + 2 local), and window w=512.
NoteA teaching simplification
Our from-scratch nsa.py aligns the block sizes (l=l'=d, non-overlapping), which makes the Eq. 9 double sum collapse to p^{\text{slc}}=p^{\text{cmp}} — the cleanest possible statement of “selection reuses compression’s scores.” We still build the general Eq. 9 (selection_importance) so you can see the faithful form, and we keep attention single-head; the paper sums importance across the query heads sharing a KV group (Eq. 10).
Walk one query through the whole thing:
Code: NSA from Scratch
The whole mechanism lives in nsa.py. First, compression — mean-pool (the default \varphi) each block into one token:
import torch
from nsa import compress_blocks
torch.manual_seed(0)
keys = torch.randn(8, 4) # 8 keys, dim 4
summaries = compress_blocks(keys, block_size=2)
print(f"{keys.shape[0]} keys → {summaries.shape[0]} block summaries")8 keys → 4 block summariesThen the crux: the block importance scores are the compression attention scores themselves. In the aligned case that reuse is an exact identity:
from nsa import selection_importance
p_cmp = torch.softmax(torch.randn(4), dim=-1) # compression attention
p_slc = selection_importance(p_cmp, cmp_block=2, cmp_stride=2, sel_block=2)
print("importance == compression scores:", torch.allclose(p_slc, p_cmp))importance == compression scores: TrueWith importance in hand, select_blocks keeps the top-n (forcing in the initial and local blocks), and the full layer ties the three branches together with a learned gate:
from nsa import NativeSparseAttention
layer = NativeSparseAttention(dim=16, block_size=4, num_selected=8, window=8)
x = torch.randn(2, 20, 16) # (batch, seq, dim)
out = layer(x)
print(f"Output shape: {out.shape}")Output shape: torch.Size([2, 20, 16])Selection Comes (Almost) for Free
A naive “read the most relevant blocks” scheme would first score every block — another full pass of query·key products. NSA avoids that entirely: the scores it needs are the softmax weights the compression branch already produced. Eq. 9 just re-buckets them from compression blocks into (possibly larger) selection blocks; when the blocks align it is the identity you just saw, and when they don’t it is a sum of a handful of numbers. Choosing where to look sharply costs no new attention.
Because the block choice is discrete (a top-n, like routing to experts in Module 11), it isn’t differentiable — but the gate, the compression weighting, and the attention over the chosen tokens all are, so the whole layer trains end to end:
from nsa import demonstrate_trainable
grads = demonstrate_trainable()
print("gradients reached every branch:",
all(grads[k] for k in ["cmp_weight", "gate", "q_proj", "o_proj"]))gradients reached every branch: TrueIt Stays Causal and Sparse
Two properties make NSA a legitimate drop-in for causal attention. First, the window branch is exactly dense attention when the window covers the whole past — NSA contains full attention as a special case, so it can only add the ability to look further for less:
from nsa import nsa_attention, full_causal_attention
q, k, v = (torch.randn(12, 8) for _ in range(3))
gate = torch.zeros(12, 3); gate[:, 2] = 1.0 # window branch only
win_only = nsa_attention(q, k, v, block_size=3, num_selected=2, window=12, gate=gate)
print("window (w≥seq) == full attention:",
torch.allclose(win_only, full_causal_attention(q, k, v), atol=1e-6))window (w≥seq) == full attention: TrueSecond, it is causal: rewriting the future never changes an earlier query’s output, because every branch only ever reads positions \le t.
layer = NativeSparseAttention(dim=8, block_size=4, num_selected=3, window=5)
seq = torch.randn(20, 8)
before = layer(seq)
seq[15:] = torch.randn(5, 8) # rewrite the future
after = layer(seq)
print("earlier outputs unchanged:", torch.equal(before[:15], after[:15]))earlier outputs unchanged: TrueAnd the payoff — where the past query reads far fewer keys than dense attention. The compression branch reads t/l summaries, selection reads n\cdot l tokens, the window reads w: a budget that grows with the sequence l-times more slowly than dense attention’s linear t.
from nsa import demonstrate_cost, nsa_selection_trace
cost_rows = demonstrate_cost()
trace = nsa_selection_trace(seq_len=20, block_size=4, num_selected=2, window=3)
ojs_define(nsaCost = cost_rows, nsaTrace = trace)Interactive: Which Blocks Does a Query Read?
Drag through the query positions and watch each query’s read pattern form — the compressed coarse view (dim), the few full-resolution blocks it selected (solid), and its sliding window (accent). Early queries see everything; later ones read a sparse, chosen slice.
The read-budget gap, plotted across sequence length: dense attention’s cost is the diagonal t; NSA’s flattens toward the compression-plus-window budget. On a 64k sequence the last query reads ~20× fewer keys.
TipTry This
- Slide
focusQueryfrom 0 to the end. Early queries have no completed blocks, so they read only their window — NSA falls back to dense local attention until there is enough past to compress and select. - Watch the selected blocks (solid) jump around as the query changes. They are chosen by content (the compression scores), not by a fixed offset — this is the difference from the sliding-window mask above.
- Read the cost chart’s crossover. Below the sparse budget the two lines sit together (NSA reads everything); past it they separate and the gap widens with length. Sparse attention only pays off once the sequence outgrows the budget.
NoteKey Insight
NSA is the fixed masks of the last section made learnable. A sliding window is NSA with only the window branch; a global-token pattern is a crude, hand-picked selection branch. By letting the query pick its blocks from scores it already had, NSA gets long-range reach at a cost that grows like t/l + n\,l + w instead of t — and, because it is trained this way from the start, without the quality hit of bolting sparsity onto a pretrained dense model.
DeepSeek Sparse Attention
NSA’s closing warning was that you pay a quality tax for “bolting sparsity onto a pretrained dense model.” DeepSeek Sparse Attention (DSA), the attention that shipped in DeepSeek-V3.2-Exp (2025), does exactly that bolting — and gets away with it. It is the third point on the sparsity spectrum this module has walked:
| pattern | granularity | trained | |
|---|---|---|---|
| Mask gallery | fixed (sliding-window / dilated) | positions | not learned |
| NSA | learned | whole blocks | sparse from scratch |
| DSA | learned | individual tokens | onto a dense model |
DSA keeps the whole expensive attention stack the frontier already uses — it was continued-trained on top of DeepSeek-V3.1-Terminus, a 685B-parameter MoE whose attention is the Multi-head Latent Attention you built earlier in this lesson — and simply teaches each query to read fewer tokens. Two cheap pieces do it:
- A lightning indexer: a lightweight, few-head, ReLU scorer that gives every past token an index score against the current query. It decides what is worth reading.
- A top-k selector: each query keeps only the k highest-scoring past tokens (DeepSeek uses k = 2048); the real attention reads only those. A fixed budget turns the main attention’s cost from O(n^2) toward O(n\cdot k).
NoteKey Insight
DSA and NSA both learn what to read, but they differ on two axes. Granularity: NSA picks whole blocks; DSA picks individual tokens, at the finest grain. Origin: NSA is trained sparse from the first step, while DSA is grafted onto a finished dense model. The graft survives because of a dense warm-up: before the top-k ever fires, the indexer is trained to imitate the dense model’s own attention, so it already points at the tokens the model was going to read anyway.
The Math: Index, Select, Attend
The lightning indexer scores query position t against past position s with a sum over H_I lightweight indexer heads (the DeepSeek-V3.2 report’s index-score equation):
I_{t,s} = \sum_{j=1}^{H_I} w_{t,j}\;\operatorname{ReLU}\!\big(q^{I}_{t,j}\cdot k^{I}_{s}\big).
Three deliberate choices make this an order of magnitude cheaper than the main attention: it uses few heads (H_I = 64 in V3.2, half the 128 main heads), the mixing weight w_{t,j} is a scalar the query carries, and the activation is ReLU, not softmax — cheap, and able to zero out irrelevant tokens. In production the whole indexer runs in FP8. It is still O(n^2) — but a tiny O(n^2), which is why only the main attention needs to be made sparse.
Given the scores, selection is a plain top-k:
\mathcal{S}_t = \operatorname{top\text{-}}k\big\{\, I_{t,s} : s \le t \,\big\}, \qquad o_t = \sum_{s \in \mathcal{S}_t} \operatorname{softmax}\!\Big(\tfrac{q_t \cdot k_s}{\sqrt{d_k}}\Big)\, v_s .
The expensive attention runs only over \mathcal{S}_t — at most k keys per query instead of all t. Walk one query through it:
Code: DSA from Scratch
The whole mechanism lives in dsa.py. First the lightning indexer — the cheap scorer that ranks the past. Notice the ReLU and the per-head, query-carried weight:
import torch
from dsa import lightning_index_scores, topk_select
torch.manual_seed(0)
q_idx = torch.randn(6, 2, 4) # 6 tokens, H_I=2 indexer heads, d_idx=4
k_idx = torch.randn(6, 4) # shared indexer keys (one vector per token)
weight = torch.randn(6, 2) # per-query, per-head mixing weight
scores = lightning_index_scores(q_idx, k_idx, weight)
print("index scores are causal (future = -inf):", bool(torch.isinf(scores[0, 1])))index scores are causal (future = -inf): TrueThe selector turns those scores into a per-query keep-mask — the top-k past tokens (always including the query itself, so the softmax is never empty):
keep = topk_select(scores, k=2)
print("tokens query 5 keeps:", torch.nonzero(keep[5]).flatten().tolist())
print("each row keeps min(k, t+1):", [int(keep[t].sum()) for t in range(6)])tokens query 5 keeps: [2, 5]
each row keeps min(k, t+1): [1, 2, 2, 2, 2, 2]And the full forward reads only the kept tokens with ordinary attention:
from dsa import DeepSeekSparseAttention
layer = DeepSeekSparseAttention(dim=16, k_top=4, n_index_heads=2, index_dim=8)
x = torch.randn(2, 20, 16) # (batch, seq, dim)
print(f"Output shape: {layer(x).shape}")Output shape: torch.Size([2, 20, 16])Why the Indexer Needs Its Own Teacher
Here is the subtlety that makes DSA more than “attention with a mask.” The indexer influences the output only through the discrete top-k selection — and a discrete choice has zero gradient almost everywhere. So the language-modeling loss, flowing back through the kept attention, trains the projections but can never teach the indexer what to select. The two must be trained by two different objectives:
from dsa import demonstrate_gradient_paths
paths = demonstrate_gradient_paths()
print("task loss → main projections:", paths["task_reaches_main"])
print("task loss → indexer: ", paths["task_reaches_indexer"]) # False!
print("imitation loss → indexer: ", paths["imitation_reaches_indexer"])task loss → main projections: True
task loss → indexer: False
imitation loss → indexer: TrueThe indexer’s own objective is imitation: on a frozen dense model, train the indexer so its softmax matches the model’s real attention distribution — a KL divergence,
\mathcal{L}^{I} = \sum_t D_{\mathrm{KL}}\!\big(p_t \,\big\|\, \operatorname{softmax}(I_{t,\cdot})\big),
where p_t is the dense model’s attention for query t (in real DSA, summed over heads and normalized). Drive that KL down and the indexer learns to put the dense model’s heavy hitters exactly where its top-k will find them.
The Dense Warm-up
That imitation step is the whole reason a pretrained dense model can be sparsified without falling apart. Run it — freeze the main model, train only the indexer, and watch the KL fall while the indexer’s top-k agreement with the dense attention climbs:
from dsa import demonstrate_warmup, demonstrate_cost, DeepSeekSparseAttention
warm = demonstrate_warmup(seq_len=24, steps=60)
# A live index-score matrix to drive the keep-map (replace -inf with null for JSON).
torch.manual_seed(0)
_layer = DeepSeekSparseAttention(dim=8, k_top=4, n_index_heads=2, index_dim=8)
_x = torch.randn(18, 8)
_scores = _layer.index_scores(_x)
dsa_score_matrix = [
[None if s == float("-inf") else float(s) for s in row]
for row in _scores.tolist()
]
ojs_define(dsaWarm = warm)
ojs_define(dsaScores = dsa_score_matrix)
ojs_define(dsaCost = demonstrate_cost())
TipTry This
The KL (dashed) falls and the agreement (solid) rises together: as the indexer learns to mimic the dense attention, its cheap top-k starts to recover the very tokens the dense model attends to. That crossing is DSA’s whole bet — once the indexer agrees with the dense model, turning on the top-k costs almost no quality, and the model can be continued-trained sparse.
It Stays Causal — and Reduces to Dense
Two properties make DSA a legitimate drop-in. First, when the budget is at least the sequence length, DSA is exactly dense attention — it can only ever skip tokens, never invent them, so a generous k recovers the full result:
from dsa import dsa_attention, full_causal_attention
q, k, v = (torch.randn(12, 8) for _ in range(3))
q_idx, k_idx, w = torch.randn(12, 2, 8), torch.randn(12, 8), torch.randn(12, 2)
dense_budget = dsa_attention(q, k, v, q_idx, k_idx, w, k_top=12) # k ≥ seq
print("k ≥ n == full attention:",
torch.allclose(dense_budget, full_causal_attention(q, k, v), atol=1e-6))k ≥ n == full attention: TrueSecond, it is causal: the indexer scores, the selection, and the attention all read only positions \le t, so rewriting the future never disturbs an earlier output:
layer = DeepSeekSparseAttention(dim=8, k_top=3, n_index_heads=2, index_dim=8)
seq = torch.randn(20, 8)
before = layer(seq)
seq[15:] = torch.randn(5, 8) # rewrite the future
after = layer(seq)
print("earlier outputs unchanged:", torch.equal(before[:15], after[:15]))earlier outputs unchanged: TrueInteractive: Which Tokens Does a Query Keep?
Drag through the query positions and the budget k. Each query’s kept tokens (accent) are chosen by the indexer’s content scores — a fine-grained, per-query slice of the past — while the rest (dim) are skipped. Turn k up and the pattern fills back in toward dense; turn it down and each query reads only its few most relevant tokens.
The read-budget gap, plotted across sequence length: dense attention’s main attention reads the diagonal t; DSA’s flattens at the fixed budget k = 2048. On a 128k sequence the last query’s main attention reads 64× fewer keys — while the lightning indexer stays cheap enough that it never becomes the bottleneck.
TipTry This
- Slide
dsaKdown to 1. Every query collapses to reading only itself — DSA with a budget of one. Slide it up past the sequence length and the map fills solid: DSA becomes dense attention, exactly the anchor you proved above. - Sweep
dsaFocusacross the queries. The kept columns move — they are chosen by the indexer’s content scores, not a fixed offset. That is the difference from the sliding-window mask earlier in this lesson. - Read the cost chart’s plateau. Dense grows with t; DSA’s main attention pins at k. The two lines sit together until the sequence outgrows the budget, then separate — sparse attention only pays once you are past k tokens.
NoteKey Insight
DSA is the production answer to the quadratic wall that keeps the frontier recipe intact. It does not replace MLA, MoE, or the pretrained weights — it inserts a cheap indexer in front of the attention and lets each query read a fixed k tokens. The trick that makes grafting sparsity onto a finished dense model work is training the indexer to imitate that model first, so the cheap selection agrees with the expensive attention it is standing in for.
Compressed Attention: CSA and HCA
Every technique so far in this section selects which past tokens to read — masks pick them by position, NSA by block, DSA by token — but each one reads its chosen tokens at full resolution, from a KV cache that is still n entries long. At a million tokens that cache is the wall: even reading a sparse k of it, you must store all n keys and values. DeepSeek-V4 (2026) attacks the store itself. The idea is compression: pool every m consecutive past tokens into a single compressed entry, so the cache the query reads from is n/m long — and keep only the most recent tokens uncompressed, in a sliding window, where full resolution still matters.
That one move splits into the two attention types V4 interleaves across its layers:
| far cache | selection | reads | |
|---|---|---|---|
| CSA — Compressed Sparse Attention | compress every m tokens | top-k compressed entries (the DSA indexer) | \min(k, n/m) + w |
| HCA — Heavily Compressed Attention | compress every m' \gg m tokens | none — read them all | n/m' + w |
CSA is DSA’s natural sequel: it reuses the same lightning indexer and top-k — but pointed at compressed entries, so both the selection and the thing selected are cheap. HCA takes the other road: compress so hard that the far cache is tiny, then just read all of it, no indexer needed. The near/far split — a sliding window of recent uncompressed tokens vs. a compressed far past — lives inside every layer; CSA vs. HCA is the choice of how hard to compress, made per layer.
NoteKey Insight
DSA made attention sparse; CSA/HCA make the cache small. Selection reads fewer of n entries; compression makes there be fewer than n entries to begin with. V4 combines both — a compressed cache you then select from (CSA) — and the payoff compounds: in the one-million-token setting, DeepSeek-V4-Pro reports needing only about 27% of the single-token inference FLOPs and 10% of the KV cache of DeepSeek-V3.2.
The Math: Compress, Select, Attend
Compression is a pooling of the far past. For a completed block of m keys, the compressed key is their average (V4 uses a learned, weighted pool; we use the uniform mean — the same stand-in NSA’s compression branch uses, and the reason m=1 is exactly the identity):
C^{K}_{i} = \frac{1}{m}\sum_{j=mi}^{m(i+1)-1} k_j, \qquad C^{V}_{i} = \frac{1}{m}\sum_{j=mi}^{m(i+1)-1} v_j .
Split each query t’s causal past into a near window (the last w positions, kept uncompressed) and a far part (everything before it, compressed). CSA scores the compressed far entries with the lightning indexer, keeps the top-k, and runs one softmax over the selected entries plus the window:
\mathcal{S}_t = \operatorname{top\text{-}}k\{\, I_{t,i} : C^{K}_{i} \text{ is far} \,\}, \qquad o_t = \operatorname{softmax}\!\Big(\tfrac{q_t \cdot \text{key}}{\sqrt{d_k}}\Big)\,\text{value}, \;\; \text{over } \{C_i : i \in \mathcal{S}_t\} \cup \{\text{window}\}.
HCA drops the \operatorname{top\text{-}}k: it compresses the far past with a much larger m' and attends densely over every compressed entry plus the window. Walk one CSA query through the pipeline:
Code: Compressed Attention from Scratch
The whole mechanism lives in csa.py, and it reuses the DSA indexer you just built — never a copy. First, compression: mean-pool complete blocks of the past into a shorter cache.
import torch
from csa import compress_kv
torch.manual_seed(0)
past = torch.randn(12, 8) # 12 far KV tokens
compressed = compress_kv(past, m=4) # → 3 compressed entries
print(f"cache {past.shape[0]} tokens → {compressed.shape[0]} entries (m=4)")
print("m=1 is the identity:", bool(torch.equal(compress_kv(past, m=1), past)))cache 12 tokens → 3 entries (m=4)
m=1 is the identity: TrueCSA compresses the far past, scores the entries, keeps the top-k, and attends over the selected entries plus the recent window:
from csa import csa_attention
q, k, v = (torch.randn(24, 8) for _ in range(3))
out = csa_attention(q, k, v, m=4, k_top=2, window=4)
print(f"CSA output shape: {out.shape}")CSA output shape: torch.Size([24, 8])HCA compresses harder and drops the selection — dense attention over a tiny far cache:
from csa import hca_attention
out_hca = hca_attention(q, k, v, m_far=8, window=4)
print(f"HCA output shape: {out_hca.shape}")HCA output shape: torch.Size([24, 8])It Reduces to Dense — and Stays Causal
Compression only coarsens the past; with no compression and no selection budget, CSA must reproduce full attention exactly. Set m = 1 (each entry is one token), k \ge n (select everything), and w = 0 (no separate window):
from csa import full_causal_attention
dense_ref = full_causal_attention(q, k, v)
csa_dense = csa_attention(q, k, v, m=1, k_top=24, window=0)
print("CSA(m=1, k≥n, w=0) == dense:",
torch.allclose(csa_dense, dense_ref, atol=1e-6))CSA(m=1, k≥n, w=0) == dense: TrueHCA reduces the same way at m' = 1. (Real V4 uses a learned weighted pool, so it only approximately recovers dense at m=1 — the mean-pool here makes the anchor exact.) And like every attention in this module, it is causal — the compression, the selection, and the softmax all read only positions \le t:
from csa import CompressedAttention
layer = CompressedAttention(dim=8, m=2, window=4, k_top=3, mode="csa")
seq = torch.randn(20, 8)
before = layer(seq)
seq[15:] = torch.randn(5, 8) # rewrite the future
after = layer(seq)
print("earlier outputs unchanged:", torch.equal(before[:15], after[:15]))earlier outputs unchanged: TrueThe Compression Dial: Cache vs Error
Compression is not free — it is a trade. As m grows, the cache shrinks like 1/m, but each compressed entry blurs more tokens together, so the attention output drifts from the dense result. Sweep the dial and watch both move:
from csa import compression_tradeoff, demonstrate_cost, csa_read_trace
tradeoff = compression_tradeoff(seq_len=64, dim=16, ratios=(1, 2, 4, 8, 16),
k_top=4, window=4)
ojs_define(compTradeoff = tradeoff)
ojs_define(compCost = demonstrate_cost())
ojs_define(compTrace = csa_read_trace(seq_len=20, m=2, k_top=3, window=4))
TipTry This
- Read the two axes against each other. At m=1 the cache is full and the error is ~0 — CSA is dense. Turn m up: the blue cache bars fall like 1/m while the orange error climbs. That gap is the whole design question — how much blur will your task tolerate to fit a million tokens?
- Compare the solid and dashed error lines. Solid is CSA (top-k over compressed), dashed is HCA (dense over compressed). When the budget k already covers every compressed entry they coincide; below that, selection lets CSA read a smaller set at the same m.
Interactive: The Two-Tier Cache
Each row is one query; each column a past position. A query reads its recent window (green, full resolution) and a selected set of compressed chunks (orange) — the far past, coarsened. Drive the compression ratio, the selection budget, and the window, and watch the read pattern and the KV-cache budget move.
TipTry This
- Set
compMto 1 andcompWto 0. Every query reads its whole causal past at full resolution — CSA collapses to dense, the anchor you proved above. Now raisecompM: the orange far cells thin out as tokens fold into chunks, and the “ratio” climbs. - Widen
compW. The green diagonal band thickens — more of the recent past is read uncompressed. This is the near tier; a window that swallows the whole sequence is dense again. - Drop
compKto 1. Each query keeps a single compressed chunk of far context plus its window — the aggressive end of CSA. HCA is the opposite corner: no selection, but the far cache compressed so hard that reading all of it is still cheap.
NoteKey Insight
CSA and HCA move the lever from which tokens to how many entries exist. A compressed cache is a lossy summary of the far past — cheaper to store and read, at the cost of blur that grows with m. V4’s bet is that most far context only needs a summary, while the recent window and a few selected chunks carry the detail — which is why it can hold a million tokens at a tenth of the cache.
Forgetting Attention: A Learned Recency Decay
Every lever so far changes which keys a query reads (GQA shares them, MLA compresses them, NSA/DSA/CSA select or shrink them) or how the sum is streamed (FlashAttention, paging). Forgetting Attention pulls an orthogonal one: it leaves the key set whole and instead decides how much each key still counts, by how long ago it was — and it lets the model learn that decay from the content itself.
You have already met two of the three corners of this idea. Back in m04, ALiBi subtracted a fixed penalty proportional to distance — a hand-set recency bias. In m22 the forget gate of Gated DeltaNet and GLA multiplied a recurrent state by a learned, per-token decay. The Forgetting Transformer (FoX; Lin et al., 2025) fills the missing corner: a learned, data-dependent decay on full softmax attention.
| decay on the scores | fixed | learned (data-dependent) |
|---|---|---|
| softmax attention | ALiBi (m04) | FoX — this section |
| linear / recurrent | RetNet decay | GLA, Gated DeltaNet (m22) |
The mechanism is one scalar per token. Each position emits a forget gate
f_t = \sigma(\mathbf{w}_f \cdot \mathbf{x}_t + b_f) \in (0, 1),
read as “how much of the running past should still be visible one step later.” A gate near 1 keeps the past; a gate near 0 wipes it. Because the gate is a function of \mathbf{x}_t, a token can decide — from its own content — to forget everything before it. ALiBi can only ever apply the same slope.
NoteKey Insight
FoX is not a new attention operator — it is ordinary softmax attention with one extra additive term on the scores. That is exactly why it inherits everything you built: it is still causal, still exact, still FlashAttention-compatible, and it needs no positional embedding at all — the learned decay is the position signal.
The Math: A Product of Gates, Added in Log-Space
The forget gate between a query at i and a key at j\le i is the running product of every gate strictly between them:
F_{ij} = \prod_{l=j+1}^{i} f_l \quad (F_{ii} = 1), \qquad D_{ij} = \log F_{ij} = \sum_{l=j+1}^{i} \log f_l .
F_{ij}\in(0,1] is the fraction of key j that survives to query i — full for the token itself, shrinking the farther back you look. Since each \log f_l<0, D_{ij} is a non-positive decay bias, and Forgetting Attention just adds it to the usual scaled scores before the softmax:
\mathbf{O} = \operatorname{softmax}\!\left(\frac{QK^\top}{\sqrt{d_k}} + D\right)V, \qquad D \text{ lower-triangular},\; D_{ij}=-\infty \text{ for } j>i .
The products look expensive, but write c_i = \sum_{l\le i}\log f_l (one cumulative sum) and every entry is a difference:
D_{ij} = c_i - c_j .
That is the whole state: one running scalar per key. And because c_i is constant along a query’s row, it cancels inside that row’s softmax — so the entire T\times T decay collapses to a single per-key bias -c_j. This shift-invariance is why FoX fuses cleanly into a FlashAttention kernel (nothing N\times N is ever built) and why the two anchors below are exact, not approximate.
Code: Forgetting Attention from Scratch
The full implementation lives in forgetting_attention.py; the core is four short functions — the gate, its cumulative sum, the decay matrix, and the two-line change to attention.
import torch
from forgetting_attention import (
forgetting_attention, standard_causal_attention,
decay_bias, constant_forget_gate, alibi_bias,
forgetting_attention_via_keybias,
)
torch.manual_seed(0)
q, k, v = (torch.randn(2, 6, 16) for _ in range(3))
f = torch.rand(2, 6) * 0.3 + 0.6 # a per-token forget gate in (0.6, 0.9)
out = forgetting_attention(q, k, v, f)
print("output shape:", tuple(out.shape)) # (2, 6, 16): same shape as plain attentionoutput shape: (2, 6, 16)Two Exact Anchors
FoX is bracketed by two things you already trust. Push the gate to 1 and the decay vanishes — you are back to the plain causal attention of m05/m09. Freeze the gate at a constant f = e^{-m} and the decay becomes D_{ij} = -m\,(i-j): the straight-line penalty of ALiBi, exactly. FoX is ALiBi with the slope set free to depend on the tokens.
# Anchor 1 — f ≡ 1 recovers standard causal attention, bit-for-bit.
ones = torch.ones(2, 6)
a1 = forgetting_attention(q, k, v, ones)
a2 = standard_causal_attention(q, k, v)
print("f=1 == plain causal attention :", torch.allclose(a1, a2, atol=1e-6))
# Anchor 2 — a constant gate f = exp(-m) is exactly ALiBi's linear bias.
m = 0.5
D = decay_bias(constant_forget_gate(6, m))
A = alibi_bias(6, m)
finite = ~torch.isinf(A)
print("constant f == ALiBi bias :", torch.allclose(D[finite], A[finite], atol=1e-6))
# The shift-invariant per-key form (-c_j) gives the identical output.
print("one-cumsum form == full matrix:",
torch.allclose(forgetting_attention(q, k, v, f),
forgetting_attention_via_keybias(q, k, v, f), atol=1e-6))f=1 == plain causal attention : True
constant f == ALiBi bias : True
one-cumsum form == full matrix: TrueDrive the gate below and watch the decay matrix F_{ij} (cell opacity = fraction of key j that survives to query i) slide between those two anchors:
TipTry This
- Leave the gate at 1.0 — the grid is a solid causal triangle: every past key counts fully. This is the plain attention from m05.
- Slide it down. The lower-left corner fades first: distant keys decay away while recent ones stay bright. The read-out shows the equivalent ALiBi slope — a constant gate is ALiBi.
- A constant gate can only tilt the whole triangle. The next widget shows what a data-dependent gate buys you.
Data-Dependence: A Forget Gate Can Reset
ALiBi’s slope is fixed for all time; FoX’s gate is read off each token, so a single position can drive f_t\to 0 and cut the past — a content-triggered segment break (a document boundary, a “forget that” instruction). The demo below builds three gate profiles over one sequence and shows the last query’s decay and the attention mass it lands on. Only the third — a mild gate with one near-zero spike — develops the hard boundary.
from forgetting_attention import demonstrate_forgetting
_demo = demonstrate_forgetting(seq_len=12, reset_at=6, seed=0)
ojs_define(foxDemo = _demo)_d = demonstrate_forgetting(seq_len=12, reset_at=6, seed=0)
print(f"mass the last query keeps *before* the reset boundary:")
print(f" mild gate : {_d['mass_before_boundary_mild']:.3f}")
print(f" reset gate : {_d['mass_before_boundary_reset']:.3f} (the spike wiped it)")mass the last query keeps *before* the reset boundary:
mild gate : 0.356
reset gate : 0.000 (the spike wiped it)
NoteKey Insight
With the fixed gates (profiles 0 and 1) the last query still spreads some weight across the whole past. With the reset gate, the near-zero spike sends D_{ij}\to-\infty across the boundary, so the query’s attention mass collapses onto the tokens after the reset — the model has learned to start a fresh context. That is the capability a data-independent slope cannot express, and it is why FoX improves long-context modeling and length extrapolation over both a vanilla Transformer and ALiBi.
The paper’s full “FoX Pro” block layers a few extra recurrent-model conveniences on top of this core — an output gate \sigma(W_g\mathbf{x}_t), output RMSNorm, and QK-Norm (all of which you have already met: QK-Norm in m05, RMSNorm in m06, output gates in m22). The forget gate above is the one essential new part; the rest is reuse.
Common Pitfalls
When implementing efficient attention, watch out for:
- Expanding K/V before caching, not after. Store the small K/V (
num_kv_heads) in the cache andrepeat_interleaveonly on the way into the multiply. Expanding first throws away the entire memory saving. - Wrong grouping order.
repeat_interleave(n_rep, dim=1)maps query heads[g·n_rep : (g+1)·n_rep]to K/V headg. Usingrepeat/tileinstead interleaves the groups differently and silently mismatches heads. num_kv_headsmust dividenum_heads. Otherwise the groups are uneven; the layer asserts this at construction.- Forgetting to rescale in the online softmax. If you skip the e^{m_{\text{old}} - m_{\text{new}}} correction when the running max grows, the earlier blocks are normalized against the wrong maximum and the result is wrong. The rescale is what keeps streaming exact.
- All-masked rows. A sliding window plus causal masking is fine (the diagonal is always attended), but an over-aggressive custom mask can leave a row with no allowed keys, producing
NaNfromsoftmaxof all-inf. - Applying RoPE to the MLA content path. The whole reason MLA can cache a tiny latent is that the content up-projections absorb into one matrix — which only works if no position rotation sits between them. Position must ride on the separate decoupled key; rotate the content keys and the absorption (and the memory win) is gone.
- Scaling MLA scores by
√d_headinstead of√(d_head + d_h^R). The query and key are[content ; rope]concatenations, so the correct denominator uses their combined width. - A block size that is too large. Paging bounds waste at
block_size − 1tokens per sequence, so a huge block (say 512) reintroduces exactly the internal fragmentation paging was meant to kill. Too small (1) wastes nothing but reads the KV in scattered single-token chunks. 16 is the usual balance. - Admitting a request you can’t finish. Continuous batching admits greedily, but a decode step can demand a new block for any running sequence. Gate admission on the pool fitting each request’s whole footprint (prompt + output), or reserve those blocks up front; otherwise the batch OOMs mid-generation. A real server instead preempts — evicting a low-priority sequence and recomputing it later — which is the harder policy this section stops short of.
- Recomputing NSA’s block-importance scores. The selection branch’s scores are the compression attention weights — reuse them (Eq. 8–9). Running a fresh query·key pass to rank blocks throws away the efficiency that motivates the design and double-counts compute.
- Compressing the current partial block. Only completed blocks (up to \lfloor (t-l)/d \rfloor) are compressed; the most recent tokens are read by the sliding window instead. Fold them into a half-empty compressed token and you both leak a padded average and starve the window branch of its purpose.
- Compressing the near window (CSA/HCA). The whole point of the sliding window is full-resolution recent context. Compress the last w tokens along with the far past and you blur exactly the tokens the local branch exists to protect — and you lose the exact reduction to dense at m=1. Compress only the far part, strictly before the window.
- Double-counting the far/near boundary. The compressed far entries and the uncompressed window must partition the causal past, not overlap. If a token appears both in a compressed chunk and in the window, the softmax reads it twice and its contribution is inflated. The far part is positions \le t-w; the window is the last w.
- Chunking the prefill too small. The token budget bounds the TBT tail, but a tiny chunk is not free: each chunk’s attention must read the whole prefix prefilled so far, so slicing a prompt into many chunks re-reads that growing prefix again and again, adding real FLOPs. The
chunked_prefill.pysimulation counts scheduled token work (the quantity the budget bounds), not this quadratic prefix re-read, so itspeak_workis a faithful stand-in for the latency the budget caps — not a full FLOP model. Real systems keep the budget moderate — large enough that the re-read overhead stays small, small enough that decodes never stall. - Applying the forget decay in probability space instead of log-space. Multiplying the normalized softmax weights by F_{ij} breaks the distribution (the rows no longer sum to one) and is not what FoX does. The decay is added to the pre-softmax scores as D=\log F, so the softmax renormalizes over the surviving mass. Add in log-space, then exponentiate once.
- Forgetting that a constant gate is already ALiBi. If your forget logits collapse to a per-head constant (e.g. the projection \mathbf{w}_f trains to zero), FoX silently degenerates to fixed ALiBi — correct, but you have thrown away the data-dependence that was the point. The gate must actually read the token; watch that \mathbf{w}_f does not vanish.
Exercises
Exercise 1: Cache size for a real model
from attention import kv_cache_bytes
# Llama-2 70B: 80 layers, 64 query heads, head_dim 128, GQA with 8 KV heads.
# How much smaller is the KV-cache than it would be under MHA, at 4096 tokens?
# Your implementation here:
# mha = kv_cache_bytes(...)
# gqa = kv_cache_bytes(...)
# print(f"GQA cache is {mha / gqa:.0f}x smaller")Exercise 2: A single dial
import torch
from attention import GroupedQueryAttention
# Build the three variants of an embed_dim=64, 8-head layer and confirm that
# only the K/V projection width changes. What is n_rep for each?
# Your implementation here:
# for kv in (8, 4, 1):
# layer = GroupedQueryAttention(embed_dim=64, num_heads=8, num_kv_heads=kv)
# ...Exercise 3: Streaming a longer sequence
import torch
from attention import online_softmax_attention
# Verify the online softmax stays exact as the block size varies. Try block sizes
# 1, 7, and larger than the sequence length; each should match the naive result.
# Your implementation here:Exercise 4: MLA vs GQA at a fixed cache budget
from attention import mla_cache_bytes, kv_cache_bytes, mla_equivalent_gqa_groups
# For DeepSeek-V2's shape (128 heads, head_dim 128, 60 layers, 8192 tokens):
# how many GQA groups would match an MLA cache with d_c=512, d_h^R=64?
# And how much smaller is MLA than full MHA?
# Your implementation here:
# groups = mla_equivalent_gqa_groups(...)
# mha = kv_cache_bytes(...)
# mla = mla_cache_bytes(...)
# print(f"MLA ≈ {groups} GQA groups, {mha / mla:.0f}x smaller than MHA")Exercise 5: Paging vs contiguous waste
from paged_attention import paged_used_slots, contiguous_reserved_slots
# Eight requests, each reaching only 50 tokens, on a server that reserves
# max_seq_len=4096 contiguously. With block_size=16, how many slots does each
# scheme use, and what fraction does the contiguous scheme waste?
# Your implementation here:
# lengths = [50] * 8
# contiguous = contiguous_reserved_slots(4096, len(lengths))
# paged = paged_used_slots(lengths, 16)
# print(f"contiguous {contiguous}, paged {paged}, wasted {1 - sum(lengths)/contiguous:.1%}")Exercise 6: The hostage cost of static batching
from continuous_batching import Request, static_batch, continuous_batch, mean_latency
# A stream of short requests behind one long generation. Sweep max_running and
# compare mean completion latency under static vs continuous batching — the gap is
# the head-of-line-blocking cost. Confirm the outputs stay identical.
# Your implementation here:
# work = [Request(0, 0, 4, 20)] + [Request(i, 0, 4, 2) for i in range(1, 8)]
# blocks = sum(r.footprint(4) for r in work)
# for mr in (2, 4, 8):
# s = static_batch(work, blocks, 4, mr)
# c = continuous_batch(work, blocks, 4, mr)
# print(mr, mean_latency(s, work), mean_latency(c, work))Exercise 7: NSA’s read budget vs dense attention
from nsa import nsa_kv_reads
# NSA's models use l=32, n=16, w=512. At what sequence length does the last
# query's NSA read budget first drop below the dense cost t? By 128k tokens, how
# many times fewer keys does it read?
# Your implementation here:
# for t in (1024, 4096, 16384, 65536, 131072):
# reads = nsa_kv_reads(t, block_size=32, num_selected=16, window=512)
# print(t, reads, f"{t / reads:.1f}x")Exercise 8: CSA vs HCA cache at a million tokens
from csa import csa_kv_entries, hca_kv_entries
# At 1M tokens, compare the KV cache CSA stores (m=8, k=512, window=512) with what
# HCA stores (m'=64, window=512), and each against dense. Which is smaller, and by
# how much? (CSA stores its whole compressed far cache; it *reads* only the top-k.)
# Your implementation here:
# csa = csa_kv_entries(1_048_576, m=8, k_top=512, window=512)
# hca = hca_kv_entries(1_048_576, m_far=64, window=512)
# print(f"dense {csa['dense']}, CSA stored {csa['stored']}, HCA stored {hca['stored']}")Exercise 9: The budget that just fits the decodes
from chunked_prefill import Request, chunked_prefill_batch, peak_iteration_work, mean_ttft
# Three requests decode (output_len 8) while one 40-token prompt is admitted.
# Sweep token_budget over 4, 5, 8, 16. What happens to peak work and mean TTFT when
# the budget barely exceeds the number of decodes (so prefill_budget = B - 3 is tiny)?
# Explain why a budget of 4 starves the prefill until the decoders finish.
# Your implementation here:
# work = [Request(i, 0, 4, 8) for i in range(3)] + [Request(3, 0, 40, 4)]
# blocks = sum(r.footprint(4) for r in work)
# for B in (4, 5, 8, 16):
# r = chunked_prefill_batch(work, blocks, 4, 4, token_budget=B)
# print(B, peak_iteration_work(r), r.makespan, round(mean_ttft(r, work), 1))Exercise 10: FoX between its two anchors
from forgetting_attention import forgetting_attention, decay_bias, constant_forget_gate, alibi_bias
import torch
# (a) Confirm the two exact anchors for a length-8 sequence:
# f ≡ 1 reproduces plain causal attention, and a constant f = exp(-m)
# reproduces the ALiBi bias for m in {0.25, 0.5, 1.0}.
# (b) Build a data-dependent gate that is 0.99 everywhere except ~0 at position 4,
# and show (via decay_bias) that a query at position 7 assigns almost no mass to
# keys 0..3. Which real behaviors (document boundaries, "ignore the above") does
# this let a model represent that a fixed slope cannot?
# Your implementation here:
# for m in (0.25, 0.5, 1.0):
# D = decay_bias(constant_forget_gate(8, m)); A = alibi_bias(8, m)
# print(m, torch.allclose(D[~torch.isinf(A)], A[~torch.isinf(A)], atol=1e-6))Summary
Key takeaways:
- Inference is memory-bound, not compute-bound — the KV-cache is re-read every decode step, so its size sets the ceiling on context length and batch.
- MHA, GQA, and MQA are one mechanism with one dial —
num_kv_headscontrols how many query heads share each key/value head; MHA and MQA are the endpoints. - The saving is real and free — the cache stores only
num_kv_headsheads;repeat_interleavere-expands them just for the multiply, leaving the attention math and (for GQA) most of the quality intact. - MLA is a different lever on the same wall — keep every head, but cache a single low-rank latent (plus one decoupled RoPE key) and reconstruct K/V from it. Its cache is worth a fraction of a GQA group; RoPE is decoupled because position rotation blocks the up-projection absorption.
- Cached decoding is exact — GQA and MLA through their caches reproduce a full forward pass to floating-point rounding, just like MHA in m08.
- FlashAttention is an exact reorganization — the online-softmax recurrence streams over key blocks and never materializes the N \times N matrix, cutting memory traffic without changing the result.
- Masks buy efficiency too — sliding-window and dilated attention attend to fewer positions, trading full coverage for cost that grows with the sequence rather than its square.
- Paging is memory management, not math — storing the cache in fixed-size blocks behind a block table cuts serving waste from ~90% to at most one partial block per request. It touches the attention result not at all, yet it is one of the biggest memory levers on a real server.
- Continuous batching is scheduling, not math — iteration-level scheduling admits and evicts requests between decode steps, so a short request never waits on a long one and freed slots refill immediately. On the hostage workload it cut mean latency 3× at higher occupancy while producing bit-for-bit identical K/V. Paging is what makes it cheap; together they are why vLLM/Orca-style servers keep the GPU saturated.
- Chunked prefill bounds the per-iteration stall — a prompt’s prefill is compute-bound and huge; sliced into budget-sized chunks with the decodes piggybacking, no single iteration balloons, so the inter-token latency (TBT) of everyone already generating stays flat. The token budget is the throughput ⇆ latency dial (tighter budget → lower TBT tail, later first token), and like paging and continuous batching it changes only scheduling, not the K/V.
- Native Sparse Attention makes the mask learnable — each query reads the past through three gated branches (coarse compression of every block, a full-resolution look at a few chosen blocks, and a local window), and the block choice reuses the compression scores so it is nearly free. It contains dense attention as a special case, stays causal, and trains sparse from scratch — reading \sim t/l + n\,l + w keys instead of t.
- DeepSeek Sparse Attention grafts sparsity onto a dense model — a cheap, few-head, ReLU lightning indexer scores every past token, a top-k selector keeps only the best k (DeepSeek: 2048), and the real attention reads only those, turning its cost from O(n^2) toward O(n\cdot k). It survives being bolted onto a pretrained model because a dense warm-up first trains the indexer to imitate the model’s own attention. Token-level and continued-trained, where NSA is block-level and from-scratch.
- Compressed attention shrinks the cache itself — where selection reads fewer of n entries, DeepSeek-V4’s CSA/HCA make there be fewer than n. Pool every m far tokens into one entry (a 1/m cache), keep the recent window uncompressed, and attend over [selected compressed chunks ++ window]. CSA selects the top-k compressed entries with the DSA indexer; HCA compresses harder (m' \gg m) and reads them all. It reduces to dense at m=1, stays causal, and trades cache for blur — the 2026 answer to the million-token wall (~10% of V3.2’s KV cache at 1M tokens).
- Forgetting Attention adds a learned recency decay to the scores — a scalar forget gate f_t=\sigma(\mathbf{w}_f\cdot\mathbf{x}_t+b_f) whose running product F_{ij}=\prod f enters as a log-space bias D_{ij}=c_i-c_j before the softmax. It is still plain causal attention (exact, FlashAttention-friendly, no positional embedding), it contains full attention (f\equiv1) and ALiBi (f\equiv e^{-m}) as exact special cases, and its data-dependence lets a single token reset the context — the missing (learned, softmax) corner of the decay-on-scores grid whose other corners are ALiBi (m04) and the linear-attention forget gate (m22).
What’s Next
You now have the efficiency toolkit — a small cache and cheap memory traffic — that long-context and fast-inference techniques are built on. Next, Module 10: Long Context uses exactly this toolkit to stretch context far past the training length (RoPE scaling, position interpolation, YaRN); from there the book turns to sparse expert models (MoE) and fast serving (quantization, speculative decoding). Each reuses the GroupedQueryAttention and online-softmax ideas you just built.
Going Deeper
Core Papers:
- Fast Transformer Decoding: One Write-Head is All You Need — Shazeer (2019), the original Multi-Query Attention.
- GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints — Ainslie et al. (2023), Grouped-Query Attention.
- DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model — DeepSeek-AI (2024), the paper that introduced Multi-head Latent Attention and the decoupled RoPE.
- DeepSeek-V3 Technical Report — DeepSeek-AI (2024), MLA at frontier scale (671B MoE).
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — Dao et al. (2022), tiling + online softmax.
- Native Sparse Attention: Hardware-Aligned and Natively Trainable Sparse Attention — Yuan et al. (DeepSeek, 2025), the compression + selection + window branches with a learned gate, trained sparse from scratch.
- DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models — DeepSeek-AI (2025), DeepSeek Sparse Attention: the lightning indexer, top-k token selection, and the dense warm-up that lets a pretrained dense model be sparsified. (Introducing DeepSeek-V3.2-Exp).
- DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligence — DeepSeek-AI (2026), Compressed Sparse Attention (compress the KV cache, then top-k select) and Heavily Compressed Attention (compress far harder, no selection), interleaved across layers; ~27% of the inference FLOPs and ~10% of the KV cache of V3.2 at one million tokens.
- Forgetting Transformer: Softmax Attention with a Forget Gate — Lin, Nikishin, He & Courville (ICLR 2025), the data-dependent forget gate on softmax attention (
forgetting_attention.py); ALiBi is its fixed, data-independent special case. - Train Short, Test Long: Attention with Linear Biases (ALiBi) — Press, Smith & Lewis (2021), the fixed distance penalty FoX generalizes (built in m04).
- Generating Long Sequences with Sparse Transformers — Child et al. (2019), the fixed strided/local sparse patterns NSA turns into a learned choice.
- Self-Attention Does Not Need O(n^2) Memory — Rabe & Staats (2021), the online-softmax memory argument.
- Efficient Memory Management for Large Language Model Serving with PagedAttention — Kwon et al. (2023), the vLLM paper: block tables, near-zero KV waste, and copy-on-write sharing.
- Orca: A Distributed Serving System for Transformer-Based Generative Models — Yu et al. (OSDI 2022), iteration-level scheduling — the continuous batching paging pairs with.
- SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills — Agrawal et al. (2023), chunked prefills + piggybacked decodes.
- Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve — Agrawal et al. (OSDI 2024), stall-free batching with a token budget — now the default scheduler in vLLM, TensorRT-LLM, and SGLang.
Practical Resources:
- Longformer: The Long-Document Transformer — Beltagy et al. (2020), sliding-window and dilated attention.
- Mistral 7B — Jiang et al. (2023), GQA + sliding-window attention in a shipped model.