From f6ed27c446829e0874281ec7744359c082188d07 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 27 Jul 2026 11:07:09 +0200
Subject: [PATCH 01/67] [PyTorch][torch.compile] Support for
 DotProductAttention

Make DotProductAttention itself traceable, so that torch.compile captures
the whole module -- input unpacking, qkv layout, backend selection and the
backend -- on the FlashAttention and UnfusedDotProductAttention backends.
FusedAttention stays an eager island, as before.

DotProductAttention.forward no longer carries @no_torch_dynamo. What that
uncovered, and the fixes:

- init_fp8_metadata unconditionally called get_fp8_recipe(), which builds a
  default Recipe when no autocast is active -- not traceable. It now returns
  early when FP8 is off everywhere, which is also what the rest of the
  function did in that case.
- get_qkv_layout detects packed q/k/v from data pointers and storage
  offsets, neither of which dynamo can trace (and neither is meaningful on
  fake tensors). While tracing, the layout is built from the format instead;
  packed inputs are declared explicitly via qkv_layer/kv_layer, and
  non-contiguous q/k/v are made contiguous so the layout stays truthful.
- get_full_cu_seqlens keyed its cache on torch.is_inference_mode_enabled(),
  a fundamental graph break; while tracing it builds the tensor directly.
- get_padding_mask read sequence lengths on the host (a device sync, and
  data-dependent under torch.compile). It is now vectorized.
- The fused sub-backend enum member does not survive a graph break: dynamo
  reconstructs the assume_constant_result value by re-emitting the call that
  produced it, so the resumed frame received the function itself instead of
  the enum -- which then reached fused_attn_fwd. get_attention_backend now
  keeps a plain int while tracing and casts back to FusedAttnBackend in
  eager mode only.
- Backend selection logging is skipped while tracing (logging.Logger methods
  graph-break, and int() of the sub-backend enum is untraceable too).

Tests: test_dpa_torch_compile covers 7 configurations (bshd/sbhd/thd,
causal/no_mask/sliding window, GQA, cross attention, packed qkv_layer)
against both backends with fullgraph=True, comparing forward and backward
against eager; plus a CUDA-graphs (reduce-overhead) test and a test that the
FusedAttention path stays correct around its graph break.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 243 ++++++++++++++++++
 .../dot_product_attention.py                  |  47 +++-
 .../attention/dot_product_attention/utils.py  |  58 ++++-
 .../pytorch/cpp_extensions/fused_attn.py      |  11 +-
 4 files changed, 330 insertions(+), 29 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index d720ca7a4a..ed2ad90c76 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -550,6 +550,249 @@ def fn(q, k, v, extra):
         assert v.grad is not None
 
 
+# ---------------------------------------------------------------------------
+# DotProductAttention under torch.compile
+# ---------------------------------------------------------------------------
+
+
+# Each config is a model configuration that both the FlashAttention and the
+# UnfusedDotProductAttention backend support, so that every one of them runs
+# against both backends.
+_DPA_COMPILE_CONFIGS = {
+    "self_bshd_causal": dict(qkv_format="bshd", attn_mask_type="causal"),
+    "self_sbhd_no_mask": dict(qkv_format="sbhd", attn_mask_type="no_mask"),
+    "self_bshd_swa": dict(qkv_format="bshd", attn_mask_type="causal", window_size=(16, 0)),
+    "gqa_bshd_causal": dict(
+        qkv_format="bshd", attn_mask_type="causal", num_heads=8, num_gqa_groups=2
+    ),
+    "cross_bshd_no_mask": dict(
+        qkv_format="bshd", attn_mask_type="no_mask", attention_type="cross", max_seqlen_kv=256
+    ),
+    "self_thd_padding_causal": dict(qkv_format="thd", attn_mask_type="padding_causal"),
+    "packed_qkv_bs3hd": dict(qkv_format="bshd", attn_mask_type="causal", packed="bs3hd"),
+}
+
+
+def _dpa_config(**overrides):
+    cfg = dict(
+        batch_size=2,
+        num_heads=4,
+        num_gqa_groups=None,
+        head_dim=64,
+        max_seqlen_q=128,
+        max_seqlen_kv=128,
+        qkv_format="bshd",
+        attn_mask_type="causal",
+        attention_type="self",
+        window_size=None,
+        packed=None,
+    )
+    cfg.update(overrides)
+    if cfg["num_gqa_groups"] is None:
+        cfg["num_gqa_groups"] = cfg["num_heads"]
+    return cfg
+
+
+def _make_dpa(cfg, dtype: torch.dtype) -> te.DotProductAttention:
+    return te.DotProductAttention(
+        num_attention_heads=cfg["num_heads"],
+        kv_channels=cfg["head_dim"],
+        num_gqa_groups=cfg["num_gqa_groups"],
+        attention_dropout=0.0,
+        qkv_format=cfg["qkv_format"],
+        attn_mask_type=cfg["attn_mask_type"],
+        window_size=cfg["window_size"],
+        attention_type=cfg["attention_type"],
+    ).to(dtype=dtype, device="cuda")
+
+
+def _make_dpa_inputs(cfg, dtype: torch.dtype):
+    """Build the (args, kwargs) that `DotProductAttention.forward` is called
+    with, plus the list of tensors whose gradients the test compares."""
+    b = cfg["batch_size"]
+    s_q, s_kv = cfg["max_seqlen_q"], cfg["max_seqlen_kv"]
+    h, g, d = cfg["num_heads"], cfg["num_gqa_groups"], cfg["head_dim"]
+    fmt = cfg["qkv_format"]
+
+    def _shape(s, heads):
+        return {
+            "bshd": (b, s, heads, d),
+            "sbhd": (s, b, heads, d),
+            "thd": (b * s, heads, d),
+        }[fmt]
+
+    def _randn(shape):
+        return torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
+
+    kwargs = {}
+    if cfg["packed"] is not None:
+        # Declarative packed QKV: q/k/v are derived from one buffer by DPA
+        # itself, and the layout comes from the declaration -- the only packed
+        # layout that torch.compile supports (pointer-based detection of packed
+        # q/k/v is untraceable).
+        assert cfg["packed"] == "bs3hd" and fmt == "bshd" and h == g
+        qkv = _randn((b, s_q, 3, h, d))
+        kwargs["qkv_layer"] = qkv
+        kwargs["qkv_interleave_dim"] = -3
+        grad_tensors = [qkv]
+        args = ()
+    else:
+        q, k, v = _randn(_shape(s_q, h)), _randn(_shape(s_kv, g)), _randn(_shape(s_kv, g))
+        grad_tensors = [q, k, v]
+        args = (q, k, v)
+
+    if fmt == "thd":
+        # All sequences have the maximum length, i.e. no padding. max_seqlen_*
+        # is passed explicitly: deriving it from cu_seqlens costs a `.item()`
+        # (a device sync, and an unbacked SymInt under torch.compile).
+        kwargs["cu_seqlens_q"] = torch.arange(
+            0, (b + 1) * s_q, step=s_q, dtype=torch.int32, device="cuda"
+        )
+        kwargs["cu_seqlens_kv"] = torch.arange(
+            0, (b + 1) * s_kv, step=s_kv, dtype=torch.int32, device="cuda"
+        )
+        kwargs["max_seqlen_q"] = s_q
+        kwargs["max_seqlen_kv"] = s_kv
+
+    return args, kwargs, grad_tensors
+
+
+def _force_dpa_backend(monkeypatch, backend: str) -> None:
+    """Restrict DotProductAttention to a single backend."""
+    from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (
+        _attention_backends,
+    )
+    from transformer_engine.pytorch.attention.dot_product_attention.utils import (
+        FlashAttentionUtils,
+    )
+
+    if backend == "flash" and not FlashAttentionUtils.is_installed:
+        pytest.skip("flash-attn is not installed")
+
+    for name, var in (
+        ("flash", "NVTE_FLASH_ATTN"),
+        ("fused", "NVTE_FUSED_ATTN"),
+        ("unfused", "NVTE_UNFUSED_ATTN"),
+    ):
+        monkeypatch.setenv(var, "1" if name == backend else "0")
+    # Backend selection is cached on the attention params only, so the env vars
+    # above are not enough to invalidate it.
+    _attention_backends["backend_selection_requires_update"] = True
+
+
+def _assert_dpa_backend(backend: str) -> None:
+    from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (
+        _attention_backends,
+    )
+
+    assert _attention_backends[f"use_{backend}_attention"], (
+        f"expected the {backend} backend to run, selected:"
+        f" flash={_attention_backends['use_flash_attention']},"
+        f" fused={_attention_backends['use_fused_attention']},"
+        f" unfused={_attention_backends['use_unfused_attention']}"
+    )
+
+
+@pytest.mark.parametrize("backend", ["flash", "unfused"])
+@pytest.mark.parametrize("config", _DPA_COMPILE_CONFIGS.keys())
+def test_dpa_torch_compile(monkeypatch, backend, config):
+    """`DotProductAttention` must trace under `torch.compile(fullgraph=True)`
+    on the FlashAttention and UnfusedDotProductAttention backends, and match
+    eager in forward and backward.
+
+    `fullgraph=True` makes the test fail on any graph break, i.e. it covers the
+    whole module: input unpacking, qkv layout, backend selection and the
+    backend itself. The FusedAttention backend is not covered -- it is an eager
+    island (`@no_torch_dynamo` on `FusedAttention.forward`)."""
+    dtype = torch.bfloat16
+    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS[config])
+    _force_dpa_backend(monkeypatch, backend)
+
+    module = _make_dpa(cfg, dtype)
+    args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
+
+    ref = module(*args, **kwargs)
+    ref.sum().backward()
+    _assert_dpa_backend(backend)
+    ref_grads = [t.grad.clone() for t in grads]
+    for t in grads:
+        t.grad = None
+
+    torch._dynamo.reset()
+    # Force backend selection to be re-run (and traced) inside the compiled
+    # region instead of being served from the cache the eager call populated.
+    _force_dpa_backend(monkeypatch, backend)
+    compiled = torch.compile(module, fullgraph=True)
+
+    out = compiled(*args, **kwargs)
+    out.sum().backward()
+    torch.cuda.synchronize()
+
+    _assert_dpa_backend(backend)
+    torch.testing.assert_close(out, ref)
+    for tensor, ref_grad in zip(grads, ref_grads):
+        assert tensor.grad is not None
+        torch.testing.assert_close(tensor.grad, ref_grad)
+
+
+@pytest.mark.parametrize("backend", ["flash", "unfused"])
+def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
+    """`mode="reduce-overhead"`: forward and backward of DotProductAttention
+    are captured into CUDA graphs and replayed on subsequent iterations."""
+    dtype = torch.bfloat16
+    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    _force_dpa_backend(monkeypatch, backend)
+
+    module = _make_dpa(cfg, dtype)
+
+    torch._dynamo.reset()
+    compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead")
+
+    for _ in range(3):
+        args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
+        out = compiled(*args, **kwargs)
+        out.sum().backward()
+        torch.cuda.synchronize()
+        assert torch.isfinite(out).all()
+        for tensor in grads:
+            assert tensor.grad is not None
+    _assert_dpa_backend(backend)
+
+
+def test_dpa_torch_compile_fused_backend(monkeypatch):
+    """The FusedAttention backend is an eager island, so it graph-breaks rather
+    than compiles; DotProductAttention must still be correct around it."""
+    dtype = torch.bfloat16
+    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    _force_dpa_backend(monkeypatch, "fused")
+
+    module = _make_dpa(cfg, dtype)
+    args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
+
+    ref = module(*args, **kwargs)
+    from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (
+        _attention_backends,
+    )
+
+    if not _attention_backends["use_fused_attention"]:
+        pytest.skip("FusedAttention does not support this config on this device")
+    ref.sum().backward()
+    ref_grads = [t.grad.clone() for t in grads]
+    for t in grads:
+        t.grad = None
+
+    torch._dynamo.reset()
+    _force_dpa_backend(monkeypatch, "fused")
+    out = torch.compile(module)(*args, **kwargs)
+    out.sum().backward()
+    torch.cuda.synchronize()
+
+    _assert_dpa_backend("fused")
+    torch.testing.assert_close(out, ref)
+    for tensor, ref_grad in zip(grads, ref_grads):
+        torch.testing.assert_close(tensor.grad, ref_grad)
+
+
 # ---------------------------------------------------------------------------
 # get_attention_backend under torch.compile
 # ---------------------------------------------------------------------------
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 4cc4cab1b8..0b21adb15c 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -42,7 +42,6 @@
     CudaRNGStatesTracker,
     graph_safe_rng_available,
 )
-from transformer_engine.pytorch.jit import no_torch_dynamo
 from transformer_engine.pytorch.graph import is_graph_capturing
 from transformer_engine.pytorch.attention.inference import InferenceParams
 
@@ -715,6 +714,22 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None:
         """
         _original_recipe = self.fp8_meta.get("recipe", None)
 
+        fp8 = FP8GlobalStateManager.is_fp8_enabled()
+        fp8_calibration = FP8GlobalStateManager.is_fp8_calibration()
+        fp8_parameters = FP8GlobalStateManager.with_fp8_parameters()
+        if not (fp8 or fp8_calibration or fp8_parameters):
+            # Nothing below has an effect when FP8 is off everywhere -- the recipe
+            # juggling ends in the "turn off and return" branch. Bailing out here
+            # also keeps DotProductAttention traceable by torch.compile: with no
+            # autocast active, get_fp8_recipe() below builds a default Recipe,
+            # whose construction is not traceable.
+            self.fast_setattr("fp8_parameters", fp8_parameters)
+            self.fast_setattr("fp8", fp8)
+            self.fast_setattr("fp8_calibration", fp8_calibration)
+            self.fp8_meta["fp8_checkpoint"] = False
+            self.fast_setattr("fp8_initialized", False)
+            return
+
         # global recipe set in autocast()
         fp8_recipe = FP8GlobalStateManager.get_fp8_recipe()
         if fp8_recipe.custom():
@@ -1090,7 +1105,6 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
-    @no_torch_dynamo(recursive=False)
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,
@@ -1820,18 +1834,23 @@ def forward(
                     _attention_backends["fused_attention_backend"] = fused_attention_backend
                     _attention_backends["use_unfused_attention"] = use_unfused_attention
                     _attention_backends["backend_selection_requires_update"] = False
-                    if use_flash_attention:
-                        self.logger.info(
-                            "Running with FlashAttention backend (version %s)",
-                            flash_attention_backend,
-                        )
-                    elif use_fused_attention:
-                        self.logger.info(
-                            "Running with FusedAttention backend (sub-backend %s)",
-                            int(fused_attention_backend),
-                        )
-                    elif use_unfused_attention:
-                        self.logger.info("Running with UnfusedDotProductAttention backend")
+                    # Backend selection is only logged in eager mode: dynamo
+                    # graph-breaks on logging.Logger methods, and even the
+                    # arguments below are untraceable (int() of the sub-backend
+                    # enum). Same restriction as in get_attention_backend.
+                    if not torch.compiler.is_compiling():
+                        if use_flash_attention:
+                            self.logger.info(
+                                "Running with FlashAttention backend (version %s)",
+                                flash_attention_backend,
+                            )
+                        elif use_fused_attention:
+                            self.logger.info(
+                                "Running with FusedAttention backend (sub-backend %s)",
+                                fused_attention_backend,
+                            )
+                        elif use_unfused_attention:
+                            self.logger.info("Running with UnfusedDotProductAttention backend")
                 else:
                     use_flash_attention = _attention_backends["use_flash_attention"]
                     flash_attention_backend = _attention_backends["flash_attention_backend"]
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 8a4c6ce286..4a8b7beb14 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -361,11 +361,17 @@ def _get_fused_attn_backend(
     *args,
 ):
     """Constant-foldable tex.get_fused_attn_backend: the result depends only on
-    the attention config, and the python-side enum keeps it traceable by
-    torch.compile (see the FusedAttnBackend docstring). Layout/bias/mask/softmax
-    are taken as their string keys and resolved to the pybind enums here, so
-    that every argument is a python literal or a python enum."""
-    return FusedAttnBackend.cast(
+    the attention config. Layout/bias/mask/softmax are taken as their string
+    keys and resolved to the pybind enums here, so that every argument is a
+    python literal or a python enum.
+
+    Returns a plain int rather than a FusedAttnBackend member: dynamo
+    reconstructs the result of an assume_constant_result call by re-emitting the
+    call, which is only valid inside the frame that made it. An int survives a
+    graph break because it is baked into the graph as a literal, while an enum
+    member comes out of the reconstruction corrupted (see the cast at the call
+    site, which restores the enum)."""
+    return int(
         tex.get_fused_attn_backend(
             is_training,
             q_type,
@@ -397,6 +403,9 @@ def get_attention_backend(
         Whether the `FusedAttention` backend has been selected.
     fused_attention_backend : FusedAttnBackend
         If `use_fused_attention = True`, one of `FusedAttention` three sub-backends, else `None`.
+        Under `torch.compile` this is the sub-backend's plain integer value instead of the enum
+        member (see the cast below); `FusedAttnBackend.cast` accepts both, as does every consumer
+        in `transformer_engine.pytorch`.
     use_unfused_attention : bool
         Whether the `UnfusedDotProductAttention` backend has been selected.
     available_backends : List[bool]
@@ -1455,11 +1464,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
             cuda_graph,
             deterministic,
         )
-        if fused_attention_backend == FusedAttnBackend["No_Backend"]:
+        if fused_attention_backend == FusedAttnBackend.No_Backend.value:
             logger.debug("Disabling FusedAttention as no backend supports the provided input")
             use_fused_attention = False
             fused_attention_backend = None
-        elif has_score_mod and fused_attention_backend != FusedAttnBackend["F16_arbitrary_seqlen"]:
+        elif (
+            has_score_mod and fused_attention_backend != FusedAttnBackend.F16_arbitrary_seqlen.value
+        ):
             logger.debug(
                 "Disabling FusedAttention for score_mod because sub-backend %s is not "
                 "F16/BF16 arbitrary-seqlen",
@@ -1509,7 +1520,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
             use_fused_attention = False
             fused_attention_backend = None
         if (
-            fused_attention_backend == FusedAttnBackend["FP8"]
+            fused_attention_backend == FusedAttnBackend.FP8.value
             and is_training
             and (device_compute_capability < (9, 0) or cudnn_version < (9, 19, 0))
         ):
@@ -1520,7 +1531,7 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
             use_fused_attention = False
             fused_attention_backend = None
         if (
-            fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]
+            fused_attention_backend == FusedAttnBackend.F16_arbitrary_seqlen.value
             and is_training
             and (
                 device_compute_capability < (9, 0)
@@ -1589,6 +1600,13 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
         use_flash_attention_4 = False
     use_flash_attention = use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4
     available_backends = [use_flash_attention, use_fused_attention, use_unfused_attention]
+    if fused_attention_backend is not None and not torch.compiler.is_compiling():
+        # The fused sub-backend is kept as a plain int while tracing: it has to
+        # survive a graph break (DotProductAttention hands it to FusedAttention,
+        # which is an eager island) and an enum member does not -- dynamo
+        # reconstructs the constant-folded member into the resumed frame as the
+        # function that produced it. Eager callers get the enum, as before.
+        fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend)
     if use_flash_attention_2:
         flash_attention_backend = FlashAttentionUtils.version
     if use_flash_attention_3:
@@ -2025,7 +2043,10 @@ def _get_cu_seqlens(batch_size, max_seqlen, device):
             device=device,
         )
 
-    if is_in_onnx_export_mode():
+    if is_in_onnx_export_mode() or torch.compiler.is_compiling():
+        # torch.is_inference_mode_enabled() (the cache key below) is a fundamental
+        # graph break under torch.compile; building the tensor is a single arange
+        # that the compiler captures into the graph anyway.
         return _get_cu_seqlens(batch_size, max_seqlen, device)
 
     is_inference = torch.is_inference_mode_enabled()
@@ -2538,7 +2559,22 @@ def run_iteratively(q, k, v):
 
         return qkv_layout
 
-    if not is_in_onnx_export_mode():
+    if torch.compiler.is_compiling():
+        # The detection in run_iteratively is untraceable by dynamo -- it reads
+        # UntypedStorage.data_ptr and Tensor.storage_offset -- and pointers are
+        # meaningless on the fake tensors used while tracing anyway. Under
+        # torch.compile, packed q/k/v must therefore be declared explicitly, via
+        # DotProductAttention's qkv_layer/kv_layer arguments (which bypass this
+        # function entirely); q/k/v passed separately are treated as three
+        # independent tensors. Views into a packed buffer are non-contiguous, so
+        # making the inputs contiguous keeps the returned layout truthful for
+        # every backend.
+        q, k, v = [x if x.is_contiguous() else x.contiguous() for x in [q, k, v]]
+        if is_same_q_kv_format:
+            qkv_layout = "_".join([qkv_format] * 3)
+        else:
+            qkv_layout = q_format + "_" + kv_format + "_" + kv_format
+    elif not is_in_onnx_export_mode():
         qkv_layout = run_iteratively(q, k, v)
     else:
         qkv_layout = "not_supported"
diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py
index 046019ee58..3a4e5aac66 100644
--- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py
+++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py
@@ -107,10 +107,13 @@ class FusedAttnBackend(IntEnum):
     ``transformer_engine_torch.NVTE_Fused_Attn_Backend`` (pybind11) enum
     value-for-value, and instances of the two enums compare equal when they
     share the same integer value. Unlike the pybind enum, a plain-python
-    ``IntEnum`` is traceable by ``torch.compile``: comparisons constant-fold
-    cleanly and instances safely cross the ``assume_constant_result`` boundary
-    in ``get_attention_backend``. Lookup by name (``FusedAttnBackend["FP8"]``)
-    works the same way as with the dict this used to be.
+    ``IntEnum`` is traceable by ``torch.compile``: comparisons against a member
+    constant-fold cleanly. Lookup by name (``FusedAttnBackend["FP8"]``) works
+    the same way as with the dict this used to be.
+
+    Members do not survive a graph break, though, so ``get_attention_backend``
+    hands out plain ints while tracing -- ``cast`` normalizes both back to a
+    member.
     """
 
     No_Backend = int(NVTE_Fused_Attn_Backend.NVTE_No_Backend)

From 3053d76f97dba1e4b1abba8ed009c9ff176a4310 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 28 Jul 2026 12:51:11 +0200
Subject: [PATCH 02/67] Delegate the no-FP8 case of DPA's init_fp8_metadata to
 the base class

The override hand-rolled the state that TransformerEngineBaseModule already
sets when FP8 is off everywhere (the three flags, fp8_checkpoint and
fp8_initialized). Call super() instead, and read the flags straight off
FP8GlobalStateManager.quantization_state, as the other hot paths do.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention.py                  | 20 ++++++-------------
 1 file changed, 6 insertions(+), 14 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 0b21adb15c..41aef33280 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -714,20 +714,12 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None:
         """
         _original_recipe = self.fp8_meta.get("recipe", None)
 
-        fp8 = FP8GlobalStateManager.is_fp8_enabled()
-        fp8_calibration = FP8GlobalStateManager.is_fp8_calibration()
-        fp8_parameters = FP8GlobalStateManager.with_fp8_parameters()
-        if not (fp8 or fp8_calibration or fp8_parameters):
-            # Nothing below has an effect when FP8 is off everywhere -- the recipe
-            # juggling ends in the "turn off and return" branch. Bailing out here
-            # also keeps DotProductAttention traceable by torch.compile: with no
-            # autocast active, get_fp8_recipe() below builds a default Recipe,
-            # whose construction is not traceable.
-            self.fast_setattr("fp8_parameters", fp8_parameters)
-            self.fast_setattr("fp8", fp8)
-            self.fast_setattr("fp8_calibration", fp8_calibration)
-            self.fp8_meta["fp8_checkpoint"] = False
-            self.fast_setattr("fp8_initialized", False)
+        qstate = FP8GlobalStateManager.quantization_state
+        if not (qstate.fp8_enabled or qstate.fp8_calibration or qstate.fp8_parameters):
+            # Without FP8 the recipe juggling below is a no-op, and its
+            # get_fp8_recipe() call would build a default Recipe -- which
+            # torch.compile cannot trace.
+            super().init_fp8_metadata(num_gemms=num_gemms)
             return
 
         # global recipe set in autocast()

From d6b285ed1cf5ff6d036acadd0cc7b517d0af2002 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 28 Jul 2026 15:42:46 +0200
Subject: [PATCH 03/67] Run DPA eagerly for undeclared packed q/k/v instead of
 guessing the layout

get_qkv_layout recognizes packed q/k/v by inspecting data pointers and storage
offsets, which dynamo cannot trace. Rather than assume such inputs are three
separate tensors -- a layout that would not describe memory -- detect the case
from what is traceable (a strided tensor that does not own its whole storage)
and run the whole module as an eager island, with a warning. Declaring the
packing via qkv_layer/kv_layer keeps the call on the compiled path.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 41 ++++++++++++++++
 .../dot_product_attention.py                  | 49 +++++++++++++++++++
 .../attention/dot_product_attention/utils.py  | 33 ++++++++-----
 3 files changed, 112 insertions(+), 11 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index ed2ad90c76..af67ea8605 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -759,6 +759,47 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     _assert_dpa_backend(backend)
 
 
+@pytest.mark.parametrize("backend", ["flash", "unfused"])
+@pytest.mark.parametrize("interleave_dim", [-3, -2])
+def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interleave_dim):
+    """Packed q/k/v passed as plain views cannot be recognized while tracing (the
+    detection reads data pointers), so DotProductAttention runs that detection
+    eagerly instead -- correct results, at the cost of a graph break. Declaring
+    the packing via qkv_layer/kv_layer keeps it on the compiled path, which
+    test_dpa_torch_compile covers."""
+    dtype = torch.bfloat16
+    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    _force_dpa_backend(monkeypatch, backend)
+
+    module = _make_dpa(cfg, dtype)
+    b, s = cfg["batch_size"], cfg["max_seqlen_q"]
+    h, d = cfg["num_heads"], cfg["head_dim"]
+    shape = (b, s, 3, h, d) if interleave_dim == -3 else (b, s, h, 3, d)
+    qkv = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
+    q, k, v = [qkv.select(interleave_dim, i) for i in range(3)]
+
+    ref = module(q, k, v)
+    ref.sum().backward()
+    ref_grad = qkv.grad.clone()
+    qkv.grad = None
+
+    torch._dynamo.reset()
+    _force_dpa_backend(monkeypatch, backend)
+    with pytest.warns(UserWarning, match="Falling back to eager execution"):
+        out = torch.compile(module)(q, k, v)
+    out.sum().backward()
+    torch.cuda.synchronize()
+
+    torch.testing.assert_close(out, ref)
+    torch.testing.assert_close(qkv.grad, ref_grad)
+
+    # The same graph break is an error when the user asked for a full graph.
+    torch._dynamo.reset()
+    _force_dpa_backend(monkeypatch, backend)
+    with pytest.raises(Exception, match="torch.compiler.disable"):
+        torch.compile(module, fullgraph=True)(q, k, v)
+
+
 def test_dpa_torch_compile_fused_backend(monkeypatch):
     """The FusedAttention backend is an eager island, so it graph-breaks rather
     than compiles; DotProductAttention must still be correct around it."""
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 41aef33280..d00d2c3849 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -4,6 +4,7 @@
 
 """Attention."""
 from contextlib import nullcontext
+from functools import wraps
 import math
 import os
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -42,6 +43,7 @@
     CudaRNGStatesTracker,
     graph_safe_rng_available,
 )
+from transformer_engine.pytorch.jit import no_torch_dynamo
 from transformer_engine.pytorch.graph import is_graph_capturing
 from transformer_engine.pytorch.attention.inference import InferenceParams
 
@@ -195,6 +197,50 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim
     return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1)
 
 
+def _needs_eager_dpa(
+    _self,
+    query_layer: Optional[torch.Tensor] = None,
+    key_layer: Optional[torch.Tensor] = None,
+    value_layer: Optional[torch.Tensor] = None,
+    *_args,
+    qkv_layer: Optional[torch.Tensor] = None,
+    kv_layer: Optional[torch.Tensor] = None,
+    **_kwargs,
+) -> bool:
+    """Whether this DotProductAttention call has to run outside the graph, i.e.
+    whether it passes packed q/k/v without declaring them."""
+    if qkv_layer is not None or kv_layer is not None:
+        return False
+    return dpa_utils.qkv_layout_needs_detection(query_layer, key_layer, value_layer)
+
+
+def _eager_under_compile_if(needs_eager: Callable[..., bool], reason: str):
+    """Decorator running the wrapped method eagerly when `needs_eager` says the
+    call is unsupported on the compiled path."""
+
+    def decorator(fn):
+        # The warning belongs inside the dynamo-disabled function: warnings.warn
+        # graph-breaks on its own, masking the break that matters.
+        @no_torch_dynamo()
+        def eager_fn(*args, **kwargs):
+            warnings.warn(
+                f"Falling back to eager execution under torch.compile: {reason} is"
+                " unsupported on the compiled path (graph-breaks under fullgraph=True).",
+                stacklevel=3,
+            )
+            return fn(*args, **kwargs)
+
+        @wraps(fn)
+        def wrapper(*args, **kwargs):
+            if torch.compiler.is_compiling() and needs_eager(*args, **kwargs):
+                return eager_fn(*args, **kwargs)
+            return fn(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 def _unpack_packed_qkv(
     qkv_layer: Optional[torch.Tensor],
     kv_layer: Optional[torch.Tensor],
@@ -1097,6 +1143,9 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
+    @_eager_under_compile_if(
+        _needs_eager_dpa, "detecting packed q/k/v that were not declared via qkv_layer/kv_layer"
+    )
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 4a8b7beb14..6ea15942d1 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -59,7 +59,7 @@
 )
 from transformer_engine.pytorch.export import is_in_onnx_export_mode
 
-from transformer_engine.pytorch.jit import jit_fuser
+from transformer_engine.pytorch.jit import jit_fuser, no_torch_dynamo
 
 # NVTE_DEBUG = 0/1 # disables/enables debug mode, default = 0
 _NVTE_DEBUG = int(os.getenv("NVTE_DEBUG", "0"))
@@ -2386,6 +2386,20 @@ def get_qkv_format(
     return qkv_format, q_format, kv_format
 
 
+def qkv_layout_needs_detection(*qkv: Optional[torch.Tensor]) -> bool:
+    """Whether the layout of these q/k/v can only be told by inspecting memory.
+
+    True for tensors that may be slices of a packed buffer, i.e. strided ones
+    that do not own their whole storage.
+    """
+    return any(
+        x is not None
+        and not x.is_contiguous()
+        and x.untyped_storage().size() != x.numel() * x.element_size()
+        for x in qkv
+    )
+
+
 def get_qkv_layout(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -2560,16 +2574,13 @@ def run_iteratively(q, k, v):
         return qkv_layout
 
     if torch.compiler.is_compiling():
-        # The detection in run_iteratively is untraceable by dynamo -- it reads
-        # UntypedStorage.data_ptr and Tensor.storage_offset -- and pointers are
-        # meaningless on the fake tensors used while tracing anyway. Under
-        # torch.compile, packed q/k/v must therefore be declared explicitly, via
-        # DotProductAttention's qkv_layer/kv_layer arguments (which bypass this
-        # function entirely); q/k/v passed separately are treated as three
-        # independent tensors. Views into a packed buffer are non-contiguous, so
-        # making the inputs contiguous keeps the returned layout truthful for
-        # every backend.
-        q, k, v = [x if x.is_contiguous() else x.contiguous() for x in [q, k, v]]
+        # run_iteratively reads data pointers and storage offsets, which dynamo
+        # cannot trace; unpacked q/k/v need no detection anyway.
+        assert not qkv_layout_needs_detection(q, k, v), (
+            "q/k/v may be views into a packed buffer, whose layout cannot be detected under"
+            " torch.compile. Pass the packed buffer explicitly via DotProductAttention's"
+            " qkv_layer/kv_layer (with qkv_interleave_dim)."
+        )
         if is_same_q_kv_format:
             qkv_layout = "_".join([qkv_format] * 3)
         else:

From 1cd6704912eff4d3fc2fe4841f49aa85bcfa431a Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 28 Jul 2026 16:23:11 +0200
Subject: [PATCH 04/67] Log backend selection in eager only; shorten comments

Log arguments are evaluated regardless of the logger, so the no-op logger
from #3189 is not enough here -- skip the whole block while tracing.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention/dot_product_attention.py       | 9 ++++-----
 .../pytorch/attention/dot_product_attention/utils.py     | 5 ++---
 2 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index d00d2c3849..fa261f0c61 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -1875,10 +1875,9 @@ def forward(
                     _attention_backends["fused_attention_backend"] = fused_attention_backend
                     _attention_backends["use_unfused_attention"] = use_unfused_attention
                     _attention_backends["backend_selection_requires_update"] = False
-                    # Backend selection is only logged in eager mode: dynamo
-                    # graph-breaks on logging.Logger methods, and even the
-                    # arguments below are untraceable (int() of the sub-backend
-                    # enum). Same restriction as in get_attention_backend.
+                    # Logged in eager only: dynamo graph-breaks on logging.Logger
+                    # methods, and log arguments are evaluated regardless of the
+                    # logger, so a no-op logger would not be enough.
                     if not torch.compiler.is_compiling():
                         if use_flash_attention:
                             self.logger.info(
@@ -1888,7 +1887,7 @@ def forward(
                         elif use_fused_attention:
                             self.logger.info(
                                 "Running with FusedAttention backend (sub-backend %s)",
-                                fused_attention_backend,
+                                int(fused_attention_backend),
                             )
                         elif use_unfused_attention:
                             self.logger.info("Running with UnfusedDotProductAttention backend")
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 6ea15942d1..84868f7740 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -2044,9 +2044,8 @@ def _get_cu_seqlens(batch_size, max_seqlen, device):
         )
 
     if is_in_onnx_export_mode() or torch.compiler.is_compiling():
-        # torch.is_inference_mode_enabled() (the cache key below) is a fundamental
-        # graph break under torch.compile; building the tensor is a single arange
-        # that the compiler captures into the graph anyway.
+        # The cache only saves an allocation, and its key -- is_inference_mode_enabled()
+        # -- is a fundamental graph break.
         return _get_cu_seqlens(batch_size, max_seqlen, device)
 
     is_inference = torch.is_inference_mode_enabled()

From 42a9e6cf209a3a10134a92bf2aec28795513e0af Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 09:04:37 +0200
Subject: [PATCH 05/67] Fix lint: unused import and keyword-arg-before-vararg

- utils.py no longer needs no_torch_dynamo: the eager island moved to
  dot_product_attention.py.
- _needs_eager_dpa reads q/k/v out of *args/**kwargs instead of naming them
  before *args, which pylint flags (W1113).

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention.py                  | 25 +++++++++----------
 .../attention/dot_product_attention/utils.py  |  2 +-
 2 files changed, 13 insertions(+), 14 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index fa261f0c61..d842a8e55b 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -197,21 +197,20 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim
     return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1)
 
 
-def _needs_eager_dpa(
-    _self,
-    query_layer: Optional[torch.Tensor] = None,
-    key_layer: Optional[torch.Tensor] = None,
-    value_layer: Optional[torch.Tensor] = None,
-    *_args,
-    qkv_layer: Optional[torch.Tensor] = None,
-    kv_layer: Optional[torch.Tensor] = None,
-    **_kwargs,
-) -> bool:
+def _needs_eager_dpa(*args, **kwargs) -> bool:
     """Whether this DotProductAttention call has to run outside the graph, i.e.
-    whether it passes packed q/k/v without declaring them."""
-    if qkv_layer is not None or kv_layer is not None:
+    whether it passes packed q/k/v without declaring them.
+
+    Takes `DotProductAttention.forward`'s arguments as they were passed, so q/k/v
+    are read from either side: `args[0]` is `self`, followed by the positional
+    query/key/value.
+    """
+    if kwargs.get("qkv_layer") is not None or kwargs.get("kv_layer") is not None:
         return False
-    return dpa_utils.qkv_layout_needs_detection(query_layer, key_layer, value_layer)
+    qkv = list(args[1:4]) + [
+        kwargs.get(name) for name in ("query_layer", "key_layer", "value_layer")
+    ]
+    return dpa_utils.qkv_layout_needs_detection(*qkv)
 
 
 def _eager_under_compile_if(needs_eager: Callable[..., bool], reason: str):
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 84868f7740..95b5a7f15f 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -59,7 +59,7 @@
 )
 from transformer_engine.pytorch.export import is_in_onnx_export_mode
 
-from transformer_engine.pytorch.jit import jit_fuser, no_torch_dynamo
+from transformer_engine.pytorch.jit import jit_fuser
 
 # NVTE_DEBUG = 0/1 # disables/enables debug mode, default = 0
 _NVTE_DEBUG = int(os.getenv("NVTE_DEBUG", "0"))

From af11f8f13b0d5e3e99aef2f6067634faf94da8cb Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 09:59:30 +0200
Subject: [PATCH 06/67] Compare against eager under CUDA graphs, and tighten
 the tolerances

The CUDA-graphs test only checked that the output was finite and that
gradients existed. A replay that reuses stale buffers -- the failure mode
that test is there to catch -- satisfies both, so it now compares against
eager on fresh inputs every iteration.

The remaining comparisons used the dtype defaults (rtol=1.6e-2 for bfloat16),
which is far looser than what these paths deliver: compiled and eager agree
bit-for-bit on every config and backend. FlashAttention now has to match
exactly, as it runs the same kernel either way; the unfused backend is
allowed a single bfloat16 rounding, since inductor may reassociate the
softmax sums.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 45 +++++++++++++++++++++++------
 1 file changed, 36 insertions(+), 9 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index af67ea8605..050f0f25a5 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -693,6 +693,19 @@ def _assert_dpa_backend(backend: str) -> None:
     )
 
 
+def _dpa_tolerances(backend: str) -> dict:
+    """How closely the compiled backend has to match eager.
+
+    FlashAttention calls the same kernel either way, so it has to match exactly.
+    The unfused backend is compiled by inductor, which may fuse the softmax
+    chain differently and reassociate its sums; allow a single bfloat16 rounding
+    (eps = 2**-8) for that, which is still 4x tighter than the dtype default.
+    """
+    if backend == "flash":
+        return {"rtol": 0.0, "atol": 0.0}
+    return {"rtol": 2**-8, "atol": 1e-5}
+
+
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
 @pytest.mark.parametrize("config", _DPA_COMPILE_CONFIGS.keys())
 def test_dpa_torch_compile(monkeypatch, backend, config):
@@ -729,10 +742,11 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
     torch.cuda.synchronize()
 
     _assert_dpa_backend(backend)
-    torch.testing.assert_close(out, ref)
+    tols = _dpa_tolerances(backend)
+    torch.testing.assert_close(out, ref, **tols)
     for tensor, ref_grad in zip(grads, ref_grads):
         assert tensor.grad is not None
-        torch.testing.assert_close(tensor.grad, ref_grad)
+        torch.testing.assert_close(tensor.grad, ref_grad, **tols)
 
 
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
@@ -748,14 +762,26 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     torch._dynamo.reset()
     compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead")
 
+    tols = _dpa_tolerances(backend)
     for _ in range(3):
+        # Fresh inputs every iteration: a replay that reuses stale buffers would
+        # still produce finite values and non-None gradients, so only comparing
+        # against eager catches it.
         args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
-        out = compiled(*args, **kwargs)
+        # CUDA graphs hand back tensors owned by their memory pool, which the
+        # next replay overwrites.
+        out = compiled(*args, **kwargs).clone()
         out.sum().backward()
         torch.cuda.synchronize()
-        assert torch.isfinite(out).all()
+        compiled_grads = [t.grad.clone() for t in grads]
         for tensor in grads:
-            assert tensor.grad is not None
+            tensor.grad = None
+
+        ref = module(*args, **kwargs)
+        ref.sum().backward()
+        torch.testing.assert_close(out, ref, **tols)
+        for compiled_grad, tensor in zip(compiled_grads, grads):
+            torch.testing.assert_close(compiled_grad, tensor.grad, **tols)
     _assert_dpa_backend(backend)
 
 
@@ -790,8 +816,9 @@ def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interlea
     out.sum().backward()
     torch.cuda.synchronize()
 
-    torch.testing.assert_close(out, ref)
-    torch.testing.assert_close(qkv.grad, ref_grad)
+    tols = _dpa_tolerances(backend)
+    torch.testing.assert_close(out, ref, **tols)
+    torch.testing.assert_close(qkv.grad, ref_grad, **tols)
 
     # The same graph break is an error when the user asked for a full graph.
     torch._dynamo.reset()
@@ -829,9 +856,9 @@ def test_dpa_torch_compile_fused_backend(monkeypatch):
     torch.cuda.synchronize()
 
     _assert_dpa_backend("fused")
-    torch.testing.assert_close(out, ref)
+    torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0)
     for tensor, ref_grad in zip(grads, ref_grads):
-        torch.testing.assert_close(tensor.grad, ref_grad)
+        torch.testing.assert_close(tensor.grad, ref_grad, rtol=0.0, atol=0.0)
 
 
 # ---------------------------------------------------------------------------

From 0b64c9695abc27aab434326c258e0edc5c4dac02 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 10:16:55 +0200
Subject: [PATCH 07/67] Support one cu_seqlens tensor for both q and kv under
 torch.compile

Self attention on thd naturally passes the same tensor as cu_seqlens_q and
cu_seqlens_kv, and flash-attn's varlen entry point forwards both to the same
autograd.Function -- which dynamo refuses to trace ("duplicate tensor input"),
failing the compilation under fullgraph=True.

Hand the varlen call a copy of one of them while tracing. The clone is
b + 1 int32 elements, only on that path, and only under torch.compile.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 43 +++++++++++++++++++
 .../dot_product_attention/backends.py         | 15 ++++---
 2 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 050f0f25a5..71eb58a10d 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -785,6 +785,49 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     _assert_dpa_backend(backend)
 
 
+@pytest.mark.parametrize("backend", ["flash", "unfused"])
+def test_dpa_torch_compile_shared_cu_seqlens(monkeypatch, backend):
+    """Self attention on `thd` naturally passes one cu_seqlens tensor for both q
+    and kv, which flash-attn hands to two inputs of the same autograd.Function --
+    something dynamo cannot trace. The other tests build two tensors, so this
+    covers the shape a caller is most likely to write."""
+    dtype = torch.bfloat16
+    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_thd_padding_causal"])
+    _force_dpa_backend(monkeypatch, backend)
+
+    module = _make_dpa(cfg, dtype)
+    b, s = cfg["batch_size"], cfg["max_seqlen_q"]
+    h, d = cfg["num_heads"], cfg["head_dim"]
+    q, k, v = [
+        torch.randn(b * s, h, d, dtype=dtype, device="cuda", requires_grad=True) for _ in range(3)
+    ]
+    cu_seqlens = torch.arange(0, (b + 1) * s, step=s, dtype=torch.int32, device="cuda")
+    kwargs = dict(
+        cu_seqlens_q=cu_seqlens,
+        cu_seqlens_kv=cu_seqlens,  # the same tensor, not a copy
+        max_seqlen_q=s,
+        max_seqlen_kv=s,
+    )
+
+    ref = module(q, k, v, **kwargs)
+    ref.sum().backward()
+    ref_grads = [t.grad.clone() for t in (q, k, v)]
+    for t in (q, k, v):
+        t.grad = None
+
+    torch._dynamo.reset()
+    _force_dpa_backend(monkeypatch, backend)
+    out = torch.compile(module, fullgraph=True)(q, k, v, **kwargs)
+    out.sum().backward()
+    torch.cuda.synchronize()
+
+    _assert_dpa_backend(backend)
+    tols = _dpa_tolerances(backend)
+    torch.testing.assert_close(out, ref, **tols)
+    for tensor, ref_grad in zip((q, k, v), ref_grads):
+        torch.testing.assert_close(tensor.grad, ref_grad, **tols)
+
+
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
 @pytest.mark.parametrize("interleave_dim", [-3, -2])
 def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interleave_dim):
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 6a60160f77..f0382be30c 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -1122,12 +1122,15 @@ def forward(
                     else:
                         func = flash_attn_with_kvcache_v3  # pylint: disable=possibly-used-before-assignment
                     if not use_flash_attn_4 and (not use_flash_attn_3 or inference_params is None):
-                        fa_optional_forward_args_thd.append(
-                            cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q
-                        )
-                        fa_optional_forward_args_thd.append(
-                            cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv
-                        )
+                        cu_q = cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q
+                        cu_kv = cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv
+                        if cu_q is cu_kv and torch.compiler.is_compiling():
+                            # Self attention passes one tensor for both, and dynamo
+                            # cannot trace autograd.Function.apply with the same
+                            # tensor on two of its inputs.
+                            cu_kv = cu_kv.clone()
+                        fa_optional_forward_args_thd.append(cu_q)
+                        fa_optional_forward_args_thd.append(cu_kv)
                         fa_optional_forward_args_thd.append(max_seqlen_q)
                         fa_optional_forward_args_thd.append(max_seqlen_kv)
                 if use_flash_attn_4:

From 7f480c6d07a9d04bca718f581e6a5e6b3e9bb113 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 10:21:49 +0200
Subject: [PATCH 08/67] Apply the shared-cu_seqlens workaround to the
 FlashAttention v4 path too

Pull the deduplication into a helper and use it wherever both cumulative
sequence lengths are handed to flash-attn: the varlen path shared by v2 and
v3 training, and v4's varlen kwargs. The v3 KV-cache path derives the kv
lengths by subtraction, so its two arguments are distinct by construction.

v4 is not installed in the environment this was tested in, so that call site
is covered by construction rather than by a test; the helper is a no-op
unless the two arguments are the same object.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention/backends.py         | 28 +++++++++++++------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index f0382be30c..2c47ff8a4b 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -793,6 +793,18 @@ def backward(
         return dq, dk, dv
 
 
+def _distinct_cu_seqlens(cu_q: torch.Tensor, cu_kv: torch.Tensor):
+    """Make q's and kv's cumulative sequence lengths distinct objects.
+
+    Self attention passes one tensor for both, and flash-attn forwards them to
+    two inputs of the same autograd.Function -- which dynamo cannot trace. In
+    eager the duplicate is harmless, so nothing is copied there.
+    """
+    if cu_q is cu_kv and torch.compiler.is_compiling():
+        return cu_q, cu_kv.clone()
+    return cu_q, cu_kv
+
+
 class FlashAttention(torch.nn.Module):
     """Dot product attention, using HazyResearch flash-attn package:
     https://github.com/Dao-AILab/flash-attention
@@ -1122,13 +1134,10 @@ def forward(
                     else:
                         func = flash_attn_with_kvcache_v3  # pylint: disable=possibly-used-before-assignment
                     if not use_flash_attn_4 and (not use_flash_attn_3 or inference_params is None):
-                        cu_q = cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q
-                        cu_kv = cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv
-                        if cu_q is cu_kv and torch.compiler.is_compiling():
-                            # Self attention passes one tensor for both, and dynamo
-                            # cannot trace autograd.Function.apply with the same
-                            # tensor on two of its inputs.
-                            cu_kv = cu_kv.clone()
+                        cu_q, cu_kv = _distinct_cu_seqlens(
+                            cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q,
+                            cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv,
+                        )
                         fa_optional_forward_args_thd.append(cu_q)
                         fa_optional_forward_args_thd.append(cu_kv)
                         fa_optional_forward_args_thd.append(max_seqlen_q)
@@ -1141,8 +1150,9 @@ def forward(
                     if inference_params is None:
                         fa_4_optional_forward_kwargs["deterministic"] = self.deterministic
                     if func is flash_attn_varlen_func_v4:
-                        fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_seqlens_q
-                        fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_seqlens_kv
+                        cu_q, cu_kv = _distinct_cu_seqlens(cu_seqlens_q, cu_seqlens_kv)
+                        fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_q
+                        fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_kv
                         fa_4_optional_forward_kwargs["max_seqlen_q"] = max_seqlen_q
                         fa_4_optional_forward_kwargs["max_seqlen_k"] = max_seqlen_kv
                     output = func(

From adc96c48a9b455678acbda96e367e8a0c63d16b2 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 10:58:31 +0200
Subject: [PATCH 09/67] Run FP8 attention eagerly under torch.compile

FP8 attention brings quantizers, fp8_meta mutation and Float8Tensor across
the graph boundary, and none of it is exercised by this PR's tests. Rather
than trace it untested, route it to the eager island that packed q/k/v
already use.

The predicate is deliberately narrow: it asks the recipe for fp8_dpa/fp8_mha
rather than whether an autocast is active, so FP8 GEMMs with attention in
high precision -- the common training setup -- stay on the compiled path.

The decorator now takes a predicate that returns the reason, so the warning
names which of the two paths was taken.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 34 ++++++++++++++
 .../dot_product_attention.py                  | 44 +++++++++++--------
 2 files changed, 60 insertions(+), 18 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 71eb58a10d..61e78e58b0 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3,6 +3,7 @@
 # See LICENSE for license information.
 
 import abc
+import warnings
 
 import pytest
 import torch
@@ -785,6 +786,39 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     _assert_dpa_backend(backend)
 
 
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("fp8_attention", [False, True], ids=["fp8_gemms_only", "fp8_attention"])
+def test_dpa_torch_compile_fp8(monkeypatch, fp8_attention):
+    """FP8 attention is not supported on the compiled path and falls back to
+    eager. FP8 elsewhere in the model with attention in high precision -- the
+    common training setup -- must stay compiled."""
+    dtype = torch.bfloat16
+    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    _force_dpa_backend(monkeypatch, "unfused")
+
+    module = _make_dpa(cfg, dtype)
+    args, kwargs, _ = _make_dpa_inputs(cfg, dtype)
+    fp8_recipe = recipe.DelayedScaling(fp8_dpa=fp8_attention)
+
+    def fn(*args, **kwargs):
+        with te.autocast(enabled=True, recipe=fp8_recipe):
+            return module(*args, **kwargs)
+
+    torch._dynamo.reset()
+    _force_dpa_backend(monkeypatch, "unfused")
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        try:
+            torch.compile(fn)(*args, **kwargs)
+        except Exception:  # pylint: disable=broad-except
+            # FP8 attention is not available on every device; what matters here
+            # is which path was taken, which the warning below reports.
+            pass
+        fell_back = any("Falling back to eager" in str(w.message) for w in caught)
+
+    assert fell_back == fp8_attention
+
+
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
 def test_dpa_torch_compile_shared_cu_seqlens(monkeypatch, backend):
     """Self attention on `thd` naturally passes one cu_seqlens tensor for both q
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index d842a8e55b..f82aa91c13 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -197,31 +197,39 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim
     return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1)
 
 
-def _needs_eager_dpa(*args, **kwargs) -> bool:
-    """Whether this DotProductAttention call has to run outside the graph, i.e.
-    whether it passes packed q/k/v without declaring them.
+def _needs_eager_dpa(*args, **kwargs) -> Optional[str]:
+    """Why this DotProductAttention call has to run outside the graph, or None.
 
     Takes `DotProductAttention.forward`'s arguments as they were passed, so q/k/v
     are read from either side: `args[0]` is `self`, followed by the positional
     query/key/value.
     """
-    if kwargs.get("qkv_layer") is not None or kwargs.get("kv_layer") is not None:
-        return False
-    qkv = list(args[1:4]) + [
-        kwargs.get(name) for name in ("query_layer", "key_layer", "value_layer")
-    ]
-    return dpa_utils.qkv_layout_needs_detection(*qkv)
+    # FP8 GEMMs with the attention itself in high precision -- the common
+    # training setup -- stay on the compiled path; only FP8 attention bails out.
+    qstate = FP8GlobalStateManager.quantization_state
+    fp8_recipe = qstate.fp8_recipe
+    if qstate.fp8_enabled and fp8_recipe is not None:
+        if fp8_recipe.fp8_dpa or fp8_recipe.fp8_mha:
+            return "FP8 attention"
+
+    if kwargs.get("qkv_layer") is None and kwargs.get("kv_layer") is None:
+        qkv = list(args[1:4]) + [
+            kwargs.get(name) for name in ("query_layer", "key_layer", "value_layer")
+        ]
+        if dpa_utils.qkv_layout_needs_detection(*qkv):
+            return "detecting packed q/k/v that were not declared via qkv_layer/kv_layer"
+    return None
 
 
-def _eager_under_compile_if(needs_eager: Callable[..., bool], reason: str):
-    """Decorator running the wrapped method eagerly when `needs_eager` says the
-    call is unsupported on the compiled path."""
+def _eager_under_compile_if(needs_eager: Callable[..., Optional[str]]):
+    """Decorator running the wrapped method eagerly whenever `needs_eager` gives
+    a reason why the call is unsupported on the compiled path."""
 
     def decorator(fn):
         # The warning belongs inside the dynamo-disabled function: warnings.warn
         # graph-breaks on its own, masking the break that matters.
         @no_torch_dynamo()
-        def eager_fn(*args, **kwargs):
+        def eager_fn(reason, *args, **kwargs):
             warnings.warn(
                 f"Falling back to eager execution under torch.compile: {reason} is"
                 " unsupported on the compiled path (graph-breaks under fullgraph=True).",
@@ -231,8 +239,10 @@ def eager_fn(*args, **kwargs):
 
         @wraps(fn)
         def wrapper(*args, **kwargs):
-            if torch.compiler.is_compiling() and needs_eager(*args, **kwargs):
-                return eager_fn(*args, **kwargs)
+            if torch.compiler.is_compiling():
+                reason = needs_eager(*args, **kwargs)
+                if reason is not None:
+                    return eager_fn(reason, *args, **kwargs)
             return fn(*args, **kwargs)
 
         return wrapper
@@ -1142,9 +1152,7 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
-    @_eager_under_compile_if(
-        _needs_eager_dpa, "detecting packed q/k/v that were not declared via qkv_layer/kv_layer"
-    )
+    @_eager_under_compile_if(_needs_eager_dpa)
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,

From 737904f09cc1957db238e73a68c7e1edeb231158 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 11:23:04 +0200
Subject: [PATCH 10/67] Drive the compile tests off ModelConfig instead of a
 hand-picked intersection

The configurations were chosen so that both backends could run all of them,
which meant nothing backend-specific was ever exercised: no biases, no
arbitrary masks, no MLA head dims, no sink softmax. Describe them with the
ModelConfig the eager attention tests already use, and let
get_available_attention_backends() decide which backend runs which, so a
configuration only one backend supports is covered rather than avoided.

Five such configurations are added. One of them immediately showed that the
tolerance for the unfused backend was too tight: an off-by-one softmax
differs from eager by two bfloat16 roundings on a single element out of
65536, because inductor reassociates the softmax sums. Only FlashAttention,
which runs the same kernel either way, keeps the exact comparison.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 208 +++++++++++++++++-----------
 1 file changed, 125 insertions(+), 83 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 61e78e58b0..77c6c31a3b 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -4,6 +4,7 @@
 
 import abc
 import warnings
+from typing import Optional
 
 import pytest
 import torch
@@ -37,7 +38,7 @@
     MXFP8Quantizer,
     NVFP4Quantizer,
 )
-from utils import recipe_id
+from utils import ModelConfig, get_available_attention_backends, recipe_id
 from transformer_engine.pytorch.attention.dot_product_attention.backends import (
     UnfusedDotProductAttention,
 )
@@ -556,93 +557,117 @@ def fn(q, k, v, extra):
 # ---------------------------------------------------------------------------
 
 
-# Each config is a model configuration that both the FlashAttention and the
-# UnfusedDotProductAttention backend support, so that every one of them runs
-# against both backends.
+# Model configurations, described with the same ModelConfig the eager attention
+# tests use. Which backends can run each of them is not hardcoded here --
+# get_available_attention_backends() answers that, so configurations only one
+# backend supports (arbitrary masks, biases, MLA head dims, ...) are covered
+# rather than avoided.
 _DPA_COMPILE_CONFIGS = {
-    "self_bshd_causal": dict(qkv_format="bshd", attn_mask_type="causal"),
-    "self_sbhd_no_mask": dict(qkv_format="sbhd", attn_mask_type="no_mask"),
-    "self_bshd_swa": dict(qkv_format="bshd", attn_mask_type="causal", window_size=(16, 0)),
-    "gqa_bshd_causal": dict(
-        qkv_format="bshd", attn_mask_type="causal", num_heads=8, num_gqa_groups=2
+    # name: (ModelConfig, qkv_format, packed layout or None)
+    "self_bshd_causal": (ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), "bshd", None),
+    "self_sbhd_no_mask": (ModelConfig(2, 128, 4, 64, attn_mask_type="no_mask"), "sbhd", None),
+    "self_bshd_swa": (
+        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", window_size=(16, 0)),
+        "bshd",
+        None,
     ),
-    "cross_bshd_no_mask": dict(
-        qkv_format="bshd", attn_mask_type="no_mask", attention_type="cross", max_seqlen_kv=256
+    "gqa_bshd_causal": (
+        ModelConfig(2, 128, 8, 64, num_gqa_groups=2, attn_mask_type="causal"),
+        "bshd",
+        None,
+    ),
+    "cross_bshd_no_mask": (
+        ModelConfig(2, 128, 4, 64, max_seqlen_kv=256, attn_mask_type="no_mask"),
+        "bshd",
+        None,
+    ),
+    "self_thd_padding_causal": (
+        ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"),
+        "thd",
+        None,
+    ),
+    "packed_qkv_bs3hd": (ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), "bshd", "bs3hd"),
+    # Configurations below are supported by one backend only, or take a code
+    # path of their own inside DotProductAttention.
+    "alibi_bshd_causal": (
+        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", attn_bias_type="alibi"),
+        "bshd",
+        None,
+    ),
+    "post_scale_bias_bshd": (
+        ModelConfig(2, 128, 4, 64, attn_bias_type="post_scale_bias", bias_shape="1hss"),
+        "bshd",
+        None,
+    ),
+    "arbitrary_mask_bshd": (ModelConfig(2, 128, 4, 64, attn_mask_type="arbitrary"), "bshd", None),
+    "mla_bshd_causal": (
+        ModelConfig(2, 128, 4, 128, head_dim_v=64, attn_mask_type="causal"),
+        "bshd",
+        None,
+    ),
+    "sink_softmax_bshd": (
+        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", softmax_type="off-by-one"),
+        "bshd",
+        None,
     ),
-    "self_thd_padding_causal": dict(qkv_format="thd", attn_mask_type="padding_causal"),
-    "packed_qkv_bs3hd": dict(qkv_format="bshd", attn_mask_type="causal", packed="bs3hd"),
 }
 
 
-def _dpa_config(**overrides):
-    cfg = dict(
-        batch_size=2,
-        num_heads=4,
-        num_gqa_groups=None,
-        head_dim=64,
-        max_seqlen_q=128,
-        max_seqlen_kv=128,
-        qkv_format="bshd",
-        attn_mask_type="causal",
-        attention_type="self",
-        window_size=None,
-        packed=None,
-    )
-    cfg.update(overrides)
-    if cfg["num_gqa_groups"] is None:
-        cfg["num_gqa_groups"] = cfg["num_heads"]
-    return cfg
+def _qkv_layout(qkv_format: str, packed: Optional[str]) -> str:
+    return packed if packed is not None else "_".join([qkv_format] * 3)
 
 
-def _make_dpa(cfg, dtype: torch.dtype) -> te.DotProductAttention:
+def _make_dpa(config: ModelConfig, qkv_format: str, dtype: torch.dtype) -> te.DotProductAttention:
     return te.DotProductAttention(
-        num_attention_heads=cfg["num_heads"],
-        kv_channels=cfg["head_dim"],
-        num_gqa_groups=cfg["num_gqa_groups"],
-        attention_dropout=0.0,
-        qkv_format=cfg["qkv_format"],
-        attn_mask_type=cfg["attn_mask_type"],
-        window_size=cfg["window_size"],
-        attention_type=cfg["attention_type"],
+        num_attention_heads=config.num_heads,
+        kv_channels=config.kv_channels,
+        num_gqa_groups=config.num_gqa_groups,
+        attention_dropout=config.dropout_p,
+        qkv_format=qkv_format,
+        attn_mask_type=config.attn_mask_type,
+        window_size=config.window_size,
+        attention_type=config.attn_type,
+        softmax_type=config.softmax_type,
     ).to(dtype=dtype, device="cuda")
 
 
-def _make_dpa_inputs(cfg, dtype: torch.dtype):
+def _make_dpa_inputs(config: ModelConfig, qkv_format: str, packed: Optional[str], dtype):
     """Build the (args, kwargs) that `DotProductAttention.forward` is called
     with, plus the list of tensors whose gradients the test compares."""
-    b = cfg["batch_size"]
-    s_q, s_kv = cfg["max_seqlen_q"], cfg["max_seqlen_kv"]
-    h, g, d = cfg["num_heads"], cfg["num_gqa_groups"], cfg["head_dim"]
-    fmt = cfg["qkv_format"]
+    b = config.batch_size
+    s_q, s_kv = config.max_seqlen_q, config.max_seqlen_kv
+    h, g = config.num_heads, config.num_gqa_groups
+    d_qk, d_v = config.head_dim_qk, config.head_dim_v
 
-    def _shape(s, heads):
+    def _shape(s, heads, head_dim):
         return {
-            "bshd": (b, s, heads, d),
-            "sbhd": (s, b, heads, d),
-            "thd": (b * s, heads, d),
-        }[fmt]
+            "bshd": (b, s, heads, head_dim),
+            "sbhd": (s, b, heads, head_dim),
+            "thd": (b * s, heads, head_dim),
+        }[qkv_format]
 
     def _randn(shape):
         return torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
 
     kwargs = {}
-    if cfg["packed"] is not None:
+    if packed is not None:
         # Declarative packed QKV: q/k/v are derived from one buffer by DPA
         # itself, and the layout comes from the declaration -- the only packed
-        # layout that torch.compile supports (pointer-based detection of packed
-        # q/k/v is untraceable).
-        assert cfg["packed"] == "bs3hd" and fmt == "bshd" and h == g
-        qkv = _randn((b, s_q, 3, h, d))
+        # layout that torch.compile supports.
+        assert packed == "bs3hd" and qkv_format == "bshd" and h == g and d_qk == d_v
+        qkv = _randn((b, s_q, 3, h, d_qk))
         kwargs["qkv_layer"] = qkv
         kwargs["qkv_interleave_dim"] = -3
         grad_tensors = [qkv]
         args = ()
     else:
-        q, k, v = _randn(_shape(s_q, h)), _randn(_shape(s_kv, g)), _randn(_shape(s_kv, g))
+        q = _randn(_shape(s_q, h, d_qk))
+        k = _randn(_shape(s_kv, g, d_qk))
+        v = _randn(_shape(s_kv, g, d_v))
         grad_tensors = [q, k, v]
         args = (q, k, v)
 
-    if fmt == "thd":
+    if qkv_format == "thd":
         # All sequences have the maximum length, i.e. no padding. max_seqlen_*
         # is passed explicitly: deriving it from cu_seqlens costs a `.item()`
         # (a device sync, and an unbacked SymInt under torch.compile).
@@ -655,9 +680,25 @@ def _randn(shape):
         kwargs["max_seqlen_q"] = s_q
         kwargs["max_seqlen_kv"] = s_kv
 
+    if config.attn_mask_type == "arbitrary":
+        kwargs["attention_mask"] = torch.zeros(b, 1, s_q, s_kv, dtype=torch.bool, device="cuda")
+    if config.attn_bias_type != "no_bias":
+        kwargs["core_attention_bias_type"] = config.attn_bias_type
+    if config.attn_bias_type == "post_scale_bias":
+        kwargs["core_attention_bias"] = torch.randn(1, h, s_q, s_kv, dtype=dtype, device="cuda")
+
     return args, kwargs, grad_tensors
 
 
+def _skip_unsupported(config: ModelConfig, qkv_layout: str, backend: str, dtype) -> None:
+    """Skip configurations the backend under test cannot run at all."""
+    available, _, _ = get_available_attention_backends(config, dtype, qkv_layout)
+    flash_supported, _, unfused_supported = available
+    supported = {"flash": flash_supported, "unfused": unfused_supported}[backend]
+    if not supported:
+        pytest.skip(f"the {backend} backend does not support this configuration")
+
+
 def _force_dpa_backend(monkeypatch, backend: str) -> None:
     """Restrict DotProductAttention to a single backend."""
     from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (
@@ -698,13 +739,13 @@ def _dpa_tolerances(backend: str) -> dict:
     """How closely the compiled backend has to match eager.
 
     FlashAttention calls the same kernel either way, so it has to match exactly.
-    The unfused backend is compiled by inductor, which may fuse the softmax
-    chain differently and reassociate its sums; allow a single bfloat16 rounding
-    (eps = 2**-8) for that, which is still 4x tighter than the dtype default.
+    The unfused backend is compiled by inductor, which fuses the softmax chain
+    differently and reassociates its sums: with an off-by-one softmax that costs
+    up to two bfloat16 roundings on individual elements, so it keeps the dtype
+    defaults. A systematic difference -- a wrong mask or scale -- lands far
+    outside them anyway.
     """
-    if backend == "flash":
-        return {"rtol": 0.0, "atol": 0.0}
-    return {"rtol": 2**-8, "atol": 1e-5}
+    return {"rtol": 0.0, "atol": 0.0} if backend == "flash" else {}
 
 
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
@@ -719,11 +760,12 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
     backend itself. The FusedAttention backend is not covered -- it is an eager
     island (`@no_torch_dynamo` on `FusedAttention.forward`)."""
     dtype = torch.bfloat16
-    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS[config])
+    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS[config]
+    _skip_unsupported(model_config, _qkv_layout(qkv_format, packed), backend, dtype)
     _force_dpa_backend(monkeypatch, backend)
 
-    module = _make_dpa(cfg, dtype)
-    args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
+    module = _make_dpa(model_config, qkv_format, dtype)
+    args, kwargs, grads = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
 
     ref = module(*args, **kwargs)
     ref.sum().backward()
@@ -755,10 +797,10 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     """`mode="reduce-overhead"`: forward and backward of DotProductAttention
     are captured into CUDA graphs and replayed on subsequent iterations."""
     dtype = torch.bfloat16
-    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _force_dpa_backend(monkeypatch, backend)
 
-    module = _make_dpa(cfg, dtype)
+    module = _make_dpa(model_config, qkv_format, dtype)
 
     torch._dynamo.reset()
     compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead")
@@ -768,7 +810,7 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
         # Fresh inputs every iteration: a replay that reuses stale buffers would
         # still produce finite values and non-None gradients, so only comparing
         # against eager catches it.
-        args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
+        args, kwargs, grads = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
         # CUDA graphs hand back tensors owned by their memory pool, which the
         # next replay overwrites.
         out = compiled(*args, **kwargs).clone()
@@ -793,11 +835,11 @@ def test_dpa_torch_compile_fp8(monkeypatch, fp8_attention):
     eager. FP8 elsewhere in the model with attention in high precision -- the
     common training setup -- must stay compiled."""
     dtype = torch.bfloat16
-    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _force_dpa_backend(monkeypatch, "unfused")
 
-    module = _make_dpa(cfg, dtype)
-    args, kwargs, _ = _make_dpa_inputs(cfg, dtype)
+    module = _make_dpa(model_config, qkv_format, dtype)
+    args, kwargs, _ = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
     fp8_recipe = recipe.DelayedScaling(fp8_dpa=fp8_attention)
 
     def fn(*args, **kwargs):
@@ -826,12 +868,12 @@ def test_dpa_torch_compile_shared_cu_seqlens(monkeypatch, backend):
     something dynamo cannot trace. The other tests build two tensors, so this
     covers the shape a caller is most likely to write."""
     dtype = torch.bfloat16
-    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_thd_padding_causal"])
+    model_config, qkv_format, _ = _DPA_COMPILE_CONFIGS["self_thd_padding_causal"]
     _force_dpa_backend(monkeypatch, backend)
 
-    module = _make_dpa(cfg, dtype)
-    b, s = cfg["batch_size"], cfg["max_seqlen_q"]
-    h, d = cfg["num_heads"], cfg["head_dim"]
+    module = _make_dpa(model_config, qkv_format, dtype)
+    b, s = model_config.batch_size, model_config.max_seqlen_q
+    h, d = model_config.num_heads, model_config.head_dim_qk
     q, k, v = [
         torch.randn(b * s, h, d, dtype=dtype, device="cuda", requires_grad=True) for _ in range(3)
     ]
@@ -871,12 +913,12 @@ def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interlea
     the packing via qkv_layer/kv_layer keeps it on the compiled path, which
     test_dpa_torch_compile covers."""
     dtype = torch.bfloat16
-    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    model_config, qkv_format, _ = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _force_dpa_backend(monkeypatch, backend)
 
-    module = _make_dpa(cfg, dtype)
-    b, s = cfg["batch_size"], cfg["max_seqlen_q"]
-    h, d = cfg["num_heads"], cfg["head_dim"]
+    module = _make_dpa(model_config, qkv_format, dtype)
+    b, s = model_config.batch_size, model_config.max_seqlen_q
+    h, d = model_config.num_heads, model_config.head_dim_qk
     shape = (b, s, 3, h, d) if interleave_dim == -3 else (b, s, h, 3, d)
     qkv = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
     q, k, v = [qkv.select(interleave_dim, i) for i in range(3)]
@@ -908,11 +950,11 @@ def test_dpa_torch_compile_fused_backend(monkeypatch):
     """The FusedAttention backend is an eager island, so it graph-breaks rather
     than compiles; DotProductAttention must still be correct around it."""
     dtype = torch.bfloat16
-    cfg = _dpa_config(**_DPA_COMPILE_CONFIGS["self_bshd_causal"])
+    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _force_dpa_backend(monkeypatch, "fused")
 
-    module = _make_dpa(cfg, dtype)
-    args, kwargs, grads = _make_dpa_inputs(cfg, dtype)
+    module = _make_dpa(model_config, qkv_format, dtype)
+    args, kwargs, grads = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
 
     ref = module(*args, **kwargs)
     from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (

From 9026a41c3303a1e602b6be6e99d2cefe5e4da64d Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 11:31:36 +0200
Subject: [PATCH 11/67] Scale the unfused comparison to the tensor, and take
 tolerances from utils

An off-by-one softmax exposed that comparing the unfused backend elementwise
does not hold up: inductor's reassociation of the softmax sums produces an
error proportional to those sums, which for gradient elements near zero is
far outside any relative tolerance -- while being one bfloat16 rounding at
the tensor's own scale.

Derive the absolute tolerance from the reference tensor's magnitude, and take
the dtype tolerances from tests/pytorch/utils.py rather than restating them.
FlashAttention and FusedAttention keep the exact comparison.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 53 ++++++++++++++++-------------
 1 file changed, 29 insertions(+), 24 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 77c6c31a3b..b88222be60 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -38,7 +38,7 @@
     MXFP8Quantizer,
     NVFP4Quantizer,
 )
-from utils import ModelConfig, get_available_attention_backends, recipe_id
+from utils import ModelConfig, dtype_tols, get_available_attention_backends, recipe_id
 from transformer_engine.pytorch.attention.dot_product_attention.backends import (
     UnfusedDotProductAttention,
 )
@@ -735,17 +735,26 @@ def _assert_dpa_backend(backend: str) -> None:
     )
 
 
-def _dpa_tolerances(backend: str) -> dict:
-    """How closely the compiled backend has to match eager.
+def _assert_matches_eager(actual, expected, backend: str, dtype: torch.dtype) -> None:
+    """Assert a compiled result matches the eager one.
 
-    FlashAttention calls the same kernel either way, so it has to match exactly.
-    The unfused backend is compiled by inductor, which fuses the softmax chain
-    differently and reassociates its sums: with an off-by-one softmax that costs
-    up to two bfloat16 roundings on individual elements, so it keeps the dtype
-    defaults. A systematic difference -- a wrong mask or scale -- lands far
-    outside them anyway.
+    FlashAttention and FusedAttention run the same kernel either way, so they
+    have to match exactly. Inductor reassociates the unfused backend's softmax
+    sums, and that error scales with the magnitude of those sums rather than
+    with each output element -- so the absolute tolerance comes from the
+    tensor's own scale, an elementwise relative tolerance being meaningless for
+    elements near zero.
     """
-    return {"rtol": 0.0, "atol": 0.0} if backend == "flash" else {}
+    if backend != "unfused":
+        torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
+        return
+    tols = dtype_tols(dtype)
+    torch.testing.assert_close(
+        actual,
+        expected,
+        rtol=tols["rtol"],
+        atol=max(tols["atol"], 1e-3 * expected.abs().max().item()),
+    )
 
 
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
@@ -785,11 +794,10 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
     torch.cuda.synchronize()
 
     _assert_dpa_backend(backend)
-    tols = _dpa_tolerances(backend)
-    torch.testing.assert_close(out, ref, **tols)
+    _assert_matches_eager(out, ref, backend, dtype)
     for tensor, ref_grad in zip(grads, ref_grads):
         assert tensor.grad is not None
-        torch.testing.assert_close(tensor.grad, ref_grad, **tols)
+        _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
 
 
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
@@ -805,7 +813,6 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     torch._dynamo.reset()
     compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead")
 
-    tols = _dpa_tolerances(backend)
     for _ in range(3):
         # Fresh inputs every iteration: a replay that reuses stale buffers would
         # still produce finite values and non-None gradients, so only comparing
@@ -822,9 +829,9 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
 
         ref = module(*args, **kwargs)
         ref.sum().backward()
-        torch.testing.assert_close(out, ref, **tols)
+        _assert_matches_eager(out, ref, backend, dtype)
         for compiled_grad, tensor in zip(compiled_grads, grads):
-            torch.testing.assert_close(compiled_grad, tensor.grad, **tols)
+            _assert_matches_eager(compiled_grad, tensor.grad, backend, dtype)
     _assert_dpa_backend(backend)
 
 
@@ -898,10 +905,9 @@ def test_dpa_torch_compile_shared_cu_seqlens(monkeypatch, backend):
     torch.cuda.synchronize()
 
     _assert_dpa_backend(backend)
-    tols = _dpa_tolerances(backend)
-    torch.testing.assert_close(out, ref, **tols)
+    _assert_matches_eager(out, ref, backend, dtype)
     for tensor, ref_grad in zip((q, k, v), ref_grads):
-        torch.testing.assert_close(tensor.grad, ref_grad, **tols)
+        _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
 
 
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
@@ -935,9 +941,8 @@ def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interlea
     out.sum().backward()
     torch.cuda.synchronize()
 
-    tols = _dpa_tolerances(backend)
-    torch.testing.assert_close(out, ref, **tols)
-    torch.testing.assert_close(qkv.grad, ref_grad, **tols)
+    _assert_matches_eager(out, ref, backend, dtype)
+    _assert_matches_eager(qkv.grad, ref_grad, backend, dtype)
 
     # The same graph break is an error when the user asked for a full graph.
     torch._dynamo.reset()
@@ -975,9 +980,9 @@ def test_dpa_torch_compile_fused_backend(monkeypatch):
     torch.cuda.synchronize()
 
     _assert_dpa_backend("fused")
-    torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0)
+    _assert_matches_eager(out, ref, "fused", dtype)
     for tensor, ref_grad in zip(grads, ref_grads):
-        torch.testing.assert_close(tensor.grad, ref_grad, rtol=0.0, atol=0.0)
+        _assert_matches_eager(tensor.grad, ref_grad, "fused", dtype)
 
 
 # ---------------------------------------------------------------------------

From bd2c5066fd871f4b3cad37166c6a441250c0c489 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 11:53:47 +0200
Subject: [PATCH 12/67] Merge the per-backend compile tests, and give padding
 masks real padding

The configurations built cu_seqlens with every sequence at the maximum length,
so a padding mask was all-False: the padding branches ran, but on data that
could not tell a difference. They now use varying sequence lengths, which
immediately showed that FlashAttention on bshd/sbhd with a padding mask does
not trace: it packs the tensors first, and get_indices built the index list
with a Python comprehension over sequence lengths held in a GPU tensor -- a
host synchronization per sequence in eager, untraceable while compiling.
It is now built on the device, verified exhaustively against the previous
implementation (10304 cases, 0 mismatches) and free of synchronization.

The separate tests for one shared cu_seqlens tensor and for the fused backend
folded into the main one: the first is a configuration, and the second is the
backend axis, which now runs flash, fused and unfused with the graph break
that FusedAttention's eager island implies.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 306 +++++++-----------
 .../attention/dot_product_attention/utils.py  |  24 +-
 2 files changed, 136 insertions(+), 194 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index b88222be60..397db816f4 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -562,68 +562,61 @@ def fn(q, k, v, extra):
 # get_available_attention_backends() answers that, so configurations only one
 # backend supports (arbitrary masks, biases, MLA head dims, ...) are covered
 # rather than avoided.
+def _cfg(model_config, qkv_format="bshd", packed=None, share_cu_seqlens=False):
+    return dict(
+        model_config=model_config,
+        qkv_format=qkv_format,
+        packed=packed,
+        share_cu_seqlens=share_cu_seqlens,
+    )
+
+
 _DPA_COMPILE_CONFIGS = {
-    # name: (ModelConfig, qkv_format, packed layout or None)
-    "self_bshd_causal": (ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), "bshd", None),
-    "self_sbhd_no_mask": (ModelConfig(2, 128, 4, 64, attn_mask_type="no_mask"), "sbhd", None),
-    "self_bshd_swa": (
-        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", window_size=(16, 0)),
-        "bshd",
-        None,
-    ),
-    "gqa_bshd_causal": (
-        ModelConfig(2, 128, 8, 64, num_gqa_groups=2, attn_mask_type="causal"),
-        "bshd",
-        None,
+    "self_bshd_causal": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal")),
+    "self_sbhd_no_mask": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="no_mask"), "sbhd"),
+    "self_bshd_swa": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal", window_size=(16, 0))),
+    "gqa_bshd_causal": _cfg(ModelConfig(2, 128, 8, 64, num_gqa_groups=2, attn_mask_type="causal")),
+    "cross_bshd_no_mask": _cfg(
+        ModelConfig(2, 128, 4, 64, max_seqlen_kv=256, attn_mask_type="no_mask")
     ),
-    "cross_bshd_no_mask": (
-        ModelConfig(2, 128, 4, 64, max_seqlen_kv=256, attn_mask_type="no_mask"),
-        "bshd",
-        None,
+    "self_bshd_padding_causal": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal")),
+    "self_thd_padding_causal": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"), "thd"
     ),
-    "self_thd_padding_causal": (
-        ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"),
-        "thd",
-        None,
+    # Self attention naturally passes one cu_seqlens tensor for both q and kv,
+    # which flash-attn hands to two inputs of the same autograd.Function.
+    "self_thd_shared_cu_seqlens": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"), "thd", share_cu_seqlens=True
     ),
-    "packed_qkv_bs3hd": (ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), "bshd", "bs3hd"),
+    "packed_qkv_bs3hd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="bs3hd"),
     # Configurations below are supported by one backend only, or take a code
     # path of their own inside DotProductAttention.
-    "alibi_bshd_causal": (
-        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", attn_bias_type="alibi"),
-        "bshd",
-        None,
+    "alibi_bshd_causal": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", attn_bias_type="alibi")
     ),
-    "post_scale_bias_bshd": (
-        ModelConfig(2, 128, 4, 64, attn_bias_type="post_scale_bias", bias_shape="1hss"),
-        "bshd",
-        None,
+    "post_scale_bias_bshd": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_bias_type="post_scale_bias", bias_shape="1hss")
     ),
-    "arbitrary_mask_bshd": (ModelConfig(2, 128, 4, 64, attn_mask_type="arbitrary"), "bshd", None),
-    "mla_bshd_causal": (
-        ModelConfig(2, 128, 4, 128, head_dim_v=64, attn_mask_type="causal"),
-        "bshd",
-        None,
-    ),
-    "sink_softmax_bshd": (
-        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", softmax_type="off-by-one"),
-        "bshd",
-        None,
+    "arbitrary_mask_bshd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="arbitrary")),
+    "mla_bshd_causal": _cfg(ModelConfig(2, 128, 4, 128, head_dim_v=64, attn_mask_type="causal")),
+    "sink_softmax_bshd": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_mask_type="causal", softmax_type="off-by-one")
     ),
 }
 
 
-def _qkv_layout(qkv_format: str, packed: Optional[str]) -> str:
-    return packed if packed is not None else "_".join([qkv_format] * 3)
+def _qkv_layout(spec: dict) -> str:
+    return spec["packed"] or "_".join([spec["qkv_format"]] * 3)
 
 
-def _make_dpa(config: ModelConfig, qkv_format: str, dtype: torch.dtype) -> te.DotProductAttention:
+def _make_dpa(spec: dict, dtype: torch.dtype) -> te.DotProductAttention:
+    config = spec["model_config"]
     return te.DotProductAttention(
         num_attention_heads=config.num_heads,
         kv_channels=config.kv_channels,
         num_gqa_groups=config.num_gqa_groups,
         attention_dropout=config.dropout_p,
-        qkv_format=qkv_format,
+        qkv_format=spec["qkv_format"],
         attn_mask_type=config.attn_mask_type,
         window_size=config.window_size,
         attention_type=config.attn_type,
@@ -631,25 +624,51 @@ def _make_dpa(config: ModelConfig, qkv_format: str, dtype: torch.dtype) -> te.Do
     ).to(dtype=dtype, device="cuda")
 
 
-def _make_dpa_inputs(config: ModelConfig, qkv_format: str, packed: Optional[str], dtype):
+def _cu_seqlens(seqlens: torch.Tensor) -> torch.Tensor:
+    cu = torch.zeros(seqlens.numel() + 1, dtype=torch.int32, device="cuda")
+    cu[1:] = torch.cumsum(seqlens, dim=0)
+    return cu
+
+
+def _make_dpa_inputs(spec: dict, dtype: torch.dtype):
     """Build the (args, kwargs) that `DotProductAttention.forward` is called
     with, plus the list of tensors whose gradients the test compares."""
+    config = spec["model_config"]
+    qkv_format, packed = spec["qkv_format"], spec["packed"]
     b = config.batch_size
     s_q, s_kv = config.max_seqlen_q, config.max_seqlen_kv
     h, g = config.num_heads, config.num_gqa_groups
     d_qk, d_v = config.head_dim_qk, config.head_dim_v
+    padded = "padding" in config.attn_mask_type
 
-    def _shape(s, heads, head_dim):
+    kwargs = {}
+    if padded:
+        # Sequences shorter than the maximum, so the padding mask is not
+        # degenerate: with all sequences full it would be all-False, and any
+        # difference in how masking is compiled would be invisible.
+        seqlens_q = torch.randint(1, s_q, [b], dtype=torch.int32, device="cuda")
+        seqlens_kv = (
+            seqlens_q
+            if config.attn_type == "self"
+            else (torch.randint(1, s_kv, [b], dtype=torch.int32, device="cuda"))
+        )
+        cu_q = _cu_seqlens(seqlens_q)
+        cu_kv = cu_q if spec["share_cu_seqlens"] else _cu_seqlens(seqlens_kv)
+        kwargs.update(cu_seqlens_q=cu_q, cu_seqlens_kv=cu_kv, max_seqlen_q=s_q, max_seqlen_kv=s_kv)
+        t_q, t_kv = int(cu_q[-1]), int(cu_kv[-1])
+    else:
+        t_q, t_kv = b * s_q, b * s_kv
+
+    def _shape(s, t, heads, head_dim):
         return {
             "bshd": (b, s, heads, head_dim),
             "sbhd": (s, b, heads, head_dim),
-            "thd": (b * s, heads, head_dim),
+            "thd": (t, heads, head_dim),
         }[qkv_format]
 
     def _randn(shape):
         return torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
 
-    kwargs = {}
     if packed is not None:
         # Declarative packed QKV: q/k/v are derived from one buffer by DPA
         # itself, and the layout comes from the declaration -- the only packed
@@ -658,27 +677,12 @@ def _randn(shape):
         qkv = _randn((b, s_q, 3, h, d_qk))
         kwargs["qkv_layer"] = qkv
         kwargs["qkv_interleave_dim"] = -3
-        grad_tensors = [qkv]
-        args = ()
+        grad_tensors, args = [qkv], ()
     else:
-        q = _randn(_shape(s_q, h, d_qk))
-        k = _randn(_shape(s_kv, g, d_qk))
-        v = _randn(_shape(s_kv, g, d_v))
-        grad_tensors = [q, k, v]
-        args = (q, k, v)
-
-    if qkv_format == "thd":
-        # All sequences have the maximum length, i.e. no padding. max_seqlen_*
-        # is passed explicitly: deriving it from cu_seqlens costs a `.item()`
-        # (a device sync, and an unbacked SymInt under torch.compile).
-        kwargs["cu_seqlens_q"] = torch.arange(
-            0, (b + 1) * s_q, step=s_q, dtype=torch.int32, device="cuda"
-        )
-        kwargs["cu_seqlens_kv"] = torch.arange(
-            0, (b + 1) * s_kv, step=s_kv, dtype=torch.int32, device="cuda"
-        )
-        kwargs["max_seqlen_q"] = s_q
-        kwargs["max_seqlen_kv"] = s_kv
+        q = _randn(_shape(s_q, t_q, h, d_qk))
+        k = _randn(_shape(s_kv, t_kv, g, d_qk))
+        v = _randn(_shape(s_kv, t_kv, g, d_v))
+        grad_tensors, args = [q, k, v], (q, k, v)
 
     if config.attn_mask_type == "arbitrary":
         kwargs["attention_mask"] = torch.zeros(b, 1, s_q, s_kv, dtype=torch.bool, device="cuda")
@@ -690,11 +694,17 @@ def _randn(shape):
     return args, kwargs, grad_tensors
 
 
-def _skip_unsupported(config: ModelConfig, qkv_layout: str, backend: str, dtype) -> None:
+def _skip_unsupported(spec: dict, backend: str, dtype) -> None:
     """Skip configurations the backend under test cannot run at all."""
-    available, _, _ = get_available_attention_backends(config, dtype, qkv_layout)
-    flash_supported, _, unfused_supported = available
-    supported = {"flash": flash_supported, "unfused": unfused_supported}[backend]
+    available, _, _ = get_available_attention_backends(
+        spec["model_config"], dtype, _qkv_layout(spec)
+    )
+    flash_supported, fused_supported, unfused_supported = available
+    supported = {
+        "flash": flash_supported,
+        "fused": fused_supported,
+        "unfused": unfused_supported,
+    }[backend]
     if not supported:
         pytest.skip(f"the {backend} backend does not support this configuration")
 
@@ -757,39 +767,22 @@ def _assert_matches_eager(actual, expected, backend: str, dtype: torch.dtype) ->
     )
 
 
-@pytest.mark.parametrize("backend", ["flash", "unfused"])
-@pytest.mark.parametrize("config", _DPA_COMPILE_CONFIGS.keys())
-def test_dpa_torch_compile(monkeypatch, backend, config):
-    """`DotProductAttention` must trace under `torch.compile(fullgraph=True)`
-    on the FlashAttention and UnfusedDotProductAttention backends, and match
-    eager in forward and backward.
-
-    `fullgraph=True` makes the test fail on any graph break, i.e. it covers the
-    whole module: input unpacking, qkv layout, backend selection and the
-    backend itself. The FusedAttention backend is not covered -- it is an eager
-    island (`@no_torch_dynamo` on `FusedAttention.forward`)."""
-    dtype = torch.bfloat16
-    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS[config]
-    _skip_unsupported(model_config, _qkv_layout(qkv_format, packed), backend, dtype)
-    _force_dpa_backend(monkeypatch, backend)
-
-    module = _make_dpa(model_config, qkv_format, dtype)
-    args, kwargs, grads = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
-
+def _compare_compiled_to_eager(
+    module, args, kwargs, grads, monkeypatch, backend: str, dtype: torch.dtype, **compile_kwargs
+) -> None:
+    """Run the module eagerly and compiled on the same inputs, and compare both
+    the output and every input gradient."""
     ref = module(*args, **kwargs)
     ref.sum().backward()
-    _assert_dpa_backend(backend)
     ref_grads = [t.grad.clone() for t in grads]
-    for t in grads:
-        t.grad = None
+    for tensor in grads:
+        tensor.grad = None
 
     torch._dynamo.reset()
     # Force backend selection to be re-run (and traced) inside the compiled
     # region instead of being served from the cache the eager call populated.
     _force_dpa_backend(monkeypatch, backend)
-    compiled = torch.compile(module, fullgraph=True)
-
-    out = compiled(*args, **kwargs)
+    out = torch.compile(module, **compile_kwargs)(*args, **kwargs)
     out.sum().backward()
     torch.cuda.synchronize()
 
@@ -800,15 +793,39 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
         _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
 
 
+@pytest.mark.parametrize("backend", ["flash", "fused", "unfused"])
+@pytest.mark.parametrize("config", _DPA_COMPILE_CONFIGS.keys())
+def test_dpa_torch_compile(monkeypatch, backend, config):
+    """`DotProductAttention` under `torch.compile` must match eager in forward
+    and backward, for every backend that supports the configuration.
+
+    FlashAttention and UnfusedDotProductAttention are compiled with
+    `fullgraph=True`, so the test fails on any graph break -- covering the whole
+    module: input unpacking, qkv layout, backend selection and the backend
+    itself. FusedAttention is an eager island (`@no_torch_dynamo` on its
+    forward), so it is expected to graph-break and only has to stay correct.
+    """
+    dtype = torch.bfloat16
+    spec = _DPA_COMPILE_CONFIGS[config]
+    _skip_unsupported(spec, backend, dtype)
+    _force_dpa_backend(monkeypatch, backend)
+
+    module = _make_dpa(spec, dtype)
+    args, kwargs, grads = _make_dpa_inputs(spec, dtype)
+    _compare_compiled_to_eager(
+        module, args, kwargs, grads, monkeypatch, backend, dtype, fullgraph=backend != "fused"
+    )
+
+
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
 def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     """`mode="reduce-overhead"`: forward and backward of DotProductAttention
     are captured into CUDA graphs and replayed on subsequent iterations."""
     dtype = torch.bfloat16
-    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+    spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _force_dpa_backend(monkeypatch, backend)
 
-    module = _make_dpa(model_config, qkv_format, dtype)
+    module = _make_dpa(spec, dtype)
 
     torch._dynamo.reset()
     compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead")
@@ -817,7 +834,7 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
         # Fresh inputs every iteration: a replay that reuses stale buffers would
         # still produce finite values and non-None gradients, so only comparing
         # against eager catches it.
-        args, kwargs, grads = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
+        args, kwargs, grads = _make_dpa_inputs(spec, dtype)
         # CUDA graphs hand back tensors owned by their memory pool, which the
         # next replay overwrites.
         out = compiled(*args, **kwargs).clone()
@@ -842,11 +859,11 @@ def test_dpa_torch_compile_fp8(monkeypatch, fp8_attention):
     eager. FP8 elsewhere in the model with attention in high precision -- the
     common training setup -- must stay compiled."""
     dtype = torch.bfloat16
-    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+    spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
     _force_dpa_backend(monkeypatch, "unfused")
 
-    module = _make_dpa(model_config, qkv_format, dtype)
-    args, kwargs, _ = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
+    module = _make_dpa(spec, dtype)
+    args, kwargs, _ = _make_dpa_inputs(spec, dtype)
     fp8_recipe = recipe.DelayedScaling(fp8_dpa=fp8_attention)
 
     def fn(*args, **kwargs):
@@ -868,48 +885,6 @@ def fn(*args, **kwargs):
     assert fell_back == fp8_attention
 
 
-@pytest.mark.parametrize("backend", ["flash", "unfused"])
-def test_dpa_torch_compile_shared_cu_seqlens(monkeypatch, backend):
-    """Self attention on `thd` naturally passes one cu_seqlens tensor for both q
-    and kv, which flash-attn hands to two inputs of the same autograd.Function --
-    something dynamo cannot trace. The other tests build two tensors, so this
-    covers the shape a caller is most likely to write."""
-    dtype = torch.bfloat16
-    model_config, qkv_format, _ = _DPA_COMPILE_CONFIGS["self_thd_padding_causal"]
-    _force_dpa_backend(monkeypatch, backend)
-
-    module = _make_dpa(model_config, qkv_format, dtype)
-    b, s = model_config.batch_size, model_config.max_seqlen_q
-    h, d = model_config.num_heads, model_config.head_dim_qk
-    q, k, v = [
-        torch.randn(b * s, h, d, dtype=dtype, device="cuda", requires_grad=True) for _ in range(3)
-    ]
-    cu_seqlens = torch.arange(0, (b + 1) * s, step=s, dtype=torch.int32, device="cuda")
-    kwargs = dict(
-        cu_seqlens_q=cu_seqlens,
-        cu_seqlens_kv=cu_seqlens,  # the same tensor, not a copy
-        max_seqlen_q=s,
-        max_seqlen_kv=s,
-    )
-
-    ref = module(q, k, v, **kwargs)
-    ref.sum().backward()
-    ref_grads = [t.grad.clone() for t in (q, k, v)]
-    for t in (q, k, v):
-        t.grad = None
-
-    torch._dynamo.reset()
-    _force_dpa_backend(monkeypatch, backend)
-    out = torch.compile(module, fullgraph=True)(q, k, v, **kwargs)
-    out.sum().backward()
-    torch.cuda.synchronize()
-
-    _assert_dpa_backend(backend)
-    _assert_matches_eager(out, ref, backend, dtype)
-    for tensor, ref_grad in zip((q, k, v), ref_grads):
-        _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
-
-
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
 @pytest.mark.parametrize("interleave_dim", [-3, -2])
 def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interleave_dim):
@@ -919,12 +894,13 @@ def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interlea
     the packing via qkv_layer/kv_layer keeps it on the compiled path, which
     test_dpa_torch_compile covers."""
     dtype = torch.bfloat16
-    model_config, qkv_format, _ = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+    spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+    config = spec["model_config"]
     _force_dpa_backend(monkeypatch, backend)
 
-    module = _make_dpa(model_config, qkv_format, dtype)
-    b, s = model_config.batch_size, model_config.max_seqlen_q
-    h, d = model_config.num_heads, model_config.head_dim_qk
+    module = _make_dpa(spec, dtype)
+    b, s = config.batch_size, config.max_seqlen_q
+    h, d = config.num_heads, config.head_dim_qk
     shape = (b, s, 3, h, d) if interleave_dim == -3 else (b, s, h, 3, d)
     qkv = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
     q, k, v = [qkv.select(interleave_dim, i) for i in range(3)]
@@ -951,40 +927,6 @@ def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interlea
         torch.compile(module, fullgraph=True)(q, k, v)
 
 
-def test_dpa_torch_compile_fused_backend(monkeypatch):
-    """The FusedAttention backend is an eager island, so it graph-breaks rather
-    than compiles; DotProductAttention must still be correct around it."""
-    dtype = torch.bfloat16
-    model_config, qkv_format, packed = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
-    _force_dpa_backend(monkeypatch, "fused")
-
-    module = _make_dpa(model_config, qkv_format, dtype)
-    args, kwargs, grads = _make_dpa_inputs(model_config, qkv_format, packed, dtype)
-
-    ref = module(*args, **kwargs)
-    from transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention import (
-        _attention_backends,
-    )
-
-    if not _attention_backends["use_fused_attention"]:
-        pytest.skip("FusedAttention does not support this config on this device")
-    ref.sum().backward()
-    ref_grads = [t.grad.clone() for t in grads]
-    for t in grads:
-        t.grad = None
-
-    torch._dynamo.reset()
-    _force_dpa_backend(monkeypatch, "fused")
-    out = torch.compile(module)(*args, **kwargs)
-    out.sum().backward()
-    torch.cuda.synchronize()
-
-    _assert_dpa_backend("fused")
-    _assert_matches_eager(out, ref, "fused", dtype)
-    for tensor, ref_grad in zip(grads, ref_grads):
-        _assert_matches_eager(tensor.grad, ref_grad, "fused", dtype)
-
-
 # ---------------------------------------------------------------------------
 # get_attention_backend under torch.compile
 # ---------------------------------------------------------------------------
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 95b5a7f15f..383afebd69 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -2005,21 +2005,21 @@ def get_indices(max_seqlen: int, cu_seqlens: torch.Tensor) -> torch.Tensor:
     tensor of shape [batch_size * max_seqlen, 1, 1] containing the indices for
     the valid tokens in a batch.
     """
+    # Built with device-side ops only: reading the sequence lengths on the host
+    # would synchronize the device once per sequence.
     bs = len(cu_seqlens) - 1
     seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
-    indices = [i * max_seqlen + ii for i, j in enumerate(seqlens) for ii in range(j)]
-    indices = torch.Tensor(indices).unsqueeze(1).unsqueeze(1).to(dtype=torch.int64, device="cuda")
-
-    num_nonzeros = indices.shape[0]
-    pad_amount = bs * max_seqlen - num_nonzeros
-    indices = F.pad(
-        input=indices,
-        pad=(0, 0, 0, 0, 0, pad_amount),
-        mode="constant",
-        value=float(bs * max_seqlen),
+    positions = torch.arange(max_seqlen, device=cu_seqlens.device)
+    valid = (positions.unsqueeze(0) < seqlens.unsqueeze(1)).flatten()
+    # Sorting the mask is stable, so the valid positions come first in order,
+    # and the invalid tail is replaced by the out-of-range padding index.
+    ordered = torch.argsort(~valid, stable=True)
+    indices = torch.where(
+        torch.arange(bs * max_seqlen, device=cu_seqlens.device) < valid.sum(),
+        ordered,
+        bs * max_seqlen,
     )
-
-    return indices
+    return indices.to(dtype=torch.int64, device="cuda").unsqueeze(1).unsqueeze(1)
 
 
 def get_full_cu_seqlens(

From 725931ec24db7425e0f008444b65b89571692a07 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 12:11:58 +0200
Subject: [PATCH 13/67] Run eagerly when max_seqlen has to be derived from
 cu_seqlens

On thd inputs without max_seqlen, DotProductAttention reads the sequence
lengths off cu_seqlens to derive it. Unlike the padding mask and the packing
indices, this one cannot be moved to the device: max_seqlen sizes tensors
downstream, so it has to be a Python int, and reading it is a device
synchronization -- a data-dependent value dynamo cannot trace. Route those
calls to the eager island, so passing max_seqlen is what keeps a call
compiled.

Arguments are now looked up by name or position, so the predicate sees them
whichever way the caller passed them.

The two eager-fallback tests merge into one: they set up different inputs but
assert the same thing -- a warning, results matching eager, and an error under
fullgraph.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 70 +++++++++++++------
 .../dot_product_attention.py                  | 44 ++++++++++--
 2 files changed, 88 insertions(+), 26 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 397db816f4..ea21cdd2ed 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -885,46 +885,76 @@ def fn(*args, **kwargs):
     assert fell_back == fp8_attention
 
 
-@pytest.mark.parametrize("backend", ["flash", "unfused"])
-@pytest.mark.parametrize("interleave_dim", [-3, -2])
-def test_dpa_torch_compile_packed_views_fall_back(monkeypatch, backend, interleave_dim):
-    """Packed q/k/v passed as plain views cannot be recognized while tracing (the
-    detection reads data pointers), so DotProductAttention runs that detection
-    eagerly instead -- correct results, at the cost of a graph break. Declaring
-    the packing via qkv_layer/kv_layer keeps it on the compiled path, which
-    test_dpa_torch_compile covers."""
-    dtype = torch.bfloat16
-    spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+def _packed_views_inputs(spec, dtype, interleave_dim=-3):
+    """Packed q/k/v handed over as plain views, without declaring the packing."""
     config = spec["model_config"]
-    _force_dpa_backend(monkeypatch, backend)
-
-    module = _make_dpa(spec, dtype)
     b, s = config.batch_size, config.max_seqlen_q
     h, d = config.num_heads, config.head_dim_qk
     shape = (b, s, 3, h, d) if interleave_dim == -3 else (b, s, h, 3, d)
     qkv = torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
     q, k, v = [qkv.select(interleave_dim, i) for i in range(3)]
+    return (q, k, v), {}, [qkv]
+
+
+def _thd_without_max_seqlen_inputs(spec, dtype):
+    """thd inputs that leave max_seqlen to be derived from cu_seqlens."""
+    args, kwargs, grads = _make_dpa_inputs(spec, dtype)
+    del kwargs["max_seqlen_q"], kwargs["max_seqlen_kv"]
+    return args, kwargs, grads
+
+
+_EAGER_FALLBACK_CASES = {
+    # name: (config name, input builder)
+    "packed_views_bs3hd": (
+        "self_bshd_causal",
+        lambda spec, dtype: _packed_views_inputs(spec, dtype, -3),
+    ),
+    "packed_views_bsh3d": (
+        "self_bshd_causal",
+        lambda spec, dtype: _packed_views_inputs(spec, dtype, -2),
+    ),
+    "thd_without_max_seqlen": ("self_thd_padding_causal", _thd_without_max_seqlen_inputs),
+}
 
-    ref = module(q, k, v)
+
+@pytest.mark.parametrize("backend", ["flash", "unfused"])
+@pytest.mark.parametrize("case", _EAGER_FALLBACK_CASES.keys())
+def test_dpa_torch_compile_eager_fallback(monkeypatch, backend, case):
+    """Calls that cannot be traced run as an eager island instead, with a
+    warning: recognizing packed q/k/v takes the data pointers dynamo cannot
+    read, and deriving max_seqlen off cu_seqlens is a device synchronization.
+    Both keep working, and both are avoidable -- by declaring the packing via
+    qkv_layer/kv_layer, or by passing max_seqlen."""
+    dtype = torch.bfloat16
+    config_name, make_inputs = _EAGER_FALLBACK_CASES[case]
+    spec = _DPA_COMPILE_CONFIGS[config_name]
+    _force_dpa_backend(monkeypatch, backend)
+
+    module = _make_dpa(spec, dtype)
+    args, kwargs, grads = make_inputs(spec, dtype)
+
+    ref = module(*args, **kwargs)
     ref.sum().backward()
-    ref_grad = qkv.grad.clone()
-    qkv.grad = None
+    ref_grads = [t.grad.clone() for t in grads]
+    for tensor in grads:
+        tensor.grad = None
 
     torch._dynamo.reset()
     _force_dpa_backend(monkeypatch, backend)
     with pytest.warns(UserWarning, match="Falling back to eager execution"):
-        out = torch.compile(module)(q, k, v)
+        out = torch.compile(module)(*args, **kwargs)
     out.sum().backward()
     torch.cuda.synchronize()
 
     _assert_matches_eager(out, ref, backend, dtype)
-    _assert_matches_eager(qkv.grad, ref_grad, backend, dtype)
+    for tensor, ref_grad in zip(grads, ref_grads):
+        _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
 
-    # The same graph break is an error when the user asked for a full graph.
+    # The same eager island is an error when the user asked for a full graph.
     torch._dynamo.reset()
     _force_dpa_backend(monkeypatch, backend)
     with pytest.raises(Exception, match="torch.compiler.disable"):
-        torch.compile(module, fullgraph=True)(q, k, v)
+        torch.compile(module, fullgraph=True)(*args, **kwargs)
 
 
 # ---------------------------------------------------------------------------
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index f82aa91c13..4e0697fcd7 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -197,12 +197,32 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim
     return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1)
 
 
+# Positions of the `DotProductAttention.forward` arguments the predicate below
+# reads, so that they are found whether they were passed by name or by position.
+_FORWARD_ARG_POSITIONS = {
+    "query_layer": 1,
+    "key_layer": 2,
+    "value_layer": 3,
+    "qkv_format": 5,
+    "max_seqlen_q": 10,
+    "max_seqlen_kv": 11,
+    "qkv_layer": 28,
+    "kv_layer": 29,
+}
+
+
+def _forward_arg(name: str, args: tuple, kwargs: dict) -> Any:
+    if name in kwargs:
+        return kwargs[name]
+    position = _FORWARD_ARG_POSITIONS[name]
+    return args[position] if len(args) > position else None
+
+
 def _needs_eager_dpa(*args, **kwargs) -> Optional[str]:
     """Why this DotProductAttention call has to run outside the graph, or None.
 
-    Takes `DotProductAttention.forward`'s arguments as they were passed, so q/k/v
-    are read from either side: `args[0]` is `self`, followed by the positional
-    query/key/value.
+    Takes `DotProductAttention.forward`'s arguments as they were passed, with
+    `args[0]` being `self`.
     """
     # FP8 GEMMs with the attention itself in high precision -- the common
     # training setup -- stay on the compiled path; only FP8 attention bails out.
@@ -212,9 +232,21 @@ def _needs_eager_dpa(*args, **kwargs) -> Optional[str]:
         if fp8_recipe.fp8_dpa or fp8_recipe.fp8_mha:
             return "FP8 attention"
 
-    if kwargs.get("qkv_layer") is None and kwargs.get("kv_layer") is None:
-        qkv = list(args[1:4]) + [
-            kwargs.get(name) for name in ("query_layer", "key_layer", "value_layer")
+    qkv_format = _forward_arg("qkv_format", args, kwargs) or args[0].qkv_format
+    if qkv_format == "thd" and (
+        _forward_arg("max_seqlen_q", args, kwargs) is None
+        or _forward_arg("max_seqlen_kv", args, kwargs) is None
+    ):
+        # Deriving it reads the sequence lengths off cu_seqlens, which is a
+        # device synchronization and a data-dependent value while tracing.
+        return "deriving max_seqlen from cu_seqlens"
+
+    if (
+        _forward_arg("qkv_layer", args, kwargs) is None
+        and _forward_arg("kv_layer", args, kwargs) is None
+    ):
+        qkv = [
+            _forward_arg(name, args, kwargs) for name in ("query_layer", "key_layer", "value_layer")
         ]
         if dpa_utils.qkv_layout_needs_detection(*qkv):
             return "detecting packed q/k/v that were not declared via qkv_layer/kv_layer"

From b2ab7ebad5034690f51ee8d6064ff0f3fbfa20fb Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 13:54:50 +0200
Subject: [PATCH 14/67] Share the run-and-compare code between the compile
 tests

Both tests ran a callable, went backward, copied the gradients out and cleared
them. That is now _run_and_capture(), with _assert_run_matches() comparing two
of its results, which halves the CUDA-graphs test.

The two stay separate tests: what CUDA graphs risk -- a replay reading stale
buffers -- does not depend on the attention configuration, so folding the mode
into the configuration matrix would multiply the cases without covering
anything new, and would report a replay bug as a failure of some unrelated
configuration.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 59 ++++++++++++++++-------------
 1 file changed, 33 insertions(+), 26 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index ea21cdd2ed..4be4ec9fea 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -767,30 +767,48 @@ def _assert_matches_eager(actual, expected, backend: str, dtype: torch.dtype) ->
     )
 
 
+def _run_and_capture(fn, args, kwargs, grads):
+    """Run `fn` and backward through its output, returning the output and a copy
+    of every input gradient.
+
+    The output is cloned because CUDA graphs hand back tensors owned by their
+    memory pool, which the next replay overwrites, and the gradients are cleared
+    so that the same inputs can be reused for the other run.
+    """
+    out = fn(*args, **kwargs).clone()
+    out.sum().backward()
+    torch.cuda.synchronize()
+    captured = []
+    for tensor in grads:
+        assert tensor.grad is not None
+        captured.append(tensor.grad.clone())
+        tensor.grad = None
+    return out, captured
+
+
+def _assert_run_matches(actual, expected, backend: str, dtype: torch.dtype) -> None:
+    """Compare an (output, gradients) pair produced by `_run_and_capture`."""
+    (out, out_grads), (ref, ref_grads) = actual, expected
+    _assert_matches_eager(out, ref, backend, dtype)
+    for out_grad, ref_grad in zip(out_grads, ref_grads):
+        _assert_matches_eager(out_grad, ref_grad, backend, dtype)
+
+
 def _compare_compiled_to_eager(
     module, args, kwargs, grads, monkeypatch, backend: str, dtype: torch.dtype, **compile_kwargs
 ) -> None:
     """Run the module eagerly and compiled on the same inputs, and compare both
     the output and every input gradient."""
-    ref = module(*args, **kwargs)
-    ref.sum().backward()
-    ref_grads = [t.grad.clone() for t in grads]
-    for tensor in grads:
-        tensor.grad = None
+    eager = _run_and_capture(module, args, kwargs, grads)
 
     torch._dynamo.reset()
     # Force backend selection to be re-run (and traced) inside the compiled
     # region instead of being served from the cache the eager call populated.
     _force_dpa_backend(monkeypatch, backend)
-    out = torch.compile(module, **compile_kwargs)(*args, **kwargs)
-    out.sum().backward()
-    torch.cuda.synchronize()
+    compiled = _run_and_capture(torch.compile(module, **compile_kwargs), args, kwargs, grads)
 
     _assert_dpa_backend(backend)
-    _assert_matches_eager(out, ref, backend, dtype)
-    for tensor, ref_grad in zip(grads, ref_grads):
-        assert tensor.grad is not None
-        _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
+    _assert_run_matches(compiled, eager, backend, dtype)
 
 
 @pytest.mark.parametrize("backend", ["flash", "fused", "unfused"])
@@ -835,20 +853,9 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
         # still produce finite values and non-None gradients, so only comparing
         # against eager catches it.
         args, kwargs, grads = _make_dpa_inputs(spec, dtype)
-        # CUDA graphs hand back tensors owned by their memory pool, which the
-        # next replay overwrites.
-        out = compiled(*args, **kwargs).clone()
-        out.sum().backward()
-        torch.cuda.synchronize()
-        compiled_grads = [t.grad.clone() for t in grads]
-        for tensor in grads:
-            tensor.grad = None
-
-        ref = module(*args, **kwargs)
-        ref.sum().backward()
-        _assert_matches_eager(out, ref, backend, dtype)
-        for compiled_grad, tensor in zip(compiled_grads, grads):
-            _assert_matches_eager(compiled_grad, tensor.grad, backend, dtype)
+        replayed = _run_and_capture(compiled, args, kwargs, grads)
+        eager = _run_and_capture(module, args, kwargs, grads)
+        _assert_run_matches(replayed, eager, backend, dtype)
     _assert_dpa_backend(backend)
 
 

From 50d19848ba97743cfe06ed619984792fbe9b3ecb Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 13:57:46 +0200
Subject: [PATCH 15/67] Use the shared run-and-compare code in the remaining
 compile tests

The eager-fallback test hand-rolled the same run, backward and gradient copy.

The backend-level unfused test only checked that the output was finite and
that gradients existed, which a CUDA-graph replay reading stale buffers also
satisfies -- the weakness the DotProductAttention CUDA-graphs test was given
an eager comparison for. It now compares against eager as well.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 28 +++++++++-------------------
 1 file changed, 9 insertions(+), 19 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 4be4ec9fea..8e3617af92 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -542,14 +542,13 @@ def fn(q, k, v, extra):
     compiled = torch.compile(fn, fullgraph=True, mode="reduce-overhead")
 
     for _ in range(3):
+        # Compared against eager rather than only checked for finiteness: a
+        # replay reusing stale buffers would pass the latter.
         q, k, v, extra, _ = _make_unfused_qkv(qkv_layout, dtype, requires_grad=True)
-        out = compiled(q, k, v, extra)
-        out.sum().backward()
-        torch.cuda.synchronize()
-        assert torch.isfinite(out).all()
-        assert q.grad is not None
-        assert k.grad is not None
-        assert v.grad is not None
+        inputs = (q, k, v, extra)
+        replayed = _run_and_capture(compiled, inputs, {}, [q, k, v])
+        eager = _run_and_capture(fn, inputs, {}, [q, k, v])
+        _assert_run_matches(replayed, eager, "unfused", dtype)
 
 
 # ---------------------------------------------------------------------------
@@ -940,22 +939,13 @@ def test_dpa_torch_compile_eager_fallback(monkeypatch, backend, case):
     module = _make_dpa(spec, dtype)
     args, kwargs, grads = make_inputs(spec, dtype)
 
-    ref = module(*args, **kwargs)
-    ref.sum().backward()
-    ref_grads = [t.grad.clone() for t in grads]
-    for tensor in grads:
-        tensor.grad = None
+    eager = _run_and_capture(module, args, kwargs, grads)
 
     torch._dynamo.reset()
     _force_dpa_backend(monkeypatch, backend)
     with pytest.warns(UserWarning, match="Falling back to eager execution"):
-        out = torch.compile(module)(*args, **kwargs)
-    out.sum().backward()
-    torch.cuda.synchronize()
-
-    _assert_matches_eager(out, ref, backend, dtype)
-    for tensor, ref_grad in zip(grads, ref_grads):
-        _assert_matches_eager(tensor.grad, ref_grad, backend, dtype)
+        fell_back = _run_and_capture(torch.compile(module), args, kwargs, grads)
+    _assert_run_matches(fell_back, eager, backend, dtype)
 
     # The same eager island is an error when the user asked for a full graph.
     torch._dynamo.reset()

From f19c52d6e58f0c4eeca4a557df39c91a718f64fc Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 14:05:10 +0200
Subject: [PATCH 16/67] Read the predicate's arguments by name, not by
 hardcoded position

The positions of forward's parameters were written out as a table of indices,
which anyone inserting a parameter into a thirty-argument signature would
silently invalidate: the predicate would go on reading whatever now sits at
index 5 or 10 and decide the fallback on it, with nothing to notice.

The decorator already has the function it wraps, so it takes the parameter
names from its signature and hands the predicate the call's arguments keyed by
name.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention.py                  | 57 +++++++------------
 1 file changed, 20 insertions(+), 37 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 4e0697fcd7..4d7c66d22d 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -5,6 +5,7 @@
 """Attention."""
 from contextlib import nullcontext
 from functools import wraps
+import inspect
 import math
 import os
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -197,32 +198,11 @@ def _trim_output(attn_out, num_attention_heads, padded_head_dim_v, orig_head_dim
     return attn_out[..., :orig_head_dim_v].reshape(*out_shape, -1)
 
 
-# Positions of the `DotProductAttention.forward` arguments the predicate below
-# reads, so that they are found whether they were passed by name or by position.
-_FORWARD_ARG_POSITIONS = {
-    "query_layer": 1,
-    "key_layer": 2,
-    "value_layer": 3,
-    "qkv_format": 5,
-    "max_seqlen_q": 10,
-    "max_seqlen_kv": 11,
-    "qkv_layer": 28,
-    "kv_layer": 29,
-}
-
-
-def _forward_arg(name: str, args: tuple, kwargs: dict) -> Any:
-    if name in kwargs:
-        return kwargs[name]
-    position = _FORWARD_ARG_POSITIONS[name]
-    return args[position] if len(args) > position else None
-
-
-def _needs_eager_dpa(*args, **kwargs) -> Optional[str]:
+def _needs_eager_dpa(call: Dict[str, Any]) -> Optional[str]:
     """Why this DotProductAttention call has to run outside the graph, or None.
 
-    Takes `DotProductAttention.forward`'s arguments as they were passed, with
-    `args[0]` being `self`.
+    `call` maps `DotProductAttention.forward`'s parameter names to the arguments
+    this call passed, including `self`.
     """
     # FP8 GEMMs with the attention itself in high precision -- the common
     # training setup -- stay on the compiled path; only FP8 attention bails out.
@@ -232,32 +212,33 @@ def _needs_eager_dpa(*args, **kwargs) -> Optional[str]:
         if fp8_recipe.fp8_dpa or fp8_recipe.fp8_mha:
             return "FP8 attention"
 
-    qkv_format = _forward_arg("qkv_format", args, kwargs) or args[0].qkv_format
+    qkv_format = call.get("qkv_format") or call["self"].qkv_format
     if qkv_format == "thd" and (
-        _forward_arg("max_seqlen_q", args, kwargs) is None
-        or _forward_arg("max_seqlen_kv", args, kwargs) is None
+        call.get("max_seqlen_q") is None or call.get("max_seqlen_kv") is None
     ):
         # Deriving it reads the sequence lengths off cu_seqlens, which is a
         # device synchronization and a data-dependent value while tracing.
         return "deriving max_seqlen from cu_seqlens"
 
-    if (
-        _forward_arg("qkv_layer", args, kwargs) is None
-        and _forward_arg("kv_layer", args, kwargs) is None
-    ):
-        qkv = [
-            _forward_arg(name, args, kwargs) for name in ("query_layer", "key_layer", "value_layer")
-        ]
+    if call.get("qkv_layer") is None and call.get("kv_layer") is None:
+        qkv = [call.get(name) for name in ("query_layer", "key_layer", "value_layer")]
         if dpa_utils.qkv_layout_needs_detection(*qkv):
             return "detecting packed q/k/v that were not declared via qkv_layer/kv_layer"
     return None
 
 
-def _eager_under_compile_if(needs_eager: Callable[..., Optional[str]]):
+def _eager_under_compile_if(needs_eager: Callable[[Dict[str, Any]], Optional[str]]):
     """Decorator running the wrapped method eagerly whenever `needs_eager` gives
-    a reason why the call is unsupported on the compiled path."""
+    a reason why the call is unsupported on the compiled path.
+
+    `needs_eager` is handed the call's arguments keyed by parameter name, taken
+    from the wrapped signature, so it does not have to care whether the caller
+    passed them by name or by position.
+    """
 
     def decorator(fn):
+        parameter_names = list(inspect.signature(fn).parameters)
+
         # The warning belongs inside the dynamo-disabled function: warnings.warn
         # graph-breaks on its own, masking the break that matters.
         @no_torch_dynamo()
@@ -272,7 +253,9 @@ def eager_fn(reason, *args, **kwargs):
         @wraps(fn)
         def wrapper(*args, **kwargs):
             if torch.compiler.is_compiling():
-                reason = needs_eager(*args, **kwargs)
+                call = dict(zip(parameter_names, args))
+                call.update(kwargs)
+                reason = needs_eager(call)
                 if reason is not None:
                     return eager_fn(reason, *args, **kwargs)
             return fn(*args, **kwargs)

From fbc7146e78d465c118ed7312d58983460cf7559f Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 14:12:38 +0200
Subject: [PATCH 17/67] Move the eager-fallback decorator to jit.py

Nothing about it is specific to attention: it takes a predicate, and runs the
wrapped method as an eager island whenever that predicate gives a reason. It
belongs next to no_torch_dynamo, which it builds on, rather than buried in the
attention module where the modules that will need it next -- MultiheadAttention
and TransformerLayer, both of which still have blockers of their own -- would
not find it.

What stays in the attention module is the predicate, which is where the
attention-specific knowledge lives.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention.py                  | 44 +------------------
 transformer_engine/pytorch/jit.py             | 42 +++++++++++++++++-
 2 files changed, 43 insertions(+), 43 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 4d7c66d22d..5f6628c951 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -4,8 +4,6 @@
 
 """Attention."""
 from contextlib import nullcontext
-from functools import wraps
-import inspect
 import math
 import os
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
@@ -44,7 +42,7 @@
     CudaRNGStatesTracker,
     graph_safe_rng_available,
 )
-from transformer_engine.pytorch.jit import no_torch_dynamo
+from transformer_engine.pytorch.jit import eager_under_compile_if
 from transformer_engine.pytorch.graph import is_graph_capturing
 from transformer_engine.pytorch.attention.inference import InferenceParams
 
@@ -227,44 +225,6 @@ def _needs_eager_dpa(call: Dict[str, Any]) -> Optional[str]:
     return None
 
 
-def _eager_under_compile_if(needs_eager: Callable[[Dict[str, Any]], Optional[str]]):
-    """Decorator running the wrapped method eagerly whenever `needs_eager` gives
-    a reason why the call is unsupported on the compiled path.
-
-    `needs_eager` is handed the call's arguments keyed by parameter name, taken
-    from the wrapped signature, so it does not have to care whether the caller
-    passed them by name or by position.
-    """
-
-    def decorator(fn):
-        parameter_names = list(inspect.signature(fn).parameters)
-
-        # The warning belongs inside the dynamo-disabled function: warnings.warn
-        # graph-breaks on its own, masking the break that matters.
-        @no_torch_dynamo()
-        def eager_fn(reason, *args, **kwargs):
-            warnings.warn(
-                f"Falling back to eager execution under torch.compile: {reason} is"
-                " unsupported on the compiled path (graph-breaks under fullgraph=True).",
-                stacklevel=3,
-            )
-            return fn(*args, **kwargs)
-
-        @wraps(fn)
-        def wrapper(*args, **kwargs):
-            if torch.compiler.is_compiling():
-                call = dict(zip(parameter_names, args))
-                call.update(kwargs)
-                reason = needs_eager(call)
-                if reason is not None:
-                    return eager_fn(reason, *args, **kwargs)
-            return fn(*args, **kwargs)
-
-        return wrapper
-
-    return decorator
-
-
 def _unpack_packed_qkv(
     qkv_layer: Optional[torch.Tensor],
     kv_layer: Optional[torch.Tensor],
@@ -1167,7 +1127,7 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
-    @_eager_under_compile_if(_needs_eager_dpa)
+    @eager_under_compile_if(_needs_eager_dpa)
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,
diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py
index 1b93b8254c..2d0ec02d8d 100644
--- a/transformer_engine/pytorch/jit.py
+++ b/transformer_engine/pytorch/jit.py
@@ -3,9 +3,11 @@
 # See LICENSE for license information.
 
 """NVFuser functions and JIT utilities"""
+import inspect
 import os
+import warnings
 from functools import wraps
-from typing import Callable, Optional, Tuple
+from typing import Any, Callable, Dict, Optional, Tuple
 import torch
 
 from .torch_version import torch_version
@@ -77,6 +79,44 @@ def no_torch_dynamo(recursive=True):  # pylint: disable=unused-argument
         return lambda func: func
 
 
+def eager_under_compile_if(needs_eager: Callable[[Dict[str, Any]], Optional[str]]):
+    """Decorator running the wrapped method eagerly whenever `needs_eager` gives
+    a reason why the call is unsupported on the compiled path.
+
+    `needs_eager` is handed the call's arguments keyed by parameter name, taken
+    from the wrapped signature, so it does not have to care whether the caller
+    passed them by name or by position. Returning None keeps the call compiled.
+    """
+
+    def decorator(fn):
+        parameter_names = list(inspect.signature(fn).parameters)
+
+        # The warning belongs inside the dynamo-disabled function: warnings.warn
+        # graph-breaks on its own, masking the break that matters.
+        @no_torch_dynamo()
+        def eager_fn(reason, *args, **kwargs):
+            warnings.warn(
+                f"Falling back to eager execution under torch.compile: {reason} is"
+                " unsupported on the compiled path (graph-breaks under fullgraph=True).",
+                stacklevel=3,
+            )
+            return fn(*args, **kwargs)
+
+        @wraps(fn)
+        def wrapper(*args, **kwargs):
+            if torch.compiler.is_compiling():
+                call = dict(zip(parameter_names, args))
+                call.update(kwargs)
+                reason = needs_eager(call)
+                if reason is not None:
+                    return eager_fn(reason, *args, **kwargs)
+            return fn(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 def set_jit_fusion_options() -> None:
     """Set PyTorch JIT layer fusion options."""
     # flags required to enable jit fusion kernels

From 6b246303b847db73c1a71312915bccec1c3514e0 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 29 Jul 2026 14:40:19 +0200
Subject: [PATCH 18/67] Skip FusedAttention rather than special-casing it in
 the compile tests

The test compiled every backend with fullgraph=True except fused, which was
passed fullgraph=False because its forward is an eager island. That put an
assumption about one backend into the comparison logic.

FusedAttention is skipped instead, next to the configurations a backend cannot
run, with the reason spelled out. Everything else is compiled the same way, and
supporting fused later is deleting the skip.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 24 ++++++++++++++----------
 1 file changed, 14 insertions(+), 10 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 8e3617af92..e32304af1b 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -694,7 +694,12 @@ def _randn(shape):
 
 
 def _skip_unsupported(spec: dict, backend: str, dtype) -> None:
-    """Skip configurations the backend under test cannot run at all."""
+    """Skip what the backend under test cannot run, or cannot compile."""
+    if backend == "fused":
+        # FusedAttention's forward carries @no_torch_dynamo, so there is nothing
+        # to compile: it runs as an eager island. Drop this skip once it traces,
+        # and the tests below cover it as they do the others.
+        pytest.skip("FusedAttention is an eager island and does not compile")
     available, _, _ = get_available_attention_backends(
         spec["model_config"], dtype, _qkv_layout(spec)
     )
@@ -813,14 +818,13 @@ def _compare_compiled_to_eager(
 @pytest.mark.parametrize("backend", ["flash", "fused", "unfused"])
 @pytest.mark.parametrize("config", _DPA_COMPILE_CONFIGS.keys())
 def test_dpa_torch_compile(monkeypatch, backend, config):
-    """`DotProductAttention` under `torch.compile` must match eager in forward
-    and backward, for every backend that supports the configuration.
-
-    FlashAttention and UnfusedDotProductAttention are compiled with
-    `fullgraph=True`, so the test fails on any graph break -- covering the whole
-    module: input unpacking, qkv layout, backend selection and the backend
-    itself. FusedAttention is an eager island (`@no_torch_dynamo` on its
-    forward), so it is expected to graph-break and only has to stay correct.
+    """`DotProductAttention` under `torch.compile(fullgraph=True)` must match
+    eager in forward and backward, for every backend that supports the
+    configuration.
+
+    `fullgraph=True` makes the test fail on any graph break, so it covers the
+    whole module: input unpacking, qkv layout, backend selection and the backend
+    itself.
     """
     dtype = torch.bfloat16
     spec = _DPA_COMPILE_CONFIGS[config]
@@ -830,7 +834,7 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
     module = _make_dpa(spec, dtype)
     args, kwargs, grads = _make_dpa_inputs(spec, dtype)
     _compare_compiled_to_eager(
-        module, args, kwargs, grads, monkeypatch, backend, dtype, fullgraph=backend != "fused"
+        module, args, kwargs, grads, monkeypatch, backend, dtype, fullgraph=True
     )
 
 

From 9a204b8e0a97237a85ee7a99ea85c7091b0fc6e4 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 11:39:18 +0200
Subject: [PATCH 19/67] Cover the compiled path around FusedAttention

FusedAttention is skipped from the backend matrix because it does not compile,
but it is the default on Hopper and Blackwell, and this change does affect it:
what used to be one eager island is now a traced prologue, a graph break at the
backend, and a traced remainder. Whatever crosses that break has to survive it,
which the sub-backend enum did not.

One test covers it, without a fullgraph and outside the matrix. Verified to
fail with the enum restored -- TypeError: int() argument must be ... not
'function', from cuDNN's entry point.

Also rename _distinct_cu_seqlens: it returns its arguments untouched in eager
and whenever they already differ, so it promised more than it does.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 24 ++++++++++++++++---
 .../dot_product_attention/backends.py         |  8 +++----
 .../dot_product_attention.py                  |  3 ---
 3 files changed, 25 insertions(+), 10 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index e32304af1b..d20a0e70dc 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -693,9 +693,10 @@ def _randn(shape):
     return args, kwargs, grad_tensors
 
 
-def _skip_unsupported(spec: dict, backend: str, dtype) -> None:
-    """Skip what the backend under test cannot run, or cannot compile."""
-    if backend == "fused":
+def _skip_unsupported(spec: dict, backend: str, dtype, compiled: bool = True) -> None:
+    """Skip what the backend under test cannot run, or -- for a test that
+    compiles it -- cannot be compiled."""
+    if compiled and backend == "fused":
         # FusedAttention's forward carries @no_torch_dynamo, so there is nothing
         # to compile: it runs as an eager island. Drop this skip once it traces,
         # and the tests below cover it as they do the others.
@@ -838,6 +839,23 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
     )
 
 
+def test_dpa_torch_compile_around_fused(monkeypatch):
+    """FusedAttention itself is an eager island, but everything around it is
+    compiled: DotProductAttention traces up to the backend call, breaks the
+    graph there and resumes afterwards. What crosses that break has to survive
+    it -- the sub-backend enum did not, and reached cuDNN as the function that
+    produced it."""
+    dtype = torch.bfloat16
+    spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+    _skip_unsupported(spec, "fused", dtype, compiled=False)
+    _force_dpa_backend(monkeypatch, "fused")
+
+    module = _make_dpa(spec, dtype)
+    args, kwargs, grads = _make_dpa_inputs(spec, dtype)
+    # No fullgraph: the graph break at the eager island is the point here.
+    _compare_compiled_to_eager(module, args, kwargs, grads, monkeypatch, "fused", dtype)
+
+
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
 def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     """`mode="reduce-overhead"`: forward and backward of DotProductAttention
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 2c47ff8a4b..04bcabe7c4 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -793,8 +793,8 @@ def backward(
         return dq, dk, dv
 
 
-def _distinct_cu_seqlens(cu_q: torch.Tensor, cu_kv: torch.Tensor):
-    """Make q's and kv's cumulative sequence lengths distinct objects.
+def _maybe_unshare_cu_seqlens(cu_q: torch.Tensor, cu_kv: torch.Tensor):
+    """Copy kv's cumulative sequence lengths when they are q's as well.
 
     Self attention passes one tensor for both, and flash-attn forwards them to
     two inputs of the same autograd.Function -- which dynamo cannot trace. In
@@ -1134,7 +1134,7 @@ def forward(
                     else:
                         func = flash_attn_with_kvcache_v3  # pylint: disable=possibly-used-before-assignment
                     if not use_flash_attn_4 and (not use_flash_attn_3 or inference_params is None):
-                        cu_q, cu_kv = _distinct_cu_seqlens(
+                        cu_q, cu_kv = _maybe_unshare_cu_seqlens(
                             cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q,
                             cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv,
                         )
@@ -1150,7 +1150,7 @@ def forward(
                     if inference_params is None:
                         fa_4_optional_forward_kwargs["deterministic"] = self.deterministic
                     if func is flash_attn_varlen_func_v4:
-                        cu_q, cu_kv = _distinct_cu_seqlens(cu_seqlens_q, cu_seqlens_kv)
+                        cu_q, cu_kv = _maybe_unshare_cu_seqlens(cu_seqlens_q, cu_seqlens_kv)
                         fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_q
                         fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_kv
                         fa_4_optional_forward_kwargs["max_seqlen_q"] = max_seqlen_q
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 5f6628c951..34aa6cf403 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -746,9 +746,6 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None:
 
         qstate = FP8GlobalStateManager.quantization_state
         if not (qstate.fp8_enabled or qstate.fp8_calibration or qstate.fp8_parameters):
-            # Without FP8 the recipe juggling below is a no-op, and its
-            # get_fp8_recipe() call would build a default Recipe -- which
-            # torch.compile cannot trace.
             super().init_fp8_metadata(num_gemms=num_gemms)
             return
 

From 2188f183a41bd0cdd6dc662d03169c10a6685c4c Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 11:47:26 +0200
Subject: [PATCH 20/67] Log backend selection through the no-op logger, as
 get_attention_backend does

Skipping the block with is_compiling() introduced a second mechanism for a
problem this module already solves one way: get_attention_backend swaps in a
no-op logger while tracing. The skip was needed while the sub-backend argument
was an enum that did not survive tracing, and is not any more.

The logger, which was private to the utils module, is now named without the
underscore, since it is used from outside it.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention.py                  | 35 ++++++++++---------
 .../attention/dot_product_attention/utils.py  |  4 +--
 2 files changed, 21 insertions(+), 18 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 34aa6cf403..d72eb6e98d 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -1854,22 +1854,25 @@ def forward(
                     _attention_backends["fused_attention_backend"] = fused_attention_backend
                     _attention_backends["use_unfused_attention"] = use_unfused_attention
                     _attention_backends["backend_selection_requires_update"] = False
-                    # Logged in eager only: dynamo graph-breaks on logging.Logger
-                    # methods, and log arguments are evaluated regardless of the
-                    # logger, so a no-op logger would not be enough.
-                    if not torch.compiler.is_compiling():
-                        if use_flash_attention:
-                            self.logger.info(
-                                "Running with FlashAttention backend (version %s)",
-                                flash_attention_backend,
-                            )
-                        elif use_fused_attention:
-                            self.logger.info(
-                                "Running with FusedAttention backend (sub-backend %s)",
-                                int(fused_attention_backend),
-                            )
-                        elif use_unfused_attention:
-                            self.logger.info("Running with UnfusedDotProductAttention backend")
+                    # logging.Logger methods graph-break under torch.compile, so
+                    # selection is only logged in eager -- as in
+                    # get_attention_backend. Note the arguments below are
+                    # evaluated either way, so they have to stay traceable.
+                    logger = (
+                        dpa_utils.no_op_logger if torch.compiler.is_compiling() else self.logger
+                    )
+                    if use_flash_attention:
+                        logger.info(
+                            "Running with FlashAttention backend (version %s)",
+                            flash_attention_backend,
+                        )
+                    elif use_fused_attention:
+                        logger.info(
+                            "Running with FusedAttention backend (sub-backend %s)",
+                            int(fused_attention_backend),
+                        )
+                    elif use_unfused_attention:
+                        logger.info("Running with UnfusedDotProductAttention backend")
                 else:
                     use_flash_attention = _attention_backends["use_flash_attention"]
                     flash_attention_backend = _attention_backends["flash_attention_backend"]
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 383afebd69..ccf46fda22 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -346,7 +346,7 @@ def error(self, *args, **kwargs):
         """No-op."""
 
 
-_no_op_logger = _NoOpLogger()
+no_op_logger = _NoOpLogger()
 
 
 @torch.compiler.assume_constant_result
@@ -461,7 +461,7 @@ def get_attention_backend(
     if torch.compiler.is_compiling():
         # logging.Logger methods graph-break under torch.compile; backend
         # selection logs are only emitted in eager mode.
-        logger = _no_op_logger
+        logger = no_op_logger
     else:
         logger = logging.getLogger("DotProductAttention")
         logger.setLevel(AttentionLogging._log_level)

From 10d141b3fc70fe5602a8a15bd84a71d25732fcef Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 11:51:46 +0200
Subject: [PATCH 21/67] Return the fused sub-backend as an int, in eager as
 well

get_attention_backend cast it back to a FusedAttnBackend member unless it was
tracing, which made its return type depend on the execution mode -- something
every reader of that function then has to keep in mind, for no benefit inside
transformer_engine: fused_attn_fwd/bwd call cast() on whatever they are given,
and every comparison against a member works with an int on either side.

It now returns an int always, and the docstring says so.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../attention/dot_product_attention/utils.py    | 17 +++++------------
 .../pytorch/cpp_extensions/fused_attn.py        |  2 +-
 2 files changed, 6 insertions(+), 13 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index ccf46fda22..9e055ef6e6 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -401,11 +401,11 @@ def get_attention_backend(
         Whether the `FlashAttention` backend has been selected.
     use_fused_attention : bool
         Whether the `FusedAttention` backend has been selected.
-    fused_attention_backend : FusedAttnBackend
-        If `use_fused_attention = True`, one of `FusedAttention` three sub-backends, else `None`.
-        Under `torch.compile` this is the sub-backend's plain integer value instead of the enum
-        member (see the cast below); `FusedAttnBackend.cast` accepts both, as does every consumer
-        in `transformer_engine.pytorch`.
+    fused_attention_backend : int
+        If `use_fused_attention = True`, the integer value of one of `FusedAttention`'s three
+        sub-backends, else `None`. It is not a `FusedAttnBackend` member because that does not
+        survive a graph break under `torch.compile`; `FusedAttnBackend.cast` turns it into one,
+        and comparing it against a member works either way.
     use_unfused_attention : bool
         Whether the `UnfusedDotProductAttention` backend has been selected.
     available_backends : List[bool]
@@ -1600,13 +1600,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
         use_flash_attention_4 = False
     use_flash_attention = use_flash_attention_2 or use_flash_attention_3 or use_flash_attention_4
     available_backends = [use_flash_attention, use_fused_attention, use_unfused_attention]
-    if fused_attention_backend is not None and not torch.compiler.is_compiling():
-        # The fused sub-backend is kept as a plain int while tracing: it has to
-        # survive a graph break (DotProductAttention hands it to FusedAttention,
-        # which is an eager island) and an enum member does not -- dynamo
-        # reconstructs the constant-folded member into the resumed frame as the
-        # function that produced it. Eager callers get the enum, as before.
-        fused_attention_backend = FusedAttnBackend.cast(fused_attention_backend)
     if use_flash_attention_2:
         flash_attention_backend = FlashAttentionUtils.version
     if use_flash_attention_3:
diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py
index 3a4e5aac66..9a33df7634 100644
--- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py
+++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py
@@ -112,7 +112,7 @@ class FusedAttnBackend(IntEnum):
     the same way as with the dict this used to be.
 
     Members do not survive a graph break, though, so ``get_attention_backend``
-    hands out plain ints while tracing -- ``cast`` normalizes both back to a
+    returns the sub-backend as a plain int and ``cast`` turns it back into a
     member.
     """
 

From 93a201d42e02f01bba21a67ad9c3e6f44f78ce05 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 11:57:12 +0200
Subject: [PATCH 22/67] Rename eager_under_compile_if to fallback_to_eager_when

The old name read as a run-on, and its '_if' suffix suggested a boolean
condition where the argument actually returns the reason. The new one reads as
a sentence where it is used and matches the warning the decorator emits.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../attention/dot_product_attention/dot_product_attention.py  | 4 ++--
 transformer_engine/pytorch/jit.py                             | 2 +-
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index d72eb6e98d..ba6c5a4769 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -42,7 +42,7 @@
     CudaRNGStatesTracker,
     graph_safe_rng_available,
 )
-from transformer_engine.pytorch.jit import eager_under_compile_if
+from transformer_engine.pytorch.jit import fallback_to_eager_when
 from transformer_engine.pytorch.graph import is_graph_capturing
 from transformer_engine.pytorch.attention.inference import InferenceParams
 
@@ -1124,7 +1124,7 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
-    @eager_under_compile_if(_needs_eager_dpa)
+    @fallback_to_eager_when(_needs_eager_dpa)
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,
diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py
index 2d0ec02d8d..f6bbf55bde 100644
--- a/transformer_engine/pytorch/jit.py
+++ b/transformer_engine/pytorch/jit.py
@@ -79,7 +79,7 @@ def no_torch_dynamo(recursive=True):  # pylint: disable=unused-argument
         return lambda func: func
 
 
-def eager_under_compile_if(needs_eager: Callable[[Dict[str, Any]], Optional[str]]):
+def fallback_to_eager_when(needs_eager: Callable[[Dict[str, Any]], Optional[str]]):
     """Decorator running the wrapped method eagerly whenever `needs_eager` gives
     a reason why the call is unsupported on the compiled path.
 

From 0866319fe05e2d17d1f768ca817919c53fa9f08d Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 12:06:44 +0200
Subject: [PATCH 23/67] Fold the eager fallback into no_torch_dynamo as a
 `when` predicate

It was a decorator of its own, which meant a second name for what is the same
thing this module already does: disable dynamo for a frame. It is now the
conditional form of no_torch_dynamo -- @no_torch_dynamo(when=predicate) -- so
there is one decorator to find and one docstring to read.

The predicate receives the call's arguments keyed by parameter name and returns
the reason this call cannot be traced, or None to have it traced.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention.py                  |  4 +-
 transformer_engine/pytorch/jit.py             | 94 +++++++++----------
 2 files changed, 48 insertions(+), 50 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index ba6c5a4769..774f6a44d1 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -42,7 +42,7 @@
     CudaRNGStatesTracker,
     graph_safe_rng_available,
 )
-from transformer_engine.pytorch.jit import fallback_to_eager_when
+from transformer_engine.pytorch.jit import no_torch_dynamo
 from transformer_engine.pytorch.graph import is_graph_capturing
 from transformer_engine.pytorch.attention.inference import InferenceParams
 
@@ -1124,7 +1124,7 @@ def get_quantizer_roles(
             ]
         return base[:num_quantizers]
 
-    @fallback_to_eager_when(_needs_eager_dpa)
+    @no_torch_dynamo(when=_needs_eager_dpa)
     def forward(
         self,
         query_layer: Optional[torch.Tensor] = None,
diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py
index f6bbf55bde..164610f5a5 100644
--- a/transformer_engine/pytorch/jit.py
+++ b/transformer_engine/pytorch/jit.py
@@ -7,7 +7,7 @@
 import os
 import warnings
 from functools import wraps
-from typing import Any, Callable, Dict, Optional, Tuple
+from typing import Callable, Optional, Tuple
 import torch
 
 from .torch_version import torch_version
@@ -51,22 +51,58 @@ def wrapper(*args, **kwargs):
 if torch.__version__ >= "2":
     import torch._dynamo
 
-    def no_torch_dynamo(recursive=True):
-        """Decorator to disable Torch Dynamo, except during ONNX export."""
+    def no_torch_dynamo(recursive=True, when=None):
+        """Decorator to disable Torch Dynamo, except during ONNX export.
 
-        def decorator(f):
+        `when` makes it conditional: it is called with the arguments of the call
+        keyed by parameter name, and returns the reason this particular call
+        cannot be traced, or None to have it traced as usual. The reason is
+        reported once per distinct message.
+        """
+
+        def _disable(f):
             # no "recursive" option in pyTorch 2.0 - it acts as if recursive was True
-            disabled_f = (
+            return (
                 torch._dynamo.disable(f, recursive=recursive)
                 if torch.__version__ >= "2.1"
                 else torch._dynamo.disable(f)
             )
 
-            @wraps(f)
-            def wrapper(*args, **kwargs):
-                if is_in_onnx_export_mode():
+        def decorator(f):
+            disabled_f = _disable(f)
+            if when is None:
+
+                @wraps(f)
+                def wrapper(*args, **kwargs):
+                    if is_in_onnx_export_mode():
+                        return f(*args, **kwargs)
+                    return disabled_f(*args, **kwargs)
+
+            else:
+                parameter_names = list(inspect.signature(f).parameters)
+
+                # The warning belongs inside the disabled function: warnings.warn
+                # graph-breaks on its own, which would mask the break that matters.
+                def report_and_run(reason, *args, **kwargs):
+                    warnings.warn(
+                        f"Falling back to eager execution under torch.compile: {reason} is"
+                        " unsupported on the compiled path (graph-breaks under fullgraph=True).",
+                        stacklevel=3,
+                    )
                     return f(*args, **kwargs)
-                return disabled_f(*args, **kwargs)
+
+                disabled_report_and_run = _disable(report_and_run)
+
+                @wraps(f)
+                def wrapper(*args, **kwargs):  # pylint: disable=function-redefined
+                    if is_in_onnx_export_mode() or not torch.compiler.is_compiling():
+                        return f(*args, **kwargs)
+                    call = dict(zip(parameter_names, args))
+                    call.update(kwargs)
+                    reason = when(call)
+                    if reason is None:
+                        return f(*args, **kwargs)
+                    return disabled_report_and_run(reason, *args, **kwargs)
 
             return wrapper
 
@@ -74,49 +110,11 @@ def wrapper(*args, **kwargs):
 
 else:
     # Fallback for PyTorch < 2.0: no-op decorator
-    def no_torch_dynamo(recursive=True):  # pylint: disable=unused-argument
+    def no_torch_dynamo(recursive=True, when=None):  # pylint: disable=unused-argument
         """No-op decorator for PyTorch < 2.0."""
         return lambda func: func
 
 
-def fallback_to_eager_when(needs_eager: Callable[[Dict[str, Any]], Optional[str]]):
-    """Decorator running the wrapped method eagerly whenever `needs_eager` gives
-    a reason why the call is unsupported on the compiled path.
-
-    `needs_eager` is handed the call's arguments keyed by parameter name, taken
-    from the wrapped signature, so it does not have to care whether the caller
-    passed them by name or by position. Returning None keeps the call compiled.
-    """
-
-    def decorator(fn):
-        parameter_names = list(inspect.signature(fn).parameters)
-
-        # The warning belongs inside the dynamo-disabled function: warnings.warn
-        # graph-breaks on its own, masking the break that matters.
-        @no_torch_dynamo()
-        def eager_fn(reason, *args, **kwargs):
-            warnings.warn(
-                f"Falling back to eager execution under torch.compile: {reason} is"
-                " unsupported on the compiled path (graph-breaks under fullgraph=True).",
-                stacklevel=3,
-            )
-            return fn(*args, **kwargs)
-
-        @wraps(fn)
-        def wrapper(*args, **kwargs):
-            if torch.compiler.is_compiling():
-                call = dict(zip(parameter_names, args))
-                call.update(kwargs)
-                reason = needs_eager(call)
-                if reason is not None:
-                    return eager_fn(reason, *args, **kwargs)
-            return fn(*args, **kwargs)
-
-        return wrapper
-
-    return decorator
-
-
 def set_jit_fusion_options() -> None:
     """Set PyTorch JIT layer fusion options."""
     # flags required to enable jit fusion kernels

From a80a41aa6d89d8922d27a4e725c6172e8b61c3ce Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 12:13:25 +0200
Subject: [PATCH 24/67] Cover the declared packed layouts beyond bs3hd

Packing was pinned to one layout by an assert in the input builder, so of the
twelve declarations DotProductAttention accepts -- qkv_layer or kv_layer,
interleaved at -3 or -2, in each of the three formats -- exactly one was
tested. The builder now takes what is packed and where it is interleaved, and
derives the shapes and the layout string the way DotProductAttention does.

Three configurations are added, so both packed inputs and both interleave
dimensions are covered.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 69 +++++++++++++++++++++++++----
 1 file changed, 60 insertions(+), 9 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index d20a0e70dc..f65bb8cb95 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -561,15 +561,34 @@ def fn(q, k, v, extra):
 # get_available_attention_backends() answers that, so configurations only one
 # backend supports (arbitrary masks, biases, MLA head dims, ...) are covered
 # rather than avoided.
-def _cfg(model_config, qkv_format="bshd", packed=None, share_cu_seqlens=False):
+def _cfg(
+    model_config,
+    qkv_format="bshd",
+    packed=None,
+    interleave_dim=-3,
+    share_cu_seqlens=False,
+):
+    """A model configuration plus what DotProductAttention is handed it as.
+
+    `packed` is "qkv" or "kv" for the declarative packed inputs, interleaved at
+    `interleave_dim`, or None for separate q/k/v.
+    """
     return dict(
         model_config=model_config,
         qkv_format=qkv_format,
         packed=packed,
+        interleave_dim=interleave_dim,
         share_cu_seqlens=share_cu_seqlens,
     )
 
 
+def _packed_layout(qkv_format: str, packed_dim: int, interleave_dim: int) -> str:
+    """bshd + 3 @ -3 -> bs3hd; bshd + 2 @ -2 -> bsh2d; as DotProductAttention
+    derives it from the declaration."""
+    position = len(qkv_format) + interleave_dim + 1
+    return qkv_format[:position] + str(packed_dim) + qkv_format[position:]
+
+
 _DPA_COMPILE_CONFIGS = {
     "self_bshd_causal": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal")),
     "self_sbhd_no_mask": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="no_mask"), "sbhd"),
@@ -587,7 +606,17 @@ def _cfg(model_config, qkv_format="bshd", packed=None, share_cu_seqlens=False):
     "self_thd_shared_cu_seqlens": _cfg(
         ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"), "thd", share_cu_seqlens=True
     ),
-    "packed_qkv_bs3hd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="bs3hd"),
+    "packed_qkv_bs3hd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="qkv"),
+    "packed_qkv_bsh3d": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="qkv", interleave_dim=-2
+    ),
+    "packed_kv_bshd_bs2hd": _cfg(ModelConfig(2, 128, 4, 64, attn_mask_type="causal"), packed="kv"),
+    "packed_kv_thd_th2d": _cfg(
+        ModelConfig(2, 128, 4, 64, attn_mask_type="padding_causal"),
+        "thd",
+        packed="kv",
+        interleave_dim=-2,
+    ),
     # Configurations below are supported by one backend only, or take a code
     # path of their own inside DotProductAttention.
     "alibi_bshd_causal": _cfg(
@@ -605,7 +634,12 @@ def _cfg(model_config, qkv_format="bshd", packed=None, share_cu_seqlens=False):
 
 
 def _qkv_layout(spec: dict) -> str:
-    return spec["packed"] or "_".join([spec["qkv_format"]] * 3)
+    qkv_format, packed = spec["qkv_format"], spec["packed"]
+    if packed is None:
+        return "_".join([qkv_format] * 3)
+    if packed == "qkv":
+        return _packed_layout(qkv_format, 3, spec["interleave_dim"])
+    return f"{qkv_format}_{_packed_layout(qkv_format, 2, spec['interleave_dim'])}"
 
 
 def _make_dpa(spec: dict, dtype: torch.dtype) -> te.DotProductAttention:
@@ -669,14 +703,31 @@ def _randn(shape):
         return torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
 
     if packed is not None:
-        # Declarative packed QKV: q/k/v are derived from one buffer by DPA
+        # Declarative packed inputs: q/k/v are derived from one buffer by DPA
         # itself, and the layout comes from the declaration -- the only packed
         # layout that torch.compile supports.
-        assert packed == "bs3hd" and qkv_format == "bshd" and h == g and d_qk == d_v
-        qkv = _randn((b, s_q, 3, h, d_qk))
-        kwargs["qkv_layer"] = qkv
-        kwargs["qkv_interleave_dim"] = -3
-        grad_tensors, args = [qkv], ()
+        assert h == g and d_qk == d_v, "packed inputs require uniform heads and head dims"
+        interleave_dim = spec["interleave_dim"]
+
+        def _packed(seqlen, tokens, packed_dim):
+            leading = {
+                "bshd": (b, seqlen),
+                "sbhd": (seqlen, b),
+                "thd": (tokens,),
+            }[qkv_format]
+            trailing = (packed_dim, h, d_qk) if interleave_dim == -3 else (h, packed_dim, d_qk)
+            return _randn(leading + trailing)
+
+        kwargs["qkv_interleave_dim"] = interleave_dim
+        if packed == "qkv":
+            qkv = _packed(s_q, t_q, 3)
+            kwargs["qkv_layer"] = qkv
+            grad_tensors, args = [qkv], ()
+        else:
+            q = _randn(_shape(s_q, t_q, h, d_qk))
+            kv = _packed(s_kv, t_kv, 2)
+            kwargs["kv_layer"] = kv
+            grad_tensors, args = [q, kv], (q,)
     else:
         q = _randn(_shape(s_q, t_q, h, d_qk))
         k = _randn(_shape(s_kv, t_kv, g, d_qk))

From 9eb1449daf1d40e349d27ad50835169f02c9e9c1 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 12:21:13 +0200
Subject: [PATCH 25/67] Let ONNX export keep its own path in get_qkv_layout

The ONNX exporter runs through dynamo, so torch.compiler.is_compiling() is
true during an export as well -- measured: inside torch.onnx.export(dynamo=True)
both it and is_in_onnx_export_mode() are set. The compile branch was written
first and took the export over: the layout came from the format instead of the
detection ONNX relies on, tensors were no longer made contiguous, and packed
views hit an assert where they used to be copied.

Export mode is the narrower context, so it is checked first. This restores the
previous behaviour for export while leaving the compiled path as it was.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/attention/dot_product_attention/utils.py   | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 9e055ef6e6..a2d260c6b2 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -2565,7 +2565,11 @@ def run_iteratively(q, k, v):
 
         return qkv_layout
 
-    if torch.compiler.is_compiling():
+    if is_in_onnx_export_mode():
+        # Checked first: the ONNX exporter runs through dynamo, so it also sets
+        # is_compiling(), and it has its own handling below.
+        qkv_layout = "not_supported"
+    elif torch.compiler.is_compiling():
         # run_iteratively reads data pointers and storage offsets, which dynamo
         # cannot trace; unpacked q/k/v need no detection anyway.
         assert not qkv_layout_needs_detection(q, k, v), (
@@ -2577,10 +2581,8 @@ def run_iteratively(q, k, v):
             qkv_layout = "_".join([qkv_format] * 3)
         else:
             qkv_layout = q_format + "_" + kv_format + "_" + kv_format
-    elif not is_in_onnx_export_mode():
-        qkv_layout = run_iteratively(q, k, v)
     else:
-        qkv_layout = "not_supported"
+        qkv_layout = run_iteratively(q, k, v)
     if qkv_layout == "not_supported":
         # force q,k,v to be contiguous and run get_layout again
         q, k, v = [x.contiguous() for x in [q, k, v]]

From 5c5b451264637a5726868043bee551be71cb88d4 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 12:31:08 +0200
Subject: [PATCH 26/67] Guard the assumption the argument binding rests on

Arguments are matched to parameter names by position, which is only right
while the signature has no *args in between -- and would misalign silently if
one were added. It is checked when the decorator is applied.

The docstring now also says that an argument the call left out is simply
absent, so a predicate reading it with .get() sees its default only as long as
that default is None.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/jit.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py
index 164610f5a5..62e6950154 100644
--- a/transformer_engine/pytorch/jit.py
+++ b/transformer_engine/pytorch/jit.py
@@ -55,9 +55,11 @@ def no_torch_dynamo(recursive=True, when=None):
         """Decorator to disable Torch Dynamo, except during ONNX export.
 
         `when` makes it conditional: it is called with the arguments of the call
-        keyed by parameter name, and returns the reason this particular call
-        cannot be traced, or None to have it traced as usual. The reason is
-        reported once per distinct message.
+        keyed by parameter name -- an argument the call left out is absent, so
+        reading it with .get() yields its default as long as that default is
+        None -- and returns the reason this particular call cannot be traced, or
+        None to have it traced as usual. The reason is reported once per
+        distinct message.
         """
 
         def _disable(f):
@@ -79,7 +81,13 @@ def wrapper(*args, **kwargs):
                     return disabled_f(*args, **kwargs)
 
             else:
-                parameter_names = list(inspect.signature(f).parameters)
+                parameters = inspect.signature(f).parameters
+                # Arguments are matched to names by position, which only holds
+                # without a *args in between.
+                assert not any(
+                    p.kind is inspect.Parameter.VAR_POSITIONAL for p in parameters.values()
+                ), f"no_torch_dynamo(when=...) does not support *args, which {f.__name__} takes"
+                parameter_names = list(parameters)
 
                 # The warning belongs inside the disabled function: warnings.warn
                 # graph-breaks on its own, which would mask the break that matters.

From 95030f9a3140e7b62067bbf903610562a55462f2 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 13:49:53 +0200
Subject: [PATCH 27/67] Check that CUDA graphs were actually captured, and
 generalize a docstring

fullgraph=True says dynamo did not break the graph, but inductor can still
decline to capture -- a mutated input, a CPU scalar -- and run the thing
normally, which would leave the CUDA-graphs test passing while measuring
nothing. Its skip counter is asserted to be zero. Measured on this branch:
nothing is skipped, the tree has one root, and interleaving the eager
reference call between replays does not disturb it.

The tolerance docstring described which backends are compared exactly by
naming them; it now states the rule -- a backend whose kernel is the same
either way -- so it still holds once FusedAttention traces.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index f65bb8cb95..6ae070b70d 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -21,6 +21,8 @@
 except ImportError:
     _opaque_available = False
 
+from torch._dynamo.utils import counters
+
 import transformer_engine.pytorch as te
 import transformer_engine_torch as tex
 from transformer_engine.common import recipe
@@ -804,12 +806,11 @@ def _assert_dpa_backend(backend: str) -> None:
 def _assert_matches_eager(actual, expected, backend: str, dtype: torch.dtype) -> None:
     """Assert a compiled result matches the eager one.
 
-    FlashAttention and FusedAttention run the same kernel either way, so they
-    have to match exactly. Inductor reassociates the unfused backend's softmax
-    sums, and that error scales with the magnitude of those sums rather than
-    with each output element -- so the absolute tolerance comes from the
-    tensor's own scale, an elementwise relative tolerance being meaningless for
-    elements near zero.
+    A backend that calls the same kernel either way has to match exactly:
+    compiling changes what surrounds it, not its arithmetic. The unfused
+    backend is instead built from PyTorch ops, which inductor fuses and
+    reassociates; that error scales with the softmax sums rather than with each
+    output element, hence an absolute tolerance taken from the tensor's scale.
     """
     if backend != "unfused":
         torch.testing.assert_close(actual, expected, rtol=0.0, atol=0.0)
@@ -918,6 +919,7 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
     module = _make_dpa(spec, dtype)
 
     torch._dynamo.reset()
+    counters.clear()
     compiled = torch.compile(module, fullgraph=True, mode="reduce-overhead")
 
     for _ in range(3):
@@ -929,6 +931,9 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
         eager = _run_and_capture(module, args, kwargs, grads)
         _assert_run_matches(replayed, eager, backend, dtype)
     _assert_dpa_backend(backend)
+    # Without this, inductor declining to capture -- a mutated input, a CPU
+    # scalar -- would leave the test passing while measuring nothing.
+    assert not counters["inductor"]["cudagraph_skips"], "inductor skipped CUDA graphs"
 
 
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)

From b0e56d450bac0300da046cc7e56ea6b098ec20b3 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 14:51:20 +0200
Subject: [PATCH 28/67] Skip lazy compilation while already tracing

Filling lazy_compile's closure cell is a side effect Dynamo rejects inside
a higher-order op, so compiling a module whose autograd.Function backward
makes the first call to a jit_fuser'd function fails with

    Unsupported: HOP: Unsafe side effect
      Attempted to mutate CellVariable()

An earlier eager run hides this by filling the cell beforehand. The nested
torch.compile is inlined into the outer graph either way, so calling the
function directly while tracing is equivalent.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/jit.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/transformer_engine/pytorch/jit.py b/transformer_engine/pytorch/jit.py
index 62e6950154..22e9c9ae3e 100644
--- a/transformer_engine/pytorch/jit.py
+++ b/transformer_engine/pytorch/jit.py
@@ -28,6 +28,8 @@ def lazy_compile(func):
     @wraps(func)
     def wrapper(*args, **kwargs):
         nonlocal compiled_func
+        if torch.compiler.is_compiling():
+            return func(*args, **kwargs)
         if compiled_func is None:
             compiled_func = torch.compile(func)
         return compiled_func(*args, **kwargs)

From 2cf695bb59a551efb8e6ad49dbe3ef85af2d54f4 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 30 Jul 2026 16:49:43 +0200
Subject: [PATCH 29/67] Run context parallel attention eagerly

The context parallel implementation uses P2P communication and its own CUDA
stream, neither of which dynamo traces, and FlashAttention reaches it from a
forward that is no longer excluded from the graph.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../attention/dot_product_attention/dot_product_attention.py   | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 774f6a44d1..0ed01a9c9e 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -210,6 +210,9 @@ def _needs_eager_dpa(call: Dict[str, Any]) -> Optional[str]:
         if fp8_recipe.fp8_dpa or fp8_recipe.fp8_mha:
             return "FP8 attention"
 
+    if call["self"].cp_group is not None:
+        return "context parallelism"
+
     qkv_format = call.get("qkv_format") or call["self"].qkv_format
     if qkv_format == "thd" and (
         call.get("max_seqlen_q") is None or call.get("max_seqlen_kv") is None

From b89446f4d04341a91708ffe5c3e8e532bfa05bff Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 13:04:46 +0200
Subject: [PATCH 30/67] Keep DotProductAttention compilable under a CUDA RNG
 states tracker

Forking the tracker swaps the global CUDA generator state, which dynamo
refuses to trace. With attention_dropout == 0 the dropout draws nothing
from the generator, so the fork is a no-op -- use nullcontext instead and
keep the call in the graph. Calls that do use dropout under a tracker run
as an eager island.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention/dot_product_attention.py     | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 0ed01a9c9e..6bb10aec56 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -225,6 +225,11 @@ def _needs_eager_dpa(call: Dict[str, Any]) -> Optional[str]:
         qkv = [call.get(name) for name in ("query_layer", "key_layer", "value_layer")]
         if dpa_utils.qkv_layout_needs_detection(*qkv):
             return "detecting packed q/k/v that were not declared via qkv_layer/kv_layer"
+
+    if call["self"].rng_states_tracker is not None and call["self"].attention_dropout > 0:
+        # Forking the tracker swaps the global CUDA generator state, which Dynamo
+        # refuses to trace. With no dropout there is nothing to fork for.
+        return "attention dropout under a CUDA RNG states tracker"
     return None
 
 
@@ -577,7 +582,10 @@ def __init__(
         else:
             self.rng_states_tracker = get_rng_state_tracker()
             set_all_rng_states(self.rng_states_tracker.get_states())
-            attention_dropout_ctx = self.rng_states_tracker.fork
+            # Forking only matters if the dropout actually draws from the generator.
+            attention_dropout_ctx = (
+                self.rng_states_tracker.fork if attention_dropout > 0 else nullcontext
+            )
 
         if softmax_scale is None:
             softmax_scale = 1.0 / math.sqrt(

From fb50154cc461711e5b9daae88ee2dcee0f3e71ca Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 14:14:50 +0200
Subject: [PATCH 31/67] Run checkpointed attention eagerly

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../attention/dot_product_attention/dot_product_attention.py   | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 6bb10aec56..d4cba79ea9 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -213,6 +213,9 @@ def _needs_eager_dpa(call: Dict[str, Any]) -> Optional[str]:
     if call["self"].cp_group is not None:
         return "context parallelism"
 
+    if call.get("checkpoint_core_attention", False):
+        return "activation checkpointing of the attention"
+
     qkv_format = call.get("qkv_format") or call["self"].qkv_format
     if qkv_format == "thd" and (
         call.get("max_seqlen_q") is None or call.get("max_seqlen_kv") is None

From 4b1ebb2806c9e8d3e5907e9d2442e2ec61934a40 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 14:48:26 +0200
Subject: [PATCH 32/67] Support KV caching under torch.compile

Wrap the KV cache kernels in custom ops: dynamo cannot trace pybind calls
into transformer_engine_torch, so a decoding step used to break the graph.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 69 +++++++++++++--
 .../dot_product_attention/backends.py         | 14 +--
 .../pytorch/attention/inference.py            | 13 +--
 .../pytorch/attention/kv_cache_ops.py         | 88 +++++++++++++++++++
 4 files changed, 167 insertions(+), 17 deletions(-)
 create mode 100644 transformer_engine/pytorch/attention/kv_cache_ops.py

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 6ae070b70d..504204a6bf 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -569,11 +569,13 @@ def _cfg(
     packed=None,
     interleave_dim=-3,
     share_cu_seqlens=False,
+    kv_cache=False,
 ):
     """A model configuration plus what DotProductAttention is handed it as.
 
     `packed` is "qkv" or "kv" for the declarative packed inputs, interleaved at
-    `interleave_dim`, or None for separate q/k/v.
+    `interleave_dim`, or None for separate q/k/v. `kv_cache` runs the call as a
+    decoding step against an InferenceParams KV cache.
     """
     return dict(
         model_config=model_config,
@@ -581,6 +583,7 @@ def _cfg(
         packed=packed,
         interleave_dim=interleave_dim,
         share_cu_seqlens=share_cu_seqlens,
+        kv_cache=kv_cache,
     )
 
 
@@ -632,6 +635,14 @@ def _packed_layout(qkv_format: str, packed_dim: int, interleave_dim: int) -> str
     "sink_softmax_bshd": _cfg(
         ModelConfig(2, 128, 4, 64, attn_mask_type="causal", softmax_type="off-by-one")
     ),
+    # A decoding step against a KV cache. FlashAttention 2 wants the non-paged
+    # cache length divisible by 256.
+    "kv_cache_bshd": _cfg(
+        ModelConfig(
+            2, 8, 4, 64, max_seqlen_kv=256, attn_mask_type="padding_causal_bottom_right"
+        ),
+        kv_cache=True,
+    ),
 }
 
 
@@ -646,6 +657,21 @@ def _qkv_layout(spec: dict) -> str:
 
 def _make_dpa(spec: dict, dtype: torch.dtype) -> te.DotProductAttention:
     config = spec["model_config"]
+    if spec["kv_cache"]:
+        # KV caching addresses the cache by layer number, and a decoding step
+        # runs in inference mode.
+        return (
+            te.DotProductAttention(
+                num_attention_heads=config.num_heads,
+                kv_channels=config.kv_channels,
+                num_gqa_groups=config.num_gqa_groups,
+                qkv_format=spec["qkv_format"],
+                attn_mask_type=config.attn_mask_type,
+                layer_number=1,
+            )
+            .to(dtype=dtype, device="cuda")
+            .eval()
+        )
     return te.DotProductAttention(
         num_attention_heads=config.num_heads,
         kv_channels=config.kv_channels,
@@ -704,6 +730,28 @@ def _shape(s, t, heads, head_dim):
     def _randn(shape):
         return torch.randn(shape, dtype=dtype, device="cuda", requires_grad=True)
 
+    if spec["kv_cache"]:
+        from collections import OrderedDict
+        from transformer_engine.pytorch.attention import InferenceParams
+
+        inference_params = InferenceParams(
+            max_batch_size=b,
+            max_sequence_length=s_kv,
+            num_heads_kv=g,
+            head_dim_k=d_qk,
+            dtype=dtype,
+            qkv_format=qkv_format,
+        )
+        inference_params.allocate_memory(1)
+        inference_params.pre_step(OrderedDict((i, s_q) for i in range(b)))
+        qkv = [
+            torch.randn(_shape(s_q, t_q, heads, d_qk), dtype=dtype, device="cuda")
+            for heads in (h, g, g)
+        ]
+        # The sequence lengths come from the cache, not from cu_seqlens, and a
+        # decoding step has no gradients to compare.
+        return tuple(qkv), {"inference_params": inference_params}, []
+
     if packed is not None:
         # Declarative packed inputs: q/k/v are derived from one buffer by DPA
         # itself, and the layout comes from the declaration -- the only packed
@@ -746,7 +794,9 @@ def _packed(seqlen, tokens, packed_dim):
     return args, kwargs, grad_tensors
 
 
-def _skip_unsupported(spec: dict, backend: str, dtype, compiled: bool = True) -> None:
+def _skip_unsupported(
+    spec: dict, backend: str, dtype, compiled: bool = True, inference_params=None
+) -> None:
     """Skip what the backend under test cannot run, or -- for a test that
     compiles it -- cannot be compiled."""
     if compiled and backend == "fused":
@@ -755,7 +805,11 @@ def _skip_unsupported(spec: dict, backend: str, dtype, compiled: bool = True) ->
         # and the tests below cover it as they do the others.
         pytest.skip("FusedAttention is an eager island and does not compile")
     available, _, _ = get_available_attention_backends(
-        spec["model_config"], dtype, _qkv_layout(spec)
+        spec["model_config"],
+        dtype,
+        _qkv_layout(spec),
+        inference_params=inference_params,
+        is_training=inference_params is None,
     )
     flash_supported, fused_supported, unfused_supported = available
     supported = {
@@ -833,7 +887,8 @@ def _run_and_capture(fn, args, kwargs, grads):
     so that the same inputs can be reused for the other run.
     """
     out = fn(*args, **kwargs).clone()
-    out.sum().backward()
+    if grads:
+        out.sum().backward()
     torch.cuda.synchronize()
     captured = []
     for tensor in grads:
@@ -881,11 +936,11 @@ def test_dpa_torch_compile(monkeypatch, backend, config):
     """
     dtype = torch.bfloat16
     spec = _DPA_COMPILE_CONFIGS[config]
-    _skip_unsupported(spec, backend, dtype)
-    _force_dpa_backend(monkeypatch, backend)
-
     module = _make_dpa(spec, dtype)
     args, kwargs, grads = _make_dpa_inputs(spec, dtype)
+    _skip_unsupported(spec, backend, dtype, inference_params=kwargs.get("inference_params"))
+    _force_dpa_backend(monkeypatch, backend)
+
     _compare_compiled_to_eager(
         module, args, kwargs, grads, monkeypatch, backend, dtype, fullgraph=True
     )
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 04bcabe7c4..825e4cd412 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -50,6 +50,10 @@
 )
 from transformer_engine.pytorch.quantization import get_fp8_torch_dtype, FP8GlobalStateManager
 from transformer_engine.pytorch.distributed import get_distributed_world_size
+from transformer_engine.pytorch.attention.kv_cache_ops import (
+    convert_bshd_to_thd,
+    convert_thd_to_bshd,
+)
 from transformer_engine.pytorch.jit import no_torch_dynamo
 from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
     attn_forward_func_with_cp,
@@ -437,7 +441,7 @@ def _forward(
 
         if qkv_format == "thd_2bshd":
             batch_size = key_layer.shape[0]
-            query_layer = tex.convert_thd_to_bshd(
+            query_layer = convert_thd_to_bshd(
                 query_layer,
                 cu_seqlens_q,
                 batch_size,
@@ -1029,7 +1033,7 @@ def forward(
                 # convert from bshd to thd_2bshd for flash_attn_varlen_func/_with_kvcache;
                 # kernel assumes tensor is contiguous
                 if isinstance(query_layer, Float8Tensor):
-                    query_layer._data = tex.convert_bshd_to_thd(
+                    query_layer._data = convert_bshd_to_thd(
                         query_layer._data,
                         cu_seqlens_q,
                         batch_size * context_len,
@@ -1038,7 +1042,7 @@ def forward(
                         query_layer, data=query_layer._data, shape=query_layer._data.shape
                     )
                 else:
-                    query_layer = tex.convert_bshd_to_thd(
+                    query_layer = convert_bshd_to_thd(
                         query_layer,
                         cu_seqlens_q,
                         batch_size * context_len,
@@ -1287,7 +1291,7 @@ def convert_to_torch_float8(tensor, dtype):
             # all KV caching cases use thd_2bshd for calculation
             # convert results back to bshd from thd_2bshd
             if isinstance(query_layer, Float8Tensor):
-                output._data = tex.convert_thd_to_bshd(
+                output._data = convert_thd_to_bshd(
                     output._data,
                     cu_seqlens_q,
                     batch_size,
@@ -1295,7 +1299,7 @@ def convert_to_torch_float8(tensor, dtype):
                 )
                 output = Float8Tensor.make_like(output, data=output._data, shape=output._data.shape)
             else:
-                output = tex.convert_thd_to_bshd(
+                output = convert_thd_to_bshd(
                     output,
                     cu_seqlens_q,
                     batch_size,
diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py
index 08e50aad8b..b42be244af 100644
--- a/transformer_engine/pytorch/attention/inference.py
+++ b/transformer_engine/pytorch/attention/inference.py
@@ -11,7 +11,10 @@
 import torch
 
 import transformer_engine_torch as tex
-from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat
+from transformer_engine.pytorch.attention.kv_cache_ops import (
+    copy_to_kv_cache,
+    QKV_FORMAT_VALUE,
+)
 
 __all__ = ["InferenceParams", "KVCacheManager", "NonPagedKVCacheManager", "PagedKVCacheManager"]
 
@@ -549,7 +552,7 @@ def step(
             batch_size = new_k.shape[1]
             ctx_len = new_k.shape[0]
 
-        tex.copy_to_kv_cache(
+        copy_to_kv_cache(
             new_k,
             new_v,
             k_cache,
@@ -557,7 +560,7 @@ def step(
             self.batch_indices,
             cu_new_seqlens,
             cu_cached_seqlens,
-            QKVFormat[qkv_format],
+            QKV_FORMAT_VALUE[qkv_format],
             batch_size,
             ctx_len,
             self.max_seqlen,
@@ -807,7 +810,7 @@ def step(
             batch_size = new_k.shape[1]
             ctx_len = new_k.shape[0]
 
-        tex.copy_to_kv_cache(
+        copy_to_kv_cache(
             new_k,
             new_v,
             k_cache,
@@ -815,7 +818,7 @@ def step(
             self.page_table,
             cu_new_seqlens,
             cu_cached_seqlens,
-            QKVFormat[qkv_format],
+            QKV_FORMAT_VALUE[qkv_format],
             batch_size,
             ctx_len,
             self.max_seqlen,
diff --git a/transformer_engine/pytorch/attention/kv_cache_ops.py b/transformer_engine/pytorch/attention/kv_cache_ops.py
new file mode 100644
index 0000000000..105e8cac3f
--- /dev/null
+++ b/transformer_engine/pytorch/attention/kv_cache_ops.py
@@ -0,0 +1,88 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""KV cache kernels wrapped as custom ops, so they don't graph-break under torch.compile."""
+
+import torch
+import transformer_engine_torch as tex
+
+from transformer_engine.pytorch.cpp_extensions.fused_attn import QKVFormat
+
+# The ops take the format's value rather than the pybind enum: converting the
+# enum inside a traced region makes dynamo recurse until it gives up.
+QKV_FORMAT_VALUE = {name: int(fmt) for name, fmt in QKVFormat.items()}
+_QKV_FORMAT_BY_VALUE = {int(fmt): fmt for fmt in QKVFormat.values()}
+
+
+@torch.library.custom_op(
+    "te_kv_cache::copy_to_kv_cache",
+    mutates_args=("k_cache", "v_cache"),
+    device_types="cuda",
+)
+def copy_to_kv_cache(
+    new_k: torch.Tensor,
+    new_v: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    page_table: torch.Tensor,
+    cu_new_lens: torch.Tensor,
+    cu_cached_lens: torch.Tensor,
+    qkv_format: int,
+    b: int,
+    max_ctx_len: int,
+    max_seq_len: int,
+    max_pages_per_seq: int,
+    is_non_paged: bool,
+) -> None:
+    """Copy new key/value tokens into the KV cache."""
+    tex.copy_to_kv_cache(
+        new_k,
+        new_v,
+        k_cache,
+        v_cache,
+        page_table,
+        cu_new_lens,
+        cu_cached_lens,
+        _QKV_FORMAT_BY_VALUE[qkv_format],
+        b,
+        max_ctx_len,
+        max_seq_len,
+        max_pages_per_seq,
+        is_non_paged,
+    )
+
+
+@copy_to_kv_cache.register_fake
+def _copy_to_kv_cache_fake(*_args, **_kwargs) -> None:
+    return None
+
+
+@torch.library.custom_op("te_kv_cache::convert_bshd_to_thd", mutates_args=(), device_types="cuda")
+def convert_bshd_to_thd(tensor: torch.Tensor, cu_seqlens: torch.Tensor, t: int) -> torch.Tensor:
+    """Convert a tensor from bshd to thd."""
+    return tex.convert_bshd_to_thd(tensor, cu_seqlens, t)
+
+
+@convert_bshd_to_thd.register_fake
+def _convert_bshd_to_thd_fake(
+    tensor: torch.Tensor, cu_seqlens: torch.Tensor, t: int
+) -> torch.Tensor:
+    del cu_seqlens
+    return tensor.new_empty((t, *tensor.shape[2:]))
+
+
+@torch.library.custom_op("te_kv_cache::convert_thd_to_bshd", mutates_args=(), device_types="cuda")
+def convert_thd_to_bshd(
+    tensor: torch.Tensor, cu_seqlens: torch.Tensor, b: int, max_seq_len: int
+) -> torch.Tensor:
+    """Convert a tensor from thd to bshd."""
+    return tex.convert_thd_to_bshd(tensor, cu_seqlens, b, max_seq_len)
+
+
+@convert_thd_to_bshd.register_fake
+def _convert_thd_to_bshd_fake(
+    tensor: torch.Tensor, cu_seqlens: torch.Tensor, b: int, max_seq_len: int
+) -> torch.Tensor:
+    del cu_seqlens
+    return tensor.new_empty((b, max_seq_len, *tensor.shape[1:]))

From eb5ad558d37a8fe822832eaae8257d529f195bbf Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 14:59:46 +0200
Subject: [PATCH 33/67] Compile the fused sbh3d QKV split

Wrap fa_prepare_fwd/bwd like the KV cache kernels, and cover the layout
that reaches them: no configuration exercised it, compiled or eager.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           |  8 ++++++
 .../{kv_cache_ops.py => custom_ops.py}        | 28 ++++++++++++++++++-
 .../dot_product_attention/backends.py         |  8 ++++--
 .../pytorch/attention/inference.py            |  2 +-
 4 files changed, 41 insertions(+), 5 deletions(-)
 rename transformer_engine/pytorch/attention/{kv_cache_ops.py => custom_ops.py} (70%)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 504204a6bf..d9a5130f40 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -635,6 +635,14 @@ def _packed_layout(qkv_format: str, packed_dim: int, interleave_dim: int) -> str
     "sink_softmax_bshd": _cfg(
         ModelConfig(2, 128, 4, 64, attn_mask_type="causal", softmax_type="off-by-one")
     ),
+    # FlashAttention takes a fused sbhd->bshd split for sbh3d with a head
+    # dimension of 128 and at least 512 tokens, and a plain transpose otherwise.
+    "packed_qkv_sbh3d_fused_split": _cfg(
+        ModelConfig(4, 128, 4, 128, attn_mask_type="no_mask"),
+        "sbhd",
+        packed="qkv",
+        interleave_dim=-2,
+    ),
     # A decoding step against a KV cache. FlashAttention 2 wants the non-paged
     # cache length divisible by 256.
     "kv_cache_bshd": _cfg(
diff --git a/transformer_engine/pytorch/attention/kv_cache_ops.py b/transformer_engine/pytorch/attention/custom_ops.py
similarity index 70%
rename from transformer_engine/pytorch/attention/kv_cache_ops.py
rename to transformer_engine/pytorch/attention/custom_ops.py
index 105e8cac3f..f364633607 100644
--- a/transformer_engine/pytorch/attention/kv_cache_ops.py
+++ b/transformer_engine/pytorch/attention/custom_ops.py
@@ -2,7 +2,7 @@
 #
 # See LICENSE for license information.
 
-"""KV cache kernels wrapped as custom ops, so they don't graph-break under torch.compile."""
+"""Attention kernels wrapped as custom ops, so they don't graph-break under torch.compile."""
 
 import torch
 import transformer_engine_torch as tex
@@ -86,3 +86,29 @@ def _convert_thd_to_bshd_fake(
 ) -> torch.Tensor:
     del cu_seqlens
     return tensor.new_empty((b, max_seq_len, *tensor.shape[1:]))
+
+
+@torch.library.custom_op("te_attention::fa_prepare_fwd", mutates_args=(), device_types="cuda")
+def fa_prepare_fwd(qkvi: torch.Tensor) -> torch.Tensor:
+    """Split interleaved sbh3d QKV into bshd q/k/v."""
+    return tex.fa_prepare_fwd(qkvi)
+
+
+@fa_prepare_fwd.register_fake
+def _fa_prepare_fwd_fake(qkvi: torch.Tensor) -> torch.Tensor:
+    # qkvi is the q view into the packed buffer, and its strides cover all of it.
+    s, b, n, h = qkvi.shape
+    return qkvi.new_empty((3, b, s, n, h))
+
+
+@torch.library.custom_op("te_attention::fa_prepare_bwd", mutates_args=(), device_types="cuda")
+def fa_prepare_bwd(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """Pack bshd gradients back into an interleaved sbh3d buffer."""
+    return tex.fa_prepare_bwd(q, k, v)
+
+
+@fa_prepare_bwd.register_fake
+def _fa_prepare_bwd_fake(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    del k, v
+    b, s, n, h = q.shape
+    return q.new_empty((s, b, n, 3 * h))
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 825e4cd412..4a6ff10b15 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -50,9 +50,11 @@
 )
 from transformer_engine.pytorch.quantization import get_fp8_torch_dtype, FP8GlobalStateManager
 from transformer_engine.pytorch.distributed import get_distributed_world_size
-from transformer_engine.pytorch.attention.kv_cache_ops import (
+from transformer_engine.pytorch.attention.custom_ops import (
     convert_bshd_to_thd,
     convert_thd_to_bshd,
+    fa_prepare_bwd,
+    fa_prepare_fwd,
 )
 from transformer_engine.pytorch.jit import no_torch_dynamo
 from transformer_engine.pytorch.attention.dot_product_attention.context_parallel import (
@@ -777,7 +779,7 @@ def forward(
         # All inputs received are non-contiguous tensors.
         # The `query_layer` tensor is used to access the
         # full memory region of the QKV tensor.
-        qkv = tex.fa_prepare_fwd(query_layer)
+        qkv = fa_prepare_fwd(query_layer)
         q, k, v = split_tensor_along_dim(qkv, 0, 3)
         query_layer = torch.squeeze(q, 0)
         key_layer = torch.squeeze(k, 0)
@@ -792,7 +794,7 @@ def backward(
         dv: torch.Tensor,
     ) -> Tuple[Union[torch.Tensor, None], ...]:
         # pylint: disable=missing-function-docstring
-        dqkv = tex.fa_prepare_bwd(dq, dk, dv)
+        dqkv = fa_prepare_bwd(dq, dk, dv)
         dq, dk, dv = split_tensor_along_dim(dqkv, -1, 3)
         return dq, dk, dv
 
diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py
index b42be244af..0db6fdca2e 100644
--- a/transformer_engine/pytorch/attention/inference.py
+++ b/transformer_engine/pytorch/attention/inference.py
@@ -11,7 +11,7 @@
 import torch
 
 import transformer_engine_torch as tex
-from transformer_engine.pytorch.attention.kv_cache_ops import (
+from transformer_engine.pytorch.attention.custom_ops import (
     copy_to_kv_cache,
     QKV_FORMAT_VALUE,
 )

From 7da1a66d053c3058313aaf05893ef8cfd818a374 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 15:23:54 +0200
Subject: [PATCH 34/67] Run FlashAttention 4 eagerly

Version 4 registers no custom ops, unlike 2 and 3: it builds its kernels
through the CUTLASS DSL as it runs. Disable dynamo on its entry points, so
only the calls that reach it become eager islands.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/attention/dot_product_attention/backends.py  | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 4a6ff10b15..a5d2b05141 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -174,11 +174,16 @@
     flash_attn_varlen_func_v4 = None
 else:
     from flash_attn.cute.interface import (  # pylint: disable=ungrouped-imports,no-name-in-module
-        flash_attn_func as flash_attn_func_v4,
-        flash_attn_varlen_func as flash_attn_varlen_func_v4,
+        flash_attn_func as _flash_attn_func_v4,
+        flash_attn_varlen_func as _flash_attn_varlen_func_v4,
         _validate_head_dims as _fa4_validate_head_dims,
     )
 
+    # Unlike versions 2 and 3, FlashAttention 4 registers no custom ops: it builds
+    # its kernels through the CUTLASS DSL as it runs. Keep it an eager island.
+    flash_attn_func_v4 = no_torch_dynamo()(_flash_attn_func_v4)
+    flash_attn_varlen_func_v4 = no_torch_dynamo()(_flash_attn_varlen_func_v4)
+
     fa_utils.v4_validate_head_dims = _fa4_validate_head_dims
     fa_utils.set_flash_attention_4_params()
 

From d1beb9b91cc16c8d80804da4c92d4a20ea2617b7 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 15:49:15 +0200
Subject: [PATCH 35/67] Pin the sequence lengths backend selection bakes in

A second sequence length makes dynamo's automatic dynamic turn max_seqlen
symbolic, and _get_fused_attn_backend, being assume_constant_result, cannot
take a symbolic argument -- so the call died once the length changed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../attention/dot_product_attention/dot_product_attention.py  | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index d4cba79ea9..383eb3a457 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -12,6 +12,7 @@
 
 import torch
 import torch.nn.functional as F
+from torch.fx.experimental.symbolic_shapes import guard_scalar
 from torch.nn.parameter import Parameter
 
 from transformer_engine.common.recipe import (
@@ -1535,6 +1536,9 @@ def forward(
                     batch_size = query_layer.shape[0]
                     max_seqlen_q = query_layer.shape[1] if max_seqlen_q is None else max_seqlen_q
                     max_seqlen_kv = key_layer.shape[1] if max_seqlen_kv is None else max_seqlen_kv
+                # Backend selection bakes these in, which it cannot do symbolically.
+                max_seqlen_q = guard_scalar(max_seqlen_q)
+                max_seqlen_kv = guard_scalar(max_seqlen_kv)
             if qkv_format == "thd":
                 assert all(
                     len(x.shape) == 3 for x in (query_layer, key_layer, value_layer)

From 4af8c30341d0bae89467b889ea1809d478c6d082 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 16:03:17 +0200
Subject: [PATCH 36/67] Capture CUDA graphs with a KV cache

The cache is mutated in place, so inductor declined to capture until its
buffers were marked static.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py               | 5 +++--
 transformer_engine/pytorch/attention/inference.py | 6 ++++++
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index d9a5130f40..d5138ec5ac 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -972,11 +972,12 @@ def test_dpa_torch_compile_around_fused(monkeypatch):
 
 
 @pytest.mark.parametrize("backend", ["flash", "unfused"])
-def test_dpa_torch_compile_cudagraphs(monkeypatch, backend):
+@pytest.mark.parametrize("config", ["self_bshd_causal", "kv_cache_bshd"])
+def test_dpa_torch_compile_cudagraphs(monkeypatch, backend, config):
     """`mode="reduce-overhead"`: forward and backward of DotProductAttention
     are captured into CUDA graphs and replayed on subsequent iterations."""
     dtype = torch.bfloat16
-    spec = _DPA_COMPILE_CONFIGS["self_bshd_causal"]
+    spec = _DPA_COMPILE_CONFIGS[config]
     _force_dpa_backend(monkeypatch, backend)
 
     module = _make_dpa(spec, dtype)
diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py
index 0db6fdca2e..ae9e02dd7d 100644
--- a/transformer_engine/pytorch/attention/inference.py
+++ b/transformer_engine/pytorch/attention/inference.py
@@ -462,6 +462,9 @@ def allocate_memory(self, layer_number):
             dtype=self.dtype,
             device=torch.cuda.current_device(),
         )
+        # Mutated in place, so CUDA graphs are only captured if the addresses are fixed.
+        torch._dynamo.mark_static_address(k_cache)
+        torch._dynamo.mark_static_address(v_cache)
         self.cache[layer_number] = (k_cache, v_cache)
 
     def pre_step(
@@ -659,6 +662,9 @@ def allocate_memory(self, layer_number):
             dtype=self.dtype,
             device=torch.cuda.current_device(),
         )
+        # Mutated in place, so CUDA graphs are only captured if the addresses are fixed.
+        torch._dynamo.mark_static_address(k_cache)
+        torch._dynamo.mark_static_address(v_cache)
         self.cache[layer_number] = (k_cache, v_cache)
 
     def print_cache(self):

From bd94fff1c155700d5d2397a6ee3fd4f020290377 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 16:16:59 +0200
Subject: [PATCH 37/67] Test generation against a KV cache

A prefill and three single-token steps, each compared to eager, with and
without CUDA graphs: the cache carries state, so a step that updates it
wrongly only shows in the step after.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py | 74 +++++++++++++++++++++++++++++
 1 file changed, 74 insertions(+)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index d5138ec5ac..ab7cddaf54 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -1000,6 +1000,80 @@ def test_dpa_torch_compile_cudagraphs(monkeypatch, backend, config):
     assert not counters["inductor"]["cudagraph_skips"], "inductor skipped CUDA graphs"
 
 
+@pytest.mark.parametrize("backend", ["flash", "unfused"])
+@pytest.mark.parametrize("paged", [False, True], ids=["non_paged", "paged"])
+@pytest.mark.parametrize("cuda_graphs", [False, True], ids=["default", "cudagraphs"])
+def test_dpa_torch_compile_kv_cache_decoding(monkeypatch, backend, paged, cuda_graphs):
+    """Generation against a KV cache, one prefill and three single-token steps.
+
+    The cache carries state from one step to the next, so a step that updates it
+    wrongly is only visible in the step after -- hence comparing every step, and
+    a cache of its own for each of the two runs.
+    """
+    from collections import OrderedDict
+    from transformer_engine.pytorch.attention import InferenceParams
+
+    dtype = torch.bfloat16
+    spec = _DPA_COMPILE_CONFIGS["kv_cache_bshd"]
+    config = spec["model_config"]
+    b, ctx_len = config.batch_size, config.max_seqlen_q
+    h, g, d = config.num_heads, config.num_gqa_groups, config.head_dim_qk
+
+    def make_cache():
+        kwargs = dict(
+            max_batch_size=b,
+            max_sequence_length=config.max_seqlen_kv,
+            num_heads_kv=g,
+            head_dim_k=d,
+            dtype=dtype,
+            qkv_format=spec["qkv_format"],
+        )
+        if paged:
+            # One page per sequence, and FlashAttention 2 wants it divisible by 256.
+            kwargs.update(is_paged=True, page_size=config.max_seqlen_kv, total_num_pages=b)
+        inference_params = InferenceParams(**kwargs)
+        inference_params.allocate_memory(1)
+        return inference_params
+
+    _skip_unsupported(spec, backend, dtype, inference_params=make_cache())
+
+    gen = torch.Generator(device="cuda").manual_seed(1234)
+    steps = [
+        [
+            torch.randn(b, seqlen, heads, d, device="cuda", dtype=dtype, generator=gen)
+            for heads in (h, g, g)
+        ]
+        for seqlen in (ctx_len, 1, 1, 1)
+    ]
+
+    _force_dpa_backend(monkeypatch, backend)
+    module = _make_dpa(spec, dtype)
+
+    def generate(fn):
+        inference_params = make_cache()
+        outputs = []
+        with torch.no_grad():
+            for args in steps:
+                inference_params.pre_step(OrderedDict((i, args[0].shape[1]) for i in range(b)))
+                _force_dpa_backend(monkeypatch, backend)
+                # Cloned because a CUDA graph replay overwrites what it returned.
+                outputs.append(fn(*args, inference_params=inference_params).clone())
+        return outputs
+
+    eager = generate(module)
+
+    torch._dynamo.reset()
+    counters.clear()
+    compile_kwargs = {"mode": "reduce-overhead"} if cuda_graphs else {}
+    compiled = generate(torch.compile(module, fullgraph=True, **compile_kwargs))
+
+    _assert_dpa_backend(backend)
+    for out, ref in zip(compiled, eager):
+        _assert_matches_eager(out, ref, backend, dtype)
+    if cuda_graphs:
+        assert not counters["inductor"]["cudagraph_skips"], "inductor skipped CUDA graphs"
+
+
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 @pytest.mark.parametrize("fp8_attention", [False, True], ids=["fp8_gemms_only", "fp8_attention"])
 def test_dpa_torch_compile_fp8(monkeypatch, fp8_attention):

From 688e10c0ac0ecb1299feebf1e3905de9fe9a684c Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 4 Aug 2026 14:20:22 +0000
Subject: [PATCH 38/67] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 tests/pytorch/test_torch_compile.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index ab7cddaf54..5876fe857f 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -646,9 +646,7 @@ def _packed_layout(qkv_format: str, packed_dim: int, interleave_dim: int) -> str
     # A decoding step against a KV cache. FlashAttention 2 wants the non-paged
     # cache length divisible by 256.
     "kv_cache_bshd": _cfg(
-        ModelConfig(
-            2, 8, 4, 64, max_seqlen_kv=256, attn_mask_type="padding_causal_bottom_right"
-        ),
+        ModelConfig(2, 8, 4, 64, max_seqlen_kv=256, attn_mask_type="padding_causal_bottom_right"),
         kv_cache=True,
     ),
 }

From 43eba11be959c3405df2d3618985233b0f170ac0 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 4 Aug 2026 21:55:38 +0200
Subject: [PATCH 39/67] Drop the imports the custom ops made unused

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/attention/dot_product_attention/backends.py          | 1 -
 transformer_engine/pytorch/attention/inference.py                | 1 -
 2 files changed, 2 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index a5d2b05141..71afed25c5 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -15,7 +15,6 @@
 
 import torch
 import torch.nn.functional as F
-import transformer_engine_torch as tex
 from transformer_engine.pytorch.utils import (
     get_device_compute_capability,
     split_tensor_along_dim,
diff --git a/transformer_engine/pytorch/attention/inference.py b/transformer_engine/pytorch/attention/inference.py
index ae9e02dd7d..f98d1693a6 100644
--- a/transformer_engine/pytorch/attention/inference.py
+++ b/transformer_engine/pytorch/attention/inference.py
@@ -10,7 +10,6 @@
 
 import torch
 
-import transformer_engine_torch as tex
 from transformer_engine.pytorch.attention.custom_ops import (
     copy_to_kv_cache,
     QKV_FORMAT_VALUE,

From 4ff5fdfed0c1cfdb6006bb3376f2d8a958a4b4bb Mon Sep 17 00:00:00 2001
From: Zhihao Wang <101526713+xiuhu17@users.noreply.github.com>
Date: Wed, 5 Aug 2026 11:33:10 -0700
Subject: [PATCH 40/67] [common] Fix UE8M0 code 0 (2^-127) and code 255 (NaN)
 expansion in ptx::exp2f (#3262)

Fix UE8M0 code 0 and 255 expansion in ptx::exp2f

UE8M0 code 0 is 2^-127 and code 255 is NaN, but the exponent-field shift
produced +0.0 and +Inf, so MXFP8 software dequantize zeroed every 1x32
block whose scale byte was 0. Mirror the special cases already present in
exp2f_rcp and add a dequantize test with planted extreme scale codes.

Signed-off-by: zhihaow6 <zhihaow6@illinois.edu>
---
 .../test_mxfp8_dequantize_extreme_scales.py   | 28 +++++++++++++++++++
 transformer_engine/common/util/ptx.cuh        |  4 +++
 2 files changed, 32 insertions(+)
 create mode 100644 tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py

diff --git a/tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py b/tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py
new file mode 100644
index 0000000000..64bfe6012e
--- /dev/null
+++ b/tests/pytorch/mxfp8/test_mxfp8_dequantize_extreme_scales.py
@@ -0,0 +1,28 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+import pytest
+import torch
+
+import transformer_engine.pytorch as te
+from transformer_engine.pytorch import MXFP8Quantizer
+
+recipe_available, reason_for_no_recipe = te.is_mxfp8_available(return_reason=True)
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_dequantize_extreme_e8m0_scale_codes() -> None:
+    """UE8M0 scale code 0 is 2^-127 and code 255 is NaN, not 0.0 and +Inf."""
+    quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, columnwise=False)
+    x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
+    qx = quantizer(x)
+    data = qx._rowwise_data.view(torch.uint8)
+    scales = qx._rowwise_scale_inv.view(torch.uint8)
+    data[0, :32] = 56
+    scales[0, 0] = 0
+    data[1, :32] = 56
+    scales[1, 0] = 255
+    y = qx.dequantize(dtype=torch.float32)
+    assert (y[0, :32] == 2.0**-127).all()
+    assert torch.isnan(y[1, :32]).all()
diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh
index 2351ddae71..83b2daa513 100644
--- a/transformer_engine/common/util/ptx.cuh
+++ b/transformer_engine/common/util/ptx.cuh
@@ -380,6 +380,10 @@ __device__ __forceinline__ bf16 exp2f_rcp<bf16>(e8m0_t biased_exp) {
 }
 
 __device__ __forceinline__ float exp2f(e8m0_t biased_exp) {
+  // Handle the special case of NaN.
+  if (biased_exp == 255) return __int_as_float(0x7fffffff);
+  // 2^-127 is subnormal, so it cannot be built by shifting into the exponent field.
+  if (biased_exp == 0) return __int_as_float(0x00400000);
   return __int_as_float(biased_exp << FP32_MANTISSA_BITS);
 }
 

From a7aec214eb5c3969984a40c3accb6d66987d8f25 Mon Sep 17 00:00:00 2001
From: Fei Wu <33940270+YangFei1990@users.noreply.github.com>
Date: Wed, 5 Aug 2026 15:56:27 -0700
Subject: [PATCH 41/67] [Pytorch] Add function to check if EP is bootstrapped
 (#3321)

expose _BOOTSTRAPPED

Signed-off-by: YangFei1990 <feiw@nvidia.com>
---
 transformer_engine/pytorch/ep.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py
index 940812b6a4..42f605367a 100644
--- a/transformer_engine/pytorch/ep.py
+++ b/transformer_engine/pytorch/ep.py
@@ -21,6 +21,7 @@
 __all__ = [
     "EpBuffer",
     "ep_bootstrap",
+    "is_ep_bootstrapped",
     "ep_finalize",
     "ep_dispatch",
     "ep_combine",
@@ -169,6 +170,11 @@ def ep_bootstrap(
         _ATEXIT_REGISTERED = True
 
 
+def is_ep_bootstrapped() -> bool:
+    """Whether EP has been initialized in this process."""
+    return _BOOTSTRAPPED
+
+
 def ep_finalize() -> None:
     """Optional explicit EP teardown; idempotent.
 

From 66b6d25484aee41ca2d4f98079e820dd010406a0 Mon Sep 17 00:00:00 2001
From: Phuong Nguyen <phuonguyen@nvidia.com>
Date: Thu, 6 Aug 2026 08:51:53 -0700
Subject: [PATCH 42/67] [JAX] Migrate EP collective-stream annotation to
 region-based compute_on (#3304)

* Migrate EP collective-stream annotation to region-based compute_on

* Constrain EP MoE output sharding to match activation input

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---
 tests/jax/test_multi_process_ep.py          | 53 +++++++++++++++++----
 tests/jax/test_te_ep_moe.py                 |  4 +-
 transformer_engine/jax/cpp_extensions/ep.py | 16 ++++++-
 transformer_engine/jax/moe.py               |  2 +
 4 files changed, 63 insertions(+), 12 deletions(-)

diff --git a/tests/jax/test_multi_process_ep.py b/tests/jax/test_multi_process_ep.py
index 4121e5d1b4..47af0b0c39 100644
--- a/tests/jax/test_multi_process_ep.py
+++ b/tests/jax/test_multi_process_ep.py
@@ -22,6 +22,7 @@
 """
 
 import os
+import re
 import sys
 import unittest
 
@@ -689,9 +690,9 @@ def run(idx, toks, w):
         "JAX/XLA lacks the gpu_stream:collective annotation (openxla/xla#39604)",
     )
     def test_z_dispatch_combine_on_collective_stream(self):
-        """Every EP FFI custom call must carry the collective-stream annotation
-        so XLA schedules them on the collective stream instead of overlapping
-        them with other collectives."""
+        """Every EP FFI custom call must run on the collective stream. compute_on
+        puts the annotation on the async wrapper XLA generates, so assert each EP
+        call is reachable from a wrapper that carries it."""
         T_dp, tokens, topk_idx, topk_w = self._make_random_inputs()
         dp_spec = PartitionSpec(("dp", "ep"), None)
         ep_spec_3d = PartitionSpec(("dp", "ep"), None, None)
@@ -719,18 +720,50 @@ def run(idx, toks, w):
 
             hlo = run.lower(topk_idx, tokens, topk_w).compile().as_text()
 
-        # Every te_ep_* FFI custom call must carry the collective-stream
-        # annotation so XLA places it on the collective stream.
-        ep_lines = [l for l in hlo.splitlines() if 'custom_call_target="te_ep_' in l]
-        self.assertTrue(ep_lines, f"no te_ep_* custom calls in compiled HLO:\n{hlo}")
+        # Parse the HLO into computations and follow the call graph: a call
+        # carrying the collective-stream annotation places its callee (and every
+        # nested callee) on the collective stream.
+        comps = {}
+        cur = None
+        for line in hlo.splitlines():
+            stripped = line.strip()
+            header = re.match(r"(?:ENTRY\s+)?(%[\w.\-]+)\s*\(", stripped)
+            if header and stripped.endswith("{"):
+                cur = header.group(1)
+                comps[cur] = []
+            elif stripped == "}":
+                cur = None
+            elif cur is not None:
+                comps[cur].append(line)
+
+        callees = lambda l: re.findall(r"(?:calls|to_apply)=(%[\w.\-]+)", l)
+        edges = {c: {x for l in ls for x in callees(l)} for c, ls in comps.items()}
+
+        collective = set()
+        for ls in comps.values():
+            for l in ls:
+                if '_xla_stream_annotation="collective"' in l.replace(" ", ""):
+                    collective.update(callees(l))
+        stack = list(collective)
+        while stack:
+            for callee in edges.get(stack.pop(), ()):
+                if callee not in collective:
+                    collective.add(callee)
+                    stack.append(callee)
+
+        ep_calls = [
+            (c, l) for c, ls in comps.items() for l in ls if 'custom_call_target="te_ep_' in l
+        ]
+        self.assertTrue(ep_calls, f"no te_ep_* custom calls in compiled HLO:\n{hlo}")
         missing = [
             l.strip()[:200]
-            for l in ep_lines
-            if '_xla_stream_annotation="collective"' not in l.replace(" ", "")
+            for c, l in ep_calls
+            if c not in collective
+            and '_xla_stream_annotation="collective"' not in l.replace(" ", "")
         ]
         self.assertFalse(
             missing,
-            "te_ep_* custom calls missing collective-stream annotation:\n" + "\n".join(missing),
+            "te_ep_* custom calls not on the collective stream:\n" + "\n".join(missing),
         )
 
     def test_z_no_unexpected_reshard_in_hlo_bwd(self):
diff --git a/tests/jax/test_te_ep_moe.py b/tests/jax/test_te_ep_moe.py
index 55ed415d21..0e9cb6598d 100644
--- a/tests/jax/test_te_ep_moe.py
+++ b/tests/jax/test_te_ep_moe.py
@@ -140,7 +140,7 @@ def _read_mp_options():
     ("exp", EP_AXIS),
     ("embed", FSDP_AXIS),
     ("mlp", None),
-    ("batch", (EP_AXIS, FSDP_AXIS)),
+    ("batch", (FSDP_AXIS, EP_AXIS)),
 )
 
 # Small shapes so the parity tests stay tight on bf16. The block still
@@ -364,6 +364,7 @@ def _make_block(
     use_expert_routing_bias=False,
     score_function="softmax",
     expert_bias_init=None,
+    input_axes=("batch", None, None),
 ):
     kwargs = dict(
         num_experts=NUM_EXPERTS,
@@ -375,6 +376,7 @@ def _make_block(
         use_expert_routing_bias=use_expert_routing_bias,
         score_function=score_function,
         dtype=DTYPE,
+        input_axes=input_axes,
     )
     # Custom expert_bias_init lets tests inject a non-zero expert_bias without
     # poking variables['params'] post-init.
diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py
index d0d472c6f2..0b18c3bbb6 100644
--- a/transformer_engine/jax/cpp_extensions/ep.py
+++ b/transformer_engine/jax/cpp_extensions/ep.py
@@ -14,6 +14,7 @@
     compound ``(dp, ep)`` axis on the leading dim.
 """
 
+import functools
 from dataclasses import dataclass
 
 import jax
@@ -34,7 +35,20 @@ def _on_collective_stream(func):
         return func
     from jax.experimental.compute_on import compute_on
 
-    return compute_on("gpu_stream:collective")(func)  # pylint: disable=not-callable
+    @functools.wraps(func)
+    def wrapper(*args, **kwargs):
+        # compute_on traces its callee and abstract-evals every argument, so it
+        # cannot take the static EpLayerConfig/PartitionSpec args directly. Wrap
+        # a nullary thunk that closes over them; the array operands are captured
+        # as consts and lifted to real operands, outputs stay on device. XLA
+        # async-wraps the resulting call onto the collective stream.
+        annotated = compute_on(
+            compute_type="gpu_stream:collective",
+            out_memory_spaces=jax.memory.Space.Device,
+        )(lambda: func(*args, **kwargs))
+        return annotated()
+
+    return wrapper
 
 
 __all__ = [
diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py
index a49d4f80c2..21ef59ba5a 100644
--- a/transformer_engine/jax/moe.py
+++ b/transformer_engine/jax/moe.py
@@ -751,6 +751,8 @@ def _body(*args):
             num_local_tokens=(B, S),
             out_partition_spec=out_partition_spec,
         )
+    # output of MLP should be sharded the same way as the activation input
+    output = with_sharding_constraint_by_logical_axes(output, input_axes)
 
     (
         casted_sorted_x_lhs_trans,

From 46e0edf9fd525dd0e89a4f9fbd96d45c9895ba06 Mon Sep 17 00:00:00 2001
From: vthumbe1503 <vthumbe@nvidia.com>
Date: Thu, 6 Aug 2026 15:14:05 -0700
Subject: [PATCH 43/67] [PyTorch]Fused ScaledActivation Kernel in TE Ops and
 reducing CPU overheads for MOE (#3238)

* initial implementation for act grouped linear fusion

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* address review comments + cleanup

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* address review comment + cleanup

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fix CI + handle activation requires_grad=false and fusion optimizations for precomputed tensor offsets in backward

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* precomputed tensor offsets in grouped linear as well

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* fix lint

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* unecsaary based on op infra

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* better name

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* activation+group_quantize fusion via ops infra

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

ugly solution for dbias fusion

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* revert to have grouped_linear + activation fusion for now

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* get rid of fused ops now

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* improve comments

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

---------

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 tests/pytorch/test_grouped_mlp.py             |  19 ++-
 transformer_engine/pytorch/csrc/extensions.h  |  25 +++
 .../pytorch/csrc/extensions/activation.cpp    | 145 +++++++++++++++++
 .../pytorch/csrc/extensions/pybind.cpp        |  21 +++
 .../pytorch/module/grouped_linear.py          |  40 +++--
 .../pytorch/ops/basic/activation.py           |  89 +++++++---
 .../pytorch/ops/basic/grouped_linear.py       | 124 ++++++++------
 .../pytorch/ops/basic/swiglu.py               | 153 +++++++++---------
 .../pytorch/ops/fused/grouped_mlp.py          |  57 +++++--
 9 files changed, 505 insertions(+), 168 deletions(-)

diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py
index d195eb2f78..28f7360641 100644
--- a/tests/pytorch/test_grouped_mlp.py
+++ b/tests/pytorch/test_grouped_mlp.py
@@ -755,6 +755,13 @@ def test_grouped_mlp(
         maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype)
         if dtype == torch.bfloat16 and not is_bf16_available():
             pytest.skip("BF16 requires SM 8.0+")
+        if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and (
+            single_grouped_weight or single_grouped_bias
+        ):
+            pytest.skip(
+                "single_grouped_weight/single_grouped_bias requires"
+                " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1"
+            )
         if single_grouped_weight and quantization != "mxfp8":
             pytest.skip("single_grouped_weight is only supported for MXFP8 quantization")
         if single_grouped_bias and not bias:
@@ -1026,8 +1033,10 @@ def _make_module():
             or (
                 quantization == "nvfp4_rht"
                 and dtype == torch.bfloat16
-                and activation == "scaled_srelu"
-                and glu_interleave_size is None
+                and (
+                    (not activation_is_glu and glu_interleave_size is None)
+                    or (activation_is_glu and glu_interleave_size == 32)
+                )
             )
         )
         if expected_grouped_mlp_fusion:
@@ -1266,6 +1275,8 @@ def test_grouped_mlp_single_weight_numerics(
     ) -> None:
         """single_grouped_weight=True/False should match exactly for fused MXFP8 grouped MLP."""
 
+        if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0":
+            pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1")
         if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
             pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
 
@@ -1584,6 +1595,8 @@ def test_grouped_mlp_overwrite_main_grad(
         that read ``.grad`` don't see stale bytes from the cached dummy).
         """
 
+        if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and single_grouped_weight:
+            pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1")
         if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
             pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
 
@@ -1715,6 +1728,8 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8(
     ) -> None:
         """Grouped MLP forward+backward should be CUDA graph capturable (MXFP8)."""
 
+        if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and single_grouped_weight:
+            pytest.skip("single_grouped_weight requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1")
         if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
             pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
         if dtype not in (torch.bfloat16, torch.float16):
diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index 95f02c64f0..8248a63680 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -290,6 +290,31 @@ py::object clamped_swiglu(const at::Tensor &input, py::handle quantizer, float l
 
 py::object clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer,
                            float limit, float alpha, float glu_linear_offset);
+
+/* Scaled activation */
+py::object scaled_swiglu(const at::Tensor &input, const at::Tensor &act_scales,
+                         py::handle quantizer, int64_t glu_interleave_size);
+
+py::object scaled_clamped_swiglu(const at::Tensor &input, const at::Tensor &act_scales,
+                                 py::handle quantizer, float limit, float alpha,
+                                 float glu_linear_offset, int64_t glu_interleave_size);
+
+py::object scaled_srelu(const at::Tensor &input, const at::Tensor &act_scales,
+                        py::handle quantizer);
+
+py::tuple scaled_dswiglu(const at::Tensor &grad, const at::Tensor &input,
+                         const at::Tensor &act_scales, py::handle quantizer,
+                         int64_t glu_interleave_size, bool compute_scale_grad);
+
+py::tuple scaled_clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input,
+                                 const at::Tensor &act_scales, py::handle quantizer, float limit,
+                                 float alpha, float glu_linear_offset, int64_t glu_interleave_size,
+                                 bool compute_scale_grad);
+
+py::tuple scaled_dsrelu(const at::Tensor &grad, const at::Tensor &input,
+                        const at::Tensor &act_scales, py::handle quantizer,
+                        bool compute_scale_grad);
+
 /***************************************************************************************************
  * LayerNorm
  **************************************************************************************************/
diff --git a/transformer_engine/pytorch/csrc/extensions/activation.cpp b/transformer_engine/pytorch/csrc/extensions/activation.cpp
index 58a8f84f85..544ff92c1b 100644
--- a/transformer_engine/pytorch/csrc/extensions/activation.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp
@@ -342,5 +342,150 @@ py::object clamped_dswiglu(const at::Tensor& grad, const at::Tensor& input, py::
                                                               glu_linear_offset);
 }
 
+/* Scaled activation helpers (activation + per-row scale via nvte_scaled_*). */
+
+template <auto act_func, typename... Args>
+at::Tensor scaled_activation_compute(const at::Tensor& input, const at::Tensor& act_scales,
+                                     int shape_divisor, Args&&... args) {
+  init_extension();
+  NVTE_CHECK(input.dim() >= 1, "scaled activation input must have at least 1 dimension");
+  NVTE_CHECK(shape_divisor > 0 && input.size(-1) % shape_divisor == 0,
+             "scaled activation input width is not compatible with activation");
+
+  auto input_tensor = input.contiguous();
+  auto scales_tensor = act_scales.contiguous().reshape({-1});
+  const int64_t rows = input_tensor.numel() / input_tensor.size(-1);
+  NVTE_CHECK(scales_tensor.numel() == rows, "scaled activation expects one scale per input row");
+
+  std::vector<int64_t> output_sizes(input_tensor.sizes().begin(), input_tensor.sizes().end());
+  output_sizes.back() /= shape_divisor;
+  auto output = at::empty(output_sizes, input_tensor.options());
+
+  const TensorWrapper& input_nvte = makeTransformerEngineTensor(input_tensor);
+  const TensorWrapper& scales_nvte = makeTransformerEngineTensor(scales_tensor);
+  const TensorWrapper& output_nvte = makeTransformerEngineTensor(output);
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+  NVTE_SCOPED_GIL_RELEASE({
+    act_func(input_nvte.data(), scales_nvte.data(), output_nvte.data(), std::forward<Args>(args)...,
+             stream);
+  });
+  return output;
+}
+
+template <auto dact_func, typename... Args>
+std::tuple<at::Tensor, at::Tensor> scaled_dactivation_compute(const at::Tensor& grad,
+                                                              const at::Tensor& input,
+                                                              const at::Tensor& act_scales,
+                                                              bool compute_scale_grad,
+                                                              Args&&... args) {
+  init_extension();
+  NVTE_CHECK(input.dim() >= 1 && grad.dim() >= 1,
+             "scaled dactivation input and grad must have at least 1 dimension");
+
+  auto grad_tensor = grad.contiguous();
+  auto input_tensor = input.contiguous();
+  auto scales_tensor = act_scales.contiguous();
+  const int64_t rows = input_tensor.numel() / input_tensor.size(-1);
+  NVTE_CHECK(scales_tensor.numel() == rows, "scaled dactivation expects one scale per input row");
+
+  auto scales_flat = scales_tensor.reshape({-1});
+  auto grad_input = at::empty_like(input_tensor);
+  auto grad_scales = compute_scale_grad ? at::empty_like(scales_tensor) : at::Tensor();
+  auto grad_scales_flat = compute_scale_grad ? grad_scales.reshape({-1}) : at::Tensor();
+
+  const TensorWrapper& grad_nvte = makeTransformerEngineTensor(grad_tensor);
+  const TensorWrapper& input_nvte = makeTransformerEngineTensor(input_tensor);
+  const TensorWrapper& scales_nvte = makeTransformerEngineTensor(scales_flat);
+  const TensorWrapper& grad_input_nvte = makeTransformerEngineTensor(grad_input);
+  std::optional<TensorWrapper> grad_scales_nvte;
+  if (compute_scale_grad) {
+    grad_scales_nvte.emplace(makeTransformerEngineTensor(grad_scales_flat));
+  }
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+  NVTE_SCOPED_GIL_RELEASE({
+    dact_func(grad_nvte.data(), input_nvte.data(), scales_nvte.data(), grad_input_nvte.data(),
+              compute_scale_grad ? grad_scales_nvte->data() : nullptr, std::forward<Args>(args)...,
+              stream);
+  });
+  return {grad_input, grad_scales};
+}
+
+py::object maybe_quantize(const at::Tensor& tensor, py::handle quantizer) {
+  if (quantizer.is_none()) {
+    return py::cast(tensor);
+  }
+  auto quantizer_cpp = convert_quantizer(quantizer);
+  const TensorWrapper& tensor_nvte = makeTransformerEngineTensor(tensor);
+  const auto shape_te = tensor_nvte.shape();
+  const std::vector<size_t> shape(shape_te.data, shape_te.data + shape_te.ndim);
+  auto fake_dtype = GetTransformerEngineDType(tensor.scalar_type());
+  auto [out_nvte, out_py] = quantizer_cpp->create_tensor(shape, fake_dtype);
+  quantizer_cpp->quantize(tensor_nvte, out_nvte);
+  return out_py;
+}
+
+template <auto act_func, typename... Args>
+py::object scaled_activation_helper(const at::Tensor& input, const at::Tensor& act_scales,
+                                    py::handle quantizer, int shape_divisor, Args&&... args) {
+  auto output = scaled_activation_compute<act_func>(input, act_scales, shape_divisor,
+                                                    std::forward<Args>(args)...);
+  return maybe_quantize(output, quantizer);
+}
+
+template <auto dact_func, typename... Args>
+py::tuple scaled_dactivation_helper(const at::Tensor& grad, const at::Tensor& input,
+                                    const at::Tensor& act_scales, py::handle quantizer,
+                                    bool compute_scale_grad, Args&&... args) {
+  auto [grad_input, grad_scales] = scaled_dactivation_compute<dact_func>(
+      grad, input, act_scales, compute_scale_grad, std::forward<Args>(args)...);
+  return py::make_tuple(maybe_quantize(grad_input, quantizer),
+                        compute_scale_grad ? py::cast(grad_scales) : py::none());
+}
+
+py::object scaled_swiglu(const at::Tensor& input, const at::Tensor& act_scales,
+                         py::handle quantizer, int64_t glu_interleave_size) {
+  return scaled_activation_helper<nvte_scaled_swiglu>(input, act_scales, quantizer,
+                                                      /*shape_divisor=*/2, glu_interleave_size);
+}
+
+py::object scaled_clamped_swiglu(const at::Tensor& input, const at::Tensor& act_scales,
+                                 py::handle quantizer, float limit, float alpha,
+                                 float glu_linear_offset, int64_t glu_interleave_size) {
+  return scaled_activation_helper<nvte_scaled_clamped_swiglu>(
+      input, act_scales, quantizer, /*shape_divisor=*/2, limit, alpha, glu_linear_offset,
+      glu_interleave_size);
+}
+
+py::object scaled_srelu(const at::Tensor& input, const at::Tensor& act_scales,
+                        py::handle quantizer) {
+  return scaled_activation_helper<nvte_scaled_srelu>(input, act_scales, quantizer,
+                                                     /*shape_divisor=*/1);
+}
+
+py::tuple scaled_dswiglu(const at::Tensor& grad, const at::Tensor& input,
+                         const at::Tensor& act_scales, py::handle quantizer,
+                         int64_t glu_interleave_size, bool compute_scale_grad) {
+  return scaled_dactivation_helper<nvte_scaled_dswiglu>(grad, input, act_scales, quantizer,
+                                                        compute_scale_grad, glu_interleave_size);
+}
+
+py::tuple scaled_clamped_dswiglu(const at::Tensor& grad, const at::Tensor& input,
+                                 const at::Tensor& act_scales, py::handle quantizer, float limit,
+                                 float alpha, float glu_linear_offset, int64_t glu_interleave_size,
+                                 bool compute_scale_grad) {
+  return scaled_dactivation_helper<nvte_scaled_clamped_dswiglu>(
+      grad, input, act_scales, quantizer, compute_scale_grad, limit, alpha, glu_linear_offset,
+      glu_interleave_size);
+}
+
+py::tuple scaled_dsrelu(const at::Tensor& grad, const at::Tensor& input,
+                        const at::Tensor& act_scales, py::handle quantizer,
+                        bool compute_scale_grad) {
+  return scaled_dactivation_helper<nvte_scaled_dsrelu>(grad, input, act_scales, quantizer,
+                                                       compute_scale_grad);
+}
+
 }  // namespace pytorch
 }  // namespace transformer_engine
diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
index c6bf7eb516..e406dc3446 100644
--- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
@@ -285,6 +285,27 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
         "Backward of SwiGLU used in GPT OSS", py::arg("grad"), py::arg("fwd_input"),
         py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f,
         py::arg("glu_linear_offset") = 1.0f);
+  /* Scaled activation */
+  m.def("scaled_swiglu", transformer_engine::pytorch::scaled_swiglu, "Scaled SwiGLU activation",
+        py::arg("input"), py::arg("act_scales"), py::arg("quantizer"),
+        py::arg("glu_interleave_size") = 0);
+  m.def("scaled_clamped_swiglu", transformer_engine::pytorch::scaled_clamped_swiglu,
+        "Scaled clamped SwiGLU activation", py::arg("input"), py::arg("act_scales"),
+        py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f,
+        py::arg("glu_linear_offset") = 1.0f, py::arg("glu_interleave_size") = 0);
+  m.def("scaled_srelu", transformer_engine::pytorch::scaled_srelu, "Scaled SReLU activation",
+        py::arg("input"), py::arg("act_scales"), py::arg("quantizer"));
+  m.def("scaled_dswiglu", transformer_engine::pytorch::scaled_dswiglu, "Scaled SwiGLU backward",
+        py::arg("grad"), py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"),
+        py::arg("glu_interleave_size") = 0, py::arg("compute_scale_grad") = true);
+  m.def("scaled_clamped_dswiglu", transformer_engine::pytorch::scaled_clamped_dswiglu,
+        "Scaled clamped SwiGLU backward", py::arg("grad"), py::arg("fwd_input"),
+        py::arg("act_scales"), py::arg("quantizer"), py::arg("limit") = 7.0f,
+        py::arg("alpha") = 1.702f, py::arg("glu_linear_offset") = 1.0f,
+        py::arg("glu_interleave_size") = 0, py::arg("compute_scale_grad") = true);
+  m.def("scaled_dsrelu", transformer_engine::pytorch::scaled_dsrelu, "Scaled SReLU backward",
+        py::arg("grad"), py::arg("fwd_input"), py::arg("act_scales"), py::arg("quantizer"),
+        py::arg("compute_scale_grad") = true);
   /* DBias + DAct fusions*/
   m.def("dbias_dgelu", transformer_engine::pytorch::dbias_dgelu, "DGeLU + DBias + Quantize",
         py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer"));
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 12497d0cb0..9860d48237 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -561,7 +561,7 @@ def _make_grouped_tensor(
         *,
         num_gemms: int,
         split_sizes: torch.Tensor,
-        base_split_offsets: torch.Tensor,
+        tensor_offsets: torch.Tensor,
         last_dim: int,
         dtype: torch.dtype,
     ) -> GroupedTensorStorage:
@@ -573,7 +573,7 @@ def _make_grouped_tensor(
             quantizer=None,
             data=data.reshape(-1),
             first_dims=split_sizes,
-            tensor_offsets=base_split_offsets * last_dim,
+            tensor_offsets=tensor_offsets,
         )
 
     @staticmethod
@@ -704,8 +704,18 @@ def _forward_grouped_tensor(
         out_features = weights[0].size(0)
         weight_requires_grad = weights[0].requires_grad
 
-        split_sizes = m_splits.to(device=device)
-        base_split_offsets = tex.splits_to_offsets(split_sizes, 1)
+        split_sizes, (
+            base_split_offsets,
+            input_tensor_offsets,
+            output_tensor_offsets,
+        ) = tex.splits_to_offsets_multi(
+            m_splits,
+            device,
+            strides=[1, in_features, out_features],
+            include_leading_zero=[True, True, True],
+            dtypes=[torch.int64, torch.int64, torch.int64],
+            bulk_allocate=True,
+        )
 
         inp_view = inp.reshape(-1, in_features)
         x = cast_if_needed(inp_view, activation_dtype)
@@ -716,13 +726,19 @@ def _forward_grouped_tensor(
                 columnwise=is_grad_enabled and weight_requires_grad,
             )
             input_quantizer.optimize_for_gemm = True
-            grouped_x = tex.group_quantize(x, input_quantizer, num_gemms, split_sizes)
+            grouped_x = tex.group_quantize(
+                x,
+                input_quantizer,
+                num_gemms,
+                split_sizes,
+                tensor_offsets=input_tensor_offsets,
+            )
         else:
             grouped_x = _GroupedLinear._make_grouped_tensor(
                 x,
                 num_gemms=num_gemms,
                 split_sizes=split_sizes,
-                base_split_offsets=base_split_offsets,
+                tensor_offsets=input_tensor_offsets,
                 last_dim=in_features,
                 dtype=activation_dtype,
             )
@@ -751,7 +767,7 @@ def _forward_grouped_tensor(
             out,
             num_gemms=num_gemms,
             split_sizes=split_sizes,
-            base_split_offsets=base_split_offsets,
+            tensor_offsets=output_tensor_offsets,
             last_dim=out_features,
             dtype=activation_dtype,
         )
@@ -796,6 +812,8 @@ def _forward_grouped_tensor(
                 *weights_to_save,
                 split_sizes,
                 base_split_offsets,
+                input_tensor_offsets,
+                output_tensor_offsets,
             )
             ctx.save_for_backward(*tensors_to_save)
             ctx.tensor_objects = tensor_objects
@@ -1217,6 +1235,8 @@ def _backward_grouped_tensor(
         weights = saved_tensors[1 : 1 + N]
         split_sizes = saved_tensors[1 + N]
         base_split_offsets = saved_tensors[2 + N]
+        input_tensor_offsets = saved_tensors[3 + N]
+        output_tensor_offsets = saved_tensors[4 + N]
 
         origin_weights = [None] * N
         main_grads = [None] * N
@@ -1253,6 +1273,7 @@ def _backward_grouped_tensor(
                     grad_output_quantizer,
                     N,
                     split_sizes,
+                    tensor_offsets=output_tensor_offsets,
                 )
             else:
                 grouped_dy = tex.group_quantize(
@@ -1260,13 +1281,14 @@ def _backward_grouped_tensor(
                     grad_output_quantizer,
                     N,
                     split_sizes,
+                    tensor_offsets=output_tensor_offsets,
                 )
         else:
             grouped_dy = _GroupedLinear._make_grouped_tensor(
                 dy_2d,
                 num_gemms=N,
                 split_sizes=split_sizes,
-                base_split_offsets=base_split_offsets,
+                tensor_offsets=output_tensor_offsets,
                 last_dim=ctx.weights_shape_0,
                 dtype=ctx.activation_dtype,
             )
@@ -1298,7 +1320,7 @@ def _backward_grouped_tensor(
                 dgrad,
                 num_gemms=N,
                 split_sizes=split_sizes,
-                base_split_offsets=base_split_offsets,
+                tensor_offsets=input_tensor_offsets,
                 last_dim=ctx.weights_shape_1,
                 dtype=ctx.activation_dtype,
             )
diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py
index f4beffe90c..8137051322 100644
--- a/transformer_engine/pytorch/ops/basic/activation.py
+++ b/transformer_engine/pytorch/ops/basic/activation.py
@@ -348,18 +348,8 @@ def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor:
         return tex.dsrelu(*args, **kwargs)
 
 
-class ScaledSReLU(BasicOperation):
-    r"""Squared ReLU with per-row post-scaling.
-
-    If the SReLU output has shape ``(d_1, ..., d_n)``, it is multiplied
-    with an extra input tensor of shape ``(d_1, ..., d_{n-1})``.
-
-    Parameters
-    ----------
-    activation_recompute_in_mlp : bool, default = ``False``
-        Enable fused grouped MLP kernels to recompute activation outputs
-        during backward when supported instead of saving them.
-    """
+class _ScaledUnary(BasicOperation, metaclass=abc.ABCMeta):
+    """Unary activation with per-row scales (fused grouped MLP middle op)."""
 
     num_extra_inputs: int = 1
 
@@ -367,6 +357,25 @@ def __init__(self, *, activation_recompute_in_mlp: bool = False) -> None:
         super().__init__()
         self.activation_recompute_in_mlp: bool = activation_recompute_in_mlp
 
+    @abc.abstractmethod
+    def _scaled_unary_forward(
+        self,
+        input_: torch.Tensor,
+        scales: torch.Tensor,
+    ) -> torch.Tensor:
+        """Apply the scaled unary activation."""
+
+    @abc.abstractmethod
+    def _scaled_unary_backward(
+        self,
+        grad_output: torch.Tensor,
+        input_: torch.Tensor,
+        scales: torch.Tensor,
+        *,
+        compute_scale_grad: bool,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Apply the scaled unary activation backward pass."""
+
     def op_forward(self, *args, **kwargs) -> None:
         raise RuntimeError(
             f"{self.__class__.__name__} operation has "
@@ -410,7 +419,7 @@ def fuser_forward(
 
         x = maybe_dequantize(input_.contiguous(), dtype)
         scales = maybe_dequantize(extra_input, dtype)
-        y = tex.srelu(x, None) * scales.unsqueeze(-1)
+        y = self._scaled_unary_forward(x, scales)
 
         ctx = basic_op_ctxs[0]
         if ctx.requires_grad:
@@ -448,21 +457,57 @@ def fuser_backward(
         scales = maybe_dequantize(scales, ctx.dtype)
         grad_output = maybe_dequantize(grad_output.contiguous(), ctx.dtype)
 
-        grad_input = None
-        if ctx.input_requires_grad:
-            grad_srelu_out = grad_output * scales.unsqueeze(-1)
-            grad_input = tex.dsrelu(grad_srelu_out, x, None)
-
-        grad_extra_input = None
-        if ctx.extra_input_requires_grad:
-            srelu_out = tex.srelu(x, None)
-            grad_extra_input = torch.linalg.vecdot(srelu_out, grad_output)
+        grad_input, grad_extra_input = self._scaled_unary_backward(
+            grad_output,
+            x,
+            scales,
+            compute_scale_grad=ctx.extra_input_requires_grad,
+        )
+        if not ctx.input_requires_grad:
+            grad_input = None
 
         clear_tensor_data(ctx.saved_tensors[0])
 
         return grad_input, [()], [(grad_extra_input,)]
 
 
+class ScaledSReLU(_ScaledUnary):
+    r"""Squared ReLU with per-row post-scaling.
+
+    If the SReLU output has shape ``(d_1, ..., d_n)``, it is multiplied
+    with an extra input tensor of shape ``(d_1, ..., d_{n-1})``.
+
+    Parameters
+    ----------
+    activation_recompute_in_mlp : bool, default = ``False``
+        Enable fused grouped MLP kernels to recompute activation outputs
+        during backward when supported instead of saving them.
+    """
+
+    def _scaled_unary_forward(
+        self,
+        input_: torch.Tensor,
+        scales: torch.Tensor,
+    ) -> torch.Tensor:
+        return tex.scaled_srelu(input_, scales, None)
+
+    def _scaled_unary_backward(
+        self,
+        grad_output: torch.Tensor,
+        input_: torch.Tensor,
+        scales: torch.Tensor,
+        *,
+        compute_scale_grad: bool,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        return tex.scaled_dsrelu(
+            grad_output,
+            input_,
+            scales,
+            None,
+            compute_scale_grad,
+        )
+
+
 class SReGLU(_ActivationOperation):
     r"""Squared Rectified Gated Linear Unit
 
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index be931829ea..c94d78fecd 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -1136,14 +1136,17 @@ def fuser_forward_save_ctx(
         # temporary workspaces freshly created in each forward pass.
         if is_cpu_offload_enabled():
             saved = tensors_to_save[0]
-            offset = 4 if self._scale_bias else 3
+            # Metadata prefix:
+            # [split_sizes, base_split_offsets, split_points,
+            #  input_tensor_offsets, output_tensor_offsets, (scales?)]
+            offset = 6 if self._scale_bias else 5
             if use_grouped_tensor_path:
-                # Layout: [split_sizes, base_split_offsets, split_points, (scales?), grouped_x, *weights]
+                # Layout: [..., grouped_x, *weights]
                 grouped_x = saved[offset]
                 if grouped_x is not None:
                     mark_activation_offload(grouped_x)
             else:
-                # Layout: [split_sizes, None, None, (scales?), *xs, *ws]
+                # Layout: [..., *xs, *ws]
                 live_xs = [t for t in saved[offset : offset + self.num_groups] if t is not None]
                 if live_xs:
                     mark_activation_offload(*live_xs)
@@ -1169,9 +1172,9 @@ def fuser_forward_save_ctx(
         ctx.weight_quantizers = weight_quantizers
         ctx.grad_output_quantizers = grad_output_quantizers
         ctx.grad_input_quantizers = None
-        # ``split_sizes`` and ``base_split_offsets`` are routed through
-        # ``save_for_backward`` (see ``_fuser_forward_split_quantize`` and
-        # ``_fuser_forward_grouped_tensor`` for the saved-tensor layout).
+        # ``split_sizes``, offset metadata, and related tensors are routed
+        # through ``save_for_backward`` (see ``_fuser_forward_split_quantize``
+        # and ``_fuser_forward_grouped_tensor`` for the saved-tensor layout).
         if torch.is_autocast_enabled():
             ctx.dtype = torch.get_autocast_dtype("cuda")
         else:
@@ -1285,12 +1288,12 @@ def _fuser_forward_split_quantize(
 
         # Build the tuple of tensors to save for backward. Layout:
         #   [split_sizes, base_split_offsets, split_points,
+        #    input_tensor_offsets, output_tensor_offsets,
         #    (scales if scale_bias), *xs, *ws]
-        # ``base_split_offsets`` and ``split_points`` are unused on the
-        # split-quantize backward path but are included as ``None`` so the
-        # saved-tensor layout matches the graph-safe
-        # ``_fuser_forward_grouped_tensor`` path (and the fused MLP forward).
-        saved: list[Optional[torch.Tensor]] = [split_sizes, None, None]
+        # Offset metadata slots are unused on the split-quantize backward path
+        # but are included as ``None`` so the saved-tensor layout matches the
+        # graph-safe ``_fuser_forward_grouped_tensor`` path.
+        saved: list[Optional[torch.Tensor]] = [split_sizes, None, None, None, None]
         if self._scale_bias:
             saved.append(scales)
         saved.extend(xs)
@@ -1312,28 +1315,36 @@ def _fuser_forward_grouped_tensor(
         device: torch.device,
         out_buffer: Optional[torch.Tensor] = None,
     ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]:
-        """Graph-safe GroupedTensor forward path (pure compute).
-        Returns ``(output, tensors_to_save)``. ``split_sizes``,
-        ``base_split_offsets`` and ``split_points`` are returned so that
-        ``fuser_forward_save_ctx`` can call ``save_for_backward`` on them.
-        """
+        """Build graph-safe grouped input storage and run grouped GEMM."""
         num_groups = self.num_groups
-        has_bias = self.has_bias
-
-        base_split_offsets = tex.splits_to_offsets(split_sizes, 1)
-        split_points = base_split_offsets[1:].to(dtype=torch.int)
-
-        # Flatten to 2D so the first dim is the total token count.
+        split_sizes, grouped_tensor_offsets = tex.splits_to_offsets_multi(
+            split_sizes,
+            device,
+            strides=[1, 1, self.in_features, self.out_features],
+            include_leading_zero=[False, True, True, True],
+            dtypes=[torch.int32, torch.int64, torch.int64, torch.int64],
+            bulk_allocate=True,
+        )
+        split_points = grouped_tensor_offsets[0]
+        base_split_offsets = grouped_tensor_offsets[1]
+        input_tensor_offsets = grouped_tensor_offsets[2]
+        output_tensor_offsets = grouped_tensor_offsets[3]
         original_shape = list(input_.size())
+        # Flatten to 2D so the first dim is the total token count.
         x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features)
         total_tokens = x.size(0)
-
-        # Build the input GroupedTensor.
+        # Build the input GroupedTensorStorage for input.
         if with_quantized_compute:
             input_quantizer = input_quantizers[0]
             input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad)
             input_quantizer.optimize_for_gemm = True
-            grouped_x = tex.group_quantize(x, input_quantizer, num_groups, split_sizes)
+            grouped_x = tex.group_quantize(
+                x,
+                input_quantizer,
+                num_groups,
+                split_sizes,
+                tensor_offsets=input_tensor_offsets,
+            )
         else:
             # No quantize: wrap the contiguous high-precision buffer.
             grouped_x = GroupedTensorStorage(
@@ -1343,8 +1354,9 @@ def _fuser_forward_grouped_tensor(
                 quantizer=None,
                 data=x.reshape(-1),
                 first_dims=split_sizes,
-                tensor_offsets=base_split_offsets * self.in_features,
+                tensor_offsets=input_tensor_offsets,
             )
+        has_bias = self.has_bias
 
         if is_cpu_offload_enabled() and grouped_x is not None:
             start_offload(grouped_x)
@@ -1379,7 +1391,7 @@ def _fuser_forward_grouped_tensor(
             quantizer=None,
             data=out.reshape(-1),
             first_dims=split_sizes,
-            tensor_offsets=base_split_offsets * self.out_features,
+            tensor_offsets=output_tensor_offsets,
         )
 
         # Bias: hand off to the grouped GEMM (graph-safe, fused). Plain bias
@@ -1414,14 +1426,23 @@ def _fuser_forward_grouped_tensor(
 
         # Build the tuple of tensors to save for backward. Layout:
         #   [split_sizes, base_split_offsets, split_points,
+        #    input_tensor_offsets, output_tensor_offsets,
         #    (scales if _scale_bias), grouped_x, *weights]
+        # ``output_tensor_offsets`` matches the linear output row layout and is
+        # reused as ``grad_output`` offsets in backward.
         if grouped_x is not None:
             # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data
             # in backward pass)
             if with_quantized_compute and grouped_x.columnwise_data is not None:
                 grouped_x.rowwise_data = None
                 grouped_x.scale_inv = None
-        saved: list[Optional[torch.Tensor]] = [split_sizes, base_split_offsets, split_points]
+        saved: list[Optional[torch.Tensor]] = [
+            split_sizes,
+            base_split_offsets,
+            split_points,
+            input_tensor_offsets,
+            output_tensor_offsets,
+        ]
         if self._scale_bias:
             saved.append(scales)
         saved.append(grouped_x)
@@ -1471,13 +1492,13 @@ def _fuser_backward_split_quantize(
 
         # Saved tensors from forward pass. Layout:
         #   [split_sizes, base_split_offsets, split_points,
+        #    input_tensor_offsets, output_tensor_offsets,
         #    (scales if _scale_bias), *xs, *ws]
-        # ``base_split_offsets`` and ``split_points`` are unused on this path
-        # but are present so the saved-tensor layout matches the graph-safe
-        # path (and the fused MLP forward).
+        # Offset metadata beyond ``split_sizes`` is unused on this path but is
+        # present so the saved-tensor layout matches the graph-safe path.
         saved_tensors = ctx.saved_tensors
         split_sizes = saved_tensors[0]
-        saved_tensors = saved_tensors[3:]
+        saved_tensors = saved_tensors[5:]
         scales = None
         if self._scale_bias:
             scales, saved_tensors = saved_tensors[0], saved_tensors[1:]
@@ -1653,24 +1674,24 @@ def _fuser_backward_grouped_tensor(
         Iterable[Iterable[Optional[torch.Tensor]]],
         Iterable[Iterable[Optional[torch.Tensor]]],
     ]:
+        """Graph-safe GroupedTensor backward path."""
         num_groups = self.num_groups
         has_bias = self.has_bias
         weights, is_dist_weight, dist_dgrad_weights = self._backward_weight_setup()
         device = weights[0].device
         dtype = ctx.dtype
-
         with_quantized_compute = bool(getattr(ctx, "with_quantized_compute", False))
 
-        # Saved tensors from forward pass
-        # Layout: [split_sizes, base_split_offsets, split_points,
-        #          (scales if _scale_bias), grouped_x, *weights]
-        # ``split_points`` is unused on this path but is present so the
-        # saved-tensor layout matches the fused MLP forward (which needs it
-        # for the cuDNN grouped GEMM kernel).
+        # Saved tensors from forward pass. Layout:
+        # [split_sizes, base_split_offsets, split_points,
+        #  input_tensor_offsets, output_tensor_offsets,
+        #  (scales if _scale_bias), grouped_x, *weights]
         saved_tensors = ctx.saved_tensors
         split_sizes = saved_tensors[0]
         base_split_offsets = saved_tensors[1]
-        saved_tensors = saved_tensors[3:]
+        input_tensor_offsets = saved_tensors[3]
+        output_tensor_offsets = saved_tensors[4]
+        saved_tensors = saved_tensors[5:]
         scales = None
         if self._scale_bias:
             scales, saved_tensors = saved_tensors[0], saved_tensors[1:]
@@ -1684,6 +1705,7 @@ def _fuser_backward_grouped_tensor(
         # to figure out total tokens.
         dy_2d = grad_output.reshape(-1, self.out_features)
         total_tokens = dy_2d.size(0)
+        grad_input_shape = list(grad_output.size())[:-1] + [self.in_features]
 
         # Build the grad_output GroupedTensor.
         # Optionally get dbias is fusion available with bgrad_group_quantize
@@ -1691,10 +1713,10 @@ def _fuser_backward_grouped_tensor(
         if with_quantized_compute:
             grad_output_quantizer = ctx.grad_output_quantizers[0]
             grad_output_quantizer.set_usage(
-                rowwise=ctx.input_requires_grad, columnwise=ctx.weight_requires_grad
+                rowwise=ctx.input_requires_grad,
+                columnwise=ctx.weight_requires_grad,
             )
             grad_output_quantizer.optimize_for_gemm = True
-
             # FP8 block scaling computes dbias in the rowwise (dgrad) pass, so only fuse
             # when dgrad is required.
             fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or (
@@ -1702,15 +1724,22 @@ def _fuser_backward_grouped_tensor(
             )
             if has_bias and not self._scale_bias and fuse_bgrad:
                 grouped_dy, dbias_packed = tex.bgrad_group_quantize(
-                    dy_2d, grad_output_quantizer, num_groups, split_sizes
+                    dy_2d,
+                    grad_output_quantizer,
+                    num_groups,
+                    split_sizes,
+                    tensor_offsets=output_tensor_offsets,
                 )
             else:
                 grouped_dy = tex.group_quantize(
-                    dy_2d, grad_output_quantizer, num_groups, split_sizes
+                    dy_2d,
+                    grad_output_quantizer,
+                    num_groups,
+                    split_sizes,
+                    tensor_offsets=output_tensor_offsets,
                 )
         else:
             dy_2d = maybe_dequantize(dy_2d, dtype)
-            # Wrap BF16/FP16 buffer as a GroupedTensor for grouped gemm
             grouped_dy = GroupedTensorStorage(
                 shape=(total_tokens, self.out_features),
                 dtype=dtype,
@@ -1718,7 +1747,7 @@ def _fuser_backward_grouped_tensor(
                 quantizer=None,
                 data=dy_2d.reshape(-1),
                 first_dims=split_sizes,
-                tensor_offsets=base_split_offsets * self.out_features,
+                tensor_offsets=output_tensor_offsets,
             )
 
         # Bias Grads compute if not already computed in bgrad_group_quantize
@@ -1745,7 +1774,6 @@ def _fuser_backward_grouped_tensor(
         # ---- dgrad GEMM ----------------------------------------------------
         grad_input = None
         if ctx.input_requires_grad:
-            grad_input_shape = list(grad_output.size())[:-1] + [self.in_features]
             grad_input = validate_or_alloc_output(
                 getattr(ctx, "dgrad_out", None), grad_input_shape, dtype, device
             )
@@ -1756,7 +1784,7 @@ def _fuser_backward_grouped_tensor(
                 quantizer=None,
                 data=grad_input.reshape(-1),
                 first_dims=split_sizes,
-                tensor_offsets=base_split_offsets * self.in_features,
+                tensor_offsets=input_tensor_offsets,
             )
             general_grouped_gemm_for_grouped_tensor(
                 dist_dgrad_weights if is_dist_weight else ws,
diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py
index 02f330ede3..fb663c0480 100644
--- a/transformer_engine/pytorch/ops/basic/swiglu.py
+++ b/transformer_engine/pytorch/ops/basic/swiglu.py
@@ -387,14 +387,21 @@ def __init__(
         self.glu_interleave_size: Optional[int] = glu_interleave_size
         self.activation_recompute_in_mlp: bool = activation_recompute_in_mlp
 
-    def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor:
+    def _scaled_glu_forward(
+        self,
+        input_: torch.Tensor,
+        scales: torch.Tensor,
+    ) -> torch.Tensor:
         raise NotImplementedError
 
-    def _glu_backward(
+    def _scaled_glu_backward(
         self,
-        grad_swiglu_out: torch.Tensor,
-        swiglu_in: torch.Tensor,
-    ) -> torch.Tensor:
+        grad_output: torch.Tensor,
+        input_: torch.Tensor,
+        scales: torch.Tensor,
+        *,
+        compute_scale_grad: bool,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         raise NotImplementedError
 
     def op_forward(self, *args, **kwargs) -> None:
@@ -442,22 +449,7 @@ def fuser_forward(
         # Make sure inputs are in correct dtype
         input_ = maybe_dequantize(input_, dtype)
         scales = maybe_dequantize(extra_input, dtype)
-
-        # Remove gate interleaving if needed
-        swiglu_in = input_
-        if self.glu_interleave_size is not None:
-            shape = swiglu_in.size()
-            swiglu_in = swiglu_in.reshape(
-                -1,
-                shape[-1] // (2 * self.glu_interleave_size),
-                2,
-                self.glu_interleave_size,
-            )
-            swiglu_in = swiglu_in.transpose(1, 2).contiguous()
-            swiglu_in = swiglu_in.view(shape)
-
-        swiglu_out = self._glu_forward(swiglu_in)
-        out = swiglu_out * scales.unsqueeze(-1)
+        out = self._scaled_glu_forward(input_, scales)
 
         # Save state for backward pass
         ctx = basic_op_ctxs[0]
@@ -469,7 +461,7 @@ def fuser_forward(
             ctx.dtype = dtype
             ctx.save_for_backward(
                 input_,
-                scales if ctx.input_requires_grad else None,
+                scales if ctx.input_requires_grad or ctx.extra_input_requires_grad else None,
             )
 
         return out, [()]
@@ -498,41 +490,14 @@ def fuser_backward(
             scales = maybe_dequantize(scales, ctx.dtype)
         grad_output = maybe_dequantize(grad_output, ctx.dtype)
 
-        # Remove gate interleaving if needed
-        swiglu_in = input_
-        if self.glu_interleave_size is not None:
-            shape = swiglu_in.size()
-            swiglu_in = swiglu_in.reshape(
-                -1,
-                shape[-1] // (2 * self.glu_interleave_size),
-                2,
-                self.glu_interleave_size,
-            )
-            swiglu_in = swiglu_in.transpose(1, 2).contiguous()
-            swiglu_in = swiglu_in.view(shape)
-
-        # Compute input grad
-        grad_input = None
-        if ctx.input_requires_grad:
-            grad_swiglu_out = grad_output * scales.unsqueeze(-1)
-            grad_swiglu_in = self._glu_backward(grad_swiglu_out, swiglu_in)
-            grad_input = grad_swiglu_in
-            if self.glu_interleave_size is not None:
-                shape = grad_input.size()
-                grad_input = grad_input.reshape(
-                    -1,
-                    2,
-                    shape[-1] // (2 * self.glu_interleave_size),
-                    self.glu_interleave_size,
-                )
-                grad_input = grad_input.transpose(1, 2).contiguous()
-                grad_input = grad_input.view(shape)
-
-        # Compute scales grad by recomputing GLU
-        grad_extra_input = None
-        if ctx.extra_input_requires_grad:
-            swiglu_out = self._glu_forward(swiglu_in)
-            grad_extra_input = torch.linalg.vecdot(swiglu_out, grad_output)
+        grad_input, grad_extra_input = self._scaled_glu_backward(
+            grad_output,
+            input_,
+            scales,
+            compute_scale_grad=ctx.extra_input_requires_grad,
+        )
+        if not ctx.input_requires_grad:
+            grad_input = None
 
         # Clear input tensor if possible
         clear_tensor_data(ctx.saved_tensors[0])  # input_
@@ -558,15 +523,34 @@ class ScaledSwiGLU(_ScaledGLU):
 
     """
 
-    def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor:
-        return tex.swiglu(swiglu_in, None)
-
-    def _glu_backward(
+    def _scaled_glu_forward(
         self,
-        grad_swiglu_out: torch.Tensor,
-        swiglu_in: torch.Tensor,
+        input_: torch.Tensor,
+        scales: torch.Tensor,
     ) -> torch.Tensor:
-        return tex.dswiglu(grad_swiglu_out, swiglu_in, None)
+        return tex.scaled_swiglu(
+            input_,
+            scales,
+            None,
+            int(self.glu_interleave_size or 0),
+        )
+
+    def _scaled_glu_backward(
+        self,
+        grad_output: torch.Tensor,
+        input_: torch.Tensor,
+        scales: torch.Tensor,
+        *,
+        compute_scale_grad: bool,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        return tex.scaled_dswiglu(
+            grad_output,
+            input_,
+            scales,
+            None,
+            int(self.glu_interleave_size or 0),
+            compute_scale_grad,
+        )
 
 
 class ScaledClampedQGeGLU(_ScaledGLU):
@@ -614,16 +598,39 @@ def __init__(
             glu_linear_offset=glu_linear_offset,
         )
 
-    def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor:
-        return self._clamped._tex_clamped_swiglu_forward(swiglu_in, None)
-
-    def _glu_backward(
+    def _scaled_glu_forward(
         self,
-        grad_swiglu_out: torch.Tensor,
-        swiglu_in: torch.Tensor,
+        input_: torch.Tensor,
+        scales: torch.Tensor,
     ) -> torch.Tensor:
-        return self._clamped._tex_clamped_dswiglu(
-            grad_swiglu_out,
-            swiglu_in,
+        clamped = self._clamped
+        return tex.scaled_clamped_swiglu(
+            input_,
+            scales,
+            None,
+            clamped.limit,
+            clamped.alpha,
+            clamped.glu_linear_offset,
+            int(self.glu_interleave_size or 0),
+        )
+
+    def _scaled_glu_backward(
+        self,
+        grad_output: torch.Tensor,
+        input_: torch.Tensor,
+        scales: torch.Tensor,
+        *,
+        compute_scale_grad: bool,
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        clamped = self._clamped
+        return tex.scaled_clamped_dswiglu(
+            grad_output,
+            input_,
+            scales,
             None,
+            clamped.limit,
+            clamped.alpha,
+            clamped.glu_linear_offset,
+            int(self.glu_interleave_size or 0),
+            compute_scale_grad,
         )
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 0113833647..db2c4e7674 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -1091,21 +1091,40 @@ def fuser_forward(
             # shared-expert path never consumes it.
             base_split_offsets = split_sizes
             fc1_x_tensor_offsets = None
+            fc1_out_tensor_offsets = None
             fc2_x_tensor_offsets = None
             fc2_out_tensor_offsets = None
         else:
+            # Bulk-allocate every grouped-tensor offset the forward and backward
+            # passes need, so the backward can reuse them from the context
+            # instead of recomputing offsets per GEMM.
             split_sizes, (
                 split_points,
                 base_split_offsets,
                 fc1_x_tensor_offsets,
+                fc1_out_tensor_offsets,
                 fc2_x_tensor_offsets,
                 fc2_out_tensor_offsets,
             ) = tex.splits_to_offsets_multi(
                 split_sizes,
                 device,
-                strides=[1, 1, fc1_weight_shape[1], fc2_weight_shape[1], fc2_weight_shape[0]],
-                include_leading_zero=[False, True, True, True, True],
-                dtypes=[torch.int32, torch.int64, torch.int64, torch.int64, torch.int64],
+                strides=[
+                    1,
+                    1,
+                    fc1_weight_shape[1],
+                    fc1_weight_shape[0],
+                    fc2_weight_shape[1],
+                    fc2_weight_shape[0],
+                ],
+                include_leading_zero=[False, True, True, True, True, True],
+                dtypes=[
+                    torch.int32,
+                    torch.int64,
+                    torch.int64,
+                    torch.int64,
+                    torch.int64,
+                    torch.int64,
+                ],
                 bulk_allocate=True,
             )
 
@@ -1743,6 +1762,10 @@ def fuser_forward(
                 split_sizes,
                 base_split_offsets,
                 split_points,
+                fc1_x_tensor_offsets,
+                fc1_out_tensor_offsets,
+                fc2_x_tensor_offsets,
+                fc2_out_tensor_offsets,
                 saved_fc1_x,
                 *fc1_weight_tensors,
                 activation_in,
@@ -1796,12 +1819,22 @@ def fuser_backward(
 
         # Saved tensors from the joint forward.
         # Layout: [split_sizes, base_split_offsets, split_points,
+        #          fc1_x_tensor_offsets, fc1_out_tensor_offsets,
+        #          fc2_x_tensor_offsets, fc2_out_tensor_offsets,
         #          grouped_fc1_x, *fc1_weights,
         #          activation_in, scales,
         #          grouped_fc2_x, *fc2_weights]
         saved_tensors = fc1_ctx.saved_tensors
-        split_sizes, base_split_offsets, split_points = saved_tensors[:3]
-        saved_tensors = saved_tensors[3:]
+        (
+            split_sizes,
+            base_split_offsets,
+            split_points,
+            fc1_x_tensor_offsets,
+            fc1_out_tensor_offsets,
+            fc2_x_tensor_offsets,
+            fc2_out_tensor_offsets,
+        ) = saved_tensors[:7]
+        saved_tensors = saved_tensors[7:]
         grouped_fc1_x, saved_tensors = saved_tensors[0], saved_tensors[1:]
         if fc1_op.single_grouped_weight:
             grouped_fc1_weight, saved_tensors = saved_tensors[0], saved_tensors[1:]
@@ -1868,6 +1901,7 @@ def fuser_backward(
                     fc2_grad_output_quantizer,
                     num_groups,
                     split_sizes,
+                    tensor_offsets=fc2_out_tensor_offsets,
                 )
             else:
                 grouped_fc2_dy = _group_quantize_for_grouped_mlp(
@@ -1875,9 +1909,7 @@ def fuser_backward(
                     fc2_grad_output_quantizer,
                     num_groups,
                     split_sizes,
-                    tensor_offsets=(
-                        None if num_groups == 1 else base_split_offsets * fc2_weight_shape[0]
-                    ),
+                    tensor_offsets=fc2_out_tensor_offsets,
                 )
 
         use_nvfp4 = (
@@ -2150,7 +2182,7 @@ def fuser_backward(
                     fc2_input_quantizer,
                     num_groups,
                     split_sizes,
-                    tensor_offsets=base_split_offsets * fc2_weight_shape[1],
+                    tensor_offsets=fc2_x_tensor_offsets,
                 )
             else:
                 sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor")
@@ -2172,7 +2204,7 @@ def fuser_backward(
                     scale_inv=None,
                     columnwise_scale_inv=fc2_x_col_scale.reshape(-1),
                     first_dims=split_sizes,
-                    tensor_offsets=base_split_offsets * fc2_weight_shape[1],
+                    tensor_offsets=fc2_x_tensor_offsets,
                     with_gemm_swizzled_scales=True,
                 )
 
@@ -2215,9 +2247,7 @@ def fuser_backward(
                     fc1_bias_grads = [dbias_2d[group_idx] for group_idx in range(num_groups)]
 
         # FC1 grad output for dgrad and wgrad GEMMs
-        fc1_dy_tensor_offsets = (
-            None if num_groups == 1 else base_split_offsets * fc1_weight_shape[0]
-        )
+        fc1_dy_tensor_offsets = fc1_out_tensor_offsets
         fc1_grad_output_quantizer = fc1_ctx.grad_output_quantizers[0]
         if use_nvfp4:
             fc1_grad_output_quantizer.set_usage(
@@ -2304,7 +2334,6 @@ def fuser_backward(
                 )
             elif use_nvfp4:
                 grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device)
-                fc1_x_tensor_offsets = base_split_offsets * fc1_weight_shape[1]
                 grouped_grad_input = GroupedTensor(
                     shape=(out_shape[0], fc1_weight_shape[1]),
                     dtype=dtype,

From dd9659563642296a00d5079558b2ead038a717f1 Mon Sep 17 00:00:00 2001
From: Chase Block <cblock@nvidia.com>
Date: Thu, 6 Aug 2026 19:25:50 -0500
Subject: [PATCH 44/67] Add support for fused Q Up-Proj GEMM/RoPE/Quant.
 (#3303)

* Add support for fused Q Up-Proj GEMM/RoPE/Quant.

This commit add support for fusing the GEMM in the Q Up Proj
step of DeepseekV3 training with the following RoPE and MXFP8
quantization operations. This uses a custom kernel from cudnn_frontend,
and supports both 16-bit projection and mxfp8 projection.

Signed-off-by: Chase Block <cblock@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove unused function from fused mla q uproj, add error handling

Signed-off-by: Chase Block <cblock@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Adjust comment SBHD/BSHD

Signed-off-by: Chase Block <cblock@nvidia.com>

* Add cudnn_frontend version check; eliminate code dup in
mxfp8_quantize_fast_path.

Signed-off-by: Chase Block <cblock@nvidia.com>

* Add missing docstrings

Signed-off-by: Chase Block <cblock@nvidia.com>

* Add unit test for q up proj fusion

Signed-off-by: Chase Block <cblock@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Add test for bwd path of fused gemm+rope+quant.

Signed-off-by: Chase Block <cblock@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Consolidate fused q proj tests and add to unit test list

Signed-off-by: Chase Block <cblock@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Remove backward test from fused mla q uproj.

These tests really belong in Megatron.

Signed-off-by: Chase Block <cblock@nvidia.com>

---------

Signed-off-by: Chase Block <cblock@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Siddhartha Raman Sundara Raman <sraman@nvidia.com>
Co-authored-by: Sudhakar Singh <sudhakars@nvidia.com>
---
 qa/L0_pytorch_unittest/test.sh                |   1 +
 .../attention/test_fused_mla_q_uproj.py       | 137 ++++++++++++++
 transformer_engine/pytorch/__init__.py        |   1 +
 .../pytorch/attention/__init__.py             |   2 +
 .../dot_product_attention/backends.py         |  14 +-
 .../dot_product_attention.py                  |  23 +++
 .../attention/dot_product_attention/utils.py  | 161 +++++++++-------
 .../pytorch/attention/fused_mla_q_uproj.py    | 175 ++++++++++++++++++
 8 files changed, 447 insertions(+), 67 deletions(-)
 create mode 100644 tests/pytorch/attention/test_fused_mla_q_uproj.py
 create mode 100644 transformer_engine/pytorch/attention/fused_mla_q_uproj.py

diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh
index 6a41515602..759432857a 100644
--- a/qa/L0_pytorch_unittest/test.sh
+++ b/qa/L0_pytorch_unittest/test.sh
@@ -60,6 +60,7 @@ NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest --tb=auto --junitxml=$X
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_flex_attention.xml $TE_PATH/tests/pytorch/attention/test_flex_attention.py || test_fail "test_flex_attention.py"
 NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic.xml $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_linear_mxfp8_attention.xml $TE_PATH/tests/pytorch/attention/test_linear_mxfp8_attention.py || test_fail "test_linear_mxfp8_attention.py"
+python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_mla_q_uproj.xml $TE_PATH/tests/pytorch/attention/test_fused_mla_q_uproj.py || test_fail "test_fused_mla_q_uproj.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_kv_cache.xml $TE_PATH/tests/pytorch/attention/test_kv_cache.py || test_fail "test_kv_cache.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cu_seqlens_cache.xml $TE_PATH/tests/pytorch/attention/test_cu_seqlens_cache.py || test_fail "test_cu_seqlens_cache.py"
 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_hf_integration.xml $TE_PATH/tests/pytorch/test_hf_integration.py || test_fail "test_hf_integration.py"
diff --git a/tests/pytorch/attention/test_fused_mla_q_uproj.py b/tests/pytorch/attention/test_fused_mla_q_uproj.py
new file mode 100644
index 0000000000..2c061607fc
--- /dev/null
+++ b/tests/pytorch/attention/test_fused_mla_q_uproj.py
@@ -0,0 +1,137 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Unit tests for FusedMLAQUpProjRopeQuant.
+
+Run:
+    pytest tests/pytorch/attention/test_fused_mla_q_uproj.py -v
+"""
+
+import pytest
+import torch
+
+import transformer_engine.pytorch  # registers transformer_engine_torch
+import transformer_engine_torch as tex
+from transformer_engine.pytorch.attention import FusedMLAQUpProjRopeQuant
+from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor
+
+# DSv3 671B MLA dims
+NUM_HEADS = 128
+HEAD_DIM_NOPE = 128
+HEAD_DIM_ROPE = 64
+HEAD_DIM = HEAD_DIM_NOPE + HEAD_DIM_ROPE  # 192
+Q_LORA_RANK = 1536
+PROJ_DIM = NUM_HEADS * HEAD_DIM  # 24576
+
+SEED = 42
+
+fused_supported, reason_not_supported = (
+    (True, "")
+    if FusedMLAQUpProjRopeQuant.is_supported()
+    else (
+        False,
+        (
+            "FusedMLAQUpProjRopeQuant.is_supported() returned False "
+            "(SM100+, cudnn-frontend >= 1.27.0, and NVTE_FUSED_MLA_Q_UPROJ=1 required)"
+        ),
+    )
+)
+
+
+def _dequantize_fused_output(query: MXFP8Tensor, s: int, b: int) -> torch.Tensor:
+    """Dequantize the rowwise fused output to bf16 [s, b, nh, head_dim].
+
+    TE's C++ dequantize kernel requires 2D layout, so reshape before calling dequantize().
+    """
+    tokens = s * b
+    q_2d = MXFP8Tensor(
+        shape=(tokens, PROJ_DIM),
+        dtype=torch.bfloat16,
+        rowwise_data=query._rowwise_data.view(tokens, PROJ_DIM),
+        rowwise_scale_inv=query._rowwise_scale_inv.view(tokens, PROJ_DIM // 32),
+        columnwise_data=None,
+        columnwise_scale_inv=None,
+        quantizer=query._quantizer,
+        requires_grad=False,
+        fp8_dtype=query._fp8_dtype,
+        with_gemm_swizzled_scales=False,
+    )
+    return q_2d.dequantize().to(torch.bfloat16).view(s, b, NUM_HEADS, HEAD_DIM)
+
+
+def _reference_q_uproj(
+    x: torch.Tensor,
+    w_mxfp8: MXFP8Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    s: int,
+    b: int,
+) -> torch.Tensor:
+    """Unfused bf16 reference: dequantize-then-GEMM + RoPE. Returns [s, b, nh, head_dim] bf16."""
+    x_dq = (
+        MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False)(x)
+        .dequantize()
+        .to(torch.bfloat16)
+    )
+    w_dq = w_mxfp8.dequantize().to(torch.bfloat16)
+    out = (x_dq @ w_dq.t()).view(s, b, NUM_HEADS, HEAD_DIM)
+
+    q_nope = out[..., :HEAD_DIM_NOPE]
+    q_rope = out[..., HEAD_DIM_NOPE:]
+    cos_ = cos[:, None, None, :].to(q_rope.dtype)
+    sin_ = sin[:, None, None, :].to(q_rope.dtype)
+    half = HEAD_DIM_ROPE // 2
+    x1, x2 = q_rope[..., 0::2], q_rope[..., 1::2]
+    q_rope_out = torch.cat(
+        [
+            x1 * cos_[..., :half] - x2 * sin_[..., :half],
+            x2 * cos_[..., half:] + x1 * sin_[..., half:],
+        ],
+        dim=-1,
+    )
+    return torch.cat([q_nope, q_rope_out], dim=-1)
+
+
+def _build_rope_tables(tokens: int, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
+    inv_freq = 1.0 / (
+        10000
+        ** (torch.arange(0, HEAD_DIM_ROPE, 2, dtype=torch.float32, device=device) / HEAD_DIM_ROPE)
+    )
+    freqs = torch.cat(
+        [torch.outer(torch.arange(tokens, device=device, dtype=torch.float32), inv_freq)] * 2,
+        dim=-1,
+    )
+    return freqs.cos().to(torch.bfloat16), freqs.sin().to(torch.bfloat16)
+
+
+@pytest.mark.skipif(not fused_supported, reason=reason_not_supported)
+@pytest.mark.parametrize("tokens", [256])
+def test_fused_mla_q_uproj(tokens: int) -> None:
+    """Forward numerics and x_saved properties for FusedMLAQUpProjRopeQuant.run().
+
+    Full forward+backward autograd testing (via _FusedMLAQUpProjFunction) lives in
+    Megatron-Core.
+    """
+    s, b = tokens, 1
+    device = torch.device("cuda")
+    torch.manual_seed(SEED)
+    torch.cuda.manual_seed(SEED)
+
+    x = torch.randn(tokens, Q_LORA_RANK, dtype=torch.bfloat16, device=device)
+    w = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False)(
+        torch.randn(PROJ_DIM, Q_LORA_RANK, dtype=torch.bfloat16, device=device)
+    )
+    cos, sin = _build_rope_tables(tokens, device)
+
+    query, x_saved = FusedMLAQUpProjRopeQuant.run(x, w, cos, sin, s, b)
+
+    # Forward numerics: FP8 GEMM + output quantize introduce ~10% relative error.
+    fused_dq = _dequantize_fused_output(query, s, b)
+    ref_dq = _reference_q_uproj(x, w, cos, sin, s, b)
+    torch.testing.assert_close(fused_dq, ref_dq, atol=0.5, rtol=0.1)
+
+    # x_saved: must be MXFP8 with only columnwise data retained for wgrad.
+    assert isinstance(x_saved, MXFP8Tensor)
+    assert x_saved._columnwise_data is not None, "x_saved must retain columnwise data for wgrad"
+    assert x_saved._rowwise_data is None, "x_saved rowwise data should be dropped after forward"
diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py
index 4c6b7fc67a..2b1803bfb2 100644
--- a/transformer_engine/pytorch/__init__.py
+++ b/transformer_engine/pytorch/__init__.py
@@ -29,6 +29,7 @@
 from transformer_engine.pytorch.module import destroy_ub
 from transformer_engine.pytorch.module import UserBufferQuantizationMode
 from transformer_engine.pytorch.attention import DotProductAttention
+from transformer_engine.pytorch.attention import FusedMLAQUpProjRopeQuant
 from transformer_engine.pytorch.attention import MultiheadAttention
 from transformer_engine.pytorch.attention import InferenceParams
 from transformer_engine.pytorch.attention import RotaryPositionEmbedding
diff --git a/transformer_engine/pytorch/attention/__init__.py b/transformer_engine/pytorch/attention/__init__.py
index c4c2aa3e72..f6e4f0b37f 100644
--- a/transformer_engine/pytorch/attention/__init__.py
+++ b/transformer_engine/pytorch/attention/__init__.py
@@ -5,12 +5,14 @@
 """Python interface for attention"""
 
 from .dot_product_attention import DotProductAttention
+from .fused_mla_q_uproj import FusedMLAQUpProjRopeQuant
 from .multi_head_attention import MultiheadAttention
 from .inference import InferenceParams
 from .rope import RotaryPositionEmbedding
 
 __all__ = [
     "DotProductAttention",
+    "FusedMLAQUpProjRopeQuant",
     "MultiheadAttention",
     "InferenceParams",
     "RotaryPositionEmbedding",
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 891a5c661d..aeac501ff8 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -1363,6 +1363,7 @@ def forward(
         deterministic,
         softmax_offset,
         fp8_output,
+        bf16_backward,
         layer_number,
         return_max_logit,
         packed_qkv=None,
@@ -1427,6 +1428,9 @@ def forward(
             #                      fp8_dtype = tex.DType.kFloat8E4M3
             if is_input_fp8:
                 q_fp8, k_fp8, v_fp8 = q, k, v
+
+                if fp8_recipe.mxfp8():
+                    qkv_scale_inv_format = "bhsd"  # Same as what combine_and_quantize would give
             else:
                 q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format = combine_and_quantize(
                     qkv_layout,
@@ -1602,6 +1606,8 @@ def forward(
 
         ctx.is_input_fp8 = is_input_fp8
         ctx.is_output_fp8 = is_output_fp8
+        # Return dQ/dK/dV in bf16 even if is_input_fp8
+        ctx.bf16_backward = bf16_backward
 
         tensors_to_save, tensor_objects = prepare_for_saving(
             *fp8_tensors,
@@ -1860,7 +1866,8 @@ def backward(ctx, d_out, *_args):
                     # dq, dk, dv:             torch.Tensor; dtype = torch.float16 or torch.bfloat16
                     dq, dk, dv = dq_, dk_, dv_
                     is_quantized_tensor = isinstance(dq_, QuantizedTensorStorage)
-                    if is_quantized_tensor and not ctx.is_input_fp8:
+
+                    if is_quantized_tensor and (not ctx.is_input_fp8 or ctx.bf16_backward):
                         # return in F16
                         dq, dk, dv = combine_and_dequantize(
                             ctx.dqkv_layout,
@@ -1869,7 +1876,7 @@ def backward(ctx, d_out, *_args):
                             dv_,
                             src_nominal_dtype=dq_.dtype,
                         )
-                    if not is_quantized_tensor and ctx.is_input_fp8:
+                    if not is_quantized_tensor and ctx.is_input_fp8 and not ctx.bf16_backward:
                         # return in FP8
                         dq, dk, dv, _, _ = combine_and_quantize(
                             ctx.dqkv_layout, dq_, dk_, dv_, ctx.dQKV_quantizer
@@ -1968,6 +1975,7 @@ def backward(ctx, d_out, *_args):
             None,
             None,  # packed_qkv
             None,  # packed_kv
+            None,
         )
 
 
@@ -2064,6 +2072,7 @@ def forward(
         score_mod_bprop_tensors: Optional[Dict[str, torch.Tensor]] = None,
         packed_qkv: Optional[torch.Tensor] = None,
         packed_kv: Optional[torch.Tensor] = None,
+        bf16_backward: bool = False,
     ) -> torch.Tensor:
         """fused attention fprop"""
         assert (
@@ -2280,6 +2289,7 @@ def forward(
                     self.deterministic,
                     softmax_offset,
                     fp8_output,
+                    bf16_backward,
                     self.layer_number,
                     self.return_max_logit,
                     packed_qkv,
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 11d6500ab7..e6f50eb0da 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -32,6 +32,7 @@
     Float8BlockScalingRecipeState,
 )
 from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage
+from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage
 from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
 from transformer_engine.pytorch.export import is_in_onnx_export_mode
 from transformer_engine.pytorch.constants import AttnMaskTypes, AttnTypes, dist_group_type, DType
@@ -1347,6 +1348,7 @@ def forward(
         inference_params: Optional[InferenceParams] = None,
         pad_between_seqs: Optional[bool] = None,
         fp8_output: Optional[bool] = False,
+        bf16_backward: Optional[bool] = False,
         num_splits: Optional[int] = 1,
         score_mod: Optional[Callable] = None,
         score_mod_bprop: Optional[Callable] = None,
@@ -1817,6 +1819,25 @@ def forward(
                     qkv_format=qkv_format,
                     inference_params=inference_params,
                 )
+            elif all(
+                isinstance(x, MXFP8TensorStorage) for x in [query_layer, key_layer, value_layer]
+            ):
+                # Pre-quantized MXFP8 q/k/v: the wrapper has no real storage, so run
+                # layout detection on the underlying rowwise data (mirrors the Float8 path).
+                (
+                    qkv_layout,
+                    query_layer._rowwise_data,
+                    key_layer._rowwise_data,
+                    value_layer._rowwise_data,
+                    q_format,
+                    kv_format,
+                ) = dpa_utils.get_qkv_layout(
+                    query_layer._rowwise_data,
+                    key_layer._rowwise_data,
+                    value_layer._rowwise_data,
+                    qkv_format=qkv_format,
+                    inference_params=inference_params,
+                )
             else:
                 (
                     qkv_layout,
@@ -2190,6 +2211,7 @@ def forward(
                         fp8_output=fp8_output,
                         packed_qkv=qkv_layer,
                         packed_kv=kv_layer,
+                        bf16_backward=bf16_backward,
                     )
                 return self.fused_attention(
                     query_layer,
@@ -2227,6 +2249,7 @@ def forward(
                     score_mod_bprop_tensors=score_mod_bprop_tensors,
                     packed_qkv=qkv_layer,
                     packed_kv=kv_layer,
+                    bf16_backward=bf16_backward,
                 )
 
             if use_unfused_attention:
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index d99ee4c8e8..6eb3ce54f1 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -2786,10 +2786,37 @@ def mxfp8_quantize_fast_path(tensor_quantizer_pairs, src_format):
     """
     if not tensor_quantizer_pairs:
         return [], src_format
+
+    fp8_tensors = mxfp8_quantize_only(tensor_quantizer_pairs, src_format)
+    mxfp8_transpose_swizzle(fp8_tensors, src_format)
+    return fp8_tensors, "bhsd"
+
+
+def mxfp8_quantize_only(tensor_quantizer_pairs, src_format):
+    """Phase 1 of mxfp8_quantize_fast_path: quantize only, no BHSD transpose or GEMM swizzle.
+
+    Returns MXFP8Tensors with data and scale_invs reshaped to src_format layout.
+    Call mxfp8_transpose_swizzle to complete the BHSD permute + swizzle when ready
+    (e.g. after pre-quantized tensors from fused kernels are also available).
+
+    Parameters
+    ----------
+    tensor_quantizer_pairs : list of (torch.Tensor, MXFP8Quantizer)
+        Same contract as mxfp8_quantize_fast_path.
+    src_format : str
+        ``"bshd"`` or ``"sbhd"``.
+
+    Returns
+    -------
+    fp8_tensors : list of MXFP8Tensor
+        Data and scale_invs in src_format layout; NOT yet BHSD-permuted or swizzled.
+    """
+    if not tensor_quantizer_pairs:
+        return []
     assert src_format in (
         "bshd",
         "sbhd",
-    ), f"mxfp8_quantize_fast_path only supports bshd/sbhd, got {src_format!r}."
+    ), f"mxfp8_quantize_only only supports bshd/sbhd, got {src_format!r}."
     _s_dim = {"bshd": 1, "sbhd": 0}
     _d_dim = {"bshd": 3, "sbhd": 3}
 
@@ -2800,45 +2827,74 @@ def mxfp8_quantize_fast_path(tensor_quantizer_pairs, src_format):
         rs_shape[_d_dim[src_format]] //= MXFP8_BLOCK_SCALING_SIZE
         cs_shape = list(original_shape)
         cs_shape[_s_dim[src_format]] //= MXFP8_BLOCK_SCALING_SIZE
-
-        # view tensor as 2D for quantization
-        # BSHD -> (B*S, H*D)
-        # SBHD -> (S, B*H*D)
         if src_format == "bshd":
-            tensor = tensor.view(*tensor.shape[:2], -1)
+            t2d = tensor.view(*tensor.shape[:2], -1)
         else:
-            tensor = tensor.view(tensor.shape[0], -1)
-
-        # quantize
+            t2d = tensor.view(tensor.shape[0], -1)
         orig_optimize = quantizer.optimize_for_gemm
         quantizer.optimize_for_gemm = False
-        fp8_tensor = quantizer(tensor)
+        fp8_2d = quantizer(t2d)
         quantizer.optimize_for_gemm = orig_optimize
-
-        # reshape rowwise/columnwise data to original shape
-        fp8_tensor._rowwise_data = (
-            fp8_tensor._rowwise_data.view(original_shape)
-            if fp8_tensor._rowwise_data is not None
-            else None
-        )
-        fp8_tensor._columnwise_data = (
-            fp8_tensor._columnwise_data.view(original_shape)
-            if fp8_tensor._columnwise_data is not None
-            else None
-        )
-        fp8_tensor._rowwise_scale_inv = (
-            fp8_tensor._rowwise_scale_inv.view(rs_shape)
-            if fp8_tensor._rowwise_scale_inv is not None
-            else None
-        )
-        fp8_tensor._columnwise_scale_inv = (
-            fp8_tensor._columnwise_scale_inv.view(cs_shape)
-            if fp8_tensor._columnwise_scale_inv is not None
-            else None
+        # Re-wrap with the original 4D SBHD/BSHD shape so that shape[-1] equals the per-head
+        # dimension (matching Q's wrapper shape) and fused_attn_bwd produces 4D dkv that
+        # matches key/value's expected gradient shape in _KFQuantizeKVForAttn.backward.
+        fp8_t = MXFP8Tensor(
+            shape=original_shape,
+            dtype=tensor.dtype,
+            rowwise_data=(
+                fp8_2d._rowwise_data.view(original_shape)
+                if fp8_2d._rowwise_data is not None
+                else None
+            ),
+            rowwise_scale_inv=(
+                fp8_2d._rowwise_scale_inv.view(rs_shape)
+                if fp8_2d._rowwise_scale_inv is not None
+                else None
+            ),
+            columnwise_data=(
+                fp8_2d._columnwise_data.view(original_shape)
+                if fp8_2d._columnwise_data is not None
+                else None
+            ),
+            columnwise_scale_inv=(
+                fp8_2d._columnwise_scale_inv.view(cs_shape)
+                if fp8_2d._columnwise_scale_inv is not None
+                else None
+            ),
+            quantizer=quantizer,
+            requires_grad=False,
+            fp8_dtype=fp8_2d._fp8_dtype,
+            with_gemm_swizzled_scales=False,
         )
-        fp8_tensors.append(fp8_tensor)
+        fp8_tensors.append(fp8_t)
+    return fp8_tensors
+
+
+def mxfp8_transpose_swizzle(fp8_tensors, src_format):
+    """Phase 2 of mxfp8_quantize_fast_path: batched BHSD-transpose + GEMM-swizzle.
+
+    For tensors whose data is already quantized (e.g. from a fused GEMM+quant kernel
+    or from mxfp8_quantize_only), permutes each tensor's scale_invs from src_format to
+    BHSD and applies the GEMM swizzle in-place.  Complements mxfp8_quantize_only to
+    allow pre-quantized tensors (like a fused-kernel Q) to be processed in the same
+    batched operation as freshly quantized K/V.
+
+    Parameters
+    ----------
+    fp8_tensors : list of MXFP8Tensor
+        Tensors with _rowwise_scale_inv / _columnwise_scale_inv in src_format layout.
+        Modified in-place: scale_invs are replaced with BHSD-permuted, swizzled versions.
+    src_format : str
+        ``"bshd"`` or ``"sbhd"``.
+    """
+    if not fp8_tensors:
+        return
+
+    assert src_format in (
+        "bshd",
+        "sbhd",
+    ), f"mxfp8_transpose_swizzle only supports bshd/sbhd, got {src_format!r}."
 
-    # ---- Pad + permute + swizzle scale_inv to BHSD ----
     rs_list = [t._rowwise_scale_inv for t in fp8_tensors]
     cs_list = [t._columnwise_scale_inv for t in fp8_tensors]
 
@@ -2872,50 +2928,25 @@ def _build_outputs(scale_list, alignment):
         buf = torch.empty(total, dtype=torch.uint8, device=device)
         return [buf[e[0] : e[0] + e[1]].view(e[2]) if e is not None else None for e in entries]
 
-    # allocate buffers with padding in mind
     rs_outs = _build_outputs(rs_list, 4)
     cs_outs = _build_outputs(cs_list, 128)
 
-    # permute scale_invs to BHSD; batched
     rs_permuted = tex.multi_tensor_transpose_to_bhsd(
-        rs_list,
-        original_format=src_format,
-        outputs=rs_outs,
+        rs_list, original_format=src_format, outputs=rs_outs
     )
     cs_permuted = tex.multi_tensor_transpose_to_bhsd(
-        cs_list,
-        original_format=src_format,
-        outputs=cs_outs,
+        cs_list, original_format=src_format, outputs=cs_outs
     )
 
-    # build output tensors
-    result = []
     for t, rp, cp in zip(fp8_tensors, rs_permuted, cs_permuted):
-        rp = rp.view(-1, rp.shape[-1]) if rp is not None else None
-        cp = cp.view(-1, cp.shape[-1]) if cp is not None else None
-        result.append(
-            MXFP8Tensor(
-                shape=t.shape,
-                dtype=t.dtype,
-                rowwise_data=t._rowwise_data,
-                rowwise_scale_inv=rp,
-                columnwise_data=t._columnwise_data,
-                columnwise_scale_inv=cp,
-                quantizer=t._quantizer,
-                requires_grad=False,
-                fp8_dtype=t._fp8_dtype,
-                with_gemm_swizzled_scales=t._with_gemm_swizzled_scales,
-            )
-        )
+        t._rowwise_scale_inv = rp.view(-1, rp.shape[-1]) if rp is not None else None
+        t._columnwise_scale_inv = cp.view(-1, cp.shape[-1]) if cp is not None else None
 
-    # swizzle in place; batched
-    tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(result, True, False)
-    tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(result, False, True)
-    for t in result:
+    tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(fp8_tensors, True, False)
+    tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(fp8_tensors, False, True)
+    for t in fp8_tensors:
         t._with_gemm_swizzled_scales = True
 
-    return result, "bhsd"
-
 
 def combine_and_quantize(
     qkv_layout,
diff --git a/transformer_engine/pytorch/attention/fused_mla_q_uproj.py b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py
new file mode 100644
index 0000000000..c176985254
--- /dev/null
+++ b/transformer_engine/pytorch/attention/fused_mla_q_uproj.py
@@ -0,0 +1,175 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Fused MLA Q up-projection + per-head RoPE + MXFP8 quantize."""
+
+from __future__ import annotations
+import functools
+import os
+from importlib.metadata import PackageNotFoundError, version as get_pkg_version
+
+import torch
+import transformer_engine_torch as tex
+from packaging.version import Version as PkgVersion
+
+from ..constants import MXFP8_BLOCK_SCALING_SIZE
+from ..quantized_tensor import QuantizedTensor
+from ..tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor
+from ..utils import get_device_compute_capability
+
+_CUDNN_FRONTEND_MIN_VERSION = "1.27.0"
+
+
+def _cudnn_frontend_version_supported() -> bool:
+    """Check that the installed nvidia-cudnn-frontend meets the minimum version."""
+    try:
+        return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion(
+            _CUDNN_FRONTEND_MIN_VERSION
+        )
+    except PackageNotFoundError:
+        return False
+
+
+class FusedMLAQUpProjRopeQuant:
+    """Wrapper for the cuDNN fused MLA Q up-proj + per-head RoPE + MXFP8 quantize kernel.
+
+    - If w is already a QuantizedTensor (primary FP8 parameter in MXFP8BlockScaling recipe),
+      this performs an MXFP8 GEMM within the fusion (and quantizes the input if necessary)
+    - Otherwise (plain BF16 weight), x and w are passed as-is to the BF16 kernel variant.
+    """
+
+    @classmethod
+    @functools.lru_cache(maxsize=None)
+    def _kernel(cls):
+        # Import directly from the subpackage to avoid depending on cudnn/__init__.py
+        # lazy-import registration (which would require overlaying cudnn/__init__.py and
+        # could revert atomicrmw fixes present in the container's version).
+        try:
+            from cudnn import gemm_proj_rope_mxfp8_wrapper_sm100
+
+            return gemm_proj_rope_mxfp8_wrapper_sm100
+        except ImportError:
+            return None
+
+    @classmethod
+    @functools.lru_cache(maxsize=None)
+    def is_supported(cls) -> bool:
+        """Whether the cuDNN FE fused gemm rope quant wrapper is available"""
+        if int(os.environ.get("NVTE_FUSED_MLA_Q_UPROJ", "1")) <= 0:
+            return False
+        if not _cudnn_frontend_version_supported():
+            return False
+        if get_device_compute_capability()[0] < 10:
+            return False
+        if cls._kernel() is None:
+            return False
+        return True
+
+    @classmethod
+    def run(
+        cls,
+        x: torch.Tensor,
+        w,  # MXFP8Tensor (primary FP8 param) or bf16 torch.Tensor
+        cos: torch.Tensor,
+        sin: torch.Tensor,
+        s: int,
+        b: int,
+    ) -> "tuple[MXFP8Tensor, torch.Tensor]":
+        """Run the fused kernel; return (Q MXFP8Tensor, activation saved for the wgrad backward).
+
+        The kernel precision is selected by the weight precision.
+        """
+
+        from cuda.bindings import driver as cuda
+
+        stream = cuda.CUstream(torch.cuda.current_stream(x.device).cuda_stream)
+        wrapper = cls._kernel()
+
+        if isinstance(w, QuantizedTensor):
+            assert isinstance(w, MXFP8Tensor), (
+                "FusedMLAQUpProjRopeQuant expects an MXFP8Tensor weight (MXFP8BlockScaling"
+                f" recipe), got {type(w).__name__}. Use the unfused path for other quantization"
+                " recipes."
+            )
+            # ---- FP8 projection: MXFP8-cast x (both usages) + reuse w's fp8 codes -> mxfp8in ----
+            # Quantize x with both rowwise (for the forward GEMM) and columnwise (for the FP8
+            # wgrad in backward, matching the unfused path).
+            x_quantizer = MXFP8Quantizer(
+                fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True
+            )
+            x_mxfp8 = x_quantizer(x)
+            x_code = x_mxfp8._rowwise_data.view(torch.float8_e4m3fn)  # [tokens, K]
+            x_scale = x_mxfp8._rowwise_scale_inv  # [tokens, K//32] uint8
+
+            # Primary FP8 parameter: already quantized; use its rowwise FP8 codes + E8M0 scales.
+            w.update_usage(rowwise_usage=True, columnwise_usage=None)
+            w_code = w._rowwise_data.view(torch.float8_e4m3fn)  # [N, K]
+            w_scale = w._rowwise_scale_inv  # [N, K//32] uint8
+
+            out = wrapper(
+                x_code,
+                w_code,
+                cos,
+                sin,
+                x_scale=x_scale,
+                w_scale=w_scale,
+                w_out_in=True,
+                stream=stream,
+            )
+
+            # Drop rowwise data now.
+            # Only columnwise x is needed for the FP8 wgrad in backward.
+            x_mxfp8.update_usage(rowwise_usage=False, columnwise_usage=True)
+            x_saved = x_mxfp8
+        else:
+            # ---- 16-bit projection: bf16 GEMM inputs -> bf16in (the projection stays bf16) ----
+            out = wrapper(x, w, cos, sin, w_out_in=True, stream=stream)
+            x_saved = x
+
+        nh = out["out_fp8_row"].shape[1]
+        d = out["out_fp8_row"].shape[2]
+        query = cls.wrap_mxfp8(
+            out["out_fp8_row"],
+            out["out_scales_row"],
+            out["out_fp8_col"],
+            out["out_scales_col"],
+            s,
+            b,
+            nh,
+            d,
+        )
+        # 2nd return is the activation to save for wgrad: MXFP8 (fp8 path) or bf16 (16-bit path).
+        return query, x_saved
+
+    @classmethod
+    def wrap_mxfp8(
+        cls,
+        fp8_row: torch.Tensor,
+        scales_row: torch.Tensor,
+        fp8_col: torch.Tensor,
+        scales_col: torch.Tensor,
+        s: int,
+        b: int,
+        nh: int,
+        d: int,
+    ) -> MXFP8Tensor:
+        """Wrap raw data and scale tensors into an MXFP8Tensor"""
+
+        blk = MXFP8_BLOCK_SCALING_SIZE
+        # Both rowwise and columnwise Q are required:
+        # - Forward QK^T uses rowwise
+        # - cuDNN backward (fused_attn_fp8_bwd_impl) requires columnwise for dK gradient
+        quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True)
+        return MXFP8Tensor(
+            shape=(s, b, nh, d),
+            dtype=torch.bfloat16,
+            rowwise_data=fp8_row.view(s, b, nh, d),
+            rowwise_scale_inv=scales_row.view(s, b, nh, d // blk),
+            columnwise_data=fp8_col.view(s, b, nh, d),
+            columnwise_scale_inv=scales_col.view(s // blk, b, nh, d),
+            quantizer=quantizer,
+            requires_grad=False,
+            fp8_dtype=tex.DType.kFloat8E4M3,
+            with_gemm_swizzled_scales=False,
+        )

From da280583e0e167a4813e1fc9d3c6ea69ff9846ff Mon Sep 17 00:00:00 2001
From: Fei Wu <33940270+YangFei1990@users.noreply.github.com>
Date: Thu, 6 Aug 2026 21:39:24 -0700
Subject: [PATCH 45/67] [Pytorch] Add support for row-wise quanted input for
 grouped gemm (#3244)

* add support for row-wise quanted input for grouped gemm

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* doc change

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* implement for backward

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* allow scaled_bias + frozen weights in prequant path

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* allow bias + frozen weight path in prequantized input

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* use .copy to create group tensor

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* move the implementation to c++ layer

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* bug fix for operation ordering

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* add comprehensive tests

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* rename to group_requantize

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* rename func with inplace tag

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* refactor the requantize function

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* minor fix for test

Signed-off-by: YangFei1990 <feiw@nvidia.com>

---------

Signed-off-by: YangFei1990 <feiw@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .../test_mxfp8_group_quantize_graph_safe.py   | 336 +++++++++++++++
 tests/pytorch/test_grouped_mlp.py             | 403 ++++++++++++++++++
 transformer_engine/pytorch/csrc/extensions.h  |   5 +
 .../pytorch/csrc/extensions/cast.cpp          |  96 +++++
 .../pytorch/csrc/extensions/pybind.cpp        |   6 +
 .../pytorch/csrc/extensions/swizzle.cpp       |  17 +
 .../pytorch/ops/basic/grouped_linear.py       |  75 +++-
 .../pytorch/ops/fused/grouped_mlp.py          | 110 ++---
 8 files changed, 985 insertions(+), 63 deletions(-)

diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
index d07953ce37..94861ea433 100644
--- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
+++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
@@ -469,3 +469,339 @@ def test_grouped_tensor_mxfp8_with_paged_stashing(
         valid_M=valid_M,
         optimize_for_gemm=optimize_for_gemm,
     )
+
+
+# ---------------------------------------------------------------------------------------------
+# Pre-quantized MXFP8 input (FP8 token dispatch)
+#
+# tex.group_requantize_inplace takes a grouped tensor that arrives
+# ALREADY rowwise-quantized (its high-precision form no longer exists), and makes it GEMM-ready
+# in both directions: the rowwise data passes through verbatim with its scales swizzled, and the
+# columnwise copy is rebuilt via dequantize + columnwise-only requantize.
+#
+# These mirror the edge-case matrices of test_grouped_tensor_mxfp8_versus_reference and
+# test_grouped_tensor_mxfp8_with_paged_stashing so the same shapes, zero-token placements and
+# uneven splits exercise this path.
+# ---------------------------------------------------------------------------------------------
+
+
+def make_prequantized_wire_tensor(x: torch.Tensor, split_section_tensor: torch.Tensor):
+    """Rowwise-only, unswizzled grouped tensor, as FP8 dispatch delivers it."""
+    wire_quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=False)
+    # Must stay unswizzled: the requantize path asserts on it and dequantize needs compact scales.
+    wire_quantizer.optimize_for_gemm = False
+    wire = fused_grouped_quantize(x, split_section_tensor, wire_quantizer)
+    assert wire.columnwise_data is None
+    assert not wire._with_gemm_swizzled_scales
+    return wire
+
+
+def make_op_quantizer(columnwise: bool):
+    """The op's input quantizer, configured the way the ops layer configures it.
+
+    ``columnwise`` mirrors ``weight_requires_grad``: it tells the helper whether a wgrad GEMM
+    will consume a columnwise copy.
+    """
+    quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=True, columnwise=columnwise)
+    quantizer.optimize_for_gemm = True
+    return quantizer
+
+
+def make_gemm_ready_tensor(x: torch.Tensor, split_section_tensor: torch.Tensor):
+    """Grouped tensor already GEMM-ready in both directions (swizzled scales)."""
+    quantizer = make_op_quantizer(columnwise=True)
+    tensor = fused_grouped_quantize(x, split_section_tensor, quantizer)
+    assert tensor.columnwise_data is not None
+    assert tensor._with_gemm_swizzled_scales
+    return tensor
+
+
+def check_prequantized_requantize_versus_reference(
+    x_dtype: torch.dtype,
+    M: int,
+    N: int,
+    split_sections: list[int],
+) -> None:
+    """Run the pre-quantized requantize path and check both directions against a reference.
+
+    The reference is derived from dequantize(wire), not from the original high-precision x: that
+    is the only data a consumer can see after dispatch, and MXFP8 rowwise requantization is
+    idempotent, so it is an exact reference rather than an approximate one.
+    """
+    device = "cuda"
+    torch.manual_seed(0)
+    torch.cuda.manual_seed(0)
+
+    # The buffer is always M rows. Paged stashing is just the case where the groups cover fewer
+    # than M of them (valid_M < M) and the tail holds garbage the kernels must leave alone; the
+    # non-paged case is the same code path with sum(split_sections) == M.
+    x = torch.randn((M, N), dtype=x_dtype, device=device)
+    split_section_tensor = torch.tensor(split_sections, dtype=torch.int64, device=device)
+    num_groups = len(split_sections)
+    # Rows the groups actually cover. Beyond this the buffers hold whatever the allocator handed
+    # out, so nothing past it may be compared.
+    valid_rows = sum(split_sections)
+
+    wire = make_prequantized_wire_tensor(x, split_section_tensor)
+
+    # Snapshot what must survive verbatim, plus the compact scales the reference swizzles.
+    rowwise_data_before = wire.rowwise_data.clone()
+    wire_splits_before = [
+        (t._rowwise_data.view(dtype=torch.uint8).clone(), t._rowwise_scale_inv.clone())
+        for t in wire.split_into_quantized_tensors()
+    ]
+
+    # Reference high-precision input: everything downstream is derived from this. Only the live
+    # rows are kept -- dequantize allocates M rows but writes only the ones the groups cover.
+    dequantized_ref = (
+        tex.group_dequantize(wire, te.DType.kBFloat16)
+        .rowwise_data.view(M, N)[:valid_rows, :]
+        .clone()
+    )
+
+    # Reference columnwise copy, quantized per group from the dequantized data.
+    colwise_quantizers = [
+        MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3, rowwise=False, columnwise=True)
+        for _ in range(num_groups)
+    ]
+    _, _, colwise_data_ref, colwise_scale_ref = reference_group_quantize(
+        dequantized_ref,
+        colwise_quantizers,
+        split_sections,
+        return_rowwise=False,
+        return_transpose=True,
+    )
+
+    # ---- the code under test ----
+    # The op's quantizer, configured as the ops layer does: columnwise_usage says a wgrad GEMM
+    # will consume the columnwise copy, so the helper builds it and switches rowwise off itself.
+    dequantized = tex.group_requantize_inplace(
+        wire,
+        make_op_quantizer(columnwise=True),
+        num_groups,
+        split_section_tensor,
+        te.DType.kBFloat16,
+        return_dequantized=True,
+    )
+
+    assert wire.columnwise_data is not None, "columnwise data must be built"
+    assert wire.columnwise_scale_inv is not None, "columnwise scales must survive the swizzle"
+
+    # The returned dequantized tensor is what bias gradients are reduced from. Compare only the
+    # live rows: both this and the reference allocate M rows but write only the covered ones, and
+    # their tails are separate uninitialized allocations.
+    torch.testing.assert_close(dequantized[:valid_rows, :], dequantized_ref, atol=0.0, rtol=0.0)
+
+    # The rowwise DATA must pass through untouched; only its scales are re-laid-out.
+    torch.testing.assert_close(wire.rowwise_data, rowwise_data_before, atol=0.0, rtol=0.0)
+
+    if valid_rows > 0:
+        # A tensor whose groups are all empty has no scales to lay out, so the swizzle is a no-op
+        # and leaves the flag unset; every other case must come back swizzled.
+        assert wire._with_gemm_swizzled_scales, "rowwise scales must be marked swizzled"
+
+    # Per-group comparison, same structure as check_grouped_tensor_mxfp8_versus_reference.
+    outputs = wire.split_into_quantized_tensors()
+    x_splits = torch.split(dequantized_ref, split_sections)
+
+    for i, out in enumerate(outputs):
+        rows_i = split_sections[i]
+        scale_before = wire_splits_before[i][1]
+        colwise_data = out._columnwise_data.view(dtype=torch.uint8)
+        colwise_scale = out._columnwise_scale_inv
+
+        if rows_i == 0:
+            # Buffers for empty groups are never written, so only shape and dtype are meaningful.
+            assert_same_shape_and_dtype(colwise_data, colwise_data_ref[i])
+            assert_same_shape_and_dtype(colwise_scale, colwise_scale_ref[i])
+            continue
+
+        # Rowwise scales: the swizzled form of the compact scales this group arrived with. The
+        # rowwise DATA is covered by the whole-buffer identity check above.
+        torch.testing.assert_close(
+            out._rowwise_scale_inv,
+            swizzle_mxfp8_scale(rows_i, N, scale_before, columnwise=False),
+            atol=0.0,
+            rtol=0.0,
+        )
+
+        # Columnwise: rebuilt from the dequantized data, and swizzled by the quantize kernel
+        # because the caller sets optimize_for_gemm.
+        torch.testing.assert_close(colwise_data, colwise_data_ref[i], atol=0.0, rtol=0.0)
+        valid_scale_shape = get_mxfp8_scale_shape_no_padding(x_splits[i].shape, True)
+        assert (
+            valid_scale_shape == colwise_scale.shape
+        ), "The columnwise scale shape is not correctly aligned"
+        torch.testing.assert_close(
+            colwise_scale,
+            swizzle_mxfp8_scale(rows_i, N, colwise_scale_ref[i], columnwise=True),
+            atol=0.0,
+            rtol=0.0,
+        )
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize(
+    "M, N",
+    [
+        # edge case, zero tokens for all
+        (0, 512),
+        # full tile cases
+        (1024, 256),
+        # larger sizes
+        (8192, 1024),
+        (16384, 8192),
+    ],
+)
+@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str)
+@pytest.mark.parametrize(
+    "edge_cases",
+    [
+        "regular",
+        "zero_tokens_front",
+        "zero_tokens_end",
+        "zero_tokens_middle",
+        "random_uneven_split",
+    ],
+)
+def test_prequantized_requantize_versus_reference(
+    x_dtype: torch.dtype,
+    M: int,
+    N: int,
+    edge_cases: str,
+) -> None:
+    split_sections = generate_split_sections(M, N, edge_cases)
+    check_prequantized_requantize_versus_reference(
+        x_dtype=x_dtype,
+        M=M,
+        N=N,
+        split_sections=split_sections,
+    )
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize(
+    "M, N",
+    [
+        # M won't be empty in paged stashing
+        (1024, 256),
+        (8192, 1024),
+        (16384, 8192),
+    ],
+)
+@pytest.mark.parametrize("x_dtype", [torch.bfloat16], ids=str)
+@pytest.mark.parametrize(
+    "edge_cases",
+    [
+        "regular",
+        "zero_tokens_all",
+        "zero_tokens_front",
+        "zero_tokens_end",
+        "zero_tokens_middle",
+        "random_uneven_split",
+    ],
+)
+def test_prequantized_requantize_with_paged_stashing(
+    x_dtype: torch.dtype,
+    M: int,
+    N: int,
+    edge_cases: str,
+) -> None:
+    # Paged stashing: the buffer holds M rows but only valid_M carry live tokens; the rest is
+    # garbage the kernels must not touch.
+    valid_M = 0 if edge_cases == "zero_tokens_all" else M // 2
+    split_sections = generate_split_sections(valid_M, N, edge_cases)
+    assert sum(split_sections) == valid_M
+
+    check_prequantized_requantize_versus_reference(
+        x_dtype=x_dtype,
+        M=M,
+        N=N,
+        split_sections=split_sections,
+    )
+
+
+def _requantize_setup(M: int = 1024, N: int = 256):
+    """Common inputs for the state-dispatch tests below."""
+    torch.manual_seed(0)
+    split_sections = [M // 4] * 4
+    x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda")
+    return x, torch.tensor(split_sections, dtype=torch.int64, device="cuda"), len(split_sections)
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_prequantized_requantize_passes_through_gemm_ready_input():
+    """A tensor already GEMM-ready in both directions is left untouched."""
+    x, splits, num_groups = _requantize_setup()
+    tensor = make_gemm_ready_tensor(x, splits)
+    rowwise_before = tensor.rowwise_data.clone()
+    columnwise_before = tensor.columnwise_data.clone()
+    scale_before = tensor.scale_inv.clone()
+
+    out = tex.group_requantize_inplace(
+        tensor, make_op_quantizer(columnwise=True), num_groups, splits, te.DType.kBFloat16
+    )
+
+    assert out is None
+    assert tensor._with_gemm_swizzled_scales
+    torch.testing.assert_close(tensor.rowwise_data, rowwise_before, atol=0.0, rtol=0.0)
+    torch.testing.assert_close(tensor.columnwise_data, columnwise_before, atol=0.0, rtol=0.0)
+    torch.testing.assert_close(tensor.scale_inv, scale_before, atol=0.0, rtol=0.0)
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_prequantized_requantize_skips_columnwise_when_not_needed():
+    """columnwise_usage=False (frozen weights) swizzles rowwise without building columnwise."""
+    x, splits, num_groups = _requantize_setup()
+    wire = make_prequantized_wire_tensor(x, splits)
+
+    out = tex.group_requantize_inplace(
+        wire, make_op_quantizer(columnwise=False), num_groups, splits, te.DType.kBFloat16
+    )
+
+    assert out is None
+    assert wire._with_gemm_swizzled_scales, "the GEMM still needs swizzled rowwise scales"
+    assert wire.columnwise_data is None, "no wgrad GEMM, so no columnwise copy should be built"
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_prequantized_requantize_rejects_dequantized_from_gemm_ready_input():
+    """Bias grads cannot be served from an already-swizzled input, so this must raise."""
+    x, splits, num_groups = _requantize_setup()
+    tensor = make_gemm_ready_tensor(x, splits)
+
+    with pytest.raises(RuntimeError, match="compact format"):
+        tex.group_requantize_inplace(
+            tensor,
+            make_op_quantizer(columnwise=True),
+            num_groups,
+            splits,
+            te.DType.kBFloat16,
+            return_dequantized=True,
+        )
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_prequantized_requantize_rejects_swizzled_without_columnwise():
+    """Swizzled rowwise scales with no columnwise copy: it can no longer be rebuilt."""
+    x, splits, num_groups = _requantize_setup()
+    wire = make_prequantized_wire_tensor(x, splits)
+    # Swizzle in place, leaving the tensor rowwise-only.
+    tex.grouped_swizzle_for_gemm(wire, True, False)
+
+    with pytest.raises(RuntimeError, match="cannot be rebuilt"):
+        tex.group_requantize_inplace(
+            wire, make_op_quantizer(columnwise=True), num_groups, splits, te.DType.kBFloat16
+        )
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+def test_prequantized_requantize_rejects_dtype_mismatch():
+    """The helper keeps the input's format; it does not convert between formats."""
+    x, splits, num_groups = _requantize_setup()
+    wire = make_prequantized_wire_tensor(x, splits)
+    mismatched = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E5M2, rowwise=True, columnwise=True)
+    mismatched.optimize_for_gemm = True
+
+    with pytest.raises(RuntimeError, match="dtype"):
+        tex.group_requantize_inplace(wire, mismatched, num_groups, splits, te.DType.kBFloat16)
diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py
index 28f7360641..d09c92ad49 100644
--- a/tests/pytorch/test_grouped_mlp.py
+++ b/tests/pytorch/test_grouped_mlp.py
@@ -15,6 +15,7 @@
 import torch
 
 import transformer_engine.pytorch as te
+from transformer_engine.pytorch.constants import TE_DType
 import transformer_engine.pytorch.ops.fused.grouped_mlp as grouped_mlp_module
 from transformer_engine.pytorch.ops.fused.grouped_mlp import (
     _cudnn_frontend_supports_grouped_gemm_srelu,
@@ -230,6 +231,23 @@ def make_reference_and_test_tensors(
     return ref, test
 
 
+class _InjectGrad(torch.autograd.Function):
+    """Replace the gradient flowing into ``x`` with ``grad``.
+
+    Mirrors how an FP8 token dispatch delivers a pre-quantized ``GroupedTensor``
+    grad output: the downstream op simply returns one as its grad input.
+    """
+
+    @staticmethod
+    def forward(ctx, x, grad):  # pylint: disable=arguments-differ
+        ctx.injected_grad = grad
+        return x
+
+    @staticmethod
+    def backward(ctx, grad_output):  # pylint: disable=arguments-differ
+        return ctx.injected_grad, None
+
+
 class TestGroupedLinearOp:
     """Tests for advanced features with grouped linear basic op"""
 
@@ -464,6 +482,189 @@ def test_grouped_linear(
                     else:
                         assert b_test.grad is None
 
+    @staticmethod
+    def _make_rowwise_mxfp8_wire_input(
+        x_hp: torch.Tensor,
+        group_size: int,
+        split_sizes: torch.Tensor,
+    ) -> "GroupedTensor":
+        """Rowwise-only MXFP8 GroupedTensor with compact scales (FP8 dispatch wire format)."""
+        wire_quantizer = MXFP8Quantizer(
+            fp8_dtype=TE_DType[torch.float8_e4m3fn], rowwise=True, columnwise=False
+        )
+        x_wire = tex.group_quantize(
+            x_hp, wire_quantizer, group_size, split_sizes.to(dtype=torch.int64)
+        )
+        # Sanity: the wire tensor is rowwise-only with unswizzled scales.
+        assert x_wire.columnwise_data is None
+        assert not x_wire._with_gemm_swizzled_scales
+        return x_wire
+
+    @pytest.mark.parametrize("weight_requires_grad", (False, True))
+    def test_grouped_linear_prequantized_mxfp8_input(
+        self,
+        *,
+        group_size: int = 4,
+        weight_shape: tuple[int, int] = (256, 256),
+        split_alignment: int = 128,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device = "cuda",
+        weight_requires_grad: bool,
+    ) -> None:
+        """Rowwise-only MXFP8 GroupedTensor input (FP8 token dispatch wire format).
+
+        The input arrives already rowwise-quantized with compact scales. The op
+        must feed the rowwise data to the forward GEMM as-is and manufacture the
+        columnwise copy for the wgrad GEMM. The reference run consumes the
+        *dequantized* wire tensor (the only data a layer can see after FP8
+        dispatch) on the normal quantize-from-BF16 path; because MXFP8 requant
+        is idempotent along the rowwise axis and both paths derive the
+        columnwise copy from the same dequantized data, the two runs must match
+        bit-for-bit.
+        """
+        if not mxfp8_available:
+            pytest.skip(reason_for_no_mxfp8)
+        maybe_skip_quantization("mxfp8", dims=weight_shape, device=device, dtype=dtype)
+
+        # Split sizes (including an empty group)
+        split_sizes = [split_alignment * i for i in range(group_size)]
+        random.shuffle(split_sizes)
+        split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device)
+
+        out_features, in_features = weight_shape
+        total_tokens = int(split_sizes.sum().item())
+        in_shape = (total_tokens, in_features)
+
+        # Wire-format input and its exact dequantization (the reference input).
+        x_hp = torch.rand(in_shape, dtype=dtype, device=device) - 0.5
+        x_wire = self._make_rowwise_mxfp8_wire_input(x_hp, group_size, split_sizes)
+        x_ref = tex.group_dequantize(x_wire, TE_DType[dtype]).rowwise_data.view(in_shape)
+
+        dy = torch.rand((total_tokens, out_features), dtype=dtype, device=device) - 0.5
+        recipe = make_recipe("mxfp8")
+        op = te.ops.GroupedLinear(
+            group_size, in_features, out_features, bias=False, device=device, dtype=dtype
+        )
+        with torch.no_grad():
+            for param in op.parameters():
+                param.requires_grad_(requires_grad=weight_requires_grad)
+
+        def _run(x):
+            with te.autocast(enabled=True, recipe=recipe):
+                y = op(x, split_sizes)
+            wgrads = []
+            if weight_requires_grad:
+                y.backward(dy)
+                for group_idx in range(group_size):
+                    weight = getattr(op, f"weight{group_idx}")
+                    wgrads.append(weight.grad.detach().clone())
+                    weight.grad = None
+            return y.detach(), wgrads
+
+        y_ref, wgrads_ref = _run(x_ref)
+        y_test, wgrads_test = _run(x_wire)
+
+        # Bit-exact match expected (identical quantized inputs and kernels).
+        torch.testing.assert_close(y_test, y_ref, rtol=0, atol=0)
+        for wgrad_test, wgrad_ref in zip(wgrads_test, wgrads_ref):
+            torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0)
+
+    @pytest.mark.parametrize("bias_mode", ("none", "plain", "scaled"))
+    @pytest.mark.parametrize("weight_requires_grad", (False, True))
+    def test_grouped_linear_prequantized_mxfp8_grad(
+        self,
+        *,
+        group_size: int = 4,
+        weight_shape: tuple[int, int] = (256, 256),
+        split_alignment: int = 128,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device = "cuda",
+        bias_mode: str,
+        weight_requires_grad: bool,
+    ) -> None:
+        """Rowwise-only MXFP8 GroupedTensor grad output (FP8 token dispatch, backward).
+
+        Mirrors ``test_grouped_linear_prequantized_mxfp8_input`` on the backward
+        side: the rowwise data feeds the dgrad GEMM and the columnwise copy is
+        manufactured for wgrad. Covers all three bias gradient sources: none,
+        ``plain`` (fused into the columnwise stage of the quantize kernel, or
+        reduced from the dequantized grad when frozen weights leave no columnwise
+        stage), and ``scaled`` (``scale_bias``, whose dbias/dscales need the
+        dequantized grad because they depend on the routing probabilities).
+
+        With frozen weights TE also requires frozen biases, so bias gradients are
+        not observable there; ``dscales`` still is, since the probabilities are an
+        input rather than a parameter.
+        """
+        if not mxfp8_available:
+            pytest.skip(reason_for_no_mxfp8)
+        maybe_skip_quantization("mxfp8", dims=weight_shape, device=device, dtype=dtype)
+
+        has_bias = bias_mode != "none"
+        use_scale_bias = bias_mode == "scaled"
+
+        # Split sizes (including an empty group)
+        split_sizes = [split_alignment * i for i in range(group_size)]
+        random.shuffle(split_sizes)
+        split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device)
+
+        out_features, in_features = weight_shape
+        total_tokens = int(split_sizes.sum().item())
+
+        x = torch.rand((total_tokens, in_features), dtype=dtype, device=device) - 0.5
+        dy_hp = torch.rand((total_tokens, out_features), dtype=dtype, device=device) - 0.5
+        probs = torch.rand((total_tokens,), dtype=dtype, device=device)
+
+        # Wire-format grad output and its exact dequantization (the reference grad).
+        dy_wire = self._make_rowwise_mxfp8_wire_input(dy_hp, group_size, split_sizes)
+        dy_ref = tex.group_dequantize(dy_wire, TE_DType[dtype]).rowwise_data.view(
+            total_tokens, out_features
+        )
+
+        recipe = make_recipe("mxfp8")
+        op = te.ops.GroupedLinear(
+            group_size,
+            in_features,
+            out_features,
+            bias=has_bias,
+            device=device,
+            dtype=dtype,
+            scale_bias=use_scale_bias,
+        )
+        if not weight_requires_grad:
+            # TE requires bias.requires_grad to match weight.requires_grad.
+            for param in op.parameters():
+                param.requires_grad_(False)
+
+        def _run(grad):
+            x_in = x.detach().clone().requires_grad_()
+            probs_in = probs.detach().clone().requires_grad_(use_scale_bias)
+            extra_inputs = (split_sizes, probs_in) if use_scale_bias else (split_sizes,)
+            with te.autocast(enabled=True, recipe=recipe):
+                y = op(x_in, *extra_inputs)
+            # Deliver ``grad`` as the op's grad output, as FP8 dispatch would.
+            _InjectGrad.apply(y, grad).backward(torch.ones_like(y))
+            grads = [("dx", x_in.grad)]
+            if use_scale_bias:
+                grads.append(("dprobs", probs_in.grad))
+            if weight_requires_grad:
+                for group_idx in range(group_size):
+                    weight = getattr(op, f"weight{group_idx}")
+                    grads.append((f"w{group_idx}", weight.grad.detach().clone()))
+                    weight.grad = None
+                    if has_bias:
+                        bias_param = getattr(op, f"bias{group_idx}")
+                        grads.append((f"b{group_idx}", bias_param.grad.detach().clone()))
+                        bias_param.grad = None
+            return grads
+
+        grads_ref = _run(dy_ref)
+        grads_test = _run(dy_wire)
+
+        # Bit-exact match expected (identical quantized grads and kernels).
+        for (_, grad_test), (_, grad_ref) in zip(grads_test, grads_ref):
+            torch.testing.assert_close(grad_test, grad_ref, rtol=0, atol=0)
+
     @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16))
     @pytest.mark.parametrize(
         "quantization",
@@ -1116,6 +1317,208 @@ def _make_module():
             assert_close(fc1.weight.grad, fc1_w_ref_grad, **tols)
             assert_close(fc2.weight.grad, fc2_w_ref_grad, **tols)
 
+    @pytest.mark.parametrize("weight_requires_grad", (False, True))
+    def test_grouped_mlp_prequantized_mxfp8_input(
+        self,
+        *,
+        group_size: int = 4,
+        hidden_size: int = 256,
+        split_alignment: int = 256,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device = "cuda",
+        weight_requires_grad: bool,
+    ) -> None:
+        """Fused grouped MLP with a rowwise-only MXFP8 GroupedTensor input.
+
+        Production (FP8 token dispatch) path: FC1 receives an already
+        rowwise-quantized input with compact scales. The fused op must feed the
+        rowwise data to the forward GEMM and manufacture FC1's columnwise copy
+        for wgrad. Compared bit-for-bit against a run on the dequantized wire
+        input (see ``test_grouped_linear_prequantized_mxfp8_input``).
+        """
+        if not mxfp8_available:
+            pytest.skip(reason_for_no_mxfp8)
+        if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
+            pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system")
+        maybe_skip_quantization(
+            "mxfp8", dims=(hidden_size, hidden_size), device=device, dtype=dtype
+        )
+
+        # Split sizes (including an empty group); sum is a multiple of 128.
+        split_sizes = [split_alignment * i for i in range(group_size)]
+        random.shuffle(split_sizes)
+        split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device)
+        total_tokens = int(split_sizes.sum().item())
+        glu_interleave_size = 32
+
+        # Wire-format FC1 input and its exact dequantization (reference input).
+        x_hp = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5
+        x_wire = TestGroupedLinearOp._make_rowwise_mxfp8_wire_input(x_hp, group_size, split_sizes)
+        x_ref = tex.group_dequantize(x_wire, TE_DType[dtype]).rowwise_data.view(
+            total_tokens, hidden_size
+        )
+
+        probs = torch.rand((total_tokens,), dtype=dtype, device=device)
+        dy = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5
+
+        recipe = make_recipe("mxfp8")
+        with te.quantized_model_init(enabled=True, recipe=recipe):
+            fc1 = te.ops.GroupedLinear(
+                group_size, hidden_size, 2 * hidden_size, bias=False, device=device, dtype=dtype
+            )
+            fc2 = te.ops.GroupedLinear(
+                group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype
+            )
+            module = te.ops.Sequential(
+                fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2
+            )
+        with torch.no_grad():
+            for param in module.parameters():
+                param.requires_grad_(requires_grad=weight_requires_grad)
+
+        def _run(x):
+            with te.autocast(enabled=True, recipe=recipe):
+                y = module(x, split_sizes, probs, split_sizes)
+            fc1_wgrads, fc2_wgrads = [], []
+            if weight_requires_grad:
+                y.backward(dy)
+                for group_idx in range(group_size):
+                    fc1_w = getattr(fc1, f"weight{group_idx}")
+                    fc2_w = getattr(fc2, f"weight{group_idx}")
+                    fc1_wgrads.append(fc1_w.grad.detach().clone())
+                    fc2_wgrads.append(fc2_w.grad.detach().clone())
+                    fc1_w.grad = None
+                    fc2_w.grad = None
+            return y.detach(), fc1_wgrads, fc2_wgrads
+
+        y_ref, fc1_wgrads_ref, fc2_wgrads_ref = _run(x_ref)
+        y_test, fc1_wgrads_test, fc2_wgrads_test = _run(x_wire)
+
+        # Confirm the CuTeDSL fused op was actually formed (not the fallback).
+        forward_ops = module._module_groups[0]._forward_ops
+        assert len(forward_ops) == 1
+        assert isinstance(forward_ops[0][0], te.ops.fused.GroupedMLP_CuTeGEMMGLU)
+
+        # Bit-exact match expected (identical quantized inputs and kernels).
+        torch.testing.assert_close(y_test, y_ref, rtol=0, atol=0)
+        for wgrad_test, wgrad_ref in zip(fc1_wgrads_test, fc1_wgrads_ref):
+            torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0)
+        for wgrad_test, wgrad_ref in zip(fc2_wgrads_test, fc2_wgrads_ref):
+            torch.testing.assert_close(wgrad_test, wgrad_ref, rtol=0, atol=0)
+
+    @pytest.mark.parametrize("bias", (False, True))
+    @pytest.mark.parametrize("weight_requires_grad", (False, True))
+    def test_grouped_mlp_prequantized_mxfp8_grad(
+        self,
+        *,
+        group_size: int = 4,
+        hidden_size: int = 256,
+        split_alignment: int = 256,
+        dtype: torch.dtype = torch.bfloat16,
+        device: torch.device = "cuda",
+        bias: bool,
+        weight_requires_grad: bool,
+    ) -> None:
+        """Fused grouped MLP with a rowwise-only MXFP8 GroupedTensor grad output.
+
+        FC2 receives the pre-quantized grad, as an FP8 token dispatch delivers it
+        on the backward pass. With ``bias`` FC2 uses ``scale_bias``, whose
+        dbias/dscales need the dequantized grad rather than the fused dbias.
+        """
+        if not mxfp8_available:
+            pytest.skip(reason_for_no_mxfp8)
+        if not te.ops.fused.GroupedMLP_CuTeGEMMGLU.is_supported():
+            pytest.skip("Fused grouped MLP (CuTeDSL) is not supported on this system")
+        if not weight_requires_grad:
+            # Independent of pre-quantization: the fused forward saves the input
+            # activations whenever anything requires grad, then asserts they carry
+            # columnwise data -- which is only built when the weights need grads.
+            pytest.skip("Fused grouped MLP does not support frozen weights")
+        maybe_skip_quantization(
+            "mxfp8", dims=(hidden_size, hidden_size), device=device, dtype=dtype
+        )
+
+        # Split sizes (including an empty group); sum is a multiple of 128.
+        split_sizes = [split_alignment * i for i in range(group_size)]
+        random.shuffle(split_sizes)
+        split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device)
+        total_tokens = int(split_sizes.sum().item())
+        glu_interleave_size = 32
+
+        x = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5
+        probs = torch.rand((total_tokens,), dtype=dtype, device=device)
+        dy_hp = torch.rand((total_tokens, hidden_size), dtype=dtype, device=device) - 0.5
+
+        # Wire-format grad output and its exact dequantization (the reference grad).
+        dy_wire = TestGroupedLinearOp._make_rowwise_mxfp8_wire_input(dy_hp, group_size, split_sizes)
+        dy_ref = tex.group_dequantize(dy_wire, TE_DType[dtype]).rowwise_data.view(
+            total_tokens, hidden_size
+        )
+
+        recipe = make_recipe("mxfp8")
+        with te.quantized_model_init(enabled=True, recipe=recipe):
+            fc1 = te.ops.GroupedLinear(
+                group_size, hidden_size, 2 * hidden_size, bias=bias, device=device, dtype=dtype
+            )
+            fc2 = te.ops.GroupedLinear(
+                group_size,
+                hidden_size,
+                hidden_size,
+                bias=bias,
+                device=device,
+                dtype=dtype,
+                scale_bias=bias,
+            )
+            module = te.ops.Sequential(
+                fc1, te.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size), fc2
+            )
+
+        # Frozen experts (weights) with a still-training bias/router is the case
+        # where the dequantized grad is needed but is not a byproduct of wgrad.
+        if not weight_requires_grad:
+            for fc in (fc1, fc2):
+                for group_idx in range(group_size):
+                    getattr(fc, f"weight{group_idx}").requires_grad_(False)
+
+        def _run(grad):
+            x_in = x.detach().clone().requires_grad_()
+            fc2_extra = (split_sizes, probs) if bias else (split_sizes,)
+            with te.autocast(enabled=True, recipe=recipe):
+                y = module(x_in, split_sizes, probs, *fc2_extra)
+            _InjectGrad.apply(y, grad).backward(torch.ones_like(y))
+            grads = [("dx", x_in.grad)]
+            for name, fc in (("fc1", fc1), ("fc2", fc2)):
+                for group_idx in range(group_size):
+                    if weight_requires_grad:
+                        weight = getattr(fc, f"weight{group_idx}")
+                        grads.append((f"{name}_w{group_idx}", weight.grad.detach().clone()))
+                        weight.grad = None
+                    if bias:
+                        bias_param = getattr(fc, f"bias{group_idx}")
+                        grads.append((f"{name}_b{group_idx}", bias_param.grad.detach().clone()))
+                        bias_param.grad = None
+            return grads
+
+        grads_ref = _run(dy_ref)
+        grads_test = _run(dy_wire)
+
+        # Confirm the CuTeDSL fused op was actually formed (not the fallback).
+        forward_ops = module._module_groups[0]._forward_ops
+        assert len(forward_ops) == 1
+        assert isinstance(forward_ops[0][0], te.ops.fused.GroupedMLP_CuTeGEMMGLU)
+
+        # Bit-exact match expected (identical quantized grads and kernels), except
+        # bias gradients: the fused kernels generate them with an accumulation
+        # that is not reproducible run to run (two runs on identical inputs differ
+        # by one BF16 ulp). Same tolerances as
+        # ``test_grouped_mlp_single_weight_numerics``.
+        bias_tols = {"rtol": 0.05, "atol": 0.015625}
+        for (name, grad_test), (_, grad_ref) in zip(grads_test, grads_ref):
+            if "_b" in name:
+                torch.testing.assert_close(grad_test, grad_ref, **bias_tols)
+            else:
+                torch.testing.assert_close(grad_test, grad_ref, rtol=0, atol=0)
+
     @pytest.mark.parametrize("bias", (False, True))
     @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list)
     @pytest.mark.parametrize(
diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index 8248a63680..1a9fabb3bb 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -392,6 +392,11 @@ py::object bgrad_group_quantize(const at::Tensor &tensor, py::handle quantizer,
                                 std::optional<at::Tensor> last_dims,
                                 std::optional<at::Tensor> tensor_offsets);
 
+py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer,
+                                    const size_t num_tensors, std::optional<at::Tensor> first_dims,
+                                    DType otype, std::optional<at::Tensor> tensor_offsets,
+                                    bool return_dequantized);
+
 std::vector<py::object> multi_tensor_quantize(const std::vector<at::Tensor> &tensor_list,
                                               std::vector<py::handle> quantizer_list);
 
diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
index 882481e0f4..5ce0261c82 100644
--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -664,6 +664,102 @@ py::object group_dequantize(const py::handle &input, transformer_engine::DType o
   return py::reinterpret_borrow<py::object>(out_py);
 }
 
+py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer,
+                                    const size_t num_tensors, std::optional<at::Tensor> first_dims,
+                                    DType otype, std::optional<at::Tensor> tensor_offsets,
+                                    bool return_dequantized) {
+  init_extension();
+
+  const bool has_rowwise =
+      !grouped_x.attr("rowwise_data").is_none() && !grouped_x.attr("scale_inv").is_none();
+  const bool has_columnwise = !grouped_x.attr("columnwise_data").is_none() &&
+                              !grouped_x.attr("columnwise_scale_inv").is_none();
+  const bool swizzled = grouped_x.attr("_with_gemm_swizzled_scales").cast<bool>();
+
+  NVTE_CHECK(has_rowwise, "Grouped input has no rowwise data and scales for the GEMM to consume.");
+
+  // The tensor's own quantization must match what the op expects on every path: even a
+  // pass-through hands its data straight to the GEMM. This keeps the input's format rather than
+  // converting between formats.
+  const auto input_quantizer = grouped_x.attr("quantizer");
+  NVTE_CHECK(!input_quantizer.is_none(), "Grouped input has no quantizer.");
+  NVTE_CHECK(Py_TYPE(input_quantizer.ptr()) == Py_TYPE(quantizer.ptr()),
+             "Grouped input and the op disagree on quantization format.");
+  NVTE_CHECK(input_quantizer.attr("dtype").cast<DType>() == quantizer.attr("dtype").cast<DType>(),
+             "Grouped input and the quantizer disagree on the FP8 dtype.");
+
+  // The columnwise copy is only worth building when a wgrad GEMM will consume it. Read this
+  // before the usage is overridden below.
+  const bool need_columnwise = quantizer.attr("columnwise_usage").cast<bool>();
+
+  if (swizzled) {
+    // Already GEMM-ready. Nothing can be derived from here, since dequantization requires scales
+    // in compact format, so the input must already carry everything that will be consumed. No
+    // quantization kernel runs, which makes this path format-agnostic.
+    NVTE_CHECK(has_columnwise || !need_columnwise,
+               "Grouped input has swizzled scales but no columnwise data for the wgrad GEMM. It "
+               "cannot be rebuilt, because dequantization requires scales in compact format.");
+    NVTE_CHECK(!return_dequantized,
+               "Cannot return a dequantized tensor for an already-swizzled grouped input: "
+               "dequantization requires scales in compact format.");
+    return py::none();
+  }
+
+  NVTE_CHECK(!has_columnwise,
+             "Grouped input already has columnwise data but unswizzled scales; requantizing "
+             "from it is not supported.");
+
+  // Everything below runs quantization kernels, so it is MXFP8-only.
+  NVTE_CHECK(detail::IsMXFP8Quantizers(quantizer.ptr()),
+             "Requantizing a grouped input is only supported for MXFP8.");
+
+  const auto logical_shape = grouped_x.attr("logical_shape").cast<py::tuple>();
+  const auto total_tokens = logical_shape[0].cast<size_t>();
+  const auto hidden_dim = logical_shape[1].cast<size_t>();
+  // Each group's token count must be a multiple of 128 too, so that every group's scales start
+  // on a swizzle-tile boundary. Those counts live on the device (host reads would break CUDA
+  // graph capture), so that half is the caller's contract rather than an assertion.
+  NVTE_CHECK(total_tokens % 128 == 0 && hidden_dim % 128 == 0,
+             "Requantizing a grouped input requires dims that are multiples of 128, but got (",
+             total_tokens, ", ", hidden_dim, ").");
+
+  // Dequantize first: it reads the rowwise scales, which the swizzle below replaces. Left
+  // undefined when nothing consumes it, which skips the pass entirely.
+  at::Tensor dequantized;
+  if (need_columnwise || return_dequantized) {
+    dequantized = group_dequantize(grouped_x, otype)
+                      .attr("rowwise_data")
+                      .cast<at::Tensor>()
+                      .view({static_cast<int64_t>(total_tokens), static_cast<int64_t>(hidden_dim)});
+  }
+
+  // Swizzle the rowwise scales before attaching any columnwise data: a rowwise-only swizzle
+  // resets columnwise_scale_inv to None, which would strand the columnwise data below with a
+  // null scale pointer.
+  grouped_swizzle_for_gemm(grouped_x, /*rowwise=*/true, /*columnwise=*/false);
+  // The swizzle hands back a 2D [num_tensors * padded_m, padded_k] scale buffer, but grouped
+  // tensors carry scales as a flat array indexed by element offsets (scale_inv_offsets), so
+  // per-group slicing breaks unless it is flattened back.
+  grouped_x.attr("scale_inv") = grouped_x.attr("scale_inv").attr("reshape")(-1);
+
+  if (need_columnwise) {
+    // Rebuild the columnwise copy the wgrad GEMM needs. It cannot be derived from the rowwise
+    // data because the two directions scale along perpendicular axes. Quantizing rowwise as well
+    // would redo work we already have, so that direction is switched off; the caller sets
+    // optimize_for_gemm, which makes the kernel emit swizzled columnwise scales directly.
+    quantizer.attr("set_usage")(py::arg("rowwise") = false, py::arg("columnwise") = true);
+    auto columnwise = group_quantize(dequantized, quantizer, num_tensors, first_dims, std::nullopt,
+                                     tensor_offsets, std::nullopt);
+    grouped_x.attr("columnwise_data") = columnwise.attr("columnwise_data");
+    grouped_x.attr("columnwise_scale_inv") = columnwise.attr("columnwise_scale_inv");
+  }
+
+  if (return_dequantized) {
+    return py::cast(dequantized);
+  }
+  return py::none();
+}
+
 namespace {
 
 void multi_tensor_quantize_impl(const std::vector<TensorWrapper> &input_list,
diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
index e406dc3446..2173dff8b2 100644
--- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
@@ -216,6 +216,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("bgrad_group_quantize", transformer_engine::pytorch::bgrad_group_quantize,
         py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
         py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none());
+  m.def("group_requantize_inplace", transformer_engine::pytorch::group_requantize_inplace,
+        "Rebuild the columnwise copy of a rowwise-prequantized MXFP8 grouped tensor and swizzle "
+        "its rowwise scales for GEMM, in place",
+        py::arg("grouped_x"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
+        py::arg("otype"), py::arg("tensor_offsets") = py::none(),
+        py::arg("return_dequantized") = false);
   m.def("bgrad_quantize", transformer_engine::pytorch::bgrad_quantize,
         "Compute bias gradient and quantize", py::arg("input"), py::arg("quantizer"));
   m.def("generic_gemm", transformer_engine::pytorch::gemm, "Compute GEMM (matrix-matrix multiply)",
diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp
index c90a7d6d0d..866d858ee3 100644
--- a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp
@@ -389,6 +389,23 @@ std::optional<SwizzledGroupedScales> maybe_swizzle_grouped_tensor(GroupedTensorW
         tensor_offsets.data_ptr, static_cast<DType>(tensor_offsets.dtype), tensor_offsets.shape);
   }
 
+  // Varying per-tensor dimensions. Leaving these unset declares the grouped tensor uniform,
+  // which selects the uniform-shape swizzle kernel.
+  const auto first_dims = input.get_first_dims();
+  if (first_dims.data_ptr != nullptr) {
+    swizzle_input.set_first_dims(first_dims.data_ptr, static_cast<DType>(first_dims.dtype),
+                                 first_dims.shape);
+    swizzle_output.set_first_dims(first_dims.data_ptr, static_cast<DType>(first_dims.dtype),
+                                  first_dims.shape);
+  }
+  const auto last_dims = input.get_last_dims();
+  if (last_dims.data_ptr != nullptr) {
+    swizzle_input.set_last_dims(last_dims.data_ptr, static_cast<DType>(last_dims.dtype),
+                                last_dims.shape);
+    swizzle_output.set_last_dims(last_dims.data_ptr, static_cast<DType>(last_dims.dtype),
+                                 last_dims.shape);
+  }
+
   // Per-tensor logical dimensions (uniform-shape grouped tensor).
   const size_t num_tensors = input.num_tensors();
   const auto logical_shape_nvte = input.logical_shape();
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index c94d78fecd..e1980d2943 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -14,7 +14,7 @@
 import torch
 
 import transformer_engine_torch as tex
-from ...constants import DType
+from ...constants import DType, TE_DType
 from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor
 from ...distributed import CudaRNGStatesTracker
 from ...module._common import WeightGradStore
@@ -1204,6 +1204,11 @@ def _fuser_forward_split_quantize(
         out_buffer: Optional[torch.Tensor] = None,
     ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]:
         """Legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow."""
+        if isinstance(input_, GroupedTensor):
+            raise NotImplementedError(
+                "Pre-quantized GroupedTensor input is only supported on the "
+                "graph-safe grouped-tensor path."
+            )
         num_groups = self.num_groups
         has_bias = self.has_bias
 
@@ -1330,14 +1335,39 @@ def _fuser_forward_grouped_tensor(
         input_tensor_offsets = grouped_tensor_offsets[2]
         output_tensor_offsets = grouped_tensor_offsets[3]
         original_shape = list(input_.size())
-        # Flatten to 2D so the first dim is the total token count.
-        x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features)
-        total_tokens = x.size(0)
-        # Build the input GroupedTensorStorage for input.
+        prequantized_input = with_quantized_compute and isinstance(input_, GroupedTensor)
         if with_quantized_compute:
             input_quantizer = input_quantizers[0]
             input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad)
             input_quantizer.optimize_for_gemm = True
+        if prequantized_input:
+            # GroupedTensor forbids reshape and is already in the canonical
+            # (total_tokens, in_features) layout; just validate the shape.
+            if input_.dim() != 2 or input_.size(-1) != self.in_features:
+                raise ValueError(
+                    "GroupedTensor input must have shape (total_tokens, "
+                    f"{self.in_features}), but got {tuple(input_.size())}."
+                )
+            total_tokens = input_.size(0)
+        else:
+            # Flatten to 2D so the first dim is the total token count.
+            x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features)
+            total_tokens = x.size(0)
+
+        # Build the input GroupedTensorStorage for input.
+        if prequantized_input:
+            # Input arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise data
+            # for the GEMM and let the helper supply whatever else the GEMMs need.
+            grouped_x = input_.copy()
+            tex.group_requantize_inplace(
+                grouped_x,
+                input_quantizer,
+                num_groups,
+                split_sizes,
+                TE_DType[dtype],
+                tensor_offsets=input_tensor_offsets,
+            )
+        elif with_quantized_compute:
             grouped_x = tex.group_quantize(
                 x,
                 input_quantizer,
@@ -1703,8 +1733,20 @@ def _fuser_backward_grouped_tensor(
 
         # Flatten grad_output to 2D (total_tokens, out_features)
         # to figure out total tokens.
-        dy_2d = grad_output.reshape(-1, self.out_features)
-        total_tokens = dy_2d.size(0)
+        prequantized_grad = with_quantized_compute and isinstance(grad_output, GroupedTensor)
+        if prequantized_grad:
+            # GroupedTensor forbids reshape and is already in the canonical
+            # (total_tokens, out_features) layout; just validate the shape.
+            if grad_output.dim() != 2 or grad_output.size(-1) != self.out_features:
+                raise ValueError(
+                    "GroupedTensor grad output must have shape (total_tokens, "
+                    f"{self.out_features}), but got {tuple(grad_output.size())}."
+                )
+            dy_2d = None
+            total_tokens = grad_output.size(0)
+        else:
+            dy_2d = grad_output.reshape(-1, self.out_features)
+            total_tokens = dy_2d.size(0)
         grad_input_shape = list(grad_output.size())[:-1] + [self.in_features]
 
         # Build the grad_output GroupedTensor.
@@ -1722,7 +1764,21 @@ def _fuser_backward_grouped_tensor(
             fuse_bgrad = isinstance(grad_output_quantizer, MXFP8Quantizer) or (
                 isinstance(grad_output_quantizer, Float8BlockQuantizer) and ctx.input_requires_grad
             )
-            if has_bias and not self._scale_bias and fuse_bgrad:
+            if prequantized_grad:
+                # Grad output arrived already quantized (e.g. FP8 token dispatch): reuse its
+                # rowwise data for the dgrad GEMM. Bias grads are reduced from the dequantized
+                # grad below, which is only kept when there is a bias.
+                grouped_dy = grad_output.copy()
+                dy_2d = tex.group_requantize_inplace(
+                    grouped_dy,
+                    grad_output_quantizer,
+                    num_groups,
+                    split_sizes,
+                    TE_DType[dtype],
+                    tensor_offsets=output_tensor_offsets,
+                    return_dequantized=has_bias,
+                )
+            elif has_bias and not self._scale_bias and fuse_bgrad:
                 grouped_dy, dbias_packed = tex.bgrad_group_quantize(
                     dy_2d,
                     grad_output_quantizer,
@@ -1764,7 +1820,8 @@ def _fuser_backward_grouped_tensor(
                     offsets=base_split_offsets,
                 )
             elif dbias_packed is None:
-                # BF16/FP16 path
+                # BF16/FP16 and pre-quantized MXFP8 paths, neither of which fuses dbias
+                # into a quantize kernel.
                 dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, num_groups)
             if self.single_grouped_bias:
                 final_bias_grads = [dbias_packed.to(dtype=dtype)]
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index db2c4e7674..76d51673f0 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -16,7 +16,7 @@
 from packaging.version import Version as PkgVersion
 
 import transformer_engine_torch as tex
-from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE
+from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE, TE_DType
 from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload
 from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor
 from ...distributed_weight import (
@@ -31,7 +31,7 @@
 from ...tensor.grouped_tensor import GroupedTensor
 from ...tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor
 from ...tensor.storage.grouped_tensor_storage import GroupedTensorStorage
-from ...triton.grouped_dbias_dscales import compute_grouped_dbias_dscales
+from ...triton.grouped_dbias_dscales import compute_grouped_dbias, compute_grouped_dbias_dscales
 from ...utils import (
     ceil_div,
     clear_tensor_data,
@@ -998,7 +998,16 @@ def fuser_forward(
         # Tensor properties
         fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features)
         fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features)
-        input_ = input_.reshape(-1, fc1_weight_shape[1])
+        if isinstance(input_, GroupedTensor):
+            # GroupedTensor forbids reshape and is already in the canonical
+            # (total_tokens, in_features) layout; just validate the shape.
+            if input_.dim() != 2 or input_.size(-1) != fc1_weight_shape[1]:
+                raise ValueError(
+                    "GroupedTensor input must have shape (total_tokens, "
+                    f"{fc1_weight_shape[1]}), but got {tuple(input_.size())}."
+                )
+        else:
+            input_ = input_.reshape(-1, fc1_weight_shape[1])
         in_shape = list(input_.size())
         if in_shape[0] % 128 != 0:
             raise ValueError(f"Unsupported input shape for fused grouped MLP ({in_shape=}).")
@@ -1218,42 +1227,18 @@ def fuser_forward(
         )
         fc1_input_quantizer.optimize_for_gemm = True
         fc1_input_quantizer.internal = True
-        input_quantizer = getattr(input_, "quantizer", None)
-        if isinstance(input_, GroupedTensor) and (
-            isinstance(fc1_input_quantizer, MXFP8Quantizer)
-            and isinstance(input_quantizer, MXFP8Quantizer)
-            or isinstance(fc1_input_quantizer, NVFP4Quantizer)
-            and isinstance(input_quantizer, NVFP4Quantizer)
-        ):
-            # GroupedTensor is a torch.Tensor subclass, so the CPU offload
-            # infrastructure's prepare_for_saving treats it as a plain tensor
-            # and does not decompose it into its component data tensors.  By
-            # repacking into a GroupedTensorStorage (not a torch.Tensor), we
-            # ensure the fuser's prepare_for_saving call correctly decomposes
-            # the activation before save_for_backward.
-            grouped_fc1_x = GroupedTensorStorage(
-                shape=input_.logical_shape,
-                dtype=input_.fake_dtype,
-                num_tensors=input_.num_tensors,
-                shapes=input_.tensor_shapes,
-                quantizer=input_.quantizer,
-                data=input_.rowwise_data,
-                columnwise_data=input_.columnwise_data,
-                scale_inv=input_.scale_inv,
-                columnwise_scale_inv=input_.columnwise_scale_inv,
-                amax=input_.amax,
-                columnwise_amax=input_.columnwise_amax,
-                scale=input_.scale,
-                first_dims=input_.first_dims,
-                last_dims=input_.last_dims,
-                tensor_offsets=input_.tensor_offsets,
-                offsets=input_.offsets,
-                scale_inv_offsets=input_.scale_inv_offsets,
-                columnwise_scale_inv_offsets=input_.columnwise_scale_inv_offsets,
-                with_gemm_swizzled_scales=input_._with_gemm_swizzled_scales,
-                row_scaled_nvfp4=input_.row_scaled_nvfp4,
-                nvfp4_use_4over6=input_.nvfp4_use_4over6,
-                nvfp4_e4m3_max=input_.nvfp4_e4m3_max,
+        if isinstance(input_, GroupedTensor):
+            # Input arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise data
+            # for the GEMM and let the helper supply whatever else the GEMMs need. An input that
+            # is already GEMM-ready in both directions passes through untouched.
+            grouped_fc1_x = input_.copy()
+            tex.group_requantize_inplace(
+                grouped_fc1_x,
+                fc1_input_quantizer,
+                num_groups,
+                split_sizes,
+                TE_DType[dtype],
+                tensor_offsets=fc1_x_tensor_offsets,
             )
         else:
             fc1_x = maybe_dequantize(input_, dtype)
@@ -1809,7 +1794,16 @@ def fuser_backward(
         # Tensor properties
         fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features)
         fc2_weight_shape = (fc2_op.out_features, fc2_op.in_features)
-        grad_output = grad_output.reshape(-1, fc2_weight_shape[0])
+        if isinstance(grad_output, GroupedTensor):
+            # GroupedTensor forbids reshape and is already in the canonical
+            # (total_tokens, out_features) layout; just validate the shape.
+            if grad_output.dim() != 2 or grad_output.size(-1) != fc2_weight_shape[0]:
+                raise ValueError(
+                    "GroupedTensor grad output must have shape (total_tokens, "
+                    f"{fc2_weight_shape[0]}), but got {tuple(grad_output.size())}."
+                )
+        else:
+            grad_output = grad_output.reshape(-1, fc2_weight_shape[0])
         out_shape = list(grad_output.size())
         num_groups = fc1_op.num_groups
         fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0
@@ -1879,20 +1873,28 @@ def fuser_backward(
         output_fc2_dbias = fc2_op.has_bias
         fc2_dbias_packed = None
         fc2_dy = None
-        grad_output_quantizer = getattr(grad_output, "quantizer", None)
-        fc2_grad_output_quantizer_matches = (
-            isinstance(fc2_grad_output_quantizer, MXFP8Quantizer)
-            and isinstance(grad_output_quantizer, MXFP8Quantizer)
-        ) or (
-            isinstance(fc2_grad_output_quantizer, NVFP4Quantizer)
-            and isinstance(grad_output_quantizer, NVFP4Quantizer)
-        )
-        if (
-            not output_fc2_dbias
-            and isinstance(grad_output, GroupedTensor)
-            and fc2_grad_output_quantizer_matches
-        ):
-            grouped_fc2_dy = grad_output
+        if isinstance(grad_output, GroupedTensor):
+            # Grad output arrived already quantized (e.g. FP8 token dispatch): reuse its rowwise
+            # data for the dgrad GEMM. Bias grads are reduced from the dequantized grad, which is
+            # only materialized when one is needed. A grad that is already GEMM-ready in both
+            # directions passes through untouched.
+            grouped_fc2_dy = grad_output.copy()
+            fc2_dy = tex.group_requantize_inplace(
+                grouped_fc2_dy,
+                fc2_grad_output_quantizer,
+                num_groups,
+                split_sizes,
+                TE_DType[dtype],
+                tensor_offsets=fc2_out_tensor_offsets,
+                return_dequantized=output_fc2_dbias or scale_bias,
+            )
+            if output_fc2_dbias and not scale_bias:
+                # This path has no quantize kernel to fuse dbias into, and the consumer below
+                # has no fallback, so reduce it here.
+                fc2_dbias_packed = compute_grouped_dbias(fc2_dy, base_split_offsets, num_groups)
+                # scale_bias is the only later consumer of the dequantized grad; drop it so the
+                # buffer is freed rather than held until backward ends.
+                fc2_dy = None
         else:
             fc2_dy = maybe_dequantize(grad_output, dtype)
             if output_fc2_dbias and not scale_bias:

From 8260f49660cbadb78bc52c90449428c51625469d Mon Sep 17 00:00:00 2001
From: fheinecke <23390735+fheinecke@users.noreply.github.com>
Date: Fri, 7 Aug 2026 00:44:50 -0500
Subject: [PATCH 46/67] Use a temp directory for CMake build (#3254)

* Use a temp directory for CMake build

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

* Add fallback to exisiting build dir behavior when using an editable install

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

---------

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
---
 build_tools/build_ext.py | 41 +++++++++++++++++++++++++++-------------
 1 file changed, 28 insertions(+), 13 deletions(-)

diff --git a/build_tools/build_ext.py b/build_tools/build_ext.py
index cbb8838b00..30ea83a55b 100644
--- a/build_tools/build_ext.py
+++ b/build_tools/build_ext.py
@@ -4,11 +4,12 @@
 
 """Installation script."""
 
+import copy
 import os
 import subprocess
 import sys
 import sysconfig
-import copy
+import tempfile
 import time
 
 from pathlib import Path
@@ -111,23 +112,37 @@ def run(self) -> None:
             for ext in self.extensions:
                 package_path = Path(self.get_ext_fullpath(ext.name))
                 install_dir = package_path.resolve().parent
-                if isinstance(ext, CMakeExtension):
-                    print(f"Building CMake extension {ext.name}")
-                    # Set up incremental builds for CMake extensions
-                    build_dir = os.getenv("NVTE_CMAKE_BUILD_DIR")
-                    if build_dir:
-                        build_dir = Path(build_dir).resolve()
-                    else:
-                        root_dir = Path(__file__).resolve().parent.parent
-                        build_dir = root_dir / "build" / "cmake"
-
-                    # Ensure the directory exists
+                if not isinstance(ext, CMakeExtension):
+                    continue
+
+                print(f"Building CMake extension {ext.name}")
+                configured_build_dir = os.getenv("NVTE_CMAKE_BUILD_DIR")
+                if not configured_build_dir and self.inplace:
+                    root_dir = Path(__file__).resolve().parent.parent
+                    configured_build_dir = root_dir / "build" / "cmake"
+
+                if configured_build_dir:
+                    # A persistent build directory enables incremental builds.
+                    build_dir = Path(configured_build_dir).resolve()
                     build_dir.mkdir(parents=True, exist_ok=True)
-
                     ext._build_cmake(
                         build_dir=build_dir,
                         install_dir=install_dir,
                     )
+                    continue
+
+                # Isolate CMake state between concurrent and successive builds.
+                build_temp = Path(self.build_temp)
+                build_temp.mkdir(parents=True, exist_ok=True)
+                with tempfile.TemporaryDirectory(
+                    prefix=f"cmake-build-{ext.name}-",
+                    dir=build_temp,
+                ) as build_dir:
+                    print(f"Building CMake extension {ext.name} in temporary directory {build_dir}")
+                    ext._build_cmake(
+                        build_dir=Path(build_dir),
+                        install_dir=install_dir,
+                    )
 
             # Build non-CMake extensions as usual
             all_extensions = self.extensions

From c4a38bc392dc23b43963465cc91b15dda5c62334 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 7 Aug 2026 16:14:53 +0200
Subject: [PATCH 47/67] Rename _maybe_unshare_cu_seqlens to _unalias_cu_seqlens

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/attention/dot_product_attention/backends.py     | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index 71afed25c5..3ccdd9d0b9 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -803,7 +803,7 @@ def backward(
         return dq, dk, dv
 
 
-def _maybe_unshare_cu_seqlens(cu_q: torch.Tensor, cu_kv: torch.Tensor):
+def _unalias_cu_seqlens(cu_q: torch.Tensor, cu_kv: torch.Tensor):
     """Copy kv's cumulative sequence lengths when they are q's as well.
 
     Self attention passes one tensor for both, and flash-attn forwards them to
@@ -1144,7 +1144,7 @@ def forward(
                     else:
                         func = flash_attn_with_kvcache_v3  # pylint: disable=possibly-used-before-assignment
                     if not use_flash_attn_4 and (not use_flash_attn_3 or inference_params is None):
-                        cu_q, cu_kv = _maybe_unshare_cu_seqlens(
+                        cu_q, cu_kv = _unalias_cu_seqlens(
                             cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q,
                             cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv,
                         )
@@ -1160,7 +1160,7 @@ def forward(
                     if inference_params is None:
                         fa_4_optional_forward_kwargs["deterministic"] = self.deterministic
                     if func is flash_attn_varlen_func_v4:
-                        cu_q, cu_kv = _maybe_unshare_cu_seqlens(cu_seqlens_q, cu_seqlens_kv)
+                        cu_q, cu_kv = _unalias_cu_seqlens(cu_seqlens_q, cu_seqlens_kv)
                         fa_4_optional_forward_kwargs["cu_seqlens_q"] = cu_q
                         fa_4_optional_forward_kwargs["cu_seqlens_k"] = cu_kv
                         fa_4_optional_forward_kwargs["max_seqlen_q"] = max_seqlen_q

From 1b495ce2616a769772b823e58c0ba84365697488 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 7 Aug 2026 16:24:37 +0200
Subject: [PATCH 48/67] Split the get_full_cu_seqlens cache bypass per reason

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/attention/dot_product_attention/utils.py     | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index a2d260c6b2..2b60c5af24 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -2036,9 +2036,12 @@ def _get_cu_seqlens(batch_size, max_seqlen, device):
             device=device,
         )
 
-    if is_in_onnx_export_mode() or torch.compiler.is_compiling():
-        # The cache only saves an allocation, and its key -- is_inference_mode_enabled()
-        # -- is a fundamental graph break.
+    if is_in_onnx_export_mode():
+        # A tensor cached by an earlier call would be baked into the exported graph.
+        return _get_cu_seqlens(batch_size, max_seqlen, device)
+    if torch.compiler.is_compiling():
+        # torch.is_inference_mode_enabled(), part of the cache key, graph-breaks, and the
+        # cache only saves one arange.
         return _get_cu_seqlens(batch_size, max_seqlen, device)
 
     is_inference = torch.is_inference_mode_enabled()

From 731f242a654bca24fea7a1adb798c144c72f267d Mon Sep 17 00:00:00 2001
From: nvegesna-netizen <nvegesna@nvidia.com>
Date: Mon, 10 Aug 2026 03:33:39 -0700
Subject: [PATCH 49/67] [PyTorch] Preserve FP8 recompute state for inner
 autocast (#3284)

* fix: preserve FP8 recompute state for inner autocast

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

* test: fix activation recompute test license header

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

* test: run FP8 recompute coverage in PyTorch QA

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

* test: compare inner FP8 autocast recompute numerics

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

* refactor: rename FP8 recompute flag to _IN_ACTIVATION_RECOMPUTE_REGION

The global no longer encodes FP8 state now that the FP8 gate lives in
is_fp8_activation_recompute_enabled(); it only marks the checkpoint
region. Rename it to match. The public getter keeps its name since it
now returns the conjunction of the region flag and the current FP8
state.

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

* refactor: rename recompute phase global for symmetry

_FP8_ACTIVATION_RECOMPUTE_PHASE carries no FP8 state either -
in_fp8_activation_recompute_phase() returns it ungated and callers
apply their own FP8 gate. Rename to _ACTIVATION_RECOMPUTE_PHASE to
match _IN_ACTIVATION_RECOMPUTE_REGION. Both public getters keep their
names.

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

* test: restore dropped FP8 recompute coverage and assert stash bookkeeping

Moving the regression into test_numerics.py silently dropped two of the
three original tests. Restore both: the non-FP8 negative case (which
needs no FP8 hardware and is the direct guard for making the region flag
FP8-agnostic) and the mixed FP8/non-FP8 region case.

The inner-autocast test could not observe a missing stash: the outer
autocast is disabled, so autocast_depth never returns to 0 with FP8
enabled, reduce_and_update_fp8_tensors is never called, and the forward
scale stays at 1.0 - making stashed and unstashed recompute identical.
Record the stash/restore of the forward scale and assert every stash is
restored exactly once, and use an inner-autocast reference so the two
runs differ only by recompute.

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

* test: make recompute stash assertions meaningful, guard bf16

The scale equality assertion was vacuous: record_restore read the scale
after the real restore had already copied the stashed value into it, so
it compared a tensor with itself, and every scale is 1.0 here anyway.
Count stash and restore per module instead, and check a module was
stashed before the real restore runs so a regression reports that rather
than a KeyError from the recompute buffer lookup.

Also skip the non-FP8 checkpoint test when bf16 is unavailable (it has no
FP8 skipif, so it would otherwise run on pre-Ampere), and correct the
is_fp8_activation_recompute_enabled docstring, which still claimed to
return a bare global.

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

* test: assert the recompute reinstalls the stashed FP8 state

Counting stashes and restores catches a missing stash but not a wrong
one. The outputs cannot catch either: the forward scale never moves here
(autocast_depth stays >= 1 so reduce_and_update_fp8_tensors never fires),
and a wrong scale would only perturb FP8 rounding regardless, since
scale_inv is derived from the same scale at cast time.

Compare the state each restore installs against a clone captured at stash
time, for both the scale and the amax history, and assert at least one
restore actually moved the live state so the comparison is not vacuous.
This is exact rather than tolerance-based and works for both reentrant
and non-reentrant checkpointing.

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

* test: drop tautological recompute state assertions

Comparing the restored state against a clone taken in the stash wrapper
could not fail: the real stash clones the same tensors at the same
instant, and the restore copies that entry back, so both sides derive
from one snapshot. It exercised the stash/restore plumbing rather than
anything this change touches.

Keep the checks that do have teeth - that some module stashed, and that
stashes and restores balance per module - and skip the bookkeeping for
modules whose recipe is not delayed scaling, since the restore site is
not gated on it while the stash site is.

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

* test: pin the changed predicate directly

Replace test_checkpoint_without_fp8_does_not_save_fp8_recompute_state,
which could not fail: with no autocast self.fp8 is False, so the stash
site is unreachable before the region flag is ever consulted, and the
assertion holds no matter what the getter returns. The mixed-region test
already makes the same negative claim in an FP8-enabled session.

Assert instead on is_fp8_activation_recompute_enabled() itself from
inside a checkpointed callable that opens its own autocast: True in both
phases inside the autocast, False outside it. The first half fails before
the fix; the second pins the FP8 term against a later over-correction
that drops it.

Also move the fp8_meta key constant up to the other module constants.

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>

---------

Signed-off-by: Nitin Vegesna <nvegesna@nvidia.com>
---
 tests/pytorch/test_numerics.py            | 176 +++++++++++++++++++++-
 transformer_engine/pytorch/distributed.py |  29 ++--
 2 files changed, 188 insertions(+), 17 deletions(-)

diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
index 8249c7fedd..e6a83d92bc 100644
--- a/tests/pytorch/test_numerics.py
+++ b/tests/pytorch/test_numerics.py
@@ -42,6 +42,10 @@
     is_nvfp4_available,
 )
 from transformer_engine.pytorch import checkpoint as te_checkpoint
+from transformer_engine.pytorch.distributed import (
+    is_fp8_activation_recompute_enabled,
+    in_fp8_activation_recompute_phase,
+)
 from transformer_engine.pytorch.cpp_extensions import general_gemm
 from transformer_engine.common import recipe
 from transformer_engine.pytorch import DType
@@ -82,6 +86,9 @@
 
 all_boolean = [True, False]
 
+# fp8_meta key written by FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute
+_FP8_RECOMPUTE_KEY = "global_fp8_buffer_pos_fwd_recompute"
+
 all_activations = [
     "gelu",
     "geglu",
@@ -648,7 +655,15 @@ def test_gpt_selective_activation_recompute(dtype, bs, model, fp8, recipe, fp8_m
 
 
 def _test_e2e_full_recompute(
-    bs, dtype, config, fp8, recipe, fp8_model_params=False, recompute=False, use_reentrant=True
+    bs,
+    dtype,
+    config,
+    fp8,
+    recipe,
+    fp8_model_params=False,
+    recompute=False,
+    use_reentrant=True,
+    inner_autocast=False,
 ):
     reset_rng_states()
     FP8GlobalStateManager.reset()
@@ -685,10 +700,17 @@ def _test_e2e_full_recompute(
         te_inp_hidden_states.retain_grad()
     te_inp_attn_mask = get_causal_attn_mask(config.max_seqlen_q)
 
-    with autocast(enabled=fp8, recipe=recipe):
+    forward = block
+    if inner_autocast:
+
+        def forward(*args, **kwargs):
+            with autocast(enabled=fp8, recipe=recipe):
+                return block(*args, **kwargs)
+
+    with autocast(enabled=fp8 and not inner_autocast, recipe=recipe):
         if recompute:
             te_out = te_checkpoint(
-                block,
+                forward,
                 te_inp_hidden_states,
                 attention_mask=te_inp_attn_mask,
                 checkpoint_core_attention=False,
@@ -697,7 +719,7 @@ def _test_e2e_full_recompute(
                 use_reentrant=use_reentrant,
             )
         else:
-            te_out = block(
+            te_out = forward(
                 te_inp_hidden_states,
                 attention_mask=te_inp_attn_mask,
                 checkpoint_core_attention=False,
@@ -787,6 +809,152 @@ def test_gpt_full_activation_recompute(
         )
 
 
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("use_reentrant", all_boolean)
+def test_gpt_full_activation_recompute_with_inner_autocast(use_reentrant, monkeypatch):
+    """Check recompute numerics when FP8 autocast starts inside the checkpointed callable."""
+    if not use_reentrant:
+        # Non-reentrant checkpoint becomes non-deterministic with bias+GELU fusion.
+        monkeypatch.setenv("NVTE_BIAS_GELU_NVFUSION", "0")
+
+    dtype = torch.bfloat16
+    fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID)
+    config = model_configs["126m"]
+
+    # Reference also opens the autocast inside the callable, so the only difference
+    # between the two runs is activation recompute.
+    outputs, names = _test_e2e_full_recompute(
+        1,
+        dtype,
+        config,
+        True,
+        fp8_recipe,
+        recompute=False,
+        use_reentrant=use_reentrant,
+        inner_autocast=True,
+    )
+
+    # Before the fix, phase 1 skipped the stash while the recompute still restored it, which
+    # surfaced as a KeyError on the recompute buffer lookup. Count both sides to pin that down.
+    stash_counts, restore_counts = {}, {}
+    stash_fn = FP8GlobalStateManager.copy_forward_fp8_meta_tensors_for_recompute
+    restore_fn = FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute
+
+    def record_stash(fp8_meta):
+        stash_fn(fp8_meta)
+        if _FP8_RECOMPUTE_KEY in fp8_meta:
+            stash_counts[id(fp8_meta)] = stash_counts.get(id(fp8_meta), 0) + 1
+
+    def record_restore(fp8_meta):
+        # The restore site is not gated on delayed scaling, but only delayed scaling stashes.
+        if not fp8_meta["recipe"].delayed():
+            restore_fn(fp8_meta)
+            return
+        key = id(fp8_meta)
+        assert key in stash_counts, "Recompute restored a scale that was never stashed"
+        restore_counts[key] = restore_counts.get(key, 0) + 1
+        restore_fn(fp8_meta)
+
+    monkeypatch.setattr(
+        FP8GlobalStateManager,
+        "copy_forward_fp8_meta_tensors_for_recompute",
+        staticmethod(record_stash),
+    )
+    monkeypatch.setattr(
+        FP8GlobalStateManager,
+        "get_old_fp8_meta_tensors_for_recompute",
+        staticmethod(record_restore),
+    )
+
+    outputs_recompute, _ = _test_e2e_full_recompute(
+        1,
+        dtype,
+        config,
+        True,
+        fp8_recipe,
+        recompute=True,
+        use_reentrant=use_reentrant,
+        inner_autocast=True,
+    )
+
+    assert stash_counts, "No FP8 module stashed a forward scale for the recompute phase"
+    assert restore_counts == stash_counts, "Stash and restore of forward scales are unbalanced"
+
+    for name, ref, test in zip(names, outputs, outputs_recompute):
+        torch.testing.assert_close(
+            test,
+            ref,
+            msg=f"Mismatch in tensor {name}",
+            rtol=0.125,
+            atol=0.0675,
+        )
+
+
+def _checkpointed_linear_backward(body, use_reentrant, *layers):
+    """Run a checkpointed callable end to end and check the gradients are finite."""
+    inp = torch.randn(16, 16, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    with torch.autocast("cuda", dtype=torch.bfloat16):
+        out = te_checkpoint(body, inp, use_reentrant=use_reentrant)
+        loss = out.float().sum()
+    loss.backward()
+    torch.cuda.synchronize()
+
+    assert torch.isfinite(loss)
+    assert inp.grad is not None and torch.isfinite(inp.grad).all()
+    for layer in layers:
+        assert layer.weight.grad is not None
+        assert torch.isfinite(layer.weight.grad).all()
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("use_reentrant", all_boolean)
+def test_checkpoint_inner_autocast_is_an_fp8_recompute_region(use_reentrant):
+    """An FP8 autocast opened inside a checkpointed callable is an FP8 recompute region."""
+    FP8GlobalStateManager.reset()
+    fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID)
+    layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda()
+
+    observed = []
+
+    def body(value):
+        outside = is_fp8_activation_recompute_enabled()
+        with autocast(enabled=True, recipe=fp8_recipe):
+            observed.append(
+                (
+                    outside,
+                    is_fp8_activation_recompute_enabled(),
+                    in_fp8_activation_recompute_phase(),
+                )
+            )
+            return layer(value)
+
+    _checkpointed_linear_backward(body, use_reentrant, layer)
+
+    # One entry for the checkpointed forward, one for the recompute during backward. The
+    # query is only an FP8 recompute region inside the autocast, in both phases.
+    assert observed == [(False, True, False), (False, True, True)]
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("use_reentrant", all_boolean)
+def test_checkpoint_with_mixed_fp8_regions_saves_only_fp8_recompute_state(use_reentrant):
+    """Only the inner FP8 region of a mixed checkpoint saves recompute metadata."""
+    FP8GlobalStateManager.reset()
+    fp8_recipe = recipe.DelayedScaling(fp8_format=recipe.Format.HYBRID)
+    non_fp8_layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda()
+    fp8_layer = Linear(16, 16, bias=False, params_dtype=torch.float32).cuda()
+
+    def body(value):
+        value = non_fp8_layer(value)
+        with autocast(enabled=True, recipe=fp8_recipe):
+            return fp8_layer(value)
+
+    _checkpointed_linear_backward(body, use_reentrant, non_fp8_layer, fp8_layer)
+
+    assert _FP8_RECOMPUTE_KEY not in non_fp8_layer.fp8_meta
+    assert _FP8_RECOMPUTE_KEY in fp8_layer.fp8_meta
+
+
 def _test_e2e_checkpointing_get_model(config, dtype):
     sigma = 0.023
     init_method = init_method_normal(sigma)
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index d1525b53f0..2dddb6587e 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -62,8 +62,8 @@
 
 _USE_REENTRANT_ACTIVATION_RECOMPUTE = True
 
-_FP8_ACTIVATION_RECOMPUTE_ENABLED = False
-_FP8_ACTIVATION_RECOMPUTE_PHASE = False
+_IN_ACTIVATION_RECOMPUTE_REGION = False
+_ACTIVATION_RECOMPUTE_PHASE = False
 
 
 _ALL_ACTIVE_RNG_STATES = {}
@@ -255,11 +255,14 @@ def __init__(self, activation_recompute: bool = False, recompute_phase: bool = F
         self.recompute_phase = recompute_phase
 
     def __enter__(self):
-        global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE
-        _FP8_ACTIVATION_RECOMPUTE_ENABLED = (
-            self.activation_recompute and FP8GlobalStateManager.is_fp8_enabled()
-        )
-        _FP8_ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase
+        global _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE
+        # Track the checkpoint region independently of the FP8 state at entry.
+        # A checkpointed callable may open its own FP8 autocast context (for
+        # example, to select precision per layer). Delayed-scaling modules in
+        # that inner context must still save their scale and amax metadata for
+        # the recompute forward.
+        _IN_ACTIVATION_RECOMPUTE_REGION = self.activation_recompute
+        _ACTIVATION_RECOMPUTE_PHASE = self.recompute_phase
 
         qstate = FP8GlobalStateManager.quantization_state
         if self.activation_recompute and not self.recompute_phase:
@@ -268,19 +271,19 @@ def __enter__(self):
             qstate.is_first_fp8_module = activation_recompute_forward._is_first_fp8_module.pop(0)
 
     def __exit__(self, *exc_details):
-        global _FP8_ACTIVATION_RECOMPUTE_ENABLED, _FP8_ACTIVATION_RECOMPUTE_PHASE
-        _FP8_ACTIVATION_RECOMPUTE_ENABLED = False
-        _FP8_ACTIVATION_RECOMPUTE_PHASE = False
+        global _IN_ACTIVATION_RECOMPUTE_REGION, _ACTIVATION_RECOMPUTE_PHASE
+        _IN_ACTIVATION_RECOMPUTE_REGION = False
+        _ACTIVATION_RECOMPUTE_PHASE = False
 
 
 def is_fp8_activation_recompute_enabled() -> bool:
-    """Return global boolean"""
-    return _FP8_ACTIVATION_RECOMPUTE_ENABLED
+    """Whether we are in an activation recompute region with FP8 currently enabled"""
+    return _IN_ACTIVATION_RECOMPUTE_REGION and FP8GlobalStateManager.is_fp8_enabled()
 
 
 def in_fp8_activation_recompute_phase() -> bool:
     """Return global boolean"""
-    return _FP8_ACTIVATION_RECOMPUTE_PHASE
+    return _ACTIVATION_RECOMPUTE_PHASE
 
 
 def _get_active_autocast_contexts():

From 32f7096a355ca803e7151da0d0f4aa65ac4ba46b Mon Sep 17 00:00:00 2001
From: Phuong Nguyen <phuonguyen@nvidia.com>
Date: Mon, 10 Aug 2026 08:44:14 -0700
Subject: [PATCH 50/67] [JAX][Fix] Gate collective-stream compute_on on import
 and silence pylint (#3329)

[JAX] Gate collective-stream compute_on on import and silence pylint not-callable

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---
 transformer_engine/jax/cpp_extensions/ep.py | 2 +-
 transformer_engine/jax/version_utils.py     | 9 ++++++++-
 2 files changed, 9 insertions(+), 2 deletions(-)

diff --git a/transformer_engine/jax/cpp_extensions/ep.py b/transformer_engine/jax/cpp_extensions/ep.py
index 0b18c3bbb6..ca70ea145c 100644
--- a/transformer_engine/jax/cpp_extensions/ep.py
+++ b/transformer_engine/jax/cpp_extensions/ep.py
@@ -42,7 +42,7 @@ def wrapper(*args, **kwargs):
         # a nullary thunk that closes over them; the array operands are captured
         # as consts and lifted to real operands, outputs stay on device. XLA
         # async-wraps the resulting call onto the collective stream.
-        annotated = compute_on(
+        annotated = compute_on(  # pylint: disable=not-callable
             compute_type="gpu_stream:collective",
             out_memory_spaces=jax.memory.Space.Device,
         )(lambda: func(*args, **kwargs))
diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py
index e797e83440..9d73ecd3b5 100644
--- a/transformer_engine/jax/version_utils.py
+++ b/transformer_engine/jax/version_utils.py
@@ -69,9 +69,16 @@ def is_triton_autotuned_alias_safe() -> bool:
 _COLLECTIVE_STREAM_MIN_JAX_VERSION = "0.10.1"
 
 
+@lru_cache(maxsize=None)
 def is_collective_stream_supported() -> bool:
     """Return True if the installed JAX supports the gpu_stream:collective annotation."""
-    return jax_version_meet_requirement(_COLLECTIVE_STREAM_MIN_JAX_VERSION)
+    if not jax_version_meet_requirement(_COLLECTIVE_STREAM_MIN_JAX_VERSION):
+        return False
+    try:
+        from jax.experimental.compute_on import compute_on  # pylint: disable=unused-import
+    except ImportError:
+        return False
+    return True
 
 
 def is_triton_extension_supported() -> bool:

From 598b9eacbe9fc34ec105cf8c12f303108ca434ca Mon Sep 17 00:00:00 2001
From: vthumbe1503 <vthumbe@nvidia.com>
Date: Mon, 10 Aug 2026 14:12:15 -0700
Subject: [PATCH 51/67] [Bug] use_single_group_runtime_offsets arg is not
 supported for ScaledSrelu (#3328)

fix cudnn version gurad

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Co-authored-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
---
 tests/pytorch/test_grouped_mlp.py                |  6 ++++--
 .../pytorch/ops/fused/grouped_mlp.py             | 16 ++++++++++------
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py
index d09c92ad49..499c22509c 100644
--- a/tests/pytorch/test_grouped_mlp.py
+++ b/tests/pytorch/test_grouped_mlp.py
@@ -1559,13 +1559,15 @@ def test_grouped_mlp_single_group_mxfp8(
         """Single-group GroupedLinear + ScaledSwiGLU + GroupedLinear with MXFP8."""
         if (
             runtime_offsets_supported
-            and not grouped_mlp_module._cudnn_frontend_supports_single_group_runtime_offsets()
+            and not grouped_mlp_module._cudnn_frontend_supports_single_group_runtime_offsets(
+                te.ops.ScaledSwiGLU
+            )
         ):
             pytest.skip("Requires cuDNN frontend >= 1.27.0")
         monkeypatch.setattr(
             grouped_mlp_module,
             "_cudnn_frontend_supports_single_group_runtime_offsets",
-            lambda: runtime_offsets_supported,
+            lambda _activation_type: runtime_offsets_supported,
         )
         self.test_grouped_mlp(
             group_size=1,
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 76d51673f0..909e5a8a9b 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -101,9 +101,13 @@ def _nvidia_cudnn_frontend_supports_wgrad() -> bool:
     return _cudnn_frontend_version_supported()
 
 
-def _cudnn_frontend_supports_single_group_runtime_offsets() -> bool:
-    """Check cuDNN FE min version for single-group runtime offsets."""
-    return _cudnn_frontend_version_at_least("1.27.0")
+def _cudnn_frontend_supports_single_group_runtime_offsets(
+    activation_type: type[FusibleOperation],
+) -> bool:
+    """Check cuDNN FE support for single-group runtime offsets."""
+    return not issubclass(activation_type, ScaledSReLU) and _cudnn_frontend_version_at_least(
+        "1.27.0"
+    )
 
 
 def _wrap_single_quantized_as_grouped(
@@ -1073,7 +1077,7 @@ def fuser_forward(
 
         activation_kernel = self.grouped_gemm_activation_kernel()
         supports_single_group_runtime_offsets = (
-            _cudnn_frontend_supports_single_group_runtime_offsets()
+            _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op))
         )
 
         # Shared experts have one dense group and all optimized kernels derive M
@@ -2032,7 +2036,7 @@ def fuser_backward(
             "use_dynamic_sched": True,
         }
         dactivation_kernel = self.grouped_gemm_dactivation_kernel()
-        if _cudnn_frontend_supports_single_group_runtime_offsets():
+        if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)):
             fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
         if self._cudnn_dact_func is not None:
             fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor
@@ -2377,7 +2381,7 @@ def fuser_backward(
                     "use_dynamic_sched": True,
                 }
                 fc1_dgrad_kernel = self.grouped_gemm_quant_kernel()
-                if _cudnn_frontend_supports_single_group_runtime_offsets():
+                if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)):
                     fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
 
                 if fc1_op.single_grouped_weight:

From f07a86029d3831d5512e26c3e0791566a72e14f0 Mon Sep 17 00:00:00 2001
From: Kaining Zhong <44538064+kainzhong@users.noreply.github.com>
Date: Mon, 10 Aug 2026 16:37:25 -0700
Subject: [PATCH 52/67] [Common] Fix pointer arithmatic to generate correct
 LDS/STS instructions (#3291)

* [Common] Fix pointer arithmatic to generate correct LDS/STS instructions

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* remove redundant comments

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

* move helper function to common.h to avoid NVRTC issue

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>

---------

Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .../common/cast/fp8/gated_fp8.cuh             |  4 +---
 .../common/cast/mxfp8/gated_mxfp8.cuh         |  4 +---
 .../common/cast/mxfp8/quantize_mxfp8.cuh      |  4 +---
 .../cast/mxfp8/specialized/quantize_mxfp8.cuh | 10 ++--------
 .../nvfp4/group_quantize_transpose_nvfp4.cuh  |  4 +---
 .../cast/nvfp4/quantize_transpose_nvfp4.cuh   |  8 ++------
 transformer_engine/common/common.h            | 20 +++++++++++++++++++
 7 files changed, 28 insertions(+), 26 deletions(-)

diff --git a/transformer_engine/common/cast/fp8/gated_fp8.cuh b/transformer_engine/common/cast/fp8/gated_fp8.cuh
index 522a9add8f..52ab6f49e9 100644
--- a/transformer_engine/common/cast/fp8/gated_fp8.cuh
+++ b/transformer_engine/common/cast/fp8/gated_fp8.cuh
@@ -68,11 +68,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
   const float scale = (scale_ptr != nullptr) ? *scale_ptr : 1;
 
   extern __shared__ char dynamic_shmem[];
-  uintptr_t base_shmem_ptr = reinterpret_cast<uintptr_t>(dynamic_shmem);
   // Manually align dynamic SHMEM per TMA requirements using padding
   // __align__(128) Does not guarantee the pointer to be aligned!
-  uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) &
-                     ~(static_cast<uintptr_t>(TMA_SHMEM_ALIGNMENT - 1));
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   constexpr size_t buff_elems = SHMEM_DIM_Y * SHMEM_DIM_X;
   constexpr size_t buff_elems_total = BUFFERS_NUM * buff_elems;
diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh
index 83b5a49cae..879c6befc3 100644
--- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh
@@ -124,11 +124,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
   __shared__ float subamax_colwise_buff[SUBAMAX_BUFF_DIM_Y][CHUNK_DIM_X];
 
   extern __shared__ char dynamic_shmem[];
-  uintptr_t base_shmem_ptr = reinterpret_cast<uintptr_t>(dynamic_shmem);
   // Manually align dynamic SHMEM per TMA requirements using padding
   // __align__(128) Does not guarantee the pointer to be aligned!
-  uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) &
-                     ~(static_cast<uintptr_t>(TMA_SHMEM_ALIGNMENT - 1));
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X;
   constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems;
diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
index 8c57f112d2..9f84195a19 100644
--- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
@@ -130,11 +130,9 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK)
   constexpr size_t out_mem_rowwise = (ROWWISE_SCALING ? buff_size_aligned_out : 0);
 
   extern __shared__ char dynamic_shmem[];
-  uintptr_t base_shmem_ptr = reinterpret_cast<uintptr_t>(dynamic_shmem);
   // Manually align dynamic SHMEM per TMA requirements using padding
   // __align__(128) Does not guarantee the pointer to be aligned!
-  uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) &
-                     ~(static_cast<uintptr_t>(TMA_SHMEM_ALIGNMENT - 1));
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned
   IType *in_sh = reinterpret_cast<IType *>(dshmem);
diff --git a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
index 81e11c40b8..0c59733d31 100644
--- a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
@@ -670,10 +670,6 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> {
                                         smem_alignment + smem_rowwise_scale + smem_colwise_reduce);
 };
 
-__device__ __forceinline__ intptr_t align_to(intptr_t x, intptr_t align) {
-  return (x + align - 1) & ~((align)-1);
-}
-
 // 32x32
 template <typename CastTraits,
           std::enable_if_t<CastTraits::isRowwise && CastTraits::isColwise, int> = 0,
@@ -706,8 +702,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(
   block_coords.x = blockIdx.x * CastTraits::blockDIM::N;
 
   extern __shared__ char smem[];
-  char *smemAligned = reinterpret_cast<char *>(
-      align_to(reinterpret_cast<intptr_t>(smem), CastTraits::smem_alignment));
+  char *smemAligned = align_up(smem, CastTraits::smem_alignment);
 
   IType *sInput = reinterpret_cast<IType *>(smemAligned);
   inputUnitType *sInputUnit = reinterpret_cast<inputUnitType *>(sInput);
@@ -1179,8 +1174,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(
   block_coords.x = blockIdx.x * CastTraits::blockDIM::N;
 
   extern __shared__ char smem[];
-  char *smemAligned = reinterpret_cast<char *>(
-      align_to(reinterpret_cast<intptr_t>(smem), CastTraits::smem_alignment));
+  char *smemAligned = align_up(smem, CastTraits::smem_alignment);
   IType *sInput = reinterpret_cast<IType *>(smemAligned);
   inputUnitType *sInputUnit = reinterpret_cast<inputUnitType *>(sInput);
 
diff --git a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh
index 91c6af26b5..0b980c5086 100644
--- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh
+++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh
@@ -262,11 +262,9 @@ __global__ void __launch_bounds__(THREADS_NUM)
   constexpr size_t out_mem_rowwise_scales = 0;
 
   extern __shared__ char dynamic_shmem[];
-  uintptr_t base_shmem_ptr = reinterpret_cast<uintptr_t>(dynamic_shmem);
   // Manually align dynamic SHMEM per TMA requirements using padding
   // __align__(128) Does not guarantee the pointer to be aligned!
-  uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) &
-                     ~(static_cast<uintptr_t>(TMA_SHMEM_ALIGNMENT - 1));
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned
   IType *in_sh = reinterpret_cast<IType *>(dshmem);
diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh
index a38a620ebe..8fc4cc7c46 100644
--- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh
+++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh
@@ -410,11 +410,9 @@ __global__ void __launch_bounds__(THREADS_NUM)
   constexpr size_t out_mem_rowwise_scales = 0;
 
   extern __shared__ char dynamic_shmem[];
-  uintptr_t base_shmem_ptr = reinterpret_cast<uintptr_t>(dynamic_shmem);
   // Manually align dynamic SHMEM per TMA requirements using padding
   // __align__(128) Does not guarantee the pointer to be aligned!
-  uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) &
-                     ~(static_cast<uintptr_t>(TMA_SHMEM_ALIGNMENT - 1));
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned
   IType *in_sh = reinterpret_cast<IType *>(dshmem);
@@ -952,11 +950,9 @@ __global__ void __launch_bounds__(THREADS_NUM)
   constexpr size_t out_mem_rowwise_scales = 0;
 
   extern __shared__ char dynamic_shmem[];
-  uintptr_t base_shmem_ptr = reinterpret_cast<uintptr_t>(dynamic_shmem);
   // Manually align dynamic SHMEM per TMA requirements using padding
   // __align__(128) Does not guarantee the pointer to be aligned!
-  uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) &
-                     ~(static_cast<uintptr_t>(TMA_SHMEM_ALIGNMENT - 1));
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned
   IType *in_sh = reinterpret_cast<IType *>(dshmem);
diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h
index eb4dcc055c..8fd7d7480c 100644
--- a/transformer_engine/common/common.h
+++ b/transformer_engine/common/common.h
@@ -1153,6 +1153,26 @@ inline bool is_aligned_ptr(const void *ptr, size_t alignment) {
   return reinterpret_cast<uintptr_t>(ptr) % alignment == 0;
 }
 
+/*! \brief Align a shared-memory base pointer up to `align` bytes.
+ *
+ * The result is derived from `p` by pointer arithmetic on purpose, without losing its
+ * identity as a pointer in between, so the address is never rounded through an integer
+ * -- in which case the compiler would lose the link back to the `extern __shared__`
+ * object, and ptxas could no longer prove the address lives in the shared window and
+ * would fall back to generic address-space accesses (`LD.E`/`ST.E`) instead of
+ * `LDS`/`STS`.
+ *
+ * `align` must be a power of two.
+ */
+__device__ __forceinline__ char *align_up(char *p, uintptr_t align) {
+  const uintptr_t misalign = reinterpret_cast<uintptr_t>(p) & (align - 1);
+  // If p is not aligned, (align - misalign) & (align - 1) is the number of bytes to fill the gap between p and
+  // the next aligned address.
+  // If p is aligned, misalign is 0 and (align - misalign) is align itself, so we use & (align - 1)
+  // to make it 0 and return p itself.
+  return p + ((align - misalign) & (align - 1));
+}
+
 inline bool is_aligned_tensor_data(const Tensor &t, size_t alignment) {
   return is_aligned_ptr(static_cast<const void *>(t.data.dptr), alignment);
 }

From 07e281f2ba93b61c5ab6145dbdaa2a768b888e19 Mon Sep 17 00:00:00 2001
From: Qiyu Wan <39144338+WanZzzzzz@users.noreply.github.com>
Date: Tue, 11 Aug 2026 04:28:21 -0700
Subject: [PATCH 53/67] [Common] MXFP8: Add swizzled-scale fast path for
 cast-only quantization (#3338)

* mxfp8: add swizzled-scale fast path for cast-only quantization

Produce GEMM-ready scales directly in the specialized cast-only kernels to avoid separate scale-swizzle launches while preserving generic fallbacks for unsupported shapes.

Signed-off-by: qiyuw <qiyuw@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* initializes padded scale entries

Signed-off-by: qiyuw <qiyuw@nvidia.com>

---------

Signed-off-by: qiyuw <qiyuw@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 tests/cpp/operator/test_cast_mxfp8.cu         | 238 +++++++
 .../test_mxfp8_quantize_swizzle_fusion.py     |  33 +
 .../common/cast/mxfp8/quantize_mxfp8.cuh      |  79 +--
 .../cast/mxfp8/specialized/quantize_mxfp8.cuh | 582 +++++++++++++++---
 4 files changed, 798 insertions(+), 134 deletions(-)

diff --git a/tests/cpp/operator/test_cast_mxfp8.cu b/tests/cpp/operator/test_cast_mxfp8.cu
index c7c778ce1e..edeb87ebe1 100644
--- a/tests/cpp/operator/test_cast_mxfp8.cu
+++ b/tests/cpp/operator/test_cast_mxfp8.cu
@@ -7,10 +7,12 @@
 #include <cuda_bf16.h>
 #include <cuda_fp8.h>
 #include <cuda_runtime.h>
+#include <cstring>
 #include <gtest/gtest.h>
 
 #include <transformer_engine/cast.h>
 #include <transformer_engine/activation.h>
+#include <transformer_engine/swizzle.h>
 #include "../test_common.h"
 #include "transformer_engine/transformer_engine.h"
 
@@ -735,3 +737,239 @@ INSTANTIATE_TEST_SUITE_P(
         ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2),
         ::testing::ValuesIn(input_scenarios)),
     test_name_generator);
+
+// ============================================================================
+// Swizzled-scales cast-only tests
+//
+// Validate the WITH_GEMM_SWIZZLED_SCALES=true code path added by the
+// CastTraitsSwizzle port. The specialized kernel dispatches to
+// CastTraitsSwizzle<..., kCacheColwise=true, kSwizzled=true> whenever the
+// output tensor has set_with_gemm_swizzled_scales(true), producing scales
+// directly in GEMM-swizzled layout.
+//
+// Reference construction: run the well-tested linear-scale path, then apply
+// nvte_swizzle_scaling_factors (independently tested by SwizzleTestSuite) to
+// transform the linear scales into the swizzled layout. Byte-compare the
+// direct swizzled cast output against this reference for both scale tensors
+// and the FP8 output data itself.
+//
+// Pass criteria per test:
+//   1. FP8 rowwise data byte-identical between linear and swizzled paths.
+//   2. FP8 colwise data byte-identical between linear and swizzled paths.
+//   3. Swizzled rowwise scale bytes match linear->swizzle reference.
+//   4. Swizzled colwise scale bytes match linear->swizzle reference.
+//   5. Rowwise-only FP8 data and swizzled scales match the same reference.
+//   6. No CUDA errors from any launch.
+// ============================================================================
+
+class SwizzledScalesFusedCastMXFP8TestSuite : public ::testing::TestWithParam<
+    std::tuple<std::vector<size_t>,
+               transformer_engine::DType,
+               transformer_engine::DType>> {};
+
+TEST_P(SwizzledScalesFusedCastMXFP8TestSuite, TestSwizzledCastMXFP8) {
+    if (getDeviceComputeCapability() < blackwellComputeCapability) {
+        GTEST_SKIP();
+    }
+
+    using namespace transformer_engine;
+    using namespace test;
+
+    const auto shape = std::get<0>(GetParam());
+    const DType itype = std::get<1>(GetParam());
+    const DType otype = std::get<2>(GetParam());
+
+    // BIDIMENSIONAL scaling needs a 2D+ input.
+    if (shape.size() < 2) {
+        GTEST_SKIP();
+    }
+
+    const size_t rows = first_dimension(shape);
+    const size_t cols = last_dimension(shape);
+    const size_t out_bytes = rows * cols;  // fp8 = 1 byte/elem
+
+    // Input filled with the same values for all casts.
+    Tensor input("input", shape, itype);
+    fillUniform(&input);
+
+    // Target: swizzled-scale cast. Dispatcher routes to
+    // CastTraitsSwizzle<..., kCacheColwise=true, kSwizzled=true> on kernel #3.
+    Tensor output_swizzled("output_swizzled", shape, otype, true, true, NVTE_MXFP8_1D_SCALING);
+    output_swizzled.set_with_gemm_swizzled_scales(true);
+    nvte_quantize(input.data(), output_swizzled.data(), 0);
+
+    // Rowwise-only activations take the pointer-based cast-only kernel. This
+    // directly exercises its WITH_GEMM_SWIZZLED_SCALES trait specialization
+    // for shapes whose column count is a multiple of 128.
+    Tensor output_rowwise_swizzled("output_rowwise_swizzled", shape, otype,
+                                   /*rowwise=*/true, /*colwise=*/false,
+                                   NVTE_MXFP8_1D_SCALING);
+    output_rowwise_swizzled.set_with_gemm_swizzled_scales(true);
+    nvte_quantize(input.data(), output_rowwise_swizzled.data(), 0);
+
+    cudaDeviceSynchronize();
+    ASSERT_EQ(cudaGetLastError(), cudaSuccess) << "swizzled-scale nvte_quantize failed";
+
+    // Reference construction: nvte_swizzle_scaling_factors accepts tensors with
+    // exactly one scale direction, so we build the rowwise and colwise references
+    // independently. Each is a single-direction linear cast followed by a
+    // single-direction swizzle transform, using the same input as the target.
+    // MXFP8 is a deterministic per-direction operation, so a rowwise-only cast
+    // produces the same rowwise fp8 bytes and scales as a rowwise+colwise cast.
+
+    // Rowwise reference.
+    Tensor linear_row("linear_row", shape, otype, /*rowwise=*/true, /*colwise=*/false,
+                      NVTE_MXFP8_1D_SCALING);
+    nvte_quantize(input.data(), linear_row.data(), 0);
+    Tensor ref_row_swz("ref_row_swz", shape, otype, /*rowwise=*/true, /*colwise=*/false,
+                       NVTE_MXFP8_1D_SCALING);
+    ref_row_swz.set_with_gemm_swizzled_scales(true);
+    if (out_bytes > 0) {
+        cudaMemcpy(ref_row_swz.rowwise_dptr(), linear_row.rowwise_dptr(),
+                   out_bytes, cudaMemcpyDeviceToDevice);
+        nvte_swizzle_scaling_factors(linear_row.data(), ref_row_swz.data(), 0);
+    }
+
+    // Colwise reference.
+    Tensor linear_col("linear_col", shape, otype, /*rowwise=*/false, /*colwise=*/true,
+                      NVTE_MXFP8_1D_SCALING);
+    nvte_quantize(input.data(), linear_col.data(), 0);
+    Tensor ref_col_swz("ref_col_swz", shape, otype, /*rowwise=*/false, /*colwise=*/true,
+                       NVTE_MXFP8_1D_SCALING);
+    ref_col_swz.set_with_gemm_swizzled_scales(true);
+    if (out_bytes > 0) {
+        cudaMemcpy(ref_col_swz.columnwise_dptr(), linear_col.columnwise_dptr(),
+                   out_bytes, cudaMemcpyDeviceToDevice);
+        nvte_swizzle_scaling_factors(linear_col.data(), ref_col_swz.data(), 0);
+    }
+
+    cudaDeviceSynchronize();
+    ASSERT_EQ(cudaGetLastError(), cudaSuccess) << "reference construction failed";
+
+    // ---- Comparisons ----
+
+    // (1) & (2): FP8 output data byte-identical to single-direction linear casts.
+    TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(otype, OutputType,
+        {
+            const uint8_t *test_row = reinterpret_cast<const uint8_t *>(
+                output_swizzled.rowwise_cpu_dptr<OutputType>());
+            const uint8_t *ref_row = reinterpret_cast<const uint8_t *>(
+                linear_row.rowwise_cpu_dptr<OutputType>());
+            for (size_t i = 0; i < out_bytes; ++i) {
+                ASSERT_EQ(test_row[i], ref_row[i])
+                    << "rowwise fp8 data mismatch at index " << i
+                    << " (swizzled=" << static_cast<int>(test_row[i])
+                    << " linear=" << static_cast<int>(ref_row[i]) << ")";
+            }
+            const uint8_t *test_col = reinterpret_cast<const uint8_t *>(
+                output_swizzled.columnwise_cpu_dptr<OutputType>());
+            const uint8_t *ref_col = reinterpret_cast<const uint8_t *>(
+                linear_col.columnwise_cpu_dptr<OutputType>());
+            for (size_t i = 0; i < out_bytes; ++i) {
+                ASSERT_EQ(test_col[i], ref_col[i])
+                    << "colwise fp8 data mismatch at index " << i
+                    << " (swizzled=" << static_cast<int>(test_col[i])
+                    << " linear=" << static_cast<int>(ref_col[i]) << ")";
+            }
+        }
+    );
+
+    // (3): Swizzled rowwise scale bytes — directly-written vs linear-then-transform.
+    {
+        const size_t n = product(output_swizzled.rowwise_scale_inv_shape());
+        const fp8e8m0 *test = output_swizzled.rowwise_cpu_scale_inv_ptr<fp8e8m0>();
+        const fp8e8m0 *ref  = ref_row_swz.rowwise_cpu_scale_inv_ptr<fp8e8m0>();
+        for (size_t i = 0; i < n; ++i) {
+            ASSERT_EQ(test[i], ref[i])
+                << "swizzled rowwise scale byte mismatch at offset " << i
+                << " (test=" << static_cast<int>(test[i])
+                << " ref=" << static_cast<int>(ref[i]) << ")";
+        }
+    }
+
+    // (4): Swizzled colwise scale bytes — exercises the uint16-pair-packed flush
+    // path added by CACHE_COLWISE_SCALE_IN_SMEM + WITH_SWIZZLED_SCALES.
+    {
+        const size_t n = product(output_swizzled.columnwise_scale_inv_shape());
+        const fp8e8m0 *test = output_swizzled.columnwise_cpu_scale_inv_ptr<fp8e8m0>();
+        const fp8e8m0 *ref  = ref_col_swz.columnwise_cpu_scale_inv_ptr<fp8e8m0>();
+        for (size_t i = 0; i < n; ++i) {
+            ASSERT_EQ(test[i], ref[i])
+                << "swizzled colwise scale byte mismatch at offset " << i
+                << " (test=" << static_cast<int>(test[i])
+                << " ref=" << static_cast<int>(ref[i]) << ")";
+        }
+    }
+
+    // (5) & (6): The rowwise-only swizzled path matches the same independently
+    // constructed linear-then-swizzle reference.
+    TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(otype, OutputType,
+        {
+            const uint8_t *test = reinterpret_cast<const uint8_t *>(
+                output_rowwise_swizzled.rowwise_cpu_dptr<OutputType>());
+            const uint8_t *ref = reinterpret_cast<const uint8_t *>(
+                linear_row.rowwise_cpu_dptr<OutputType>());
+            for (size_t i = 0; i < out_bytes; ++i) {
+                ASSERT_EQ(test[i], ref[i])
+                    << "rowwise-only fp8 data mismatch at index " << i;
+            }
+        }
+    );
+    {
+        const size_t n = product(output_rowwise_swizzled.rowwise_scale_inv_shape());
+        const fp8e8m0 *test =
+            output_rowwise_swizzled.rowwise_cpu_scale_inv_ptr<fp8e8m0>();
+        const fp8e8m0 *ref = ref_row_swz.rowwise_cpu_scale_inv_ptr<fp8e8m0>();
+        for (size_t i = 0; i < n; ++i) {
+            ASSERT_EQ(test[i], ref[i])
+                << "rowwise-only swizzled scale mismatch at offset " << i;
+        }
+    }
+
+}
+
+std::string swizzled_test_name_generator(
+    const testing::TestParamInfo<SwizzledScalesFusedCastMXFP8TestSuite::ParamType>& info) {
+    std::string name;
+    const auto &shape = std::get<0>(info.param);
+    for (size_t i = 0; i < shape.size(); ++i) {
+        if (i > 0) name += "x";
+        name += std::to_string(shape[i]);
+    }
+    name += "X" + test::typeName(std::get<1>(info.param))
+          + "X" + test::typeName(std::get<2>(info.param));
+    return name;
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    OperatorTest_FusedCastMXFP8_SwizzledCastOnly,
+    SwizzledScalesFusedCastMXFP8TestSuite,
+    ::testing::Values(
+        // 1. Aligned, small — sanity.
+        std::make_tuple(std::vector<size_t>{128, 128}, DType::kBFloat16, DType::kFloat8E4M3),
+        // 2. Aligned, small — second dtype pair.
+        std::make_tuple(std::vector<size_t>{128, 128}, DType::kFloat32, DType::kFloat8E5M2),
+        // 3. Aligned, medium.
+        std::make_tuple(std::vector<size_t>{256, 384}, DType::kBFloat16, DType::kFloat8E4M3),
+        // 4. Odd number of scale rows (96/32 = 3) - exercises colwise flush
+        //    odd-row scalar tail.
+        std::make_tuple(std::vector<size_t>{96, 512}, DType::kBFloat16, DType::kFloat8E4M3),
+        // 5. cols/32 not a multiple of 4 - exercises rowwise flush scalar tail
+        //    (need cols multiple of 8 for TMA-alignment on 16-bit input).
+        std::make_tuple(std::vector<size_t>{256, 160}, DType::kBFloat16, DType::kFloat8E4M3),
+        // 6. Odd colwise scale rows + rowwise tail cols - combined tail paths.
+        std::make_tuple(std::vector<size_t>{96, 160}, DType::kBFloat16, DType::kFloat8E4M3),
+        // 7. Small but tile-aligned.
+        std::make_tuple(std::vector<size_t>{32, 64}, DType::kBFloat16, DType::kFloat8E4M3),
+        // 8. Minimum aligned size - single 32x32 tile.
+        std::make_tuple(std::vector<size_t>{32, 32}, DType::kBFloat16, DType::kFloat8E4M3),
+        // 9. Multi-CTA at scale — stresses gmem writes across many CTAs.
+        std::make_tuple(std::vector<size_t>{4096, 32768}, DType::kBFloat16, DType::kFloat8E4M3),
+        // 10. 4D input — matches the existing suite's rank coverage.
+        std::make_tuple(std::vector<size_t>{16, 8, 4, 512}, DType::kBFloat16, DType::kFloat8E4M3),
+        // 11. Third input dtype.
+        std::make_tuple(std::vector<size_t>{128, 128}, DType::kFloat16, DType::kFloat8E4M3),
+        // 12. Larger fp32.
+        std::make_tuple(std::vector<size_t>{1024, 1024}, DType::kFloat32, DType::kFloat8E4M3)
+    ),
+    swizzled_test_name_generator);
diff --git a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py
index 127b487650..16a2d75de6 100644
--- a/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py
+++ b/tests/pytorch/mxfp8/test_mxfp8_quantize_swizzle_fusion.py
@@ -129,3 +129,36 @@ def test_mxfp8_quantize_swizzle_fusion(
         return_rowwise=return_rowwise,
         return_transpose=return_transpose,
     )
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize("M, N", [(96, 160), (4096, 576), (4096, 2112)])
+def test_mxfp8_bidirectional_swizzled_row_scale_padding(M: int, N: int) -> None:
+    """The specialized bidirectional kernel must not overwrite padded row scales."""
+    x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda")
+    quantizer = MXFP8Quantizer(
+        fp8_dtype=te.DType.kFloat8E4M3,
+        rowwise=True,
+        columnwise=True,
+    )
+    quantizer.optimize_for_gemm = True
+    scale = quantizer(x)._rowwise_scale_inv.view(torch.uint8)
+
+    scale_rows = torch.arange(M, device=scale.device, dtype=torch.int64).view(-1, 1)
+    scale_cols = torch.arange(N // 32, device=scale.device, dtype=torch.int64).view(1, -1)
+    num_tiles_x = math.ceil(N / 128)
+    scale_indices = (
+        ((scale_rows // 128) * num_tiles_x + scale_cols // 4) * (128 * 4)
+        + (scale_rows % 32) * 16
+        + ((scale_rows % 128) // 32) * 4
+        + scale_cols % 4
+    )
+    valid_mask = torch.zeros(scale.numel(), dtype=torch.bool, device=scale.device)
+    valid_mask[scale_indices.view(-1)] = True
+
+    torch.testing.assert_close(
+        scale.view(-1)[~valid_mask],
+        torch.zeros_like(scale.view(-1)[~valid_mask]),
+        atol=0,
+        rtol=0,
+    )
diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
index 9f84195a19..3900f072bb 100644
--- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh
@@ -646,6 +646,30 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop,
   float *const amax_ptr = reinterpret_cast<float *>(output->amax.dptr);
   const float *noop_ptr = reinterpret_cast<const float *>(noop->data.dptr);
 
+  // Clear padding before either the generic or specialized kernel writes
+  // directly into the GEMM-swizzled scale layout.
+  if (with_gemm_swizzled_scales && (cols % 128 != 0 || rows % 128 != 0)) {
+    constexpr size_t zero_threads = 256;
+    if (use_rowwise_scaling) {
+      const size_t size_bytes = output->scale_inv.buffer_size_bytes();
+      if (size_bytes > 0) {
+        const size_t zero_blocks = DIVUP(size_bytes, zero_threads);
+        zero_scales_kernel<<<zero_blocks, zero_threads, 0, stream>>>(
+            reinterpret_cast<uint8_t *>(output->scale_inv.dptr), size_bytes, noop_ptr);
+        NVTE_CHECK_CUDA(cudaGetLastError());
+      }
+    }
+    if (use_colwise_scaling) {
+      const size_t size_bytes = output->columnwise_scale_inv.buffer_size_bytes();
+      if (size_bytes > 0) {
+        const size_t zero_blocks = DIVUP(size_bytes, zero_threads);
+        zero_scales_kernel<<<zero_blocks, zero_threads, 0, stream>>>(
+            reinterpret_cast<uint8_t *>(output->columnwise_scale_inv.dptr), size_bytes, noop_ptr);
+        NVTE_CHECK_CUDA(cudaGetLastError());
+      }
+    }
+  }
+
   TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(
       input.dtype(), IType,
       TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(
@@ -667,17 +691,26 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop,
                    bidimensional_traits::blockDIM::M) <= max_grid_dim_y;
 
               const bool is_full_rowwise_chunk = (cols % 128 == 0);
+              const bool has_full_bidimensional_chunks =
+                  (rows % bidimensional_traits::colChunkElems == 0) &&
+                  (cols % bidimensional_traits::rowChunkElems == 0);
+              // Both rowwise and bidimensional cast-only kernels select their
+              // scale layout from WITH_GEMM_SWIZZLED_SCALES.
               const bool scaling_type_has_specialized_support =
                   (scaling_type == ScalingType::ROWWISE && is_full_rowwise_chunk &&
                    rowwise_specialized_grid_fits) ||
-                  (scaling_type == ScalingType::BIDIMENSIONAL &&
+                  (scaling_type == ScalingType::BIDIMENSIONAL && has_full_bidimensional_chunks &&
                    bidimensional_specialized_grid_fits);
 
-              if (specialized::hasSpec<IS_DBIAS, IS_DACT, IS_ACT, IType, OType>() &&
-                  !WITH_GEMM_SWIZZLED_SCALES && scaling_type_has_specialized_support) {
+              // Specialized cast-only kernels do not consume the device noop flag.
+              // Preserve cached outputs by keeping noop-aware calls on the generic path.
+              if (noop_ptr == nullptr &&
+                  specialized::hasSpec<IS_DBIAS, IS_DACT, IS_ACT, IType, OType>() &&
+                  scaling_type_has_specialized_support) {
                 switch (scaling_type) {
                   case ScalingType::ROWWISE: {
-                    using traits = specialized::CastTraits<IType, OType, true, false>;
+                    using traits = specialized::CastTraits<IType, OType, true, false,
+                                                           WITH_GEMM_SWIZZLED_SCALES>;
                     auto kernel = specialized::quantize_mxfp8_kernel_cast_only<traits>;
 
                     NVTE_CHECK_CUDA(cudaFuncSetAttribute(
@@ -696,7 +729,11 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop,
                     break;
                   }
                   case ScalingType::BIDIMENSIONAL: {
-                    using traits = specialized::CastTraits<IType, OType, true, true>;
+                    using traits =
+                        specialized::CastTraitsSwizzle<IType, OType,
+                                                       /*NumStages=*/2, /*IterM=*/1, /*IterN=*/4,
+                                                       /*kCacheColwise=*/WITH_GEMM_SWIZZLED_SCALES,
+                                                       /*kSwizzled=*/WITH_GEMM_SWIZZLED_SCALES>;
                     auto kernel = specialized::quantize_mxfp8_kernel_cast_only<traits>;
 
                     NVTE_CHECK_CUDA(cudaFuncSetAttribute(
@@ -786,38 +823,6 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop,
 
               const size_t dshmem_size = in_mem + out_mem + TMA_SHMEM_ALIGNMENT;
 
-              // Zero out swizzled scales if padding is needed
-              /// TODO (tmoon) Handle this within the cast kernel
-              if (with_gemm_swizzled_scales) {
-                constexpr size_t TILE_DIM_X = 128;  // Tile dim in data buffer
-                constexpr size_t TILE_DIM_Y = 128;
-                if (cols % TILE_DIM_X != 0 || rows % TILE_DIM_Y != 0) {
-                  // Use a noop-aware zero kernel so that the clear is skipped
-                  // when quantization is a noop (e.g. FP8 weight caching).
-                  constexpr size_t zero_threads = 256;
-                  if (use_rowwise_scaling) {
-                    const size_t size_bytes = output->scale_inv.buffer_size_bytes();
-                    if (size_bytes > 0) {
-                      const size_t zero_blocks = DIVUP(size_bytes, zero_threads);
-                      zero_scales_kernel<<<zero_blocks, zero_threads, 0, stream>>>(
-                          reinterpret_cast<uint8_t *>(output->scale_inv.dptr), size_bytes,
-                          noop_ptr);
-                      NVTE_CHECK_CUDA(cudaGetLastError());
-                    }
-                  }
-                  if (use_colwise_scaling) {
-                    const size_t size_bytes = output->columnwise_scale_inv.buffer_size_bytes();
-                    if (size_bytes > 0) {
-                      const size_t zero_blocks = DIVUP(size_bytes, zero_threads);
-                      zero_scales_kernel<<<zero_blocks, zero_threads, 0, stream>>>(
-                          reinterpret_cast<uint8_t *>(output->columnwise_scale_inv.dptr),
-                          size_bytes, noop_ptr);
-                      NVTE_CHECK_CUDA(cudaGetLastError());
-                    }
-                  }
-                }
-              }
-
               switch (scaling_type) {
                 case ScalingType::ROWWISE: {
                   auto kernel = quantize_mxfp8_kernel<IS_DBIAS, IS_DACT, IS_ACT, ParamOP, OP, IType,
diff --git a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
index 0c59733d31..507122cf8d 100644
--- a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh
@@ -14,8 +14,9 @@
 #include <cstdlib>
 
 #include "../../../util/ptx.cuh"
+#include "../swizzle.cuh"  // gemm_swizzled_scale_idx (parent dir, GEMM scale swizzle)
 #include "state_counter.cuh"
-#include "swizzle.cuh"
+#include "swizzle.cuh"  // specialized/swizzle.cuh (TMA input bank-conflict swizzle)
 
 namespace transformer_engine {
 namespace dispatch {
@@ -24,6 +25,10 @@ namespace quantize_kernel {
 namespace specialized {
 
 namespace ptx = transformer_engine::ptx;
+
+// Bring in the GEMM-swizzled scale index helper (from ../swizzle.cuh).
+// Used only when the kernel is instantiated with CastTraits::_with_swizzled_scales=true.
+using transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx;
 namespace {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
 
@@ -122,19 +127,20 @@ struct Layout {
   static constexpr int32_t num = M * N;
 };
 
-template <typename IType, typename OType, bool rowwise, bool colwise>
+template <typename IType, typename OType, bool rowwise, bool colwise,
+          bool with_swizzled_scales = false>
 struct CastTraits;
 
 // 1x32
-template <typename _IType, typename _OType>
-struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false> {
+template <typename _IType, typename _OType, bool _kSwizzled>
+struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false, _kSwizzled> {
   static constexpr bool isRowwise = true;
   static constexpr bool isColwise = false;
   using IType = _IType;
   using OType = _OType;
 
   static constexpr int32_t chunkElems = 32;
-  using threadLayout = Layout<1, 32>;
+  using threadLayout = Layout<1, THREADS_PER_WARP>;
   static constexpr int32_t numThreadsPerChunk = 1;
   static constexpr int32_t warpDimM = threadLayout::M;
   static constexpr int32_t warpDimN = threadLayout::N * chunkElems;
@@ -151,14 +157,22 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/false> {
   using iterLayout = Layout<1, 1>;
   static constexpr int32_t blockDimM = iterLayout::M * blockIterDimM;
   static constexpr int32_t blockDimN = iterLayout::N * blockIterDimN;
+  static constexpr int32_t rowwiseScaleStride = blockDimN / chunkElems;
+  using PreferredDataType = std::conditional_t<
+      rowwiseScaleStride % 16 == 0, uint4,
+      std::conditional_t<
+          rowwiseScaleStride % 8 == 0, uint2,
+          std::conditional_t<rowwiseScaleStride % 4 == 0, uint32_t,
+                             std::conditional_t<rowwiseScaleStride % 2 == 0, uint16_t, uint8_t>>>>;
 
   static constexpr int32_t numStages = 1;
   static constexpr int32_t numPrefetch = numStages - 1;
 
   static constexpr bool _use_cvt_4x = true;
   static constexpr bool _cache_rowwise_scale_in_smem = true;
+  static constexpr bool _with_swizzled_scales = _kSwizzled;
 
-  static constexpr int32_t numThreads = warpLayout::num * 32;
+  static constexpr int32_t numThreads = warpLayout::num * THREADS_PER_WARP;
 
   static constexpr size_t smem_rowwise_scale =
       _cache_rowwise_scale_in_smem ? (blockDimM * (blockDimN / chunkElems) * sizeof(e8m0_t)) : 0ul;
@@ -498,13 +512,8 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re
     block_coords.y = blockIdx.y * CastTraits::blockDimM;
     block_coords.x = blockIdx.x * CastTraits::blockDimN;
 
-    constexpr int32_t stride_in_smem = CastTraits::blockDimN / CastTraits::chunkElems;
-    using PreferredDataType = std::conditional_t<
-        stride_in_smem % 16 == 0, uint4,
-        std::conditional_t<
-            stride_in_smem % 8 == 0, uint2,
-            std::conditional_t<stride_in_smem % 4 == 0, uint32_t,
-                               std::conditional_t<stride_in_smem % 2 == 0, uint16_t, uint8_t>>>>;
+    constexpr int32_t stride_in_smem = CastTraits::rowwiseScaleStride;
+    using PreferredDataType = typename CastTraits::PreferredDataType;
 
     int2 end_coords;
     end_coords.y = std::min(block_coords.y + CastTraits::blockDimM, rows);
@@ -514,7 +523,40 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re
     valid_coords.y = end_coords.y - block_coords.y;
     valid_coords.x = end_coords.x - (block_coords.x / CastTraits::chunkElems);
 
-    if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) {
+    if constexpr (CastTraits::_with_swizzled_scales) {
+      // Four adjacent rowwise scale columns are contiguous in the GEMM scale
+      // layout, so write them as one uint32_t whenever possible.
+      constexpr int32_t cols_per_group = 4;
+      const int32_t groups_per_row = valid_coords.x / cols_per_group;
+      const int32_t total_groups = valid_coords.y * groups_per_row;
+      const int32_t base_col = block_coords.x / CastTraits::chunkElems;
+      const size_t num_tiles_x = DIVUP(cols, static_cast<int32_t>(128));
+
+      for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups;
+           i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
+        const int32_t row = i / groups_per_row;
+        const int32_t col = (i % groups_per_row) * cols_per_group;
+        const uint32_t value =
+            *reinterpret_cast<const uint32_t *>(&sRowwiseScale[row * stride_in_smem + col]);
+        const size_t idx =
+            gemm_swizzled_scale_idx(block_coords.y + row, base_col + col, num_tiles_x);
+        *reinterpret_cast<uint32_t *>(&scales_rowwise[idx]) = value;
+      }
+
+      const int32_t remaining_start = groups_per_row * cols_per_group;
+      const int32_t remaining_per_row = valid_coords.x - remaining_start;
+      if (remaining_per_row > 0) {
+        const int32_t total_remaining = valid_coords.y * remaining_per_row;
+        for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining;
+             i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
+          const int32_t row = i / remaining_per_row;
+          const int32_t col = remaining_start + (i % remaining_per_row);
+          const size_t idx =
+              gemm_swizzled_scale_idx(block_coords.y + row, base_col + col, num_tiles_x);
+          scales_rowwise[idx] = sRowwiseScale[row * stride_in_smem + col];
+        }
+      }
+    } else if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) {
       using DataType = int32_t;
       constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t);
       constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group;
@@ -527,8 +569,9 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re
           reinterpret_cast<DataType *>(scales_rowwise + block_coords.y * scale_stride_rowwise +
                                        block_coords.x / CastTraits::chunkElems);
 
-      for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row);
-           i += CastTraits::warpLayout::num * 32) {
+      for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP;
+           i < (valid_coords.y * num_threads_per_row);
+           i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
         int32_t row = i / num_threads_per_row;
         int32_t col = i % num_threads_per_row;
         gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col];
@@ -546,8 +589,9 @@ __global__ void quantize_mxfp8_kernel_cast_only(typename CastTraits::IType *__re
           reinterpret_cast<DataType *>(scales_rowwise + block_coords.y * scale_stride_rowwise +
                                        block_coords.x / CastTraits::chunkElems);
 
-      for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row);
-           i += CastTraits::warpLayout::num * 32) {
+      for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP;
+           i < (valid_coords.y * num_threads_per_row);
+           i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
         int32_t row = i / num_threads_per_row;
         int32_t col = i % num_threads_per_row;
         gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col];
@@ -577,11 +621,12 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> {
   static constexpr int32_t rowChunkElems = 32;
   static constexpr int32_t colChunkElems = 32;
 
-  using rowThreadLayout = Layout<32, 1>;                                   // 32x1
+  using rowThreadLayout = Layout<THREADS_PER_WARP, 1>;                     // 32x1
   using colThreadLayout = Layout<rowThreadLayout::N, rowThreadLayout::M>;  // 1x32
   static_assert(rowThreadLayout::num == colThreadLayout::num,
                 "rowThreadLayout::num must be equal to colThreadLayout::num");
-  static_assert(rowThreadLayout::num == 32, "rowThreadLayout::num must be 32");
+  static_assert(rowThreadLayout::num == THREADS_PER_WARP,
+                "rowThreadLayout::num must match the warp size");
 
   using rowWarpDim = Layout<rowThreadLayout::M, rowThreadLayout::N * rowChunkElems>;
   using colWarpDim = Layout<colThreadLayout::M * colChunkElems, colThreadLayout::N>;
@@ -603,6 +648,13 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> {
 
   using iterLayout = Layout<1, 4>;
   using blockDIM = Layout<iterLayout::M * blockIterDim::M, iterLayout::N * blockIterDim::N>;
+  static constexpr int32_t rowwiseScaleStride = blockDIM::N / rowChunkElems;
+  using PreferredDataType = std::conditional_t<
+      rowwiseScaleStride % 16 == 0, uint4,
+      std::conditional_t<
+          rowwiseScaleStride % 8 == 0, uint2,
+          std::conditional_t<rowwiseScaleStride % 4 == 0, uint32_t,
+                             std::conditional_t<rowwiseScaleStride % 2 == 0, uint16_t, uint8_t>>>>;
 
   static constexpr int32_t numStages = 2;
 
@@ -636,7 +688,7 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> {
                 "It requires aligned smem pointer");
 
   static constexpr int32_t numWarps = warpLayout::num + 2 * (int32_t)_use_warp_specialization;
-  static constexpr int32_t numThreads = numWarps * 32;
+  static constexpr int32_t numThreads = numWarps * THREADS_PER_WARP;
   static_assert(numThreads <= 1024, "numThreads must be less than or equal to 1024");
 
   static constexpr size_t smemInputPerWarp = warpDim::num * sizeof(IType);
@@ -660,7 +712,9 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> {
   static constexpr bool _need_smem_for_colwise_reduce =
       _colwise_source_coming_from_rowwise;  // && _colwise_reduce_max != ColwiseReduceMax::Redux;
   static constexpr size_t smem_colwise_reduce =
-      _need_smem_for_colwise_reduce ? 32 * warpLayout::num * sizeof(ColwiseReduceDataType) : 0ul;
+      _need_smem_for_colwise_reduce
+          ? THREADS_PER_WARP * warpLayout::num * sizeof(ColwiseReduceDataType)
+          : 0ul;
 
   static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul;
   static constexpr size_t smem = _reuse_input_out_smem
@@ -670,6 +724,139 @@ struct CastTraits<_IType, _OType, /*rowwise=*/true, /*colwise=*/true> {
                                         smem_alignment + smem_rowwise_scale + smem_colwise_reduce);
 };
 
+// Standalone trait for the non-warp-specialized rowwise+colwise cast_only kernel.
+// Exposes numStages, iterM, iterN, and the two colwise-scale features as
+// caller-controllable template axes. Both colwise flags default to true,
+// giving callers the swizzled + colwise-scale-cached fast path by default.
+//
+// This trait duck-types the same interface CastTraits<_, _, true, true> exposes,
+// so it drops into quantize_mxfp8_kernel_cast_only<traits, ...> without any
+// changes to the kernel signature. Kernel #1 (rowwise-only) and Kernel #2
+// (warp-specialized row+col) don't accept it: kernel #1 requires isColwise=false,
+// kernel #2 requires _use_warp_specialization=true - both are wrong here.
+template <typename _IType, typename _OType, int _NumStages = 2, int _IterM = 1, int _IterN = 4,
+          bool _kCacheColwise = true, bool _kSwizzled = true>
+struct CastTraitsSwizzle {
+  static constexpr bool isRowwise = true;
+  static constexpr bool isColwise = true;
+  using IType = _IType;
+  using OType = _OType;
+
+  static constexpr int32_t rowChunkElems = 32;
+  static constexpr int32_t colChunkElems = 32;
+
+  using rowThreadLayout = Layout<THREADS_PER_WARP, 1>;                     // 32x1
+  using colThreadLayout = Layout<rowThreadLayout::N, rowThreadLayout::M>;  // 1x32
+  static_assert(rowThreadLayout::num == colThreadLayout::num,
+                "rowThreadLayout::num must be equal to colThreadLayout::num");
+  static_assert(rowThreadLayout::num == THREADS_PER_WARP,
+                "rowThreadLayout::num must match the warp size");
+
+  using rowWarpDim = Layout<rowThreadLayout::M, rowThreadLayout::N * rowChunkElems>;
+  using colWarpDim = Layout<colThreadLayout::M * colChunkElems, colThreadLayout::N>;
+  using warpDim =
+      Layout<std::max(rowWarpDim::M, colWarpDim::M), std::max(rowWarpDim::N, colWarpDim::N)>;
+
+  static constexpr bool _tma_swizzle = true;
+  using warpLayout = Layout<1, 2>;
+  static_assert(_tma_swizzle ? (warpLayout::N == 2) : true);
+  static constexpr CUtensorMapSwizzle input_swizzle_pattern =
+      _tma_swizzle ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_128B
+                   : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE;
+
+  static constexpr CUtensorMapSwizzle output_swizzle_pattern =
+      _tma_swizzle ? CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_64B
+                   : CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE;
+
+  using blockIterDim = Layout<warpLayout::M * warpDim::M, warpLayout::N * warpDim::N>;
+
+  using iterLayout = Layout<_IterM, _IterN>;
+  using blockDIM = Layout<iterLayout::M * blockIterDim::M, iterLayout::N * blockIterDim::N>;
+  static constexpr int32_t rowwiseScaleStride = blockDIM::N / rowChunkElems;
+  using PreferredDataType = std::conditional_t<
+      rowwiseScaleStride % 16 == 0, uint4,
+      std::conditional_t<
+          rowwiseScaleStride % 8 == 0, uint2,
+          std::conditional_t<rowwiseScaleStride % 4 == 0, uint32_t,
+                             std::conditional_t<rowwiseScaleStride % 2 == 0, uint16_t, uint8_t>>>>;
+
+  static constexpr int32_t numStages = _NumStages;
+
+  using inputUnitType = uint4;
+  static constexpr int32_t rowNumElemsPerUnit = sizeof(inputUnitType) / sizeof(IType);
+  static constexpr int32_t rowNumUnitsPerChunk = rowChunkElems / rowNumElemsPerUnit;
+  using inputElemSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 3, 3>, swz::Linear>;
+  using inputUnitSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<3, 0, 3>, swz::Linear>;
+
+  using colIndexSwz = swz::Swizzle<5, 0, 5>;
+
+  using rowOutputUnitType = uint4;
+  static constexpr int32_t rowNumOutUnitsPerChunk =
+      rowChunkElems * sizeof(OType) / sizeof(rowOutputUnitType);
+  static constexpr int32_t rowOutNumElemsPerUnit = sizeof(rowOutputUnitType) / sizeof(OType);
+
+  using rowOutputChunkSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 0, 3>, swz::Linear>;
+  using colOutputSwz = std::conditional_t<_tma_swizzle, swz::Swizzle<2, 4, 3>, swz::Linear>;
+
+  static constexpr bool _use_cvt_4x = true;
+  static constexpr bool _use_warp_specialization = false;
+  static constexpr bool _need_wait_group = iterLayout::num > numStages;
+  static constexpr bool _reuse_input_out_smem = false;
+  static_assert(_reuse_input_out_smem == false, "Just don't use it");
+  static constexpr bool _cache_rowwise_scale_in_smem = true;
+
+  static constexpr bool _colwise_source_coming_from_rowwise = true;
+  static constexpr ColwiseReduceMax _colwise_reduce_max = ColwiseReduceMax::Redux;
+  static_assert(_colwise_reduce_max != ColwiseReduceMax::RedAsync,
+                "It requires aligned smem pointer");
+
+  // The two colwise-scale features exposed as caller-controllable trait axes.
+  // Both default to true so callers get the vectorized swizzled path by default.
+  static constexpr bool _cache_colwise_scale_in_smem = _kCacheColwise;
+  static constexpr bool _with_swizzled_scales = _kSwizzled;
+
+  static constexpr int32_t numWarps = warpLayout::num + 2 * (int32_t)_use_warp_specialization;
+  static constexpr int32_t numThreads = numWarps * THREADS_PER_WARP;
+  static_assert(numThreads <= 1024, "numThreads must be less than or equal to 1024");
+
+  static constexpr size_t smemInputPerWarp = warpDim::num * sizeof(IType);
+  static constexpr size_t smemInputPerBlock = smemInputPerWarp * warpLayout::num;
+
+  static constexpr size_t smemRowwiseOutputPerWarp = warpDim::num * sizeof(OType);
+  static constexpr size_t smemRowwiseOutputPerBlock = smemRowwiseOutputPerWarp * warpLayout::num;
+
+  static constexpr size_t smemColwiseOutputPerWarp = warpDim::num * sizeof(OType);
+  static constexpr size_t smemColwiseOutputPerBlock = smemColwiseOutputPerWarp * warpLayout::num;
+
+  static constexpr size_t smemInput = smemInputPerBlock * numStages;
+  static constexpr size_t smemRowwiseOutput = smemRowwiseOutputPerBlock * numStages;
+  static constexpr size_t smemColwiseOutput = smemColwiseOutputPerBlock * numStages;
+
+  static constexpr size_t smem_rowwise_scale =
+      _cache_rowwise_scale_in_smem ? (blockDIM::M * (blockDIM::N / rowChunkElems) * sizeof(e8m0_t))
+                                   : 0ul;
+
+  // Extra shmem for cached colwise scales - only when the flag is on.
+  static constexpr size_t smem_colwise_scale =
+      _cache_colwise_scale_in_smem ? (blockDIM::M / colChunkElems) * blockDIM::N * sizeof(e8m0_t)
+                                   : 0ul;
+
+  using ColwiseReduceDataType = float;
+  static constexpr bool _need_smem_for_colwise_reduce = _colwise_source_coming_from_rowwise;
+  static constexpr size_t smem_colwise_reduce =
+      _need_smem_for_colwise_reduce
+          ? THREADS_PER_WARP * warpLayout::num * sizeof(ColwiseReduceDataType)
+          : 0ul;
+
+  static constexpr size_t smem_alignment = _tma_swizzle ? 1024ul : 128ul;
+  static constexpr size_t smem =
+      _reuse_input_out_smem
+          ? (std::max(smemInput, smemColwiseOutput) + smemRowwiseOutput + smem_alignment +
+             smem_rowwise_scale + smem_colwise_scale + smem_colwise_reduce)
+          : (smemInput + smemRowwiseOutput + smemColwiseOutput + smem_alignment +
+             smem_rowwise_scale + smem_colwise_scale + smem_colwise_reduce);
+};
+
 // 32x32
 template <typename CastTraits,
           std::enable_if_t<CastTraits::isRowwise && CastTraits::isColwise, int> = 0,
@@ -723,12 +910,12 @@ __global__ void quantize_mxfp8_kernel_cast_only(
     if constexpr (CastTraits::_need_smem_for_colwise_reduce) {
       sColwiseReduce = reinterpret_cast<ColwiseReduceDataType *>(
           sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t));
-      sColwiseReduce += warpId * 32;
+      sColwiseReduce += warpId * THREADS_PER_WARP;
     }
   } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) {
     sColwiseReduce = reinterpret_cast<ColwiseReduceDataType *>(
         sColOutput + CastTraits::blockIterDim::num * CastTraits::numStages);
-    sColwiseReduce += warpId * 32;
+    sColwiseReduce += warpId * THREADS_PER_WARP;
   }
 
   // TODO: maybe we can assign a different barrier for each warp
@@ -739,8 +926,8 @@ __global__ void quantize_mxfp8_kernel_cast_only(
 #pragma unroll
     for (int32_t i = 0; i < CastTraits::numStages; i++) {
       ptx::mbarrier_init(&ldg_producer[i], 1);
-      ptx::mbarrier_init(&ldg_consumer[i], CastTraits::warpLayout::num * 32);
-      ptx::mbarrier_init(&stg_producer[i], CastTraits::warpLayout::num * 32);
+      ptx::mbarrier_init(&ldg_consumer[i], CastTraits::warpLayout::num * THREADS_PER_WARP);
+      ptx::mbarrier_init(&stg_producer[i], CastTraits::warpLayout::num * THREADS_PER_WARP);
       ptx::mbarrier_init(&stg_consumer[i], 1);
     }
     ptx::fence_mbarrier_init_release_cluster();
@@ -1080,15 +1267,10 @@ __global__ void quantize_mxfp8_kernel_cast_only(
     }
 
     if constexpr (CastTraits::_cache_rowwise_scale_in_smem) {
-      ptx::numbered_barrier_sync(CastTraits::warpLayout::num * 32, 0u);
+      ptx::numbered_barrier_sync(CastTraits::warpLayout::num * THREADS_PER_WARP, 0u);
 
-      constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems;
-      using PreferredDataType = std::conditional_t<
-          stride_in_smem % 16 == 0, uint4,
-          std::conditional_t<
-              stride_in_smem % 8 == 0, uint2,
-              std::conditional_t<stride_in_smem % 4 == 0, uint32_t,
-                                 std::conditional_t<stride_in_smem % 2 == 0, uint16_t, uint8_t>>>>;
+      constexpr int32_t stride_in_smem = CastTraits::rowwiseScaleStride;
+      using PreferredDataType = typename CastTraits::PreferredDataType;
 
       int2 end_coords;
       end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows);
@@ -1112,8 +1294,9 @@ __global__ void quantize_mxfp8_kernel_cast_only(
             reinterpret_cast<DataType *>(scales_rowwise + block_coords.y * scale_stride_rowwise +
                                          block_coords.x / CastTraits::rowChunkElems);
 
-        for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row);
-             i += CastTraits::warpLayout::num * 32) {
+        for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP;
+             i < (valid_coords.y * num_threads_per_row);
+             i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
           int32_t row = i / num_threads_per_row;
           int32_t col = i % num_threads_per_row;
           gScales[row * gmem_stride_in_group + col] =
@@ -1132,8 +1315,9 @@ __global__ void quantize_mxfp8_kernel_cast_only(
             reinterpret_cast<DataType *>(scales_rowwise + block_coords.y * scale_stride_rowwise +
                                          block_coords.x / CastTraits::rowChunkElems);
 
-        for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row);
-             i += CastTraits::warpLayout::num * 32) {
+        for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP;
+             i < (valid_coords.y * num_threads_per_row);
+             i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
           int32_t row = i / num_threads_per_row;
           int32_t col = i % num_threads_per_row;
           gScales[row * gmem_stride_in_group + col] =
@@ -1185,15 +1369,20 @@ __global__ void quantize_mxfp8_kernel_cast_only(
   // colwise output will reuse input buffer
   OType *sColOutput;
   e8m0_t *sRowwiseScale = nullptr;
+  e8m0_t *sColwiseScale = nullptr;
   ColwiseReduceDataType *sColwiseReduce = nullptr;
   if constexpr (CastTraits::_reuse_input_out_smem) {
     sColOutput = reinterpret_cast<OType *>(sInput);
     if constexpr (CastTraits::_cache_rowwise_scale_in_smem) {
       sRowwiseScale = reinterpret_cast<e8m0_t *>(sRowOutput + CastTraits::blockIterDim::num *
                                                                   CastTraits::numStages);
+      if constexpr (CastTraits::_cache_colwise_scale_in_smem) {
+        sColwiseScale = sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t);
+      }
       if constexpr (CastTraits::_need_smem_for_colwise_reduce) {
         sColwiseReduce = reinterpret_cast<ColwiseReduceDataType *>(
-            sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t));
+            sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t) +
+            CastTraits::smem_colwise_scale);
       }
     } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) {
       sColwiseReduce = reinterpret_cast<ColwiseReduceDataType *>(
@@ -1205,9 +1394,13 @@ __global__ void quantize_mxfp8_kernel_cast_only(
     if constexpr (CastTraits::_cache_rowwise_scale_in_smem) {
       sRowwiseScale = reinterpret_cast<e8m0_t *>(sColOutput + CastTraits::blockIterDim::num *
                                                                   CastTraits::numStages);
+      if constexpr (CastTraits::_cache_colwise_scale_in_smem) {
+        sColwiseScale = sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t);
+      }
       if constexpr (CastTraits::_need_smem_for_colwise_reduce) {
         sColwiseReduce = reinterpret_cast<ColwiseReduceDataType *>(
-            sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t));
+            sRowwiseScale + CastTraits::smem_rowwise_scale / sizeof(e8m0_t) +
+            CastTraits::smem_colwise_scale);
       }
     } else if constexpr (CastTraits::_need_smem_for_colwise_reduce) {
       sColwiseReduce = reinterpret_cast<ColwiseReduceDataType *>(
@@ -1217,7 +1410,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(
   rowOutputUnitType *sColOutputUnit = reinterpret_cast<rowOutputUnitType *>(sColOutput);
 
   if constexpr (CastTraits::_need_smem_for_colwise_reduce) {
-    sColwiseReduce += warpId * 32;
+    sColwiseReduce += warpId * THREADS_PER_WARP;
   }
 
   __shared__ uint64_t producer[CastTraits::numStages];
@@ -1237,7 +1430,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(
     }
     if constexpr (CastTraits::_colwise_source_coming_from_rowwise &&
                   CastTraits::_colwise_reduce_max == ColwiseReduceMax::RedAsync) {
-      ptx::mbarrier_init(colwise_reduce_barrier, 32);
+      ptx::mbarrier_init(colwise_reduce_barrier, THREADS_PER_WARP);
     }
 
     ptx::fence_mbarrier_init_release_cluster();
@@ -1257,18 +1450,35 @@ __global__ void quantize_mxfp8_kernel_cast_only(
                                (threadIdx.x % CastTraits::rowThreadLayout::N) *
                                    (CastTraits::rowChunkElems / CastTraits::rowNumElemsPerUnit);
 
-  size_t rowwise_scale_base_offset =
-      (block_coords.y + warp_coords.y + (threadIdx.x / CastTraits::rowThreadLayout::N)) *
-          static_cast<size_t>(scale_stride_rowwise) +
+  // Scale coordinates in absolute (compact) scale-tensor space. Shared by both
+  // compact-layout offsets and by the CastTraits::_with_swizzled_scales branches below.
+  const int32_t row_scale_row_base =
+      block_coords.y + warp_coords.y + (threadIdx.x / CastTraits::rowThreadLayout::N);
+  const int32_t row_scale_col_base =
       (block_coords.x + warp_coords.x +
        (threadIdx.x % CastTraits::rowThreadLayout::N) * CastTraits::rowChunkElems) /
-          CastTraits::rowChunkElems;
+      CastTraits::rowChunkElems;
+  const int32_t col_scale_row_base =
+      (block_coords.y + warp_coords.y +
+       (threadIdx.x / CastTraits::colThreadLayout::N) * CastTraits::colChunkElems) /
+      CastTraits::colChunkElems;
+  const int32_t col_scale_col_base =
+      block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N);
+
+  size_t rowwise_scale_base_offset =
+      static_cast<size_t>(row_scale_row_base) * static_cast<size_t>(scale_stride_rowwise) +
+      row_scale_col_base;
   size_t colwise_scale_base_offset =
-      ((block_coords.y + warp_coords.y +
-        (threadIdx.x / CastTraits::colThreadLayout::N) * CastTraits::colChunkElems) /
-       CastTraits::colChunkElems) *
-          static_cast<size_t>(scale_stride_colwise) +
-      (block_coords.x + warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N));
+      static_cast<size_t>(col_scale_row_base) * static_cast<size_t>(scale_stride_colwise) +
+      col_scale_col_base;
+
+  // Precomputed swizzle constants (each swizzle tile is 128 rows x 4 cols in scale space).
+  // Rowwise scale tensor has DIVUP(cols, 128) tiles across;
+  // colwise scale tensor has DIVUP(rows, 128) tiles across (X/Y axes are transposed).
+  const size_t row_swz_num_tiles_X =
+      CastTraits::_with_swizzled_scales ? DIVUP(cols, static_cast<int32_t>(128)) : 0;
+  const size_t col_swz_num_tiles_X =
+      CastTraits::_with_swizzled_scales ? DIVUP(rows, static_cast<int32_t>(128)) : 0;
 
   constexpr int32_t rowwise_scale_stride_in_smem =
       CastTraits::blockDIM::N / CastTraits::rowChunkElems;
@@ -1395,6 +1605,12 @@ __global__ void quantize_mxfp8_kernel_cast_only(
                   iter_m * CastTraits::blockIterDim::M * rowwise_scale_stride_in_smem +
                   iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems);
               sRowwiseScale[rowwise_scale_offset] = row_biased_exponent;
+            } else if constexpr (CastTraits::_with_swizzled_scales) {
+              int32_t abs_row = row_scale_row_base + iter_m * CastTraits::blockIterDim::M;
+              int32_t abs_col = row_scale_col_base +
+                                iter_n * (CastTraits::blockIterDim::N / CastTraits::rowChunkElems);
+              size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X);
+              scales_rowwise[idx] = row_biased_exponent;
             } else {
               size_t rowwise_scale_offset =
                   rowwise_scale_base_offset +
@@ -1410,12 +1626,32 @@ __global__ void quantize_mxfp8_kernel_cast_only(
             e8m0_t col_biased_exponent = to_e8m0<OType>(col_amax);
             float col_scale_inverse = ptx::exp2f_rcp<float>(col_biased_exponent);
             sColwiseReduce[threadIdx.x] = col_scale_inverse;
-            size_t colwise_scale_offset =
-                colwise_scale_base_offset +
-                iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems) *
-                    static_cast<size_t>(scale_stride_colwise) +
-                iter_n * CastTraits::blockIterDim::N;
-            scales_colwise[colwise_scale_offset] = col_biased_exponent;
+            if constexpr (CastTraits::_cache_colwise_scale_in_smem) {
+              // Cache in shmem; end-of-kernel flush handles gmem indexing.
+              int32_t smem_row = (warp_coords.y + (threadIdx.x / CastTraits::colThreadLayout::N) *
+                                                      CastTraits::colChunkElems) /
+                                     CastTraits::colChunkElems +
+                                 iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems);
+              int32_t smem_col = warp_coords.x + (threadIdx.x % CastTraits::colThreadLayout::N) +
+                                 iter_n * CastTraits::blockIterDim::N;
+              sColwiseScale[smem_row * CastTraits::blockDIM::N + smem_col] = col_biased_exponent;
+            } else if constexpr (CastTraits::_with_swizzled_scales) {
+              int32_t abs_row = col_scale_row_base +
+                                iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems);
+              int32_t abs_col = col_scale_col_base + iter_n * CastTraits::blockIterDim::N;
+              // Colwise scale tensor's X/Y axes are transposed vs rowwise
+              // (see col_swz_num_tiles_X = DIVUP(rows, 128)), so pass
+              // (abs_col, abs_row) - abs_col is the swizzle "row" dim.
+              size_t idx = gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X);
+              scales_colwise[idx] = col_biased_exponent;
+            } else {
+              size_t colwise_scale_offset =
+                  colwise_scale_base_offset +
+                  iter_m * (CastTraits::blockIterDim::M / CastTraits::colChunkElems) *
+                      static_cast<size_t>(scale_stride_colwise) +
+                  iter_n * CastTraits::blockIterDim::N;
+              scales_colwise[colwise_scale_offset] = col_biased_exponent;
+            }
             __syncwarp();
           }
         }
@@ -1540,58 +1776,210 @@ __global__ void quantize_mxfp8_kernel_cast_only(
 
   if constexpr (CastTraits::_cache_rowwise_scale_in_smem) {
     constexpr int32_t stride_in_smem = CastTraits::blockDIM::N / CastTraits::rowChunkElems;
-    using PreferredDataType = std::conditional_t<
-        stride_in_smem % 16 == 0, uint4,
-        std::conditional_t<
-            stride_in_smem % 8 == 0, uint2,
-            std::conditional_t<stride_in_smem % 4 == 0, uint32_t,
-                               std::conditional_t<stride_in_smem % 2 == 0, uint16_t, uint8_t>>>>;
 
     int2 end_coords;
     end_coords.y = std::min(block_coords.y + CastTraits::blockDIM::M, rows);
-    end_coords.x = std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems,
-                            scale_stride_rowwise);
+    if constexpr (CastTraits::_with_swizzled_scales) {
+      end_coords.x =
+          std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems,
+                   DIVUP(cols, static_cast<int32_t>(CastTraits::rowChunkElems)));
+    } else {
+      // The compact layout's padded entries are consumed by a later swizzle.
+      end_coords.x =
+          std::min((block_coords.x + CastTraits::blockDIM::N) / CastTraits::rowChunkElems,
+                   scale_stride_rowwise);
+    }
     int2 valid_coords;
     valid_coords.y = end_coords.y - block_coords.y;
     valid_coords.x = end_coords.x - (block_coords.x / CastTraits::rowChunkElems);
 
-    if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) {
-      using DataType = int32_t;
-      constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t);
-      constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group;
-
-      int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group);
-      int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group;
-
-      DataType *sScales = reinterpret_cast<DataType *>(sRowwiseScale);
-      DataType *gScales =
-          reinterpret_cast<DataType *>(scales_rowwise + block_coords.y * scale_stride_rowwise +
-                                       block_coords.x / CastTraits::rowChunkElems);
+    if constexpr (CastTraits::_with_swizzled_scales) {
+      // Swizzled flush: within a 128x4 swizzle tile, 4 consecutive column entries
+      // for the same row live at 4 consecutive gmem bytes. Group by 4 so each
+      // thread writes a uint32_t when col%4 == 0, then scalar tail for remainder.
+      constexpr int32_t cols_per_group = 4;
+      const int32_t groups_per_row = valid_coords.x / cols_per_group;
+      const int32_t total_groups = valid_coords.y * groups_per_row;
+      const int32_t base_col = block_coords.x / CastTraits::rowChunkElems;
+
+      for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups;
+           i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
+        int32_t row = i / groups_per_row;
+        int32_t group = i % groups_per_row;
+        int32_t col = group * cols_per_group;
+
+        uint32_t val4 =
+            *reinterpret_cast<const uint32_t *>(&sRowwiseScale[row * stride_in_smem + col]);
+
+        int32_t abs_row = block_coords.y + row;
+        int32_t abs_col = base_col + col;
+        size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X);
+        *reinterpret_cast<uint32_t *>(&scales_rowwise[idx]) = val4;
+      }
 
-      for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row);
-           i += CastTraits::warpLayout::num * 32) {
-        int32_t row = i / num_threads_per_row;
-        int32_t col = i % num_threads_per_row;
-        gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col];
+      // Tail (valid_coords.x % 4 != 0)
+      const int32_t remaining_start = groups_per_row * cols_per_group;
+      const int32_t remaining_per_row = valid_coords.x - remaining_start;
+      if (remaining_per_row > 0) {
+        const int32_t total_remaining = valid_coords.y * remaining_per_row;
+        for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining;
+             i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
+          int32_t row = i / remaining_per_row;
+          int32_t col = remaining_start + (i % remaining_per_row);
+          e8m0_t val = sRowwiseScale[row * stride_in_smem + col];
+          int32_t abs_row = block_coords.y + row;
+          int32_t abs_col = base_col + col;
+          size_t idx = gemm_swizzled_scale_idx(abs_row, abs_col, row_swz_num_tiles_X);
+          scales_rowwise[idx] = val;
+        }
       }
     } else {
-      using DataType = PreferredDataType;
-      constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t);
-      constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group;
+      using PreferredDataType = typename CastTraits::PreferredDataType;
 
-      int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group);
-      int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group;
+      if (scale_stride_rowwise % sizeof(PreferredDataType) != 0) {
+        using DataType = int32_t;
+        constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t);
+        constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group;
 
-      DataType *sScales = reinterpret_cast<DataType *>(sRowwiseScale);
-      DataType *gScales =
-          reinterpret_cast<DataType *>(scales_rowwise + block_coords.y * scale_stride_rowwise +
-                                       block_coords.x / CastTraits::rowChunkElems);
+        int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group);
+        int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group;
 
-      for (int32_t i = threadIdx.x + warpId * 32; i < (valid_coords.y * num_threads_per_row);
-           i += CastTraits::warpLayout::num * 32) {
-        int32_t row = i / num_threads_per_row;
-        int32_t col = i % num_threads_per_row;
-        gScales[row * gmem_stride_in_group + col] = sScales[row * num_groups_per_row_in_smem + col];
+        DataType *sScales = reinterpret_cast<DataType *>(sRowwiseScale);
+        DataType *gScales =
+            reinterpret_cast<DataType *>(scales_rowwise + block_coords.y * scale_stride_rowwise +
+                                         block_coords.x / CastTraits::rowChunkElems);
+
+        for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP;
+             i < (valid_coords.y * num_threads_per_row);
+             i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
+          int32_t row = i / num_threads_per_row;
+          int32_t col = i % num_threads_per_row;
+          gScales[row * gmem_stride_in_group + col] =
+              sScales[row * num_groups_per_row_in_smem + col];
+        }
+      } else {
+        using DataType = PreferredDataType;
+        constexpr int32_t num_elems_per_group = sizeof(DataType) / sizeof(e8m0_t);
+        constexpr int32_t num_groups_per_row_in_smem = stride_in_smem / num_elems_per_group;
+
+        int32_t num_threads_per_row = (valid_coords.x / num_elems_per_group);
+        int32_t gmem_stride_in_group = scale_stride_rowwise / num_elems_per_group;
+
+        DataType *sScales = reinterpret_cast<DataType *>(sRowwiseScale);
+        DataType *gScales =
+            reinterpret_cast<DataType *>(scales_rowwise + block_coords.y * scale_stride_rowwise +
+                                         block_coords.x / CastTraits::rowChunkElems);
+
+        for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP;
+             i < (valid_coords.y * num_threads_per_row);
+             i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
+          int32_t row = i / num_threads_per_row;
+          int32_t col = i % num_threads_per_row;
+          gScales[row * gmem_stride_in_group + col] =
+              sScales[row * num_groups_per_row_in_smem + col];
+        }
+      }
+    }
+  }
+
+  // Cached colwise scale flush (swizzled path only). Same barrier semantics as
+  // the rowwise flush above: the last-iter __syncthreads already ordered every
+  // in-loop sColwiseScale byte store before this block reads them.
+  if constexpr (CastTraits::_cache_colwise_scale_in_smem && CastTraits::_with_swizzled_scales) {
+    const int32_t scale_row_base = block_coords.y / CastTraits::colChunkElems;
+    // DIVUP so a partial last block (e.g. rows=993, colChunkElems=32 -> last CTA
+    // has 1 valid input row) still emits its scale row. Truncating divison here
+    // drops the last partial scale row, causing the fast path to diverge from
+    // nvte_swizzle_scaling_factors on non-multiple-of-32-rows inputs.
+    const int32_t valid_rows =
+        DIVUP(std::min(block_coords.y + CastTraits::blockDIM::M, rows) - block_coords.y,
+              static_cast<int32_t>(CastTraits::colChunkElems));
+    const int32_t valid_cols =
+        std::min(block_coords.x + CastTraits::blockDIM::N, cols) - block_coords.x;
+
+    // In GEMM swizzle, contiguous bytes run over four scale-row indices for a
+    // fixed logical column. A 64-row CTA owns two such rows; pack them as a
+    // uint16 store only when the pair stays inside the same 4-row swizzle group.
+    const int32_t row_pairs = valid_rows / 2;
+    const int32_t total_pairs = row_pairs * valid_cols;
+    for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_pairs;
+         i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
+      int32_t row = (i / valid_cols) * 2;
+      int32_t col = i % valid_cols;
+      const int32_t abs_col = block_coords.x + col;
+      const int32_t abs_row = scale_row_base + row;
+      e8m0_t val0 = sColwiseScale[row * CastTraits::blockDIM::N + col];
+      e8m0_t val1 = sColwiseScale[(row + 1) * CastTraits::blockDIM::N + col];
+      size_t idx = gemm_swizzled_scale_idx(abs_col, abs_row, col_swz_num_tiles_X);
+      if (((abs_row & 3) != 3) &&
+          ((reinterpret_cast<uintptr_t>(&scales_colwise[idx]) & (alignof(uint16_t) - 1)) == 0)) {
+        uint16_t val2 = static_cast<uint16_t>(val0) | (static_cast<uint16_t>(val1) << 8);
+        *reinterpret_cast<uint16_t *>(&scales_colwise[idx]) = val2;
+      } else {
+        scales_colwise[idx] = val0;
+        size_t idx1 = gemm_swizzled_scale_idx(abs_col, abs_row + 1, col_swz_num_tiles_X);
+        scales_colwise[idx1] = val1;
+      }
+    }
+    // Odd-row tail.
+    if ((valid_rows & 1) != 0) {
+      const int32_t row = valid_rows - 1;
+      for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < valid_cols;
+           i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
+        const int32_t col = i;
+        e8m0_t val = sColwiseScale[row * CastTraits::blockDIM::N + col];
+        size_t idx = gemm_swizzled_scale_idx(block_coords.x + col, scale_row_base + row,
+                                             col_swz_num_tiles_X);
+        scales_colwise[idx] = val;
+      }
+    }
+  }
+
+  // Cached colwise scale flush (non-swizzled linear layout).
+  // Rows in scales_colwise are indexed by (block_coords.y / colChunkElems),
+  // columns are logical input columns; 4 adjacent columns for the same scale
+  // row live at 4 consecutive gmem bytes, so pack as uint32 stores.
+  if constexpr (CastTraits::_cache_colwise_scale_in_smem && !CastTraits::_with_swizzled_scales) {
+    const int32_t scale_row_base = block_coords.y / CastTraits::colChunkElems;
+    // DIVUP so a partial last block (e.g. rows=993, colChunkElems=32 -> last CTA
+    // has 1 valid input row) still emits its scale row. Truncating divison here
+    // drops the last partial scale row, causing the fast path to diverge from
+    // nvte_swizzle_scaling_factors on non-multiple-of-32-rows inputs.
+    const int32_t valid_rows =
+        DIVUP(std::min(block_coords.y + CastTraits::blockDIM::M, rows) - block_coords.y,
+              static_cast<int32_t>(CastTraits::colChunkElems));
+    const int32_t valid_cols =
+        std::min(block_coords.x + CastTraits::blockDIM::N, cols) - block_coords.x;
+
+    constexpr int32_t cols_per_group = 4;
+    const int32_t groups_per_row = valid_cols / cols_per_group;
+    const int32_t total_groups = valid_rows * groups_per_row;
+    for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_groups;
+         i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
+      int32_t row = i / groups_per_row;
+      int32_t group = i % groups_per_row;
+      int32_t col = group * cols_per_group;
+      uint32_t val4 =
+          *reinterpret_cast<const uint32_t *>(&sColwiseScale[row * CastTraits::blockDIM::N + col]);
+      size_t idx =
+          static_cast<size_t>(scale_row_base + row) * static_cast<size_t>(scale_stride_colwise) +
+          static_cast<size_t>(block_coords.x + col);
+      *reinterpret_cast<uint32_t *>(&scales_colwise[idx]) = val4;
+    }
+    // Column tail (valid_cols not multiple of 4).
+    const int32_t remaining_start = groups_per_row * cols_per_group;
+    if (remaining_start < valid_cols) {
+      const int32_t remaining_cols = valid_cols - remaining_start;
+      const int32_t total_remaining = valid_rows * remaining_cols;
+      for (int32_t i = threadIdx.x + warpId * THREADS_PER_WARP; i < total_remaining;
+           i += CastTraits::warpLayout::num * THREADS_PER_WARP) {
+        int32_t row = i / remaining_cols;
+        int32_t col = remaining_start + (i % remaining_cols);
+        e8m0_t val = sColwiseScale[row * CastTraits::blockDIM::N + col];
+        size_t idx =
+            static_cast<size_t>(scale_row_base + row) * static_cast<size_t>(scale_stride_colwise) +
+            static_cast<size_t>(block_coords.x + col);
+        scales_colwise[idx] = val;
       }
     }
   }
@@ -1599,7 +1987,7 @@ __global__ void quantize_mxfp8_kernel_cast_only(
   ptx::cp_async_bulk_wait_group_read<0>();
 
 #endif  // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
-}
+}  // NOLINT(readability/fn_size)
 
 }  // namespace specialized
 }  // namespace quantize_kernel

From b54bb1e5fdec06bbce633a696020a5b207c6b868 Mon Sep 17 00:00:00 2001
From: fheinecke <23390735+fheinecke@users.noreply.github.com>
Date: Tue, 11 Aug 2026 17:08:29 -0500
Subject: [PATCH 54/67] Enable runtime resolution of CUDA header path for NVRTC
 (#3252)

* Enable runtime resolution of CUDA header path for NVRTC

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Update transformer_engine/common/util/cuda_runtime.cpp

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: fheinecke <23390735+fheinecke@users.noreply.github.com>

* Fix resolution not working properly when package is installed outside of site-packages directory

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

* fix cuda include dir resolution when installing as an editable package

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

* fix linter failure

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>

---------

Signed-off-by: Fred Heinecke <fheinecke@nvidia.com>
Signed-off-by: fheinecke <23390735+fheinecke@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
---
 transformer_engine/common/CMakeLists.txt      |  1 +
 transformer_engine/common/__init__.py         | 46 ++++++++---
 .../common/util/cuda_runtime.cpp              | 79 +++++++++++++++++++
 3 files changed, 117 insertions(+), 9 deletions(-)

diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt
index 9734e24fbe..6ceadc7405 100644
--- a/transformer_engine/common/CMakeLists.txt
+++ b/transformer_engine/common/CMakeLists.txt
@@ -371,6 +371,7 @@ target_link_libraries(transformer_engine PUBLIC
                       CUDA::cublas
                       CUDA::cudart
                       CUDNN::cudnn_all)
+target_link_libraries(transformer_engine PRIVATE ${CMAKE_DL_LIBS})
 
 target_include_directories(transformer_engine PRIVATE
                            ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES})
diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py
index 231680321e..6b276d22a8 100644
--- a/transformer_engine/common/__init__.py
+++ b/transformer_engine/common/__init__.py
@@ -235,24 +235,49 @@ def _get_sys_extension() -> str:
     raise RuntimeError(f"Unsupported operating system ({system})")
 
 
+def _cuda_runtime_major(cuda_runtime: ctypes.CDLL) -> Optional[int]:
+    """Return the major version of the CUDA runtime loaded with Transformer Engine."""
+
+    runtime_version = ctypes.c_int()
+    get_runtime_version = cuda_runtime.cudaRuntimeGetVersion
+    get_runtime_version.argtypes = [ctypes.POINTER(ctypes.c_int)]
+    get_runtime_version.restype = ctypes.c_int
+    if get_runtime_version(ctypes.byref(runtime_version)) != 0 or runtime_version.value <= 0:
+        return None
+    return runtime_version.value // 1000
+
+
 @functools.lru_cache(maxsize=None)
-def _nvidia_cudart_include_dir() -> str:
+def _nvidia_cudart_include_dir(cuda_major_version: int) -> str:
     """Returns the include directory for cuda_runtime.h if exists in python environment."""
 
+    # This is primarily here to support editable installs. cuda_runtime.cpp handles the
+    # resolution for install via wheel or when using the shared library ABI directly.
+
     try:
         import nvidia
     except ModuleNotFoundError:
         return ""
 
-    # Installing some nvidia-* packages, like nvshmem, create nvidia name, so "import nvidia"
-    # above doesn't throw. However, they don't set "__file__" attribute.
+    # NVIDIA packages may use either a regular package or a namespace package spread
+    # across multiple package roots.
     if nvidia.__file__ is not None:
-        nvidia_root = Path(nvidia.__file__).parent
+        nvidia_roots = (Path(nvidia.__file__).parent,)
     else:
-        nvidia_root = Path(nvidia.__path__[0])  # namespace package
+        nvidia_roots = tuple(Path(path) for path in nvidia.__path__)
+
+    layouts = [f"cu{cuda_major_version}"]
+    if cuda_major_version == 12:
+        layouts.append("cuda_runtime")
 
-    include_dir = nvidia_root / "cuda_runtime"
-    return str(include_dir) if include_dir.exists() else ""
+    for layout in layouts:
+        for nvidia_root in nvidia_roots:
+            cuda_root = nvidia_root / layout
+            if (cuda_root / "cuda_runtime.h").is_file() or (
+                cuda_root / "include" / "cuda_runtime.h"
+            ).is_file():
+                return str(cuda_root)
+    return ""
 
 
 @functools.lru_cache(maxsize=None)
@@ -382,5 +407,8 @@ def _load_core_library():
     _TE_LIB_CTYPES = _load_core_library()
 
     # Needed to find the correct headers for NVRTC kernels.
-    if not os.getenv("NVTE_CUDA_INCLUDE_DIR") and _nvidia_cudart_include_dir():
-        os.environ["NVTE_CUDA_INCLUDE_DIR"] = _nvidia_cudart_include_dir()
+    _cuda_major_version = _cuda_runtime_major(_TE_LIB_CTYPES)
+    if not os.getenv("NVTE_CUDA_INCLUDE_DIR") and _cuda_major_version is not None:
+        cuda_include_dir = _nvidia_cudart_include_dir(_cuda_major_version)
+        if cuda_include_dir:
+            os.environ["NVTE_CUDA_INCLUDE_DIR"] = cuda_include_dir
diff --git a/transformer_engine/common/util/cuda_runtime.cpp b/transformer_engine/common/util/cuda_runtime.cpp
index 504d761bb1..5690ef5e3c 100644
--- a/transformer_engine/common/util/cuda_runtime.cpp
+++ b/transformer_engine/common/util/cuda_runtime.cpp
@@ -7,6 +7,7 @@
 #include "../util/cuda_runtime.h"
 
 #include <cublasLt.h>
+#include <dlfcn.h>
 
 #include <filesystem>
 #include <fstream>
@@ -26,6 +27,83 @@ namespace {
 // String with build-time CUDA include path
 #include "string_path_cuda_include.h"
 
+// Get the runtime directory of the shared library that contains this code
+std::filesystem::path shared_library_directory() {
+  static const char library_anchor = 0;
+  Dl_info library_info{};
+  if (dladdr(static_cast<const void *>(&library_anchor), &library_info) == 0 ||
+      library_info.dli_fname == nullptr) {
+    return {};
+  }
+
+  std::filesystem::path library_path = library_info.dli_fname;
+  if (library_path.is_relative()) {
+    std::error_code error;
+    library_path = std::filesystem::absolute(library_path, error);
+    if (error) {
+      return {};
+    }
+  }
+
+  return library_path.parent_path();
+}
+
+std::string runtime_cuda_major_version() {
+  int runtime_version = 0;
+  // Header discovery is best-effort, so do not throw if the runtime cannot
+  // report its version.
+  if (cudaRuntimeGetVersion(&runtime_version) != cudaSuccess || runtime_version <= 0) {
+    return {};
+  }
+
+  return std::to_string(runtime_version / 1000);
+}
+
+std::filesystem::path python_cuda_directory() {
+  using Path = std::filesystem::path;
+
+  // Find the Python package root from the installed Transformer Engine package. Do not
+  // assume that the root is named site-packages or dist-packages since valid installs
+  // may use an arbitrary target directory.
+  Path te_package_directory = shared_library_directory();
+  while (true) {
+    if (te_package_directory.filename() == "transformer_engine") {
+      break;
+    }
+
+    const Path parent = te_package_directory.parent_path();
+    if (parent == te_package_directory) {
+      // Root directory reached
+      return {};
+    }
+
+    te_package_directory = parent;
+  }
+
+  const Path nvidia_directory = te_package_directory.parent_path() / "nvidia";
+  const auto cuda_major_version = runtime_cuda_major_version();
+  if (cuda_major_version.empty()) {
+    return {};
+  }
+
+  std::error_code error;
+  const Path cuda_directory = nvidia_directory / ("cu" + cuda_major_version);
+  if (std::filesystem::is_directory(cuda_directory, error)) {
+    return cuda_directory;
+  }
+
+  // CUDA 12 Python wheels use the older nvidia/cuda_runtime layout.
+  if (cuda_major_version == "12") {
+    error.clear();
+    const Path legacy_cuda_directory = nvidia_directory / "cuda_runtime";
+    if (std::filesystem::is_directory(legacy_cuda_directory, error)) {
+      return legacy_cuda_directory;
+    }
+  }
+
+  return {};
+}
+
 }  // namespace
 
 int num_devices() {
@@ -152,6 +230,7 @@ const std::string &include_directory(bool required) {
     std::vector<std::pair<std::string, Path>> search_paths = {{"NVTE_CUDA_INCLUDE_DIR", ""},
                                                               {"CUDA_HOME", ""},
                                                               {"CUDA_DIR", ""},
+                                                              {"", python_cuda_directory()},
                                                               {"", string_path_cuda_include},
                                                               {"", "/usr/local/cuda"}};
     for (auto &[env, p] : search_paths) {

From 56a05e7f5ce8d600ace27cb807482cccf5bf809c Mon Sep 17 00:00:00 2001
From: Phuong Nguyen <phuonguyen@nvidia.com>
Date: Tue, 11 Aug 2026 19:05:33 -0700
Subject: [PATCH 55/67] [Common][PyTorch] EP dispatch with unfused MXFP8
 quantization (#3270)

* EP Dispatch with MXFP8

* Separate EP dispatch-forward and combine-backward quant recipes

* Simplify EP quantize path and guard zero-copy pool alloc under CUDA graph capture

*  Add eager-mode coverage for MXFP8 EP combine backward

* Supply persistent symm-mem recv buffers for zero-copy 1F1B CUDA-graph capture

* Enable MXFP8 EP combine backward under zero-copy

* Release symm-mem pool in ep_finalize

* Skip MXFP8 EP tests on pre-Blackwell devices

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---
 tests/cpp_distributed/test_ep.cu              | 137 ++++++
 tests/cpp_distributed/test_ep_common.h        |   4 +-
 tests/pytorch/distributed/run_ep.py           | 219 +++++++++-
 tests/pytorch/distributed/run_test_ep.sh      |  12 +-
 transformer_engine/common/ep/ep_backend.cpp   |  70 ++-
 .../include/transformer_engine/comm_window.h  |   3 +
 .../common/include/transformer_engine/ep.h    |   6 +
 transformer_engine/pytorch/csrc/extensions.h  |   8 +-
 .../pytorch/csrc/extensions/ep.cpp            | 134 +++++-
 transformer_engine/pytorch/distributed.py     |  45 +-
 transformer_engine/pytorch/ep.py              | 402 +++++++++++++++---
 11 files changed, 937 insertions(+), 103 deletions(-)

diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu
index 47732196ed..be6dbed3ab 100644
--- a/tests/cpp_distributed/test_ep.cu
+++ b/tests/cpp_distributed/test_ep.cu
@@ -97,6 +97,26 @@ static std::vector<float> expected_recv_values_sorted(
   return vals;
 }
 
+// Sorted multiset of per-token identity bytes each recv_rank should receive,
+// derived from the deterministic routing map (same id_of stamp as the test).
+static std::vector<uint8_t> expected_recv_ids_sorted(
+    int recv_rank, int num_processes, int num_tokens, int top_k,
+    int num_experts, int num_local_experts) {
+  int base = recv_rank * num_local_experts;
+  std::vector<uint8_t> ids;
+  for (int src = 0; src < num_processes; ++src) {
+    auto idx = routing_balanced(src, num_tokens, top_k, num_experts, num_local_experts);
+    for (int t = 0; t < num_tokens; ++t)
+      for (int k = 0; k < top_k; ++k) {
+        int64_t e = idx[t * top_k + k];
+        if (e >= base && e < base + num_local_experts)
+          ids.push_back(static_cast<uint8_t>((src * num_tokens + t + 1) & 0xFF));
+      }
+  }
+  std::sort(ids.begin(), ids.end());
+  return ids;
+}
+
 // 2^-5 relative tolerance for BF16 (matches mantissa precision with margin),
 // plus a small atol floor for near-zero expected values.
 static constexpr float kBf16Rtol = 1.0f / 32.0f;
@@ -368,6 +388,123 @@ TYPED_TEST(EPDispatchTest, PrepareAndDispatch) {
   NVTE_CHECK_CUDA(cudaStreamDestroy(stream));
 }
 
+// =============================================================================
+// EPDispatchScaledTest: MXFP8 dispatch routes each token's e8m0 scale row in
+// lockstep with its e4m3 data row.
+// =============================================================================
+
+// Each source token stamps a nonzero identity byte across its whole data row
+// and its whole scale row. Dispatch is a pure permutation, so every filled recv
+// slot must carry the same id in both buffers; a zero id means scales were not
+// routed.
+TEST_F(EpOpTestBase, MXFP8DispatchScales) {
+  if (g_sm_major < 9) GTEST_SKIP() << "EP requires SM_90+";
+  // MXFP8 dispatch is supported on all HT EM modes (local_permute, local_dup,
+  // nvlink_dup); the mode is selected by the NCCL_EP_HT_EM_* env at group init.
+  // MXFP8 e8m0 scale bytes/token = hidden/32; NCCL EP HT requires 16B alignment.
+  if (hidden_dim_ % 512 != 0)
+    GTEST_SKIP() << "MXFP8 dispatch needs hidden_dim % 512 == 0 (16B scale alignment)";
+  const int scale_cols = hidden_dim_ / 32;
+  using Shape = std::vector<size_t>;
+
+  EPBuffers<nv_bfloat16> buf;
+  buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_,
+            ep_size_, max_tokens_per_rank_);
+  EPTensors<nv_bfloat16> t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_);
+
+  auto h_idx = routing_balanced(g_process_id, num_tokens_, top_k_,
+                                num_experts_, num_local_experts_);
+  std::vector<float> h_w(num_tokens_ * top_k_, 1.0f / top_k_);
+  NVTE_CHECK_CUDA(cudaMemcpy(buf.topk_idx.get(), h_idx.data(),
+                        h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice));
+  NVTE_CHECK_CUDA(cudaMemcpy(buf.topk_weights.get(), h_w.data(),
+                        h_w.size() * sizeof(float), cudaMemcpyHostToDevice));
+
+  auto id_of = [&](int tok) {
+    return static_cast<uint8_t>((g_process_id * num_tokens_ + tok + 1) & 0xFF);
+  };
+  std::vector<uint8_t> h_data(num_tokens_ * hidden_dim_);
+  std::vector<uint8_t> h_scale(num_tokens_ * scale_cols);
+  for (int tok = 0; tok < num_tokens_; ++tok) {
+    const uint8_t id = id_of(tok);
+    std::fill_n(&h_data[tok * hidden_dim_], hidden_dim_, id);
+    std::fill_n(&h_scale[tok * scale_cols], scale_cols, id);
+  }
+
+  DevBuf<uint8_t> d_data(num_tokens_ * hidden_dim_), d_scale(num_tokens_ * scale_cols);
+  DevBuf<uint8_t> d_recv_data(buf.recv_capacity * hidden_dim_);
+  DevBuf<uint8_t> d_recv_scale(buf.recv_capacity * scale_cols);
+  NVTE_CHECK_CUDA(cudaMemcpy(d_data.get(), h_data.data(), h_data.size(), cudaMemcpyHostToDevice));
+  NVTE_CHECK_CUDA(cudaMemcpy(d_scale.get(), h_scale.data(), h_scale.size(), cudaMemcpyHostToDevice));
+  NVTE_CHECK_CUDA(cudaMemset(d_recv_data.get(), 0, d_recv_data.bytes()));
+  NVTE_CHECK_CUDA(cudaMemset(d_recv_scale.get(), 0, d_recv_scale.bytes()));
+
+  TensorWrapper tokens(NVTE_MXFP8_1D_SCALING);
+  tokens.set_rowwise_data(d_data.get(), DType::kFloat8E4M3,
+                          Shape{(size_t)num_tokens_, (size_t)hidden_dim_});
+  tokens.set_rowwise_scale_inv(d_scale.get(), DType::kFloat8E8M0,
+                               Shape{(size_t)num_tokens_, (size_t)scale_cols});
+  TensorWrapper recv_tokens(NVTE_MXFP8_1D_SCALING);
+  recv_tokens.set_rowwise_data(d_recv_data.get(), DType::kFloat8E4M3,
+                               Shape{buf.recv_capacity, (size_t)hidden_dim_});
+  recv_tokens.set_rowwise_scale_inv(d_recv_scale.get(), DType::kFloat8E8M0,
+                                    Shape{buf.recv_capacity, (size_t)scale_cols});
+
+  cudaStream_t stream;
+  NVTE_CHECK_CUDA(cudaStreamCreate(&stream));
+  ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(),
+                                  t.recv_tokens_per_expert.data(), nullptr, &t.layer_cfg_, stream));
+  ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(),
+                                   tokens.data(), NVTECommWindow{}, t.topk_weights.data(),
+                                   NVTECommWindow{}, recv_tokens.data(), NVTECommWindow{},
+                                   t.recv_topk_weights.data(), NVTECommWindow{}, stream));
+  NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));
+
+  std::vector<int32_t> counts(num_local_experts_);
+  NVTE_CHECK_CUDA(cudaMemcpy(counts.data(), buf.recv_tokens_per_expert.get(),
+                        num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost));
+  auto exp_counts = expected_recv_tokens_per_expert(g_process_id, g_num_processes, num_tokens_,
+                                          top_k_, num_experts_, num_local_experts_);
+  int total_recv = 0;
+  for (int e = 0; e < num_local_experts_; ++e) {
+    EXPECT_EQ(counts[e], exp_counts[e]) << "local expert " << e;
+    total_recv += exp_counts[e];
+  }
+  ASSERT_LE(total_recv, static_cast<int>(buf.recv_capacity));
+
+  std::vector<uint8_t> r_data(buf.recv_capacity * hidden_dim_);
+  std::vector<uint8_t> r_scale(buf.recv_capacity * scale_cols);
+  NVTE_CHECK_CUDA(cudaMemcpy(r_data.data(), d_recv_data.get(), r_data.size(), cudaMemcpyDeviceToHost));
+  NVTE_CHECK_CUDA(cudaMemcpy(r_scale.data(), d_recv_scale.get(), r_scale.size(), cudaMemcpyDeviceToHost));
+
+  std::vector<uint8_t> got_data_ids, got_scale_ids;
+  got_data_ids.reserve(total_recv);
+  got_scale_ids.reserve(total_recv);
+  size_t slot = 0;
+  for (int e = 0; e < num_local_experts_; ++e) {
+    for (int i = 0; i < counts[e]; ++i, ++slot) {
+      got_data_ids.push_back(r_data[slot * hidden_dim_]);
+      got_scale_ids.push_back(r_scale[slot * scale_cols]);
+    }
+  }
+  std::sort(got_data_ids.begin(), got_data_ids.end());
+  std::sort(got_scale_ids.begin(), got_scale_ids.end());
+
+  auto exp_ids = expected_recv_ids_sorted(g_process_id, g_num_processes, num_tokens_,
+                                          top_k_, num_experts_, num_local_experts_);
+  ASSERT_EQ(got_data_ids.size(), exp_ids.size());
+  ASSERT_EQ(got_scale_ids.size(), exp_ids.size());
+  for (size_t i = 0; i < exp_ids.size(); ++i) {
+    EXPECT_EQ(got_data_ids[i], exp_ids[i]) << "recv data id mismatch at sorted index " << i;
+    EXPECT_EQ(got_scale_ids[i], exp_ids[i]) << "recv scale id mismatch at sorted index " << i;
+  }
+
+  if (g_process_id == 0)
+    printf("  MXFP8DispatchScales: passed (recv=%d, data + scales match routing map)\n", total_recv);
+
+  NVTE_CHECK_CUDA(cudaStreamDestroy(stream));
+}
+
 // =============================================================================
 // EPCombineTest: round-trip identity expert -> result == top_k * tokens.
 // =============================================================================
diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h
index 7cf6017090..37f324eabb 100644
--- a/tests/cpp_distributed/test_ep_common.h
+++ b/tests/cpp_distributed/test_ep_common.h
@@ -7,7 +7,7 @@
 /*
  * Shared TE EP test infrastructure. Include once per TU; ep_bootstrap() in
  * each test binary's main() populates process-level globals.
- * Defaults: 4 experts/rank, hidden_dim=256, max_tokens_per_rank=64.
+ * Defaults: 4 experts/rank, hidden_dim=512, max_tokens_per_rank=64.
  */
 #pragma once
 
@@ -48,7 +48,7 @@ static int         g_num_processes       = -1;
 static int         g_sm_major            = -1;   // set by ep_bootstrap; -1 until then
 static int         g_ep_size             = -1;
 static int         g_num_experts         = -1;
-static int         g_hidden_dim          = 256;
+static int         g_hidden_dim          = 512;
 static int         g_max_tokens_per_rank = 64;
 static NVTEDType   g_max_token_dtype     = kNVTEFloat32;  // staging-buffer sizing
 static bool        g_ep_initialized      = false;
diff --git a/tests/pytorch/distributed/run_ep.py b/tests/pytorch/distributed/run_ep.py
index 067f75d725..6e09ed316f 100644
--- a/tests/pytorch/distributed/run_ep.py
+++ b/tests/pytorch/distributed/run_ep.py
@@ -11,6 +11,7 @@
 import torch
 import torch.distributed as dist
 
+from transformer_engine.common.recipe import MXFP8BlockScaling
 from transformer_engine.pytorch.ep import (
     EpBuffer,
     ep_bootstrap,
@@ -19,12 +20,12 @@
     ep_dispatch,
     ep_combine,
     symm_mem_alloc,
+    release_symm_mem_pool,
     is_symm_backed,
     _ep_combine_raw,
     _ep_dispatch_raw,
 )
 
-
 ZERO_COPY = os.environ.get("NVTE_EP_ZERO_COPY", "0") == "1"
 EAGER = os.environ.get("NVTE_EP_EAGER", "0") == "1"
 OVERFLOW = os.environ.get("NVTE_EP_OVERFLOW", "0") == "1"
@@ -32,11 +33,13 @@
 # Must come after the transformer_engine import so libtransformer_engine.so is loaded.
 import transformer_engine_torch as tex  # noqa: F401
 
-
 NUM_LOCAL_EXPERTS = 2
-HIDDEN_DIM = 32
+# MXFP8 dispatch needs HIDDEN_DIM % 512 == 0 and TOKENS_PER_RANK % 32 == 0. Defaults
+# satisfy both so the MXFP8 tests run by default; override via NVTE_EP_HIDDEN_DIM /
+# NVTE_EP_TOKENS_PER_RANK.
+HIDDEN_DIM = int(os.environ.get("NVTE_EP_HIDDEN_DIM", "512"))
 TOP_K = 2
-TOKENS_PER_RANK = 4
+TOKENS_PER_RANK = int(os.environ.get("NVTE_EP_TOKENS_PER_RANK", "32"))
 
 
 def _zero_copy_test_include(fn):
@@ -57,6 +60,18 @@ def _overflow_test_include(fn):
     return fn
 
 
+# MXFP8 grouped dispatch needs a per-expert alignment of 128, but the EP backend caches a single
+# alignment per process, so alignment=128 tests cannot share a process with the alignment=0 tests.
+# They run in a dedicated pass (NVTE_EP_MXFP8_PASS=1) instead.
+MXFP8_PASS = os.environ.get("NVTE_EP_MXFP8_PASS", "0") == "1"
+
+
+def _mxfp8_align_test(fn):
+    """Mark a test that dispatches with alignment=128; runs only in the MXFP8 pass."""
+    fn._mxfp8_align_test = True
+    return fn
+
+
 class _StageToSymm(torch.autograd.Function):
     """Identity op that stages ``src`` into a symm-mem buffer; grad passes through.
     Lets a test feed a symm-mem-backed, autograd-tracked tensor into ep_combine.
@@ -121,6 +136,16 @@ def _make_identity_inputs(rank, ep_size, device="cuda"):
     )
 
 
+def _degroup_mxfp8(recv_grouped, valid_counts=None):
+    """Dequantize a per-expert MXFP8 GroupedTensor to a dense tensor in expert-major order.
+    With ``valid_counts`` keep only the first ``valid_counts[e]`` rows of each padded expert
+    slot; otherwise return every (padded) row."""
+    parts = recv_grouped.split_into_quantized_tensors()
+    if valid_counts is None:
+        return torch.cat([p.dequantize() for p in parts], dim=0)
+    return torch.cat([p.dequantize()[:v] for p, v in zip(parts, valid_counts)], dim=0)
+
+
 class _Cfg:
     rank: int
     world_size: int
@@ -171,6 +196,16 @@ def setUpClass(cls):
         )
 
     def setUp(self):
+        # alignment=128 MXFP8 tests run only in the dedicated MXFP8 pass; everything else skips
+        # there (and the MXFP8 tests skip outside it) since the backend pins one alignment/process.
+        is_mxfp8_align = getattr(getattr(self, self._testMethodName), "_mxfp8_align_test", False)
+        if MXFP8_PASS and not is_mxfp8_align:
+            self.skipTest("only alignment=128 MXFP8 tests run in the MXFP8 pass")
+        if not MXFP8_PASS and is_mxfp8_align:
+            self.skipTest("alignment=128 MXFP8 tests run in the dedicated MXFP8 pass")
+        # MXFP8 quantization requires Blackwell (SM 10.0) or newer.
+        if is_mxfp8_align and torch.cuda.get_device_capability() < (10, 0):
+            self.skipTest("MXFP8 EP tests require Blackwell (SM 10.0) or newer")
         # Only the zero-copy-capable tests run in the zero-copy pass.
         if ZERO_COPY and not getattr(
             getattr(self, self._testMethodName), "_zero_copy_test_include", False
@@ -185,7 +220,13 @@ def setUp(self):
         ):
             self.skipTest("not exercised in overflow mode")
 
-    def _make_buffer(self, alignment=0, top_k=TOP_K):
+    def _make_buffer(
+        self,
+        alignment=0,
+        top_k=TOP_K,
+        dispatch_fwd_quant_recipe=None,
+        combine_bwd_quant_recipe=None,
+    ):
         return EpBuffer(
             top_k=top_k,
             max_tokens_per_rank=TOKENS_PER_RANK,
@@ -193,13 +234,17 @@ def _make_buffer(self, alignment=0, top_k=TOP_K):
             num_local_experts=NUM_LOCAL_EXPERTS,
             recv_capacity_per_rank=None if EAGER else self.cfg.recv_capacity_per_rank,
             alignment=alignment,
+            dispatch_fwd_quant_recipe=dispatch_fwd_quant_recipe,
+            combine_bwd_quant_recipe=combine_bwd_quant_recipe,
         )
 
     def _expert_out(self, expert_out):
         """Stage the combine input into symm-mem under zero-copy (combine requires it)."""
         if not ZERO_COPY:
             return expert_out
-        symm_buf = symm_mem_alloc(tuple(expert_out.shape), expert_out.dtype, self.ep_group)
+        symm_buf = symm_mem_alloc(
+            tuple(expert_out.shape), expert_out.dtype, self.ep_group, use_pool=True
+        )
         return _StageToSymm.apply(expert_out, symm_buf)
 
     def _stage_grad_symm(self, x, symm_buf=None):
@@ -372,6 +417,70 @@ def test_dispatch_autograd(self):
                     tokens_p.grad.float(), tokens.float() * float(TOP_K), atol=5e-2, rtol=5e-2
                 )
 
+    # MXFP8 dispatch
+
+    def _mxfp8_quantizer(self):
+        from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer
+
+        return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False)
+
+    def _require_mxfp8_shapes(self):
+        if HIDDEN_DIM % 512 != 0 or TOKENS_PER_RANK % 32 != 0:
+            self.skipTest(
+                "MXFP8 needs HIDDEN_DIM % 512 == 0 and TOKENS_PER_RANK % 32 == 0 "
+                "(set NVTE_EP_HIDDEN_DIM / NVTE_EP_TOKENS_PER_RANK)"
+            )
+
+    def _assert_mxfp8_matches_bf16(self, recv_mx, tokens, topk_idx, w, tc):
+        """Dequantized MXFP8 recv matches a bf16 dispatch of the same tokens. Both share the
+        alignment=128 padded expert-major layout, so compare the full prefix [0:sum(padded)]."""
+        ref_tokens = self._mxfp8_quantizer().quantize(tokens).dequantize()
+        ref_recv, _rw, _tc = ep_dispatch(self._make_buffer(alignment=128), ref_tokens, topk_idx, w)
+        torch.cuda.synchronize()
+        n = int(tc.sum())
+        torch.testing.assert_close(
+            _degroup_mxfp8(recv_mx).float(), ref_recv.float()[:n], atol=1e-2, rtol=1e-2
+        )
+
+    @_eager_test_include
+    @_zero_copy_test_include
+    @_mxfp8_align_test
+    def test_dispatch_mxfp8(self):
+        """MXFP8 dispatch quantizes bf16 tokens internally; recv (a per-expert GroupedTensor)
+        dequantized matches a bf16 dispatch of the same tokens. Under zero-copy the recv data and
+        scales are symm-mem backed."""
+        self._require_mxfp8_shapes()
+        topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+        buf = self._make_buffer(dispatch_fwd_quant_recipe=MXFP8BlockScaling(), alignment=128)
+        recv_mx, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w)
+        if ZERO_COPY:
+            self.assertTrue(is_symm_backed(recv_mx.rowwise_data))
+            self.assertTrue(is_symm_backed(recv_mx.scale_inv))
+        self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc)
+
+    @_zero_copy_test_include
+    @_mxfp8_align_test
+    def test_caller_provides_dispatch_recv_mxfp8(self):
+        """One caller-supplied buffer holds the recv data followed by the e8m0 scales; ep_dispatch
+        slices it and the returned GroupedTensor views the data and scale regions of that buffer."""
+        self._require_mxfp8_shapes()
+        from transformer_engine.pytorch.constants import MXFP8_BLOCK_SCALING_SIZE
+
+        rc = self.cfg.recv_capacity_per_rank
+        cols = HIDDEN_DIM // MXFP8_BLOCK_SCALING_SIZE
+        nbytes = rc * (HIDDEN_DIM + cols)  # fp8 data + e8m0 scales, one byte per element
+        if ZERO_COPY:
+            recv_buf = symm_mem_alloc((nbytes,), torch.uint8, self.ep_group)
+        else:
+            recv_buf = torch.empty(nbytes, dtype=torch.uint8, device=self.cfg.device)
+        buf = self._make_buffer(dispatch_fwd_quant_recipe=MXFP8BlockScaling(), alignment=128)
+        topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+        recv_mx, _rw, tc = ep_dispatch(buf, tokens, topk_idx, w, recv_tokens=recv_buf)
+        # the returned GroupedTensor views the caller buffer's data then scale regions
+        self.assertEqual(recv_mx.rowwise_data.data_ptr(), recv_buf.data_ptr())
+        self.assertEqual(recv_mx.scale_inv.data_ptr(), recv_buf.data_ptr() + rc * HIDDEN_DIM)
+        self._assert_mxfp8_matches_bf16(recv_mx, tokens, topk_idx, w, tc)
+
     @_zero_copy_test_include
     def test_caller_provides_dispatch_recv_tokens(self):
         """Caller-supplied recv_tokens (symm-mem-backed in zero-copy): ep_dispatch
@@ -419,6 +528,85 @@ def test_caller_provides_grad_expert_out(self):
         # the caller-owned buffer was used as the combine-bwd scatter target
         self.assertGreater(gbuf.abs().sum().item(), 0.0)
 
+    @_zero_copy_test_include
+    @_mxfp8_align_test
+    def test_combine_bwd_mxfp8_caller_grad_out(self):
+        """MXFP8 combine backward into a single caller buffer sliced into data + e8m0 scales: the
+        returned per-expert GroupedTensor views those regions and, dequantized, matches a bf16
+        combine backward reference on the same routing. Under zero-copy the caller buffer and combine
+        input are symm-mem backed."""
+        self._require_mxfp8_shapes()
+        from transformer_engine.pytorch.constants import MXFP8_BLOCK_SCALING_SIZE
+
+        rc = self.cfg.recv_capacity_per_rank
+        cols = HIDDEN_DIM // MXFP8_BLOCK_SCALING_SIZE
+        topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+        eo_vals = (
+            torch.linspace(-0.5, 0.5, rc * HIDDEN_DIM, device=self.cfg.device)
+            .reshape(rc, HIDDEN_DIM)
+            .to(torch.bfloat16)
+        )
+        # MXFP8 combine backward writes into one caller buffer (data then e8m0 scales)
+        buf_mx = self._make_buffer(combine_bwd_quant_recipe=MXFP8BlockScaling(), alignment=128)
+        _recv, _rw, tc = ep_dispatch(buf_mx, tokens, topk_idx, w)  # seeds the routing
+        nbytes = rc * (HIDDEN_DIM + cols)
+        if ZERO_COPY:
+            grad_buf = symm_mem_alloc((nbytes,), torch.uint8, self.ep_group)
+        else:
+            grad_buf = torch.empty(nbytes, dtype=torch.uint8, device=self.cfg.device)
+        src_mx = eo_vals.detach().clone().requires_grad_(True)
+        out_mx = ep_combine(buf_mx, self._expert_out(src_mx), grad_out=grad_buf)
+        (0.5 * (out_mx.float() ** 2).sum()).backward()
+        g_mx = src_mx.grad  # per-expert GroupedTensor viewing grad_buf
+        self.assertEqual(g_mx.rowwise_data.data_ptr(), grad_buf.data_ptr())
+        self.assertEqual(g_mx.scale_inv.data_ptr(), grad_buf.data_ptr() + rc * HIDDEN_DIM)
+        # bf16 reference combine backward on the same routing
+        buf_bf = self._make_buffer(alignment=128)
+        ep_dispatch(buf_bf, tokens, topk_idx, w)
+        src_bf = eo_vals.detach().clone().requires_grad_(True)
+        out_bf = ep_combine(buf_bf, self._expert_out(src_bf))
+        (0.5 * (out_bf.float() ** 2).sum()).backward()
+        torch.cuda.synchronize()
+        n = int(tc.sum())
+        torch.testing.assert_close(
+            _degroup_mxfp8(g_mx).float(), src_bf.grad.float()[:n], atol=5e-2, rtol=5e-2
+        )
+
+    @_eager_test_include
+    @_zero_copy_test_include
+    @_mxfp8_align_test
+    def test_combine_bwd_mxfp8(self):
+        """MXFP8 combine backward with an internally allocated grad target: the returned per-expert
+        GroupedTensor, dequantized, matches a bf16 combine backward reference on the same routing.
+        Under zero-copy the combine input is symm-mem backed.
+        """
+        self._require_mxfp8_shapes()
+        topk_idx, tokens, w = _make_identity_inputs(self.cfg.rank, self.cfg.ep_size)
+        buf_mx = self._make_buffer(combine_bwd_quant_recipe=MXFP8BlockScaling(), alignment=128)
+        _recv, _rw, tc = ep_dispatch(buf_mx, tokens, topk_idx, w)  # seeds the routing
+        # Combine input rows match the recv total (per-step in eager, capacity otherwise).
+        rows = int(buf_mx.total_recv_tokens.item()) if EAGER else self.cfg.recv_capacity_per_rank
+        eo_vals = (
+            torch.linspace(-0.5, 0.5, rows * HIDDEN_DIM, device=self.cfg.device)
+            .reshape(rows, HIDDEN_DIM)
+            .to(torch.bfloat16)
+        )
+        src_mx = eo_vals.detach().clone().requires_grad_(True)
+        out_mx = ep_combine(buf_mx, self._expert_out(src_mx))
+        (0.5 * (out_mx.float() ** 2).sum()).backward()
+        g_mx = src_mx.grad  # per-expert GroupedTensor
+        # bf16 reference combine backward on the same routing
+        buf_bf = self._make_buffer(alignment=128)
+        ep_dispatch(buf_bf, tokens, topk_idx, w)
+        src_bf = eo_vals.detach().clone().requires_grad_(True)
+        out_bf = ep_combine(buf_bf, self._expert_out(src_bf))
+        (0.5 * (out_bf.float() ** 2).sum()).backward()
+        torch.cuda.synchronize()
+        n = int(tc.sum())
+        torch.testing.assert_close(
+            _degroup_mxfp8(g_mx).float(), src_bf.grad.float()[:n], atol=5e-2, rtol=5e-2
+        )
+
     @_zero_copy_test_include
     def test_zero_copy_pool_auto_alloc(self):
         """Zero-copy with recv/grad left None: ep_dispatch/ep_combine allocate their IO
@@ -527,18 +715,29 @@ def _run_1f1b(self, capture):
 
         recv = [None, None, None]
         # Per-microbatch grad-staging buffers, symm-mem under zero-copy and
-        # pre-allocated so nothing is allocated/freed mid-interleave. The recv
-        # outputs are owned by each EpBuffer (symm-mem under zero-copy).
+        # pre-allocated so nothing is allocated/freed mid-interleave.
         recv_w = [None, None, None]
         rc = self.cfg.recv_capacity_per_rank
         if ZERO_COPY:
             gbuf_t = [symm_mem_alloc((rc, H), torch.bfloat16, self.ep_group) for _ in scales]
             gbuf_w = [symm_mem_alloc((rc,), torch.float32, self.ep_group) for _ in scales]
+            # Persistent symm-mem recv buffers per microbatch: leaving recv None
+            # pool-allocates, which is not CUDA-graph capturable.
+            rbuf_t = [symm_mem_alloc((rc, H), torch.bfloat16, self.ep_group) for _ in scales]
+            rbuf_w = [symm_mem_alloc((rc,), torch.float32, self.ep_group) for _ in scales]
         else:
             gbuf_t = gbuf_w = [None, None, None]
+            rbuf_t = rbuf_w = [None, None, None]
 
         def fwd(k):
-            rt, rw, _ = ep_dispatch(buffers[k], tokens_p[k], idx, w)
+            rt, rw, _ = ep_dispatch(
+                buffers[k],
+                tokens_p[k],
+                idx,
+                w,
+                recv_tokens=rbuf_t[k],
+                recv_topk_weights=rbuf_w[k],
+            )
             recv[k] = self._stage_grad_symm(rt, gbuf_t[k])
             recv_w[k] = self._stage_grad_symm(rw, gbuf_w[k])
 
@@ -631,5 +830,7 @@ def _init_distributed():
     result = runner.run(suite)
     dist.barrier()
     ep_finalize()
+    # Deregister symm-mem windows while the comm is still valid.
+    release_symm_mem_pool()
     dist.destroy_process_group()
     sys.exit(0 if result.wasSuccessful() else 1)
diff --git a/tests/pytorch/distributed/run_test_ep.sh b/tests/pytorch/distributed/run_test_ep.sh
index 10e814eb80..5788c2ac9f 100755
--- a/tests/pytorch/distributed/run_test_ep.sh
+++ b/tests/pytorch/distributed/run_test_ep.sh
@@ -47,11 +47,13 @@ run_pass() {
   local zc="$2"
   local eager="${3:-0}"
   local overflow="${4:-0}"
+  local mxfp8="${5:-0}"
   local log="stdout_ep_${label}.txt"
   echo "=== Running ${SCRIPT} [${label}] on ${NUM_RANKS} GPUs (timeout=${TEST_TIMEOUT_S}s) ==="
   # setsid + kill-after so SIGKILL takes down the whole process group, not just torchrun.
-  NVTE_EP_ZERO_COPY="${zc}" NVTE_EP_EAGER="${eager}" NVTE_EP_OVERFLOW="${overflow}" setsid timeout --foreground \
-    --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \
+  NVTE_EP_ZERO_COPY="${zc}" NVTE_EP_EAGER="${eager}" NVTE_EP_OVERFLOW="${overflow}" \
+    NVTE_EP_MXFP8_PASS="${mxfp8}" \
+    setsid timeout --foreground --kill-after=10 --signal=TERM "${TEST_TIMEOUT_S}" \
     torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_RANKS}" \
     "${SCRIPT}" 2>&1 | tee "${log}"
   local rc=${PIPESTATUS[0]}
@@ -72,5 +74,11 @@ run_pass "default" 0
 run_pass "zero_copy" 1
 run_pass "eager" 0 1
 run_pass "overflow" 0 0 1
+# MXFP8 grouped dispatch pins the per-expert alignment to 128, which the backend caches
+# process-wide, so its tests get their own passes (normal + zero-copy + eager IO). mxfp8 is the
+# 5th arg; eager is the 3rd.
+run_pass "mxfp8" 0 0 0 1
+run_pass "mxfp8_zero_copy" 1 0 0 1
+run_pass "mxfp8_eager" 0 1 0 1
 
 exit $RET
diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp
index 2725008257..b726784d03 100644
--- a/transformer_engine/common/ep/ep_backend.cpp
+++ b/transformer_engine/common/ep/ep_backend.cpp
@@ -47,26 +47,50 @@ ncclDataType_t te_dtype_to_nccl_dtype(NVTEDType dtype) {
       return ncclFloat8e4m3;
     case kNVTEFloat8E5M2:
       return ncclFloat8e5m2;
+    case kNVTEFloat8E8M0:
+      return ncclUint8;
     default:
       NVTE_ERROR("Unsupported NVTEDType for NCCL dtype conversion: ", static_cast<int>(dtype));
   }
   return ncclFloat32;  // unreachable
 }
 
-// shape_out is caller-owned; desc.sizes aliases shape_out.data and must
-// outlive the NCCL EP call.
+// Which part of a TE tensor an NCCL descriptor points at: the data payload, or
+// (for block-scaled tensors) the rowwise scale-inverse that rides alongside it.
+enum class DescSource { kData, kScaleInv };
+
+// Build an NCCL descriptor for a TE tensor's data (kData) or its rowwise
+// scale-inverse (kScaleInv). shape_out is caller-owned; desc.sizes aliases
+// shape_out.data and must outlive the NCCL EP call. Uses the matching window
+// field (win.window / win.scale_window) when set, else the raw pointer.
 inline ncclEpTensor_t make_nccl_ep_tensor(const NVTETensor t, NVTEShape& shape_out,
-                                          const NVTECommWindow& win = {}) {
-  shape_out = nvte_tensor_shape(t);
+                                          const NVTECommWindow& win = {},
+                                          DescSource source = DescSource::kData) {
   ncclEpTensor_t desc = NCCL_EP_TENSOR_INIT;
+  void* raw_ptr = nullptr;
+  ncclWindow_t win_hdl = nullptr;
+  uint64_t win_offset = 0;
+  if (source == DescSource::kData) {
+    shape_out = nvte_tensor_shape(t);
+    desc.datatype = te_dtype_to_nccl_dtype(nvte_tensor_type(t));
+    raw_ptr = nvte_tensor_data(t);
+    win_hdl = win.window;
+    win_offset = win.offset;
+  } else {
+    const SimpleTensor& si = convertNVTETensorCheck(t)->scale_inv;
+    shape_out = nvte_make_shape(si.shape.data(), si.shape.size());
+    desc.datatype = te_dtype_to_nccl_dtype(static_cast<NVTEDType>(si.dtype));
+    raw_ptr = si.dptr;
+    win_hdl = win.scale_window;
+    win_offset = win.scale_offset;
+  }
   desc.ndim = shape_out.ndim;
   desc.sizes = shape_out.data;
-  desc.datatype = te_dtype_to_nccl_dtype(nvte_tensor_type(t));
-  if (win.window != nullptr) {
-    desc.win_hdl = win.window;
-    desc.win_offset = win.offset;
+  if (win_hdl != nullptr) {
+    desc.win_hdl = win_hdl;
+    desc.win_offset = win_offset;
   } else {
-    desc.data = nvte_tensor_data(t);
+    desc.data = raw_ptr;
     NVTE_CHECK(desc.data != nullptr || nvte_tensor_numel(t) == 0,
                "non-empty tensor data must not be null");
   }
@@ -412,16 +436,44 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE
         make_nccl_ep_tensor(recv_topk_weights, recv_topk_weights_shape, recv_topk_weights_win);
   }
 
+  // Block-scaled (e.g. MXFP8): route the per-token scale-inverse alongside the
+  // data. High-precision (bf16/fp16/fp32) and per-tensor FP8 payloads carry the
+  // default delayed scaling mode and skip this. Keys on is_block_scaling so
+  // NVFP4 can reuse this path later.
+  const NVTEScalingMode tokens_scaling_mode = nvte_tensor_scaling_mode(tokens);
+  const bool is_scaled = is_block_scaling(tokens_scaling_mode);
+  NVTEShape scales_in_shape, scales_out_shape;
+  ncclEpTensor_t nccl_scales_in = NCCL_EP_TENSOR_INIT, nccl_scales_out = NCCL_EP_TENSOR_INIT;
+  if (is_scaled) {
+    NVTE_CHECK(is_mxfp8_scaling(tokens_scaling_mode),
+               "EP dispatch supports MXFP8 block scaling only; got scaling mode ",
+               static_cast<int>(tokens_scaling_mode));
+    NVTE_CHECK(nvte_tensor_scaling_mode(recv_tokens) == tokens_scaling_mode,
+               "recv_tokens scaling mode must match tokens scaling mode");
+    nccl_scales_in =
+        make_nccl_ep_tensor(tokens, scales_in_shape, tokens_win, DescSource::kScaleInv);
+    nccl_scales_out =
+        make_nccl_ep_tensor(recv_tokens, scales_out_shape, recv_tokens_win, DescSource::kScaleInv);
+  } else {
+    NVTE_CHECK(!is_fp8_dtype(static_cast<DType>(tok_dtype)),
+               "EP dispatch of FP8 tokens requires a block scaling mode (e.g. MXFP8); "
+               "per-tensor (delayed) FP8 scaling is not supported");
+  }
+
   ncclEpDispatchInputs_t in_struct = NCCL_EP_DISPATCH_INPUTS_INIT;
   in_struct.tokens = &nccl_tokens_in;
   in_struct.topk_weights = is_forward ? &nccl_topk_weights_in : nullptr;
+  in_struct.scales = is_scaled ? &nccl_scales_in : nullptr;
 
   ncclEpDispatchOutputs_t out_struct = NCCL_EP_DISPATCH_OUTPUTS_INIT;
   out_struct.tokens = &nccl_tokens_out;
   out_struct.topk_weights = is_forward ? &nccl_topk_weights_out : nullptr;
+  out_struct.scales = is_scaled ? &nccl_scales_out : nullptr;
 
   ncclEpDispatchConfig_t dispatch_cfg = NCCL_EP_DISPATCH_CONFIG_INIT;
   dispatch_cfg.pass_direction = is_forward ? NCCL_EP_FWD_PASS : NCCL_EP_BWD_PASS;
+  // Block-scaled payloads forward the per-token scale-inverse; select the matching recipe.
+  dispatch_cfg.quant_recipe = is_scaled ? NCCL_EP_DISP_QUANT_FWD : NCCL_EP_DISP_QUANT_NONE;
 
   std::lock_guard<std::mutex> lock(mutex_);
   NVTE_CHECK(initialized_, "EPBackend not initialized");
diff --git a/transformer_engine/common/include/transformer_engine/comm_window.h b/transformer_engine/common/include/transformer_engine/comm_window.h
index 424c350bbd..ed67774493 100644
--- a/transformer_engine/common/include/transformer_engine/comm_window.h
+++ b/transformer_engine/common/include/transformer_engine/comm_window.h
@@ -26,6 +26,9 @@ struct ncclWindow_vidmem;
 typedef struct {
   struct ncclWindow_vidmem* window; /*!< NCCL window, or NULL to use the raw data pointer. */
   uint64_t offset;                  /*!< Byte offset of the payload within window. */
+  struct ncclWindow_vidmem*
+      scale_window;      /*!< Window for a block-scaled tensor's scale-inverse, or NULL for raw. */
+  uint64_t scale_offset; /*!< Byte offset of the scale-inverse within scale_window. */
 } NVTECommWindow;
 
 #ifdef __cplusplus
diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h
index a4800944b5..31417fbff1 100644
--- a/transformer_engine/common/include/transformer_engine/ep.h
+++ b/transformer_engine/common/include/transformer_engine/ep.h
@@ -150,6 +150,12 @@ void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor recv
  *  *_win arguments enable zero-copy via symmem windows; pass NVTECommWindow{}
  *  when unused. Requires a prior nvte_ep_prepare on this handle_mem.
  *
+ *  tokens/recv_tokens may be high-precision (bf16/fp16) or FP8:
+ *  for the latter, set rowwise data and rowwise scale-inverse (unswizzled
+ *  [T, hidden_dim/block]) on the tensor and the scales are routed alongside the
+ *  data. tokens and recv_tokens must share a scaling mode. For now, only MXFP8
+ *  (NVTE_MXFP8_1D_SCALING, e4m3 data + e8m0 scales) is supported.
+ *
  *  \param[in]     handle_mem             uint8 routing-state buffer (from prepare).
  *  \param[in]     topk_idx               [T, top_k] int64 sparse routing indices.
  *  \param[in]     tokens                 [T, hidden_dim] input tokens.
diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index 1a9fabb3bb..be92651656 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -721,14 +721,18 @@ void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens_pe
                 at::Tensor total_recv_tokens);
 
 void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens,
-                 at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights);
+                 at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights,
+                 std::optional<at::Tensor> tokens_scale_inv = std::nullopt,
+                 std::optional<at::Tensor> recv_scale_inv = std::nullopt);
 
 void ep_combine(at::Tensor handle_mem, at::Tensor expert_out, at::Tensor result);
 
 void ep_dispatch_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor g_recv_topk_weights,
                      at::Tensor grad_tokens, at::Tensor grad_topk_weights);
 
-void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out);
+void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out,
+                    std::optional<at::Tensor> grad_scale_inv = std::nullopt,
+                    std::optional<at::Tensor> grad_expert_out_scale_inv = std::nullopt);
 
 // Registers the EP pybind functions on `m`. Defined under NVTE_WITH_NCCL_EP.
 void register_ep_bindings(pybind11::module_ &m);
diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp
index 9bf821721f..cef489dbbb 100644
--- a/transformer_engine/pytorch/csrc/extensions/ep.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp
@@ -73,8 +73,14 @@ NVTECommWindow maybe_make_window(const at::Tensor& t) {
   NVTE_CHECK(nccl_sm != nullptr,
              "Symm-mem backend mismatch: expected NCCLSymmetricMemory. Set the backend to "
              "\"NCCL\" before allocating EP payload buffers.");
-  return NVTECommWindow{static_cast<ncclWindow_t>(nccl_sm->get_window()),
-                        static_cast<uint64_t>(nccl_sm->get_offset())};
+  // rendezvous resolves ``t`` by its storage base, so get_offset() is the allocation's offset in
+  // the NCCL window. Add ``t``'s own storage offset so a slice/view of a symm-mem allocation
+  // (e.g. the scale region carved from a shared recv buffer) resolves to its true position in the
+  // window rather than the allocation base.
+  const uint64_t offset =
+      static_cast<uint64_t>(nccl_sm->get_offset()) +
+      static_cast<uint64_t>(t.storage_offset()) * static_cast<uint64_t>(t.element_size());
+  return NVTECommWindow{static_cast<ncclWindow_t>(nccl_sm->get_window()), offset};
 #else
   (void)t;
   return kNoWindow;
@@ -114,6 +120,34 @@ DType check_topk_idx_dtype(at::Tensor topk_idx) {
 
 using Shape = std::vector<size_t>;
 
+// EP block scaling supports only E4M3 MXFP8 today. A future block-scaled recipe (e.g. NVFP4)
+// would also set is_scaled but carry a non-FP8 token dtype, so key the guard on the FP8 dtype.
+bool is_mxfp8_scaled(bool is_scaled, const at::Tensor& data) {
+  return is_scaled && data.scalar_type() == at::kFloat8_e4m3fn;
+}
+
+// Validate the scale-inverse pair of a block-scaled EP op and return the scale column count
+// (hidden/block). Both scales must be 2D contiguous with matching cols dividing H and numels
+// equal to their row counts times cols; the recv scale must be symm-mem-backed under zero-copy.
+size_t check_mxfp8_scale_pair(const at::Tensor& send_scale, const at::Tensor& recv_scale,
+                              size_t send_rows, size_t recv_rows, size_t H, const char* recv_name) {
+  NVTE_CHECK(send_scale.dim() >= 2 && recv_scale.dim() >= 2,
+             "scale-inverses must be at least 2D [., H/block]");
+  NVTE_CHECK(send_scale.is_contiguous() && recv_scale.is_contiguous(),
+             "scale-inverses must be contiguous");
+  const size_t sc_cols = static_cast<size_t>(send_scale.size(-1));
+  NVTE_CHECK(sc_cols > 0 && H % sc_cols == 0, "scale cols (", sc_cols,
+             ") must be a non-zero divisor of hidden (", H, ")");
+  NVTE_CHECK(static_cast<size_t>(recv_scale.size(-1)) == sc_cols,
+             "recv scale cols must match send scale cols");
+  NVTE_CHECK(static_cast<size_t>(send_scale.numel()) == send_rows * sc_cols,
+             "send scale numel must equal rows * cols");
+  NVTE_CHECK(static_cast<size_t>(recv_scale.numel()) == recv_rows * sc_cols,
+             "recv scale numel must equal rows * cols");
+  check_symm_mem_required(recv_scale, recv_name);
+  return sc_cols;
+}
+
 }  // namespace
 
 bool ep_get_zero_copy() { return g_zero_copy_enabled.load(std::memory_order_relaxed); }
@@ -218,8 +252,13 @@ void ep_prepare(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens_pe
                   total_recv_tokens_te.data(), &layer_cfg, stream);
 }
 
+// tokens_scale_inv / recv_scale_inv are set only for block-scaled dispatch (for
+// now MXFP8): tokens/recv_tokens carry e4m3 data and the scale tensors carry the
+// unswizzled e8m0 scale-inverses [T, H/block]. Both null => bf16/fp16/fp32.
 void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens,
-                 at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights) {
+                 at::Tensor topk_weights, at::Tensor recv_tokens, at::Tensor recv_topk_weights,
+                 std::optional<at::Tensor> tokens_scale_inv,
+                 std::optional<at::Tensor> recv_scale_inv) {
   auto stream = at::cuda::getCurrentCUDAStream().stream();
   NVTE_CHECK(tokens.dim() >= 2, "tokens must be at least 2D [..., H]");
   NVTE_CHECK(topk_idx.dim() >= 2, "topk_idx must be at least 2D [..., top_k]");
@@ -250,16 +289,39 @@ void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens,
   check_symm_mem_required(recv_tokens, "recv_tokens");
   check_symm_mem_required(recv_topk_weights, "recv_topk_weights");
 
+  // Block-scaled dispatch: tokens carry e4m3 data and the scale tensors carry
+  // unswizzled e8m0 scale-inverses [T, H/block]. Scales ride in the tensor; the
+  // backend keys on the tensor's scaling mode. is_mxfp8 is split from is_scaled
+  // so future block-scaled recipes can reuse the scale-routing plumbing while
+  // building their own TE tensors; only MXFP8 is supported for now.
+  const bool is_scaled = tokens_scale_inv.has_value();
+  const bool is_mxfp8 = is_mxfp8_scaled(is_scaled, tokens);
+  size_t sc_cols = 0;
+  if (is_scaled) {
+    NVTE_CHECK(recv_scale_inv.has_value(),
+               "recv_scale_inv must be provided together with tokens_scale_inv");
+    NVTE_CHECK(is_mxfp8, "EP dispatch currently supports only E4M3 MXFP8 block scaling");
+    sc_cols = check_mxfp8_scale_pair(*tokens_scale_inv, *recv_scale_inv, T_flat, recv_pr, H,
+                                     "recv_scale_inv");
+  }
+
   auto tok_dtype = GetTransformerEngineDType(tokens.scalar_type());
   auto handle_mem_te = makeTransformerEngineTensor(
       handle_mem.data_ptr(), Shape{static_cast<size_t>(handle_mem.numel())}, DType::kByte);
   auto topk_idx_te =
       makeTransformerEngineTensor(topk_idx.data_ptr(), Shape{T_flat, topk_n}, idx_dtype);
-  auto tokens_te = makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype);
+  auto tokens_te =
+      is_mxfp8 ? makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype,
+                                             nullptr, nullptr, tokens_scale_inv->data_ptr(),
+                                             Shape{T_flat, sc_cols}, NVTE_MXFP8_1D_SCALING)
+               : makeTransformerEngineTensor(tokens.data_ptr(), Shape{T_flat, H}, tok_dtype);
   auto topk_w_te =
       makeTransformerEngineTensor(topk_weights.data_ptr(), Shape{T_flat, topk_n}, DType::kFloat32);
   auto recv_tokens_te =
-      makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype);
+      is_mxfp8 ? makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype,
+                                             nullptr, nullptr, recv_scale_inv->data_ptr(),
+                                             Shape{recv_pr, sc_cols}, NVTE_MXFP8_1D_SCALING)
+               : makeTransformerEngineTensor(recv_tokens.data_ptr(), Shape{recv_pr, H}, tok_dtype);
   auto recv_topk_w_te =
       makeTransformerEngineTensor(recv_topk_weights.data_ptr(), Shape{recv_pr}, DType::kFloat32);
 
@@ -269,6 +331,17 @@ void ep_dispatch(at::Tensor handle_mem, at::Tensor topk_idx, at::Tensor tokens,
   NVTECommWindow topk_w_win = maybe_make_window(topk_weights);
   NVTECommWindow recv_tokens_win = maybe_make_window(recv_tokens);
   NVTECommWindow recv_topk_w_win = maybe_make_window(recv_topk_weights);
+  // Block-scaled zero-copy: the scale-inverse rides on the data tensor's window.
+  // Send scales (tokens_scale_inv) stay staged like the send data; recv scales
+  // are the one-sided write target and must be symm-mem-backed under zero-copy.
+  if (is_scaled) {
+    const NVTECommWindow tsi_win = maybe_make_window(*tokens_scale_inv);
+    const NVTECommWindow rsi_win = maybe_make_window(*recv_scale_inv);
+    tokens_win.scale_window = tsi_win.window;
+    tokens_win.scale_offset = tsi_win.offset;
+    recv_tokens_win.scale_window = rsi_win.window;
+    recv_tokens_win.scale_offset = rsi_win.offset;
+  }
   nvte_ep_dispatch(handle_mem_te.data(), topk_idx_te.data(), tokens_te.data(), tokens_win,
                    topk_w_te.data(), topk_w_win, recv_tokens_te.data(), recv_tokens_win,
                    recv_topk_w_te.data(), recv_topk_w_win, stream);
@@ -344,7 +417,9 @@ void ep_dispatch_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor g_recv_t
                        g_recv_w_win, grad_tokens_te.data(), grad_topk_w_te.data(), stream);
 }
 
-void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out) {
+void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expert_out,
+                    std::optional<at::Tensor> grad_scale_inv,
+                    std::optional<at::Tensor> grad_expert_out_scale_inv) {
   auto stream = at::cuda::getCurrentCUDAStream().stream();
   NVTE_CHECK(grad.dim() >= 2, "grad must be at least 2D [..., H]");
   NVTE_CHECK(grad_expert_out.dim() >= 2, "grad_expert_out must be at least 2D [..., recv_pr, H]");
@@ -363,17 +438,51 @@ void ep_combine_bwd(at::Tensor handle_mem, at::Tensor grad, at::Tensor grad_expe
   // EpBuffer-owned scatter target and must be symm-mem in zero-copy mode.
   check_symm_mem_required(grad_expert_out, "grad_expert_out");
 
+  // Block-scaled (MXFP8) backward: grad/grad_expert_out carry e4m3 data and the scale
+  // tensors carry the unswizzled e8m0 scale-inverses [., H/block]; the reverse-direction
+  // dispatch forwards them like the forward path. is_mxfp8 is split from is_scaled so
+  // future block-scaled recipes can reuse the plumbing; only MXFP8 is supported for now.
+  const bool is_scaled = grad_scale_inv.has_value();
+  const bool is_mxfp8 = is_mxfp8_scaled(is_scaled, grad);
+  size_t sc_cols = 0;
+  if (is_scaled) {
+    NVTE_CHECK(grad_expert_out_scale_inv.has_value(),
+               "grad_expert_out_scale_inv must be provided together with grad_scale_inv");
+    NVTE_CHECK(is_mxfp8, "EP combine backward currently supports only E4M3 MXFP8 block scaling");
+    sc_cols = check_mxfp8_scale_pair(*grad_scale_inv, *grad_expert_out_scale_inv, T_flat, recv_pr,
+                                     H, "grad_expert_out_scale_inv");
+  }
+
   auto g_dtype = GetTransformerEngineDType(grad.scalar_type());
   auto handle_mem_te = makeTransformerEngineTensor(
       handle_mem.data_ptr(), Shape{static_cast<size_t>(handle_mem.numel())}, DType::kByte);
-  auto grad_te = makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype);
+  auto grad_te = is_mxfp8
+                     ? makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype,
+                                                   nullptr, nullptr, grad_scale_inv->data_ptr(),
+                                                   Shape{T_flat, sc_cols}, NVTE_MXFP8_1D_SCALING)
+                     : makeTransformerEngineTensor(grad.data_ptr(), Shape{T_flat, H}, g_dtype);
   auto grad_expert_out_te =
-      makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype);
+      is_mxfp8
+          ? makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype,
+                                        nullptr, nullptr, grad_expert_out_scale_inv->data_ptr(),
+                                        Shape{recv_pr, sc_cols}, NVTE_MXFP8_1D_SCALING)
+          : makeTransformerEngineTensor(grad_expert_out.data_ptr(), Shape{recv_pr, H}, g_dtype);
 
   // grad is autograd-allocated (staged); grad_expert_out resolves to a symm-mem
   // window in zero-copy mode, else kNoWindow for the staged path.
   NVTECommWindow grad_win = maybe_make_window(grad);
   NVTECommWindow grad_expert_out_win = maybe_make_window(grad_expert_out);
+  // Block-scaled zero-copy: the scale-inverse rides on the data tensor's window,
+  // mirroring the forward dispatch. Send scales stay staged like the send data;
+  // recv scales are the one-sided write target and must be symm-mem-backed.
+  if (is_scaled) {
+    const NVTECommWindow gsi_win = maybe_make_window(*grad_scale_inv);
+    const NVTECommWindow gesi_win = maybe_make_window(*grad_expert_out_scale_inv);
+    grad_win.scale_window = gsi_win.window;
+    grad_win.scale_offset = gsi_win.offset;
+    grad_expert_out_win.scale_window = gesi_win.window;
+    grad_expert_out_win.scale_offset = gesi_win.offset;
+  }
   nvte_ep_combine_bwd(handle_mem_te.data(), grad_te.data(), grad_win, grad_expert_out_te.data(),
                       grad_expert_out_win, stream);
 }
@@ -396,11 +505,16 @@ void register_ep_bindings(pybind11::module_& m) {
         py::arg("tokens_per_expert"), py::arg("top_k"),
         py::arg("dispatch_output_per_expert_alignment"), py::arg("total_recv_tokens"),
         py::call_guard<py::gil_scoped_release>());
-  m.def("ep_dispatch", &ep_dispatch, "EP dispatch", py::call_guard<py::gil_scoped_release>());
+  m.def("ep_dispatch", &ep_dispatch, "EP dispatch", py::arg("handle_mem"), py::arg("topk_idx"),
+        py::arg("tokens"), py::arg("topk_weights"), py::arg("recv_tokens"),
+        py::arg("recv_topk_weights"), py::arg("tokens_scale_inv") = std::nullopt,
+        py::arg("recv_scale_inv") = std::nullopt, py::call_guard<py::gil_scoped_release>());
   m.def("ep_combine", &ep_combine, "EP combine", py::call_guard<py::gil_scoped_release>());
   m.def("ep_dispatch_bwd", &ep_dispatch_bwd, "EP dispatch backward",
         py::call_guard<py::gil_scoped_release>());
-  m.def("ep_combine_bwd", &ep_combine_bwd, "EP combine backward",
+  m.def("ep_combine_bwd", &ep_combine_bwd, "EP combine backward", py::arg("handle_mem"),
+        py::arg("grad"), py::arg("grad_expert_out"), py::arg("grad_scale_inv") = std::nullopt,
+        py::arg("grad_expert_out_scale_inv") = std::nullopt,
         py::call_guard<py::gil_scoped_release>());
 }
 
diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py
index 2dddb6587e..8605a4746b 100644
--- a/transformer_engine/pytorch/distributed.py
+++ b/transformer_engine/pytorch/distributed.py
@@ -3,6 +3,7 @@
 # See LICENSE for license information.
 
 """Methods needed for distributed training (DP/TP)."""
+
 from __future__ import annotations
 
 from collections.abc import Iterable
@@ -50,7 +51,6 @@
 from .tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage
 from ..debug.pytorch.debug_quantization import DebugQuantizedTensor
 
-
 __all__ = ["checkpoint", "CudaRNGStatesTracker"]
 
 
@@ -1885,6 +1885,11 @@ def get_symmetric_memory_tensor(tensor_numel, tensor_dtype, tensor_device, tp_gr
 
 _SYMM_MEM_POOL = None
 _SYMM_MEM_POOL_BACKEND = None
+# Device the pool was created for; the torch symm_mem._symm_mem_pools cache is keyed by it.
+_SYMM_MEM_POOL_DEVICE = None
+# True when the pool was created via torch's get_mem_pool, which caches it in the private
+# symm_mem._symm_mem_pools dict; release then has to drop that cached reference.
+_SYMM_MEM_POOL_TORCH_CACHED = False
 
 
 def _get_symm_mem_pool(device: torch.device, backend: str = "NCCL"):
@@ -1893,12 +1898,14 @@ def _get_symm_mem_pool(device: torch.device, backend: str = "NCCL"):
     backend arg, so the (process-global) backend is always set before the pool is created. The
     collective rendezvous cost is amortized across allocations (paid per new segment, not per buffer).
     """
-    global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND
+    global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND, _SYMM_MEM_POOL_DEVICE, _SYMM_MEM_POOL_TORCH_CACHED
     if _SYMM_MEM_POOL is None:
         symm_mem.set_backend(backend)
         _SYMM_MEM_POOL_BACKEND = backend
+        _SYMM_MEM_POOL_DEVICE = device
         if hasattr(symm_mem, "get_mem_pool"):
             _SYMM_MEM_POOL = symm_mem.get_mem_pool(device)
+            _SYMM_MEM_POOL_TORCH_CACHED = True
         elif hasattr(torch.cuda, "MemPool") and hasattr(symm_mem, "get_mempool_allocator"):
             _SYMM_MEM_POOL = torch.cuda.MemPool(symm_mem.get_mempool_allocator(device))
         else:
@@ -1914,6 +1921,40 @@ def _get_symm_mem_pool(device: torch.device, backend: str = "NCCL"):
     return _SYMM_MEM_POOL
 
 
+def release_symm_mem_pool() -> None:
+    """Free the process-wide symm-mem pool's segments, deregistering their NCCL windows.
+
+    Call before ``dist.destroy_process_group()``: the pool's windows are registered on
+    the group's NCCL comm, which becomes invalid once the group is destroyed. No-op if
+    no pool was created.
+    """
+    global _SYMM_MEM_POOL, _SYMM_MEM_POOL_BACKEND, _SYMM_MEM_POOL_DEVICE, _SYMM_MEM_POOL_TORCH_CACHED
+    if _SYMM_MEM_POOL is None:
+        return
+    # The torch symm_mem._symm_mem_pools cache is keyed by the pool's creation device.
+    device = _SYMM_MEM_POOL_DEVICE
+    _SYMM_MEM_POOL = None
+    _SYMM_MEM_POOL_BACKEND = None
+    _SYMM_MEM_POOL_DEVICE = None
+    torch_cached = _SYMM_MEM_POOL_TORCH_CACHED
+    _SYMM_MEM_POOL_TORCH_CACHED = False
+    # A pool from torch's get_mem_pool is also cached in the private module dict
+    # symm_mem._symm_mem_pools; drop that reference so the segments' refcount reaches zero
+    # and their NCCL windows deregister. Fail loudly if this internal has changed shape,
+    # otherwise the windows would silently leak past destroy_process_group().
+    if torch_cached:
+        pools = getattr(symm_mem, "_symm_mem_pools", None)
+        if not isinstance(pools, dict) or device not in pools:
+            raise RuntimeError(
+                "torch symmetric-memory pool cache (symm_mem._symm_mem_pools) is missing or "
+                "has changed layout; cannot release the pooled segments and their NCCL windows "
+                "would leak past destroy_process_group(). This torch version needs an updated "
+                "release_symm_mem_pool()."
+            )
+        pools.pop(device, None)
+    torch.cuda.empty_cache()
+
+
 def symm_mem_alloc(
     shape,
     dtype: torch.dtype,
diff --git a/transformer_engine/pytorch/ep.py b/transformer_engine/pytorch/ep.py
index 42f605367a..b0c66df4cc 100644
--- a/transformer_engine/pytorch/ep.py
+++ b/transformer_engine/pytorch/ep.py
@@ -7,7 +7,7 @@
 
 import atexit
 import warnings
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 import torch
 import torch.distributed as dist
@@ -15,8 +15,13 @@
 import transformer_engine_torch as tex
 
 from .cpu_offload import mark_not_offload
-from .distributed import symm_mem_alloc
+from .distributed import symm_mem_alloc, release_symm_mem_pool
+from .quantized_tensor import QuantizedTensor
 
+# Type-hint-only import; keeps the ``Recipe`` annotation without a runtime import of
+# common.recipe (the concrete recipe classes are imported lazily where used).
+if TYPE_CHECKING:
+    from ..common.recipe import Recipe
 
 __all__ = [
     "EpBuffer",
@@ -26,6 +31,7 @@
     "ep_dispatch",
     "ep_combine",
     "symm_mem_alloc",
+    "release_symm_mem_pool",
     "is_symm_backed",
 ]
 
@@ -116,7 +122,9 @@ def ep_bootstrap(
 
     ``zero_copy`` opts the EP group into the symm-mem zero-copy IO path; pass
     ``True`` only when payload tensors are allocated via ``symm_mem_alloc``.
-    Requires ``recv_capacity_per_rank``.
+    Requires ``recv_capacity_per_rank``. To capture a CUDA graph, supply
+    persistent recv_tokens / grad_out buffers to dispatch/combine; the pool-based
+    auto-allocation used when they are omitted is not CUDA-graph capturable.
 
     ``num_topk`` is the per-token top-k; it sizes NCCL EP internal buffers.
 
@@ -181,12 +189,16 @@ def ep_finalize() -> None:
     An atexit handler covers normal interpreter shutdown, so most users do not
     need to call this. Call it explicitly only before
     ``dist.destroy_process_group()``, since the borrowed NCCL comm becomes
-    invalid once the PG is destroyed.
+    invalid once the PG is destroyed. This also releases the symm-mem pool, so
+    a caller that used ``symm_mem_alloc(use_pool=True)`` does not need a separate
+    ``release_symm_mem_pool()`` before destroying the PG.
     """
     global _BOOTSTRAPPED, _EP_GROUP, _EAGER
     if not _BOOTSTRAPPED:
         return
     try:
+        # Deregister pooled symm-mem windows while the group's comm is still valid.
+        release_symm_mem_pool()
         tex.ep_finalize()
     finally:
         _BOOTSTRAPPED = False
@@ -237,6 +249,8 @@ class EpBuffer:
         "eager",
         "total_recv_tokens",
         "_host_total_recv_tokens",
+        "dispatch_fwd_quant_recipe",
+        "combine_bwd_quant_recipe",
     )
 
     def __init__(
@@ -249,6 +263,8 @@ def __init__(
         alignment: int = 0,
         payload_dtype: torch.dtype = torch.bfloat16,
         device: Optional[torch.device] = None,
+        dispatch_fwd_quant_recipe: Optional["Recipe"] = None,
+        combine_bwd_quant_recipe: Optional["Recipe"] = None,
     ) -> None:
         if not _BOOTSTRAPPED:
             raise RuntimeError("EpBuffer requires ep_bootstrap() to be called first.")
@@ -274,6 +290,8 @@ def __init__(
         self.payload_dtype = payload_dtype
         self.device = device
         self.zero_copy = bool(tex.ep_get_zero_copy())
+        self.dispatch_fwd_quant_recipe = dispatch_fwd_quant_recipe
+        self.combine_bwd_quant_recipe = combine_bwd_quant_recipe
 
         size_bytes = tex.ep_handle_mem_size(self.top_k, self.alignment)
         self.handle_mem = torch.empty(int(size_bytes), dtype=torch.uint8, device=device)
@@ -319,7 +337,7 @@ def _(*_args, **_kw):
 
 @torch.library.custom_op(
     f"{_LIB}::dispatch",
-    mutates_args=("recv_tokens", "recv_topk_weights"),
+    mutates_args=("recv_tokens", "recv_topk_weights", "recv_scale_inv"),
     device_types="cuda",
 )
 def _dispatch_op(
@@ -329,8 +347,19 @@ def _dispatch_op(
     topk_weights: torch.Tensor,
     recv_tokens: torch.Tensor,
     recv_topk_weights: torch.Tensor,
+    tokens_scale_inv: Optional[torch.Tensor] = None,
+    recv_scale_inv: Optional[torch.Tensor] = None,
 ) -> None:
-    tex.ep_dispatch(handle_mem, topk_idx, tokens, topk_weights, recv_tokens, recv_topk_weights)
+    tex.ep_dispatch(
+        handle_mem,
+        topk_idx,
+        tokens,
+        topk_weights,
+        recv_tokens,
+        recv_topk_weights,
+        tokens_scale_inv,
+        recv_scale_inv,
+    )
 
 
 @_dispatch_op.register_fake
@@ -378,15 +407,17 @@ def _(*_args, **_kw):
 
 @torch.library.custom_op(
     f"{_LIB}::combine_bwd",
-    mutates_args=("grad_expert_out",),
+    mutates_args=("grad_expert_out", "grad_expert_out_scale_inv"),
     device_types="cuda",
 )
 def _combine_bwd_op(
     handle_mem: torch.Tensor,
     grad: torch.Tensor,
     grad_expert_out: torch.Tensor,
+    grad_scale_inv: Optional[torch.Tensor] = None,
+    grad_expert_out_scale_inv: Optional[torch.Tensor] = None,
 ) -> None:
-    tex.ep_combine_bwd(handle_mem, grad, grad_expert_out)
+    tex.ep_combine_bwd(handle_mem, grad, grad_expert_out, grad_scale_inv, grad_expert_out_scale_inv)
 
 
 @_combine_bwd_op.register_fake
@@ -448,32 +479,84 @@ class _EpDispatch(torch.autograd.Function):
     def forward(  # type: ignore[override]
         ctx,
         handle_mem: torch.Tensor,
-        recv_tokens: torch.Tensor,
-        recv_topk_weights: torch.Tensor,
+        recv_tokens: Optional[torch.Tensor],
+        recv_topk_weights: Optional[torch.Tensor],
         topk_idx: torch.Tensor,
         tokens: torch.Tensor,
         topk_weights: torch.Tensor,
+        tokens_scale_inv: Optional[torch.Tensor] = None,
+        token_counts: Optional[torch.Tensor] = None,
+        num_recv_tokens: Optional[int] = None,
+        payload_dtype: torch.dtype = torch.bfloat16,
     ):
-        """Dispatch fwd; prepare must have run into ``handle_mem`` beforehand."""
+        """Dispatch fwd; prepare must have run into ``handle_mem`` beforehand. When scales are set
+        (MXFP8 for now), ``tokens`` is the quantized tensor kept as the autograd operand so grad
+        reaches the pre-quant input. Recv outputs are carved/allocated here: a caller may supply
+        ``recv_tokens`` / ``recv_topk_weights``, else they are sized to ``num_recv_tokens``."""
+        is_scaled = tokens_scale_inv is not None
+        tokens_data = tokens._rowwise_data if isinstance(tokens, QuantizedTensor) else tokens
+        assert tokens_data.dim() == 2, "EP dispatch tokens must be 2D [num_tokens, hidden]"
+        hidden = tokens_data.shape[-1]
+        device = tokens_data.device
+        zero_copy = tex.ep_get_zero_copy()
+
+        recv_scale_inv = None
+        if is_scaled:
+            if tokens._fp8_dtype != tex.DType.kFloat8E4M3:
+                raise NotImplementedError("EP dispatch supports only E4M3 MXFP8 tokens for now.")
+            # recv data + scales share one buffer (data then scales); carve or allocate it here.
+            recv_tokens, recv_scale_inv = _scale_alloc_io(
+                recv_tokens,
+                num_recv_tokens,
+                hidden,
+                tokens_scale_inv.shape[-1],
+                tokens_data.dtype,
+                tokens_scale_inv.dtype,
+                device,
+                zero_copy,
+            )
+            # Reinterpret byte-backed FP8 data as the fp8 dtype so the backend sees a scaled tensor.
+            dispatch_tokens = tokens_data.view(torch.float8_e4m3fn)
+            dispatch_recv = recv_tokens.view(torch.float8_e4m3fn)
+        else:
+            if recv_tokens is None:
+                recv_tokens = _alloc_io((num_recv_tokens, hidden), payload_dtype, device, zero_copy)
+            dispatch_tokens = tokens_data
+            dispatch_recv = recv_tokens
+        if recv_topk_weights is None:
+            recv_topk_weights = _alloc_io((num_recv_tokens,), torch.float32, device, zero_copy)
         torch.ops.transformer_engine_ep.dispatch(
             handle_mem,
             topk_idx,
-            tokens,
+            dispatch_tokens,
             topk_weights,
-            recv_tokens,
+            dispatch_recv,
             recv_topk_weights,
+            tokens_scale_inv,
+            recv_scale_inv,
         )
         ctx.save_for_backward(handle_mem)
         ctx.tokens_shape = tokens.shape
-        ctx.tokens_dtype = tokens.dtype
         ctx.topk_weights_shape = topk_weights.shape
-        ctx.tokens_T_flat = tokens.numel() // tokens.shape[-1]
+        ctx.num_tokens = tokens_data.shape[0]
         ctx.topk_T_flat = topk_weights.numel() // topk_weights.shape[-1]
         ctx.top_k = topk_weights.shape[-1]
-        ctx.hidden_dim = tokens.shape[-1]
+        ctx.hidden_dim = hidden
         # Detach so the long-lived buffers aren't tracked as differentiable outputs;
-        # autograd re-attaches grad_fn pointing back at this Function.
-        return recv_tokens.detach(), recv_topk_weights.detach()
+        # autograd re-attaches grad_fn pointing back at this Function. For scaled inputs
+        # the expert-major recv data + scales are wrapped into a per-expert GroupedTensor
+        # so downstream grouped GEMM and autograd see a proper quantized grouped tensor.
+        if is_scaled:
+            recv_out = _make_grouped_mxfp8(
+                recv_tokens.view(tokens._rowwise_data.dtype),
+                recv_scale_inv,
+                token_counts,
+                tokens._fp8_dtype,
+                tokens.dtype,
+            )
+        else:
+            recv_out = recv_tokens.detach()
+        return recv_out, recv_topk_weights.detach()
 
     @staticmethod
     def backward(ctx, g_recv_tokens, g_recv_topk_weights):  # type: ignore[override]
@@ -482,8 +565,10 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights):  # type: ignore[override]
         device = handle_mem.device
         g_recv_tokens = g_recv_tokens.contiguous()
         g_recv_topk_weights = g_recv_topk_weights.contiguous()
+        # Dispatch grad follows the recv grad's (high-precision) dtype; the quantizer's STE
+        # owns the fp8 boundary for scaled inputs.
         grad_tokens = torch.empty(
-            ctx.tokens_T_flat, ctx.hidden_dim, dtype=ctx.tokens_dtype, device=device
+            ctx.num_tokens, ctx.hidden_dim, dtype=g_recv_tokens.dtype, device=device
         )
         grad_topk_weights = torch.empty(
             ctx.topk_T_flat, ctx.top_k, dtype=torch.float32, device=device
@@ -502,21 +587,21 @@ def backward(ctx, g_recv_tokens, g_recv_topk_weights):  # type: ignore[override]
             None,  # topk_idx
             grad_tokens.view(ctx.tokens_shape),
             grad_topk_weights.view(ctx.topk_weights_shape),
+            None,  # tokens_scale_inv (scales; non-differentiable)
+            None,  # token_counts (per-expert counts; non-differentiable)
+            None,  # num_recv_tokens (sizing scalar)
+            None,  # payload_dtype (sizing scalar)
         )
 
 
 class _EpCombine(torch.autograd.Function):
-    """Autograd combine.
+    """Autograd combine; bwd scatters the expert_out grad into ``grad_out``. When the caller
+    supplies it that buffer is used as-is; otherwise it is allocated in the backward from the
+    symm-mem pool in zero-copy mode (one-sided target) or a plain tensor in normal mode (keeps
+    allocation torch.compile / CUDA-graph safe and lets autograd own the grad's lifetime).
 
-    bwd scatters the expert_out grad into ``grad_out``. When the caller supplies it
-    (mcore-managed mode) that buffer is used as-is; otherwise it is allocated on the
-    fly here — from the symm-mem pool in zero-copy mode (one-sided target), or a plain
-    tensor in normal mode (keeps allocation torch.compile / CUDA-graph safe and lets
-    autograd own the grad's lifetime).
-
-    ``grad_out`` is the backward's scatter target (an output it writes, never reads),
-    so it is stashed as a plain ctx attribute rather than via save_for_backward, which
-    would version-track a tensor we mutate.
+    ``grad_out`` is a write-only scatter target, so it is stashed as a plain ctx attribute rather
+    than via save_for_backward, which would version-track a tensor we mutate.
     """
 
     @staticmethod
@@ -527,55 +612,100 @@ def forward(  # type: ignore[override]
         hidden_dim: int,
         grad_out: Optional[torch.Tensor],
         expert_out: torch.Tensor,
+        bwd_quant_recipe=None,
+        token_counts: Optional[torch.Tensor] = None,
     ):
-        """Combine fwd; stashes the bwd grad target or expert_out shape to size it."""
+        """Combine fwd; stashes the bwd grad target or expert_out shape to size it. When
+        ``bwd_quant_recipe`` is set, the backward sends the result-grad as MXFP8."""
         device = expert_out.device
         result = torch.empty(num_local_tokens, hidden_dim, dtype=expert_out.dtype, device=device)
         torch.ops.transformer_engine_ep.combine(handle_mem, expert_out, result)
         ctx.save_for_backward(handle_mem)
         ctx.grad_out = grad_out
-        if grad_out is None:
-            ctx.expert_out_shape = expert_out.shape
-            ctx.expert_out_dtype = expert_out.dtype
-            ctx.device = device
+        ctx.bwd_quant_recipe = bwd_quant_recipe
+        ctx.token_counts = token_counts
+        ctx.expert_out_shape = expert_out.shape
+        ctx.expert_out_dtype = expert_out.dtype
+        ctx.device = device
         return result
 
     @staticmethod
     def backward(ctx, g_result):  # type: ignore[override]
-        """Combine bwd; scatters the result grad into the grad target."""
+        """Combine bwd; scatters the result-grad to expert positions. High-precision sends the grad
+        as-is; a quantized recipe (MXFP8 today) quantizes it and returns the expert_out grad as a
+        per-expert GroupedTensor."""
         if not g_result.is_contiguous():
             g_result = g_result.contiguous()
         (handle_mem,) = ctx.saved_tensors
-        grad_expert_out = ctx.grad_out
-        if grad_expert_out is None:
-            grad_expert_out = _alloc_io(
-                ctx.expert_out_shape, ctx.expert_out_dtype, ctx.device, tex.ep_get_zero_copy()
+
+        if ctx.bwd_quant_recipe is None:
+            grad_expert_out = ctx.grad_out
+            if grad_expert_out is None:
+                grad_expert_out = _alloc_io(
+                    ctx.expert_out_shape, ctx.expert_out_dtype, ctx.device, tex.ep_get_zero_copy()
+                )
+            torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out)
+        else:
+            mx, g_scale_inv = _quantize_mxfp8(g_result)
+            g_data = mx._rowwise_data
+            recv_pr, hidden = ctx.expert_out_shape[0], ctx.expert_out_shape[-1]
+            ge_data, ge_scale_inv = _scale_alloc_io(
+                ctx.grad_out,
+                recv_pr,
+                hidden,
+                g_scale_inv.shape[-1],
+                g_data.dtype,
+                g_scale_inv.dtype,
+                ctx.device,
+                tex.ep_get_zero_copy(),
+            )
+            # The backend keys on the fp8 scaling mode; reinterpret the byte-backed data as fp8.
+            torch.ops.transformer_engine_ep.combine_bwd(
+                handle_mem,
+                g_data.view(torch.float8_e4m3fn),
+                ge_data.view(torch.float8_e4m3fn),
+                g_scale_inv,
+                ge_scale_inv,
             )
-        torch.ops.transformer_engine_ep.combine_bwd(handle_mem, g_result, grad_expert_out)
+            grad_expert_out = _make_grouped_mxfp8(
+                ge_data, ge_scale_inv, ctx.token_counts, mx._fp8_dtype, ctx.expert_out_dtype
+            )
+
         return (
             None,  # handle_mem
             None,  # num_local_tokens
             None,  # hidden_dim
             None,  # grad_out
             grad_expert_out,
+            None,  # bwd_quant_recipe
+            None,  # token_counts
         )
 
 
 # Public high-level wrappers
 
 
-# NCCL EP currently only supports bfloat16 payload tensors.
+# NCCL EP inputs are bfloat16; MXFP8 is applied internally via the buffer's dispatch_fwd_quant_recipe.
 def _require_bf16(name: str, t: torch.Tensor) -> None:
     if t.dtype is not torch.bfloat16:
         raise NotImplementedError(
-            f"NCCL EP currently supports only bfloat16 payloads; got {name}.dtype={t.dtype}."
+            "NCCL EP currently supports only bfloat16 or MXFP8 payloads; got"
+            f" {name}.dtype={t.dtype}."
         )
 
 
 def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tensor:
     """Allocate a dispatch/combine IO tensor the caller did not supply: from the symm-mem pool in
-    zero-copy mode (auto-registered segment, lifecycle managed by torch refcount), else plain."""
+    zero-copy mode (auto-registered segment, lifecycle managed by torch refcount), else plain.
+
+    The zero-copy pool path is not CUDA-graph capturable; supply persistent recv_tokens / grad_out
+    buffers to capture a graph."""
     if zero_copy:
+        if torch.cuda.is_current_stream_capturing():
+            raise RuntimeError(
+                "EP zero-copy pool allocation is not CUDA-graph capturable; supply persistent "
+                "recv_tokens / grad_out buffers (allocated once via symm_mem_alloc) before capture."
+            )
         t = symm_mem_alloc(shape, dtype, _EP_GROUP, device=device, use_pool=True)
         # symm-mem storage is non-resizable; exempt it from CPU activation offloading (which
         # releases via storage.resize_(0)). Matters for bf16 recv_tokens (the saved activation).
@@ -584,6 +714,103 @@ def _alloc_io(shape, dtype: torch.dtype, device, zero_copy: bool) -> torch.Tenso
     return torch.empty(*shape, dtype=dtype, device=device)
 
 
+def _quantize_mxfp8(x: torch.Tensor):
+    """Quantize a high-precision tensor to MXFP8 and return ``(quantized_tensor, scale_inv)`` where
+    ``scale_inv`` is the compact ``[T, H/block]`` scale the EP backend routes. The quantized tensor
+    is returned so callers can keep it as the autograd operand; its ``_rowwise_data`` is the fp8
+    payload and ``scale_inv.shape[-1]`` the scale-column count. EP routes and returns E4M3 data in
+    both directions, so quantize to E4M3 regardless of pass. Strips the GEMM scale row padding to
+    the compact ``[T, H/block]`` layout; requires a 16-byte-aligned scale row."""
+    from .constants import MXFP8_BLOCK_SCALING_SIZE
+    from .tensor.mxfp8_tensor import MXFP8Quantizer
+
+    mx = MXFP8Quantizer(tex.DType.kFloat8E4M3, rowwise=True, columnwise=False).quantize(x)
+    if mx._with_gemm_swizzled_scales:
+        raise RuntimeError(
+            "internal MXFP8 quantization produced swizzled scales; EP dispatch needs compact."
+        )
+    data = mx._rowwise_data
+    scale_inv = mx._rowwise_scale_inv
+    if data is None or scale_inv is None:
+        raise ValueError("MXFP8 tokens must carry rowwise data and scale_inv for EP dispatch.")
+    t_flat = x.shape[0]
+    hidden = x.shape[-1]
+    cols = hidden // MXFP8_BLOCK_SCALING_SIZE
+    # The backend forwards each token's scale row with a 16-byte-aligned store, so the row
+    # (cols * dtype bytes) must be a multiple of 16.
+    scale_row_bytes = cols * scale_inv.element_size()
+    if scale_row_bytes % 16 != 0:
+        raise ValueError(
+            f"MXFP8 dispatch requires a 16-byte-aligned scale row; hidden={hidden} gives "
+            f"{scale_row_bytes} bytes. Use a hidden size that is a multiple of "
+            f"{16 * MXFP8_BLOCK_SCALING_SIZE}."
+        )
+    # scale_inv is 2D [round_up(T, 128), cols]; drop the row padding to the logical [T, H/block]
+    # the backend expects. cols is a multiple of 4 (16-byte row), so no column padding and the
+    # slice stays contiguous; assert rather than force a copy.
+    scale_inv = scale_inv[:t_flat, :cols]
+    if not scale_inv.is_contiguous():
+        raise ValueError(
+            "MXFP8 dispatch requires compact contiguous scales [T, H/block]; got a "
+            f"non-contiguous [{t_flat}, {cols}] slice."
+        )
+    return mx, scale_inv
+
+
+def _scale_alloc_io(buf, rows, data_cols, scale_cols, data_dtype, scale_dtype, device, zero_copy):
+    """Block-scaled output data + scale buffers, each ``rows`` tall, laid out back-to-back
+    (``[rows, data_cols]`` data of ``data_dtype`` then ``[rows, scale_cols]`` scales of
+    ``scale_dtype``). Carve both from a single caller ``buf`` when it is large enough, so one
+    symm-mem window backs both views; else allocate them (symm-mem pool under zero-copy, else
+    plain). Recipe-agnostic: byte sizes come from the element sizes."""
+    data_bytes = rows * data_cols * torch.empty((), dtype=data_dtype).element_size()
+    scale_bytes = rows * scale_cols * torch.empty((), dtype=scale_dtype).element_size()
+    if buf is not None:
+        # Reinterpret in place; a non-contiguous buf would force a copy and leave the caller's
+        # buffer unwritten, so require contiguous and view rather than reshape.
+        if not buf.is_contiguous():
+            raise ValueError("scaled output buffer must be contiguous.")
+        flat = buf.view(-1).view(torch.uint8)
+        if flat.numel() < data_bytes + scale_bytes:
+            raise ValueError(
+                f"scaled output buffer too small: need {data_bytes + scale_bytes} bytes "
+                f"(data + scales), got {flat.numel()}."
+            )
+        data = flat[:data_bytes].view(data_dtype).reshape(rows, data_cols)
+        scale_inv = (
+            flat[data_bytes : data_bytes + scale_bytes].view(scale_dtype).reshape(rows, scale_cols)
+        )
+        return data, scale_inv
+    data = _alloc_io((rows, data_cols), data_dtype, device, zero_copy)
+    scale_inv = _alloc_io((rows, scale_cols), scale_dtype, device, zero_copy)
+    return data, scale_inv
+
+
+def _make_grouped_mxfp8(data, scale_inv, token_counts, fp8_dtype, fake_dtype):
+    """Wrap expert-major MXFP8 recv data + compact e8m0 scales as a per-expert ``GroupedTensor``.
+
+    ``token_counts`` (int64 [num_local_experts]) is the padded per-expert row counts (128-aligned),
+    used as the group sizes. Grouping is device-side (first_dims/tensor_offsets), so the counts never
+    sync to host; the outer shape is the static recv capacity, bounded per expert by first_dims.
+    """
+    from .tensor.grouped_tensor import GroupedTensor
+    from .tensor.mxfp8_tensor import MXFP8Quantizer
+
+    assert data.dim() == 2, "recv data must be 2D [capacity_rows, hidden]"
+    capacity_rows, hidden = data.shape
+    quantizer = MXFP8Quantizer(fp8_dtype, rowwise=True, columnwise=False)
+    return GroupedTensor(
+        shape=(capacity_rows, hidden),
+        dtype=fake_dtype,
+        num_tensors=token_counts.numel(),
+        quantizer=quantizer,
+        data=data.reshape(-1).detach(),
+        scale_inv=scale_inv.reshape(-1).detach(),
+        first_dims=token_counts,
+        tensor_offsets=tex.splits_to_offsets(token_counts, hidden),
+    )
+
+
 def ep_dispatch(
     buffer: EpBuffer,
     tokens: torch.Tensor,
@@ -593,39 +820,59 @@ def ep_dispatch(
     recv_tokens: Optional[torch.Tensor] = None,
     recv_topk_weights: Optional[torch.Tensor] = None,
 ):
-    """Prepare + dispatch with autograd. topk_idx must be int32 or int64.
-
-    ``recv_tokens`` / ``recv_topk_weights`` are the dispatch recv outputs: pass caller-owned buffers
-    (mcore-managed mode; in zero-copy they must be symm-mem-backed) or leave them None to allocate on
-    the fly (zero-copy: symm-mem pool; normal: plain). In eager mode the recv outputs are sized to
-    this step's recv-token total and must not be caller-supplied. Returns (recv_tokens,
-    recv_topk_weights, tokens_per_expert); tokens_per_expert is non-diff. See ``buffer.total_recv_tokens`` for
-    the per-step recv total.
+    """Prepare + dispatch with autograd. ``tokens`` is bfloat16; ``topk_idx`` is int32 or int64.
+
+    When the buffer's ``dispatch_fwd_quant_recipe`` is set (``MXFP8BlockScaling`` only for now), tokens
+    are quantized internally and recv is returned as a per-expert ``GroupedTensor``; otherwise recv
+    stays bfloat16. A pre-quantized ``tokens`` is not accepted.
+
+    ``recv_tokens`` / ``recv_topk_weights`` are the recv outputs: pass caller-owned buffers
+    (symm-mem-backed under zero-copy) or leave them None to allocate. For MXFP8 the recv data and
+    scales share ``recv_tokens`` (data then scales), so size it to at least
+    ``recv_capacity_per_rank * (hidden + hidden/block)`` bytes. Eager mode sizes the recv outputs
+    per step and forbids caller-supplied buffers. Under zero-copy, leaving them None allocates from
+    the symm-mem pool, which is not CUDA-graph capturable; pass persistent buffers to capture a graph.
+
+    Returns (recv_tokens, recv_topk_weights, tokens_per_expert); tokens_per_expert is non-diff. See
+    ``buffer.total_recv_tokens`` for the per-step recv total.
     """
-    _require_bf16("tokens", tokens)
     if topk_weights.dtype is not torch.float32:
         raise TypeError(
             f"topk_weights must be float32; got dtype={topk_weights.dtype}. "
             "Cast with topk_weights.float() before calling."
         )
+    if isinstance(tokens, QuantizedTensor):
+        raise NotImplementedError(
+            "NCCL EP dispatch takes a bfloat16 input and quantizes internally when the buffer's "
+            "dispatch_fwd_quant_recipe is set; a pre-quantized tensor is not accepted."
+        )
+    _require_bf16("tokens", tokens)
     if buffer.eager and (recv_tokens is not None or recv_topk_weights is not None):
         raise ValueError(
             "eager mode sizes the recv outputs from the per-step recv-token total "
             "and cannot use caller-supplied recv_tokens / recv_topk_weights"
         )
+
     # Prepare (routing AllGather) up front so the recv outputs can be sized; in
     # eager mode ep_prepare also host-syncs this step's recv-token total.
     tokens_per_expert = ep_prepare(buffer, topk_idx)
-    rows = buffer._host_total_recv_tokens if buffer.eager else buffer.recv_capacity_per_rank
-    if recv_tokens is None:
-        recv_tokens = _alloc_io(
-            (rows, buffer.hidden_dim),
-            buffer.payload_dtype,
-            buffer.device,
-            buffer.zero_copy,
-        )
-    if recv_topk_weights is None:
-        recv_topk_weights = _alloc_io((rows,), torch.float32, buffer.device, buffer.zero_copy)
+    num_recv_tokens = (
+        buffer._host_total_recv_tokens if buffer.eager else buffer.recv_capacity_per_rank
+    )
+
+    tokens_scale_inv = None
+    if buffer.dispatch_fwd_quant_recipe is not None:
+        from ..common.recipe import MXFP8BlockScaling
+
+        if not isinstance(buffer.dispatch_fwd_quant_recipe, MXFP8BlockScaling):
+            raise NotImplementedError(
+                "EP block-scaled dispatch supports MXFP8BlockScaling only; got "
+                f"{type(buffer.dispatch_fwd_quant_recipe).__name__}."
+            )
+        # Quantize here (not in forward) so the quantized tensor stays the autograd operand and grad
+        # reaches the pre-quant input; forward then carves the recv buffers and routes.
+        tokens, tokens_scale_inv = _quantize_mxfp8(tokens)
+
     recv_tokens, recv_topk_weights = _EpDispatch.apply(
         buffer.handle_mem,
         recv_tokens,
@@ -633,6 +880,10 @@ def ep_dispatch(
         topk_idx,
         tokens,
         topk_weights,
+        tokens_scale_inv,
+        tokens_per_expert,
+        num_recv_tokens,
+        buffer.payload_dtype,
     )
     return recv_tokens, recv_topk_weights, tokens_per_expert
 
@@ -646,11 +897,14 @@ def ep_combine(
 ):
     """Combine with autograd; caller pre-applies topk weighting.
 
-    ``expert_out`` is the combine input (always caller-supplied; in zero-copy it must be symm-mem-
-    backed). ``grad_out`` is the backward's grad target: pass a caller-owned buffer (mcore-managed
-    mode) or leave it None to allocate on the fly in the backward (zero-copy: symm-mem pool; normal:
-    plain). In eager mode the grad target is sized per step and must not be caller-supplied. Result
-    shape is (num_local_tokens, buffer.hidden_dim); defaults to buffer.max_tokens_per_rank rows.
+    ``expert_out`` is the combine input (symm-mem-backed under zero-copy). ``grad_out`` is the
+    backward's grad target: pass a caller-owned buffer or leave it None to allocate. For MXFP8 the
+    grad data and scales share ``grad_out`` (data then scales), so size it to at least
+    ``recv_capacity_per_rank * (hidden + hidden/block)`` bytes (non-zero-copy only). Eager mode sizes
+    the grad target per step and forbids a caller-supplied buffer. Under zero-copy, leaving it None
+    allocates from the symm-mem pool, which is not CUDA-graph capturable; pass a persistent buffer to
+    capture a graph. Result shape is (num_local_tokens, hidden_dim); num_local_tokens defaults to
+    buffer.max_tokens_per_rank.
     """
     _require_bf16("expert_out", expert_out)
     if buffer.eager and grad_out is not None:
@@ -660,10 +914,24 @@ def ep_combine(
         )
     if num_local_tokens is None:
         num_local_tokens = buffer.max_tokens_per_rank
+    # When combine_bwd_quant_recipe is set the combine backward sends the result-grad over the
+    # wire as MXFP8 and returns the expert_out grad as a GroupedTensor.
+    bwd_quant_recipe = None
+    if buffer.combine_bwd_quant_recipe is not None:
+        from ..common.recipe import MXFP8BlockScaling
+
+        if not isinstance(buffer.combine_bwd_quant_recipe, MXFP8BlockScaling):
+            raise NotImplementedError(
+                "EP combine backward supports MXFP8BlockScaling only; got "
+                f"{type(buffer.combine_bwd_quant_recipe).__name__}."
+            )
+        bwd_quant_recipe = buffer.combine_bwd_quant_recipe
     return _EpCombine.apply(
         buffer.handle_mem,
         num_local_tokens,
         buffer.hidden_dim,
         grad_out,
         expert_out,
+        bwd_quant_recipe,
+        buffer.tokens_per_expert,
     )

From 734ce31e9f9d03bc6e528b53510d3971653f58af Mon Sep 17 00:00:00 2001
From: Aditya Singh <60082699+adityasingh2400@users.noreply.github.com>
Date: Tue, 11 Aug 2026 21:14:07 -0700
Subject: [PATCH 56/67] [PyTorch] Advance FusedAdam step counter for empty
 param groups (#3318)

FusedAdam.step() skipped a param group with no parameters before touching
its step counter, so a group that is empty on one data-parallel rank and
populated on another stopped counting on the empty ranks. Since step is
stored in param_groups it is checkpointed, and a rank that loads its shard
from a rank where the group was empty resumes with a stale step and a wrong
bias correction.

Move the counter update above the empty-group skip. Empty groups have no
parameter to read a device from, so the capturable tensor now falls back to
the device of the optimizer scratch buffer.

Fixes #1986

Signed-off-by: Aditya Singh <adisin650@gmail.com>
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
---
 tests/pytorch/test_fused_optimizer.py         | 25 +++++++++++++++++
 .../pytorch/optimizers/fused_adam.py          | 28 ++++++++++++++-----
 2 files changed, 46 insertions(+), 7 deletions(-)

diff --git a/tests/pytorch/test_fused_optimizer.py b/tests/pytorch/test_fused_optimizer.py
index 6832ef89dd..50dc8e724c 100644
--- a/tests/pytorch/test_fused_optimizer.py
+++ b/tests/pytorch/test_fused_optimizer.py
@@ -166,6 +166,31 @@ def test_frozen_model(self):
 
             torch.testing.assert_close(ref_param, tst_param)
 
+    @pytest.mark.parametrize("capturable", [False, True])
+    def test_empty_param_group_advances_step(self, capturable):
+        # An empty param group must advance its step counter like a populated one.
+        # The same group can be empty on one data-parallel rank and populated on
+        # another, and "step" is part of the checkpoint, so a group that stops
+        # counting makes a resumed run apply a stale bias correction.
+        param = torch.nn.Parameter(torch.rand(4, dtype=torch.float, device="cuda"))
+        tst_optim = self.fused_optim(
+            [{"params": [param]}, {"params": []}], capturable=capturable, **self.options
+        )
+
+        num_steps = 3
+        for _ in range(num_steps):
+            param.grad = torch.rand_like(param)
+            tst_optim.step()
+
+        populated_step, empty_step = (int(g["step"]) for g in tst_optim.param_groups)
+        assert populated_step == num_steps
+        assert empty_step == populated_step
+
+        # The counter is checkpointed through the param groups, so the empty group
+        # has to round trip with the same value as the populated one.
+        checkpoint = tst_optim.state_dict()
+        assert int(checkpoint["param_groups"][1]["step"]) == num_steps
+
     def test_empty_param_at_end_of_group(self):
         tensors = [
             torch.ones(4, dtype=torch.float, device="cuda"),
diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py
index b48953e8e3..15290d7723 100644
--- a/transformer_engine/pytorch/optimizers/fused_adam.py
+++ b/transformer_engine/pytorch/optimizers/fused_adam.py
@@ -557,12 +557,11 @@ def step(self, closure=None, grad_scaler=None):
             loss = closure()
 
         for group in self.param_groups:
-            if len(group["params"]) == 0:
-                continue
-            device = group["params"][0].device
-            bias_correction = 1 if group["bias_correction"] else 0
-            beta1, beta2 = group["betas"]
-
+            # Advance the step counter before skipping empty groups. A param group can be
+            # empty on some data-parallel ranks and populated on others, so incrementing
+            # only for populated groups desynchronizes "step" across the ranks that share
+            # an optimizer state shard. A rank then resumes from a checkpoint written by a
+            # rank where the group was empty and applies a stale bias correction.
             # assume same step across group now to simplify things
             # per parameter step can be easily support by making it tensor, or pass list into kernel
             if "step" in group:
@@ -570,10 +569,25 @@ def step(self, closure=None, grad_scaler=None):
                     1 if not self.capturable else (self._dummy_overflow_buf != 1).to(torch.int)
                 )
             else:
+                # Empty groups have no parameter to take the device from, so fall back to
+                # the device of the optimizer's own scratch buffer.
+                step_device = (
+                    group["params"][0].device
+                    if len(group["params"]) > 0
+                    else self._dummy_overflow_buf.device
+                )
                 group["step"] = (
-                    1 if not self.capturable else torch.tensor([1], dtype=torch.int, device=device)
+                    1
+                    if not self.capturable
+                    else torch.tensor([1], dtype=torch.int, device=step_device)
                 )
 
+            if len(group["params"]) == 0:
+                continue
+            device = group["params"][0].device
+            bias_correction = 1 if group["bias_correction"] else 0
+            beta1, beta2 = group["betas"]
+
             # create lists for multi-tensor apply
             p_main_of_fp8_model = []
             p_main_of_f16_model = []

From fc1fc0e787fc4fe3b50c71e4a9e353c0593615ff Mon Sep 17 00:00:00 2001
From: David Kogan <38084383+davidkny22@users.noreply.github.com>
Date: Wed, 12 Aug 2026 00:20:38 -0400
Subject: [PATCH 57/67] [Common] Fix NVFP4 stochastic rounding on architectures
 without cvt.rs (#3281)

* Fix NVFP4 stochastic rounding on architectures without cvt.rs

The four e2m1 stochastic rounding helpers gate on
ARCH_HAS_STOCHASTIC_ROUNDING, which is sm_100 and sm_103, where the
cvt.rs instruction exists. The other branch is NVTE_DEVICE_ERROR, which
is printf plus assert(0), and release builds define NDEBUG, so on
sm_120 and sm_121 the quantize returns all-zero FP4 data while printing
once per thread. Replace that branch with a software e2m1 stochastic
rounder that takes 8 random bits per element, the same as the
instruction's rbits operand, and follows cvt.satfinite at the edges.
Architectures with cvt.rs keep the asm path.

Signed-off-by: David Kogan <davidkny22@gmail.com>

* Make the software e2m1 stochastic rounder branchless

The rounding ladder becomes a floor in units of the grid step at the
jittered value, and the isnan branch folds into the final sign select,
since fminf carries a NaN through to max_norm and the select keeps it
positive, matching cvt.satfinite. Output is bit-identical to the
previous form across all 2^32 float bit patterns at 16 rbits values and
across every float with magnitude in [1, 8) at all 256 rbits values. An
isolated microbenchmark on sm_121 runs 1.34x faster.

Signed-off-by: David Kogan <davidkny22@gmail.com>

---------

Signed-off-by: David Kogan <davidkny22@gmail.com>
---
 tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py | 36 ++++++++++
 ...quantize_transpose_vector_blockwise_fp4.cu | 10 +--
 transformer_engine/common/util/ptx.cuh        | 66 ++++++++++++++++---
 3 files changed, 98 insertions(+), 14 deletions(-)

diff --git a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py
index b14eeb815b..2ed496c730 100755
--- a/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py
+++ b/tests/pytorch/nvfp4/test_nvfp4_sr_quantize.py
@@ -431,3 +431,39 @@ def test_group_stochastic_rounding_quantization_versus_reference(
         num_splits=num_splits,
         use_tex_split_quantize=use_tex_split_quantize,
     )
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str)
+@pytest.mark.parametrize("use_2D", [False, True], ids=str)
+def test_stochastic_rounding_picks_an_adjacent_fp4_value(
+    x_dtype: torch.dtype, use_2D: bool
+) -> None:
+    device = "cuda"
+    torch.manual_seed(seed)
+    M, N = 512, 1024
+    x = torch.randn((M, N), dtype=x_dtype, device=device)
+    amax = torch.max(torch.abs(x)).float()
+
+    q, s, q_t, s_t = quantize_fp4(x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=False)
+    q_redraw, _, q_t_redraw, _ = quantize_fp4(
+        x, use_stochastic_rounding=True, use_2D=use_2D, use_RHT=False
+    )
+    assert not torch.equal(q, q_redraw), "two stochastic rounding draws are identical"
+    assert not torch.equal(q_t, q_t_redraw), "two stochastic rounding draws are identical"
+
+    def check_adjacent(qx: torch.Tensor, sx: torch.Tensor, reference: torch.Tensor) -> None:
+        # The kernel rounds reference / block_scale onto the E2M1 grid, so every
+        # output has to be one of the two grid values that bracket it, i.e. within
+        # one grid step. An all-zero output fails this everywhere the input is not
+        # already near zero.
+        block_scale = sx.repeat_interleave(16, dim=1).view(torch.float8_e4m3fn).to(torch.float32)
+        block_scale = block_scale[: reference.shape[0], : reference.shape[1]] * (amax / (6.0 * 448))
+        exact = (reference.float() / block_scale).clamp(-6.0, 6.0)
+        magnitude = exact.abs()
+        step = torch.where(magnitude >= 4.0, 2.0, torch.where(magnitude >= 2.0, 1.0, 0.5))
+        gap = (fp4_to_fp32(unpack_fp4(qx)) - exact).abs()
+        assert torch.all(gap <= step * 1.0001), f"largest gap to the exact value is {gap.max()}"
+
+    check_adjacent(q, s, x)
+    check_adjacent(q_t, s_t, x.t().contiguous())
diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
index d5f2fa9a2c..7605798758 100644
--- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
+++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
@@ -269,11 +269,11 @@ __device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x_with_stochastic_ro
         : "f"(in01.y), "f"(in01.x), "f"(in23.y), "f"(in23.x), "r"(rbits));
     return *reinterpret_cast<__nv_fp4x4_e2m1*>(&out_4x);
   } else {
-    NVTE_DEVICE_ERROR(
-        "FP4 cvt.rs PTX instructions are architecture-specific. "
-        "Try recompiling with sm_XXXa instead of sm_XXX.");
-    uint16_t dummy = 0;
-    return *reinterpret_cast<__nv_fp4x4_e2m1*>(&dummy);
+    const float q0 = ptx::stochastic_round_fp4_e2m1(in01.x, rbits);
+    const float q1 = ptx::stochastic_round_fp4_e2m1(in01.y, rbits >> 8);
+    const float q2 = ptx::stochastic_round_fp4_e2m1(in23.x, rbits >> 16);
+    const float q3 = ptx::stochastic_round_fp4_e2m1(in23.y, rbits >> 24);
+    return __nv_fp4x4_e2m1(make_float4(q0, q1, q2, q3));
   }
 }
 
diff --git a/transformer_engine/common/util/ptx.cuh b/transformer_engine/common/util/ptx.cuh
index 83b2daa513..8dbb8b3044 100644
--- a/transformer_engine/common/util/ptx.cuh
+++ b/transformer_engine/common/util/ptx.cuh
@@ -604,6 +604,29 @@ __device__ __forceinline__ void mul_cvt_4x(fp4e2m1x4 &out, const Tx2 &in01, cons
   out = fp4e2m1x4(make_float4(x0, x1, x2, x3));
 }
 
+// Software stochastic rounding onto the e2m1 grid, for architectures without
+// cvt.rs. Takes 8 random bits per element, which is what
+// cvt.rs.satfinite.e2m1x4.f32 takes from its rbits operand. Follows satfinite
+// for the edges: NaN becomes positive MAX_NORM and larger magnitudes clamp to
+// MAX_NORM with the sign kept. The result is exactly representable in e2m1, so
+// packing it is lossless.
+__device__ __forceinline__ float stochastic_round_fp4_e2m1(const float x, const uint32_t rbits8) {
+  constexpr float max_norm = 6.0f;
+  const float u = static_cast<float>(rbits8 & 0xFFu) * (1.0f / 256.0f);
+  const float a = fabsf(x);
+  // Grid step at |x|: 0.5 below 2, 1 in [2, 4), 2 in [4, 6].
+  const float step = (a >= 4.0f) ? 2.0f : ((a >= 2.0f) ? 1.0f : 0.5f);
+  const float t = fmaf(u, step, a);
+  // The jitter can carry t across one region boundary, so the rounding step
+  // derives from t. Flooring in units of that step lands on the e2m1 grid in
+  // every region, and the saturation also maps a NaN t to max_norm, which the
+  // sign select keeps positive.
+  const float step_t = (t >= 4.0f) ? 2.0f : ((t >= 2.0f) ? 1.0f : 0.5f);
+  const float inv_step_t = (t >= 4.0f) ? 0.5f : ((t >= 2.0f) ? 1.0f : 2.0f);
+  const float q = fminf(floorf(t * inv_step_t) * step_t, max_norm);
+  return copysignf(q, (x != x) ? 1.0f : x);
+}
+
 __device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_stochastic_rounding(
     const uint64_t in_4x, const float2 scale, const uint32_t rbits) {
   uint16_t out_4x = 0;
@@ -637,9 +660,14 @@ __device__ __forceinline__ fp4e2m1x4 mul_cvt_bf16_to_fp4_4x_with_stochastic_roun
         : "=h"(out_4x)
         : "l"(in_4x), "l"(reinterpret_cast<const uint64_t &>(scale)), "r"(rbits));
   } else {
-    NVTE_DEVICE_ERROR(
-        "FP4 cvt PTX instructions are architecture-specific. "
-        "Try recompiling with sm_XXXa instead of sm_XXX.");
+    // mul.f32x2 above applies scale.x to the even elements and scale.y to the odd ones.
+    const bf16 *vals = reinterpret_cast<const bf16 *>(&in_4x);
+    const float q0 = stochastic_round_fp4_e2m1(static_cast<float>(vals[0]) * scale.x, rbits);
+    const float q1 = stochastic_round_fp4_e2m1(static_cast<float>(vals[1]) * scale.y, rbits >> 8);
+    const float q2 = stochastic_round_fp4_e2m1(static_cast<float>(vals[2]) * scale.x, rbits >> 16);
+    const float q3 = stochastic_round_fp4_e2m1(static_cast<float>(vals[3]) * scale.y, rbits >> 24);
+    const fp4e2m1x4 packed(make_float4(q0, q1, q2, q3));
+    out_4x = *reinterpret_cast<const uint16_t *>(&packed);
   }
   return *reinterpret_cast<fp4e2m1x4 *>(&out_4x);
 }
@@ -729,9 +757,12 @@ __device__ __forceinline__ fp4e2m1x4 mul_cvt_fp32_to_fp4_4x_with_stochastic_roun
           "l"(reinterpret_cast<const uint64_t &>(in23)),
           "l"(reinterpret_cast<const uint64_t &>(scale)), "r"(rbits));
   } else {
-    NVTE_DEVICE_ERROR(
-        "FP4 cvt PTX instructions are architecture-specific. "
-        "Try recompiling with sm_XXXa instead of sm_XXX.");
+    const float q0 = stochastic_round_fp4_e2m1(in01.x * scale.x, rbits);
+    const float q1 = stochastic_round_fp4_e2m1(in01.y * scale.y, rbits >> 8);
+    const float q2 = stochastic_round_fp4_e2m1(in23.x * scale.x, rbits >> 16);
+    const float q3 = stochastic_round_fp4_e2m1(in23.y * scale.y, rbits >> 24);
+    const fp4e2m1x4 packed(make_float4(q0, q1, q2, q3));
+    out_4x = *reinterpret_cast<const uint16_t *>(&packed);
   }
   return *reinterpret_cast<fp4e2m1x4 *>(&out_4x);
 }
@@ -954,9 +985,26 @@ __device__ __forceinline__ uint32_t mul_cvt_bf16_to_fp4_8x_stochastic_rounding(
       NVTE_DEVICE_ERROR("Not supported scaling coefficient type.");
     }
   } else {
-    NVTE_DEVICE_ERROR(
-        "FP4 cvt PTX instructions are architecture-specific. "
-        "Try recompiling with sm_XXXa instead of sm_XXX.");
+    constexpr bool known_coeff = std::is_same<SCALING_COEFFICIENT_TYPE, bf16>::value ||
+                                 std::is_same<SCALING_COEFFICIENT_TYPE, float>::value;
+    if constexpr (known_coeff) {
+      const float coeff = static_cast<float>(scaling_coefficient);
+      const bf16 *vals03 = reinterpret_cast<const bf16 *>(&in03);
+      const bf16 *vals47 = reinterpret_cast<const bf16 *>(&in47);
+      float q[8];
+#pragma unroll
+      for (int i = 0; i < 4; ++i) {
+        q[i] = stochastic_round_fp4_e2m1(static_cast<float>(vals03[i]) * coeff, rbits03 >> (8 * i));
+        q[i + 4] =
+            stochastic_round_fp4_e2m1(static_cast<float>(vals47[i]) * coeff, rbits47 >> (8 * i));
+      }
+      const fp4e2m1x4 lo(make_float4(q[0], q[1], q[2], q[3]));
+      const fp4e2m1x4 hi(make_float4(q[4], q[5], q[6], q[7]));
+      out_8x = static_cast<uint32_t>(*reinterpret_cast<const uint16_t *>(&lo)) |
+               (static_cast<uint32_t>(*reinterpret_cast<const uint16_t *>(&hi)) << 16);
+    } else {
+      NVTE_DEVICE_ERROR("Not supported scaling coefficient type.");
+    }
   }
   return out_8x;
 }

From 1ad30c686f74de2354d6beb70ea1dfff3371906d Mon Sep 17 00:00:00 2001
From: Fei Wu <33940270+YangFei1990@users.noreply.github.com>
Date: Wed, 12 Aug 2026 09:54:19 -0700
Subject: [PATCH 58/67] [Pytorch] Fix swizzle scale output shape in variable
 shape case (#3344)

* fix swizzle scale output shape in variable shape case

Signed-off-by: YangFei1990 <feiw@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: YangFei1990 <feiw@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .../test_mxfp8_group_quantize_graph_safe.py   | 87 +++++++++++++++++++
 .../pytorch/csrc/extensions/swizzle.cpp       | 20 +++--
 2 files changed, 102 insertions(+), 5 deletions(-)

diff --git a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
index 94861ea433..23662eb0c3 100644
--- a/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
+++ b/tests/pytorch/mxfp8/test_mxfp8_group_quantize_graph_safe.py
@@ -60,6 +60,13 @@ def generate_split_sections(M: int, N: int, edge_cases: str) -> list[int]:
         split_sections = [avg_split] * (num_chunks - 2) + [0] + [avg_split * 2]
     elif edge_cases == "random_uneven_split":
         split_sections = generate_random_multiples_sum(M, num_chunks, least_multiple)
+    elif edge_cases == "imbalanced_avg_misaligned":
+        # Three groups whose 128-aligned sizes average to a non-128-aligned value, so any
+        # buffer sized from num_groups * round_up(M // num_groups, 128) diverges from the
+        # per-group padded sum. "random_uneven_split" cannot catch that: with 4 chunks the
+        # average only depends on M, which is always 128 * num_chunks aligned here.
+        assert M >= 3 * least_multiple, "M too small for the imbalanced case"
+        split_sections = [least_multiple, least_multiple, M - 2 * least_multiple]
     else:
         raise ValueError(f"Invalid edge case: {edge_cases}")
 
@@ -545,6 +552,7 @@ def check_prequantized_requantize_versus_reference(
     wire = make_prequantized_wire_tensor(x, split_section_tensor)
 
     # Snapshot what must survive verbatim, plus the compact scales the reference swizzles.
+    rowwise_scale_shape_before = wire.scale_inv.shape
     rowwise_data_before = wire.rowwise_data.clone()
     wire_splits_before = [
         (t._rowwise_data.view(dtype=torch.uint8).clone(), t._rowwise_scale_inv.clone())
@@ -587,6 +595,13 @@ def check_prequantized_requantize_versus_reference(
     assert wire.columnwise_data is not None, "columnwise data must be built"
     assert wire.columnwise_scale_inv is not None, "columnwise scales must survive the swizzle"
 
+    # The rowwise swizzle is size-preserving; the per-group content checks below cannot catch a
+    # wrongly sized buffer because they read through offsets derived from the splits, so the
+    # capacity must be checked explicitly.
+    assert (
+        wire.scale_inv.shape == rowwise_scale_shape_before
+    ), "the swizzled rowwise scale buffer must keep the compact buffer's capacity"
+
     # The returned dequantized tensor is what bias gradients are reduced from. Compare only the
     # live rows: both this and the reference allocate M rows but write only the covered ones, and
     # their tails are separate uninitialized allocations.
@@ -662,6 +677,7 @@ def check_prequantized_requantize_versus_reference(
         "zero_tokens_end",
         "zero_tokens_middle",
         "random_uneven_split",
+        "imbalanced_avg_misaligned",
     ],
 )
 def test_prequantized_requantize_versus_reference(
@@ -699,6 +715,7 @@ def test_prequantized_requantize_versus_reference(
         "zero_tokens_end",
         "zero_tokens_middle",
         "random_uneven_split",
+        "imbalanced_avg_misaligned",
     ],
 )
 def test_prequantized_requantize_with_paged_stashing(
@@ -805,3 +822,73 @@ def test_prequantized_requantize_rejects_dtype_mismatch():
 
     with pytest.raises(RuntimeError, match="dtype"):
         tex.group_requantize_inplace(wire, mismatched, num_groups, splits, te.DType.kBFloat16)
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize("columnwise", [False, True], ids=["rowwise", "columnwise"])
+@pytest.mark.parametrize(
+    "split_sections",
+    [
+        # Imbalanced groups whose mean is not 128-aligned: sizing the swizzled output from
+        # num_tensors * round_up(mean, 128) would give 3 * 512 = 1536 rows instead of the
+        # input's 1280.
+        [128, 128, 1024],
+        # Same property with a zero-token group in the mix.
+        [0, 256, 128, 1152],
+    ],
+    ids=["imbalanced", "imbalanced_with_empty"],
+)
+def test_grouped_swizzle_variable_shape_preserves_scale_capacity(
+    columnwise: bool, split_sections: list[int]
+):
+    """Swizzling a variable-shape grouped tensor must keep the scale buffer's exact capacity.
+
+    The swizzle kernel walks input and output with identical per-group padded strides, so the
+    operation is size-preserving. Sizing the output from the per-tensor average instead breaks
+    any consumer that derives per-group offsets from the split sizes, because the average of
+    128-aligned group sizes is generally not 128-aligned. The content checks alone cannot catch
+    a wrong allocation (per-group offsets are derived from the splits, not the buffer), so the
+    shape assertion is the actual regression check.
+    """
+    torch.manual_seed(0)
+    N = 256
+    M = sum(split_sections)
+    x = torch.randn((M, N), dtype=torch.bfloat16, device="cuda")
+    splits = torch.tensor(split_sections, dtype=torch.int64, device="cuda")
+
+    quantizer = MXFP8Quantizer(
+        fp8_dtype=te.DType.kFloat8E4M3, rowwise=not columnwise, columnwise=columnwise
+    )
+    # Compact (unswizzled) scales, as they arrive over the wire.
+    quantizer.optimize_for_gemm = False
+    tensor = fused_grouped_quantize(x, splits, quantizer)
+    assert not tensor._with_gemm_swizzled_scales
+
+    scale_attr = "columnwise_scale_inv" if columnwise else "scale_inv"
+    compact_shape = getattr(tensor, scale_attr).shape
+    compact_groups = [
+        (t._columnwise_scale_inv if columnwise else t._rowwise_scale_inv).clone()
+        for t in tensor.split_into_quantized_tensors()
+    ]
+
+    tex.grouped_swizzle_for_gemm(tensor, not columnwise, columnwise)
+
+    assert tensor._with_gemm_swizzled_scales
+    assert (
+        getattr(tensor, scale_attr).shape == compact_shape
+    ), "grouped swizzle must preserve the scale buffer's shape for variable-shape tensors"
+
+    # The swizzled content of each group must match a per-group dense swizzle of the compact
+    # scales it arrived with.
+    for rows_i, group_before, out in zip(
+        split_sections, compact_groups, tensor.split_into_quantized_tensors()
+    ):
+        if rows_i == 0:
+            continue
+        out_scale = out._columnwise_scale_inv if columnwise else out._rowwise_scale_inv
+        torch.testing.assert_close(
+            out_scale,
+            swizzle_mxfp8_scale(rows_i, N, group_before, columnwise=columnwise),
+            atol=0.0,
+            rtol=0.0,
+        )
diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp
index 866d858ee3..59e3a63b29 100644
--- a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp
@@ -415,11 +415,21 @@ std::optional<SwizzledGroupedScales> maybe_swizzle_grouped_tensor(GroupedTensorW
   const size_t per_tensor_last_dim = logical_shape_nvte.data[logical_shape_nvte.ndim - 1];
   constexpr size_t kMxfp8BlockSize = 32;
 
-  // Output is always allocated in the per-tensor padded ("swizzle-ready") layout
-  // so the cuDNN grouped GEMM consumer sees the correct stride between experts.
-  // The swizzle kernel itself handles converting from the kernel-emitted compact
-  // layout (per-tensor first dim is the unpadded value) to this padded layout.
-  auto compute_padded_grouped_scale_shape = [&](bool rowwise) {
+  const bool variable_shape = first_dims.data_ptr != nullptr || last_dims.data_ptr != nullptr;
+
+  // Output is allocated in the layout the swizzle kernel writes so its consumer sees the
+  // correct stride between experts.
+  auto compute_padded_grouped_scale_shape = [&](bool rowwise) -> std::vector<size_t> {
+    if (variable_shape) {
+      // Grouped variable-shape scale storage is a concatenation of per-tensor padded
+      // regions whose sizes live on the device. The swizzle kernel walks input and output
+      // with identical per-group strides, so swizzling is size-preserving and the output
+      // needs exactly the input's shape. The uniform-average formula below would be wrong
+      // here: the average of per-group sizes is generally not tile-aligned, so its
+      // rounded-up total matches neither the kernel's walk nor the input's capacity.
+      const auto &scales = rowwise ? row_scales : col_scales;
+      return nvte_shape_to_vector(scales.shape);
+    }
     const size_t m = rowwise ? per_tensor_first_dim : per_tensor_last_dim;
     const size_t k = rowwise ? per_tensor_last_dim : per_tensor_first_dim;
     const size_t padded_m = ceildiv(m, size_t{128}) * 128;

From 779f17d37b615380348d696f6af36d4a1d1a6208 Mon Sep 17 00:00:00 2001
From: Phuong Nguyen <phuonguyen@nvidia.com>
Date: Wed, 12 Aug 2026 10:43:49 -0700
Subject: [PATCH 59/67] [Common] Upgrade nccl-extensions commit (#3349)

bump nccl-extensions commit

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---
 3rdparty/nccl-extensions | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/3rdparty/nccl-extensions b/3rdparty/nccl-extensions
index 705ca8eb38..9f47d6eb3b 160000
--- a/3rdparty/nccl-extensions
+++ b/3rdparty/nccl-extensions
@@ -1 +1 @@
-Subproject commit 705ca8eb38297f3a8c4af6adf4185176a1116cb9
+Subproject commit 9f47d6eb3b60962d8157a579b4caaaa4ae6b19f4

From 2d80391f77e542c06d0281688ddd17b6e34adec8 Mon Sep 17 00:00:00 2001
From: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com>
Date: Wed, 12 Aug 2026 11:54:14 -0700
Subject: [PATCH 60/67] Improve device-init grouped linear module with single
 grouped weight support  (#3224)

* improve device grouped linear

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* proper support for grouped bias

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* fixes and improvements

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

# Conflicts:
#	tests/pytorch/test_grouped_mlp.py

* fix unit test failure

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix group linear UT

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix group mlp UT by disabling wrong fallback

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* lint

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix CI errors

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* rename toggle resolve comments

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* don't always set optimize_for_gemm=True

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* test other fp8 recipes on hopper

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix test failure

Signed-off-by: zhongboz <zhongboz@nvidia.com>

* relax atol for bf16 sum errors

Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix CI

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* continue to fix edge cases

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* unify grouped tensor check

Signed-off-by: zhongboz <zhongboz@nvidia.com>

* special case handling for all empty moe inputs for fp8 CS

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* add another specific case guard

Signed-off-by: zhongboz <zhongboz@nvidia.com>

* fix cublas setup meta data issue for empty inputs

Signed-off-by: zhongboz <zhongboz@nvidia.com>

* mcore integration fixes

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* chore fix after rebase

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* chore

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* chore fix after rebase

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* resolve comments, enable save original input

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* fix cudnn version gurad

Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>

* comments

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* update benchmark script

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
Signed-off-by: zhongboz <zhongboz@nvidia.com>
Signed-off-by: Varun Thumbe <vthumbe@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Varun Thumbe <vthumbe@nvidia.com>
---
 ...py => benchmark_graph_safe_grouped_mlp.py} |   8 +-
 .../linear/benchmark_grouped_gemm_kernels.py  | 581 +++++++++++
 ...mark_grouped_linear_cublas_grouped_gemm.py | 520 ++++++++++
 tests/pytorch/test_grouped_linear.py          | 973 ++++++++++++++++--
 tests/pytorch/test_grouped_mlp.py             |  56 +-
 tests/pytorch/test_grouped_tensor.py          | 291 +++++-
 tests/pytorch/test_sanity.py                  |  58 +-
 .../cast/mxfp8/group_quantize_mxfp8.cuh       |  19 +-
 .../common/gemm/cublaslt_grouped_gemm.cu      | 102 +-
 transformer_engine/pytorch/csrc/extensions.h  |   2 +-
 .../pytorch/csrc/extensions/cast.cpp          |  40 +-
 .../pytorch/csrc/extensions/pybind.cpp        |   2 +-
 .../pytorch/csrc/type_converters.cpp          |  15 +-
 transformer_engine/pytorch/module/__init__.py |   2 +-
 .../pytorch/module/grouped_linear.py          | 676 ++++++++----
 .../pytorch/ops/basic/__init__.py             |   2 +-
 .../pytorch/ops/basic/grouped_linear.py       | 195 ++--
 .../pytorch/ops/fused/grouped_mlp.py          |   3 +-
 18 files changed, 3051 insertions(+), 494 deletions(-)
 rename benchmarks/linear/{benchmark_graph_safe_grouped_linear.py => benchmark_graph_safe_grouped_mlp.py} (97%)
 create mode 100644 benchmarks/linear/benchmark_grouped_gemm_kernels.py
 create mode 100644 benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py

diff --git a/benchmarks/linear/benchmark_graph_safe_grouped_linear.py b/benchmarks/linear/benchmark_graph_safe_grouped_mlp.py
similarity index 97%
rename from benchmarks/linear/benchmark_graph_safe_grouped_linear.py
rename to benchmarks/linear/benchmark_graph_safe_grouped_mlp.py
index d8230c38fe..00f7f516c8 100644
--- a/benchmarks/linear/benchmark_graph_safe_grouped_linear.py
+++ b/benchmarks/linear/benchmark_graph_safe_grouped_mlp.py
@@ -13,21 +13,21 @@
 
 Example:
 
-    python benchmarks/linear/benchmark_graph_safe_grouped_linear.py
+    python benchmarks/linear/benchmark_graph_safe_grouped_mlp.py
 
 Forward-only:
 
-    python benchmarks/linear/benchmark_graph_safe_grouped_linear.py --fwd-only
+    python benchmarks/linear/benchmark_graph_safe_grouped_mlp.py --fwd-only
 
 Nsight Systems:
 
     (optionally: unset DEBUGINFOD_URLS)
 
     nsys profile \
-        --output=./benchmarks/linear/graph_safe_grouped_linear_mxfp8 \
+        --output=./benchmarks/linear/graph_safe_grouped_mlp_mxfp8 \
         --force-overwrite true \
         --trace=cuda,nvtx,cudnn,cublas \
-        python benchmarks/linear/benchmark_graph_safe_grouped_linear.py --profile
+        python benchmarks/linear/benchmark_graph_safe_grouped_mlp.py --profile
 """
 
 # Match the Qwen MXFP8 SFT launch toggles before importing TE.
diff --git a/benchmarks/linear/benchmark_grouped_gemm_kernels.py b/benchmarks/linear/benchmark_grouped_gemm_kernels.py
new file mode 100644
index 0000000000..88b3f8a466
--- /dev/null
+++ b/benchmarks/linear/benchmark_grouped_gemm_kernels.py
@@ -0,0 +1,581 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Benchmark multi-stream and cuBLASLt grouped GEMM kernels.
+
+Unlike ``benchmark_grouped_linear_cublas_grouped_gemm.py``, this benchmark does
+not construct a GroupedLinear module or invoke autograd. Inputs are allocated
+and, for MXFP8, quantized and GEMM-swizzled before timing. The timed region only
+calls TE's low-level grouped GEMM wrappers.
+
+The default shape is Qwen3.5-397B-A17B with sequence length 4096, top-10
+routing, and EP32. Each rank owns 16 experts and processes 2560 rows per expert.
+All six expert GEMMs are measured independently:
+
+============== ====== ====== ====== ====== ==================================
+Projection     Pass   Layout M      K      N
+============== ====== ====== ====== ====== ==================================
+FC1            fwd    TN     2560   4096   2048
+FC1            dgrad  NN     2560   4096   2048
+FC1            wgrad  NT     2560   4096   2048
+FC2            fwd    TN     2560   1024   4096
+FC2            dgrad  NN     2560   1024   4096
+FC2            wgrad  NT     2560   1024   4096
+============== ====== ====== ====== ====== ==================================
+
+The layouts describe TE's operands for each expert:
+
+* Forward (TN): A=``weight[N,K]``, B=``input[M,K]``;
+  ``input @ weight.T -> output[M,N]``.
+* Dgrad (NN): A=``weight[N,K]``, B=``dy[M,N]``;
+  ``dy @ weight -> dx[M,K]``.
+* Wgrad (NT): A=``input[M,K]``, B=``dy[M,N]``;
+  ``dy.T @ input -> dweight[N,K]``.
+
+The multi-stream path receives lists of discrete tensors. The cuBLASLt grouped
+path receives packed GroupedTensor activations and discrete weights/wgrads,
+matching GroupedLinear with discrete parameters.
+
+Examples
+--------
+Run the full BF16 and MXFP8 matrix:
+
+    python benchmarks/linear/benchmark_grouped_gemm_kernels.py
+
+Run only MXFP8 FC1 wgrad:
+
+    python benchmarks/linear/benchmark_grouped_gemm_kernels.py \
+        --precision mxfp8 --projection fc1 --gemm wgrad
+"""
+
+import argparse
+import os
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import pandas as pd
+import torch
+import torch.utils.benchmark as benchmark
+
+from transformer_engine.common.recipe import MXFP8BlockScaling, Recipe
+from transformer_engine.pytorch import MXFP8Quantizer
+from transformer_engine.pytorch.cpp_extensions import (
+    general_grouped_gemm,
+    general_grouped_gemm_for_grouped_tensor,
+)
+from transformer_engine.pytorch.module import is_module_grouped_tensor_path_supported
+from transformer_engine.pytorch.module.base import (
+    _2X_ACC_DGRAD,
+    _2X_ACC_FPROP,
+    _2X_ACC_WGRAD,
+)
+from transformer_engine.pytorch.quantization import FP8GlobalStateManager
+from transformer_engine.pytorch.tensor import GroupedTensor, GroupedTensorStorage
+import transformer_engine_torch as tex
+
+
+QWEN_NUM_EXPERTS = 512
+QWEN_TOP_K = 10
+QWEN_SEQUENCE_LENGTH = 4096
+QWEN_EXPERT_PARALLEL_SIZE = 32
+QWEN_HIDDEN_SIZE = 4096
+QWEN_MOE_INTERMEDIATE_SIZE = 1024
+
+
+@dataclass(frozen=True)
+class GemmSpec:
+    """One expert GEMM in the Qwen routed MLP."""
+
+    projection: str
+    gemm: str
+    layout: str
+    k: int
+    n: int
+
+
+GEMM_SPECS = tuple(
+    GemmSpec(projection, gemm, layout, k, n)
+    for projection, k, n in (
+        ("fc1", QWEN_HIDDEN_SIZE, 2 * QWEN_MOE_INTERMEDIATE_SIZE),
+        ("fc2", QWEN_MOE_INTERMEDIATE_SIZE, QWEN_HIDDEN_SIZE),
+    )
+    for gemm, layout in (("fwd", "TN"), ("dgrad", "NN"), ("wgrad", "NT"))
+)
+
+
+@dataclass
+class PreparedGemm:
+    """Preallocated operands and outputs for both grouped GEMM paths."""
+
+    spec: GemmSpec
+    split_sizes: list[int]
+    multistream_a: list[Any]
+    multistream_b: list[Any]
+    multistream_out: list[torch.Tensor]
+    grouped_a: Any
+    grouped_b: GroupedTensorStorage
+    grouped_out: Any
+
+
+def _make_grouped_bf16(
+    packed: Optional[torch.Tensor],
+    split_sizes: list[int],
+    last_dim: int,
+) -> GroupedTensorStorage:
+    """Create packed GroupedTensor storage and optionally initialize its data."""
+    first_dims = torch.tensor(split_sizes, dtype=torch.int64, device="cuda")
+    grouped = GroupedTensor.make_grouped_tensor(
+        num_tensors=len(split_sizes),
+        first_dims=first_dims,
+        last_dims=None,
+        logical_first_dim=sum(split_sizes),
+        logical_last_dim=last_dim,
+        quantizer=None,
+        device=torch.device("cuda"),
+        dtype=torch.bfloat16,
+    )
+    if packed is not None:
+        grouped.rowwise_data.view(-1).copy_(packed.reshape(-1))
+    return grouped
+
+
+def _new_mxfp8_quantizer(*, rowwise: bool, columnwise: bool) -> MXFP8Quantizer:
+    """Construct a quantizer that emits GEMM-ready scales before timing."""
+    quantizer = MXFP8Quantizer(
+        fp8_dtype=tex.DType.kFloat8E4M3,
+        rowwise=rowwise,
+        columnwise=columnwise,
+    )
+    quantizer.optimize_for_gemm = True
+    return quantizer
+
+
+def _quantize_discrete_mxfp8(
+    tensors: list[torch.Tensor],
+    *,
+    rowwise: bool,
+    columnwise: bool,
+) -> list[Any]:
+    """Quantize per-expert tensors outside the timed region."""
+    quantizer = _new_mxfp8_quantizer(rowwise=rowwise, columnwise=columnwise)
+    return [quantizer(tensor) for tensor in tensors]
+
+
+def _quantize_grouped_mxfp8(
+    packed: torch.Tensor,
+    split_sizes: list[int],
+    *,
+    rowwise: bool,
+    columnwise: bool,
+) -> GroupedTensorStorage:
+    """Quantize one packed activation into GEMM-ready grouped storage."""
+    quantizer = _new_mxfp8_quantizer(rowwise=rowwise, columnwise=columnwise)
+    first_dims = torch.tensor(split_sizes, dtype=torch.int64, device="cuda")
+    return tex.group_quantize(
+        packed,
+        quantizer,
+        len(split_sizes),
+        first_dims,
+    )
+
+
+def _make_high_precision_operands(
+    spec: GemmSpec,
+    split_sizes: list[int],
+) -> tuple[list[torch.Tensor], Optional[torch.Tensor], torch.Tensor]:
+    """Create canonical BF16 operands shared by both precision paths."""
+    num_experts = len(split_sizes)
+    total_rows = sum(split_sizes)
+
+    if spec.gemm in ("fwd", "dgrad"):
+        weights = [
+            torch.randn(spec.n, spec.k, dtype=torch.bfloat16, device="cuda")
+            for _ in range(num_experts)
+        ]
+        b_width = spec.k if spec.gemm == "fwd" else spec.n
+        packed_b = torch.randn(total_rows, b_width, dtype=torch.bfloat16, device="cuda")
+        return weights, None, packed_b
+
+    packed_x = torch.randn(total_rows, spec.k, dtype=torch.bfloat16, device="cuda")
+    packed_dy = torch.randn(total_rows, spec.n, dtype=torch.bfloat16, device="cuda")
+    return [], packed_x, packed_dy
+
+
+def _operand_usage(spec: GemmSpec) -> tuple[tuple[bool, bool], tuple[bool, bool]]:
+    """Return rowwise/columnwise MXFP8 storage used by A and B."""
+    if spec.gemm == "fwd":
+        return (True, False), (True, False)
+    if spec.gemm == "dgrad":
+        return (False, True), (True, False)
+    return (False, True), (False, True)
+
+
+def _allocate_outputs(
+    spec: GemmSpec,
+    split_sizes: list[int],
+) -> tuple[list[torch.Tensor], Any]:
+    """Allocate output storage matching each low-level API's native contract."""
+    num_experts = len(split_sizes)
+    if spec.gemm == "wgrad":
+        multistream_out = [
+            torch.empty(spec.n, spec.k, dtype=torch.bfloat16, device="cuda")
+            for _ in range(num_experts)
+        ]
+        grouped_out = [
+            torch.empty(spec.n, spec.k, dtype=torch.bfloat16, device="cuda")
+            for _ in range(num_experts)
+        ]
+        return multistream_out, grouped_out
+
+    out_features = spec.n if spec.gemm == "fwd" else spec.k
+    multistream_out = [
+        torch.empty(sum(split_sizes), out_features, dtype=torch.bfloat16, device="cuda")
+    ]
+    grouped_out = _make_grouped_bf16(None, split_sizes, out_features)
+    return multistream_out, grouped_out
+
+
+def _prepare_gemm(
+    spec: GemmSpec,
+    split_sizes: list[int],
+    precision: str,
+) -> PreparedGemm:
+    """Prepare one GEMM without leaving quantization inside the timed region."""
+    hp_a, packed_a, packed_b = _make_high_precision_operands(spec, split_sizes)
+    a_usage, b_usage = _operand_usage(spec)
+
+    if spec.gemm in ("fwd", "dgrad"):
+        if precision == "mxfp8":
+            multistream_a = _quantize_discrete_mxfp8(
+                hp_a,
+                rowwise=a_usage[0],
+                columnwise=a_usage[1],
+            )
+        else:
+            multistream_a = hp_a
+        grouped_a = multistream_a
+    elif precision == "mxfp8":
+        assert packed_a is not None
+        grouped_a = _quantize_grouped_mxfp8(
+            packed_a,
+            split_sizes,
+            rowwise=a_usage[0],
+            columnwise=a_usage[1],
+        )
+    else:
+        assert packed_a is not None
+        grouped_a = _make_grouped_bf16(packed_a, split_sizes, spec.k)
+    if spec.gemm == "wgrad":
+        multistream_a = list(grouped_a.split_into_quantized_tensors())
+
+    if precision == "mxfp8":
+        grouped_b = _quantize_grouped_mxfp8(
+            packed_b,
+            split_sizes,
+            rowwise=b_usage[0],
+            columnwise=b_usage[1],
+        )
+    else:
+        b_width = spec.k if spec.gemm == "fwd" else spec.n
+        grouped_b = _make_grouped_bf16(packed_b, split_sizes, b_width)
+    # Give multi-stream zero-copy member views into the exact same packed operand used by
+    # cuBLASLt. This prevents quantization or input-data differences from entering the result.
+    multistream_b = list(grouped_b.split_into_quantized_tensors())
+
+    multistream_out, grouped_out = _allocate_outputs(spec, split_sizes)
+    prepared = PreparedGemm(
+        spec=spec,
+        split_sizes=split_sizes,
+        multistream_a=multistream_a,
+        multistream_b=multistream_b,
+        multistream_out=multistream_out,
+        grouped_a=grouped_a,
+        grouped_b=grouped_b,
+        grouped_out=grouped_out,
+    )
+    if precision == "mxfp8":
+        _assert_mxfp8_operands_are_gemm_ready(prepared)
+    return prepared
+
+
+def _iter_operands(operand: Any):
+    """Yield discrete members or one packed operand."""
+    if isinstance(operand, (list, tuple)):
+        yield from operand
+    else:
+        yield operand
+
+
+def _assert_mxfp8_operands_are_gemm_ready(prepared: PreparedGemm) -> None:
+    """Ensure scale swizzling cannot leak into the GEMM timing."""
+    operands = (
+        prepared.multistream_a,
+        prepared.multistream_b,
+        prepared.grouped_a,
+        prepared.grouped_b,
+    )
+    for operand_group in operands:
+        for operand in _iter_operands(operand_group):
+            assert getattr(
+                operand, "_with_gemm_swizzled_scales", False
+            ), "MXFP8 operand was not prepared with GEMM-swizzled scales"
+
+
+def _use_split_accumulator(spec: GemmSpec) -> bool:
+    """Match GroupedLinear's accumulator policy for each training GEMM."""
+    if spec.gemm == "fwd":
+        return _2X_ACC_FPROP
+    if spec.gemm == "dgrad":
+        return _2X_ACC_DGRAD
+    return _2X_ACC_WGRAD
+
+
+def _run_multistream(prepared: PreparedGemm, iterations: int) -> None:
+    """Launch the multi-stream grouped GEMM repeatedly."""
+    spec = prepared.spec
+    for _ in range(iterations):
+        general_grouped_gemm(
+            prepared.multistream_a,
+            prepared.multistream_b,
+            prepared.multistream_out,
+            [None] * len(prepared.split_sizes),
+            torch.bfloat16,
+            layout=spec.layout,
+            m_splits=prepared.split_sizes,
+            grad=spec.gemm != "fwd",
+            single_output=spec.gemm != "wgrad",
+            use_split_accumulator=_use_split_accumulator(spec),
+        )
+
+
+def _run_cublas_grouped(prepared: PreparedGemm, iterations: int) -> None:
+    """Launch the device-described cuBLASLt grouped GEMM repeatedly."""
+    for _ in range(iterations):
+        general_grouped_gemm_for_grouped_tensor(
+            prepared.grouped_a,
+            prepared.grouped_b,
+            prepared.grouped_out,
+            layout=prepared.spec.layout,
+            use_split_accumulator=_use_split_accumulator(prepared.spec),
+        )
+
+
+def _canonical_output(prepared: PreparedGemm, *, grouped: bool) -> torch.Tensor:
+    """Return either path's output in one comparable packed tensor."""
+    output = prepared.grouped_out if grouped else prepared.multistream_out
+    if prepared.spec.gemm == "wgrad":
+        return torch.stack(output, dim=0)
+    if grouped:
+        out_features = prepared.spec.n if prepared.spec.gemm == "fwd" else prepared.spec.k
+        return output.rowwise_data.view(sum(prepared.split_sizes), out_features)
+    return output[0]
+
+
+def _validate_outputs(prepared: PreparedGemm, precision: str) -> None:
+    """Check that execution-path selection does not change GEMM numerics unexpectedly."""
+    _run_multistream(prepared, 1)
+    _run_cublas_grouped(prepared, 1)
+    if precision == "mxfp8":
+        tolerances = {"rtol": 0.125, "atol": 0.0675}
+    else:
+        tolerances = {"rtol": 1e-2, "atol": 1e-2}
+    torch.testing.assert_close(
+        _canonical_output(prepared, grouped=True),
+        _canonical_output(prepared, grouped=False),
+        **tolerances,
+    )
+
+
+def _benchmark_path(
+    prepared: PreparedGemm,
+    *,
+    path: str,
+    iterations_per_run: int,
+    warmup_iterations: int,
+    min_run_time: float,
+    profile: bool,
+    label: str,
+) -> float:
+    """Return milliseconds for one grouped GEMM launch."""
+    run = _run_multistream if path == "multistream" else _run_cublas_grouped
+    run(prepared, warmup_iterations)
+    torch.cuda.synchronize()
+
+    if profile:
+        torch.cuda.nvtx.range_push(label)
+    timing = benchmark.Timer(
+        stmt="run(prepared, iterations_per_run)",
+        globals={
+            "run": run,
+            "prepared": prepared,
+            "iterations_per_run": iterations_per_run,
+        },
+        num_threads=1,
+    ).blocked_autorange(min_run_time=min_run_time)
+    if profile:
+        torch.cuda.nvtx.range_pop()
+    return timing.median * 1000 / iterations_per_run
+
+
+def _gemm_tflops(spec: GemmSpec, total_rows: int, time_ms: float) -> float:
+    """Compute aggregate GEMM throughput across all local experts."""
+    flops = 2 * total_rows * spec.k * spec.n
+    return flops / time_ms / 1e9
+
+
+def _shape_description(spec: GemmSpec, rows_per_expert: str) -> str:
+    """Format the actual per-expert operands passed to TE."""
+    if spec.gemm == "fwd":
+        return (
+            f"A[{spec.n},{spec.k}] B[{rows_per_expert},{spec.k}] -> D[{rows_per_expert},{spec.n}]"
+        )
+    if spec.gemm == "dgrad":
+        return (
+            f"A[{spec.n},{spec.k}] B[{rows_per_expert},{spec.n}] -> D[{rows_per_expert},{spec.k}]"
+        )
+    return f"A[{rows_per_expert},{spec.k}] B[{rows_per_expert},{spec.n}] -> D[{spec.n},{spec.k}]"
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--precision", choices=("all", "bf16", "mxfp8"), default="all")
+    parser.add_argument("--projection", choices=("all", "fc1", "fc2"), default="all")
+    parser.add_argument("--gemm", choices=("all", "fwd", "dgrad", "wgrad"), default="all")
+    parser.add_argument("--iterations-per-run", type=int, default=1000)
+    parser.add_argument("--warmup-iterations", type=int, default=500)
+    parser.add_argument("--min-run-time", type=float, default=5.0)
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument(
+        "--m-splits",
+        type=str,
+        default=None,
+        help="Optional comma-separated per-expert rows; defaults to Qwen3.5-397B EP32.",
+    )
+    parser.add_argument("--skip-correctness", action="store_true")
+    parser.add_argument("--profile", action="store_true", help="Add per-path NVTX ranges.")
+    parser.add_argument("--output-csv", type=str, default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for this benchmark.")
+    if os.getenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "0") == "1":
+        raise RuntimeError(
+            "Unset NVTE_USE_CUTLASS_GROUPED_GEMM: this benchmark requires the "
+            "multi-stream cuBLAS baseline."
+        )
+    if args.iterations_per_run < 1 or args.warmup_iterations < 1:
+        raise ValueError("iterations-per-run and warmup-iterations must be positive.")
+
+    if args.m_splits is None:
+        num_local_experts = QWEN_NUM_EXPERTS // QWEN_EXPERT_PARALLEL_SIZE
+        rows_per_expert = QWEN_SEQUENCE_LENGTH * QWEN_TOP_K // num_local_experts
+        split_sizes = [rows_per_expert] * num_local_experts
+    else:
+        split_sizes = [int(value) for value in args.m_splits.split(",") if value]
+        if not split_sizes or any(value <= 0 for value in split_sizes):
+            raise ValueError("m_splits must contain positive integers.")
+        num_local_experts = len(split_sizes)
+
+    total_rows = sum(split_sizes)
+    precisions = ("bf16", "mxfp8") if args.precision == "all" else (args.precision,)
+    specs = [
+        spec
+        for spec in GEMM_SPECS
+        if (args.projection == "all" or spec.projection == args.projection)
+        and (args.gemm == "all" or spec.gemm == args.gemm)
+    ]
+    recipes: dict[str, Optional[Recipe]] = {
+        "bf16": None,
+        "mxfp8": MXFP8BlockScaling(),
+    }
+    mxfp8_available, reason_for_no_mxfp8 = FP8GlobalStateManager.is_mxfp8_available()
+
+    uniform_rows = str(split_sizes[0]) if len(set(split_sizes)) == 1 else "variable"
+    print("Qwen3.5-397B-A17B grouped GEMM kernel benchmark")
+    print(f"  local experts: {num_local_experts}")
+    print(f"  total rows: {total_rows}")
+    print(f"  m_splits: {split_sizes}")
+    print("  quantization is outside the timed region")
+    print()
+    for spec in specs:
+        print(
+            f"  {spec.projection} {spec.gemm:5s} {spec.layout}: "
+            f"{_shape_description(spec, uniform_rows)}"
+        )
+    print()
+
+    rows = []
+    for precision in precisions:
+        recipe = recipes[precision]
+        if precision == "mxfp8" and not mxfp8_available:
+            print(f"Skipping MXFP8: {reason_for_no_mxfp8}")
+            continue
+        if not is_module_grouped_tensor_path_supported(recipe, torch.bfloat16):
+            print(
+                f"Skipping {precision}: cuBLASLt grouped GEMM is unsupported on this "
+                "GPU or cuBLASLt version."
+            )
+            continue
+
+        for spec in specs:
+            torch.manual_seed(args.seed)
+            torch.cuda.manual_seed(args.seed)
+            prepared = _prepare_gemm(spec, split_sizes, precision)
+            if not args.skip_correctness:
+                _validate_outputs(prepared, precision)
+
+            timings = {}
+            for path in ("multistream", "cublas_grouped_gemm"):
+                label = f"{precision}_{spec.projection}_{spec.gemm}_{path}"
+                timings[path] = _benchmark_path(
+                    prepared,
+                    path=path,
+                    iterations_per_run=args.iterations_per_run,
+                    warmup_iterations=args.warmup_iterations,
+                    min_run_time=args.min_run_time,
+                    profile=args.profile,
+                    label=label,
+                )
+
+            speedup = timings["multistream"] / timings["cublas_grouped_gemm"]
+            for path, time_ms in timings.items():
+                rows.append(
+                    {
+                        "precision": precision,
+                        "projection": spec.projection,
+                        "gemm": spec.gemm,
+                        "layout": spec.layout,
+                        "execution_path": path,
+                        "num_local_experts": num_local_experts,
+                        "total_rows": total_rows,
+                        "k": spec.k,
+                        "n": spec.n,
+                        "time_ms": time_ms,
+                        "tflops": _gemm_tflops(spec, total_rows, time_ms),
+                        "speedup_vs_multistream": (
+                            speedup if path == "cublas_grouped_gemm" else 1.0
+                        ),
+                    }
+                )
+
+            print(
+                f"{precision:6s} {spec.projection} {spec.gemm:5s}: "
+                f"multistream={timings['multistream']:.3f} ms, "
+                f"cuBLAS grouped={timings['cublas_grouped_gemm']:.3f} ms, "
+                f"speedup={speedup:.3f}x"
+            )
+
+    results = pd.DataFrame(rows)
+    print()
+    print(results.to_string(index=False))
+    if args.output_csv is not None:
+        results.to_csv(args.output_csv, index=False)
+        print(f"\nWrote {args.output_csv}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py b/benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py
new file mode 100644
index 0000000000..1b400c0c24
--- /dev/null
+++ b/benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py
@@ -0,0 +1,520 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Compare GroupedLinear's multi-stream and cuBLASLt grouped GEMM paths.
+
+The default problem is the routed-expert shape from Qwen3.5-397B-A17B with
+sequence length 4096, top-10 routing, and expert parallelism 32:
+
+    512 global experts / EP32 = 16 local experts
+    4096 tokens * top-10 / 16 local experts = 2560 rows per expert
+
+The 2560-row split is already 256-aligned, so both paths execute identical
+GEMM work. FC1 and FC2 are benchmarked independently:
+
+    FC1: 16 x (M=2560, K=4096, N=2048)
+    FC2: 16 x (M=2560, K=1024, N=4096)
+
+Both paths use discrete parameters initialized with identical values:
+
+* ``use_grouped_tensor=False`` receives CPU splits and launches the
+  multi-stream grouped GEMM implementation.
+* ``use_grouped_tensor=True`` receives CUDA int64 splits and launches the
+  cuBLASLt grouped GEMM implementation.
+
+The MXFP8 cases use BF16 primary parameters. Each timing bundle refreshes the
+MXFP8 weight cache on its first microbatch and reuses it on later microbatches.
+
+Examples
+--------
+Run the complete BF16 and MXFP8 comparison:
+
+    python benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py
+
+Run only MXFP8 FC1 forward and backward:
+
+    python benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py \
+        --precision mxfp8 --projection fc1 --mode fwd_bwd
+
+Profile one microbatch per timing invocation:
+
+    nsys profile \
+        --output=grouped_linear_cublas_grouped_gemm \
+        --force-overwrite=true \
+        --trace=cuda,nvtx,cublas \
+        python benchmarks/linear/benchmark_grouped_linear_cublas_grouped_gemm.py \
+            --precision mxfp8 --projection fc1 --mode fwd_bwd \
+            --num-microbatches 1 --profile
+"""
+
+import argparse
+from contextlib import nullcontext
+from dataclasses import dataclass
+from typing import Optional
+
+import pandas as pd
+import torch
+import torch.utils.benchmark as benchmark
+
+import transformer_engine.pytorch as te
+from transformer_engine.common.recipe import MXFP8BlockScaling, Recipe
+from transformer_engine.pytorch.module import (
+    GroupedLinear,
+    is_module_grouped_tensor_path_supported,
+)
+from transformer_engine.pytorch.quantization import FP8GlobalStateManager
+
+
+QWEN_NUM_EXPERTS = 512
+QWEN_TOP_K = 10
+QWEN_SEQUENCE_LENGTH = 4096
+QWEN_EXPERT_PARALLEL_SIZE = 32
+QWEN_HIDDEN_SIZE = 4096
+QWEN_MOE_INTERMEDIATE_SIZE = 1024
+
+
+@dataclass(frozen=True)
+class Projection:
+    """GroupedLinear dimensions for one Qwen routed-expert projection."""
+
+    name: str
+    in_features: int
+    out_features: int
+
+
+PROJECTIONS = {
+    "fc1": Projection(
+        name="fc1",
+        in_features=QWEN_HIDDEN_SIZE,
+        out_features=2 * QWEN_MOE_INTERMEDIATE_SIZE,
+    ),
+    "fc2": Projection(
+        name="fc2",
+        in_features=QWEN_MOE_INTERMEDIATE_SIZE,
+        out_features=QWEN_HIDDEN_SIZE,
+    ),
+}
+
+
+@dataclass(frozen=True)
+class ExecutionPath:
+    """GroupedLinear path and the corresponding m_splits representation."""
+
+    name: str
+    use_grouped_tensor: bool
+
+
+EXECUTION_PATHS = (
+    ExecutionPath(name="multistream", use_grouped_tensor=False),
+    ExecutionPath(name="cublas_grouped_gemm", use_grouped_tensor=True),
+)
+
+
+def _quantization_context(recipe: Optional[Recipe]):
+    """Construct a fresh quantization context for one invocation."""
+    if recipe is None:
+        return nullcontext()
+    return te.autocast(enabled=True, recipe=recipe)
+
+
+def _make_m_splits(
+    *,
+    use_grouped_tensor: bool,
+    split_sizes: list[int],
+) -> torch.Tensor:
+    """Use the split representation consumed natively by each execution path."""
+    device = "cuda" if use_grouped_tensor else "cpu"
+    return torch.tensor(split_sizes, dtype=torch.int64, device=device)
+
+
+def _build_layer(
+    *,
+    projection: Projection,
+    num_local_experts: int,
+    use_grouped_tensor: bool,
+) -> GroupedLinear:
+    """Construct a discrete-parameter GroupedLinear module."""
+    return GroupedLinear(
+        num_local_experts,
+        projection.in_features,
+        projection.out_features,
+        bias=False,
+        params_dtype=torch.bfloat16,
+        device="cuda",
+        use_grouped_tensor=use_grouped_tensor,
+    )
+
+
+def _run_microbatches(
+    layer: GroupedLinear,
+    x: torch.Tensor,
+    m_splits: torch.Tensor,
+    grad_output: torch.Tensor,
+    *,
+    recipe: Optional[Recipe],
+    mode: str,
+    num_microbatches: int,
+) -> torch.Tensor:
+    """Run one timed bundle while preserving realistic MXFP8 weight caching."""
+    layer.zero_grad(set_to_none=True)
+    x.grad = None
+
+    if mode == "fwd":
+        with torch.no_grad(), _quantization_context(recipe):
+            for microbatch in range(num_microbatches):
+                output = layer(
+                    x,
+                    m_splits,
+                    is_first_microbatch=(microbatch == 0),
+                )
+        return output
+
+    with _quantization_context(recipe):
+        for microbatch in range(num_microbatches):
+            output = layer(
+                x,
+                m_splits,
+                is_first_microbatch=(microbatch == 0),
+            )
+            output.backward(grad_output)
+    return output
+
+
+def _run_correctness_step(
+    layer: GroupedLinear,
+    x: torch.Tensor,
+    m_splits: torch.Tensor,
+    grad_output: torch.Tensor,
+    *,
+    recipe: Optional[Recipe],
+    mode: str,
+) -> tuple[torch.Tensor, Optional[torch.Tensor], list[Optional[torch.Tensor]]]:
+    """Run one microbatch and retain outputs and gradients for path parity."""
+    output = _run_microbatches(
+        layer,
+        x,
+        m_splits,
+        grad_output,
+        recipe=recipe,
+        mode=mode,
+        num_microbatches=1,
+    )
+    input_grad = None if x.grad is None else x.grad.detach().clone()
+    parameter_grads = [
+        None if param.grad is None else param.grad.detach().clone() for param in layer.parameters()
+    ]
+    return output.detach().clone(), input_grad, parameter_grads
+
+
+def _validate_path_parity(
+    *,
+    multistream_layer: GroupedLinear,
+    grouped_layer: GroupedLinear,
+    multistream_x: torch.Tensor,
+    grouped_x: torch.Tensor,
+    multistream_splits: torch.Tensor,
+    grouped_splits: torch.Tensor,
+    grad_output: torch.Tensor,
+    recipe: Optional[Recipe],
+    mode: str,
+) -> None:
+    """Require the two execution paths to produce compatible results."""
+    reference = _run_correctness_step(
+        multistream_layer,
+        multistream_x,
+        multistream_splits,
+        grad_output,
+        recipe=recipe,
+        mode=mode,
+    )
+    actual = _run_correctness_step(
+        grouped_layer,
+        grouped_x,
+        grouped_splits,
+        grad_output,
+        recipe=recipe,
+        mode=mode,
+    )
+
+    tolerances = {"rtol": 1e-2, "atol": 1e-2}
+
+    torch.testing.assert_close(actual[0], reference[0], **tolerances)
+    if mode == "fwd":
+        return
+
+    assert actual[1] is not None and reference[1] is not None
+    torch.testing.assert_close(actual[1], reference[1], **tolerances)
+    assert len(actual[2]) == len(reference[2])
+    for actual_grad, reference_grad in zip(actual[2], reference[2]):
+        assert actual_grad is not None and reference_grad is not None
+        torch.testing.assert_close(actual_grad, reference_grad, **tolerances)
+
+
+def _benchmark_path(
+    *,
+    layer: GroupedLinear,
+    x: torch.Tensor,
+    m_splits: torch.Tensor,
+    grad_output: torch.Tensor,
+    recipe: Optional[Recipe],
+    mode: str,
+    num_microbatches: int,
+    warmup_steps: int,
+    min_run_time: float,
+    profile: bool,
+    label: str,
+) -> float:
+    """Benchmark one execution path and return milliseconds per microbatch."""
+    _run_microbatches(
+        layer,
+        x,
+        m_splits,
+        grad_output,
+        recipe=recipe,
+        mode=mode,
+        num_microbatches=warmup_steps,
+    )
+    torch.cuda.synchronize()
+
+    if profile:
+        torch.cuda.nvtx.range_push(label)
+
+    timing = benchmark.Timer(
+        stmt=(
+            "_run_microbatches(layer, x, m_splits, grad_output, recipe=recipe, "
+            "mode=mode, num_microbatches=num_microbatches)"
+        ),
+        globals={
+            "_run_microbatches": _run_microbatches,
+            "layer": layer,
+            "x": x,
+            "m_splits": m_splits,
+            "grad_output": grad_output,
+            "recipe": recipe,
+            "mode": mode,
+            "num_microbatches": num_microbatches,
+        },
+        num_threads=1,
+    ).blocked_autorange(min_run_time=min_run_time)
+
+    if profile:
+        torch.cuda.nvtx.range_pop()
+
+    return timing.median * 1000 / num_microbatches
+
+
+def _gemm_tflops(
+    *,
+    total_rows: int,
+    projection: Projection,
+    mode: str,
+    time_ms: float,
+) -> float:
+    """Report GEMM FLOP/s while timing also includes quantization and Python overhead."""
+    flops = 2 * total_rows * projection.in_features * projection.out_features
+    if mode == "fwd_bwd":
+        flops *= 3
+    return flops / time_ms / 1e9
+
+
+def _parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--precision",
+        choices=("all", "bf16", "mxfp8"),
+        default="all",
+    )
+    parser.add_argument(
+        "--projection",
+        choices=("all", "fc1", "fc2"),
+        default="all",
+    )
+    parser.add_argument(
+        "--mode",
+        choices=("fwd", "fwd_bwd"),
+        default="fwd_bwd",
+    )
+    parser.add_argument("--num-microbatches", type=int, default=16)
+    parser.add_argument("--warmup-steps", type=int, default=500)
+    parser.add_argument("--min-run-time", type=float, default=30.0)
+    parser.add_argument("--seed", type=int, default=1234)
+    parser.add_argument(
+        "--m-splits",
+        type=str,
+        default=None,
+        help="Optional comma-separated per-expert rows; defaults to Qwen3.5-397B EP32.",
+    )
+    parser.add_argument(
+        "--skip-correctness",
+        action="store_true",
+        help="Skip output, input-gradient, and parameter-gradient parity checks.",
+    )
+    parser.add_argument("--profile", action="store_true", help="Add per-path NVTX ranges.")
+    parser.add_argument("--output-csv", type=str, default=None)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = _parse_args()
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required for this benchmark.")
+    if args.num_microbatches < 1 or args.warmup_steps < 1:
+        raise ValueError("num_microbatches and warmup_steps must both be positive.")
+
+    if args.m_splits is None:
+        num_local_experts = QWEN_NUM_EXPERTS // QWEN_EXPERT_PARALLEL_SIZE
+        rows_per_expert = QWEN_SEQUENCE_LENGTH * QWEN_TOP_K // num_local_experts
+        split_sizes = [rows_per_expert] * num_local_experts
+    else:
+        split_sizes = [int(value) for value in args.m_splits.split(",") if value]
+        if not split_sizes or any(value < 0 for value in split_sizes):
+            raise ValueError("m_splits must contain non-negative integers.")
+        num_local_experts = len(split_sizes)
+
+    total_rows = sum(split_sizes)
+    if total_rows == 0:
+        raise ValueError("At least one routed token row is required for this benchmark.")
+
+    precision_names = ("bf16", "mxfp8") if args.precision == "all" else (args.precision,)
+    projection_names = ("fc1", "fc2") if args.projection == "all" else (args.projection,)
+
+    recipes: dict[str, Optional[Recipe]] = {
+        "bf16": None,
+        "mxfp8": MXFP8BlockScaling(),
+    }
+    mxfp8_available, reason_for_no_mxfp8 = FP8GlobalStateManager.is_mxfp8_available()
+
+    print("Qwen3.5-397B-A17B GroupedLinear benchmark")
+    print(f"  local experts: {num_local_experts}")
+    print(f"  total routed rows: {total_rows}")
+    print(f"  m_splits: {split_sizes}")
+    print(f"  mode: {args.mode}")
+    print("  primary parameter dtype: BF16")
+    print(f"  microbatches per timing invocation: {args.num_microbatches}")
+    print()
+
+    rows = []
+    for precision_name in precision_names:
+        recipe = recipes[precision_name]
+        if precision_name == "mxfp8" and not mxfp8_available:
+            print(f"Skipping MXFP8: {reason_for_no_mxfp8}")
+            continue
+        if not is_module_grouped_tensor_path_supported(recipe, torch.bfloat16):
+            print(
+                f"Skipping {precision_name}: the cuBLASLt grouped-tensor path is unsupported "
+                "on this GPU or cuBLASLt version."
+            )
+            continue
+
+        for projection_name in projection_names:
+            projection = PROJECTIONS[projection_name]
+            torch.manual_seed(args.seed)
+            torch.cuda.manual_seed(args.seed)
+
+            layers = {
+                path.name: _build_layer(
+                    projection=projection,
+                    num_local_experts=num_local_experts,
+                    use_grouped_tensor=path.use_grouped_tensor,
+                )
+                for path in EXECUTION_PATHS
+            }
+            layers["cublas_grouped_gemm"].load_state_dict(layers["multistream"].state_dict())
+
+            base_x = torch.randn(
+                total_rows,
+                projection.in_features,
+                dtype=torch.bfloat16,
+                device="cuda",
+            )
+            inputs = {
+                path.name: base_x.detach().clone().requires_grad_(args.mode == "fwd_bwd")
+                for path in EXECUTION_PATHS
+            }
+            grad_output = torch.randn(
+                total_rows,
+                projection.out_features,
+                dtype=torch.bfloat16,
+                device="cuda",
+            )
+            splits = {
+                path.name: _make_m_splits(
+                    use_grouped_tensor=path.use_grouped_tensor,
+                    split_sizes=split_sizes,
+                )
+                for path in EXECUTION_PATHS
+            }
+
+            if not args.skip_correctness:
+                _validate_path_parity(
+                    multistream_layer=layers["multistream"],
+                    grouped_layer=layers["cublas_grouped_gemm"],
+                    multistream_x=inputs["multistream"],
+                    grouped_x=inputs["cublas_grouped_gemm"],
+                    multistream_splits=splits["multistream"],
+                    grouped_splits=splits["cublas_grouped_gemm"],
+                    grad_output=grad_output,
+                    recipe=recipe,
+                    mode=args.mode,
+                )
+
+            timings = {}
+            for path in EXECUTION_PATHS:
+                label = f"{precision_name}_{projection_name}_{args.mode}_{path.name}"
+                timing_ms = _benchmark_path(
+                    layer=layers[path.name],
+                    x=inputs[path.name],
+                    m_splits=splits[path.name],
+                    grad_output=grad_output,
+                    recipe=recipe,
+                    mode=args.mode,
+                    num_microbatches=args.num_microbatches,
+                    warmup_steps=args.warmup_steps,
+                    min_run_time=args.min_run_time,
+                    profile=args.profile,
+                    label=label,
+                )
+                timings[path.name] = timing_ms
+
+            speedup = timings["multistream"] / timings["cublas_grouped_gemm"]
+            for path in EXECUTION_PATHS:
+                timing_ms = timings[path.name]
+                rows.append(
+                    {
+                        "precision": precision_name,
+                        "projection": projection_name,
+                        "mode": args.mode,
+                        "execution_path": path.name,
+                        "num_local_experts": num_local_experts,
+                        "total_rows": total_rows,
+                        "in_features": projection.in_features,
+                        "out_features": projection.out_features,
+                        "time_ms": timing_ms,
+                        "gemm_tflops": _gemm_tflops(
+                            total_rows=total_rows,
+                            projection=projection,
+                            mode=args.mode,
+                            time_ms=timing_ms,
+                        ),
+                        "speedup_vs_multistream": (1.0 if path.name == "multistream" else speedup),
+                    }
+                )
+
+            print(
+                f"{precision_name:6s} {projection_name} {args.mode}: "
+                f"multistream={timings['multistream']:.3f} ms, "
+                f"cuBLAS grouped={timings['cublas_grouped_gemm']:.3f} ms, "
+                f"speedup={speedup:.3f}x"
+            )
+
+    results = pd.DataFrame(rows)
+    print()
+    print(results.to_string(index=False))
+    if args.output_csv is not None:
+        results.to_csv(args.output_csv, index=False)
+        print(f"\nWrote {args.output_csv}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py
index 10d0dbd227..1dd2b16ade 100644
--- a/tests/pytorch/test_grouped_linear.py
+++ b/tests/pytorch/test_grouped_linear.py
@@ -14,6 +14,8 @@
 import transformer_engine.pytorch as te
 from transformer_engine.common import recipe
 from transformer_engine.pytorch import (
+    Float8BlockQuantizer,
+    Float8CurrentScalingQuantizer,
     Float8Quantizer,
     Fp8Padding,
     Fp8Unpadding,
@@ -30,6 +32,11 @@
     general_grouped_gemm,
     general_grouped_gemm_for_grouped_tensor,
 )
+from transformer_engine.pytorch.constants import TE_DType
+from transformer_engine.pytorch.module.grouped_linear import (
+    _GroupedLinear,
+    is_module_grouped_tensor_path_supported,
+)
 from transformer_engine.pytorch.quantization import (
     FP8GlobalStateManager,
     get_align_size_for_quantization,
@@ -1154,46 +1161,124 @@ def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate, use_b
             torch.testing.assert_close(o, o_ref, **tols)
 
 
+@pytest.mark.parametrize("use_bias_scale", [False, True])
+def test_grouped_gemm_grouped_tensor_zero_work_bias(use_bias_scale) -> None:
+    """A grouped bias operation is a no-op when every group has zero rows.
+
+    Zero-sized CUDA tensors may legally have a null data pointer. Exercise both bias entry points
+    so neither the ordinary nor scaled path mistakes that pointer for missing output storage.
+    This BF16 case runs on both Hopper and Blackwell when grouped cuBLASLt GEMM is available.
+    """
+    if not is_module_grouped_tensor_path_supported(
+        None,
+        torch.bfloat16,
+    ):
+        pytest.skip("BF16 GroupedTensor GEMM is unavailable.")
+
+    num_groups = 4
+    in_features = 256
+    out_features = 256
+    m_sizes = [0] * num_groups
+    dtype = torch.bfloat16
+    device = torch.device("cuda")
+
+    weights = [
+        torch.randn(out_features, in_features, dtype=dtype, device=device)
+        for _ in range(num_groups)
+    ]
+    biases = [torch.randn(1, out_features, dtype=dtype, device=device) for _ in range(num_groups)]
+    grouped_weights = _make_grouped_tensor_uniform(
+        num_groups, out_features, in_features, device, dtype
+    )
+    grouped_input = _make_grouped_tensor_from_splits(m_sizes, in_features, device, dtype)
+    grouped_output = _make_grouped_tensor_from_splits(m_sizes, out_features, device, dtype)
+    grouped_bias = _make_grouped_tensor_uniform(num_groups, 1, out_features, device, dtype)
+    _pack_grouped_tensor(grouped_weights, weights)
+    _pack_grouped_tensor(grouped_bias, biases)
+
+    bias_scale = torch.empty(0, dtype=torch.float32, device=device) if use_bias_scale else None
+    general_grouped_gemm_for_grouped_tensor(
+        grouped_weights,
+        grouped_input,
+        grouped_output,
+        layout="TN",
+        bias=grouped_bias,
+        bias_scale=bias_scale,
+    )
+    torch.cuda.synchronize()
+
+    assert grouped_output.rowwise_data.numel() == 0
+
+
 @pytest.mark.parametrize("layout", ["TN", "NN", "NT"])
 @pytest.mark.parametrize("accumulate", [False, True])
-@pytest.mark.parametrize("quant_type", ["bf16", "mxfp8"])
+@pytest.mark.parametrize(
+    "quant_type", ["bf16", "fp8_current_scaling", "mxfp8", "fp8_block_scaling"]
+)
 def test_grouped_gemm_grouped_tensor_zero_work(layout, accumulate, quant_type) -> None:
     """Grouped GEMM with all-zero split sizes (zero total work).
 
     For wgrad (NT layout) the output should be zero when not accumulating,
     or unchanged when accumulating with beta=1.
     """
-    if torch.cuda.get_device_capability() < (10, 0):
-        pytest.skip("Grouped GEMM requires Blackwell (SM100) or newer.")
     if not is_bf16_available():
         pytest.skip("bfloat16 is required for grouped GEMM test.")
-    if quant_type == "mxfp8" and not mxfp8_available:
-        pytest.skip(reason_for_no_mxfp8)
 
     z = 4
     k, n = 256, 256
     dtype = torch.bfloat16
     device = torch.device("cuda")
-    use_mxfp8 = quant_type == "mxfp8"
+
+    test_recipe = {
+        "bf16": None,
+        "fp8_current_scaling": recipe.Float8CurrentScaling(),
+        "mxfp8": recipe.MXFP8BlockScaling(),
+        "fp8_block_scaling": recipe.Float8BlockScaling(),
+    }[quant_type]
+    if not is_module_grouped_tensor_path_supported(
+        test_recipe,
+        dtype,
+    ):
+        pytest.skip(f"{quant_type} grouped-tensor GEMM is unavailable")
 
     transa = layout[0] == "T"
     transb = layout[1] == "T"
     zero_first_dims = torch.zeros(z, dtype=torch.int64, device=device)
 
+    def _make_quantizer(fp8_dtype, rowwise, columnwise):
+        if quant_type == "fp8_current_scaling":
+            quantizer = Float8CurrentScalingQuantizer(fp8_dtype=fp8_dtype, device=device)
+            quantizer.set_usage(rowwise=rowwise, columnwise=columnwise)
+        elif quant_type == "mxfp8":
+            quantizer = MXFP8Quantizer(
+                fp8_dtype=fp8_dtype,
+                rowwise=rowwise,
+                columnwise=columnwise,
+            )
+        elif quant_type == "fp8_block_scaling":
+            quantizer = Float8BlockQuantizer(
+                fp8_dtype=fp8_dtype,
+                rowwise=rowwise,
+                columnwise=columnwise,
+                force_pow_2_scales=False,
+                amax_epsilon=0.0,
+                block_scaling_dim=1,
+            )
+        else:
+            raise ValueError(f"Unsupported quantized zero-work test type {quant_type}")
+        quantizer.optimize_for_gemm = True
+        return quantizer
+
     def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a):
         """Create a GroupedTensor with non-zero logical_shape but zero first_dims."""
         buf = torch.randn(0, logical_last_dim, dtype=dtype, device=device)
-        if use_mxfp8:
+        if test_recipe is not None:
             if is_a:
                 rowwise, columnwise = transa, not transa
             else:
                 rowwise, columnwise = not transb, transb
-            quantizer = MXFP8Quantizer(
-                fp8_dtype=tex.DType.kFloat8E4M3,
-                rowwise=rowwise,
-                columnwise=columnwise,
-            )
-            quantizer.optimize_for_gemm = True
+            fp8_dtype = TE_DType[torch.float8_e4m3fn]
+            quantizer = _make_quantizer(fp8_dtype, rowwise, columnwise)
             return tex.group_quantize(buf, quantizer, z, zero_first_dims)
         return GroupedTensor.make_grouped_tensor(
             num_tensors=z,
@@ -1208,12 +1293,18 @@ def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a):
 
     if layout in ("TN", "NN"):
         weight_tensors = [torch.randn(n, k, dtype=dtype, device=device) for _ in range(z)]
-        if use_mxfp8:
-            grouped_A = _make_grouped_tensor_quantized_mxfp8(
-                weight_tensors,
+        if test_recipe is not None:
+            grouped_weight = torch.cat(weight_tensors, dim=0)
+            weight_quantizer = _make_quantizer(
+                TE_DType[torch.float8_e4m3fn],
                 rowwise=transa,
                 columnwise=not transa,
-                device=device,
+            )
+            grouped_A = tex.group_quantize(
+                grouped_weight,
+                weight_quantizer,
+                z,
+                torch.full((z,), n, dtype=torch.int64, device=device),
             )
         else:
             grouped_A = _make_grouped_tensor_uniform(z, n, k, device, dtype)
@@ -1506,16 +1597,638 @@ def test_fp8_grouped_gemm(shape, accumulate):
 
 @pytest.fixture(autouse=True)
 def _reset_fp8_state(monkeypatch):
-    monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "0")
+    monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False)
     yield
     FP8GlobalStateManager.reset()
     monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False)
 
 
+@pytest.mark.parametrize(
+    "m_splits,exception",
+    [([256, 256], ValueError), (torch.tensor([256, 256]), ValueError)],
+    ids=["python-list", "cpu-tensor"],
+)
+def test_single_grouped_weight_rejects_host_m_splits(monkeypatch, m_splits, exception):
+    """A single parent parameter must never fall back to host-split per-expert GEMMs."""
+    if not is_module_grouped_tensor_path_supported(
+        None,
+        torch.bfloat16,
+    ):
+        pytest.skip("Native GroupedTensor GEMM is unavailable on this system.")
+
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+    grouped_linear = GroupedLinear(
+        2,
+        64,
+        64,
+        bias=False,
+        params_dtype=torch.bfloat16,
+        device="cuda",
+        single_grouped_weight=True,
+        use_grouped_tensor=True,
+    )
+    x = torch.randn(512, 64, dtype=torch.bfloat16, device="cuda")
+    with pytest.raises(exception, match="requires.*CUDA"):
+        grouped_linear(x, m_splits)
+
+
+@pytest.mark.parametrize(
+    "fp8_recipe",
+    [
+        pytest.param(None, id="bf16"),
+        pytest.param(
+            recipe.Float8CurrentScaling(),
+            marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8),
+            id="fp8-current-scaling",
+        ),
+        pytest.param(
+            recipe.MXFP8BlockScaling(),
+            marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8),
+            id="mxfp8",
+        ),
+        pytest.param(
+            recipe.Float8BlockScaling(),
+            marks=pytest.mark.skipif(
+                not _fp8_block_scaling_available,
+                reason=_reason_for_no_fp8_block_scaling,
+            ),
+            id="fp8-block-scaling",
+        ),
+    ],
+)
+def test_single_grouped_weight_matches_discrete_grouped_tensor_path(monkeypatch, fp8_recipe):
+    """Match single and discrete weights while both use CUDA m_splits and grouped GEMM."""
+    if not is_module_grouped_tensor_path_supported(
+        fp8_recipe,
+        torch.bfloat16,
+    ):
+        pytest.skip("Recipe is not supported with a single grouped weight on this system.")
+
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+    FP8GlobalStateManager.reset()
+
+    num_gemms = 3
+    in_features = 256
+    out_features = 256
+    m_splits = torch.tensor([256, 512, 256], dtype=torch.int64, device="cuda")
+    total_tokens = int(m_splits.sum())
+    weights = torch.randn(
+        num_gemms,
+        out_features,
+        in_features,
+        dtype=torch.bfloat16,
+        device="cuda",
+    )
+
+    discrete = GroupedLinear(
+        num_gemms,
+        in_features,
+        out_features,
+        bias=False,
+        params_dtype=torch.bfloat16,
+        device="cuda",
+        use_grouped_tensor=True,
+    )
+    single = GroupedLinear(
+        num_gemms,
+        in_features,
+        out_features,
+        bias=False,
+        params_dtype=torch.bfloat16,
+        device="cuda",
+        single_grouped_weight=True,
+        use_grouped_tensor=True,
+    )
+    with torch.no_grad():
+        for idx in range(num_gemms):
+            getattr(discrete, f"weight{idx}").copy_(weights[idx])
+        single.weight.rowwise_data.view_as(weights).copy_(weights)
+
+    x = torch.randn(total_tokens, in_features, dtype=torch.bfloat16, device="cuda")
+    dy = torch.randn(total_tokens, out_features, dtype=torch.bfloat16, device="cuda")
+    x_discrete = x.detach().clone().requires_grad_(True)
+    x_single = x.detach().clone().requires_grad_(True)
+
+    with autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe):
+        y_discrete = discrete(x_discrete, m_splits)
+        y_single = single(x_single, m_splits)
+    y_discrete.backward(dy)
+    y_single.backward(dy)
+
+    tolerances = dict(rtol=1e-2, atol=5e-3)
+    torch.testing.assert_close(y_single.float(), y_discrete.float(), **tolerances)
+    torch.testing.assert_close(x_single.grad.float(), x_discrete.grad.float(), **tolerances)
+    discrete_wgrad = torch.stack(
+        [getattr(discrete, f"weight{idx}").grad for idx in range(num_gemms)]
+    )
+    torch.testing.assert_close(single.weight.grad.float(), discrete_wgrad.float(), **tolerances)
+
+
+@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8)
+def test_single_grouped_weight_mxfp8_workspace_cache(monkeypatch):
+    """BF16 primary weights update one persistent MXFP8 grouped workspace per iteration."""
+    mxfp8_recipe = recipe.MXFP8BlockScaling()
+    if not is_module_grouped_tensor_path_supported(
+        mxfp8_recipe,
+        torch.bfloat16,
+    ):
+        pytest.skip("MXFP8 single-weight GroupedTensor path is unavailable on this system.")
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+    FP8GlobalStateManager.reset()
+    grouped_linear = GroupedLinear(
+        2,
+        256,
+        256,
+        bias=False,
+        params_dtype=torch.bfloat16,
+        device="cuda",
+        single_grouped_weight=True,
+        use_grouped_tensor=True,
+    )
+
+    x = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+    m_splits = torch.tensor([256, 256], dtype=torch.int64, device="cuda")
+
+    with autocast(enabled=True, recipe=mxfp8_recipe):
+        grouped_linear(x, m_splits, is_first_microbatch=True)
+        workspace = grouped_linear._fp8_workspaces["weight"]
+        assert isinstance(workspace, GroupedTensor)
+        pointers = (
+            workspace.rowwise_data.data_ptr(),
+            workspace.columnwise_data.data_ptr(),
+            workspace.scale_inv.data_ptr(),
+            workspace.columnwise_scale_inv.data_ptr(),
+        )
+        old_data = workspace.rowwise_data.clone()
+
+        with torch.no_grad():
+            grouped_linear.weight.rowwise_data.add_(1)
+        grouped_linear(x, m_splits, is_first_microbatch=False)
+        assert torch.equal(workspace.rowwise_data, old_data)
+
+        grouped_linear(x, m_splits, is_first_microbatch=True)
+        assert pointers == (
+            workspace.rowwise_data.data_ptr(),
+            workspace.columnwise_data.data_ptr(),
+            workspace.scale_inv.data_ptr(),
+            workspace.columnwise_scale_inv.data_ptr(),
+        )
+        assert not torch.equal(workspace.rowwise_data, old_data)
+
+
+@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8)
+@pytest.mark.parametrize("fp8_recipe", [recipe.MXFP8BlockScaling()], ids=recipe_id)
+def test_single_grouped_weight_with_disabled_weight_preswizzle(monkeypatch, fp8_recipe):
+    """Grouped weight preparation preserves a disabled preswizzle decision."""
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+    FP8GlobalStateManager.reset()
+    with quantized_model_init(enabled=True, recipe=fp8_recipe):
+        grouped_linear = GroupedLinear(
+            2,
+            256,
+            256,
+            bias=False,
+            params_dtype=torch.bfloat16,
+            device="cuda",
+            single_grouped_weight=True,
+            use_grouped_tensor=True,
+        )
+
+    weight_quantizer = grouped_linear.weight.quantizer
+    assert weight_quantizer is not None
+    preswizzle = grouped_linear._enable_weight_preswizzle(
+        weight_quantizer,
+        grouped_linear.weight,
+    )
+    assert preswizzle is False
+    weight_quantizer.optimize_for_gemm = preswizzle
+    grouped_weight, new_workspaces = _GroupedLinear._prepare_weights_for_grouped_tensor_gemm(
+        (grouped_linear.weight,),
+        [weight_quantizer],
+        [None],
+        num_gemms=grouped_linear.num_gemms,
+        single_grouped_weight=True,
+        with_quantized_compute=True,
+        columnwise_usage=True,
+        activation_dtype=torch.bfloat16,
+        is_first_microbatch=True,
+        skip_fp8_weight_update=None,
+        cache_weight=True,
+    )
+
+    assert weight_quantizer.optimize_for_gemm is False
+    assert len(new_workspaces) == 1
+    assert new_workspaces[0] is None
+    assert grouped_weight is grouped_linear.weight
+    assert hasattr(grouped_weight, "_with_gemm_swizzled_scales")
+    assert grouped_weight._with_gemm_swizzled_scales is False
+
+
+@pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8)
+def test_single_grouped_primary_mxfp8_bypasses_weight_workspace(monkeypatch):
+    """An MXFP8 primary grouped parameter is already GEMM-ready and is not requantized."""
+    mxfp8_recipe = recipe.MXFP8BlockScaling()
+    if not is_module_grouped_tensor_path_supported(
+        mxfp8_recipe,
+        torch.bfloat16,
+    ):
+        pytest.skip("MXFP8 single-weight GroupedTensor path is unavailable on this system.")
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+    FP8GlobalStateManager.reset()
+    with quantized_model_init(enabled=True, recipe=mxfp8_recipe):
+        grouped_linear = GroupedLinear(
+            2,
+            256,
+            256,
+            bias=False,
+            params_dtype=torch.bfloat16,
+            device="cuda",
+            single_grouped_weight=True,
+            use_grouped_tensor=True,
+        )
+
+    x = torch.randn(512, 256, dtype=torch.bfloat16, device="cuda")
+    m_splits = torch.tensor([256, 256], dtype=torch.int64, device="cuda")
+    with torch.no_grad(), autocast(enabled=True, recipe=mxfp8_recipe):
+        grouped_linear(x, m_splits, is_first_microbatch=True)
+    assert grouped_linear.weight.quantizer is not None
+    assert "weight" not in grouped_linear._fp8_workspaces
+
+
 def _clone_outputs(outputs):
     return [None if out is None else out.detach().clone() for out in outputs]
 
 
+def _grouped_linear_weight_params(module):
+    if module.single_grouped_weight:
+        return [module.weight]
+    return [getattr(module, f"weight{i}") for i in range(module.num_gemms)]
+
+
+def _grouped_linear_bias_params(module):
+    if not module.use_bias:
+        return []
+    if module.single_grouped_bias:
+        return [module.bias]
+    return [getattr(module, f"bias{i}") for i in range(module.num_gemms)]
+
+
+def _run_grouped_parameter_layout(
+    *,
+    use_grouped_tensor,
+    fp8_recipe,
+    single_grouped_weight,
+    single_grouped_bias,
+    use_bias,
+    delay_wgrad_compute,
+    fuse_wgrad_accumulation,
+    x_base,
+    dy,
+    weights,
+    biases,
+    m_splits,
+    save_original_input=False,
+):
+    """Run one layout and return all numerically observable forward/backward results."""
+    FP8GlobalStateManager.reset()
+    num_gemms, out_features, in_features = weights.shape
+    module = GroupedLinear(
+        num_gemms,
+        in_features,
+        out_features,
+        bias=use_bias,
+        params_dtype=torch.bfloat16,
+        device="cuda",
+        fuse_wgrad_accumulation=fuse_wgrad_accumulation,
+        delay_wgrad_compute=delay_wgrad_compute,
+        single_grouped_weight=single_grouped_weight,
+        single_grouped_bias=single_grouped_bias,
+        use_grouped_tensor=use_grouped_tensor,
+        save_original_input=save_original_input,
+    )
+
+    with torch.no_grad():
+        if single_grouped_weight:
+            module.weight.rowwise_data.view_as(weights).copy_(weights)
+        else:
+            for i in range(num_gemms):
+                getattr(module, f"weight{i}").copy_(weights[i])
+        if use_bias:
+            if single_grouped_bias:
+                module.bias.rowwise_data.view_as(biases).copy_(biases)
+            else:
+                for i in range(num_gemms):
+                    getattr(module, f"bias{i}").copy_(biases[i])
+
+    flat_main_grad = None
+    initial_main_grad = None
+    weight_params = _grouped_linear_weight_params(module)
+    if fuse_wgrad_accumulation:
+        # MCore owns one flat FP32 grad buffer. Discrete parameters receive per-expert
+        # views, while a single grouped parameter receives one view over the full range.
+        flat_main_grad = torch.full(
+            (weights.numel(),),
+            0.25,
+            dtype=torch.float32,
+            device="cuda",
+        )
+        packed_main_grad = flat_main_grad.view_as(weights)
+        main_grad_views = (
+            [packed_main_grad]
+            if single_grouped_weight
+            else [packed_main_grad[i] for i in range(num_gemms)]
+        )
+        for param, main_grad in zip(weight_params, main_grad_views):
+            param.main_grad = main_grad
+            param.overwrite_main_grad = False
+            param.zero_out_wgrad = False
+            param.grad_added_to_main_grad = False
+            assert (
+                param.main_grad.untyped_storage().data_ptr()
+                == flat_main_grad.untyped_storage().data_ptr()
+            )
+        initial_main_grad = flat_main_grad.clone()
+
+    x = x_base.detach().clone().requires_grad_(True)
+    m_splits_arg = (
+        torch.tensor(m_splits, dtype=torch.int64, device="cuda") if use_grouped_tensor else m_splits
+    )
+    with autocast(enabled=fp8_recipe is not None, recipe=fp8_recipe):
+        y = module(
+            x,
+            m_splits_arg,
+            is_first_microbatch=False if fuse_wgrad_accumulation else None,
+        )
+    y.backward(dy)
+
+    if fuse_wgrad_accumulation and delay_wgrad_compute:
+        torch.testing.assert_close(flat_main_grad, initial_main_grad, rtol=0, atol=0)
+
+    # The grouped-tensor path computes dbias during the main backward even when dW is delayed.
+    if use_bias and use_grouped_tensor:
+        assert all(param.grad is not None for param in _grouped_linear_bias_params(module))
+
+    if delay_wgrad_compute:
+        module.backward_dw()
+
+    if fuse_wgrad_accumulation:
+        assert not torch.equal(flat_main_grad, initial_main_grad)
+        for param in weight_params:
+            assert param.grad_added_to_main_grad
+            assert (
+                param.main_grad.untyped_storage().data_ptr()
+                == flat_main_grad.untyped_storage().data_ptr()
+            )
+        packed_wgrad = flat_main_grad.view_as(weights)
+    elif single_grouped_weight:
+        packed_wgrad = module.weight.grad.view_as(weights)
+    else:
+        packed_wgrad = torch.stack([param.grad for param in weight_params])
+
+    packed_dbias = None
+    if use_bias:
+        bias_params = _grouped_linear_bias_params(module)
+        if single_grouped_bias:
+            packed_dbias = bias_params[0].grad.view_as(biases)
+        else:
+            packed_dbias = torch.stack([param.grad for param in bias_params])
+
+    return {
+        "output": y.detach().clone(),
+        "dgrad": x.grad.detach().clone(),
+        "wgrad": packed_wgrad.detach().clone(),
+        "dbias": None if packed_dbias is None else packed_dbias.detach().clone(),
+    }
+
+
+_GROUPED_PARAMETER_LAYOUTS = [
+    pytest.param(False, False, False, id="no-bias-discrete-weight"),
+    pytest.param(False, True, False, id="no-bias-single-weight"),
+    pytest.param(True, False, False, id="bias-discrete-weight-discrete-bias"),
+    pytest.param(True, True, False, id="bias-single-weight-discrete-bias"),
+    pytest.param(True, False, True, id="bias-discrete-weight-single-bias"),
+    pytest.param(True, True, True, id="bias-single-weight-single-bias"),
+]
+
+
+@pytest.mark.parametrize(
+    "use_bias,single_grouped_weight,single_grouped_bias", _GROUPED_PARAMETER_LAYOUTS
+)
+@pytest.mark.parametrize(
+    "fp8_recipe",
+    [
+        pytest.param(None, id="bf16"),
+        pytest.param(
+            recipe.Float8CurrentScaling(),
+            marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8),
+            id="fp8-current-scaling",
+        ),
+        pytest.param(
+            recipe.MXFP8BlockScaling(),
+            marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8),
+            id="mxfp8",
+        ),
+        pytest.param(
+            recipe.Float8BlockScaling(),
+            marks=pytest.mark.skipif(
+                not _fp8_block_scaling_available,
+                reason=_reason_for_no_fp8_block_scaling,
+            ),
+            id="fp8-block-scaling",
+        ),
+    ],
+)
+@pytest.mark.parametrize("delay_wgrad_compute", _ALL_BOOLEAN)
+@pytest.mark.parametrize("fuse_wgrad_accumulation", _ALL_BOOLEAN)
+def test_grouped_parameter_layout_matches_cpu_m_splits(
+    monkeypatch,
+    use_bias,
+    single_grouped_weight,
+    single_grouped_bias,
+    fp8_recipe,
+    delay_wgrad_compute,
+    fuse_wgrad_accumulation,
+):
+    """Match CUDA m_splits and all meaningful parameter layouts against the legacy path."""
+    if not is_module_grouped_tensor_path_supported(
+        fp8_recipe,
+        torch.bfloat16,
+    ):
+        pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.")
+    FP8GlobalStateManager.reset()
+
+    torch.manual_seed(1234)
+    # MXFP8 and the grouped FP8 recipes require aligned expert problems. Use the same
+    # 256-aligned shapes for every recipe so all precision modes exercise one layout.
+    num_gemms = 2
+    in_features = 256
+    out_features = 256
+    m_splits = [256, 512]
+    total_tokens = sum(m_splits)
+    x_base = (0.1 * torch.randn(total_tokens, in_features, device="cuda")).to(torch.bfloat16)
+    dy = (0.1 * torch.randn(total_tokens, out_features, device="cuda")).to(torch.bfloat16)
+    weights = (0.1 * torch.randn(num_gemms, out_features, in_features, device="cuda")).to(
+        torch.bfloat16
+    )
+    biases = None
+    if use_bias:
+        biases = (0.1 * torch.randn(num_gemms, out_features, device="cuda")).to(torch.bfloat16)
+
+    # The CPU m_splits baseline is explicitly the legacy, discrete-parameter contract.
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0")
+    reference = _run_grouped_parameter_layout(
+        use_grouped_tensor=False,
+        fp8_recipe=fp8_recipe,
+        single_grouped_weight=False,
+        single_grouped_bias=False,
+        use_bias=use_bias,
+        delay_wgrad_compute=delay_wgrad_compute,
+        fuse_wgrad_accumulation=fuse_wgrad_accumulation,
+        x_base=x_base,
+        dy=dy,
+        weights=weights,
+        biases=biases,
+        m_splits=m_splits,
+    )
+
+    # Enable single parameters only for the CUDA m_splits target. The explicit layout flags
+    # below still decide whether this particular case uses discrete or grouped parameters.
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+    result = _run_grouped_parameter_layout(
+        use_grouped_tensor=True,
+        fp8_recipe=fp8_recipe,
+        single_grouped_weight=single_grouped_weight,
+        single_grouped_bias=single_grouped_bias,
+        use_bias=use_bias,
+        delay_wgrad_compute=delay_wgrad_compute,
+        fuse_wgrad_accumulation=fuse_wgrad_accumulation,
+        x_base=x_base,
+        dy=dy,
+        weights=weights,
+        biases=biases,
+        m_splits=m_splits,
+    )
+
+    tolerances = dict(rtol=1e-2, atol=5e-3)
+
+    for name in ("output", "dgrad", "wgrad", "dbias"):
+        if reference[name] is None:
+            assert result[name] is None
+        else:
+            torch.testing.assert_close(
+                result[name].float(),
+                reference[name].float(),
+                **tolerances,
+                msg=f"Mismatch for {name}",
+            )
+
+
+@pytest.mark.parametrize(
+    "fp8_recipe",
+    [
+        pytest.param(None, id="bf16"),
+        pytest.param(
+            recipe.Float8CurrentScaling(),
+            marks=pytest.mark.skipif(not _fp8_available, reason=_reason_for_no_fp8),
+            id="fp8-current-scaling",
+        ),
+        pytest.param(
+            recipe.MXFP8BlockScaling(),
+            marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8),
+            id="mxfp8",
+        ),
+        pytest.param(
+            recipe.Float8BlockScaling(),
+            marks=pytest.mark.skipif(
+                not _fp8_block_scaling_available,
+                reason=_reason_for_no_fp8_block_scaling,
+            ),
+            id="fp8-block-scaling",
+        ),
+        pytest.param(
+            recipe.NVFP4BlockScaling(disable_stochastic_rounding=True),
+            marks=pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4),
+            id="nvfp4",
+        ),
+    ],
+)
+@pytest.mark.parametrize("single_grouped_weight", _ALL_BOOLEAN)
+def test_grouped_tensor_save_original_input_matches_saved_grouped_input(
+    monkeypatch,
+    fp8_recipe,
+    single_grouped_weight,
+):
+    """Saving raw input must preserve native grouped forward, dgrad, and wgrad numerics."""
+    if not is_module_grouped_tensor_path_supported(fp8_recipe, torch.bfloat16):
+        pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.")
+    if single_grouped_weight and fp8_recipe is not None and fp8_recipe.nvfp4():
+        pytest.skip(
+            "NVFP4 grouped GEMM with single_grouped_weight is not supported yet; "
+            "only discrete weights are supported."
+        )
+
+    def reject_split_fallback(*_args, **_kwargs):
+        pytest.fail("save_original_input unexpectedly selected the split-quantize path")
+
+    monkeypatch.setattr(
+        "transformer_engine.pytorch.module.grouped_linear._split_quantize",
+        reject_split_fallback,
+    )
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+
+    torch.manual_seed(1234)
+    num_gemms = 2
+    in_features = 256
+    out_features = 256
+    m_splits = [256, 512]
+    total_tokens = sum(m_splits)
+    x_base = (0.1 * torch.randn(total_tokens, in_features, device="cuda")).to(torch.bfloat16)
+    dy = (0.1 * torch.randn(total_tokens, out_features, device="cuda")).to(torch.bfloat16)
+    weights = (0.1 * torch.randn(num_gemms, out_features, in_features, device="cuda")).to(
+        torch.bfloat16
+    )
+
+    saved_grouped = _run_grouped_parameter_layout(
+        use_grouped_tensor=True,
+        fp8_recipe=fp8_recipe,
+        single_grouped_weight=single_grouped_weight,
+        single_grouped_bias=False,
+        use_bias=False,
+        delay_wgrad_compute=False,
+        fuse_wgrad_accumulation=False,
+        x_base=x_base,
+        dy=dy,
+        weights=weights,
+        biases=None,
+        m_splits=m_splits,
+        save_original_input=False,
+    )
+    saved_original = _run_grouped_parameter_layout(
+        use_grouped_tensor=True,
+        fp8_recipe=fp8_recipe,
+        single_grouped_weight=single_grouped_weight,
+        single_grouped_bias=False,
+        use_bias=False,
+        delay_wgrad_compute=False,
+        fuse_wgrad_accumulation=False,
+        x_base=x_base,
+        dy=dy,
+        weights=weights,
+        biases=None,
+        m_splits=m_splits,
+        save_original_input=True,
+    )
+
+    for name in ("output", "dgrad", "wgrad"):
+        torch.testing.assert_close(
+            saved_original[name].float(),
+            saved_grouped[name].float(),
+            rtol=1e-2,
+            atol=5e-3,
+            msg=f"Mismatch for {name}",
+        )
+
+
 def _run_grouped_linear_path(
     *,
     enable_grouped_tensor_path: bool,
@@ -1609,33 +2322,16 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy(
     fp8_recipe, bias, fp8_model_params, delay_wgrad_compute, monkeypatch
 ):
     use_fp8 = fp8_recipe is not None
-    device_capability = torch.cuda.get_device_capability()
-    if not (9, 0) <= device_capability <= (11, 0):
-        pytest.skip(
-            "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)."
-        )
-    # MXFP8/NVFP4 grouped quantization kernels require Blackwell; FP8 per-tensor
-    # current scaling runs on Hopper and Blackwell; FP8 block scaling is Hopper-only.
-    is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling()
-    is_block_scaling = use_fp8 and fp8_recipe.float8_block_scaling()
-    if is_block_scaling and not (9, 0) <= device_capability < (10, 0):
-        pytest.skip("Fused grouped FP8 block-scaling requires Hopper (SM90).")
-    if use_fp8 and not is_current_scaling and not is_block_scaling and device_capability < (10, 0):
-        pytest.skip(
-            "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)."
-        )
-    cublaslt_version = tex.get_cublasLt_version()
-    if device_capability < (10, 0) and cublaslt_version < 130400:
-        pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.")
-    if is_current_scaling and device_capability < (10, 0) and cublaslt_version < 130500:
-        pytest.skip("FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+.")
-    if cublaslt_version < 130300:
-        pytest.skip("Grouped GEMM requires cuBLAS 13.3+.")
+    dtype = torch.bfloat16
+    if not is_module_grouped_tensor_path_supported(
+        fp8_recipe,
+        dtype,
+    ):
+        pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.")
 
     if fp8_model_params and not use_fp8:
         pytest.skip("fp8_model_params requires FP8")
 
-    dtype = torch.bfloat16
     num_gemms = 3
     in_features = 128
     out_features = 128
@@ -1692,8 +2388,11 @@ def test_grouped_linear_grouped_tensor_path_matches_legacy(
 
 
 def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monkeypatch):
-    if torch.cuda.get_device_capability() < (10, 0):
-        pytest.skip("GroupedTensor grouped GEMM path requires SM100+")
+    if not is_module_grouped_tensor_path_supported(
+        None,
+        torch.bfloat16,
+    ):
+        pytest.skip("BF16 GroupedTensor path is unavailable on this system.")
 
     monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1")
 
@@ -1722,6 +2421,102 @@ def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monk
     grouped_linear.backward_dw()
 
 
+def test_grouped_linear_returns_single_grouped_bias_parameter(monkeypatch):
+    """return_bias preserves the grouped parent and accumulates dbias into it.
+
+    This mirrors how MCore applies a returned MoE bias::
+
+        x -> GroupedLinear (bias not applied) -> output
+                                                   +
+        grouped bias [2, 128]
+                 |
+                 +-> repeat_interleave([256, 256])
+                       -> per-token bias [512, 128]
+                              |
+                              * routing probabilities
+                              |
+                              +-> loss.backward() -> grouped bias.grad
+
+    Since the loss sums every output feature, each feature of expert ``i`` receives
+    ``sum(probs_for_expert_i)``. The identity assertion also ensures that TE returns the
+    registered grouped parent rather than copied or split bias tensors.
+    """
+    if not is_module_grouped_tensor_path_supported(
+        None,
+        torch.bfloat16,
+    ):
+        pytest.skip("BF16 GroupedTensor path is unavailable on this system.")
+    monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+
+    dtype = torch.bfloat16
+    num_gemms = 2
+    in_features = 128
+    out_features = 128
+    m_splits = torch.tensor([256, 256], dtype=torch.int64, device="cuda")
+    total_tokens = 512
+    grouped_linear = GroupedLinear(
+        num_gemms,
+        in_features,
+        out_features,
+        bias=True,
+        return_bias=True,
+        params_dtype=dtype,
+        device="cuda",
+        single_grouped_weight=False,
+        single_grouped_bias=True,
+        use_grouped_tensor=True,
+    )
+
+    x = torch.randn(
+        total_tokens,
+        in_features,
+        dtype=dtype,
+        device="cuda",
+        requires_grad=True,
+    )
+
+    probs = torch.cat(
+        (
+            torch.full((256,), 0.25, dtype=dtype, device="cuda"),
+            torch.full((256,), 0.5, dtype=dtype, device="cuda"),
+        )
+    )
+    output, returned_bias = grouped_linear(x, m_splits)
+
+    assert returned_bias is grouped_linear.bias
+    assert returned_bias.shape == (num_gemms, out_features)
+
+    bias_per_token = torch.repeat_interleave(
+        returned_bias,
+        m_splits,
+        dim=0,
+        output_size=total_tokens,
+    )
+    biased_output = (output + bias_per_token * probs.reshape(-1, 1)).to(output.dtype)
+    biased_output.sum().backward()
+
+    # For token t and output feature j, the externally applied bias contributes
+    # bias[expert(t), j] * probs[t] to the summed loss. Therefore every bias feature for
+    # expert e receives sum(probs[t]) over that expert's token range. m_splits places the
+    # first 256 tokens on expert 0 and the remaining 256 tokens on expert 1.
+    expected_dbias = (
+        torch.stack(
+            (probs[:256].float().sum(), probs[256:].float().sum()),
+        )
+        .unsqueeze(-1)
+        .expand(num_gemms, out_features)
+    )
+    assert grouped_linear.bias.grad is not None
+    # Megatron applies the packed bias in BF16. Compare its gradient against an independently
+    # accumulated FP32 reference with a tolerance appropriate for a 256-element BF16 reduction.
+    torch.testing.assert_close(
+        grouped_linear.bias.grad.float(),
+        expected_dbias,
+        rtol=5e-2,
+        atol=5e-3,
+    )
+
+
 @pytest.mark.parametrize("use_fused_path", [False, True], ids=["legacy", "grouped_tensor"])
 @pytest.mark.parametrize("supply", ["out", "dgrad_out", "both"])
 def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatch):
@@ -1730,18 +2525,11 @@ def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatc
     Checks that a supplied buffer is written in place (bit-for-bit vs internal allocation)
     and, on the fused path with a padded input, that only the valid rows are touched.
     """
-    if use_fused_path:
-        device_capability = torch.cuda.get_device_capability()
-        if not (9, 0) <= device_capability <= (11, 0):
-            pytest.skip(
-                "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell"
-                " (SM10x and SM110)."
-            )
-        cublaslt_version = tex.get_cublasLt_version()
-        if device_capability < (10, 0) and cublaslt_version < 130400:
-            pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.")
-        if cublaslt_version < 130300:
-            pytest.skip("Grouped GEMM requires cuBLAS 13.3+.")
+    if use_fused_path and not is_module_grouped_tensor_path_supported(
+        None,
+        torch.bfloat16,
+    ):
+        pytest.skip("BF16 GroupedTensor path is unavailable on this system.")
 
     monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1" if use_fused_path else "0")
     give_out = supply in ("out", "both")
@@ -1819,19 +2607,16 @@ def test_grouped_linear_caller_output_buffers(use_fused_path, supply, monkeypatc
 
 
 @pytest.mark.skipif(not _nvfp4_available, reason=_reason_for_no_nvfp4)
-def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch):
-    """Non-RHT NVFP4 falls back to the legacy path; check it stays numerically correct.
-
-    Graph-safe grouped quantization currently requires RHT, so requesting NVFP4 with
-    ``disable_rht=True`` while the fused grouped-tensor path is enabled falls back to the
-    legacy path internally. We verify the output and gradients against a reference built from
-    per-GEMM ``te.Linear`` modules that share the same weights and use the same NVFP4 recipe;
-    the grouped GEMM should match the loop of single GEMMs.
+def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4():
+    """Non-RHT NVFP4 falls back to split-quantize for discrete parameters.
+
+    Graph-safe grouped quantization currently requires RHT. Discrete parameters have a valid
+    split-quantize fallback, so enabling the grouped-tensor path is a preference rather than a
+    hard requirement for this parameter layout.
     """
     if torch.cuda.get_device_capability() < (10, 0):
         pytest.skip("NVFP4 GroupedTensor grouped GEMM path requires SM100+")
 
-    monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1")
     FP8GlobalStateManager.reset()
 
     dtype = torch.bfloat16
@@ -1854,7 +2639,6 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch):
         disable_stochastic_rounding=True,
     )
 
-    # Grouped path: fused path enabled, but non-RHT NVFP4 falls back to legacy internally.
     grouped_linear = GroupedLinear(
         num_gemms,
         in_features,
@@ -1862,6 +2646,7 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch):
         bias=False,
         params_dtype=dtype,
         device="cuda",
+        use_grouped_tensor=True,
     )
     with torch.no_grad():
         for i in range(num_gemms):
@@ -1872,7 +2657,6 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch):
         y = grouped_linear(x, m_splits)
     y.backward(dy)
 
-    # Reference: one te.Linear per GEMM sharing the same weights and NVFP4 recipe.
     ref_linears = torch.nn.ModuleList(
         [
             Linear(in_features, out_features, bias=False, params_dtype=dtype, device="cuda")
@@ -1890,7 +2674,6 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch):
         )
     y_ref.backward(dy)
 
-    # cuBLAS grouped GEMM should match the loop of single GEMMs bit-for-bit.
     tols = dict(rtol=0, atol=0)
     torch.testing.assert_close(y.float(), y_ref.float(), **tols)
     torch.testing.assert_close(x.grad.float(), x_ref.grad.float(), **tols)
@@ -1902,6 +2685,33 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch):
         )
 
 
+def test_grouped_linear_delay_wgrad_rejects_implicit_fallback(monkeypatch):
+    """Delayed wgrad reports when a grouped-tensor request used the legacy path."""
+    monkeypatch.setattr(
+        "transformer_engine.pytorch.module.grouped_linear.is_module_grouped_tensor_path_supported",
+        lambda *_args, **_kwargs: False,
+    )
+    grouped_linear = GroupedLinear(
+        2,
+        64,
+        64,
+        bias=True,
+        params_dtype=torch.bfloat16,
+        device="cuda",
+        delay_wgrad_compute=True,
+        use_grouped_tensor=True,
+    )
+    x = torch.randn(16, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+    m_splits = torch.tensor([8, 8], dtype=torch.int64, device="cuda")
+
+    grouped_linear(x, m_splits).sum().backward()
+    with pytest.raises(
+        RuntimeError,
+        match="implicit fallback is unsupported with delay_wgrad_compute=True",
+    ):
+        grouped_linear.backward_dw()
+
+
 @pytest.mark.parametrize(
     "fp8_recipe",
     [
@@ -1931,33 +2741,16 @@ def test_grouped_linear_grouped_tensor_path_skips_non_rht_nvfp4(monkeypatch):
 def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch):
     """Fused GroupedTensor GEMM path should be CUDA graph capturable."""
     use_fp8 = fp8_recipe is not None
-    device_capability = torch.cuda.get_device_capability()
-    if not (9, 0) <= device_capability <= (11, 0):
-        pytest.skip(
-            "GroupedTensor grouped GEMM path requires Hopper (SM90) or Blackwell (SM10x and SM110)."
-        )
-    # MXFP8/NVFP4 grouped quantization kernels require Blackwell; FP8 per-tensor
-    # current scaling runs on Hopper and Blackwell; FP8 block scaling is Hopper-only.
-    is_current_scaling = use_fp8 and fp8_recipe.float8_current_scaling()
-    is_block_scaling = use_fp8 and fp8_recipe.float8_block_scaling()
-    if is_block_scaling and not (9, 0) <= device_capability < (10, 0):
-        pytest.skip("Fused grouped FP8 block-scaling requires Hopper (SM90).")
-    if use_fp8 and not is_current_scaling and not is_block_scaling and device_capability < (10, 0):
-        pytest.skip(
-            "Quantized GroupedTensor grouped GEMM path (MXFP8/NVFP4) requires Blackwell (SM100+)."
-        )
-    cublaslt_version = tex.get_cublasLt_version()
-    if device_capability < (10, 0) and cublaslt_version < 130400:
-        pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.")
-    if is_current_scaling and device_capability < (10, 0) and cublaslt_version < 130500:
-        pytest.skip("FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+.")
-    if cublaslt_version < 130300:
-        pytest.skip("Grouped GEMM requires cuBLAS 13.3+.")
+    dtype = torch.bfloat16
+    if not is_module_grouped_tensor_path_supported(
+        fp8_recipe,
+        dtype,
+    ):
+        pytest.skip("Recipe is not supported by the module GroupedTensor path on this system.")
 
     monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1")
     FP8GlobalStateManager.reset()
 
-    dtype = torch.bfloat16
     device = "cuda"
     num_gemms = 3
     in_features = 128
diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py
index 499c22509c..c6d73cba63 100644
--- a/tests/pytorch/test_grouped_mlp.py
+++ b/tests/pytorch/test_grouped_mlp.py
@@ -24,6 +24,7 @@
 from transformer_engine.pytorch.ops.basic.grouped_linear import (
     OUTPUT_BUFFER_KEY,
     GRAD_INPUT_BUFFER_KEY,
+    is_op_fuser_grouped_tensor_path_supported,
 )
 from transformer_engine.pytorch import (
     QuantizedTensor,
@@ -278,6 +279,26 @@ def test_meta_single_grouped_weight_with_delayed_wgrad(self, monkeypatch) -> Non
 
         assert grouped_weight.skip_backward_post_hook
 
+    def test_single_grouped_bias_uses_registered_packed_storage(self, monkeypatch) -> None:
+        """The grouped bias compute view must alias the registered trainable parent."""
+        monkeypatch.setenv("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "1")
+        op = te.ops.GroupedLinear(
+            2,
+            128,
+            128,
+            bias=True,
+            device="cuda",
+            dtype=torch.bfloat16,
+            single_grouped_bias=True,
+        )
+
+        bias_packed = op._get_packed_bias_tensor(torch.bfloat16)
+
+        assert bias_packed.shape == (op.num_groups, op.out_features)
+        assert bias_packed.untyped_storage().data_ptr() == op.bias.rowwise_data.data_ptr()
+        assert op.bias.requires_grad
+        assert dict(op.named_parameters())["bias"] is op.bias
+
     @pytest.mark.parametrize("bias", (False, True))
     @pytest.mark.parametrize("dtype", _dtypes)
     @pytest.mark.parametrize("quantization", _quantization_list)
@@ -338,6 +359,16 @@ def test_grouped_linear(
 
         if single_grouped_bias and not bias:
             pytest.skip("single_grouped_bias requires bias=True")
+        recipe = make_recipe(quantization)
+        compute_recipe = recipe if quantized_compute else None
+        if (
+            single_grouped_weight or single_grouped_bias
+        ) and not is_op_fuser_grouped_tensor_path_supported(
+            compute_recipe,
+            dtype,
+        ):
+            # Single grouped parameters intentionally have no split-quantize fallback.
+            pytest.skip("Single grouped parameters require the native grouped-tensor path")
         if single_grouped_weight and quantized_weight and quantization in ("fp8_delayed_scaling"):
             pytest.skip(
                 "single_grouped_weight does not support FP8 delayed scaling "
@@ -393,7 +424,6 @@ def test_grouped_linear(
             y_ref.backward(dy_ref)
 
         # Construct fusible operation
-        recipe = make_recipe(quantization)
         with te.quantized_model_init(enabled=quantized_weight, recipe=recipe):
             op = te.ops.GroupedLinear(
                 group_size,
@@ -707,22 +737,6 @@ def test_grouped_linear_cuda_graph_safe(
                 "single_grouped_weight/single_grouped_bias requires"
                 " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1"
             )
-        device_capability = torch.cuda.get_device_capability()
-        if device_capability < (9, 0):
-            pytest.skip(
-                "Grouped GEMM CUDA-graph-safe path requires Hopper (SM90) or Blackwell (SM100+)"
-            )
-        # BF16/FP16 and FP8 per-tensor current scaling run on the Hopper grouped GEMM path,
-        # but MXFP8/NVFP4 grouped quantization kernels require Blackwell (SM100+).
-        requires_blackwell = quantization is not None and quantization != "fp8_current_scaling"
-        if requires_blackwell and device_capability < (10, 0):
-            pytest.skip("MXFP8/NVFP4 grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)")
-        # Grouped GEMM on Hopper requires cuBLAS 13.4+; Blackwell requires cuBLAS 13.3+.
-        cublaslt_version = tex.get_cublasLt_version()
-        if device_capability < (10, 0) and cublaslt_version < 130400:
-            pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.")
-        if cublaslt_version < 130300:
-            pytest.skip("Grouped GEMM requires cuBLAS 13.3+.")
         if quantization is None and quantized_weight:
             pytest.skip("quantized_weight requires a quantization recipe")
         if (
@@ -740,6 +754,13 @@ def test_grouped_linear_cuda_graph_safe(
                 "only discrete weights (single_grouped_weight=False) are supported."
             )
 
+        recipe = make_recipe(quantization)
+        if not is_op_fuser_grouped_tensor_path_supported(
+            recipe,
+            dtype,
+        ):
+            pytest.skip("Configuration falls back to the non-CUDA-graph-safe path")
+
         single_grouped_bias = bias and single_grouped_weight
 
         # Split sizes (statically pinned for graph capture)
@@ -752,7 +773,6 @@ def test_grouped_linear_cuda_graph_safe(
         in_shape = (num_active_tokens + token_padding, in_features)
         out_shape = (in_shape[0], out_features)
 
-        recipe = make_recipe(quantization)
         with te.quantized_model_init(enabled=quantized_weight, recipe=recipe):
             op = te.ops.GroupedLinear(
                 group_size,
diff --git a/tests/pytorch/test_grouped_tensor.py b/tests/pytorch/test_grouped_tensor.py
index 4dd52ee2bd..ac8c493290 100644
--- a/tests/pytorch/test_grouped_tensor.py
+++ b/tests/pytorch/test_grouped_tensor.py
@@ -577,7 +577,7 @@ def test_quantize_grouped_mxfp8(self, shape_case: str, output_dbias: bool) -> No
 
         if output_dbias:
             expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors])
-            assert torch.allclose(dbias, expected_dbias)
+            torch.testing.assert_close(dbias, expected_dbias, rtol=1e-5, atol=4e-3)
 
     @pytest.mark.parametrize("output_dbias", [False, True])
     @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
@@ -805,6 +805,152 @@ def _run(inp):
         assert torch.equal(static_output.rowwise_data, expected.rowwise_data)
         assert torch.equal(static_output.scale_inv, expected.scale_inv)
 
+    @pytest.mark.parametrize(
+        "quantization",
+        [
+            pytest.param(
+                "fp8_current_scaling",
+                marks=pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8),
+            ),
+            pytest.param(
+                "fp8_blockwise",
+                marks=pytest.mark.skipif(
+                    not fp8_block_scaling_grouped_available,
+                    reason=reason_for_no_fp8_block_scaling_grouped,
+                ),
+            ),
+        ],
+    )
+    def test_group_quantize_reuses_destination_and_honors_noop_fp8(self, quantization: str) -> None:
+        """Check pointer-stable weight-cache updates and no-op CUDA graph replays."""
+        num_tensors = 2
+        shape = (num_tensors * 256, 256)
+        if quantization == "fp8_current_scaling":
+            quantizer = Float8CurrentScalingQuantizer(
+                fp8_dtype=tex.DType.kFloat8E4M3,
+                device="cuda",
+                force_pow_2_scales=False,
+                amax_epsilon=0.0,
+            )
+            quantizer.set_usage(rowwise=True, columnwise=True)
+        else:
+            quantizer = Float8BlockQuantizer(
+                fp8_dtype=tex.DType.kFloat8E4M3,
+                rowwise=True,
+                columnwise=True,
+                force_pow_2_scales=False,
+                amax_epsilon=0.0,
+                block_scaling_dim=1,
+            )
+
+        source = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+        output = tex.group_quantize(source, quantizer, num_tensors, None)
+        output_buffers = tuple(
+            buffer
+            for buffer in (
+                output.rowwise_data,
+                output.columnwise_data,
+                output.scale_inv,
+                output.columnwise_scale_inv,
+                output.amax,
+                output.columnwise_amax,
+                output.scale,
+            )
+            if buffer is not None
+        )
+        assert output_buffers
+        pointers = tuple(buffer.data_ptr() for buffer in output_buffers)
+
+        # Re-quantization must update the existing allocations because a captured GEMM keeps
+        # their raw addresses rather than looking them up again through the Python object.
+        updated_source = source + 1
+        updated = tex.group_quantize(
+            updated_source,
+            quantizer,
+            num_tensors,
+            None,
+            output=output,
+        )
+        reference = tex.group_quantize(updated_source, quantizer, num_tensors, None)
+        assert updated is output
+        assert pointers == tuple(buffer.data_ptr() for buffer in output_buffers)
+        reference_buffers = tuple(
+            buffer
+            for buffer in (
+                reference.rowwise_data,
+                reference.columnwise_data,
+                reference.scale_inv,
+                reference.columnwise_scale_inv,
+                reference.amax,
+                reference.columnwise_amax,
+                reference.scale,
+            )
+            if buffer is not None
+        )
+        assert len(output_buffers) == len(reference_buffers)
+        for output_buffer, reference_buffer in zip(output_buffers, reference_buffers, strict=True):
+            assert torch.equal(output_buffer, reference_buffer)
+
+        # A nonzero device flag must preserve every cached payload and metadata buffer.
+        old_buffers = tuple(buffer.clone() for buffer in output_buffers)
+        noop = torch.ones(1, dtype=torch.float32, device="cuda")
+        tex.group_quantize(
+            updated_source * 2,
+            quantizer,
+            num_tensors,
+            None,
+            noop_flag=noop,
+            output=output,
+        )
+        for output_buffer, old_buffer in zip(output_buffers, old_buffers, strict=True):
+            assert torch.equal(output_buffer, old_buffer)
+
+        # Capture one pointer-stable update. Replays change only source contents and the device
+        # no-op value, matching TE's CUDA graph microbatch weight-cache protocol.
+        graph_source = updated_source.clone()
+        noop.zero_()
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            tex.group_quantize(
+                graph_source,
+                quantizer,
+                num_tensors,
+                None,
+                noop_flag=noop,
+                output=output,
+            )
+
+        graph_source.copy_(source - 1)
+        noop.zero_()
+        graph.replay()
+        torch.cuda.synchronize()
+        reference = tex.group_quantize(graph_source, quantizer, num_tensors, None)
+        reference_buffers = tuple(
+            buffer
+            for buffer in (
+                reference.rowwise_data,
+                reference.columnwise_data,
+                reference.scale_inv,
+                reference.columnwise_scale_inv,
+                reference.amax,
+                reference.columnwise_amax,
+                reference.scale,
+            )
+            if buffer is not None
+        )
+        assert len(output_buffers) == len(reference_buffers)
+        for output_buffer, reference_buffer in zip(output_buffers, reference_buffers, strict=True):
+            assert torch.equal(output_buffer, reference_buffer)
+
+        old_buffers = tuple(buffer.clone() for buffer in output_buffers)
+        graph_source.copy_(source + 2)
+        noop.fill_(1)
+        graph.replay()
+        torch.cuda.synchronize()
+        for output_buffer, old_buffer in zip(output_buffers, old_buffers, strict=True):
+            assert torch.equal(output_buffer, old_buffer)
+
     @pytest.mark.parametrize("mode", ["rowwise", "columnwise", "both"])
     @pytest.mark.parametrize(
         "shape_case",
@@ -1121,7 +1267,7 @@ def test_quantize_grouped_fp8_blockwise(
 
         if output_dbias:
             expected_dbias = torch.stack([t.sum(dim=0) for t in input_tensors])
-            assert torch.allclose(dbias, expected_dbias)
+            torch.testing.assert_close(dbias, expected_dbias, rtol=1e-5, atol=4e-3)
 
     @pytest.mark.parametrize(
         "shape",
@@ -1204,6 +1350,147 @@ def test_group_dequantize_cudagraph_capturable(self) -> None:
         for exp, got in zip(expected_tensors, static_tensors):
             assert torch.equal(got, exp)
 
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    def test_group_quantize_reuses_destination_and_honors_noop(self) -> None:
+        """Verify the in-place and graph-safe contracts of grouped MXFP8 quantization.
+
+        A captured training graph records the addresses of all MXFP8 weight storage, not just
+        the Python ``GroupedTensor`` object. Re-quantizing an updated BF16 weight must therefore
+        refresh the existing FP8 payloads and scales without replacing any allocation. During
+        graph replay, ``noop_flag`` must additionally allow the captured quantization kernel to
+        execute as a no-op on microbatches that should reuse the cached weight.
+        """
+        num_tensors = 2
+        # Treat two contiguous [256, 256] matrices as one grouped BF16 source. These dimensions
+        # satisfy the MXFP8 block-alignment requirements in both rowwise and columnwise layouts.
+        shape = (num_tensors * 256, 256)
+        quantizer = MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)
+
+        # Forward GEMM consumes rowwise weight storage, while backward dgrad consumes columnwise
+        # storage. Both representations and both scale tensors must be updated and kept stable.
+        quantizer.set_usage(rowwise=True, columnwise=True)
+        # Store the quantized output in GEMM-ready (swizzled) scale layout. This makes the test
+        # cover the exact destination format used by cached MXFP8 model weights.
+        quantizer.optimize_for_gemm = True
+
+        # The first call owns allocation: it creates the persistent destination that subsequent
+        # optimizer updates and CUDA graph replays must modify in place.
+        source = torch.randn(shape, dtype=torch.bfloat16, device="cuda")
+        output = tex.group_quantize(source, quantizer, num_tensors, None)
+
+        # CUDA graphs retain these raw addresses. Object identity alone is insufficient because
+        # replacing one internal payload or scale allocation would leave the graph with a stale
+        # pointer even if ``output`` remained the same Python object.
+        pointers = (
+            output.rowwise_data.data_ptr(),
+            output.columnwise_data.data_ptr(),
+            output.scale_inv.data_ptr(),
+            output.columnwise_scale_inv.data_ptr(),
+        )
+
+        # Eager in-place update: quantize new BF16 values into the previously allocated output.
+        # An independently allocated quantization provides the numerical reference.
+        updated_source = source + 1
+        updated = tex.group_quantize(
+            updated_source,
+            quantizer,
+            num_tensors,
+            None,
+            output=output,
+        )
+        reference = tex.group_quantize(updated_source, quantizer, num_tensors, None)
+
+        # The API must return the caller-provided destination and preserve every captured address.
+        assert updated is output
+        assert pointers == (
+            output.rowwise_data.data_ptr(),
+            output.columnwise_data.data_ptr(),
+            output.scale_inv.data_ptr(),
+            output.columnwise_scale_inv.data_ptr(),
+        )
+
+        # In-place quantization must still produce exactly the same FP8 bytes and scales as a
+        # normal allocating call, for both GEMM orientations.
+        assert torch.equal(output.rowwise_data, reference.rowwise_data)
+        assert torch.equal(output.columnwise_data, reference.columnwise_data)
+        assert torch.equal(output.scale_inv, reference.scale_inv)
+        assert torch.equal(output.columnwise_scale_inv, reference.columnwise_scale_inv)
+
+        # Eager no-op: a nonzero device flag tells the kernel to preserve the cached destination.
+        # Clone all four buffers so this checks payloads and metadata independently.
+        old_buffers = (
+            output.rowwise_data.clone(),
+            output.columnwise_data.clone(),
+            output.scale_inv.clone(),
+            output.columnwise_scale_inv.clone(),
+        )
+        noop = torch.ones(1, dtype=torch.float32, device="cuda")
+
+        # The source is deliberately different. If the kernel ignores ``noop_flag``, at least
+        # one payload or scale tensor below will change and expose the failure.
+        tex.group_quantize(
+            updated_source * 2,
+            quantizer,
+            num_tensors,
+            None,
+            noop_flag=noop,
+            output=output,
+        )
+        assert torch.equal(output.rowwise_data, old_buffers[0])
+        assert torch.equal(output.columnwise_data, old_buffers[1])
+        assert torch.equal(output.scale_inv, old_buffers[2])
+        assert torch.equal(output.columnwise_scale_inv, old_buffers[3])
+
+        # Capture one fixed launch. ``graph_source``, ``noop``, and ``output`` keep stable device
+        # addresses; replay changes only their contents. This is how weight caching works when
+        # different microbatches replay the same CUDA graph.
+        graph_source = updated_source.clone()
+        # zero means "perform quantization"; nonzero means "leave destination untouched".
+        noop.zero_()
+        torch.cuda.synchronize()
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            tex.group_quantize(
+                graph_source,
+                quantizer,
+                num_tensors,
+                None,
+                noop_flag=noop,
+                output=output,
+            )
+
+        # Replay with new source contents and noop=0. The captured kernel must refresh the same
+        # destination allocations, and the result must equal a fresh eager quantization.
+        graph_source.copy_(source - 1)
+        noop.zero_()
+        graph.replay()
+        torch.cuda.synchronize()
+        reference = tex.group_quantize(graph_source, quantizer, num_tensors, None)
+        assert torch.equal(output.rowwise_data, reference.rowwise_data)
+        assert torch.equal(output.columnwise_data, reference.columnwise_data)
+        assert torch.equal(output.scale_inv, reference.scale_inv)
+        assert torch.equal(output.columnwise_scale_inv, reference.columnwise_scale_inv)
+
+        # Snapshot the successfully refreshed cache before testing the captured no-op branch.
+        old_buffers = (
+            output.rowwise_data.clone(),
+            output.columnwise_data.clone(),
+            output.scale_inv.clone(),
+            output.columnwise_scale_inv.clone(),
+        )
+
+        # Replay the exact same captured graph with different source data but noop=1. The launch
+        # still occurs, which is required for CUDA graph structural consistency, but the kernel
+        # must not mutate any cached FP8 payload or scale buffer.
+        graph_source.copy_(source + 2)
+        noop.fill_(1)
+        graph.replay()
+        torch.cuda.synchronize()
+        assert torch.equal(output.rowwise_data, old_buffers[0])
+        assert torch.equal(output.columnwise_data, old_buffers[1])
+        assert torch.equal(output.scale_inv, old_buffers[2])
+        assert torch.equal(output.columnwise_scale_inv, old_buffers[3])
+
     @pytest.mark.parametrize("block_scaling_dim", [1, 2], ids=["1D", "2D"])
     @pytest.mark.parametrize("direction", ["rowwise", "columnwise"])
     @pytest.mark.skipif(
diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py
index de62e56e53..c9b620fa1e 100644
--- a/tests/pytorch/test_sanity.py
+++ b/tests/pytorch/test_sanity.py
@@ -37,6 +37,7 @@
 from transformer_engine.common import recipe
 from transformer_engine.pytorch.cpp_extensions import general_gemm
 from transformer_engine.pytorch.tensor.utils import replace_raw_data
+from transformer_engine.pytorch.module import is_module_grouped_tensor_path_supported
 from utils import ModelConfig, recipe_id, skip_unsupported_backward_override
 
 # Only run FP8 tests on supported devices.
@@ -86,11 +87,11 @@ def is_fp8_supported(config: ModelConfig):
 
 
 def nvfp4_vanilla():
-    nvfp4_recipe = recipe.NVFP4BlockScaling()
-    nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams()
-    nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams()
-    nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams()
-    return nvfp4_recipe
+    return recipe.NVFP4BlockScaling(
+        disable_rht=True,
+        disable_stochastic_rounding=True,
+        disable_2d_quantization=True,
+    )
 
 
 def nvfp4_row_scaled():
@@ -123,7 +124,7 @@ def nvfp4_4over6():
 if mxfp8_available:
     fp8_recipes.append(recipe.MXFP8BlockScaling())
 if nvfp4_available:
-    fp8_recipes.append(nvfp4_vanilla())  # TODO: fix check for this
+    fp8_recipes.append(nvfp4_vanilla())
     fp8_recipes.append(nvfp4_4over6())
 if fp8_block_scaling_available:
     fp8_recipes.append(recipe.Float8BlockScaling())
@@ -563,6 +564,7 @@ def test_sanity_linear_with_zero_tokens(
         out = te_linear(inp_hidden_states)
     loss = out.sum()
     loss.backward()
+    torch.cuda.synchronize()
     assert out.shape == (num_tokens, ffn_hidden_size)
 
 
@@ -603,17 +605,34 @@ def test_sanity_grouped_linear(
     if fp8_recipe is not None:
         fp8_recipe = copy.deepcopy(fp8_recipe)
         fp8_recipe.backward_override = backward_override
+    if single_param and not is_module_grouped_tensor_path_supported(
+        fp8_recipe,
+        dtype,
+    ):
+        pytest.skip("Single grouped parameters require the native grouped-tensor path")
+    if single_param:
+        # Single grouped parameters intentionally have no split-quantize fallback, so this
+        # test must satisfy the native grouped kernels' shape contract. MCore pads each
+        # expert's token count to 256; TE requires at least 128-row alignment. Weight K must
+        # be 64-aligned.
+        tokens_per_nonempty_expert = bs * config.max_seqlen_q
+        if tokens_per_nonempty_expert % 128 != 0:
+            pytest.skip("Single grouped parameters require each nonempty m_split to be 128-aligned")
+        k_alignment = 64
+        if config.hidden_size % k_alignment != 0:
+            pytest.skip(f"Single grouped parameters require GEMM K to be {k_alignment}-aligned")
 
     if fp8_recipe is not None:
         if not is_fp8_supported(config):
             pytest.skip("Model config does not support FP8")
         if fp8_recipe.nvfp4():
-            if not getattr(fp8_recipe, "row_scaled_activation", False):
-                pytest.skip("NVFP4 not supported for grouped linear")
-            if single_param:
-                pytest.skip("Row-scaled NVFP4 does not support GroupedTensor grouped linear")
-            if dtype == torch.float16:
-                pytest.skip("FP16 output for NVFP4 not supported")
+            if dtype != torch.bfloat16:
+                pytest.skip("NVFP4 GroupedLinear requires BF16")
+            if single_param and not fp8_model_params:
+                pytest.skip(
+                    "NVFP4 single grouped BF16 primary weights require unsupported non-RHT "
+                    "grouped weight quantization; enable quantized model initialization"
+                )
 
     use_fp8 = fp8_recipe is not None
     with quantized_model_init(enabled=use_fp8 and fp8_model_params, recipe=fp8_recipe):
@@ -625,6 +644,7 @@ def test_sanity_grouped_linear(
             params_dtype=dtype,
             single_grouped_weight=single_param,
             single_grouped_bias=single_param,
+            use_grouped_tensor=single_param,
         ).cuda()
 
     # Verify grouped linear exposes a single grouped weight parameter(and bias when applicable).
@@ -644,11 +664,25 @@ def test_sanity_grouped_linear(
         m_splits[-1] = 0
     elif empty_split == "middle":
         m_splits[num_gemms // 2] = 0
+    if single_param:
+        m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cuda")
+
+    if NVTE_TEST_NVINSPECT_ENABLED and single_param:
+        # DebugQuantizer operates on per-GEMM tensors, while single grouped parameters
+        # intentionally have no split-quantize fallback.
+        with pytest.raises(
+            RuntimeError,
+            match="TE debug features do not support single grouped parameters",
+        ):
+            with autocast(enabled=use_fp8, recipe=fp8_recipe):
+                te_grouped_linear(inp_hidden_states, m_splits)
+        return
 
     with autocast(enabled=use_fp8, recipe=fp8_recipe):
         out = te_grouped_linear(inp_hidden_states, m_splits)
     loss = out.sum()
     loss.backward()
+    torch.cuda.synchronize()
     assert out.shape == (num_tokens, ffn_hidden_size)
 
 
diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh
index 14832573d7..bc00c5824f 100644
--- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh
@@ -572,9 +572,18 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel
     const size_t scale_stride_colwise = DIVUP_TO_MULTIPLE(cols, scale_alignment_X_colwise);
 
     const size_t tensor_base = current_block.tensor_base;
-    const size_t tensor_base_for_scales = (is_single_tensor && num_tensors > 1)
-                                              ? static_cast<size_t>(offsets_ptr[tensor_id])
-                                              : tensor_base;
+    size_t tensor_base_for_scales = tensor_base;
+    size_t tensor_rows_for_scales = rows;
+    if constexpr (WITH_GEMM_SWIZZLED_SCALES && SHAPE_REP == ShapeRepresentation::SAME_BOTH_DIMS) {
+      // The payload is one tall tensor, but GEMM scales are swizzled independently per member.
+      // Uniform groups omit first_dims/tensor_offsets, so derive the member geometry directly.
+      tensor_rows_for_scales = first_logical_dim / num_tensors;
+      tensor_base_for_scales = tensor_id * tensor_rows_for_scales * cols;
+    }
+    if constexpr (WITH_GEMM_SWIZZLED_SCALES &&
+                  SHAPE_REP == ShapeRepresentation::VARYING_FIRST_DIM) {
+      tensor_base_for_scales = static_cast<size_t>(offsets_ptr[tensor_id]);
+    }
     const size_t block_id_Y = current_block.block_id_Y;
     const size_t block_id_X = current_block.block_id_X;
     const size_t block_offset_Y = current_block.block_offset_Y;
@@ -680,8 +689,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel
         process_colwise_stage<IS_DBIAS, IS_DACT, IS_ACT, ParamOP, OP, IType, OType, ROWWISE_SCALING,
                               WITH_GEMM_SWIZZLED_SCALES>(
             buff, stage, tid_X_colwise, scales_offset_Y_colwise, scales_offset_X_colwise,
-            scale_stride_colwise, tensor_base_for_scales, rows, cols, sIn_ptr, sActIn_ptr,
-            sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise);
+            scale_stride_colwise, tensor_base_for_scales, tensor_rows_for_scales, cols, sIn_ptr,
+            sActIn_ptr, sCachedAct_ptr, sOutColwise_ptr, scales_colwise, partial_dbias_colwise);
       }
 
       if constexpr (ROWWISE_SCALING) {
diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu
index 3997e5249d..3663106d50 100644
--- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu
+++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu
@@ -1324,7 +1324,8 @@ __global__ void setup_grouped_gemm_kernel(
     char *a_base, char *b_base, char *c_base, char *d_base, TensorShapeInfo A_meta,
     TensorShapeInfo B_meta, TensorShapeInfo C_meta, TensorShapeInfo D_meta, size_t a_bits_per_elem,
     size_t b_bits_per_elem, size_t c_elem_size, size_t d_elem_size, float *alpha_ptr,
-    float *beta_ptr, bool use_per_group_alpha_beta,
+    float *beta_ptr, bool use_per_group_alpha_beta, bool a_is_discrete, bool c_is_discrete,
+    bool d_is_discrete,
     // Scale inputs: for tensor scaling, pass float* and set mxfp8_base to nullptr
     // For MXFP8, pass nullptr for tensor_scale and set mxfp8_base
     float *a_scale_base, float *b_scale_base, bool a_rowwise, bool b_rowwise,
@@ -1339,12 +1340,9 @@ __global__ void setup_grouped_gemm_kernel(
   if (idx >= num_tensors) return;
 
   // Get dimensions for this tensor (from array or uniform value)
-  const bool has_a_multi_tensor = (a_base == nullptr);
-  const bool has_c_multi_tensor = (c_base == nullptr);
-  const bool has_d_multi_tensor = (d_base == nullptr);
   int64_t a_first = 0;
   int64_t a_last = 0;
-  if (!has_a_multi_tensor) {
+  if (!a_is_discrete) {
     a_first = A_meta.first_dims ? A_meta.first_dims[idx] : A_meta.uniform_first;
     a_last = A_meta.last_dims ? A_meta.last_dims[idx] : A_meta.uniform_last;
   }
@@ -1354,24 +1352,25 @@ __global__ void setup_grouped_gemm_kernel(
   int64_t d_last = D_meta.last_dims ? D_meta.last_dims[idx] : D_meta.uniform_last;
 
   // Compute offsets (from explicit array, cumulative from per-tensor dims, or uniform)
-  int64_t a_offset = has_a_multi_tensor ? 0 : compute_grouped_tensor_offset(A_meta, idx);
+  int64_t a_offset = a_is_discrete ? 0 : compute_grouped_tensor_offset(A_meta, idx);
   int64_t b_offset = compute_grouped_tensor_offset(B_meta, idx);
   int64_t c_offset = compute_grouped_tensor_offset(C_meta, idx);
   int64_t d_offset = compute_grouped_tensor_offset(D_meta, idx);
 
   // Compute data pointers
-  A_ptrs[idx] = has_a_multi_tensor ? a_multi_tensor_args.data_ptrs[idx]
-                                   : (a_base + (a_offset * a_bits_per_elem) / 8);
-  B_ptrs[idx] = b_base + (b_offset * b_bits_per_elem) / 8;
-  C_ptrs[idx] =
-      has_c_multi_tensor ? c_multi_tensor_args.data_ptrs[idx] : (c_base + c_offset * c_elem_size);
-  D_ptrs[idx] =
-      has_d_multi_tensor ? d_multi_tensor_args.data_ptrs[idx] : (d_base + d_offset * d_elem_size);
+  A_ptrs[idx] = a_is_discrete
+                    ? a_multi_tensor_args.data_ptrs[idx]
+                    : (a_base == nullptr ? nullptr : a_base + (a_offset * a_bits_per_elem) / 8);
+  B_ptrs[idx] = b_base == nullptr ? nullptr : b_base + (b_offset * b_bits_per_elem) / 8;
+  C_ptrs[idx] = c_is_discrete ? c_multi_tensor_args.data_ptrs[idx]
+                              : (c_base == nullptr ? nullptr : c_base + c_offset * c_elem_size);
+  D_ptrs[idx] = d_is_discrete ? d_multi_tensor_args.data_ptrs[idx]
+                              : (d_base == nullptr ? nullptr : d_base + d_offset * d_elem_size);
 
   // Compute storage dimensions for cuBLAS matrix layouts from logical dims.
   // Rowwise and MXFP8 columnwise storage use logical row-major layout, viewed as
   // column-major rows=last, cols=first. Transposed columnwise storage reverses this.
-  if (has_a_multi_tensor) {
+  if (a_is_discrete) {
     a_rows[idx] = a_multi_tensor_args.rows[idx];
     a_cols[idx] = a_multi_tensor_args.cols[idx];
   } else if (a_storage_transposed) {
@@ -1388,7 +1387,7 @@ __global__ void setup_grouped_gemm_kernel(
     b_rows[idx] = static_cast<int>(b_last);
     b_cols[idx] = static_cast<int>(b_first);
   }
-  if (has_d_multi_tensor) {
+  if (d_is_discrete) {
     d_rows[idx] = d_multi_tensor_args.rows[idx];
     d_cols[idx] = d_multi_tensor_args.cols[idx];
   } else {
@@ -1403,7 +1402,7 @@ __global__ void setup_grouped_gemm_kernel(
   if (use_per_group_alpha_beta) {
     float a_amax_val = 0.0f;
     bool has_a_amax = false;
-    if (has_a_multi_tensor) {
+    if (a_is_discrete) {
       auto *a_amax_p = static_cast<float *>(a_multi_tensor_args.amax_ptrs[idx]);
       if (a_amax_p != nullptr) {
         a_amax_val = *a_amax_p;
@@ -1471,10 +1470,12 @@ __global__ void setup_grouped_gemm_kernel(
     }
   };
 
-  if (a_scale_base) {
+  if (a_is_discrete) {
+    a_scale_inv_ptrs[idx] = a_multi_tensor_args.scale_inv_ptrs[idx];
+  } else if (a_scale_base) {
     fill_scale_ptr(a_scale_inv_ptrs, a_scale_base, A_meta, a_rowwise, a_scaling_mode);
   } else {
-    a_scale_inv_ptrs[idx] = a_multi_tensor_args.scale_inv_ptrs[idx];
+    a_scale_inv_ptrs[idx] = nullptr;
   }
   if (b_scale_base) {
     fill_scale_ptr(b_scale_inv_ptrs, b_scale_base, B_meta, b_rowwise, b_scaling_mode);
@@ -1491,7 +1492,7 @@ inline void launch_grouped_gemm_setup(
     const transformer_engine::Tensor *beta_tensor, bool use_per_group_alpha_beta,
     size_t num_tensors, cudaStream_t stream,
     const MultiTensorGroupGemmInputArgs &a_multi_tensor_args, const NVTETensor *C_list,
-    const NVTETensor *D_list, char *a_base, transformer_engine::DType c_dtype,
+    const NVTETensor *D_list, bool a_is_discrete, char *a_base, transformer_engine::DType c_dtype,
     transformer_engine::DType d_dtype) {
   // Use logical shape info from selection; storage transposes are tracked separately.
   TensorShapeInfo A_meta = A_sel.logical_tensor_shape;
@@ -1499,31 +1500,31 @@ inline void launch_grouped_gemm_setup(
   TensorShapeInfo C_meta{};
   TensorShapeInfo D_meta{};
 
-  const bool has_d_multi_tensor = (D_list != nullptr);
-  const bool has_c_multi_tensor = (C_list != nullptr) || has_d_multi_tensor;
+  const bool d_is_discrete = (D_list != nullptr);
+  const bool c_is_discrete = (C_list != nullptr) || d_is_discrete;
   MultiTensorGroupGemmOutputArgs c_multi_tensor_args{};
   MultiTensorGroupGemmOutputArgs d_multi_tensor_args{};
-  if (has_d_multi_tensor) {
+  if (d_is_discrete) {
     d_multi_tensor_args =
         build_grouped_gemm_multi_out_args(D_list, num_tensors, num_tensors, d_dtype, "D");
   }
   if (C_list != nullptr) {
     c_multi_tensor_args =
         build_grouped_gemm_multi_out_args(C_list, num_tensors, num_tensors, d_dtype, "C");
-  } else if (has_d_multi_tensor) {
+  } else if (d_is_discrete) {
     c_multi_tensor_args = d_multi_tensor_args;
   }
 
   char *c_base = nullptr;
   char *d_base = nullptr;
 
-  if (!has_c_multi_tensor) {
+  if (!c_is_discrete) {
     NVTE_CHECK(C != nullptr && D != nullptr,
                "Grouped GEMM: C/D grouped tensors are required when no C list is provided");
     C_meta = TensorShapeInfo::create_shape_info_for_C(C, D);
     c_base = static_cast<char *>(C->data.dptr);
   }
-  if (!has_d_multi_tensor) {
+  if (!d_is_discrete) {
     NVTE_CHECK(D != nullptr,
                "Grouped GEMM: D grouped tensor is required when no D list is provided");
     D_meta = TensorShapeInfo::from_tensor(D);
@@ -1544,8 +1545,8 @@ inline void launch_grouped_gemm_setup(
   const bool b_rowwise = B_sel.rowwise;
 
   // NVFP4 alpha needs A's amax from either A_sel.amax (grouped) or amax_ptrs (discrete).
-  const bool a_has_amax = (A_sel.amax != nullptr) ||
-                          (A_sel.dptr == nullptr && a_multi_tensor_args.amax_ptrs[0] != nullptr);
+  const bool a_has_amax =
+      a_is_discrete ? a_multi_tensor_args.amax_ptrs[0] != nullptr : A_sel.amax != nullptr;
   const bool needs_nvfp4_alpha = transformer_engine::is_nvfp_scaling(A_sel.scaling_mode) &&
                                  a_has_amax && (B_sel.amax != nullptr);
 
@@ -1554,11 +1555,12 @@ inline void launch_grouped_gemm_setup(
       ws.d_rows, ws.d_cols, ws.alpha_ptrs, ws.beta_ptrs, ws.a_scale_inv_ptrs, ws.b_scale_inv_ptrs,
       A_sel.dptr, B_sel.dptr, c_base, d_base, A_meta, B_meta, C_meta, D_meta, a_bits_per_elem,
       b_bits_per_elem, c_elem_size, d_elem_size, static_cast<float *>(alpha_tensor->data.dptr),
-      static_cast<float *>(beta_tensor->data.dptr), use_per_group_alpha_beta,
-      reinterpret_cast<float *>(A_sel.scale_inv), reinterpret_cast<float *>(B_sel.scale_inv),
-      a_rowwise, b_rowwise, A_sel.storage_transposed, B_sel.storage_transposed, A_sel.scaling_mode,
-      B_sel.scaling_mode, num_tensors, a_multi_tensor_args, c_multi_tensor_args,
-      d_multi_tensor_args, A_sel.amax ? static_cast<float *>(A_sel.amax) : nullptr,
+      static_cast<float *>(beta_tensor->data.dptr), use_per_group_alpha_beta, a_is_discrete,
+      c_is_discrete, d_is_discrete, reinterpret_cast<float *>(A_sel.scale_inv),
+      reinterpret_cast<float *>(B_sel.scale_inv), a_rowwise, b_rowwise, A_sel.storage_transposed,
+      B_sel.storage_transposed, A_sel.scaling_mode, B_sel.scaling_mode, num_tensors,
+      a_multi_tensor_args, c_multi_tensor_args, d_multi_tensor_args,
+      A_sel.amax ? static_cast<float *>(A_sel.amax) : nullptr,
       B_sel.amax ? static_cast<float *>(B_sel.amax) : nullptr,
       needs_nvfp4_alpha ? ws.nvfp4_computed_alpha : nullptr);
 
@@ -1633,8 +1635,8 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT
   MultiTensorGroupGemmInputArgs a_multi_tensor_args{};
   launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor,
                             beta_tensor, use_per_group_alpha_beta, num_tensors, stream,
-                            a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, A_sel.dptr,
-                            inputC->dtype(), outputD->dtype());
+                            a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr,
+                            /*a_is_discrete=*/false, A_sel.dptr, inputC->dtype(), outputD->dtype());
 
   // Compute average dimensions for heuristics
   // K dimension: if transa, K is A's last dim; if not, K is A's first dim
@@ -1788,8 +1790,8 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num
 
   launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor,
                             beta_tensor, use_per_group_alpha_beta, num_tensors, stream,
-                            a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, nullptr,
-                            inputC->dtype(), outputD->dtype());
+                            a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr,
+                            /*a_is_discrete=*/true, nullptr, inputC->dtype(), outputD->dtype());
 
   GroupedGemmConfig gemm_config;
   gemm_config.use_split_accumulator = config_.use_split_accumulator;
@@ -1873,8 +1875,8 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa,
   MultiTensorGroupGemmInputArgs a_multi_tensor_args{};
   launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, /*C=*/nullptr, /*D=*/nullptr,
                             alpha_tensor, beta_tensor, use_per_group_alpha_beta, num_tensors,
-                            stream, a_multi_tensor_args, C_list, D_list, A_sel.dptr, d_dtype,
-                            d_dtype);
+                            stream, a_multi_tensor_args, C_list, D_list,
+                            /*a_is_discrete=*/false, A_sel.dptr, d_dtype, d_dtype);
 
   GroupedGemmConfig gemm_config;
   gemm_config.use_split_accumulator = config_.use_split_accumulator;
@@ -1909,8 +1911,6 @@ void launch_grouped_bias_add(const transformer_engine::GroupedTensor *outputD,
   NVTE_CHECK(outputD->num_tensors >= 1, api_name, ": number of tensors must be at least 1");
   NVTE_CHECK(outputD->num_tensors == bias_tensor->num_tensors, api_name,
              ": output and bias must have the same number of tensors");
-  NVTE_CHECK(outputD->has_data(), api_name, ": output is missing row-wise data");
-  NVTE_CHECK(bias_tensor->has_data(), api_name, ": bias is missing row-wise data");
   NVTE_CHECK(outputD->dtype() == bias_tensor->dtype(), api_name,
              ": output and bias must have matching dtypes");
   NVTE_CHECK(bias_tensor->all_same_first_dim(), api_name,
@@ -1921,15 +1921,23 @@ void launch_grouped_bias_add(const transformer_engine::GroupedTensor *outputD,
   NVTE_CHECK(outputD->get_common_last_dim() == bias_tensor->get_common_last_dim(), api_name,
              ": output and bias last dims must match");
 
+  const int num_tensors = static_cast<int>(outputD->num_tensors);
+  NVTE_CHECK(num_tensors <= kMaxGroups, api_name, " supports at most ", kMaxGroups,
+             " tensors, got ", num_tensors);
+  const int total_rows = static_cast<int>(outputD->logical_shape.data[0]);
+  // A valid zero-sized CUDA allocation may have a null data pointer.
+  if (total_rows == 0) {
+    return;
+  }
+
+  NVTE_CHECK(outputD->has_data(), api_name, ": output is missing row-wise data");
+  NVTE_CHECK(bias_tensor->has_data(), api_name, ": bias is missing row-wise data");
+
   const TensorShapeInfo d_meta = TensorShapeInfo::from_tensor(outputD);
 
   const DType dtype = outputD->dtype();
   constexpr int kThreads = 128;
 
-  const int num_tensors = static_cast<int>(outputD->num_tensors);
-  NVTE_CHECK(num_tensors <= kMaxGroups, api_name, " supports at most ", kMaxGroups,
-             " tensors, got ", num_tensors);
-  const int total_rows = static_cast<int>(outputD->logical_shape.data[0]);
   const int n = static_cast<int>(outputD->get_common_last_dim());
 
   const size_t elem_size = typeToSize(dtype);
@@ -1997,8 +2005,6 @@ void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGrou
   const GroupedTensor *bias_tensor = convertNVTEGroupedTensorCheck(bias);
   const Tensor *scale_tensor = convertNVTETensorCheck(scale);
 
-  NVTE_CHECK(scale_tensor->data.dptr != nullptr,
-             "Grouped scaled bias add: scale tensor must not be null");
   NVTE_CHECK(scale_tensor->dtype() == DType::kFloat32,
              "Grouped scaled bias add: scale must be float32");
   NVTE_CHECK(scale_tensor->data.shape.size() == 1,
@@ -2007,6 +2013,10 @@ void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGrou
   const size_t total_rows = static_cast<size_t>(outputD->logical_shape.data[0]);
   NVTE_CHECK(scale_tensor->data.shape[0] == total_rows, "Grouped scaled bias add: scale size (",
              scale_tensor->data.shape[0], ") must equal total rows (", total_rows, ")");
+  if (total_rows > 0) {
+    NVTE_CHECK(scale_tensor->data.dptr != nullptr,
+               "Grouped scaled bias add: scale tensor must not be null");
+  }
 
   const float *scale_ptr = static_cast<const float *>(scale_tensor->data.dptr);
   launch_grouped_bias_add(outputD, bias_tensor, scale_ptr, true, stream);
diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h
index be92651656..9114b1e453 100644
--- a/transformer_engine/pytorch/csrc/extensions.h
+++ b/transformer_engine/pytorch/csrc/extensions.h
@@ -375,7 +375,7 @@ py::object dequantize(const py::handle &input, DType otype);
 py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors,
                           std::optional<at::Tensor> first_dims, std::optional<at::Tensor> last_dims,
                           std::optional<at::Tensor> tensor_offsets,
-                          std::optional<at::Tensor> noop_flag);
+                          std::optional<at::Tensor> noop_flag, const py::object &output);
 
 py::object nvfp4_group_quantize_with_amax(const at::Tensor &tensor, py::handle quantizer,
                                           const size_t num_tensors,
diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp
index 5ce0261c82..89c0418f01 100644
--- a/transformer_engine/pytorch/csrc/extensions/cast.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp
@@ -277,7 +277,7 @@ void compute_grouped_fp8_current_scaling_amax_and_scale(
 py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const size_t num_tensors,
                           std::optional<at::Tensor> first_dims, std::optional<at::Tensor> last_dims,
                           std::optional<at::Tensor> tensor_offsets,
-                          std::optional<at::Tensor> noop_flag) {
+                          std::optional<at::Tensor> noop_flag, const py::object &output) {
   using namespace transformer_engine::pytorch::detail;
   init_extension();
 
@@ -306,11 +306,34 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const
                                         GetTransformerEngineDType(tensor.scalar_type()),
                                         std::vector<size_t>{static_cast<size_t>(tensor.numel())});
 
-  // Create output GroupedTensor.
-  auto [grouped_output_tensor_cpp, grouped_output_py] = quantizer_cpp->create_grouped_tensor(
-      num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()),
-      py::reinterpret_borrow<py::object>(quantizer), first_dims, last_dims, tensor_offsets,
-      logical_first_dim, logical_last_dim);
+  // Create a GroupedTensor or reuse an existing destination. Reusing the destination is
+  // required for weight caching and CUDA graph replay because captured GEMMs retain the
+  // original data and scale pointers.
+  py::object grouped_output_py;
+  auto grouped_output_tensor_cpp = [&]() -> GroupedTensorWrapper {
+    if (!output.is_none()) {
+      NVTE_CHECK(!first_dims.has_value() && !last_dims.has_value() && !tensor_offsets.has_value(),
+                 "group_quantize: output reuse currently requires uniform tensor shapes.");
+      NVTE_CHECK(output.attr("num_tensors").cast<size_t>() == num_tensors,
+                 "group_quantize: output has a different number of tensors.");
+      NVTE_CHECK(output.attr("logical_shape").cast<std::vector<size_t>>() == logical_shape,
+                 "group_quantize: output has an incompatible logical shape.");
+      py::object output_quantizer = output.attr("quantizer");
+      NVTE_CHECK(!output_quantizer.is_none(),
+                 "group_quantize: output must have quantized storage.");
+      NVTE_CHECK(output_quantizer.is(quantizer),
+                 "group_quantize: output must have been created by the same quantizer.");
+      grouped_output_py = output;
+      return GroupedTensorFromPyTorchGroupedTensor(output);
+    }
+
+    auto result = quantizer_cpp->create_grouped_tensor(
+        num_tensors, logical_shape, GetTransformerEngineDType(tensor.scalar_type()),
+        py::reinterpret_borrow<py::object>(quantizer), first_dims, last_dims, tensor_offsets,
+        logical_first_dim, logical_last_dim);
+    grouped_output_py = std::move(result.second);
+    return std::move(result.first);
+  }();
 
   // dispatch to scaling methods
   enum class GroupedQuantizationMode {
@@ -381,6 +404,9 @@ py::object group_quantize(const at::Tensor &tensor, py::handle quantizer, const
       QuantizationConfigWrapper quant_config_cpp;
       quant_config_cpp.set_force_pow_2_scales(fp8_block_quantizer_cpp->force_pow_2_scales);
       quant_config_cpp.set_amax_epsilon(fp8_block_quantizer_cpp->amax_epsilon);
+      if (noop_flag_cpp.has_value()) {
+        quant_config_cpp.set_noop_tensor(noop_flag_cpp->data());
+      }
       NVTE_SCOPED_GIL_RELEASE({
         nvte_group_quantize(grouped_input_tensor.data(), grouped_output_tensor_cpp.data(),
                             quant_config_cpp, at::cuda::getCurrentCUDAStream());
@@ -749,7 +775,7 @@ py::object group_requantize_inplace(py::handle grouped_x, py::handle quantizer,
     // optimize_for_gemm, which makes the kernel emit swizzled columnwise scales directly.
     quantizer.attr("set_usage")(py::arg("rowwise") = false, py::arg("columnwise") = true);
     auto columnwise = group_quantize(dequantized, quantizer, num_tensors, first_dims, std::nullopt,
-                                     tensor_offsets, std::nullopt);
+                                     tensor_offsets, std::nullopt, py::none());
     grouped_x.attr("columnwise_data") = columnwise.attr("columnwise_data");
     grouped_x.attr("columnwise_scale_inv") = columnwise.attr("columnwise_scale_inv");
   }
diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
index 2173dff8b2..7c8793b0f6 100644
--- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp
@@ -209,7 +209,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("group_quantize", transformer_engine::pytorch::group_quantize, py::arg("tensor"),
         py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims"),
         py::arg("last_dims") = py::none(), py::arg("tensor_offsets") = py::none(),
-        py::arg("noop_flag") = py::none());
+        py::arg("noop_flag") = py::none(), py::arg("output") = py::none());
   transformer_engine::pytorch::bind_quantize_with_amax_extensions(m);
   m.def("group_dequantize", transformer_engine::pytorch::group_dequantize,
         "Dequantize group tensor", py::arg("input"), py::arg("otype"));
diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp
index ddb85808a5..7a3b284aca 100644
--- a/transformer_engine/pytorch/csrc/type_converters.cpp
+++ b/transformer_engine/pytorch/csrc/type_converters.cpp
@@ -219,12 +219,23 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) {
   }
   auto ret = GroupedTensorWrapper(num_tensors, logical_shape, scaling_mode);
 
+  auto get_initialized_storage_shape = [](const at::Tensor &data) {
+    auto shape = getTensorShape(data);
+    if (data.numel() == 0) {
+      // PyTorch may use a null pointer for a valid zero-sized allocation. TE Common reserves
+      // {nullptr, {0}} for uninitialized storage, so use an orientation-neutral 2D empty shape
+      // to distinguish explicitly provided rowwise or columnwise storage from missing storage.
+      shape = {0, 0};
+    }
+    return shape;
+  };
+
   // Rowwise data
   if (!tensor.attr("rowwise_data").is_none()) {
     const auto &data = tensor.attr("rowwise_data").cast<at::Tensor>();
     DType data_dtype =
         quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype;
-    ret.set_rowwise_data(data.data_ptr(), data_dtype, getTensorShape(data));
+    ret.set_rowwise_data(data.data_ptr(), data_dtype, get_initialized_storage_shape(data));
   } else if (quantizer_dtype != DType::kNumTypes) {
     ret.set_rowwise_data(nullptr, quantizer_dtype, std::vector<size_t>{0});
   }
@@ -234,7 +245,7 @@ GroupedTensorWrapper GroupedTensorFromPyTorchGroupedTensor(py::handle tensor) {
     const auto &data = tensor.attr("columnwise_data").cast<at::Tensor>();
     DType data_dtype =
         quantizer.is_none() ? GetTransformerEngineDType(data.scalar_type()) : quantizer_dtype;
-    ret.set_columnwise_data(data.data_ptr(), data_dtype, getTensorShape(data));
+    ret.set_columnwise_data(data.data_ptr(), data_dtype, get_initialized_storage_shape(data));
   } else if (quantizer_dtype != DType::kNumTypes) {
     ret.set_columnwise_data(nullptr, quantizer_dtype, std::vector<size_t>{0});
   }
diff --git a/transformer_engine/pytorch/module/__init__.py b/transformer_engine/pytorch/module/__init__.py
index 3cf15efc11..98ab745448 100644
--- a/transformer_engine/pytorch/module/__init__.py
+++ b/transformer_engine/pytorch/module/__init__.py
@@ -5,7 +5,7 @@
 """Module level PyTorch APIs"""
 from .layernorm_linear import LayerNormLinear
 from .linear import Linear
-from .grouped_linear import GroupedLinear
+from .grouped_linear import GroupedLinear, is_module_grouped_tensor_path_supported
 from .layernorm_mlp import LayerNormMLP
 from .layernorm import LayerNorm
 from .rmsnorm import RMSNorm
diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py
index 9860d48237..fe283cf1a3 100644
--- a/transformer_engine/pytorch/module/grouped_linear.py
+++ b/transformer_engine/pytorch/module/grouped_linear.py
@@ -82,7 +82,6 @@
 from ...debug.pytorch.debug_quantization import DebugQuantizer
 from ...debug.pytorch.debug_state import TEDebugState
 
-
 _NATIVE_SPLIT_QUANTIZER_TYPES = frozenset(
     {
         Float8Quantizer,
@@ -446,7 +445,73 @@ def _split_quantize_and_bias(
     return outputs, grad_biases
 
 
-__all__ = ["GroupedLinear"]
+__all__ = ["GroupedLinear", "is_module_grouped_tensor_path_supported"]
+
+
+def is_module_grouped_tensor_path_supported(
+    recipe: Optional[Recipe],
+    dtype: torch.dtype,
+) -> bool:
+    """Whether the module grouped-tensor path supports this recipe and dtype.
+
+    The grouped-tensor path dispatches to ``general_grouped_gemm_for_grouped_tensor``
+    and does not inspect split values because they may reside in a CUDA tensor.
+    Inspecting them on the host would add synchronization and break CUDA Graph safety.
+
+    Supported Compute Capability (CC) and precisions:
+
+    * Hopper (CC 9.0): BF16/FP16, FP8 per-tensor current scaling, and FP8
+      block scaling.
+    * Blackwell (CC 10.x and 11.0): BF16/FP16, FP8 per-tensor current scaling,
+      MXFP8, and NVFP4 with RHT.
+    * Custom recipes are unsupported because they may assign different
+      quantizers to input, weight, and grad-output roles. This predicate
+      currently supports only built-in recipes with known uniform layouts.
+    * FP8 delayed scaling is unsupported because the required grouped
+      quantization kernels are unavailable.
+    * FP8 block scaling is unsupported by this path on Blackwell because it
+      does not implement the legacy path's MXFP8-broadcast emulation.
+    * Grouped GEMM requires cuBLASLt 13.3+, with 13.4+ required on Hopper,
+      13.5+ required for FP8 per-tensor current scaling on Hopper, and 13.6+
+      required for FP8 block scaling on Hopper.
+    * FP32 is unsupported by the cuBLASLt grouped GEMM.
+
+    Runtime-only restrictions such as debug mode, CPU offloading, calibration,
+    output quantization, and backend selection are checked separately by
+    ``GroupedLinear``.
+    """
+    if dtype not in (torch.bfloat16, torch.float16):
+        return False
+
+    device_capability = get_device_compute_capability()
+    if not (9, 0) <= device_capability <= (11, 0):
+        return False
+    cublaslt_version = tex.get_cublasLt_version()
+    if cublaslt_version < 130300:
+        return False
+    if device_capability < (10, 0) and cublaslt_version < 130400:
+        return False
+
+    if recipe is None:
+        return True
+    if recipe.custom():
+        return False
+    if recipe.backward_override is not None:
+        return False
+    if recipe.float8_current_scaling():
+        return device_capability >= (10, 0) or cublaslt_version >= 130500
+    if recipe.float8_block_scaling():
+        # cuBLASLt 13.6 fixes Hopper grouped GEMM algo selection for block-scaled FP8.
+        return device_capability < (10, 0) and cublaslt_version >= 130600
+    if recipe.mxfp8():
+        return device_capability >= (10, 0)
+    if recipe.nvfp4():
+        return (
+            device_capability >= (10, 0)
+            and not recipe.disable_rht
+            and not recipe.row_scaled_activation
+        )
+    return False
 
 
 class _GroupedLinear(torch.autograd.Function):
@@ -464,97 +529,6 @@ def _maybe_dequantize(
             return tensor.dequantize(dtype=dtype)
         return cast_if_needed(tensor, dtype)
 
-    @staticmethod
-    def _is_grouped_tensor_path_supported(
-        *,
-        fp8: bool,
-        fp8_calibration: bool,
-        debug: bool,
-        cpu_offloading: bool,
-        backward_override: Optional[str],
-        save_original_input: bool,
-        activation_dtype: torch.dtype,
-        input_quantizers: List[Optional[Quantizer]],
-        output_quantizers: List[Optional[Quantizer]],
-    ) -> bool:
-        """Whether to use cuBLASLt grouped GEMM through GroupedTensor metadata.
-
-        There are no checks whether split sizes are supported. Splits
-        may be in a CUDA tensor, so checking would hurt performance
-        and be incompatible with CUDA Graphs.
-
-        Supported Compute Capability (CC) and precisions:
-        * Hopper (CC 9.0): BF16/FP16, FP8 per-tensor current scaling, and FP8
-          block scaling (1D/2D, including power-of-2 scales).
-        * Blackwell (CC 10.x and 11.0): BF16/FP16/MXFP8/NVFP4 with RHT and FP8
-          per-tensor current scaling.
-        FP8 delayed scaling is not supported because the corresponding grouped
-        quantization kernels are missing. FP8 block scaling on Blackwell (SM100 and
-        SM110) raises instead of falling back: the fused path is Hopper-only and has
-        no MXFP8-broadcast emulation. Architectures outside the fused-path window
-        (e.g. SM120) fall back to the legacy path like every other recipe.
-        Grouped GEMM requires cuBLAS 13.3+ (13.4+ on Hopper, 13.5+ for FP8
-        per-tensor current scaling on Hopper); otherwise the legacy path is used.
-        Non-RHT NVFP4 falls back to the legacy path because graph-safe grouped quantization
-        currently requires RHT.
-
-        Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it would
-        trigger a fatal error in the cuBLASLt grouped GEMM check.
-        """
-        # 1. Filter by environment variable
-        if not bool(int(os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0"))):
-            return False
-        # 2. Filter out advanced features
-        if (
-            debug
-            or cpu_offloading
-            or fp8_calibration
-            or backward_override is not None
-            or save_original_input
-        ):
-            return False
-        # 3. Filter by compute capability and cuBLAS version
-        device_capability = get_device_compute_capability()
-        if not (9, 0) <= device_capability <= (11, 0):
-            return False
-        cublaslt_version = tex.get_cublasLt_version()
-        if cublaslt_version < 130300:
-            return False
-        if device_capability < (10, 0) and cublaslt_version < 130400:
-            return False
-        # 4. Output quantization is not supported.
-        if any(q is not None for q in output_quantizers):
-            return False
-        # 5. Filter by quantization recipes.
-        if fp8:
-            if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers):
-                # FP8 per-tensor scaling grouped GEMM on Hopper requires cuBLAS 13.5+.
-                if device_capability < (10, 0) and cublaslt_version < 130500:
-                    return False
-                return True
-            if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers):
-                # Grouped FP8 block-scaling quantize kernels and cuBLASLt grouped GEMM
-                # scale modes are Hopper-only, and the fused path has no MXFP8-broadcast
-                # emulation. On Blackwell (SM100/SM110, the only other arch that reaches
-                # this branch) fail loudly rather than silently falling back to the
-                # unfused path the user explicitly opted out of.
-                if get_device_compute_capability() >= (10, 0):
-                    raise RuntimeError(
-                        "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM=1 does not support the"
-                        " FP8 block-scaling recipe on Blackwell GPUs: the fused grouped"
-                        " FP8 block-scaling path is Hopper-only. Unset"
-                        " NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM to use the unfused"
-                        " path (emulated via MXFP8 GEMM on Blackwell)."
-                    )
-                return True
-            # MXFP8 and NVFP4 require Blackwell+.
-            if not (10, 0) <= device_capability <= (11, 0):
-                return False
-            return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) or all(
-                isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers
-            )
-        return activation_dtype in (torch.bfloat16, torch.float16)
-
     @staticmethod
     def _make_grouped_tensor(
         data: torch.Tensor,
@@ -604,14 +578,101 @@ def _prepare_weights_for_grouped_tensor_gemm(
         weight_quantizers: List[Optional[Quantizer]],
         weight_workspaces: List[Optional[QuantizedTensorStorage]],
         *,
+        num_gemms: int,
+        single_grouped_weight: bool,
         with_quantized_compute: bool,
         columnwise_usage: bool,
         activation_dtype: torch.dtype,
         is_first_microbatch: Optional[bool],
         skip_fp8_weight_update: Optional[torch.Tensor],
         cache_weight: bool,
-    ) -> Tuple[List[torch.Tensor], List[Optional[QuantizedTensorStorage]]]:
-        """Prepare discrete weight tensors for GroupedTensor GEMM."""
+    ) -> Tuple[
+        Union[GroupedTensorStorage, List[torch.Tensor]],
+        List[Optional[QuantizedTensorStorage]],
+    ]:
+        """Prepare a grouped parameter or discrete weights for GroupedTensor GEMM."""
+        if single_grouped_weight:
+            weight = weights[0]
+            if not isinstance(weight, GroupedTensorStorage):
+                raise TypeError(
+                    "single_grouped_weight requires the weight parameter to be a GroupedTensor."
+                )
+
+            new_workspaces: List[Optional[QuantizedTensorStorage]] = [None]
+            if weight.quantizer is not None:
+                if not with_quantized_compute:
+                    raise RuntimeError(
+                        "Quantized single grouped weights require quantized grouped GEMM compute."
+                    )
+                return weight, new_workspaces
+
+            if not with_quantized_compute:
+                if weight.rowwise_data is None:
+                    raise RuntimeError("Single grouped weight has no rowwise storage.")
+                if weight.rowwise_data.dtype == activation_dtype:
+                    return weight, new_workspaces
+                data = weight.rowwise_data.to(dtype=activation_dtype)
+                return (
+                    GroupedTensorStorage(
+                        shape=weight.logical_shape,
+                        dtype=activation_dtype,
+                        num_tensors=num_gemms,
+                        shapes=weight.tensor_shapes,
+                        quantizer=None,
+                        data=data,
+                    ),
+                    new_workspaces,
+                )
+
+            weight_quantizer = weight_quantizers[0]
+            if weight_quantizer is None:
+                raise RuntimeError("Quantized grouped compute requires a weight quantizer.")
+            weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)
+            # forward() already applied _enable_weight_preswizzle(); preserve that decision
+            # because not every quantizer and weight shape supports fused quantize-swizzle.
+
+            workspace = weight_workspaces[0] if weight_workspaces else None
+
+            if workspace is not None and (
+                workspace.quantizer is not weight_quantizer
+                or (columnwise_usage and workspace.columnwise_data is None)
+            ):
+                workspace = None
+
+            if weight.rowwise_data is None:
+                raise RuntimeError("Single grouped weight has no rowwise storage to quantize.")
+            source = weight.rowwise_data.view(weight.logical_shape)
+            update_workspace = is_first_microbatch is None or is_first_microbatch
+            if workspace is None:
+                if cache_weight:
+                    # Match quantize_weight(): persistent workspaces must be Tensor subclasses
+                    # so autograd can save them without decomposing their storage metadata.
+                    saved_internal = weight_quantizer.internal
+                    weight_quantizer.internal = False
+                grouped_weight = tex.group_quantize(
+                    source,
+                    weight_quantizer,
+                    num_gemms,
+                    None,
+                )
+                if cache_weight:
+                    weight_quantizer.internal = saved_internal
+            elif skip_fp8_weight_update is not None or update_workspace:
+                grouped_weight = tex.group_quantize(
+                    source,
+                    weight_quantizer,
+                    num_gemms,
+                    None,
+                    noop_flag=skip_fp8_weight_update,
+                    output=workspace,
+                )
+            else:
+                grouped_weight = workspace
+
+            if cache_weight:
+                new_workspaces[0] = grouped_weight
+            return grouped_weight, new_workspaces
+
         weights_for_gemm: List[torch.Tensor] = []
         new_workspaces: List[Optional[QuantizedTensorStorage]] = [None] * len(weights)
         if not with_quantized_compute:
@@ -651,13 +712,10 @@ def _validate_or_alloc_output(
         """
         if buffer is None:
             return torch.empty((rows, cols), dtype=dtype, device=device)
-        if buffer.dim() != 2:
-            raise ValueError(f"Output buffer must be 2D, got {buffer.dim()}D.")
-        if buffer.size(0) != rows:
-            raise ValueError(f"Output buffer rows {buffer.size(0)} must match input rows {rows}.")
-        if buffer.size(1) != cols:
+        expected_shape = (rows, cols)
+        if buffer.shape != expected_shape:
             raise ValueError(
-                f"Output buffer last dim {buffer.size(1)} does not match required {cols}."
+                f"Output buffer shape {tuple(buffer.shape)} must match required {expected_shape}."
             )
         if buffer.dtype != dtype:
             raise ValueError(f"Output buffer dtype {buffer.dtype} does not match required {dtype}.")
@@ -671,6 +729,43 @@ def _validate_or_alloc_output(
             raise ValueError("Output buffer must not require gradient.")
         return buffer
 
+    @staticmethod
+    def _prepare_bias_for_grouped_tensor_gemm(
+        biases: Tuple[torch.Tensor, ...],
+        *,
+        single_grouped_bias: bool,
+        num_gemms: int,
+        out_features: int,
+        dtype: torch.dtype,
+    ) -> GroupedTensorStorage:
+        """Prepare grouped or discrete bias storage for grouped GEMM."""
+        if not single_grouped_bias:
+            return _GroupedLinear._make_grouped_bias(
+                biases,
+                num_gemms=num_gemms,
+                out_features=out_features,
+                dtype=dtype,
+            )
+
+        bias = biases[0]
+        if not isinstance(bias, GroupedTensorStorage):
+            raise TypeError("single_grouped_bias requires a GroupedTensor parameter.")
+        bias_data = bias.rowwise_data
+        if bias_data.dtype != dtype:
+            bias_data = bias_data.to(dtype=dtype)
+
+        # The parameter exposes a packed [num_gemms, out_features] tensor, but its grouped
+        # members are 1D vectors. The grouped bias-add kernel consumes those same bytes as
+        # num_gemms row matrices with shape [1, out_features].
+        return GroupedTensorStorage(
+            shape=(num_gemms, out_features),
+            dtype=dtype,
+            num_tensors=num_gemms,
+            shapes=[(1, out_features)] * num_gemms,
+            quantizer=None,
+            data=bias_data.reshape(-1),
+        )
+
     @staticmethod
     def _forward_grouped_tensor(
         ctx,
@@ -692,6 +787,9 @@ def _forward_grouped_tensor(
         weight_workspaces: List[Optional[QuantizedTensorStorage]],
         cache_weight: bool,
         skip_fp8_weight_update: Optional[torch.Tensor],
+        save_original_input: bool,
+        single_grouped_weight: bool,
+        single_grouped_bias: bool,
         weights: Tuple[torch.Tensor, ...],
         biases: Tuple[torch.Tensor, ...],
         out: Optional[torch.Tensor] = None,
@@ -701,8 +799,9 @@ def _forward_grouped_tensor(
         num_gemms = len(m_splits)
         device = inp.device
         in_features = weights[0].size(-1)
-        out_features = weights[0].size(0)
+        out_features = weights[0].size(-2)
         weight_requires_grad = weights[0].requires_grad
+        save_original_input = save_original_input and weight_requires_grad
 
         split_sizes, (
             base_split_offsets,
@@ -723,7 +822,7 @@ def _forward_grouped_tensor(
             input_quantizer = input_quantizers[0]
             input_quantizer.set_usage(
                 rowwise=True,
-                columnwise=is_grad_enabled and weight_requires_grad,
+                columnwise=(is_grad_enabled and weight_requires_grad and not save_original_input),
             )
             input_quantizer.optimize_for_gemm = True
             grouped_x = tex.group_quantize(
@@ -748,6 +847,8 @@ def _forward_grouped_tensor(
             weights,
             weight_quantizers,
             weight_workspaces,
+            num_gemms=num_gemms,
+            single_grouped_weight=single_grouped_weight,
             with_quantized_compute=fp8,
             columnwise_usage=columnwise_usage,
             activation_dtype=activation_dtype,
@@ -774,8 +875,9 @@ def _forward_grouped_tensor(
 
         grouped_bias = None
         if use_bias:
-            grouped_bias = _GroupedLinear._make_grouped_bias(
+            grouped_bias = _GroupedLinear._prepare_bias_for_grouped_tensor_gemm(
                 biases,
+                single_grouped_bias=single_grouped_bias,
                 num_gemms=num_gemms,
                 out_features=out_features,
                 dtype=activation_dtype,
@@ -797,18 +899,25 @@ def _forward_grouped_tensor(
         )
 
         if is_grad_enabled:
+            input_to_save = grouped_x
             if weight_requires_grad:
-                # (For FP8 per tensor current scaling on Hopper --> Free Rowwise Data
-                # in backward pass)
-                if fp8 and grouped_x.columnwise_data is not None:
+                if save_original_input:
+                    # Save the high-precision input and reconstruct the grouped columnwise
+                    # operand in backward instead of retaining a second quantized copy.
+                    input_to_save = inp
+                elif fp8 and grouped_x.columnwise_data is not None:
+                    # Wgrad only consumes the columnwise representation.
                     grouped_x.rowwise_data = None
                     grouped_x.scale_inv = None
             else:
-                grouped_x = None
+                input_to_save = None
+
+            weights_to_save = [weights_for_gemm] if single_grouped_weight else weights_for_gemm
+            if not inp.requires_grad:
+                weights_to_save = [None] * len(weights_to_save)
 
-            weights_to_save = weights_for_gemm if inp.requires_grad else [None] * num_gemms
             tensors_to_save, tensor_objects = prepare_for_saving(
-                grouped_x,
+                input_to_save,
                 *weights_to_save,
                 split_sizes,
                 base_split_offsets,
@@ -818,7 +927,7 @@ def _forward_grouped_tensor(
             ctx.save_for_backward(*tensors_to_save)
             ctx.tensor_objects = tensor_objects
 
-            ctx.use_grouped_tensor_path = True
+            ctx.grouped_tensor_supported = True
             ctx.weight_quantizers = weight_quantizers
             ctx.weights_shape_0 = out_features
             ctx.weights_shape_1 = in_features
@@ -826,16 +935,18 @@ def _forward_grouped_tensor(
             ctx.grad_output_quantizers = grad_output_quantizers
             ctx.grad_weight_quantizers = grad_weight_quantizers
             ctx.weights_requires_grad = weight_requires_grad
+            ctx.single_grouped_weight = single_grouped_weight
+            ctx.single_grouped_bias = single_grouped_bias
             if fuse_wgrad_accumulation and ctx.weights_requires_grad:
                 ctx.origin_weight_refs = [weakref.ref(w) for w in weights]
                 ctx.origin_weights_overwrite_main_grad = getattr(
                     weights[0], "overwrite_main_grad", False
                 )
                 if hasattr(weights[0], "__fsdp_param__"):
-                    ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)]
+                    ctx.main_grad_funcs = [weight.get_main_grad for weight in weights]
                 else:
                     ctx.main_grad_funcs = [
-                        lambda j=i: weights[j].main_grad for i in range(num_gemms)
+                        lambda j=i: weights[j].main_grad for i in range(len(weights))
                     ]
             ctx.device = device
             ctx.dgrad_out = dgrad_out
@@ -859,7 +970,7 @@ def _forward_grouped_tensor(
                 )
             ctx.wgrad_store = wgrad_store
             ctx.debug = False
-            ctx.save_original_input = False
+            ctx.save_original_input = save_original_input
             ctx.input_quantizers = input_quantizers
 
         return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces
@@ -903,19 +1014,22 @@ def forward(
             delayed_scaling_input_quantizer,
             unsafe_requantization_input_quantizer,
             debug,
+            single_grouped_weight,
+            single_grouped_bias,
+            use_grouped_tensor,
         ) = non_tensor_args
-        if fp8:
-            backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override
-        else:
-            backward_override = None
+        recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None
+        backward_override = recipe.backward_override if recipe is not None else None
         if backward_override == "high_precision":
             save_original_input = True
         elif backward_override == "dequantized":
             save_original_input = False
 
         num_gemms = len(m_splits)
-        weights = weights_and_biases[:num_gemms]
-        biases = weights_and_biases[num_gemms:]
+        num_weight_args = 1 if single_grouped_weight else num_gemms
+        num_bias_args = 1 if single_grouped_bias else num_gemms
+        weights = weights_and_biases[:num_weight_args]
+        biases = weights_and_biases[num_weight_args : num_weight_args + num_bias_args]
         device = inp.device
         weight_requires_grad = weights[0].requires_grad
 
@@ -993,17 +1107,46 @@ def forward(
                 f"weight tensor (shape={tuple(weights[0].size())})"
             )
 
-        if _GroupedLinear._is_grouped_tensor_path_supported(
-            fp8=fp8,
-            fp8_calibration=fp8_calibration,
-            debug=debug,
-            cpu_offloading=cpu_offloading,
-            backward_override=backward_override,
-            save_original_input=save_original_input,
-            activation_dtype=activation_dtype,
-            input_quantizers=input_quantizers,
-            output_quantizers=output_quantizers,
+        grouped_tensor_supported = False
+        if use_grouped_tensor and not (
+            fp8_calibration
+            or debug
+            or cpu_offloading
+            or any(q is not None for q in output_quantizers)
         ):
+            if (
+                fp8
+                and recipe.float8_block_scaling()
+                and (10, 0) <= get_device_compute_capability() <= (11, 0)
+            ):
+                raise RuntimeError(
+                    "use_grouped_tensor=True does not support the FP8 block-scaling recipe on "
+                    "Blackwell GPUs: the native grouped FP8 block-scaling path is Hopper-only. "
+                    "Set use_grouped_tensor=False, or unset "
+                    "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM if it enabled this path, to use "
+                    "the MXFP8-emulated path on Blackwell."
+                )
+            grouped_tensor_supported = is_module_grouped_tensor_path_supported(
+                recipe,
+                activation_dtype,
+            )
+        if (
+            use_grouped_tensor
+            and not grouped_tensor_supported
+            and (single_grouped_weight or single_grouped_bias)
+        ):
+            raise RuntimeError(
+                "Single grouped parameters require the native grouped-tensor path, but the active "
+                "device, cuBLASLt version, quantization recipe, or GroupedLinear feature "
+                "configuration does not support it. Disable single_grouped_weight and "
+                "single_grouped_bias to allow the split-quantize fallback."
+            )
+        if grouped_tensor_supported:
+            if m_splits.device.type != "cuda":
+                raise ValueError(
+                    "The native grouped_tensor path requires CUDA m_splits. Pass a CUDA int64 "
+                    "tensor, or set use_grouped_tensor=False."
+                )
             return _GroupedLinear._forward_grouped_tensor(
                 ctx,
                 inp=inp,
@@ -1023,6 +1166,9 @@ def forward(
                 weight_workspaces=weight_workspaces,
                 cache_weight=cache_weight,
                 skip_fp8_weight_update=skip_fp8_weight_update,
+                save_original_input=save_original_input,
+                single_grouped_weight=single_grouped_weight,
+                single_grouped_bias=single_grouped_bias,
                 weights=weights,
                 biases=biases,
                 out=out,
@@ -1114,7 +1260,7 @@ def forward(
             mark_not_offload(*weights_fp8, *weights)
 
         if is_grad_enabled:
-            ctx.use_grouped_tensor_path = False
+            ctx.grouped_tensor_supported = False
             ctx.weight_quantizers = weight_quantizers
             ctx.weights_shape_1 = weights[0].shape[1]
 
@@ -1231,15 +1377,47 @@ def _backward_grouped_tensor(
         """Backward path paired with ``_forward_grouped_tensor``."""
         saved_tensors = restore_from_func_ctx(ctx)
         N = ctx.num_gemms
-        grouped_x = saved_tensors[0]
-        weights = saved_tensors[1 : 1 + N]
-        split_sizes = saved_tensors[1 + N]
-        base_split_offsets = saved_tensors[2 + N]
-        input_tensor_offsets = saved_tensors[3 + N]
-        output_tensor_offsets = saved_tensors[4 + N]
-
-        origin_weights = [None] * N
-        main_grads = [None] * N
+        saved_input = saved_tensors[0]
+        if ctx.single_grouped_weight:
+            weights_for_gemm = saved_tensors[1]
+            weight_tensors = [weights_for_gemm]
+            split_sizes = saved_tensors[2]
+            base_split_offsets = saved_tensors[3]
+            input_tensor_offsets = saved_tensors[4]
+            output_tensor_offsets = saved_tensors[5]
+        else:
+            weight_tensors = saved_tensors[1 : 1 + N]
+            weights_for_gemm = weight_tensors
+            split_sizes = saved_tensors[1 + N]
+            base_split_offsets = saved_tensors[2 + N]
+            input_tensor_offsets = saved_tensors[3 + N]
+            output_tensor_offsets = saved_tensors[4 + N]
+
+        if ctx.save_original_input:
+            x = cast_if_needed(
+                saved_input.reshape(-1, ctx.weights_shape_1),
+                ctx.activation_dtype,
+            )
+            if ctx.fp8:
+                input_quantizer = ctx.input_quantizers[0]
+                input_quantizer.set_usage(rowwise=False, columnwise=True)
+                input_quantizer.optimize_for_gemm = True
+                grouped_x = tex.group_quantize(x, input_quantizer, N, split_sizes)
+            else:
+                grouped_x = _GroupedLinear._make_grouped_tensor(
+                    x,
+                    num_gemms=N,
+                    split_sizes=split_sizes,
+                    tensor_offsets=input_tensor_offsets,
+                    last_dim=ctx.weights_shape_1,
+                    dtype=ctx.activation_dtype,
+                )
+        else:
+            grouped_x = saved_input
+
+        num_weight_args = 1 if ctx.single_grouped_weight else N
+        origin_weights = [None] * num_weight_args
+        main_grads = [None] * num_weight_args
         if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad:
             origin_weight_refs = ctx.origin_weight_refs
             ctx.origin_weight_refs = None
@@ -1293,11 +1471,16 @@ def _backward_grouped_tensor(
                 dtype=ctx.activation_dtype,
             )
 
-        grad_biases = [None] * N
         if ctx.use_bias:
             if dbias_packed is None:
                 dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, N)
-            grad_biases = [dbias_packed[i].to(dtype=ctx.activation_dtype) for i in range(N)]
+            if ctx.single_grouped_bias:
+                grad_bias_args = [dbias_packed.to(dtype=ctx.activation_dtype)]
+            else:
+                grad_bias_args = [dbias_packed[i].to(dtype=ctx.activation_dtype) for i in range(N)]
+        else:
+            num_bias_args = 1 if ctx.single_grouped_bias else N
+            grad_bias_args = [None] * num_bias_args
 
         dgrad = None
         if ctx.requires_dgrad:
@@ -1306,7 +1489,7 @@ def _backward_grouped_tensor(
                 recipe = ctx.fp8_recipe
                 if hasattr(recipe, "fp8_gemm_dgrad"):
                     dgrad_gemm_use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator
-            for weight in weights:
+            for weight in weight_tensors:
                 if isinstance(weight, QuantizedTensorStorage):
                     weight.update_usage(columnwise_usage=True)
             dgrad = _GroupedLinear._validate_or_alloc_output(
@@ -1325,7 +1508,7 @@ def _backward_grouped_tensor(
                 dtype=ctx.activation_dtype,
             )
             general_grouped_gemm_for_grouped_tensor(
-                weights,
+                weights_for_gemm,
                 grouped_dy,
                 grouped_dgrad,
                 layout="NN",
@@ -1346,16 +1529,42 @@ def _backward_grouped_tensor(
                 if hasattr(recipe, "fp8_gemm_wgrad"):
                     wgrad_gemm_use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator
             if ctx.fuse_wgrad_accumulation:
-                wgrad_list = main_grads
+                if ctx.single_grouped_weight:
+                    main_grad = main_grads[0]
+                    grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data(
+                        num_tensors=N,
+                        tensor_shape=(ctx.weights_shape_0, ctx.weights_shape_1),
+                        rowwise_data=main_grad.view(-1),
+                        dtype=main_grad.dtype,
+                    )
+                    wgrad_output = grouped_wgrad
+                    wgrad_list = [main_grad]
+                else:
+                    wgrad_output = main_grads
+                    wgrad_list = main_grads
             else:
-                wgrad_packed = torch.empty(
-                    N,
-                    ctx.weights_shape_0,
-                    ctx.weights_shape_1,
-                    dtype=ctx.activation_dtype,
-                    device=ctx.device,
-                )
-                wgrad_list = [wgrad_packed[i] for i in range(N)]
+                if ctx.single_grouped_weight:
+                    grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes(
+                        num_tensors=N,
+                        shapes=[(ctx.weights_shape_0, ctx.weights_shape_1)] * N,
+                        quantizer=None,
+                        device=ctx.device,
+                        dtype=ctx.activation_dtype,
+                    )
+                    wgrad_output = grouped_wgrad
+                    wgrad_list = [
+                        grouped_wgrad.rowwise_data.view(N, ctx.weights_shape_0, ctx.weights_shape_1)
+                    ]
+                else:
+                    wgrad_packed = torch.empty(
+                        N,
+                        ctx.weights_shape_0,
+                        ctx.weights_shape_1,
+                        dtype=ctx.activation_dtype,
+                        device=ctx.device,
+                    )
+                    wgrad_output = [wgrad_packed[i] for i in range(N)]
+                    wgrad_list = wgrad_output
 
             accumulate = (
                 accumulate_wgrad_into_param_main_grad
@@ -1375,9 +1584,9 @@ def grouped_gemm_wgrad(inputmats, grad_output_mats, grad_weights):
                 return None, [None] * N, None
 
             if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute():
-                ctx.wgrad_store.put([grouped_x, grouped_dy, wgrad_list], grouped_gemm_wgrad)
+                ctx.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], grouped_gemm_wgrad)
             else:
-                grouped_gemm_wgrad(grouped_x, grouped_dy, wgrad_list)
+                grouped_gemm_wgrad(grouped_x, grouped_dy, wgrad_output)
 
             def handle_custom_ddp_from_mcore(weight, main_grad, wgrad):
                 if ctx.weights_requires_grad:
@@ -1405,10 +1614,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad):
                 for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list)
             ]
         else:
-            wgrad_list = [None] * N
-
-        if not ctx.use_bias:
-            grad_biases = [None] * N
+            wgrad_list = [None] * num_weight_args
 
         if ctx.reduce_and_update_bwd_fp8_tensors:
             FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False)
@@ -1419,7 +1625,7 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad):
             None,  # out
             None,  # dgrad_out
             *wgrad_list,
-            *grad_biases,
+            *grad_bias_args,
         )
 
     @staticmethod
@@ -1428,7 +1634,7 @@ def backward(
     ) -> Tuple[Union[torch.Tensor, None], ...]:
         # pylint: disable=missing-function-docstring
         with get_nvtx_range_context("_GroupedLinear_backward"):
-            if ctx.use_grouped_tensor_path:
+            if ctx.grouped_tensor_supported:
                 return _GroupedLinear._backward_grouped_tensor(ctx, grad_output)
 
             saved_tensors = restore_from_func_ctx(ctx)
@@ -1727,7 +1933,9 @@ class GroupedLinear(TransformerEngineBaseModule):
                  when set to ``True``, this module will not apply the additive bias itself, but
                  instead return the bias value during the forward pass together with the
                  output of the linear transformation :math:`y = xA^T`. This is useful when
-                 the bias addition can be fused to subsequent operations.
+                 the bias addition can be fused to subsequent operations. A single grouped
+                 bias is returned as its packed ``GroupedTensor`` parameter; discrete biases
+                 are returned as a list of per-GEMM tensors.
     params_dtype : torch.dtype, default = torch.get_default_dtype()
                   it controls the type used to allocate the initial parameters. Useful when
                   the model is trained with lower precision and the original FP32 parameters
@@ -1752,6 +1960,13 @@ class GroupedLinear(TransformerEngineBaseModule):
                        EXPERIMENTAL and subject to change. Gated by the
                        ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` environment variable: if the env var
                        is not set this argument is forced to ``False`` with a warning.
+    use_grouped_tensor : bool or None, default = None
+                       Prefer the native GroupedTensor grouped GEMM path. Discrete parameters
+                       fall back to split-quantize when the path is unsupported. Single grouped
+                       parameters require the native path and raise instead of falling back.
+                       The native path requires CUDA ``m_splits``. ``None`` preserves the deprecated
+                       ``NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM`` environment-variable
+                       selection for compatibility. New callers should pass a boolean explicitly.
 
     Notes
     -----
@@ -1785,6 +2000,7 @@ def __init__(
         single_grouped_weight: bool = False,
         single_grouped_bias: bool = False,
         name: Optional[str] = None,
+        use_grouped_tensor: Optional[bool] = None,
     ) -> None:
         super().__init__(name)
 
@@ -1800,11 +2016,37 @@ def __init__(
         self.ub_overlap_ag = ub_overlap_ag
         self.ub_name = ub_name
         self.save_original_input = save_original_input
+        if use_grouped_tensor is None:
+            use_grouped_tensor_env = os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM")
+            if use_grouped_tensor_env is not None:
+                warnings.warn(
+                    "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM is deprecated and will be "
+                    "removed in a future release. Pass use_grouped_tensor=True or "
+                    "use_grouped_tensor=False to GroupedLinear instead.",
+                    FutureWarning,
+                    stacklevel=2,
+                )
+            else:
+                use_grouped_tensor_env = "0"
+            use_grouped_tensor = bool(int(use_grouped_tensor_env))
+        if not isinstance(use_grouped_tensor, bool):
+            raise TypeError(
+                f"use_grouped_tensor must be a bool or None, got {type(use_grouped_tensor)}."
+            )
+        self.use_grouped_tensor = use_grouped_tensor
         single_grouped_weight, single_grouped_bias = resolve_grouped_linear_single_param_flags(
             single_grouped_weight, single_grouped_bias
         )
         self.single_grouped_weight = single_grouped_weight
         self.single_grouped_bias = single_grouped_bias
+        if self.use_bias and self.single_grouped_weight and not self.single_grouped_bias:
+            warnings.warn(
+                "GroupedLinear has single_grouped_weight=True and bias=True, but "
+                "single_grouped_bias=False. This requires packing the per-GEMM biases on every "
+                "forward; enable single_grouped_bias to keep both parameters grouped.",
+                UserWarning,
+                stacklevel=2,
+            )
         if ub_overlap_rs or ub_overlap_ag:
             raise ValueError("GroupedLinear doesn't support Userbuffer overlap.")
         self.init_method = init_method
@@ -2022,7 +2264,7 @@ def make_grouped_weights(self, defer_init=False) -> None:
             if weight_quantizers and weight_quantizers[0] is not None
             else None
         )
-        if recipe is not None and (recipe.delayed() or recipe.float8_current_scaling()):
+        if recipe is not None and recipe.delayed():
             self.set_tensor_parallel_attributes(defer_init=defer_init)
             return
 
@@ -2310,16 +2552,13 @@ def forward(
         is_grad_enabled = torch.is_grad_enabled()
         num_gemms = self.num_gemms
 
-        if FP8GlobalStateManager.fp8_graph_capturing():
-            skip_fp8_weight_update = (
-                FP8GlobalStateManager.quantization_state.skip_fp8_weight_update_tensor
-            )
-        else:
-            skip_fp8_weight_update = None
-        if skip_fp8_weight_update is not None:
-            is_first_microbatch = False
-
         # Make sure splits are in expected format
+        if (self.single_grouped_weight or self.single_grouped_bias) and not self.use_grouped_tensor:
+            raise RuntimeError(
+                "single_grouped_weight and single_grouped_bias require "
+                "use_grouped_tensor=True; the split-quantize path only supports discrete "
+                "parameters."
+            )
         if not isinstance(m_splits, torch.Tensor):
             # Convert list of ints to tensor for backward compatibility
             m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cpu")
@@ -2348,6 +2587,7 @@ def forward(
         try:
             weight_tensors = self._get_weight_tensors()
             bias_tensors = self._get_bias_tensors()
+            use_grouped_bias = self.use_bias and self.single_grouped_bias
 
             quantizers = self._get_quantizers() if not debug else self._get_debug_quantizers()
 
@@ -2355,6 +2595,13 @@ def forward(
                 if self.no_debug_features_active(list(chain(*quantizers))):
                     debug = False
                     quantizers = self._get_quantizers()
+            if debug and (self.single_grouped_weight or self.single_grouped_bias):
+                raise RuntimeError(
+                    "TE debug features do not support single grouped parameters. DebugQuantizer "
+                    "uses the split-quantize path, which only supports discrete parameters. "
+                    "Disable single_grouped_weight and single_grouped_bias, or disable TE debug "
+                    "features for this GroupedLinear."
+                )
 
             (
                 input_quantizers,
@@ -2380,11 +2627,14 @@ def forward(
                 autograd_ctx = [None]
 
             cache_weight = is_first_microbatch is not None
-            weight_workspaces = (
-                [self._fp8_workspaces.get(f"weight{i}") for i in range(num_gemms)]
-                if cache_weight
-                else [None] * num_gemms
-            )
+            if self.single_grouped_weight:
+                weight_workspaces = [self._fp8_workspaces.get("weight")] if cache_weight else [None]
+            else:
+                weight_workspaces = (
+                    [self._fp8_workspaces.get(f"weight{i}") for i in range(num_gemms)]
+                    if cache_weight
+                    else [None] * num_gemms
+                )
 
             non_tensor_args = (
                 self.apply_bias,
@@ -2410,6 +2660,9 @@ def forward(
                 self._delayed_scaling_input_quantizer,
                 self._unsafe_requantization_input_quantizer,
                 debug,
+                self.single_grouped_weight,
+                use_grouped_bias,
+                self.use_grouped_tensor,
             )
             out, new_workspaces = linear_fn(
                 *autograd_ctx,
@@ -2427,12 +2680,15 @@ def forward(
                     if ws is not None:
                         if isinstance(ws, torch.Tensor):
                             ws = ws.detach()
-                        self._fp8_workspaces[f"weight{i}"] = ws
+                        key = "weight" if self.single_grouped_weight else f"weight{i}"
+                        self._fp8_workspaces[key] = ws
 
         finally:
             self.end_forward()
 
         if self.return_bias:
+            if use_grouped_bias:
+                return out, bias_tensors[0]
             return out, [cast_if_needed(b, self.activation_dtype) for b in bias_tensors]
         return out
 
@@ -2447,31 +2703,34 @@ def backward_dw(self):
             return
         with get_nvtx_range_context("_GroupedLinear_wgrad"):
             (_, grad_biases_, _), tensor_list = self.wgrad_store.pop()
-            wgrad_list = tensor_list[2]
+            wgrad_output = tensor_list[2]
             weight_params = self._get_weight_tensors()
             if not self.fuse_wgrad_accumulation:
-                for i in range(self.num_gemms):
-                    weight_params[i].grad = wgrad_list[i].to(weight_params[i].dtype)
+                if self.single_grouped_weight:
+                    weight_params[0].grad = wgrad_output.rowwise_data.view(
+                        self.num_gemms, self.out_features, self.in_features
+                    ).to(weight_params[0].dtype)
+                else:
+                    for i in range(self.num_gemms):
+                        weight_params[i].grad = wgrad_output[i].to(weight_params[i].dtype)
             has_grad_biases = [
                 grad_bias is not None and grad_bias.numel() != 0 for grad_bias in grad_biases_
             ]
             if self.use_bias and any(has_grad_biases):
-                grouped_bias = getattr(self, "bias", None)
-                if grouped_bias is not None:
-                    if not all(has_grad_biases):
-                        raise RuntimeError("Expected all grouped bias gradients to be present.")
-                    gstack = torch.stack(grad_biases_, dim=0).to(grouped_bias.dtype)
-                    if grouped_bias.grad is None:
-                        grouped_bias.grad = gstack
-                    else:
-                        grouped_bias.grad.add_(gstack)
-                else:
-                    bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)]
-                    for i in range(self.num_gemms):
-                        if has_grad_biases[i] and bias_params[i].grad is None:
-                            bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype)
+                if self.use_grouped_tensor:
+                    raise RuntimeError(
+                        "GroupedLinear(use_grouped_tensor=True) fell back to the split-quantize "
+                        "path, which produced per-expert bias gradients during delayed wgrad. "
+                        "This implicit fallback is unsupported with delay_wgrad_compute=True. "
+                        "Use a configuration supported by the grouped-tensor path, or set "
+                        "use_grouped_tensor=False to select the legacy path explicitly."
+                    )
+                bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)]
+                for i in range(self.num_gemms):
+                    if has_grad_biases[i] and bias_params[i].grad is None:
+                        bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype)
             del grad_biases_
-            del wgrad_list
+            del wgrad_output
             del tensor_list
             self._trigger_wgrad_accumulation_and_reduce_hooks()
 
@@ -2514,10 +2773,7 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage
         """Get the weight tensors of the module."""
         grouped_weight = getattr(self, "weight", None)
         if grouped_weight is not None:
-            weight_tensors = grouped_weight.quantized_tensors
-            if weight_tensors is None:
-                # TODO(ksivaman): Remove this after GEMM integration.
-                weight_tensors = grouped_weight.split_into_quantized_tensors()
+            weight_tensors = [grouped_weight]
         else:
             weight_tensors = [getattr(self, f"weight{i}") for i in range(self.num_gemms)]
         if not self.fp8 and any(isinstance(w, QuantizedTensorStorage) for w in weight_tensors):
@@ -2532,13 +2788,23 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage
         return weight_tensors
 
     def _get_bias_tensors(self) -> List[torch.Tensor]:
-        """Per-GEMM bias tensors (views into grouped storage when ``single_grouped_bias``)."""
+        """Get bias parameters in their registered grouped or per-GEMM layout.
+
+        A single grouped bias remains one packed GroupedTensor. When return_bias=True,
+        an upper-level framework such as MCore must apply that packed bias accordingly;
+        Discrete bias parameters retain the existing list-of-per-GEMM contract.
+
+        Example with 2 experts, 128 output features:
+
+            single grouped bias:
+              GroupedTensor shape = [2, 128] -> [grouped_bias]
+
+            discrete biases:
+              bias0 [128] + bias1 [128] -> [bias0, bias1]
+        """
         grouped_bias = getattr(self, "bias", None)
         if grouped_bias is not None:
-            parts = grouped_bias.quantized_tensors
-            if parts is None:
-                parts = grouped_bias.split_into_quantized_tensors()
-            return [p.reshape(-1) for p in parts]
+            return [grouped_bias]
         return [getattr(self, f"bias{i}") for i in range(self.num_gemms)]
 
     def _get_weight_quantizers(self) -> List[Quantizer]:
diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py
index 6def36ffc7..caab9c9c8d 100644
--- a/transformer_engine/pytorch/ops/basic/__init__.py
+++ b/transformer_engine/pytorch/ops/basic/__init__.py
@@ -24,7 +24,7 @@
 from .bias import Bias
 from .constant_scale import ConstantScale
 from .dropout import Dropout
-from .grouped_linear import GroupedLinear
+from .grouped_linear import GroupedLinear, is_op_fuser_grouped_tensor_path_supported
 from .identity import Identity
 from .l2normalization import L2Normalization
 from .layer_norm import LayerNorm
diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py
index e1980d2943..1a932fb44d 100644
--- a/transformer_engine/pytorch/ops/basic/grouped_linear.py
+++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py
@@ -14,6 +14,7 @@
 import torch
 
 import transformer_engine_torch as tex
+from transformer_engine.common.recipe import Recipe
 from ...constants import DType, TE_DType
 from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor
 from ...distributed import CudaRNGStatesTracker
@@ -24,14 +25,12 @@
     _2X_ACC_WGRAD,
 )
 from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload, start_offload
-from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe
+from ...quantization import FP8GlobalStateManager, QuantizerRole
 from ...quantized_tensor import QuantizedTensorStorage
 from ...tensor import (
     Float8BlockQuantizer,
-    Float8CurrentScalingQuantizer,
     MXFP8Quantizer,
     MXFP8Tensor,
-    NVFP4Quantizer,
     Quantizer,
 )
 from ...utils import (
@@ -72,6 +71,67 @@
 GRAD_INPUT_BUFFER_KEY = "grad_input"
 
 
+def is_op_fuser_grouped_tensor_path_supported(
+    recipe: Optional[Recipe],
+    dtype: torch.dtype,
+) -> bool:
+    """Whether the op-fuser grouped-tensor path supports this recipe and dtype.
+
+    * The graph-safe path dispatches to ``general_grouped_gemm_for_grouped_tensor``,
+      which is backed by ``nvte_grouped_gemm_with_discrete_inputA`` in the common
+      library.
+    * MXFP8 and NVFP4 are supported on Blackwell GPUs with Compute Capability
+      (CC) 10.x and 11.0. NVFP4 requires RHT because graph-safe grouped
+      quantization currently requires it.
+    * FP8 per-tensor current scaling uses grouped current-scaling quantization
+      through ``tex.group_quantize`` and cuBLASLt grouped GEMM with per-batch
+      scalar FP8 scaling. It is supported on Hopper and Blackwell, with
+      cuBLASLt 13.5+ required on Hopper.
+    * FP8 block scaling uses the grouped-tensor path only on Hopper with
+      cuBLASLt 13.6+. On other architectures or older cuBLAS versions it
+      falls back to the split-quantize path for discrete parameters.
+    * Custom recipes are unsupported because they may assign different
+      quantizers to input, weight, and grad-output roles. This predicate
+      currently supports only built-in recipes with known uniform layouts.
+    * Other quantization recipes, including FP8 delayed scaling, fall back to
+      split quantization because their grouped quantization kernels are missing.
+    * Unquantized BF16/FP16 compute is supported on Hopper and Blackwell. FP32
+      is excluded because cuBLASLt grouped GEMM does not support it.
+    * Single grouped parameters have no split-quantize fallback, so callers
+      must reject them when this function returns ``False``.
+    """
+    if dtype not in (torch.bfloat16, torch.float16):
+        return False
+
+    device_capability = get_device_compute_capability()
+    if not (9, 0) <= device_capability <= (11, 0):
+        return False
+    cublaslt_version = tex.get_cublasLt_version()
+    if cublaslt_version < 130300:
+        return False
+    if device_capability < (10, 0) and cublaslt_version < 130400:
+        return False
+
+    if recipe is None:
+        return True
+    if recipe.custom():
+        return False
+    if recipe.float8_current_scaling():
+        return device_capability >= (10, 0) or cublaslt_version >= 130500
+    if recipe.float8_block_scaling():
+        # cuBLASLt 13.6 fixes Hopper grouped GEMM algo selection for block-scaled FP8.
+        return device_capability < (10, 0) and cublaslt_version >= 130600
+    if recipe.mxfp8():
+        return device_capability >= (10, 0)
+    if recipe.nvfp4():
+        return (
+            device_capability >= (10, 0)
+            and not recipe.disable_rht
+            and not recipe.row_scaled_activation
+        )
+    return False
+
+
 class GroupedLinear(BasicOperation):
     r"""Apply multiple linear transformations: :math:``y_i = x_i W_i^T + b_i``
 
@@ -316,17 +376,28 @@ def backward_dw(self) -> None:
                 w.grad = grad_weights[group_idx].to(w.dtype)
         self._trigger_wgrad_accumulation_and_reduce_hooks()
 
-    def _get_bias_tensors(self, dtype: torch.dtype) -> list[torch.Tensor]:
-        """Retrieve per-group bias tensors in the given dtype."""
+    def _get_discrete_bias_tensors(self, dtype: torch.dtype) -> list[torch.Tensor]:
+        """Retrieve discrete per-group bias parameters in the given dtype."""
         if self.single_grouped_bias:
-            bias_parts = self.bias.quantized_tensors
-            if bias_parts is None:
-                bias_parts = self.bias.split_into_quantized_tensors()
-            return [maybe_dequantize(p.reshape(-1), dtype) for p in bias_parts]
+            raise RuntimeError(
+                "Discrete bias tensors were requested for a single grouped bias parameter."
+            )
         return [
             maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(self.num_groups)
         ]
 
+    def _get_packed_bias_tensor(self, dtype: torch.dtype) -> torch.Tensor:
+        """Return all per-group biases as one dense tensor with shape [num_groups, out_features].
+
+        A single grouped bias is already stored in this layout, so return a view of the
+        registered parent parameter instead of splitting it into members and stacking it again.
+        Discrete biases require a stack because they are independent parameters.
+        """
+        if self.single_grouped_bias:
+            bias_data = self.bias.rowwise_data.view(self.num_groups, self.out_features)
+            return bias_data if bias_data.dtype == dtype else bias_data.to(dtype=dtype)
+        return torch.stack(self._get_discrete_bias_tensors(dtype), dim=0)
+
     def num_quantizers(self, mode: str) -> int:
         if mode == "forward":
             return 2 * self.num_groups
@@ -791,68 +862,6 @@ def op_backward(self, *args, **kwargs):
             "It overrides `fuser_backward` instead of `op_backward`."
         )
 
-    @staticmethod
-    def _is_graph_safe_path_supported(
-        *,
-        with_quantized_compute: bool,
-        input_quantizers: Sequence[Optional[Quantizer]],
-        dtype: torch.dtype,
-        single_grouped_weight: bool,
-    ) -> bool:
-        """Whether the graph-safe grouped-tensor flow can be used.
-
-        * The graph-safe path dispatches to ``general_grouped_gemm_for_grouped_tensor``,
-          which is backed by ``nvte_grouped_gemm_with_discrete_inputA`` in the common
-          library. This filter mirrors cuBLASLt grouped GEMM's architecture
-          requirement without duplicating its cuBLAS version checks.
-        * Quantized compute supports MXFP8 and NVFP4 on Blackwell GPUs with Compute Capability (CC)
-          10.x and 11.0. NVFP4 requires RHT because graph-safe grouped quantization currently
-          requires RHT. NVFP4 is additionally restricted to discrete weights: with
-          ``single_grouped_weight=True`` the weight quantizer is non-RHT and cannot use the
-          graph-safe grouped quantize kernel, so we fall back to the split-quantize flow.
-        * FP8 per-tensor current scaling is backed by grouped current-scaling quantization
-          (``tex.group_quantize``) and cuBLASLt grouped GEMM with per-batch scalar FP8 scaling,
-          which are supported on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0).
-        * FP8 block scaling uses the grouped-tensor path on Hopper (CC 9.0) with cuBLAS 13.4+;
-          it is Hopper-only (no MXFP8-broadcast emulation), so elsewhere it falls back.
-        * Every other quantization recipe (fp8 delayed scaling, ...) falls back to the legacy flow
-          because the corresponding grouped quantization kernels are missing.
-        * Unquantized compute supports BF16/FP16 on Hopper (CC 9.0) and Blackwell (CC 10.x and 11.0)
-          -- FP32 is excluded because the cuBLASLt grouped GEMM doesn't support it.
-        * Input/weight/grad_output quantizers are assumed to be of the same type, otherwise it
-          would trigger a fatal error in the cuBLASLt grouped GEMM check.
-        """
-        if not (9, 0) <= get_device_compute_capability() <= (11, 0):
-            return False
-        if with_quantized_compute:
-            # FP8 per-tensor current scaling runs on the Hopper and Blackwell grouped GEMM
-            # path; the compute-capability range was already checked above. On Hopper it
-            # requires cuBLAS 13.5+; fall back to the legacy flow on older cuBLAS.
-            if all(isinstance(q, Float8CurrentScalingQuantizer) for q in input_quantizers):
-                if (
-                    get_device_compute_capability() < (10, 0)
-                    and tex.get_cublasLt_version() < 130500
-                ):
-                    return False
-                return True
-            if all(isinstance(q, Float8BlockQuantizer) for q in input_quantizers):
-                # Grouped FP8 block scaling is Hopper-only and needs cuBLAS 13.4+; elsewhere
-                # fall back to the split-quantize (MXFP8-emulated) flow.
-                if get_device_compute_capability() >= (10, 0):
-                    return False
-                return tex.get_cublasLt_version() >= 130400
-            # MXFP8 and NVFP4 grouped quantization kernels require Blackwell.
-            if not (10, 0) <= get_device_compute_capability() <= (11, 0):
-                return False
-            if all(isinstance(q, MXFP8Quantizer) for q in input_quantizers):
-                return True
-            # NVFP4 graph-safe grouped quantization requires RHT and only supports
-            # discrete weights; otherwise fall back to the split-quantize flow.
-            if all(isinstance(q, NVFP4Quantizer) and q.with_rht for q in input_quantizers):
-                return not single_grouped_weight
-            return False
-        return dtype in (torch.bfloat16, torch.float16)
-
     def _get_grouped_weight_for_gemm(
         self,
         weight_param: GroupedTensor,
@@ -974,16 +983,7 @@ def _get_grouped_bias_for_gemm(
             return None
         num_groups = self.num_groups
 
-        if self.single_grouped_bias:
-            # Already a contiguous (num_groups * out_features) buffer.
-            bias_data = self.bias.rowwise_data
-            if bias_data.dtype != dtype:
-                bias_data = bias_data.to(dtype=dtype)
-        else:
-            bias_list = [
-                maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(num_groups)
-            ]
-            bias_data = torch.stack(bias_list, dim=0).contiguous()
+        bias_data = self._get_packed_bias_tensor(dtype)
 
         return GroupedTensorStorage(
             shape=(num_groups, self.out_features),
@@ -1052,16 +1052,22 @@ def fuser_forward(
         out_buffer = basic_op_kwargs[0].get(OUTPUT_BUFFER_KEY)
 
         # Dispatch: graph-safe GroupedTensor flow whenever it can be used.
-        # See ``_is_graph_safe_path_supported`` for the gating rationale --
+        # See ``is_op_fuser_grouped_tensor_path_supported`` for the gating rationale --
         # in short it requires Hopper (SM90+) plus a supported dtype /
         # quantization recipe. Otherwise we fall back to the legacy
         # ``tex.split_quantize`` + ``general_grouped_gemm`` flow.
-        use_grouped_tensor_path = self._is_graph_safe_path_supported(
-            with_quantized_compute=with_quantized_compute,
-            input_quantizers=input_quantizers,
-            dtype=dtype,
-            single_grouped_weight=self.single_grouped_weight,
+        recipe = FP8GlobalStateManager.get_fp8_recipe() if with_quantized_compute else None
+        use_grouped_tensor_path = is_op_fuser_grouped_tensor_path_supported(
+            recipe,
+            dtype,
         )
+        if (self.single_grouped_weight or self.single_grouped_bias) and not use_grouped_tensor_path:
+            raise RuntimeError(
+                "Single grouped parameters require the native grouped-tensor GroupedLinear path, "
+                "which is unavailable for the current device, dtype, or quantization recipe. "
+                "Disable single_grouped_weight/single_grouped_bias or use a supported grouped-"
+                "tensor configuration."
+            )
 
         if use_grouped_tensor_path:
             out, tensors_to_save = self._fuser_forward_grouped_tensor(
@@ -1215,16 +1221,11 @@ def _fuser_forward_split_quantize(
         # Need CPU split sizes for split_quantize / general_grouped_gemm.
         split_sizes_int = [int(s) for s in split_sizes.tolist()]
 
-        # Extract params
-        if self.single_grouped_weight:
-            weights = self.weight.quantized_tensors
-            if weights is None:
-                weights = self.weight.split_into_quantized_tensors()
-        else:
-            weights = self._forward_weight_list()  # materialized when distributed
+        # Single grouped parameters are rejected before entering this legacy path.
+        weights = self._forward_weight_list()  # materialized when distributed
         bs = None
         if has_bias:
-            bs = self._get_bias_tensors(dtype)
+            bs = self._get_discrete_bias_tensors(dtype)
 
         ws = self._get_discrete_weights_for_gemm(
             weights,
@@ -1556,7 +1557,7 @@ def _fuser_backward_split_quantize(
             offsets = torch.zeros(num_groups + 1, dtype=torch.int64, device=device)
             offsets[1:] = split_sizes.cumsum(0)
             if self._scale_bias:
-                bias_packed = torch.stack(self._get_bias_tensors(ctx.dtype))
+                bias_packed = self._get_packed_bias_tensor(ctx.dtype)
                 scales_f32 = scales.to(dtype=torch.float32)
                 dbias_packed, grad_scales = compute_grouped_dbias_dscales(
                     dy_2d,
@@ -1811,7 +1812,7 @@ def _fuser_backward_grouped_tensor(
         grad_scales: Optional[torch.Tensor] = None
         if has_bias:
             if self._scale_bias:
-                bias_packed = torch.stack(self._get_bias_tensors(dtype))
+                bias_packed = self._get_packed_bias_tensor(dtype)
                 scales_f32 = scales.to(dtype=torch.float32)
                 dbias_packed, grad_scales = compute_grouped_dbias_dscales(
                     dy_2d,
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 909e5a8a9b..3582d35ccd 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -2217,8 +2217,7 @@ def fuser_backward(
         fc2_bias_grads: Optional[list[Optional[torch.Tensor]]] = None
         fc2_bias_grad_packed: Optional[torch.Tensor] = None
         if scale_bias:
-            fc2_biases = fc2_op._get_bias_tensors(dtype)
-            bias_packed = torch.stack(fc2_biases)
+            bias_packed = fc2_op._get_packed_bias_tensor(dtype)
             fc2_dbias_packed_result, grad_scales = compute_grouped_dbias_dscales(
                 fc2_dy,
                 scales_f32,

From 67c26e252efc971e15a949945cb2095e6ad5f480 Mon Sep 17 00:00:00 2001
From: Andrew White <andrewh@cdw.com>
Date: Thu, 13 Aug 2026 11:53:34 -0500
Subject: [PATCH 61/67] fix: allgather striped window_size (5,0) id says (8,0)
 (#3374)

Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
---
 tests/jax/test_distributed_fused_attn.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/jax/test_distributed_fused_attn.py b/tests/jax/test_distributed_fused_attn.py
index 9678195942..c1cc2b8282 100644
--- a/tests/jax/test_distributed_fused_attn.py
+++ b/tests/jax/test_distributed_fused_attn.py
@@ -539,7 +539,7 @@ def check_has_backend_for_mask(mask_type):
         "window_size",
         [
             pytest.param((-1, -1), id="window_size(-1, -1)"),
-            pytest.param((5, 0), id="window_size(8, 0)"),
+            pytest.param((5, 0), id="window_size(5, 0)"),
         ],
     )
     @pytest.mark.parametrize(

From 910be8b2f80d4ecd9399ab4c30cce36536d87b17 Mon Sep 17 00:00:00 2001
From: Andrew White <andrewh@cdw.com>
Date: Thu, 13 Aug 2026 11:56:41 -0500
Subject: [PATCH 62/67] fix: remove redundant second export of XLA_FLAGS in L0
 (#3377)

Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
---
 qa/L0_jax_unittest/test.sh | 1 -
 1 file changed, 1 deletion(-)

diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh
index 0be19bba67..3e4742d8ae 100644
--- a/qa/L0_jax_unittest/test.sh
+++ b/qa/L0_jax_unittest/test.sh
@@ -40,7 +40,6 @@ pip3 install -r $TE_PATH/examples/jax/encoder/requirements.txt || error_exit "Fa
 export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops"
 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py"
 # Test without custom calls
-export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops"
 NVTE_JAX_CUSTOM_CALLS="false" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder_without_custom_call.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py without custom calls"
 
 # Exercise the docs/examples/jax tutorials. The multi-GPU tests are

From 900ac82aab483feb7a0cbd9a27ae9da4a515e007 Mon Sep 17 00:00:00 2001
From: Andrew White <andrewh@cdw.com>
Date: Thu, 13 Aug 2026 12:06:29 -0500
Subject: [PATCH 63/67] fix: remove unused `import time` (#3364)

Squashed to single commit for review.

Original PR: https://github.com/andrewwhitecdw/TransformerEngine/pull/14

Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
---
 tests/pytorch/attention/test_attention_with_cp.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py
index 2c0a5d9217..3c86f99cef 100644
--- a/tests/pytorch/attention/test_attention_with_cp.py
+++ b/tests/pytorch/attention/test_attention_with_cp.py
@@ -9,7 +9,6 @@
 import subprocess
 import sys
 import threading
-import time
 import pathlib
 import logging
 import copy

From 7d0365b35859f825e70215b95c77cbf8b09ebe7c Mon Sep 17 00:00:00 2001
From: Sudhakar Singh <sudhakars@nvidia.com>
Date: Thu, 13 Aug 2026 10:13:31 -0700
Subject: [PATCH 64/67] Stabilize FlashAttention CI environments (#3341)

* Gate FA4 and stabilize attention test imports

FA4 can be installed on SM8x even though its current implementation rejects those GPUs. Disable selection and skip dedicated FA4 tests there so A100 and L40 use supported attention backends.

FA4 and CUTLASS can also expose a generic utils package on sys.path. Prepend the Transformer Engine test helper directory in the context-parallel test so collection resolves the intended utilities.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Isolate FlashAttention CI backends

Moving images can install FA4 alongside older FlashAttention generations, which mixes a shared Python namespace and can make context-parallel reference runs compile an unsupported backend. Isolate the L3 version matrix, keep current CP comparisons on FA2/FA3, and temporarily reject symmetric D512 FA4 on Blackwell until upstream kernel support is complete.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Prepend PyTorch test utility imports

FA4 and its CUTLASS dependency expose a top-level utils module after Transformer Engine imports. Appending the test root can therefore bind these late imports to the installed module and fail collection. Give the repository helper precedence in the four test files that exhibited this ordering.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Recognize FA3 sliding-window CP support

The all-gather and a2a guards use the FA2 package version check to recognize FlashAttention support, so an isolated FA3 run is rejected even though FA3 implements sliding-window attention. Accept the explicit FA3 backend in both guards.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Keep CP backend selection at CI boundaries

The CP runner must honor an explicit generation selected by its caller, particularly the existing B200 L3 FA4 lane. Remove its internal V4 override, restore the L3 SM100 selection changed in 0f6c71eb, and disable V4 only for the L1 suite that still targets FA2/FA3. This keeps per-generation L3 isolation intact without making the shared runner silently override directed coverage.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Align FlashAttention CI coverage by architecture

Keep L0 on the mature FA2 path while L3 owns newer-generation coverage. Restrict H100 L3 to FA3 and B200 L3 to non-CP FA4 so unsupported H100 FA4 kernels and mislabeled Blackwell CP results do not obscure the intended signal. Make FA4-specific tests honor backend enablement to prevent silent fallback under an FA4 label.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

* Guard FA4 against incompatible CUTLASS installs

Package metadata can report FA4 present even when a later dependency install leaves its transitive CUTLASS stack unusable. Reject the known b24/CUTLASS combination below the stable 4.6.2 release and treat a nested interface ImportError as an unavailable optional backend so unrelated Transformer Engine imports can continue.

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>

---------

Signed-off-by: Sudhakar Singh <sudhakars@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 qa/L0_pytorch_debug_unittest/test.sh          |  5 ++
 qa/L0_pytorch_unittest/test.sh                |  5 ++
 qa/L1_pytorch_distributed_unittest/test.sh    | 10 ++--
 qa/L3_pytorch_FA_versions_test/test.sh        | 48 ++++++++++++++-----
 tests/pytorch/attention/test_attention.py     | 33 +++++++------
 .../attention/test_attention_with_cp.py       |  7 ++-
 tests/pytorch/attention/test_kv_cache.py      |  3 +-
 tests/pytorch/distributed/test_fusible_ops.py |  3 +-
 .../test_fusible_ops_with_userbuffers.py      |  3 +-
 tests/pytorch/distributed/test_sanity.py      |  3 +-
 .../dot_product_attention/backends.py         | 34 ++++++++++---
 .../dot_product_attention/context_parallel.py |  7 ++-
 .../attention/dot_product_attention/utils.py  | 19 ++++++--
 13 files changed, 129 insertions(+), 51 deletions(-)

diff --git a/qa/L0_pytorch_debug_unittest/test.sh b/qa/L0_pytorch_debug_unittest/test.sh
index 3efa462628..36efe485f5 100644
--- a/qa/L0_pytorch_debug_unittest/test.sh
+++ b/qa/L0_pytorch_debug_unittest/test.sh
@@ -22,6 +22,11 @@ FAILED_CASES=""
 : ${XML_LOG_DIR:=/logs}
 mkdir -p "$XML_LOG_DIR"
 
+# L0 keeps one mature FlashAttention generation; L3 owns newer-generation coverage.
+export NVTE_FLASH_ATTN_V2=1
+export NVTE_FLASH_ATTN_V3=0
+export NVTE_FLASH_ATTN_V4=0
+
 # Config with the dummy feature which prevents nvinspect from being disabled.
 # Nvinspect will be disabled if no feature is active.
 : ${NVTE_TEST_NVINSPECT_DUMMY_CONFIG_FILE:=$TE_PATH/tests/pytorch/debug/test_configs/dummy_feature.yaml}
diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh
index 759432857a..ee3ef89e66 100644
--- a/qa/L0_pytorch_unittest/test.sh
+++ b/qa/L0_pytorch_unittest/test.sh
@@ -22,6 +22,11 @@ set -x
 : ${XML_LOG_DIR:=/logs}
 mkdir -p "$XML_LOG_DIR"
 
+# L0 keeps one mature FlashAttention generation; L3 owns newer-generation coverage.
+export NVTE_FLASH_ATTN_V2=1
+export NVTE_FLASH_ATTN_V3=0
+export NVTE_FLASH_ATTN_V4=0
+
 pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
 
 NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py"
diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh
index 90ff6ba2fd..ec19492ee7 100644
--- a/qa/L1_pytorch_distributed_unittest/test.sh
+++ b/qa/L1_pytorch_distributed_unittest/test.sh
@@ -24,20 +24,22 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest"
 
 # Run CP tests (deterministic + non-deterministic) first so they can be parallelized.
 # Each needs 4 GPUs, so >=8 GPUs allows them to run concurrently on disjoint GPU sets.
+# Main's CP implementation supports FA2/FA3. Keep FA4 disabled at the suite
+# boundary so both the reference and CP halves use the same backend generation.
 NUM_GPUS=$(python3 -c "import torch; print(torch.cuda.device_count())")
 echo "Detected $NUM_GPUS GPU(s)"
 if [ "$NUM_GPUS" -ge 8 ]; then
     echo "Running CP tests in parallel: non-deterministic on GPUs 0-3, deterministic on GPUs 4-7"
-    CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py &
+    CUDA_VISIBLE_DEVICES=0,1,2,3 NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py &
     PID_CP_NONDET=$!
-    CUDA_VISIBLE_DEVICES=4,5,6,7 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py &
+    CUDA_VISIBLE_DEVICES=4,5,6,7 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py &
     PID_CP_DET=$!
     wait $PID_CP_NONDET || test_fail "test_attention_with_cp.py"
     wait $PID_CP_DET || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py"
 else
     echo "Running CP tests sequentially: need >=8 GPUs for parallel execution"
-    python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py"
-    NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py"
+    NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py"
+    NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FLASH_ATTN_V4=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py"
 fi
 
 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/distributed/test_sanity.py || test_fail "test_sanity.py"
diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh
index 468ef04d76..047358b301 100644
--- a/qa/L3_pytorch_FA_versions_test/test.sh
+++ b/qa/L3_pytorch_FA_versions_test/test.sh
@@ -31,37 +31,57 @@ export NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1
 # Iterate over Flash Attention versions
 sm_arch=`python3 -c "import torch; sm = torch.cuda.get_device_capability(0); print(sm[0]*10+sm[1])"`
 export FLASH_ATTN_CUDA_ARCHS=$sm_arch
-# CP tests are expensive and run only once per arch:
-#   - sm90 (H100):  FA3 (3.0.0b1) - context_parallel.py only supports FA3 on Hopper
-#   - sm>90 (B200): latest FA4    - FA3 is not built/installed for sm>90
-# Non-CP tests still run for every FA version in the array.
+# Run one architecture-owned FlashAttention generation. CP remains FA3-only
+# until the production selector and runner support FA4 CP end to end.
+CP_FA_VERSION=""
 if [ $sm_arch -gt 90 ]
 then
-  FA_versions=(2.8.3 4.0.0b11)
-  CP_FA_VERSION="${FA_versions[-1]}"
+  FA_versions=(4.0.0b11)
 elif [ $sm_arch -eq 90 ]
 then
-  FA_versions=(2.8.3 3.0.0b1 4.0.0b11)
+  FA_versions=(3.0.0b1)
   CP_FA_VERSION="3.0.0b1"
+else
+  error_exit "No L3 FlashAttention generation is defined for sm${sm_arch}"
 fi
 
 for fa_version in "${FA_versions[@]}"
 do
 
+  # The FA distributions share the flash_attn namespace. Keep exactly one
+  # installed so import-time discovery and the iteration label cannot disagree.
+  pip3 uninstall -y flash-attn flash-attn-3 flash-attn-4 \
+    || error_exit "Failed to isolate Flash Attention $fa_version"
+  export NVTE_FLASH_ATTN_V2=0
+  export NVTE_FLASH_ATTN_V3=0
+  export NVTE_FLASH_ATTN_V4=0
+
   # Build Flash Attention
   if [ "${fa_version}" \< "3.0.0" ]
   then
-    pip3 install flash-attn==${fa_version} --no-build-isolation
+    export NVTE_FLASH_ATTN_V2=1
+    pip3 install flash-attn==${fa_version} --no-build-isolation \
+      || error_exit "Failed to install Flash Attention $fa_version"
   elif [[ "${fa_version}" == 4.* ]]
   then
-    pip3 install flash-attn-4==${fa_version} nvidia-cutlass-dsl[cu13]==4.4.2 --no-build-isolation
+    export NVTE_FLASH_ATTN_V4=1
+    # FA4 is intentionally last in every version array. Its b11 test pin needs
+    # CUTLASS DSL 4.4.2, so replace the image-matched stack only for this final
+    # iteration; later iterations would otherwise need that stack restored.
+    pip3 uninstall -y nvidia-cutlass-dsl nvidia-cutlass-dsl-libs-base \
+      nvidia-cutlass-dsl-libs-cu12 nvidia-cutlass-dsl-libs-cu13 \
+      || error_exit "Failed to isolate CUTLASS DSL for Flash Attention $fa_version"
+    pip3 install flash-attn-4==${fa_version} nvidia-cutlass-dsl[cu13]==4.4.2 \
+      --no-build-isolation || error_exit "Failed to install Flash Attention $fa_version"
   else
+    export NVTE_FLASH_ATTN_V3=1
     # FA3 source build (~20 min). Skip if FA3 is already installed.
     if python3 -c "import flash_attn_3" 2>/dev/null; then
       echo "FA3 already installed (from base image); skipping source build"
     else
       git clone https://github.com/Dao-AILab/flash-attention.git
-      cd flash-attention/hopper && python setup.py install
+      cd flash-attention/hopper && python setup.py install \
+        || error_exit "Failed to install Flash Attention $fa_version"
       cd ../../
     fi
   fi
@@ -82,7 +102,7 @@ do
 
   # test_attention.py reloads its own trusted delayed-scaling FP8 checkpoint,
   # whose legacy extra state requires an explicit pickle opt-in.
-  if [ "$fa_version" = "$CP_FA_VERSION" ]; then
+  if [ -n "$CP_FA_VERSION" ] && [ "$fa_version" = "$CP_FA_VERSION" ]; then
     echo "Running CP tests with FA $fa_version (CP version for sm$sm_arch)"
     if [ "$NUM_GPUS" -ge 5 ]; then
       CP_NUM_GPUS=$(( NUM_GPUS - 1 > 4 ? 4 : NUM_GPUS - 1 ))
@@ -107,7 +127,11 @@ do
       NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s --junitxml=$XML_CP $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py (FA $fa_version)"
     fi
   else
-    echo "Skipping CP tests for FA $fa_version (CP only runs with FA $CP_FA_VERSION on sm$sm_arch)"
+    if [ -n "$CP_FA_VERSION" ]; then
+      echo "Skipping CP tests for FA $fa_version (CP uses FA $CP_FA_VERSION on sm$sm_arch)"
+    else
+      echo "CP tests are not scheduled for the FA generation on sm$sm_arch"
+    fi
     NVTE_TORCH_COMPILE=0 NVTE_ALLOW_UNSAFE_PICKLE_EXTRA_STATE=1 python3 -m pytest -v -s --junitxml=$XML_ATTN $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py (FA $fa_version)"
   fi
 done
diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py
index 82791084d8..c4f852f112 100644
--- a/tests/pytorch/attention/test_attention.py
+++ b/tests/pytorch/attention/test_attention.py
@@ -341,9 +341,18 @@ def test_dpa_num_splits(dtype, model_configs, model):
 }
 
 
-@pytest.mark.skipif(
-    not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required."
+fa4_enabled = bool(int(os.getenv("NVTE_FLASH_ATTN", "1"))) and bool(
+    int(os.getenv("NVTE_FLASH_ATTN_V4", "1"))
+)
+requires_fa4 = pytest.mark.skipif(
+    not fa4_enabled
+    or not FlashAttentionUtils.v4_is_installed
+    or device_compute_capability < (9, 0),
+    reason="Enabled Flash-attn v4 and compute capability >= SM90 are required.",
 )
+
+
+@requires_fa4
 @pytest.mark.parametrize("dtype", param_types_lean)
 @pytest.mark.parametrize("model_configs", [model_configs_fa4_base])
 @pytest.mark.parametrize("model", model_configs_fa4_base.keys())
@@ -362,9 +371,7 @@ def test_dpa_fa4_base(dtype, model_configs, model):
 }
 
 
-@pytest.mark.skipif(
-    not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required."
-)
+@requires_fa4
 @pytest.mark.skipif(
     device_compute_capability not in ((10, 0), (10, 3)),
     reason="FA4 head_dim=256 dedicated kernel is SM100/103-only.",
@@ -442,9 +449,7 @@ def test_dpa_d256(dtype, model_configs, model, qkv_layout):
 }
 
 
-@pytest.mark.skipif(
-    not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required."
-)
+@requires_fa4
 @pytest.mark.parametrize("dtype", param_types_lean)
 @pytest.mark.parametrize("model_configs", [model_configs_fa4_mla])
 @pytest.mark.parametrize("model", model_configs_fa4_mla.keys())
@@ -466,9 +471,7 @@ def test_dpa_fa4_mla(dtype, model_configs, model):
 }
 
 
-@pytest.mark.skipif(
-    not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required."
-)
+@requires_fa4
 @pytest.mark.parametrize("dtype", param_types_lean)
 @pytest.mark.parametrize("model_configs", [model_configs_fa4_swa])
 @pytest.mark.parametrize("model", model_configs_fa4_swa.keys())
@@ -489,9 +492,7 @@ def test_dpa_fa4_sliding_window(dtype, model_configs, model, qkv_layout):
 }
 
 
-@pytest.mark.skipif(
-    not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required."
-)
+@requires_fa4
 @pytest.mark.parametrize("dtype", param_types_lean)
 @pytest.mark.parametrize("model_configs", [model_configs_fa4_varlen])
 @pytest.mark.parametrize("model", model_configs_fa4_varlen.keys())
@@ -514,9 +515,7 @@ def test_dpa_fa4_varlen(dtype, model_configs, model, qkv_layout):
 }
 
 
-@pytest.mark.skipif(
-    not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required."
-)
+@requires_fa4
 @pytest.mark.parametrize("dtype", param_types_lean)
 @pytest.mark.parametrize("model_configs", [model_configs_fa4_mask])
 @pytest.mark.parametrize("model", model_configs_fa4_mask.keys())
diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py
index 3c86f99cef..d7eb16b862 100644
--- a/tests/pytorch/attention/test_attention_with_cp.py
+++ b/tests/pytorch/attention/test_attention_with_cp.py
@@ -28,7 +28,7 @@
 from transformer_engine.pytorch.attention.dot_product_attention.utils import FlashAttentionUtils
 
 _current_file = pathlib.Path(__file__).resolve()
-sys.path.append(str(_current_file.parent.parent))
+sys.path = [str(_current_file.parent.parent)] + sys.path
 from utils import ModelConfig, get_available_attention_backends
 
 pytest_logging_level = logging.getLevelName(logging.root.level)
@@ -299,7 +299,10 @@ def _submit(pool: PoolWorker, **kwargs) -> None:
     qkv_formats = ["sbhd", "thd"]
 
 
-@pytest.mark.skipif(not FlashAttentionUtils.v2_plus, reason="Flash-attn 2.0+ is required.")
+@pytest.mark.skipif(
+    not (FlashAttentionUtils.v2_plus or FlashAttentionUtils.v3_is_installed),
+    reason="Flash-attn v2 or v3 is required.",
+)
 @pytest.mark.skipif(get_device_compute_capability() < (8, 0), reason="CP tests require sm80+.")
 @pytest.mark.parametrize("dtype", dtypes)
 @pytest.mark.parametrize("model", model_configs_flash_attn.keys())
diff --git a/tests/pytorch/attention/test_kv_cache.py b/tests/pytorch/attention/test_kv_cache.py
index cdd98d2445..2a857a10dc 100644
--- a/tests/pytorch/attention/test_kv_cache.py
+++ b/tests/pytorch/attention/test_kv_cache.py
@@ -33,7 +33,8 @@
 )
 
 _current_file = pathlib.Path(__file__).resolve()
-sys.path.append(str(_current_file.parent.parent))
+# Prepend so installed packages with a top-level utils module cannot shadow the test helpers.
+sys.path = [str(_current_file.parent.parent)] + sys.path
 from utils import (
     ModelConfig,
     reset_rng_states,
diff --git a/tests/pytorch/distributed/test_fusible_ops.py b/tests/pytorch/distributed/test_fusible_ops.py
index 9d08a81a8f..d733286093 100644
--- a/tests/pytorch/distributed/test_fusible_ops.py
+++ b/tests/pytorch/distributed/test_fusible_ops.py
@@ -32,7 +32,8 @@
 
 # Import utility functions
 _current_file = pathlib.Path(__file__).resolve()
-sys.path.append(str(_current_file.parent.parent))
+# Prepend so installed packages with a top-level utils module cannot shadow the test helpers.
+sys.path = [str(_current_file.parent.parent)] + sys.path
 from utils import dtype_tols, make_recipe, quantization_tols
 
 
diff --git a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py
index eb43ba7e75..07dffebf5f 100644
--- a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py
+++ b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py
@@ -37,7 +37,8 @@
 
 # Import utility functions
 _current_file = pathlib.Path(__file__).resolve()
-sys.path.append(str(_current_file.parent.parent))
+# Prepend so installed packages with a top-level utils module cannot shadow the test helpers.
+sys.path = [str(_current_file.parent.parent)] + sys.path
 from utils import dtype_tols, make_recipe, run_distributed, str_to_dtype
 
 # Check if FP8 is supported
diff --git a/tests/pytorch/distributed/test_sanity.py b/tests/pytorch/distributed/test_sanity.py
index 2e7a63e0a2..1b8bf5890f 100644
--- a/tests/pytorch/distributed/test_sanity.py
+++ b/tests/pytorch/distributed/test_sanity.py
@@ -19,7 +19,8 @@
 from transformer_engine.common import recipe
 
 _current_file = pathlib.Path(__file__).resolve()
-sys.path.append(str(_current_file.parent.parent))
+# Prepend so installed packages with a top-level utils module cannot shadow the test helpers.
+sys.path = [str(_current_file.parent.parent)] + sys.path
 from utils import ModelConfig
 
 model_configs = {
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
index aeac501ff8..aa3bc15e84 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py
@@ -167,14 +167,34 @@
     flash_attn_func_v4 = None
     flash_attn_varlen_func_v4 = None
 else:
-    from flash_attn.cute.interface import (  # pylint: disable=ungrouped-imports,no-name-in-module
-        flash_attn_func as flash_attn_func_v4,
-        flash_attn_varlen_func as flash_attn_varlen_func_v4,
-        _validate_head_dims as _fa4_validate_head_dims,
-    )
+    try:
+        cutlass_dsl_version = PkgVersion(get_pkg_version("nvidia-cutlass-dsl"))
+
+        # FA4 4.0.0b24 requires CUTLASS DSL 4.6.2 or newer.
+        if fa_utils.fa4_version == PkgVersion("4.0.0b24") and cutlass_dsl_version < PkgVersion(
+            "4.6.2"
+        ):
+            raise ImportError(
+                "flash-attn-4 4.0.0b24 requires nvidia-cutlass-dsl>=4.6.2; "
+                f"found {cutlass_dsl_version}"
+            )
 
-    fa_utils.v4_validate_head_dims = _fa4_validate_head_dims
-    fa_utils.set_flash_attention_4_params()
+        from flash_attn.cute.interface import (  # pylint: disable=ungrouped-imports,no-name-in-module
+            flash_attn_func as flash_attn_func_v4,
+            flash_attn_varlen_func as flash_attn_varlen_func_v4,
+            _validate_head_dims as _fa4_validate_head_dims,
+        )
+    except ImportError as exc:
+        flash_attn_func_v4 = None
+        flash_attn_varlen_func_v4 = None
+        warnings.warn(
+            f"FlashAttention 4 is installed but cannot be loaded: {exc}",
+            RuntimeWarning,
+            stacklevel=2,
+        )
+    else:
+        fa_utils.v4_validate_head_dims = _fa4_validate_head_dims
+        fa_utils.set_flash_attention_4_params()
 
 # Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16
 _dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1"
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
index b444f034c8..ea89ca97eb 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py
@@ -3087,10 +3087,12 @@ def forward(
             window_size == (-1, 0)
             or window_size == (-1, -1)
             or use_fused_attention
+            or use_flash_attn_3
             or fa_utils.v2_3_plus
         ), (
             "cp_comm_type='all_gather' only supports SWA through FusedAttention or FlashAttention"
-            f" >= 2.3. Found {use_fused_attention=} and {fa_utils.v2_3_plus=}."
+            f" >= 2.3. Found {use_fused_attention=}, {use_flash_attn_3=}, "
+            f"and {fa_utils.v2_3_plus=}."
         )
         assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, (
             "cp_comm_type='all_gather' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q ="
@@ -4231,10 +4233,11 @@ def forward(
             window_size == (-1, 0)
             or window_size == (-1, -1)
             or use_fused_attention
+            or use_flash_attn_3
             or fa_utils.v2_3_plus
         ), (
             "cp_comm_type='a2a' only supports SWA through FusedAttention or FlashAttention >= 2.3."
-            f" Found {use_fused_attention=} and {fa_utils.v2_3_plus=}."
+            f" Found {use_fused_attention=}, {use_flash_attn_3=}, and {fa_utils.v2_3_plus=}."
         )
         assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, (
             "cp_comm_type='a2a' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q ="
diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 6eb3ce54f1..e9fc03b174 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -563,10 +563,10 @@ def _disable_all_flash_attention() -> None:
         if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed:
             logger.debug("Disabling FlashAttention 3 for compute capability != sm90")
         use_flash_attention_3 = False
-    # FA4 supports SM80, SM90, SM100, SM120
-    if device_compute_capability < (8, 0):
+    # FA4 does not currently support SM8x.
+    if device_compute_capability < (9, 0):
         if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed:
-            logger.debug("Disabling FlashAttention 4 for compute capability < sm80")
+            logger.debug("Disabling FlashAttention 4 for compute capability < sm90")
         use_flash_attention_4 = False
     # On SM90, prefer FA3 over FA4 when FA3 is available.
     # FA3 is more mature on Hopper; FA4's SM90 backward has limitations
@@ -996,6 +996,19 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt
                 device_compute_capability[0] * 10 + device_compute_capability[1],
             )
             use_flash_attention_4 = False
+        # FA4's validator currently accepts symmetric (512, 512) on SM100/SM110,
+        # but the generic forward kernel exceeds its TMEM allocation for that shape.
+        # Preserve the supported asymmetric (64, 512) MLA path while D512 support
+        # is completed upstream.
+        if (
+            use_flash_attention_4
+            and (10, 0) <= device_compute_capability < (12, 0)
+            and head_dim_qk == head_dim_v == 512
+        ):
+            logger.debug(
+                "Disabling FlashAttention 4 for unsupported symmetric head_dim=512 on SM100/SM110."
+            )
+            use_flash_attention_4 = False
         # flash-attn-4 4.0.0b11 validates (256, 256) on SM100, but its dedicated
         # hd256 kernel diverges from the reference for cross-attention/decode-like
         # shapes such as sq=1, skv=2048. Keep FA4 enabled for the self-attention

From 76ebc3c5c809a3c6e49945eae786df74b87be9f8 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 13 Aug 2026 21:21:51 +0200
Subject: [PATCH 65/67] Skip get_fp8_recipe only while tracing in
 init_fp8_metadata

The unconditional early return skipped main's new CustomRecipe
local-recipes inference when quantization is off.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../dot_product_attention/dot_product_attention.py          | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index ae6f735e88..35a71d0768 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -921,8 +921,12 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None:
         """
         _original_recipe = self.fp8_meta.get("recipe", None)
 
+        # get_fp8_recipe() may build a default Recipe, which asserts it is not tracing.
+        # With quantization off the base class does all that is needed.
         qstate = FP8GlobalStateManager.quantization_state
-        if not (qstate.fp8_enabled or qstate.fp8_calibration or qstate.fp8_parameters):
+        if torch.compiler.is_compiling() and not (
+            qstate.fp8_enabled or qstate.fp8_calibration or qstate.fp8_parameters
+        ):
             super().init_fp8_metadata(num_gemms=num_gemms)
             return
 

From 075bce897be0671b5410a7a546a979e716bf17c9 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 17 Aug 2026 12:11:49 +0200
Subject: [PATCH 66/67] Skip the max_seqlen bake-in under ONNX export

guard_scalar specializes user-declared dynamic sequence dims in the
exported model; the exporter skips backend selection, which the bake-in
is for.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../attention/dot_product_attention/dot_product_attention.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
index 35a71d0768..d5adbbcadf 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py
@@ -1775,8 +1775,9 @@ def forward(
                     max_seqlen_q = query_layer.shape[1] if max_seqlen_q is None else max_seqlen_q
                     max_seqlen_kv = key_layer.shape[1] if max_seqlen_kv is None else max_seqlen_kv
                 # Backend selection bakes these in, which it cannot do symbolically.
-                max_seqlen_q = guard_scalar(max_seqlen_q)
-                max_seqlen_kv = guard_scalar(max_seqlen_kv)
+                if not is_in_onnx_export_mode():
+                    max_seqlen_q = guard_scalar(max_seqlen_q)
+                    max_seqlen_kv = guard_scalar(max_seqlen_kv)
             if qkv_format == "thd":
                 assert all(
                     len(x.shape) == 3 for x in (query_layer, key_layer, value_layer)

From b3e9949e85c9f78d041febe3ce78064a9df09e2d Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 17 Aug 2026 12:11:49 +0200
Subject: [PATCH 67/67] Make undeclared non-contiguous q/k/v contiguous under
 compile

The declared layout is only truthful for contiguous tensors. Eager
repairs such inputs via its detect-and-retry loop; mirror that in the
compiled branch instead of declaring a layout the memory contradicts.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/attention/dot_product_attention/utils.py             | 1 +
 1 file changed, 1 insertion(+)

diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
index 0893dfec4b..b5e5633e14 100644
--- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py
+++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py
@@ -2598,6 +2598,7 @@ def run_iteratively(q, k, v):
             " torch.compile. Pass the packed buffer explicitly via DotProductAttention's"
             " qkv_layer/kv_layer (with qkv_interleave_dim)."
         )
+        q, k, v = [x if x.is_contiguous() else x.contiguous() for x in (q, k, v)]
         if is_same_q_kv_format:
             qkv_layout = "_".join([qkv_format] * 3)
         else: