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

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions