From fa5881331e579ede745bd868fba989f21723d1d3 Mon Sep 17 00:00:00 2001 From: David Kogan Date: Mon, 27 Jul 2026 04:08:06 -0400 Subject: [PATCH 1/2] Add architecture gate to the NVFP4 split_quantize RHT path split_quantize with RHT-enabled NVFP4 quantizers dispatches unconditionally to the grouped Hadamard transform kernels, which are SM100 only, so GroupedLinear under the NVFP4 recipe fails at runtime on sm_120/sm_121. Apply the same architecture band the single-tensor path uses and fall back to per-tensor quantize outside it. Signed-off-by: David Kogan --- .../nvfp4/test_nvfp4_group_quantize.py | 22 +++++++++++++++++++ .../pytorch/csrc/extensions/cast.cpp | 20 +++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index 148ad83772..f89ba85e3f 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -252,3 +252,25 @@ def test_rht_with_quantization_block_tiling_versus_reference( with_random_sign_mask=with_random_sign_mask, optimize_for_gemm=optimize_for_gemm, ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize("quantize_mode", ["rowwise_only", "both_directions", "columnwise_only"]) +def test_rht_split_quantize_matches_per_tensor_reference(quantize_mode: str) -> None: + # split_quantize sends RHT-enabled NVFP4 quantizers to the grouped Hadamard + # transform kernels, which are implemented for the SM100 family only. On + # other architectures it falls back to quantizing each split on its own. + # Both routes have to give the same result as per-tensor quantization. + split_sections = [128, 128, 128, 128] + return_rowwise = quantize_mode in ("rowwise_only", "both_directions") + return_transpose = quantize_mode in ("columnwise_only", "both_directions") + + check_group_quantization_nvfp4_versus_reference( + x_dtype=torch.bfloat16, + M=sum(split_sections), + N=256, + return_rowwise=return_rowwise, + return_transpose=return_transpose, + split_sections=split_sections, + with_rht=True, + ) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 8b1cd384aa..60391b6bf0 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -19,6 +19,7 @@ #include "../extensions.h" #include "common.h" #include "common/common.h" +#include "common/util/cuda_runtime.h" #include "common/util/system.h" #include "pybind.h" #include "transformer_engine/multi_tensor.h" @@ -1631,6 +1632,25 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, // CUDA stream auto stream = at::cuda::getCurrentCUDAStream(); + // The RHT helper dispatches to grouped Hadamard transform kernels that are + // implemented for the SM100 family only. On other architectures, where + // NVFP4Quantizer::is_eligible_for_rht_cast_fusion is false as well, quantize + // each split on its own instead. That takes the generic unfused RHT path. + const int sm = transformer_engine::cuda::sm_arch(); + const bool grouped_rht_supported = sm >= 100 && sm <= 110; + if (quantizer.with_rht && !grouped_rht_supported) { + NVTE_CHECK(input.dtype() == DType::kBFloat16, "RHT is only supported for bfloat16 input"); + NVTE_SCOPED_GIL_RELEASE({ + for (size_t i = 0; i < num_tensors; ++i) { + if (input_list[i].numel() == 0) { + continue; + } + quantizers[i]->quantize(input_list[i], output_list[i], std::nullopt); + } + }); + return; + } + // Perform multi-tensor quantization NVTE_SCOPED_GIL_RELEASE({ if (quantizer.with_rht) { // Quantize row-wise data, RHT+quantize column-wise data From 5fb5bfecf51736e80207d5c3c9caff93c484182e Mon Sep 17 00:00:00 2001 From: David Kogan Date: Tue, 28 Jul 2026 10:01:18 -0400 Subject: [PATCH 2/2] Move the fallback inside the existing GIL release block Folds the architecture choice into the with_rht branch of the existing NVTE_SCOPED_GIL_RELEASE block instead of adding a second scope and an early return. Drops the duplicated bfloat16 check. Signed-off-by: David Kogan --- .../pytorch/csrc/extensions/cast.cpp | 31 +++++++++---------- 1 file changed, 14 insertions(+), 17 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 60391b6bf0..882481e0f4 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -1632,33 +1632,30 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, // CUDA stream auto stream = at::cuda::getCurrentCUDAStream(); - // The RHT helper dispatches to grouped Hadamard transform kernels that are - // implemented for the SM100 family only. On other architectures, where + // The grouped Hadamard transform kernels are implemented for the SM100 family + // only. On other architectures, where // NVFP4Quantizer::is_eligible_for_rht_cast_fusion is false as well, quantize // each split on its own instead. That takes the generic unfused RHT path. const int sm = transformer_engine::cuda::sm_arch(); const bool grouped_rht_supported = sm >= 100 && sm <= 110; - if (quantizer.with_rht && !grouped_rht_supported) { - NVTE_CHECK(input.dtype() == DType::kBFloat16, "RHT is only supported for bfloat16 input"); - NVTE_SCOPED_GIL_RELEASE({ - for (size_t i = 0; i < num_tensors; ++i) { - if (input_list[i].numel() == 0) { - continue; - } - quantizers[i]->quantize(input_list[i], output_list[i], std::nullopt); - } - }); - return; - } // Perform multi-tensor quantization NVTE_SCOPED_GIL_RELEASE({ if (quantizer.with_rht) { // Quantize row-wise data, RHT+quantize column-wise data // Check that config is supported NVTE_CHECK(input.dtype() == DType::kBFloat16, "RHT is only supported for bfloat16 input"); - // Fuse the rowwise and colwise into one when the kernel is ready - split_quantize_nvfp4_impl_with_rht_helper(input, input_list, output_list, split_sections, - quantizers, stream); + if (grouped_rht_supported) { + // Fuse the rowwise and colwise into one when the kernel is ready + split_quantize_nvfp4_impl_with_rht_helper(input, input_list, output_list, split_sections, + quantizers, stream); + } else { + for (size_t i = 0; i < num_tensors; ++i) { + if (input_list[i].numel() == 0) { + continue; + } + quantizers[i]->quantize(input_list[i], output_list[i], std::nullopt); + } + } } else { // NVFP4 quantize // Fuse the rowwise and colwise into one when the kernel is ready split_quantize_nvfp4_impl_helper(input, input_list, output_list, split_sections, quantizers,