[PyTorch] Enable e5m3 fused GEMM kernels from cuDNN

tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py

@@ -690,3 +690,112 @@ def test_nvfp4_row_scaled_gemm_matches_emulated(
         use_4over6=use_4over6,
         nvfp4_4over6_err_mode=nvfp4_4over6_err_mode,
     )
+
+
+def _check_ue5m3_gemm_versus_dequantized(
+    M, K, N, x_columnwise, w_columnwise, disable_second_level_scale
+):
+    """Run an NVFP4/UE5M3 GEMM and compare against a dequantized FP32 reference."""
+    if M % 256 != 0:
+        pytest.skip(
+            "cuDNN's grouped GEMM pads every group to 256 rows, so the UE5M3 path (which "
+            "routes there while cuBLAS lacks UE5M3 kernels) requires M % 256 == 0."
+        )
+    torch.manual_seed(0)
+    device, dtype, out_dtype = "cuda", torch.bfloat16, torch.bfloat16
+    x_shape = (K, M) if x_columnwise else (M, K)
+    w_shape = (K, N) if w_columnwise else (N, K)
+    x = torch.randn(x_shape, dtype=dtype, device=device)
+    w = torch.randn(w_shape, dtype=dtype, device=device)
+
+    common = dict(
+        fp4_dtype=tex.DType.kFloat4E2M1,
+        scale_dtype=tex.DType.kFloat8UE5M3,
+        rowwise=True,
+        columnwise=True,
+        with_amax_reduction=False,
+        amax_reduction_group=None,
+        with_rht=False,
+        with_post_rht_amax=False,
+    )
+    # disable_second_level_scale is given per operand, as (x, w).
+    xq = NVFP4Quantizer(**common, disable_second_level_scale=disable_second_level_scale[0])
+    wq = NVFP4Quantizer(**common, disable_second_level_scale=disable_second_level_scale[1])
+    x_q = xq.update_quantized(x, xq.make_empty(x_shape, dtype=dtype, device=device))
+    w_q = wq.update_quantized(w, wq.make_empty(w_shape, dtype=dtype, device=device))
+
+    if disable_second_level_scale[0]:
+        assert x_q._amax_rowwise is None, "disable_second_level_scale should drop the amax"
+    if disable_second_level_scale[1]:
+        assert w_q._amax_rowwise is None, "disable_second_level_scale should drop the amax"
+
+    # Reference: dequantize the orientation each operand is actually read in.
+    x_ref = _dequantize_nvfp4_usage(x_q, columnwise=x_columnwise)
+    w_ref = _dequantize_nvfp4_usage(w_q, columnwise=w_columnwise)
+    # _dequantize_nvfp4_usage returns each operand canonically as (rows, K), so
+    # the reference is the same expression for every layout.
+    ref = x_ref @ w_ref.t()
+
+    if x_columnwise:
+        x_q.update_usage(rowwise_usage=False)
+    if w_columnwise:
+        w_q.update_usage(rowwise_usage=False)
+    transa, transb = not w_columnwise, x_columnwise
+    layout = ("T" if transa else "N") + ("T" if transb else "N")
+    y = general_gemm(w_q, x_q, out_dtype=out_dtype, layout=layout)[0]
+
+    # Both sides see identically quantized operands, so quantization error cancels and
+    # only accumulation order and the bf16 output rounding differ. One bf16 ulp is
+    # already ~4e-3 relative, which no elementwise tolerance survives, so compare the
+    # whole result instead.
+    rel_err = (y.float() - ref).norm() / ref.norm()
+    assert rel_err < 5e-3, f"relative error {rel_err:.2e} is too large"
+
+ue5m3_available, reason_for_no_ue5m3 = te.is_fp8_ue5m3_available(return_reason=True)
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.skipif(not ue5m3_available, reason=reason_for_no_ue5m3)
+@pytest.mark.parametrize(
+    "M, K, N",
+    [
+        (256, 128, 256),
+        (256, 256, 256),
+        (256, 1024, 256),
+        (1024, 1024, 1024),
+        (4096, 512, 3072),
+        (112, 128, 96),
+        (304, 640, 304),
+        (1008, 3072, 992),
+        (256, 64, 256),
+        (128, 128, 112),
+    ],
+)
+@pytest.mark.parametrize(
+    "x_columnwise, w_columnwise",
+    [
+        (False, False),  # TN -- w rowwise, x rowwise   (fprop)
+        (False, True),  # NN -- w colwise, x rowwise   (dgrad)
+        (True, True),  # NT -- w colwise, x colwise   (wgrad)
+    ], ids=["FF", "FT", "TT"]
+)
+@pytest.mark.parametrize(
+    "disable_second_level_scale", [
+        (True, False),
+    ], ids=["TF"]
+)
+def test_nvfp4_ue5m3_gemm_versus_reference(
+    M: int,
+    K: int,
+    N: int,
+    x_columnwise: bool,
+    w_columnwise: bool,
+    disable_second_level_scale: bool,
+):
+    """NVFP4 GEMM with UE5M3 block scales, with and without second-level scaling.
+
+    UE5M3's wider range is what makes dropping the per-tensor global scale
+    viable, so both configurations must match the dequantized reference.
+    """
+    _check_ue5m3_gemm_versus_dequantized(
+        M, K, N, x_columnwise, w_columnwise, disable_second_level_scale
+    )

tests/pytorch/test_fusible_ops.py

@@ -82,6 +82,8 @@
 if nvfp4_available:
     _quantization_list.append("nvfp4")
     _quantization_list.append("nvfp4_4over6")
+    if fp8_ue5m3_available:
+        _quantization_list.append("nvfp4_rht_ue5m3")
 if fp8_block_scaling_available:
     _quantization_list.append("fp8_block_scaling")
 
@@ -136,6 +138,11 @@ def maybe_skip_quantization(
         elif quantization in nvfp4_variant_names:
             if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0:
                 pytest.skip("NVFP4 GEMMs require dims that are divisible by 16")
+            if (
+                quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3")
+                and (math.prod(dims[:-1]) % 64 != 0 or dims[-1] % 64 != 0)
+            ):
+                pytest.skip("cuDNN FE NVFP4-UE5M3 GEMMs produce incorrect values with 32x32 tensors")
 
     # Check dtype
     if dtype is not None:
@@ -3588,7 +3595,12 @@ def test_grouped_mlp(
 
         # Skip invalid configurations
         with_quantization = quantization is not None
-        maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype)
+        maybe_skip_quantization(
+            quantization,
+            dims=in_shape,
+            device=device,
+            dtype=dtype,
+        )
         if with_quantization and dtype not in (torch.bfloat16, torch.float16):
             pytest.skip("Quantized group GEMM is only supported with BF16/FP16")
         if activation == "scaled_srelu" and quantization == "nvfp4_rht" and bias:

tests/pytorch/test_grouped_mlp.py

@@ -43,6 +43,7 @@
     dtype_tols,
     make_recipe,
     MegatronTrainingHelper,
+    nvfp4_variant_names,
     quantization_tols,
     reset_rng_states,
 )
@@ -51,6 +52,7 @@
 fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True)
 mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True)
 nvfp4_available, reason_for_no_nvfp4 = te.is_nvfp4_available(return_reason=True)
+fp8_ue5m3_available, reason_for_no_fp8_ue5m3 = te.is_fp8_ue5m3_available(return_reason=True)
 
 # Supported data types
 _dtypes: list[torch.dtype] = [torch.float32, torch.float16]
@@ -73,6 +75,8 @@
     _grouped_mlp_quantization_list.append("mxfp8")
 if nvfp4_available:
     _grouped_mlp_quantization_list.append("nvfp4_rht")
+    if fp8_ue5m3_available:
+        _grouped_mlp_quantization_list.append("nvfp4_rht_ue5m3")
 
 
 @pytest.fixture(autouse=True, scope="function")
@@ -102,11 +106,10 @@ def maybe_skip_quantization(
         pytest.skip(reason_for_no_fp8)
     if quantization == "mxfp8" and not mxfp8_available:
         pytest.skip(reason_for_no_mxfp8)
-    if (
-        quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht")
-        and not nvfp4_available
-    ):
+    if quantization in nvfp4_variant_names and not nvfp4_available:
         pytest.skip(reason_for_no_nvfp4)
+    if quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3") and not fp8_ue5m3_available:
+        pytest.skip(reason_for_no_fp8_ue5m3)
 
     # Check dims
     if dims is not None:
@@ -118,16 +121,18 @@ def maybe_skip_quantization(
         elif quantization == "mxfp8":
             if math.prod(dims[:-1]) % 32 != 0 or dims[-1] % 32 != 0:
                 pytest.skip("MXFP8 GEMMs require dims that are divisible by 32")
-        elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"):
+        elif quantization in nvfp4_variant_names:
             if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0:
                 pytest.skip("NVFP4 GEMMs require dims that are divisible by 16")
+            if (
+                quantization in ("nvfp4_ue5m3", "nvfp4_rht_ue5m3")
+                and (math.prod(dims[:-1]) % 64 != 0 or dims[-1] % 64 != 0)
+            ):
+                pytest.skip("cuDNN FE NVFP4-UE5M3 GEMMs produce incorrect values with 32x32 tensors")
 
     # Check dtype
     if dtype is not None:
-        if (
-            quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht")
-            and dtype != torch.bfloat16
-        ):
+        if quantization in nvfp4_variant_names and dtype != torch.bfloat16:
             pytest.skip("NVFP4 quantization is only supported with BF16 data")
 
 
@@ -183,17 +188,27 @@ def make_reference_and_test_tensors(
         test = quantizer(test)
     elif quantization == "mxfp8":
         test = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)(test)
-    elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht"):
+    elif quantization in (
+        "nvfp4",
+        "nvfp4_row_scaled",
+        "nvfp4_rht",
+        "nvfp4_ue5m3",
+        "nvfp4_rht_ue5m3",
+    ):
         tensor_type = "input"
         if quantizer_role is not None:
             tensor_type = quantizer_role.tensor_type
-        with_rht = quantization == "nvfp4_rht" and tensor_type != "weight"
+        with_rht = quantization in ("nvfp4_rht", "nvfp4_rht_ue5m3") and tensor_type != "weight"
+        scale_dtype = (
+            te.DType.kFloat8UE5M3 if quantization == "nvfp4_rht_ue5m3" else te.DType.kFloat8E4M3
+        )
         test = NVFP4Quantizer(
+            scale_dtype=scale_dtype,
             with_rht=with_rht,
             with_post_rht_amax=with_rht,
             with_2d_quantization=False,
             stochastic_rounding=False,
-            with_random_sign_mask=False,
+            with_random_sign_mask=with_rht,
         )(test)
     elif quantization == "nvfp4_4over6":
         tensor_type = "input"