pub const AXPY_W: usize = 4;
The sub-blocks of outputs and the row groups a block of the staged per-head matmul takes.