Decoder Block 6: Online Attention, the Capstone

Lesson

This is the capstone, and it earns the promise stage 5 made: the materialized score row goes away. The driver is the whole decoder block now, built from the path's own graded solutions: the packed QKV projection from stage 3 (one launch, 768 slab rows), the in-place RoPE from stage 4, and the trusted tail through output projection, residual, second RMSNorm, SwiGLU and down projection. This is the same pre-norm decoder block Llama-family models run per layer at inference time. One seam in the middle is yours: runCausalAttention in attention.cu, called once per case, timed with CUDA events against stage 5's exact kernel running in the same run on the same packed tensor.

The starter IS that stage 5 kernel behind the new host seam, and it measures about 1.0x by construction. Its cost is architectural, not a tuning miss: one block per query row means every row re-reads its whole K and V prefix from the memory system. At 2048 tokens that is roughly 2048 traversals of the same KV head per head, and no amount of polishing the inner loop changes that multiplier.

The online version changes the multiplier. Keep a TILE of query rows resident in one block, stage each K/V tile in shared memory once, and let all 32 rows consume it before it is replaced. K and V then cross the memory system once per query tile instead of once per query row, a 32x cut in attention's dominant traffic. The price is that a row no longer sees its whole score line at once, which is exactly what the Reduction path's online-softmax stage prepared you for. Per row you carry three things: the running maximum m, the running denominator l, and the 64-value unnormalized output accumulator. When a new tile arrives with tile maximum m_tile:

m_new = max(m_old, m_tile)
alpha = exp(m_old - m_new)
l_new = alpha * l_old + sum_j exp(score_j - m_new)
acc = alpha * acc + sum_j exp(score_j - m_new) * V[j]

One division by l at the very end normalizes the whole row. Forget the alpha on acc or on l and the block is wrong the moment the running maximum moves between tiles, which the odd length 33 already catches: its rows past 32 see two tiles. This rescale discipline is the algorithmic core of FlashAttention; what you are writing is the teaching-scale version of the idea that made long-context inference affordable.

The structure of the run enforces the rest. The driver allocates no attention workspace of any kind, so there is nowhere to put a T x T score tensor; the baseline needs none and neither does your kernel. Keep as many clear kernels as you like behind the host seam; the events bracket everything you launch. The gate asks for 1.8x. The starter measures about 1.0x, the reference online solution about 2.7x on the calibration GPU, and the judged evidence prints all three numbers: ATTN BASELINE MS, ATTN USER MS, and SPEEDUP. Correctness is judged at every value: the small lengths against the double reference that has judged this path since stage 1, and the profile length against the trusted baseline path, attention output first and final block output second, so a failure names the earliest wrong tensor.

This is teaching attention, not a production FlashAttention implementation. Compare with the materialized baseline in the same run, with identical shape, mask and precision.

What you will learn

• Carry online-softmax state (max, denominator, accumulator) across streamed K/V tiles.
• Rescale prior contributions by exp(m_old - m_new) whenever the running maximum moves.
• Cut attention's dominant traffic by reusing each staged K/V tile across a resident query tile.
• Separate model mathematics from implementation choice across the whole decoder block.