LESSON 11 · 2026.04.22 · L4

Rewriting vLLM's Paged Attention in Triton — and repeating vLLM's mistake

A story about how independently converging on the same refactor is the evidence the design is right.

Prologue — the one line you slip into contiguous attention

In Lesson 09 I wrote a (B, H, N, d) 4-D MHA + causal flash attention in Triton. That kernel runs at 78–90 % of torch.nn.functional.scaled_dot_product_attention (which, under the hood, is Tri Dao's FA-2 CUDA). By the bench alone, "a good kernel."

The problem: LLM serving (vLLM / SGLang / TensorRT-LLM) doesn't use a contiguous KV cache. Sequences vary wildly in length and churn in and out, so pre-allocating burns 70 %+ of GPU memory to fragmentation. The fix vLLM pushed through at SOSP '23 is simple:

Split the KV cache into a pool of fixed-size blocks, and use a per-sequence block_table to reference them indirectly.

(before)
K: (B, H, N, d)                             ← contiguous per sequence

(paged)
K_cache:     (num_blocks, block_size, H_kv, d)   ← block pool
block_table: (B, max_blocks_per_seq)             ← seq → physical block id
context_lens:(B,)                                 ← valid length per seq

What changes in the attention kernel is exactly one line:

# before: load K[b, h, start_n:end_n, :] contiguously
# paged: look up block_table[b, logical_blk] → phys_blk →
#        load K_cache[phys_blk, :, kv_head, :]

This essay measures "does slipping that one line in actually yield vLLM paged attention?" It does — but I was wrong twice on the way. That's the essay.

Phase 1–2 · correctness passed too easily (this was the trap)

Order of work:

1. Build a PyTorch reference (Python loop through block_table, gather, then standard attention) as the oracle.
2. Triton kernel: initially grid = (B, H_q) — the natural reduction of Lesson 09's (cdiv(N, BM), H, B) with "decode, so N=1."
3. GQA support: add GQA_GROUP_SIZE = H_q // H_kv as a constexpr and kv_head = pid_h // GQA_GROUP_SIZE. Four lines changed total.

Correctness bench (16 shapes × 2 dtypes = 32/32 PASS):

Sub-1e-3 on fp16, 1e-7 on fp32. Everything any model shop would accept.

At this point the judgement was: "It works. We're done." That was the trap.

Phase 3 · the speed bench exposed a structural bug

Had I stopped at correctness, I would have shipped a kernel 2–3× slower than SDPA on LLaMA-3-8B.

MHA parity. From GQA onward, the gap is roughly linear in GROUP_SIZE.

• group=4 → +16–37 %
• group=8 → +117–217 %
• group=32 → +1316 % (13× slower)

That linearity is the diagnosis. Not a random regression — structural.

The cause — (B, H_q) grid reloads KV GROUP times

A grid of (B, H_q) means one program per (batch, query head). The GROUP_SIZE query heads in a GQA group share the same KV head, but each program independently walks block_table and reloads the same K/V block from DRAM again. Redundant DRAM traffic scales with GROUP_SIZE.

Why isn't SDPA slow? Contiguous KV lets the L2 prefetcher absorb the redundant loads. SDPA on MQA reaches 542 GB/s on L4 — 1.8× the DRAM peak (300 GB/s). That throughput is impossible from DRAM alone. L2 is absorbing more than half.

Our paged kernel breaks the L2 prefetch pattern because of block_table indirection, so the redundant loads all go to DRAM.

Phase 3.5 · I only changed the grid (and another bug popped up)

The fix is clear: change the grid to (B, H_kv) and handle the GROUP_SIZE query heads of the GQA group at once inside the program. K/V blocks get loaded once per program.

grid = (B, H_kv)                              # program count ÷ GROUP
q = tl.load(q_ptrs)                            # (GROUP, HEAD_DIM) — 2D tile
# inside the block loop:
scores = tl.dot(q_scaled, tl.trans(k))         # (GROUP, BLOCK)
acc += tl.dot(p.to(v.dtype), v)                # (GROUP, HEAD)

The change itself is ~20 lines. Bench:

LLaMA-3-8B at -14 % — our Triton kernel beats cuDNN / FA-2 at a shape common in production.

But fp32 correctness broke

• fp16: all shapes still pass ✅
• fp32 MQA: max diff 4.1e-04 ← it was 3.6e-07 before. Three orders of magnitude worse.

Confused. Why would fp32 break from just changing the grid?

The real cause — tl.dot(fp32, fp32) defaults to TF32 on sm_80+

On Ampere and later, Triton silently downgrades fp32 × fp32 in tl.dot to TF32 (10-bit mantissa). Without specifying input_precision, the default is TF32. On MQA's (GROUP=16, BLOCK=16, HEAD=64) score tile, 80–100 summation steps accumulate 10-bit truncation error and bias the softmax-max boundary by 4e-4.

Fix:

if IS_FP32:
    # 3-pass TF32 stack (2 low-bit corrections) to reconstruct IEEE — 3× slower
    scores = tl.dot(q_scaled, tl.trans(k), input_precision="ieee")
else:
    # fp16/bf16 MMA — default is already IEEE fp16
    scores = tl.dot(q_scaled, tl.trans(k)).to(tl.float32)

fp32 max diff: 4.1e-04 → 3.6e-07 recovered. No fp16 speed cost.

Why didn't this show up in Phase 3?

Phase 3 computed scores via manual broadcast (tl.sum(q * k)) — a pure fp32 path. No TF32 detour. The bug first showed up the moment Phase 3.5 introduced tl.dot.

Phase 4 · reading vLLM's source, I realized I had reproduced vLLM's history in miniature

I read the vLLM source only after finishing Phase 3.5 — deliberately, to see whether an independent design would converge with vLLM's.

Files I read:

Finding 1 — my Phase 3.5 matches vLLM's current Triton kernel axis-for-axis

The only real difference: vLLM's axis-0 is (batch × query block) because they handle prefill in the same kernel (variable q_len); mine is pure batch because decode pins q_len=1. Same structure, different scope.

Finding 2 — vLLM itself went through the same refactor

paged_attention_v1.cu:86:

dim3 grid(num_heads, num_seqs, 1);

That is a per-query-head grid — identical in design to my Phase 3. vLLM's 2023 ship.

Why it was fine then:

• Most of the LLM market was MHA (H_kv == H_q). Group redundancy didn't structurally exist.
• The few GQA models had small KVs that L2 could absorb.
• Triton 2.x's MMA couldn't yet compete with CUDA, so CUDA was the answer.

As LLaMA-2-chat, LLaMA-3, and Mistral shipped with GQA, the per-query-head grid became the bottleneck. vLLM moved to Triton and restructured to (q_block, H_kv) — exactly the refactor I did from Phase 3 → 3.5.

Finding 3 — one thing I did that vLLM didn't

vLLM doesn't specify precision on tl.dot. Production runs fp16/bf16 only, so it doesn't matter. But if someone flows fp32 through that path on sm_80+, they get the same 4e-4 error I caught. My IS_FP32 branching + input_precision="ieee" is only relevant in a lesson context, but still caught it.

Finding 4 — what I didn't do: split-k over the ctx axis

The source of my residual +85 % MQA gap. SDPA hits 698 GB/s on this shape (2.3× DRAM) — L2 absorbs more than half. 1 KV head × 4k tokens × 128 dim × 2 B fp16 = 1 MB fits trivially into L4's 48 MB L2 and is shared by 32 query heads.

My paged kernel can't replicate that L2 reuse because of block_table indirection — structural. The grid fix alone can't close it. vLLM's v2 partitions the ctx axis and uses a reduce kernel to recombine softmax. Pushed to Lesson 12.

Final numbers (L4 sm_89, fp16, warmup=50, iters=200)

Correctness: 32/32 PASS. A 275-line Triton kernel that beats SDPA (= Tri Dao FA-2 CUDA) by 14 % on LLaMA-3-8B B=8 ctx=2k. Parity on LLaMA-70B. MQA at +85 % — a residual gap that split-k will close.

Three things that stick

(1) Correctness passing doesn't mean "correct"

If I had called it done at 32/32 PASS, I would have shipped a kernel 2–13× slower on GQA shapes. This trap doesn't get caught without the SDPA-gap column in the speed bench. Report allclose and gap together.

(2) Bugs hide behind bugs

The grid bug (Phase 3) and the TF32 bug (Phase 3.5) were independent and only surfaced sequentially. After a big refactor, rerun correctness — always. "Fixed one, so we're safe" is exactly wrong in this situation.

(3) Converging independently is the evidence of design

I read vLLM's source only after finishing Phase 3.5, and found an axis-for-axis match. That's not me being clever; it's "the right answer is one" when the tool (Triton) + HW (sm_80+) + workload (GQA) are fixed. Recording this convergence is actually the credible story — "the ecosystem already did this refactor and I reproduced it in miniature."