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_heads dial.
  • 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, plus how continuous batching rides on it.
  • Read the common attention mask shapes: full, causal, sliding-window, dilated.

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 KVCache we 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
  1. 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.
  2. 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  <- smaller

The 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=4

Just 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:

  1. Decoding through the cache is exact. Running a layer token-by-token with a KVCache returns the identical result to one full forward pass.
  2. 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.79e-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.79e-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| = 1.79e-07
  (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| = 1.79e-07  (≈ 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-07

Now 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
  1. 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.
  2. 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.
  3. 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_size tokens (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 n tokens 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
  1. 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.
  2. 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.
  3. Shrink block_size in 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 -> physical

The 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:  True

Now 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: True

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

Sharing Blocks: Copy-on-Write

Because a sequence is just a list of physical block numbers, two sequences can point at the same block. When many requests share a prefix — a common system prompt, or several samples from one prompt — the shared blocks are cached once. A per-block reference count lets a sharer keep reading until it writes into a shared block, at which point that one block is copied (copy-on-write), exactly like fork() in an OS.

cache = PagedKVCache(num_blocks=16, block_size=4, num_kv_heads=1, head_dim=2)
prompt = cache.add_sequence()
for _ in range(8):                       # an 8-token shared prompt (2 blocks)
    cache.append(prompt, torch.ones(1, 2), torch.ones(1, 2))

used_before = cache.allocator.num_used
sample = cache.share_prefix(prompt, num_prefix_blocks=2)   # a second sample reuses it

print(f"Blocks used before sharing: {used_before}")
print(f"Blocks used after sharing:  {cache.allocator.num_used}  (no new blocks)")
print(f"Same physical prefix block: {cache.block_table(sample)[0] == cache.block_table(prompt)[0]}")
Blocks used before sharing: 2
Blocks used after sharing:  2  (no new blocks)
Same physical prefix block: True

Continuous Batching

Paging fixes space; continuous batching (Orca, Yu et al., 2022) fixes time. A static batch runs every request to completion together, so one long generation holds the whole batch hostage and finished requests leave their slots idle. Iteration-level scheduling instead admits and evicts requests between decode steps: the moment a sequence emits its end token, its blocks return to the pool and a waiting request takes the freed capacity — on the very next step. Paging is what makes that cheap: a request joins or leaves by editing a block table, never by moving a contiguous cache. Together they are why modern servers (vLLM, TensorRT-LLM, TGI) keep the GPU saturated.

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

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| = 3.58e-07
Match: True

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

Common Pitfalls

When implementing efficient attention, watch out for:

  1. Expanding K/V before caching, not after. Store the small K/V (num_kv_heads) in the cache and repeat_interleave only on the way into the multiply. Expanding first throws away the entire memory saving.
  2. Wrong grouping order. repeat_interleave(n_rep, dim=1) maps query heads [g·n_rep : (g+1)·n_rep] to K/V head g. Using repeat/tile instead interleaves the groups differently and silently mismatches heads.
  3. num_kv_heads must divide num_heads. Otherwise the groups are uneven; the layer asserts this at construction.
  4. 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.
  5. 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 NaN from softmax of all -inf.
  6. 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.
  7. Scaling MLA scores by √d_head instead of √(d_head + d_h^R). The query and key are [content ; rope] concatenations, so the correct denominator uses their combined width.
  8. A block size that is too large. Paging bounds waste at block_size − 1 tokens 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.

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%}")

Summary

Key takeaways:

  1. 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.
  2. MHA, GQA, and MQA are one mechanism with one dial — num_kv_heads controls how many query heads share each key/value head; MHA and MQA are the endpoints.
  3. The saving is real and free — the cache stores only num_kv_heads heads; repeat_interleave re-expands them just for the multiply, leaving the attention math and (for GQA) most of the quality intact.
  4. 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.
  5. Cached decoding is exact — GQA and MLA through their caches reproduce a full forward pass to floating-point rounding, just like MHA in m08.
  6. 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.
  7. 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.
  8. 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, and lets continuous batching swap requests in and out between decode steps. Neither touches the attention result, yet together they are the biggest throughput lever on a real server.

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:

Practical Resources: