LESSON 12 · 2026.04.22 · L4 · FINAL

vLLM's Paged Attention V2 in Triton — closing the MQA SM-occupancy hole with split-k

"When a parameter sweep doesn't move the metric, suspect the architecture," rewound one more time.

The most uncomfortable result from the previous chapter (Paged Attention) was a single line:

At MQA (B=16, H=32, H_kv=1, ctx=4k) our paged kernel is +85 % slower than SDPA — a structural flaw.

Lesson 12 opened for that one line.

(Note: "+85 %" was logged at the end of Phase 4. At the start of this session, re-running the same shape with block_size pinned at 16 gave +646 %. Block size changes the gap by that much. Either way the conclusion is the same: MQA is structurally slow.)

1. The flaw in one line

In MQA, the single-pass grid is (B, H_kv) = (16, 1) = 16 programs. L4 has 58 SMs. 16 / 58 ≈ 28 % occupancy. The other 42 SMs sit idle.

How do you wake those idle SMs? That's the entire session. Sweeping block size 8–128 barely moves the number. This isn't a parameter-tuning problem — it's architecture. vLLM's v1 → v2 transition is exactly their answer to this.

2. vLLM v2 — slice the ctx axis into two kernels

The crux of paged_attention_v2.cu:

1. Forward kernel: grid = (num_heads, num_seqs, max_num_partitions). A third axis for the ctx-partition index. PARTITION_SIZE = 512.
2. Reduce kernel: grid = (num_heads, num_seqs). Fold the (max, lse, partial_out) per partition using online softmax.

Ported to Triton, two new kernels:

• paged_attention_split_kernel — grid (B, H_kv, SEGMENTS). Each program processes only its segment's PARTITION_SIZE=512 tokens. Writes unnormalized (m_i, l_i, acc) to scratch.
• paged_attention_reduce_kernel — grid (B, H_q). Reads the SEGMENTS axis and recombines with alpha = exp(m_s - m_global). Writes the final normalized output.

Lesson 11's single-pass kernel stays put, alongside the new ones. Reasons:

• For dense shapes (B·H_kv ≥ half the SMs) single-pass is faster. No reason to pay the reduce kernel's launch overhead.
• Auto-dispatch picks per shape. Principle.

3. "Segment-level recombine" is a formula you already know

The split-k recombine isn't actually new. Lesson 09 used the same formula — there as "block-level recombine," here as "segment-level recombine."

m_global = max(m_s over s)
alpha_s  = exp(m_s - m_global)
l_global = sum(alpha_s * l_s)
acc_global = sum(alpha_s * acc_s)      # rescale each segment's acc, then add
out = acc_global / l_global

Literally the Lesson-09 online softmax lifted one level outward. Knowing that formula means the reduce kernel is 10 lines.

4. "Invalid segment" is mathematically the same as "empty segment"

The split kernel's grid is (B, H_kv, SEGMENTS), and SEGMENTS is sized to the longest sequence. What about sequences short enough that some segment is empty?

An empty segment:

• Inner loop runs 0 times → m_i, l_i, acc stay at their init (-inf, 0, 0).
• Those values are written to scratch as-is.
• In reduce: alpha = exp(-inf - m_global) = 0, so alpha * l = 0, alpha * acc = 0 — auto-ignored.

No separate "empty mask" logic needed. When the sentinel value is mathematically equivalent to "ignore," the code collapses. That was Lesson 12's nicest moment.

5. The trap — tl.arange requires power-of-2 ranges

The reduce kernel loads the SEGMENTS axis with tl.arange(0, SEGMENTS), and Triton complains:

CompilationError: arange's range must be a power of 2

Test case: ctx=513, partition_size=32 → SEGMENTS = ceil(513/32) = 17. Not pow2.

Naive fix: pad scratch to next_pow2. Wastes memory + the forward grid grows too.

Proper fix: only the kernel constexpr uses pow2; scratch stays at the actual size.

# kernel side
offs_s = tl.arange(0, SEGMENTS_P2)     # pow2, 17 → 32
mask_s = offs_s < SEGMENTS              # actual
m_s = tl.load(ptr + offs_s * stride, mask=mask_s, other=-float("inf"))

Pad lanes load as -inf → the recombine naturally contributes 0. Same sentinel trick, used twice (empty segment + pow2 padding).

After this patch: 32/32 PASS. Single-pass still 32/32 (no regression).

6. The first heuristic was wrong

First shot at auto-dispatch:

use_split_k = (B*H_kv < 0.75 * SM_COUNT) and (segments >= 2)

Bench reveals it's too permissive:

LLaMA-70B (B=4, H_kv=8) tripped the heuristic with B*H_kv = 32, but because GROUP=8, the single-pass kernel's inner body already amortizes KV loads across 8 query heads. Split-k's extra parallelism wins nothing — only adds the reduce kernel's launch overhead.

Second attempt:

use_split_k = (B*H_kv < 0.5 * SM_COUNT) and (segments >= 4)

• B*H_kv < 29: LLaMA-70B B=4 (32) drops out, only MQA B=16 (16) survives.
• segments >= 4: ctx=1k fails (2 segments), ctx=4k passes (8 segments).

Result: 9 of 10 shapes, auto picks the right path. The one miss (llama7b-B1-ctx4k) leaves SK 17 % ahead that auto doesn't take — "don't regress" beats "catch every win" is the principle.

7. Final numbers

MQA B=16 H_kv=1 ctx=4k bs=16 fp16, L4:

The paged kernel itself is 1.68× faster. The structural fix works. The gap to SDPA is still large — because this part can't be closed with split-k.

8. Why we can't catch SDPA — L2 locality

SDPA 0.044 ms. bytes_moved = 2·B·H·d + 2·B·H_kv·ctx·d = 2·16·32·128 + 2·16·1·4096·128 = 16.4 MB. Effective BW = 16.4e6 / 44e-6 = 761 GB/s.

L4's DRAM peak = 300 GB/s. 761 is physically impossible from DRAM alone.

Answer: L2 hit. K/V total is 16 MB; L4's L2 is 48 MB — fully cache-resident. SDPA reads the repeated accesses out of L2, so it looks L2-bound rather than DRAM-bound.

Our paged kernel:

• block_table lookups consume extra L2.
• Block-level gather disrupts L2 spatial locality.
• Measured BW: 165 GB/s (close to DRAM).

SDPA eats from L2, paged eats from DRAM. Split-k fixes the SM-occupancy problem, not the L2 residency problem. The remaining MQA gap is L2-aware prefetch / pinned block layout territory — next session.

9. Generalizing — "when a sweep doesn't move the metric, suspect the architecture"

One level deeper than Lesson 10's "if something's slow, there's a reason — nail the reason with a metric."

Lesson 11 Phase 4's observation: block_size 8, 16, 32, 64, 128 — any of them, MQA bottoms at 0.203 ms at bs=32. Parameter space doesn't close the hole. That was the decisive signal: "this isn't a parameter problem."

vLLM probably hit the same wall at v1 → v2. That's why their answer wasn't "tune partition_size," it was "add a new grid axis." Ours has to match.

10. Next-session candidates

1. L2-aware paged: how do we mimic MQA's L2 residency? Pinned block layout? Per-SM L1 staging?
2. Unified kernel: like vLLM's kernel_unified_attention_2d, prefill + decode in one kernel. Collapses serving-code branches.
3. OSS contribution: both above could actually be vLLM-worthy PRs.