From 32de01986b2c02524bcd406934c0042f279c0a10 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 14 May 2026 03:59:02 -0700 Subject: [PATCH 1/7] Support MXFP8 2D quantization Signed-off-by: kunlunl --- qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/test_mxfp8_2d_quantize.py | 485 ++++++++++++++++++ .../common/cast/dispatch/quantize.cuh | 4 +- .../common/cast/mxfp8/quantize_mxfp8.cuh | 118 +++-- transformer_engine/common/common.h | 4 +- .../transformer_engine/transformer_engine.h | 9 + transformer_engine/common/recipe/__init__.py | 4 + .../common/transformer_engine.cpp | 6 + transformer_engine/pytorch/csrc/common.h | 2 + transformer_engine/pytorch/csrc/quantizer.cpp | 4 + transformer_engine/pytorch/quantization.py | 18 +- .../pytorch/tensor/mxfp8_tensor.py | 4 + 12 files changed, 617 insertions(+), 42 deletions(-) create mode 100644 tests/pytorch/test_mxfp8_2d_quantize.py diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index c35dc4c063..2a916ab585 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -59,6 +59,7 @@ fi python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_checkpoint.xml $TE_PATH/tests/pytorch/test_checkpoint.py || test_fail "test_checkpoint.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_router.xml $TE_PATH/tests/pytorch/test_fused_router.py || test_fail "test_fused_router.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_partial_cast.xml $TE_PATH/tests/pytorch/test_partial_cast.py || test_fail "test_partial_cast.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mxfp8_2d_quantize.xml $TE_PATH/tests/pytorch/test_mxfp8_2d_quantize.py || test_fail "test_mxfp8_2d_quantize.py" # Disable autotuning to make unittests faster. In addition, disable TF32 path to fully align with the pytorch reference implementation's precision NVTE_DISABLE_TRITON_AUTOTUNING=1 NVIDIA_TF32_OVERRIDE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mhc.xml $TE_PATH/tests/pytorch/test_mhc.py || test_fail "test_mhc.py" diff --git a/tests/pytorch/test_mxfp8_2d_quantize.py b/tests/pytorch/test_mxfp8_2d_quantize.py new file mode 100644 index 0000000000..d5baf213c9 --- /dev/null +++ b/tests/pytorch/test_mxfp8_2d_quantize.py @@ -0,0 +1,485 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Tests for MXFP8 2D quantization.""" + +import pytest +import torch + +import transformer_engine.pytorch as te +import transformer_engine_torch as tex +from transformer_engine.common.recipe import MXFP8BlockScaling +from transformer_engine.pytorch import MXFP8Quantizer +from transformer_engine.pytorch.quantization import ( + MXFP8BlockScalingRecipeState, + QuantizerRole, +) + + +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +MXFP8_BLOCK_SIZE = 32 +FP8_E4M3_MAX = 448.0 +MXFP8_TEST_SHAPES = [ + (64, 64), + (128, 128), + (256, 256), + (256, 1024), + (1024, 256), + (256, 288), + (320, 320), + (352, 256), + (2048, 2048), + (1024, 2048), + (2048, 1024), +] +MXFP8_TEST_DTYPES = [torch.float32, torch.bfloat16] + + +def _valid_rowwise_scale(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """Return the logical, unpadded rowwise scale region.""" + return scale[:rows, : (cols + MXFP8_BLOCK_SIZE - 1) // MXFP8_BLOCK_SIZE] + + +def _valid_columnwise_scale(scale: torch.Tensor, rows: int, cols: int) -> torch.Tensor: + """Return the logical, unpadded columnwise scale region.""" + return scale[: (rows + MXFP8_BLOCK_SIZE - 1) // MXFP8_BLOCK_SIZE, :cols] + + +def _float_to_e8m0(amax: torch.Tensor) -> torch.Tensor: + """Convert amax values to E8M0 scale bytes with the same ceil policy as TE.""" + val = (amax.to(torch.float32) / FP8_E4M3_MAX).contiguous() + val_u32 = val.view(torch.int32) + exponent = ((val_u32 >> 23) & 0xFF).to(torch.int32) + mantissa = val_u32 & 0x7FFFFF + + round_up = (mantissa > 0) & (exponent != 254) & ~( + (exponent == 0) & (mantissa <= 0x400000) + ) + exponent = exponent + round_up.to(torch.int32) + exponent = torch.where(val == 0, torch.zeros_like(exponent), exponent) + + return exponent.to(torch.uint8) + + +def _e8m0_to_scale_inv(e8m0: torch.Tensor) -> torch.Tensor: + """Convert E8M0 scale bytes back to scale-inverse values.""" + return torch.pow(2.0, e8m0.to(torch.float32) - 127) + + +def _mxfp8_2d_quantize_reference( + x: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Reference MXFP8 2D quantization using one scale per 32x32 block.""" + rows, cols = x.shape + assert rows % MXFP8_BLOCK_SIZE == 0 + assert cols % MXFP8_BLOCK_SIZE == 0 + + block_rows = rows // MXFP8_BLOCK_SIZE + block_cols = cols // MXFP8_BLOCK_SIZE + x_blocks = x.view( + block_rows, + MXFP8_BLOCK_SIZE, + block_cols, + MXFP8_BLOCK_SIZE, + ).permute(0, 2, 1, 3) + + block_amax = torch.amax(torch.abs(x_blocks.to(torch.float32)), dim=(-1, -2)) + block_scale_e8m0 = _float_to_e8m0(block_amax) + block_scale_inv = _e8m0_to_scale_inv(block_scale_e8m0) + + x_scaled = x_blocks.to(torch.float32) / block_scale_inv[:, :, None, None] + x_quantized = x_scaled.to(torch.float8_e4m3fn) + rowwise_data = x_quantized.permute(0, 2, 1, 3).reshape(rows, cols) + rowwise_scale = block_scale_e8m0.repeat_interleave(MXFP8_BLOCK_SIZE, dim=0) + columnwise_scale = block_scale_e8m0.repeat_interleave(MXFP8_BLOCK_SIZE, dim=1) + + return rowwise_data, rowwise_scale, columnwise_scale + + +def _quantize( + quantizer: MXFP8Quantizer, + x: torch.Tensor, + use_preallocated_output: bool, +) -> torch.Tensor: + """Quantize with either the C++ allocator or an explicitly preallocated output.""" + if not use_preallocated_output: + return quantizer(x) + + out = quantizer.make_empty( + x.shape, + dtype=x.dtype, + device=x.device, + requires_grad=False, + ) + return quantizer.update_quantized(x, out) + + +def _assert_rowwise_scales_are_2d(scales: torch.Tensor, rows: int, cols: int) -> None: + """Check that each 32x32 block uses one rowwise scale for all rows.""" + valid = _valid_rowwise_scale(scales, rows, cols) + block_rows = (rows + MXFP8_BLOCK_SIZE - 1) // MXFP8_BLOCK_SIZE + block_cols = valid.shape[1] + for block_row in range(block_rows): + row_start = block_row * MXFP8_BLOCK_SIZE + row_end = min(row_start + MXFP8_BLOCK_SIZE, rows) + for block_col in range(block_cols): + block_scales = valid[row_start:row_end, block_col] + torch.testing.assert_close( + block_scales, + block_scales[0].expand_as(block_scales), + atol=0, + rtol=0, + ) + + +def _assert_bidirectional_scales_are_2d( + rowwise_scales: torch.Tensor, + columnwise_scales: torch.Tensor, + rows: int, + cols: int, +) -> None: + """Check that rowwise and columnwise metadata agree per 32x32 block.""" + rowwise_valid = _valid_rowwise_scale(rowwise_scales, rows, cols) + columnwise_valid = _valid_columnwise_scale(columnwise_scales, rows, cols) + block_rows = columnwise_valid.shape[0] + block_cols = rowwise_valid.shape[1] + + for block_row in range(block_rows): + row_start = block_row * MXFP8_BLOCK_SIZE + row_end = min(row_start + MXFP8_BLOCK_SIZE, rows) + for block_col in range(block_cols): + col_start = block_col * MXFP8_BLOCK_SIZE + col_end = min(col_start + MXFP8_BLOCK_SIZE, cols) + rowwise_block_scales = rowwise_valid[row_start:row_end, block_col] + columnwise_block_scales = columnwise_valid[block_row, col_start:col_end] + torch.testing.assert_close( + rowwise_block_scales, + rowwise_block_scales[0].expand_as(rowwise_block_scales), + atol=0, + rtol=0, + ) + torch.testing.assert_close( + columnwise_block_scales, + columnwise_block_scales[0].expand_as(columnwise_block_scales), + atol=0, + rtol=0, + ) + torch.testing.assert_close( + rowwise_block_scales[0], + columnwise_block_scales[0], + atol=0, + rtol=0, + ) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("columnwise", [False, True], ids=["rowwise_only", "bidirectional"]) +def test_mxfp8_2d_quantize_scales_match_known_block_amax(columnwise: bool) -> None: + """Check exact 2D MXFP8 scale bytes on a hand-built 2x2 block matrix. + + The input has one known amax per 32x32 block. Each amax is chosen so the + expected E8M0 scale byte is known exactly, which makes this test independent + of the random-input comparisons below. Rowwise-only mode should emit only + rowwise tensors, while bidirectional mode should emit matching rowwise and + columnwise scale metadata for the same 2D blocks. + """ + rows, cols = 64, 64 + x = torch.zeros((rows, cols), dtype=torch.float32, device="cuda") + block_exponents = torch.tensor([[-2, -1], [0, 1]], device="cuda") + expected_block_scales = (block_exponents + 127).to(torch.uint8) + + for block_row in range(block_exponents.shape[0]): + for block_col in range(block_exponents.shape[1]): + row_start = block_row * MXFP8_BLOCK_SIZE + col_start = block_col * MXFP8_BLOCK_SIZE + amax = 448.0 * (2.0 ** int(block_exponents[block_row, block_col].item())) + x[ + row_start : row_start + MXFP8_BLOCK_SIZE, + col_start : col_start + MXFP8_BLOCK_SIZE, + ] = amax * 0.5 + x[row_start, col_start] = amax + + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=columnwise, + with_2d_quantization=True, + ) + out = quantizer(x) + + expected_rowwise_scales = expected_block_scales.repeat_interleave(MXFP8_BLOCK_SIZE, dim=0) + torch.testing.assert_close( + _valid_rowwise_scale(out._rowwise_scale_inv, rows, cols), + expected_rowwise_scales, + atol=0, + rtol=0, + ) + + if columnwise: + expected_columnwise_scales = expected_block_scales.repeat_interleave( + MXFP8_BLOCK_SIZE, dim=1 + ) + torch.testing.assert_close( + _valid_columnwise_scale(out._columnwise_scale_inv, rows, cols), + expected_columnwise_scales, + atol=0, + rtol=0, + ) + else: + assert out._columnwise_data is None + assert out._columnwise_scale_inv is None + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("shape", MXFP8_TEST_SHAPES) +@pytest.mark.parametrize("dtype", MXFP8_TEST_DTYPES, ids=str) +@pytest.mark.parametrize("columnwise", [False, True], ids=["rowwise_only", "bidirectional"]) +@pytest.mark.parametrize( + "use_preallocated_output", + [False, True], + ids=["cpp_allocator", "preallocated_output"], +) +def test_mxfp8_2d_quantize_matches_torch_reference( + shape: tuple[int, int], + dtype: torch.dtype, + columnwise: bool, + use_preallocated_output: bool, +) -> None: + """Compare random-input MXFP8 2D data and scales against a PyTorch reference.""" + rows, cols = shape + torch.manual_seed(9012) + x = torch.randn(shape, dtype=dtype, device="cuda") + + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=columnwise, + with_2d_quantization=True, + ) + out = _quantize(quantizer, x, use_preallocated_output) + + ref_data, ref_rowwise_scale, ref_columnwise_scale = _mxfp8_2d_quantize_reference(x) + ref_data_uint8 = ref_data.view(torch.uint8) + + assert out._rowwise_data is not None + assert out._rowwise_scale_inv is not None + torch.testing.assert_close( + out._rowwise_data.view(torch.uint8), + ref_data_uint8, + atol=0, + rtol=0, + ) + torch.testing.assert_close( + _valid_rowwise_scale(out._rowwise_scale_inv, rows, cols), + ref_rowwise_scale, + atol=0, + rtol=0, + ) + + if columnwise: + assert out._columnwise_data is not None + assert out._columnwise_scale_inv is not None + torch.testing.assert_close( + out._columnwise_data.view(torch.uint8), + ref_data_uint8, + atol=0, + rtol=0, + ) + torch.testing.assert_close( + _valid_columnwise_scale(out._columnwise_scale_inv, rows, cols), + ref_columnwise_scale, + atol=0, + rtol=0, + ) + else: + assert out._columnwise_data is None + assert out._columnwise_scale_inv is None + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("shape", MXFP8_TEST_SHAPES) +@pytest.mark.parametrize("dtype", MXFP8_TEST_DTYPES, ids=str) +@pytest.mark.parametrize( + "use_preallocated_output", + [False, True], + ids=["cpp_allocator", "preallocated_output"], +) +def test_mxfp8_2d_quantize_rowwise_only_matches_bidirectional( + shape: tuple[int, int], + dtype: torch.dtype, + use_preallocated_output: bool, +) -> None: + """2D MXFP8 must support inference-style rowwise-only weight quantization.""" + rows, cols = shape + torch.manual_seed(1234) + x = torch.randn(shape, dtype=dtype, device="cuda") + + rowwise_only_quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + with_2d_quantization=True, + ) + bidirectional_quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + with_2d_quantization=True, + ) + + rowwise_only = _quantize(rowwise_only_quantizer, x, use_preallocated_output) + bidirectional = _quantize(bidirectional_quantizer, x, use_preallocated_output) + + assert rowwise_only._rowwise_data is not None + assert rowwise_only._rowwise_scale_inv is not None + assert rowwise_only._columnwise_data is None + assert rowwise_only._columnwise_scale_inv is None + + torch.testing.assert_close( + rowwise_only._rowwise_data.view(torch.uint8), + bidirectional._rowwise_data.view(torch.uint8), + atol=0, + rtol=0, + ) + torch.testing.assert_close( + _valid_rowwise_scale(rowwise_only._rowwise_scale_inv, rows, cols), + _valid_rowwise_scale(bidirectional._rowwise_scale_inv, rows, cols), + atol=0, + rtol=0, + ) + _assert_rowwise_scales_are_2d(rowwise_only._rowwise_scale_inv, rows, cols) + + +@pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) +@pytest.mark.parametrize("shape", MXFP8_TEST_SHAPES) +@pytest.mark.parametrize("dtype", MXFP8_TEST_DTYPES, ids=str) +@pytest.mark.parametrize( + "use_preallocated_output", + [False, True], + ids=["cpp_allocator", "preallocated_output"], +) +def test_mxfp8_2d_quantize_bidirectional_scales_match( + shape: tuple[int, int], + dtype: torch.dtype, + use_preallocated_output: bool, +) -> None: + """Rowwise and columnwise scale metadata should encode the same 32x32 block scales.""" + rows, cols = shape + torch.manual_seed(5678) + x = torch.randn(shape, dtype=dtype, device="cuda") + + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + with_2d_quantization=True, + ) + out = _quantize(quantizer, x, use_preallocated_output) + + assert out._rowwise_scale_inv is not None + assert out._columnwise_scale_inv is not None + _assert_bidirectional_scales_are_2d( + out._rowwise_scale_inv, + out._columnwise_scale_inv, + rows, + cols, + ) + + +def test_mxfp8_recipe_default_2d_quantization_disabled() -> None: + """MXFP8 2D quantization is opt-in.""" + mxfp8_recipe = MXFP8BlockScaling() + assert mxfp8_recipe.enable_2d_quantization is False + + state = MXFP8BlockScalingRecipeState( + recipe=mxfp8_recipe, + mode="forward", + num_quantizers=3, + roles=[ + QuantizerRole(module_type="linear", tensor_type="input"), + QuantizerRole(module_type="linear", tensor_type="weight"), + QuantizerRole(module_type="linear", tensor_type="output"), + ], + ) + assert [q.with_2d_quantization for q in state.make_quantizers()] == [ + False, + False, + False, + ] + + +def test_mxfp8_recipe_state_uses_2d_only_for_forward_weights() -> None: + """Only forward weight quantizers should inherit MXFP8 2D quantization.""" + recipe = MXFP8BlockScaling(enable_2d_quantization=True) + roles = [ + QuantizerRole(module_type="linear", tensor_type="input"), + QuantizerRole(module_type="linear", tensor_type="weight"), + QuantizerRole(module_type="linear", tensor_type="output"), + ] + state = MXFP8BlockScalingRecipeState( + recipe=recipe, mode="forward", num_quantizers=3, roles=roles + ) + quantizers = state.make_quantizers() + + assert [q.with_2d_quantization for q in quantizers] == [False, True, False] + + backward_state = MXFP8BlockScalingRecipeState( + recipe=recipe, + mode="backward", + num_quantizers=2, + roles=[ + QuantizerRole(module_type="linear", tensor_type="grad_output"), + QuantizerRole(module_type="linear", tensor_type="grad_input"), + ], + ) + assert [q.with_2d_quantization for q in backward_state.make_quantizers()] == [ + False, + False, + ] + + +def test_mxfp8_recipe_state_2d_requires_explicit_weight_role() -> None: + """MXFP8 2D should not enable itself for unknown positional slots.""" + recipe = MXFP8BlockScaling(enable_2d_quantization=True) + state = MXFP8BlockScalingRecipeState(recipe=recipe, mode="forward", num_quantizers=3) + assert [q.with_2d_quantization for q in state.make_quantizers()] == [ + False, + False, + False, + ] + + +def test_mxfp8_recipe_state_2d_ignores_non_linear_roles() -> None: + """MXFP8 2D is limited to Linear-style weight quantizers.""" + recipe = MXFP8BlockScaling(enable_2d_quantization=True) + roles = [ + QuantizerRole(module_type="dpa", tensor_type="qkv"), + QuantizerRole(module_type="dpa", tensor_type="weight"), + QuantizerRole(module_type="", tensor_type="weight"), + ] + state = MXFP8BlockScalingRecipeState( + recipe=recipe, + mode="forward", + num_quantizers=len(roles), + roles=roles, + ) + assert [q.with_2d_quantization for q in state.make_quantizers()] == [ + False, + False, + False, + ] + + +def test_mxfp8_quantizer_copy_preserves_2d_flag() -> None: + """MXFP8Quantizer.copy should preserve the 2D quantization setting.""" + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=False, + with_2d_quantization=True, + ) + copied = quantizer.copy() + assert copied.rowwise_usage is True + assert copied.columnwise_usage is False + assert copied.with_2d_quantization is True diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 123362ce10..303ed00e60 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -85,7 +85,7 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, Tensor *dummy_workspace_tensor = nullptr; mxfp8::quantize( *input_tensor, dummy_input_tensor, noop_tensor, output_tensor, dummy_dbias_tensor, - dummy_workspace_tensor, stream); + dummy_workspace_tensor, quant_config_cpp.mxfp8_2d_quantization, stream); break; } case NVTE_NVFP4_1D_SCALING: { @@ -233,7 +233,7 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens case NVTE_MXFP8_1D_SCALING: { mxfp8::quantize( *grad_tensor, input_tensor, noop_tensor, output_tensor, dbias_tensor, workspace_tensor, - stream); + quant_config_cpp.mxfp8_2d_quantization, stream); break; } case NVTE_NVFP4_1D_SCALING: { diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 1549a292d8..1d4489d3e7 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -45,7 +45,7 @@ constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM_X; // 4 = 128 template + size_t CHUNK_DIM_X, size_t THREADS_PER_CHUNK, bool kIs2DBlockScaling> __global__ void __launch_bounds__(THREADS_PER_CHUNK) quantize_mxfp8_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_act_input, @@ -264,6 +264,13 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } } + if constexpr (kIs2DBlockScaling) { +#pragma unroll + for (int i = 16; i > 0; i /= 2) { + thread_amax = fmaxf(thread_amax, __shfl_xor_sync(0xffffffff, thread_amax, i)); + } + } + // 2. Compute E8M0 scaling factor const e8m0_t biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); @@ -415,8 +422,26 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } // 2. Compute E8M0 scaling factor - const e8m0_t biased_exponent = - ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + e8m0_t biased_exponent; + if constexpr (kIs2DBlockScaling) { + __shared__ e8m0_t block_scales_2d[THREADS_X]; + __shared__ float block_amax_2d[THREADS_X * THREADS_Y]; + block_amax_2d[tid_X_rowwise * THREADS_Y + tid_Y_rowwise] = thread_amax; + __syncthreads(); + if (tid_Y_rowwise == 0) { + float amax_2d = 0.0f; +#pragma unroll + for (int i = 0; i < THREADS_Y; ++i) { + amax_2d = fmaxf(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); + } + block_scales_2d[tid_X_rowwise] = + ptx::float_to_e8m0(amax_2d * Quantized_Limits::max_norm_rcp); + } + __syncthreads(); + biased_exponent = block_scales_2d[tid_X_rowwise]; + } else { + biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + } const int stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; const int stage_scales_offset_X = scales_offset_X_rowwise; size_t scale_idx; @@ -558,7 +583,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) template void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, // TODO (ksivamani) - Tensor *output, Tensor *dbias, Tensor *workspace, cudaStream_t stream) { + Tensor *output, Tensor *dbias, Tensor *workspace, const bool use_2d_quantization, + cudaStream_t stream) { using namespace quantize_kernel; checkCuDriverContext(stream); @@ -664,7 +690,8 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, bidimensional_specialized_grid_fits); if (specialized::hasSpec() && - !WITH_GEMM_SWIZZLED_SCALES && scaling_type_has_specialized_support) { + !WITH_GEMM_SWIZZLED_SCALES && !use_2d_quantization && + scaling_type_has_specialized_support) { switch (scaling_type) { case ScalingType::ROWWISE: { using traits = specialized::CastTraits; @@ -795,45 +822,60 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, switch (scaling_type) { case ScalingType::ROWWISE: { - auto kernel = quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, - workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_2d_quantization, kIs2DBlockScaling, { + auto kernel = + quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, + noop_ptr, workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise); + }); break; } case ScalingType::COLWISE: { - auto kernel = quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, - workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_2d_quantization, kIs2DBlockScaling, { + auto kernel = + quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, + noop_ptr, workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise); + }); break; } case ScalingType::BIDIMENSIONAL: { - auto kernel = quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, - workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_2d_quantization, kIs2DBlockScaling, { + auto kernel = + quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, + noop_ptr, workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise); + }); break; } } NVTE_CHECK_CUDA(cudaGetLastError()); diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 12479f2a9c..b580e0c7f6 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -477,6 +477,7 @@ struct QuantizationConfig { bool nvfp4_2d_quantization = false; bool stochastic_rounding = false; bool use_fast_math = false; + bool mxfp8_2d_quantization = false; static constexpr size_t attr_sizes[] = { sizeof(uint8_t), // force_pow_2_scales @@ -486,7 +487,8 @@ struct QuantizationConfig { sizeof(NVTETensor), // rng_seed and offset sizeof(uint8_t), // nvfp4_2d_quantization sizeof(uint8_t), // stochastic_rounding - sizeof(uint8_t) // use_fast_math + sizeof(uint8_t), // use_fast_math + sizeof(uint8_t) // mxfp8_2d_quantization }; }; diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 045ae88893..ac62f02de1 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -381,6 +381,8 @@ enum NVTEQuantizationConfigAttribute { * inconsistently between kernels. */ kNVTEQuantizationConfigUseFastMath = 7, + /*! Whether to use 2D block scaling for MXFP8 */ + kNVTEQuantizationConfigMXFP82DQuantization = 8, kNVTEQuantizationConfigNumAttributes }; @@ -1304,6 +1306,13 @@ class QuantizationConfigWrapper { &val, sizeof(val)); } + /*! \brief Set whether to use 2D block scaling for MXFP8 */ + void set_mxfp8_2d_quantization(bool mxfp8_2d_quantization) { + const auto val = static_cast(mxfp8_2d_quantization); + nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigMXFP82DQuantization, + &val, sizeof(val)); + } + /*! \brief Set whether to use stochastic rounding */ void set_stochastic_rounding(bool stochastic_rounding) { const auto val = static_cast(stochastic_rounding); diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index b773a81d1b..d0f4561e3e 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -359,8 +359,11 @@ class MXFP8BlockScaling(Recipe): `high_precision` keeps original high-precision operands for backward, and `dequantized` dequantizes saved operands to the active high-precision compute dtype (e.g. BF16/FP16/FP32) for backward. + enable_2d_quantization : bool, default = False + If set to `True`, 2D block scaling is used for weight tensors. """ + enable_2d_quantization: bool = os.getenv("NVTE_MXFP8_ENABLE_2D_QUANTIZATION", "0") == "1" margin: int = 0 fp8_format: Format = Format.E4M3 fp8_dpa: bool = False @@ -378,6 +381,7 @@ def _make_repr(self) -> str: f"recipe_type={self.__class__.__name__}, " f"margin={self.margin}, " f"format={str(self.fp8_format).split('.')[1]}, " + f"enable_2d_quantization={self.enable_2d_quantization}, " f"backward_override={self.backward_override}" ) diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 1a52d76019..f43a966a19 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -1049,6 +1049,9 @@ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigUseFastMath: bool_to_uint8(config_.use_fast_math, buf); break; + case kNVTEQuantizationConfigMXFP82DQuantization: + bool_to_uint8(config_.mxfp8_2d_quantization, buf); + break; default: NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); } @@ -1104,6 +1107,9 @@ void nvte_set_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigUseFastMath: uint8_to_bool(buf, config_.use_fast_math); break; + case kNVTEQuantizationConfigMXFP82DQuantization: + uint8_to_bool(buf, config_.mxfp8_2d_quantization); + break; default: NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); } diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 94350da1e6..8364e46fec 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -293,6 +293,8 @@ class Float8BlockQuantizer : public Quantizer { class MXFP8Quantizer : public Quantizer { public: + bool with_2d_quantization = false; + explicit MXFP8Quantizer(const py::handle& quantizer); NVTEScalingMode get_scaling_mode() const override { return NVTE_MXFP8_1D_SCALING; } diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 7045995dd7..f5aace4b5a 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1364,6 +1364,7 @@ std::vector Float8BlockQuantizer::get_scale_shape(const std::vectordtype = quantizer.attr("dtype").cast(); + this->with_2d_quantization = quantizer.attr("with_2d_quantization").cast(); } void MXFP8Quantizer::set_quantization_params(TensorWrapper* tensor) const {} @@ -1687,6 +1688,9 @@ void MXFP8Quantizer::quantize(const TensorWrapper& input, TensorWrapper& out, if (noop_flag) { quant_config.set_noop_tensor(noop_flag->data()); } + if (this->with_2d_quantization) { + quant_config.set_mxfp8_2d_quantization(true); + } NVTE_SCOPED_GIL_RELEASE({ nvte_quantize_v2(input.data(), out.data(), quant_config, at::cuda::getCurrentCUDAStream()); }); diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 0c40723517..e503a9fea9 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1506,7 +1506,23 @@ def make_quantizers(self) -> list: # TODO(ksivamani); Find better design for this, adding here to avoid circular import. from .tensor.mxfp8_tensor import MXFP8Quantizer - return [MXFP8Quantizer(self.dtype) for i in range(self.num_quantizers)] + if self.mode not in ("forward", "backward"): + raise RuntimeError(f"Unexpected recipe mode ({self.mode})") + + if self.mode == "backward" or not self.recipe.enable_2d_quantization: + return [MXFP8Quantizer(self.dtype) for i in range(self.num_quantizers)] + + def _use_2d_quantization(idx: int) -> bool: + role = self._slot_role(idx) + return role.module_type in ("linear", "grouped_linear") and role.tensor_type == "weight" + + return [ + MXFP8Quantizer( + self.dtype, + with_2d_quantization=_use_2d_quantization(idx), + ) + for idx in range(self.num_quantizers) + ] class Float8BlockScalingRecipeState(RecipeState): diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 134f8b5a61..e2b368732b 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -34,6 +34,7 @@ class MXFP8Quantizer(Quantizer): """ dtype: TE_DType + with_2d_quantization: bool def __init__( self, @@ -41,9 +42,11 @@ def __init__( *, rowwise: bool = True, columnwise: bool = True, + with_2d_quantization: bool = False, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) self.dtype = fp8_dtype + self.with_2d_quantization = with_2d_quantization def copy(self) -> MXFP8Quantizer: """Create shallow copy""" @@ -52,6 +55,7 @@ def copy(self) -> MXFP8Quantizer: fp8_dtype=self.dtype, rowwise=self.rowwise_usage, columnwise=self.columnwise_usage, + with_2d_quantization=self.with_2d_quantization, ) quantizer.internal = self.internal quantizer.optimize_for_gemm = self.optimize_for_gemm From 93e34deee8f0c43113320404de4bf9148d694b1d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 14 May 2026 10:59:32 +0000 Subject: [PATCH 2/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- tests/pytorch/test_mxfp8_2d_quantize.py | 8 +- .../common/cast/mxfp8/quantize_mxfp8.cuh | 90 +++++++++---------- 2 files changed, 46 insertions(+), 52 deletions(-) diff --git a/tests/pytorch/test_mxfp8_2d_quantize.py b/tests/pytorch/test_mxfp8_2d_quantize.py index d5baf213c9..3da1e20d09 100644 --- a/tests/pytorch/test_mxfp8_2d_quantize.py +++ b/tests/pytorch/test_mxfp8_2d_quantize.py @@ -53,9 +53,7 @@ def _float_to_e8m0(amax: torch.Tensor) -> torch.Tensor: exponent = ((val_u32 >> 23) & 0xFF).to(torch.int32) mantissa = val_u32 & 0x7FFFFF - round_up = (mantissa > 0) & (exponent != 254) & ~( - (exponent == 0) & (mantissa <= 0x400000) - ) + round_up = (mantissa > 0) & (exponent != 254) & ~((exponent == 0) & (mantissa <= 0x400000)) exponent = exponent + round_up.to(torch.int32) exponent = torch.where(val == 0, torch.zeros_like(exponent), exponent) @@ -197,7 +195,9 @@ def test_mxfp8_2d_quantize_scales_match_known_block_amax(columnwise: bool) -> No x[ row_start : row_start + MXFP8_BLOCK_SIZE, col_start : col_start + MXFP8_BLOCK_SIZE, - ] = amax * 0.5 + ] = ( + amax * 0.5 + ) x[row_start, col_start] = amax quantizer = MXFP8Quantizer( diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 1d4489d3e7..a98e36dc76 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -822,60 +822,54 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, switch (scaling_type) { case ScalingType::ROWWISE: { - TRANSFORMER_ENGINE_SWITCH_CONDITION( - use_2d_quantization, kIs2DBlockScaling, { - auto kernel = - quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, - noop_ptr, workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); - }); + TRANSFORMER_ENGINE_SWITCH_CONDITION(use_2d_quantization, kIs2DBlockScaling, { + auto kernel = + quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, + workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise); + }); break; } case ScalingType::COLWISE: { - TRANSFORMER_ENGINE_SWITCH_CONDITION( - use_2d_quantization, kIs2DBlockScaling, { - auto kernel = - quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, - noop_ptr, workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); - }); + TRANSFORMER_ENGINE_SWITCH_CONDITION(use_2d_quantization, kIs2DBlockScaling, { + auto kernel = + quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, + workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise); + }); break; } case ScalingType::BIDIMENSIONAL: { - TRANSFORMER_ENGINE_SWITCH_CONDITION( - use_2d_quantization, kIs2DBlockScaling, { - auto kernel = - quantize_mxfp8_kernel; - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, - noop_ptr, workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); - }); + TRANSFORMER_ENGINE_SWITCH_CONDITION(use_2d_quantization, kIs2DBlockScaling, { + auto kernel = + quantize_mxfp8_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, + workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, + scale_stride_colwise); + }); break; } } NVTE_CHECK_CUDA(cudaGetLastError()); From 30ad3a452b4513b7fe8f712668c52bba7a7f1182 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Thu, 14 May 2026 07:13:13 -0700 Subject: [PATCH 3/7] Fix comments Signed-off-by: kunlunl --- tests/pytorch/test_mxfp8_2d_quantize.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_mxfp8_2d_quantize.py b/tests/pytorch/test_mxfp8_2d_quantize.py index 3da1e20d09..5f5e2b6d85 100644 --- a/tests/pytorch/test_mxfp8_2d_quantize.py +++ b/tests/pytorch/test_mxfp8_2d_quantize.py @@ -387,8 +387,11 @@ def test_mxfp8_2d_quantize_bidirectional_scales_match( ) -def test_mxfp8_recipe_default_2d_quantization_disabled() -> None: +def test_mxfp8_recipe_default_2d_quantization_disabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: """MXFP8 2D quantization is opt-in.""" + monkeypatch.setenv("NVTE_MXFP8_ENABLE_2D_QUANTIZATION", "0") mxfp8_recipe = MXFP8BlockScaling() assert mxfp8_recipe.enable_2d_quantization is False From 9142f1d69d407a994e3065af6e2c023df908d1fe Mon Sep 17 00:00:00 2001 From: kunlunl Date: Fri, 15 May 2026 10:08:23 -0700 Subject: [PATCH 4/7] Fix comments Signed-off-by: kunlunl --- tests/cpp/operator/test_cast_mxfp8.cu | 254 ++++++++++++++++++ tests/pytorch/test_mxfp8_2d_quantize.py | 10 +- .../common/cast/mxfp8/quantize_mxfp8.cuh | 23 +- transformer_engine/common/recipe/__init__.py | 6 +- transformer_engine/pytorch/quantization.py | 2 +- 5 files changed, 277 insertions(+), 18 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8.cu b/tests/cpp/operator/test_cast_mxfp8.cu index c7c778ce1e..35f0707167 100644 --- a/tests/cpp/operator/test_cast_mxfp8.cu +++ b/tests/cpp/operator/test_cast_mxfp8.cu @@ -4,6 +4,8 @@ * See LICENSE for license information. ************************************************************************/ +#include + #include #include #include @@ -36,6 +38,12 @@ enum ActivationType { SReLU }; +enum MXFP82DScalingDirection { + RowwiseOnly, + ColwiseOnly, + Bidirectional +}; + template void compute_ref(const ProcessingMethod processing_method, float (*OP)(const float), @@ -165,6 +173,70 @@ void compute_ref(const ProcessingMethod processing_method, } } +template +void compute_ref_2d_quantize(const bool rowwise, + const bool colwise, + const InputType* input, + OutputType* output_rowwise, + OutputType* output_colwise, + fp8e8m0* output_scales_rowwise, + fp8e8m0* output_scales_colwise, + const size_t rows, + const size_t cols, + const size_t scales_stride_rowwise, + const size_t scales_stride_colwise) { + const size_t tile_size_Y = 32; + const size_t tile_size_X = 32; + const size_t tiles_num_Y = (rows + tile_size_Y - 1) / tile_size_Y; + const size_t tiles_num_X = (cols + tile_size_X - 1) / tile_size_X; + + #pragma omp parallel for collapse(2) proc_bind(spread) + for (size_t tile_Y = 0; tile_Y < tiles_num_Y; ++tile_Y) { + for (size_t tile_X = 0; tile_X < tiles_num_X; ++tile_X) { + const size_t i_min = tile_Y * tile_size_Y; + const size_t i_max = std::min(i_min + tile_size_Y, rows); + const size_t j_min = tile_X * tile_size_X; + const size_t j_max = std::min(j_min + tile_size_X, cols); + + float block_amax = 0.0f; + for (size_t i = i_min; i < i_max; ++i) { + for (size_t j = j_min; j < j_max; ++j) { + const size_t idx = i * cols + j; + block_amax = std::max(block_amax, std::abs(static_cast(input[idx]))); + } + } + + const fp8e8m0 biased_exponent = + float_to_e8m0(block_amax * Quantized_Limits::max_reciprocal()); + const float scale_reciprocal = exp2f_rcp(biased_exponent); + + if (rowwise) { + for (size_t i = i_min; i < i_max; ++i) { + output_scales_rowwise[i * scales_stride_rowwise + tile_X] = biased_exponent; + for (size_t j = j_min; j < j_max; ++j) { + const size_t idx = i * cols + j; + output_rowwise[idx] = + static_cast(static_cast(input[idx]) * + scale_reciprocal); + } + } + } + + if (colwise) { + for (size_t j = j_min; j < j_max; ++j) { + output_scales_colwise[tile_Y * scales_stride_colwise + j] = biased_exponent; + for (size_t i = i_min; i < i_max; ++i) { + const size_t idx = i * cols + j; + output_colwise[idx] = + static_cast(static_cast(input[idx]) * + scale_reciprocal); + } + } + } + } + } +} + /** * Scaling along single dimension (either rows or columns) * Produces one set of output data and the corresponding data of the fused operation (dbias): @@ -521,6 +593,106 @@ void performTest_x2(const ProcessingMethod processing_method, } } +template +void performTest_2d_quantize(const std::vector& shape, + const MXFP82DScalingDirection scaling_direction, + InputsFillCase fill_case) { + using namespace test; + using EncodingType = fp32; + DType itype = TypeInfo::dtype; + DType otype = TypeInfo::dtype; + + if (shape.size() < 2) { + GTEST_SKIP(); + } + + const size_t rows = first_dimension(shape); + const size_t cols = last_dimension(shape); + + const bool rowwise = scaling_direction != MXFP82DScalingDirection::ColwiseOnly; + const bool colwise = scaling_direction != MXFP82DScalingDirection::RowwiseOnly; + + const std::array scale_dims_rowwise = get_scale_tensor_dims(rows, cols, 1, 32); + const std::array scale_dims_colwise = get_scale_tensor_dims(rows, cols, 32, 1); + + const size_t unpadded_blocks_Y_rowwise = scale_dims_rowwise[0]; + const size_t unpadded_blocks_X_rowwise = scale_dims_rowwise[1]; + const size_t blocks_Y_rowwise = scale_dims_rowwise[2]; + const size_t blocks_X_rowwise = scale_dims_rowwise[3]; + const size_t scales_stride_rowwise = blocks_X_rowwise; + + const size_t unpadded_blocks_Y_colwise = scale_dims_colwise[0]; + const size_t unpadded_blocks_X_colwise = scale_dims_colwise[1]; + const size_t blocks_Y_colwise = scale_dims_colwise[2]; + const size_t blocks_X_colwise = scale_dims_colwise[3]; + const size_t scales_stride_colwise = blocks_X_colwise; + + Tensor input("input", shape, itype); + Tensor output("output", shape, otype, rowwise, colwise, NVTE_MXFP8_1D_SCALING); + + std::unique_ptr ref_output_rowwise = std::make_unique(rows * cols); + std::unique_ptr ref_output_colwise = std::make_unique(rows * cols); + std::unique_ptr ref_scales_rowwise = + std::make_unique(blocks_Y_rowwise * blocks_X_rowwise); + std::unique_ptr ref_scales_colwise = + std::make_unique(blocks_Y_colwise * blocks_X_colwise); + std::fill_n(ref_scales_rowwise.get(), blocks_Y_rowwise * blocks_X_rowwise, 0); + std::fill_n(ref_scales_colwise.get(), blocks_Y_colwise * blocks_X_colwise, 0); + + fillCase(&input, fill_case); + + QuantizationConfigWrapper quant_config; + quant_config.set_mxfp8_2d_quantization(true); + nvte_quantize_v2(input.data(), output.data(), quant_config, 0); + + cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + compute_ref_2d_quantize( + rowwise, + colwise, + input.rowwise_cpu_dptr(), + ref_output_rowwise.get(), + ref_output_colwise.get(), + ref_scales_rowwise.get(), + ref_scales_colwise.get(), + rows, + cols, + scales_stride_rowwise, + scales_stride_colwise); + + const size_t scale_diff_abs_tolerance = 0; + const double abs_tolerable_mismatches_limit = 0.0; + const double rel_tolerable_mismatches_limit = 0.0; + + auto [atol, rtol] = getTolerances(otype); + + if (rowwise) { + size_t mismatches_scales_rowwise = 0; + compare_scaling_factors("scales_rowwise", output.rowwise_cpu_scale_inv_ptr(), + ref_scales_rowwise.get(), unpadded_blocks_Y_rowwise, + unpadded_blocks_X_rowwise, scales_stride_rowwise, + mismatches_scales_rowwise, + scale_diff_abs_tolerance, + abs_tolerable_mismatches_limit, + rel_tolerable_mismatches_limit); + compareResults("output_rowwise", output, ref_output_rowwise.get(), true, atol, rtol, true); + } + + if (colwise) { + size_t mismatches_scales_colwise = 0; + compare_scaling_factors("scales_colwise", output.columnwise_cpu_scale_inv_ptr(), + ref_scales_colwise.get(), unpadded_blocks_Y_colwise, + unpadded_blocks_X_colwise, scales_stride_colwise, + mismatches_scales_colwise, + scale_diff_abs_tolerance, + abs_tolerable_mismatches_limit, + rel_tolerable_mismatches_limit); + compareResults("output_colwise", output, ref_output_colwise.get(), false, atol, rtol, true); + } +} + std::vector> matrix_sizes = { {1, 16}, {16, 48}, @@ -532,6 +704,18 @@ std::vector> matrix_sizes = { {8192, 7168}, }; +std::vector> matrix_sizes_2d_quantize = { + {1, 16}, + {16, 48}, + {65, 80}, + {127, 400}, + {128, 128}, + {993, 512}, + {8, 32, 1024}, + {16, 8, 4, 512}, + {8192, 7168}, +}; + std::vector> block_sizes = { {1, 32}, {32, 1}, @@ -554,6 +738,12 @@ std::vector processing_methods = { ProcessingMethod::CAST_ACT, }; +std::vector scaling_directions_2d_quantize = { + MXFP82DScalingDirection::RowwiseOnly, + MXFP82DScalingDirection::ColwiseOnly, + MXFP82DScalingDirection::Bidirectional, +}; + // Only GeLU activation tests are supported std::vector Activation_types = { ActivationType::Identity, @@ -573,6 +763,13 @@ class FusedCastMXFP8TestSuite : public ::testing::TestWithParam transformer_engine::DType, InputsFillCase>> {}; +class CastMXFP82DQuantizationTestSuite : public ::testing::TestWithParam + , + transformer_engine::DType, + transformer_engine::DType, + InputsFillCase>> {}; + TEST_P(FusedCastMXFP8TestSuite, TestFusedCastMXFP8) { // Skip tests for pre-Blackwell architectures if (getDeviceComputeCapability() < blackwellComputeCapability) { @@ -653,6 +850,29 @@ TEST_P(FusedCastMXFP8TestSuite, TestFusedCastMXFP8) { } } +TEST_P(CastMXFP82DQuantizationTestSuite, TestCastMXFP82DQuantization) { + // Skip tests for pre-Blackwell architectures + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + using namespace transformer_engine; + using namespace test; + + const auto scaling_direction = std::get<0>(GetParam()); + const auto matrix_size = std::get<1>(GetParam()); + const DType input_type = std::get<2>(GetParam()); + const DType output_type = std::get<3>(GetParam()); + const InputsFillCase fill_case = std::get<4>(GetParam()); + + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(input_type, InputType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(output_type, OutputType, + performTest_2d_quantize( + matrix_size, scaling_direction, fill_case); + ); + ); +} + std::string to_string(const ProcessingMethod method) { switch (method) { case ProcessingMethod::CAST_ONLY: return "CAST_ONLY"; @@ -676,6 +896,15 @@ std::string to_string(const ActivationType Act_type) { } } +std::string to_string(const MXFP82DScalingDirection scaling_direction) { + switch (scaling_direction) { + case MXFP82DScalingDirection::RowwiseOnly: return "RowwiseOnly"; + case MXFP82DScalingDirection::ColwiseOnly: return "ColwiseOnly"; + case MXFP82DScalingDirection::Bidirectional: return "Bidirectional"; + default: return ""; + } +} + std::string test_name_generator( const testing::TestParamInfo& info) { std::string name = to_string(std::get<0>(info.param)) + "X" + @@ -692,6 +921,19 @@ std::string test_name_generator( return name; } +std::string mxfp8_2d_quantization_test_name_generator( + const testing::TestParamInfo& info) { + std::string name = to_string(std::get<0>(info.param)); + const auto& shape = std::get<1>(info.param); + for ( const auto& s: shape) { + name += "X" + std::to_string(s); + } + name += "X" + test::typeName(std::get<2>(info.param)) + + "X" + test::typeName(std::get<3>(info.param)) + + "X" + test::caseName(std::get<4>(info.param)); + return name; +} + } // namespace // Test cases with only cast kernels @@ -708,6 +950,18 @@ INSTANTIATE_TEST_SUITE_P( ::testing::ValuesIn(input_scenarios)), test_name_generator); +// Test cases for MXFP8 2D block scaling through the common C++ API. +INSTANTIATE_TEST_SUITE_P( + OperatorTest_CastMXFP8_2DQuantization, + CastMXFP82DQuantizationTestSuite, + ::testing::Combine( + ::testing::ValuesIn(scaling_directions_2d_quantize), + ::testing::ValuesIn(matrix_sizes_2d_quantize), + ::testing::Values(DType::kFloat32, DType::kBFloat16), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), + ::testing::Values(InputsFillCase::uniform)), + mxfp8_2d_quantization_test_name_generator); + // Test cases with varying matrix shapes and block shapes INSTANTIATE_TEST_SUITE_P( OperatorTest_FusedCastMXFP8_Sizes, diff --git a/tests/pytorch/test_mxfp8_2d_quantize.py b/tests/pytorch/test_mxfp8_2d_quantize.py index 5f5e2b6d85..f750eb3991 100644 --- a/tests/pytorch/test_mxfp8_2d_quantize.py +++ b/tests/pytorch/test_mxfp8_2d_quantize.py @@ -23,14 +23,11 @@ MXFP8_TEST_SHAPES = [ (64, 64), (128, 128), - (256, 256), (256, 1024), (1024, 256), (256, 288), (320, 320), (352, 256), - (2048, 2048), - (1024, 2048), (2048, 1024), ] MXFP8_TEST_DTYPES = [torch.float32, torch.bfloat16] @@ -75,6 +72,7 @@ def _mxfp8_2d_quantize_reference( block_rows = rows // MXFP8_BLOCK_SIZE block_cols = cols // MXFP8_BLOCK_SIZE + x_blocks = x.view( block_rows, MXFP8_BLOCK_SIZE, @@ -453,10 +451,11 @@ def test_mxfp8_recipe_state_2d_requires_explicit_weight_role() -> None: ] -def test_mxfp8_recipe_state_2d_ignores_non_linear_roles() -> None: - """MXFP8 2D is limited to Linear-style weight quantizers.""" +def test_mxfp8_recipe_state_2d_ignores_unsupported_roles() -> None: + """MXFP8 2D is limited to regular Linear weight quantizers.""" recipe = MXFP8BlockScaling(enable_2d_quantization=True) roles = [ + QuantizerRole(module_type="grouped_linear", tensor_type="weight"), QuantizerRole(module_type="dpa", tensor_type="qkv"), QuantizerRole(module_type="dpa", tensor_type="weight"), QuantizerRole(module_type="", tensor_type="weight"), @@ -471,6 +470,7 @@ def test_mxfp8_recipe_state_2d_ignores_non_linear_roles() -> None: False, False, False, + False, ] diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index a98e36dc76..5deca0bb8d 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -265,10 +265,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) } if constexpr (kIs2DBlockScaling) { -#pragma unroll - for (int i = 16; i > 0; i /= 2) { - thread_amax = fmaxf(thread_amax, __shfl_xor_sync(0xffffffff, thread_amax, i)); - } + thread_amax = warp_reduce_max_broadcast(thread_amax); } // 2. Compute E8M0 scaling factor @@ -424,18 +421,26 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) // 2. Compute E8M0 scaling factor e8m0_t biased_exponent; if constexpr (kIs2DBlockScaling) { + using AMax2DType = std::conditional_t< + NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v), IType, float>; __shared__ e8m0_t block_scales_2d[THREADS_X]; - __shared__ float block_amax_2d[THREADS_X * THREADS_Y]; - block_amax_2d[tid_X_rowwise * THREADS_Y + tid_Y_rowwise] = thread_amax; + __shared__ AMax2DType block_amax_2d[THREADS_X * THREADS_Y]; + block_amax_2d[tid_X_rowwise * THREADS_Y + tid_Y_rowwise] = + static_cast(thread_amax); __syncthreads(); if (tid_Y_rowwise == 0) { - float amax_2d = 0.0f; + AMax2DType amax_2d = static_cast(0.0f); #pragma unroll for (int i = 0; i < THREADS_Y; ++i) { - amax_2d = fmaxf(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); + if constexpr (std::is_same_v) { + amax_2d = fmaxf(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); + } else { + amax_2d = __hmax(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); + } } block_scales_2d[tid_X_rowwise] = - ptx::float_to_e8m0(amax_2d * Quantized_Limits::max_norm_rcp); + ptx::float_to_e8m0(static_cast(amax_2d) * + Quantized_Limits::max_norm_rcp); } __syncthreads(); biased_exponent = block_scales_2d[tid_X_rowwise]; diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index d0f4561e3e..311615259b 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -363,12 +363,12 @@ class MXFP8BlockScaling(Recipe): If set to `True`, 2D block scaling is used for weight tensors. """ - enable_2d_quantization: bool = os.getenv("NVTE_MXFP8_ENABLE_2D_QUANTIZATION", "0") == "1" margin: int = 0 fp8_format: Format = Format.E4M3 fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) + enable_2d_quantization: bool = os.getenv("NVTE_MXFP8_ENABLE_2D_QUANTIZATION", "0") == "1" def __post_init__(self) -> None: assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." @@ -381,8 +381,8 @@ def _make_repr(self) -> str: f"recipe_type={self.__class__.__name__}, " f"margin={self.margin}, " f"format={str(self.fp8_format).split('.')[1]}, " - f"enable_2d_quantization={self.enable_2d_quantization}, " - f"backward_override={self.backward_override}" + f"backward_override={self.backward_override}, " + f"enable_2d_quantization={self.enable_2d_quantization}" ) diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index e503a9fea9..2282f93243 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1514,7 +1514,7 @@ def make_quantizers(self) -> list: def _use_2d_quantization(idx: int) -> bool: role = self._slot_role(idx) - return role.module_type in ("linear", "grouped_linear") and role.tensor_type == "weight" + return role.module_type == "linear" and role.tensor_type == "weight" return [ MXFP8Quantizer( From 5ef3cfea0a992950a2d26ff215783ceb9b229de6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 15 May 2026 17:09:37 +0000 Subject: [PATCH 5/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../common/cast/mxfp8/quantize_mxfp8.cuh | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 5deca0bb8d..4625ac564a 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -421,8 +421,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) // 2. Compute E8M0 scaling factor e8m0_t biased_exponent; if constexpr (kIs2DBlockScaling) { - using AMax2DType = std::conditional_t< - NO_ACTIVATIONS && (!IS_DBIAS) && (!std::is_same_v), IType, float>; + using AMax2DType = + std::conditional_t), + IType, float>; __shared__ e8m0_t block_scales_2d[THREADS_X]; __shared__ AMax2DType block_amax_2d[THREADS_X * THREADS_Y]; block_amax_2d[tid_X_rowwise * THREADS_Y + tid_Y_rowwise] = @@ -438,9 +439,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) amax_2d = __hmax(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); } } - block_scales_2d[tid_X_rowwise] = - ptx::float_to_e8m0(static_cast(amax_2d) * - Quantized_Limits::max_norm_rcp); + block_scales_2d[tid_X_rowwise] = ptx::float_to_e8m0( + static_cast(amax_2d) * Quantized_Limits::max_norm_rcp); } __syncthreads(); biased_exponent = block_scales_2d[tid_X_rowwise]; From 80a51afe58614c2ae6ce9237eb59bd84745e3780 Mon Sep 17 00:00:00 2001 From: kunlunl Date: Mon, 10 Aug 2026 04:40:46 -0700 Subject: [PATCH 6/7] Support explicit 2D MXFP8 grouped quantization Signed-off-by: kunlunl --- tests/cpp/operator/test_cast_mxfp8_grouped.cu | 76 +++++++++++++--- .../test_mxfp8_group_quantize_graph_safe.py | 26 ++++++ tests/pytorch/test_grouped_linear.py | 13 ++- tests/pytorch/test_mxfp8_2d_quantize.py | 50 +++++------ .../cast/mxfp8/group_quantize_mxfp8.cuh | 86 ++++++++++++++----- transformer_engine/common/recipe/__init__.py | 2 +- .../pytorch/csrc/extensions/cast.cpp | 2 + transformer_engine/pytorch/quantization.py | 5 +- 8 files changed, 195 insertions(+), 65 deletions(-) diff --git a/tests/cpp/operator/test_cast_mxfp8_grouped.cu b/tests/cpp/operator/test_cast_mxfp8_grouped.cu index 80b80da0d8..43ddb66984 100644 --- a/tests/cpp/operator/test_cast_mxfp8_grouped.cu +++ b/tests/cpp/operator/test_cast_mxfp8_grouped.cu @@ -46,6 +46,7 @@ enum ShapeRepresentation { template void compute_ref(const ProcessingMethod processing_method, float (*OP)(const float), + const bool use_2d_quantization, const bool rowwise, const bool colwise, const InputType* input, @@ -117,13 +118,29 @@ void compute_ref(const ProcessingMethod processing_method, } } - if (rowwise) { + float block_amax_2d = 0.0f; + if (use_2d_quantization) { for (size_t i = i_min; i < i_max; ++i) { - float block_amax = 0.0f; - for (size_t j = j_min; j < j_max; ++j) { - const size_t cache_idx = (i - i_min) * tile_size_X + (j - j_min); - block_amax = std::max(block_amax, std::abs(cache_buffer[cache_idx])); + const size_t cache_idx = + (i - i_min) * tile_size_X + (j - j_min); + block_amax_2d = + std::max(block_amax_2d, std::abs(cache_buffer[cache_idx])); + } + } + } + + if (rowwise) { + for (size_t i = i_min; i < i_max; ++i) { + float block_amax = block_amax_2d; + + if (!use_2d_quantization) { + for (size_t j = j_min; j < j_max; ++j) { + const size_t cache_idx = + (i - i_min) * tile_size_X + (j - j_min); + block_amax = + std::max(block_amax, std::abs(cache_buffer[cache_idx])); + } } const fp8e8m0 biased_exponent = float_to_e8m0(block_amax * Quantized_Limits::max_reciprocal()); @@ -140,11 +157,15 @@ void compute_ref(const ProcessingMethod processing_method, } if (colwise) { for (size_t j = j_min; j < j_max; ++j) { - float block_amax = 0.0f; - - for (size_t i = i_min; i < i_max; ++i) { - const size_t cache_idx = (i - i_min) * tile_size_X + (j - j_min); - block_amax = std::max(block_amax, std::abs(cache_buffer[cache_idx])); + float block_amax = block_amax_2d; + + if (!use_2d_quantization) { + for (size_t i = i_min; i < i_max; ++i) { + const size_t cache_idx = + (i - i_min) * tile_size_X + (j - j_min); + block_amax = + std::max(block_amax, std::abs(cache_buffer[cache_idx])); + } } const fp8e8m0 biased_exponent = float_to_e8m0(block_amax * Quantized_Limits::max_reciprocal()); @@ -241,7 +262,8 @@ void performTest(const ProcessingMethod processing_method, const std::vector& last_dims_h, const std::vector& offsets_h, const bool rowwise, - const bool colwise) { + const bool colwise, + const bool use_2d_quantization = false) { using namespace test; DType itype = TypeInfo::dtype; @@ -491,7 +513,7 @@ void performTest(const ProcessingMethod processing_method, InputType* const ref_output_dbias_ptr = ref_output_dbias.data() + dbias_offset; compute_ref( - processing_method, OP, rowwise, colwise, in_ptr, grad_ptr, + processing_method, OP, use_2d_quantization, rowwise, colwise, in_ptr, grad_ptr, out_data_rowwise_ptr, out_data_colwise_ptr, out_scales_rowwise_ptr, out_scales_colwise_ptr, ref_output_dbias_ptr, M, K, @@ -500,6 +522,7 @@ void performTest(const ProcessingMethod processing_method, } QuantizationConfigWrapper quant_config; + quant_config.set_mxfp8_2d_quantization(use_2d_quantization); // GPU Tensor workspace; @@ -803,6 +826,35 @@ TEST_P(GroupedFusedCastMXFP8TestSuite, Test) { ); } +TEST(OperatorTest_GroupedFusedCastMXFP8, Test2DQuantization) { + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + constexpr size_t num_tensors = 2; + const std::vector logical_shape = {256, 128}; + const std::vector first_dims = {128, 128}; + const std::vector last_dims = {128, 128}; + const std::vector offsets = {0, 128 * 128, 2 * 128 * 128}; + + for (const auto scaling_direction : scaling_directions) { + const bool rowwise = scaling_direction != ScalingDirection::COLWISE; + const bool colwise = scaling_direction != ScalingDirection::ROWWISE; + performTest( + ProcessingMethod::CAST_ONLY, + &identity, + ShapeRepresentation::SAME_BOTH_DIMS, + num_tensors, + logical_shape, + first_dims, + last_dims, + offsets, + rowwise, + colwise, + /*use_2d_quantization=*/true); + } +} + std::string to_string(const ProcessingMethod method) { switch (method) { case ProcessingMethod::CAST_ONLY: return "CAST_ONLY"; diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py index d07953ce37..734792dd50 100644 --- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py +++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py @@ -137,6 +137,7 @@ def check_grouped_tensor_mxfp8_versus_reference( return_transpose: bool, split_sections: list[int], optimize_for_gemm: bool = False, + with_2d_quantization: bool = False, ) -> None: te_dtype = te.DType.kFloat8E4M3 @@ -159,6 +160,7 @@ def check_grouped_tensor_mxfp8_versus_reference( fp8_dtype=te_dtype, rowwise=return_rowwise, columnwise=return_transpose, + with_2d_quantization=with_2d_quantization, ) for _ in range(len(split_sections)) ] @@ -330,6 +332,30 @@ def check_grouped_tensor_mxfp8_with_paged_stashing( torch.testing.assert_close(x_sx_t_i, x_sx_t_ref_i, atol=0.0, rtol=0.0) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) +@pytest.mark.parametrize( + "optimize_for_gemm", [True, False], ids=["optimize_for_gemm", "no_optimize_for_gemm"] +) +def test_grouped_tensor_mxfp8_2d_quantization_versus_reference( + quantize_mode: str, + optimize_for_gemm: bool, +) -> None: + """Grouped MXFP8 should match independent 2D quantization of each tensor.""" + return_rowwise = quantize_mode != "columnwise_only" + return_transpose = quantize_mode != "rowwise_only" + check_grouped_tensor_mxfp8_versus_reference( + x_dtype=torch.bfloat16, + M=1024, + N=256, + return_rowwise=return_rowwise, + return_transpose=return_transpose, + split_sections=[256, 256, 256, 256], + optimize_for_gemm=optimize_for_gemm, + with_2d_quantization=True, + ) + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( "M, N", diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py index 10d0dbd227..0139634cb0 100644 --- a/tests/pytorch/test_grouped_linear.py +++ b/tests/pytorch/test_grouped_linear.py @@ -1589,6 +1589,10 @@ def _run_grouped_linear_path( recipe.MXFP8BlockScaling(), marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), ), + pytest.param( + recipe.MXFP8BlockScaling(enable_2d_quantization=True), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + ), pytest.param( recipe.NVFP4BlockScaling(disable_stochastic_rounding=True), marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4), @@ -1600,7 +1604,14 @@ def _run_grouped_linear_path( ), ), ], - ids=["bf16", "fp8_current_scaling", "mxfp8", "nvfp4", "fp8_block_scaling"], + ids=[ + "bf16", + "fp8_current_scaling", + "mxfp8", + "mxfp8_2d", + "nvfp4", + "fp8_block_scaling", + ], ) @pytest.mark.parametrize("bias", _ALL_BOOLEAN) @pytest.mark.parametrize("fp8_model_params", _ALL_BOOLEAN) diff --git a/tests/pytorch/test_mxfp8_2d_quantize.py b/tests/pytorch/test_mxfp8_2d_quantize.py index f750eb3991..a0159720c5 100644 --- a/tests/pytorch/test_mxfp8_2d_quantize.py +++ b/tests/pytorch/test_mxfp8_2d_quantize.py @@ -385,31 +385,6 @@ def test_mxfp8_2d_quantize_bidirectional_scales_match( ) -def test_mxfp8_recipe_default_2d_quantization_disabled( - monkeypatch: pytest.MonkeyPatch, -) -> None: - """MXFP8 2D quantization is opt-in.""" - monkeypatch.setenv("NVTE_MXFP8_ENABLE_2D_QUANTIZATION", "0") - mxfp8_recipe = MXFP8BlockScaling() - assert mxfp8_recipe.enable_2d_quantization is False - - state = MXFP8BlockScalingRecipeState( - recipe=mxfp8_recipe, - mode="forward", - num_quantizers=3, - roles=[ - QuantizerRole(module_type="linear", tensor_type="input"), - QuantizerRole(module_type="linear", tensor_type="weight"), - QuantizerRole(module_type="linear", tensor_type="output"), - ], - ) - assert [q.with_2d_quantization for q in state.make_quantizers()] == [ - False, - False, - False, - ] - - def test_mxfp8_recipe_state_uses_2d_only_for_forward_weights() -> None: """Only forward weight quantizers should inherit MXFP8 2D quantization.""" recipe = MXFP8BlockScaling(enable_2d_quantization=True) @@ -451,11 +426,31 @@ def test_mxfp8_recipe_state_2d_requires_explicit_weight_role() -> None: ] -def test_mxfp8_recipe_state_2d_ignores_unsupported_roles() -> None: - """MXFP8 2D is limited to regular Linear weight quantizers.""" +def test_mxfp8_recipe_state_uses_2d_for_grouped_linear_weights() -> None: + """GroupedLinear weight quantizers should inherit MXFP8 2D quantization.""" recipe = MXFP8BlockScaling(enable_2d_quantization=True) roles = [ + QuantizerRole(module_type="grouped_linear", tensor_type="input"), QuantizerRole(module_type="grouped_linear", tensor_type="weight"), + QuantizerRole(module_type="grouped_linear", tensor_type="output"), + ] + state = MXFP8BlockScalingRecipeState( + recipe=recipe, + mode="forward", + num_quantizers=len(roles), + roles=roles, + ) + assert [q.with_2d_quantization for q in state.make_quantizers()] == [ + False, + True, + False, + ] + + +def test_mxfp8_recipe_state_2d_ignores_unsupported_roles() -> None: + """MXFP8 2D is limited to supported Linear weight quantizers.""" + recipe = MXFP8BlockScaling(enable_2d_quantization=True) + roles = [ QuantizerRole(module_type="dpa", tensor_type="qkv"), QuantizerRole(module_type="dpa", tensor_type="weight"), QuantizerRole(module_type="", tensor_type="weight"), @@ -470,7 +465,6 @@ def test_mxfp8_recipe_state_2d_ignores_unsupported_roles() -> None: False, False, False, - False, ] diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 14832573d7..129292f353 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -79,7 +79,7 @@ constexpr uint THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM_X; // 4 = 128 / template + bool WITH_GEMM_SWIZZLED_SCALES, bool kIs2DBlockScaling> __device__ __forceinline__ void process_colwise_stage( const size_t buff, const int stage, const size_t tid_X_colwise, const size_t scales_offset_Y_colwise, const size_t scales_offset_X_colwise, @@ -163,9 +163,13 @@ __device__ __forceinline__ void process_colwise_stage( "+r"(reinterpret_cast(thread_amax_2x)) : "r"(src_smem_ptr), "r"(IN_SHMEM_STRIDE)); } - const float thread_amax = + float thread_amax = static_cast(__hmax(__habs(thread_amax_2x.x), __habs(thread_amax_2x.y))); + if constexpr (kIs2DBlockScaling) { + thread_amax = warp_reduce_max_broadcast(thread_amax); + } + const e8m0_t biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); // OOB padded region needs to be zeroed out. @@ -238,6 +242,10 @@ __device__ __forceinline__ void process_colwise_stage( } } + if constexpr (kIs2DBlockScaling) { + thread_amax = warp_reduce_max_broadcast(thread_amax); + } + const e8m0_t biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); // OOB padded region needs to be zeroed out. @@ -262,7 +270,7 @@ __device__ __forceinline__ void process_colwise_stage( template + bool WITH_GEMM_SWIZZLED_SCALES, bool kIs2DBlockScaling> __device__ __forceinline__ void process_rowwise_stage( const size_t buff, const size_t stage_offset_Y, const size_t thread_offset_Y_rowwise, const size_t thread_offset_X_rowwise, const int bank_group, @@ -290,6 +298,8 @@ __device__ __forceinline__ void process_rowwise_stage( auto &sOutRowwise = *reinterpret_cast(sOutRowwise_ptr); const size_t i = thread_offset_Y_rowwise; + const size_t tid_Y_rowwise = thread_offset_Y_rowwise; + const size_t tid_X_rowwise = thread_offset_X_rowwise / SCALE_DIM_X; float thread_amax = 0.0f; float rInCompute[SCALE_DIM_X]; @@ -388,8 +398,33 @@ __device__ __forceinline__ void process_rowwise_stage( } } - const e8m0_t biased_exponent = - ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + e8m0_t biased_exponent; + if constexpr (kIs2DBlockScaling) { + using AMax2DType = std::conditional_t; + __shared__ e8m0_t block_scales_2d[THREADS_X]; + __shared__ AMax2DType block_amax_2d[THREADS_X * THREADS_Y]; + block_amax_2d[tid_X_rowwise * THREADS_Y + tid_Y_rowwise] = + static_cast(thread_amax); + __syncthreads(); + if (tid_Y_rowwise == 0) { + AMax2DType amax_2d = static_cast(0.0f); +#pragma unroll + for (int i = 0; i < THREADS_Y; ++i) { + if constexpr (std::is_same_v) { + amax_2d = fmaxf(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); + } else { + amax_2d = __hmax(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); + } + } + block_scales_2d[tid_X_rowwise] = ptx::float_to_e8m0( + static_cast(amax_2d) * Quantized_Limits::max_norm_rcp); + } + __syncthreads(); + biased_exponent = block_scales_2d[tid_X_rowwise]; + } else { + biased_exponent = + ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + } const size_t stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; const size_t stage_scales_offset_X = scales_offset_X_rowwise; @@ -448,7 +483,8 @@ __device__ __forceinline__ void process_rowwise_stage( template + ScalingType SCALING_TYPE, bool WITH_GEMM_SWIZZLED_SCALES, bool kIs2DBlockScaling, + ShapeRepresentation SHAPE_REP> __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel( const __grid_constant__ CUtensorMap tensor_map_input_static, const __grid_constant__ CUtensorMap tensor_map_act_input_static, @@ -678,7 +714,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel const size_t buff = buff_in; if constexpr (COLWISE_SCALING) { process_colwise_stage( + WITH_GEMM_SWIZZLED_SCALES, kIs2DBlockScaling>( buff, stage, tid_X_colwise, scales_offset_Y_colwise, scales_offset_X_colwise, scale_stride_colwise, tensor_base_for_scales, rows, cols, sIn_ptr, sActIn_ptr, sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise); @@ -686,7 +722,7 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel if constexpr (ROWWISE_SCALING) { process_rowwise_stage( + WITH_GEMM_SWIZZLED_SCALES, kIs2DBlockScaling>( buff, stage_offset_Y, thread_offset_Y_rowwise, thread_offset_X_rowwise, bank_group, scales_offset_Y_rowwise, scales_offset_X_rowwise, scale_stride_rowwise, rowwise_scale_is_within_bounds, cols, sIn_ptr, sActIn_ptr, sCachedAct_ptr, @@ -836,6 +872,8 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations const size_t block_size = THREADS_PER_CHUNK; const bool with_gemm_swizzled_scales = output->with_gemm_swizzled_scales; + const bool use_2d_quantization = + quant_config != nullptr && quant_config->mxfp8_2d_quantization; // Logical shape of a tensor with varying all dims is [1, M*K] if (shape_rep != ShapeRepresentation::VARYING_BOTH_DIMS) { @@ -972,20 +1010,24 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations last_dims_ptr, use_rowwise_scaling, use_colwise_scaling, IS_DACT); } - auto kernel = - group_quantize_mxfp8_kernel; - - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size)); - - kernel<<>>( - tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, - tensor_map_output_colwise, num_tensors, first_logical_dim, - last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, - scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, workspace_ptr, - amax_ptr, work_blocks_X, work_blocks_Y); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + use_2d_quantization, kIs2DBlockScaling, { + auto kernel = group_quantize_mxfp8_kernel< + IS_DBIAS, IS_DACT, IS_ACT, ParamOP, OP, IType, OType, + SCALING_TYPE, WITH_GEMM_SWIZZLED_SCALES, kIs2DBlockScaling, + SHAPE_REP>; + + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size)); + + kernel<<>>( + tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise, + tensor_map_output_colwise, num_tensors, first_logical_dim, + last_logical_dim, offsets_ptr, first_dims_ptr, last_dims_ptr, + scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, workspace_ptr, + amax_ptr, work_blocks_X, work_blocks_Y); + }); if constexpr (IS_DBIAS) { common::grouped_reduce_dbias( diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index ad7c8234ed..a6575b4d1e 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -370,7 +370,7 @@ class MXFP8BlockScaling(Recipe): fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) - enable_2d_quantization: bool = os.getenv("NVTE_MXFP8_ENABLE_2D_QUANTIZATION", "0") == "1" + enable_2d_quantization: bool = False def __post_init__(self) -> None: assert self.fp8_format != Format.E5M2, "Pure E5M2 training is not supported." diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8b1cd384aa..9e7411d27a 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -364,10 +364,12 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const break; } case GroupedQuantizationMode::MXFP8_GROUPED_QUANTIZE: { + auto *mxfp8_quantizer_cpp = static_cast(quantizer_cpp.get()); QuantizationConfigWrapper quant_config_cpp; if (noop_flag_cpp.has_value()) { quant_config_cpp.set_noop_tensor(noop_flag_cpp->data()); } + quant_config_cpp.set_mxfp8_2d_quantization(mxfp8_quantizer_cpp->with_2d_quantization); NVTE_SCOPED_GIL_RELEASE({ nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(), quant_config_cpp, at::cuda::getCurrentCUDAStream()); diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 103f11653c..1b05658c42 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1528,7 +1528,10 @@ def make_quantizers(self) -> list: def _use_2d_quantization(idx: int) -> bool: role = self._slot_role(idx) - return role.module_type == "linear" and role.tensor_type == "weight" + return ( + role.module_type in ("linear", "grouped_linear") + and role.tensor_type == "weight" + ) return [ MXFP8Quantizer( From 2e1dfd2a98496a7906097345f35163b39ff328ae Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 10 Aug 2026 11:45:25 +0000 Subject: [PATCH 7/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../cast/mxfp8/group_quantize_mxfp8.cuh | 22 +++++++++---------- transformer_engine/pytorch/quantization.py | 5 +---- 2 files changed, 11 insertions(+), 16 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index 129292f353..a7b4d32301 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -403,8 +403,7 @@ __device__ __forceinline__ void process_rowwise_stage( using AMax2DType = std::conditional_t; __shared__ e8m0_t block_scales_2d[THREADS_X]; __shared__ AMax2DType block_amax_2d[THREADS_X * THREADS_Y]; - block_amax_2d[tid_X_rowwise * THREADS_Y + tid_Y_rowwise] = - static_cast(thread_amax); + block_amax_2d[tid_X_rowwise * THREADS_Y + tid_Y_rowwise] = static_cast(thread_amax); __syncthreads(); if (tid_Y_rowwise == 0) { AMax2DType amax_2d = static_cast(0.0f); @@ -416,14 +415,13 @@ __device__ __forceinline__ void process_rowwise_stage( amax_2d = __hmax(amax_2d, block_amax_2d[tid_X_rowwise * THREADS_Y + i]); } } - block_scales_2d[tid_X_rowwise] = ptx::float_to_e8m0( - static_cast(amax_2d) * Quantized_Limits::max_norm_rcp); + block_scales_2d[tid_X_rowwise] = + ptx::float_to_e8m0(static_cast(amax_2d) * Quantized_Limits::max_norm_rcp); } __syncthreads(); biased_exponent = block_scales_2d[tid_X_rowwise]; } else { - biased_exponent = - ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); + biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); } const size_t stage_scales_offset_Y = scales_offset_Y_rowwise + stage_offset_Y; const size_t stage_scales_offset_X = scales_offset_X_rowwise; @@ -872,8 +870,7 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations const size_t block_size = THREADS_PER_CHUNK; const bool with_gemm_swizzled_scales = output->with_gemm_swizzled_scales; - const bool use_2d_quantization = - quant_config != nullptr && quant_config->mxfp8_2d_quantization; + const bool use_2d_quantization = quant_config != nullptr && quant_config->mxfp8_2d_quantization; // Logical shape of a tensor with varying all dims is [1, M*K] if (shape_rep != ShapeRepresentation::VARYING_BOTH_DIMS) { @@ -1012,10 +1009,11 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations TRANSFORMER_ENGINE_SWITCH_CONDITION( use_2d_quantization, kIs2DBlockScaling, { - auto kernel = group_quantize_mxfp8_kernel< - IS_DBIAS, IS_DACT, IS_ACT, ParamOP, OP, IType, OType, - SCALING_TYPE, WITH_GEMM_SWIZZLED_SCALES, kIs2DBlockScaling, - SHAPE_REP>; + auto kernel = + group_quantize_mxfp8_kernel; NVTE_CHECK_CUDA(cudaFuncSetAttribute( kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 1b05658c42..1b7afad4e5 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1528,10 +1528,7 @@ def make_quantizers(self) -> list: def _use_2d_quantization(idx: int) -> bool: role = self._slot_role(idx) - return ( - role.module_type in ("linear", "grouped_linear") - and role.tensor_type == "weight" - ) + return role.module_type in ("linear", "grouped_linear") and role.tensor_type == "weight" return [ MXFP8Quantizer(