Fix SM89 FP8 blockwise scale indexing under threadblock swizzling

[Noperi0r]

Problem

MmaMultistageBlockwise derives the ScaleA and ScaleB row indices from the raw CUDA block coordinates, blockIdx.x and blockIdx.y.

With GemmIdentityThreadblockSwizzle<N> and N > 1, the physical launch coordinates do not necessarily match the logical GEMM tile coordinates returned by the swizzle and used by the operand iterators and epilogue. The blockwise mainloop can therefore select scale rows that do not correspond to the M and N tiles being processed. This produces incorrect output and can read beyond exact-size scale tensors.

Fix

Pass the post-swizzle logical M and N tile coordinates from GemmUniversalBlockwise to MmaMultistageBlockwise, and use those coordinates when calculating the blockwise scale row indices instead of reading blockIdx directly.

Add an SM89 FP8 blockwise regression test covering:

• GemmIdentityThreadblockSwizzle<1> as a positive control.
• GemmIdentityThreadblockSwizzle<2> as the swizzled case that exposes the incorrect indexing.

The regression uses a 256x256x128 problem with a 64x128x128 threadblock shape, creating multiple logical tiles along both M and N. A and B are filled with ones, while each 128-row M scale block and 128-column N scale block uses a distinct power-of-two value, so the expected result can be computed directly for every output element.

The scale tensors are allocated at their exact required size, and the full output matrix is compared exactly. This makes an incorrect scale-row selection observable and prevents it from being hidden by unused padding.

Validation

Tested on an NVIDIA GeForce RTX 4060 Laptop GPU (Ada, SM89) with CUDA 12.8 under WSL2 Ubuntu 24.04.

• New SM89 blockwise unit-test target: 2/2 tests passed.
• Before the fix, the Identity<2> case produced 49,152 output mismatches.
• After the fix, the same case matched the expected output exactly.
• Compute Sanitizer memcheck on the fixed Identity<2> case reported 0 errors.
• Additional local checks passed with Identity<4>, multiple non-split K scale blocks, and the public ColumnMajor adapter.
• Example 94 SM89 FP8 blockwise verification passed.
• Existing adjacent SM89 unit tests: 34/34 passed.
• git diff --check passed.

[Noperi0r]

Hi @hwu36, could you please help route this PR to the appropriate reviewer when you have a chance? It fixes a correctness issue in the SM89 FP8 CUTLASS 2.x GemmBlockwise path and includes a focused regression test and local Compute Sanitizer validation. Thanks.