diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 4e64215cdf5..1f82c1abd46 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/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: diff --git a/transformer_engine/common/cast/mxfp8/specialized/cast_bidim.cuh b/transformer_engine/common/cast/mxfp8/specialized/cast_bidim.cuh index 014d3a632ee..0c870d91fb2 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/cast_bidim.cuh +++ b/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;