[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 SM100-family architecture gate around grouped NVFP4 RHT quantization.

• Uses per-tensor unfused RHT quantization when grouped kernels are unavailable.
• Adds numerical coverage for rowwise, columnwise, and bidirectional split quantization.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/csrc/extensions/cast.cpp	Gates grouped RHT dispatch by GPU architecture and falls back to ordinary per-tensor quantization on unsupported devices.
tests/pytorch/nvfp4/test_nvfp4_group_quantize.py	Adds an independent per-tensor numerical reference test for all supported output-direction modes.

Flowchart

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

Loading

_{Reviews (2): Last reviewed commit: "Move the fallback inside the existing GI..." | Re-trigger Greptile}

[davidkny22]

Moved it inside the existing NVTE_SCOPED_GIL_RELEASE block. That also drops the duplicated bfloat16 check and the early return, so the diff is smaller.

Rebuilt on sm_121a: the added test passes, test_nvfp4_group_quantize.py is unchanged at 120 failed and 453 passed, and the reproducer still passes.

The failing CI jobs look unrelated to this change. They all stop at fatbinary fatal: Could not open input file 'flash_attn.compute_75.cubin' in common/fused_attn, and the same failure is on the other open PRs including ones that only touch common.

[ptrendx]

/te-ci pytorch