From f40be70f9f4ecc828588b9a257cb1390d4d4d52e Mon Sep 17 00:00:00 2001 From: tony <864832769@qq.com> Date: Mon, 3 Aug 2026 22:01:11 +0800 Subject: [PATCH 01/12] fix(nvfp4): match RHT reference on SM120 and SM121 Signed-off-by: tony <864832769@qq.com> --- transformer_engine/pytorch/csrc/quantizer.cpp | 31 ++++++++++++++++--- 1 file changed, 26 insertions(+), 5 deletions(-) diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 4fff3f92de..1b79a8e60d 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2465,11 +2465,32 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper( out_columnwise_amax.shape); // Invoking fallback RHT kernel unfused. - NVTE_SCOPED_GIL_RELEASE({ - // Perform the RHT(input.t), and write to rht_output_cpp.columnwise. - nvte_hadamard_transform(input.data(), rht_output_t_cpp.data(), 0, - this->rht_matrix_random_sign_mask_t, stream); - }); + const int sm_arch = transformer_engine::cuda::sm_arch(); + if (sm_arch == 120 || sm_arch == 121) { + // Match the PyTorch reference arithmetic on GeForce Blackwell. The BF16 MMA + // Hadamard kernel uses a different accumulation order and can differ by one + // BF16 ULP, which is enough to change exact FP4 rounding at ties. + NVTE_CHECK(this->rht_matrix.defined() && this->rht_matrix.numel() > 0, + "RHT matrix is not available."); + const auto [rows, cols] = get_2d_dims(input.shape()); + const auto torch_dtype = GetATenDType(input.dtype()); + auto options = at::TensorOptions().dtype(torch_dtype).device(torch::kCUDA); + auto input_torch = at::from_blob( + input.get_rowwise_data().data_ptr, + {static_cast(rows), static_cast(cols)}, [](void*) {}, options); + auto output_torch = at::from_blob( + rht_output_t_cpp.get_rowwise_data().data_ptr, + {static_cast(cols), static_cast(rows)}, [](void*) {}, options); + auto rht_matrix = this->rht_matrix.to(torch_dtype); + at::matmul_out(output_torch.view({-1, 16}), + input_torch.transpose(0, 1).contiguous().view({-1, 16}), rht_matrix); + } else { + NVTE_SCOPED_GIL_RELEASE({ + // Perform the RHT(input.t), and write to rht_output_cpp.columnwise. + nvte_hadamard_transform(input.data(), rht_output_t_cpp.data(), 0, + this->rht_matrix_random_sign_mask_t, stream); + }); + } // Quantize kernel will treat everything as rowwise input/output, which is // intended. From caa2476cb8269a88fa4ac65bdd1869658e56b740 Mon Sep 17 00:00:00 2001 From: tony <864832769@qq.com> Date: Mon, 3 Aug 2026 22:12:36 +0800 Subject: [PATCH 02/12] fix(nvfp4): recompute post-RHT amax from fallback output Signed-off-by: tony <864832769@qq.com> --- transformer_engine/pytorch/csrc/common.h | 2 +- transformer_engine/pytorch/csrc/quantizer.cpp | 10 ++++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index aa0e0c87fe..79a7c47f63 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -421,7 +421,7 @@ class NVFP4Quantizer : public Quantizer { TensorWrapper& rht_output_t_cpp, QuantizationConfigWrapper& quant_config, QuantizationConfigWrapper& quant_config_columnwise, - cudaStream_t stream); + bool compute_amax, cudaStream_t stream); }; std::unique_ptr convert_quantizer(py::handle quantizer); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 1b79a8e60d..a6b80f35dd 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2409,7 +2409,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( void NVFP4Quantizer::quantize_with_rht_unfused_helper( const TensorWrapper& input, TensorWrapper& out, TensorWrapper& rht_output_t_cpp, QuantizationConfigWrapper& quant_config, QuantizationConfigWrapper& quant_config_columnwise, - cudaStream_t stream) { + bool compute_amax, cudaStream_t stream) { // The kernels invoked below reject swizzled-SF output, so trip a clear // error here before reaching them. NVTE_CHECK(!out.get_with_gemm_swizzled_scales(), @@ -2484,6 +2484,12 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper( auto rht_matrix = this->rht_matrix.to(torch_dtype); at::matmul_out(output_torch.view({-1, 16}), input_torch.transpose(0, 1).contiguous().view({-1, 16}), rht_matrix); + if (compute_amax) { + auto amax_options = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + auto columnwise_amax = at::from_blob(out_columnwise_amax.data_ptr, {1}, [](void*) {}, + amax_options); + columnwise_amax.copy_(output_torch.abs().amax().to(torch::kFloat32).view({1})); + } } else { NVTE_SCOPED_GIL_RELEASE({ // Perform the RHT(input.t), and write to rht_output_cpp.columnwise. @@ -2722,7 +2728,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou rht_output_t_cpp.set_rowwise_data(rht_output_t.data_ptr(), input.dtype(), std::vector{cols, rows}); this->quantize_with_rht_unfused_helper(input, out, rht_output_t_cpp, quant_config, - columnwise_quant_config_to_use, stream); + columnwise_quant_config_to_use, compute_amax, stream); } } else { NVTE_SCOPED_GIL_RELEASE({ nvte_quantize_v2(input.data(), out.data(), quant_config, stream); }); From 1381291cbbf07f97dd19e255a203ed5fd554a447 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 3 Aug 2026 14:13:38 +0000 Subject: [PATCH 03/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/csrc/quantizer.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index a6b80f35dd..e2d5731500 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2486,8 +2486,8 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper( input_torch.transpose(0, 1).contiguous().view({-1, 16}), rht_matrix); if (compute_amax) { auto amax_options = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); - auto columnwise_amax = at::from_blob(out_columnwise_amax.data_ptr, {1}, [](void*) {}, - amax_options); + auto columnwise_amax = + at::from_blob(out_columnwise_amax.data_ptr, {1}, [](void*) {}, amax_options); columnwise_amax.copy_(output_torch.abs().amax().to(torch::kFloat32).view({1})); } } else { From 9142e86686ad98157d9352488527696012877506 Mon Sep 17 00:00:00 2001 From: tony <864832769@qq.com> Date: Mon, 3 Aug 2026 22:44:37 +0800 Subject: [PATCH 04/12] fix(nvfp4): reduce ATen RHT amax across ranks Signed-off-by: tony <864832769@qq.com> --- transformer_engine/pytorch/csrc/common.h | 2 +- transformer_engine/pytorch/csrc/quantizer.cpp | 43 +++++++++++++------ 2 files changed, 32 insertions(+), 13 deletions(-) diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 79a7c47f63..aa0e0c87fe 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -421,7 +421,7 @@ class NVFP4Quantizer : public Quantizer { TensorWrapper& rht_output_t_cpp, QuantizationConfigWrapper& quant_config, QuantizationConfigWrapper& quant_config_columnwise, - bool compute_amax, cudaStream_t stream); + cudaStream_t stream); }; std::unique_ptr convert_quantizer(py::handle quantizer); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index e2d5731500..0a579faf0c 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2409,7 +2409,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( void NVFP4Quantizer::quantize_with_rht_unfused_helper( const TensorWrapper& input, TensorWrapper& out, TensorWrapper& rht_output_t_cpp, QuantizationConfigWrapper& quant_config, QuantizationConfigWrapper& quant_config_columnwise, - bool compute_amax, cudaStream_t stream) { + cudaStream_t stream) { // The kernels invoked below reject swizzled-SF output, so trip a clear // error here before reaching them. NVTE_CHECK(!out.get_with_gemm_swizzled_scales(), @@ -2484,12 +2484,6 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper( auto rht_matrix = this->rht_matrix.to(torch_dtype); at::matmul_out(output_torch.view({-1, 16}), input_torch.transpose(0, 1).contiguous().view({-1, 16}), rht_matrix); - if (compute_amax) { - auto amax_options = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); - auto columnwise_amax = - at::from_blob(out_columnwise_amax.data_ptr, {1}, [](void*) {}, amax_options); - columnwise_amax.copy_(output_torch.abs().amax().to(torch::kFloat32).view({1})); - } } else { NVTE_SCOPED_GIL_RELEASE({ // Perform the RHT(input.t), and write to rht_output_cpp.columnwise. @@ -2639,10 +2633,35 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou // 1. Rowwise amax = amax for input // 2. Columnwise amax = amax for RHT(input.t) if (compute_amax) { - NVTE_SCOPED_GIL_RELEASE({ - nvte_hadamard_transform_amax(input.data(), out.data(), 0, - this->rht_matrix_random_sign_mask_t, stream); - }); + const int sm_arch = transformer_engine::cuda::sm_arch(); + if (sm_arch == 120 || sm_arch == 121) { + const auto torch_dtype = GetATenDType(input.dtype()); + auto options = at::TensorOptions().dtype(torch_dtype).device(torch::kCUDA); + auto input_torch = at::from_blob( + input.get_rowwise_data().data_ptr, + {static_cast(rows), static_cast(cols)}, [](void*) {}, options); + auto amax_options = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + auto copy_amax = [&amax_options](const at::Tensor& tensor, void* dst) { + if (dst != nullptr) { + auto amax = at::from_blob(dst, {1}, [](void*) {}, amax_options); + amax.copy_(tensor.abs().amax().to(torch::kFloat32).view({1})); + } + }; + copy_amax(input_torch, out.get_amax().data_ptr); + if (out.get_columnwise_amax().data_ptr != nullptr) { + NVTE_CHECK(this->rht_matrix.defined() && this->rht_matrix.numel() > 0, + "RHT matrix is not available."); + auto rht_matrix = this->rht_matrix.to(torch_dtype); + auto rht_output = at::matmul( + input_torch.transpose(0, 1).contiguous().view({-1, 16}), rht_matrix); + copy_amax(rht_output, out.get_columnwise_amax().data_ptr); + } + } else { + NVTE_SCOPED_GIL_RELEASE({ + nvte_hadamard_transform_amax(input.data(), out.data(), 0, + this->rht_matrix_random_sign_mask_t, stream); + }); + } } } else { // raise error since it's not supported yet @@ -2728,7 +2747,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou rht_output_t_cpp.set_rowwise_data(rht_output_t.data_ptr(), input.dtype(), std::vector{cols, rows}); this->quantize_with_rht_unfused_helper(input, out, rht_output_t_cpp, quant_config, - columnwise_quant_config_to_use, compute_amax, stream); + columnwise_quant_config_to_use, stream); } } else { NVTE_SCOPED_GIL_RELEASE({ nvte_quantize_v2(input.data(), out.data(), quant_config, stream); }); From 70a5da7d7803a9aeb3e15fb83af6a1865e160ecb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 3 Aug 2026 14:46:12 +0000 Subject: [PATCH 05/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/pytorch/csrc/quantizer.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 0a579faf0c..cc472c5575 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2652,8 +2652,8 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou NVTE_CHECK(this->rht_matrix.defined() && this->rht_matrix.numel() > 0, "RHT matrix is not available."); auto rht_matrix = this->rht_matrix.to(torch_dtype); - auto rht_output = at::matmul( - input_torch.transpose(0, 1).contiguous().view({-1, 16}), rht_matrix); + auto rht_output = + at::matmul(input_torch.transpose(0, 1).contiguous().view({-1, 16}), rht_matrix); copy_amax(rht_output, out.get_columnwise_amax().data_ptr); } } else { From ea87075716fc8a795a1e00364ebc237e8f7694b6 Mon Sep 17 00:00:00 2001 From: tony <864832769@qq.com> Date: Mon, 3 Aug 2026 23:19:31 +0800 Subject: [PATCH 06/12] test(nvfp4): cover SM12x RHT amax semantics Signed-off-by: tony <864832769@qq.com> --- tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py | 110 ++++++++++++++++++++ 1 file changed, 110 insertions(+) create mode 100644 tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py b/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py new file mode 100644 index 0000000000..0ad0d50e3f --- /dev/null +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py @@ -0,0 +1,110 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Regression tests for the unfused NVFP4 RHT path on SM120/SM121.""" + +import os + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +import transformer_engine.pytorch as te +from transformer_engine.pytorch import NVFP4Quantizer +from transformer_engine.pytorch.custom_recipes import utils +from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef + + +recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) + + +def _is_sm12x(device: int = 0) -> bool: + return torch.cuda.get_device_capability(device) in ((12, 0), (12, 1)) + + +def _reference_post_rht_amax(x: torch.Tensor, with_random_sign_mask: bool) -> torch.Tensor: + quantizer = NVFP4QuantizerRef( + dtype=utils.Fp4Formats.E2M1, + rowwise=False, + columnwise=True, + pow_2_scales=False, + eps=0.0, + quant_tile_shape=(1, 16), + with_rht=True, + with_random_sign_mask=with_random_sign_mask, + ) + transformed = quantizer._apply_rht(x.t().contiguous()) + return transformed.abs().amax().to(torch.float32).view(1) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("with_random_sign_mask", [False, True]) +def test_sm12x_post_rht_amax_matches_aten_output(with_random_sign_mask: bool) -> None: + """The stored post-RHT amax must describe the ATen-generated RHT tensor.""" + + if not _is_sm12x(): + pytest.skip("Test targets the SM120/SM121 ATen RHT fallback") + + torch.manual_seed(1234) + x = torch.randn((128, 128), device="cuda", dtype=torch.bfloat16) + expected_amax = _reference_post_rht_amax(x, with_random_sign_mask) + + quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=False, + columnwise=True, + with_amax_reduction=False, + with_rht=True, + with_post_rht_amax=True, + with_random_sign_mask=with_random_sign_mask, + ) + out = quantizer(x) + + torch.testing.assert_close(out._amax_columnwise, expected_amax, atol=0.0, rtol=0.0) + + +def _distributed_amax_worker(rank: int, world_size: int, init_file: str) -> None: + torch.cuda.set_device(rank) + dist.init_process_group( + backend="nccl", + init_method=f"file://{init_file}", + rank=rank, + world_size=world_size, + ) + try: + torch.manual_seed(4321 + rank) + x = torch.randn((128, 128), device=f"cuda:{rank}", dtype=torch.bfloat16) + x.mul_(rank + 1) + expected_amax = _reference_post_rht_amax(x, with_random_sign_mask=True) + + quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=False, + columnwise=True, + with_amax_reduction=True, + amax_reduction_group=dist.group.WORLD, + with_rht=True, + with_post_rht_amax=True, + with_random_sign_mask=True, + ) + out = quantizer(x) + + dist.all_reduce(expected_amax, op=dist.ReduceOp.MAX) + torch.testing.assert_close(out._amax_columnwise, expected_amax, atol=0.0, rtol=0.0) + finally: + dist.destroy_process_group() + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="Two CUDA devices are required") +def test_sm12x_post_rht_amax_reduction_is_global(tmp_path) -> None: + """ATen post-RHT amax must be computed before the distributed MAX reduction.""" + + if not all(_is_sm12x(device) for device in range(2)): + pytest.skip("Test targets the SM120/SM121 ATen RHT fallback") + + init_file = os.fspath(tmp_path / "nvfp4_rht_amax_init") + mp.spawn(_distributed_amax_worker, args=(2, init_file), nprocs=2, join=True) From 7a0a0d09dc597eed00957e15977f57f74f6df4ab Mon Sep 17 00:00:00 2001 From: tony <864832769@qq.com> Date: Tue, 4 Aug 2026 00:03:11 +0800 Subject: [PATCH 07/12] perf(nvfp4): reuse SM12x ATen RHT output Signed-off-by: tony <864832769@qq.com> --- transformer_engine/pytorch/csrc/common.h | 2 +- transformer_engine/pytorch/csrc/quantizer.cpp | 50 ++++++++++++------- 2 files changed, 33 insertions(+), 19 deletions(-) diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index aa0e0c87fe..9f4da7e1e0 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -421,7 +421,7 @@ class NVFP4Quantizer : public Quantizer { TensorWrapper& rht_output_t_cpp, QuantizationConfigWrapper& quant_config, QuantizationConfigWrapper& quant_config_columnwise, - cudaStream_t stream); + cudaStream_t stream, bool rht_output_is_ready); }; std::unique_ptr convert_quantizer(py::handle quantizer); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index cc472c5575..9512038953 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -2409,7 +2409,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( void NVFP4Quantizer::quantize_with_rht_unfused_helper( const TensorWrapper& input, TensorWrapper& out, TensorWrapper& rht_output_t_cpp, QuantizationConfigWrapper& quant_config, QuantizationConfigWrapper& quant_config_columnwise, - cudaStream_t stream) { + cudaStream_t stream, bool rht_output_is_ready) { // The kernels invoked below reject swizzled-SF output, so trip a clear // error here before reaching them. NVTE_CHECK(!out.get_with_gemm_swizzled_scales(), @@ -2464,9 +2464,13 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper( static_cast(out_columnwise_amax.dtype), out_columnwise_amax.shape); - // Invoking fallback RHT kernel unfused. + // Invoking fallback RHT kernel unfused unless the SM120/SM121 ATen result was already + // materialized while computing post-RHT amax. const int sm_arch = transformer_engine::cuda::sm_arch(); - if (sm_arch == 120 || sm_arch == 121) { + if (rht_output_is_ready) { + NVTE_CHECK(sm_arch == 120 || sm_arch == 121, + "A precomputed unfused RHT output is only supported on SM120/SM121."); + } else if (sm_arch == 120 || sm_arch == 121) { // Match the PyTorch reference arithmetic on GeForce Blackwell. The BF16 MMA // Hadamard kernel uses a different accumulation order and can differ by one // BF16 ULP, which is enough to change exact FP4 rounding at ties. @@ -2482,8 +2486,9 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper( rht_output_t_cpp.get_rowwise_data().data_ptr, {static_cast(cols), static_cast(rows)}, [](void*) {}, options); auto rht_matrix = this->rht_matrix.to(torch_dtype); - at::matmul_out(output_torch.view({-1, 16}), - input_torch.transpose(0, 1).contiguous().view({-1, 16}), rht_matrix); + auto output_view = output_torch.view({-1, 16}); + auto input_view = input_torch.transpose(0, 1).contiguous().view({-1, 16}); + at::matmul_out(output_view, input_view, rht_matrix); } else { NVTE_SCOPED_GIL_RELEASE({ // Perform the RHT(input.t), and write to rht_output_cpp.columnwise. @@ -2581,6 +2586,19 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou input.dtype() == DType::kBFloat16 && NVFP4Quantizer::is_eligible_for_rht_cast_fusion(convertShape(input.shape())); + // The unfused path needs an intermediate in transposed layout. Allocate it before amax + // computation so SM120/SM121 can materialize the ATen RHT once, derive post-RHT amax from + // that exact buffer, and reuse it for quantization after any distributed amax reduction. + at::Tensor rht_output_t; + TensorWrapper rht_output_t_cpp; + bool rht_output_is_ready = false; + if (this->with_rht && !eligible_for_rht_cast_fusion) { + rht_output_t = + allocateTorchTensor(static_cast(cols), static_cast(rows), input.dtype()); + rht_output_t_cpp.set_rowwise_data(rht_output_t.data_ptr(), input.dtype(), + std::vector{cols, rows}); + } + // Stochastic rounding // When both rowwise and columnwise quantization are used with RHT, // we need separate RNG states for each to ensure they use different random numbers. @@ -2652,9 +2670,14 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou NVTE_CHECK(this->rht_matrix.defined() && this->rht_matrix.numel() > 0, "RHT matrix is not available."); auto rht_matrix = this->rht_matrix.to(torch_dtype); - auto rht_output = - at::matmul(input_torch.transpose(0, 1).contiguous().view({-1, 16}), rht_matrix); + auto rht_output = at::from_blob( + rht_output_t_cpp.get_rowwise_data().data_ptr, + {static_cast(cols), static_cast(rows)}, [](void*) {}, options); + auto rht_output_view = rht_output.view({-1, 16}); + auto input_view = input_torch.transpose(0, 1).contiguous().view({-1, 16}); + at::matmul_out(rht_output_view, input_view, rht_matrix); copy_amax(rht_output, out.get_columnwise_amax().data_ptr); + rht_output_is_ready = true; } } else { NVTE_SCOPED_GIL_RELEASE({ @@ -2736,18 +2759,9 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou // are separate kernel launches auto& columnwise_quant_config_to_use = need_separate_columnwise_rng ? quant_config_columnwise : quant_config; - // unfused path also needs memory allocation for intermediate buffer for RHT output - at::Tensor rht_output_t; // The RHT(x_t) output, in columnwise layout - // This wrapper is going to be passed as input to the quantization kernel. - TensorWrapper rht_output_t_cpp; // Wrapper to contain the RHT(x) and RHT(x_t) outputs - rht_output_t = - allocateTorchTensor(static_cast(cols), static_cast(rows), input.dtype()); - // NOTE (frsun): This is non-intuitive, we are writing the - // result of transposed RHT to the output of rowwise. - rht_output_t_cpp.set_rowwise_data(rht_output_t.data_ptr(), input.dtype(), - std::vector{cols, rows}); this->quantize_with_rht_unfused_helper(input, out, rht_output_t_cpp, quant_config, - columnwise_quant_config_to_use, stream); + columnwise_quant_config_to_use, stream, + rht_output_is_ready); } } else { NVTE_SCOPED_GIL_RELEASE({ nvte_quantize_v2(input.data(), out.data(), quant_config, stream); }); From 434f24295688b5d4fd944e9a641ada956929cbf7 Mon Sep 17 00:00:00 2001 From: tony <864832769@qq.com> Date: Tue, 4 Aug 2026 03:52:27 +0800 Subject: [PATCH 08/12] feat(nvfp4): fuse SM12x RHT quantization with TMA Signed-off-by: tony <864832769@qq.com> --- tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py | 67 ++++++- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 185 ++++++++++++++---- ...cast_col_hadamard_transform_cast_fusion.cu | 18 ++ transformer_engine/pytorch/csrc/quantizer.cpp | 30 ++- 4 files changed, 257 insertions(+), 43 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py b/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py index 0ad0d50e3f..15404f4737 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py @@ -2,7 +2,7 @@ # # See LICENSE for license information. -"""Regression tests for the unfused NVFP4 RHT path on SM120/SM121.""" +"""Regression tests for the no-TMEM fused NVFP4 RHT path on SM120/SM121.""" import os @@ -43,7 +43,7 @@ def _reference_post_rht_amax(x: torch.Tensor, with_random_sign_mask: bool) -> to @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @pytest.mark.parametrize("with_random_sign_mask", [False, True]) def test_sm12x_post_rht_amax_matches_aten_output(with_random_sign_mask: bool) -> None: - """The stored post-RHT amax must describe the ATen-generated RHT tensor.""" + """The fused post-RHT amax must describe the ATen-equivalent RHT tensor.""" if not _is_sm12x(): pytest.skip("Test targets the SM120/SM121 ATen RHT fallback") @@ -66,6 +66,69 @@ def test_sm12x_post_rht_amax_matches_aten_output(with_random_sign_mask: bool) -> torch.testing.assert_close(out._amax_columnwise, expected_amax, atol=0.0, rtol=0.0) +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +@pytest.mark.parametrize("shape", [(128, 128), (256, 256)]) +@pytest.mark.parametrize("rowwise", [False, True]) +@pytest.mark.parametrize("with_random_sign_mask", [False, True]) +def test_sm12x_fused_rht_codes_and_scales_match_aten_pipeline( + shape: tuple[int, int], rowwise: bool, with_random_sign_mask: bool +) -> None: + """Fused RHT codes/scales must match ATen RHT followed by the tuned quantizer.""" + + if not _is_sm12x(): + pytest.skip("Test targets the SM120/SM121 no-TMEM fused RHT path") + + torch.manual_seed(2026) + x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + fused_quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=rowwise, + columnwise=True, + with_amax_reduction=False, + with_rht=True, + with_post_rht_amax=True, + with_random_sign_mask=with_random_sign_mask, + ) + fused = fused_quantizer(x) + + reference_rht = NVFP4QuantizerRef( + dtype=utils.Fp4Formats.E2M1, + rowwise=False, + columnwise=True, + pow_2_scales=False, + eps=0.0, + quant_tile_shape=(1, 16), + with_rht=True, + with_random_sign_mask=with_random_sign_mask, + ) + transformed = reference_rht._apply_rht(x.t().contiguous()) + plain_quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=True, + columnwise=False, + with_amax_reduction=False, + with_rht=False, + ) + expected_columnwise = plain_quantizer(transformed) + + torch.testing.assert_close( + fused._columnwise_data.view(torch.uint8), + expected_columnwise._rowwise_data.view(torch.uint8), + atol=0.0, + rtol=0.0, + ) + torch.testing.assert_close( + fused._columnwise_scale_inv, + expected_columnwise._rowwise_scale_inv, + atol=0.0, + rtol=0.0, + ) + torch.testing.assert_close( + fused._amax_columnwise, expected_columnwise._amax_rowwise, atol=0.0, rtol=0.0 + ) + + def _distributed_amax_worker(rank: int, world_size: int, init_file: str) -> None: torch.cuda.set_device(rank) dist.init_process_group( diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index a38a620ebe..f240b3b112 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -317,11 +317,13 @@ constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM; // 8 = 128 / template + bool ROW_SCALED_NVFP4, bool RETURN_ROWWISE = true, + bool APPLY_COLUMNWISE_RHT = false> __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, const __grid_constant__ CUtensorMap tensor_map_output_t, + const __grid_constant__ CUtensorMap tensor_map_rht, nvfp4_scale_t *const scales_ptr, nvfp4_scale_t *const scales_t_ptr, const float *noop, const float *const amax_rowwise_ptr, @@ -408,6 +410,8 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_data = buff_size_aligned_out; constexpr size_t out_mem_colwise_data = buff_size_aligned_out; constexpr size_t out_mem_rowwise_scales = 0; + constexpr size_t out_mem_colwise_scales = + (CHUNK_DIM_Y * CHUNK_DIM_X) / SCALE_DIM * sizeof(nvfp4_scale_t); extern __shared__ char dynamic_shmem[]; uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); @@ -425,6 +429,9 @@ __global__ void __launch_bounds__(THREADS_NUM) dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); nvfp4_scale_t *out_colwise_scales_sh = reinterpret_cast( dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); + IType *rht_sh = reinterpret_cast( + dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + + out_mem_rowwise_scales + out_mem_colwise_scales); IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; @@ -448,9 +455,16 @@ __global__ void __launch_bounds__(THREADS_NUM) // Initialize shared memory barrier with the number of threads participating in the barrier. #pragma nv_diag_suppress static_var_with_dynamic_init __shared__ alignas(8) uint64_t mbar[STAGES]; + __shared__ alignas(8) uint64_t mbar_rht[1]; initialize_barriers(mbar, is_master_thread); + if constexpr (APPLY_COLUMNWISE_RHT) { + initialize_barriers<1, THREADS_NUM>(mbar_rht, is_master_thread); + copy_2d_to_shared(rht_sh, &tensor_map_rht, 0, 0, 16 * 16 * sizeof(IType), &mbar_rht[0], + is_master_thread); + } + copy_2d_to_shared(&in_sh[0], &tensor_map_input, block_offset_X, block_offset_Y, shmem_buff_size, &mbar[0], is_master_thread); @@ -483,6 +497,9 @@ __global__ void __launch_bounds__(THREADS_NUM) // Wait for the data to have arrived ptx::mbarrier_wait_parity(&mbar[stage], 0); + if constexpr (APPLY_COLUMNWISE_RHT) { + ptx::mbarrier_wait_parity(&mbar_rht[0], 0); + } float block_amax = 0.0f; @@ -544,6 +561,47 @@ __global__ void __launch_bounds__(THREADS_NUM) in_compute_colwise[i] = elt; } } + if constexpr (APPLY_COLUMNWISE_RHT) { + float input_vec[SCALE_DIM]; +#pragma unroll + for (int i = 0; i < SCALE_DIM; ++i) { + input_vec[i] = static_cast(in_colwise_IType[i]); + } + block_amax = 0.0f; +#pragma unroll + for (int j = 0; j < SCALE_DIM; ++j) { + float value = 0.0f; + // ATen selects a reduced-precision BF16 GEMM for at most 1024 + // K=16 vectors on SM12x. It forms two K=8 FP32 partials, rounds each + // partial to BF16, and then adds them. Larger inputs use a normal + // FP32 K=16 reduction. Match both orders at the fused-kernel boundary. + if (rows * cols <= 1024 * SCALE_DIM) { + float partial[2] = {0.0f, 0.0f}; +#pragma unroll + for (int group = 0; group < 2; ++group) { +#pragma unroll + for (int k = 0; k < SCALE_DIM / 2; ++k) { + const int rht_k = group * (SCALE_DIM / 2) + k; + partial[group] = + fmaf(input_vec[rht_k], + static_cast(rht_sh[rht_k * SCALE_DIM + j]), partial[group]); + } + partial[group] = static_cast(static_cast(partial[group])); + } + value = partial[0] + partial[1]; + } else { +#pragma unroll + for (int k = 0; k < SCALE_DIM; ++k) { + value = + fmaf(input_vec[k], static_cast(rht_sh[k * SCALE_DIM + j]), value); + } + } + in_colwise_IType[j] = static_cast(value); + value = static_cast(in_colwise_IType[j]); + in_compute_colwise[j] = value; + block_amax = fmaxf(block_amax, fabsf(value)); + } + } // 2. Compute E4M3 scaling factor const nvfp4_scale_t S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc_colwise); @@ -562,18 +620,38 @@ __global__ void __launch_bounds__(THREADS_NUM) // 3. Scale elements fp4e2m1x4 regs[SCALE_DIM / 4]; + if constexpr (APPLY_COLUMNWISE_RHT && NO_ACTIVATIONS_NOT_FP32_INPUT) { + uint32_t *regs_8x = reinterpret_cast(regs); #pragma unroll - for (int e = 0; e < SCALE_DIM / 4; ++e) { - const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); - if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { - const uint64_t elts = *reinterpret_cast(&in_colwise_IType[4 * e]); - regs[e] = ptx::mul_cvt_bf16_to_fp4_4x( - elts, block_scale_inverse_2x, rbits); - } else { - const float2 in01 = *reinterpret_cast(&in_compute_colwise[4 * e]); - const float2 in23 = *reinterpret_cast(&in_compute_colwise[4 * e + 2]); - regs[e] = ptx::mul_cvt_fp32_to_fp4_4x( - in01, in23, block_scale_inverse_2x, rbits); + for (int e = 0; e < SCALE_DIM / 8; ++e) { + const uint64_t elts03 = + *reinterpret_cast(&in_colwise_IType[8 * e]); + const uint64_t elts47 = + *reinterpret_cast(&in_colwise_IType[8 * e + 4]); + if constexpr (USE_STOCHASTIC_ROUNDING) { + const uint32_t rbits03 = get_rbits(rng, random_uint4, rnd_idx); + const uint32_t rbits47 = get_rbits(rng, random_uint4, rnd_idx); + regs_8x[e] = ptx::mul_cvt_bf16_to_fp4_8x_stochastic_rounding( + elts03, elts47, block_scale_inverse, rbits03, rbits47); + } else { + regs_8x[e] = ptx::mul_cvt_bf16_to_fp4_8x_round_to_nearest( + elts03, elts47, block_scale_inverse); + } + } + } else { +#pragma unroll + for (int e = 0; e < SCALE_DIM / 4; ++e) { + const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); + if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { + const uint64_t elts = *reinterpret_cast(&in_colwise_IType[4 * e]); + regs[e] = ptx::mul_cvt_bf16_to_fp4_4x( + elts, block_scale_inverse_2x, rbits); + } else { + const float2 in01 = *reinterpret_cast(&in_compute_colwise[4 * e]); + const float2 in23 = *reinterpret_cast(&in_compute_colwise[4 * e + 2]); + regs[e] = ptx::mul_cvt_fp32_to_fp4_4x( + in01, in23, block_scale_inverse_2x, rbits); + } } } @@ -601,7 +679,7 @@ __global__ void __launch_bounds__(THREADS_NUM) } // ROWWISE scaling - { + if constexpr (RETURN_ROWWISE) { const size_t stage_rowwise_scales_offset_Y = stage * BUFF_DIM_Y; #pragma unroll for (size_t it = 0; it < ITERATIONS_NORMAL; ++it) { @@ -818,9 +896,11 @@ __global__ void __launch_bounds__(THREADS_NUM) const size_t global_offset_Y_t = block_offset_Y_t; const size_t global_offset_X_t = block_offset_X_t + stage_offset_Y; - ptx::cp_async_bulk_tensor_2d_shared_to_global( - reinterpret_cast(&tensor_map_output), global_offset_X, global_offset_Y, - reinterpret_cast(&out_data_sh[buff_offset_out])); + if constexpr (RETURN_ROWWISE) { + ptx::cp_async_bulk_tensor_2d_shared_to_global( + reinterpret_cast(&tensor_map_output), global_offset_X, + global_offset_Y, reinterpret_cast(&out_data_sh[buff_offset_out])); + } if constexpr (RETURN_TRANSPOSE) { ptx::cp_async_bulk_tensor_2d_shared_to_global( @@ -853,6 +933,9 @@ __global__ void __launch_bounds__(THREADS_NUM) } destroy_barriers(mbar, is_master_thread); + if constexpr (APPLY_COLUMNWISE_RHT) { + destroy_barriers<1>(mbar_rht, is_master_thread); + } #else NVTE_DEVICE_ERROR("sm_100 or higher is required."); #endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) @@ -1418,9 +1501,10 @@ __global__ void __launch_bounds__(THREADS_NUM) #endif // FP4_TYPE_SUPPORTED } // namespace quantize_transpose_kernel -template +template void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, - const QuantizationConfig *quant_config, cudaStream_t stream) { + const QuantizationConfig *quant_config, cudaStream_t stream, + const Tensor *rht_matrix = nullptr) { #if FP4_TYPE_SUPPORTED using namespace quantize_transpose_kernel; using namespace ptx; @@ -1437,6 +1521,15 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, // Columnwise-only (no rowwise output) is supported on the optimized 2D path; the rowwise pass // and its store are gated out via the RETURN_ROWWISE template bool. const bool return_rowwise = output->has_data(); + if constexpr (apply_columnwise_rht) { + NVTE_CHECK(rht_matrix != nullptr, "Columnwise RHT requires an RHT matrix."); + NVTE_CHECK(return_transpose, "Columnwise RHT requires columnwise output."); + NVTE_CHECK(rht_matrix->dtype() == DType::kBFloat16, + "Columnwise RHT matrix must be BF16."); + NVTE_CHECK(rht_matrix->dim() == 2 && rht_matrix->shape()[0] == 16 && + rht_matrix->shape()[1] == 16, + "Columnwise RHT matrix must have shape [16, 16]."); + } if (!use_2d_quantization && (input.dtype() == DType::kBFloat16)) { quantize_transpose_tuned_1D(input, noop, output, quant_config, stream); @@ -1453,7 +1546,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, CheckOutputTensor(*output, "output", false); NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); - NVTE_CHECK(return_rowwise || (return_transpose && use_2d_quantization), + NVTE_CHECK(return_rowwise || (return_transpose && (use_2d_quantization || apply_columnwise_rht)), "NVFP4 optimized kernel supports rowwise output (1D or 2D), or columnwise-only output " "with 2D quantization."); if (return_rowwise) { @@ -1523,6 +1616,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, alignas(64) CUtensorMap tensor_map_input{}; alignas(64) CUtensorMap tensor_map_output{}; alignas(64) CUtensorMap tensor_map_output_transpose{}; + alignas(64) CUtensorMap tensor_map_rht{}; create_2D_tensor_map(tensor_map_input, input.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, cols, 0, sizeof(IType) * 8); @@ -1535,6 +1629,10 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, create_2D_tensor_map(tensor_map_output_transpose, output->columnwise_data, cols, rows, BUFF_DIM_X, BUFF_DIM_Y, rows, 0, 4); } + if constexpr (apply_columnwise_rht) { + create_2D_tensor_map(tensor_map_rht, rht_matrix->data, 16, 16, 16, 16, 16, 0, + sizeof(IType) * 8); + } constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; constexpr size_t buff_size_aligned_in = @@ -1542,6 +1640,8 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, constexpr size_t buff_size_aligned_out = DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); constexpr size_t buff_size_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * sizeof(nvfp4_scale_t); + constexpr size_t rht_mem = + DIVUP_TO_MULTIPLE(16 * 16 * sizeof(IType), TMA_SHMEM_ALIGNMENT); constexpr size_t in_mem = buff_size_aligned_in; @@ -1551,7 +1651,8 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, constexpr size_t out_mem = out_data_mem + out_data_transpose_mem; - constexpr size_t dshmem_size = in_mem + out_mem + out_scales_transpose_mem + TMA_SHMEM_ALIGNMENT; + const size_t dshmem_size = in_mem + out_mem + out_scales_transpose_mem + + (apply_columnwise_rht ? rht_mem : 0) + TMA_SHMEM_ALIGNMENT; TRANSFORMER_ENGINE_SWITCH_CONDITION( use_stochastic_rounding, USE_STOCHASTIC_ROUNDING, @@ -1559,13 +1660,21 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, TRANSFORMER_ENGINE_SWITCH_CONDITION(row_scaled_nvfp4, ROW_SCALED_NVFP4, { TRANSFORMER_ENGINE_SWITCH_CONDITION(return_rowwise, RETURN_ROWWISE, { TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { - // The 1D kernel always produces rowwise output (no RETURN_ROWWISE); the dispatch only - // routes columnwise-only requests here when use_2d_quantization is true. - auto kernel = quantize_transpose_nvfp4_kernel; - - if constexpr (use_2d_quantization) { + if constexpr (apply_columnwise_rht) { + auto kernel = quantize_transpose_nvfp4_kernel< + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, + RETURN_TRANSPOSE, ROW_SCALED_NVFP4, RETURN_ROWWISE, + /*APPLY_COLUMNWISE_RHT=*/true>; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size); + kernel<<>>( + tensor_map_input, tensor_map_output, tensor_map_output_transpose, tensor_map_rht, + scales_ptr, scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, + rows, cols, scale_stride, scale_stride_transpose, rng_state); + } else if constexpr (use_2d_quantization) { + auto kernel = quantize_transpose_nvfp4_2D_kernel< + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, + RETURN_ROWWISE, RETURN_TRANSPOSE, false>; if (with_gemm_swizzled_scales) { kernel = quantize_transpose_nvfp4_2D_kernel< COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, @@ -1575,13 +1684,23 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, RETURN_ROWWISE, RETURN_TRANSPOSE, /*WITH_GEMM_SWIZZLED_SCALES=*/false>; } + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size); + kernel<<>>( + tensor_map_input, tensor_map_output, tensor_map_output_transpose, scales_ptr, + scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, + scale_stride, scale_stride_transpose, rng_state); + } else { + auto kernel = quantize_transpose_nvfp4_kernel< + COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, + RETURN_TRANSPOSE, ROW_SCALED_NVFP4>; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size); + kernel<<>>( + tensor_map_input, tensor_map_output, tensor_map_output_transpose, tensor_map_rht, + scales_ptr, scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, + rows, cols, scale_stride, scale_stride_transpose, rng_state); } - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); - kernel<<>>( - tensor_map_input, tensor_map_output, tensor_map_output_transpose, scales_ptr, - scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, - scale_stride, scale_stride_transpose, rng_state); }); }); });); diff --git a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu index 479922a9bf..84d9088337 100644 --- a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu @@ -18,10 +18,12 @@ #include #include "common/common.h" +#include "common/cast/core/common.cuh" #include "common/util/cuda_runtime.h" #include "common/util/curanddx.hpp" #include "common/util/ptx.cuh" #include "common/utils.cuh" +#include "common/cast/nvfp4/quantize_transpose_nvfp4.cuh" #include "customized_pipeline.cuh" #include "cutlass/arch/barrier.h" #include "cutlass/arch/reg_reconfig.h" @@ -1314,6 +1316,22 @@ void hadamard_transform_cast_fusion(const Tensor &input_, Tensor &output_, NVTE_CHECK(m % hadamard_dimension == 0, "num_rows must be divisible by hadamard_dimension"); + // SM120/121 do not provide TMEM. Reuse the NVFP4 1D TMA pipeline and perform the + // 16-point columnwise RHT in registers, while leaving the SM100/110 UMMA/TMEM path below intact. + const int sm_arch = transformer_engine::cuda::sm_arch(transformer_engine::cuda::current_device()); + if (sm_arch == 120 || sm_arch == 121) { + Tensor noop; + if (output_.columnwise_data.dptr != nullptr) { + dispatch::nvfp4::quantize_transpose(input_, &noop, &output_, &quant_config, + stream, &hadamard_matrix_); + } else { + // RHT only affects the columnwise result. Keep rowwise-only quantization on + // the regular 1D scaling path rather than accidentally selecting 2D scaling. + dispatch::nvfp4::quantize_transpose(input_, &noop, &output_, &quant_config, stream); + } + return; + } + int k_tile_size = 1024; // Honor the output tensor's GEMM-swizzled-scales flag: when set, emit diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 9512038953..a20fe4320b 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1910,8 +1910,10 @@ bool NVFP4Quantizer::is_eligible_for_rht_cast_fusion(const std::vector& bool for_grouped_kernel) { const auto [rows, cols] = get_2d_dims(shape); const size_t row_align = for_grouped_kernel ? 128 : 64; - return rows % row_align == 0 && cols % 128 == 0 && transformer_engine::cuda::sm_arch() >= 100 && - transformer_engine::cuda::sm_arch() <= 110; + const int sm_arch = transformer_engine::cuda::sm_arch(); + const bool supported_arch = + (sm_arch >= 100 && sm_arch <= 110) || (!for_grouped_kernel && (sm_arch == 120 || sm_arch == 121)); + return rows % row_align == 0 && cols % 128 == 0 && supported_arch; } bool NVFP4Quantizer::is_eligible_for_2d_swizzle_fusion(const std::vector& shape) { @@ -1930,6 +1932,12 @@ bool nvfp4_emits_gemm_swizzled_scales(const NVFP4Quantizer& q, const std::vector return false; } if (q.with_rht) { + const int sm_arch = transformer_engine::cuda::sm_arch(); + // The SM12x no-TMEM RHT path reuses the compact-scale 1D TMA quantizer. + // Keep the existing post-quantize swizzle until that path gains an in-kernel swizzle epilogue. + if (sm_arch == 120 || sm_arch == 121) { + return false; + } return NVFP4Quantizer::is_eligible_for_rht_cast_fusion(shape); } // Plain (non-RHT) 2D quantize kernel can bake the swizzled layout for aligned shapes. @@ -2670,14 +2678,20 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou NVTE_CHECK(this->rht_matrix.defined() && this->rht_matrix.numel() > 0, "RHT matrix is not available."); auto rht_matrix = this->rht_matrix.to(torch_dtype); - auto rht_output = at::from_blob( - rht_output_t_cpp.get_rowwise_data().data_ptr, - {static_cast(cols), static_cast(rows)}, [](void*) {}, options); - auto rht_output_view = rht_output.view({-1, 16}); auto input_view = input_torch.transpose(0, 1).contiguous().view({-1, 16}); - at::matmul_out(rht_output_view, input_view, rht_matrix); + at::Tensor rht_output; + if (eligible_for_rht_cast_fusion) { + rht_output = at::matmul(input_view, rht_matrix) + .view({static_cast(cols), static_cast(rows)}); + } else { + rht_output = at::from_blob( + rht_output_t_cpp.get_rowwise_data().data_ptr, + {static_cast(cols), static_cast(rows)}, [](void*) {}, options); + auto rht_output_view = rht_output.view({-1, 16}); + at::matmul_out(rht_output_view, input_view, rht_matrix); + rht_output_is_ready = true; + } copy_amax(rht_output, out.get_columnwise_amax().data_ptr); - rht_output_is_ready = true; } } else { NVTE_SCOPED_GIL_RELEASE({ From fea70f2b9ba40c8e241a997d81bdddd4ce4dd302 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 3 Aug 2026 19:53:46 +0000 Subject: [PATCH 09/12] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 60 +++++++++---------- ...cast_col_hadamard_transform_cast_fusion.cu | 6 +- transformer_engine/pytorch/csrc/quantizer.cpp | 4 +- 3 files changed, 33 insertions(+), 37 deletions(-) diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index f240b3b112..dff7370713 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -317,8 +317,7 @@ constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM; // 8 = 128 / template + bool ROW_SCALED_NVFP4, bool RETURN_ROWWISE = true, bool APPLY_COLUMNWISE_RHT = false> __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, @@ -429,9 +428,9 @@ __global__ void __launch_bounds__(THREADS_NUM) dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); nvfp4_scale_t *out_colwise_scales_sh = reinterpret_cast( dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales); - IType *rht_sh = reinterpret_cast( - dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + - out_mem_rowwise_scales + out_mem_colwise_scales); + IType *rht_sh = + reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + + out_mem_rowwise_scales + out_mem_colwise_scales); IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; @@ -583,8 +582,8 @@ __global__ void __launch_bounds__(THREADS_NUM) for (int k = 0; k < SCALE_DIM / 2; ++k) { const int rht_k = group * (SCALE_DIM / 2) + k; partial[group] = - fmaf(input_vec[rht_k], - static_cast(rht_sh[rht_k * SCALE_DIM + j]), partial[group]); + fmaf(input_vec[rht_k], static_cast(rht_sh[rht_k * SCALE_DIM + j]), + partial[group]); } partial[group] = static_cast(static_cast(partial[group])); } @@ -592,8 +591,7 @@ __global__ void __launch_bounds__(THREADS_NUM) } else { #pragma unroll for (int k = 0; k < SCALE_DIM; ++k) { - value = - fmaf(input_vec[k], static_cast(rht_sh[k * SCALE_DIM + j]), value); + value = fmaf(input_vec[k], static_cast(rht_sh[k * SCALE_DIM + j]), value); } } in_colwise_IType[j] = static_cast(value); @@ -624,18 +622,16 @@ __global__ void __launch_bounds__(THREADS_NUM) uint32_t *regs_8x = reinterpret_cast(regs); #pragma unroll for (int e = 0; e < SCALE_DIM / 8; ++e) { - const uint64_t elts03 = - *reinterpret_cast(&in_colwise_IType[8 * e]); - const uint64_t elts47 = - *reinterpret_cast(&in_colwise_IType[8 * e + 4]); + const uint64_t elts03 = *reinterpret_cast(&in_colwise_IType[8 * e]); + const uint64_t elts47 = *reinterpret_cast(&in_colwise_IType[8 * e + 4]); if constexpr (USE_STOCHASTIC_ROUNDING) { const uint32_t rbits03 = get_rbits(rng, random_uint4, rnd_idx); const uint32_t rbits47 = get_rbits(rng, random_uint4, rnd_idx); regs_8x[e] = ptx::mul_cvt_bf16_to_fp4_8x_stochastic_rounding( elts03, elts47, block_scale_inverse, rbits03, rbits47); } else { - regs_8x[e] = ptx::mul_cvt_bf16_to_fp4_8x_round_to_nearest( - elts03, elts47, block_scale_inverse); + regs_8x[e] = ptx::mul_cvt_bf16_to_fp4_8x_round_to_nearest(elts03, elts47, + block_scale_inverse); } } } else { @@ -1524,11 +1520,10 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, if constexpr (apply_columnwise_rht) { NVTE_CHECK(rht_matrix != nullptr, "Columnwise RHT requires an RHT matrix."); NVTE_CHECK(return_transpose, "Columnwise RHT requires columnwise output."); - NVTE_CHECK(rht_matrix->dtype() == DType::kBFloat16, - "Columnwise RHT matrix must be BF16."); - NVTE_CHECK(rht_matrix->dim() == 2 && rht_matrix->shape()[0] == 16 && - rht_matrix->shape()[1] == 16, - "Columnwise RHT matrix must have shape [16, 16]."); + NVTE_CHECK(rht_matrix->dtype() == DType::kBFloat16, "Columnwise RHT matrix must be BF16."); + NVTE_CHECK( + rht_matrix->dim() == 2 && rht_matrix->shape()[0] == 16 && rht_matrix->shape()[1] == 16, + "Columnwise RHT matrix must have shape [16, 16]."); } if (!use_2d_quantization && (input.dtype() == DType::kBFloat16)) { @@ -1640,8 +1635,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, constexpr size_t buff_size_aligned_out = DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); constexpr size_t buff_size_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * sizeof(nvfp4_scale_t); - constexpr size_t rht_mem = - DIVUP_TO_MULTIPLE(16 * 16 * sizeof(IType), TMA_SHMEM_ALIGNMENT); + constexpr size_t rht_mem = DIVUP_TO_MULTIPLE(16 * 16 * sizeof(IType), TMA_SHMEM_ALIGNMENT); constexpr size_t in_mem = buff_size_aligned_in; @@ -1661,10 +1655,11 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, TRANSFORMER_ENGINE_SWITCH_CONDITION(return_rowwise, RETURN_ROWWISE, { TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { if constexpr (apply_columnwise_rht) { - auto kernel = quantize_transpose_nvfp4_kernel< - COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, - RETURN_TRANSPOSE, ROW_SCALED_NVFP4, RETURN_ROWWISE, - /*APPLY_COLUMNWISE_RHT=*/true>; + auto kernel = + quantize_transpose_nvfp4_kernel; cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); kernel<<>>( @@ -1672,9 +1667,10 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, scales_ptr, scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, scale_stride, scale_stride_transpose, rng_state); } else if constexpr (use_2d_quantization) { - auto kernel = quantize_transpose_nvfp4_2D_kernel< - COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, - RETURN_ROWWISE, RETURN_TRANSPOSE, false>; + auto kernel = + quantize_transpose_nvfp4_2D_kernel; if (with_gemm_swizzled_scales) { kernel = quantize_transpose_nvfp4_2D_kernel< COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, @@ -1691,9 +1687,9 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, scale_stride, scale_stride_transpose, rng_state); } else { - auto kernel = quantize_transpose_nvfp4_kernel< - COMPUTE_ACTIVATIONS, ParamOP, OP, IType, USE_STOCHASTIC_ROUNDING, - RETURN_TRANSPOSE, ROW_SCALED_NVFP4>; + auto kernel = quantize_transpose_nvfp4_kernel; cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); kernel<<>>( diff --git a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu index 84d9088337..36cabd9658 100644 --- a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu @@ -17,13 +17,13 @@ #include #include -#include "common/common.h" #include "common/cast/core/common.cuh" +#include "common/cast/nvfp4/quantize_transpose_nvfp4.cuh" +#include "common/common.h" #include "common/util/cuda_runtime.h" #include "common/util/curanddx.hpp" #include "common/util/ptx.cuh" #include "common/utils.cuh" -#include "common/cast/nvfp4/quantize_transpose_nvfp4.cuh" #include "customized_pipeline.cuh" #include "cutlass/arch/barrier.h" #include "cutlass/arch/reg_reconfig.h" @@ -1323,7 +1323,7 @@ void hadamard_transform_cast_fusion(const Tensor &input_, Tensor &output_, Tensor noop; if (output_.columnwise_data.dptr != nullptr) { dispatch::nvfp4::quantize_transpose(input_, &noop, &output_, &quant_config, - stream, &hadamard_matrix_); + stream, &hadamard_matrix_); } else { // RHT only affects the columnwise result. Keep rowwise-only quantization on // the regular 1D scaling path rather than accidentally selecting 2D scaling. diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index a20fe4320b..84f408fc47 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1911,8 +1911,8 @@ bool NVFP4Quantizer::is_eligible_for_rht_cast_fusion(const std::vector& const auto [rows, cols] = get_2d_dims(shape); const size_t row_align = for_grouped_kernel ? 128 : 64; const int sm_arch = transformer_engine::cuda::sm_arch(); - const bool supported_arch = - (sm_arch >= 100 && sm_arch <= 110) || (!for_grouped_kernel && (sm_arch == 120 || sm_arch == 121)); + const bool supported_arch = (sm_arch >= 100 && sm_arch <= 110) || + (!for_grouped_kernel && (sm_arch == 120 || sm_arch == 121)); return rows % row_align == 0 && cols % 128 == 0 && supported_arch; } From f9e28f43d7375a8bbe8b24721c27c20e47ea0afb Mon Sep 17 00:00:00 2001 From: tony <864832769@qq.com> Date: Tue, 4 Aug 2026 10:21:37 +0800 Subject: [PATCH 10/12] perf(nvfp4): use warp MMA for SM12x fused RHT Signed-off-by: tony <864832769@qq.com> --- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 151 ++++++++++++++---- 1 file changed, 118 insertions(+), 33 deletions(-) diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index dff7370713..324f19b119 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -411,6 +412,9 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_scales = 0; constexpr size_t out_mem_colwise_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / SCALE_DIM * sizeof(nvfp4_scale_t); + constexpr size_t rht_matrix_mem = 16 * 16 * sizeof(IType); + constexpr size_t rht_result_mem = BUFF_DIM_Y * BUFF_IN_DIM_X * sizeof(IType); + constexpr size_t rht_mma_a_mem = (THREADS_NUM / THREADS_PER_WARP) * 16 * 16 * sizeof(IType); extern __shared__ char dynamic_shmem[]; uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); @@ -431,6 +435,12 @@ __global__ void __launch_bounds__(THREADS_NUM) IType *rht_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales + out_mem_colwise_scales); + IType *rht_result_sh = + reinterpret_cast(reinterpret_cast(rht_sh) + rht_matrix_mem); + IType *rht_mma_a_sh = + reinterpret_cast(reinterpret_cast(rht_result_sh) + rht_result_mem); + float *rht_mma_acc_sh = + reinterpret_cast(reinterpret_cast(rht_mma_a_sh) + rht_mma_a_mem); IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; @@ -498,6 +508,103 @@ __global__ void __launch_bounds__(THREADS_NUM) ptx::mbarrier_wait_parity(&mbar[stage], 0); if constexpr (APPLY_COLUMNWISE_RHT) { ptx::mbarrier_wait_parity(&mbar_rht[0], 0); + IType *stage_rht_result_sh = rht_result_sh; + + // SM120/121 have legacy warp MMA but no TMEM. Form sixteen 16x16 tiles from + // the TMA input stage, transpose each tile into WMMA A layout, and compute + // A @ H in registers. Four warps independently process four tiles each. + constexpr int kWarps = THREADS_NUM / THREADS_PER_WARP; + constexpr int kTilesX = BUFF_IN_DIM_X / SCALE_DIM; + constexpr int kTilesY = BUFF_DIM_Y / SCALE_DIM; + constexpr int kMmaTiles = kTilesX * kTilesY; + const int warp = threadIdx.x / THREADS_PER_WARP; + const int lane = threadIdx.x % THREADS_PER_WARP; + IType *warp_a = rht_mma_a_sh + warp * SCALE_DIM * SCALE_DIM; + float *warp_acc = rht_mma_acc_sh + warp * SCALE_DIM * SCALE_DIM; + const bool reduced_precision_rht = rows * cols <= 1024 * SCALE_DIM; + + for (int tile = warp; tile < kMmaTiles; tile += kWarps) { + const int tile_y = tile / kTilesX; + const int tile_x = tile % kTilesX; + nvcuda::wmma::fragment + a_frag; + nvcuda::wmma::fragment + b_frag; + nvcuda::wmma::fragment acc_frag; + nvcuda::wmma::load_matrix_sync(b_frag, reinterpret_cast<__nv_bfloat16 *>(rht_sh), + SCALE_DIM); + + if (!reduced_precision_rht) { + // cuBLAS uses two K=8 MMA steps with an FP32 accumulator for larger + // RHT batches. Mask the opposite half of K in each step to reproduce + // that accumulation order with the K=16 WMMA primitive available here. + nvcuda::wmma::fill_fragment(acc_frag, 0.0f); + for (int mma_pass = 0; mma_pass < 2; ++mma_pass) { +#pragma unroll + for (int idx = lane; idx < SCALE_DIM * SCALE_DIM; idx += THREADS_PER_WARP) { + const int vector = idx / SCALE_DIM; + const int k = idx % SCALE_DIM; + warp_a[idx] = k / (SCALE_DIM / 2) == mma_pass + ? in_sh[buff_offset_in + (tile_y * SCALE_DIM + k) * BUFF_IN_DIM_X + + tile_x * SCALE_DIM + vector] + : static_cast(0.0f); + } + __syncwarp(); + nvcuda::wmma::load_matrix_sync(a_frag, reinterpret_cast<__nv_bfloat16 *>(warp_a), + SCALE_DIM); + nvcuda::wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag); + } + nvcuda::wmma::store_matrix_sync(warp_acc, acc_frag, SCALE_DIM, + nvcuda::wmma::mem_row_major); + __syncwarp(); +#pragma unroll + for (int idx = lane; idx < SCALE_DIM * SCALE_DIM; idx += THREADS_PER_WARP) { + const int vector = idx / SCALE_DIM; + const int out = idx % SCALE_DIM; + IType *dst = &stage_rht_result_sh[(tile_y * SCALE_DIM + out) * BUFF_IN_DIM_X + + tile_x * SCALE_DIM + vector]; + *dst = static_cast(warp_acc[idx]); + } + __syncwarp(); + } else { + // Small ATen GEMMs round two independent K=8 partials to BF16 before + // adding them. Keep the same boundary while still using MMA for both halves. + for (int mma_pass = 0; mma_pass < 2; ++mma_pass) { +#pragma unroll + for (int idx = lane; idx < SCALE_DIM * SCALE_DIM; idx += THREADS_PER_WARP) { + const int vector = idx / SCALE_DIM; + const int k = idx % SCALE_DIM; + warp_a[idx] = k / (SCALE_DIM / 2) == mma_pass + ? in_sh[buff_offset_in + (tile_y * SCALE_DIM + k) * BUFF_IN_DIM_X + + tile_x * SCALE_DIM + vector] + : static_cast(0.0f); + } + __syncwarp(); + nvcuda::wmma::load_matrix_sync(a_frag, reinterpret_cast<__nv_bfloat16 *>(warp_a), + SCALE_DIM); + nvcuda::wmma::fill_fragment(acc_frag, 0.0f); + nvcuda::wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag); + nvcuda::wmma::store_matrix_sync(warp_acc, acc_frag, SCALE_DIM, + nvcuda::wmma::mem_row_major); + __syncwarp(); +#pragma unroll + for (int idx = lane; idx < SCALE_DIM * SCALE_DIM; idx += THREADS_PER_WARP) { + const int vector = idx / SCALE_DIM; + const int out = idx % SCALE_DIM; + IType *dst = &stage_rht_result_sh[(tile_y * SCALE_DIM + out) * BUFF_IN_DIM_X + + tile_x * SCALE_DIM + vector]; + const IType rounded_partial = static_cast(warp_acc[idx]); + *dst = mma_pass == 0 ? rounded_partial + : static_cast(static_cast(*dst) + + static_cast(rounded_partial)); + } + __syncwarp(); + } + } + } + __syncthreads(); } float block_amax = 0.0f; @@ -561,41 +668,14 @@ __global__ void __launch_bounds__(THREADS_NUM) } } if constexpr (APPLY_COLUMNWISE_RHT) { - float input_vec[SCALE_DIM]; -#pragma unroll - for (int i = 0; i < SCALE_DIM; ++i) { - input_vec[i] = static_cast(in_colwise_IType[i]); - } + IType *stage_rht_result_sh = rht_result_sh; block_amax = 0.0f; #pragma unroll for (int j = 0; j < SCALE_DIM; ++j) { - float value = 0.0f; - // ATen selects a reduced-precision BF16 GEMM for at most 1024 - // K=16 vectors on SM12x. It forms two K=8 FP32 partials, rounds each - // partial to BF16, and then adds them. Larger inputs use a normal - // FP32 K=16 reduction. Match both orders at the fused-kernel boundary. - if (rows * cols <= 1024 * SCALE_DIM) { - float partial[2] = {0.0f, 0.0f}; -#pragma unroll - for (int group = 0; group < 2; ++group) { -#pragma unroll - for (int k = 0; k < SCALE_DIM / 2; ++k) { - const int rht_k = group * (SCALE_DIM / 2) + k; - partial[group] = - fmaf(input_vec[rht_k], static_cast(rht_sh[rht_k * SCALE_DIM + j]), - partial[group]); - } - partial[group] = static_cast(static_cast(partial[group])); - } - value = partial[0] + partial[1]; - } else { -#pragma unroll - for (int k = 0; k < SCALE_DIM; ++k) { - value = fmaf(input_vec[k], static_cast(rht_sh[k * SCALE_DIM + j]), value); - } - } - in_colwise_IType[j] = static_cast(value); - value = static_cast(in_colwise_IType[j]); + const int rht_offset = in_thread_offset_Y + j; + in_colwise_IType[j] = + stage_rht_result_sh[rht_offset * BUFF_IN_DIM_X + in_thread_offset_X]; + float value = static_cast(in_colwise_IType[j]); in_compute_colwise[j] = value; block_amax = fmaxf(block_amax, fabsf(value)); } @@ -1635,7 +1715,12 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, constexpr size_t buff_size_aligned_out = DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); constexpr size_t buff_size_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * sizeof(nvfp4_scale_t); - constexpr size_t rht_mem = DIVUP_TO_MULTIPLE(16 * 16 * sizeof(IType), TMA_SHMEM_ALIGNMENT); + constexpr size_t rht_matrix_mem = 16 * 16 * sizeof(IType); + constexpr size_t rht_result_mem = BUFF_DIM_Y * BUFF_DIM_X * sizeof(IType); + constexpr size_t rht_mma_a_mem = (THREADS_NUM / THREADS_PER_WARP) * 16 * 16 * sizeof(IType); + constexpr size_t rht_mma_acc_mem = (THREADS_NUM / THREADS_PER_WARP) * 16 * 16 * sizeof(float); + const size_t rht_mem = DIVUP_TO_MULTIPLE( + rht_matrix_mem + rht_result_mem + rht_mma_a_mem + rht_mma_acc_mem, TMA_SHMEM_ALIGNMENT); constexpr size_t in_mem = buff_size_aligned_in; From 36e1179f4f47981c086630030cf3a4057c779fe0 Mon Sep 17 00:00:00 2001 From: tony <864832769@qq.com> Date: Tue, 4 Aug 2026 11:09:15 +0800 Subject: [PATCH 11/12] fix(nvfp4): align SM12x fused RHT with native K16 MMA Signed-off-by: tony <864832769@qq.com> --- .../nvfp4/test_nvfp4_rht_quantize_exact.py | 30 +++++ tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py | 101 ++++++++-------- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 86 ++++--------- transformer_engine/pytorch/csrc/common.h | 2 +- transformer_engine/pytorch/csrc/quantizer.cpp | 113 ++++-------------- 5 files changed, 127 insertions(+), 205 deletions(-) diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py index a65010ef02..0bc1b6c337 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py @@ -9,6 +9,8 @@ # Due to the structure of NVFP4Quantizer, we need to test the RHT functionality # together with the quantization functionality. +import os + import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer @@ -127,6 +129,34 @@ def check_quantization_nvfp4_versus_reference( ref_quantizer._apply_rht(x.t().contiguous()) if with_rht else x.t().contiguous() ) ref_amax_colwise_t = torch.max(torch.abs(x_t_for_amax)).to(torch.float32).view(1) + + # SM120/121 uses TE's native single-K=16 MMA Hadamard arithmetic. cuBLAS may + # choose a different reduction order (and applies the random-sign matrix in + # the ATen reference orientation), so use the unfused TE kernel as the exact + # reference for the fused TE kernel on these architectures. + if torch.cuda.get_device_capability() in ((12, 0), (12, 1)): + env_name = "NVTE_NVFP4_DISABLE_RHT_CAST_FUSION" + old_value = os.environ.get(env_name) + os.environ[env_name] = "1" + try: + native_quantizer = NVFP4Quantizer( + fp4_dtype=te_dtype, + rowwise=False, + columnwise=True, + with_amax_reduction=False, + with_rht=True, + with_post_rht_amax=True, + with_random_sign_mask=with_random_sign_mask, + ) + native = native_quantizer(x) + finally: + if old_value is None: + os.environ.pop(env_name) + else: + os.environ[env_name] = old_value + qx_t_ref = unpack_fp4(native._columnwise_data.view(dtype=torch.uint8)) + sx_t_ref = native._columnwise_scale_inv.view(dtype=torch.uint8) + ref_amax_colwise_t = native._amax_columnwise else: qx_t_ref = None sx_t_ref = None diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py b/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py index 15404f4737..9152c5dc13 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_sm12x.py @@ -13,9 +13,6 @@ import transformer_engine.pytorch as te from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes import utils -from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef - recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) @@ -24,34 +21,51 @@ def _is_sm12x(device: int = 0) -> bool: return torch.cuda.get_device_capability(device) in ((12, 0), (12, 1)) -def _reference_post_rht_amax(x: torch.Tensor, with_random_sign_mask: bool) -> torch.Tensor: - quantizer = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, - rowwise=False, - columnwise=True, - pow_2_scales=False, - eps=0.0, - quant_tile_shape=(1, 16), - with_rht=True, - with_random_sign_mask=with_random_sign_mask, - ) - transformed = quantizer._apply_rht(x.t().contiguous()) - return transformed.abs().amax().to(torch.float32).view(1) +def _native_unfused_columnwise(x: torch.Tensor, with_random_sign_mask: bool): + """Run TE's native K=16 RHT with cast fusion disabled.""" + + env_name = "NVTE_NVFP4_DISABLE_RHT_CAST_FUSION" + old_value = os.environ.get(env_name) + os.environ[env_name] = "1" + try: + quantizer = NVFP4Quantizer( + fp4_dtype=te.DType.kFloat4E2M1, + rowwise=False, + columnwise=True, + with_rht=True, + with_post_rht_amax=True, + with_random_sign_mask=with_random_sign_mask, + ) + return quantizer(x) + finally: + if old_value is None: + os.environ.pop(env_name) + else: + os.environ[env_name] = old_value + + +def _unpack_fp4(x: torch.Tensor) -> torch.Tensor: + unpacked = x.view(torch.uint8).repeat_interleave(2, dim=-1) + unpacked[..., 0::2] &= 0x0F + unpacked[..., 1::2] >>= 4 + return unpacked @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") @pytest.mark.parametrize("with_random_sign_mask", [False, True]) -def test_sm12x_post_rht_amax_matches_aten_output(with_random_sign_mask: bool) -> None: - """The fused post-RHT amax must describe the ATen-equivalent RHT tensor.""" +def test_sm12x_post_rht_amax_matches_native_k16(with_random_sign_mask: bool) -> None: + """The fused post-RHT amax must match TE's native K=16 MMA RHT.""" if not _is_sm12x(): - pytest.skip("Test targets the SM120/SM121 ATen RHT fallback") + pytest.skip("Test targets the SM120/SM121 no-TMEM fused RHT path") torch.manual_seed(1234) x = torch.randn((128, 128), device="cuda", dtype=torch.bfloat16) - expected_amax = _reference_post_rht_amax(x, with_random_sign_mask) + torch.manual_seed(5678) + expected = _native_unfused_columnwise(x, with_random_sign_mask) + torch.manual_seed(5678) quantizer = NVFP4Quantizer( fp4_dtype=te.DType.kFloat4E2M1, rowwise=False, @@ -63,7 +77,7 @@ def test_sm12x_post_rht_amax_matches_aten_output(with_random_sign_mask: bool) -> ) out = quantizer(x) - torch.testing.assert_close(out._amax_columnwise, expected_amax, atol=0.0, rtol=0.0) + torch.testing.assert_close(out._amax_columnwise, expected._amax_columnwise, atol=0.0, rtol=0.0) @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @@ -71,16 +85,18 @@ def test_sm12x_post_rht_amax_matches_aten_output(with_random_sign_mask: bool) -> @pytest.mark.parametrize("shape", [(128, 128), (256, 256)]) @pytest.mark.parametrize("rowwise", [False, True]) @pytest.mark.parametrize("with_random_sign_mask", [False, True]) -def test_sm12x_fused_rht_codes_and_scales_match_aten_pipeline( - shape: tuple[int, int], rowwise: bool, with_random_sign_mask: bool +@pytest.mark.parametrize("seed", [1234, 2026]) +def test_sm12x_fused_rht_codes_and_scales_match_native_k16( + shape: tuple[int, int], rowwise: bool, with_random_sign_mask: bool, seed: int ) -> None: - """Fused RHT codes/scales must match ATen RHT followed by the tuned quantizer.""" + """Fused RHT codes/scales must match TE's native K=16 MMA RHT path.""" if not _is_sm12x(): pytest.skip("Test targets the SM120/SM121 no-TMEM fused RHT path") - torch.manual_seed(2026) + torch.manual_seed(seed) x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) + torch.manual_seed(5678) fused_quantizer = NVFP4Quantizer( fp4_dtype=te.DType.kFloat4E2M1, rowwise=rowwise, @@ -92,40 +108,23 @@ def test_sm12x_fused_rht_codes_and_scales_match_aten_pipeline( ) fused = fused_quantizer(x) - reference_rht = NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, - rowwise=False, - columnwise=True, - pow_2_scales=False, - eps=0.0, - quant_tile_shape=(1, 16), - with_rht=True, - with_random_sign_mask=with_random_sign_mask, - ) - transformed = reference_rht._apply_rht(x.t().contiguous()) - plain_quantizer = NVFP4Quantizer( - fp4_dtype=te.DType.kFloat4E2M1, - rowwise=True, - columnwise=False, - with_amax_reduction=False, - with_rht=False, - ) - expected_columnwise = plain_quantizer(transformed) + torch.manual_seed(5678) + expected = _native_unfused_columnwise(x, with_random_sign_mask) torch.testing.assert_close( - fused._columnwise_data.view(torch.uint8), - expected_columnwise._rowwise_data.view(torch.uint8), + _unpack_fp4(fused._columnwise_data), + _unpack_fp4(expected._columnwise_data), atol=0.0, rtol=0.0, ) torch.testing.assert_close( fused._columnwise_scale_inv, - expected_columnwise._rowwise_scale_inv, + expected._columnwise_scale_inv, atol=0.0, rtol=0.0, ) torch.testing.assert_close( - fused._amax_columnwise, expected_columnwise._amax_rowwise, atol=0.0, rtol=0.0 + fused._amax_columnwise, expected._amax_columnwise, atol=0.0, rtol=0.0 ) @@ -141,8 +140,10 @@ def _distributed_amax_worker(rank: int, world_size: int, init_file: str) -> None torch.manual_seed(4321 + rank) x = torch.randn((128, 128), device=f"cuda:{rank}", dtype=torch.bfloat16) x.mul_(rank + 1) - expected_amax = _reference_post_rht_amax(x, with_random_sign_mask=True) + torch.manual_seed(5678) + expected_amax = _native_unfused_columnwise(x, with_random_sign_mask=True)._amax_columnwise + torch.manual_seed(5678) quantizer = NVFP4Quantizer( fp4_dtype=te.DType.kFloat4E2M1, rowwise=False, @@ -167,7 +168,7 @@ def test_sm12x_post_rht_amax_reduction_is_global(tmp_path) -> None: """ATen post-RHT amax must be computed before the distributed MAX reduction.""" if not all(_is_sm12x(device) for device in range(2)): - pytest.skip("Test targets the SM120/SM121 ATen RHT fallback") + pytest.skip("Test targets the SM120/SM121 no-TMEM fused RHT path") init_file = os.fspath(tmp_path / "nvfp4_rht_amax_init") mp.spawn(_distributed_amax_worker, args=(2, init_file), nprocs=2, join=True) diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 324f19b119..38961acaf4 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -521,7 +521,6 @@ __global__ void __launch_bounds__(THREADS_NUM) const int lane = threadIdx.x % THREADS_PER_WARP; IType *warp_a = rht_mma_a_sh + warp * SCALE_DIM * SCALE_DIM; float *warp_acc = rht_mma_acc_sh + warp * SCALE_DIM * SCALE_DIM; - const bool reduced_precision_rht = rows * cols <= 1024 * SCALE_DIM; for (int tile = warp; tile < kMmaTiles; tile += kWarps) { const int tile_y = tile / kTilesX; @@ -535,74 +534,29 @@ __global__ void __launch_bounds__(THREADS_NUM) nvcuda::wmma::fragment acc_frag; nvcuda::wmma::load_matrix_sync(b_frag, reinterpret_cast<__nv_bfloat16 *>(rht_sh), SCALE_DIM); - - if (!reduced_precision_rht) { - // cuBLAS uses two K=8 MMA steps with an FP32 accumulator for larger - // RHT batches. Mask the opposite half of K in each step to reproduce - // that accumulation order with the K=16 WMMA primitive available here. - nvcuda::wmma::fill_fragment(acc_frag, 0.0f); - for (int mma_pass = 0; mma_pass < 2; ++mma_pass) { -#pragma unroll - for (int idx = lane; idx < SCALE_DIM * SCALE_DIM; idx += THREADS_PER_WARP) { - const int vector = idx / SCALE_DIM; - const int k = idx % SCALE_DIM; - warp_a[idx] = k / (SCALE_DIM / 2) == mma_pass - ? in_sh[buff_offset_in + (tile_y * SCALE_DIM + k) * BUFF_IN_DIM_X + - tile_x * SCALE_DIM + vector] - : static_cast(0.0f); - } - __syncwarp(); - nvcuda::wmma::load_matrix_sync(a_frag, reinterpret_cast<__nv_bfloat16 *>(warp_a), - SCALE_DIM); - nvcuda::wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag); - } - nvcuda::wmma::store_matrix_sync(warp_acc, acc_frag, SCALE_DIM, - nvcuda::wmma::mem_row_major); - __syncwarp(); -#pragma unroll - for (int idx = lane; idx < SCALE_DIM * SCALE_DIM; idx += THREADS_PER_WARP) { - const int vector = idx / SCALE_DIM; - const int out = idx % SCALE_DIM; - IType *dst = &stage_rht_result_sh[(tile_y * SCALE_DIM + out) * BUFF_IN_DIM_X + - tile_x * SCALE_DIM + vector]; - *dst = static_cast(warp_acc[idx]); - } - __syncwarp(); - } else { - // Small ATen GEMMs round two independent K=8 partials to BF16 before - // adding them. Keep the same boundary while still using MMA for both halves. - for (int mma_pass = 0; mma_pass < 2; ++mma_pass) { #pragma unroll - for (int idx = lane; idx < SCALE_DIM * SCALE_DIM; idx += THREADS_PER_WARP) { - const int vector = idx / SCALE_DIM; - const int k = idx % SCALE_DIM; - warp_a[idx] = k / (SCALE_DIM / 2) == mma_pass - ? in_sh[buff_offset_in + (tile_y * SCALE_DIM + k) * BUFF_IN_DIM_X + - tile_x * SCALE_DIM + vector] - : static_cast(0.0f); - } - __syncwarp(); - nvcuda::wmma::load_matrix_sync(a_frag, reinterpret_cast<__nv_bfloat16 *>(warp_a), - SCALE_DIM); - nvcuda::wmma::fill_fragment(acc_frag, 0.0f); - nvcuda::wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag); - nvcuda::wmma::store_matrix_sync(warp_acc, acc_frag, SCALE_DIM, - nvcuda::wmma::mem_row_major); - __syncwarp(); + for (int idx = lane; idx < SCALE_DIM * SCALE_DIM; idx += THREADS_PER_WARP) { + const int vector = idx / SCALE_DIM; + const int k = idx % SCALE_DIM; + warp_a[idx] = in_sh[buff_offset_in + (tile_y * SCALE_DIM + k) * BUFF_IN_DIM_X + + tile_x * SCALE_DIM + vector]; + } + __syncwarp(); + nvcuda::wmma::load_matrix_sync(a_frag, reinterpret_cast<__nv_bfloat16 *>(warp_a), + SCALE_DIM); + nvcuda::wmma::fill_fragment(acc_frag, 0.0f); + nvcuda::wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag); + nvcuda::wmma::store_matrix_sync(warp_acc, acc_frag, SCALE_DIM, nvcuda::wmma::mem_row_major); + __syncwarp(); #pragma unroll - for (int idx = lane; idx < SCALE_DIM * SCALE_DIM; idx += THREADS_PER_WARP) { - const int vector = idx / SCALE_DIM; - const int out = idx % SCALE_DIM; - IType *dst = &stage_rht_result_sh[(tile_y * SCALE_DIM + out) * BUFF_IN_DIM_X + - tile_x * SCALE_DIM + vector]; - const IType rounded_partial = static_cast(warp_acc[idx]); - *dst = mma_pass == 0 ? rounded_partial - : static_cast(static_cast(*dst) + - static_cast(rounded_partial)); - } - __syncwarp(); - } + for (int idx = lane; idx < SCALE_DIM * SCALE_DIM; idx += THREADS_PER_WARP) { + const int vector = idx / SCALE_DIM; + const int out = idx % SCALE_DIM; + IType *dst = &stage_rht_result_sh[(tile_y * SCALE_DIM + out) * BUFF_IN_DIM_X + + tile_x * SCALE_DIM + vector]; + *dst = static_cast(warp_acc[idx]); } + __syncwarp(); } __syncthreads(); } diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 9f4da7e1e0..aa0e0c87fe 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -421,7 +421,7 @@ class NVFP4Quantizer : public Quantizer { TensorWrapper& rht_output_t_cpp, QuantizationConfigWrapper& quant_config, QuantizationConfigWrapper& quant_config_columnwise, - cudaStream_t stream, bool rht_output_is_ready); + cudaStream_t stream); }; std::unique_ptr convert_quantizer(py::handle quantizer); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 84f408fc47..865d0f9432 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1908,6 +1908,9 @@ void NVFP4Quantizer::set_quantization_params(TensorWrapper* tensor) const { bool NVFP4Quantizer::is_eligible_for_rht_cast_fusion(const std::vector& shape, bool for_grouped_kernel) { + if (transformer_engine::getenv("NVTE_NVFP4_DISABLE_RHT_CAST_FUSION")) { + return false; + } const auto [rows, cols] = get_2d_dims(shape); const size_t row_align = for_grouped_kernel ? 128 : 64; const int sm_arch = transformer_engine::cuda::sm_arch(); @@ -2417,7 +2420,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( void NVFP4Quantizer::quantize_with_rht_unfused_helper( const TensorWrapper& input, TensorWrapper& out, TensorWrapper& rht_output_t_cpp, QuantizationConfigWrapper& quant_config, QuantizationConfigWrapper& quant_config_columnwise, - cudaStream_t stream, bool rht_output_is_ready) { + cudaStream_t stream) { // The kernels invoked below reject swizzled-SF output, so trip a clear // error here before reaching them. NVTE_CHECK(!out.get_with_gemm_swizzled_scales(), @@ -2472,38 +2475,12 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper( static_cast(out_columnwise_amax.dtype), out_columnwise_amax.shape); - // Invoking fallback RHT kernel unfused unless the SM120/SM121 ATen result was already - // materialized while computing post-RHT amax. - const int sm_arch = transformer_engine::cuda::sm_arch(); - if (rht_output_is_ready) { - NVTE_CHECK(sm_arch == 120 || sm_arch == 121, - "A precomputed unfused RHT output is only supported on SM120/SM121."); - } else if (sm_arch == 120 || sm_arch == 121) { - // Match the PyTorch reference arithmetic on GeForce Blackwell. The BF16 MMA - // Hadamard kernel uses a different accumulation order and can differ by one - // BF16 ULP, which is enough to change exact FP4 rounding at ties. - NVTE_CHECK(this->rht_matrix.defined() && this->rht_matrix.numel() > 0, - "RHT matrix is not available."); - const auto [rows, cols] = get_2d_dims(input.shape()); - const auto torch_dtype = GetATenDType(input.dtype()); - auto options = at::TensorOptions().dtype(torch_dtype).device(torch::kCUDA); - auto input_torch = at::from_blob( - input.get_rowwise_data().data_ptr, - {static_cast(rows), static_cast(cols)}, [](void*) {}, options); - auto output_torch = at::from_blob( - rht_output_t_cpp.get_rowwise_data().data_ptr, - {static_cast(cols), static_cast(rows)}, [](void*) {}, options); - auto rht_matrix = this->rht_matrix.to(torch_dtype); - auto output_view = output_torch.view({-1, 16}); - auto input_view = input_torch.transpose(0, 1).contiguous().view({-1, 16}); - at::matmul_out(output_view, input_view, rht_matrix); - } else { - NVTE_SCOPED_GIL_RELEASE({ - // Perform the RHT(input.t), and write to rht_output_cpp.columnwise. - nvte_hadamard_transform(input.data(), rht_output_t_cpp.data(), 0, - this->rht_matrix_random_sign_mask_t, stream); - }); - } + // Invoking fallback RHT kernel unfused. + NVTE_SCOPED_GIL_RELEASE({ + // Perform the RHT(input.t), and write to rht_output_cpp.columnwise. + nvte_hadamard_transform(input.data(), rht_output_t_cpp.data(), 0, + this->rht_matrix_random_sign_mask_t, stream); + }); // Quantize kernel will treat everything as rowwise input/output, which is // intended. @@ -2594,19 +2571,6 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou input.dtype() == DType::kBFloat16 && NVFP4Quantizer::is_eligible_for_rht_cast_fusion(convertShape(input.shape())); - // The unfused path needs an intermediate in transposed layout. Allocate it before amax - // computation so SM120/SM121 can materialize the ATen RHT once, derive post-RHT amax from - // that exact buffer, and reuse it for quantization after any distributed amax reduction. - at::Tensor rht_output_t; - TensorWrapper rht_output_t_cpp; - bool rht_output_is_ready = false; - if (this->with_rht && !eligible_for_rht_cast_fusion) { - rht_output_t = - allocateTorchTensor(static_cast(cols), static_cast(rows), input.dtype()); - rht_output_t_cpp.set_rowwise_data(rht_output_t.data_ptr(), input.dtype(), - std::vector{cols, rows}); - } - // Stochastic rounding // When both rowwise and columnwise quantization are used with RHT, // we need separate RNG states for each to ensure they use different random numbers. @@ -2659,46 +2623,10 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou // 1. Rowwise amax = amax for input // 2. Columnwise amax = amax for RHT(input.t) if (compute_amax) { - const int sm_arch = transformer_engine::cuda::sm_arch(); - if (sm_arch == 120 || sm_arch == 121) { - const auto torch_dtype = GetATenDType(input.dtype()); - auto options = at::TensorOptions().dtype(torch_dtype).device(torch::kCUDA); - auto input_torch = at::from_blob( - input.get_rowwise_data().data_ptr, - {static_cast(rows), static_cast(cols)}, [](void*) {}, options); - auto amax_options = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); - auto copy_amax = [&amax_options](const at::Tensor& tensor, void* dst) { - if (dst != nullptr) { - auto amax = at::from_blob(dst, {1}, [](void*) {}, amax_options); - amax.copy_(tensor.abs().amax().to(torch::kFloat32).view({1})); - } - }; - copy_amax(input_torch, out.get_amax().data_ptr); - if (out.get_columnwise_amax().data_ptr != nullptr) { - NVTE_CHECK(this->rht_matrix.defined() && this->rht_matrix.numel() > 0, - "RHT matrix is not available."); - auto rht_matrix = this->rht_matrix.to(torch_dtype); - auto input_view = input_torch.transpose(0, 1).contiguous().view({-1, 16}); - at::Tensor rht_output; - if (eligible_for_rht_cast_fusion) { - rht_output = at::matmul(input_view, rht_matrix) - .view({static_cast(cols), static_cast(rows)}); - } else { - rht_output = at::from_blob( - rht_output_t_cpp.get_rowwise_data().data_ptr, - {static_cast(cols), static_cast(rows)}, [](void*) {}, options); - auto rht_output_view = rht_output.view({-1, 16}); - at::matmul_out(rht_output_view, input_view, rht_matrix); - rht_output_is_ready = true; - } - copy_amax(rht_output, out.get_columnwise_amax().data_ptr); - } - } else { - NVTE_SCOPED_GIL_RELEASE({ - nvte_hadamard_transform_amax(input.data(), out.data(), 0, - this->rht_matrix_random_sign_mask_t, stream); - }); - } + NVTE_SCOPED_GIL_RELEASE({ + nvte_hadamard_transform_amax(input.data(), out.data(), 0, + this->rht_matrix_random_sign_mask_t, stream); + }); } } else { // raise error since it's not supported yet @@ -2773,9 +2701,18 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou // are separate kernel launches auto& columnwise_quant_config_to_use = need_separate_columnwise_rng ? quant_config_columnwise : quant_config; + // unfused path also needs memory allocation for intermediate buffer for RHT output + at::Tensor rht_output_t; // The RHT(x_t) output, in columnwise layout + // This wrapper is going to be passed as input to the quantization kernel. + TensorWrapper rht_output_t_cpp; // Wrapper to contain the RHT(x) and RHT(x_t) outputs + rht_output_t = + allocateTorchTensor(static_cast(cols), static_cast(rows), input.dtype()); + // NOTE (frsun): This is non-intuitive, we are writing the + // result of transposed RHT to the output of rowwise. + rht_output_t_cpp.set_rowwise_data(rht_output_t.data_ptr(), input.dtype(), + std::vector{cols, rows}); this->quantize_with_rht_unfused_helper(input, out, rht_output_t_cpp, quant_config, - columnwise_quant_config_to_use, stream, - rht_output_is_ready); + columnwise_quant_config_to_use, stream); } } else { NVTE_SCOPED_GIL_RELEASE({ nvte_quantize_v2(input.data(), out.data(), quant_config, stream); }); From 9923d1cd4b090c94214fa5855b38546b0ef2ba88 Mon Sep 17 00:00:00 2001 From: tony <864832769@qq.com> Date: Tue, 4 Aug 2026 11:30:00 +0800 Subject: [PATCH 12/12] perf(nvfp4): quantize SM12x RHT fragments in registers Signed-off-by: tony <864832769@qq.com> --- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 130 +++++++++--------- 1 file changed, 64 insertions(+), 66 deletions(-) diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 38961acaf4..532c760bf0 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -14,12 +14,12 @@ #include #include #include -#include #include #include #include "../../common.h" +#include "../../hadamard_transform/hadamard_transform_utils.cuh" #include "../../util/math.h" #include "../../util/ptx.cuh" #include "../../utils.cuh" @@ -31,6 +31,16 @@ namespace transformer_engine { namespace dispatch { namespace nvfp4 { +__device__ __forceinline__ void load_matrix_b_16x16_from_shared(uint32_t &b0, uint32_t &b1, + uint32_t &b2, uint32_t &b3, + const void *addr, uint32_t stride) { + asm volatile( + "wmma.load.b.sync.aligned.row.m16n16k16.shared::cta.bf16 " + "{%0,%1,%2,%3}, [%4], %5;\n" + : "=r"(b0), "=r"(b1), "=r"(b2), "=r"(b3) + : "l"(addr), "r"(stride)); +} + namespace rowwise_amax_kernel { using namespace ptx; @@ -412,9 +422,6 @@ __global__ void __launch_bounds__(THREADS_NUM) constexpr size_t out_mem_rowwise_scales = 0; constexpr size_t out_mem_colwise_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / SCALE_DIM * sizeof(nvfp4_scale_t); - constexpr size_t rht_matrix_mem = 16 * 16 * sizeof(IType); - constexpr size_t rht_result_mem = BUFF_DIM_Y * BUFF_IN_DIM_X * sizeof(IType); - constexpr size_t rht_mma_a_mem = (THREADS_NUM / THREADS_PER_WARP) * 16 * 16 * sizeof(IType); extern __shared__ char dynamic_shmem[]; uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); @@ -435,12 +442,6 @@ __global__ void __launch_bounds__(THREADS_NUM) IType *rht_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data + out_mem_rowwise_scales + out_mem_colwise_scales); - IType *rht_result_sh = - reinterpret_cast(reinterpret_cast(rht_sh) + rht_matrix_mem); - IType *rht_mma_a_sh = - reinterpret_cast(reinterpret_cast(rht_result_sh) + rht_result_mem); - float *rht_mma_acc_sh = - reinterpret_cast(reinterpret_cast(rht_mma_a_sh) + rht_mma_a_mem); IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer constexpr size_t shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; @@ -508,63 +509,77 @@ __global__ void __launch_bounds__(THREADS_NUM) ptx::mbarrier_wait_parity(&mbar[stage], 0); if constexpr (APPLY_COLUMNWISE_RHT) { ptx::mbarrier_wait_parity(&mbar_rht[0], 0); - IType *stage_rht_result_sh = rht_result_sh; // SM120/121 have legacy warp MMA but no TMEM. Form sixteen 16x16 tiles from - // the TMA input stage, transpose each tile into WMMA A layout, and compute - // A @ H in registers. Four warps independently process four tiles each. + // the TMA input stage and compute A @ H in registers. The col-major A load + // consumes the transposed TMA layout directly, avoiding an operand staging copy. constexpr int kWarps = THREADS_NUM / THREADS_PER_WARP; constexpr int kTilesX = BUFF_IN_DIM_X / SCALE_DIM; constexpr int kTilesY = BUFF_DIM_Y / SCALE_DIM; constexpr int kMmaTiles = kTilesX * kTilesY; const int warp = threadIdx.x / THREADS_PER_WARP; - const int lane = threadIdx.x % THREADS_PER_WARP; - IType *warp_a = rht_mma_a_sh + warp * SCALE_DIM * SCALE_DIM; - float *warp_acc = rht_mma_acc_sh + warp * SCALE_DIM * SCALE_DIM; + uint32_t b_frag[4]; + load_matrix_b_16x16_from_shared(b_frag[0], b_frag[1], b_frag[2], b_frag[3], rht_sh, + SCALE_DIM); for (int tile = warp; tile < kMmaTiles; tile += kWarps) { const int tile_y = tile / kTilesX; const int tile_x = tile % kTilesX; - nvcuda::wmma::fragment - a_frag; - nvcuda::wmma::fragment - b_frag; - nvcuda::wmma::fragment acc_frag; - nvcuda::wmma::load_matrix_sync(b_frag, reinterpret_cast<__nv_bfloat16 *>(rht_sh), - SCALE_DIM); + uint32_t a_frag[4]; + uint32_t c_frag[4]; + uint32_t unused_amax = 0; + IType *tile_in = + &in_sh[buff_offset_in + tile_y * SCALE_DIM * BUFF_IN_DIM_X + tile_x * SCALE_DIM]; + load_matrix_16x16_from_shared(a_frag[0], a_frag[1], a_frag[2], a_frag[3], tile_in, + BUFF_IN_DIM_X); + mma_m16_n16_k16_b16_b16_b16_noacc( + a_frag[0], a_frag[1], a_frag[2], a_frag[3], b_frag[0], b_frag[1], b_frag[2], b_frag[3], + c_frag[0], c_frag[1], c_frag[2], c_frag[3], unused_amax); + + // A WMMA lane owns two adjacent pairs in each of two rows. Quantize + // those fragments in registers and write both packed FP4 pairs + // directly to the transposed-output staging buffer. + const int lane = threadIdx.x % THREADS_PER_WARP; + const int lane_in_quad = lane & 3; + const int row_in_half = lane >> 2; #pragma unroll - for (int idx = lane; idx < SCALE_DIM * SCALE_DIM; idx += THREADS_PER_WARP) { - const int vector = idx / SCALE_DIM; - const int k = idx % SCALE_DIM; - warp_a[idx] = in_sh[buff_offset_in + (tile_y * SCALE_DIM + k) * BUFF_IN_DIM_X + - tile_x * SCALE_DIM + vector]; - } - __syncwarp(); - nvcuda::wmma::load_matrix_sync(a_frag, reinterpret_cast<__nv_bfloat16 *>(warp_a), - SCALE_DIM); - nvcuda::wmma::fill_fragment(acc_frag, 0.0f); - nvcuda::wmma::mma_sync(acc_frag, a_frag, b_frag, acc_frag); - nvcuda::wmma::store_matrix_sync(warp_acc, acc_frag, SCALE_DIM, nvcuda::wmma::mem_row_major); - __syncwarp(); -#pragma unroll - for (int idx = lane; idx < SCALE_DIM * SCALE_DIM; idx += THREADS_PER_WARP) { - const int vector = idx / SCALE_DIM; - const int out = idx % SCALE_DIM; - IType *dst = &stage_rht_result_sh[(tile_y * SCALE_DIM + out) * BUFF_IN_DIM_X + - tile_x * SCALE_DIM + vector]; - *dst = static_cast(warp_acc[idx]); + for (int row_half = 0; row_half < 2; ++row_half) { + const uint32_t c_lo = c_frag[row_half * 2]; + const uint32_t c_hi = c_frag[row_half * 2 + 1]; + __nv_bfloat162 lo = *reinterpret_cast(&c_lo); + __nv_bfloat162 hi = *reinterpret_cast(&c_hi); + float row_amax = + fmaxf(fmaxf(fabsf(__bfloat162float(lo.x)), fabsf(__bfloat162float(lo.y))), + fmaxf(fabsf(__bfloat162float(hi.x)), fabsf(__bfloat162float(hi.y)))); + row_amax = fmaxf(row_amax, __shfl_xor_sync(0xffffffff, row_amax, 1)); + row_amax = fmaxf(row_amax, __shfl_xor_sync(0xffffffff, row_amax, 2)); + const nvfp4_scale_t scale = compute_decoding_scaling_factor(row_amax, S_enc_colwise); + const float scale_inverse = fminf(1.0f / (static_cast(scale) * S_dec_colwise), + detail::TypeExtrema::max); + const float2 scale_inverse_2x{scale_inverse, scale_inverse}; + const uint64_t values = static_cast(c_lo) | (static_cast(c_hi) << 32); + const uint32_t rbits = get_rbits(rng, random_uint4, rnd_idx); + const fp4e2m1x4 packed = + ptx::mul_cvt_bf16_to_fp4_4x(values, scale_inverse_2x, rbits); + const uint16_t packed_bits = *reinterpret_cast(&packed); + const int output_row = tile_x * SCALE_DIM + row_in_half + row_half * 8; + uint8_t *output_bytes = reinterpret_cast(out_t_data_sh) + buff_offset_out_t + + output_row * BUFF_OUT_T_DIM_X + tile_y * (SCALE_DIM / 2); + output_bytes[lane_in_quad] = static_cast(packed_bits); + output_bytes[lane_in_quad + 4] = static_cast(packed_bits >> 8); + if (lane_in_quad == 0) { + const size_t scale_idx = + output_row * SCALES_PER_CHUNK_Y + stage * ITERATIONS_TRANSPOSE + tile_y; + out_colwise_scales_sh[scale_idx] = scale; + } } - __syncwarp(); } - __syncthreads(); } float block_amax = 0.0f; // COLWISE scaling - if constexpr (RETURN_TRANSPOSE) { + if constexpr (RETURN_TRANSPOSE && !APPLY_COLUMNWISE_RHT) { #pragma unroll for (size_t it = 0; it < ITERATIONS_TRANSPOSE; ++it) { const size_t in_thread_offset_Y = 0 + it * SCALE_DIM; @@ -621,19 +636,6 @@ __global__ void __launch_bounds__(THREADS_NUM) in_compute_colwise[i] = elt; } } - if constexpr (APPLY_COLUMNWISE_RHT) { - IType *stage_rht_result_sh = rht_result_sh; - block_amax = 0.0f; -#pragma unroll - for (int j = 0; j < SCALE_DIM; ++j) { - const int rht_offset = in_thread_offset_Y + j; - in_colwise_IType[j] = - stage_rht_result_sh[rht_offset * BUFF_IN_DIM_X + in_thread_offset_X]; - float value = static_cast(in_colwise_IType[j]); - in_compute_colwise[j] = value; - block_amax = fmaxf(block_amax, fabsf(value)); - } - } // 2. Compute E4M3 scaling factor const nvfp4_scale_t S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc_colwise); @@ -1670,11 +1672,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); constexpr size_t buff_size_scales = (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * sizeof(nvfp4_scale_t); constexpr size_t rht_matrix_mem = 16 * 16 * sizeof(IType); - constexpr size_t rht_result_mem = BUFF_DIM_Y * BUFF_DIM_X * sizeof(IType); - constexpr size_t rht_mma_a_mem = (THREADS_NUM / THREADS_PER_WARP) * 16 * 16 * sizeof(IType); - constexpr size_t rht_mma_acc_mem = (THREADS_NUM / THREADS_PER_WARP) * 16 * 16 * sizeof(float); - const size_t rht_mem = DIVUP_TO_MULTIPLE( - rht_matrix_mem + rht_result_mem + rht_mma_a_mem + rht_mma_acc_mem, TMA_SHMEM_ALIGNMENT); + constexpr size_t rht_mem = DIVUP_TO_MULTIPLE(rht_matrix_mem, TMA_SHMEM_ALIGNMENT); constexpr size_t in_mem = buff_size_aligned_in;