GroupedLinear under the NVFP4 recipe fails on sm_120/sm_121. split_quantize with RHT-enabled NVFP4 quantizers dispatches unconditionally to the grouped Hadamard-transform kernels (nvte_group_hadamard_transform_*), which are implemented with SM100-family instructions (tcgen05/TMEM) and have no sm_120/121 implementation. There is no architecture check before the dispatch, so the call 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 already handles this: NVFP4Quantizer checks is_eligible_for_rht_cast_fusion (arch band 100..110) and falls back to the generic unfused RHT kernels, so plain Linear works on sm_120/121. Only the split/grouped entry point lacks the check.
Repro: script below. GB10 (DGX Spark), sm_121a, CUDA 13.2, driver 580.95.05, TE main @ 9d92fa0.
Proposed fix: add the same arch gate to split_quantize_nvfp4_impl and fall back to per-tensor quantize on non-SM100 archs. I have this working on sm_121 hardware and can open a PR.
test_split_quantize_rht_sm12x.py
# Reproducer: split_quantize with RHT-enabled NVFP4 quantizers on sm_120/121
# (the GroupedLinear input/grad-output quantization path).
# Before the fix: dispatches to grouped Hadamard-transform kernels implemented
# with SM100-family instructions -> runtime failure on sm_120/121.
# After the fix: falls back to per-tensor quantization (generic unfused RHT
# path); outputs are well-formed and rowwise dequant matches the input at the
# NVFP4 quantization floor.
import sys
import torch
import transformer_engine.pytorch # noqa: F401 (loads the torch extension)
import transformer_engine_torch as tex
from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
def main():
torch.manual_seed(0)
n_splits, rows_per_split, cols = 4, 128, 256
x = torch.randn(n_splits * rows_per_split, cols, dtype=torch.bfloat16, device="cuda")
quantizers = [
NVFP4Quantizer(rowwise=True, columnwise=True, with_rht=True, with_post_rht_amax=True)
for _ in range(n_splits)
]
outs = tex.split_quantize(x, [rows_per_split] * n_splits, quantizers)
torch.cuda.synchronize()
ok = True
for i, y in enumerate(outs):
d = y.dequantize(dtype=torch.float32)
ref = x[i * rows_per_split:(i + 1) * rows_per_split].float()
cos = torch.nn.functional.cosine_similarity(
d.flatten().unsqueeze(0), ref.flatten().unsqueeze(0)
).item()
nr = (d.norm() / (ref.norm() + 1e-30)).item()
print(f"split {i}: rowwise dequant cos {cos:.6f} norm-ratio {nr:.6f}")
ok &= cos > 0.98 and abs(nr - 1.0) < 0.05
print("PASS" if ok else "FAIL")
return 0 if ok else 1
if __name__ == "__main__":
try:
sys.exit(main())
except Exception as e:
print(f"ERROR: {type(e).__name__}: {e}")
sys.exit(2)
GroupedLinearunder the NVFP4 recipe fails on sm_120/sm_121.split_quantizewith RHT-enabled NVFP4 quantizers dispatches unconditionally to the grouped Hadamard-transform kernels (nvte_group_hadamard_transform_*), which are implemented with SM100-family instructions (tcgen05/TMEM) and have no sm_120/121 implementation. There is no architecture check before the dispatch, so the call fails:The single-tensor path already handles this:
NVFP4Quantizerchecksis_eligible_for_rht_cast_fusion(arch band 100..110) and falls back to the generic unfused RHT kernels, so plainLinearworks on sm_120/121. Only the split/grouped entry point lacks the check.Repro: script below. GB10 (DGX Spark), sm_121a, CUDA 13.2, driver 580.95.05, TE main @ 9d92fa0.
Proposed fix: add the same arch gate to
split_quantize_nvfp4_impland fall back to per-tensor quantize on non-SM100 archs. I have this working on sm_121 hardware and can open a PR.test_split_quantize_rht_sm12x.py