Article. Learning FlashAttention the hard way. When a loop is secretly associative?

I'm writing a short series of tutorials on FlashAttention, the algorithm that largely powers modern LLMs. The core idea is to spot the associative structure hiding in the loop. Once you see it, you get two things: you can fuse the passes (in the case of attention, it is huge, because it avoids ever materializing the score matrix) and split and recombine the work across GPU threads or any parallel processor. This post is about recognizing that structure. The stable softmax operation on a vector x (part of the attention kernel) is computed as follows: m = max_j x_j softmax(x)_i = exp(x_i - m) / Σ_j exp(x_j - m) Naively, you would write something like: # pass 1 — running max m = -inf for j in 0..N: m = max(m, x[j]) # pass 2 — denominator, needs the final m d = 0 for j in 0..N: d += exp(x[j] - m) # pass 3 — normalize for j in 0..N: y[j] = exp(x[j] - m) / d The trick is not to wait for the final max before accumulating the denominator. Carry both in a small state (m, d) and rescale d whenever the max moves. Here is what the online version of the algorithm above looks like: # one pass — carry (m, d) together m, d = -inf, 0 for j in 0..N: m_new = max(m, x[j]) d = d * exp(m - m_new) + exp(x[j] - m_new) # rescale, then add m = m_new # normalize (unchanged) for j in 0..N: y[j] = exp(x[j] - m) / d It turns out this is not a one-off trick. There is a whole class of "secretly associative" loops that you can parallelize by introducing the right carrier state. The tutorial goes into detail, shows a few examples of secretly associative operations, provides the algebraic formulation for those, and gives some tools to help you recognize secretly associative loops. Overview: Safe softmax, Welford's variance, and FlashAttention belong to the same class of secretly-associative operations The twisted monoid via transport of structure, why the max-rescale coupling doesn't break associativity Third Homomorphism Theorem as a test for whether any loop is secretly associative Numerical analys