Attention softmax fusion

What fused has to mean

Fusing the softmax into the QK stage keeps the score tile in registers from the contraction through the row max, the exponential and the division. It does not remove the probabilities from memory unless the PV matmul is fused too, because a separate op has to read them from somewhere. Fusing one boundary removes half the traffic and leaves an S by S tensor behind; the claim only holds when all three ops share a tile.

The tile loop

For each row block: load a Q tile of S_t by D, stream K tiles, produce a score tile of S_t by S_t in registers, reduce along rows for the max, subtract and exponentiate, reduce for the sum, divide, then contract with the V tile and accumulate into an output tile of S_t by D.

If the columns are tiled as well, a later tile can raise the row maximum, and the accumulator and the sum both need the online correction: multiply by exp(m_old - m_new) before adding the new tile’s contribution. That recurrence is the same one the GPU track’s softmax kernel uses, and it is the part of attention that is genuinely about order.

Keep the intermediate out of memory

The score tile has to remain a vector value. If bufferization runs before the tile exists, the intermediate becomes a memref of S by S, and the fusion becomes a story about IR that never reached the loop. Tile first, vectorise second, bufferize last, then grep the final IR for an allocation the size of the score matrix.

Things that break

Fusion attempted after bufferization, which leaves the allocation in place and a diagram claiming otherwise. A stack allocation that dies at 4096 columns. A fully masked tile whose maximum is negative infinity, poisoning the running max with NaN. A fused kernel compared against a materialised baseline at one short sequence length, where the traffic difference is too small to show.