From bdac43a876d99a4ca91d97c93b3766ebdaef9752 Mon Sep 17 00:00:00 2001
From: Aristotle Martin <Aristotle.Martin@amd.com>
Date: Fri, 30 Oct 2026 14:40:26 +0000
Subject: [PATCH 01/11] Add gfx1250 MXFP4 GEMM support to FlyDSL backend

---
 tests/pytorch/test_gemm_backends.py           | 107 ++-
 .../pytorch/cpp_extensions/gemm.py            | 140 ++--
 .../pytorch/flydsl_kernels/__init__.py        |   2 +-
 .../pytorch/flydsl_kernels/gemm/__init__.py   |   2 +-
 .../flydsl_kernels/gemm/gemm_wrappers.py      | 144 +++-
 .../pytorch/flydsl_kernels/gemm/mxfp4_gemm.py | 619 ++++++++++++++++++
 6 files changed, 952 insertions(+), 62 deletions(-)
 create mode 100644 transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py

diff --git a/tests/pytorch/test_gemm_backends.py b/tests/pytorch/test_gemm_backends.py
index fedfdea54b..e85ab201f3 100644
--- a/tests/pytorch/test_gemm_backends.py
+++ b/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 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,9 @@
 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 / TT
+with FP32 / BF16 / FP16 output against a torch.matmul reference only: the native
+C++ MXFP4 GEMM does not run on gfx1250. 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 +50,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 +97,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,8 +173,20 @@ 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
+]
+
 LAYOUTS = ["TN", "NN", "NT"]
 
+# MXFP4 also supports TT (every layout runs the same TN kernel).
+MXFP4_LAYOUTS = ["TN", "NN", "NT", "TT"]
+
+MXFP4_OUT_DTYPES = [torch.float32, torch.bfloat16, torch.float16]
+
 # Same-format FP8 combos, shared by both backends.
 FP8_FORMAT_COMBOS = [
     (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3),
@@ -255,6 +278,8 @@ def get_shapes(layout, M, K, N):
         return (K, M), (N, K)
     if layout == "NT":  # A not transposed, B transposed
         return (K, M), (K, N)
+    if layout == "TT":  # both transposed
+        return (M, K), (K, N)
     raise ValueError(f"Unsupported layout: {layout}")
 
 
@@ -267,6 +292,8 @@ def compute_pytorch_reference(A_ref, B_ref, layout):
         return torch.matmul(B_ref, A_ref)
     if layout == "NT":
         return torch.matmul(B_ref.T, A_ref)
+    if layout == "TT":
+        return torch.matmul(B_ref.T, A_ref.T)
     raise ValueError(f"Unsupported layout: {layout}")
 
 
@@ -334,6 +361,41 @@ 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.
+
+    MXFP4 rowwise and columnwise data are quantized separately (columnwise is the
+    rowwise quantization of the transposed tensor), so the reference must use the
+    buffer the GEMM reads: rowwise for a transposed operand, columnwise otherwise.
+    """
+    # 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 +1156,45 @@ 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", MXFP4_LAYOUTS)
+@pytest.mark.parametrize("out_dtype", MXFP4_OUT_DTYPES, 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.
+
+    That is the type ``general_gemm`` catches to fall back. The native MXFP4
+    GEMM does not run on gfx1250, so the fallback itself is not exercised.
+    """
+    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 --------------
 
 
diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py
index ddb9531380..760714b0b2 100644
--- a/transformer_engine/pytorch/cpp_extensions/gemm.py
+++ b/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,43 @@ 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 if NVTE_GEMM_BACKEND=FLYDSL, 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)
+            and get_gemm_backend() == "flydsl"
+        ):
+            result = _try_flydsl_gemm(
+                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 result is not None:
+                out, bias_grad, gelu_input, extra_output = result
+                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 +709,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:
diff --git a/transformer_engine/pytorch/flydsl_kernels/__init__.py b/transformer_engine/pytorch/flydsl_kernels/__init__.py
index 0057be1e27..d464e3cad1 100644
--- a/transformer_engine/pytorch/flydsl_kernels/__init__.py
+++ b/transformer_engine/pytorch/flydsl_kernels/__init__.py
@@ -1,4 +1,4 @@
 # Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
 #
 # See LICENSE for license information.
-"""FlyDSL kernels for ROCm gfx950 GEMM replacement."""
+"""FlyDSL kernels for ROCm gfx950 and gfx1250 (MXFP4) GEMM replacement."""
diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py
index aa20712c72..c47a43d070 100644
--- a/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py
+++ b/transformer_engine/pytorch/flydsl_kernels/gemm/__init__.py
@@ -2,7 +2,7 @@
 #
 # See LICENSE for license information.
 
-"""FlyDSL GEMM kernels (dense, non-grouped) for BF16/FP16/FP32/FP8/MXFP8."""
+"""FlyDSL dense GEMM kernels: BF16/FP16/FP32/FP8/MXFP8 (gfx950), MXFP4 (gfx1250)."""
 
 # Supported flydsl release series. Single source of truth for the version
 # requirement, since flydsl is user-installed (not a declared dependency). A
diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py
index 4a1dbb09a4..38a25ca054 100644
--- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py
+++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py
@@ -15,6 +15,7 @@
 from .fp32_gemm import fp32_matmul
 from .fp8_gemm import fp8_matmul
 from .mxfp8_gemm import mxfp8_matmul
+from .mxfp4_gemm import mxfp4_matmul
 
 # TODO: Some backend-independent GEMM wrapper utilities overlap with the
 # Triton GEMM backend (PR#667), including operand classification, logical
@@ -172,6 +173,16 @@ def _classify_input(t):
     except ImportError:
         pass
 
+    try:
+        from transformer_engine.pytorch.tensor.storage.mxfp4_tensor_storage import (
+            MXFP4TensorStorage,
+        )
+
+        if isinstance(t, MXFP4TensorStorage):
+            return "mxfp4", t
+    except ImportError:
+        pass
+
     try:
         from transformer_engine.pytorch.quantized_tensor import (
             QuantizedTensorStorage,
@@ -181,8 +192,8 @@ def _classify_input(t):
             raise ValueError(
                 "The FlyDSL GEMM backend does not support "
                 f"{type(t).__name__}. Only Float8Tensor / "
-                "Float8TensorStorage and MXFP8Tensor / "
-                "MXFP8TensorStorage are implemented."
+                "Float8TensorStorage, MXFP8Tensor / "
+                "MXFP8TensorStorage and MXFP4TensorStorage are implemented."
             )
     except ImportError:
         pass
@@ -1008,6 +1019,97 @@ def _run_mxfp8(
     return D, aux
 
 
+def _select_mxfp4_storage_for_layout(A, transa, B, transb):
+    """Select TE MXFP4 data and scales for a layout.
+
+    MXFP4 columnwise storage is a materialized transpose (as for tensor-wise FP8),
+    so every layout reduces to the TN kernel without transpose reads:
+
+        TN: B.rowwise    [M,K], A.rowwise    [N,K]
+        NN: B.rowwise    [M,K], A.columnwise [N,K]
+        NT: B.columnwise [M,K], A.columnwise [N,K]
+        TT: B.columnwise [M,K], A.rowwise    [N,K]
+    """
+
+    def select(t, rowwise, name):
+        orientation = "rowwise" if rowwise else "columnwise"
+        data = t._rowwise_data if rowwise else t._columnwise_data
+        scale = t._rowwise_scale_inv if rowwise else t._columnwise_scale_inv
+        if data is None or scale is None:
+            raise FlyDSLUnsupportedError(
+                f"FlyDSL MXFP4 requires existing {name} {orientation} data and scales"
+            )
+        if t._shuffle_rowwise_data if rowwise else t._shuffle_columnwise_data:
+            raise FlyDSLUnsupportedError(
+                f"FlyDSL MXFP4 requires plain FP4 data; {name} {orientation} data is shuffled"
+            )
+        if t._with_gemm_swizzled_scales:
+            raise FlyDSLUnsupportedError(
+                f"FlyDSL MXFP4 requires plain E8M0 scales; {name} scales are GEMM-swizzled"
+            )
+        if t._fp4_dtype != tex.DType.kFloat4E2M1:
+            raise FlyDSLUnsupportedError(f"FlyDSL MXFP4 requires E2M1 {name}, got {t._fp4_dtype}")
+        flatten = _flatten_rowwise if rowwise else _flatten_columnwise
+        return flatten(data.view(torch.uint8), f"{name}.{orientation}"), scale, data
+
+    A_data, A_scale, _ = select(A, bool(transa), "A")
+    B_data, B_scale, B_payload = select(B, not bool(transb), "B")
+    return A_data, A_scale, B_data, B_scale, B_payload
+
+
+def _run_mxfp4(
+    A,
+    transa,
+    B,
+    transb,
+    D,
+    *,
+    output_dtype: torch.dtype,
+):
+    """Dispatch MXFP4 (gfx1250) through the single TN kernel for TN/NN/NT/TT."""
+    layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}"
+    A_data, A_scale, B_data, B_scale, B_payload = _select_mxfp4_storage_for_layout(
+        A, transa, B, transb
+    )
+
+    # Kernel operand ownership is always swapped relative to TE
+    a_flydsl, a_scale = B_data, B_scale
+    b_flydsl, b_scale = A_data, A_scale
+    m, n = a_flydsl.shape[0], b_flydsl.shape[0]
+    if a_flydsl.shape[1] != b_flydsl.shape[1]:
+        raise FlyDSLUnsupportedError(
+            f"FlyDSL MXFP4 {layout} selected incompatible payloads: "
+            f"a={tuple(a_flydsl.shape)} and b={tuple(b_flydsl.shape)}"
+        )
+    if a_flydsl.device != b_flydsl.device:
+        raise ValueError(
+            f"FlyDSL MXFP4 {layout} operands must be on the same device, got "
+            f"{a_flydsl.device} and {b_flydsl.device}"
+        )
+
+    if D is not None:
+        output_shape = torch.Size(D.shape)
+    elif transb:
+        output_shape = torch.Size((m, n))
+    else:
+        output_shape = torch.Size((*B_payload.shape[:-1], n))
+    if _product(output_shape) != m * n:
+        raise FlyDSLUnsupportedError(
+            f"FlyDSL MXFP4 {layout} logical output shape "
+            f"{tuple(output_shape)} does not match kernel output {(m, n)}"
+        )
+
+    D = _validate_or_allocate_output(
+        D,
+        shape=output_shape,
+        dtype=output_dtype,
+        device=a_flydsl.device,
+        backend_name=f"MXFP4 {layout}",
+    )
+    mxfp4_matmul(a_flydsl, a_scale, b_flydsl, b_scale, D.view(m, n))
+    return D
+
+
 def _select_fp8_storage_for_layout(A, transa, B, transb):
     """Select existing TE FP8 backings and normalize to one core contract.
 
@@ -1242,9 +1344,10 @@ def te_generic_gemm_flydsl(
       - NN: transa=False, transb=False
       - NT: transa=False, transb=True
 
-    TT is intentionally rejected.
+    TT is intentionally rejected, except for MXFP4.
 
     Supported dtypes:
+      - MXFP4 E2M1 A/B (gfx1250) with FP16, BF16, or FP32 output
       - MXFP8 E4M3/E5M2 A/B combinations with FP16, BF16, or FP32 output
       - tensor-wise E4M3/E5M2 FP8 A/B combinations with FP16, BF16, or FP32 output
       - BF16 input with FP16, BF16, or FP32 output
@@ -1273,12 +1376,12 @@ def te_generic_gemm_flydsl(
     ):
         raise FlyDSLUnsupportedError("FlyDSL GEMM does not support comm+GEMM overlap")
 
-    if transa and transb:
-        raise FlyDSLUnsupportedError("FlyDSL GEMM does not support transa=True, transb=True (TT)")
-
     a_kind, _ = _classify_input(A)
     b_kind, _ = _classify_input(B)
 
+    if transa and transb and a_kind != "mxfp4":
+        raise FlyDSLUnsupportedError("FlyDSL GEMM does not support transa=True, transb=True (TT)")
+
     _validate_common_epilogue(
         quantizer=quantizer,
         bias=bias,
@@ -1289,6 +1392,35 @@ def te_generic_gemm_flydsl(
         beta=beta,
     )
 
+    if a_kind == "mxfp4" or b_kind == "mxfp4":
+        if a_kind != b_kind:
+            raise ValueError("Mixed MXFP4 and non-MXFP4 FlyDSL GEMM inputs are not supported")
+        if bias is not None:
+            raise FlyDSLUnsupportedError("FlyDSL MXFP4 GEMM fused bias is not implemented")
+        if gelu:
+            raise FlyDSLUnsupportedError("FlyDSL MXFP4 GEMM fused GELU is not implemented")
+
+        mxfp4_output_dtypes = {
+            None: torch.bfloat16,
+            tex.DType.kFloat16: torch.float16,
+            tex.DType.kBFloat16: torch.bfloat16,
+            tex.DType.kFloat32: torch.float32,
+        }
+        if output_dtype not in mxfp4_output_dtypes:
+            raise FlyDSLUnsupportedError(
+                f"FlyDSL MXFP4 supports FP16, BF16, or FP32 output, got {output_dtype}"
+            )
+
+        D = _run_mxfp4(
+            A,
+            transa,
+            B,
+            transb,
+            D,
+            output_dtype=mxfp4_output_dtypes[output_dtype],
+        )
+        return D, None, None, None
+
     if a_kind == "mxfp8" or b_kind == "mxfp8":
         # Validate both are MXFP8
         if a_kind != b_kind:
diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
new file mode 100644
index 0000000000..9a38f27683
--- /dev/null
+++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
@@ -0,0 +1,619 @@
+# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+# License for AMD contributions = MIT. See LICENSE for more information
+
+"""FlyDSL MXFP4 TN GEMM for gfx1250.
+
+D[M, N] = a[M, K] @ b[N, K]^T with K-contiguous packed FP4 operands and E8M0 scales along K. TE's
+MXFP4 columnwise storage is a materialized transpose, so this one kernel serves TN/NN/NT/TT.
+Tiles are staged with TDM into a 3-stage LDS pipeline and C is written back through LDS by TDM.
+"""
+
+import functools
+
+import torch
+
+import flydsl.compiler as flyc
+import flydsl.expr as fx
+from flydsl.expr import const_expr, range_constexpr
+
+from .exceptions import FlyDSLUnsupportedError
+
+MXFP4_BLOCK_SIZE = 32
+SCALE_WORD_K = 4 * MXFP4_BLOCK_SIZE
+BLOCK_M = 256
+BLOCK_N = 256
+BLOCK_K = 128
+NUM_WAVES = 8
+WAVE_32 = 32
+NUM_THREADS = NUM_WAVES * WAVE_32
+NUM_STAGES = 3  # LDS pages per operand
+# K-loop tile j is loaded by wave group j % 2, two waves per operand
+TDM_GROUP_WAVES = NUM_WAVES // 2
+TDM_WAVES = 2
+PRO_SCALE_WAVES = 2
+LDS_B_START = 64 * 1024
+LDS_BYTES = 320 * 1024
+C_HALF_ROWS = BLOCK_M // 2  # C is stored in two halves
+C_WAVE_ROWS = C_HALF_ROWS // NUM_WAVES
+SCALE_TDM_COLS = 32
+# Logical 16 B chunk c of row r is stored at chunk c ^ ((r >> 2) & 3) to avoid LDS bank conflicts
+SWZ_CHUNKS = 4
+SCHED_TAIL_MMAS = 8
+# Each lane's scale words are made contiguous so they load as ds_load_b128
+SCALE_LANE_ROWS_A = 32
+SCALE_LANE_ROWS_B = 64
+OUT_KINDS = {torch.float32: 0, torch.bfloat16: 1, torch.float16: 2}
+
+
+def _lds_bytes(K):
+    k_iters = K // BLOCK_K
+    return max(
+        LDS_B_START + NUM_STAGES * BLOCK_N * BLOCK_K // 2 + 2 * k_iters * BLOCK_M * 4,
+        C_HALF_ROWS * BLOCK_N * 4,
+    )
+
+
+MAX_K = max(k for k in range(2 * BLOCK_K, 64 * 1024 + 1, BLOCK_K) if _lds_bytes(k) <= LDS_BYTES)
+
+
+# Scale words (one BLOCK_K tile of a row) stored [rows // block_rows][K_iters][block_rows],
+# viewed by logical (row, k_tile)
+def iter_major_scale_words(t, rows, words_per_row, block_rows, lane_rows):
+    src = fx.get_iter(t)
+    ptr = fx.recast_iter(fx.PointerType.get(fx.Int32.ir_type, src.memspace, 4), src)
+    reps = block_rows // lane_rows
+    return fx.make_view(
+        ptr,
+        fx.make_layout(
+            (((lane_rows, reps), rows // block_rows), words_per_row),
+            (((reps, 1), words_per_row * block_rows), block_rows),
+        ),
+    )
+
+
+# The same words as 2D rows of SCALE_TDM_COLS for TDM; row block b is row-tile b
+def iter_major_scale_rows(t, rows, K_iters, block_rows):
+    src = fx.get_iter(t)
+    ptr = fx.recast_iter(fx.PointerType.get(fx.Int32.ir_type, src.memspace, 4), src)
+    tdm_rows = rows // block_rows * (K_iters * block_rows // SCALE_TDM_COLS)
+    return fx.make_view(ptr, fx.make_layout((tdm_rows, SCALE_TDM_COLS), (SCALE_TDM_COLS, 1)))
+
+
+@flyc.kernel
+def gemm_kernel(
+    A: fx.Tensor,
+    B: fx.Tensor,
+    A_s: fx.Tensor,
+    B_s: fx.Tensor,
+    C: fx.Tensor,
+    M: fx.Int32,
+    N: fx.Int32,
+    K: fx.Int32,
+    K_iters: fx.Constexpr[int],
+    OUT_KIND: fx.Constexpr[int],
+):
+
+    def as_words(t, rows, words_per_row):
+        src = fx.get_iter(t)
+        ptr = fx.recast_iter(fx.PointerType.get(fx.Int32.ir_type, src.memspace, 4), src)
+        return fx.make_view(ptr, fx.make_layout((rows, words_per_row), (words_per_row, 1)))
+
+    out_fx = (fx.Float32, fx.BFloat16, fx.Float16)[OUT_KIND]
+    out_bytes = 4 if OUT_KIND == 0 else 2
+    c_stage_bytes = C_HALF_ROWS * BLOCK_N * out_bytes
+
+    tid = fx.thread_idx.x
+    bid = fx.block_idx.x
+    blocks_n = N // BLOCK_N
+    bid_m = bid // blocks_n
+    bid_n = bid % blocks_n
+
+    # wave-uniform, so TDM descriptors derived from it stay in SGPRs
+    lane = tid % 32
+    wave = fx.Int32(fx.rocdl.readfirstlane(fx.Int32.ir_type, tid // 32))
+    l16, khalf = lane % 16, lane // 16
+    wm, wn = wave % 2, wave // 2
+
+    sA_rows = iter_major_scale_rows(A_s, M, K_iters, BLOCK_M)
+    sB_rows = iter_major_scale_rows(B_s, N, K_iters, BLOCK_N)
+    k_words = K // 8
+    A_words = as_words(A, M, k_words)
+    B_words = as_words(B, N, k_words)
+
+    # physical chunks of this lane's logical chunks khalf and khalf + 2 (same for all its rows)
+    swz = (l16 // 4) % SWZ_CHUNKS
+    chunk_lo = khalf ^ swz
+    chunk_hi = chunk_lo ^ 2
+    a_lds_lo = (wm * 16 + l16) * 16 + 4 * chunk_lo
+    a_lds_hi = (wm * 16 + l16) * 16 + 4 * chunk_hi
+    b_lds_lo = (wn * 16 + l16) * 16 + 4 * chunk_lo
+    b_lds_hi = (wn * 16 + l16) * 16 + 4 * chunk_hi
+    as_lds_base = (wm * 16 + l16) * (BLOCK_M // SCALE_LANE_ROWS_A)
+    bs_lds_base = (wn * 16 + l16) * (BLOCK_N // SCALE_LANE_ROWS_B)
+
+    # tile k lives in LDS page k % NUM_STAGES
+    smem = fx.SharedAllocator(static=False)
+    lds_A = smem.allocate(fx.Array[fx.Int32, NUM_STAGES * BLOCK_M * 16]).peek()
+    # B starts at 64 KB so all its pages are reachable with immediate offsets from one base register
+    smem.allocate(fx.Array[fx.Int32, LDS_B_START // 4 - NUM_STAGES * BLOCK_M * 16])
+    lds_B = smem.allocate(fx.Array[fx.Int32, NUM_STAGES * BLOCK_N * 16]).peek()
+    box_A = fx.make_layout((BLOCK_M, 16), (16, 1))
+    sA = fx.make_view(
+        lds_A.ptr, fx.make_layout(((BLOCK_M, 16), NUM_STAGES), ((16, 1), BLOCK_M * 16))
+    )
+    box_B = fx.make_layout((BLOCK_N, 16), (16, 1))
+    sB = fx.make_view(
+        lds_B.ptr, fx.make_layout(((BLOCK_N, 16), NUM_STAGES), ((16, 1), BLOCK_N * 16))
+    )
+
+    # all of the block's scales, loaded once
+    sc_rows_A = K_iters * BLOCK_M // SCALE_TDM_COLS
+    sc_rows_B = K_iters * BLOCK_N // SCALE_TDM_COLS
+    lds_As = smem.allocate(fx.Array[fx.Int32, K_iters * BLOCK_M]).peek()
+    lds_Bs = smem.allocate(fx.Array[fx.Int32, K_iters * BLOCK_N]).peek()
+    box_As = fx.make_layout((sc_rows_A, SCALE_TDM_COLS), (SCALE_TDM_COLS, 1))
+    sAs = fx.make_view(
+        lds_As.ptr,
+        fx.make_layout(((sc_rows_A, SCALE_TDM_COLS), 1), ((SCALE_TDM_COLS, 1), K_iters * BLOCK_M)),
+    )
+    box_Bs = fx.make_layout((sc_rows_B, SCALE_TDM_COLS), (SCALE_TDM_COLS, 1))
+    sBs = fx.make_view(
+        lds_Bs.ptr,
+        fx.make_layout(((sc_rows_B, SCALE_TDM_COLS), 1), ((SCALE_TDM_COLS, 1), K_iters * BLOCK_N)),
+    )
+
+    # C is staged from LDS offset 0 after the K loop
+    if smem.allocated_bytes < c_stage_bytes:
+        smem.allocate(c_stage_bytes - smem.allocated_bytes)
+    lds_C_ptr = fx.recast_iter(out_fx, smem.base_ptr)
+
+    sA_base_ptr = fx.get_iter(sA)
+    sB_base_ptr = fx.get_iter(sB)
+    sAs_base_ptr = fx.get_iter(sAs)
+    sBs_base_ptr = fx.get_iter(sBs)
+
+    # K-loop TDM atoms over this wave's half tile; later K tiles are reached with imm_offset
+    tdm_wave = wave % TDM_WAVES
+    HALF_ROWS = BLOCK_M // TDM_WAVES
+    half_box = fx.make_layout((HALF_ROWS, 16), (16, 1))
+    gA_half = fx.make_view(
+        fx.add_offset(fx.get_iter(A_words), (bid_m * BLOCK_M + tdm_wave * HALF_ROWS) * k_words),
+        half_box,
+    )
+    gB_half = fx.make_view(
+        fx.add_offset(fx.get_iter(B_words), (bid_n * BLOCK_N + tdm_wave * HALF_ROWS) * k_words),
+        half_box,
+    )
+    tdm_A = fx.rocdl.make_tdm_atom(
+        gA_half, [None, None], strides=[fx.Int64(k_words), None], num_warps=1
+    )
+    tdm_B = fx.rocdl.make_tdm_atom(
+        gB_half, [None, None], strides=[fx.Int64(k_words), None], num_warps=1
+    )
+
+    def lds_half(base_ptr, page):
+        return fx.make_view(
+            fx.add_offset(base_ptr, page * BLOCK_M * 16 + tdm_wave * HALF_ROWS * 16), half_box
+        )
+
+    # prologue atoms: whole tiles and scales
+    tdm_A1w, cA1w = fx.rocdl.cdna5.make_tiled_tdm_atom(
+        fx.rocdl.TensorLoad(), A_words, box_A, (BLOCK_M, 16), num_warps=1
+    )
+    blk_cA1w = fx.zipped_divide(cA1w, (BLOCK_M, 16))[None, (bid_m, None)]
+    tAsA1w, tAgA1w = fx.rocdl.cdna5.tdm_partition(tdm_A1w, 0, fx.make_layout(1, 1), sA, blk_cA1w)
+    tdm_B1w, cB1w = fx.rocdl.cdna5.make_tiled_tdm_atom(
+        fx.rocdl.TensorLoad(), B_words, box_B, (BLOCK_N, 16), num_warps=1
+    )
+    blk_cB1w = fx.zipped_divide(cB1w, (BLOCK_N, 16))[None, (bid_n, None)]
+    tBsB1w, tBgB1w = fx.rocdl.cdna5.tdm_partition(tdm_B1w, 0, fx.make_layout(1, 1), sB, blk_cB1w)
+
+    sc_wave = wave % PRO_SCALE_WAVES
+    tdm_As, cAs = fx.rocdl.cdna5.make_tiled_tdm_atom(
+        fx.rocdl.TensorLoad(),
+        sA_rows,
+        box_As,
+        (sc_rows_A, SCALE_TDM_COLS),
+        num_warps=PRO_SCALE_WAVES,
+    )
+    blk_cAs = fx.zipped_divide(cAs, (sc_rows_A, SCALE_TDM_COLS))[None, (None, 0)]
+    tAssAs, tAsgAs = fx.rocdl.cdna5.tdm_partition(
+        tdm_As, sc_wave, fx.make_layout(PRO_SCALE_WAVES, 1), sAs, blk_cAs
+    )
+
+    tdm_Bs, cBs = fx.rocdl.cdna5.make_tiled_tdm_atom(
+        fx.rocdl.TensorLoad(),
+        sB_rows,
+        box_Bs,
+        (sc_rows_B, SCALE_TDM_COLS),
+        num_warps=PRO_SCALE_WAVES,
+    )
+    blk_cBs = fx.zipped_divide(cBs, (sc_rows_B, SCALE_TDM_COLS))[None, (None, 0)]
+    tBssBs, tBsgBs = fx.rocdl.cdna5.tdm_partition(
+        tdm_Bs, sc_wave, fx.make_layout(PRO_SCALE_WAVES, 1), sBs, blk_cBs
+    )
+
+    # C store: each wave writes C_WAVE_ROWS rows of each half
+    c_wave_box = fx.make_layout((C_WAVE_ROWS, BLOCK_N), (BLOCK_N, 1))
+    c_wave_row = bid_m * BLOCK_M + wave * C_WAVE_ROWS
+    gC_wave = [
+        fx.make_view(
+            fx.add_offset(fx.get_iter(C), (c_wave_row + h * C_HALF_ROWS) * N + bid_n * BLOCK_N),
+            c_wave_box,
+        )
+        for h in range(2)
+    ]
+    tdm_C = fx.rocdl.make_tdm_atom(
+        gC_wave[0], [None, None], strides=[fx.Int64(N), None], num_warps=1
+    )
+    lds_C_wave = fx.make_view(fx.add_offset(lds_C_ptr, wave * C_WAVE_ROWS * BLOCK_N), c_wave_box)
+
+    bC = fx.slice(fx.zipped_divide(C, (BLOCK_M, BLOCK_N)), (None, (bid_m, bid_n)))
+
+    mma_atom = fx.make_mma_atom(fx.rocdl.WMMAScale(16, 16, 128, fx.Float4E2M1FN))
+    tiled_mma = fx.make_tiled_mma(mma_atom, fx.make_layout((2, 4, 1), (1, 2, 0)))
+    thr_mma = tiled_mma.thr_slice(tid)
+
+    copy_atom_i32 = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32)
+    copy_atom_c = fx.make_copy_atom(
+        fx.UniversalCopy32b() if OUT_KIND == 0 else fx.UniversalCopy16b(), out_fx
+    )
+
+    # double-buffered register fragments: tile k is read into set k % 2
+    frag_A = [fx.make_rmem_tensor(fx.make_layout((8, 8, 1), (1, 8, 0)), fx.Int32) for _ in range(2)]
+    frag_B = [fx.make_rmem_tensor(fx.make_layout((8, 4, 1), (1, 8, 0)), fx.Int32) for _ in range(2)]
+    frag_sA = [fx.make_rmem_tensor((1, 8, 1), fx.Int32) for _ in range(2)]
+    frag_sB = [fx.make_rmem_tensor((1, 4, 1), fx.Int32) for _ in range(2)]
+    frag_acc = thr_mma.make_fragment_C(bC)
+
+    # tiles issued per wave class, to derive each class's tensorcnt wait for tile t
+    WAVE_CLASSES = ((0, 2), (2, 4), (4, NUM_WAVES))
+    tdm_issued = [[0], [1], [0]]  # prologue ops
+
+    def wait_tile(t):
+        counts = [sum(1 for x in tiles if x > t) for tiles in tdm_issued]
+        if counts[0] == counts[1] == counts[2]:
+            fx.rocdl.s_wait_tensorcnt(counts[0])
+        else:
+            if wave < WAVE_CLASSES[0][1]:
+                fx.rocdl.s_wait_tensorcnt(counts[0])
+            else:
+                if wave < WAVE_CLASSES[1][1]:
+                    fx.rocdl.s_wait_tensorcnt(counts[1])
+                else:
+                    fx.rocdl.s_wait_tensorcnt(counts[2])
+
+    def issue_tdm(j):
+        page = j % NUM_STAGES
+        group = j % 2
+        k_off = fx.Int64(j * BLOCK_K // 2)
+        if wave // TDM_GROUP_WAVES == group:
+            if wave % TDM_GROUP_WAVES < TDM_WAVES:
+                fx.copy(tdm_A, gA_half, lds_half(sA_base_ptr, page), imm_offset=k_off)
+            else:
+                fx.copy(tdm_B, gB_half, lds_half(sB_base_ptr, page), imm_offset=k_off)
+        for c in ((0, 1) if group == 0 else (2,)):
+            tdm_issued[c].append(j)
+
+    # one TDM op per wave: tiles 0 and 1 and the scales
+    def issue_prologue_tdm():
+        if wave < 4:
+            if wave == 0:
+                fx.copy(tdm_A1w, tAgA1w[None, 0], tAsA1w[None, 0])
+            else:
+                if wave == 1:
+                    fx.copy(tdm_B1w, tBgB1w[None, 0], tBsB1w[None, 0])
+                else:
+                    if wave == 2:
+                        fx.copy(tdm_A1w, tAgA1w[None, 1], tAsA1w[None, 1])
+                    else:
+                        fx.copy(tdm_B1w, tBgB1w[None, 1], tBsB1w[None, 1])
+        else:
+            if wave < 4 + PRO_SCALE_WAVES:
+                fx.copy(tdm_As, tAsgAs[None, bid_m], tAssAs[None, 0])
+            else:
+                fx.copy(tdm_Bs, tBsgBs[None, bid_n], tBssBs[None, 0])
+
+    def lds_to_regs(k):
+        page, r = k % NUM_STAGES, k % 2
+        # one copy per swizzled chunk; registers stay in logical chunk order
+        for h, a_off, b_off in ((0, a_lds_lo, b_lds_lo), (1, a_lds_hi, b_lds_hi)):
+            src_A = fx.make_view(
+                fx.add_offset(sA_base_ptr, page * BLOCK_M * 16 + a_off),
+                fx.make_layout((4, 8), (1, 32 * 16)),
+            )
+            src_B = fx.make_view(
+                fx.add_offset(sB_base_ptr, page * BLOCK_N * 16 + b_off),
+                fx.make_layout((4, 4), (1, 64 * 16)),
+            )
+            fx.copy(
+                copy_atom_i32,
+                src_A,
+                fx.make_view(
+                    fx.add_offset(fx.get_iter(frag_A[r]), 4 * h), fx.make_layout((4, 8), (1, 8))
+                ),
+            )
+            fx.copy(
+                copy_atom_i32,
+                src_B,
+                fx.make_view(
+                    fx.add_offset(fx.get_iter(frag_B[r]), 4 * h), fx.make_layout((4, 4), (1, 8))
+                ),
+            )
+        s_off = k
+        src_sA = fx.make_view(
+            fx.add_offset(sAs_base_ptr, s_off * BLOCK_M + as_lds_base),
+            fx.make_layout((4, 2), (1, 4)),
+        )
+        src_sB = fx.make_view(
+            fx.add_offset(sBs_base_ptr, s_off * BLOCK_N + bs_lds_base),
+            fx.make_layout((4, 1), (1, 4)),
+        )
+        fx.copy(
+            copy_atom_i32,
+            src_sA,
+            fx.make_view(fx.get_iter(frag_sA[r]), fx.make_layout((4, 2), (1, 4))),
+        )
+        fx.copy(
+            copy_atom_i32,
+            src_sB,
+            fx.make_view(fx.get_iter(frag_sB[r]), fx.make_layout((4, 1), (1, 4))),
+        )
+
+    A_REPS, B_REPS = 8, 4
+    n_dsrd = 2 * (A_REPS + B_REPS) + (A_REPS + B_REPS) // 2
+    n_mma = A_REPS * B_REPS
+
+    # interleave tile k+1's LDS reads with gemm(k)'s WMMAs
+    def hot_loop_scheduler():
+        n_spread = n_mma - SCHED_TAIL_MMAS
+        issued = 0
+        for i in range_constexpr(n_spread):
+            want = (i + 1) * n_dsrd // n_spread
+            if want > issued:
+                fx.rocdl.sched_dsrd(want - issued)
+                issued = want
+            fx.rocdl.sched_mfma(1)
+        fx.rocdl.sched_mfma(SCHED_TAIL_MMAS)
+        fx.rocdl.sched_barrier(0)
+
+    def gemm_step(k):
+        r = k % 2
+        fx.gemm(mma_atom, frag_acc, [frag_A[r], frag_sA[r]], [frag_B[r], frag_sB[r]], frag_acc)
+
+    frag_acc.fill(0)
+
+    issue_prologue_tdm()
+    wait_tile(0)
+    fx.barrier()
+    lds_to_regs(0)
+
+    for k in range_constexpr(K_iters):
+        if k + 1 < K_iters:
+            # keep WMMAs from being scheduled across the barrier
+            fx.rocdl.sched_barrier(0)
+            wait_tile(k + 1)
+            fx.rocdl.s_wait_dscnt(0)
+            fx.barrier()
+            fx.rocdl.sched_barrier(0)
+            if k == 0 and NUM_STAGES - 1 < K_iters:
+                issue_tdm(NUM_STAGES - 1)
+            if k + NUM_STAGES < K_iters:
+                issue_tdm(k + NUM_STAGES)
+            lds_to_regs(k + 1)
+            gemm_step(k)
+            hot_loop_scheduler()
+        else:
+            gemm_step(k)
+
+    if const_expr(OUT_KIND == 0):
+        frag_c = frag_acc
+    else:
+        frag_c = fx.make_fragment_like(frag_acc, out_fx)
+        frag_c.store(frag_acc.load().to(out_fx))
+
+    # accumulator (v, j, i) is tile row wm*16 + 8*khalf + v + 32*i, column wn*16 + l16 + 64*j
+    fx.rocdl.s_wait_dscnt(0)
+    fx.barrier()
+    lane_c = (wm * 16 + 8 * khalf) * BLOCK_N + wn * 16 + l16
+    for h in range_constexpr(2):
+        if h > 0:
+            fx.rocdl.s_wait_tensorcnt(0)
+            fx.barrier()
+        # sched_barrier(0x7FF) keeps the stores from being merged into ds_store_2addr
+        for ii in range_constexpr(4):
+            for j in range_constexpr(4):
+                for v in range_constexpr(8):
+                    src_c = fx.make_view(
+                        fx.add_offset(fx.get_iter(frag_c), v + 8 * j + 32 * (4 * h + ii)),
+                        fx.make_layout(1, 1),
+                    )
+                    dst_c = fx.make_view(
+                        fx.add_offset(lds_C_ptr, lane_c + (32 * ii + v) * BLOCK_N + 64 * j),
+                        fx.make_layout(1, 1),
+                    )
+                    fx.copy(copy_atom_c, src_c, dst_c)
+                    fx.rocdl.sched_barrier(0x7FF)
+        fx.rocdl.s_wait_dscnt(0)
+        fx.barrier()
+        fx.copy(tdm_C, lds_C_wave, gC_wave[h])
+    fx.rocdl.s_wait_tensorcnt(0)
+
+
+@flyc.jit
+def launch_gemm(
+    A: fx.Tensor,
+    B: fx.Tensor,
+    A_s: fx.Tensor,
+    B_s: fx.Tensor,
+    C: fx.Tensor,
+    M: fx.Int32,
+    N: fx.Int32,
+    K: fx.Int32,
+    K_iters: fx.Constexpr[int],
+    OUT_KIND: fx.Constexpr[int],
+    stream: fx.Stream = fx.Stream(None),
+):
+    grid_tot = (M // BLOCK_M) * (N // BLOCK_N)
+    gemm_kernel(A, B, A_s, B_s, C, M, N, K, K_iters, OUT_KIND).launch(
+        grid=(grid_tot, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream
+    )
+
+
+@flyc.kernel
+def shuffle_scales_kernel(
+    scales: fx.Tensor,
+    scales_shuf: fx.Tensor,
+    rows: fx.Int32,
+    words_per_row: fx.Int32,
+    src_row_words: fx.Int32,
+    block_rows: fx.Constexpr[int],
+    lane_rows: fx.Constexpr[int],
+):
+    # row-major [rows][src_row_words] (TE's padded buffer) -> iteration-major, one thread per word
+    tid = fx.thread_idx.x
+    bid_r = fx.block_idx.x
+    bid_k = fx.block_idx.y
+
+    src = fx.get_iter(scales)
+    src_ptr = fx.recast_iter(fx.PointerType.get(fx.Int32.ir_type, src.memspace, 4), src)
+    src_words = fx.make_view(src_ptr, fx.make_layout((rows, words_per_row), (src_row_words, 1)))
+    dst_words = iter_major_scale_words(scales_shuf, rows, words_per_row, block_rows, lane_rows)
+
+    src_tile = fx.slice(fx.zipped_divide(src_words, (block_rows, 1)), (None, (bid_r, bid_k)))
+    dst_tile = fx.slice(fx.zipped_divide(dst_words, (block_rows, 1)), (None, (bid_r, bid_k)))
+
+    copy_atom_i32 = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32)
+    tile_mn, tv_layout = fx.make_layout_tv(
+        fx.make_layout((block_rows, 1), (1, 1)), fx.make_layout((1, 1), (1, 1))
+    )
+    thr_copy = fx.make_tiled_copy(copy_atom_i32, tv_layout, tile_mn).get_slice(tid)
+    part_src = thr_copy.partition_S(src_tile)
+    part_dst = thr_copy.partition_D(dst_tile)
+    frag = fx.make_fragment_like(part_src)
+    fx.copy(copy_atom_i32, part_src, frag)
+    fx.copy(copy_atom_i32, frag, part_dst)
+
+
+@flyc.jit
+def shuffle_scales(
+    scales: fx.Tensor,
+    scales_shuf: fx.Tensor,
+    rows: fx.Int32,
+    K: fx.Int32,
+    src_row_words: fx.Int32,
+    block_rows: fx.Constexpr[int],
+    lane_rows: fx.Constexpr[int],
+    stream: fx.Stream = fx.Stream(None),
+):
+    words_per_row = K // SCALE_WORD_K
+    grid_r = rows // block_rows
+    shuffle_scales_kernel(
+        scales, scales_shuf, rows, words_per_row, src_row_words, block_rows, lane_rows
+    ).launch(grid=(grid_r, words_per_row, 1), block=(block_rows, 1, 1), stream=stream)
+
+
+@functools.lru_cache(maxsize=None)
+def _swizzle_index(rows: int, K: int, device: torch.device) -> torch.Tensor:
+    swz = (torch.arange(rows, device=device) >> 2) & (SWZ_CHUNKS - 1)
+    src_chunk = torch.arange(SWZ_CHUNKS, device=device)[None, :] ^ swz[:, None]
+    return src_chunk[:, None, :, None].expand(rows, K // BLOCK_K, SWZ_CHUNKS, 4)
+
+
+def swizzle_operand(q, rows, K):
+    """Copy of packed FP4 operand q [rows, K/2] with the LDS chunk swizzle applied."""
+    words = q.view(torch.int32).view(rows, K // BLOCK_K, SWZ_CHUNKS, 4)
+    return words.gather(2, _swizzle_index(rows, K, q.device)).view(torch.uint8).view(rows, K // 2)
+
+
+def _shuffle_scale_words(scale, rows, K, block_rows, lane_rows, stream):
+    """Reorder TE E8M0 scales (possibly padded) into the kernel's iteration-major order."""
+    if scale.dtype != torch.uint8 or scale.ndim != 2:
+        raise FlyDSLUnsupportedError(
+            f"FlyDSL MXFP4 expects 2D uint8 E8M0 scales, got {scale.dtype} {tuple(scale.shape)}"
+        )
+    if scale.shape[0] < rows or scale.shape[1] < K // MXFP4_BLOCK_SIZE:
+        raise FlyDSLUnsupportedError(
+            f"FlyDSL MXFP4 scales {tuple(scale.shape)} do not cover "
+            f"({rows}, {K // MXFP4_BLOCK_SIZE})"
+        )
+    if scale.stride(1) != 1 or scale.stride(0) % 4 != 0 or scale.data_ptr() % 4 != 0:
+        raise FlyDSLUnsupportedError(
+            f"FlyDSL MXFP4 needs word-aligned row-major scales, got strides {scale.stride()}"
+        )
+    out = torch.empty(rows, K // MXFP4_BLOCK_SIZE, dtype=torch.uint8, device=scale.device)
+    shuffle_scales(scale, out, rows, K, scale.stride(0) // 4, block_rows, lane_rows, stream=stream)
+    return out
+
+
+def check_mxfp4_shape(m, n, k):
+    """Raise FlyDSLUnsupportedError unless the kernel supports C[m, n] with reduction depth k."""
+    if m % BLOCK_M or n % BLOCK_N:
+        raise FlyDSLUnsupportedError(
+            f"FlyDSL MXFP4 needs M, N multiples of {BLOCK_M}, got M={m}, N={n}"
+        )
+    if k % BLOCK_K or k < 2 * BLOCK_K or k > MAX_K:
+        raise FlyDSLUnsupportedError(
+            f"FlyDSL MXFP4 needs K a multiple of {BLOCK_K} in [{2 * BLOCK_K}, {MAX_K}], got K={k}"
+        )
+
+
+def mxfp4_matmul(
+    a: torch.Tensor,
+    a_scale: torch.Tensor,
+    b: torch.Tensor,
+    b_scale: torch.Tensor,
+    D: torch.Tensor,
+    stream=None,
+):
+    """D[M, N] = dequant(a, a_scale) @ dequant(b, b_scale)^T.
+
+    a [M, K/2] and b [N, K/2] packed FP4; scales E8M0 covering [M, K/32] and [N, K/32];
+    D float32, bfloat16 or float16.
+    """
+    if a.ndim != 2 or b.ndim != 2 or D.ndim != 2:
+        raise ValueError(
+            f"FlyDSL MXFP4 expects rank-2 a, b, D; got "
+            f"{tuple(a.shape)}, {tuple(b.shape)}, {tuple(D.shape)}"
+        )
+    m, k2 = a.shape
+    n, kb2 = b.shape
+    k = 2 * k2
+    if kb2 != k2:
+        raise ValueError(f"Incompatible MXFP4 operands: a{tuple(a.shape)} and b{tuple(b.shape)}")
+    if tuple(D.shape) != (m, n):
+        raise ValueError(f"D shape {tuple(D.shape)} != expected {(m, n)}")
+    if D.dtype not in OUT_KINDS:
+        raise FlyDSLUnsupportedError(
+            f"FlyDSL MXFP4 output must be float32, bfloat16 or float16, got {D.dtype}"
+        )
+    if not (a.is_contiguous() and b.is_contiguous() and D.is_contiguous()):
+        raise FlyDSLUnsupportedError("FlyDSL MXFP4 requires contiguous a, b and D")
+    if len({a.device, b.device, a_scale.device, b_scale.device, D.device}) != 1:
+        raise ValueError("a, b, scales and D must be on the same device")
+    check_mxfp4_shape(m, n, k)
+
+    if stream is None:
+        stream = torch.cuda.current_stream(a.device)
+    a_u8 = a.view(torch.uint8)
+    b_u8 = b.view(torch.uint8)
+    a_scale_k = _shuffle_scale_words(a_scale, m, k, BLOCK_M, SCALE_LANE_ROWS_A, stream)
+    b_scale_k = _shuffle_scale_words(b_scale, n, k, BLOCK_N, SCALE_LANE_ROWS_B, stream)
+    with torch.cuda.stream(stream):
+        a_k = swizzle_operand(a_u8, m, k)
+        b_k = swizzle_operand(b_u8, n, k)
+    launch_gemm(
+        a_k, b_k, a_scale_k, b_scale_k, D, m, n, k, k // BLOCK_K, OUT_KINDS[D.dtype], stream=stream
+    )
+    return D
+
+
+__all__ = [
+    "BLOCK_M",
+    "BLOCK_N",
+    "BLOCK_K",
+    "MAX_K",
+    "MXFP4_BLOCK_SIZE",
+    "check_mxfp4_shape",
+    "mxfp4_matmul",
+]

From 512b0fff59749d1d808b9d7cf28fdf90d34759f5 Mon Sep 17 00:00:00 2001
From: Aristotle Martin <Aristotle.Martin@amd.com>
Date: Fri, 30 Oct 2026 15:19:22 +0000
Subject: [PATCH 02/11] clean up comments

---
 tests/pytorch/test_gemm_backends.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/tests/pytorch/test_gemm_backends.py b/tests/pytorch/test_gemm_backends.py
index e85ab201f3..bc488808ef 100644
--- a/tests/pytorch/test_gemm_backends.py
+++ b/tests/pytorch/test_gemm_backends.py
@@ -15,7 +15,7 @@
   (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 ``test_flydsl_*_mxfp4`` tests run on
-  gfx1250 instead (the only arch the FlyDSL MXFP4 path supports).
+  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
@@ -34,8 +34,7 @@
 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. MXFP4 (gfx1250) covers TN / NN / NT / TT
-with FP32 / BF16 / FP16 output against a torch.matmul reference only: the native
-C++ MXFP4 GEMM does not run on gfx1250. FlyDSL kernels require tile-aligned
+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.
 """

From ac84339865e8bf0437e6d860fae390e0293bf2fb Mon Sep 17 00:00:00 2001
From: Aristotle Martin <Aristotle.Martin@amd.com>
Date: Mon, 2 Nov 2026 18:56:56 +0000
Subject: [PATCH 03/11] restructure scale pipeline to remove K size limitation

---
 .../pytorch/flydsl_kernels/gemm/mxfp4_gemm.py | 184 +++++++++++-------
 1 file changed, 112 insertions(+), 72 deletions(-)

diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
index 9a38f27683..523c05ecd4 100644
--- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
+++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
@@ -31,11 +31,12 @@
 TDM_GROUP_WAVES = NUM_WAVES // 2
 TDM_WAVES = 2
 PRO_SCALE_WAVES = 2
+# Scales stream through SCALE_STAGES LDS pages of SCALE_GROUP K tiles
+SCALE_GROUP = 8
+SCALE_STAGES = 2
 LDS_B_START = 64 * 1024
-LDS_BYTES = 320 * 1024
 C_HALF_ROWS = BLOCK_M // 2  # C is stored in two halves
 C_WAVE_ROWS = C_HALF_ROWS // NUM_WAVES
-SCALE_TDM_COLS = 32
 # Logical 16 B chunk c of row r is stored at chunk c ^ ((r >> 2) & 3) to avoid LDS bank conflicts
 SWZ_CHUNKS = 4
 SCHED_TAIL_MMAS = 8
@@ -45,17 +46,6 @@
 OUT_KINDS = {torch.float32: 0, torch.bfloat16: 1, torch.float16: 2}
 
 
-def _lds_bytes(K):
-    k_iters = K // BLOCK_K
-    return max(
-        LDS_B_START + NUM_STAGES * BLOCK_N * BLOCK_K // 2 + 2 * k_iters * BLOCK_M * 4,
-        C_HALF_ROWS * BLOCK_N * 4,
-    )
-
-
-MAX_K = max(k for k in range(2 * BLOCK_K, 64 * 1024 + 1, BLOCK_K) if _lds_bytes(k) <= LDS_BYTES)
-
-
 # Scale words (one BLOCK_K tile of a row) stored [rows // block_rows][K_iters][block_rows],
 # viewed by logical (row, k_tile)
 def iter_major_scale_words(t, rows, words_per_row, block_rows, lane_rows):
@@ -71,14 +61,6 @@ def iter_major_scale_words(t, rows, words_per_row, block_rows, lane_rows):
     )
 
 
-# The same words as 2D rows of SCALE_TDM_COLS for TDM; row block b is row-tile b
-def iter_major_scale_rows(t, rows, K_iters, block_rows):
-    src = fx.get_iter(t)
-    ptr = fx.recast_iter(fx.PointerType.get(fx.Int32.ir_type, src.memspace, 4), src)
-    tdm_rows = rows // block_rows * (K_iters * block_rows // SCALE_TDM_COLS)
-    return fx.make_view(ptr, fx.make_layout((tdm_rows, SCALE_TDM_COLS), (SCALE_TDM_COLS, 1)))
-
-
 @flyc.kernel
 def gemm_kernel(
     A: fx.Tensor,
@@ -114,8 +96,6 @@ def as_words(t, rows, words_per_row):
     l16, khalf = lane % 16, lane // 16
     wm, wn = wave % 2, wave // 2
 
-    sA_rows = iter_major_scale_rows(A_s, M, K_iters, BLOCK_M)
-    sB_rows = iter_major_scale_rows(B_s, N, K_iters, BLOCK_N)
     k_words = K // 8
     A_words = as_words(A, M, k_words)
     B_words = as_words(B, N, k_words)
@@ -146,21 +126,10 @@ def as_words(t, rows, words_per_row):
         lds_B.ptr, fx.make_layout(((BLOCK_N, 16), NUM_STAGES), ((16, 1), BLOCK_N * 16))
     )
 
-    # all of the block's scales, loaded once
-    sc_rows_A = K_iters * BLOCK_M // SCALE_TDM_COLS
-    sc_rows_B = K_iters * BLOCK_N // SCALE_TDM_COLS
-    lds_As = smem.allocate(fx.Array[fx.Int32, K_iters * BLOCK_M]).peek()
-    lds_Bs = smem.allocate(fx.Array[fx.Int32, K_iters * BLOCK_N]).peek()
-    box_As = fx.make_layout((sc_rows_A, SCALE_TDM_COLS), (SCALE_TDM_COLS, 1))
-    sAs = fx.make_view(
-        lds_As.ptr,
-        fx.make_layout(((sc_rows_A, SCALE_TDM_COLS), 1), ((SCALE_TDM_COLS, 1), K_iters * BLOCK_M)),
-    )
-    box_Bs = fx.make_layout((sc_rows_B, SCALE_TDM_COLS), (SCALE_TDM_COLS, 1))
-    sBs = fx.make_view(
-        lds_Bs.ptr,
-        fx.make_layout(((sc_rows_B, SCALE_TDM_COLS), 1), ((SCALE_TDM_COLS, 1), K_iters * BLOCK_N)),
-    )
+    num_groups = (K_iters + SCALE_GROUP - 1) // SCALE_GROUP
+    tail_tiles = K_iters - (num_groups - 1) * SCALE_GROUP
+    lds_As = smem.allocate(fx.Array[fx.Int32, SCALE_STAGES * SCALE_GROUP * BLOCK_M]).peek()
+    lds_Bs = smem.allocate(fx.Array[fx.Int32, SCALE_STAGES * SCALE_GROUP * BLOCK_N]).peek()
 
     # C is staged from LDS offset 0 after the K loop
     if smem.allocated_bytes < c_stage_bytes:
@@ -169,8 +138,8 @@ def as_words(t, rows, words_per_row):
 
     sA_base_ptr = fx.get_iter(sA)
     sB_base_ptr = fx.get_iter(sB)
-    sAs_base_ptr = fx.get_iter(sAs)
-    sBs_base_ptr = fx.get_iter(sBs)
+    sAs_base_ptr = lds_As.ptr
+    sBs_base_ptr = lds_Bs.ptr
 
     # K-loop TDM atoms over this wave's half tile; later K tiles are reached with imm_offset
     tdm_wave = wave % TDM_WAVES
@@ -208,29 +177,54 @@ def lds_half(base_ptr, page):
     blk_cB1w = fx.zipped_divide(cB1w, (BLOCK_N, 16))[None, (bid_n, None)]
     tBsB1w, tBgB1w = fx.rocdl.cdna5.tdm_partition(tdm_B1w, 0, fx.make_layout(1, 1), sB, blk_cB1w)
 
+    # scale group g of a row block is contiguous: 1-D TDM, bumped to group g with imm_offset
+    def scale_ptr(t):
+        src = fx.get_iter(t)
+        return fx.recast_iter(fx.PointerType.get(fx.Int32.ir_type, src.memspace, 4), src)
+
+    sA_blk = fx.add_offset(scale_ptr(A_s), bid_m * K_iters * BLOCK_M)
+    sB_blk = fx.add_offset(scale_ptr(B_s), bid_n * K_iters * BLOCK_N)
+
+    def scale_box(tiles, block_rows):
+        return fx.make_layout(tiles * block_rows, 1)
+
+    # prologue: scale groups 0 and 1, half per wave
     sc_wave = wave % PRO_SCALE_WAVES
-    tdm_As, cAs = fx.rocdl.cdna5.make_tiled_tdm_atom(
-        fx.rocdl.TensorLoad(),
-        sA_rows,
-        box_As,
-        (sc_rows_A, SCALE_TDM_COLS),
-        num_warps=PRO_SCALE_WAVES,
+    pro_tiles = min(SCALE_STAGES * SCALE_GROUP, K_iters)
+    pro_half_A = pro_tiles * BLOCK_M // PRO_SCALE_WAVES
+    pro_half_B = pro_tiles * BLOCK_N // PRO_SCALE_WAVES
+    gAs_pro = fx.make_view(
+        fx.add_offset(sA_blk, sc_wave * pro_half_A), fx.make_layout(pro_half_A, 1)
+    )
+    gBs_pro = fx.make_view(
+        fx.add_offset(sB_blk, sc_wave * pro_half_B), fx.make_layout(pro_half_B, 1)
+    )
+    tdm_As_pro = fx.rocdl.make_tdm_atom(gAs_pro, [None], num_warps=1)
+    tdm_Bs_pro = fx.rocdl.make_tdm_atom(gBs_pro, [None], num_warps=1)
+    lAs_pro = fx.make_view(
+        fx.add_offset(sAs_base_ptr, sc_wave * pro_half_A), fx.make_layout(pro_half_A, 1)
     )
-    blk_cAs = fx.zipped_divide(cAs, (sc_rows_A, SCALE_TDM_COLS))[None, (None, 0)]
-    tAssAs, tAsgAs = fx.rocdl.cdna5.tdm_partition(
-        tdm_As, sc_wave, fx.make_layout(PRO_SCALE_WAVES, 1), sAs, blk_cAs
+    lBs_pro = fx.make_view(
+        fx.add_offset(sBs_base_ptr, sc_wave * pro_half_B), fx.make_layout(pro_half_B, 1)
     )
 
-    tdm_Bs, cBs = fx.rocdl.cdna5.make_tiled_tdm_atom(
-        fx.rocdl.TensorLoad(),
-        sB_rows,
-        box_Bs,
-        (sc_rows_B, SCALE_TDM_COLS),
-        num_warps=PRO_SCALE_WAVES,
+    gAs_grp = fx.make_view(sA_blk, scale_box(SCALE_GROUP, BLOCK_M))
+    gBs_grp = fx.make_view(sB_blk, scale_box(SCALE_GROUP, BLOCK_N))
+    tdm_As_grp = fx.rocdl.make_tdm_atom(gAs_grp, [None], num_warps=1)
+    tdm_Bs_grp = fx.rocdl.make_tdm_atom(gBs_grp, [None], num_warps=1)
+    gAs_tail = fx.make_view(sA_blk, scale_box(tail_tiles, BLOCK_M))
+    gBs_tail = fx.make_view(sB_blk, scale_box(tail_tiles, BLOCK_N))
+    tdm_As_tail = fx.rocdl.make_tdm_atom(gAs_tail, [None], num_warps=1)
+    tdm_Bs_tail = fx.rocdl.make_tdm_atom(gBs_tail, [None], num_warps=1)
+
+    # whole tiles, for K-loop iterations that also load a scale group
+    gA_full = fx.make_view(fx.add_offset(fx.get_iter(A_words), bid_m * BLOCK_M * k_words), box_A)
+    gB_full = fx.make_view(fx.add_offset(fx.get_iter(B_words), bid_n * BLOCK_N * k_words), box_B)
+    tdm_A_full = fx.rocdl.make_tdm_atom(
+        gA_full, [None, None], strides=[fx.Int64(k_words), None], num_warps=1
     )
-    blk_cBs = fx.zipped_divide(cBs, (sc_rows_B, SCALE_TDM_COLS))[None, (None, 0)]
-    tBssBs, tBsgBs = fx.rocdl.cdna5.tdm_partition(
-        tdm_Bs, sc_wave, fx.make_layout(PRO_SCALE_WAVES, 1), sBs, blk_cBs
+    tdm_B_full = fx.rocdl.make_tdm_atom(
+        gB_full, [None, None], strides=[fx.Int64(k_words), None], num_warps=1
     )
 
     # C store: each wave writes C_WAVE_ROWS rows of each half
@@ -283,15 +277,57 @@ def wait_tile(t):
                 else:
                     fx.rocdl.s_wait_tensorcnt(counts[2])
 
-    def issue_tdm(j):
+    # With scale_group g the group's waves load the A tile, A scale group g, the B tile and
+    # B scale group g, still one TDM op per wave; the scale op shares tile j's wait slot
+    def issue_tdm(j, scale_group=None):
         page = j % NUM_STAGES
         group = j % 2
         k_off = fx.Int64(j * BLOCK_K // 2)
-        if wave // TDM_GROUP_WAVES == group:
-            if wave % TDM_GROUP_WAVES < TDM_WAVES:
-                fx.copy(tdm_A, gA_half, lds_half(sA_base_ptr, page), imm_offset=k_off)
-            else:
-                fx.copy(tdm_B, gB_half, lds_half(sB_base_ptr, page), imm_offset=k_off)
+        if scale_group is None:
+            if wave // TDM_GROUP_WAVES == group:
+                if wave % TDM_GROUP_WAVES < TDM_WAVES:
+                    fx.copy(tdm_A, gA_half, lds_half(sA_base_ptr, page), imm_offset=k_off)
+                else:
+                    fx.copy(tdm_B, gB_half, lds_half(sB_base_ptr, page), imm_offset=k_off)
+        else:
+            g = scale_group
+            spage = g % SCALE_STAGES
+            tail = g == num_groups - 1 and tail_tiles < SCALE_GROUP
+            atom_As, atom_Bs = (tdm_As_tail, tdm_Bs_tail) if tail else (tdm_As_grp, tdm_Bs_grp)
+            g_As, g_Bs = (gAs_tail, gBs_tail) if tail else (gAs_grp, gBs_grp)
+            tiles = tail_tiles if tail else SCALE_GROUP
+            l_As = fx.make_view(
+                fx.add_offset(sAs_base_ptr, spage * SCALE_GROUP * BLOCK_M),
+                scale_box(tiles, BLOCK_M),
+            )
+            l_Bs = fx.make_view(
+                fx.add_offset(sBs_base_ptr, spage * SCALE_GROUP * BLOCK_N),
+                scale_box(tiles, BLOCK_N),
+            )
+            s_offA = fx.Int64(g * SCALE_GROUP * BLOCK_M * 4)
+            s_offB = fx.Int64(g * SCALE_GROUP * BLOCK_N * 4)
+            if wave // TDM_GROUP_WAVES == group:
+                gw = wave % TDM_GROUP_WAVES
+                if gw < 2:
+                    if gw == 0:
+                        fx.copy(
+                            tdm_A_full,
+                            gA_full,
+                            fx.make_view(fx.add_offset(sA_base_ptr, page * BLOCK_M * 16), box_A),
+                            imm_offset=k_off,
+                        )
+                    else:
+                        fx.copy(atom_As, g_As, l_As, imm_offset=s_offA)
+                else:
+                    if gw == 2:
+                        fx.copy(
+                            tdm_B_full,
+                            gB_full,
+                            fx.make_view(fx.add_offset(sB_base_ptr, page * BLOCK_N * 16), box_B),
+                            imm_offset=k_off,
+                        )
+                    else:
+                        fx.copy(atom_Bs, g_Bs, l_Bs, imm_offset=s_offB)
         for c in ((0, 1) if group == 0 else (2,)):
             tdm_issued[c].append(j)
 
@@ -310,9 +346,9 @@ def issue_prologue_tdm():
                         fx.copy(tdm_B1w, tBgB1w[None, 1], tBsB1w[None, 1])
         else:
             if wave < 4 + PRO_SCALE_WAVES:
-                fx.copy(tdm_As, tAsgAs[None, bid_m], tAssAs[None, 0])
+                fx.copy(tdm_As_pro, gAs_pro, lAs_pro)
             else:
-                fx.copy(tdm_Bs, tBsgBs[None, bid_n], tBssBs[None, 0])
+                fx.copy(tdm_Bs_pro, gBs_pro, lBs_pro)
 
     def lds_to_regs(k):
         page, r = k % NUM_STAGES, k % 2
@@ -340,7 +376,7 @@ def lds_to_regs(k):
                     fx.add_offset(fx.get_iter(frag_B[r]), 4 * h), fx.make_layout((4, 4), (1, 8))
                 ),
             )
-        s_off = k
+        s_off = (k // SCALE_GROUP) % SCALE_STAGES * SCALE_GROUP + k % SCALE_GROUP
         src_sA = fx.make_view(
             fx.add_offset(sAs_base_ptr, s_off * BLOCK_M + as_lds_base),
             fx.make_layout((4, 2), (1, 4)),
@@ -399,7 +435,12 @@ def gemm_step(k):
             if k == 0 and NUM_STAGES - 1 < K_iters:
                 issue_tdm(NUM_STAGES - 1)
             if k + NUM_STAGES < K_iters:
-                issue_tdm(k + NUM_STAGES)
+                # scale group g >= SCALE_STAGES: issued at k = SCALE_GROUP * (g - 1)
+                g_next = k // SCALE_GROUP + 1
+                if k % SCALE_GROUP == 0 and SCALE_STAGES <= g_next < num_groups:
+                    issue_tdm(k + NUM_STAGES, scale_group=g_next)
+                else:
+                    issue_tdm(k + NUM_STAGES)
             lds_to_regs(k + 1)
             gemm_step(k)
             hot_loop_scheduler()
@@ -552,9 +593,9 @@ def check_mxfp4_shape(m, n, k):
         raise FlyDSLUnsupportedError(
             f"FlyDSL MXFP4 needs M, N multiples of {BLOCK_M}, got M={m}, N={n}"
         )
-    if k % BLOCK_K or k < 2 * BLOCK_K or k > MAX_K:
+    if k % BLOCK_K or k < 2 * BLOCK_K:
         raise FlyDSLUnsupportedError(
-            f"FlyDSL MXFP4 needs K a multiple of {BLOCK_K} in [{2 * BLOCK_K}, {MAX_K}], got K={k}"
+            f"FlyDSL MXFP4 needs K a multiple of {BLOCK_K}, at least {2 * BLOCK_K}, got K={k}"
         )
 
 
@@ -573,7 +614,7 @@ def mxfp4_matmul(
     """
     if a.ndim != 2 or b.ndim != 2 or D.ndim != 2:
         raise ValueError(
-            f"FlyDSL MXFP4 expects rank-2 a, b, D; got "
+            "FlyDSL MXFP4 expects rank-2 a, b, D; got "
             f"{tuple(a.shape)}, {tuple(b.shape)}, {tuple(D.shape)}"
         )
     m, k2 = a.shape
@@ -612,7 +653,6 @@ def mxfp4_matmul(
     "BLOCK_M",
     "BLOCK_N",
     "BLOCK_K",
-    "MAX_K",
     "MXFP4_BLOCK_SIZE",
     "check_mxfp4_shape",
     "mxfp4_matmul",

From 126f1dd02e3dba776c51756fdb762a570ea776ca Mon Sep 17 00:00:00 2001
From: Aristotle Martin <Aristotle.Martin@amd.com>
Date: Tue, 3 Nov 2026 20:44:51 +0000
Subject: [PATCH 04/11] address review comments

---
 tests/pytorch/test_gemm_backends.py           | 18 +++------------
 .../flydsl_kernels/gemm/gemm_wrappers.py      | 17 ++++----------
 .../pytorch/flydsl_kernels/gemm/mxfp4_gemm.py | 22 +++++++++----------
 3 files changed, 18 insertions(+), 39 deletions(-)

diff --git a/tests/pytorch/test_gemm_backends.py b/tests/pytorch/test_gemm_backends.py
index bc488808ef..b6e924f126 100644
--- a/tests/pytorch/test_gemm_backends.py
+++ b/tests/pytorch/test_gemm_backends.py
@@ -33,7 +33,7 @@
 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. MXFP4 (gfx1250) covers TN / NN / NT / TT
+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.
@@ -181,9 +181,6 @@ def _device_capability():
 
 LAYOUTS = ["TN", "NN", "NT"]
 
-# MXFP4 also supports TT (every layout runs the same TN kernel).
-MXFP4_LAYOUTS = ["TN", "NN", "NT", "TT"]
-
 MXFP4_OUT_DTYPES = [torch.float32, torch.bfloat16, torch.float16]
 
 # Same-format FP8 combos, shared by both backends.
@@ -277,8 +274,6 @@ def get_shapes(layout, M, K, N):
         return (K, M), (N, K)
     if layout == "NT":  # A not transposed, B transposed
         return (K, M), (K, N)
-    if layout == "TT":  # both transposed
-        return (M, K), (K, N)
     raise ValueError(f"Unsupported layout: {layout}")
 
 
@@ -291,8 +286,6 @@ def compute_pytorch_reference(A_ref, B_ref, layout):
         return torch.matmul(B_ref, A_ref)
     if layout == "NT":
         return torch.matmul(B_ref.T, A_ref)
-    if layout == "TT":
-        return torch.matmul(B_ref.T, A_ref.T)
     raise ValueError(f"Unsupported layout: {layout}")
 
 
@@ -375,12 +368,7 @@ def _dequantize_mxfp4(data, scale):
 
 
 def create_mxfp4_tensors(M, K, N, layout):
-    """Create MXFP4 inputs and references dequantized from the buffers ``layout`` consumes.
-
-    MXFP4 rowwise and columnwise data are quantized separately (columnwise is the
-    rowwise quantization of the transposed tensor), so the reference must use the
-    buffer the GEMM reads: rowwise for a transposed operand, columnwise otherwise.
-    """
+    """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)
@@ -1157,7 +1145,7 @@ def test_flydsl_mxfp8_unsupported_shape_falls_back():
 
 @requires_flydsl_mxfp4_backend
 @pytest.mark.parametrize("M, K, N", FLYDSL_MXFP4_SHAPES)
-@pytest.mark.parametrize("layout", MXFP4_LAYOUTS)
+@pytest.mark.parametrize("layout", LAYOUTS)
 @pytest.mark.parametrize("out_dtype", MXFP4_OUT_DTYPES, 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."""
diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py
index 38a25ca054..9be2c8512a 100644
--- a/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py
+++ b/transformer_engine/pytorch/flydsl_kernels/gemm/gemm_wrappers.py
@@ -1020,16 +1020,7 @@ def _run_mxfp8(
 
 
 def _select_mxfp4_storage_for_layout(A, transa, B, transb):
-    """Select TE MXFP4 data and scales for a layout.
-
-    MXFP4 columnwise storage is a materialized transpose (as for tensor-wise FP8),
-    so every layout reduces to the TN kernel without transpose reads:
-
-        TN: B.rowwise    [M,K], A.rowwise    [N,K]
-        NN: B.rowwise    [M,K], A.columnwise [N,K]
-        NT: B.columnwise [M,K], A.columnwise [N,K]
-        TT: B.columnwise [M,K], A.rowwise    [N,K]
-    """
+    """Select TE MXFP4 data and scales for a layout."""
 
     def select(t, rowwise, name):
         orientation = "rowwise" if rowwise else "columnwise"
@@ -1066,7 +1057,7 @@ def _run_mxfp4(
     *,
     output_dtype: torch.dtype,
 ):
-    """Dispatch MXFP4 (gfx1250) through the single TN kernel for TN/NN/NT/TT."""
+    """Dispatch MXFP4 (gfx1250) through the single TN kernel for TN/NN/NT."""
     layout = f"{'T' if transa else 'N'}{'T' if transb else 'N'}"
     A_data, A_scale, B_data, B_scale, B_payload = _select_mxfp4_storage_for_layout(
         A, transa, B, transb
@@ -1344,7 +1335,7 @@ def te_generic_gemm_flydsl(
       - NN: transa=False, transb=False
       - NT: transa=False, transb=True
 
-    TT is intentionally rejected, except for MXFP4.
+    TT is intentionally rejected.
 
     Supported dtypes:
       - MXFP4 E2M1 A/B (gfx1250) with FP16, BF16, or FP32 output
@@ -1379,7 +1370,7 @@ def te_generic_gemm_flydsl(
     a_kind, _ = _classify_input(A)
     b_kind, _ = _classify_input(B)
 
-    if transa and transb and a_kind != "mxfp4":
+    if transa and transb:
         raise FlyDSLUnsupportedError("FlyDSL GEMM does not support transa=True, transb=True (TT)")
 
     _validate_common_epilogue(
diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
index 523c05ecd4..24e7b820f9 100644
--- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
+++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
@@ -4,7 +4,7 @@
 """FlyDSL MXFP4 TN GEMM for gfx1250.
 
 D[M, N] = a[M, K] @ b[N, K]^T with K-contiguous packed FP4 operands and E8M0 scales along K. TE's
-MXFP4 columnwise storage is a materialized transpose, so this one kernel serves TN/NN/NT/TT.
+MXFP4 columnwise storage is a materialized transpose, so this one kernel serves TN/NN/NT.
 Tiles are staged with TDM into a 3-stage LDS pipeline and C is written back through LDS by TDM.
 """
 
@@ -93,7 +93,7 @@ def as_words(t, rows, words_per_row):
     # wave-uniform, so TDM descriptors derived from it stay in SGPRs
     lane = tid % 32
     wave = fx.Int32(fx.rocdl.readfirstlane(fx.Int32.ir_type, tid // 32))
-    l16, khalf = lane % 16, lane // 16
+    lane_mod16, khalf = lane % 16, lane // 16
     wm, wn = wave % 2, wave // 2
 
     k_words = K // 8
@@ -101,15 +101,15 @@ def as_words(t, rows, words_per_row):
     B_words = as_words(B, N, k_words)
 
     # physical chunks of this lane's logical chunks khalf and khalf + 2 (same for all its rows)
-    swz = (l16 // 4) % SWZ_CHUNKS
+    swz = (lane_mod16 // 4) % SWZ_CHUNKS
     chunk_lo = khalf ^ swz
     chunk_hi = chunk_lo ^ 2
-    a_lds_lo = (wm * 16 + l16) * 16 + 4 * chunk_lo
-    a_lds_hi = (wm * 16 + l16) * 16 + 4 * chunk_hi
-    b_lds_lo = (wn * 16 + l16) * 16 + 4 * chunk_lo
-    b_lds_hi = (wn * 16 + l16) * 16 + 4 * chunk_hi
-    as_lds_base = (wm * 16 + l16) * (BLOCK_M // SCALE_LANE_ROWS_A)
-    bs_lds_base = (wn * 16 + l16) * (BLOCK_N // SCALE_LANE_ROWS_B)
+    a_lds_lo = (wm * 16 + lane_mod16) * 16 + 4 * chunk_lo
+    a_lds_hi = (wm * 16 + lane_mod16) * 16 + 4 * chunk_hi
+    b_lds_lo = (wn * 16 + lane_mod16) * 16 + 4 * chunk_lo
+    b_lds_hi = (wn * 16 + lane_mod16) * 16 + 4 * chunk_hi
+    as_lds_base = (wm * 16 + lane_mod16) * (BLOCK_M // SCALE_LANE_ROWS_A)
+    bs_lds_base = (wn * 16 + lane_mod16) * (BLOCK_N // SCALE_LANE_ROWS_B)
 
     # tile k lives in LDS page k % NUM_STAGES
     smem = fx.SharedAllocator(static=False)
@@ -453,10 +453,10 @@ def gemm_step(k):
         frag_c = fx.make_fragment_like(frag_acc, out_fx)
         frag_c.store(frag_acc.load().to(out_fx))
 
-    # accumulator (v, j, i) is tile row wm*16 + 8*khalf + v + 32*i, column wn*16 + l16 + 64*j
+    # accumulator (v, j, i) is tile row wm*16 + 8*khalf + v + 32*i, column wn*16 + lane_mod16 + 64*j
     fx.rocdl.s_wait_dscnt(0)
     fx.barrier()
-    lane_c = (wm * 16 + 8 * khalf) * BLOCK_N + wn * 16 + l16
+    lane_c = (wm * 16 + 8 * khalf) * BLOCK_N + wn * 16 + lane_mod16
     for h in range_constexpr(2):
         if h > 0:
             fx.rocdl.s_wait_tensorcnt(0)

From a0a43bd114d73017f842af054fad9c9508c328ff Mon Sep 17 00:00:00 2001
From: Aristotle Martin <Aristotle.Martin@amd.com>
Date: Wed, 4 Nov 2026 12:32:38 +0000
Subject: [PATCH 05/11] address review comment

---
 tests/pytorch/test_gemm_backends.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/pytorch/test_gemm_backends.py b/tests/pytorch/test_gemm_backends.py
index b6e924f126..29180549a7 100644
--- a/tests/pytorch/test_gemm_backends.py
+++ b/tests/pytorch/test_gemm_backends.py
@@ -181,8 +181,6 @@ def _device_capability():
 
 LAYOUTS = ["TN", "NN", "NT"]
 
-MXFP4_OUT_DTYPES = [torch.float32, torch.bfloat16, torch.float16]
-
 # Same-format FP8 combos, shared by both backends.
 FP8_FORMAT_COMBOS = [
     (tex.DType.kFloat8E4M3, tex.DType.kFloat8E4M3),
@@ -1146,7 +1144,9 @@ def test_flydsl_mxfp8_unsupported_shape_falls_back():
 @requires_flydsl_mxfp4_backend
 @pytest.mark.parametrize("M, K, N", FLYDSL_MXFP4_SHAPES)
 @pytest.mark.parametrize("layout", LAYOUTS)
-@pytest.mark.parametrize("out_dtype", MXFP4_OUT_DTYPES, ids=["fp32", "bf16", "fp16"])
+@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)

From f7098494acbf5df2634525f16f2f54dc68d1de71 Mon Sep 17 00:00:00 2001
From: Aristotle Martin <Aristotle.Martin@amd.com>
Date: Thu, 5 Nov 2026 02:16:26 +0000
Subject: [PATCH 06/11] Enable MXFP4 recipe on gfx1250 via FlyDSL

---
 tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py  |  7 ++++++
 .../pytorch/mxfp4/test_mxfp4_module_exact.py  | 18 +++++++++++--
 tests/pytorch/test_gemm_backends.py           |  3 +--
 .../triton_kernels/test_grouped_gemm_mxfp4.py |  4 +++
 .../pytorch/cpp_extensions/gemm.py            | 22 +++++++---------
 transformer_engine/pytorch/quantization.py    | 25 ++++++++++++++-----
 6 files changed, 56 insertions(+), 23 deletions(-)

diff --git a/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py b/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py
index 116a566bf0..f084270318 100644
--- a/tests/pytorch/mxfp4/test_mxfp4_gemm_exact.py
+++ b/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"])
diff --git a/tests/pytorch/mxfp4/test_mxfp4_module_exact.py b/tests/pytorch/mxfp4/test_mxfp4_module_exact.py
index d3d19784df..1adca56459 100644
--- a/tests/pytorch/mxfp4/test_mxfp4_module_exact.py
+++ b/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,
diff --git a/tests/pytorch/test_gemm_backends.py b/tests/pytorch/test_gemm_backends.py
index 29180549a7..68218ed346 100644
--- a/tests/pytorch/test_gemm_backends.py
+++ b/tests/pytorch/test_gemm_backends.py
@@ -1167,8 +1167,7 @@ def test_flydsl_vs_pytorch_mxfp4(M, K, N, layout, out_dtype):
 def test_flydsl_mxfp4_unsupported_shape_raises():
     """MXFP4 configs FlyDSL cannot run raise FlyDSLUnsupportedError.
 
-    That is the type ``general_gemm`` catches to fall back. The native MXFP4
-    GEMM does not run on gfx1250, so the fallback itself is not exercised.
+    ``general_gemm`` raises it, as gfx1250 has no other MXFP4 GEMM.
     """
     from transformer_engine.pytorch.flydsl_kernels.gemm import (
         FlyDSLUnsupportedError,
diff --git a/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py b/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py
index f255b6ad0b..d1bc63a46f 100644
--- a/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py
+++ b/tests/pytorch/triton_kernels/test_grouped_gemm_mxfp4.py
@@ -7,6 +7,7 @@
 import torch
 
 from transformer_engine.pytorch.quantization import check_mxfp4_support
+from transformer_engine.pytorch.utils import get_device_compute_capability
 
 _MXFP4_OK, _MXFP4_REASON = check_mxfp4_support()
 
@@ -25,6 +26,9 @@
 
 pytestmark = [
     pytest.mark.skipif(not _MXFP4_OK, reason=f"MXFP4 unsupported: {_MXFP4_REASON}"),
+    pytest.mark.skipif(
+        _MXFP4_OK and get_device_compute_capability() != (9, 5), reason="requires gfx950"
+    ),
 ]
 
 DTYPE = torch.bfloat16
diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py
index 760714b0b2..92dac22680 100644
--- a/transformer_engine/pytorch/cpp_extensions/gemm.py
+++ b/transformer_engine/pytorch/cpp_extensions/gemm.py
@@ -590,15 +590,13 @@ def general_gemm(
     # Use bfloat16 as default bias_dtype
     bias_dtype = TE_DType[torch.bfloat16 if bias is None else bias.dtype]
 
-    # MXFP4 GEMM: FlyDSL on gfx1250 if NVTE_GEMM_BACKEND=FLYDSL, else 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)
-            and get_gemm_backend() == "flydsl"
-        ):
-            result = _try_flydsl_gemm(
+        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,
@@ -622,11 +620,9 @@ def general_gemm(
                 alpha=alpha,
                 beta=beta,
             )
-            if result is not None:
-                out, bias_grad, gelu_input, extra_output = result
-                if debug_quantizer is not None:
-                    out = debug_quantizer.process_gemm_output(out)
-                return out, bias_grad, gelu_input, extra_output
+            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,
diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py
index 176ee9e071..8e42ae47ea 100644
--- a/transformer_engine/pytorch/quantization.py
+++ b/transformer_engine/pytorch/quantization.py
@@ -218,9 +218,15 @@ def _compute_mxfp4_support() -> Tuple[bool, str]:
     """Return if mxfp4 support is available"""
     if IS_HIP_EXTENSION:
         gpu_arch = get_device_compute_capability()
-        if gpu_arch in ((9, 5),):  # TODO: enable for gfx1250 when GEMM is available
+        if gpu_arch == (9, 5):
             return True, ""
-        return False, "Device arch gfx95x is required for MXFP4 execution."
+        if gpu_arch == (12, 5):
+            try:
+                from .flydsl_kernels.gemm import te_generic_gemm_flydsl  # noqa: F401
+            except ImportError as e:
+                return False, f"MXFP4 on gfx1250 requires the FlyDSL GEMM backend: {e}"
+            return True, ""
+        return False, "Device arch gfx950 or gfx1250 is required for MXFP4 execution."
     return False, "Only ROCm supports MXFP4"
 
 
@@ -1815,19 +1821,26 @@ def make_quantizers(self) -> list:
         # scales are plain (VEC32_UE8M0) by default, but can optionally be pre-swizzled into
         # the 32x8 tile order that hipBLASLt's BLK32_UE8M0_32_8_EXT (mode 1001) reads in place
         # -- opt in via the recipe's use_swizzled_scales flag. The AITER a4w4 backend always
-        # needs the 16x16 weight shuffle and swizzled scales.
-        use_hipblaslt = bool(int(os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0")))
+        # needs the 16x16 weight shuffle and swizzled scales. FlyDSL (gfx1250) takes plain data
+        # and scales.
+        use_flydsl = IS_HIP_EXTENSION and get_device_compute_capability() == (12, 5)
+        use_hipblaslt = not use_flydsl and bool(
+            int(os.environ.get("NVTE_ROCM_USE_HIPBLASLT_MXFP4", "0"))
+        )
         use_swizzled = use_hipblaslt and self.recipe.use_swizzled_scales
         # AITER path swizzles scales; hipBLASLt path swizzles only when the recipe opts in.
         # FP4 data shuffle stays off on the hipBLASLt path regardless
-        swizzled_scales = use_swizzled if use_hipblaslt else True
+        if use_flydsl:
+            swizzled_scales = False
+        else:
+            swizzled_scales = use_swizzled if use_hipblaslt else True
 
         if self.mode == "forward":
 
             def _make_quantizer(idx: int):
                 is_activation = idx % 3 == 0
                 is_weight = idx % 3 == 1
-                if use_hipblaslt:
+                if use_hipblaslt or use_flydsl:
                     shuffle_rowwise_data = False
                     shuffle_columnwise_data = False
                 elif is_activation:

From ba96d7fde0a43a12db77f968fd7a1aa92a15eb18 Mon Sep 17 00:00:00 2001
From: Aristotle Martin <Aristotle.Martin@amd.com>
Date: Thu, 5 Nov 2026 02:40:01 +0000
Subject: [PATCH 07/11] FlyDSL MXFP4: replace operand swizzle pass with TDM LDS
 padding

---
 .../pytorch/flydsl_kernels/gemm/mxfp4_gemm.py | 92 +++++++++----------
 1 file changed, 42 insertions(+), 50 deletions(-)

diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
index 24e7b820f9..da9d04270a 100644
--- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
+++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
@@ -8,8 +8,6 @@
 Tiles are staged with TDM into a 3-stage LDS pipeline and C is written back through LDS by TDM.
 """
 
-import functools
-
 import torch
 
 import flydsl.compiler as flyc
@@ -37,8 +35,9 @@
 LDS_B_START = 64 * 1024
 C_HALF_ROWS = BLOCK_M // 2  # C is stored in two halves
 C_WAVE_ROWS = C_HALF_ROWS // NUM_WAVES
-# Logical 16 B chunk c of row r is stored at chunk c ^ ((r >> 2) & 3) to avoid LDS bank conflicts
-SWZ_CHUNKS = 4
+# A/B LDS rows: 16 words padded to 20 by the TDM, so ds_load_b128 reads avoid bank conflicts
+ROW_WORDS = 16
+LDS_ROW_WORDS = 20
 SCHED_TAIL_MMAS = 8
 # Each lane's scale words are made contiguous so they load as ds_load_b128
 SCALE_LANE_ROWS_A = 32
@@ -100,31 +99,32 @@ def as_words(t, rows, words_per_row):
     A_words = as_words(A, M, k_words)
     B_words = as_words(B, N, k_words)
 
-    # physical chunks of this lane's logical chunks khalf and khalf + 2 (same for all its rows)
-    swz = (lane_mod16 // 4) % SWZ_CHUNKS
-    chunk_lo = khalf ^ swz
-    chunk_hi = chunk_lo ^ 2
-    a_lds_lo = (wm * 16 + lane_mod16) * 16 + 4 * chunk_lo
-    a_lds_hi = (wm * 16 + lane_mod16) * 16 + 4 * chunk_hi
-    b_lds_lo = (wn * 16 + lane_mod16) * 16 + 4 * chunk_lo
-    b_lds_hi = (wn * 16 + lane_mod16) * 16 + 4 * chunk_hi
+    # this lane's 16 B chunks khalf and khalf + 2
+    a_lds_lo = (wm * 16 + lane_mod16) * LDS_ROW_WORDS + 4 * khalf
+    a_lds_hi = a_lds_lo + 8
+    b_lds_lo = (wn * 16 + lane_mod16) * LDS_ROW_WORDS + 4 * khalf
+    b_lds_hi = b_lds_lo + 8
     as_lds_base = (wm * 16 + lane_mod16) * (BLOCK_M // SCALE_LANE_ROWS_A)
     bs_lds_base = (wn * 16 + lane_mod16) * (BLOCK_N // SCALE_LANE_ROWS_B)
 
     # tile k lives in LDS page k % NUM_STAGES
     smem = fx.SharedAllocator(static=False)
-    lds_A = smem.allocate(fx.Array[fx.Int32, NUM_STAGES * BLOCK_M * 16]).peek()
+    PAGE_A, PAGE_B = BLOCK_M * LDS_ROW_WORDS, BLOCK_N * LDS_ROW_WORDS
+    lds_A = smem.allocate(fx.Array[fx.Int32, NUM_STAGES * PAGE_A]).peek()
     # B starts at 64 KB so all its pages are reachable with immediate offsets from one base register
-    smem.allocate(fx.Array[fx.Int32, LDS_B_START // 4 - NUM_STAGES * BLOCK_M * 16])
-    lds_B = smem.allocate(fx.Array[fx.Int32, NUM_STAGES * BLOCK_N * 16]).peek()
-    box_A = fx.make_layout((BLOCK_M, 16), (16, 1))
+    smem.allocate(fx.Array[fx.Int32, LDS_B_START // 4 - NUM_STAGES * PAGE_A])
+    lds_B = smem.allocate(fx.Array[fx.Int32, NUM_STAGES * PAGE_B]).peek()
+    box_A = fx.make_layout((BLOCK_M, ROW_WORDS), (LDS_ROW_WORDS, 1))
     sA = fx.make_view(
-        lds_A.ptr, fx.make_layout(((BLOCK_M, 16), NUM_STAGES), ((16, 1), BLOCK_M * 16))
+        lds_A.ptr,
+        fx.make_layout(((BLOCK_M, ROW_WORDS), NUM_STAGES), ((LDS_ROW_WORDS, 1), PAGE_A)),
     )
-    box_B = fx.make_layout((BLOCK_N, 16), (16, 1))
+    box_B = fx.make_layout((BLOCK_N, ROW_WORDS), (LDS_ROW_WORDS, 1))
     sB = fx.make_view(
-        lds_B.ptr, fx.make_layout(((BLOCK_N, 16), NUM_STAGES), ((16, 1), BLOCK_N * 16))
+        lds_B.ptr,
+        fx.make_layout(((BLOCK_N, ROW_WORDS), NUM_STAGES), ((LDS_ROW_WORDS, 1), PAGE_B)),
     )
+    pad = dict(pad_interval=ROW_WORDS, pad_amount=LDS_ROW_WORDS - ROW_WORDS)
 
     num_groups = (K_iters + SCALE_GROUP - 1) // SCALE_GROUP
     tail_tiles = K_iters - (num_groups - 1) * SCALE_GROUP
@@ -154,15 +154,16 @@ def as_words(t, rows, words_per_row):
         half_box,
     )
     tdm_A = fx.rocdl.make_tdm_atom(
-        gA_half, [None, None], strides=[fx.Int64(k_words), None], num_warps=1
+        gA_half, [None, None], strides=[fx.Int64(k_words), None], num_warps=1, **pad
     )
     tdm_B = fx.rocdl.make_tdm_atom(
-        gB_half, [None, None], strides=[fx.Int64(k_words), None], num_warps=1
+        gB_half, [None, None], strides=[fx.Int64(k_words), None], num_warps=1, **pad
     )
 
     def lds_half(base_ptr, page):
         return fx.make_view(
-            fx.add_offset(base_ptr, page * BLOCK_M * 16 + tdm_wave * HALF_ROWS * 16), half_box
+            fx.add_offset(base_ptr, page * PAGE_A + tdm_wave * HALF_ROWS * LDS_ROW_WORDS),
+            half_box,
         )
 
     # prologue atoms: whole tiles and scales
@@ -221,10 +222,10 @@ def scale_box(tiles, block_rows):
     gA_full = fx.make_view(fx.add_offset(fx.get_iter(A_words), bid_m * BLOCK_M * k_words), box_A)
     gB_full = fx.make_view(fx.add_offset(fx.get_iter(B_words), bid_n * BLOCK_N * k_words), box_B)
     tdm_A_full = fx.rocdl.make_tdm_atom(
-        gA_full, [None, None], strides=[fx.Int64(k_words), None], num_warps=1
+        gA_full, [None, None], strides=[fx.Int64(k_words), None], num_warps=1, **pad
     )
     tdm_B_full = fx.rocdl.make_tdm_atom(
-        gB_full, [None, None], strides=[fx.Int64(k_words), None], num_warps=1
+        gB_full, [None, None], strides=[fx.Int64(k_words), None], num_warps=1, **pad
     )
 
     # C store: each wave writes C_WAVE_ROWS rows of each half
@@ -313,7 +314,7 @@ def issue_tdm(j, scale_group=None):
                         fx.copy(
                             tdm_A_full,
                             gA_full,
-                            fx.make_view(fx.add_offset(sA_base_ptr, page * BLOCK_M * 16), box_A),
+                            fx.make_view(fx.add_offset(sA_base_ptr, page * PAGE_A), box_A),
                             imm_offset=k_off,
                         )
                     else:
@@ -323,7 +324,7 @@ def issue_tdm(j, scale_group=None):
                         fx.copy(
                             tdm_B_full,
                             gB_full,
-                            fx.make_view(fx.add_offset(sB_base_ptr, page * BLOCK_N * 16), box_B),
+                            fx.make_view(fx.add_offset(sB_base_ptr, page * PAGE_B), box_B),
                             imm_offset=k_off,
                         )
                     else:
@@ -352,15 +353,14 @@ def issue_prologue_tdm():
 
     def lds_to_regs(k):
         page, r = k % NUM_STAGES, k % 2
-        # one copy per swizzled chunk; registers stay in logical chunk order
         for h, a_off, b_off in ((0, a_lds_lo, b_lds_lo), (1, a_lds_hi, b_lds_hi)):
             src_A = fx.make_view(
-                fx.add_offset(sA_base_ptr, page * BLOCK_M * 16 + a_off),
-                fx.make_layout((4, 8), (1, 32 * 16)),
+                fx.add_offset(sA_base_ptr, page * PAGE_A + a_off),
+                fx.make_layout((4, 8), (1, 32 * LDS_ROW_WORDS)),
             )
             src_B = fx.make_view(
-                fx.add_offset(sB_base_ptr, page * BLOCK_N * 16 + b_off),
-                fx.make_layout((4, 4), (1, 64 * 16)),
+                fx.add_offset(sB_base_ptr, page * PAGE_B + b_off),
+                fx.make_layout((4, 4), (1, 64 * LDS_ROW_WORDS)),
             )
             fx.copy(
                 copy_atom_i32,
@@ -554,19 +554,6 @@ def shuffle_scales(
     ).launch(grid=(grid_r, words_per_row, 1), block=(block_rows, 1, 1), stream=stream)
 
 
-@functools.lru_cache(maxsize=None)
-def _swizzle_index(rows: int, K: int, device: torch.device) -> torch.Tensor:
-    swz = (torch.arange(rows, device=device) >> 2) & (SWZ_CHUNKS - 1)
-    src_chunk = torch.arange(SWZ_CHUNKS, device=device)[None, :] ^ swz[:, None]
-    return src_chunk[:, None, :, None].expand(rows, K // BLOCK_K, SWZ_CHUNKS, 4)
-
-
-def swizzle_operand(q, rows, K):
-    """Copy of packed FP4 operand q [rows, K/2] with the LDS chunk swizzle applied."""
-    words = q.view(torch.int32).view(rows, K // BLOCK_K, SWZ_CHUNKS, 4)
-    return words.gather(2, _swizzle_index(rows, K, q.device)).view(torch.uint8).view(rows, K // 2)
-
-
 def _shuffle_scale_words(scale, rows, K, block_rows, lane_rows, stream):
     """Reorder TE E8M0 scales (possibly padded) into the kernel's iteration-major order."""
     if scale.dtype != torch.uint8 or scale.ndim != 2:
@@ -636,15 +623,20 @@ def mxfp4_matmul(
 
     if stream is None:
         stream = torch.cuda.current_stream(a.device)
-    a_u8 = a.view(torch.uint8)
-    b_u8 = b.view(torch.uint8)
     a_scale_k = _shuffle_scale_words(a_scale, m, k, BLOCK_M, SCALE_LANE_ROWS_A, stream)
     b_scale_k = _shuffle_scale_words(b_scale, n, k, BLOCK_N, SCALE_LANE_ROWS_B, stream)
-    with torch.cuda.stream(stream):
-        a_k = swizzle_operand(a_u8, m, k)
-        b_k = swizzle_operand(b_u8, n, k)
     launch_gemm(
-        a_k, b_k, a_scale_k, b_scale_k, D, m, n, k, k // BLOCK_K, OUT_KINDS[D.dtype], stream=stream
+        a.view(torch.uint8),
+        b.view(torch.uint8),
+        a_scale_k,
+        b_scale_k,
+        D,
+        m,
+        n,
+        k,
+        k // BLOCK_K,
+        OUT_KINDS[D.dtype],
+        stream=stream,
     )
     return D
 

From 506d2cd1203e2e5e030ae1f8e8c3a167ab4d7918 Mon Sep 17 00:00:00 2001
From: Aristotle Martin <Aristotle.Martin@amd.com>
Date: Thu, 5 Nov 2026 02:55:45 +0000
Subject: [PATCH 08/11] FlyDSL MXFP4: read TE's row-major scales directly, drop
 per-call scale shuffle

---
 .../pytorch/flydsl_kernels/gemm/mxfp4_gemm.py | 225 ++++++------------
 1 file changed, 77 insertions(+), 148 deletions(-)

diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
index da9d04270a..093472f7da 100644
--- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
+++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
@@ -17,7 +17,6 @@
 from .exceptions import FlyDSLUnsupportedError
 
 MXFP4_BLOCK_SIZE = 32
-SCALE_WORD_K = 4 * MXFP4_BLOCK_SIZE
 BLOCK_M = 256
 BLOCK_N = 256
 BLOCK_K = 128
@@ -39,27 +38,12 @@
 ROW_WORDS = 16
 LDS_ROW_WORDS = 20
 SCHED_TAIL_MMAS = 8
-# Each lane's scale words are made contiguous so they load as ds_load_b128
-SCALE_LANE_ROWS_A = 32
-SCALE_LANE_ROWS_B = 64
+# Scales are TE's row-major words (one per row and K tile), kept in LDS as [row][SCALE_STAGES *
+# SCALE_GROUP] with an odd row pitch so the per-row reads avoid bank conflicts
+LDS_SCALE_PITCH = 17
 OUT_KINDS = {torch.float32: 0, torch.bfloat16: 1, torch.float16: 2}
 
 
-# Scale words (one BLOCK_K tile of a row) stored [rows // block_rows][K_iters][block_rows],
-# viewed by logical (row, k_tile)
-def iter_major_scale_words(t, rows, words_per_row, block_rows, lane_rows):
-    src = fx.get_iter(t)
-    ptr = fx.recast_iter(fx.PointerType.get(fx.Int32.ir_type, src.memspace, 4), src)
-    reps = block_rows // lane_rows
-    return fx.make_view(
-        ptr,
-        fx.make_layout(
-            (((lane_rows, reps), rows // block_rows), words_per_row),
-            (((reps, 1), words_per_row * block_rows), block_rows),
-        ),
-    )
-
-
 @flyc.kernel
 def gemm_kernel(
     A: fx.Tensor,
@@ -70,6 +54,8 @@ def gemm_kernel(
     M: fx.Int32,
     N: fx.Int32,
     K: fx.Int32,
+    As_sw: fx.Int32,
+    Bs_sw: fx.Int32,
     K_iters: fx.Constexpr[int],
     OUT_KIND: fx.Constexpr[int],
 ):
@@ -104,8 +90,8 @@ def as_words(t, rows, words_per_row):
     a_lds_hi = a_lds_lo + 8
     b_lds_lo = (wn * 16 + lane_mod16) * LDS_ROW_WORDS + 4 * khalf
     b_lds_hi = b_lds_lo + 8
-    as_lds_base = (wm * 16 + lane_mod16) * (BLOCK_M // SCALE_LANE_ROWS_A)
-    bs_lds_base = (wn * 16 + lane_mod16) * (BLOCK_N // SCALE_LANE_ROWS_B)
+    as_lds_base = (wm * 16 + lane_mod16) * LDS_SCALE_PITCH
+    bs_lds_base = (wn * 16 + lane_mod16) * LDS_SCALE_PITCH
 
     # tile k lives in LDS page k % NUM_STAGES
     smem = fx.SharedAllocator(static=False)
@@ -128,8 +114,8 @@ def as_words(t, rows, words_per_row):
 
     num_groups = (K_iters + SCALE_GROUP - 1) // SCALE_GROUP
     tail_tiles = K_iters - (num_groups - 1) * SCALE_GROUP
-    lds_As = smem.allocate(fx.Array[fx.Int32, SCALE_STAGES * SCALE_GROUP * BLOCK_M]).peek()
-    lds_Bs = smem.allocate(fx.Array[fx.Int32, SCALE_STAGES * SCALE_GROUP * BLOCK_N]).peek()
+    lds_As = smem.allocate(fx.Array[fx.Int32, BLOCK_M * LDS_SCALE_PITCH]).peek()
+    lds_Bs = smem.allocate(fx.Array[fx.Int32, BLOCK_N * LDS_SCALE_PITCH]).peek()
 
     # C is staged from LDS offset 0 after the K loop
     if smem.allocated_bytes < c_stage_bytes:
@@ -178,45 +164,47 @@ def lds_half(base_ptr, page):
     blk_cB1w = fx.zipped_divide(cB1w, (BLOCK_N, 16))[None, (bid_n, None)]
     tBsB1w, tBgB1w = fx.rocdl.cdna5.tdm_partition(tdm_B1w, 0, fx.make_layout(1, 1), sB, blk_cB1w)
 
-    # scale group g of a row block is contiguous: 1-D TDM, bumped to group g with imm_offset
+    # scale group g of a row block is a [rows][SCALE_GROUP] box at column g * SCALE_GROUP
     def scale_ptr(t):
         src = fx.get_iter(t)
         return fx.recast_iter(fx.PointerType.get(fx.Int32.ir_type, src.memspace, 4), src)
 
-    sA_blk = fx.add_offset(scale_ptr(A_s), bid_m * K_iters * BLOCK_M)
-    sB_blk = fx.add_offset(scale_ptr(B_s), bid_n * K_iters * BLOCK_N)
-
-    def scale_box(tiles, block_rows):
-        return fx.make_layout(tiles * block_rows, 1)
+    sA_blk = fx.add_offset(scale_ptr(A_s), bid_m * BLOCK_M * As_sw)
+    sB_blk = fx.add_offset(scale_ptr(B_s), bid_n * BLOCK_N * Bs_sw)
+
+    def scale_atom(view, sw, tiles, width):
+        return fx.rocdl.make_tdm_atom(
+            view,
+            [None, tiles],
+            strides=[fx.Int64(sw), None],
+            num_warps=1,
+            pad_interval=width,
+            pad_amount=LDS_SCALE_PITCH - width,
+        )
 
-    # prologue: scale groups 0 and 1, half per wave
+    # prologue: scale groups 0 and 1, half the rows per wave
     sc_wave = wave % PRO_SCALE_WAVES
-    pro_tiles = min(SCALE_STAGES * SCALE_GROUP, K_iters)
-    pro_half_A = pro_tiles * BLOCK_M // PRO_SCALE_WAVES
-    pro_half_B = pro_tiles * BLOCK_N // PRO_SCALE_WAVES
-    gAs_pro = fx.make_view(
-        fx.add_offset(sA_blk, sc_wave * pro_half_A), fx.make_layout(pro_half_A, 1)
-    )
-    gBs_pro = fx.make_view(
-        fx.add_offset(sB_blk, sc_wave * pro_half_B), fx.make_layout(pro_half_B, 1)
-    )
-    tdm_As_pro = fx.rocdl.make_tdm_atom(gAs_pro, [None], num_warps=1)
-    tdm_Bs_pro = fx.rocdl.make_tdm_atom(gBs_pro, [None], num_warps=1)
-    lAs_pro = fx.make_view(
-        fx.add_offset(sAs_base_ptr, sc_wave * pro_half_A), fx.make_layout(pro_half_A, 1)
-    )
-    lBs_pro = fx.make_view(
-        fx.add_offset(sBs_base_ptr, sc_wave * pro_half_B), fx.make_layout(pro_half_B, 1)
-    )
-
-    gAs_grp = fx.make_view(sA_blk, scale_box(SCALE_GROUP, BLOCK_M))
-    gBs_grp = fx.make_view(sB_blk, scale_box(SCALE_GROUP, BLOCK_N))
-    tdm_As_grp = fx.rocdl.make_tdm_atom(gAs_grp, [None], num_warps=1)
-    tdm_Bs_grp = fx.rocdl.make_tdm_atom(gBs_grp, [None], num_warps=1)
-    gAs_tail = fx.make_view(sA_blk, scale_box(tail_tiles, BLOCK_M))
-    gBs_tail = fx.make_view(sB_blk, scale_box(tail_tiles, BLOCK_N))
-    tdm_As_tail = fx.rocdl.make_tdm_atom(gAs_tail, [None], num_warps=1)
-    tdm_Bs_tail = fx.rocdl.make_tdm_atom(gBs_tail, [None], num_warps=1)
+    PRO_ROWS = BLOCK_M // PRO_SCALE_WAVES
+    PRO_COLS = SCALE_STAGES * SCALE_GROUP
+    pro_box = fx.make_layout((PRO_ROWS, PRO_COLS), (PRO_COLS, 1))
+    gAs_pro = fx.make_view(fx.add_offset(sA_blk, sc_wave * PRO_ROWS * As_sw), pro_box)
+    gBs_pro = fx.make_view(fx.add_offset(sB_blk, sc_wave * PRO_ROWS * Bs_sw), pro_box)
+    tdm_As_pro = scale_atom(gAs_pro, As_sw, K_iters, PRO_COLS)
+    tdm_Bs_pro = scale_atom(gBs_pro, Bs_sw, K_iters, PRO_COLS)
+    pro_lds_off = sc_wave * PRO_ROWS * LDS_SCALE_PITCH
+    lAs_pro = fx.make_view(fx.add_offset(sAs_base_ptr, pro_lds_off), pro_box)
+    lBs_pro = fx.make_view(fx.add_offset(sBs_base_ptr, pro_lds_off), pro_box)
+
+    grp_box = fx.make_layout((BLOCK_M, SCALE_GROUP), (SCALE_GROUP, 1))
+    gAs_grp = fx.make_view(sA_blk, grp_box)
+    gBs_grp = fx.make_view(sB_blk, grp_box)
+    tdm_As_grp = scale_atom(gAs_grp, As_sw, None, SCALE_GROUP)
+    tdm_Bs_grp = scale_atom(gBs_grp, Bs_sw, None, SCALE_GROUP)
+    tail_col = (num_groups - 1) * SCALE_GROUP
+    gAs_tail = fx.make_view(fx.add_offset(sA_blk, tail_col), grp_box)
+    gBs_tail = fx.make_view(fx.add_offset(sB_blk, tail_col), grp_box)
+    tdm_As_tail = scale_atom(gAs_tail, As_sw, tail_tiles, SCALE_GROUP)
+    tdm_Bs_tail = scale_atom(gBs_tail, Bs_sw, tail_tiles, SCALE_GROUP)
 
     # whole tiles, for K-loop iterations that also load a scale group
     gA_full = fx.make_view(fx.add_offset(fx.get_iter(A_words), bid_m * BLOCK_M * k_words), box_A)
@@ -250,6 +238,7 @@ def scale_box(tiles, block_rows):
     thr_mma = tiled_mma.thr_slice(tid)
 
     copy_atom_i32 = fx.make_copy_atom(fx.UniversalCopy128b(), fx.Int32)
+    copy_atom_s32 = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32)
     copy_atom_c = fx.make_copy_atom(
         fx.UniversalCopy32b() if OUT_KIND == 0 else fx.UniversalCopy16b(), out_fx
     )
@@ -296,17 +285,9 @@ def issue_tdm(j, scale_group=None):
             tail = g == num_groups - 1 and tail_tiles < SCALE_GROUP
             atom_As, atom_Bs = (tdm_As_tail, tdm_Bs_tail) if tail else (tdm_As_grp, tdm_Bs_grp)
             g_As, g_Bs = (gAs_tail, gBs_tail) if tail else (gAs_grp, gBs_grp)
-            tiles = tail_tiles if tail else SCALE_GROUP
-            l_As = fx.make_view(
-                fx.add_offset(sAs_base_ptr, spage * SCALE_GROUP * BLOCK_M),
-                scale_box(tiles, BLOCK_M),
-            )
-            l_Bs = fx.make_view(
-                fx.add_offset(sBs_base_ptr, spage * SCALE_GROUP * BLOCK_N),
-                scale_box(tiles, BLOCK_N),
-            )
-            s_offA = fx.Int64(g * SCALE_GROUP * BLOCK_M * 4)
-            s_offB = fx.Int64(g * SCALE_GROUP * BLOCK_N * 4)
+            l_As = fx.make_view(fx.add_offset(sAs_base_ptr, spage * SCALE_GROUP), grp_box)
+            l_Bs = fx.make_view(fx.add_offset(sBs_base_ptr, spage * SCALE_GROUP), grp_box)
+            s_offA = s_offB = fx.Int64(0 if tail else g * SCALE_GROUP * 4)
             if wave // TDM_GROUP_WAVES == group:
                 gw = wave % TDM_GROUP_WAVES
                 if gw < 2:
@@ -376,28 +357,29 @@ def lds_to_regs(k):
                     fx.add_offset(fx.get_iter(frag_B[r]), 4 * h), fx.make_layout((4, 4), (1, 8))
                 ),
             )
-        s_off = (k // SCALE_GROUP) % SCALE_STAGES * SCALE_GROUP + k % SCALE_GROUP
-        src_sA = fx.make_view(
-            fx.add_offset(sAs_base_ptr, s_off * BLOCK_M + as_lds_base),
-            fx.make_layout((4, 2), (1, 4)),
-        )
-        src_sB = fx.make_view(
-            fx.add_offset(sBs_base_ptr, s_off * BLOCK_N + bs_lds_base),
-            fx.make_layout((4, 1), (1, 4)),
-        )
+        s_col = (k // SCALE_GROUP) % SCALE_STAGES * SCALE_GROUP + k % SCALE_GROUP
         fx.copy(
-            copy_atom_i32,
-            src_sA,
-            fx.make_view(fx.get_iter(frag_sA[r]), fx.make_layout((4, 2), (1, 4))),
-        )
-        fx.copy(
-            copy_atom_i32,
-            src_sB,
-            fx.make_view(fx.get_iter(frag_sB[r]), fx.make_layout((4, 1), (1, 4))),
+            copy_atom_s32,
+            fx.make_view(
+                fx.add_offset(sAs_base_ptr, as_lds_base + s_col),
+                fx.make_layout((1, 8), (1, 32 * LDS_SCALE_PITCH)),
+            ),
+            fx.make_view(fx.get_iter(frag_sA[r]), fx.make_layout((1, 8), (1, 1))),
         )
+        # B rows are 64 apart; sched_barrier keeps them from merging into ds_load_2addr_stride64
+        for j in range_constexpr(4):
+            fx.copy(
+                copy_atom_s32,
+                fx.make_view(
+                    fx.add_offset(sBs_base_ptr, bs_lds_base + s_col + j * 64 * LDS_SCALE_PITCH),
+                    fx.make_layout(1, 1),
+                ),
+                fx.make_view(fx.add_offset(fx.get_iter(frag_sB[r]), j), fx.make_layout(1, 1)),
+            )
+            fx.rocdl.sched_barrier(0x7FF)
 
     A_REPS, B_REPS = 8, 4
-    n_dsrd = 2 * (A_REPS + B_REPS) + (A_REPS + B_REPS) // 2
+    n_dsrd = 3 * (A_REPS + B_REPS)
     n_mma = A_REPS * B_REPS
 
     # interleave tile k+1's LDS reads with gemm(k)'s WMMAs
@@ -491,71 +473,20 @@ def launch_gemm(
     M: fx.Int32,
     N: fx.Int32,
     K: fx.Int32,
+    As_sw: fx.Int32,
+    Bs_sw: fx.Int32,
     K_iters: fx.Constexpr[int],
     OUT_KIND: fx.Constexpr[int],
     stream: fx.Stream = fx.Stream(None),
 ):
     grid_tot = (M // BLOCK_M) * (N // BLOCK_N)
-    gemm_kernel(A, B, A_s, B_s, C, M, N, K, K_iters, OUT_KIND).launch(
+    gemm_kernel(A, B, A_s, B_s, C, M, N, K, As_sw, Bs_sw, K_iters, OUT_KIND).launch(
         grid=(grid_tot, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream
     )
 
 
-@flyc.kernel
-def shuffle_scales_kernel(
-    scales: fx.Tensor,
-    scales_shuf: fx.Tensor,
-    rows: fx.Int32,
-    words_per_row: fx.Int32,
-    src_row_words: fx.Int32,
-    block_rows: fx.Constexpr[int],
-    lane_rows: fx.Constexpr[int],
-):
-    # row-major [rows][src_row_words] (TE's padded buffer) -> iteration-major, one thread per word
-    tid = fx.thread_idx.x
-    bid_r = fx.block_idx.x
-    bid_k = fx.block_idx.y
-
-    src = fx.get_iter(scales)
-    src_ptr = fx.recast_iter(fx.PointerType.get(fx.Int32.ir_type, src.memspace, 4), src)
-    src_words = fx.make_view(src_ptr, fx.make_layout((rows, words_per_row), (src_row_words, 1)))
-    dst_words = iter_major_scale_words(scales_shuf, rows, words_per_row, block_rows, lane_rows)
-
-    src_tile = fx.slice(fx.zipped_divide(src_words, (block_rows, 1)), (None, (bid_r, bid_k)))
-    dst_tile = fx.slice(fx.zipped_divide(dst_words, (block_rows, 1)), (None, (bid_r, bid_k)))
-
-    copy_atom_i32 = fx.make_copy_atom(fx.UniversalCopy32b(), fx.Int32)
-    tile_mn, tv_layout = fx.make_layout_tv(
-        fx.make_layout((block_rows, 1), (1, 1)), fx.make_layout((1, 1), (1, 1))
-    )
-    thr_copy = fx.make_tiled_copy(copy_atom_i32, tv_layout, tile_mn).get_slice(tid)
-    part_src = thr_copy.partition_S(src_tile)
-    part_dst = thr_copy.partition_D(dst_tile)
-    frag = fx.make_fragment_like(part_src)
-    fx.copy(copy_atom_i32, part_src, frag)
-    fx.copy(copy_atom_i32, frag, part_dst)
-
-
-@flyc.jit
-def shuffle_scales(
-    scales: fx.Tensor,
-    scales_shuf: fx.Tensor,
-    rows: fx.Int32,
-    K: fx.Int32,
-    src_row_words: fx.Int32,
-    block_rows: fx.Constexpr[int],
-    lane_rows: fx.Constexpr[int],
-    stream: fx.Stream = fx.Stream(None),
-):
-    words_per_row = K // SCALE_WORD_K
-    grid_r = rows // block_rows
-    shuffle_scales_kernel(
-        scales, scales_shuf, rows, words_per_row, src_row_words, block_rows, lane_rows
-    ).launch(grid=(grid_r, words_per_row, 1), block=(block_rows, 1, 1), stream=stream)
-
-
-def _shuffle_scale_words(scale, rows, K, block_rows, lane_rows, stream):
-    """Reorder TE E8M0 scales (possibly padded) into the kernel's iteration-major order."""
+def _scale_row_words(scale, rows, K):
+    """Row stride in words of TE's row-major E8M0 scales (possibly padded), read in place."""
     if scale.dtype != torch.uint8 or scale.ndim != 2:
         raise FlyDSLUnsupportedError(
             f"FlyDSL MXFP4 expects 2D uint8 E8M0 scales, got {scale.dtype} {tuple(scale.shape)}"
@@ -569,9 +500,7 @@ def _shuffle_scale_words(scale, rows, K, block_rows, lane_rows, stream):
         raise FlyDSLUnsupportedError(
             f"FlyDSL MXFP4 needs word-aligned row-major scales, got strides {scale.stride()}"
         )
-    out = torch.empty(rows, K // MXFP4_BLOCK_SIZE, dtype=torch.uint8, device=scale.device)
-    shuffle_scales(scale, out, rows, K, scale.stride(0) // 4, block_rows, lane_rows, stream=stream)
-    return out
+    return scale.stride(0) // 4
 
 
 def check_mxfp4_shape(m, n, k):
@@ -623,17 +552,17 @@ def mxfp4_matmul(
 
     if stream is None:
         stream = torch.cuda.current_stream(a.device)
-    a_scale_k = _shuffle_scale_words(a_scale, m, k, BLOCK_M, SCALE_LANE_ROWS_A, stream)
-    b_scale_k = _shuffle_scale_words(b_scale, n, k, BLOCK_N, SCALE_LANE_ROWS_B, stream)
     launch_gemm(
         a.view(torch.uint8),
         b.view(torch.uint8),
-        a_scale_k,
-        b_scale_k,
+        a_scale,
+        b_scale,
         D,
         m,
         n,
         k,
+        _scale_row_words(a_scale, m, k),
+        _scale_row_words(b_scale, n, k),
         k // BLOCK_K,
         OUT_KINDS[D.dtype],
         stream=stream,

From c27c0e335e2adfbd0d5345f93c9c71ec4a31fa42 Mon Sep 17 00:00:00 2001
From: Aristotle Martin <Aristotle.Martin@amd.com>
Date: Thu, 5 Nov 2026 02:56:06 +0000
Subject: [PATCH 09/11] Add FlyDSL MXFP4 test shapes with K > 2048 (streamed
 and partial scale groups)

---
 tests/pytorch/test_gemm_backends.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/pytorch/test_gemm_backends.py b/tests/pytorch/test_gemm_backends.py
index 68218ed346..eaa50c1a26 100644
--- a/tests/pytorch/test_gemm_backends.py
+++ b/tests/pytorch/test_gemm_backends.py
@@ -177,6 +177,8 @@ def _device_capability():
     (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"]

From 825ed444fa35a0801de3367f3fa1f99875356a85 Mon Sep 17 00:00:00 2001
From: Aristotle Martin <Aristotle.Martin@amd.com>
Date: Thu, 5 Nov 2026 03:19:03 +0000
Subject: [PATCH 10/11] FlyDSL MXFP4: use 64-bit offsets for A, B and C;
 bound-check 32-bit scale offsets

---
 .../pytorch/flydsl_kernels/gemm/mxfp4_gemm.py | 30 ++++++++++++++-----
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
index 093472f7da..fbebd0ad4f 100644
--- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
+++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
@@ -82,6 +82,8 @@ def as_words(t, rows, words_per_row):
     wm, wn = wave % 2, wave // 2
 
     k_words = K // 8
+    # global offsets in 64 bits: rows * k_words and rows * N exceed 2^31 for large shapes
+    k_words64, N64 = fx.Int64(k_words), fx.Int64(N)
     A_words = as_words(A, M, k_words)
     B_words = as_words(B, N, k_words)
 
@@ -132,11 +134,15 @@ def as_words(t, rows, words_per_row):
     HALF_ROWS = BLOCK_M // TDM_WAVES
     half_box = fx.make_layout((HALF_ROWS, 16), (16, 1))
     gA_half = fx.make_view(
-        fx.add_offset(fx.get_iter(A_words), (bid_m * BLOCK_M + tdm_wave * HALF_ROWS) * k_words),
+        fx.add_offset(
+            fx.get_iter(A_words), fx.Int64(bid_m * BLOCK_M + tdm_wave * HALF_ROWS) * k_words64
+        ),
         half_box,
     )
     gB_half = fx.make_view(
-        fx.add_offset(fx.get_iter(B_words), (bid_n * BLOCK_N + tdm_wave * HALF_ROWS) * k_words),
+        fx.add_offset(
+            fx.get_iter(B_words), fx.Int64(bid_n * BLOCK_N + tdm_wave * HALF_ROWS) * k_words64
+        ),
         half_box,
     )
     tdm_A = fx.rocdl.make_tdm_atom(
@@ -169,6 +175,7 @@ def scale_ptr(t):
         src = fx.get_iter(t)
         return fx.recast_iter(fx.PointerType.get(fx.Int32.ir_type, src.memspace, 4), src)
 
+    # scale offsets stay 32-bit (rows * K / 128 words; the host checks they fit)
     sA_blk = fx.add_offset(scale_ptr(A_s), bid_m * BLOCK_M * As_sw)
     sB_blk = fx.add_offset(scale_ptr(B_s), bid_n * BLOCK_N * Bs_sw)
 
@@ -207,8 +214,12 @@ def scale_atom(view, sw, tiles, width):
     tdm_Bs_tail = scale_atom(gBs_tail, Bs_sw, tail_tiles, SCALE_GROUP)
 
     # whole tiles, for K-loop iterations that also load a scale group
-    gA_full = fx.make_view(fx.add_offset(fx.get_iter(A_words), bid_m * BLOCK_M * k_words), box_A)
-    gB_full = fx.make_view(fx.add_offset(fx.get_iter(B_words), bid_n * BLOCK_N * k_words), box_B)
+    gA_full = fx.make_view(
+        fx.add_offset(fx.get_iter(A_words), fx.Int64(bid_m * BLOCK_M) * k_words64), box_A
+    )
+    gB_full = fx.make_view(
+        fx.add_offset(fx.get_iter(B_words), fx.Int64(bid_n * BLOCK_N) * k_words64), box_B
+    )
     tdm_A_full = fx.rocdl.make_tdm_atom(
         gA_full, [None, None], strides=[fx.Int64(k_words), None], num_warps=1, **pad
     )
@@ -221,14 +232,15 @@ def scale_atom(view, sw, tiles, width):
     c_wave_row = bid_m * BLOCK_M + wave * C_WAVE_ROWS
     gC_wave = [
         fx.make_view(
-            fx.add_offset(fx.get_iter(C), (c_wave_row + h * C_HALF_ROWS) * N + bid_n * BLOCK_N),
+            fx.add_offset(
+                fx.get_iter(C),
+                fx.Int64(c_wave_row + h * C_HALF_ROWS) * N64 + fx.Int64(bid_n * BLOCK_N),
+            ),
             c_wave_box,
         )
         for h in range(2)
     ]
-    tdm_C = fx.rocdl.make_tdm_atom(
-        gC_wave[0], [None, None], strides=[fx.Int64(N), None], num_warps=1
-    )
+    tdm_C = fx.rocdl.make_tdm_atom(gC_wave[0], [None, None], strides=[N64, None], num_warps=1)
     lds_C_wave = fx.make_view(fx.add_offset(lds_C_ptr, wave * C_WAVE_ROWS * BLOCK_N), c_wave_box)
 
     bC = fx.slice(fx.zipped_divide(C, (BLOCK_M, BLOCK_N)), (None, (bid_m, bid_n)))
@@ -500,6 +512,8 @@ def _scale_row_words(scale, rows, K):
         raise FlyDSLUnsupportedError(
             f"FlyDSL MXFP4 needs word-aligned row-major scales, got strides {scale.stride()}"
         )
+    if rows * (scale.stride(0) // 4) > 2**31 - 1:
+        raise FlyDSLUnsupportedError(f"FlyDSL MXFP4 scale offsets exceed 32 bits ({rows} rows)")
     return scale.stride(0) // 4
 
 

From 637f2bde33a078c118d44cd2f78f308d42a9988f Mon Sep 17 00:00:00 2001
From: Aristotle Martin <Aristotle.Martin@amd.com>
Date: Thu, 5 Nov 2026 04:33:40 +0000
Subject: [PATCH 11/11] FlyDSL MXFP4: replace fully unrolled K loop with a
 runtime loop

---
 .../pytorch/flydsl_kernels/gemm/mxfp4_gemm.py | 525 +++++++++++-------
 1 file changed, 311 insertions(+), 214 deletions(-)

diff --git a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
index fbebd0ad4f..f6e72cd740 100644
--- a/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
+++ b/transformer_engine/pytorch/flydsl_kernels/gemm/mxfp4_gemm.py
@@ -28,8 +28,7 @@
 TDM_GROUP_WAVES = NUM_WAVES // 2
 TDM_WAVES = 2
 PRO_SCALE_WAVES = 2
-# Scales stream through SCALE_STAGES LDS pages of SCALE_GROUP K tiles
-SCALE_GROUP = 8
+SCALE_GROUP = 6
 SCALE_STAGES = 2
 LDS_B_START = 64 * 1024
 C_HALF_ROWS = BLOCK_M // 2  # C is stored in two halves
@@ -38,9 +37,7 @@
 ROW_WORDS = 16
 LDS_ROW_WORDS = 20
 SCHED_TAIL_MMAS = 8
-# Scales are TE's row-major words (one per row and K tile), kept in LDS as [row][SCALE_STAGES *
-# SCALE_GROUP] with an odd row pitch so the per-row reads avoid bank conflicts
-LDS_SCALE_PITCH = 17
+K_LOOP_BODY = SCALE_GROUP
 OUT_KINDS = {torch.float32: 0, torch.bfloat16: 1, torch.float16: 2}
 
 
@@ -56,7 +53,8 @@ def gemm_kernel(
     K: fx.Int32,
     As_sw: fx.Int32,
     Bs_sw: fx.Int32,
-    K_iters: fx.Constexpr[int],
+    K_SHORT: fx.Constexpr[int],
+    K_REM: fx.Constexpr[int],
     OUT_KIND: fx.Constexpr[int],
 ):
 
@@ -82,18 +80,17 @@ def as_words(t, rows, words_per_row):
     wm, wn = wave % 2, wave // 2
 
     k_words = K // 8
+    K_iters = K // BLOCK_K
     # global offsets in 64 bits: rows * k_words and rows * N exceed 2^31 for large shapes
     k_words64, N64 = fx.Int64(k_words), fx.Int64(N)
     A_words = as_words(A, M, k_words)
     B_words = as_words(B, N, k_words)
 
-    # this lane's 16 B chunks khalf and khalf + 2
+    # this lane's 16 B chunk khalf (khalf + 2 is 8 words on)
     a_lds_lo = (wm * 16 + lane_mod16) * LDS_ROW_WORDS + 4 * khalf
-    a_lds_hi = a_lds_lo + 8
     b_lds_lo = (wn * 16 + lane_mod16) * LDS_ROW_WORDS + 4 * khalf
-    b_lds_hi = b_lds_lo + 8
-    as_lds_base = (wm * 16 + lane_mod16) * LDS_SCALE_PITCH
-    bs_lds_base = (wn * 16 + lane_mod16) * LDS_SCALE_PITCH
+    as_lds_base = (wm * 16 + lane_mod16) * SCALE_GROUP
+    bs_lds_base = (wn * 16 + lane_mod16) * SCALE_GROUP
 
     # tile k lives in LDS page k % NUM_STAGES
     smem = fx.SharedAllocator(static=False)
@@ -114,10 +111,9 @@ def as_words(t, rows, words_per_row):
     )
     pad = dict(pad_interval=ROW_WORDS, pad_amount=LDS_ROW_WORDS - ROW_WORDS)
 
-    num_groups = (K_iters + SCALE_GROUP - 1) // SCALE_GROUP
-    tail_tiles = K_iters - (num_groups - 1) * SCALE_GROUP
-    lds_As = smem.allocate(fx.Array[fx.Int32, BLOCK_M * LDS_SCALE_PITCH]).peek()
-    lds_Bs = smem.allocate(fx.Array[fx.Int32, BLOCK_N * LDS_SCALE_PITCH]).peek()
+    PAGE_S = BLOCK_M * SCALE_GROUP
+    lds_As = smem.allocate(fx.Array[fx.Int32, SCALE_STAGES * PAGE_S]).peek()
+    lds_Bs = smem.allocate(fx.Array[fx.Int32, SCALE_STAGES * PAGE_S]).peek()
 
     # C is staged from LDS offset 0 after the K loop
     if smem.allocated_bytes < c_stage_bytes:
@@ -158,7 +154,7 @@ def lds_half(base_ptr, page):
             half_box,
         )
 
-    # prologue atoms: whole tiles and scales
+    # prologue atoms: whole tiles 0 and 1
     tdm_A1w, cA1w = fx.rocdl.cdna5.make_tiled_tdm_atom(
         fx.rocdl.TensorLoad(), A_words, box_A, (BLOCK_M, 16), num_warps=1
     )
@@ -179,39 +175,29 @@ def scale_ptr(t):
     sA_blk = fx.add_offset(scale_ptr(A_s), bid_m * BLOCK_M * As_sw)
     sB_blk = fx.add_offset(scale_ptr(B_s), bid_n * BLOCK_N * Bs_sw)
 
-    def scale_atom(view, sw, tiles, width):
-        return fx.rocdl.make_tdm_atom(
-            view,
-            [None, tiles],
-            strides=[fx.Int64(sw), None],
-            num_warps=1,
-            pad_interval=width,
-            pad_amount=LDS_SCALE_PITCH - width,
-        )
-
-    # prologue: scale groups 0 and 1, half the rows per wave
-    sc_wave = wave % PRO_SCALE_WAVES
-    PRO_ROWS = BLOCK_M // PRO_SCALE_WAVES
-    PRO_COLS = SCALE_STAGES * SCALE_GROUP
-    pro_box = fx.make_layout((PRO_ROWS, PRO_COLS), (PRO_COLS, 1))
-    gAs_pro = fx.make_view(fx.add_offset(sA_blk, sc_wave * PRO_ROWS * As_sw), pro_box)
-    gBs_pro = fx.make_view(fx.add_offset(sB_blk, sc_wave * PRO_ROWS * Bs_sw), pro_box)
-    tdm_As_pro = scale_atom(gAs_pro, As_sw, K_iters, PRO_COLS)
-    tdm_Bs_pro = scale_atom(gBs_pro, Bs_sw, K_iters, PRO_COLS)
-    pro_lds_off = sc_wave * PRO_ROWS * LDS_SCALE_PITCH
-    lAs_pro = fx.make_view(fx.add_offset(sAs_base_ptr, pro_lds_off), pro_box)
-    lBs_pro = fx.make_view(fx.add_offset(sBs_base_ptr, pro_lds_off), pro_box)
-
     grp_box = fx.make_layout((BLOCK_M, SCALE_GROUP), (SCALE_GROUP, 1))
+
     gAs_grp = fx.make_view(sA_blk, grp_box)
     gBs_grp = fx.make_view(sB_blk, grp_box)
-    tdm_As_grp = scale_atom(gAs_grp, As_sw, None, SCALE_GROUP)
-    tdm_Bs_grp = scale_atom(gBs_grp, Bs_sw, None, SCALE_GROUP)
-    tail_col = (num_groups - 1) * SCALE_GROUP
-    gAs_tail = fx.make_view(fx.add_offset(sA_blk, tail_col), grp_box)
-    gBs_tail = fx.make_view(fx.add_offset(sB_blk, tail_col), grp_box)
-    tdm_As_tail = scale_atom(gAs_tail, As_sw, tail_tiles, SCALE_GROUP)
-    tdm_Bs_tail = scale_atom(gBs_tail, Bs_sw, tail_tiles, SCALE_GROUP)
+    tdm_As_grp = fx.rocdl.make_tdm_atom(
+        gAs_grp, [None, None], strides=[fx.Int64(As_sw), None], num_warps=1
+    )
+    tdm_Bs_grp = fx.rocdl.make_tdm_atom(
+        gBs_grp, [None, None], strides=[fx.Int64(Bs_sw), None], num_warps=1
+    )
+
+    def copy_scale_group(blk, sw, lds_base, g, full=False, page=None):
+        page = g % SCALE_STAGES if page is None else page
+        lds = fx.make_view(fx.add_offset(lds_base, page * PAGE_S), grp_box)
+        if const_expr(full):
+            atom, view = (tdm_As_grp, gAs_grp) if blk is sA_blk else (tdm_Bs_grp, gBs_grp)
+            fx.copy(atom, view, lds, imm_offset=fx.Int64(g) * (SCALE_GROUP * 4))
+        else:
+            view = fx.make_view(fx.add_offset(blk, g * SCALE_GROUP), grp_box)
+            atom = fx.rocdl.make_tdm_atom(
+                view, [None, K_iters - g * SCALE_GROUP], strides=[fx.Int64(sw), None], num_warps=1
+            )
+            fx.copy(atom, view, lds)
 
     # whole tiles, for K-loop iterations that also load a scale group
     gA_full = fx.make_view(
@@ -227,22 +213,6 @@ def scale_atom(view, sw, tiles, width):
         gB_full, [None, None], strides=[fx.Int64(k_words), None], num_warps=1, **pad
     )
 
-    # C store: each wave writes C_WAVE_ROWS rows of each half
-    c_wave_box = fx.make_layout((C_WAVE_ROWS, BLOCK_N), (BLOCK_N, 1))
-    c_wave_row = bid_m * BLOCK_M + wave * C_WAVE_ROWS
-    gC_wave = [
-        fx.make_view(
-            fx.add_offset(
-                fx.get_iter(C),
-                fx.Int64(c_wave_row + h * C_HALF_ROWS) * N64 + fx.Int64(bid_n * BLOCK_N),
-            ),
-            c_wave_box,
-        )
-        for h in range(2)
-    ]
-    tdm_C = fx.rocdl.make_tdm_atom(gC_wave[0], [None, None], strides=[N64, None], num_warps=1)
-    lds_C_wave = fx.make_view(fx.add_offset(lds_C_ptr, wave * C_WAVE_ROWS * BLOCK_N), c_wave_box)
-
     bC = fx.slice(fx.zipped_divide(C, (BLOCK_M, BLOCK_N)), (None, (bid_m, bid_n)))
 
     mma_atom = fx.make_mma_atom(fx.rocdl.WMMAScale(16, 16, 128, fx.Float4E2M1FN))
@@ -255,197 +225,317 @@ def scale_atom(view, sw, tiles, width):
         fx.UniversalCopy32b() if OUT_KIND == 0 else fx.UniversalCopy16b(), out_fx
     )
 
-    # double-buffered register fragments: tile k is read into set k % 2
-    frag_A = [fx.make_rmem_tensor(fx.make_layout((8, 8, 1), (1, 8, 0)), fx.Int32) for _ in range(2)]
-    frag_B = [fx.make_rmem_tensor(fx.make_layout((8, 4, 1), (1, 8, 0)), fx.Int32) for _ in range(2)]
-    frag_sA = [fx.make_rmem_tensor((1, 8, 1), fx.Int32) for _ in range(2)]
-    frag_sB = [fx.make_rmem_tensor((1, 4, 1), fx.Int32) for _ in range(2)]
-    frag_acc = thr_mma.make_fragment_C(bC)
-
-    # tiles issued per wave class, to derive each class's tensorcnt wait for tile t
-    WAVE_CLASSES = ((0, 2), (2, 4), (4, NUM_WAVES))
-    tdm_issued = [[0], [1], [0]]  # prologue ops
-
-    def wait_tile(t):
-        counts = [sum(1 for x in tiles if x > t) for tiles in tdm_issued]
-        if counts[0] == counts[1] == counts[2]:
-            fx.rocdl.s_wait_tensorcnt(counts[0])
-        else:
-            if wave < WAVE_CLASSES[0][1]:
-                fx.rocdl.s_wait_tensorcnt(counts[0])
+    def new_frags():
+        return (
+            fx.make_rmem_tensor(fx.make_layout((8, 8, 1), (1, 8, 0)), fx.Int32),
+            fx.make_rmem_tensor(fx.make_layout((8, 4, 1), (1, 8, 0)), fx.Int32),
+            fx.make_rmem_tensor(fx.make_layout((1, 8, 1), (0, 1, 0)), fx.Int32),
+            fx.make_rmem_tensor(fx.make_layout((1, 4, 1), (0, 1, 0)), fx.Int32),
+        )
+
+    A_REPS, B_REPS = 8, 4
+    acc_tile = fx.make_layout((8, 1, 1), (1, 0, 0))
+    scale_tile = fx.make_layout((1, 1, 1), (0, 0, 0))
+    frag_acc = [
+        [fx.make_rmem_tensor(acc_tile, fx.Float32) for _ in range(B_REPS)] for _ in range(A_REPS)
+    ]
+
+    def frag_view(t, off, layout):
+        return fx.make_view(fx.add_offset(fx.get_iter(t), off), layout)
+
+    def frag_view_ptr(ptr, off):
+        return fx.make_view(fx.add_offset(ptr, off), fx.make_layout(1, 1))
+
+    # one TDM op per wave: tiles 0 and 1 and the scales
+    def issue_prologue_tdm():
+        if wave < 4:
+            if wave == 0:
+                fx.copy(tdm_A1w, tAgA1w[None, 0], tAsA1w[None, 0])
             else:
-                if wave < WAVE_CLASSES[1][1]:
-                    fx.rocdl.s_wait_tensorcnt(counts[1])
-                else:
-                    fx.rocdl.s_wait_tensorcnt(counts[2])
-
-    # With scale_group g the group's waves load the A tile, A scale group g, the B tile and
-    # B scale group g, still one TDM op per wave; the scale op shares tile j's wait slot
-    def issue_tdm(j, scale_group=None):
-        page = j % NUM_STAGES
-        group = j % 2
-        k_off = fx.Int64(j * BLOCK_K // 2)
-        if scale_group is None:
-            if wave // TDM_GROUP_WAVES == group:
-                if wave % TDM_GROUP_WAVES < TDM_WAVES:
-                    fx.copy(tdm_A, gA_half, lds_half(sA_base_ptr, page), imm_offset=k_off)
+                if wave == 1:
+                    fx.copy(tdm_B1w, tBgB1w[None, 0], tBsB1w[None, 0])
                 else:
-                    fx.copy(tdm_B, gB_half, lds_half(sB_base_ptr, page), imm_offset=k_off)
+                    if wave == 2:
+                        fx.copy(tdm_A1w, tAgA1w[None, 1], tAsA1w[None, 1])
+                    else:
+                        fx.copy(tdm_B1w, tBgB1w[None, 1], tBsB1w[None, 1])
         else:
-            g = scale_group
-            spage = g % SCALE_STAGES
-            tail = g == num_groups - 1 and tail_tiles < SCALE_GROUP
-            atom_As, atom_Bs = (tdm_As_tail, tdm_Bs_tail) if tail else (tdm_As_grp, tdm_Bs_grp)
-            g_As, g_Bs = (gAs_tail, gBs_tail) if tail else (gAs_grp, gBs_grp)
-            l_As = fx.make_view(fx.add_offset(sAs_base_ptr, spage * SCALE_GROUP), grp_box)
-            l_Bs = fx.make_view(fx.add_offset(sBs_base_ptr, spage * SCALE_GROUP), grp_box)
-            s_offA = s_offB = fx.Int64(0 if tail else g * SCALE_GROUP * 4)
-            if wave // TDM_GROUP_WAVES == group:
-                gw = wave % TDM_GROUP_WAVES
+            if wave < 6:
+                if wave == 4:
+                    copy_scale_group(sA_blk, As_sw, sAs_base_ptr, 0)
+                else:
+                    if const_expr(K_SHORT == 0 or K_SHORT > SCALE_GROUP):
+                        copy_scale_group(sA_blk, As_sw, sAs_base_ptr, 1)
+            else:
+                if wave == 6:
+                    copy_scale_group(sB_blk, Bs_sw, sBs_base_ptr, 0)
+                else:
+                    if const_expr(K_SHORT == 0 or K_SHORT > SCALE_GROUP):
+                        copy_scale_group(sB_blk, Bs_sw, sBs_base_ptr, 1)
+
+    n_dsrd = 3 * (A_REPS + B_REPS)
+    n_mma = A_REPS * B_REPS
+
+    # interleave tile k+1's LDS reads with gemm(k)'s WMMAs
+    def hot_loop_scheduler():
+        n_spread = n_mma - SCHED_TAIL_MMAS
+        issued = 0
+        for i in range_constexpr(n_spread):
+            want = (i + 1) * n_dsrd // n_spread
+            if want > issued:
+                fx.rocdl.sched_dsrd(want - issued)
+                issued = want
+            fx.rocdl.sched_mfma(1)
+        fx.rocdl.sched_mfma(SCHED_TAIL_MMAS)
+        fx.rocdl.sched_barrier(0)
+
+    def gemm_step(src):
+        frag_A, frag_B, frag_sA, frag_sB = src
+        for i in range_constexpr(A_REPS):
+            for j in range_constexpr(B_REPS):
+                fx.gemm(
+                    mma_atom,
+                    frag_acc[i][j],
+                    [frag_view(frag_A, 8 * i, acc_tile), frag_view(frag_sA, i, scale_tile)],
+                    [frag_view(frag_B, 8 * j, acc_tile), frag_view(frag_sB, j, scale_tile)],
+                    frag_acc[i][j],
+                )
+
+    def rt_ptrs0():
+        return [fx.get_iter(v) for v in (gA_half, gB_half, gA_full, gB_full)] + [sA_blk, sB_blk]
+
+    def rt_advance(ptrs, tiles):
+        tw = tiles * (BLOCK_K // 8)
+        return [fx.add_offset(p, tw) for p in ptrs[:4]] + [
+            fx.add_offset(p, tiles) for p in ptrs[4:]
+        ]
+
+    def rt_views(ptrs):
+        layouts = (half_box, half_box, box_A, box_B, grp_box, grp_box)
+        return tuple(fx.make_view(p, l) for p, l in zip(ptrs, layouts))
+
+    def issue_tdm_rt(kb, vw, ks, scales, exact):
+        vA_half, vB_half, vA_full, vB_full, vAs, vBs = vw
+        js = ks + NUM_STAGES
+        page, group = js % NUM_STAGES, js % 2
+        k_off = fx.Int64(js * (BLOCK_K // 2))
+        if wave // TDM_GROUP_WAVES == group:
+            gw = wave % TDM_GROUP_WAVES
+            if const_expr(scales):
+                gi = ks // SCALE_GROUP + 1
+                gpage = (kb // SCALE_GROUP + gi) % SCALE_STAGES
+                s_off = fx.Int64(gi * SCALE_GROUP * 4)
                 if gw < 2:
                     if gw == 0:
                         fx.copy(
                             tdm_A_full,
-                            gA_full,
+                            vA_full,
                             fx.make_view(fx.add_offset(sA_base_ptr, page * PAGE_A), box_A),
                             imm_offset=k_off,
                         )
                     else:
-                        fx.copy(atom_As, g_As, l_As, imm_offset=s_offA)
+                        if const_expr(exact):
+                            g = kb // SCALE_GROUP + gi
+                            copy_scale_group(sA_blk, As_sw, sAs_base_ptr, g, page=gpage)
+                        else:
+                            lds = fx.make_view(fx.add_offset(sAs_base_ptr, gpage * PAGE_S), grp_box)
+                            fx.copy(tdm_As_grp, vAs, lds, imm_offset=s_off)
                 else:
                     if gw == 2:
                         fx.copy(
                             tdm_B_full,
-                            gB_full,
+                            vB_full,
                             fx.make_view(fx.add_offset(sB_base_ptr, page * PAGE_B), box_B),
                             imm_offset=k_off,
                         )
                     else:
-                        fx.copy(atom_Bs, g_Bs, l_Bs, imm_offset=s_offB)
-        for c in ((0, 1) if group == 0 else (2,)):
-            tdm_issued[c].append(j)
-
-    # one TDM op per wave: tiles 0 and 1 and the scales
-    def issue_prologue_tdm():
-        if wave < 4:
-            if wave == 0:
-                fx.copy(tdm_A1w, tAgA1w[None, 0], tAsA1w[None, 0])
+                        if const_expr(exact):
+                            g = kb // SCALE_GROUP + gi
+                            copy_scale_group(sB_blk, Bs_sw, sBs_base_ptr, g, page=gpage)
+                        else:
+                            lds = fx.make_view(fx.add_offset(sBs_base_ptr, gpage * PAGE_S), grp_box)
+                            fx.copy(tdm_Bs_grp, vBs, lds, imm_offset=s_off)
             else:
-                if wave == 1:
-                    fx.copy(tdm_B1w, tBgB1w[None, 0], tBsB1w[None, 0])
+                if gw < TDM_WAVES:
+                    fx.copy(tdm_A, vA_half, lds_half(sA_base_ptr, page), imm_offset=k_off)
                 else:
-                    if wave == 2:
-                        fx.copy(tdm_A1w, tAgA1w[None, 1], tAsA1w[None, 1])
-                    else:
-                        fx.copy(tdm_B1w, tBgB1w[None, 1], tBsB1w[None, 1])
-        else:
-            if wave < 4 + PRO_SCALE_WAVES:
-                fx.copy(tdm_As_pro, gAs_pro, lAs_pro)
-            else:
-                fx.copy(tdm_Bs_pro, gBs_pro, lBs_pro)
-
-    def lds_to_regs(k):
-        page, r = k % NUM_STAGES, k % 2
-        for h, a_off, b_off in ((0, a_lds_lo, b_lds_lo), (1, a_lds_hi, b_lds_hi)):
-            src_A = fx.make_view(
-                fx.add_offset(sA_base_ptr, page * PAGE_A + a_off),
-                fx.make_layout((4, 8), (1, 32 * LDS_ROW_WORDS)),
-            )
-            src_B = fx.make_view(
-                fx.add_offset(sB_base_ptr, page * PAGE_B + b_off),
-                fx.make_layout((4, 4), (1, 64 * LDS_ROW_WORDS)),
-            )
+                    fx.copy(tdm_B, vB_half, lds_half(sB_base_ptr, page), imm_offset=k_off)
+
+    def lds_to_regs_rt(r, dst, bases):
+        a_lo, b_lo, s_a, s_b, s_d = bases
+        page = r % NUM_STAGES
+        frag_A, frag_B, frag_sA, frag_sB = dst
+        for h in range_constexpr(2):
             fx.copy(
                 copy_atom_i32,
-                src_A,
                 fx.make_view(
-                    fx.add_offset(fx.get_iter(frag_A[r]), 4 * h), fx.make_layout((4, 8), (1, 8))
+                    fx.add_offset(sA_base_ptr, a_lo + (page * PAGE_A + 8 * h)),
+                    fx.make_layout((4, 8), (1, 32 * LDS_ROW_WORDS)),
                 ),
+                frag_view(frag_A, 4 * h, fx.make_layout((4, 8), (1, 8))),
             )
             fx.copy(
                 copy_atom_i32,
-                src_B,
                 fx.make_view(
-                    fx.add_offset(fx.get_iter(frag_B[r]), 4 * h), fx.make_layout((4, 4), (1, 8))
+                    fx.add_offset(sB_base_ptr, b_lo + (page * PAGE_B + 8 * h)),
+                    fx.make_layout((4, 4), (1, 64 * LDS_ROW_WORDS)),
                 ),
+                frag_view(frag_B, 4 * h, fx.make_layout((4, 4), (1, 8))),
             )
-        s_col = (k // SCALE_GROUP) % SCALE_STAGES * SCALE_GROUP + k % SCALE_GROUP
-        fx.copy(
-            copy_atom_s32,
-            fx.make_view(
-                fx.add_offset(sAs_base_ptr, as_lds_base + s_col),
-                fx.make_layout((1, 8), (1, 32 * LDS_SCALE_PITCH)),
-            ),
-            fx.make_view(fx.get_iter(frag_sA[r]), fx.make_layout((1, 8), (1, 1))),
-        )
-        # B rows are 64 apart; sched_barrier keeps them from merging into ds_load_2addr_stride64
-        for j in range_constexpr(4):
+        if const_expr(r // SCALE_GROUP % SCALE_STAGES == 1):
+            s_a, s_b = s_a + s_d, s_b + s_d
+        col = r % SCALE_GROUP
+        for i in range_constexpr(A_REPS):
             fx.copy(
                 copy_atom_s32,
-                fx.make_view(
-                    fx.add_offset(sBs_base_ptr, bs_lds_base + s_col + j * 64 * LDS_SCALE_PITCH),
-                    fx.make_layout(1, 1),
-                ),
-                fx.make_view(fx.add_offset(fx.get_iter(frag_sB[r]), j), fx.make_layout(1, 1)),
+                frag_view_ptr(sAs_base_ptr, s_a + (col + i * 32 * SCALE_GROUP)),
+                frag_view(frag_sA, i, fx.make_layout(1, 1)),
+            )
+            fx.rocdl.sched_barrier(0x7FF)
+        for j in range_constexpr(B_REPS):
+            fx.copy(
+                copy_atom_s32,
+                frag_view_ptr(sBs_base_ptr, s_b + (col + j * 64 * SCALE_GROUP)),
+                frag_view(frag_sB, j, fx.make_layout(1, 1)),
             )
             fx.rocdl.sched_barrier(0x7FF)
 
-    A_REPS, B_REPS = 8, 4
-    n_dsrd = 3 * (A_REPS + B_REPS)
-    n_mma = A_REPS * B_REPS
+    def wait_tile_rt(ks, pending):
+        if const_expr(pending == 0):
+            fx.rocdl.s_wait_tensorcnt(0)
+        else:
+            if wave // TDM_GROUP_WAVES == (ks + 2) % 2:
+                fx.rocdl.s_wait_tensorcnt(1)
+            else:
+                fx.rocdl.s_wait_tensorcnt(0)
 
-    # interleave tile k+1's LDS reads with gemm(k)'s WMMAs
-    def hot_loop_scheduler():
-        n_spread = n_mma - SCHED_TAIL_MMAS
-        issued = 0
-        for i in range_constexpr(n_spread):
-            want = (i + 1) * n_dsrd // n_spread
-            if want > issued:
-                fx.rocdl.sched_dsrd(want - issued)
-                issued = want
-            fx.rocdl.sched_mfma(1)
-        fx.rocdl.sched_mfma(SCHED_TAIL_MMAS)
+    def k_step_rt(ks, cur, bases, issue=None, scales=False, exact=False, pending=1):
         fx.rocdl.sched_barrier(0)
+        wait_tile_rt(ks, pending)
+        fx.rocdl.s_wait_dscnt(0)
+        fx.barrier()
+        fx.rocdl.sched_barrier(0)
+        if const_expr(issue is not None):
+            issue_tdm_rt(issue[0], issue[1], ks, scales, exact)
+        nxt = new_frags()
+        lds_to_regs_rt(ks + 1, nxt, bases)
+        gemm_step(cur)
+        hot_loop_scheduler()
+        return nxt
+
+    def carry_values(fr):
+        frag_A, frag_B, frag_sA, frag_sB = fr
+        return (
+            [frag_view(frag_A, 8 * i, acc_tile).load() for i in range_constexpr(A_REPS)]
+            + [frag_view(frag_B, 8 * j, acc_tile).load() for j in range_constexpr(B_REPS)]
+            + [frag_sA.load(), frag_sB.load()]
+            + [acc.load() for acc_row in frag_acc for acc in acc_row]
+        )
 
-    def gemm_step(k):
-        r = k % 2
-        fx.gemm(mma_atom, frag_acc, [frag_A[r], frag_sA[r]], [frag_B[r], frag_sB[r]], frag_acc)
-
-    frag_acc.fill(0)
+    def from_carry(vals):
+        fr = new_frags()
+        frag_A, frag_B, frag_sA, frag_sB = fr
+        for i in range_constexpr(A_REPS):
+            frag_view(frag_A, 8 * i, acc_tile).store(vals[i])
+        for j in range_constexpr(B_REPS):
+            frag_view(frag_B, 8 * j, acc_tile).store(vals[A_REPS + j])
+        frag_sA.store(vals[A_REPS + B_REPS])
+        frag_sB.store(vals[A_REPS + B_REPS + 1])
+        accs = vals[A_REPS + B_REPS + 2 :]
+        for i in range_constexpr(A_REPS):
+            for j in range_constexpr(B_REPS):
+                frag_acc[i][j].store(accs[i * B_REPS + j])
+        return fr
+
+    for acc_row in frag_acc:
+        for acc in acc_row:
+            acc.fill(0)
 
     issue_prologue_tdm()
-    wait_tile(0)
+    if wave // TDM_WAVES == 1:
+        fx.rocdl.s_wait_tensorcnt(1)
+    else:
+        fx.rocdl.s_wait_tensorcnt(0)
     fx.barrier()
-    lds_to_regs(0)
-
-    for k in range_constexpr(K_iters):
-        if k + 1 < K_iters:
-            # keep WMMAs from being scheduled across the barrier
-            fx.rocdl.sched_barrier(0)
-            wait_tile(k + 1)
-            fx.rocdl.s_wait_dscnt(0)
-            fx.barrier()
-            fx.rocdl.sched_barrier(0)
-            if k == 0 and NUM_STAGES - 1 < K_iters:
-                issue_tdm(NUM_STAGES - 1)
-            if k + NUM_STAGES < K_iters:
-                # scale group g >= SCALE_STAGES: issued at k = SCALE_GROUP * (g - 1)
-                g_next = k // SCALE_GROUP + 1
-                if k % SCALE_GROUP == 0 and SCALE_STAGES <= g_next < num_groups:
-                    issue_tdm(k + NUM_STAGES, scale_group=g_next)
-                else:
-                    issue_tdm(k + NUM_STAGES)
-            lds_to_regs(k + 1)
-            gemm_step(k)
-            hot_loop_scheduler()
-        else:
-            gemm_step(k)
+    bases = (a_lds_lo, b_lds_lo, as_lds_base, bs_lds_base, fx.Int32(PAGE_S))
+    cur = new_frags()
+    lds_to_regs_rt(0, cur, bases)
+
+    fx.rocdl.sched_barrier(0)
+    fx.rocdl.s_wait_tensorcnt(0)
+    fx.rocdl.s_wait_dscnt(0)
+    fx.barrier()
+    fx.rocdl.sched_barrier(0)
+    views0 = rt_views(rt_ptrs0())
+    for ks in range_constexpr(-1, 1):
+        if const_expr(K_SHORT == 0 or ks + NUM_STAGES < K_SHORT):
+            issue_tdm_rt(0, views0, ks, False, False)
+    nxt = new_frags()
+    lds_to_regs_rt(1, nxt, bases)
+    gemm_step(cur)
+    hot_loop_scheduler()
+    cur = nxt
+
+    if const_expr(K_SHORT == 0):
+        # kb & 2^30 is zero, but keeps the carried LDS bases from being hoisted
+        n_body = K_iters // K_LOOP_BODY - 2
+        init = carry_values(cur) + list(bases) + rt_ptrs0()
+        for i, state in range(0, n_body, 1, init=init):
+            nc = len(state) - 11
+            cur = from_carry(state[:nc])
+            a_lo, b_lo, s_a, s_b, s_d = (fx.Int32(v) for v in state[nc : nc + 5])
+            bases = (a_lo, b_lo, s_a, s_b, s_d)
+            ptrs = list(state[nc + 5 :])
+            kb = fx.Int32(i) * K_LOOP_BODY
+            issue = (kb, rt_views(ptrs))
+            for u in range_constexpr(K_LOOP_BODY):
+                at_group = (1 + u) % SCALE_GROUP == 0
+                cur = k_step_rt(1 + u, cur, bases, issue, scales=at_group)
+            zero = kb & (1 << 30)
+            bases_next = [a_lo + zero, b_lo + zero, s_a + s_d, s_b + s_d, -s_d]
+            results = yield carry_values(cur) + bases_next + rt_advance(ptrs, K_LOOP_BODY)
+        nc = len(results) - 11
+        cur = from_carry(results[:nc])
+        bases = tuple(fx.Int32(v) for v in results[nc : nc + 5])
+        kb = n_body * K_LOOP_BODY
+        issue = (kb, rt_views(list(results[nc + 5 :])))
+        R = 2 * SCALE_GROUP + K_REM - 4
+    else:
+        issue = (0, views0)
+        R = max(K_SHORT - 4, 0)
+    for u in range_constexpr(R):
+        at_group = (1 + u) % SCALE_GROUP == 0 and u < R - 3
+        cur = k_step_rt(1 + u, cur, bases, issue, scales=at_group, exact=u >= R - 8)
+    ks_last = R + 3 if K_SHORT == 0 else K_SHORT - 1
+    for ks in range_constexpr(max(ks_last - 2, 1), ks_last):
+        cur = k_step_rt(ks, cur, bases, pending=int(ks < ks_last - 1))
+    gemm_step(cur)
+    x = bases[0] - wm * 16 * LDS_ROW_WORDS
+    lane_mod16, khalf = x // LDS_ROW_WORDS, x % LDS_ROW_WORDS // 4
+
+    # C store: each wave writes C_WAVE_ROWS rows of each half
+    c_wave_box = fx.make_layout((C_WAVE_ROWS, BLOCK_N), (BLOCK_N, 1))
+    c_wave_row = bid_m * BLOCK_M + wave * C_WAVE_ROWS
+    gC_wave = [
+        fx.make_view(
+            fx.add_offset(
+                fx.get_iter(C),
+                fx.Int64(c_wave_row + h * C_HALF_ROWS) * N64 + fx.Int64(bid_n * BLOCK_N),
+            ),
+            c_wave_box,
+        )
+        for h in range(2)
+    ]
+    tdm_C = fx.rocdl.make_tdm_atom(gC_wave[0], [None, None], strides=[N64, None], num_warps=1)
+    lds_C_wave = fx.make_view(fx.add_offset(lds_C_ptr, wave * C_WAVE_ROWS * BLOCK_N), c_wave_box)
 
     if const_expr(OUT_KIND == 0):
         frag_c = frag_acc
     else:
-        frag_c = fx.make_fragment_like(frag_acc, out_fx)
-        frag_c.store(frag_acc.load().to(out_fx))
+        frag_c = [[fx.make_fragment_like(acc, out_fx) for acc in acc_row] for acc_row in frag_acc]
+        for i in range_constexpr(A_REPS):
+            for j in range_constexpr(B_REPS):
+                frag_c[i][j].store(frag_acc[i][j].load().to(out_fx))
 
     # accumulator (v, j, i) is tile row wm*16 + 8*khalf + v + 32*i, column wn*16 + lane_mod16 + 64*j
     fx.rocdl.s_wait_dscnt(0)
@@ -459,10 +549,7 @@ def gemm_step(k):
         for ii in range_constexpr(4):
             for j in range_constexpr(4):
                 for v in range_constexpr(8):
-                    src_c = fx.make_view(
-                        fx.add_offset(fx.get_iter(frag_c), v + 8 * j + 32 * (4 * h + ii)),
-                        fx.make_layout(1, 1),
-                    )
+                    src_c = frag_view(frag_c[4 * h + ii][j], v, fx.make_layout(1, 1))
                     dst_c = fx.make_view(
                         fx.add_offset(lds_C_ptr, lane_c + (32 * ii + v) * BLOCK_N + 64 * j),
                         fx.make_layout(1, 1),
@@ -487,12 +574,13 @@ def launch_gemm(
     K: fx.Int32,
     As_sw: fx.Int32,
     Bs_sw: fx.Int32,
-    K_iters: fx.Constexpr[int],
+    K_SHORT: fx.Constexpr[int],
+    K_REM: fx.Constexpr[int],
     OUT_KIND: fx.Constexpr[int],
     stream: fx.Stream = fx.Stream(None),
 ):
     grid_tot = (M // BLOCK_M) * (N // BLOCK_N)
-    gemm_kernel(A, B, A_s, B_s, C, M, N, K, As_sw, Bs_sw, K_iters, OUT_KIND).launch(
+    gemm_kernel(A, B, A_s, B_s, C, M, N, K, As_sw, Bs_sw, K_SHORT, K_REM, OUT_KIND).launch(
         grid=(grid_tot, 1, 1), block=(NUM_THREADS, 1, 1), stream=stream
     )
 
@@ -517,6 +605,15 @@ def _scale_row_words(scale, rows, K):
     return scale.stride(0) // 4
 
 
+def k_loop_constexprs(K):
+    """(K_SHORT, K_REM): the K tile count if the kernel needs no K loop, else K tiles mod
+    K_LOOP_BODY, so that kernels are compiled per remainder, not per K."""
+    k_iters = K // BLOCK_K
+    if k_iters < 2 * K_LOOP_BODY:
+        return k_iters, 0
+    return 0, k_iters % K_LOOP_BODY
+
+
 def check_mxfp4_shape(m, n, k):
     """Raise FlyDSLUnsupportedError unless the kernel supports C[m, n] with reduction depth k."""
     if m % BLOCK_M or n % BLOCK_N:
@@ -577,7 +674,7 @@ def mxfp4_matmul(
         k,
         _scale_row_words(a_scale, m, k),
         _scale_row_words(b_scale, n, k),
-        k // BLOCK_K,
+        *k_loop_constexprs(k),
         OUT_KINDS[D.dtype],
         stream=stream,
     )