MXFP4 Workgroup Reordering Based on SP3 Kernel

[suryajasper]

MXFP4 Workgroup Reordering

Remaps the workgroup grid by creating groups of 32 tiles along the N dimension. This way, the first 32 workgroups in launch order are assigned along the N dimension on the same M row (0), then the next 32 to the next M row (1), and so on until wrapping around back to the first row.

Through this grouping, consecutive workgroups are able to reuse the same A row slices rather than spreading over a larger, sparse footprint. This improves L2 cache reuse since workgroups running close together in launch order work on the same row of output tiles.

The index remapping is calculated through two cases:

• Main case: Creates full groups of 32 consecutive tiles along N dimension
• Tail case: If the number of tiles along the N dimension is not a multiple of our group size (32), after forming as many full groups as possible, we are left with a small rectangle in the launch grid, which we fill with the tailing workgroups in row-major order.