From Naive Implementation to Auto-Tuning

In the previous chapter, we derived the math behind Flash Attention (Tiling + Online Softmax). Now it is time to turn the math into code.

Core Loop Structure

Three questions sit between the formula and a working kernel, and it pays to settle them first.

Which loop is parallel, and which is serial. The derivation was written as "Outer Q, Inner K": the outer loop walks blocks of $Q$, the inner loop walks blocks of $K, V$. In Triton you will never see that outer loop, because the SPMD model parallelizes it implicitly: each Program claims one $Q$ block, and tl.program_id(0) is its index. The only loop written as a for is the inner one. Miss this and you will spend the whole kernel hunting for an outer loop that does not exist.

Which data stays put across iterations. A $Q$ block is loaded once and reused for the entire inner loop; $K$ and $V$ are re-loaded every round. That distinction decides whether a tl.load belongs inside or outside the loop, and it is precisely why Flash Attention drops HBM traffic from $O(N^2)$ to $O(N^2 d / M)$.

How much state the accumulator carries. Naive softmax can compute a full row and normalize at the end. Tiled softmax cannot: each new $K$ block may raise the running maximum, which means everything accumulated so far has to be rescaled. So alongside the output accumulator acc, we carry a running max and a running sum — three pieces of state that survive across iterations.

With those three settled, the kernel below is just a line-by-line translation into Triton: