LESSON 05 · 2026.04.18 · T4

Softmax & Fusion — the mathematical half of Flash Attention

Fuse three kernels into one and you get exactly 2× speedup. And online softmax, born next to it, becomes the heart of Flash Attention.

Three versions

• v1 naive — 3 kernels sequential (max → sum → divide). HBM: 4 trips/elem
• v2 fused — 1 kernel, row cached in shared memory. HBM: 2 trips/elem
• v3 online — 1 kernel + normalize pass. Any N. HBM: 3 trips/elem

Results

Lesson 1 · The promise of fusion is observable

Theoretical HBM trips: v1 = 4, v2 = 2 → v2 should be 2× faster. Observed: 2.02× (N=1024), 1.86× (N=2048), 1.92× (N=4096). Very close to theory. This is why more than half of LLM-inference optimization is fusion.

Lesson 2 · Occupancy cliff @ N=8192

v2 smem usage = N × 4 bytes. At N=8192, that's 32 KB. Against T4's 64 KB smem per SM, only 2 blocks reside → threads/SM drops from 1024 to 512 → 50% occupancy.

It's not a bandwidth shortage — it's a shortage of warps to hide latency. This is the generic trap of shared-memory-heavy kernels.

Lesson 3 · L2 absorbs v3's "extra read"

In theory v3 should be 1.5× slower than v2 because of 3 trips. Yet at N=1024, v3 is actually slightly faster. Reason: a row is 4 KB → fits comfortably in L2 (4 MB) → pass 1's input hits L2 in pass 2. Effective 356 GB/s (111% of theoretical) is the evidence. As N grows, L2 gets evicted and the benefit fades, v2 retakes the lead.

Lesson 4 · v3's real value — the online update formula

v3 isn't faster than v2 at our sizes. But:

• Supports arbitrary N (v2 fails for N > 12288)
• Its online update formula is the output-update formula of Flash Attention

new_max = max(m1, m2)
new_sum = s1 * exp(m1 - new_max) + s2 * exp(m2 - new_max)

FA layers tiled matmul fusion on top of this so the intermediate matrix (P = softmax(Q@K^T)) never lands in HBM. Implementing v3 = understanding the mathematical half of FA. The other half is attention-specific tiling — next lesson.

Regime map

small N (smem slack)     v2       fusion clean win (2×)
mid N (L2 hits)          v2 ≈ v3  L2 absorbs v3's 3rd read
large N (smem saturated) v3 / FA  v2's occupancy collapses
very large N (attention) FA       intermediate matrix cannot materialize