[PyTorch] Add architecture gate to NVFP4 split_quantize RHT path

[davidkny22]

Description

GroupedLinear under the NVFP4 recipe fails at runtime on sm_120/sm_121. split_quantize with RHT-enabled quantizers goes straight to split_quantize_nvfp4_impl_with_rht_helper, which calls the grouped Hadamard transform kernels. Those are SM100 only, so the launch fails:

RuntimeError: .../group_row_cast_col_hadamard_transform_cast_fusion.cu:1276 in function group_row_col_rht_gemm_ntt_w_sfc: CUDA Error: invalid argument

The single-tensor path does not have this problem. NVFP4Quantizer checks is_eligible_for_rht_cast_fusion, which gates on 100 <= sm_arch <= 110, and uses the unfused RHT kernels otherwise. This adds the same check to split_quantize_nvfp4_impl and quantizes the splits per tensor when it fails. That is slower than a fused grouped kernel would be, but it makes the configuration work.

I repeated the band instead of calling is_eligible_for_rht_cast_fusion because that function also checks shape alignment, and reusing it would change SM100 behavior. Happy to factor the band into a shared helper if you prefer.

Fixes #3219

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• transformer_engine/pytorch/csrc/extensions/cast.cpp: architecture gate and per-tensor fallback in split_quantize_nvfp4_impl.
• tests/pytorch/nvfp4/test_nvfp4_group_quantize.py: test that RHT split_quantize matches per-tensor quantization, on either dispatch route.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

No documentation changes needed. On the last box, see below.

Testing done

GB10 (DGX Spark), sm_121a, CUDA 13.2, driver 580.95.05, main @ f1d5f8d built with NVTE_CUDA_ARCHS=121a. cpplint, black and clang-format are clean on the changed files.

The new test fails on all three parametrizations without the fix and passes with it. It is not gated on architecture, so it exercises the grouped kernels on SM100 and the fallback elsewhere.

The rest of tests/pytorch/nvfp4/test_nvfp4_group_quantize.py goes from 274 failures to 120. The 154 that start passing are the RHT cases with optimize_for_gemm=False, which is what this patch covers.

The 120 that remain fail before the patch too, for a different reason. They are all optimize_for_gemm=True: swizzled scale factor emission is gated on the same architecture band, so the outputs here carry compact scale factors, but the test swizzles the reference whenever optimize_for_gemm is set.

tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py is 16 failures and 284 passes both before and after, so the single-tensor path is unaffected. Those 16 are small numeric mismatches already on main.

The reproducer in #3219 fails before the patch and passes after: per-split rowwise dequant cos 0.995463 to 0.995507, norm ratio 0.999135 to 1.001009.

cc @zhongbozhu @cael-ling

[greptile-apps]

Greptile Summary

Adds an architecture-gated fallback for NVFP4 split quantization with RHT.

• Uses grouped Hadamard kernels on SM100–SM110.
• Quantizes each split independently on other architectures.
• Adds coverage comparing split and per-tensor quantization for rowwise, columnwise, and bidirectional modes.

Confidence Score: 5/5

The PR appears safe to merge, with the fallback correctly populating each preallocated split through the existing per-tensor NVFP4 quantization path.

The architecture check matches the existing fused-RHT eligibility band, while unsupported architectures avoid the grouped kernel and retain the expected output storage and metadata semantics.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/extensions/cast.cpp	Adds a hardware gate that preserves grouped RHT dispatch on supported architectures and uses the existing per-tensor quantization path elsewhere.
tests/pytorch/nvfp4/test_nvfp4_group_quantize.py	Adds architecture-independent coverage for all rowwise and columnwise output combinations of RHT split quantization.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
  A["NVFP4 split_quantize with RHT"] --> B{"SM architecture 100–110?"}
  B -->|Yes| C["Grouped RHT quantization"]
  B -->|No| D["Quantize each split independently"]
  C --> E["Return split NVFP4 tensors"]
  D --> E

Loading

_{Reviews (1): Last reviewed commit: "Add architecture gate to the NVFP4 split..." | Re-trigger Greptile}