[gfx1250] Add MXFP4 GEMM support to FlyDSL backend

tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py

@@ -18,9 +18,14 @@
 import transformer_engine_torch as tex
 from transformer_engine.pytorch.tensor.mxfp4_tensor import MXFP4Quantizer
 from transformer_engine.pytorch.custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef
+from transformer_engine.pytorch.utils import get_device_compute_capability
 
 
 recipe_available, reason_for_no_recipe = te.is_mxfp4_available(return_reason=True)
+requires_gfx950 = pytest.mark.skipif(
+    not recipe_available or get_device_compute_capability() != (9, 5),
+    reason="AITER and hipBLASLt MXFP4 GEMMs require gfx950",
+)
 
 try:
     import aiter  # noqa: F401
@@ -122,6 +127,7 @@ def check_mxfp4_gemm_versus_reference(
 
 # AITER a4w4 backend: shuffled FP4 data + swizzled scales, TN only. Requires the aiter package.
 @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@requires_gfx950
 @pytest.mark.skipif(not _aiter_available, reason="aiter package not available")
 @pytest.mark.parametrize("M, K, N", _SHAPES)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16], ids=str)
@@ -146,6 +152,7 @@ def test_mxfp4_gemm_aiter_versus_reference(
 
 # hipBLASLt F4F4 backend: plain FP4 data, plain or pre-swizzled UE8M0 scales
 @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@requires_gfx950
 @pytest.mark.parametrize("M, K, N", _SHAPES)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16], ids=str)
 @pytest.mark.parametrize("layout", ["TN", "NN", "NT", "TT"])

tests/pytorch/mxfp4/test_mxfp4_module_exact.py

@@ -19,6 +19,7 @@
 )
 from transformer_engine.pytorch.custom_recipes.quantization_mxfp4 import MXFP4QuantizerRef
 from transformer_engine.pytorch.quantization import FP8GlobalStateManager
+from transformer_engine.pytorch.utils import get_device_compute_capability
 
 
 recipe_available, reason_for_no_recipe = te.is_mxfp4_available(return_reason=True)
@@ -30,6 +31,16 @@
 except ImportError:
     _aiter_available = False
 
+_is_gfx1250 = recipe_available and get_device_compute_capability() == (12, 5)
+requires_mxfp4_gemm_backend = pytest.mark.skipif(
+    not (_is_gfx1250 or _aiter_available), reason="aiter package not available"
+)
+
+
+def skip_if_unsupported_shape(in_features: int, out_features: int):
+    if _is_gfx1250 and (in_features % 256 or out_features % 256):
+        pytest.skip("FlyDSL MXFP4 GEMM needs in_features and out_features divisible by 256")
+
 
 class GetRecipes:
     @staticmethod
@@ -244,7 +255,7 @@ def check_mxfp4_module_versus_reference(
 
 
 @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
-@pytest.mark.skipif(not _aiter_available, reason="aiter package not available")
+@requires_mxfp4_gemm_backend
 @pytest.mark.parametrize(
     "in_features, out_features",
     [
@@ -270,6 +281,7 @@ def test_mxfp4_linear_versus_reference(
     """Test MXFP4 Linear module against reference implementation."""
     if x_dtype != torch.bfloat16:
         pytest.skip("MXFP4 quantization is only supported for bfloat16 input")
+    skip_if_unsupported_shape(in_features, out_features)
 
     check_mxfp4_module_versus_reference(
         module_class=te.Linear,
@@ -412,12 +424,13 @@ def check_mxfp4_layernorm_linear_versus_reference(
 
 
 @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
-@pytest.mark.skipif(not _aiter_available, reason="aiter package not available")
+@requires_mxfp4_gemm_backend
 @pytest.mark.parametrize(
     "in_features, out_features",
     [
         (128, 256),
         (256, 128),
+        (512, 512),
     ],
 )
 @pytest.mark.parametrize("bias", [False], ids=["no_bias"])
@@ -434,6 +447,7 @@ def test_mxfp4_layernorm_linear_versus_reference(
     num_steps: int,
     use_hadamard: bool,
 ):
+    skip_if_unsupported_shape(in_features, out_features)
     check_mxfp4_layernorm_linear_versus_reference(
         in_features=in_features,
         out_features=out_features,

tests/pytorch/test_gemm_backends.py

@@ -14,7 +14,8 @@
 - The ``test_triton_*`` family runs when ``pytorch-triton-rocm`` is importable
   (the Triton GEMM path has no arch gate, so it runs on gfx942 and gfx950).
 - The ``test_flydsl_*`` family runs when the version-gated ``flydsl`` package is
-  importable and the device is gfx950 (the only arch the FlyDSL path supports).
+  importable and the device is gfx950; the ``test_flydsl_*_mxfp4`` tests run on
+  gfx1250 instead (the only arch the FlyDSL MXFP4 path currently supports).
 
 Triton coverage: fp32 / fp16 / bf16 / same-format FP8 / mixed FP8 (skipped for a
 compiler bug) / MXFP8 across TN / NN / NT, plus bias / bias-grad epilogues and a
@@ -32,7 +33,8 @@
 config: fp32 / fp16 / bf16, same- and mixed-format tensor-wise FP8, same- and
 mixed-format MXFP8, TN / NN / NT, batched multidim FP8, the fused forward BIAS
 epilogue, and the fused GELU_AUX / GELU_AUX_BIAS epilogues. BGRADB and DGELU are
-not implemented on any FlyDSL path yet. FlyDSL kernels require tile-aligned
+not implemented on any FlyDSL path yet. MXFP4 (gfx1250) covers TN / NN / NT
+with FP32 / BF16 / FP16 output against a torch.matmul reference only. FlyDSL kernels require tile-aligned
 launch dimensions, so its shapes are aligned to the 256x256x128 kernel contract
 rather than reusing the odd-sized Triton edge-mask cases.
 """
@@ -47,6 +49,7 @@
 from transformer_engine.pytorch import Float8Tensor, is_mxfp8_available
 from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm
 from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer
+from transformer_engine.pytorch.tensor.mxfp4_tensor import MXFP4Quantizer
 from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor
 from transformer_engine.pytorch import torch_version
 from transformer_engine.pytorch.utils import get_gemm_backend
@@ -93,7 +96,14 @@ def _device_capability():
     reason="FlyDSL GEMM tests require the flydsl package on gfx950",
 )
 
-if _flydsl_available:
+# FlyDSL MXFP4 is gfx1250-only (matching the C++ dispatch).
+_flydsl_mxfp4_available = _CAP == (12, 5) and importlib.util.find_spec("flydsl") is not None
+requires_flydsl_mxfp4_backend = pytest.mark.skipif(
+    not _flydsl_mxfp4_available,
+    reason="FlyDSL MXFP4 GEMM tests require the flydsl package on gfx1250",
+)
+
+if _flydsl_available or _flydsl_mxfp4_available:
     pytest.importorskip("flydsl", reason="FlyDSL package is not installed")
 
 # gfx942's Triton fp32 matmul has a stable numerical divergence from
@@ -162,6 +172,15 @@ def _device_capability():
     (512, 512, 1024),  # M != N
 ]
 
+FLYDSL_MXFP4_SHAPES = [
+    (512, 512, 512),
+    (512, 1024, 512),
+    (512, 512, 1024),  # M != N
+    (512, 640, 512),  # odd K tile count
+    (512, 4096, 512),  # K > 2048: scales streamed in the K loop
+    (512, 3200, 512),  # streamed scales with a partial last scale group
+]
+
 LAYOUTS = ["TN", "NN", "NT"]
 
 # Same-format FP8 combos, shared by both backends.
@@ -334,6 +353,36 @@ def create_mxfp8_tensors(
     return A_mxfp8, B_mxfp8, A_mxfp8.dequantize(), B_mxfp8.dequantize()
 
 
+_E2M1_MAGNITUDES = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0]
+
+
+def _dequantize_mxfp4(data, scale):
+    """Dequantize packed FP4 ``data`` [rows, K/2] with E8M0 ``scale`` (padded) to FP32."""
+    data = data.view(torch.uint8)
+    rows, k = data.shape[0], data.shape[1] * 2
+    codes = torch.stack([data & 0xF, data >> 4], dim=-1).reshape(rows, k).long()
+    values = torch.tensor(_E2M1_MAGNITUDES, device=data.device)[codes & 7]
+    values = torch.where(codes & 8 != 0, -values, values)
+    scales = torch.exp2(scale[:rows, : k // 32].float() - 127.0)
+    return values * scales.repeat_interleave(32, dim=1)
+
+
+def create_mxfp4_tensors(M, K, N, layout):
+    """Create MXFP4 inputs and references dequantized from the buffers ``layout`` consumes."""
+    # the MXFP4 quantizer takes BF16 input
+    A_bf16, B_bf16 = create_regular_tensors(M, K, N, layout, dtype=torch.bfloat16)
+    quantizer = MXFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1, rowwise=True, columnwise=True)
+    A_mxfp4, B_mxfp4 = quantizer(A_bf16), quantizer(B_bf16)
+
+    def reference(t, rowwise):
+        if rowwise:
+            return _dequantize_mxfp4(t._rowwise_data, t._rowwise_scale_inv)
+        return _dequantize_mxfp4(t._columnwise_data, t._columnwise_scale_inv).T
+
+    transa, transb = layout[0] == "T", layout[1] == "T"
+    return A_mxfp4, B_mxfp4, reference(A_mxfp4, transa), reference(B_mxfp4, not transb)
+
+
 # ==============================================================================
 # Triton backend (NVTE_GEMM_BACKEND=TRITON)
 # ==============================================================================
@@ -1094,6 +1143,46 @@ def test_flydsl_mxfp8_unsupported_shape_falls_back():
     assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2)
 
 
+@requires_flydsl_mxfp4_backend
+@pytest.mark.parametrize("M, K, N", FLYDSL_MXFP4_SHAPES)
+@pytest.mark.parametrize("layout", LAYOUTS)
+@pytest.mark.parametrize(
+    "out_dtype", [torch.float32, torch.bfloat16, torch.float16], ids=["fp32", "bf16", "fp16"]
+)
+def test_flydsl_vs_pytorch_mxfp4(M, K, N, layout, out_dtype):
+    """Test MXFP4 FlyDSL GEMMs (gfx1250) against an FP32 PyTorch reference."""
+    torch.manual_seed(42)
+
+    A_mxfp4, B_mxfp4, A_deq, B_deq = create_mxfp4_tensors(M, K, N, layout)
+
+    output = flydsl_call_gemm(A_mxfp4, B_mxfp4, layout, out_dtype=out_dtype, use_flydsl=True)
+    expected = compute_pytorch_reference(A_deq, B_deq, layout)
+
+    assert output.dtype == out_dtype
+    if out_dtype == torch.float32:
+        assert_gemm_close(output, expected, atol=1e-3, rtol=1e-4)
+    else:
+        assert_gemm_close(output, expected, atol=5e-3, rtol=1e-2)
+
+
+@requires_flydsl_mxfp4_backend
+def test_flydsl_mxfp4_unsupported_shape_raises():
+    """MXFP4 configs FlyDSL cannot run raise FlyDSLUnsupportedError.
+
+    ``general_gemm`` raises it, as gfx1250 has no other MXFP4 GEMM.
+    """
+    from transformer_engine.pytorch.flydsl_kernels.gemm import (
+        FlyDSLUnsupportedError,
+        te_generic_gemm_flydsl,
+    )
+
+    torch.manual_seed(42)
+    # M=128 is not a multiple of the 256-wide kernel tile
+    A_mxfp4, B_mxfp4, _, _ = create_mxfp4_tensors(128, 512, 256, "TN")
+    with pytest.raises(FlyDSLUnsupportedError):
+        te_generic_gemm_flydsl(A_mxfp4, True, B_mxfp4, False, None, None, tex.DType.kBFloat16)
+
+
 # --- Approach 2: FlyDSL vs native C++ ``generic_gemm`` reference --------------
 
 

transformer_engine/pytorch/cpp_extensions/gemm.py

@@ -437,6 +437,56 @@ def _validate_native_gemm_output_quantizer(quantization_params):
 _flydsl_import_warned = False
 
 
+def _try_flydsl_gemm(*args, **kwargs):
+    """Run te_generic_gemm_flydsl; return None if the caller should fall back."""
+    try:
+        # Lazy import keeps FlyDSL off the normal Transformer Engine
+        # import path. It is done inside the try so a wheel built without
+        # flydsl (NVTE_GEMM_BACKEND!=FLYDSL at build time) degrades to the
+        # default backend instead of raising a bare ImportError.
+        from ..flydsl_kernels.gemm import (
+            FlyDSLUnsupportedError,
+            te_generic_gemm_flydsl,
+        )
+
+        return te_generic_gemm_flydsl(
+            *args,
+            **kwargs,
+        )
+    except ImportError as exc:
+        # NVTE_GEMM_BACKEND=FLYDSL was requested but the flydsl package is
+        # missing or too old (see flydsl_kernels.gemm._MIN_FLYDSL). This
+        # is a misconfiguration, not an unsupported GEMM config, so always
+        # warn (once) regardless of the opt-in fallback flag before
+        # degrading to the default backend.
+        global _flydsl_import_warned
+        if not _flydsl_import_warned:
+            _flydsl_import_warned = True
+            warnings.warn(
+                "[FLYDSL WARNING]: NVTE_GEMM_BACKEND=FLYDSL but the flydsl "
+                "package is unavailable; falling back to the default backend. "
+                f"Install a supported version (e.g. `pip install flydsl`) to "
+                f"enable it. Reason: {exc}",
+                UserWarning,
+                stacklevel=3,
+            )
+        return None
+    except FlyDSLUnsupportedError as exc:
+        warn_fallback = os.environ.get(
+            "NVTE_FLYDSL_GEMM_WARN_FALLBACK",
+            "0",
+        ).lower() not in ("", "0", "false", "no", "off")
+
+        if warn_fallback:
+            warnings.warn(
+                "[FLYDSL WARNING]: FlyDSL GEMM does not support this configuration; "
+                f"falling back to the default backend. Reason: {exc}",
+                UserWarning,
+                stacklevel=3,
+            )
+        return None
+
+
 def general_gemm(
     A: torch.Tensor,
     B: torch.Tensor,
@@ -540,9 +590,39 @@ def general_gemm(
     # Use bfloat16 as default bias_dtype
     bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype]
 
-    # MXFP4 GEMM: route to AITER a4w4 ASM kernels, unless the hipBLASLt backend is
-    # opted in via NVTE_ROCM_USE_HIPBLASLT_MXFP4
+    # MXFP4 GEMM: FlyDSL on gfx1250 (no fallback), else route to AITER a4w4 ASM kernels,
+    # unless the hipBLASLt backend is opted in via NVTE_ROCM_USE_HIPBLASLT_MXFP4
     if isinstance(A, MXFP4TensorStorage) or isinstance(B, MXFP4TensorStorage):
+        if IS_HIP_EXTENSION and get_device_compute_capability() == (12, 5):
+            from ..flydsl_kernels.gemm import te_generic_gemm_flydsl
+
+            out, bias_grad, gelu_input, extra_output = te_generic_gemm_flydsl(
+                A,
+                transa,
+                B,
+                transb,
+                out,
+                quantization_params,
+                TE_DType[out_dtype] if out_dtype is not None else None,
+                bias,
+                bias_dtype,
+                gelu,
+                gelu_in,
+                grad,
+                workspace,
+                workspace.shape[0],
+                accumulate,
+                use_split_accumulator,
+                comm_overlap=ub,
+                comm_type=ub_type,
+                extra_output=extra_output,
+                bulk_overlap=bulk_overlap,
+                alpha=alpha,
+                beta=beta,
+            )
+            if debug_quantizer is not None:
+                out = debug_quantizer.process_gemm_output(out)
+            return out, bias_grad, gelu_input, extra_output
         if not bool(int(os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0"))):
             result = mxfp4_gemm(
                 A,
@@ -625,55 +705,9 @@ def general_gemm(
             and get_device_compute_capability() == (9, 5)
             and gemm_backend == "flydsl"
         )
-        if use_gemm_flydsl:
-            try:
-                # Lazy import keeps FlyDSL off the normal Transformer Engine
-                # import path. It is done inside the try so a wheel built without
-                # flydsl (NVTE_GEMM_BACKEND!=FLYDSL at build time) degrades to the
-                # default backend instead of raising a bare ImportError.
-                from ..flydsl_kernels.gemm import (
-                    FlyDSLUnsupportedError,
-                    te_generic_gemm_flydsl,
-                )
-
-                out, bias_grad, gelu_input, extra_output = te_generic_gemm_flydsl(
-                    *args,
-                    **kwargs,
-                )
-            except ImportError as exc:
-                # NVTE_GEMM_BACKEND=FLYDSL was requested but the flydsl package is
-                # missing or too old (see flydsl_kernels.gemm._MIN_FLYDSL). This
-                # is a misconfiguration, not an unsupported GEMM config, so always
-                # warn (once) regardless of the opt-in fallback flag before
-                # degrading to the default backend.
-                global _flydsl_import_warned
-                if not _flydsl_import_warned:
-                    _flydsl_import_warned = True
-                    warnings.warn(
-                        "[FLYDSL WARNING]: NVTE_GEMM_BACKEND=FLYDSL but the flydsl "
-                        "package is unavailable; falling back to the default backend. "
-                        f"Install a supported version (e.g. `pip install flydsl`) to "
-                        f"enable it. Reason: {exc}",
-                        UserWarning,
-                        stacklevel=2,
-                    )
-
-                out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs)
-            except FlyDSLUnsupportedError as exc:
-                warn_fallback = os.environ.get(
-                    "NVTE_FLYDSL_GEMM_WARN_FALLBACK",
-                    "0",
-                ).lower() not in ("", "0", "false", "no", "off")
-
-                if warn_fallback:
-                    warnings.warn(
-                        "[FLYDSL WARNING]: FlyDSL GEMM does not support this configuration; "
-                        f"falling back to the default backend. Reason: {exc}",
-                        UserWarning,
-                        stacklevel=2,
-                    )
-
-                out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs)
+        result = _try_flydsl_gemm(*args, **kwargs) if use_gemm_flydsl else None
+        if result is not None:
+            out, bias_grad, gelu_input, extra_output = result
         else:
             out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs)
     else: