Skip to content

GroupedLinear NVFP4 with RHT fails at runtime on sm_120/sm_121 #3219

Description

@davidkny22

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)

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions