Probability-by-V matmul

The second contraction

O[i, d] = sum over j of P[i, j] * V[j, d]. P is the softmax output tile from the previous node, still in registers. V is read tile by tile. The reduction runs over the sequence dimension, which is the same index the softmax normalised, so the tile you just produced is exactly the left operand.

What matching the reference means

A dense implementation on the same inputs, after the same 1/sqrt(D) scale: a scalar triple loop, or the framework expression in f64. State the tolerance as a function of the sequence length rather than as one number, because f32 accumulation error grows with the count of terms. A relative 1e-5 is plausible at 512 columns and tight at 8192; a single bound for every shape is a test chosen at its easiest point.

Bits depend on accumulation order

vector.contract becomes vfmacc when the target allows contraction, and a fused multiply-add rounds once where a multiply and an add round twice. Two builds of the same sum therefore disagree in the last bits with neither being wrong. Record whether the kernel contracts, and keep the reference uncontracted: if both contract, a difference in tile order disappears and a real bug hides inside it.

The mask is part of the contract

Causal attention zeroes the upper triangle before the softmax, and that zeroing changes the denominator, not just the numerator. A PV matmul tested only on unmasked input can be wrong and still match. Test one masked shape and one row that is entirely masked.

Things that break

Comparing against a reference that forgot the scale, then loosening the tolerance until the test passes. A tolerance wide enough to accept a missing mask. Accumulating in f16 without saying so, which is a different accuracy claim rather than a faster version of this one.