From 41e075b5001185e37e1817a3ffdac9fd7fa809fa Mon Sep 17 00:00:00 2001 From: ceo Date: Tue, 27 Oct 2026 10:56:16 +0900 Subject: [PATCH 1/2] [Common, PyTorch] Write both NVFP4 amax outputs in the reduction Plain NVFP4 quantization with row-wise and column-wise outputs reduced the amax into one buffer and then copied it into the other with a separate device-to-device copy. Let the amax kernels zero and update an optional second output, used for the column-wise amax of NVFP4 outputs, and drop the copy. Other callers of the amax reduction keep one output. Co-Authored-By: Claude Opus 5.5 Signed-off-by: ceo --- tests/pytorch/nvfp4/test_nvfp4_dual_amax.py | 40 +++++++++++++++++++ .../include/transformer_engine/recipe.h | 7 +++- .../common/recipe/current_scaling.cu | 40 +++++++++++++------ transformer_engine/pytorch/csrc/quantizer.cpp | 14 ++----- 4 files changed, 75 insertions(+), 26 deletions(-) create mode 100644 tests/pytorch/nvfp4/test_nvfp4_dual_amax.py diff --git a/tests/pytorch/nvfp4/test_nvfp4_dual_amax.py b/tests/pytorch/nvfp4/test_nvfp4_dual_amax.py new file mode 100644 index 00000000000..2404f42fe58 --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_dual_amax.py @@ -0,0 +1,40 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Check both NVFP4 amax outputs, including reuse with a decreasing maximum.""" + +import pytest +import torch + +import transformer_engine.pytorch as te +from transformer_engine.pytorch import NVFP4Quantizer + + +@pytest.mark.parametrize("direction", ["row", "column", "both"]) +@pytest.mark.parametrize("shape", [(128, 128), (160, 192)]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("with_2d_quantization", [False, True]) +def test_nvfp4_amax_matches_input_after_reuse(direction, shape, dtype, with_2d_quantization): + """The second amax must reset and track the same input maximum as the first.""" + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + available, reason = te.is_nvfp4_available(return_reason=True) + if not available: + pytest.skip(reason) + q = NVFP4Quantizer( + rowwise=direction != "column", + columnwise=direction != "row", + with_2d_quantization=with_2d_quantization, + ) + q.optimize_for_gemm = False + out = q.make_empty(shape, dtype=dtype) + # The maximum decreases, reaches zero and rises again; stale atomic-max state fails. + for magnitude in (16.0, 0.5, 0.0, 2.0): + x = torch.full(shape, -magnitude, device="cuda", dtype=dtype) + q.update_quantized(x, out) + expected = x.float().abs().amax().reshape(1) + for name in ("_amax_rowwise", "_amax_columnwise"): + actual = getattr(out, name) + if actual is not None: + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/transformer_engine/common/include/transformer_engine/recipe.h b/transformer_engine/common/include/transformer_engine/recipe.h index b3feeba6b75..4f70f153019 100644 --- a/transformer_engine/common/include/transformer_engine/recipe.h +++ b/transformer_engine/common/include/transformer_engine/recipe.h @@ -76,7 +76,9 @@ void nvte_delayed_scaling_recipe_amax_and_scale_update_after_reduction( /*! \brief Compute an FP8 tensor's amax. * * The amax (maximum absolute value) of the input tensor is computed - * and written to the amax buffer of the output tensor. + * and written to the amax buffer of the output tensor. For an NVFP4 + * output with distinct row-wise and column-wise amax buffers, the + * amax is written to both. * * \param[in] input Input tensor. Must be unquantized. * \param[in,out] output Output tensor. Must be an FP8 tensor with per-tensor scaling. @@ -88,7 +90,8 @@ void nvte_compute_amax(const NVTETensor input, NVTETensor output, cudaStream_t s * * The amax (maximum absolute value) of the input tensor is computed * and written to the amax buffer of the output tensor, using the provided - * quantization configuration. + * quantization configuration. For an NVFP4 output with distinct row-wise + * and column-wise amax buffers, the amax is written to both. * One useful config is the noop tensor, which is needed by cuda graph. * * \param[in] input Input tensor. Must be unquantized. diff --git a/transformer_engine/common/recipe/current_scaling.cu b/transformer_engine/common/recipe/current_scaling.cu index 3ba51dba754..2db3e12d222 100644 --- a/transformer_engine/common/recipe/current_scaling.cu +++ b/transformer_engine/common/recipe/current_scaling.cu @@ -21,16 +21,20 @@ namespace { constexpr int amax_kernel_threads = 512; -__launch_bounds__(1) __global__ void zero_amax_kernel(float *amax_ptr, const float *noop_ptr) { +__launch_bounds__(1) __global__ + void zero_amax_kernel(float *amax_ptr, float *amax2_ptr, const float *noop_ptr) { if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) { return; } *amax_ptr = 0; + if (amax2_ptr != nullptr) { + *amax2_ptr = 0; + } } template __launch_bounds__(amax_kernel_threads) __global__ - void amax_kernel(const InputType *input, float *amax, const size_t N, + void amax_kernel(const InputType *input, float *amax, float *amax2, const size_t N, const size_t num_aligned_elements, const float *noop_ptr) { if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) { return; @@ -66,14 +70,17 @@ __launch_bounds__(amax_kernel_threads) __global__ max = reduce_max(max, warp_id); if (threadIdx.x == 0) { atomicMaxFloat(amax, max); + if (amax2 != nullptr) { + atomicMaxFloat(amax2, max); + } } } template void launch_amax_kernel(const InputType *input, float *amax, const size_t N, const float *noop_ptr, - cudaStream_t stream) { - // Zero out amax so we can update with atomic max - zero_amax_kernel<<<1, 1, 0, stream>>>(amax, noop_ptr); + cudaStream_t stream, float *amax2 = nullptr) { + // Zero out amax (and the optional second output) so we can update with atomic max + zero_amax_kernel<<<1, 1, 0, stream>>>(amax, amax2, noop_ptr); NVTE_CHECK_CUDA(cudaGetLastError()); // Return immediately if tensor is empty @@ -94,18 +101,18 @@ void launch_amax_kernel(const InputType *input, float *amax, const size_t N, con // Launch kernel switch (align) { case Alignment::SAME_ALIGNED: - amax_kernel - <<>>(input, amax, N, num_aligned_elements, noop_ptr); + amax_kernel<<>>( + input, amax, amax2, N, num_aligned_elements, noop_ptr); break; case Alignment::SAME_UNALIGNED: - amax_kernel - <<>>(input, amax, N, num_aligned_elements, noop_ptr); + amax_kernel<<>>( + input, amax, amax2, N, num_aligned_elements, noop_ptr); break; case Alignment::DIFFERENT: { // This case is a logic error, since there is only one pointer (input) // in the alignment check. Still safe to process without vectorization. amax_kernel<1, true, InputType> - <<>>(input, amax, N, N, noop_ptr); + <<>>(input, amax, amax2, N, N, noop_ptr); break; } } @@ -170,10 +177,17 @@ void compute_amax_impl(const NVTETensor input_, const NVTETensor output_, cudaSt // Compute amax float *amax_ptr = reinterpret_cast( (output.amax.dptr != nullptr) ? output.amax.dptr : output.columnwise_amax.dptr); + // NVFP4 outputs with a distinct column-wise amax get the same value written there as well, + // which replaces a separate device-to-device copy by the caller. + float *amax2_ptr = nullptr; + if (output.scaling_mode == NVTE_NVFP4_1D_SCALING && output.columnwise_amax.dptr != nullptr && + output.columnwise_amax.dptr != amax_ptr) { + amax2_ptr = reinterpret_cast(output.columnwise_amax.dptr); + } TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( - input.data.dtype, IType, constexpr int nvec = 32 / sizeof(IType); launch_amax_kernel( - reinterpret_cast(input.data.dptr), amax_ptr, input.data.numel(), noop_ptr, - stream);); // NOLINT(*) + input.data.dtype, IType, constexpr int nvec = 32 / sizeof(IType); + launch_amax_kernel(reinterpret_cast(input.data.dptr), amax_ptr, + input.data.numel(), noop_ptr, stream, amax2_ptr);); // NOLINT(*) } void group_compute_amax_impl(const NVTEGroupedTensor input_, NVTEGroupedTensor output_, diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index ba0410bcacb..72f7b477dcf 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2730,21 +2730,13 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou void* amax_ptr = rowwise_amax_ptr != nullptr ? rowwise_amax_ptr : columnwise_amax_ptr; NVTE_CHECK(amax_ptr != nullptr, "Could not find amax pointer"); - // Compute amax of input tensor + // Compute amax of input tensor. When both the row-wise and the column-wise amax are + // allocated, the amax kernel updates both (identical values), so no device-to-device copy + // is needed afterwards. out.set_amax(amax_ptr, DType::kFloat32, std::vector{1}); NVTE_SCOPED_GIL_RELEASE( { nvte_compute_amax_with_config(input.data(), out.data(), quant_config, stream); }); out.set_amax(rowwise_amax_ptr, DType::kFloat32, std::vector{1}); - - // Make sure row-wise and column-wise amaxes match - if (rowwise_amax_ptr != amax_ptr && rowwise_amax_ptr != nullptr) { - NVTE_CHECK_CUDA(cudaMemcpyAsync(rowwise_amax_ptr, amax_ptr, sizeof(float), - cudaMemcpyDeviceToDevice, stream)); - } - if (columnwise_amax_ptr != amax_ptr && columnwise_amax_ptr != nullptr) { - NVTE_CHECK_CUDA(cudaMemcpyAsync(columnwise_amax_ptr, amax_ptr, sizeof(float), - cudaMemcpyDeviceToDevice, stream)); - } } } From 73cc88c3af99322414e21282f4744e3f0da44c32 Mon Sep 17 00:00:00 2001 From: ceo Date: Tue, 10 Nov 2026 07:16:18 +0900 Subject: [PATCH 2/2] [Common, PyTorch] Address dual amax review feedback Name the zeroing kernel outputs rowwise and colwise. Fold decreasing-amax output reuse checks into the existing NVFP4 quantization tests, including columnwise-only 1D and 2D coverage. Signed-off-by: ceo --- tests/pytorch/nvfp4/test_nvfp4_dual_amax.py | 40 ---------- .../nvfp4/test_nvfp4_quantize_exact.py | 77 +++++++++++++++---- .../common/recipe/current_scaling.cu | 8 +- 3 files changed, 65 insertions(+), 60 deletions(-) delete mode 100644 tests/pytorch/nvfp4/test_nvfp4_dual_amax.py diff --git a/tests/pytorch/nvfp4/test_nvfp4_dual_amax.py b/tests/pytorch/nvfp4/test_nvfp4_dual_amax.py deleted file mode 100644 index 2404f42fe58..00000000000 --- a/tests/pytorch/nvfp4/test_nvfp4_dual_amax.py +++ /dev/null @@ -1,40 +0,0 @@ -# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# -# See LICENSE for license information. - -"""Check both NVFP4 amax outputs, including reuse with a decreasing maximum.""" - -import pytest -import torch - -import transformer_engine.pytorch as te -from transformer_engine.pytorch import NVFP4Quantizer - - -@pytest.mark.parametrize("direction", ["row", "column", "both"]) -@pytest.mark.parametrize("shape", [(128, 128), (160, 192)]) -@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) -@pytest.mark.parametrize("with_2d_quantization", [False, True]) -def test_nvfp4_amax_matches_input_after_reuse(direction, shape, dtype, with_2d_quantization): - """The second amax must reset and track the same input maximum as the first.""" - if not torch.cuda.is_available(): - pytest.skip("CUDA is required") - available, reason = te.is_nvfp4_available(return_reason=True) - if not available: - pytest.skip(reason) - q = NVFP4Quantizer( - rowwise=direction != "column", - columnwise=direction != "row", - with_2d_quantization=with_2d_quantization, - ) - q.optimize_for_gemm = False - out = q.make_empty(shape, dtype=dtype) - # The maximum decreases, reaches zero and rises again; stale atomic-max state fails. - for magnitude in (16.0, 0.5, 0.0, 2.0): - x = torch.full(shape, -magnitude, device="cuda", dtype=dtype) - q.update_quantized(x, out) - expected = x.float().abs().amax().reshape(1) - for name in ("_amax_rowwise", "_amax_columnwise"): - actual = getattr(out, name) - if actual is not None: - torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 51b5b46c158..1a5015b6762 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -123,6 +123,7 @@ def check_quantization_nvfp4_versus_reference( nvfp4_4over6_err_mode: str = "MAE", nvfp4_4over6_err_use_fast_math: bool = False, scale_dtype: te.DType = te.DType.kFloat8E4M3, + reuse_output: bool = False, ) -> None: scale_dtype = te.DType.cast(scale_dtype) if nvfp4_e4m3_max not in NVFP4_SUPPORTED_SCALE_MAX_VALUES: @@ -252,6 +253,20 @@ def check_quantization_nvfp4_versus_reference( torch.testing.assert_close(qx_amax, ref_amax, atol=0.0, rtol=0.0) + if reuse_output: + assert not use_cpp_allocator + assert not row_scaled_nvfp4 + # Reusing an output whose maximum decreases verifies that every amax buffer is reset. + for magnitude in (16.0, 0.5, 0.0, 2.0): + x = torch.full((M, N), -magnitude, dtype=x_dtype, device=device) + nvfp4_quantizer.update_quantized(x, x_nvfp4_sut) + expected_amax = x.float().abs().amax().reshape(1) + torch.testing.assert_close(x_nvfp4_sut._amax_rowwise, expected_amax, atol=0.0, rtol=0.0) + if return_transpose: + torch.testing.assert_close( + x_nvfp4_sut._amax_columnwise, expected_amax, atol=0.0, rtol=0.0 + ) + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize("return_transpose", [False, True], ids=["rowwise", "with_columnwise"]) @@ -359,6 +374,12 @@ def test_quantization_block_tiling_versus_reference( row_scaled_nvfp4: bool, nvfp4_4over6_config: NVFP44Over6TestConfig, ) -> None: + reuse_output = ( + not use_cpp_allocator + and not row_scaled_nvfp4 + and not nvfp4_4over6_config.use_4over6 + and (M, N) in ((128, 128), (256, 272)) + ) check_quantization_nvfp4_versus_reference( x_dtype=x_dtype, M=M, @@ -372,6 +393,7 @@ def test_quantization_block_tiling_versus_reference( nvfp4_e4m3_max=nvfp4_4over6_config.e4m3_max, nvfp4_4over6_err_mode=nvfp4_4over6_config.err_mode, nvfp4_4over6_err_use_fast_math=nvfp4_4over6_config.err_use_fast_math, + reuse_output=reuse_output, ) @@ -807,29 +829,42 @@ def test_nvfp4_quantization_noncontiguous_inputs( @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( - "M, N", + "with_2d_quantization, M, N", [ - # Aligned tiles - (128, 128), - (256, 256), - (512, 512), - (2048, 2048), - # Padded tiles (non-multiple of kTileDim=128) - (256, 272), - (304, 304), - (320, 256), + # Existing 2D coverage: aligned and padded tiles. + (True, 128, 128), + (True, 256, 256), + (True, 512, 512), + (True, 2048, 2048), + (True, 256, 272), + (True, 304, 304), + (True, 320, 256), + # Focused 1D coverage for aligned and padded output reuse. + (False, 128, 128), + (False, 256, 272), + ], + ids=[ + "2d-128x128", + "2d-256x256", + "2d-512x512", + "2d-2048x2048", + "2d-256x272", + "2d-304x304", + "2d-320x256", + "1d-128x128", + "1d-256x272", ], ) @pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) -def test_nvfp4_2d_columnwise_only_matches_both_directions( +def test_nvfp4_columnwise_only_matches_both_directions( x_dtype: torch.dtype, + with_2d_quantization: bool, M: int, N: int, ): - """Bitwise check: 2D NVFP4 with columnwise-only must produce the same - columnwise data/scales as the columnwise half of (rowwise + columnwise) 2D. + """Columnwise-only NVFP4 must match the columnwise half of both-directions. - Covers both kernels depending on the (dtype, shape) routing: + The 2D cases cover both kernels depending on the (dtype, shape) routing: - bf16 with rows % 32 == 0 and cols % 32 == 0 routes to the optimized ``quantize_transpose_nvfp4_2D_kernel`` (instantiated with RETURN_ROWWISE=false), validating that gating the rowwise pass/store leaves the shared @@ -853,7 +888,7 @@ def _make_quantizer(*, rowwise: bool, columnwise: bool) -> NVFP4Quantizer: amax_reduction_group=None, with_rht=False, with_post_rht_amax=False, - with_2d_quantization=True, + with_2d_quantization=with_2d_quantization, row_scaled_nvfp4=False, ) @@ -861,7 +896,7 @@ def _make_quantizer(*, rowwise: bool, columnwise: bool) -> NVFP4Quantizer: q_both = _make_quantizer(rowwise=True, columnwise=True) out_both = q_both(x) - # SUT: produce columnwise only (the path that hits the new amax-only pass). + # SUT: produce columnwise only. q_col_only = _make_quantizer(rowwise=False, columnwise=True) out_col_only = q_col_only(x) @@ -901,5 +936,15 @@ def _make_quantizer(*, rowwise: bool, columnwise: bool) -> NVFP4Quantizer: out_col_only._amax_columnwise, out_both._amax_columnwise, atol=0, rtol=0 ) + if (M, N) in ((128, 128), (256, 272)): + # Reusing a columnwise-only output exercises the path where its amax is primary. + for magnitude in (16.0, 0.5, 0.0, 2.0): + x = torch.full((M, N), -magnitude, dtype=x_dtype, device=device) + q_col_only.update_quantized(x, out_col_only) + expected_amax = x.float().abs().amax().reshape(1) + torch.testing.assert_close( + out_col_only._amax_columnwise, expected_amax, atol=0.0, rtol=0.0 + ) + # Sanity: column-only path must not allocate a rowwise output. assert out_col_only._rowwise_data is None diff --git a/transformer_engine/common/recipe/current_scaling.cu b/transformer_engine/common/recipe/current_scaling.cu index 2db3e12d222..8bc7e166e6e 100644 --- a/transformer_engine/common/recipe/current_scaling.cu +++ b/transformer_engine/common/recipe/current_scaling.cu @@ -22,13 +22,13 @@ namespace { constexpr int amax_kernel_threads = 512; __launch_bounds__(1) __global__ - void zero_amax_kernel(float *amax_ptr, float *amax2_ptr, const float *noop_ptr) { + void zero_amax_kernel(float *rowwise_amax_ptr, float *colwise_amax_ptr, const float *noop_ptr) { if (noop_ptr != nullptr && noop_ptr[0] == 1.0f) { return; } - *amax_ptr = 0; - if (amax2_ptr != nullptr) { - *amax2_ptr = 0; + *rowwise_amax_ptr = 0; + if (colwise_amax_ptr != nullptr) { + *colwise_amax_ptr = 0; } }