Attention block decomposition

Three shapes, two reductions

Q, K and V are [B, H, S, D]. Scores are Q times K transposed, scaled: [S, S], with the contraction over D. Probabilities are a row-wise softmax over the last dimension. Output is P times V: [S, D], with the contraction over S.

The two contractions reduce different dimensions, which is why the block is not one matmul with a function in the middle.

Layout decides the boundary

Q and K are both [S, D] row-major, so the QK contraction reads contiguous data from both operands. That survives vectorisation when the indexing maps say so, and it is why the QK side needs no copy. The PV side does: P[i, j] steps contiguously in j, while V[j, d] steps a whole row. Either V is transposed once, or the microkernel iterates so the reduction stays innermost in both.

The fusion boundary sits at the score matrix. Materialised, one head at S=4096, D=64, f32 writes and reads 4 * S * S * 4 bytes, about 268 MB, against 2 * S * S * D flops, about 2.1 GFLOP. That is roughly 8 flop/byte, below the ridge point of any CPU here, so the kernel is memory-bound and the arithmetic is not what you are optimising.

Write the decomposition first

Build the three ops, check the block against a dense reference, then decide where to fuse. The decomposition is also the ablation plan for the benchmark node: every boundary you remove is a variant you can time.

Things that break

Assuming the layout is free and inserting a transpose on every call, which is a memory pass you then measure as attention. Applying the mask after the softmax, which changes the normalisation rather than the selection. Fusing at one boundary and calling the block fused when the score matrix still reaches memory as P.