Module gemm 

FP32 matrix products in the layout of a PyTorch Linear: y = x wᵀ + b, with x as [m, k] and w as [n, k], both row major.

The weights are packed once into panels of 16 rows, [n / 16][k][16] with the last panel padded with zeros, and a micro kernel keeps a 6 by 16 block of outputs in registers while it walks k, one broadcast of x and two vector loads of the panel per step. Every output is summed the same way wherever it lands: in order over k with fused multiply adds in f32, moved to an f64 sum every 64 steps, then rounded to f32 and given its bias. Tiling and threading only decide which outputs run side by side, so the result is the same bit for bit for any thread count and any split, and the same on x86 with FMA as on ARM.

On macOS the GEMM goes to Accelerate instead, whose sgemm runs on the AMX units at two to four times what the NEON kernel reaches. Its sums are in an order of its own that changes with the number of rows in the call, so it is always called on blocks of exactly 64 rows, the last one padded with zeros. With the row count fixed, a row’s result does not depend on the rows around it or on its place in the block, so the bits still do not depend on the batch or the split. They are not the same bits as on other machines, which the parity tests allow for.

dot, which attention uses for its scores, sums in eight lanes instead and is not meant to match the GEMM bit for bit.

Structs

Gemm

One GEMM with its epilogue, y = ep(x wᵀ + b), split into tiles that any thread may run. Every output element is computed the same way whatever the split, so the split is free to follow the thread count.

Constants

NR

Rows of w per panel, two vectors.

Functions

dot

dot(a, b) in the order described in the module docs.

linear

y = x wᵀ + b with x as [m, k], w as [n, k], b as [n] and y as [m, n]. This packs w on every call, so a caller with a fixed weight should pack it once and run a Gemm.

pack

Lays out w, [n, k] row major, the way Gemm reads it: panels of 16 rows, or as it is on macOS, where Accelerate reads it.

scratch_len

Floats of scratch each task of a GEMM with these sizes needs.