[PyTorch] Add architecture gate to NVFP4 split_quantize RHT path

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,
+    )

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