[Common] Honor MXFP8 scale strides in specialized casts

tests/pytorch/test_quantized_tensor.py

@@ -1013,6 +1013,56 @@ def test_quantize_dequantize_autograd(
         ), "Gradient did not flow back to the input through quantize -> dequantize"
 
 
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10,
+    reason="Register-resident MXFP8 quantization requires Blackwell or newer",
+)
+@pytest.mark.parametrize("fp8_dtype", _fp8_dtypes)
+@pytest.mark.parametrize("cols", [1024, 2048, 4096])
+@pytest.mark.parametrize("extra", [(0, 0), (4, 0), (0, 16), (4, 16)])
+def test_mxfp8_quantize_with_pitched_scales(
+    fp8_dtype: te.DType, cols: int, extra: Tuple[int, int]
+) -> None:
+    """Quantization honors the scale strides supplied by the output tensor."""
+    torch.manual_seed(1234)
+    rows = 64
+    x = torch.randn((rows, cols), dtype=torch.bfloat16, device="cuda")
+    quantizer = MXFP8Quantizer(fp8_dtype=fp8_dtype)
+    quantizer.optimize_for_gemm = False
+    ref = quantizer(x)
+    out = MXFP8Tensor(
+        shape=x.shape,
+        dtype=x.dtype,
+        device=x.device,
+        rowwise_data=torch.empty_like(ref._rowwise_data),
+        columnwise_data=torch.empty_like(ref._columnwise_data),
+        rowwise_scale_inv=torch.full(
+            (rows, cols // 32 + extra[0]), 0xAA, dtype=torch.uint8, device=x.device
+        ),
+        columnwise_scale_inv=torch.full(
+            (rows // 32, cols + extra[1]), 0xAA, dtype=torch.uint8, device=x.device
+        ),
+        fp8_dtype=fp8_dtype,
+        quantizer=quantizer,
+        with_gemm_swizzled_scales=False,
+    )
+    out.quantize_(x)
+    torch.testing.assert_close(out._rowwise_data, ref._rowwise_data, rtol=0, atol=0)
+    torch.testing.assert_close(out._columnwise_data, ref._columnwise_data, rtol=0, atol=0)
+    torch.testing.assert_close(
+        out._rowwise_scale_inv[:, : cols // 32],
+        ref._rowwise_scale_inv[:rows, : cols // 32],
+        rtol=0,
+        atol=0,
+    )
+    torch.testing.assert_close(
+        out._columnwise_scale_inv[:, :cols],
+        ref._columnwise_scale_inv[: rows // 32, :cols],
+        rtol=0,
+        atol=0,
+    )
+
+
 @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
 class TestMXFP8Tensor:
 

transformer_engine/common/cast/mxfp8/specialized/cast_bidim.cuh

@@ -571,7 +571,10 @@ void launch_tiled(const void *input, void *output_rowwise, void *scales_rowwise,
   // Keep the small cases; everything else takes the runtime path.  All column
   // counts work either way: `default` is the runtime instantiation, not a
   // rejection.
-  switch (cols) {
+  const bool has_padded_scale_strides =
+      scale_stride_rowwise == DIVUP_TO_MULTIPLE(cols / kBlockElems, 4) &&
+      scale_stride_colwise == DIVUP_TO_MULTIPLE(cols, 128);
+  switch (has_padded_scale_strides ? cols : 0) {
     case 2048:
       NVTE_LAUNCH_BIDIM(2048);
       break;