QK-transpose matmul

One contraction, two maps

scores[i, j] = sum over d of Q[i, d] * K[j, d]. As a linalg.generic that is three maps: (i, d), (j, d) and (i, j). Both operands are indexed by the reduction dimension in the same position, which is exactly what a pair of row-major [S, D] buffers wants.

linalg.matmul is written with maps (m, k), (k, n) and (m, n). Its second operand is indexed by the reduction dimension first, so feeding it K means handing it K transposed. Reusing the matmul signature buys a transpose; writing the map costs nothing.

What the vectoriser does with the map

Vectorising the contraction turns each operand access into a transfer whose map is the original map restricted to the tile. When that map permutes the innermost dimension, the transfer is a strided read: vlse32 rather than vle32, and a gather when the stride is not constant. On many implementations a strided load moves one element per memory transaction, so the QK stage slows down without becoming wrong, which is the kind of regression a correctness test never catches.

Three ways out, in order

Write the maps so the contraction reads K as it lies. If the operand is a static weight, transpose it once outside the timed region and say so in the rig line. If neither is possible, change the tile order so the K tile loads by row and the microkernel walks down columns.

Things that break

A transpose inserted inside the call path, so the benchmark times it and the table charges the cost to attention. A permutation that survives vectorisation and then scalarises at the transfer, which looks like a slow kernel and is really a slow load. A batch dimension mistaken for the reduction, which produces the right shape and the wrong sum.