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..882481e0f4 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,14 +1632,30 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, // CUDA stream auto stream = at::cuda::getCurrentCUDAStream(); + // 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; + // 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,