From cb69309b91c47bd98776cea60ec482b32d8af3bc Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 5 Aug 2026 22:21:38 +0200
Subject: [PATCH 01/27] [PyTorch] Register an op's forward and backward without
 autograd glue

register_custom_op now defines an operation's forward and backward as two
independent two-tier custom ops and hands back both, leaving autograd to the
caller. That is what lets a pipeline-level autograd.Function decide how the
two are wired, and so group the forward and backward passes differently --
which is what ops.OperationFuser does.

The variant that wires autograd itself keeps the old behaviour under
register_custom_op_with_autograd, and is now built on the same registration:
the pair is the primitive, autograd is what the other one adds. About two
thirds of the two bodies were the same code before.

BasicOperation gains the plumbing an operation needs to opt in: declare two
argument containers and implement four compute classmethods, and
__init_subclass__ registers the custom ops while op_forward / op_backward are
written once in the base. compile_unsupported_reason lets an operation say why
it cannot be compiled -- it sits here rather than on the args, as Linear has
it, because in ops/ the compile boundary is the fuser group, not the
operation.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/dynamo/__init__.py |   3 +-
 .../pytorch/dynamo/custom_op.py               | 343 +++++++++++++-----
 transformer_engine/pytorch/module/linear.py   |   4 +-
 transformer_engine/pytorch/ops/op.py          | 204 ++++++++++-
 4 files changed, 459 insertions(+), 95 deletions(-)

diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py
index e42eb8f9f6..083a2aa1fb 100644
--- a/transformer_engine/pytorch/dynamo/__init__.py
+++ b/transformer_engine/pytorch/dynamo/__init__.py
@@ -6,7 +6,7 @@
 
 from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer
 from .tensor_spec import TensorSpec, to_tensor_spec
-from .custom_op import register_custom_op, TensorOrQuantized
+from .custom_op import register_custom_op, register_custom_op_with_autograd, TensorOrQuantized
 
 __all__ = [
     "register_value_opaque_quantizer",
@@ -14,5 +14,6 @@
     "TensorSpec",
     "to_tensor_spec",
     "register_custom_op",
+    "register_custom_op_with_autograd",
     "TensorOrQuantized",
 ]
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 00846d615a..455a1d8150 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -6,7 +6,10 @@
 
 Registers TE modules' eager forward/backward as ``torch.library`` custom ops so
 ``torch.compile(fullgraph=True)`` traces them as single graph nodes.
-``register_custom_op`` is the entry point; ``module/linear.py`` is the first user.
+``register_custom_op_with_autograd`` is the entry point for a module that
+wires autograd on the op itself (``module/linear.py`` is the first user);
+``register_custom_op`` hands back the forward and backward ops separately, for a
+caller that drives autograd at a higher level (``ops/fuser.py``).
 
 A TE forward/backward implementation takes one dataclass argument
 (``fwd_arg_type`` / ``bwd_arg_type``, e.g. ``LinearFwdArgs``) whose fields mix
@@ -41,9 +44,9 @@
     only when its value is trivial (``None`` / all-``None``) at call time.
 
 What runs where. Each op registers a data-free fake (``register_fake``) so it
-traces under ``torch.compile`` without allocating. ``register_custom_op`` returns
-``forward_fn`` -- the drop-in for the eager ``autograd.Function.apply``. A forward
-call through it:
+traces under ``torch.compile`` without allocating.
+``register_custom_op_with_autograd`` returns ``forward_fn`` -- the drop-in for
+the eager ``autograd.Function.apply``. A forward call through it:
 
   * runs the fake ``fwd_fake_impl`` on ``TensorSpec`` descriptors (data-free; see
     ``tensor_spec.py``) and parses its result into an ``_OutputPlan`` -- the
@@ -930,7 +933,7 @@ def _slice_user_grads(
 
 
 # --------------------------------------------------------------------------- #
-# Op registration
+# Op registration: base and wrapper ops, autograd wiring
 # --------------------------------------------------------------------------- #
 
 
@@ -1166,7 +1169,233 @@ def _all_quantized_tensor_subclasses() -> List[type]:
     return found
 
 
+@dataclasses.dataclass(frozen=True)
+class _OpPair:
+    """One registered forward/backward pair, and what a caller needs to drive it."""
+
+    fwd_plan: _ArgPlan
+    bwd_plan: _ArgPlan
+    base_fwd_def: Any
+    base_bwd_op: Any
+    wrapper_fwd_def: Any
+    wrapper_fwd_op: Any
+    wrapper_bwd_op: Any
+
+    def call_forward(
+        self, fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], fwd_args: Any
+    ) -> Tuple[_OutputPlan, List[torch.Tensor]]:
+        """Run the forward op on ``fwd_args``: its output plan and flat payload."""
+        spec_obj = _spec_view(fwd_args, self.fwd_plan.tensor_field_names())
+        out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj))
+        kwargs = self.fwd_plan.pack(fwd_args)
+        payload = self.wrapper_fwd_op(*[kwargs[name] for name in self.fwd_plan.slot_names])
+        return out_plan, payload
+
+
+def _register_two_tier_pair(
+    *,
+    op_name: str,
+    fwd_arg_type: type,
+    fwd_impl: Callable[[Any], Any],
+    fwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
+    bwd_arg_type: type,
+    bwd_impl: Callable[[Any], Any],
+    bwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
+    num_grad_inputs: int,
+) -> _OpPair:
+    """Define an operation's forward and backward as two-tier custom ops.
+
+    Everything that is common to :func:`register_custom_op` and
+    :func:`register_custom_op_with_autograd`: the arg plans, the base kernels,
+    the wrapper ops that flatten ``QuantizedTensor`` subclass inputs, and the
+    passthrough registrations. Autograd is deliberately not touched here -- that
+    is what the two entry points differ on.
+    """
+    wrapper_fwd_name = op_name
+    wrapper_bwd_name = f"{op_name}_backward"
+    base_fwd_name = f"{op_name}_base"
+    base_bwd_name = f"{wrapper_bwd_name}_base"
+    subclass_list = _all_quantized_tensor_subclasses()
+
+    fwd_plan = _parse_arg_type(fwd_arg_type)
+    bwd_plan = _parse_arg_type(bwd_arg_type)
+
+    fwd_schema = f"{fwd_plan.schema_str} -> Tensor[]"
+    bwd_schema = f"{bwd_plan.schema_str} -> Tensor[]"
+
+    base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}"
+
+    base_fwd_def = _register_base_op(
+        op_name=base_fwd_name,
+        schema_str=fwd_schema,
+        plan=fwd_plan,
+        impl=fwd_impl,
+        fake_impl=fwd_fake_impl,
+        pack_result=_pack_fwd_result,
+    )
+    _register_base_op(
+        op_name=base_bwd_name,
+        schema_str=bwd_schema,
+        plan=bwd_plan,
+        impl=bwd_impl,
+        fake_impl=bwd_fake_impl,
+        pack_result=lambda g: _pack_bwd_result(g, num_grad_inputs, base_bwd_qualname),
+    )
+
+    base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name)
+    base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name)
+
+    fwd_slot_offsets = fwd_plan.tensor_or_quantized_offsets()
+    bwd_slot_offsets = bwd_plan.tensor_or_quantized_offsets()
+
+    wrapper_fwd_def = _register_wrapper_op(
+        wrapper_op_name=wrapper_fwd_name,
+        schema_str=fwd_schema,
+        base_op=base_fwd_op,
+        slot_offsets=fwd_slot_offsets,
+        subclasses=subclass_list,
+    )
+    # Pass-through: a subclass input reaches the base op through the dispatch
+    # rule below, never through the wrapper body.
+    wrapper_bwd_def = _register_wrapper_op(
+        wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op
+    )
+    wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name)
+    wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name)
+
+    _fwd_rule = _make_dispatch_rule(
+        _make_slot_forwarder(base_fwd_op, fwd_slot_offsets, subclass_list)
+    )
+    _bwd_rule = _make_dispatch_rule(
+        _make_slot_forwarder(base_bwd_op, bwd_slot_offsets, subclass_list)
+    )
+
+    for sub in subclass_list:
+        wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule)
+        wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule)
+
+    for op in (wrapper_fwd_op, wrapper_bwd_op, base_fwd_op, base_bwd_op):
+        _quantized_tensor_passthrough_ops.add(op.default)
+
+    return _OpPair(
+        fwd_plan=fwd_plan,
+        bwd_plan=bwd_plan,
+        base_fwd_def=base_fwd_def,
+        base_bwd_op=base_bwd_op,
+        wrapper_fwd_def=wrapper_fwd_def,
+        wrapper_fwd_op=wrapper_fwd_op,
+        wrapper_bwd_op=wrapper_bwd_op,
+    )
+
+
+# --------------------------------------------------------------------------- #
+# Op registration: the forward/backward pair, and the autograd-wired variant
+# --------------------------------------------------------------------------- #
+
+
 def register_custom_op(
+    *,
+    op_name: str,
+    fwd_arg_type: type,
+    fwd_impl: Callable[[Any], Any],
+    fwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
+    bwd_arg_type: type,
+    bwd_impl: Callable[[Any], Any],
+    bwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
+    num_grad_inputs: int,
+) -> Optional[Tuple[Callable[[Any], Any], Callable[[Any], Any]]]:
+    """Register an op's forward and backward as two independent custom ops.
+
+    Autograd is the caller's: it decides how the two are wired, which is what
+    lets a pipeline-level ``torch.autograd.Function`` -- traced by Dynamo as a
+    higher-order op -- group the forward and backward passes differently, as
+    ``ops.OperationFuser`` does. :func:`register_custom_op_with_autograd` builds
+    on this and wires them the usual way instead.
+
+    Both ops are two-tier, so ``QuantizedTensor`` subclass inputs pass through
+    without dequantization.
+
+    Contracts, mirroring :func:`register_custom_op_with_autograd`:
+
+    * ``fwd_impl(fwd_args) -> (*user_outputs, tensors_to_save, ctx_attrs)``
+    * ``fwd_fake_impl`` -- its data-free twin over :class:`TensorSpec`
+    * ``bwd_impl(bwd_args) -> tuple`` of ``num_grad_inputs`` gradients
+    * ``bwd_fake_impl`` -- its data-free twin
+
+    Returns ``(forward_fn, backward_fn)``:
+
+    * ``forward_fn(fwd_args) -> (outputs, saved_tensors, ctx_attrs)`` --
+      ``outputs`` is a single value or a tuple, mirroring ``fwd_impl``'s user
+      outputs; ``saved_tensors`` is the reassembled ``tensors_to_save`` tuple,
+      which the caller is expected to persist (e.g. ``ctx.save_for_backward``).
+    * ``backward_fn(bwd_args) -> tuple`` of gradients.
+
+    Returns ``None`` if registration fails (recorded once), so callers can fall
+    back to eager rather than breaking import.
+    """
+    try:
+        return _register_custom_op_impl(
+            op_name=op_name,
+            fwd_arg_type=fwd_arg_type,
+            fwd_impl=fwd_impl,
+            fwd_fake_impl=fwd_fake_impl,
+            bwd_arg_type=bwd_arg_type,
+            bwd_impl=bwd_impl,
+            bwd_fake_impl=bwd_fake_impl,
+            num_grad_inputs=num_grad_inputs,
+        )
+    except (ImportError, AttributeError, RuntimeError, TypeError) as e:
+        record_compile_disabled(
+            f"could not register the autograd-free custom ops '{op_name}' ({type(e).__name__}: {e})"
+        )
+        return None
+
+
+def _register_custom_op_impl(
+    *,
+    op_name: str,
+    fwd_arg_type: type,
+    fwd_impl: Callable[[Any], Any],
+    fwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
+    bwd_arg_type: type,
+    bwd_impl: Callable[[Any], Any],
+    bwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
+    num_grad_inputs: int,
+) -> Tuple[Callable[[Any], Any], Callable[[Any], Any]]:
+    """Body of :func:`register_custom_op`; see it for semantics."""
+    pair = _register_two_tier_pair(
+        op_name=op_name,
+        fwd_arg_type=fwd_arg_type,
+        fwd_impl=fwd_impl,
+        fwd_fake_impl=fwd_fake_impl,
+        bwd_arg_type=bwd_arg_type,
+        bwd_impl=bwd_impl,
+        bwd_fake_impl=bwd_fake_impl,
+        num_grad_inputs=num_grad_inputs,
+    )
+
+    def forward_fn(fwd_args):
+        out_plan, payload = pair.call_forward(fwd_fake_impl, fwd_args)
+        outputs = out_plan.user_outputs(payload)
+        saved = out_plan.saved_tensors(payload)
+        return (
+            (outputs[0] if len(outputs) == 1 else tuple(outputs)),
+            tuple(saved),
+            out_plan.ctx_attrs,
+        )
+
+    def backward_fn(bwd_args):
+        # Unlike the forward payload, each grad occupies exactly one slot
+        # (``_pack_bwd_result`` materializes a TensorSpec grad), so there is
+        # nothing to reassemble.
+        kwargs = pair.bwd_plan.pack(bwd_args)
+        payload = pair.wrapper_bwd_op(*[kwargs[name] for name in pair.bwd_plan.slot_names])
+        return tuple(_decode_none(t) for t in payload)
+
+    return forward_fn, backward_fn
+
+
+def register_custom_op_with_autograd(
     *,
     op_name: str,
     input_tensors_for_grad: List[str],
@@ -1231,7 +1460,7 @@ def register_custom_op(
     ``torch.compile`` (a graph break) rather than breaking import.
     """
     try:
-        return _register_custom_op_impl(
+        return _register_custom_op_with_autograd_impl(
             op_name=op_name,
             input_tensors_for_grad=input_tensors_for_grad,
             fwd_arg_type=fwd_arg_type,
@@ -1249,7 +1478,7 @@ def register_custom_op(
         return None
 
 
-def _register_custom_op_impl(
+def _register_custom_op_with_autograd_impl(
     *,
     op_name: str,
     input_tensors_for_grad: List[str],
@@ -1261,7 +1490,7 @@ def _register_custom_op_impl(
     fwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
     bwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
 ) -> Callable[..., Any]:
-    """Body of :func:`register_custom_op`; see it for semantics."""
+    """Body of :func:`register_custom_op_with_autograd`; see it for semantics."""
     # Existence check at the API boundary: every ``input_tensors_for_grad`` name
     # must be an actual field of ``fwd_arg_type`` (differentiability -- whether
     # that field can carry a gradient -- is checked later, in
@@ -1271,96 +1500,32 @@ def _register_custom_op_impl(
     if missing:
         raise ValueError(f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}")
 
-    wrapper_fwd_name = op_name
-    wrapper_bwd_name = f"{op_name}_backward"
-    base_fwd_name = f"{op_name}_base"
-    base_bwd_name = f"{wrapper_bwd_name}_base"
-    subclass_list = _all_quantized_tensor_subclasses()
-
-    fwd_plan = _parse_arg_type(fwd_arg_type)
-    bwd_plan = _parse_arg_type(bwd_arg_type)
-
-    num_grad_inputs = len(input_tensors_for_grad)
-    grad_targets = fwd_plan.resolve_grad_targets(input_tensors_for_grad)
-
-    fwd_schema = f"{fwd_plan.schema_str} -> Tensor[]"
-    bwd_schema = f"{bwd_plan.schema_str} -> Tensor[]"
-
-    base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}"
-
-    base_fwd_def = _register_base_op(
-        op_name=base_fwd_name,
-        schema_str=fwd_schema,
-        plan=fwd_plan,
-        impl=fwd_impl,
-        fake_impl=fwd_fake_impl,
-        pack_result=_pack_fwd_result,
-    )
-    _register_base_op(
-        op_name=base_bwd_name,
-        schema_str=bwd_schema,
-        plan=bwd_plan,
-        impl=bwd_impl,
-        fake_impl=bwd_fake_impl,
-        pack_result=lambda g: _pack_bwd_result(g, num_grad_inputs, base_bwd_qualname),
-    )
-
-    base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name)
-    base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name)
-
-    fwd_slot_offsets = fwd_plan.tensor_or_quantized_offsets()
-    bwd_slot_offsets = bwd_plan.tensor_or_quantized_offsets()
-
-    wrapper_fwd_def = _register_wrapper_op(
-        wrapper_op_name=wrapper_fwd_name,
-        schema_str=fwd_schema,
-        base_op=base_fwd_op,
-        slot_offsets=fwd_slot_offsets,
-        subclasses=subclass_list,
-    )
-    # Pass-through: a subclass input reaches the base op through the dispatch
-    # rule below, never through the wrapper body.
-    wrapper_bwd_def = _register_wrapper_op(
-        wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op
+    pair = _register_two_tier_pair(
+        op_name=op_name,
+        fwd_arg_type=fwd_arg_type,
+        fwd_impl=fwd_impl,
+        fwd_fake_impl=fwd_fake_impl,
+        bwd_arg_type=bwd_arg_type,
+        bwd_impl=bwd_impl,
+        bwd_fake_impl=bwd_fake_impl,
+        num_grad_inputs=len(input_tensors_for_grad),
     )
 
     autograd_common = {
-        "fwd_plan": fwd_plan,
-        "bwd_plan": bwd_plan,
-        "grad_targets": grad_targets,
+        "fwd_plan": pair.fwd_plan,
+        "bwd_plan": pair.bwd_plan,
+        "grad_targets": pair.fwd_plan.resolve_grad_targets(input_tensors_for_grad),
         "setup_context_user": setup_context,
         "fwd_fake_impl": fwd_fake_impl,
     }
-    wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name)
-    wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name)
-
-    _register_autograd_for_op(fwd_op=base_fwd_def, bwd_op=base_bwd_op, **autograd_common)
-    _register_autograd_for_op(fwd_op=wrapper_fwd_def, bwd_op=wrapper_bwd_op, **autograd_common)
-
-    _fwd_rule = _make_dispatch_rule(
-        _make_slot_forwarder(base_fwd_op, fwd_slot_offsets, subclass_list)
-    )
-    _bwd_rule = _make_dispatch_rule(
-        _make_slot_forwarder(base_bwd_op, bwd_slot_offsets, subclass_list)
+    _register_autograd_for_op(fwd_op=pair.base_fwd_def, bwd_op=pair.base_bwd_op, **autograd_common)
+    _register_autograd_for_op(
+        fwd_op=pair.wrapper_fwd_def, bwd_op=pair.wrapper_bwd_op, **autograd_common
     )
 
-    for sub in subclass_list:
-        wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule)
-        wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule)
-
-    _quantized_tensor_passthrough_ops.add(wrapper_fwd_op.default)
-    _quantized_tensor_passthrough_ops.add(wrapper_bwd_op.default)
-    _quantized_tensor_passthrough_ops.add(base_fwd_op.default)
-    _quantized_tensor_passthrough_ops.add(base_bwd_op.default)
-
     def forward_fn(fwd_args):
-        spec_obj = _spec_view(fwd_args, fwd_plan.tensor_field_names())
-        out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj))
-        kwargs = fwd_plan.pack(fwd_args)
-        flat_in = [kwargs[name] for name in fwd_plan.slot_names]
-        result = wrapper_fwd_op(*flat_in)
-
-        outputs = out_plan.user_outputs(result)
+        out_plan, payload = pair.call_forward(fwd_fake_impl, fwd_args)
+        outputs = out_plan.user_outputs(payload)
         if len(outputs) == 1:
             return outputs[0]
         return tuple(outputs)
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 55fc69ef7f..3b3c99facd 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -88,7 +88,7 @@
 from ..dynamo import (
     TensorSpec,
     TensorOrQuantized,
-    register_custom_op,
+    register_custom_op_with_autograd,
     is_value_opaque_quantizer,
 )
 from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer
@@ -1788,7 +1788,7 @@ def _linear_backward_fake(
 
 
 # Custom op used under ``torch.compile``.
-_linear_op = register_custom_op(
+_linear_op = register_custom_op_with_autograd(
     op_name="linear",
     input_tensors_for_grad=["weight", "inp", "bias"],
     fwd_arg_type=LinearFwdArgs,
diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py
index d057d46816..0fc6cb5c99 100644
--- a/transformer_engine/pytorch/ops/op.py
+++ b/transformer_engine/pytorch/ops/op.py
@@ -9,7 +9,7 @@
 from collections.abc import Iterable, Sequence
 import dataclasses
 import pickle
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import torch
 
@@ -22,6 +22,7 @@
     autocast,
 )
 from ..tensor import Quantizer
+from ..dynamo import is_value_opaque_quantizer, register_custom_op
 
 
 @dataclasses.dataclass
@@ -186,6 +187,45 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta):
     # Number of extra tensor outputs
     num_extra_outputs: int = 0
 
+    # torch.compile support. An operation opts in by declaring the two arg
+    # containers and implementing the four compute classmethods below; the base
+    # class then registers its custom ops and drives them from op_forward /
+    # op_backward, so no operation writes that plumbing itself.
+    fwd_args_type: Optional[type] = None
+    bwd_args_type: Optional[type] = None
+    # Gradients returned by backward_compute: the input's, then any parameters'.
+    num_grad_inputs: int = 1
+    # (forward_fn, backward_fn) pair, or None if the operation cannot be compiled.
+    compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None
+
+    def __init_subclass__(cls, **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
+        if cls.fwd_args_type is None or cls.bwd_args_type is None:
+            return
+        if getattr(cls.forward_compute, "__isabstractmethod__", False):
+            return
+        for name, arg_type in (
+            ("fwd_args_type", cls.fwd_args_type),
+            ("bwd_args_type", cls.bwd_args_type),
+        ):
+            # The op schema is built from the container's fields, so this is the
+            # framework's actual requirement -- check it where it is declared.
+            if not dataclasses.is_dataclass(arg_type):
+                raise TypeError(f"{cls.__name__}.{name} must be a dataclass")
+        # One registration per class. The compute halves are bound here, so a
+        # subclass that only swaps kernels (the activations) still gets its own
+        # op without repeating any of this.
+        cls.compile_ops = register_custom_op(
+            op_name=cls.__name__.lower(),
+            fwd_arg_type=cls.fwd_args_type,
+            fwd_impl=cls.forward_compute,
+            fwd_fake_impl=cls.forward_fake,
+            bwd_arg_type=cls.bwd_args_type,
+            bwd_impl=cls.backward_compute,
+            bwd_fake_impl=cls.backward_fake,
+            num_grad_inputs=cls.num_grad_inputs,
+        )
+
     def __init__(self) -> None:
         super().__init__()
 
@@ -274,6 +314,93 @@ def set_extra_output_channel(
         self._extra_output_to_caller[index] = output_to_caller
         return self
 
+    # ------------------------------------------------------------------ #
+    # Compute halves. Classmethods, not free functions: they belong to the
+    # operation, and binding to the class is what lets a family of operations
+    # share one implementation while dispatching to per-class kernels.
+    # ------------------------------------------------------------------ #
+
+    @classmethod
+    def forward_compute(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]:
+        """Pure forward: ``(output, tensors_to_save, ctx_attrs)``.
+
+        Takes everything through ``args``; must not read ``self`` or global
+        state, both of which are invisible to the compiler at this point.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def forward_fake(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]:
+        """Allocation-free twin of :meth:`forward_compute` over ``TensorSpec``.
+
+        Runs as a meta kernel, outside the traced frame, and more than once per
+        compile, so it must be a pure function of ``args`` -- a read of global
+        state here is unguarded and can silently disagree with the real impl.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def backward_compute(cls, args: Any) -> tuple:
+        """Pure backward: ``num_grad_inputs`` gradients."""
+        raise NotImplementedError
+
+    @classmethod
+    def backward_fake(cls, args: Any) -> tuple:
+        """Allocation-free twin of :meth:`backward_compute`."""
+        raise NotImplementedError
+
+    def compile_unsupported_reason(self) -> Optional[str]:
+        """Why this operation cannot go through its custom op, or ``None``.
+
+        Asked per operation, but acted on per fuser group: a pipeline compiles
+        as a whole, so one unsupported operation sends the whole group to eager.
+        Recipe-level limits are not checked here -- they belong to whoever reads
+        the recipe, which is the fuser.
+        """
+        if self.compile_ops is None:
+            return f"{self.__class__.__name__} without compute halves"
+        for mode in ("forward", "backward"):
+            for index in range(self.num_quantizers(mode)):
+                quantizer = self.get_quantizer(mode, index)
+                if quantizer is not None and not is_value_opaque_quantizer(quantizer):
+                    # Delayed scaling holds live scale/amax tensors, so its
+                    # quantizer cannot be specialized on and would be baked into
+                    # the graph as a stale constant.
+                    return (
+                        f"{type(quantizer).__name__} (not a torch.compile value-opaque quantizer)"
+                    )
+        return None
+
+    def resolve_fwd_args(
+        self,
+        input_: torch.Tensor,
+        *,
+        requires_grad: bool,
+        prev_op_grad_output_quantizer: Optional[Quantizer] = None,
+        next_op_input_quantizer: Optional[Quantizer] = None,
+    ) -> Any:
+        """Gather the forward's inputs into a flat, ``self``-free container.
+
+        This is where module config and global state are read, so it belongs in
+        the traced region where Dynamo guards those reads -- never inside the
+        custom op.
+        """
+        raise NotImplementedError
+
+    def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> Any:
+        """Rebuild the backward's inputs from the forward's saved state."""
+        raise NotImplementedError
+
+    def saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple:
+        """Tensors to persist, given what the forward handed back.
+
+        An operation whose backward needs its input but whose forward does not
+        produce a distinct tensor for it overrides this; a custom op may not
+        return one of its own inputs.
+        """
+        del input_
+        return saved
+
     @property
     def is_fused_op(self) -> bool:
         return False
@@ -508,7 +635,6 @@ def _load_fp8_metas(self, fp8_metas: Optional[dict[str, Any]]) -> None:
                 self._fp8_metas[mode][fp8_meta_key].scale.copy_(scale)
                 self._fp8_metas[mode][fp8_meta_key].amax_history.copy_(amax_history)
 
-    @abc.abstractmethod
     def op_forward(
         self,
         ctx: OperationContext,
@@ -520,6 +646,10 @@ def op_forward(
     ) -> torch.Tensor:
         """Forward pass
 
+        Operations that declare the compute halves inherit this: it resolves the
+        arguments, runs the forward, and records what the backward will need. The
+        rest override it.
+
         Parameters
         ----------
         ctx: OperationContext
@@ -537,8 +667,63 @@ def op_forward(
             Output tensor
 
         """
+        if self.fwd_args_type is None:
+            raise NotImplementedError(
+                f"{self.__class__.__name__} implements neither op_forward nor the compute halves"
+            )
+        if kwargs:
+            raise ValueError(f"{self.__class__.__name__} forward does not expect keyword arguments")
+        args = self.resolve_fwd_args(
+            input_,
+            requires_grad=ctx.requires_grad,
+            prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
+            next_op_input_quantizer=next_op_input_quantizer,
+        )
+        output, saved, ctx_attrs = self.forward_compute(args)
+        if ctx.requires_grad:
+            ctx.save_for_backward(*self.saved_for_backward(saved, input_))
+            for name, value in ctx_attrs.items():
+                setattr(ctx, name, value)
+        return output
+
+    def compiled_op_forward(
+        self,
+        ctx: OperationContext,
+        input_: torch.Tensor,
+        *,
+        prev_op_grad_output_quantizer: Optional[Quantizer],
+        next_op_input_quantizer: Optional[Quantizer],
+    ) -> torch.Tensor:
+        """:meth:`op_forward` routed through this operation's custom op.
+
+        Same bookkeeping, but the computation crosses an op boundary so Dynamo
+        sees one graph node instead of tracing into the kernels.
+        """
+        args = self.resolve_fwd_args(
+            input_,
+            requires_grad=ctx.requires_grad,
+            prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
+            next_op_input_quantizer=next_op_input_quantizer,
+        )
+        output, saved, ctx_attrs = self.compile_ops[0](args)
+        if ctx.requires_grad:
+            ctx.save_for_backward(*self.saved_for_backward(saved, input_))
+            for name, value in ctx_attrs.items():
+                setattr(ctx, name, value)
+        return output
+
+    def compiled_op_backward(
+        self,
+        ctx: OperationContext,
+        grad_output: torch.Tensor,
+    ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]:
+        """:meth:`op_backward` routed through this operation's custom op."""
+        grads = self.compile_ops[1](self.resolve_bwd_args(ctx, grad_output))
+        grad_input = grads[0]
+        if grad_input is None:
+            grad_input = grad_output
+        return grad_input, tuple(grads[1:])
 
-    @abc.abstractmethod
     def op_backward(
         self,
         ctx: OperationContext,
@@ -546,6 +731,8 @@ def op_backward(
     ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]:
         """Backward pass
 
+        Counterpart to the inherited :meth:`op_forward`.
+
         Parameters
         ----------
         ctx: OperationContext
@@ -561,6 +748,17 @@ def op_backward(
             Loss gradients w.r.t. parameters
 
         """
+        if self.bwd_args_type is None:
+            raise NotImplementedError(
+                f"{self.__class__.__name__} implements neither op_backward nor the compute halves"
+            )
+        grads = self.backward_compute(self.resolve_bwd_args(ctx, grad_output))
+        grad_input = grads[0]
+        if grad_input is None:
+            # "The incoming gradient, unchanged": a custom op may not return one
+            # of its own inputs, so the compute half hands back None instead.
+            grad_input = grad_output
+        return grad_input, tuple(grads[1:])
 
     def fuser_forward(
         self,

From 419c4e2065c1b655c92f4b6c21b1234438ffaa1b Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 5 Aug 2026 22:21:51 +0200
Subject: [PATCH 02/27] [PyTorch] Compile an OperationFuser group holding one
 operation

A group whose operations declare their compute halves now runs through their
custom ops under torch.compile(fullgraph=True). The pipeline-level
autograd.Function is traced as a higher-order op, which is what will later let
its forward and backward walk different op groupings.

Four side effects reached outside the higher-order op's scope and had to go:

- OperationContext objects are created in the forward, but the backward is a
  separate subgraph, so writing to them there mutates an enclosing scope; the
  backward copies them into its own scope instead;
- requires_grad_ on an output, which AOTAutograd's functionalization drops
  anyway -- autograd marks the outputs of an apply() itself;
- _do_not_clear on inputs and outputs.

They are gated on being traced rather than on using the custom ops. Under
fullgraph there is no leaving the graph, so an unsupported operation does not
fall back: the pipeline is traced either way and only the choice of
implementation changes.

The gate reports why a group runs eagerly through warn_compile_eager_fallback,
which is safe to call from the traced region.

Sequential builds its module groups outside the forward pass, since that
constructs nn.Modules.

Tested with a test-only operation, so the fuser's path does not depend on
which real operations happen to declare their halves.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py          | 140 +++++++++++++++++++
 transformer_engine/pytorch/ops/fuser.py      | 133 ++++++++++++++----
 transformer_engine/pytorch/ops/sequential.py |  17 ++-
 3 files changed, 257 insertions(+), 33 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index f8e09d5ce1..534e45e394 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -4,6 +4,7 @@
 
 import abc
 import contextlib
+import dataclasses
 import os
 import re
 import sys
@@ -37,6 +38,7 @@
 from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
 from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole
 from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear
+from transformer_engine.pytorch.ops.op import BasicOperation
 from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer
 from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer
 from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer
@@ -2339,3 +2341,141 @@ def fn(inp):
             "Unexpected recompilation(s) across different batch sizes: "
             f"{unique_graphs_after - unique_graphs_baseline} extra graph(s) compiled"
         )
+
+
+# --------------------------------------------------------------------------- #
+# transformer_engine.pytorch.ops under torch.compile
+# --------------------------------------------------------------------------- #
+
+
+@dataclasses.dataclass(slots=True)
+class _ScaleFwdArgs:
+    """Flat, ``self``-free inputs to the test operation's forward."""
+
+    input_: torch.Tensor
+    scale: torch.Tensor
+
+
+@dataclasses.dataclass(slots=True)
+class _ScaleBwdArgs:
+    """Flat inputs to the test operation's backward."""
+
+    grad_output: torch.Tensor = None
+    saved_input: torch.Tensor = None
+    scale: torch.Tensor = None
+
+
+class _ScaleOp(BasicOperation):
+    """Test-only operation: multiply by a learnable scalar.
+
+    Exists so the fuser's compiled path can be exercised without depending on
+    which real operations happen to declare their compute halves. It is the
+    smallest operation that still has a parameter gradient and a saved tensor.
+    """
+
+    fwd_args_type = _ScaleFwdArgs
+    bwd_args_type = _ScaleBwdArgs
+    num_grad_inputs = 2  # grad input, grad scale
+
+    def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> None:
+        super().__init__()
+        self.scale = torch.nn.Parameter(torch.full((), 2.0, device=device, dtype=dtype))
+
+    @classmethod
+    def forward_compute(cls, args):
+        return args.input_ * args.scale, (), {}
+
+    @classmethod
+    def forward_fake(cls, args):
+        x = args.input_
+        return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), (), {}
+
+    @classmethod
+    def backward_compute(cls, args):
+        dy = args.grad_output
+        return dy * args.scale, (dy * args.saved_input).sum()
+
+    @classmethod
+    def backward_fake(cls, args):
+        dy = args.grad_output
+        return (
+            TensorSpec(shape=tuple(dy.shape), dtype=dy.dtype, device=dy.device),
+            TensorSpec(shape=(), dtype=dy.dtype, device=dy.device),
+        )
+
+    def saved_for_backward(self, saved, input_):
+        # The forward produces no distinct tensor for its input, and a custom op
+        # may not return one of its own inputs.
+        del saved
+        return (input_,)
+
+    def resolve_fwd_args(
+        self,
+        input_,
+        *,
+        requires_grad,
+        prev_op_grad_output_quantizer=None,
+        next_op_input_quantizer=None,
+    ):
+        del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer
+        return _ScaleFwdArgs(input_=input_, scale=self.scale)
+
+    def resolve_bwd_args(self, ctx, grad_output):
+        (x,) = ctx.saved_tensors
+        return _ScaleBwdArgs(grad_output=grad_output, saved_input=x, scale=self.scale)
+
+
+def _assert_sequential_matches_eager(model, compiled, base):
+    """Run a Sequential eagerly and compiled on identical inputs; compare both
+    the output and every parameter gradient."""
+    inp_eager = base.detach().clone().requires_grad_(True)
+    model.zero_grad(set_to_none=True)
+    out_eager = model(inp_eager)
+    out_eager.sum().backward()
+    ref_out = out_eager.detach().clone()
+    ref_igrad = inp_eager.grad.detach().clone()
+    ref_pgrads = [p.grad.detach().clone() for p in model.parameters()]
+
+    inp_compiled = base.detach().clone().requires_grad_(True)
+    model.zero_grad(set_to_none=True)
+    out_compiled = compiled(inp_compiled).clone()
+    out_compiled.sum().backward()
+
+    torch.testing.assert_close(out_compiled, ref_out)
+    torch.testing.assert_close(inp_compiled.grad, ref_igrad)
+    for got, expected in zip(model.parameters(), ref_pgrads):
+        torch.testing.assert_close(got.grad, expected)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_te_ops_single_op_group_compiles():
+    """``fullgraph=True`` over an ``OperationFuser`` group holding one operation.
+
+    The pipeline-level ``autograd.Function`` is traced as a higher-order op and
+    calls the operation's custom ops inside, so forward and backward both end up
+    in the graph.
+    """
+    torch._dynamo.reset()
+    model = te.ops.Sequential(_ScaleOp())
+    compiled = torch.compile(model, fullgraph=True)
+    base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
+    _assert_sequential_matches_eager(model, compiled, base)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_te_ops_unsupported_group_still_compiles_eagerly():
+    """An operation without the compute halves runs its eager implementation.
+
+    Note that this is not a fallback: under ``fullgraph=True`` there is no
+    leaving the graph, so the pipeline is traced either way and only the choice
+    of implementation changes. That is why the tracing constraints -- no
+    mutation of anything from an enclosing scope -- have to hold on both paths.
+    """
+    torch._dynamo.reset()
+    op = te.ops.Identity()
+    assert op.compile_unsupported_reason() is not None
+
+    model = te.ops.Sequential(op)
+    compiled = torch.compile(model, fullgraph=True)
+    base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
+    _assert_sequential_matches_eager(model, compiled, base)
diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py
index fd66529ba8..4478509b1d 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 from collections.abc import Callable, Iterable, Sequence
+import copy
 import itertools
 from typing import Any, Optional, TypeAlias
 
@@ -13,6 +14,7 @@
 
 from ..quantization import FP8GlobalStateManager, Recipe, DelayedScaling
 from ..quantized_tensor import prepare_for_saving, restore_from_func_ctx
+from ..utils import warn_compile_eager_fallback
 from .op import (
     BasicOperation,
     FusibleOperation,
@@ -66,6 +68,7 @@ def forward(
         fuser: OperationFuser,
         basic_op_kwargs: list[dict[str, Any]],
         set_output_requires_grad: bool,
+        use_compiled: bool,
         *params_and_extra_inputs: torch.Tensor,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Forward pass
@@ -82,6 +85,9 @@ def forward(
             Keyword arguments to BasicOperation
         set_output_requires_grad: bool
             Whether to set ``requires_grad`` flags on returned tensors
+        use_compiled: bool
+            Whether to call the operations' custom ops instead of their eager
+            implementations. Decided once per group by ``OperationFuser``.
         *params_and_extra_inputs: torch.Tensor
             Other tensor inputs to include in autograd graph. Consists
             of parameter tensors, followed by extra operation inputs.
@@ -98,9 +104,14 @@ def forward(
         # Operation autograd contexts
         basic_op_ctxs = [OperationContext() for _ in range(fuser._num_basic_ops)]
 
-        # Mark input tensors as not deletable in backward
-        for tensor in (input_,) + params_and_extra_inputs:
-            tensor._do_not_clear = True
+        # Mark input tensors as not deletable in backward. Skipped whenever this
+        # is being traced -- not merely when the custom ops are used: these
+        # tensors are created outside this function, and a higher-order op may
+        # not mutate anything from an enclosing scope. Under fullgraph there is
+        # no falling back out of the graph, so the constraint holds either way.
+        if not torch.compiler.is_compiling():
+            for tensor in (input_,) + params_and_extra_inputs:
+                tensor._do_not_clear = True
 
         # Place user provided extra inputs into their basic-op slots. Slots bound to
         # internal channels are filled lazily as their producers execute.
@@ -153,14 +164,23 @@ def forward(
             if next_op is not None:
                 next_op_input_quantizer = next_op.get_input_quantizer()
 
-            x, fused_op_extra_outputs = op.fuser_forward(
-                [basic_op_ctxs[idx] for idx in basic_op_idxs],
-                x,
-                basic_op_extra_inputs=extra_inputs,
-                prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
-                next_op_input_quantizer=next_op_input_quantizer,
-                basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs],
-            )
+            if use_compiled:
+                x = op.compiled_op_forward(
+                    basic_op_ctxs[basic_op_idxs[0]],
+                    x,
+                    prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
+                    next_op_input_quantizer=next_op_input_quantizer,
+                )
+                fused_op_extra_outputs = [()]
+            else:
+                x, fused_op_extra_outputs = op.fuser_forward(
+                    [basic_op_ctxs[idx] for idx in basic_op_idxs],
+                    x,
+                    basic_op_extra_inputs=extra_inputs,
+                    prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
+                    next_op_input_quantizer=next_op_input_quantizer,
+                    basic_op_kwargs=[basic_op_kwargs[idx] for idx in basic_op_idxs],
+                )
             if len(fused_op_extra_outputs) != len(basic_op_idxs):
                 raise RuntimeError(
                     f"Expected {type(op).__name__} to generate extra outputs for "
@@ -227,9 +247,13 @@ def forward(
             func_ctx.save_for_backward(*tensors_to_save)
             func_ctx.tensor_objects = tensor_objects
 
-            # Whether to perform recipe update in backward pass
+            # Whether to perform recipe update in backward pass. Skipped under
+            # compile: this reads and flips global FP8 state, and delayed
+            # scaling -- the only recipe it serves -- is gated out anyway.
             is_first_module = False
-            if fuser.first_op_requiring_backward < fuser._num_basic_ops:
+            if not torch.compiler.is_compiling() and (
+                fuser.first_op_requiring_backward < fuser._num_basic_ops
+            ):
                 is_first_module = FP8GlobalStateManager.is_first_fp8_module()
 
             # Other context
@@ -244,15 +268,20 @@ def forward(
             func_ctx.basic_op_extra_output_consumers = fuser._basic_op_extra_output_consumers
             func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources
             func_ctx.is_first_module = is_first_module
-
-        # Mark output tensors as not deletable in backward
-        for tensor in itertools.chain(
-            (x,),
-            (y for ys in extra_outputs for y in ys if y is not None),
-        ):
-            tensor._do_not_clear = True
-
-        if set_output_requires_grad:
+            func_ctx.use_compiled = use_compiled
+
+        # Mark output tensors as not deletable in backward (eager only; see above)
+        if not torch.compiler.is_compiling():
+            for tensor in itertools.chain(
+                (x,),
+                (y for ys in extra_outputs for y in ys if y is not None),
+            ):
+                tensor._do_not_clear = True
+
+        # Autograd marks the outputs of an ``apply`` itself, so this is only
+        # needed on the eager path -- and AOTAutograd's functionalization drops
+        # a requires_grad_() applied to a graph output anyway.
+        if set_output_requires_grad and not torch.compiler.is_compiling():
             x.requires_grad_(fuser.first_op_requiring_backward < fuser._num_basic_ops)
 
         if extra_outputs_flat:
@@ -277,7 +306,12 @@ def backward(
         # Restore saved tensors
         saved_tensors = restore_from_func_ctx(func_ctx)
 
-        # Unflatten list of saved tensors
+        # Unflatten list of saved tensors. Under compile the contexts were
+        # created in the forward, which is a different subgraph, so writing to
+        # them here would be a side effect on an enclosing scope; copy them into
+        # this one instead. The copy carries the attributes the forward set.
+        if torch.compiler.is_compiling():
+            basic_op_ctxs = [copy.copy(ctx) for ctx in basic_op_ctxs]
         for ctx in basic_op_ctxs:
             ctx.saved_tensors = saved_tensors[slice(*ctx._saved_tensors_range)]
             ctx._saved_tensors_range = None
@@ -327,14 +361,22 @@ def backward(
                                 channel_grad if output_grad is None else output_grad + channel_grad
                             )
             grad_extra_outputs = [basic_op_grad_extra_outputs[idx] for idx in basic_op_idxs]
-            dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward(
-                [basic_op_ctxs[idx] for idx in basic_op_idxs],
-                dx,
-                basic_op_grad_extra_outputs=grad_extra_outputs,
-            )
+            if func_ctx.use_compiled:
+                dx, grad_params_one = op.compiled_op_backward(basic_op_ctxs[basic_op_idxs[0]], dx)
+                fused_op_grad_params = [grad_params_one]
+                fused_op_grad_extra_inputs = [()]
+            else:
+                dx, fused_op_grad_params, fused_op_grad_extra_inputs = op.fuser_backward(
+                    [basic_op_ctxs[idx] for idx in basic_op_idxs],
+                    dx,
+                    basic_op_grad_extra_outputs=grad_extra_outputs,
+                )
             for idx, dparams in zip(basic_op_idxs, fused_op_grad_params):
                 grad_params[idx] = dparams
-                basic_op_ctxs[idx].saved_tensors = None
+                # Dropping the reference frees the activation early; on the
+                # compiled path the graph owns that lifetime instead.
+                if not torch.compiler.is_compiling():
+                    basic_op_ctxs[idx].saved_tensors = None
             for idx, dxs in zip(basic_op_idxs, fused_op_grad_extra_inputs):
                 grad_extra_inputs[idx] = dxs
                 for input_idx, grad in enumerate(dxs):
@@ -392,6 +434,7 @@ def backward(
             None,  # fuser
             None,  # basic_op_kwargs
             None,  # set_output_requires_grad
+            None,  # use_compiled
             *grad_params_flat,
             *grad_extra_inputs_flat,
         )
@@ -691,6 +734,35 @@ def maybe_fuse_ops(
         else:
             self._last_amax_history_len = 0
 
+    def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) -> Optional[str]:
+        """Why this group may not run through its operations' custom ops."""
+        if len(self._forward_ops) != self._num_basic_ops:
+            # A fused op covers several basic ops; only single-op groups so far.
+            return "a fused operation"
+        if any(kwargs for kwargs in basic_op_kwargs):
+            return "operation keyword arguments are not supported"
+        for op in self._basic_ops:
+            if op.num_extra_inputs or op.num_extra_outputs:
+                return f"{type(op).__name__} with extra tensor inputs or outputs"
+            reason = op.compile_unsupported_reason()
+            if reason is not None:
+                return reason
+        return None
+
+    def _use_compiled(self, basic_op_kwargs: list[dict[str, Any]]) -> bool:
+        """Whether this group runs through its operations' custom ops.
+
+        Decided once for the whole group: a pipeline compiles as a whole, so one
+        unsupported operation sends all of them to eager.
+        """
+        if not torch.compiler.is_compiling():
+            return False
+        reason = self._compile_unsupported_reason(basic_op_kwargs)
+        if reason is None:
+            return True
+        warn_compile_eager_fallback(reason)
+        return False
+
     def __call__(
         self,
         input: torch.Tensor,  # pylint: disable=redefined-builtin
@@ -733,11 +805,14 @@ def __call__(
         # Note: We call forward directly when is_grad_enabled=False,
         # which can expose non-leaf tensors to the inner ops. Avoid
         # problems in this case by passing set_output_requires_grad=False.
+        use_compiled = self._use_compiled(basic_op_kwargs)
+
         args = (
             input,
             self,
             basic_op_kwargs,
             is_grad_enabled,  # set_output_requires_grad
+            use_compiled,
             *self._flat_basic_op_params,
             *extra_inputs,
         )
diff --git a/transformer_engine/pytorch/ops/sequential.py b/transformer_engine/pytorch/ops/sequential.py
index cb5dfecb9f..b8724ca460 100644
--- a/transformer_engine/pytorch/ops/sequential.py
+++ b/transformer_engine/pytorch/ops/sequential.py
@@ -179,9 +179,7 @@ def forward(
             or grouped MLP.
         """
 
-        # Create module groups if needed
-        if self._module_groups is None:
-            self._module_groups = self._make_module_groups(self._modules.values())
+        module_groups = self._get_module_groups()
 
         # Route op kwargs to each module group's basic ops
         group_op_kwargs = self._resolve_op_kwargs(op_kwargs)
@@ -189,7 +187,7 @@ def forward(
         # Forward pass for each module group
         x = input
         extra_outputs: list[torch.Tensor] = []
-        for group_idx, module_group in enumerate(self._module_groups):
+        for group_idx, module_group in enumerate(module_groups):
             if isinstance(module_group, OperationFuser):
                 xs, extra_inputs = (
                     (x,) + extra_inputs[: module_group.num_extra_inputs],
@@ -208,6 +206,17 @@ def forward(
             return (x,) + tuple(extra_outputs)
         return x
 
+    def _get_module_groups(self) -> list[OperationFuser | torch.nn.Module]:
+        """Module groups, built once.
+
+        Kept out of the forward pass: building them constructs ``OperationFuser``
+        and fused-operation objects, and an ``nn.Module`` cannot be constructed
+        inside a traced region.
+        """
+        if self._module_groups is None:
+            self._module_groups = self._make_module_groups(self._modules.values())
+        return self._module_groups
+
     def _resolve_op_kwargs(
         self,
         op_kwargs: Optional[dict[torch.nn.Module | int, dict[str, Any]]],

From 7cf976f2ea9f030a01ede5e79ebe3eeb15181e24 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Thu, 6 Aug 2026 19:41:39 +0200
Subject: [PATCH 03/27] [PyTorch] Accept an operation's declared forward kwargs
 under compile

An operation lists the forward kwargs it takes in fwd_kwarg_names. They are
resolved into its args container like any other config, in the traced Python
where Dynamo guards them, so they reach the custom op through the existing
schema -- a value is guarded, a tensor is lifted into the graph, and a quantized
one crosses as its inner buffers.

An undeclared kwarg still sends the whole group to eager. That is not a schema
limitation, as the old message implied: the kwargs that remain are the grouped
operations' preallocated buffers, which the op writes to, and a custom op may
not mutate a tensor from an enclosing scope.

A kwarg carries no gradient. This matches the eager path, where kwargs never
entered the autograd graph either, and is why only read-only ones are accepted.

The fuser test helper now builds a separate model for the eager and the compiled
pass. Previously both shared one model and the eager pass ran first to produce
the reference, so the compiled pass was always traced on a model whose module
groups, fusions and pre_first_fuser_forward had already run. Those paths are now
traced as well.

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     | 199 ++++++++++++++++++++----
 transformer_engine/pytorch/ops/fuser.py |  11 +-
 transformer_engine/pytorch/ops/op.py    |  22 ++-
 3 files changed, 200 insertions(+), 32 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 534e45e394..40c29411d9 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -9,6 +9,7 @@
 import re
 import sys
 import warnings
+from typing import Union
 
 import pytest
 import torch
@@ -43,7 +44,11 @@
 from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer
 from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer
 from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
-from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, Quantizer
+from transformer_engine.pytorch.quantized_tensor import (
+    QuantizedTensor,
+    QuantizedTensorStorage,
+    Quantizer,
+)
 from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec
 from transformer_engine.pytorch import (
     is_fp8_available,
@@ -2425,26 +2430,146 @@ def resolve_bwd_args(self, ctx, grad_output):
         return _ScaleBwdArgs(grad_output=grad_output, saved_input=x, scale=self.scale)
 
 
-def _assert_sequential_matches_eager(model, compiled, base):
+@dataclasses.dataclass(slots=True)
+class _ScaleKwargsFwdArgs:
+    """Flat inputs to the kwarg-taking test operation's forward."""
+
+    input_: torch.Tensor
+    scale: torch.Tensor
+    extra_scale: float
+    offset: Union[torch.Tensor, QuantizedTensorStorage]
+
+
+@dataclasses.dataclass(slots=True)
+class _ScaleKwargsBwdArgs:
+    """Flat inputs to the kwarg-taking test operation's backward."""
+
+    grad_output: torch.Tensor = None
+    saved_input: torch.Tensor = None
+    scale: torch.Tensor = None
+    extra_scale: float = 1.0
+
+
+class _ScaleWithKwargsOp(BasicOperation):
+    """Test-only operation taking forward kwargs: a value and a tensor.
+
+    ``offset`` is declared as tensor-or-quantized, so a quantized kwarg crosses
+    the op boundary as its inner buffers. Neither kwarg carries a gradient --
+    that is what "read-only" means here.
+    """
+
+    fwd_args_type = _ScaleKwargsFwdArgs
+    bwd_args_type = _ScaleKwargsBwdArgs
+    num_grad_inputs = 2  # grad input, grad scale
+    fwd_kwarg_names = ("extra_scale", "offset")
+
+    def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> None:
+        super().__init__()
+        self.scale = torch.nn.Parameter(torch.full((), 2.0, device=device, dtype=dtype))
+
+    @classmethod
+    def forward_compute(cls, args):
+        offset = args.offset
+        if isinstance(offset, QuantizedTensor):
+            offset = offset.dequantize()
+        out = args.input_ * args.scale * args.extra_scale + offset
+        return out, (), {"extra_scale": args.extra_scale}
+
+    @classmethod
+    def forward_fake(cls, args):
+        x = args.input_
+        return (
+            TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device),
+            (),
+            {"extra_scale": args.extra_scale},
+        )
+
+    @classmethod
+    def backward_compute(cls, args):
+        dy = args.grad_output
+        return (
+            dy * args.scale * args.extra_scale,
+            (dy * args.saved_input).sum() * args.extra_scale,
+        )
+
+    @classmethod
+    def backward_fake(cls, args):
+        dy = args.grad_output
+        return (
+            TensorSpec(shape=tuple(dy.shape), dtype=dy.dtype, device=dy.device),
+            TensorSpec(shape=(), dtype=dy.dtype, device=dy.device),
+        )
+
+    def saved_for_backward(self, saved, input_):
+        del saved
+        return (input_,)
+
+    def resolve_fwd_args(
+        self,
+        input_,
+        *,
+        requires_grad,
+        prev_op_grad_output_quantizer=None,
+        next_op_input_quantizer=None,
+        extra_scale=1.0,
+        offset=None,
+    ):
+        del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer
+        if offset is None:
+            offset = torch.zeros((), device=input_.device, dtype=input_.dtype)
+        return _ScaleKwargsFwdArgs(
+            input_=input_,
+            scale=self.scale,
+            extra_scale=extra_scale,
+            offset=offset,
+        )
+
+    def resolve_bwd_args(self, ctx, grad_output):
+        (x,) = ctx.saved_tensors
+        return _ScaleKwargsBwdArgs(
+            grad_output=grad_output,
+            saved_input=x,
+            scale=self.scale,
+            extra_scale=ctx.extra_scale,
+        )
+
+
+def _assert_sequential_matches_eager(make_model, base, op_kwargs_seq=(None,)):
     """Run a Sequential eagerly and compiled on identical inputs; compare both
-    the output and every parameter gradient."""
-    inp_eager = base.detach().clone().requires_grad_(True)
-    model.zero_grad(set_to_none=True)
-    out_eager = model(inp_eager)
-    out_eager.sum().backward()
-    ref_out = out_eager.detach().clone()
-    ref_igrad = inp_eager.grad.detach().clone()
-    ref_pgrads = [p.grad.detach().clone() for p in model.parameters()]
+    the output and every parameter gradient.
 
-    inp_compiled = base.detach().clone().requires_grad_(True)
-    model.zero_grad(set_to_none=True)
-    out_compiled = compiled(inp_compiled).clone()
-    out_compiled.sum().backward()
+    Each pass gets its own freshly built model, so the compiled one is traced on
+    a first run: nothing has built the module groups, resolved the fusions or run
+    ``pre_first_fuser_forward`` on it beforehand. ``make_model`` must therefore
+    build deterministically identical models.
+
+    Several ``op_kwargs`` are run in order on the same pair of models, which is
+    what exercises Dynamo's guards on a kwarg value.
+    """
+    eager_model = make_model()
+    compiled_model = make_model()
+    compiled = torch.compile(compiled_model, fullgraph=True)
 
-    torch.testing.assert_close(out_compiled, ref_out)
-    torch.testing.assert_close(inp_compiled.grad, ref_igrad)
-    for got, expected in zip(model.parameters(), ref_pgrads):
-        torch.testing.assert_close(got.grad, expected)
+    for op_kwargs in op_kwargs_seq:
+        call_kwargs = {} if op_kwargs is None else {"op_kwargs": op_kwargs}
+
+        inp_eager = base.detach().clone().requires_grad_(True)
+        eager_model.zero_grad(set_to_none=True)
+        out_eager = eager_model(inp_eager, **call_kwargs)
+        out_eager.sum().backward()
+        ref_out = out_eager.detach().clone()
+        ref_igrad = inp_eager.grad.detach().clone()
+        ref_pgrads = [p.grad.detach().clone() for p in eager_model.parameters()]
+
+        inp_compiled = base.detach().clone().requires_grad_(True)
+        compiled_model.zero_grad(set_to_none=True)
+        out_compiled = compiled(inp_compiled, **call_kwargs).clone()
+        out_compiled.sum().backward()
+
+        torch.testing.assert_close(out_compiled, ref_out)
+        torch.testing.assert_close(inp_compiled.grad, ref_igrad)
+        for got, expected in zip(compiled_model.parameters(), ref_pgrads):
+            torch.testing.assert_close(got.grad, expected)
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@@ -2456,10 +2581,8 @@ def test_te_ops_single_op_group_compiles():
     in the graph.
     """
     torch._dynamo.reset()
-    model = te.ops.Sequential(_ScaleOp())
-    compiled = torch.compile(model, fullgraph=True)
     base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
-    _assert_sequential_matches_eager(model, compiled, base)
+    _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp()), base)
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
@@ -2472,10 +2595,34 @@ def test_te_ops_unsupported_group_still_compiles_eagerly():
     mutation of anything from an enclosing scope -- have to hold on both paths.
     """
     torch._dynamo.reset()
-    op = te.ops.Identity()
-    assert op.compile_unsupported_reason() is not None
+    assert te.ops.Identity().compile_unsupported_reason() is not None
 
-    model = te.ops.Sequential(op)
-    compiled = torch.compile(model, fullgraph=True)
     base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
-    _assert_sequential_matches_eager(model, compiled, base)
+    _assert_sequential_matches_eager(lambda: te.ops.Sequential(te.ops.Identity()), base)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_te_ops_forward_kwargs_compile():
+    """Forward kwargs reach the operation through its custom op.
+
+    Covers both kinds at once: a value, which Dynamo guards on -- hence the
+    second call with a different one -- and a tensor, quantized here, which
+    crosses the op boundary as its inner buffers.
+    """
+    torch._dynamo.reset()
+    quantizer = Float8CurrentScalingQuantizer(
+        fp8_dtype=tex.DType.kFloat8E4M3,
+        device=torch.device("cuda"),
+    )
+    offset = quantizer(torch.randn(64, dtype=torch.bfloat16, device="cuda"))
+
+    base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
+    _assert_sequential_matches_eager(
+        lambda: te.ops.Sequential(_ScaleWithKwargsOp()),
+        base,
+        op_kwargs_seq=(
+            {0: {"extra_scale": 3.0, "offset": offset}},
+            {0: {"extra_scale": 5.0, "offset": offset}},
+        ),
+    )
diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py
index 4478509b1d..25f14311c4 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -170,6 +170,7 @@ def forward(
                     x,
                     prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
                     next_op_input_quantizer=next_op_input_quantizer,
+                    **basic_op_kwargs[basic_op_idxs[0]],
                 )
                 fused_op_extra_outputs = [()]
             else:
@@ -739,8 +740,14 @@ def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) ->
         if len(self._forward_ops) != self._num_basic_ops:
             # A fused op covers several basic ops; only single-op groups so far.
             return "a fused operation"
-        if any(kwargs for kwargs in basic_op_kwargs):
-            return "operation keyword arguments are not supported"
+        for op, kwargs in zip(self._basic_ops, basic_op_kwargs):
+            # A kwarg an operation declares is resolved into its args container
+            # like any other config. Anything else -- notably the preallocated
+            # buffers of the grouped operations -- is written to by the op, and a
+            # custom op may not mutate a tensor from an enclosing scope.
+            unsupported = sorted(name for name in kwargs if name not in op.fwd_kwarg_names)
+            if unsupported:
+                return f"{type(op).__name__} does not support keyword arguments {unsupported}"
         for op in self._basic_ops:
             if op.num_extra_inputs or op.num_extra_outputs:
                 return f"{type(op).__name__} with extra tensor inputs or outputs"
diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py
index 0fc6cb5c99..7b876089a1 100644
--- a/transformer_engine/pytorch/ops/op.py
+++ b/transformer_engine/pytorch/ops/op.py
@@ -195,6 +195,9 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta):
     bwd_args_type: Optional[type] = None
     # Gradients returned by backward_compute: the input's, then any parameters'.
     num_grad_inputs: int = 1
+    # Forward kwargs this operation accepts, resolved into fwd_args_type like
+    # any other config. A kwarg carries no gradient and must not be mutated.
+    fwd_kwarg_names: tuple[str, ...] = ()
     # (forward_fn, backward_fn) pair, or None if the operation cannot be compiled.
     compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None
 
@@ -378,12 +381,15 @@ def resolve_fwd_args(
         requires_grad: bool,
         prev_op_grad_output_quantizer: Optional[Quantizer] = None,
         next_op_input_quantizer: Optional[Quantizer] = None,
+        **kwargs: Any,
     ) -> Any:
         """Gather the forward's inputs into a flat, ``self``-free container.
 
         This is where module config and global state are read, so it belongs in
         the traced region where Dynamo guards those reads -- never inside the
-        custom op.
+        custom op. ``kwargs`` are the caller's forward kwargs, restricted to
+        ``fwd_kwarg_names``; an operation declaring them supplies their defaults
+        here, since a kwarg may be absent.
         """
         raise NotImplementedError
 
@@ -671,13 +677,17 @@ def op_forward(
             raise NotImplementedError(
                 f"{self.__class__.__name__} implements neither op_forward nor the compute halves"
             )
-        if kwargs:
-            raise ValueError(f"{self.__class__.__name__} forward does not expect keyword arguments")
+        unsupported = sorted(name for name in kwargs if name not in self.fwd_kwarg_names)
+        if unsupported:
+            raise ValueError(
+                f"{self.__class__.__name__} forward does not accept keyword arguments {unsupported}"
+            )
         args = self.resolve_fwd_args(
             input_,
             requires_grad=ctx.requires_grad,
             prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
             next_op_input_quantizer=next_op_input_quantizer,
+            **kwargs,
         )
         output, saved, ctx_attrs = self.forward_compute(args)
         if ctx.requires_grad:
@@ -693,17 +703,21 @@ def compiled_op_forward(
         *,
         prev_op_grad_output_quantizer: Optional[Quantizer],
         next_op_input_quantizer: Optional[Quantizer],
+        **kwargs: Any,
     ) -> torch.Tensor:
         """:meth:`op_forward` routed through this operation's custom op.
 
         Same bookkeeping, but the computation crosses an op boundary so Dynamo
-        sees one graph node instead of tracing into the kernels.
+        sees one graph node instead of tracing into the kernels. ``kwargs`` are
+        not validated here -- the fuser's gate already rejected a group whose
+        kwargs an operation does not declare.
         """
         args = self.resolve_fwd_args(
             input_,
             requires_grad=ctx.requires_grad,
             prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
             next_op_input_quantizer=next_op_input_quantizer,
+            **kwargs,
         )
         output, saved, ctx_attrs = self.compile_ops[0](args)
         if ctx.requires_grad:

From 4f773fca0b05b1437a82a825d0640c6e33350360 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 7 Aug 2026 16:25:23 +0200
Subject: [PATCH 04/27] [PyTorch] Restrict compiled forward kwargs to tensors

A value kwarg does not survive a second call. The other fields of an args
container are read off the module and are constant across calls, so they are
baked into the graph; a kwarg changes per call, and on the second value Dynamo
hands over a symbolic scalar, which OpaqueValueBundle cannot carry -- it fails
with AsPythonConstantNotImplementedError, not with a graph break. Measured on
int and float alike; specialize_float=True cures only the float, and is global.

The gate now takes tensor kwargs only, so a value sends the group to eager
deterministically instead of failing on its second call. A 0-d tensor is the way
to pass a scalar: it is a graph input, so it recompiles for no value at all.

The test carries a quantized offset that changes on every call and confirms no
recompilation, then adds a value kwarg to cover the gated path. That last call
keeps its offset unquantized on purpose: a gated group runs the eager
implementation, which is traced directly rather than hidden behind a custom op,
and dequantize() graph-breaks there.

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     | 21 ++++++++++++++-------
 transformer_engine/pytorch/ops/fuser.py | 16 ++++++++++++----
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 40c29411d9..5be47f5e87 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2604,25 +2604,32 @@ def test_te_ops_unsupported_group_still_compiles_eagerly():
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_te_ops_forward_kwargs_compile():
-    """Forward kwargs reach the operation through its custom op.
+    """A tensor forward kwarg reaches the operation through its custom op.
 
-    Covers both kinds at once: a value, which Dynamo guards on -- hence the
-    second call with a different one -- and a tensor, quantized here, which
-    crosses the op boundary as its inner buffers.
+    The tensor is quantized, so it crosses the op boundary as its inner buffers,
+    and it changes between calls, which a graph input absorbs without a
+    recompilation. The last call adds a value kwarg: that one is gated onto the
+    eager implementation, since Dynamo turns a changed scalar into a symbol that
+    cannot be carried as opaque config.
     """
     torch._dynamo.reset()
     quantizer = Float8CurrentScalingQuantizer(
         fp8_dtype=tex.DType.kFloat8E4M3,
         device=torch.device("cuda"),
     )
-    offset = quantizer(torch.randn(64, dtype=torch.bfloat16, device="cuda"))
+
+    def offset(value):
+        return quantizer(torch.full((64,), value, dtype=torch.bfloat16, device="cuda"))
 
     base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
     _assert_sequential_matches_eager(
         lambda: te.ops.Sequential(_ScaleWithKwargsOp()),
         base,
         op_kwargs_seq=(
-            {0: {"extra_scale": 3.0, "offset": offset}},
-            {0: {"extra_scale": 5.0, "offset": offset}},
+            {0: {"offset": offset(0.5)}},
+            {0: {"offset": offset(1.5)}},
+            # No quantized offset here: this call runs the eager implementation,
+            # which is traced directly, and dequantize() is not traceable.
+            {0: {"extra_scale": 3.0}},
         ),
     )
diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py
index 25f14311c4..01e944248a 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -740,14 +740,22 @@ def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) ->
         if len(self._forward_ops) != self._num_basic_ops:
             # A fused op covers several basic ops; only single-op groups so far.
             return "a fused operation"
-        for op, kwargs in zip(self._basic_ops, basic_op_kwargs):
+        for op, kwargs in zip(self._basic_ops, basic_op_kwargs, strict=True):
             # A kwarg an operation declares is resolved into its args container
             # like any other config. Anything else -- notably the preallocated
             # buffers of the grouped operations -- is written to by the op, and a
             # custom op may not mutate a tensor from an enclosing scope.
-            unsupported = sorted(name for name in kwargs if name not in op.fwd_kwarg_names)
-            if unsupported:
-                return f"{type(op).__name__} does not support keyword arguments {unsupported}"
+            undeclared = sorted(name for name in kwargs if name not in op.fwd_kwarg_names)
+            if undeclared:
+                return f"{type(op).__name__} with undeclared keyword arguments {undeclared}"
+            # Only tensors. The other fields of an args container are values read
+            # off the module, constant across calls and baked into the graph; a
+            # kwarg changes per call, and on the second value Dynamo hands over a
+            # symbolic scalar, which cannot go into an opaque value bundle. Pass a
+            # 0-d tensor instead -- it is a graph input, so it does not recompile.
+            values = sorted(name for name, v in kwargs.items() if not isinstance(v, torch.Tensor))
+            if values:
+                return f"{type(op).__name__} with non-tensor keyword arguments {values}"
         for op in self._basic_ops:
             if op.num_extra_inputs or op.num_extra_outputs:
                 return f"{type(op).__name__} with extra tensor inputs or outputs"

From 97e4764974d772c367b564002806de055532de52 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 4 Sep 2026 16:28:21 +0200
Subject: [PATCH 05/27] [PyTorch] Refine fusible custom-op integration

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 83 ++++++++++++++-----
 transformer_engine/pytorch/dynamo/__init__.py |  8 +-
 .../pytorch/dynamo/custom_op.py               | 50 +++++++----
 transformer_engine/pytorch/ops/fuser.py       | 37 +++++----
 transformer_engine/pytorch/ops/op.py          | 35 +++-----
 transformer_engine/pytorch/ops/sequential.py  | 17 +---
 6 files changed, 142 insertions(+), 88 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 5be47f5e87..c02863a1a0 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -39,6 +39,7 @@
 from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
 from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole
 from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear
+from transformer_engine.pytorch.ops.fuser import OperationFuser
 from transformer_engine.pytorch.ops.op import BasicOperation
 from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer
 from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer
@@ -49,7 +50,7 @@
     QuantizedTensorStorage,
     Quantizer,
 )
-from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec
+from transformer_engine.pytorch.dynamo import ForwardResult, TensorSpec, to_tensor_spec
 from transformer_engine.pytorch import (
     is_fp8_available,
     is_mxfp8_available,
@@ -2388,12 +2389,12 @@ def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16)
 
     @classmethod
     def forward_compute(cls, args):
-        return args.input_ * args.scale, (), {}
+        return ForwardResult(args.input_ * args.scale)
 
     @classmethod
     def forward_fake(cls, args):
         x = args.input_
-        return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), (), {}
+        return ForwardResult(TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device))
 
     @classmethod
     def backward_compute(cls, args):
@@ -2408,11 +2409,9 @@ def backward_fake(cls, args):
             TensorSpec(shape=(), dtype=dy.dtype, device=dy.device),
         )
 
-    def saved_for_backward(self, saved, input_):
-        # The forward produces no distinct tensor for its input, and a custom op
-        # may not return one of its own inputs.
-        del saved
-        return (input_,)
+    def setup_context(self, ctx, args, aux):
+        del aux
+        ctx.save_for_backward(args.input_, args.scale)
 
     def resolve_fwd_args(
         self,
@@ -2426,8 +2425,23 @@ def resolve_fwd_args(
         return _ScaleFwdArgs(input_=input_, scale=self.scale)
 
     def resolve_bwd_args(self, ctx, grad_output):
-        (x,) = ctx.saved_tensors
-        return _ScaleBwdArgs(grad_output=grad_output, saved_input=x, scale=self.scale)
+        x, scale = ctx.saved_tensors
+        return _ScaleBwdArgs(grad_output=grad_output, saved_input=x, scale=scale)
+
+
+class _BackwardScalePair(te.ops.FusedOperation):
+    """Backward-only fusion for the compile gate test."""
+
+    def fuser_backward(self, basic_op_ctxs, grad_output, **unused):
+        dx, grad_params_1 = self.basic_ops[1].op_backward(basic_op_ctxs[1], grad_output)
+        dx, grad_params_0 = self.basic_ops[0].op_backward(basic_op_ctxs[0], dx)
+        return dx, [grad_params_0, grad_params_1], [(), ()]
+
+
+def _fuse_backward_scale_pair(ops, **unused):
+    if len(ops) == 2 and all(isinstance(op, _ScaleOp) for op in ops):
+        return [_BackwardScalePair(ops)]
+    return ops
 
 
 @dataclasses.dataclass(slots=True)
@@ -2473,16 +2487,12 @@ def forward_compute(cls, args):
         if isinstance(offset, QuantizedTensor):
             offset = offset.dequantize()
         out = args.input_ * args.scale * args.extra_scale + offset
-        return out, (), {"extra_scale": args.extra_scale}
+        return ForwardResult(out)
 
     @classmethod
     def forward_fake(cls, args):
         x = args.input_
-        return (
-            TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device),
-            (),
-            {"extra_scale": args.extra_scale},
-        )
+        return ForwardResult(TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device))
 
     @classmethod
     def backward_compute(cls, args):
@@ -2500,9 +2510,10 @@ def backward_fake(cls, args):
             TensorSpec(shape=(), dtype=dy.dtype, device=dy.device),
         )
 
-    def saved_for_backward(self, saved, input_):
-        del saved
-        return (input_,)
+    def setup_context(self, ctx, args, aux):
+        del aux
+        ctx.save_for_backward(args.input_, args.scale)
+        ctx.extra_scale = args.extra_scale
 
     def resolve_fwd_args(
         self,
@@ -2525,11 +2536,11 @@ def resolve_fwd_args(
         )
 
     def resolve_bwd_args(self, ctx, grad_output):
-        (x,) = ctx.saved_tensors
+        x, scale = ctx.saved_tensors
         return _ScaleKwargsBwdArgs(
             grad_output=grad_output,
             saved_input=x,
-            scale=self.scale,
+            scale=scale,
             extra_scale=ctx.extra_scale,
         )
 
@@ -2585,6 +2596,36 @@ def test_te_ops_single_op_group_compiles():
     _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp()), base)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_te_ops_backward_fusion_uses_eager_implementations():
+    """A backward fusion prevents the group from using basic-op custom ops."""
+    te.ops.register_backward_fusion(_fuse_backward_scale_pair, prepend=True)
+    try:
+        base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
+        with pytest.warns(UserWarning, match="backward fusion"):
+            _assert_sequential_matches_eager(
+                lambda: te.ops.Sequential(_ScaleOp(), _ScaleOp()), base
+            )
+    finally:
+        OperationFuser.backward_fusion_functions.remove(_fuse_backward_scale_pair)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("compile_model", [False, True], ids=["eager", "compiled"])
+def test_te_ops_setup_context_saves_parameter(compile_model):
+    """Backward observes mutation of a tensor used by the forward."""
+    op = _ScaleOp()
+    model = te.ops.Sequential(op)
+    if compile_model:
+        model = torch.compile(model, fullgraph=True)
+    x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+    y = model(x)
+    with torch.no_grad():
+        op.scale.add_(1)
+    with pytest.raises(RuntimeError, match="modified by an inplace operation"):
+        y.sum().backward()
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 def test_te_ops_unsupported_group_still_compiles_eagerly():
     """An operation without the compute halves runs its eager implementation.
diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py
index 083a2aa1fb..88a2d716a7 100644
--- a/transformer_engine/pytorch/dynamo/__init__.py
+++ b/transformer_engine/pytorch/dynamo/__init__.py
@@ -6,13 +6,19 @@
 
 from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer
 from .tensor_spec import TensorSpec, to_tensor_spec
-from .custom_op import register_custom_op, register_custom_op_with_autograd, TensorOrQuantized
+from .custom_op import (
+    ForwardResult,
+    register_custom_op,
+    register_custom_op_with_autograd,
+    TensorOrQuantized,
+)
 
 __all__ = [
     "register_value_opaque_quantizer",
     "is_value_opaque_quantizer",
     "TensorSpec",
     "to_tensor_spec",
+    "ForwardResult",
     "register_custom_op",
     "register_custom_op_with_autograd",
     "TensorOrQuantized",
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 455a1d8150..1f5df798f6 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -14,6 +14,8 @@
 A TE forward/backward implementation takes one dataclass argument
 (``fwd_arg_type`` / ``bwd_arg_type``, e.g. ``LinearFwdArgs``) whose fields mix
 tensors, quantized tensors, quantizers, process groups and plain Python values.
+The autograd-free forward returns ``ForwardResult(output, aux)``; the
+autograd-wired API keeps its saved-tensor and context-metadata contract.
 
 A ``torch.library`` custom op is narrower: it only accepts flat schema slots
 (tensors plus opaque objects) and returns a flat ``Tensor[]``.
@@ -114,6 +116,15 @@
 
 _TE_OP_NAMESPACE = "transformer_engine_compile"
 
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class ForwardResult:
+    """Output and fresh auxiliary tensors produced by an autograd-free forward."""
+
+    output: Any
+    aux: tuple = ()
+
+
 # Annotation for an op arg field that may hold a plain tensor, a quantized
 # tensor subclass or a *bare* ``QuantizedTensorStorage`` (the internal-quantizer
 # optimization). Matched exactly by ``_TensorOrQuantizedAdapter``.
@@ -1315,19 +1326,18 @@ def register_custom_op(
     Both ops are two-tier, so ``QuantizedTensor`` subclass inputs pass through
     without dequantization.
 
-    Contracts, mirroring :func:`register_custom_op_with_autograd`:
+    Callable contracts:
 
-    * ``fwd_impl(fwd_args) -> (*user_outputs, tensors_to_save, ctx_attrs)``
+    * ``fwd_impl(fwd_args) -> ForwardResult(output, aux)``
     * ``fwd_fake_impl`` -- its data-free twin over :class:`TensorSpec`
     * ``bwd_impl(bwd_args) -> tuple`` of ``num_grad_inputs`` gradients
     * ``bwd_fake_impl`` -- its data-free twin
 
     Returns ``(forward_fn, backward_fn)``:
 
-    * ``forward_fn(fwd_args) -> (outputs, saved_tensors, ctx_attrs)`` --
-      ``outputs`` is a single value or a tuple, mirroring ``fwd_impl``'s user
-      outputs; ``saved_tensors`` is the reassembled ``tensors_to_save`` tuple,
-      which the caller is expected to persist (e.g. ``ctx.save_for_backward``).
+    * ``forward_fn(fwd_args) -> (output, aux)`` -- ``aux`` contains only fresh
+      tensors produced by the custom op. The caller decides which tensors and
+      metadata to persist for backward.
     * ``backward_fn(bwd_args) -> tuple`` of gradients.
 
     Returns ``None`` if registration fails (recorded once), so callers can fall
@@ -1363,11 +1373,25 @@ def _register_custom_op_impl(
     num_grad_inputs: int,
 ) -> Tuple[Callable[[Any], Any], Callable[[Any], Any]]:
     """Body of :func:`register_custom_op`; see it for semantics."""
+
+    def adapt_forward(impl):
+        def wrapped(args):
+            result = impl(args)
+            if not isinstance(result, ForwardResult):
+                raise TypeError(
+                    f"autograd-free fwd impl must return ForwardResult, got {type(result).__name__}"
+                )
+            return result.output, result.aux, None
+
+        return wrapped
+
+    adapted_fwd_impl = adapt_forward(fwd_impl)
+    adapted_fwd_fake_impl = adapt_forward(fwd_fake_impl)
     pair = _register_two_tier_pair(
         op_name=op_name,
         fwd_arg_type=fwd_arg_type,
-        fwd_impl=fwd_impl,
-        fwd_fake_impl=fwd_fake_impl,
+        fwd_impl=adapted_fwd_impl,
+        fwd_fake_impl=adapted_fwd_fake_impl,
         bwd_arg_type=bwd_arg_type,
         bwd_impl=bwd_impl,
         bwd_fake_impl=bwd_fake_impl,
@@ -1375,14 +1399,10 @@ def _register_custom_op_impl(
     )
 
     def forward_fn(fwd_args):
-        out_plan, payload = pair.call_forward(fwd_fake_impl, fwd_args)
+        out_plan, payload = pair.call_forward(adapted_fwd_fake_impl, fwd_args)
         outputs = out_plan.user_outputs(payload)
-        saved = out_plan.saved_tensors(payload)
-        return (
-            (outputs[0] if len(outputs) == 1 else tuple(outputs)),
-            tuple(saved),
-            out_plan.ctx_attrs,
-        )
+        aux = out_plan.saved_tensors(payload)
+        return outputs[0], tuple(aux)
 
     def backward_fn(bwd_args):
         # Unlike the forward payload, each grad occupies exactly one slot
diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py
index 01e944248a..87a3f18d96 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -68,7 +68,7 @@ def forward(
         fuser: OperationFuser,
         basic_op_kwargs: list[dict[str, Any]],
         set_output_requires_grad: bool,
-        use_compiled: bool,
+        use_custom_ops: bool,
         *params_and_extra_inputs: torch.Tensor,
     ) -> torch.Tensor | tuple[torch.Tensor, ...]:
         """Forward pass
@@ -85,9 +85,9 @@ def forward(
             Keyword arguments to BasicOperation
         set_output_requires_grad: bool
             Whether to set ``requires_grad`` flags on returned tensors
-        use_compiled: bool
-            Whether to call the operations' custom ops instead of their eager
-            implementations. Decided once per group by ``OperationFuser``.
+        use_custom_ops: bool
+            Whether to call the operations' custom ops instead of tracing their
+            eager implementations. Decided once per group by ``OperationFuser``.
         *params_and_extra_inputs: torch.Tensor
             Other tensor inputs to include in autograd graph. Consists
             of parameter tensors, followed by extra operation inputs.
@@ -164,7 +164,7 @@ def forward(
             if next_op is not None:
                 next_op_input_quantizer = next_op.get_input_quantizer()
 
-            if use_compiled:
+            if use_custom_ops:
                 x = op.compiled_op_forward(
                     basic_op_ctxs[basic_op_idxs[0]],
                     x,
@@ -269,7 +269,7 @@ def forward(
             func_ctx.basic_op_extra_output_consumers = fuser._basic_op_extra_output_consumers
             func_ctx.basic_op_extra_input_sources = fuser._basic_op_extra_input_sources
             func_ctx.is_first_module = is_first_module
-            func_ctx.use_compiled = use_compiled
+            func_ctx.use_custom_ops = use_custom_ops
 
         # Mark output tensors as not deletable in backward (eager only; see above)
         if not torch.compiler.is_compiling():
@@ -362,7 +362,7 @@ def backward(
                                 channel_grad if output_grad is None else output_grad + channel_grad
                             )
             grad_extra_outputs = [basic_op_grad_extra_outputs[idx] for idx in basic_op_idxs]
-            if func_ctx.use_compiled:
+            if func_ctx.use_custom_ops:
                 dx, grad_params_one = op.compiled_op_backward(basic_op_ctxs[basic_op_idxs[0]], dx)
                 fused_op_grad_params = [grad_params_one]
                 fused_op_grad_extra_inputs = [()]
@@ -435,7 +435,7 @@ def backward(
             None,  # fuser
             None,  # basic_op_kwargs
             None,  # set_output_requires_grad
-            None,  # use_compiled
+            None,  # use_custom_ops
             *grad_params_flat,
             *grad_extra_inputs_flat,
         )
@@ -735,11 +735,16 @@ def maybe_fuse_ops(
         else:
             self._last_amax_history_len = 0
 
-    def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) -> Optional[str]:
+    def _custom_ops_unsupported_reason(
+        self, basic_op_kwargs: list[dict[str, Any]]
+    ) -> Optional[str]:
         """Why this group may not run through its operations' custom ops."""
-        if len(self._forward_ops) != self._num_basic_ops:
-            # A fused op covers several basic ops; only single-op groups so far.
-            return "a fused operation"
+        for mode, ops in (("forward", self._forward_ops), ("backward", self._backward_ops)):
+            if len(ops) != self._num_basic_ops or any(
+                op is not self._basic_ops[idx] or basic_op_idxs != [idx]
+                for idx, (op, basic_op_idxs) in enumerate(ops)
+            ):
+                return f"a {mode} fusion"
         for op, kwargs in zip(self._basic_ops, basic_op_kwargs, strict=True):
             # A kwarg an operation declares is resolved into its args container
             # like any other config. Anything else -- notably the preallocated
@@ -764,7 +769,7 @@ def _compile_unsupported_reason(self, basic_op_kwargs: list[dict[str, Any]]) ->
                 return reason
         return None
 
-    def _use_compiled(self, basic_op_kwargs: list[dict[str, Any]]) -> bool:
+    def _use_custom_ops(self, basic_op_kwargs: list[dict[str, Any]]) -> bool:
         """Whether this group runs through its operations' custom ops.
 
         Decided once for the whole group: a pipeline compiles as a whole, so one
@@ -772,7 +777,7 @@ def _use_compiled(self, basic_op_kwargs: list[dict[str, Any]]) -> bool:
         """
         if not torch.compiler.is_compiling():
             return False
-        reason = self._compile_unsupported_reason(basic_op_kwargs)
+        reason = self._custom_ops_unsupported_reason(basic_op_kwargs)
         if reason is None:
             return True
         warn_compile_eager_fallback(reason)
@@ -820,14 +825,14 @@ def __call__(
         # Note: We call forward directly when is_grad_enabled=False,
         # which can expose non-leaf tensors to the inner ops. Avoid
         # problems in this case by passing set_output_requires_grad=False.
-        use_compiled = self._use_compiled(basic_op_kwargs)
+        use_custom_ops = self._use_custom_ops(basic_op_kwargs)
 
         args = (
             input,
             self,
             basic_op_kwargs,
             is_grad_enabled,  # set_output_requires_grad
-            use_compiled,
+            use_custom_ops,
             *self._flat_basic_op_params,
             *extra_inputs,
         )
diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py
index 7b876089a1..6c5b997ee2 100644
--- a/transformer_engine/pytorch/ops/op.py
+++ b/transformer_engine/pytorch/ops/op.py
@@ -22,7 +22,7 @@
     autocast,
 )
 from ..tensor import Quantizer
-from ..dynamo import is_value_opaque_quantizer, register_custom_op
+from ..dynamo import ForwardResult, is_value_opaque_quantizer, register_custom_op
 
 
 @dataclasses.dataclass
@@ -324,8 +324,8 @@ def set_extra_output_channel(
     # ------------------------------------------------------------------ #
 
     @classmethod
-    def forward_compute(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]:
-        """Pure forward: ``(output, tensors_to_save, ctx_attrs)``.
+    def forward_compute(cls, args: Any) -> ForwardResult:
+        """Forward computation over explicit arguments.
 
         Takes everything through ``args``; must not read ``self`` or global
         state, both of which are invisible to the compiler at this point.
@@ -333,7 +333,7 @@ def forward_compute(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]:
         raise NotImplementedError
 
     @classmethod
-    def forward_fake(cls, args: Any) -> tuple[Any, tuple, dict[str, Any]]:
+    def forward_fake(cls, args: Any) -> ForwardResult:
         """Allocation-free twin of :meth:`forward_compute` over ``TensorSpec``.
 
         Runs as a meta kernel, outside the traced frame, and more than once per
@@ -397,15 +397,10 @@ def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) ->
         """Rebuild the backward's inputs from the forward's saved state."""
         raise NotImplementedError
 
-    def saved_for_backward(self, saved: tuple, input_: torch.Tensor) -> tuple:
-        """Tensors to persist, given what the forward handed back.
-
-        An operation whose backward needs its input but whose forward does not
-        produce a distinct tensor for it overrides this; a custom op may not
-        return one of its own inputs.
-        """
-        del input_
-        return saved
+    def setup_context(self, ctx: OperationContext, args: Any, aux: tuple) -> None:
+        """Prepare backward state from the original arguments and fresh auxiliary tensors."""
+        del args
+        ctx.save_for_backward(*aux)
 
     @property
     def is_fused_op(self) -> bool:
@@ -689,12 +684,10 @@ def op_forward(
             next_op_input_quantizer=next_op_input_quantizer,
             **kwargs,
         )
-        output, saved, ctx_attrs = self.forward_compute(args)
+        result = self.forward_compute(args)
         if ctx.requires_grad:
-            ctx.save_for_backward(*self.saved_for_backward(saved, input_))
-            for name, value in ctx_attrs.items():
-                setattr(ctx, name, value)
-        return output
+            self.setup_context(ctx, args, result.aux)
+        return result.output
 
     def compiled_op_forward(
         self,
@@ -719,11 +712,9 @@ def compiled_op_forward(
             next_op_input_quantizer=next_op_input_quantizer,
             **kwargs,
         )
-        output, saved, ctx_attrs = self.compile_ops[0](args)
+        output, aux = self.compile_ops[0](args)
         if ctx.requires_grad:
-            ctx.save_for_backward(*self.saved_for_backward(saved, input_))
-            for name, value in ctx_attrs.items():
-                setattr(ctx, name, value)
+            self.setup_context(ctx, args, aux)
         return output
 
     def compiled_op_backward(
diff --git a/transformer_engine/pytorch/ops/sequential.py b/transformer_engine/pytorch/ops/sequential.py
index b8724ca460..cb5dfecb9f 100644
--- a/transformer_engine/pytorch/ops/sequential.py
+++ b/transformer_engine/pytorch/ops/sequential.py
@@ -179,7 +179,9 @@ def forward(
             or grouped MLP.
         """
 
-        module_groups = self._get_module_groups()
+        # Create module groups if needed
+        if self._module_groups is None:
+            self._module_groups = self._make_module_groups(self._modules.values())
 
         # Route op kwargs to each module group's basic ops
         group_op_kwargs = self._resolve_op_kwargs(op_kwargs)
@@ -187,7 +189,7 @@ def forward(
         # Forward pass for each module group
         x = input
         extra_outputs: list[torch.Tensor] = []
-        for group_idx, module_group in enumerate(module_groups):
+        for group_idx, module_group in enumerate(self._module_groups):
             if isinstance(module_group, OperationFuser):
                 xs, extra_inputs = (
                     (x,) + extra_inputs[: module_group.num_extra_inputs],
@@ -206,17 +208,6 @@ def forward(
             return (x,) + tuple(extra_outputs)
         return x
 
-    def _get_module_groups(self) -> list[OperationFuser | torch.nn.Module]:
-        """Module groups, built once.
-
-        Kept out of the forward pass: building them constructs ``OperationFuser``
-        and fused-operation objects, and an ``nn.Module`` cannot be constructed
-        inside a traced region.
-        """
-        if self._module_groups is None:
-            self._module_groups = self._make_module_groups(self._modules.values())
-        return self._module_groups
-
     def _resolve_op_kwargs(
         self,
         op_kwargs: Optional[dict[torch.nn.Module | int, dict[str, Any]]],

From b86d8f9c7ab1506c185f802a88bfd16f409addf3 Mon Sep 17 00:00:00 2001
From: William Yang <77467499+wilyan09007@users.noreply.github.com>
Date: Fri, 4 Sep 2026 11:35:54 -0400
Subject: [PATCH 06/27] [PyTorch] Decline fused grouped MLP when the backward
 format is not E4M3 (#3352)

* [PyTorch] Decline fused grouped MLP when the backward format is not E4M3

The fused grouped MLP packs the incoming activation gradient by reinterpreting
its storage as E4M3, conditioned only on NVFP4 and never on the FP8 format.
Under MXFP8BlockScaling(fp8_format=Format.HYBRID) the backward quantizers emit
E5M2, so those bytes are read as the wrong format rather than converted, and
every gradient out of the fusion is wrong. The forward pass is unaffected, so
this shows up as a model that trains too slowly instead of one that fails.

Fall back to the unfused ops when the recipe's backward format is not E4M3, and
raise instead of reinterpreting if such a gradient reaches the kernel path.

Signed-off-by: William <wilyan090@gmail.com>

* [PyTorch] Gate the grouped MLP backward format check on MXFP8

fp8_format describes the FP8 formats of an MXFP8 recipe. NVFP4BlockScaling
carries one too, pinned to E4M3, but its gradients are quantized to FP4 and the
value says nothing about them, so testing it for an NVFP4 recipe reached the
right answer for the wrong reason. Restrict the check to recipes where it means
something.

The runtime check at the pack site is unchanged: it reads the grad output
quantizer's own dtype on the non-NVFP4 branch, not the recipe.

Signed-off-by: William <wilyan090@gmail.com>

* Test grouped MLP format fusion with real ops

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

---------

Signed-off-by: William <wilyan090@gmail.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Co-authored-by: Przemek Tredak <ptredak@nvidia.com>
---
 tests/pytorch/test_grouped_mlp.py             | 40 +++++++++++++++++++
 .../pytorch/ops/fused/grouped_mlp.py          | 17 +++++++-
 2 files changed, 55 insertions(+), 2 deletions(-)

diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py
index d48f7afae6..1d173c76d8 100644
--- a/tests/pytorch/test_grouped_mlp.py
+++ b/tests/pytorch/test_grouped_mlp.py
@@ -1067,6 +1067,46 @@ def train_step(
 class TestGroupedMLPFusedOp:
     """Tests for grouped MLP fused op"""
 
+    def test_fusion_requires_supported_grad_output_format(self, monkeypatch) -> None:
+        """Fuse E4M3 MXFP8 and NVFP4, but decline MXFP8 with an E5M2 backward."""
+        from transformer_engine.common.recipe import Format, MXFP8BlockScaling, NVFP4BlockScaling
+
+        fused_op_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMGLU
+        monkeypatch.setattr(fused_op_cls, "is_supported", classmethod(lambda cls: True))
+
+        fc1 = te.ops.GroupedLinear(1, 64, 128, bias=False, device="cuda")
+        activation = te.ops.ScaledSwiGLU(glu_interleave_size=32)
+        fc2 = te.ops.GroupedLinear(1, 64, 64, bias=False, device="cuda")
+        ops = [fc1, activation, fc2]
+
+        def fuse(recipe):
+            return grouped_mlp_module.fuse_grouped_mlp_ops(
+                ops,
+                recipe=recipe,
+                fused_op_cls=fused_op_cls,
+            )
+
+        def assert_fused(recipe):
+            fused_ops = fuse(recipe)
+            assert len(fused_ops) == 1
+            fused_op = fused_ops[0]
+            assert isinstance(fused_op, fused_op_cls)
+            assert list(fused_op.basic_ops) == ops
+
+        hybrid = MXFP8BlockScaling(fp8_format=Format.HYBRID)
+        assert fuse(hybrid) is ops
+
+        e4m3 = MXFP8BlockScaling(fp8_format=Format.E4M3)
+        assert_fused(e4m3)
+
+        # NVFP4 quantizes gradients to FP4, so the FP8 format must not gate it. Forcing the
+        # lookup to E5M2 is what an NVFP4 recipe would hit if the check were not MXFP8-only.
+        monkeypatch.setattr(
+            grouped_mlp_module, "get_fp8_torch_dtype", lambda *_, **__: torch.float8_e5m2
+        )
+        nvfp4 = NVFP4BlockScaling(disable_rht=False)
+        assert_fused(nvfp4)
+
     @pytest.mark.parametrize("bias", (False, True))
     @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list)
     @pytest.mark.parametrize("single_grouped_weight", (False, True))
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 61f80b9d9f..2128bae1ec 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -17,7 +17,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, TE_DType
+from ...constants import DType, 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 (
@@ -27,7 +27,7 @@
     finalize_weight_grads,
 )
 from ...module.base import _2X_ACC_WGRAD
-from ...quantization import Recipe
+from ...quantization import Recipe, get_fp8_torch_dtype
 from ...tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer
 from ...tensor.grouped_tensor import GroupedTensor
 from ...tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor
@@ -877,6 +877,12 @@ def fuse_grouped_mlp_ops(
     # NVFP4 fused grouped MLP uses graph-safe grouped quantize, which currently requires RHT.
     if recipe.nvfp4() and recipe.disable_rht:
         return ops
+    # The fused MXFP8 backward reinterprets the grad output's storage as E4M3, so an E5M2
+    # backward format would have its gradients misread rather than converted. This declines
+    # MXFP8 with Format.HYBRID. fp8_format does not describe NVFP4 gradients, so NVFP4 is
+    # excluded from the check rather than relying on its value.
+    if recipe.mxfp8() and get_fp8_torch_dtype(recipe, fprop_tensor=False) != torch.float8_e4m3fn:
+        return ops
     if activation_op_types is None:
         activation_op_types = [ScaledSwiGLU, ScaledClampedQGeGLU]
         if _cudnn_frontend_supports_grouped_gemm_situglu():
@@ -2026,6 +2032,13 @@ def fuser_backward(
             or isinstance(fc1_weight_param, NVFP4Tensor)
             or isinstance(fc2_weight_param, NVFP4Tensor)
         )
+        if not use_nvfp4 and fc2_grad_output_quantizer.dtype != DType.kFloat8E4M3:
+            # The pack below reinterprets the grad output's storage as E4M3 rather than
+            # converting it, so anything else would be read as the wrong format.
+            raise RuntimeError(
+                "Fused grouped MLP backward requires an E4M3 grad output, but the recipe "
+                f"produced {fc2_grad_output_quantizer.dtype}."
+            )
         data_dtype = torch.float4_e2m1fn_x2 if use_nvfp4 else torch.float8_e4m3fn
         scale_view_dtype = torch.float8_e4m3fn if use_nvfp4 else torch.float8_e8m0fnu
         sf_vec_size = NVFP4_BLOCK_SCALING_SIZE if use_nvfp4 else MXFP8_BLOCK_SCALING_SIZE

From e61b66d6efd6c0f34c517cc333f57ce853bf4633 Mon Sep 17 00:00:00 2001
From: Phuong Nguyen <phuonguyen@nvidia.com>
Date: Fri, 4 Sep 2026 09:21:15 -0700
Subject: [PATCH 07/27] [PyTorch] Fix: Resolve EP symm-mem window offset for
 both old and new torch version  (#3466)

[PyTorch] Resolve EP symm-mem window offset against both torch symm-mem layouts

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>
---
 .../pytorch/csrc/extensions/ep.cpp            | 23 +++++++++++++++----
 1 file changed, 19 insertions(+), 4 deletions(-)

diff --git a/transformer_engine/pytorch/csrc/extensions/ep.cpp b/transformer_engine/pytorch/csrc/extensions/ep.cpp
index c74d4ddb6d..97bc70bddd 100644
--- a/transformer_engine/pytorch/csrc/extensions/ep.cpp
+++ b/transformer_engine/pytorch/csrc/extensions/ep.cpp
@@ -59,6 +59,21 @@ std::atomic<bool> g_zero_copy_enabled{false};
 // is not symm-mem-backed; the backend treats it as "no window, use staged copy".
 constexpr NVTECommWindow kNoWindow = {nullptr, 0};
 
+#ifdef NCCL_HAS_SYMMEM_SUPPORT
+// Offset of a symm-mem allocation relative to the start of its NCCL window.
+// Newer torch places the signal pad at the front of the allocation and exposes
+// get_window_offset() for it; on older torch the window starts at the buffer
+// base, where get_offset() is already window-relative.
+template <typename T>
+auto symm_mem_window_offset(T* sm, int) -> decltype(sm->get_window_offset()) {
+  return sm->get_window_offset();
+}
+template <typename T>
+size_t symm_mem_window_offset(T* sm, ...) {
+  return sm->get_offset();
+}
+#endif
+
 // Resolve ``t`` to an NCCL symm-mem window for the zero-copy one-sided path.
 // Returns ``kNoWindow`` when symm-mem support isn't compiled in, zero-copy is
 // disabled, no group is set, or ``t`` isn't symm-mem-backed; callers pass the
@@ -78,11 +93,11 @@ 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.");
-  // NCCL EP consumes window-relative offsets (the NCCL window starts at the signal pad,
-  // not at the buffer base). get_window_offset() = buffer_offset + get_offset(); add
-  // ``t``'s own storage offset for slice/view positioning.
+  // NCCL EP consumes window-relative offsets. 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.
   const uint64_t offset =
-      static_cast<uint64_t>(nccl_sm->get_window_offset()) +
+      static_cast<uint64_t>(symm_mem_window_offset(nccl_sm, 0)) +
       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

From 1f25d6ab1f30938d143bbdaf88a27c9fc045159e Mon Sep 17 00:00:00 2001
From: Wei Wang <143543872+nWEIdia@users.noreply.github.com>
Date: Fri, 4 Sep 2026 13:31:48 -0700
Subject: [PATCH 08/27] [NCCL][PyTorch] Fix
 test-fusible-ops-file-rendezvous-bug (#3478)

* world_group() used a single hardcoded init_method="file:///tmp/rdzv",
shared across every world_size this test parametrizes over
([device_count(), 1, 2]), and neither world_group() nor
test_distributed_fuser_ops ever calls destroy_process_group() or
removes the file afterwards. A world_size=N run's leftover FileStore
content can then corrupt a differently-shaped world_size=M run that
reuses the same path in the same pytest session, surfacing as a
confusing NCCL bootstrap failure:

  torch.distributed.DistBackendError: NCCL error ...
  ncclOsSocketPollConnect: connect to <self> ... Connection refused,
  exceeded error retry count after 35 attempts

instead of a clear rendezvous error. Confirmed 100% deterministic:
running the full file fresh (no pre-existing /tmp/rdzv) still fails
test_distributed_fuser_ops[2] every time, because the [4] parametrization
(which runs first) leaves /tmp/rdzv behind for [2] to trip over.

This was previously masked by older bundled NCCL (2.30.7), which
apparently tolerated the stale/mismatched FileStore well enough to
still succeed; a newer NCCL (2.31.2) surfaces it as a hard failure.
See the investigation writeup for the full comparison:

Fix: key the rendezvous path by world_size
(file:///tmp/rdzv_test_fusible_ops_{world_size}), and defensively
remove any pre-existing file at that path in test_distributed_fuser_ops
before launching each subprocess job, to also cover a leftover file
from an earlier crashed run of the same world_size.

Verified on a GB200 node
NCCL 2.31.2, the container that previously failed 2 of 3 parametrizations):
all 3 world_size parametrizations now pass, including two runs of the
full file back to back with no manual cleanup in between.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Signed-off-by: Wei Wang <weiwan@nvidia.com>

* Use unique rendezvous files in fusible ops tests

Signed-off-by: Przemek Tredak <ptredak@nvidia.com>

---------

Signed-off-by: Wei Wang <weiwan@nvidia.com>
Signed-off-by: Przemek Tredak <ptredak@nvidia.com>
Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
Co-authored-by: Przemek Tredak <ptredak@nvidia.com>
---
 tests/pytorch/distributed/test_fusible_ops.py        | 12 +++++++-----
 .../distributed/test_fusible_ops_with_userbuffers.py |  8 ++++++--
 2 files changed, 13 insertions(+), 7 deletions(-)

diff --git a/tests/pytorch/distributed/test_fusible_ops.py b/tests/pytorch/distributed/test_fusible_ops.py
index d733286093..c314bf5bb6 100644
--- a/tests/pytorch/distributed/test_fusible_ops.py
+++ b/tests/pytorch/distributed/test_fusible_ops.py
@@ -12,6 +12,7 @@
 import pathlib
 import subprocess
 import sys
+import tempfile
 from typing import Optional
 
 import pytest
@@ -58,7 +59,8 @@ def world_group() -> torch.distributed.ProcessGroup:
     torch.cuda.set_device(rank)
     group = torch.distributed.init_process_group(
         "nccl",
-        init_method="file:///tmp/rdzv",
+        # Each parallel job must use a fresh FileStore shared by only its ranks.
+        init_method=f"file://{os.environ['NVTE_TEST_RDZV_PATH']}",
         world_size=world_size,
         rank=rank,
     )
@@ -1053,10 +1055,10 @@ def test_distributed_fuser_ops(world_size: int) -> None:
         current_file,
         "--parallel",
     ]
-    result = subprocess.run(
-        command,
-        check=True,
-    )
+    with tempfile.TemporaryDirectory(prefix="te-test-fusible-ops-") as temp_dir:
+        env = dict(os.environ)
+        env["NVTE_TEST_RDZV_PATH"] = str(pathlib.Path(temp_dir) / "rdzv")
+        subprocess.run(command, check=True, env=env)
 
 
 def main() -> None:
diff --git a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py
index 07dffebf5f..38f49a96cb 100644
--- a/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py
+++ b/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py
@@ -12,6 +12,7 @@
 import pathlib
 import subprocess
 import sys
+import tempfile
 
 import pytest
 import torch
@@ -107,7 +108,8 @@ def world_group() -> torch.distributed.ProcessGroup:
     torch.cuda.set_device(local_rank)
     group = torch.distributed.init_process_group(
         "nccl",
-        init_method="file:///tmp/rdzv",
+        # Each parallel job must use a fresh FileStore shared by only its ranks.
+        init_method=f"file://{os.environ['NVTE_TEST_RDZV_PATH']}",
         world_size=world_size,
         rank=rank,
         device_id=torch.device(f"cuda:{local_rank}"),
@@ -471,7 +473,9 @@ def test_fuser_ops_with_userbuffers(
     env["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0"
 
     # Launch parallel job
-    run_distributed(command, env=env)
+    with tempfile.TemporaryDirectory(prefix="te-test-fusible-ops-userbuffers-") as temp_dir:
+        env["NVTE_TEST_RDZV_PATH"] = str(pathlib.Path(temp_dir) / "rdzv")
+        run_distributed(command, env=env)
 
 
 def main() -> None:

From 80a89adc272887fba3999a6a4b434347327d2c74 Mon Sep 17 00:00:00 2001
From: Zhongbo Zhu <42691305+zhongbozhu@users.noreply.github.com>
Date: Fri, 4 Sep 2026 13:47:15 -0700
Subject: [PATCH 09/27] Reduce Grouped MLP Fuser CPU Overhead (#3410)

* Reduce grouped MLP fuser CPU overhead

Reuse fused operation plans when full activation recompute changes grad mode, and avoid redundant CUDA current-device discovery for grouped MLP stream lookups.

Co-authored-by: Ting-Yang Kao <tingyangk@nvidia.com>
Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>

* resolve comments

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

* fix cutedsl wgrad crash

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

* resolve comments

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

---------

Signed-off-by: Zhongbo Zhu <zhongboz@nvidia.com>
Signed-off-by: tingyangk <tingyangk@nvidia.com>
Co-authored-by: Ting-Yang Kao <tingyangk@nvidia.com>
---
 tests/pytorch/test_fusible_ops.py             | 167 ++++++++++++++++++
 .../pytorch/ops/fused/grouped_mlp.py          |   4 +-
 transformer_engine/pytorch/ops/fuser.py       |  84 +++++----
 3 files changed, 219 insertions(+), 36 deletions(-)

diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py
index 6cd1fc3065..2adce717c9 100644
--- a/tests/pytorch/test_fusible_ops.py
+++ b/tests/pytorch/test_fusible_ops.py
@@ -146,6 +146,173 @@ def maybe_skip_quantization(
             pytest.skip("NVFP4 quantization is only supported with BF16 data")
 
 
+def test_operation_fuser_caches_plans_by_grad_requirement(monkeypatch) -> None:
+    """Cache and restore fusion plans for checkpoint forward and recompute."""
+
+    # Count fusion-plan construction without depending on any particular real
+    # fusion implementation. Each distinct fusion configuration invokes this
+    # hook once, while a cache hit must bypass it entirely.
+    fusion_calls = 0
+
+    def track_fusion(ops, *, recipe):  # pylint: disable=unused-argument
+        nonlocal fusion_calls
+        fusion_calls += 1
+        # Preserve the operation list so this hook observes plan construction
+        # without changing the topology under test.
+        return ops
+
+    # The fusion registries are class attributes shared by every OperationFuser.
+    # pytest's monkeypatch fixture restores all three after the test, preventing
+    # this synthetic fusion function from leaking into other tests. Keep only a
+    # joint forward-backward fusion hook so each plan build has one countable
+    # callback and no registered TE fusion can affect the result.
+    monkeypatch.setattr(OperationFuser, "forward_backward_fusion_functions", [track_fusion])
+    monkeypatch.setattr(OperationFuser, "forward_fusion_functions", [])
+    monkeypatch.setattr(OperationFuser, "backward_fusion_functions", [])
+
+    # One Identity op is enough to exercise the cache. With one basic op,
+    # first_op_requiring_backward has an intentionally simple interpretation:
+    #   0: backward starts at the Identity op;
+    #   1: the boundary is past the only op, so no backward work is required.
+    fuser = OperationFuser([te_ops.Identity()])
+    x = torch.ones(1, requires_grad=True)
+    # maybe_fuse_ops expects one extra-input collection per basic op. Identity
+    # has no extra inputs, so its collection is an empty tuple.
+    extra_inputs = [()]
+
+    # Phase 1: the original checkpointed forward runs with grad disabled. This
+    # is the first invocation, so the fuser must construct and cache the no-grad
+    # configuration. The runtime backward boundary is past the only op.
+    fuser.maybe_fuse_ops(False, None, x, extra_inputs)
+    assert fusion_calls == 1
+    assert fuser.first_op_requiring_backward == 1
+    no_grad_forward_ops = fuser._forward_ops
+    no_grad_backward_ops = fuser._backward_ops
+
+    # Phase 2: backward replays the checkpointed region with grad enabled. The
+    # backward boundary is part of the fusion key, allowing future fusion rules
+    # to choose a training-specific topology. The first grad-enabled invocation
+    # therefore constructs and caches a second configuration.
+    fuser.maybe_fuse_ops(True, None, x, extra_inputs)
+    assert fusion_calls == 2
+    assert fuser.first_op_requiring_backward == 0
+    grad_forward_ops = fuser._forward_ops
+    grad_backward_ops = fuser._backward_ops
+    assert grad_forward_ops is not no_grad_forward_ops
+    assert grad_backward_ops is not no_grad_backward_ops
+
+    # Phase 3: the next checkpointed forward must select the exact no-grad lists
+    # cached in phase 1. Before the cache was added, every boundary transition
+    # rebuilt the fused operations and called track_fusion again.
+    fuser.maybe_fuse_ops(False, None, x, extra_inputs)
+    assert fusion_calls == 2
+    assert fuser.first_op_requiring_backward == 1
+    assert fuser._forward_ops is no_grad_forward_ops
+    assert fuser._backward_ops is no_grad_backward_ops
+
+    # Phase 4: another recomputation must likewise restore the grad-enabled
+    # lists from phase 2. The full alternating sequence has built only the two
+    # configurations represented by its two fusion keys.
+    fuser.maybe_fuse_ops(True, None, x, extra_inputs)
+    assert fusion_calls == 2
+    assert fuser.first_op_requiring_backward == 0
+    assert fuser._forward_ops is grad_forward_ops
+    assert fuser._backward_ops is grad_backward_ops
+
+
+def test_operation_fuser_resets_recipe_state_independently_from_plan_cache(monkeypatch) -> None:
+    """Track recipe-state resets independently from fusion-plan construction."""
+
+    fusion_calls = 0
+
+    def track_fusion(ops, *, recipe):  # pylint: disable=unused-argument
+        nonlocal fusion_calls
+        fusion_calls += 1
+        return ops
+
+    # Replace the process-wide fusion registries so one callback corresponds to
+    # one plan construction. monkeypatch restores the registries after the test.
+    monkeypatch.setattr(OperationFuser, "forward_backward_fusion_functions", [track_fusion])
+    monkeypatch.setattr(OperationFuser, "forward_fusion_functions", [])
+    monkeypatch.setattr(OperationFuser, "backward_fusion_functions", [])
+
+    op = te_ops.Identity()
+    reset_recipes = []
+    first_forward_calls = 0
+
+    def track_recipe_reset(*, recipe):
+        reset_recipes.append(recipe)
+
+    def track_first_forward():
+        nonlocal first_forward_calls
+        first_forward_calls += 1
+
+    # Identity has no quantizers, so replace its state hooks with counters. This
+    # keeps the test CPU-only and isolates OperationFuser's reset decisions.
+    monkeypatch.setattr(op, "reset_recipe_state", track_recipe_reset)
+    monkeypatch.setattr(op, "pre_first_fuser_forward", track_first_forward)
+
+    fuser = OperationFuser([op])
+    x = torch.ones(1)
+    extra_inputs = [()]
+
+    current_scaling = transformer_engine.common.recipe.Float8CurrentScaling(backward_override=None)
+    fuser.maybe_fuse_ops(False, current_scaling, x, extra_inputs)
+    assert reset_recipes == [current_scaling]
+    assert first_forward_calls == 1
+    assert fusion_calls == 1
+
+    # A fresh but equivalent recipe does not invalidate state or the plan.
+    equivalent_current_scaling = transformer_engine.common.recipe.Float8CurrentScaling(
+        backward_override=None
+    )
+    fuser.maybe_fuse_ops(False, equivalent_current_scaling, x, extra_inputs)
+    assert reset_recipes == [current_scaling]
+    assert first_forward_calls == 1
+    assert fusion_calls == 1
+
+    # Backward override affects both recipe state and fusion topology, so it
+    # triggers one reset and constructs a distinct cached plan.
+    overridden_current_scaling = transformer_engine.common.recipe.Float8CurrentScaling(
+        backward_override="high_precision"
+    )
+    fuser.maybe_fuse_ops(False, overridden_current_scaling, x, extra_inputs)
+    assert reset_recipes == [current_scaling, overridden_current_scaling]
+    assert first_forward_calls == 1
+    assert fusion_calls == 2
+
+    delayed_scaling = transformer_engine.common.recipe.DelayedScaling(
+        amax_history_len=8,
+        backward_override=None,
+    )
+    fuser.maybe_fuse_ops(False, delayed_scaling, x, extra_inputs)
+    assert reset_recipes == [current_scaling, overridden_current_scaling, delayed_scaling]
+    assert first_forward_calls == 1
+    assert fusion_calls == 3
+
+    # Amax history length only affects delayed-scaling recipe state. Reset that
+    # state, but restore the existing DelayedScaling fusion plan from the cache.
+    resized_delayed_scaling = transformer_engine.common.recipe.DelayedScaling(
+        amax_history_len=16,
+        backward_override=None,
+    )
+    fuser.maybe_fuse_ops(False, resized_delayed_scaling, x, extra_inputs)
+    assert reset_recipes == [
+        current_scaling,
+        overridden_current_scaling,
+        delayed_scaling,
+        resized_delayed_scaling,
+    ]
+    assert first_forward_calls == 1
+    assert fusion_calls == 3
+
+    # Repeating the exact recipe parameters performs neither operation again.
+    fuser.maybe_fuse_ops(False, resized_delayed_scaling, x, extra_inputs)
+    assert len(reset_recipes) == 4
+    assert first_forward_calls == 1
+    assert fusion_calls == 3
+
+
 @torch.no_grad()
 def make_reference_and_test_tensors(
     shape: int | Iterable[int],
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 2128bae1ec..66c5bbb196 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -1404,7 +1404,7 @@ def fuser_forward(
 
         alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device)
         norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device)
-        current_stream = torch.cuda.current_stream().cuda_stream
+        current_stream = torch.cuda.current_stream(device.index).cuda_stream
 
         fc1_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc1_op)
         fc2_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc2_op)
@@ -2095,7 +2095,7 @@ def fuser_backward(
         # Kernel scaling factors
         alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device)
         norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device)
-        current_stream = torch.cuda.current_stream().cuda_stream
+        current_stream = torch.cuda.current_stream(device.index).cuda_stream
 
         unit_activation_scale = bool(getattr(fc1_ctx, "unit_activation_scale", False))
         scales_f32 = None
diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py
index fd66529ba8..2500002700 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -11,7 +11,7 @@
 
 import torch
 
-from ..quantization import FP8GlobalStateManager, Recipe, DelayedScaling
+from ..quantization import FP8GlobalStateManager, Recipe
 from ..quantized_tensor import prepare_for_saving, restore_from_func_ctx
 from .op import (
     BasicOperation,
@@ -48,6 +48,8 @@ def _is_graph_capturing() -> bool:
 OperationFusionFunction: TypeAlias = (
     "Callable[tuple[list[FusibleOperation], ...], list[FusibleOperation]]"
 )
+_FusedOpList: TypeAlias = list[tuple[FusibleOperation, list[int]]]
+_FusionParams: TypeAlias = tuple[type, int, Optional[str]]
 
 
 class _OperationFuserAutogradFunction(torch.autograd.Function):
@@ -535,15 +537,22 @@ def __init__(
             op._lock_extra_tensor_channels()
 
         # Ops for forward and backward pass, will be populated in maybe_fuse_ops
-        self._forward_ops: list[tuple[FusibleOperation, list[int]]]
-        self._backward_ops: list[tuple[FusibleOperation, list[int]]]
+        self._forward_ops: _FusedOpList
+        self._backward_ops: _FusedOpList
+
+        # Fused operation configurations are reusable wrappers around the basic
+        # ops, so cache each configuration by the state that selected it.
+        self._fused_ops_cache: dict[_FusionParams, tuple[_FusedOpList, _FusedOpList]] = {}
 
         # Cache and detect change of state relevant for fusing operations
         self.recipe_type = None
-        self.first_op_requiring_backward = 0
         self.backward_override = None
         self._last_amax_history_len = 0
 
+        # Runtime backward boundary. Full activation recompute alternates this
+        # between the checkpointed forward and the grad-enabled recomputation.
+        self.first_op_requiring_backward = 0
+
         # Flatten list of parameters
         self._basic_op_params = [list(op.parameters()) for op in self._basic_ops]
         self._basic_op_num_params = list(map(len, self._basic_op_params))
@@ -626,36 +635,48 @@ def maybe_fuse_ops(
                     first_op_requiring_backward = op_idx
                     break
 
-        # Early exit if fusion parameters haven't changed
-        need_reset = False
+        # Update the runtime backward boundary on every invocation, including
+        # paths that reuse a cached fused operation configuration.
+        self.first_op_requiring_backward = first_op_requiring_backward
+
+        # Check if recipe parameters don't match cached values. In this case,
+        # the recipe state in the basic ops might be invalid, so reset it.
         recipe_type = type(recipe)
+        need_to_reset_recipe_state = self.recipe_type != recipe_type
+
         backward_override = recipe.backward_override if recipe is not None else None
-        fusion_params = (recipe_type, first_op_requiring_backward, backward_override)
-        if fusion_params != (
-            self.recipe_type,
-            self.first_op_requiring_backward,
-            self.backward_override,
-        ):
-            # Recipe type, backward override, or grad requirements have changed
-            need_reset = True
-        elif (
+        if backward_override != self.backward_override:
+            self.backward_override = backward_override
+            need_to_reset_recipe_state = True
+
+        if (
             recipe is not None
             and recipe.delayed()
             and self._last_amax_history_len != recipe.amax_history_len
         ):
-            # FP8 delayed scaling has changed amax history length
-            need_reset = True
-        if not need_reset:
-            return
-
-        # Reset recipe state
-        for op in self._basic_ops:
-            op.reset_recipe_state(recipe=recipe)
+            self._last_amax_history_len = recipe.amax_history_len
+            need_to_reset_recipe_state = True
 
-        # Check if this is the first iteration
-        if self.recipe_type is None:
+        if need_to_reset_recipe_state:
             for op in self._basic_ops:
-                op.pre_first_fuser_forward()
+                op.reset_recipe_state(recipe=recipe)
+
+            # Check if this is the first iteration
+            if self.recipe_type is None:
+                for op in self._basic_ops:
+                    op.pre_first_fuser_forward()
+
+            self.recipe_type = recipe_type
+
+        # Training and inference may support different fusions. Keep the
+        # backward boundary in the key, but pay construction cost only once for
+        # each configuration. Full recompute therefore builds at most one
+        # no-grad plan and one grad-enabled plan for a stable recipe.
+        fusion_params = (recipe_type, first_op_requiring_backward, backward_override)
+        cached_ops = self._fused_ops_cache.get(fusion_params)
+        if cached_ops is not None:
+            self._forward_ops, self._backward_ops = cached_ops
+            return
 
         # Apply joint forward-backward fusions first
         joint_ops = OperationFuser._apply_fusions(
@@ -682,14 +703,9 @@ def maybe_fuse_ops(
             self._basic_ops,
         )
 
-        # Save current fusion params
-        self.recipe_type, self.first_op_requiring_backward, self.backward_override = fusion_params
-
-        # Save amax history length
-        if isinstance(recipe, DelayedScaling):
-            self._last_amax_history_len = recipe.amax_history_len
-        else:
-            self._last_amax_history_len = 0
+        # The FusedOperation contract excludes parameters and per-invocation
+        # state, so the mapped lists can be selected directly on cache hits.
+        self._fused_ops_cache[fusion_params] = (self._forward_ops, self._backward_ops)
 
     def __call__(
         self,

From d1e9c33449ef583937532a003aa26abac2fe8981 Mon Sep 17 00:00:00 2001
From: Zhiyu Li <zhiyul@NVIDIA.com>
Date: Fri, 4 Sep 2026 14:22:00 -0700
Subject: [PATCH 10/27] Use cuDNN's deterministic dprob in the fused grouped
 MLP (#3407)

* [PyTorch] Ask cuDNN for a deterministic dprob under NVTE_ALLOW_NONDETERMINISTIC_ALGO=0

The cuDNN grouped-GEMM dactivation backward that the CuTe DSL fused grouped MLP calls
accumulates the scale gradient (dprob) with cross-CTA atomic adds, so its floating-point
summation order follows the tile scheduler and varies run to run. Until now there was no
way to switch that off, and NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 did not reach it: the run
trained fine and was silently not reproducible.

cuDNN frontend 1.28.0 (NVIDIA/cudnn-frontend#521) added a `deterministic` argument to
grouped_gemm_dsrelu_wrapper_sm100 that parks each N-subtile's partial result in its own
slot and sums the slots in a canonical order, for dprob and for dbias. Pass it from the
TE flag.

Passed as True or not at all, never as False. The wrapper's own default is None, which
follows torch.use_deterministic_algorithms; sending an explicit False would override that
and take determinism away from a caller who asked torch for it without setting the TE
variable.

The capability is reported per subclass rather than per environment variable, because
grouped_gemm_dglu_wrapper_sm100 has no equivalent argument -- a GLU activation stays
non-deterministic however new the installed front-end is. That case, and an SReLU op on a
front-end older than 1.28.0, warn instead, once per distinct reason since the remedies
differ. The warning is raised from where dprob is actually produced: with a unit
activation scale the epilogue never runs its atomic accumulation, so there is nothing to
make deterministic and nothing to warn about.

Tests: TestGroupedMLPDeterminism covers the env-var parse, that only the SReLU op reports
the capability and that it tracks the front-end version (no GPU or cuDNN needed for
either), that the warning fires once per reason, and an MXFP8 end-to-end run under
determinism for both SwiGLU and SReLU that checks numerics and pins which of the two arms
warns.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Honor torch.use_deterministic_algorithms too, not just the env variable

_deterministic_algorithms_required() copied the narrow check from
transformer_engine.pytorch.triton.grouped_dbias_dscales, which reads
NVTE_ALLOW_NONDETERMINISTIC_ALGO and nothing else. DotProductAttention takes the union
instead -- the variable OR torch.use_deterministic_algorithms -- and that is the right
precedent here.

The two knobs answer different questions. The variable is set once in a job launcher,
applies uniformly across ranks, and is the only one TE's C++ layer can read. The torch
flag is the framework standard, is togglable at runtime, and is what a user who wants
reproducibility usually reaches for; most have never heard of the variable.

Keying on the variable alone left the torch flag half-honored. The SReLU path happened to
come out right, but by delegation rather than by decision: TE passed nothing and the
wrapper's own default read torch.are_deterministic_algorithms_enabled(). The GLU path did
not -- TE stayed silent about an atomic dprob it cannot fix, for a user who had asked
torch for reproducibility. That silence is the exact failure mode the warning exists to
prevent, so it was the one case that most needed to warn.

Passing the argument only as True, never as False, now needs a different justification
than the one the first commit gave: with the union in place the two are equivalent, since
the wrapper's default reads the same torch flag TE just read. The reason that survives is
narrower and firmer -- the argument does not exist on the dGLU wrapper or on a front-end
older than 1.28.0, where passing it at all, even as False, is a TypeError.

Tests: the env-var parametrization becomes the two-knob truth table, including the row
that motivates the change (torch flag set, NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 -- the
variable's default is the absence of a request, not a request for non-determinism, so the
torch flag still wins). A fixture restores the process-global torch flag.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Test that dprob is actually bit-exact, not just within tolerance

Review caught that nothing in the suite tested the property this change exists for. The
end-to-end test runs the op once and checks numerics against a reference with
rtol=0.125 / atol=0.25; reordering the same atomic adds moves dprob by about an ulp, so a
run that is silently not reproducible passes it comfortably. The tolerance check proves
the deterministic path is correct, which is worth keeping, but it cannot prove the path is
deterministic.

Add a second run. Same module, same inputs, grads cleared between passes, probs.grad
compared with torch.equal.

Three things the test has to get right to be worth having:

* hidden_size 1024, not the 128 used elsewhere. dprob's reduction is over that extent and
  the tile is 256 wide, so 128 gives a single N-tile, one writer per token, and nothing to
  reorder -- the assertion would hold by construction and test nothing.
* No bias. With an FC2 scale_bias the scale gradient is finished by the Triton grouped
  dbias/dscales kernel, which refuses to run under determinism, and probs.grad would stop
  being the dprob under test.
* An assertion that the fusion happened, since dprob only comes from the cuDNN epilogue on
  the fused path.

Skipped rather than xfailed on a front-end older than 1.28.0: there the kernel has no
deterministic mode and is expected to vary, which is not a failure of this change. Weight
gradients are deliberately left out of the comparison -- the CuTe DSL wgrad kernel has its
own K-split atomics that this PR does not address.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Probe the dsrelu wrapper's signature instead of the frontend version

_cudnn_frontend_supports_deterministic_dprob() gated on
_cudnn_frontend_version_at_least("1.28.0"). That check is too coarse to answer the
question it is asked, and would have raised at runtime on a build TE is actually run
against.

#521 merged after v1.27.0 was tagged, so `deterministic` ships in 1.28.0. But
cudnn-frontend's develop branch has called itself 1.28.0 since shortly after that tag --
eleven days before the merge. Any front-end built from develop in that window reports
1.28.0 and does not accept the argument, so the version check passes, TE adds
`deterministic=True` to the call, and the backward dies with

    TypeError: grouped_gemm_dsrelu_wrapper_sm100() got an unexpected keyword argument
    'deterministic'

This is not hypothetical, and not new. The same coarseness already bit
use_single_group_runtime_offsets: a cuDNN reporting 1.27.0 that did not implement 1.27.0's
arguments failed the identical way, in fuser_forward, before any backward code ran.
Version numbers describe a release; they do not describe whatever happens to be installed.

Ask the function instead. `"deterministic" in inspect.signature(...).parameters` is exact,
cannot drift, and needs no maintenance when the release lands. The import is wrapped the
way _grouped_gemm_dsrelu_backward_supported() already wraps it, so a missing cuDNN answers
False rather than raising. Cached, since the call site runs every backward.

This also removes the version constant from the code path entirely -- 1.28.0 now appears
only in user-facing text, where a release number is the useful thing to say.

Tests: a smoke test that the probe returns a bool without raising, with or without cuDNN
installed, since reading a signature has more ways to fail than comparing two version
strings. It deliberately does not assert which answer -- that depends on the installed
front-end, and pinning it would only restate the implementation.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Revert the unrelated nccl-extensions submodule bump

`git add -u` in the previous commit swept in a local 3rdparty/nccl-extensions pointer
change that has nothing to do with this PR. Restore it to main's commit so the branch
touches only the three files it means to.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Raise instead of warning, and cut the change down to what it needs

Review asked for two things on the unsupported path: make it an error rather than a
warning, and stop branching on self._cudnn_dact_func to pick a message. Both are right,
and taking them removes most of the machinery this PR had accumulated.

Raising matches what TE already does elsewhere: the Triton grouped dbias/dscales kernel
refuses to run under determinism rather than running non-deterministically. It also matches
what the variable documents -- "only deterministic algorithms are allowed" is not "prefer
deterministic algorithms". A silently non-reproducible run is the failure this PR exists to
prevent, so continuing past a request TE cannot honor was the wrong default. Checked that
no existing determinism test hits this path: test_hybrid_quantization sets the variable for
an attention recipe, and test_fusible_ops_with_userbuffers for linear ops.

One message, no branch. The two cases did have different remedies, which is why the branch
was there, but a single sentence states both facts -- "needs the scaled-SReLU activation
and nvidia-cudnn-frontend 1.28.0 or later" -- without telling a SwiGLU user to go upgrade.

What that let me delete:

* _warn_nondeterministic_cudnn_dprob and its per-reason lru_cache, the two reason strings
  and the branch selecting them: 30 lines at the call site and above it, down to a single
  raise.
* _cudnn_frontend_supports_deterministic_dprob as a standalone function. The probe now
  lives in GroupedMLP_CuTeGEMMUnary.grouped_gemm_dactivation_is_deterministic(), which
  reaches the wrapper through grouped_gemm_dactivation_kernel() -- the import and its
  ImportError handling already existed there, so folding it in dropped a duplicate import
  and an indirection.
* The warn-once cache-clearing fixture in the tests, and the two tests that existed only to
  cover the warning.

Tests: test_deterministic_dactivation_is_numerically_correct becomes
test_determinism_either_runs_or_refuses -- it expects RuntimeError where the request cannot
be honored and runs the full numerical check where it can, so both arms assert something
either way. The bit-exactness and two-knob tests are unchanged in substance.

Net: transformer_engine/pytorch/ops/fused/grouped_mlp.py goes from +106 to +68, all of it
addition, no line of pre-existing code touched.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Apply suggestion from @vthumbe1503

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

* Match the feature-detection idiom main just landed

The SiTU-GLU merge (#3402) brought _cudnn_frontend_supports_grouped_gemm_situglu() into
this file, which asks inspect.signature(wrapper).parameters for the arguments it needs
rather than comparing frontend versions -- the same conclusion this branch reached
independently, now the house style.

Two things to match. Guard the signature call with `except (TypeError, ValueError)`: a
callable that is not introspectable answers "no" instead of raising out of a backward pass.
I had left this out on the grounds that the wrapper is a plain undecorated function, which
is true today but is not a property this code controls. And say "feature-detect" in the
docstring summary, as the neighbor does.

Also dropped the sentence about use_single_group_runtime_offsets from the docstring. The
neighbor now demonstrates the pattern in the same file, so the cautionary tale is no longer
what makes the choice legible.

`import inspect` came in with the merge, so this branch no longer adds it.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Cut the comments down to the file's own register

The new code carried multi-paragraph docstrings into a file whose 44 functions have a
median docstring of one line. Measured before and after:

  grouped_mlp.py     _deterministic_algorithms_required           10 -> 3 lines
                     grouped_gemm_dactivation_is_deterministic (base)  5 -> 1
                     grouped_gemm_dactivation_is_deterministic (unary) 7 -> 1
  test_grouped_mlp.py  four new tests                            4-6 -> 1-4
                       four inline comment blocks                2-3 -> 1 each

Before this, the three new functions were the 2nd, 3rd and 5th longest docstrings in
grouped_mlp.py; only fuse_grouped_mlp_ops, which has a full Parameters block, was longer.
In the test file, 63 pre-existing tests have a median docstring of zero lines.

Most of what came out was rationale, not explanation: why the union matches
DotProductAttention, why feature detection beats a version compare, which cuDNN release
window motivated it. That belongs in the commits that made those choices, where it already
is, and it reads as noise next to _cudnn_frontend_supports_grouped_gemm_situglu -- the
neighbor doing the very same feature detection in a one-line docstring with no rationale
at all.

What stayed is what the code cannot say itself: that the check sits inside the
non-unit-scale branch because a unit scale produces no dprob; that hidden_size must exceed
one N-tile or the bit-exactness test is vacuous; that bias would reroute probs.grad through
Triton; that weight grads are excluded because wgrad has its own atomics. Each is now one
line.

No behavior change -- comments, docstrings and one local variable's reading order only.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Flatten the determinism check and the runs-or-refuses test

Structural cleanups from the review pass.

grouped_mlp.py: the check was nested two deep inside `if not unit_activation_scale`, and
assigned deterministic_dactivation only to immediately test its own assignment. Hoisted to
two flat statements right after unit_activation_scale is computed. `not
unit_activation_scale and _deterministic_algorithms_required()` now says in the expression
what the comment had to say in prose, and the separate `= False` initializer is gone. The
local itself stays -- the kwargs dict is built about sixty lines further down.

Also shortened the error: the tile-scheduler detail was not actionable, and "this
activation's cuDNN dactivation kernel" is more accurate than naming the grouped-GEMM
backward, since which kernel it is depends on the activation.

test_grouped_mlp.py: fused_cls was derived from `activation` by a five-line conditional
inside the test; it is now the second half of the parametrize pair. That also fixes the
skip guard, which asked GroupedMLP_CuTeGEMMGLU.is_supported() on both parametrizations
including the SReLU one -- the sibling test three functions down already gets this right.
The _run closure existed only so an if/else could call it twice; a contextlib.nullcontext
/ pytest.raises choice removes the closure and the branch. nullcontext is used in ten test
files here, so it is the local idiom rather than a new one.

Not taken: dropping the `isinstance(..., bool)` assertion. It looks vacuous but it is the
only coverage of the ImportError branch in the capability probe, which is the branch that
runs on every machine without cuDNN -- including CI.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Close the second dprob producer, and stop discarding fp64 test tensors

Two findings from the review pass.

dprob has two producers in this backward, and the check only covered one. The cuDNN
epilogue produces grad_scales at fuser_backward, and when scale_bias is set
compute_grouped_dbias_dscales accumulates into it further down -- the Triton kernel that
grouped_dbias_dscales.py documents as nondeterministic atomic adds. That kernel's own guard
reads NVTE_ALLOW_NONDETERMINISTIC_ALGO and nothing else.

So the hole opened exactly where this branch widened the trigger. With
torch.use_deterministic_algorithms(True) and the variable unset -- the case the union
exists to start honoring -- SReLU on a 1.28.0 front-end with scale_bias passed the new
check, set deterministic=True, raised nothing, and then routed dprob through the
nondeterministic path anyway. Env-var users were never exposed: the Triton guard fires for
them. It was reachable only via the torch flag, which is to say only through what this
branch added. The test picked bias=False and so never crossed it.

scale_bias is computed ~130 lines earlier in the same scope, so the fix is to require both
producers rather than one. Still one condition and one message, per review -- the message
now lists all three requirements instead of two.

Separately, the bit-exactness test built its tensors with make_reference_and_test_tensors
and discarded the reference every time. That helper allocates an fp64 CPU companion,
quantizes and dequantizes for MXFP8 representability, then copies back D2H with an implicit
sync -- about 16 MB of host allocation across the two (1024, 1024) calls, for a test that
compares run 1 against run 2 and never against a reference. Twelve of the file's other
fifteen uses keep the reference; this one had no use for it. Plain uniform_ tensors instead.
Also dropped a .item() sync for a token count already known in Python.

Not taken, with reasons:
* Hoisting _deterministic_algorithms_required into pytorch/utils.py so the Triton guard
  reads the same union. That is the deeper fix and it is correct, but broadening that guard
  changes behavior for callers this PR does not touch (ops/basic/grouped_linear.py,
  module/grouped_linear.py) -- users who set only the torch flag would start seeing
  RuntimeError where they now get silent nondeterminism. Worth doing deliberately, not as a
  side effect of this branch.
* Extracting the signature-probe shared with _cudnn_frontend_supports_grouped_gemm_situglu.
  The overlap is about four lines and the two are not interchangeable; refactoring working
  code outside the diff to save them is not this PR's job.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Add the regression test for the scale_bias hole

The previous commit fixed a real bug and shipped it with no test. Every test in the class
used bias=False and every end-to-end one set the env var, so neither half of the bug was
reachable: not scale_bias, and not the torch-flag-only trigger.

Both halves are load-bearing. With the env var the Triton kernel raises on its own, so an
env-var test would have passed before the fix as well as after and pinned nothing. Only
torch.use_deterministic_algorithms with the variable unset reaches the state where this
op's check said yes and the Triton reduction then ran nondeterministically.

warn_only=True so torch's own enforcement cannot raise first and be mistaken for TE's
refusal. are_deterministic_algorithms_enabled() still reports True in that mode -- the
separate is_deterministic_algorithms_warn_only_enabled() getter exists precisely because
the two are independent -- so the predicate under test sees what it should.

Not executed: no GPU or torch on the machine this was written on. Formatting and syntax
only, like the rest of the branch.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Make the bit-exactness test capable of failing

Followed cudnn-frontend#521's own test work and found this test had the flaw its commit
88c7fab was written to fix, at the same config.

That commit measured 16 launches per shape and found that at l=4 / [256]*4 / n=512 the
NONDETERMINISTIC dprob is already bit-stable: the assertion cannot fail there, so a pass
certifies nothing. It varies 15/15 at l=8 / [1024]*8 / n=2048. This test used
l=4 / [256]*4 / n=1024 -- the vacuous shape, one power of two along n. Moved to the shape
that actually varies.

n > 256 was necessary but not sufficient, which is what the old comment got wrong. Spanning
several N-tiles exercises the within-CTA subtile ordering; making the cross-CTA reduction
unstable needs the larger token count and expert count too.

Also took the rest of #521's discipline for these comparisons:

* Repeat rather than compare a pair. The order determinism removes is set by the tile
  scheduler, so two runs can match by luck. Four by default, NVTE_TEST_DETERMINISM_REPEATS
  to raise it, matching that file's DETERMINISM_REPEATS.
* Compare bytes, not values. torch.equal treats +0.0 and -0.0 as equal, and a change in
  reduction order produces exactly that; upstream's bitwise_bits views as uint8 for the
  same reason.
* Assert the output is finite first, so a NaN run cannot be read as a determinism result.

Not copied: asserting that the nondeterministic path *does* vary. It is the thing that
makes the config meaningful, but as an assertion it is timing-dependent and would flake.
Upstream settled this by measuring once and pinning the config; the comment now cites that
measurement so the next person does not shrink the shape back.

Not run yet: job 535935 is building the previous revision of this test.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Pick the bit-exactness config by measuring it, not by borrowing one

Measured on GB300 across five shapes, determinism off, 8 launches each (job 538058),
counting how many runs differ from run 0:

  l=8  tok/grp=1024 n=2048   2/7   max|d| 5.96e-08
  l=8  tok/grp=1024 n=4096   2/7   max|d| 9.54e-07
  l=16 tok/grp=1024 n=2048   7/7   max|d| 1.19e-07
  l=8  tok/grp=2048 n=2048   5/7   max|d| 7.63e-06
  l=4  tok/grp=512  n=8192   6/7   max|d| 1.91e-06

Moved to l=16, the only shape where every run differs, so the assertion cannot pass by
luck. The previous choice, l=8, varies 2/7 -- an eight-run sample calls it stable often
enough to be a poor detector, and an earlier control run (537313) did exactly that and
reported 0/7 at this shape. I took that single sample as proof the config was vacuous and
said so; it was a sampling artifact, and the shape does vary, just weakly.

That earlier shape came from cudnn-frontend#521's own measurement, which was taken on its
direct wrapper test. It does not transfer to TE's path -- different scheduler settings,
different quantization -- so borrowing the number was the mistake underneath both errors.
This config is measured through the fused grouped MLP itself.

Two things the same job settled that are worth recording:

* Without #521 the values genuinely move: 6e-08 to 8e-06 absolute across these shapes.
  Small, but nonzero every time, and the reason the refusal exists rather than a warning.
* The refusal cannot be exercised against the stock 1.27.0 frontend on this image at all.
  TE's forward passes prob_tensor=None because _cudnn_frontend_version_at_least("1.27.0")
  reports optional-prob support that a stock 1.27.0 does not implement, so the op dies in
  fuser_forward with "prob_tensor is required" before any determinism code runs. Same
  version-gate-too-coarse failure this PR avoids for its own argument, on a gate it does
  not own.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

* Make the scale_bias half of the dprob check readable

The condition was written as `not (A and not B)` with `scale_bias` as B, which gives a
reader no way to tell why an FC2 bias flag decides whether a scale gradient is
reproducible -- and the call that makes it relevant is ~250 lines further down.

Same logic, named and nested: `dprob_is_deterministic` says what the conjunction means,
and the comment names the mechanism instead of gesturing at it. dprob is finished by two
kernels, not one -- the cuDNN dactivation epilogue writes it, and then, when scale_bias is
set, fuser_backward hands it to compute_grouped_dbias_dscales as the `dscales` accumulator,
which atomically adds into it. Its docstring is explicit: "Both outputs use fp32 atomic
adds, so pre-populated tensors are accumulated into."

No logic change; the truth table is identical.

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>

---------

Signed-off-by: Zhiyu Li <zhiyul@nvidia.com>
Signed-off-by: vthumbe1503 <vthumbe@nvidia.com>
Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
---
 tests/pytorch/test_grouped_mlp.py             | 189 ++++++++++++++++++
 .../pytorch/ops/fused/grouped_mlp.py          |  54 +++++
 2 files changed, 243 insertions(+)

diff --git a/tests/pytorch/test_grouped_mlp.py b/tests/pytorch/test_grouped_mlp.py
index 1d173c76d8..15c1ff6c51 100644
--- a/tests/pytorch/test_grouped_mlp.py
+++ b/tests/pytorch/test_grouped_mlp.py
@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
+import contextlib
 import functools
 import os
 import math
@@ -2845,6 +2846,194 @@ def train_step(
                 assert_close(graph_grad, param.grad, **tols)
 
 
+class TestGroupedMLPDeterminism:
+    """Determinism coverage for the CuTe DSL fused grouped MLP.
+
+    Only the dSReLU wrapper can make ``dprob`` bit-exact, and only from cuDNN FE 1.28.0 on.
+    Anything else must refuse a determinism request rather than run non-deterministically.
+    """
+
+    @pytest.fixture
+    def _restore_torch_determinism(self):
+        """``use_deterministic_algorithms`` is process-global, so put it back."""
+        previous = torch.are_deterministic_algorithms_enabled()
+        yield
+        torch.use_deterministic_algorithms(previous)
+
+    @pytest.mark.parametrize(
+        "allow_nondeterministic,torch_flag,expected",
+        (
+            (None, False, False),  # default: non-deterministic algorithms are allowed
+            ("1", False, False),
+            ("0", False, True),  # the TE variable alone
+            (None, True, True),  # the torch flag alone, which TE must not ignore
+            ("1", True, True),  # ... including when the TE variable says otherwise
+            ("0", True, True),
+        ),
+    )
+    def test_either_knob_requests_determinism(
+        self,
+        monkeypatch,
+        _restore_torch_determinism,
+        *,
+        allow_nondeterministic: Optional[str],
+        torch_flag: bool,
+        expected: bool,
+    ) -> None:
+        """``=1`` is the absence of a request, not a request for non-determinism."""
+        if allow_nondeterministic is None:
+            monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False)
+        else:
+            monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", allow_nondeterministic)
+        torch.use_deterministic_algorithms(torch_flag)
+        assert grouped_mlp_module._deterministic_algorithms_required() is expected
+
+    def test_only_the_srelu_path_can_be_deterministic(self) -> None:
+        """The capability belongs to the wrapper, not the environment. Needs no GPU."""
+        glu = grouped_mlp_module.GroupedMLP_CuTeGEMMGLU
+        unary = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary
+        assert glu.grouped_gemm_dactivation_is_deterministic() is False
+        assert isinstance(unary.grouped_gemm_dactivation_is_deterministic(), bool)
+
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    @pytest.mark.parametrize(
+        "activation,fused_cls",
+        (
+            ("scaled_srelu", grouped_mlp_module.GroupedMLP_CuTeGEMMUnary),
+            ("scaled_swiglu", grouped_mlp_module.GroupedMLP_CuTeGEMMGLU),
+        ),
+    )
+    def test_determinism_either_runs_or_refuses(
+        self, monkeypatch, *, activation, fused_cls
+    ) -> None:
+        """A request TE cannot honor must fail loudly; one it can must still be correct."""
+        if not fused_cls.is_supported():
+            pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
+
+        monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0")
+        expectation = (
+            contextlib.nullcontext()
+            if fused_cls.grouped_gemm_dactivation_is_deterministic()
+            else pytest.raises(RuntimeError, match="dprob")
+        )
+        with expectation:
+            TestGroupedMLPFusedOp().test_grouped_mlp(
+                bias=False,
+                hidden_size=128,
+                quantization="mxfp8",
+                single_grouped_weight=False,
+                activation=activation,
+            )
+
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    def test_scale_bias_refuses_under_the_torch_flag(
+        self, monkeypatch, _restore_torch_determinism
+    ) -> None:
+        """``scale_bias`` finishes ``dprob`` in a Triton kernel that reads only the env var.
+
+        So the torch flag alone is the combination that used to pass this op's own check and
+        then reduce nondeterministically anyway, on a front-end new enough to say yes.
+        """
+        fused_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary
+        if not fused_cls.is_supported():
+            pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
+
+        monkeypatch.delenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", raising=False)
+        # warn_only so torch's own enforcement cannot raise first and mask what TE does.
+        torch.use_deterministic_algorithms(True, warn_only=True)
+        with pytest.raises(RuntimeError, match="dprob"):
+            TestGroupedMLPFusedOp().test_grouped_mlp(
+                bias=True,
+                hidden_size=128,
+                quantization="mxfp8",
+                single_grouped_weight=False,
+                activation="scaled_srelu",
+            )
+
+    @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
+    def test_dprob_is_bit_exact_across_runs(self, monkeypatch) -> None:
+        """Repeated identical runs must give a bit-identical ``dprob``.
+
+        An ulp of reordering passes every tolerance in this file, so only an exact
+        comparison across runs can see it.
+        """
+        fused_cls = grouped_mlp_module.GroupedMLP_CuTeGEMMUnary
+        if not fused_cls.is_supported():
+            pytest.skip("MXFP8 fused grouped MLP is not supported on this system")
+        if not fused_cls.grouped_gemm_dactivation_is_deterministic():
+            pytest.skip("dSReLU determinism needs cuDNN frontend 1.28.0 or later")
+
+        monkeypatch.setenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "0")
+
+        device = torch.device("cuda")
+        dtype = torch.bfloat16
+        # Measured on GB300, determinism off, 8 launches per shape (job 538058): this shape
+        # gives 7/7 runs differing from run 0, so the assertion below can actually fail.
+        # Shapes matter more than they look -- l=8 with the same n and tokens/group varies
+        # only 2/7, which an 8-run sample reports as stable often enough to be useless, and
+        # cudnn-frontend#521 measured its own l=4 / [256]*4 / n=512 as never varying.
+        group_size = 16
+        hidden_size = 2048
+        tokens_per_group = 1024
+        split_sizes = torch.tensor([tokens_per_group] * group_size, dtype=torch.int, device=device)
+        num_tokens = tokens_per_group * group_size
+
+        recipe = make_recipe("mxfp8")
+
+        # Plain random tensors, not make_reference_and_test_tensors: this test compares two
+        # runs against each other, never against a reference, so the fp64 companion and the
+        # MXFP8 representability round-trip would both be allocated and thrown away.
+        def _rand(*shape, requires_grad=True) -> torch.Tensor:
+            out = torch.empty(shape, dtype=dtype, device=device).uniform_(-0.25, 0.25)
+            return out.requires_grad_() if requires_grad else out
+
+        x = _rand(num_tokens, hidden_size)
+        dy = _rand(num_tokens, hidden_size, requires_grad=False)
+        probs = _rand(num_tokens)
+
+        # No bias, or probs.grad comes from the Triton dbias kernel instead of cuDNN.
+        with te.quantized_model_init(enabled=True, recipe=recipe):
+            module = te.ops.Sequential(
+                te.ops.GroupedLinear(
+                    group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype
+                ),
+                te.ops.ScaledSReLU(),
+                te.ops.GroupedLinear(
+                    group_size, hidden_size, hidden_size, bias=False, device=device, dtype=dtype
+                ),
+            )
+
+        def _run() -> torch.Tensor:
+            x.grad = None
+            probs.grad = None
+            with te.autocast(enabled=True, recipe=recipe):
+                y = module(x, split_sizes, probs, split_sizes)
+            y.backward(dy)
+            return probs.grad.detach().clone()
+
+        runs = [_run()]
+        # Without the fusion there is no cuDNN dprob and the comparison proves nothing.
+        forward_ops = module._module_groups[0]._forward_ops
+        assert len(forward_ops) == 1
+        assert isinstance(forward_ops[0][0], fused_cls)
+        # More than two, as cudnn-frontend#521 does: the cross-CTA order that determinism
+        # removes is set by the scheduler, so two runs can agree by luck.
+        runs += [_run() for _ in range(int(os.getenv("NVTE_TEST_DETERMINISM_REPEATS", "4")) - 1)]
+        torch.cuda.synchronize()
+
+        assert torch.isfinite(runs[0]).all(), "dprob is not finite; the comparison would be moot"
+        # Bytes, not values: torch.equal calls +0.0 and -0.0 equal, and a change in reduction
+        # order can produce exactly that. Weight grads are excluded from the comparison --
+        # the CuTe DSL wgrad kernel has its own K-split atomics, which this change leaves.
+        for index, later in enumerate(runs[1:], start=1):
+            assert torch.equal(
+                runs[0].contiguous().view(torch.uint8), later.contiguous().view(torch.uint8)
+            ), (
+                f"dprob differs between run 0 and run {index} under determinism; max |delta| ="
+                f" {(runs[0].float() - later.float()).abs().max().item()}"
+            )
+
+
 def test_grouped_gemm_quant_cute_matches_mxfp8_quantized() -> None:
     if not mxfp8_available:
         pytest.skip(reason_for_no_mxfp8)
diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 66c5bbb196..58f22c874c 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -170,6 +170,17 @@ def _cudnn_frontend_supports_single_group_runtime_offsets(
     ) and _cudnn_frontend_version_at_least("1.27.0")
 
 
+def _deterministic_algorithms_required() -> bool:
+    """Whether bit-exact reproducibility was asked for. Same union as ``DotProductAttention``.
+
+    Uncached: both knobs can change during the process.
+    """
+    return (
+        not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1")))
+        or torch.are_deterministic_algorithms_enabled()
+    )
+
+
 def _wrap_single_quantized_as_grouped(
     tensor: torch.Tensor,
     quantized: MXFP8Tensor | NVFP4Tensor | NVFP4TensorStorage,
@@ -956,6 +967,11 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable:
         """Fused kernel for grouped GEMM, activation backward, and scale grad."""
         raise NotImplementedError
 
+    @classmethod
+    def grouped_gemm_dactivation_is_deterministic(cls) -> bool:
+        """Whether this op's dactivation kernel can produce a bit-exact ``dprob``."""
+        return False
+
     @classmethod
     @functools.lru_cache(maxsize=None)
     def grouped_gemm_quant_kernel(cls) -> Callable:
@@ -2098,6 +2114,28 @@ def fuser_backward(
         current_stream = torch.cuda.current_stream(device.index).cuda_stream
 
         unit_activation_scale = bool(getattr(fc1_ctx, "unit_activation_scale", False))
+        # A unit activation scale produces no dprob, so there is nothing to make deterministic.
+        deterministic_dactivation = (
+            not unit_activation_scale and _deterministic_algorithms_required()
+        )
+        if deterministic_dactivation:
+            # Two kernels write dprob and both have to be exact. The cuDNN dactivation
+            # epilogue produces it below; then, when scale_bias is set, it is passed to
+            # compute_grouped_dbias_dscales as the ``dscales`` accumulator and atomically
+            # added into (see triton/grouped_dbias_dscales.py). That Triton kernel is never
+            # deterministic, so scale_bias rules out a bit-exact dprob on its own.
+            dprob_is_deterministic = (
+                self.grouped_gemm_dactivation_is_deterministic() and not scale_bias
+            )
+            if not dprob_is_deterministic:
+                raise RuntimeError(
+                    "Deterministic execution was requested"
+                    " (NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 or"
+                    " torch.use_deterministic_algorithms), but the scale gradient (dprob) is"
+                    " accumulated with nondeterministic atomics on this configuration."
+                    " A bit-exact dprob requires the scaled-SReLU activation,"
+                    " nvidia-cudnn-frontend 1.28.0 or later, and an FC2 without scale_bias."
+                )
         scales_f32 = None
         scales_tensor = None
         dscales_tensor = None
@@ -2152,6 +2190,9 @@ def fuser_backward(
             "use_dynamic_sched": True,
         }
         dactivation_kernel = self.grouped_gemm_dactivation_kernel()
+        if deterministic_dactivation:
+            # Never passed to a wrapper that would reject it -- the check above raises first.
+            fc2_dactivation_kwargs["deterministic"] = True
         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:
@@ -2682,6 +2723,19 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable:
 
         return grouped_gemm_dsrelu_wrapper_sm100
 
+    @classmethod
+    @functools.lru_cache(maxsize=None)
+    def grouped_gemm_dactivation_is_deterministic(cls) -> bool:
+        """Feature-detect the dSReLU wrapper's ``deterministic`` argument (cuDNN FE 1.28.0+)."""
+        try:
+            kernel = cls.grouped_gemm_dactivation_kernel()
+        except ImportError:
+            return False
+        try:
+            return "deterministic" in inspect.signature(kernel).parameters
+        except (TypeError, ValueError):
+            return False
+
 
 def fuse_ops(
     ops: list[FusibleOperation],

From 29229b5a030d1672bb79af61a735caf42ce1effc Mon Sep 17 00:00:00 2001
From: Robin Zhang <robinz@nvidia.com>
Date: Sat, 5 Sep 2026 06:51:41 +0800
Subject: [PATCH 11/27] [PyTorch] Reduce CUDA graph memory retention (#3427)

* [PyTorch] Release warmup outputs after their last use

Signed-off-by: Robin Zhang <robinz@nvidia.com>

* [PyTorch] Release buffer-reuse capture temporaries

Signed-off-by: Robin Zhang <robinz@nvidia.com>

* [PyTorch] Release per-callable state on reset

Signed-off-by: Robin Zhang <robinz@nvidia.com>

* [PyTorch] Bundle per-callable lifecycle helpers

Signed-off-by: Robin Zhang <robinz@nvidia.com>

---------

Signed-off-by: Robin Zhang <robinz@nvidia.com>
---
 tests/pytorch/test_cuda_graphs.py   | 202 ++++++++++++++++++++++++++++
 transformer_engine/pytorch/graph.py | 129 +++++++++++++++---
 2 files changed, 310 insertions(+), 21 deletions(-)

diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py
index 5a848dc0e8..8f85f57f32 100644
--- a/tests/pytorch/test_cuda_graphs.py
+++ b/tests/pytorch/test_cuda_graphs.py
@@ -5,6 +5,8 @@
 from typing import Callable, Dict, Iterable, List, Tuple, Union
 import pytest
 import copy
+import gc
+import weakref
 
 import torch
 from transformer_engine.pytorch import (
@@ -994,6 +996,206 @@ def hook(module: torch.nn.Module) -> None:
     ]
 
 
+def test_ordered_warmup_releases_consumed_outputs() -> None:
+    """Ordered warmup should only retain outputs until their corresponding backward."""
+
+    class OutputLifetimeModule(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.previous_output = None
+
+        def forward(self, input_: torch.Tensor) -> torch.Tensor:
+            is_warmup = not torch.cuda.is_current_stream_capturing()
+            if is_warmup and self.previous_output is not None:
+                assert self.previous_output() is None
+            output = input_ * 2
+            if is_warmup:
+                self.previous_output = weakref.ref(output)
+            return output
+
+    module = OutputLifetimeModule()
+    sample_args = tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in range(2))
+    graphed_callables = make_graphed_callables(
+        (module,),
+        sample_args,
+        num_warmup_iters=2,
+        _order=[1, -1, 1, -1],
+        _num_layers_per_chunk=[1],
+    )
+    assert module.previous_output is not None
+    assert module.previous_output() is None
+    reset_graphs(graphed_callables)
+
+
+def test_unordered_warmup_releases_consumed_outputs() -> None:
+    """Unordered warmup should release each output after its corresponding backward."""
+
+    class OutputLifetimeModule(torch.nn.Module):
+        def __init__(self, output_refs: list, module_idx: int) -> None:
+            super().__init__()
+            self.output_refs = output_refs
+            self.module_idx = module_idx
+            self.capture_started = False
+
+        def forward(self, input_: torch.Tensor) -> torch.Tensor:
+            output = input_ * 2
+            if torch.cuda.is_current_stream_capturing():
+                self.capture_started = True
+            else:
+                self.output_refs[self.module_idx] = weakref.ref(output)
+            return output
+
+    output_refs = [None, None]
+    modules = tuple(OutputLifetimeModule(output_refs, module_idx) for module_idx in range(2))
+
+    def first_module_backward_pre_hook(_module: torch.nn.Module) -> None:
+        if not modules[0].capture_started:
+            assert output_refs[1] is not None
+            assert output_refs[1]() is None
+
+    graphed_callables = make_graphed_callables(
+        modules,
+        tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in modules),
+        num_warmup_iters=2,
+        capture_time_hooks=[
+            {"backward_pre_hooks": {0: first_module_backward_pre_hook}},
+            None,
+        ],
+    )
+    assert all(output_ref is not None and output_ref() is None for output_ref in output_refs)
+    reset_graphs(graphed_callables)
+
+
+def test_inference_warmup_does_not_retain_outputs() -> None:
+    """Inference warmup should release outputs as soon as each forward returns."""
+
+    class OutputLifetimeModule(torch.nn.Module):
+        def __init__(self, previous_output: list) -> None:
+            super().__init__()
+            self.previous_output = previous_output
+
+        def forward(self, input_: torch.Tensor) -> torch.Tensor:
+            is_warmup = not torch.cuda.is_current_stream_capturing()
+            if is_warmup and self.previous_output[0] is not None:
+                assert self.previous_output[0]() is None
+            output = input_ * 2
+            if is_warmup:
+                self.previous_output[0] = weakref.ref(output)
+            return output
+
+    previous_output = [None]
+    modules = tuple(OutputLifetimeModule(previous_output).eval() for _ in range(2))
+    graphed_callables = make_graphed_callables(
+        modules,
+        tuple((torch.ones(4, 8, device="cuda"),) for _ in modules),
+        num_warmup_iters=2,
+    )
+    assert previous_output[0] is not None
+    assert previous_output[0]() is None
+    reset_graphs(graphed_callables)
+
+
+def test_reused_capture_buffers_release_outputs_after_backward() -> None:
+    """Capture locals must not keep weak-refed output buffers alive."""
+
+    class OutputLifetimeModule(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.previous_capture_output = None
+
+        def forward(self, input_: torch.Tensor) -> torch.Tensor:
+            if (
+                torch.cuda.is_current_stream_capturing()
+                and self.previous_capture_output is not None
+            ):
+                assert self.previous_capture_output() is None
+            output = input_ * 2
+            if torch.cuda.is_current_stream_capturing():
+                self.previous_capture_output = weakref.ref(output)
+            return output
+
+    module = OutputLifetimeModule()
+    sample_args = tuple((torch.ones(4, 8, device="cuda", requires_grad=True),) for _ in range(2))
+    graphed_callables = make_graphed_callables(
+        (module,),
+        sample_args,
+        _order=[1, -1, 1, -1],
+        _num_layers_per_chunk=[1],
+        _reuse_graph_input_output_buffers=True,
+    )
+    assert module.previous_capture_output is not None
+    assert module.previous_capture_output() is None
+    reset_graphs(graphed_callables)
+
+
+def test_reset_releases_only_the_selected_callable() -> None:
+    """Reset releases one callable's graph state without retaining its peers."""
+
+    class CaptureOutputModule(torch.nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.capture_output = None
+
+        def forward(self, input_: torch.Tensor) -> torch.Tensor:
+            output = input_ * 2
+            if torch.cuda.is_current_stream_capturing():
+                self.capture_output = weakref.ref(output)
+            return output
+
+    modules = tuple(CaptureOutputModule().cuda() for _ in range(2))
+    graphed_callables = make_graphed_callables(
+        modules,
+        tuple((torch.ones(4, device="cuda", requires_grad=True),) for _ in modules),
+    )
+    capture_outputs = tuple(module.capture_output for module in modules)
+    assert all(output is not None and output() is not None for output in capture_outputs)
+
+    graphed_callables[0].reset()
+    graphed_callables[0].reset()
+    gc.collect()
+    assert capture_outputs[0]() is None
+    assert capture_outputs[1]() is not None
+
+    output = graphed_callables[1](torch.randn(4, device="cuda", requires_grad=True))
+    output.sum().backward()
+    del output
+    graphed_callables[1].reset()
+    gc.collect()
+    assert capture_outputs[1]() is None
+
+
+@pytest.mark.parametrize("with_order", (False, True))
+def test_reset_rejects_all_replay_entry_points(with_order: bool) -> None:
+    """Reset is idempotent and terminal for forward and backward replay."""
+
+    class TestModule(torch.nn.Module):
+        def forward(self, input_: torch.Tensor) -> torch.Tensor:
+            return input_ * 2
+
+    module = TestModule().cuda()
+    sample_input = torch.ones(4, device="cuda", requires_grad=True)
+    graph_options = {}
+    if with_order:
+        graph_options = {"_order": [1, -1], "_num_layers_per_chunk": [1]}
+    graphed_callable = make_graphed_callables(module, (sample_input,), **graph_options)
+    output = graphed_callable(torch.randn_like(sample_input, requires_grad=True))
+    torch.cuda.synchronize()
+
+    graphed_callable.reset()
+    graphed_callable.reset()
+    if not with_order:
+        # The eager fallback for a different training state is invalid after reset too.
+        graphed_callable.eval()
+
+    error = "has been reset and can no longer be used"
+    with pytest.raises(RuntimeError, match=error):
+        graphed_callable(torch.randn_like(sample_input, requires_grad=True))
+    with pytest.raises(RuntimeError, match=error):
+        graphed_callable.backward_dw()
+    with pytest.raises(RuntimeError, match=error):
+        output.sum().backward()
+
+
 @pytest.mark.parametrize("with_order", (False, True))
 def test_make_graphed_callables_with_capture_time_hooks(with_order: bool) -> None:
     """Test capture-time hooks around warmup and graph capture."""
diff --git a/transformer_engine/pytorch/graph.py b/transformer_engine/pytorch/graph.py
index b298b3d8ff..04fa56721d 100644
--- a/transformer_engine/pytorch/graph.py
+++ b/transformer_engine/pytorch/graph.py
@@ -8,7 +8,7 @@
 import gc
 import warnings
 from math import ceil
-from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union
+from typing import Any, Callable, Dict, List, NamedTuple, Optional, Tuple, TypeVar, Union
 
 import torch
 from torch.utils._pytree import tree_flatten as _tree_flatten
@@ -44,6 +44,13 @@
 )
 
 
+class _GraphedCallableHelpers(NamedTuple):
+    """Lifecycle helpers owned by one graphed callable invocation."""
+
+    ensure_not_reset: Callable[[], None]
+    release_static_state: Callable[[], None]
+
+
 def set_capture_start() -> None:
     """Record beginning of `make_graphed_callables`."""
     global _IS_GRAPH_CAPTURING
@@ -633,13 +640,17 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx):
                 warmup_outputs = []
                 for func_idx, func in zip(warmup_func_idx, warmup_func):
                     outputs = _run_warmup_forward(func_idx, func, func_idx)
-                    warmup_outputs.append((func_idx, func, outputs))
-                if is_training:
-                    for func_idx, func, outputs in reversed(warmup_outputs):
-                        _run_warmup_backward(func_idx, func, outputs, warmup_iter, func_idx)
+                    if is_training:
+                        warmup_outputs.append((func_idx, func, outputs))
+                    else:
+                        del outputs
+                while warmup_outputs:
+                    func_idx, func, outputs = warmup_outputs.pop()
+                    _run_warmup_backward(func_idx, func, outputs, warmup_iter, func_idx)
+                    del outputs
             else:
                 # Follow _order exactly, mirroring the capture phase.
-                per_fwd_outputs = {}  # per_callable_fwd_idx -> flattened outputs
+                per_fwd_outputs = {}  # per_callable_fwd_idx -> outstanding flattened outputs
                 fwd_idx = [0] * num_model_chunks
                 bwd_idx = [0] * num_model_chunks
                 for c_id in _order:
@@ -653,7 +664,10 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx):
                             ) + (fwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no)
                             func = callables[callable_idx]
                             outputs = _run_warmup_forward(per_callable_fwd_idx, func, callable_idx)
-                            per_fwd_outputs[per_callable_fwd_idx] = outputs
+                            if is_training:
+                                per_fwd_outputs[per_callable_fwd_idx] = outputs
+                            else:
+                                del outputs
                         fwd_idx[m_chunk] += 1
                     elif ceil(c_id) == c_id:
                         # Backward pass for chunk -c_id.
@@ -665,10 +679,11 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx):
                                     _prefix_num_layers[m_chunk] * num_microbatches
                                 ) + (bwd_idx[m_chunk] * _num_layers_per_chunk[m_chunk] + l_no)
                                 func = callables[callable_idx]
-                                outputs = per_fwd_outputs[per_callable_bwd_idx]
+                                outputs = per_fwd_outputs.pop(per_callable_bwd_idx)
                                 _run_warmup_backward(
                                     per_callable_bwd_idx, func, outputs, warmup_iter, callable_idx
                                 )
+                                del outputs
                             bwd_idx[m_chunk] += 1
 
         if post_warmup_hook is not None:
@@ -729,6 +744,7 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx):
                     per_callable_static_outputs[per_callable_fwd_idx] = tuple(flatten_outputs)
                     per_callable_output_unflatten_spec[per_callable_fwd_idx] = spec
                     graph_callables[per_callable_fwd_idx] = func
+                    del outputs, flatten_outputs
                 fwd_idx[m_chunk] += 1
             else:
                 # Capture backward graph for model chunk c_id, microbatch bwd_idx[-c_id-1]
@@ -917,6 +933,11 @@ def _run_warmup_backward(func_idx, func, outputs, warmup_iter, callable_idx):
                                     per_callable_static_grad_inputs[idx]
                                 )
                             previous_chunk_last_callable_bwd_idx = per_callable_bwd_idx
+
+                        # The per-callable containers now own all tensors that must survive
+                        # capture. Drop local strong references so weak-refed graph buffers can
+                        # be returned to the shared CUDA graph pool before the next capture.
+                        del static_outputs, static_grad_inputs, grad_inputs
                 if ceil(c_id) == c_id:
                     bwd_idx[m_chunk] += 1
     else:
@@ -1028,12 +1049,22 @@ def make_graphed_autograd_function(
         static_grad_inputs,
         returned_param_grad_clone_slots,
     ):
+        is_reset = False
+
+        def ensure_not_reset():
+            """Reject replay after this callable's graph state has been released."""
+            if is_reset:
+                raise RuntimeError(
+                    "This graphed callable has been reset and can no longer be used."
+                )
+
         class Graphed(torch.autograd.Function):
             """Autograd function for graph replay."""
 
             @staticmethod
             def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *inputs):
                 # pylint: disable=missing-function-docstring
+                ensure_not_reset()
 
                 # Set flag for whether to update FP8 weight updates
                 ctx.is_first_module = FP8GlobalStateManager.is_first_fp8_module()
@@ -1071,6 +1102,7 @@ def forward(ctx, skip_fp8_weight_update, cuda_graph_stream, cuda_graph_event, *i
             @torch.autograd.function.once_differentiable
             def backward(ctx, *grads):
                 # pylint: disable=missing-function-docstring
+                ensure_not_reset()
 
                 # Replay backward graph
                 if len(grads) != len(static_grad_outputs):
@@ -1119,6 +1151,7 @@ def backward(ctx, *grads):
                 return (None, None, None) + tuple(grad_inputs)
 
         def functionalized(*user_args, **user_kwargs):
+            ensure_not_reset()
 
             # Decide whether to update FP8 weights
             skip_fp8_weight_update = None
@@ -1170,16 +1203,43 @@ def functionalized(*user_args, **user_kwargs):
             )
             return _tree_unflatten(out, output_unflatten_spec)
 
-        return functionalized
-
-    def make_graphed_attribute_functions(graph_idx):
-        # Get te modules for current graph
+        def release_static_state():
+            """Release per-callable state captured by replay closures."""
+            nonlocal fwd_graph, bwd_graph, is_reset
+            nonlocal module_params
+            nonlocal static_input_surface, static_outputs
+            nonlocal static_grad_outputs, static_grad_inputs
+
+            is_reset = True
+
+            # Drop the per-callable references that can own graph-pool storage.
+            fwd_graph = None
+            bwd_graph = None
+            module_params = ()
+            static_input_surface = ()
+            static_outputs = ()
+            static_grad_outputs = ()
+            static_grad_inputs = ()
+
+        helpers = _GraphedCallableHelpers(
+            ensure_not_reset=ensure_not_reset,
+            release_static_state=release_static_state,
+        )
+        return functionalized, helpers
+
+    def make_graphed_attribute_functions(graph_idx, helpers):
+        # Snapshot per-callable state so returned closures do not retain the outer lists.
+        fwd_graph = fwd_graphs[graph_idx]
+        bwd_graph = bwd_graphs[graph_idx]
+        bwd_dw_graph = bwd_dw_graphs[graph_idx]
+        need_bwd_dw = need_bwd_dw_graph.get(graph_idx, False)
         te_modules = visited_te_modules.get(graph_idx, set())
 
         # Attach backward_dw as an attribute to the graphed callable.
         def backward_dw():
-            if need_bwd_dw_graph.get(graph_idx, False):
-                bwd_dw_graphs[graph_idx].replay()
+            helpers.ensure_not_reset()
+            if need_bwd_dw:
+                bwd_dw_graph.replay()
 
                 # Trigger the grad accumulation hook for wgrad graphs.
                 for module in te_modules:
@@ -1191,16 +1251,24 @@ def backward_dw():
 
         # Attach reset as an attribute to the graphed callable.
         def reset():
-            fwd_graphs[graph_idx].reset()
-            bwd_graphs[graph_idx].reset()
-            bwd_dw_graphs[graph_idx].reset()
+            nonlocal fwd_graph, bwd_graph, bwd_dw_graph, te_modules
+
+            for graph in (fwd_graph, bwd_graph, bwd_dw_graph):
+                if graph is not None:
+                    graph.reset()
+
+            fwd_graph = None
+            bwd_graph = None
+            bwd_dw_graph = None
+            te_modules = ()
+            helpers.release_static_state()
 
         return backward_dw, reset
 
     # Put together the final graphed callables
     ret = []
     for i in range(len(sample_args)):
-        graphed = make_graphed_autograd_function(
+        graphed, helpers = make_graphed_autograd_function(
             fwd_graphs[i],
             bwd_graphs[i],
             per_callable_module_params[i],
@@ -1218,8 +1286,17 @@ def reset():
         te_modules = visited_te_modules.get(i, set())
         if isinstance(func, torch.nn.Module):
 
-            def make_graphed_forward(func, graph_training_state, graphed, orig_fwd, te_modules):
+            def make_graphed_forward(
+                func,
+                graph_training_state,
+                graphed,
+                orig_fwd,
+                te_modules,
+                helpers,
+            ):
                 def new_fwd(*user_args, **user_kwargs):
+                    helpers.ensure_not_reset()
+
                     # If the module's training-or-eval state matches what we graphed,
                     # run the graph, otherwise run the original forward method
                     if func.training == graph_training_state:
@@ -1264,7 +1341,14 @@ def new_fwd(*user_args, **user_kwargs):
 
                 return new_fwd
 
-            forward = make_graphed_forward(func, func.training, graphed, func.forward, te_modules)
+            forward = make_graphed_forward(
+                func,
+                func.training,
+                graphed,
+                func.forward,
+                te_modules,
+                helpers,
+            )
             if _order is None:
                 func.forward = forward
                 ret.append(func)
@@ -1273,7 +1357,10 @@ def new_fwd(*user_args, **user_kwargs):
         else:
             ret.append(graphed)
 
-        backward_dw_func, reset_func = make_graphed_attribute_functions(i)
+        backward_dw_func, reset_func = make_graphed_attribute_functions(
+            i,
+            helpers,
+        )
         setattr(ret[-1], "backward_dw", backward_dw_func)
         setattr(ret[-1], "reset", reset_func)
 

From 846379d3c4e8dee0d12fe8066c8c31dce3ef3ee2 Mon Sep 17 00:00:00 2001
From: Ravi Ghadia <40660742+ghadiaravi13@users.noreply.github.com>
Date: Fri, 4 Sep 2026 17:52:54 -0500
Subject: [PATCH 12/27] Relax runtime checks for activation recompute into
 Warnings (#3436)

* Remove redundant runtime checks for activation recompute in MLP from _ScaledUnary class in activation.py

Signed-off-by: Ravi Ghadia <rghadia@nvidia.com>

* Add warning for activation recompute in MLP outside fused path in _ScaledUnary class

Signed-off-by: Ravi Ghadia <rghadia@nvidia.com>

* Add test for Scaled SReLU activation recompute warning outside fused MLP path

Signed-off-by: Ravi Ghadia <rghadia@nvidia.com>

* Enhance activation recompute warning in ScaledSReLU: Update test to verify multiple warnings during backward passes and refactor warning mechanism to ensure it survives Dynamo tracing.

Signed-off-by: Ravi Ghadia <rghadia@nvidia.com>

* Remove redundant logic to only warn once

Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Remove unhelpful test

Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------

Signed-off-by: Ravi Ghadia <rghadia@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>
---
 .../pytorch/ops/basic/activation.py           | 28 ++++++++++---------
 1 file changed, 15 insertions(+), 13 deletions(-)

diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py
index 5c33c08b44..26d5261b13 100644
--- a/transformer_engine/pytorch/ops/basic/activation.py
+++ b/transformer_engine/pytorch/ops/basic/activation.py
@@ -16,7 +16,7 @@
 from ...constants import DType
 from ...cpu_offload import is_cpu_offload_enabled, mark_activation_offload
 from ...tensor.float8_tensor import Float8CurrentScalingQuantizer, Quantizer
-from ...utils import clear_tensor_data
+from ...utils import _compile_safe_warn, clear_tensor_data
 from ..op import BasicOperation, OperationContext
 from .._common import maybe_dequantize
 
@@ -404,14 +404,14 @@ def fuser_forward(
         next_op_input_quantizer: Optional[Quantizer],  # pylint: disable=unused-argument
         basic_op_kwargs: list[dict[str, Any]],  # pylint: disable=unused-argument
     ) -> tuple[torch.Tensor, Sequence[Sequence[torch.Tensor]]]:
+        extra_input = basic_op_extra_inputs[0][0]
+
         if self.activation_recompute_in_mlp:
-            raise RuntimeError(
-                f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the "
-                "fused grouped MLP path."
+            _compile_safe_warn(
+                f"{self.__class__.__name__}(activation_recompute_in_mlp=True) is only supported "
+                "in the fused grouped MLP path."
             )
 
-        extra_input = basic_op_extra_inputs[0][0]
-
         if torch.is_autocast_enabled():
             dtype = torch.get_autocast_dtype("cuda")
         elif isinstance(input_, torch.Tensor):
@@ -447,18 +447,18 @@ def fuser_backward(
     ]:
         del basic_op_grad_extra_outputs
 
-        if self.activation_recompute_in_mlp:
-            raise RuntimeError(
-                f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the "
-                "fused grouped MLP path."
-            )
-
         ctx = basic_op_ctxs[0]
         x, scales = ctx.saved_tensors
         x = maybe_dequantize(x.contiguous(), ctx.dtype)
         scales = maybe_dequantize(scales, ctx.dtype)
         grad_output = maybe_dequantize(grad_output.contiguous(), ctx.dtype)
 
+        if self.activation_recompute_in_mlp:
+            _compile_safe_warn(
+                f"{self.__class__.__name__}(activation_recompute_in_mlp=True) is only supported "
+                "in the fused grouped MLP path."
+            )
+
         grad_input, grad_extra_input = self._scaled_unary_backward(
             grad_output,
             x,
@@ -483,7 +483,9 @@ class ScaledSReLU(_ScaledUnary):
     ----------
     activation_recompute_in_mlp : bool, default = ``False``
         Enable fused grouped MLP kernels to recompute activation outputs
-        during backward when supported instead of saving them.
+        during backward when supported instead of saving them. Outside the
+        fused grouped MLP path this option has no effect and a warning is
+        emitted.
     """
 
     def _scaled_unary_forward(

From a30aee5b5536d9e0edf8645eda6eabb66238e108 Mon Sep 17 00:00:00 2001
From: Andrew White <andrewh@cdw.com>
Date: Fri, 4 Sep 2026 19:06:45 -0500
Subject: [PATCH 13/27] fix: Unreachable backend check after earlier backend
 skip (#3368)

* fix: remove unreachable fused-attn backend skip and add regression test

Root cause: FusedAttnRunner._check_configs skipped with 'Unsupported
inputs combination or device compute capability.' unless the backend was
NVTE_F16_arbitrary_seqlen. The later elif re-testing
self.backend != NVTE_F16_arbitrary_seqlen could therefore never be
reached, so its skip message ('B1SS, BHSS and 11SS bias shapes are only
supported for the F16_arbitrary_seqlen backend') was dead code.

Fix: drop the dead elif arm. The padding-mask skip in the sibling if arm
is retained. A regression test locks the remaining behavior: a
non-1HSS post-scale-bias config (BiasShape._B1SS) that passes
_check_configs selects NVTE_F16_arbitrary_seqlen, proving the removal is
behaviorally invisible and the earlier guard is the sole gate.

Testing: not run locally - the JAX test stack (jax,
transformer_engine_jax) is not installed on this machine. The suite runs
in NVIDIA TransformerEngine CI. Contribution: tests/jax tests target
real GPU/cuDNN fused-attention kernels and skip otherwise; the new test
will follow that path via CI.

Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>

* Remove unnecessary test

Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com>

---------

Signed-off-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com>
Co-authored-by: andrewwhitecdw <andrewwhitecdw@users.noreply.github.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
---
 tests/jax/test_fused_attn.py | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py
index b6fc8f7794..1ca4121a4f 100644
--- a/tests/jax/test_fused_attn.py
+++ b/tests/jax/test_fused_attn.py
@@ -612,11 +612,6 @@ def _check_configs(self):
                 pytest.skip(
                     "B1SS, BHSS and 11SS bias shapes are only supported for non-padding mask"
                 )
-            elif self.backend != NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen:
-                pytest.skip(
-                    "B1SS, BHSS and 11SS bias shapes are only supported for "
-                    "the F16_arbitrary_seqlen backend."
-                )
 
     def _setup_inputs(self):
         self._check_configs()

From 5f6105b900778068144f1b24e25a2e84066983cb Mon Sep 17 00:00:00 2001
From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>
Date: Sat, 5 Sep 2026 15:13:46 +0200
Subject: [PATCH 14/27] [Common] Preserve shared-memory pointer provenance in
 TMA kernels (#3482)

* Use the shmem alignment operator consistently

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

---------

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 transformer_engine/common/cast/core/grouped_tma.cuh    |  6 ------
 .../fp8_blockwise/group_quantize_fp8_blockwise.cuh     | 10 +++++-----
 .../common/cast/mxfp8/group_quantize_mxfp8.cuh         |  4 ++--
 .../common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh    |  4 ++--
 .../specialized/quantize_transpose_nvfp4_tuned_1D.cuh  |  4 ++--
 5 files changed, 11 insertions(+), 17 deletions(-)

diff --git a/transformer_engine/common/cast/core/grouped_tma.cuh b/transformer_engine/common/cast/core/grouped_tma.cuh
index 61218d654a..8603fd1fd2 100644
--- a/transformer_engine/common/cast/core/grouped_tma.cuh
+++ b/transformer_engine/common/cast/core/grouped_tma.cuh
@@ -53,12 +53,6 @@ inline bool dimensions_supported_by_TMA(const Tensor *const t) {
   return cols % alignment_requirement == 0;
 }
 
-__device__ __forceinline__ unsigned char *align_smem_ptr_per_TMA_requirements(unsigned char *p) {
-  size_t addr = reinterpret_cast<size_t>(p);
-  addr = (addr + TMA_SHMEM_ALIGNMENT - 1) & ~(TMA_SHMEM_ALIGNMENT - 1);
-  return reinterpret_cast<unsigned char *>(addr);
-}
-
 // Copies the base tensor map to shmem, modifies the copy, stores the modified tensor map at index
 __device__ __forceinline__ void modify_base_tensor_map(const CUtensorMap base_tensor_map,
                                                        CUtensorMap *global_tensor_map,
diff --git a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh
index 203f569471..31feaf833d 100644
--- a/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh
+++ b/transformer_engine/common/cast/fp8_blockwise/group_quantize_fp8_blockwise.cuh
@@ -314,9 +314,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock, 4) group_block_scaled_2d_tma
 
   // Dynamic smem holds the IType input tile (TMA dest, must be 128 B aligned).
   // warp_amaxes and tma_mbar are static smem.
-  extern __shared__ unsigned char smem_raw_2d_tma[];
-  IType(*smem_in)[kTileDim] = reinterpret_cast<IType(*)[kTileDim]>(
-      common::align_smem_ptr_per_TMA_requirements(smem_raw_2d_tma));
+  extern __shared__ char smem_raw_2d_tma[];
+  IType(*smem_in)[kTileDim] =
+      reinterpret_cast<IType(*)[kTileDim]>(align_up(smem_raw_2d_tma, TMA_SHMEM_ALIGNMENT));
 
   __shared__ CType warp_amaxes[kNumWarps];
   __shared__ size_t warp_offset_partials[kNumWarps];
@@ -603,8 +603,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) group_block_scaled_1d_tma_ke
 
   // Dynamic smem: IType[kTileDim][kTileDim], 128 B aligned for TMA. Static smem
   // (smem_T when CW, tma_mbar) lives outside the dynamic region.
-  extern __shared__ unsigned char smem_raw_1d_tma[];
-  unsigned char* smem_base = common::align_smem_ptr_per_TMA_requirements(smem_raw_1d_tma);
+  extern __shared__ char smem_raw_1d_tma[];
+  char* smem_base = align_up(smem_raw_1d_tma, TMA_SHMEM_ALIGNMENT);
   IType(*smem)[kTileDim] = reinterpret_cast<IType(*)[kTileDim]>(smem_base);
 
   __shared__ uint64_t tma_mbar;
diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh
index 980a77db0a..b0383d95f3 100644
--- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh
@@ -704,8 +704,8 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_
   constexpr size_t out_mem_rowwise = (ROWWISE_SCALING ? buff_size_aligned_out : 0);
 
   // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned
-  extern __shared__ unsigned char dynamic_shmem[];
-  unsigned char *dshmem = align_smem_ptr_per_TMA_requirements(dynamic_shmem);
+  extern __shared__ char dynamic_shmem[];
+  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 *sIn_ptr = reinterpret_cast<IType *>(dshmem);
diff --git a/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh
index 24f84fa359..878fc93107 100644
--- a/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/group_scaled_swiglu_mxfp8.cuh
@@ -209,8 +209,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_scaled_swiglu_mxfp8_k
         DIVUP_TO_MULTIPLE(CHUNK_DIM_Y * sizeof(float), TMA_SHMEM_ALIGNMENT);
 
     // shmem layout: [act input][gate input][colwise output][prob]
-    extern __shared__ unsigned char dynamic_shmem[];
-    unsigned char *dshmem = align_smem_ptr_per_TMA_requirements(dynamic_shmem);
+    extern __shared__ char dynamic_shmem[];
+    char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
     IType *sInAct_ptr = reinterpret_cast<IType *>(dshmem);
     IType *sInGate_ptr = reinterpret_cast<IType *>(dshmem + buff_size_aligned_in);
diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh
index cdd0d4916a..58af1f7938 100644
--- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh
+++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh
@@ -410,8 +410,8 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D
       TunableConfig::CHUNK_DIM_Y * SCALES_PER_CHUNK_X * sizeof(nvfp4_scale_t), TMA_SHMEM_ALIGNMENT);
 
   // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned
-  extern __shared__ unsigned char dynamic_shmem[];
-  unsigned char *dshmem = common::align_smem_ptr_per_TMA_requirements(dynamic_shmem);
+  extern __shared__ char dynamic_shmem[];
+  char *dshmem = align_up(dynamic_shmem, TMA_SHMEM_ALIGNMENT);
 
   IType *sIn_ptr = reinterpret_cast<IType *>(dshmem);
   fp4e2m1x2 *sOut_ptr = reinterpret_cast<fp4e2m1x2 *>(dshmem + in_mem);

From dc8909da2e2c9cebba3f2583ea2a47b531ff1bcb Mon Sep 17 00:00:00 2001
From: Evgeny Tsykunov <etsykunov@nvidia.com>
Date: Tue, 8 Sep 2026 12:24:29 +0200
Subject: [PATCH 15/27] Update Linear docstring referring torch.Linear (#3491)

* Warn Linear argument documentation

Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Update docstring, remove warning

Signed-off-by: Evgeny <etsykunov@nvidia.com>

---------

Signed-off-by: Evgeny <etsykunov@nvidia.com>
---
 docs/api/pytorch.rst                        | 2 +-
 transformer_engine/pytorch/module/linear.py | 8 +++++++-
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst
index 54981c9086..739a9864b9 100644
--- a/docs/api/pytorch.rst
+++ b/docs/api/pytorch.rst
@@ -6,7 +6,7 @@
 PyTorch
 =======
 
-.. autoapiclass:: transformer_engine.pytorch.Linear(in_features, out_features, bias=True, **kwargs)
+.. autoapiclass:: transformer_engine.pytorch.Linear(in_features, out_features, **kwargs)
   :members: forward, set_tensor_parallel_group
 
 .. autoapiclass:: transformer_engine.pytorch.GroupedLinear(in_features, out_features, bias=True, **kwargs)
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 55fc69ef7f..94de69e975 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -1901,7 +1901,13 @@ def _linear_eager(
 class Linear(TransformerEngineBaseModule):
     """Applies a linear transformation to the incoming data :math:`y = xA^T + b`
 
-    On NVIDIA GPUs it is a drop-in replacement for ``torch.nn.Linear``.
+    On NVIDIA GPUs, this module implements the same linear transformation as
+    ``torch.nn.Linear``.
+
+    .. note::
+
+       Its constructor signature differs from ``torch.nn.Linear``. Pass optional
+       arguments, including ``bias``, by keyword.
 
     Parameters
     ----------

From f2bec2314c10754d4ea5a1fd0cb0a8fa1442bf6d Mon Sep 17 00:00:00 2001
From: Evgeny Tsykunov <etsykunov@nvidia.com>
Date: Tue, 8 Sep 2026 12:25:51 +0200
Subject: [PATCH 16/27] [PyTorch] Document fine-grained quantization recipes
 (#3336)

* Fine-grained recipe docs

Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Update docs/features/low_precision_training/fine_grained_quantization/pytorch_fine_grained_quantization_example.py

Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Signed-off-by: Evgeny Tsykunov <e.tsykunov@gmail.com>

* resolve comments

Signed-off-by: Evgeny <etsykunov@nvidia.com>

* Rework heterogeneous quantization docs into mixed-format quantization; add per-recipe Quantizer sections

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Align mixed-format quantization diagrams with shared diagram-colors.css and dark mode

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Rename mixed-format quantization docs to fine-grained quantization recipes

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Simplify quantizer factory paragraph

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

* Generalize the fallback-path description

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

---------

Signed-off-by: Evgeny <etsykunov@nvidia.com>
Signed-off-by: Evgeny Tsykunov <e.tsykunov@gmail.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 docs/_static/css/diagram-colors.css           |  15 +
 docs/api/pytorch.rst                          |  18 +
 ...torch_fine_grained_quantization_example.py | 130 +++++++
 .../fine_grained_quantization.rst             | 364 ++++++++++++++++++
 .../img/fine_grained_assignments.svg          |  95 +++++
 .../img/fine_grained_linear_mapping.svg       |  59 +++
 .../img/hybrid_columnwise_source.svg          |  74 ++++
 .../img/hybrid_quantizer.svg                  |  58 +++
 .../fp8_blockwise_scaling.rst                 |  31 ++
 .../fp8_current_scaling.rst                   |  53 ++-
 .../fp8_delayed_scaling.rst                   |  67 +++-
 .../features/low_precision_training/index.rst |   1 +
 .../introduction/introduction.rst             |  68 ++++
 .../low_precision_training/mxfp8/mxfp8.rst    |  53 ++-
 .../low_precision_training/nvfp4/nvfp4.rst    |  54 +++
 .../performance_considerations.rst            |  56 ++-
 transformer_engine/common/recipe/__init__.py  |  12 +-
 .../pytorch/tensor/hybrid_tensor.py           |   4 +
 .../pytorch/tensor/identity_tensor.py         |  19 +
 19 files changed, 1224 insertions(+), 7 deletions(-)
 create mode 100644 docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py
 create mode 100644 docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst
 create mode 100644 docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg
 create mode 100644 docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg
 create mode 100644 docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg
 create mode 100644 docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg

diff --git a/docs/_static/css/diagram-colors.css b/docs/_static/css/diagram-colors.css
index f5dc7da4dd..9ee5827bd1 100644
--- a/docs/_static/css/diagram-colors.css
+++ b/docs/_static/css/diagram-colors.css
@@ -279,3 +279,18 @@ html[data-theme="dark"] .subtitle,
 html[data-theme="dark"] .memory-label { fill: #e0e0e0; }
 html[data-theme="dark"] .connector { stroke: #bdbdbd; }
 
+
+/* fine_grained_quantization diagrams */
+html[data-theme="dark"] .fmt-mxfp8 { fill: #10375c; stroke: #64b5f6; }
+html[data-theme="dark"] .fmt-nvfp4 { fill: #5c3a10; stroke: #ffb74d; }
+html[data-theme="dark"] .fmt-bf16 { fill: #1e4620; stroke: #81c784; }
+html[data-theme="dark"] .fmt-mxfp8-text { fill: #90caf9; }
+html[data-theme="dark"] .fmt-nvfp4-text { fill: #ffcc80; }
+html[data-theme="dark"] .fmt-bf16-text { fill: #a5d6a7; }
+html[data-theme="dark"] .source { fill: #3a2f5c; stroke: #b39ddb; }
+html[data-theme="dark"] .quantizer { fill: #1e4620; stroke: #81c784; }
+html[data-theme="dark"] .representation { fill: #10375c; stroke: #64b5f6; }
+html[data-theme="dark"] .dequantize { fill: #5c3a10; stroke: #ffb74d; }
+html[data-theme="dark"] .rowlabel,
+html[data-theme="dark"] .legend,
+html[data-theme="dark"] .op { fill: #e0e0e0; }
diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst
index 739a9864b9..a6afa2d0cc 100644
--- a/docs/api/pytorch.rst
+++ b/docs/api/pytorch.rst
@@ -112,6 +112,12 @@ Communication-computation overlap
   :members: FP8, NONE
 
 
+Fine-grained quantization recipes
+---------------------------------
+
+.. autoapiclass:: transformer_engine.pytorch.QuantizerRole(module_type="", tensor_type="", name="")
+
+
 Quantized tensors
 -----------------
 
@@ -129,6 +135,10 @@ Quantized tensors
 
 .. autoapiclass:: transformer_engine.pytorch.NVFP4TensorStorage(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, amax_rowwise, amax_columnwise, fp4_dtype, quantizer)
 
+.. autoapiclass:: transformer_engine.pytorch.HybridQuantizedTensorStorage(*, rowwise_storage, columnwise_storage, quantizer, fake_dtype=None)
+
+.. autoapiclass:: transformer_engine.pytorch.IdentityTensorStorage(*, hp_data, fake_dtype=None, quantizer=None)
+
 .. autoapiclass:: transformer_engine.pytorch.Float8Tensor(shape, dtype, data, fp8_scale_inv, fp8_dtype, requires_grad=False, data_transpose=None, quantizer=None)
 
 .. autoapiclass:: transformer_engine.pytorch.MXFP8Tensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, fp8_dtype, quantizer)
@@ -137,6 +147,10 @@ Quantized tensors
 
 .. autoapiclass:: transformer_engine.pytorch.NVFP4Tensor(rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv, amax_rowwise, amax_columnwise, fp4_dtype, quantizer)
 
+.. autoapiclass:: transformer_engine.pytorch.HybridQuantizedTensor(shape, dtype, *, rowwise_storage, columnwise_storage, quantizer, requires_grad=False, device=None)
+
+.. autoapiclass:: transformer_engine.pytorch.IdentityTensor(shape, dtype, *, hp_data, quantizer=None, requires_grad=False, device=None)
+
 Quantizers
 ----------
 
@@ -153,6 +167,10 @@ Quantizers
 
 .. autoapiclass:: transformer_engine.pytorch.NVFP4Quantizer(fp4_dtype, *, rowwise=True, columnwise=True, **kwargs)
 
+.. autoapiclass:: transformer_engine.pytorch.HybridQuantizer(*, rowwise_quantizer, columnwise_quantizer, columnwise_source="original")
+
+.. autoapiclass:: transformer_engine.pytorch.IdentityQuantizer(*, dtype=None, rowwise=True, columnwise=True)
+
 Tensor saving and restoring functions
 -------------------------------------
 
diff --git a/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py b/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py
new file mode 100644
index 0000000000..b511fc41dc
--- /dev/null
+++ b/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py
@@ -0,0 +1,130 @@
+#!/usr/bin/env python3
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Runnable fine-grained quantization recipe example.
+
+The factory assigns one precision to each ``demo.fc1`` Linear GEMM:
+
+* fprop: ``weight.row(MXFP8) x input.row(MXFP8)``
+* dgrad: ``weight.col(NVFP4) x grad_output.row(NVFP4)``
+* wgrad: ``input.col(original BF16) x grad_output.col(original BF16)``
+
+``demo.fc2`` runs every GEMM in high precision. ``demo.output`` is not
+special-cased and therefore exercises the MXFP8 base-factory fallback.
+
+Run from the Transformer Engine repository root::
+
+    python docs/examples/fine_grained_quantization/\
+        pytorch_fine_grained_quantization_example.py
+"""
+
+from __future__ import annotations
+
+import torch
+import transformer_engine.pytorch as te
+
+
+def require_supported_hardware() -> None:
+    """Fail early with TE's reason when either required format is unavailable."""
+
+    if not torch.cuda.is_available():
+        raise SystemExit("This example requires a CUDA-capable NVIDIA GPU.")
+
+    failures = []
+    for name, check in (
+        ("MXFP8", te.is_mxfp8_available),
+        ("NVFP4", te.is_nvfp4_available),
+    ):
+        available, reason = check(return_reason=True)
+        if not available:
+            failures.append(f"{name}: {reason}")
+    if failures:
+        raise SystemExit("Required formats are unavailable: " + "; ".join(failures))
+
+
+require_supported_hardware()
+
+# START_FINE_GRAINED_QUANTIZATION_EXAMPLE
+
+from typing import Optional
+
+import torch
+
+import transformer_engine.pytorch as te
+from transformer_engine.common.recipe import CustomRecipe
+from transformer_engine.pytorch.custom_recipes.quantizer_factories import (
+    mxfp8_factory,
+    nvfp4_factory,
+)
+
+
+THREE_FORMAT_MODULE = "demo.fc1"
+HIGH_PRECISION_MODULE = "demo.fc2"
+BASE_FACTORY = mxfp8_factory
+
+
+def quantizer_factory(role: Optional[te.QuantizerRole]):
+    """Return a fresh quantizer for every role, including ``None``.
+
+    ``BASE_FACTORY`` makes the factory total: unknown roles, future role values,
+    and untargeted modules all retain valid MXFP8 behavior.
+    """
+
+    if role is not None and role.name == THREE_FORMAT_MODULE:
+        # Constructing fresh child quantizers for every call is recommended.
+        if role.tensor_type == "input":
+            # Wgrad retains the original BF16 input.
+            return te.HybridQuantizer(
+                rowwise_quantizer=mxfp8_factory(role),
+                columnwise_quantizer=te.IdentityQuantizer(),
+                columnwise_source="original",
+            )
+        if role.tensor_type == "weight":
+            # Dgrad uses NVFP4 quantized from the dequantized MXFP8 fprop weight.
+            return te.HybridQuantizer(
+                rowwise_quantizer=mxfp8_factory(role),
+                columnwise_quantizer=nvfp4_factory(role),
+                columnwise_source="rowwise_dequantized",
+            )
+        if role.tensor_type == "grad_output":
+            # Dgrad uses NVFP4 while wgrad retains the original BF16 gradient.
+            return te.HybridQuantizer(
+                rowwise_quantizer=nvfp4_factory(role),
+                columnwise_quantizer=te.IdentityQuantizer(),
+                columnwise_source="original",
+            )
+
+    if role is not None and role.name == HIGH_PRECISION_MODULE:
+        return te.IdentityQuantizer()
+
+    return BASE_FACTORY(role)
+
+
+linear_options = {"bias": False, "params_dtype": torch.bfloat16, "device": "cuda"}
+model = torch.nn.Sequential(
+    te.Linear(128, 256, name=THREE_FORMAT_MODULE, **linear_options),
+    torch.nn.GELU(),
+    te.Linear(256, 256, name=HIGH_PRECISION_MODULE, **linear_options),
+    torch.nn.GELU(),
+    te.Linear(256, 128, name="demo.output", **linear_options),
+)
+inputs = torch.randn(64, 128, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+recipe = CustomRecipe(qfactory=quantizer_factory)
+
+with te.autocast(enabled=True, recipe=recipe):
+    outputs = model(inputs)
+
+loss = outputs.float().square().mean()
+loss.backward()
+
+# END_FINE_GRAINED_QUANTIZATION_EXAMPLE
+
+gradients = [inputs.grad, *(parameter.grad for parameter in model.parameters())]
+assert all(gradient is not None for gradient in gradients)
+assert all(torch.isfinite(gradient).all() for gradient in gradients)
+
+print(f"GPU: {torch.cuda.get_device_name()}")
+print(f"TE Linear names: {[model[index].name for index in (0, 2, 4)]}")
+print(f"loss: {loss.item():.6f}; forward and backward completed")
diff --git a/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst
new file mode 100644
index 0000000000..301b655b3f
--- /dev/null
+++ b/docs/features/low_precision_training/fine_grained_quantization/fine_grained_quantization.rst
@@ -0,0 +1,364 @@
+..
+    Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+    See LICENSE for license information.
+
+.. _fine-grained-quantization-recipes:
+.. _heterogeneous-quantization-recipes:
+
+Fine-grained quantization recipes
+=================================
+
+Standard TE recipes quantize the whole model the same way. That is often too
+coarse: one sensitive layer may need BF16 while the rest runs in MXFP8, or a
+gradient GEMM may tolerate a cheaper format than the forward pass. Fine-grained
+recipes lift this restriction: you write a small factory function that picks a
+quantizer for each slot TE asks about, and pass it via
+:class:`~transformer_engine.common.recipe.CustomRecipe` to the usual
+:class:`~transformer_engine.pytorch.autocast`.
+"Fine-grained" refers to the granularity of that choice (per module, tensor
+role, and GEMM direction), not to the block size of the scaling factors.
+
+.. warning::
+
+   Fine-grained recipes are currently available only in the PyTorch API of
+   TE.
+
+.. warning::
+
+   Fine-grained recipes and their construction APIs are experimental: API,
+   validation, and kernel coverage may change without notice. This guide does
+   not define a supported recipe or an expected accuracy/performance ordering.
+
+
+Example: mixing MXFP8, NVFP4, and BF16
+--------------------------------------
+
+The `runnable example <https://github.com/NVIDIA/TransformerEngine/blob/main/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py>`__
+makes the following assignments:
+
+.. raw:: html
+   :file: img/fine_grained_assignments.svg
+
+*Figure 1. Precision assignments per module and GEMM used throughout this
+guide.*
+
+A minimal factory implementing these assignments, plugged into the standard
+TE autocast path:
+
+.. tabs::
+
+   .. tab:: PyTorch
+
+      .. code-block:: python
+
+         import transformer_engine.pytorch as te
+         from transformer_engine.common.recipe import CustomRecipe
+         from transformer_engine.pytorch.custom_recipes.quantizer_factories import (
+             mxfp8_factory,
+             nvfp4_factory,
+         )
+
+
+         def quantizer_factory(role):
+             if role is not None and role.name == "demo.fc1":
+                 if role.tensor_type == "input":
+                     # wgrad keeps the original BF16 input
+                     return te.HybridQuantizer(
+                         rowwise_quantizer=mxfp8_factory(role),
+                         columnwise_quantizer=te.IdentityQuantizer(),
+                         columnwise_source="original",
+                     )
+                 if role.tensor_type == "weight":
+                     # fprop in MXFP8, dgrad in NVFP4
+                     return te.HybridQuantizer(
+                         rowwise_quantizer=mxfp8_factory(role),
+                         columnwise_quantizer=nvfp4_factory(role),
+                         columnwise_source="rowwise_dequantized",
+                     )
+                 if role.tensor_type == "grad_output":
+                     # dgrad in NVFP4, wgrad keeps the original BF16 gradient
+                     return te.HybridQuantizer(
+                         rowwise_quantizer=nvfp4_factory(role),
+                         columnwise_quantizer=te.IdentityQuantizer(),
+                         columnwise_source="original",
+                     )
+             if role is not None and role.name == "demo.fc2":
+                 return te.IdentityQuantizer()  # whole module stays in BF16
+             return mxfp8_factory(role)  # every other TE module in MXFP8
+
+
+         recipe = CustomRecipe(qfactory=quantizer_factory)
+
+         with te.autocast(enabled=True, recipe=recipe):
+             output = model(inputs)
+
+The complete, runnable version is available
+`on GitHub <https://github.com/NVIDIA/TransformerEngine/blob/main/docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py>`__
+(requires Blackwell or later); run it from the repository root after
+installing TE:
+
+.. code-block:: bash
+
+   python docs/examples/fine_grained_quantization/pytorch_fine_grained_quantization_example.py
+
+CustomRecipe and quantizer factory
+----------------------------------
+
+:class:`~transformer_engine.common.recipe.CustomRecipe` is used like any
+other TE recipe (``DelayedScaling``, ``MXFP8BlockScaling``, ...), but carries
+no quantization logic of its own: TE asks your ``qfactory`` for a quantizer
+whenever a module needs one.
+
+Each TE module defines an ordered role list for the forward and backward
+quantizer slots it needs. When module recipe state is initialized or rebuilt,
+a ``CustomRecipe`` calls ``qfactory(role)`` once for every slot in that list.
+It does not call the factory on every unchanged forward.
+
+.. tabs::
+
+   .. tab:: PyTorch
+
+      .. code-block:: python
+
+         # QuantizerRole describes the slot being configured (fields below):
+         #
+         #   @dataclasses.dataclass(frozen=True)
+         #   class QuantizerRole:
+         #       module_type: str = ""
+         #       tensor_type: str = ""
+         #       name: str = ""
+
+
+         def quantizer_factory(role: Optional[te.QuantizerRole]):
+             # construct a fresh quantizer on every call
+             ...
+             # Boundary slots may pass role=None or a role with empty fields, so
+             # always end with a default that covers every remaining role.
+             return mxfp8_factory(role)
+
+
+         # The factory plugs into the standard TE autocast path:
+         recipe = CustomRecipe(qfactory=quantizer_factory)
+
+         with te.autocast(enabled=True, recipe=recipe):
+             output = model(inputs)
+
+      **Module type**
+
+      The kind of TE module that owns the slot, filled in by TE itself:
+
+      * ``"linear"`` — ``Linear``, ``LayerNormLinear``, ``fc1``/``fc2`` in
+        ``LayerNormMLP``, ``qkv``/``proj`` in ``MultiheadAttention``;
+      * ``"grouped_linear"`` — ``GroupedLinear``;
+      * ``"dpa"`` — ``DotProductAttention``.
+
+      **Tensor type**
+
+      Which tensor of that module the quantizer will process, also filled in
+      by TE. For ``"linear"`` and ``"grouped_linear"``:
+
+      * ``"input"`` — the activation (fprop, wgrad);
+      * ``"weight"`` — (fprop, dgrad);
+      * ``"grad_output"`` — the incoming gradient (dgrad, wgrad).
+
+      For ``"dpa"``:
+
+      * ``"qkv"`` — the query/key/value tensor;
+      * ``"s"`` — the softmax output;
+      * ``"do"`` — the output gradient;
+      * ``"dp"`` — the gradient of ``"s"``.
+
+      **Name**
+
+      The identity of one concrete module instance, supplied by the caller:
+      ``te.Linear(..., name="decoder.39.fc2")``. Composite TE modules may
+      append suffixes such as ``.fc1``, ``.fc2``, and ``.proj``.
+
+      The role vocabulary is experimental and may grow between releases —
+      one more reason to end the factory with a total default. Treat the role
+      strings as selectors, not a fixed enumeration. Prefer a module-level
+      function for the factory itself, so that launchers and checkpointing
+      setups can import or pickle it.
+
+      TE provides factories for its native quantizers in
+      ``transformer_engine.pytorch.custom_recipes.quantizer_factories``
+      (``mxfp8_factory``, ``nvfp4_factory``, ...). They can be used as
+      defaults or to construct ``HybridQuantizer`` children. Additional
+      specialized recipes are available in
+      ``transformer_engine.pytorch.custom_recipes.quantizer_factory_zoo``.
+
+      The factory is not limited to TE-native quantizers: it may return your
+      own :class:`~transformer_engine.pytorch.Quantizer` subclass, and custom
+      quantizers can also serve as ``HybridQuantizer`` children. The GEMMs
+      still need to receive representations in formats they support.
+
+HybridQuantizer
+---------------
+
+During training, each tensor of a ``Linear`` or ``GroupedLinear`` layer feeds
+two different GEMMs: its rowwise representation feeds one, its columnwise
+representation the other (the exact operand layout is described in the
+:doc:`Introduction <../introduction/introduction>`). Since those two GEMMs may
+want different formats, the tensor needs a quantizer per direction:
+:class:`~transformer_engine.pytorch.HybridQuantizer` composes a rowwise and a
+columnwise quantizer, and its output,
+:class:`~transformer_engine.pytorch.HybridQuantizedTensor`, composes the
+corresponding representations.
+
+.. tabs::
+
+   .. tab:: PyTorch
+
+      The following is pseudocode illustrating the composition:
+
+      .. code-block:: text
+
+         quantizer = te.HybridQuantizer(
+             rowwise_quantizer=MXFP8Quantizer(fp8_dtype=DType.kFloat8E4M3),
+             columnwise_quantizer=NVFP4Quantizer(),
+             columnwise_source="original",  # or "rowwise_dequantized"
+         )
+
+         # Quantization yields a HybridQuantizedTensor whose rowwise
+         # representation is MXFP8 and columnwise representation is NVFP4;
+         # each GEMM consumes the representation it needs.
+         qtensor = quantizer(tensor)
+
+.. raw:: html
+   :file: img/hybrid_quantizer.svg
+
+*Figure 2. HybridQuantizer composes a rowwise and a columnwise quantizer; each
+representation of the result feeds a different GEMM.*
+
+Choosing the columnwise source
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+``columnwise_source`` is a separate numerical recipe choice that controls the
+source for the columnwise representation:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Value
+     - Columnwise source
+   * - ``"original"``
+     - The original high-precision tensor.
+   * - ``"rowwise_dequantized"``
+     - Dequantized rowwise representation.
+
+.. raw:: html
+   :file: img/hybrid_columnwise_source.svg
+
+*Figure 3. The columnwise representation can be derived from the original
+high-precision tensor or from the dequantized rowwise representation.*
+
+For forward inputs and weights, ``"rowwise_dequantized"`` derives the backward
+representation from the value consumed in the forward direction. This
+can improve forward/backward numerical consistency and may affect convergence.
+It does not recover information discarded by rowwise quantization.
+``"original"`` instead derives both representations from the original tensor.
+Choose the provenance as part of the numerical recipe.
+
+IdentityQuantizer
+-----------------
+
+:class:`~transformer_engine.pytorch.IdentityQuantizer` stores its input in the
+held compute dtype, typically BF16, FP16, or FP32. It can keep a complete slot
+in high precision or act as one child of a ``HybridQuantizer``:
+
+.. tabs::
+
+   .. tab:: PyTorch
+
+      .. code-block:: python
+
+         # whole slot in high precision (e.g. a module kept in BF16)
+         quantizer = te.IdentityQuantizer()
+
+         # one direction in high precision, the other quantized
+         quantizer = te.HybridQuantizer(
+             rowwise_quantizer=mxfp8_factory(role),
+             columnwise_quantizer=te.IdentityQuantizer(),
+             columnwise_source="rowwise_dequantized",
+         )
+
+Note that in the second example the columnwise direction is high precision but
+holds the value reconstructed from MXFP8, not the original input — see
+`Choosing the columnwise source`_ above.
+
+Example: one format per GEMM
+----------------------------
+
+A natural way to design a recipe is to pick one format for each GEMM. To
+translate that into quantizers, look at what each GEMM consumes — both of its
+operands must be in that GEMM's format:
+
+* **fprop** consumes ``input.rowwise`` and ``weight.rowwise``;
+* **dgrad** consumes ``grad_output.rowwise`` and ``weight.columnwise``;
+* **wgrad** consumes ``input.columnwise`` and ``grad_output.columnwise``.
+
+Reading the same table per tensor gives the ``HybridQuantizer`` for each role.
+For the example assignments (fprop in MXFP8, dgrad in NVFP4, wgrad in BF16):
+
+.. code-block:: text
+
+   input       = HybridQuantizer(rowwise=MXFP8,  columnwise=BF16)    # fprop | wgrad
+   weight      = HybridQuantizer(rowwise=MXFP8,  columnwise=NVFP4)   # fprop | dgrad
+   grad_output = HybridQuantizer(rowwise=NVFP4,  columnwise=BF16)    # dgrad | wgrad
+
+.. raw:: html
+   :file: img/fine_grained_linear_mapping.svg
+
+*Figure 4. Each GEMM consumes one representation of each of its two operand
+tensors; giving both operands the same format sets that GEMM's precision.*
+
+If two directions use the same quantizer configuration, a plain quantizer may
+replace the corresponding hybrid; one factory may return both plain and hybrid
+quantizers.
+The two operands of each GEMM still need a combination supported by that GEMM
+backend. TE may reject incompatible quantizer pairs or unsupported layouts.
+
+.. note::
+
+   On supported hardware these recipes run TE's regular quantized kernels: the
+   tensors are quantized on the GPU and the GEMMs execute in the selected
+   low-precision formats. TE does not fall back to fake quantization
+   (quantize-dequantize followed by a high-precision GEMM).
+
+
+Validating and optimizing a recipe
+----------------------------------
+
+The factory API can express more recipes than TE has kernels for, so any
+assignment lands in one of three buckets:
+
+* **Fast** — quantization hits TE's fused kernels and every GEMM runs a
+  native low-precision implementation.
+* **Correct but potentially unoptimized** — the recipe executes, but some
+  selected paths may not have fused or optimized implementations in the
+  current TE release. For example, ``HybridQuantizer`` may produce its rowwise
+  and columnwise representations in separate kernel launches; future releases
+  may fuse this work.
+* **Rejected** — the two operands of some GEMM end up in a combination of
+  formats or layouts that no GEMM backend supports, and TE raises an error.
+  This can happen with plain and hybrid quantizers alike.
+
+Before adopting a recipe for a real workload, check that:
+
+* it executes at all on the target GPU, software version, and modules;
+* it runs on optimized kernels rather than fallback paths;
+* accuracy and convergence hold on the target model and distributed setup;
+* throughput and memory actually improve on the target workload.
+
+The unoptimized paths are still useful: accuracy and convergence experiments can run
+on them before dedicated kernels exist, so the precision of each GEMM can be
+treated as an accuracy/performance trade-off to explore.
+
+API reference
+-------------
+
+See the :doc:`PyTorch API <../../../api/pytorch>` for ``QuantizerRole``,
+``HybridQuantizer``, ``IdentityQuantizer``, and their returned tensor types.
+See the :doc:`Common API <../../../api/common>` for ``CustomRecipe``.
diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg
new file mode 100644
index 0000000000..b9aa48333d
--- /dev/null
+++ b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_assignments.svg
@@ -0,0 +1,95 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 660" width="100%" style="max-width: 900px;" role="img" aria-labelledby="title description">
+  <title id="title">Precision assignment by tensor role and module</title>
+  <desc id="description">Each tensor role provides a rowwise and a columnwise representation, consumed by fprop, dgrad, and wgrad GEMMs. demo.fc1: input is MXFP8 rowwise and original BF16 columnwise; weight is MXFP8 rowwise and NVFP4 columnwise; grad_output is NVFP4 rowwise and original BF16 columnwise. demo.fc2 keeps every tensor in BF16. Other TE modules use MXFP8 everywhere.</desc>
+  <defs>
+    <style>
+      @import url("../_static/css/diagram-colors.css");
+      .label { font: 600 14px 'Segoe UI', Arial, sans-serif; fill: #212121; text-anchor: middle; }
+      .rowlabel { font: 13px 'Segoe UI', Arial, sans-serif; fill: #424242; text-anchor: end; }
+      .format { font: 600 14px 'Segoe UI', Arial, sans-serif; text-anchor: middle; }
+      .legend { font: 13px 'Segoe UI', Arial, sans-serif; fill: #424242; text-anchor: start; }
+      .fmt-mxfp8 { fill: #e3f2fd; stroke: #1976d2; stroke-width: 2; }
+      .fmt-nvfp4 { fill: #fff3e0; stroke: #f57c00; stroke-width: 2; }
+      .fmt-bf16 { fill: #e8f5e9; stroke: #388e3c; stroke-width: 2; }
+      .fmt-mxfp8-text { fill: #1565c0; }
+      .fmt-nvfp4-text { fill: #e65100; }
+      .fmt-bf16-text { fill: #2e7d32; }
+    </style>
+  </defs>
+
+  <text x="500" y="28" class="title">Precision assignment by tensor role and module</text>
+
+  <text x="335" y="70" class="label">demo.fc1</text>
+  <text x="585" y="70" class="label">demo.fc2</text>
+  <text x="835" y="70" class="label">Other TE modules</text>
+
+  <g aria-label="Input tensor">
+    <rect x="40" y="84" width="920" height="150" rx="7" class="panel"/>
+    <text x="66" y="159" class="section-title" transform="rotate(-90 66 159)">input</text>
+    <text x="200" y="150" class="rowlabel">rowwise (fprop)</text>
+    <rect x="220" y="122" width="230" height="44" rx="6" class="fmt-mxfp8"/>
+    <text x="335" y="150" class="format fmt-mxfp8-text">MXFP8</text>
+    <rect x="470" y="122" width="230" height="44" rx="6" class="fmt-bf16"/>
+    <text x="585" y="150" class="format fmt-bf16-text">BF16</text>
+    <rect x="720" y="122" width="230" height="44" rx="6" class="fmt-mxfp8"/>
+    <text x="835" y="150" class="format fmt-mxfp8-text">MXFP8</text>
+
+    <text x="200" y="206" class="rowlabel">columnwise (wgrad)</text>
+    <rect x="220" y="178" width="230" height="44" rx="6" class="fmt-bf16"/>
+    <text x="335" y="206" class="format fmt-bf16-text">BF16 (original)</text>
+    <rect x="470" y="178" width="230" height="44" rx="6" class="fmt-bf16"/>
+    <text x="585" y="206" class="format fmt-bf16-text">BF16</text>
+    <rect x="720" y="178" width="230" height="44" rx="6" class="fmt-mxfp8"/>
+    <text x="835" y="206" class="format fmt-mxfp8-text">MXFP8</text>
+  </g>
+
+  <g aria-label="Weight tensor">
+    <rect x="40" y="250" width="920" height="150" rx="7" class="panel"/>
+    <text x="66" y="325" class="section-title" transform="rotate(-90 66 325)">weight</text>
+    <text x="200" y="316" class="rowlabel">rowwise (fprop)</text>
+    <rect x="220" y="288" width="230" height="44" rx="6" class="fmt-mxfp8"/>
+    <text x="335" y="316" class="format fmt-mxfp8-text">MXFP8</text>
+    <rect x="470" y="288" width="230" height="44" rx="6" class="fmt-bf16"/>
+    <text x="585" y="316" class="format fmt-bf16-text">BF16</text>
+    <rect x="720" y="288" width="230" height="44" rx="6" class="fmt-mxfp8"/>
+    <text x="835" y="316" class="format fmt-mxfp8-text">MXFP8</text>
+
+    <text x="200" y="372" class="rowlabel">columnwise (dgrad)</text>
+    <rect x="220" y="344" width="230" height="44" rx="6" class="fmt-nvfp4"/>
+    <text x="335" y="372" class="format fmt-nvfp4-text">NVFP4</text>
+    <rect x="470" y="344" width="230" height="44" rx="6" class="fmt-bf16"/>
+    <text x="585" y="372" class="format fmt-bf16-text">BF16</text>
+    <rect x="720" y="344" width="230" height="44" rx="6" class="fmt-mxfp8"/>
+    <text x="835" y="372" class="format fmt-mxfp8-text">MXFP8</text>
+  </g>
+
+  <g aria-label="Grad output tensor">
+    <rect x="40" y="416" width="920" height="150" rx="7" class="panel"/>
+    <text x="66" y="491" class="section-title" transform="rotate(-90 66 491)">grad_output</text>
+    <text x="200" y="482" class="rowlabel">rowwise (dgrad)</text>
+    <rect x="220" y="454" width="230" height="44" rx="6" class="fmt-nvfp4"/>
+    <text x="335" y="482" class="format fmt-nvfp4-text">NVFP4</text>
+    <rect x="470" y="454" width="230" height="44" rx="6" class="fmt-bf16"/>
+    <text x="585" y="482" class="format fmt-bf16-text">BF16</text>
+    <rect x="720" y="454" width="230" height="44" rx="6" class="fmt-mxfp8"/>
+    <text x="835" y="482" class="format fmt-mxfp8-text">MXFP8</text>
+
+    <text x="200" y="538" class="rowlabel">columnwise (wgrad)</text>
+    <rect x="220" y="510" width="230" height="44" rx="6" class="fmt-bf16"/>
+    <text x="335" y="538" class="format fmt-bf16-text">BF16 (original)</text>
+    <rect x="470" y="510" width="230" height="44" rx="6" class="fmt-bf16"/>
+    <text x="585" y="538" class="format fmt-bf16-text">BF16</text>
+    <rect x="720" y="510" width="230" height="44" rx="6" class="fmt-mxfp8"/>
+    <text x="835" y="538" class="format fmt-mxfp8-text">MXFP8</text>
+  </g>
+
+  <g aria-label="Legend">
+    <rect x="220" y="600" width="26" height="18" rx="4" class="fmt-mxfp8"/>
+    <text x="254" y="614" class="legend">MXFP8</text>
+    <rect x="360" y="600" width="26" height="18" rx="4" class="fmt-nvfp4"/>
+    <text x="394" y="614" class="legend">NVFP4</text>
+    <rect x="500" y="600" width="26" height="18" rx="4" class="fmt-bf16"/>
+    <text x="534" y="614" class="legend">BF16 (high precision)</text>
+  </g>
+</svg>
diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg
new file mode 100644
index 0000000000..0e17ffdfa4
--- /dev/null
+++ b/docs/features/low_precision_training/fine_grained_quantization/img/fine_grained_linear_mapping.svg
@@ -0,0 +1,59 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 270" width="100%" style="max-width: 900px;" role="img" aria-labelledby="title description">
+  <title id="title">Per-GEMM formats and the operands each GEMM consumes</title>
+  <desc id="description">Three GEMM cards. Fprop consumes input.rowwise and weight.rowwise, both MXFP8. Dgrad consumes weight.columnwise and grad_output.rowwise, both NVFP4. Wgrad consumes input.columnwise and grad_output.columnwise, both BF16.</desc>
+  <defs>
+    <style>
+      @import url("../_static/css/diagram-colors.css");
+      .label { font: 600 14px 'Segoe UI', Arial, sans-serif; fill: #212121; text-anchor: middle; }
+      .format { font: 600 13px 'Segoe UI', Arial, sans-serif; text-anchor: middle; }
+      .op { font: 600 16px 'Segoe UI', Arial, sans-serif; fill: #616161; text-anchor: middle; }
+      .fmt-mxfp8 { fill: #e3f2fd; stroke: #1976d2; stroke-width: 2; }
+      .fmt-nvfp4 { fill: #fff3e0; stroke: #f57c00; stroke-width: 2; }
+      .fmt-bf16 { fill: #e8f5e9; stroke: #388e3c; stroke-width: 2; }
+      .fmt-mxfp8-text { fill: #1565c0; }
+      .fmt-nvfp4-text { fill: #e65100; }
+      .fmt-bf16-text { fill: #2e7d32; }
+      .legend { font: 13px 'Segoe UI', Arial, sans-serif; fill: #424242; text-anchor: start; }
+    </style>
+  </defs>
+
+  <g aria-label="Fprop GEMM">
+    <rect x="30" y="30" width="300" height="220" rx="7" class="panel"/>
+    <text x="180" y="58" class="label">fprop</text>
+    <text x="180" y="78" class="small-text">format: MXFP8</text>
+    <rect x="55" y="95" width="250" height="50" rx="6" class="fmt-mxfp8"/>
+    <text x="180" y="116" class="small-text">input.rowwise</text>
+    <text x="180" y="134" class="format fmt-mxfp8-text">MXFP8</text>
+    <text x="180" y="167" class="op">&#215;</text>
+    <rect x="55" y="180" width="250" height="50" rx="6" class="fmt-mxfp8"/>
+    <text x="180" y="201" class="small-text">weight.rowwise</text>
+    <text x="180" y="219" class="format fmt-mxfp8-text">MXFP8</text>
+  </g>
+
+  <g aria-label="Dgrad GEMM">
+    <rect x="350" y="30" width="300" height="220" rx="7" class="panel"/>
+    <text x="500" y="58" class="label">dgrad</text>
+    <text x="500" y="78" class="small-text">format: NVFP4</text>
+    <rect x="375" y="95" width="250" height="50" rx="6" class="fmt-nvfp4"/>
+    <text x="500" y="116" class="small-text">grad_output.rowwise</text>
+    <text x="500" y="134" class="format fmt-nvfp4-text">NVFP4</text>
+    <text x="500" y="167" class="op">&#215;</text>
+    <rect x="375" y="180" width="250" height="50" rx="6" class="fmt-nvfp4"/>
+    <text x="500" y="201" class="small-text">weight.columnwise</text>
+    <text x="500" y="219" class="format fmt-nvfp4-text">NVFP4</text>
+  </g>
+
+  <g aria-label="Wgrad GEMM">
+    <rect x="670" y="30" width="300" height="220" rx="7" class="panel"/>
+    <text x="820" y="58" class="label">wgrad</text>
+    <text x="820" y="78" class="small-text">format: BF16</text>
+    <rect x="695" y="95" width="250" height="50" rx="6" class="fmt-bf16"/>
+    <text x="820" y="116" class="small-text">input.columnwise</text>
+    <text x="820" y="134" class="format fmt-bf16-text">BF16</text>
+    <text x="820" y="167" class="op">&#215;</text>
+    <rect x="695" y="180" width="250" height="50" rx="6" class="fmt-bf16"/>
+    <text x="820" y="201" class="small-text">grad_output.columnwise</text>
+    <text x="820" y="219" class="format fmt-bf16-text">BF16</text>
+  </g>
+</svg>
diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg
new file mode 100644
index 0000000000..ebc077eb04
--- /dev/null
+++ b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_columnwise_source.svg
@@ -0,0 +1,74 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 430" width="100%" style="max-width: 900px;" role="img" aria-labelledby="title description">
+  <title id="title">Hybrid quantizer columnwise source choices</title>
+  <desc id="description">With original provenance, both quantizers consume the original high-precision tensor. With rowwise-dequantized provenance, the columnwise quantizer consumes the dequantized rowwise representation.</desc>
+  <defs>
+    <style>
+      @import url("../_static/css/diagram-colors.css");
+      .source { fill: #ede7f6; stroke: #673ab7; stroke-width: 2; }
+      .quantizer { fill: #e8f5e9; stroke: #388e3c; stroke-width: 2; }
+      .representation { fill: #e3f2fd; stroke: #1976d2; stroke-width: 2; }
+      .dequantize { fill: #fff3e0; stroke: #f57c00; stroke-width: 2; }
+    </style>
+    <marker id="arrowhead" markerWidth="8" markerHeight="8" refX="7" refY="3" orient="auto">
+      <polygon points="0 0, 8 3, 0 6" fill="#616161"/>
+    </marker>
+  </defs>
+
+  <text x="500" y="28" class="title">Choosing the columnwise source</text>
+
+  <g aria-label="Original source">
+    <rect x="20" y="50" width="460" height="350" rx="8" class="panel"/>
+    <text x="250" y="82" class="section-title">columnwise_source="original"</text>
+
+    <rect x="150" y="105" width="200" height="48" rx="6" class="source"/>
+    <text x="250" y="134" class="text">High-precision tensor</text>
+
+    <path d="M 250 153 L 250 178 L 130 178 L 130 205" class="arrow"/>
+    <path d="M 250 153 L 250 178 L 370 178 L 370 205" class="arrow"/>
+    <text x="250" y="198" class="small-text">same original source</text>
+
+    <rect x="55" y="210" width="150" height="50" rx="6" class="quantizer"/>
+    <text x="130" y="240" class="text">Rowwise quantizer</text>
+    <rect x="295" y="210" width="150" height="50" rx="6" class="quantizer"/>
+    <text x="370" y="240" class="text">Columnwise quantizer</text>
+
+    <path d="M 130 260 L 130 300" class="arrow"/>
+    <path d="M 370 260 L 370 300" class="arrow"/>
+    <rect x="55" y="305" width="150" height="55" rx="6" class="representation"/>
+    <text x="130" y="329" class="text">Rowwise</text>
+    <text x="130" y="347" class="small-text">representation</text>
+    <rect x="295" y="305" width="150" height="55" rx="6" class="representation"/>
+    <text x="370" y="329" class="text">Columnwise</text>
+    <text x="370" y="347" class="small-text">representation</text>
+  </g>
+
+  <g aria-label="Rowwise-dequantized source">
+    <rect x="520" y="50" width="460" height="350" rx="8" class="panel"/>
+    <text x="750" y="82" class="section-title">columnwise_source="rowwise_dequantized"</text>
+
+    <rect x="555" y="95" width="180" height="44" rx="6" class="source"/>
+    <text x="645" y="122" class="text">High-precision tensor</text>
+    <path d="M 645 139 L 645 160" class="arrow"/>
+
+    <rect x="570" y="165" width="150" height="46" rx="6" class="quantizer"/>
+    <text x="645" y="193" class="text">Rowwise quantizer</text>
+    <path d="M 645 211 L 645 232" class="arrow"/>
+
+    <rect x="570" y="237" width="150" height="52" rx="6" class="representation"/>
+    <text x="645" y="259" class="text">Rowwise</text>
+    <text x="645" y="277" class="small-text">representation</text>
+    <path d="M 720 263 L 775 263" class="arrow"/>
+
+    <rect x="780" y="237" width="150" height="52" rx="6" class="dequantize"/>
+    <text x="855" y="269" class="text">Dequantize</text>
+    <path d="M 855 289 L 855 320" class="arrow"/>
+
+    <rect x="780" y="325" width="150" height="46" rx="6" class="quantizer"/>
+    <text x="855" y="353" class="text">Columnwise quantizer</text>
+    <path d="M 780 348 L 725 348" class="arrow"/>
+    <rect x="570" y="325" width="150" height="46" rx="6" class="representation"/>
+    <text x="645" y="344" class="text">Columnwise</text>
+    <text x="645" y="361" class="small-text">representation</text>
+  </g>
+</svg>
diff --git a/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg
new file mode 100644
index 0000000000..6e542306c7
--- /dev/null
+++ b/docs/features/low_precision_training/fine_grained_quantization/img/hybrid_quantizer.svg
@@ -0,0 +1,58 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 1000 300" width="100%" style="max-width: 900px;" role="img" aria-labelledby="title description">
+  <title id="title">HybridQuantizer data flow</title>
+  <desc id="description">A high-precision tensor enters a HybridQuantizer whose rowwise child is an MXFP8 quantizer and columnwise child is an NVFP4 quantizer. The result is a HybridQuantizedTensor with an MXFP8 rowwise representation and an NVFP4 columnwise representation, each consumed by a different GEMM.</desc>
+  <defs>
+    <style>
+      @import url("../_static/css/diagram-colors.css");
+      .label { font: 600 14px 'Segoe UI', Arial, sans-serif; fill: #212121; text-anchor: middle; }
+      .format { font: 600 13px 'Segoe UI', Arial, sans-serif; text-anchor: middle; }
+      .fmt-mxfp8 { fill: #e3f2fd; stroke: #1976d2; stroke-width: 2; }
+      .fmt-nvfp4 { fill: #fff3e0; stroke: #f57c00; stroke-width: 2; }
+      .fmt-bf16 { fill: #e8f5e9; stroke: #388e3c; stroke-width: 2; }
+      .fmt-mxfp8-text { fill: #1565c0; }
+      .fmt-nvfp4-text { fill: #e65100; }
+      .fmt-bf16-text { fill: #2e7d32; }
+    </style>
+    <marker id="arrowhead" markerWidth="8" markerHeight="8" refX="7" refY="3" orient="auto">
+      <polygon points="0 0, 8 3, 0 6" fill="#616161"/>
+    </marker>
+  </defs>
+
+  <g aria-label="High-precision input">
+    <rect x="30" y="115" width="160" height="70" rx="7" class="fmt-bf16"/>
+    <text x="110" y="145" class="label fmt-bf16-text">tensor</text>
+    <text x="110" y="165" class="small-text">high precision (BF16)</text>
+  </g>
+
+  <line x1="190" y1="150" x2="250" y2="150" class="arrow"/>
+
+  <g aria-label="HybridQuantizer">
+    <rect x="255" y="60" width="280" height="180" rx="7" class="panel"/>
+    <text x="395" y="88" class="label">HybridQuantizer</text>
+    <rect x="275" y="105" width="240" height="50" rx="6" class="fmt-mxfp8"/>
+    <text x="395" y="126" class="small-text">rowwise_quantizer</text>
+    <text x="395" y="144" class="format fmt-mxfp8-text">MXFP8Quantizer</text>
+    <rect x="275" y="170" width="240" height="50" rx="6" class="fmt-nvfp4"/>
+    <text x="395" y="191" class="small-text">columnwise_quantizer</text>
+    <text x="395" y="209" class="format fmt-nvfp4-text">NVFP4Quantizer</text>
+  </g>
+
+  <line x1="535" y1="150" x2="595" y2="150" class="arrow"/>
+
+  <g aria-label="HybridQuantizedTensor">
+    <rect x="600" y="60" width="250" height="180" rx="7" class="panel"/>
+    <text x="725" y="88" class="label">HybridQuantizedTensor</text>
+    <rect x="620" y="105" width="210" height="50" rx="6" class="fmt-mxfp8"/>
+    <text x="725" y="126" class="small-text">rowwise</text>
+    <text x="725" y="144" class="format fmt-mxfp8-text">MXFP8</text>
+    <rect x="620" y="170" width="210" height="50" rx="6" class="fmt-nvfp4"/>
+    <text x="725" y="191" class="small-text">columnwise</text>
+    <text x="725" y="209" class="format fmt-nvfp4-text">NVFP4</text>
+  </g>
+
+  <line x1="830" y1="130" x2="900" y2="130" class="arrow"/>
+  <line x1="830" y1="195" x2="900" y2="195" class="arrow"/>
+  <text x="935" y="134" class="small-text">GEMM 1</text>
+  <text x="935" y="199" class="small-text">GEMM 2</text>
+</svg>
diff --git a/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst b/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst
index 48d17db8d5..557a4b09e8 100644
--- a/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst
+++ b/docs/features/low_precision_training/fp8_blockwise_scaling/fp8_blockwise_scaling.rst
@@ -180,6 +180,37 @@ Blackwell and later (SM >= 10.0) – the recipe is emulated with MXFP8. Note tha
 
 ----
 
+
+Quantizer
+---------
+
+.. tabs::
+
+   .. tab:: PyTorch
+
+      Blockwise scaling uses
+      :class:`~transformer_engine.pytorch.Float8BlockQuantizer`. Each block of
+      the tensor gets its own power-of-two scale: ``block_scaling_dim=1``
+      scales 1x128 blocks, ``block_scaling_dim=2`` (the default) scales
+      128x128 blocks. This recipe is not available in TE/JAX.
+
+      .. code-block:: python
+
+         import torch
+         import transformer_engine.pytorch as te
+
+         tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16)
+
+         quantizer = te.Float8BlockQuantizer(
+             fp8_dtype=te.DType.kFloat8E4M3,
+             rowwise=True,
+             columnwise=True,
+             block_scaling_dim=1,
+         )
+
+         qtensor = quantizer(tensor)
+         roundtrip = qtensor.dequantize()
+
 Developer Notes
 ---------------
 
diff --git a/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst b/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst
index cac3792194..2436a07566 100644
--- a/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst
+++ b/docs/features/low_precision_training/fp8_current_scaling/fp8_current_scaling.rst
@@ -164,6 +164,57 @@ Here's how to use FP8 Current Scaling recipe in PyTorch and JAX:
 
 ----
 
+
+Quantizer
+---------
+
+.. tabs::
+
+   .. tab:: PyTorch
+
+      Current scaling uses
+      :class:`~transformer_engine.pytorch.Float8CurrentScalingQuantizer`. It
+      needs no external state: at each call it computes the amax of the input
+      tensor, derives the scale from it, and then quantizes.
+
+      .. code-block:: python
+
+         import torch
+         import transformer_engine.pytorch as te
+
+         tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16)
+
+         quantizer = te.Float8CurrentScalingQuantizer(
+             fp8_dtype=te.DType.kFloat8E4M3,
+             device="cuda",
+         )
+
+         qtensor = quantizer(tensor)
+         roundtrip = qtensor.dequantize()
+
+   .. tab:: JAX
+
+      Current scaling uses ``CurrentScaleQuantizer``. At each call it computes
+      the amax of the input tensor, derives the scale from it, and then
+      quantizes.
+
+      .. code-block:: python
+
+         import jax.numpy as jnp
+         from transformer_engine.jax.quantize import (
+             QuantizerFactory, ScalingMode, QuantizeLayout,
+         )
+
+         x = jnp.ones((256, 256), dtype=jnp.bfloat16)
+
+         quantizer = QuantizerFactory.create(
+             scaling_mode=ScalingMode.CURRENT_TENSOR_SCALING,
+             q_dtype=jnp.float8_e4m3fn,
+             q_layout=QuantizeLayout.ROWWISE,
+         )
+         qtensor = quantizer.quantize(x)
+         roundtrip = qtensor.dequantize()
+
 Developer Notes
 ---------------
 
@@ -177,4 +228,4 @@ On Blackwell and later, rowwise and columnwise tensors share the same memory lay
 so all-gather of columnwise tensors is directly supported.
 
 For Hopper and Ada, all-gather of transposed FP8 tensors is not supported. 
-The rowwise tensor is gathered first, then transposed to columnwise format.
\ No newline at end of file
+The rowwise tensor is gathered first, then transposed to columnwise format.
diff --git a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst
index d39787f6f5..99a379eed1 100644
--- a/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst
+++ b/docs/features/low_precision_training/fp8_delayed_scaling/fp8_delayed_scaling.rst
@@ -160,4 +160,69 @@ However, amax reduction works slightly differently in different frameworks.
 Supported devices
 -----------------
 
-Ada and later (SM 8.9+)
\ No newline at end of file
+Ada and later (SM 8.9+)
+
+Quantizer
+---------
+
+.. tabs::
+
+   .. tab:: PyTorch
+
+      Delayed scaling uses
+      :class:`~transformer_engine.pytorch.Float8Quantizer`. It does not
+      compute the scaling factor from the current tensor: one-element
+      ``scale`` and ``amax`` buffers are supplied at construction.
+      Quantization applies the given scale and records the tensor's amax into
+      the ``amax`` buffer.
+
+      During training both buffers are views into the recipe state: ``scale``
+      into its per-quantizer scale vector, ``amax`` into the current row of
+      its ``(amax_history_len, num_quantizers)`` amax history. At the end of
+      each step the recipe state computes a new scale from the history (its
+      max or most recent entry, per ``amax_compute_algo``), rolls the history
+      by one slot, and zeroes the current row — all in place, so the views
+      held by the quantizer stay valid for the whole training run.
+
+      .. code-block:: python
+
+         import torch
+         import transformer_engine.pytorch as te
+
+         tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16)
+
+         quantizer = te.Float8Quantizer(
+             scale=torch.ones(1, device="cuda"),
+             amax=torch.zeros(1, device="cuda"),
+             fp8_dtype=te.DType.kFloat8E4M3,
+         )
+
+         qtensor = quantizer(tensor)
+         roundtrip = qtensor.dequantize()
+
+   .. tab:: JAX
+
+      Delayed scaling uses ``DelayedScaleQuantizer``. The ``scale`` and the
+      ``amax_history`` (1024 entries by default) are fields of the quantizer
+      itself, carried through JAX transformations as its pytree state. Each
+      ``quantize()`` call applies the current ``scale``, then updates the
+      state: the tensor's amax is written into the history, a new scale is
+      computed from the history (max or most-recent entry, per
+      ``amax_compute_algo``), and the history is rolled by one slot.
+
+      .. code-block:: python
+
+         import jax.numpy as jnp
+         from transformer_engine.jax.quantize import (
+             QuantizerFactory, ScalingMode, QuantizeLayout,
+         )
+
+         x = jnp.ones((256, 256), dtype=jnp.bfloat16)
+
+         quantizer = QuantizerFactory.create(
+             scaling_mode=ScalingMode.DELAYED_TENSOR_SCALING,
+             q_dtype=jnp.float8_e4m3fn,
+             q_layout=QuantizeLayout.ROWWISE,
+         )
+         qtensor = quantizer.quantize(x)
+         roundtrip = qtensor.dequantize()
diff --git a/docs/features/low_precision_training/index.rst b/docs/features/low_precision_training/index.rst
index 0a798f1364..b9649c00a4 100644
--- a/docs/features/low_precision_training/index.rst
+++ b/docs/features/low_precision_training/index.rst
@@ -15,4 +15,5 @@ Low precision training
    fp8_blockwise_scaling/fp8_blockwise_scaling.rst
    mxfp8/mxfp8.rst
    nvfp4/nvfp4.rst
+   fine_grained_quantization/fine_grained_quantization.rst
    speedups.rst
diff --git a/docs/features/low_precision_training/introduction/introduction.rst b/docs/features/low_precision_training/introduction/introduction.rst
index fba7796ece..2255308b04 100644
--- a/docs/features/low_precision_training/introduction/introduction.rst
+++ b/docs/features/low_precision_training/introduction/introduction.rst
@@ -283,3 +283,71 @@ so GEMM with tensors ``A`` and ``B`` returns ``B * A^T``.
    :file: img/fp8_linear_flow.svg
 
 *Figure 4: Forward pass of a Linear layer with low precision data flow.*
+
+Quantizers
+----------
+
+Every recipe implements its quantization logic in a **quantizer** — an object
+that converts a high-precision tensor into a quantized one. TE modules create
+and use quantizers internally according to the active recipe, but a quantizer
+can also be used directly:
+
+.. tabs::
+
+   .. tab:: PyTorch
+
+      .. code-block:: python
+
+         import torch
+         import transformer_engine.pytorch as te
+
+         tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16)
+
+         quantizer = te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)
+         qtensor = quantizer(tensor)            # quantize
+         roundtrip = qtensor.dequantize()       # back to high precision
+
+      The main parts of the interface are:
+
+      * ``quantize(tensor)`` — quantizes a high-precision tensor and returns a
+        ``QuantizedTensor``; calling the quantizer (``quantizer(tensor)``) is
+        a shorthand;
+      * ``update_quantized(src, dst)`` — quantizes ``src`` in place into an
+        already-allocated quantized tensor ``dst``;
+      * ``make_empty(shape)`` — allocates an uninitialized quantized tensor to
+        be filled later;
+      * ``rowwise_usage`` / ``columnwise_usage`` — flags selecting which of
+        the two GEMM-oriented representations the produced tensor holds;
+      * the returned ``QuantizedTensor`` supports ``dequantize()`` back to
+        high precision.
+
+   .. tab:: JAX
+
+      .. code-block:: python
+
+         import jax.numpy as jnp
+         from transformer_engine.jax.quantize import (
+             QuantizerFactory, ScalingMode, QuantizeLayout,
+         )
+
+         x = jnp.ones((256, 256), dtype=jnp.bfloat16)
+
+         quantizer = QuantizerFactory.create(
+             scaling_mode=ScalingMode.MXFP8_1D_SCALING,
+             q_dtype=jnp.float8_e4m3fn,
+             q_layout=QuantizeLayout.ROWWISE,
+         )
+         qtensor = quantizer.quantize(x)
+         roundtrip = qtensor.dequantize()
+
+      The main parts of the interface are:
+
+      * ``quantize(x, is_rowwise=..., is_colwise=...)`` — quantizes a tensor
+        and returns a ``ScaledTensor`` holding the requested representations
+        (the default comes from the quantizer's ``q_layout``);
+      * the returned ``ScaledTensor`` supports ``dequantize()`` back to high
+        precision;
+      * quantizers are registered pytrees, so they can be passed through JAX
+        transformations.
+
+Each recipe section ends with a short description of that recipe's quantizer.
diff --git a/docs/features/low_precision_training/mxfp8/mxfp8.rst b/docs/features/low_precision_training/mxfp8/mxfp8.rst
index 1fbcc43af9..1827d42cb6 100644
--- a/docs/features/low_precision_training/mxfp8/mxfp8.rst
+++ b/docs/features/low_precision_training/mxfp8/mxfp8.rst
@@ -152,6 +152,57 @@ SM 10.0, SM 10.3
 
 ----
 
+
+Quantizer
+---------
+
+.. tabs::
+
+   .. tab:: PyTorch
+
+      MXFP8 uses :class:`~transformer_engine.pytorch.MXFP8Quantizer`. Every
+      32-element block shares one power-of-two (E8M0) scale, computed from the
+      block's amax at quantization time; no external state is needed.
+
+      .. code-block:: python
+
+         import torch
+         import transformer_engine.pytorch as te
+
+         tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16)
+
+         quantizer = te.MXFP8Quantizer(fp8_dtype=te.DType.kFloat8E4M3)
+
+         qtensor = quantizer(tensor)
+         roundtrip = qtensor.dequantize()
+
+   .. tab:: JAX
+
+      MXFP8 uses ``BlockScaleQuantizer`` — the JAX quantizer for block-based
+      scaling, selected by ``ScalingMode.MXFP8_1D_SCALING``. Instead of one
+      scale per tensor, the tensor is split along the quantization axis into
+      32-element blocks and each block gets its own power-of-two (E8M0) scale,
+      computed from that block's amax at quantization time. Because the scale
+      is derived from the current data, no external state (scale buffers or
+      amax history) is needed.
+
+      .. code-block:: python
+
+         import jax.numpy as jnp
+         from transformer_engine.jax.quantize import (
+             QuantizerFactory, ScalingMode, QuantizeLayout,
+         )
+
+         x = jnp.ones((256, 256), dtype=jnp.bfloat16)
+
+         quantizer = QuantizerFactory.create(
+             scaling_mode=ScalingMode.MXFP8_1D_SCALING,
+             q_dtype=jnp.float8_e4m3fn,
+             q_layout=QuantizeLayout.ROWWISE,
+         )
+         qtensor = quantizer.quantize(x)
+         roundtrip = qtensor.dequantize()
+
 Developer Notes
 ---------------
 
@@ -210,4 +261,4 @@ All-gather of columnwise tensors
 All-gather of columnwise tensors is supported and necessary because:
 
 - columnwise quantized tensors cannot be computed from rowwise quantized ones,
-- gathering high-precision tensors is avoided in most cases for performance reasons.
\ No newline at end of file
+- gathering high-precision tensors is avoided in most cases for performance reasons.
diff --git a/docs/features/low_precision_training/nvfp4/nvfp4.rst b/docs/features/low_precision_training/nvfp4/nvfp4.rst
index 900789b0d3..26d798651d 100644
--- a/docs/features/low_precision_training/nvfp4/nvfp4.rst
+++ b/docs/features/low_precision_training/nvfp4/nvfp4.rst
@@ -250,6 +250,60 @@ Supported devices
 
 ----
 
+
+Quantizer
+---------
+
+.. tabs::
+
+   .. tab:: PyTorch
+
+      NVFP4 uses :class:`~transformer_engine.pytorch.NVFP4Quantizer`. It
+      implements the two-level scaling described above: an FP8 (E4M3) scale
+      per 16-element block plus one FP32 scale per tensor. Further keyword
+      options select the recipe variations from this page (random Hadamard
+      transforms, stochastic rounding, 2D weight scaling); they are internal
+      knobs and may change without notice.
+
+      .. code-block:: python
+
+         import torch
+         import transformer_engine.pytorch as te
+
+         tensor = torch.randn(256, 256, device="cuda", dtype=torch.bfloat16)
+
+         quantizer = te.NVFP4Quantizer(
+             fp4_dtype=te.DType.kFloat4E2M1,
+             rowwise=True,
+             columnwise=True,
+         )
+
+         qtensor = quantizer(tensor)
+         roundtrip = qtensor.dequantize()
+
+   .. tab:: JAX
+
+      NVFP4 uses its own ``NVFP4Quantizer``, with the same two-level scaling.
+      ``ScalingMode.NVFP4_1D_SCALING`` selects per-block scaling only,
+      ``ScalingMode.NVFP4_2D_SCALING`` adds 2D weight scaling.
+
+      .. code-block:: python
+
+         import jax.numpy as jnp
+         from transformer_engine.jax.quantize import (
+             QuantizerFactory, ScalingMode, QuantizeLayout,
+         )
+
+         x = jnp.ones((256, 256), dtype=jnp.bfloat16)
+
+         quantizer = QuantizerFactory.create(
+             scaling_mode=ScalingMode.NVFP4_1D_SCALING,
+             q_dtype=jnp.float4_e2m1fn,
+             q_layout=QuantizeLayout.ROWWISE,
+         )
+         qtensor = quantizer.quantize(x)
+         roundtrip = qtensor.dequantize()
+
 Developer Notes
 ---------------
 
diff --git a/docs/features/low_precision_training/performance_considerations/performance_considerations.rst b/docs/features/low_precision_training/performance_considerations/performance_considerations.rst
index 2c21799dd6..afa5d16c86 100644
--- a/docs/features/low_precision_training/performance_considerations/performance_considerations.rst
+++ b/docs/features/low_precision_training/performance_considerations/performance_considerations.rst
@@ -143,6 +143,61 @@ Transformer Engine chooses the best possible fusion internally taking the recipe
 
 *Figure 3: Three scenarios of producing quantized tensors in rowwise and columnwise usages.*
 
+**Usages in the quantizer API**
+
+The usages are visible directly in the quantizer API:
+
+.. tabs::
+
+   .. tab:: PyTorch
+
+      At quantization time, the quantizer's ``rowwise_usage`` and
+      ``columnwise_usage`` flags select which representations ``quantize()``
+      produces; when both are set, the representations are computed together
+      in one fused kernel (scenario 1 above).
+
+      After quantization, ``update_usage()`` on the quantized tensor removes a
+      representation or, when supported by the format, generates a missing one.
+      Passing ``rowwise_usage=False`` after the forward pass frees the rowwise
+      data while keeping the columnwise data for backward. Some formats also
+      support ``columnwise_usage=True`` to create the columnwise representation
+      from the data already present (e.g. by a transpose on Hopper — scenario 3
+      above); unsupported requests raise an error. Arguments left as ``None``
+      preserve the current state.
+
+      .. code-block:: python
+
+         quantizer = te.MXFP8Quantizer(
+             fp8_dtype=te.DType.kFloat8E4M3,
+             rowwise=True,
+             columnwise=True,
+         )
+
+         qtensor = quantizer(tensor)   # both representations, one fused kernel
+
+         qtensor.update_usage(rowwise_usage=False)  # drop rowwise, keep columnwise
+
+   .. tab:: JAX
+
+      The usages are selected when the tensor is quantized: the quantizer's
+      ``q_layout`` (``QuantizeLayout.ROWWISE``, ``COLWISE``, or
+      ``ROWWISE_COLWISE``) sets the default, and ``quantize()`` accepts
+      ``is_rowwise``/``is_colwise`` overrides. Requesting both usages returns
+      a ``ScaledTensor2x`` holding the two representations. There is no
+      in-place ``update_usage()``: JAX arrays are immutable, so a
+      representation is not added or dropped later — unneeded ones are simply
+      not requested and get dropped by XLA's dead-code elimination.
+
+      .. code-block:: python
+
+         quantizer = QuantizerFactory.create(
+             scaling_mode=ScalingMode.MXFP8_1D_SCALING,
+             q_dtype=jnp.float8_e4m3fn,
+             q_layout=QuantizeLayout.ROWWISE_COLWISE,
+         )
+
+         qtensor = quantizer.quantize(x)  # ScaledTensor2x, both representations
+         rowwise_only = quantizer.quantize(x, is_rowwise=True, is_colwise=False)
 
 
 Memory usage
@@ -470,4 +525,3 @@ Actual behavior depends on the recipe and module configuration.
 *Figure 5: All-gather of quantized tensors for input and gradient tensors. 
 This is one possible scenario — actual behavior varies depending on the recipe and module configuration.*
 
-
diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py
index 5d5ce1f6cf..128e8280bb 100644
--- a/transformer_engine/common/recipe/__init__.py
+++ b/transformer_engine/common/recipe/__init__.py
@@ -639,13 +639,18 @@ class CustomRecipe(Recipe):
     ----------
     qfactory : Callable
         Factory callable that returns a quantizer instance *or* a
-        ``QuantizerRequest`` subclass for a given ``QuantizerRole``.
+        ``QuantizerRequest`` subclass for a given optional ``QuantizerRole``.
         The callable is invoked as::
 
             qfactory(
-                role: QuantizerRole,
+                role: Optional[QuantizerRole],
             ) -> Union[Quantizer, QuantizerRequest]
 
+        Boundary slots may provide ``None`` or a role with empty fields. The
+        factory must return a valid object for every call. Return an
+        ``IdentityQuantizer`` for an intentional high-precision slot instead
+        of returning ``None``.
+
         ``QuantizerRole`` is a frozen dataclass with the following fields:
 
         - ``module_type`` (str): module type (empty string when not set), e.g.
@@ -663,7 +668,8 @@ class CustomRecipe(Recipe):
 
         See ``transformer_engine.pytorch.quantization.QuantizerRole``
         and ``transformer_engine.pytorch.quantization.DelayedScalingRequest``
-        for full documentation.
+        for API details. See :ref:`heterogeneous-quantization-recipes` for
+        construction rules and direction mapping.
 
     backward_override : {None, 'high_precision', 'dequantized'}, default = None
         Backward precision mode. None does not modify backward behavior,
diff --git a/transformer_engine/pytorch/tensor/hybrid_tensor.py b/transformer_engine/pytorch/tensor/hybrid_tensor.py
index 8df2ec8b4b..26d0798b92 100644
--- a/transformer_engine/pytorch/tensor/hybrid_tensor.py
+++ b/transformer_engine/pytorch/tensor/hybrid_tensor.py
@@ -19,6 +19,10 @@
 class HybridQuantizer(Quantizer):
     """Quantizer that composes rowwise and columnwise representations.
 
+    .. warning::
+        **EXPERIMENTAL**: ``HybridQuantizer`` is under active development and
+        its API is subject to change without notice.
+
     When both representations are requested, applies ``rowwise_quantizer`` to
     produce the rowwise representation and ``columnwise_quantizer`` to produce
     the columnwise representation. The results are wrapped in a
diff --git a/transformer_engine/pytorch/tensor/identity_tensor.py b/transformer_engine/pytorch/tensor/identity_tensor.py
index 8310afc653..ec171564fe 100644
--- a/transformer_engine/pytorch/tensor/identity_tensor.py
+++ b/transformer_engine/pytorch/tensor/identity_tensor.py
@@ -26,6 +26,10 @@
 class IdentityQuantizer(Quantizer):
     """Quantizer that produces a high-precision passthrough representation.
 
+    .. warning::
+        **EXPERIMENTAL**: ``IdentityQuantizer`` is under active development and
+        its API is subject to change without notice.
+
     Returns an :class:`IdentityTensorStorage` (or :class:`IdentityTensor`)
     holding the tensor directly, without a low-precision encoding.
     ``general_gemm`` materializes it as a plain tensor, so a GEMM consumes it
@@ -174,6 +178,21 @@ class IdentityTensor(IdentityTensorStorage, QuantizedTensor):
 
     Presents as a standard tensor of its nominal dtype; internally it just
     holds data directly in that dtype, without a low-precision encoding.
+
+    Parameters
+    ----------
+    shape : iterable of int
+        Tensor dimensions.
+    dtype : torch.dtype
+        Logical tensor datatype.
+    hp_data : torch.Tensor
+        Held high-precision data.
+    quantizer : IdentityQuantizer, optional
+        Quantizer that produced the tensor.
+    requires_grad : bool, default = False
+        Whether to compute gradients for this tensor.
+    device : torch.device, optional
+        Device containing the tensor.
     """
 
     def __repr__(self, *, tensor_contents=None):

From cd245040282ba06d98afc3aa09ef7e71d06ac30a Mon Sep 17 00:00:00 2001
From: Evgeny Tsykunov <etsykunov@nvidia.com>
Date: Tue, 8 Sep 2026 12:29:02 +0200
Subject: [PATCH 17/27] Fix hybrid fp8 test guard (#3492)

Signed-off-by: Evgeny <etsykunov@nvidia.com>
---
 tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py | 4 ++++
 tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py      | 4 ++++
 2 files changed, 8 insertions(+)

diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py
index 08e762045a..fc482ce5b2 100644
--- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py
+++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py
@@ -1733,6 +1733,10 @@ def test_fused_adam_hybrid_scale_uniform_across_shards(hybrid_recipe_name):
     ), f"missing hybrid current-scaling directions: {checked}"
 
 
+@pytest.mark.skipif(
+    not te.is_fp8_available(),
+    reason=te.is_fp8_available(return_reason=True)[1],
+)
 def test_fused_adam_hybrid_identity_fp8_master_weights():
     """FSDP2 + FusedAdam with Hybrid(FP8 current rowwise, Identity columnwise).
 
diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py
index 39d825e701..c2a9df765e 100644
--- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py
+++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py
@@ -510,6 +510,10 @@ def _hybrid_param_count():
     _check_fp8_fsdp2_allgather(model, tols=dict(atol=5e-4, rtol=5e-3))
 
 
+@pytest.mark.skipif(
+    not te.is_fp8_available(),
+    reason=te.is_fp8_available(return_reason=True)[1],
+)
 def test_distributed_hybrid_identity_all():
     """FSDP2 training/all-gather with an all-Identity CustomRecipe.
 

From 2281cd5391ffa1c88df70d7b9d2149aee94bd25e Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 8 Sep 2026 15:44:10 +0200
Subject: [PATCH 18/27] [PyTorch] Register custom ops one at a time

The forward/backward pair was the registration primitive, returned as an
_OpPair carrying every object the function had created, because the
autograd-wired variant finished the registration outside it. Forward and
backward are registered almost identically, so the primitive is now one op:
_register_op returns a _RegisteredOp with the plan and the base/wrapper
handles, and the two entry points register the pair as two calls. This also
makes a forward-only op possible later without a placeholder backward.

ForwardResult is gone: the autograd-free forward returns a plain
(output, aux) tuple, symmetric with the backward.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           |  10 +-
 transformer_engine/pytorch/dynamo/__init__.py |   2 -
 .../pytorch/dynamo/custom_op.py               | 281 +++++++++---------
 transformer_engine/pytorch/ops/op.py          |  12 +-
 4 files changed, 145 insertions(+), 160 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index c02863a1a0..0004d86508 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -50,7 +50,7 @@
     QuantizedTensorStorage,
     Quantizer,
 )
-from transformer_engine.pytorch.dynamo import ForwardResult, TensorSpec, to_tensor_spec
+from transformer_engine.pytorch.dynamo import TensorSpec, to_tensor_spec
 from transformer_engine.pytorch import (
     is_fp8_available,
     is_mxfp8_available,
@@ -2389,12 +2389,12 @@ def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16)
 
     @classmethod
     def forward_compute(cls, args):
-        return ForwardResult(args.input_ * args.scale)
+        return args.input_ * args.scale, ()
 
     @classmethod
     def forward_fake(cls, args):
         x = args.input_
-        return ForwardResult(TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device))
+        return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), ()
 
     @classmethod
     def backward_compute(cls, args):
@@ -2487,12 +2487,12 @@ def forward_compute(cls, args):
         if isinstance(offset, QuantizedTensor):
             offset = offset.dequantize()
         out = args.input_ * args.scale * args.extra_scale + offset
-        return ForwardResult(out)
+        return out, ()
 
     @classmethod
     def forward_fake(cls, args):
         x = args.input_
-        return ForwardResult(TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device))
+        return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), ()
 
     @classmethod
     def backward_compute(cls, args):
diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py
index 88a2d716a7..042c209c1f 100644
--- a/transformer_engine/pytorch/dynamo/__init__.py
+++ b/transformer_engine/pytorch/dynamo/__init__.py
@@ -7,7 +7,6 @@
 from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer
 from .tensor_spec import TensorSpec, to_tensor_spec
 from .custom_op import (
-    ForwardResult,
     register_custom_op,
     register_custom_op_with_autograd,
     TensorOrQuantized,
@@ -18,7 +17,6 @@
     "is_value_opaque_quantizer",
     "TensorSpec",
     "to_tensor_spec",
-    "ForwardResult",
     "register_custom_op",
     "register_custom_op_with_autograd",
     "TensorOrQuantized",
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 1f5df798f6..d35958109b 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -14,8 +14,8 @@
 A TE forward/backward implementation takes one dataclass argument
 (``fwd_arg_type`` / ``bwd_arg_type``, e.g. ``LinearFwdArgs``) whose fields mix
 tensors, quantized tensors, quantizers, process groups and plain Python values.
-The autograd-free forward returns ``ForwardResult(output, aux)``; the
-autograd-wired API keeps its saved-tensor and context-metadata contract.
+The autograd-free forward returns an ``(output, aux)`` tuple; the autograd-wired
+API keeps its saved-tensor and context-metadata contract.
 
 A ``torch.library`` custom op is narrower: it only accepts flat schema slots
 (tensors plus opaque objects) and returns a flat ``Tensor[]``.
@@ -117,14 +117,6 @@
 _TE_OP_NAMESPACE = "transformer_engine_compile"
 
 
-@dataclasses.dataclass(frozen=True, slots=True)
-class ForwardResult:
-    """Output and fresh auxiliary tensors produced by an autograd-free forward."""
-
-    output: Any
-    aux: tuple = ()
-
-
 # Annotation for an op arg field that may hold a plain tensor, a quantized
 # tensor subclass or a *bare* ``QuantizedTensorStorage`` (the internal-quantizer
 # optimization). Matched exactly by ``_TensorOrQuantizedAdapter``.
@@ -944,7 +936,7 @@ def _slice_user_grads(
 
 
 # --------------------------------------------------------------------------- #
-# Op registration: base and wrapper ops, autograd wiring
+# Op registration: base and wrapper ops, autograd wiring, one op
 # --------------------------------------------------------------------------- #
 
 
@@ -1181,126 +1173,119 @@ def _all_quantized_tensor_subclasses() -> List[type]:
 
 
 @dataclasses.dataclass(frozen=True)
-class _OpPair:
-    """One registered forward/backward pair, and what a caller needs to drive it."""
-
-    fwd_plan: _ArgPlan
-    bwd_plan: _ArgPlan
-    base_fwd_def: Any
-    base_bwd_op: Any
-    wrapper_fwd_def: Any
-    wrapper_fwd_op: Any
-    wrapper_bwd_op: Any
-
-    def call_forward(
-        self, fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], fwd_args: Any
-    ) -> Tuple[_OutputPlan, List[torch.Tensor]]:
-        """Run the forward op on ``fwd_args``: its output plan and flat payload."""
-        spec_obj = _spec_view(fwd_args, self.fwd_plan.tensor_field_names())
-        out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj))
-        kwargs = self.fwd_plan.pack(fwd_args)
-        payload = self.wrapper_fwd_op(*[kwargs[name] for name in self.fwd_plan.slot_names])
-        return out_plan, payload
+class _RegisteredOp:
+    """One registered custom op: its arg plan and the base / wrapper definitions."""
+
+    plan: _ArgPlan
+    base_def: Any
+    base_op: Any
+    wrapper_def: Any
+    wrapper_op: Any
 
+    def __call__(self, args: Any) -> List[torch.Tensor]:
+        """Pack the args dataclass into slots and call the wrapper op."""
+        kwargs = self.plan.pack(args)
+        return self.wrapper_op(*[kwargs[name] for name in self.plan.slot_names])
 
-def _register_two_tier_pair(
+
+def _register_op(
     *,
-    op_name: str,
-    fwd_arg_type: type,
-    fwd_impl: Callable[[Any], Any],
-    fwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
-    bwd_arg_type: type,
-    bwd_impl: Callable[[Any], Any],
-    bwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
-    num_grad_inputs: int,
-) -> _OpPair:
-    """Define an operation's forward and backward as two-tier custom ops.
-
-    Everything that is common to :func:`register_custom_op` and
-    :func:`register_custom_op_with_autograd`: the arg plans, the base kernels,
-    the wrapper ops that flatten ``QuantizedTensor`` subclass inputs, and the
-    passthrough registrations. Autograd is deliberately not touched here -- that
-    is what the two entry points differ on.
+    name: str,
+    arg_type: type,
+    impl: Callable[[Any], Any],
+    fake_impl: Callable[[Any], Any],
+    pack_result: Callable[[Any], List[torch.Tensor]],
+    flatten_in_body: bool,
+) -> _RegisteredOp:
+    """Define one two-tier custom op: the base kernel, the wrapper op that lets
+    ``QuantizedTensor`` subclasses be inputs, and the passthrough registrations.
+
+    ``flatten_in_body`` also flattens subclass inputs inside the wrapper body,
+    not only through the ``register_torch_dispatch`` rules.
     """
-    wrapper_fwd_name = op_name
-    wrapper_bwd_name = f"{op_name}_backward"
-    base_fwd_name = f"{op_name}_base"
-    base_bwd_name = f"{wrapper_bwd_name}_base"
-    subclass_list = _all_quantized_tensor_subclasses()
-
-    fwd_plan = _parse_arg_type(fwd_arg_type)
-    bwd_plan = _parse_arg_type(bwd_arg_type)
-
-    fwd_schema = f"{fwd_plan.schema_str} -> Tensor[]"
-    bwd_schema = f"{bwd_plan.schema_str} -> Tensor[]"
-
-    base_bwd_qualname = f"{_TE_OP_NAMESPACE}::{base_bwd_name}"
-
-    base_fwd_def = _register_base_op(
-        op_name=base_fwd_name,
-        schema_str=fwd_schema,
-        plan=fwd_plan,
-        impl=fwd_impl,
-        fake_impl=fwd_fake_impl,
-        pack_result=_pack_fwd_result,
+    plan = _parse_arg_type(arg_type)
+    schema = f"{plan.schema_str} -> Tensor[]"
+    subclasses = _all_quantized_tensor_subclasses()
+    slot_offsets = plan.tensor_or_quantized_offsets()
+    namespace = getattr(torch.ops, _TE_OP_NAMESPACE)
+
+    base_def = _register_base_op(
+        op_name=f"{name}_base",
+        schema_str=schema,
+        plan=plan,
+        impl=impl,
+        fake_impl=fake_impl,
+        pack_result=pack_result,
     )
-    _register_base_op(
-        op_name=base_bwd_name,
-        schema_str=bwd_schema,
-        plan=bwd_plan,
-        impl=bwd_impl,
-        fake_impl=bwd_fake_impl,
-        pack_result=lambda g: _pack_bwd_result(g, num_grad_inputs, base_bwd_qualname),
+    base_op = getattr(namespace, f"{name}_base")
+    wrapper_def = _register_wrapper_op(
+        wrapper_op_name=name,
+        schema_str=schema,
+        base_op=base_op,
+        slot_offsets=slot_offsets if flatten_in_body else (),
+        subclasses=subclasses if flatten_in_body else (),
+    )
+    wrapper_op = getattr(namespace, name)
+
+    rule = _make_dispatch_rule(_make_slot_forwarder(base_op, slot_offsets, subclasses))
+    for sub in subclasses:
+        wrapper_def.register_torch_dispatch(sub, rule)
+    _quantized_tensor_passthrough_ops.update((base_op.default, wrapper_op.default))
+
+    return _RegisteredOp(
+        plan=plan,
+        base_def=base_def,
+        base_op=base_op,
+        wrapper_def=wrapper_def,
+        wrapper_op=wrapper_op,
     )
 
-    base_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_fwd_name)
-    base_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), base_bwd_name)
-
-    fwd_slot_offsets = fwd_plan.tensor_or_quantized_offsets()
-    bwd_slot_offsets = bwd_plan.tensor_or_quantized_offsets()
 
-    wrapper_fwd_def = _register_wrapper_op(
-        wrapper_op_name=wrapper_fwd_name,
-        schema_str=fwd_schema,
-        base_op=base_fwd_op,
-        slot_offsets=fwd_slot_offsets,
-        subclasses=subclass_list,
-    )
-    # Pass-through: a subclass input reaches the base op through the dispatch
-    # rule below, never through the wrapper body.
-    wrapper_bwd_def = _register_wrapper_op(
-        wrapper_op_name=wrapper_bwd_name, schema_str=bwd_schema, base_op=base_bwd_op
+def _register_forward_op(
+    *, name: str, arg_type: type, impl: Callable[[Any], Any], fake_impl: Callable[[Any], Any]
+) -> _RegisteredOp:
+    return _register_op(
+        name=name,
+        arg_type=arg_type,
+        impl=impl,
+        fake_impl=fake_impl,
+        pack_result=_pack_fwd_result,
+        flatten_in_body=True,
     )
-    wrapper_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_fwd_name)
-    wrapper_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), wrapper_bwd_name)
 
-    _fwd_rule = _make_dispatch_rule(
-        _make_slot_forwarder(base_fwd_op, fwd_slot_offsets, subclass_list)
-    )
-    _bwd_rule = _make_dispatch_rule(
-        _make_slot_forwarder(base_bwd_op, bwd_slot_offsets, subclass_list)
-    )
 
-    for sub in subclass_list:
-        wrapper_fwd_def.register_torch_dispatch(sub, _fwd_rule)
-        wrapper_bwd_def.register_torch_dispatch(sub, _bwd_rule)
-
-    for op in (wrapper_fwd_op, wrapper_bwd_op, base_fwd_op, base_bwd_op):
-        _quantized_tensor_passthrough_ops.add(op.default)
-
-    return _OpPair(
-        fwd_plan=fwd_plan,
-        bwd_plan=bwd_plan,
-        base_fwd_def=base_fwd_def,
-        base_bwd_op=base_bwd_op,
-        wrapper_fwd_def=wrapper_fwd_def,
-        wrapper_fwd_op=wrapper_fwd_op,
-        wrapper_bwd_op=wrapper_bwd_op,
+def _register_backward_op(
+    *,
+    name: str,
+    arg_type: type,
+    impl: Callable[[Any], Any],
+    fake_impl: Callable[[Any], Any],
+    num_grad_inputs: int,
+) -> _RegisteredOp:
+    # Pass-through body: a subclass input reaches the base op through the
+    # dispatch rule, never through the wrapper body.
+    qualname = f"{_TE_OP_NAMESPACE}::{name}_base"
+    return _register_op(
+        name=name,
+        arg_type=arg_type,
+        impl=impl,
+        fake_impl=fake_impl,
+        pack_result=lambda g: _pack_bwd_result(g, num_grad_inputs, qualname),
+        flatten_in_body=False,
     )
 
 
+def _run_forward(
+    fwd_op: _RegisteredOp, fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], fwd_args: Any
+) -> Tuple[_OutputPlan, List[torch.Tensor]]:
+    """Run the forward op on ``fwd_args``: its output plan and flat payload."""
+    spec_obj = _spec_view(fwd_args, fwd_op.plan.tensor_field_names())
+    out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj))
+    return out_plan, fwd_op(fwd_args)
+
+
 # --------------------------------------------------------------------------- #
-# Op registration: the forward/backward pair, and the autograd-wired variant
+# Op registration: the autograd-free pair, and the autograd-wired variant
 # --------------------------------------------------------------------------- #
 
 
@@ -1328,7 +1313,7 @@ def register_custom_op(
 
     Callable contracts:
 
-    * ``fwd_impl(fwd_args) -> ForwardResult(output, aux)``
+    * ``fwd_impl(fwd_args) -> (output, aux)`` -- ``aux`` is a tuple of fresh tensors
     * ``fwd_fake_impl`` -- its data-free twin over :class:`TensorSpec`
     * ``bwd_impl(bwd_args) -> tuple`` of ``num_grad_inputs`` gradients
     * ``bwd_fake_impl`` -- its data-free twin
@@ -1377,29 +1362,33 @@ def _register_custom_op_impl(
     def adapt_forward(impl):
         def wrapped(args):
             result = impl(args)
-            if not isinstance(result, ForwardResult):
+            if not isinstance(result, tuple) or len(result) != 2:
                 raise TypeError(
-                    f"autograd-free fwd impl must return ForwardResult, got {type(result).__name__}"
+                    "autograd-free fwd impl must return an (output, aux) tuple, got"
+                    f" {type(result).__name__}"
                 )
-            return result.output, result.aux, None
+            output, aux = result
+            return output, tuple(aux), None
 
         return wrapped
 
-    adapted_fwd_impl = adapt_forward(fwd_impl)
     adapted_fwd_fake_impl = adapt_forward(fwd_fake_impl)
-    pair = _register_two_tier_pair(
-        op_name=op_name,
-        fwd_arg_type=fwd_arg_type,
-        fwd_impl=adapted_fwd_impl,
-        fwd_fake_impl=adapted_fwd_fake_impl,
-        bwd_arg_type=bwd_arg_type,
-        bwd_impl=bwd_impl,
-        bwd_fake_impl=bwd_fake_impl,
+    fwd_op = _register_forward_op(
+        name=op_name,
+        arg_type=fwd_arg_type,
+        impl=adapt_forward(fwd_impl),
+        fake_impl=adapted_fwd_fake_impl,
+    )
+    bwd_op = _register_backward_op(
+        name=f"{op_name}_backward",
+        arg_type=bwd_arg_type,
+        impl=bwd_impl,
+        fake_impl=bwd_fake_impl,
         num_grad_inputs=num_grad_inputs,
     )
 
     def forward_fn(fwd_args):
-        out_plan, payload = pair.call_forward(adapted_fwd_fake_impl, fwd_args)
+        out_plan, payload = _run_forward(fwd_op, adapted_fwd_fake_impl, fwd_args)
         outputs = out_plan.user_outputs(payload)
         aux = out_plan.saved_tensors(payload)
         return outputs[0], tuple(aux)
@@ -1408,9 +1397,7 @@ def backward_fn(bwd_args):
         # Unlike the forward payload, each grad occupies exactly one slot
         # (``_pack_bwd_result`` materializes a TensorSpec grad), so there is
         # nothing to reassemble.
-        kwargs = pair.bwd_plan.pack(bwd_args)
-        payload = pair.wrapper_bwd_op(*[kwargs[name] for name in pair.bwd_plan.slot_names])
-        return tuple(_decode_none(t) for t in payload)
+        return tuple(_decode_none(t) for t in bwd_op(bwd_args))
 
     return forward_fn, backward_fn
 
@@ -1520,31 +1507,31 @@ def _register_custom_op_with_autograd_impl(
     if missing:
         raise ValueError(f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}")
 
-    pair = _register_two_tier_pair(
-        op_name=op_name,
-        fwd_arg_type=fwd_arg_type,
-        fwd_impl=fwd_impl,
-        fwd_fake_impl=fwd_fake_impl,
-        bwd_arg_type=bwd_arg_type,
-        bwd_impl=bwd_impl,
-        bwd_fake_impl=bwd_fake_impl,
+    fwd_op = _register_forward_op(
+        name=op_name, arg_type=fwd_arg_type, impl=fwd_impl, fake_impl=fwd_fake_impl
+    )
+    bwd_op = _register_backward_op(
+        name=f"{op_name}_backward",
+        arg_type=bwd_arg_type,
+        impl=bwd_impl,
+        fake_impl=bwd_fake_impl,
         num_grad_inputs=len(input_tensors_for_grad),
     )
 
     autograd_common = {
-        "fwd_plan": pair.fwd_plan,
-        "bwd_plan": pair.bwd_plan,
-        "grad_targets": pair.fwd_plan.resolve_grad_targets(input_tensors_for_grad),
+        "fwd_plan": fwd_op.plan,
+        "bwd_plan": bwd_op.plan,
+        "grad_targets": fwd_op.plan.resolve_grad_targets(input_tensors_for_grad),
         "setup_context_user": setup_context,
         "fwd_fake_impl": fwd_fake_impl,
     }
-    _register_autograd_for_op(fwd_op=pair.base_fwd_def, bwd_op=pair.base_bwd_op, **autograd_common)
+    _register_autograd_for_op(fwd_op=fwd_op.base_def, bwd_op=bwd_op.base_op, **autograd_common)
     _register_autograd_for_op(
-        fwd_op=pair.wrapper_fwd_def, bwd_op=pair.wrapper_bwd_op, **autograd_common
+        fwd_op=fwd_op.wrapper_def, bwd_op=bwd_op.wrapper_op, **autograd_common
     )
 
     def forward_fn(fwd_args):
-        out_plan, payload = pair.call_forward(fwd_fake_impl, fwd_args)
+        out_plan, payload = _run_forward(fwd_op, fwd_fake_impl, fwd_args)
         outputs = out_plan.user_outputs(payload)
         if len(outputs) == 1:
             return outputs[0]
diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py
index 6c5b997ee2..e353d860ee 100644
--- a/transformer_engine/pytorch/ops/op.py
+++ b/transformer_engine/pytorch/ops/op.py
@@ -22,7 +22,7 @@
     autocast,
 )
 from ..tensor import Quantizer
-from ..dynamo import ForwardResult, is_value_opaque_quantizer, register_custom_op
+from ..dynamo import is_value_opaque_quantizer, register_custom_op
 
 
 @dataclasses.dataclass
@@ -324,7 +324,7 @@ def set_extra_output_channel(
     # ------------------------------------------------------------------ #
 
     @classmethod
-    def forward_compute(cls, args: Any) -> ForwardResult:
+    def forward_compute(cls, args: Any) -> tuple[Any, tuple]:
         """Forward computation over explicit arguments.
 
         Takes everything through ``args``; must not read ``self`` or global
@@ -333,7 +333,7 @@ def forward_compute(cls, args: Any) -> ForwardResult:
         raise NotImplementedError
 
     @classmethod
-    def forward_fake(cls, args: Any) -> ForwardResult:
+    def forward_fake(cls, args: Any) -> tuple[Any, tuple]:
         """Allocation-free twin of :meth:`forward_compute` over ``TensorSpec``.
 
         Runs as a meta kernel, outside the traced frame, and more than once per
@@ -684,10 +684,10 @@ def op_forward(
             next_op_input_quantizer=next_op_input_quantizer,
             **kwargs,
         )
-        result = self.forward_compute(args)
+        output, aux = self.forward_compute(args)
         if ctx.requires_grad:
-            self.setup_context(ctx, args, result.aux)
-        return result.output
+            self.setup_context(ctx, args, aux)
+        return output
 
     def compiled_op_forward(
         self,

From b97f5ab97f7d3a2647b79f619c889df0a1a78a5c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Micha=C5=82=20Marcinkiewicz?=
 <43240942+mmarcinkiewicz@users.noreply.github.com>
Date: Tue, 8 Sep 2026 20:35:25 +0200
Subject: [PATCH 19/27] [PyTorch] Use torch's register_custom_class API for
 opaque quantizers when available (#3495)

[PyTorch] Use torch's register_custom_class API for opaque quantizers

PyTorch renamed the opaque-object entry points: register_opaque_type ->
register_custom_class, is_opaque_value_type -> is_opaque_constant_type,
typ="value" -> typ="constant". The old names remain as thin wrappers that
log a deprecation warning on every call, so each TE import now emits 21
warning lines (5 quantizer classes registered, 11 type checks) -- once per
process, including every torch.compile worker. A 16-rank training job logs
them from ~320 processes.

Switch quantizer_opaque.py, custom_op.py and the torch.compile test to the
new names. Both import sites already sit behind try/except that records the
torch.compile custom-op path as disabled, so a PyTorch that predates the
rename (torch < 2.14) keeps importing and running eagerly with the compile
path off -- the same behaviour as a PyTorch without the opaque-object API
at all.

Signed-off-by: Michal Marcinkiewicz <michalm@nvidia.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
---
 tests/pytorch/test_torch_compile.py            |  6 +++---
 transformer_engine/pytorch/dynamo/custom_op.py | 18 +++++++++---------
 .../pytorch/dynamo/quantizer_opaque.py         | 10 +++++-----
 3 files changed, 17 insertions(+), 17 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index f8e09d5ce1..e1c6eaf599 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -22,7 +22,7 @@
     from torch._opaque_base import OpaqueBaseMeta
     from torch._library.opaque_object import (
         get_opaque_type_name,
-        register_opaque_type,
+        register_custom_class,
         MemberType,
     )
 
@@ -215,9 +215,9 @@ def __fx_repr__(self):
                 {"ToyQuantizer": ToyQuantizer},
             )
 
-    register_opaque_type(
+    register_custom_class(
         ToyQuantizer,
-        typ="value",
+        typ="constant",
         members={
             "__setattr__": MemberType.USE_REAL,
             "set_usage": MemberType.USE_REAL,
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 00846d615a..f887a48d95 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -180,7 +180,7 @@ def is_simple_value(cls, value: Any) -> bool:
             return True
         if isinstance(value, type):
             return True
-        if _is_opaque_value_type is not None and _is_opaque_value_type(type(value)):
+        if _is_opaque_constant_type is not None and _is_opaque_constant_type(type(value)):
             return True
         if isinstance(value, dict):
             return all(isinstance(k, str) and cls.is_simple_value(v) for k, v in value.items())
@@ -223,7 +223,7 @@ def _fmt_simple(cls, value: Any) -> str:
         if isinstance(value, tuple):
             body = ", ".join(cls._fmt_simple(v) for v in value)
             return f"({body},)" if len(value) == 1 else f"({body})"
-        if _is_opaque_value_type(type(value)):
+        if _is_opaque_constant_type(type(value)):
             return value.__fx_repr__()[0]
         # repr(float('inf')) is 'inf', which is not an evaluable literal.
         if isinstance(value, float) and not math.isfinite(value):
@@ -285,7 +285,7 @@ def _collect(value: Any) -> None:
                 return
             if isinstance(value, OpaqueValueBundle.PRIMITIVE_TYPES):
                 return
-            if _is_opaque_value_type(type(value)):
+            if _is_opaque_constant_type(type(value)):
                 _, extra = value.__fx_repr__()
                 globals_.update(extra)
 
@@ -297,18 +297,18 @@ def _collect(value: Any) -> None:
 try:
     from torch._library.opaque_object import (
         get_opaque_type_name,
-        is_opaque_value_type as _is_opaque_value_type,
-        register_opaque_type,
+        is_opaque_constant_type as _is_opaque_constant_type,
+        register_custom_class,
     )
 
-    register_opaque_type(OpaqueValueBundle, typ="value")
+    register_custom_class(OpaqueValueBundle, typ="constant")
     _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle)
 # Older torch without opaque_object support.
 except Exception as e:  # pylint: disable=broad-exception-caught  # pragma: no cover
     record_compile_disabled(
         f"could not register OpaqueValueBundle as an opaque type ({e}); use a newer PyTorch build"
     )
-    _is_opaque_value_type = None
+    _is_opaque_constant_type = None
     _OPAQUE_VALUE_BUNDLE_TYPE_NAME = None
 
 try:
@@ -461,8 +461,8 @@ def _is_simple_annot(annot: Any) -> bool:
         return True
     if (
         isinstance(annot, type)
-        and _is_opaque_value_type is not None
-        and _is_opaque_value_type(annot)
+        and _is_opaque_constant_type is not None
+        and _is_opaque_constant_type(annot)
     ):
         return True
     if get_origin(annot) in (tuple, list):
diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
index dc689258d1..da595b1070 100644
--- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py
+++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
@@ -108,15 +108,15 @@ def register_value_opaque_quantizer(cls: type) -> None:
                 "the field in ``_rebuild_derived_state`` instead."
             )
     cls._value_field_names = tuple(fields)
-    # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the
+    # ``register_custom_class`` requires ``__fx_repr__`` to already exist on the
     # class, so attach it before registering.
     if "__fx_repr__" not in cls.__dict__:
         cls.__fx_repr__ = _quantizer_fx_repr
 
     try:
         from torch._library.opaque_object import (  # pylint: disable=import-outside-toplevel
-            register_opaque_type,
-            is_opaque_value_type,
+            register_custom_class,
+            is_opaque_constant_type,
         )
     except (ImportError, AttributeError) as e:
         # Older PyTorch without the opaque-object API: eager value semantics
@@ -127,8 +127,8 @@ def register_value_opaque_quantizer(cls: type) -> None:
         return
 
     try:
-        if not is_opaque_value_type(cls):
-            register_opaque_type(cls, typ="value")
+        if not is_opaque_constant_type(cls):
+            register_custom_class(cls, typ="constant")
     except (RuntimeError, TypeError) as e:
         # Keep TE importable: neither the opaque-type query nor the registration
         # must crash the import, e.g. on PyTorch versions with only partial /

From 7c6fa5986d01b8175986ed4d6959c8b8afbee946 Mon Sep 17 00:00:00 2001
From: Xianduo Li <30922914+lxd-cumt@users.noreply.github.com>
Date: Wed, 9 Sep 2026 09:20:23 +0800
Subject: [PATCH 20/27] Add entrypoint for flagos multi-backend plugin system
 (#3401)

* Add NVTE_PLUGIN support for plugin system on main branch

Port plugin system changes from backend_rc2.14:
- common/__init__.py: load plugin after framework extension init
- dot_product_attention.py: override FlashAttention and get_attention_backend when plugin is active

Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>

* Fix lint

Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com>

---------

Signed-off-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>
Signed-off-by: Przemyslaw Tredak <ptrendx@gmail.com>
Co-authored-by: Xianduo Li <lixianduo@mail.nankai.edu.cn>
Co-authored-by: Przemyslaw Tredak <ptredak@nvidia.com>
Co-authored-by: Przemyslaw Tredak <ptrendx@gmail.com>
---
 transformer_engine/common/__init__.py         | 26 +++++++++++++++++++
 .../dot_product_attention.py                  | 12 +++++++++
 2 files changed, 38 insertions(+)

diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py
index e47a6f11b2..155ceb9baf 100644
--- a/transformer_engine/common/__init__.py
+++ b/transformer_engine/common/__init__.py
@@ -16,6 +16,7 @@
 import sys
 import sysconfig
 from typing import Optional, Tuple
+import warnings
 
 
 @functools.lru_cache(maxsize=None)
@@ -193,6 +194,31 @@ def load_framework_extension(framework: str) -> None:
     sys.modules[module_name] = solib
     spec.loader.exec_module(solib)
 
+    # Plugin system: set NVTE_PLUGIN=<module_name> to let plugin stub take over
+    # transformer_engine_torch and register original pybind as _nv for CUDA backend.
+    # Only applies to the PyTorch extension — JAX has no plugin stub.
+    _nvte_plugin = os.environ.get("NVTE_PLUGIN")
+    if _nvte_plugin and framework == "torch":
+        _original_module = sys.modules.get(module_name)
+        try:
+            # Register _nv alias BEFORE importing the plugin, because the
+            # plugin module may import transformer_engine_torch_nv at top level.
+            sys.modules[module_name + "_nv"] = solib
+            _plugin = importlib.import_module(_nvte_plugin)
+            _plugin.load_plugins()
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            # Rollback to pre-plugin state if plugin failed to fully initialize
+            sys.modules.pop(module_name + "_nv", None)
+            if _original_module is not None:
+                sys.modules[module_name] = _original_module
+            else:
+                sys.modules.pop(module_name, None)
+            warnings.warn(
+                f"NVTE_PLUGIN={_nvte_plugin} but plugin loading failed: {e}",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+
 
 def sanity_checks_for_pypi_installation() -> None:
     """Ensure that package is installed correctly if using PyPI."""
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 de082c8fae..cd087fd642 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
@@ -408,6 +408,18 @@ def _delayed_scaling_recipe() -> Optional[DelayedScaling]:
 _dpa_fp8ds_reduce_amax = os.getenv("NVTE_DPA_FP8DS_REDUCE_AMAX", "1") == "1"
 
 
+# Plugin system: override FlashAttention and get_attention_backend if enabled
+if os.environ.get("NVTE_PLUGIN"):
+    import transformer_engine_torch as tex
+
+    _FlashAttentionNative = FlashAttention
+    FlashAttention = getattr(tex, "flash_attention", _FlashAttentionNative)
+    _plugin_get_attention_backend = getattr(tex, "get_attention_backend", None)
+    if _plugin_get_attention_backend is not None:
+        dpa_utils._original_get_attention_backend = dpa_utils.get_attention_backend
+        dpa_utils.get_attention_backend = _plugin_get_attention_backend
+
+
 __all__ = ["DotProductAttention"]
 
 

From 39a54018b1cec075bb3be6a0582ec2a761fea12a Mon Sep 17 00:00:00 2001
From: Wentao Guo <47335620+GarlGuo@users.noreply.github.com>
Date: Tue, 8 Sep 2026 22:45:07 -0400
Subject: [PATCH 21/27] [PyTorch] Fix NaN expert weight gradients at num_groups
 == 1 with SReLU (#3485)

[PyTorch] Take the num_groups == 1 dense shortcuts only when the kernels agree

At num_groups == 1 the fused grouped MLP may quantize and scale-swizzle the token
buffer as ONE dense tensor over `tensor.shape[0]` rows and read it back the same
way. That is coherent only if the cuDNN kernels also write their intermediates
densely over `tensor.shape[0]`, which is what `use_single_group_runtime_offsets`
asks of them -- they then derive M from the runtime tensor shapes rather than from
the offsets.

The two halves of that decision are currently gated differently. The request to
the kernels is already guarded, because it can be refused:

    if supports_single_group_runtime_offsets:
        fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1

but its siblings -- `use_single_group_weight_swizzle`, `use_single_discrete_weight`,
`use_single_group_dense_fc2`, `use_single_group_dense_dgrad`, the dense quantize
and the plain wgrad GEMM -- are gated on `num_groups == 1` alone. So when the
specialization is unavailable, `ScaledSReLU` at any version or any activation on a
cuDNN frontend older than 1.27.0, the kernels pack their output to
`sum(split_sizes)` rows while TE still reads everything densely.

Because the columnwise swizzled MXFP8 scale layout is `[k/128][m/128][32][4][4]`,
with the m-tile as an INNER stride, an m-extent mismatch misindexes every k-tile
past the first: consumer k-tile j reads producer k-tile j*T_c/T_p. The HEAD rows
-- the ones that do belong to the group -- pick up the wrong scale factors, and
high k-tiles read scale memory the producer never wrote. The data is read
correctly; only the scales move, so the error is order 1 and often NaN. Crucially
this is value-independent: zeroing the padded tail does not prevent it, because
the defect is in the indexing rather than in the padding's contents, so a caller
cannot work around it.

Callers may legally pass a token buffer padded past `sum(split_sizes)` -- an MoE
expert-parallel fixed-capacity receive buffer always does, and
`test_grouped_linear_cuda_graph_safe` documents the padded tail as "intentionally
outside every group" and "uninitialized".

Measured on B300 with MXFP8 and a native ScaledSReLU at num_groups == 1, input
padded 512 -> 768 rows, FC1 main_grad against the unpadded result:

    before   zeroed tail    406528 / 1048576 non-finite
    before   poisoned tail  406528 / 1048576 non-finite   (identical -- not the values)
    after    either tail    0.0000e+00, bit-identical
    control  num_groups=2   clean before and after

Every corrected arm also sits at 1.70e-02 against an unfused reference, the MXFP8
fused-vs-unfused noise floor.

Gate the whole shortcut family on the same predicate the kernels use, applying the
guard TE already writes at the sites that were missing it. No new helper: the
binding calls `_cudnn_frontend_supports_single_group_runtime_offsets` directly,
as this file already does at three other sites, and the flag is named
`use_dense_single_group` to match its neighbours. It deliberately does not claim
the group spans the buffer -- that is unknowable here -- only that producer and
consumer will agree on the extent. It is a host-side check
on an activation type and a package version, so it costs nothing and stays CUDA
graph capturable, and it leaves the GLU activations -- which do get the
specialization -- byte-for-byte on their existing fast path. That matters:
removing the shortcuts outright costs a shared expert +6% to +30% of the fused MLP
step (+5-10% at hidden 4096 / ffn 14336, rising to +28-30% at 1024/4096 with few
tokens), measured by CUDA-graph replay against num_groups >= 2 null controls, and
Megatron's FusedSharedExpertMLP builds exactly this path.

In the two grouped-quantize helpers the condition is just `use_dense_single_group`:
`num_groups == 1` is already folded into that flag at both binding sites, and the
`isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer))` test was dead because
`fuse_grouped_mlp_ops` only builds this fusion for an MXFP8 or NVFP4 recipe.

All the shortcuts have to move together: quantization layout and kernel indexing
must agree, so enabling a subset pairs a dense consumer with a packed producer and
yields NaN rather than a slowdown. Hence one flag threaded explicitly through the
grouped-quantize helpers and `_compute_grad_params` rather than a predicate per
site.

What this does NOT do: where the specialization is available, producer and
consumer agree and the residual exposure is that the plain wgrad GEMM still
contracts over `logical_shape[0]` rows, summing any padding past
`sum(split_sizes)` into the weight gradient. That part is value-dependent -- a
zeroed tail contributes an all-zero outer product and is bit-exact, verified over
10 launches -- so a caller that zeroes its padding, or passes none, is unaffected.
Closing that gap needs either a device-side masked zeroing of TE's own
intermediate buffers or an explicit per-op dense-single-group contract; neither is
attempted here.

The condition cannot be evaluated at run time: `split_sizes` is device-resident,
so testing `sum(split_sizes) == tensor.shape[0]` forces a device-to-host sync, and
that test exists to keep this flow sync-free and graph-capturable.

tests/pytorch/test_grouped_mlp.py is at baseline: 6 failed, 48 passed, the 6 being
pre-existing nvfp4_rht and cuBLAS-version failures also present on main.

Signed-off-by: Wentao Guo <wg0420@princeton.edu>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: vthumbe1503 <vthumbe@nvidia.com>
---
 .../pytorch/ops/fused/grouped_mlp.py          | 48 +++++++++++++------
 1 file changed, 34 insertions(+), 14 deletions(-)

diff --git a/transformer_engine/pytorch/ops/fused/grouped_mlp.py b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
index 58f22c874c..bd8b82d3fa 100644
--- a/transformer_engine/pytorch/ops/fused/grouped_mlp.py
+++ b/transformer_engine/pytorch/ops/fused/grouped_mlp.py
@@ -239,10 +239,11 @@ def _group_quantize_for_grouped_mlp(
     split_sizes: Optional[torch.Tensor],
     *,
     tensor_offsets: Optional[torch.Tensor] = None,
+    use_dense_single_group: bool = False,
 ) -> GroupedTensor:
     """Quantize into grouped storage."""
 
-    if num_groups != 1 or not isinstance(quantizer, (MXFP8Quantizer, NVFP4Quantizer)):
+    if not use_dense_single_group:
         return tex.group_quantize(
             tensor,
             quantizer,
@@ -270,6 +271,7 @@ def _group_quantize_with_amax_for_grouped_mlp(
     columnwise_amax: torch.Tensor,
     *,
     tensor_offsets: Optional[torch.Tensor] = None,
+    use_dense_single_group: bool = False,
 ) -> GroupedTensor:
     """Quantize with precomputed NVFP4 amaxes into grouped storage."""
     if not isinstance(quantizer, NVFP4Quantizer):
@@ -279,9 +281,10 @@ def _group_quantize_with_amax_for_grouped_mlp(
             num_groups,
             split_sizes,
             tensor_offsets=tensor_offsets,
+            use_dense_single_group=use_dense_single_group,
         )
 
-    if num_groups != 1:
+    if not use_dense_single_group:
         return tex.nvfp4_group_quantize_with_amax(
             tensor,
             quantizer,
@@ -690,6 +693,7 @@ def _compute_grad_params(
     scale_view_dtype,
     sf_vec_size,
     offsets,
+    use_dense_single_group,
 ):
     """Compute weight gradients and build grad_params for a GroupedLinear layer.
     Returns the grad_params list in parameter registration order.
@@ -757,7 +761,7 @@ def _compute_grad_params(
                 "distributed-weight fused grouped-MLP requires delay_wgrad_compute=False."
             )
         if (
-            num_groups == 1
+            use_dense_single_group
             and isinstance(grouped_x, (GroupedTensor, GroupedTensorStorage))
             and isinstance(grouped_dy, (GroupedTensor, GroupedTensorStorage))
             and isinstance(grouped_x.quantizer, (MXFP8Quantizer, NVFP4Quantizer))
@@ -1125,6 +1129,9 @@ def fuser_forward(
             raise ValueError(f"Unsupported input shape for fused grouped MLP ({in_shape=}).")
 
         num_groups = fc1_op.num_groups
+        use_dense_single_group = num_groups == 1 and (
+            _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op))
+        )
         fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0
         fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0
 
@@ -1195,7 +1202,7 @@ def fuser_forward(
         # Older cuDNN frontends do not expose this specialization, so use the
         # live generic offset calculation rather than caching CUDA metadata.
         use_offsetless_metadata = (
-            num_groups == 1
+            use_dense_single_group
             and unit_activation_scale
             and isinstance(fc1_input_quantizer, MXFP8Quantizer)
             and supports_single_group_runtime_offsets
@@ -1276,6 +1283,7 @@ def fuser_forward(
                     fc1_weight_quantizer,
                     num_groups,
                     None,
+                    use_dense_single_group=use_dense_single_group,
                 )
         else:
             fc1_weights = [getattr(fc1_op, f"weight{idx}") for idx in range(num_groups)]
@@ -1313,6 +1321,7 @@ def fuser_forward(
                     fc2_weight_quantizer,
                     num_groups,
                     None,
+                    use_dense_single_group=use_dense_single_group,
                 )
         else:
             fc2_weights = [getattr(fc2_op, f"weight{idx}") for idx in range(num_groups)]
@@ -1360,6 +1369,7 @@ def fuser_forward(
                 num_groups,
                 split_sizes,
                 tensor_offsets=fc1_x_tensor_offsets,
+                use_dense_single_group=use_dense_single_group,
             )
 
         use_nvfp4 = isinstance(fc1_input_quantizer, NVFP4Quantizer) or isinstance(
@@ -1495,7 +1505,7 @@ def fuser_forward(
             fc1_activation_kwargs["norm_const_tensor"] = fc1_norm_const_tensor
             fc1_activation_kwargs["discrete_col_sfd"] = not use_nvfp4
             if supports_single_group_runtime_offsets:
-                fc1_activation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
+                fc1_activation_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group
         if self._pass_geglu_runtime_params:
             fc1_activation_kwargs.update(
                 linear_offset=self._cudnn_linear_offset,
@@ -1514,7 +1524,7 @@ def fuser_forward(
         if fc1_op.single_grouped_weight:
             # Clone and swizzle scales for GEMM.
             fc1_weight_for_gemm = grouped_fc1_weight.copy()
-            use_single_group_weight_swizzle = num_groups == 1
+            use_single_group_weight_swizzle = use_dense_single_group
             if use_single_group_weight_swizzle:
                 fc1_weight_single = _single_quantized_tensor_from_grouped(fc1_weight_for_gemm)
                 fc1_weight_single._columnwise_data = None
@@ -1551,7 +1561,7 @@ def fuser_forward(
             fc1_activation_kwargs["b_tensor"] = fc1_w_data
             fc1_activation_kwargs["sfb_tensor"] = fc1_w_scales
         else:
-            use_single_discrete_weight = num_groups == 1
+            use_single_discrete_weight = use_dense_single_group
             if use_single_discrete_weight:
                 fc1_weight_single = grouped_fc1_weight[0]
                 original_rowwise_scale = fc1_weight_single._rowwise_scale_inv
@@ -1653,6 +1663,7 @@ def fuser_forward(
                     fc1_kernel_out["amax_tensor"].view(-1),
                     fc1_kernel_out["post_rht_amax_tensor"].view(-1),
                     tensor_offsets=fc2_x_tensor_offsets,
+                    use_dense_single_group=use_dense_single_group,
                 )
             else:
                 grouped_fc2_x = _group_quantize_for_grouped_mlp(
@@ -1661,11 +1672,12 @@ def fuser_forward(
                     num_groups,
                     split_sizes,
                     tensor_offsets=fc2_x_tensor_offsets,
+                    use_dense_single_group=use_dense_single_group,
                 )
 
             fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device)
             if (
-                num_groups == 1
+                use_dense_single_group
                 and grouped_fc2_x.columnwise_data is not None
                 and grouped_fc2_x.columnwise_scale_inv is not None
             ):
@@ -1723,7 +1735,7 @@ def fuser_forward(
                 with_gemm_swizzled_scales=True,
             )
 
-            use_single_group_dense_fc2 = num_groups == 1
+            use_single_group_dense_fc2 = use_dense_single_group
             fc2_out_buf = validate_or_alloc_output(output_buffer, fc2_out_shape, dtype, device)
             if use_single_group_dense_fc2:
                 fc2_out = _single_group_fc2_gemm(
@@ -1764,7 +1776,7 @@ def fuser_forward(
                 }
                 fc2_quant_kernel = self.grouped_gemm_quant_kernel()
                 if supports_single_group_runtime_offsets:
-                    fc2_quant_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
+                    fc2_quant_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group
 
                 if fc2_op.single_grouped_weight:
                     # Clone and swizzle scales for GEMM
@@ -1935,6 +1947,9 @@ def fuser_backward(
             grad_output = grad_output.reshape(-1, fc2_weight_shape[0])
         out_shape = list(grad_output.size())
         num_groups = fc1_op.num_groups
+        use_dense_single_group = num_groups == 1 and (
+            _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op))
+        )
         fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0
         fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0
         device = fc1_weight_param.device
@@ -2041,6 +2056,7 @@ def fuser_backward(
                     num_groups,
                     split_sizes,
                     tensor_offsets=fc2_out_tensor_offsets,
+                    use_dense_single_group=use_dense_single_group,
                 )
 
         use_nvfp4 = (
@@ -2194,7 +2210,7 @@ def fuser_backward(
             # Never passed to a wrapper that would reject it -- the check above raises first.
             fc2_dactivation_kwargs["deterministic"] = True
         if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)):
-            fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
+            fc2_dactivation_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group
         if self._cudnn_dact_func is not None:
             fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor
             fc2_dactivation_kwargs["act_func"] = self._cudnn_dact_func
@@ -2248,7 +2264,7 @@ def fuser_backward(
             fc2_dactivation_kwargs["b_tensor"] = fc2_w_data
             fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales
         else:
-            use_single_discrete_weight = num_groups == 1
+            use_single_discrete_weight = use_dense_single_group
             if use_single_discrete_weight:
                 fc2_weight_single = grouped_fc2_weight[0]
                 original_rowwise_data = fc2_weight_single._rowwise_data
@@ -2353,6 +2369,7 @@ def fuser_backward(
                     num_groups,
                     split_sizes,
                     tensor_offsets=fc2_x_tensor_offsets,
+                    use_dense_single_group=use_dense_single_group,
                 )
             else:
                 sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor")
@@ -2430,6 +2447,7 @@ def fuser_backward(
                 num_groups,
                 split_sizes,
                 tensor_offsets=fc1_dy_tensor_offsets,
+                use_dense_single_group=use_dense_single_group,
             )
         else:
             grouped_fc1_dy = GroupedTensor(
@@ -2466,6 +2484,7 @@ def fuser_backward(
             scale_view_dtype=scale_view_dtype,
             sf_vec_size=sf_vec_size,
             offsets=split_points,
+            use_dense_single_group=use_dense_single_group,
         )
 
         # Clear FC2 input tensor if possible
@@ -2491,7 +2510,7 @@ def fuser_backward(
             if is_distributed_weight(fc1_leader):
                 grouped_fc1_weight = materialize_weight_for_backward(fc1_leader)
 
-            use_single_group_dense_dgrad = num_groups == 1
+            use_single_group_dense_dgrad = use_dense_single_group
             if use_single_group_dense_dgrad:
                 grad_input = validate_or_alloc_output(grad_input_buffer, in_shape, dtype, device)
                 _single_group_dgrad_gemm(
@@ -2545,7 +2564,7 @@ def fuser_backward(
                 }
                 fc1_dgrad_kernel = self.grouped_gemm_quant_kernel()
                 if _cudnn_frontend_supports_single_group_runtime_offsets(type(activation_op)):
-                    fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = num_groups == 1
+                    fc1_dgrad_kwargs["use_single_group_runtime_offsets"] = use_dense_single_group
 
                 if fc1_op.single_grouped_weight:
                     # Clone and swizzle scales for GEMM
@@ -2622,6 +2641,7 @@ def fuser_backward(
             scale_view_dtype=scale_view_dtype,
             sf_vec_size=sf_vec_size,
             offsets=split_points,
+            use_dense_single_group=use_dense_single_group,
         )
 
         # Clear FC1 input tensor if possible

From 7ba2f9f17a9016c8b991c0702870b49a7d0f35f8 Mon Sep 17 00:00:00 2001
From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com>
Date: Wed, 9 Sep 2026 11:10:32 +0200
Subject: [PATCH 22/27] [Common] Fix Grouped MXFP8 work mapping and TMA
 synchronization (#3483)

* Added device check that first dim is a multiple of 128

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fixed tensormap release-acquire protocol

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

* Fixed error description

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fix grouped quantize TMA descriptor updater launch

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

* Fix shared-memory race in Grouped MXFP8 direct mapper

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>

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

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

---------

Signed-off-by: Oleg Goncharov <ogoncharov@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
 .../common/cast/core/grouped_tma.cuh          | 67 ++++++++++++-------
 .../cast/mxfp8/group_dequantize_mxfp8.cuh     |  5 +-
 .../cast/mxfp8/group_quantize_mxfp8.cuh       | 60 ++++++++++++-----
 3 files changed, 86 insertions(+), 46 deletions(-)

diff --git a/transformer_engine/common/cast/core/grouped_tma.cuh b/transformer_engine/common/cast/core/grouped_tma.cuh
index 8603fd1fd2..a919b0a5ee 100644
--- a/transformer_engine/common/cast/core/grouped_tma.cuh
+++ b/transformer_engine/common/cast/core/grouped_tma.cuh
@@ -53,45 +53,59 @@ inline bool dimensions_supported_by_TMA(const Tensor *const t) {
   return cols % alignment_requirement == 0;
 }
 
-// Copies the base tensor map to shmem, modifies the copy, stores the modified tensor map at index
+// Copy the base tensor map to shared memory, modify it, and publish it to global memory. This
+// function must be called convergently by every thread in a one-warp CTA.
 __device__ __forceinline__ void modify_base_tensor_map(const CUtensorMap base_tensor_map,
                                                        CUtensorMap *global_tensor_map,
                                                        const uintptr_t global_data_ptr,
                                                        const size_t global_dim_Y,
                                                        const size_t global_dim_X,
                                                        const size_t data_type_size_bytes) {
-  __shared__ CUtensorMap shared_tensor_map;
-  shared_tensor_map = base_tensor_map;  // Copy the base tensor map into shmem
+  __shared__ alignas(128) CUtensorMap shared_tensor_map;
   constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY;
   if constexpr (is_blackwell) {
-    const size_t global_stride_bytes = global_dim_X * data_type_size_bytes;
-    if (global_stride_bytes % TMA_GMEM_ALIGNMENT != 0) {
-      NVTE_DEVICE_ERROR("Shape not supported. Data stride must be 16B aligned.");
-    }
-    if (global_data_ptr % TMA_GMEM_ALIGNMENT != 0) {
-      NVTE_DEVICE_ERROR("Tensor data pointer must be 16B aligned");
+    const uint32_t shared_tensor_map_ptr = __cvta_generic_to_shared(&shared_tensor_map);
+    if (threadIdx.x == 0) {
+      shared_tensor_map = base_tensor_map;
+
+      const size_t global_stride_bytes = global_dim_X * data_type_size_bytes;
+      if (global_stride_bytes % TMA_GMEM_ALIGNMENT != 0) {
+        NVTE_DEVICE_ERROR("Shape not supported. Data stride must be 16B aligned.");
+      }
+      if (global_data_ptr % TMA_GMEM_ALIGNMENT != 0) {
+        NVTE_DEVICE_ERROR("Tensor data pointer must be 16B aligned");
+      }
+
+      asm volatile(
+          "tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;\n\t"
+          "tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 1, %2;\n\t"
+          "tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %3;\n\t"
+          "tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %4;\n"
+          :
+          : "r"(shared_tensor_map_ptr), "l"(global_data_ptr),
+            "r"(static_cast<uint32_t>(global_dim_Y)), "r"(static_cast<uint32_t>(global_dim_X)),
+            "l"(static_cast<uint64_t>(global_stride_bytes))
+          : "memory");
     }
 
+    // tensormap.cp_fenceproxy is a warp-collective instruction. Besides copying the complete
+    // 128-byte descriptor, its GPU-scope release makes the update visible to tensor-map proxy
+    // accesses that perform a matching acquire in a consumer CTA.
+    __syncwarp();
+    const uintptr_t global_tensor_map_ptr = reinterpret_cast<uintptr_t>(global_tensor_map);
     asm volatile(
-        "{\n\t"
-        ".reg.b64 tensor_map_ptr; \n\t"
-        "mov.b64 tensor_map_ptr, %0; \n\t"
-        "tensormap.replace.tile.global_address.b1024.b64  [tensor_map_ptr], %1; \n\t"
-        "tensormap.replace.tile.global_dim.b1024.b32  [tensor_map_ptr], 1, %2; \n\t"  // DIM Y
-        "tensormap.replace.tile.global_dim.b1024.b32  [tensor_map_ptr], 0, %3; \n\t"  // DIM X
-        "tensormap.replace.tile.global_stride.b1024.b64  [tensor_map_ptr], 0, %4; \n"
-        "}\n" ::"l"(reinterpret_cast<uintptr_t>(&shared_tensor_map)),
-        "l"(global_data_ptr), "r"(static_cast<uint32_t>(global_dim_Y)),
-        "r"(static_cast<uint32_t>(global_dim_X)), "l"(static_cast<uint64_t>(global_stride_bytes))
+        "tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned "
+        "[%0], [%1], 128;"
+        :
+        : "l"(global_tensor_map_ptr), "r"(shared_tensor_map_ptr)
         : "memory");
-    *global_tensor_map = shared_tensor_map;
   } else {
     NVTE_DEVICE_ERROR("tensormap.replace is architecture-specific. ");
   }
 }
 
 template <typename IType, typename OType>
-__global__ void __launch_bounds__(1)
+__global__ void __launch_bounds__(THREADS_PER_WARP)
     update_tma_descriptors(const __grid_constant__ CUtensorMap base_tensor_map_input,
                            const __grid_constant__ CUtensorMap base_tensor_map_act_input,
                            const __grid_constant__ CUtensorMap base_tensor_map_output_rowwise,
@@ -112,9 +126,11 @@ __global__ void __launch_bounds__(1)
   const size_t cols = get_tensor_cols_num(tensor_id, shape_rep, last_logical_dim, last_dims_ptr);
 
   const size_t offset_elts = offsets_ptr[tensor_id];
-  g_tensor_maps.rows[tensor_id] = rows;
-  g_tensor_maps.cols[tensor_id] = cols;
-  g_tensor_maps.offsets[tensor_id] = offset_elts;
+  if (threadIdx.x == 0) {
+    g_tensor_maps.rows[tensor_id] = rows;
+    g_tensor_maps.cols[tensor_id] = cols;
+    g_tensor_maps.offsets[tensor_id] = offset_elts;
+  }
 
   // Zero-sized groups: skip TMA descriptor update. The main kernel already returns
   // early for rows==0 or cols==0, but creating a TMA descriptor with a zero dimension
@@ -156,7 +172,8 @@ __global__ void __launch_bounds__(1)
 
 __device__ __forceinline__ void fence_acquire_tensormap(const CUtensorMap *tensor_map) {
 #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
-  asm volatile("fence.proxy.tensormap::generic.acquire.cta [%0], 128;" ::"l"(tensor_map));
+  // The descriptor updater and consumer execute in different CTAs, so CTA scope is insufficient.
+  asm volatile("fence.proxy.tensormap::generic.acquire.gpu [%0], 128;" ::"l"(tensor_map));
 #else
   NVTE_DEVICE_ERROR("fence_acquire_tensormap is only supported on SM 9.0+.");
 #endif  // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)
diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
index dad8d18d6f..141efabbaf 100644
--- a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
@@ -89,7 +89,6 @@ __global__ void update_tma_descriptors(const __grid_constant__ CUtensorMap base_
                                        const int64_t *const __restrict__ offsets_ptr,
                                        const int64_t *const __restrict__ first_dims_ptr,
                                        const int64_t *const __restrict__ last_dims_ptr) {
-  const bool leading_thread = (threadIdx.x == 0);
   const size_t tensor_id = blockIdx.x;
 
   const size_t rows =
@@ -105,7 +104,7 @@ __global__ void update_tma_descriptors(const __grid_constant__ CUtensorMap base_
     return;
   }
 
-  if (leading_thread && (tensor_id < num_tensors)) {
+  if (tensor_id < num_tensors) {
     {
       const uintptr_t global_data_ptr = reinterpret_cast<uintptr_t>(input_data_ptr + offset_elts);
       modify_base_tensor_map(base_tensor_map_input, &g_tensor_maps_input[tensor_id],
@@ -469,7 +468,7 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output,
             const IType *const input_dptr = reinterpret_cast<const IType *>(input_data.dptr);
             OType *const output_dptr = reinterpret_cast<OType *>(output->data.dptr);
 
-            update_tma_descriptors<IType, OType><<<num_tensors, 32, 0, stream>>>(
+            update_tma_descriptors<IType, OType><<<num_tensors, THREADS_PER_WARP, 0, stream>>>(
                 tensor_map_input, tensor_map_output, input_dptr, output_dptr, shape_rep,
                 num_tensors, first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr,
                 last_dims_ptr);
diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh
index b0383d95f3..bd8c3052f6 100644
--- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh
@@ -133,6 +133,13 @@ struct LaunchConfig {
   dim3 grid;
 };
 
+struct alignas(8) DirectVaryingFirstMapperStorage {
+  size_t active_elements;
+  size_t tensor_id;
+  size_t rows;
+  size_t tensor_start_offset;
+};
+
 template <typename CastTraits>
 LaunchConfig get_launch_config(const size_t first_logical_dim, const size_t last_logical_dim,
                                const size_t elts_total, const size_t num_tensors) {
@@ -659,6 +666,22 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_
   constexpr bool is_single_tensor = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM);
 
   const bool leading_thread = (threadIdx.x == 0);
+  // Keep mapper metadata separate from dynamic shared memory. The latter is the destination of
+  // asynchronous TMA loads and may be overwritten before every warp has consumed the metadata.
+  __shared__ DirectVaryingFirstMapperStorage direct_mapper_storage;
+
+  if constexpr (use_direct_varying_first_mapper) {
+    static_assert(CastTraits::THREADS_PER_CHUNK >= MAX_SUPPORTED_TENSOR_DESCRIPTORS,
+                  "The first CTA must have enough threads to validate every tensor.");
+    // The direct mapper relies on every tensor boundary being TILE_DIM_Y-aligned. The legacy
+    // mapper validated this while decoding each job, but the direct path no longer calls it.
+    // Validate all per-tensor row counts once, using the first CTA. num_tensors is bounded by
+    // MAX_SUPPORTED_TENSOR_DESCRIPTORS (64), which is smaller than the CTA size.
+    if (blockIdx.x == 0 && threadIdx.x < num_tensors) {
+      get_tensor_rows_num<ShapeRepresentation::VARYING_FIRST_DIM>(threadIdx.x, first_logical_dim,
+                                                                  first_dims_ptr, num_tensors);
+    }
+  }
 
   // Decode the linear direct-mapper grid once per CTA. Valid CUDA grid extents fit in uint, which
   // also keeps this one-time coordinate calculation in 32-bit arithmetic.
@@ -719,26 +742,26 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_
 
   if constexpr (use_direct_varying_first_mapper) {
     // logical_shape may describe graph-safe capacity beyond the active tensors. Resolve the
-    // active tail once per CTA and reject it before initializing TMA barriers. Reuse the same
-    // temporary storage for the exceptional colwise-swizzled tensor metadata.
-    size_t *const mapper_storage = reinterpret_cast<size_t *>(dshmem);
+    // active tail once per CTA and reject it before initializing TMA barriers. Cache the
+    // exceptional colwise-swizzled tensor metadata in dedicated static shared memory.
     const size_t block_offset_Y = direct_block_id_Y * CHUNK_DIM_Y;
     const size_t tensor_offset = block_offset_Y * last_logical_dim;
     if (leading_thread) {
       const size_t active_elements = static_cast<size_t>(offsets_ptr[num_tensors]);
-      mapper_storage[0] = active_elements;
+      direct_mapper_storage.active_elements = active_elements;
       if constexpr (WITH_GEMM_SWIZZLED_SCALES && COLWISE_SCALING) {
         if (tensor_offset < active_elements) {
           const size_t mapped_tensor_id =
               find_tensor_from_offsets(offsets_ptr, num_tensors, tensor_offset);
-          mapper_storage[1] = mapped_tensor_id;
-          mapper_storage[2] = static_cast<size_t>(first_dims_ptr[mapped_tensor_id]);
-          mapper_storage[3] = static_cast<size_t>(offsets_ptr[mapped_tensor_id]);
+          direct_mapper_storage.tensor_id = mapped_tensor_id;
+          direct_mapper_storage.rows = static_cast<size_t>(first_dims_ptr[mapped_tensor_id]);
+          direct_mapper_storage.tensor_start_offset =
+              static_cast<size_t>(offsets_ptr[mapped_tensor_id]);
         }
       }
     }
     __syncthreads();
-    if (tensor_offset >= mapper_storage[0]) {
+    if (tensor_offset >= direct_mapper_storage.active_elements) {
       return;
     }
   }
@@ -839,10 +862,10 @@ __global__ void __launch_bounds__(CastTraits::THREADS_PER_CHUNK) group_quantize_
       if constexpr (WITH_GEMM_SWIZZLED_SCALES && COLWISE_SCALING) {
         // Colwise GEMM-swizzled scale indices restart at each tensor and depend on M_i.
         // The leading thread decoded this exceptional metadata before barrier initialization.
-        size_t *const mapper_storage = reinterpret_cast<size_t *>(dshmem);
-        tensor_id = mapper_storage[1];
-        rows = mapper_storage[2];
-        tensor_start_offset = mapper_storage[3];
+        tensor_id = direct_mapper_storage.tensor_id;
+        rows = direct_mapper_storage.rows;
+        tensor_start_offset = direct_mapper_storage.tensor_start_offset;
+        tensor_offset_Y = block_offset_Y - tensor_start_offset / cols;
       }
     } else {
       block_id_Y = current_block_id / fixed_blocks_X;
@@ -1289,12 +1312,13 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations
                               use_colwise_scaling
                                   ? reinterpret_cast<OType *>(output->columnwise_data.dptr)
                                   : nullptr;
-                          update_tma_descriptors<IType, OType><<<num_tensors, 1, 0, stream>>>(
-                              tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise,
-                              tensor_map_output_colwise, input_dptr, act_input_dptr,
-                              output_rowwise_dptr, output_colwise_dptr, shape_rep, num_tensors,
-                              first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr,
-                              last_dims_ptr, use_rowwise_scaling, use_colwise_scaling, IS_DACT);
+                          update_tma_descriptors<IType, OType>
+                              <<<num_tensors, THREADS_PER_WARP, 0, stream>>>(
+                                  tensor_map_input, tensor_map_act_input, tensor_map_output_rowwise,
+                                  tensor_map_output_colwise, input_dptr, act_input_dptr,
+                                  output_rowwise_dptr, output_colwise_dptr, shape_rep, num_tensors,
+                                  first_logical_dim, last_logical_dim, offsets_ptr, first_dims_ptr,
+                                  last_dims_ptr, use_rowwise_scaling, use_colwise_scaling, IS_DACT);
                         }
 
                         TRANSFORMER_ENGINE_SWITCH_CONDITION(

From 1b178bc14a1fd9ef200b5fabe17345eb3b936c03 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 9 Sep 2026 15:37:36 +0200
Subject: [PATCH 23/27] [PyTorch] Derive an op's compile declarations; gate
 groups to one op

num_grad_inputs leaves BasicOperation: on the autograd-free path it only
validated the length of backward_compute's result, and a default of 1 was
wrong for every op with a parameter. register_custom_op takes it as an
optional check. fwd_kwarg_names is read off resolve_fwd_args's keyword-only
parameters instead of being declared twice.

The fuser gate now matches the title: a group of several operations runs the
eager implementations. Multi-op groups without fusions were slipping through.

The saved_tensors reset in backward no longer skips compile; the contexts are
copies local to that subgraph, so the write is allowed, and gating it only
suggested a constraint that is not there.

Drop test_te_ops_setup_context_saves_parameter.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 28 ++++++-------------
 .../pytorch/dynamo/custom_op.py               | 19 +++++++------
 transformer_engine/pytorch/ops/fuser.py       |  7 ++---
 transformer_engine/pytorch/ops/op.py          | 23 +++++++++++----
 4 files changed, 40 insertions(+), 37 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 0004d86508..4342e1a330 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2381,7 +2381,6 @@ class _ScaleOp(BasicOperation):
 
     fwd_args_type = _ScaleFwdArgs
     bwd_args_type = _ScaleBwdArgs
-    num_grad_inputs = 2  # grad input, grad scale
 
     def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> None:
         super().__init__()
@@ -2474,8 +2473,6 @@ class _ScaleWithKwargsOp(BasicOperation):
 
     fwd_args_type = _ScaleKwargsFwdArgs
     bwd_args_type = _ScaleKwargsBwdArgs
-    num_grad_inputs = 2  # grad input, grad scale
-    fwd_kwarg_names = ("extra_scale", "offset")
 
     def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> None:
         super().__init__()
@@ -2596,6 +2593,15 @@ def test_te_ops_single_op_group_compiles():
     _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp()), base)
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_te_ops_multi_op_group_uses_eager_implementations():
+    """A group of several operations is gated onto the eager implementations."""
+    torch._dynamo.reset()
+    base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
+    with pytest.warns(UserWarning, match="several operations"):
+        _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp(), _ScaleOp()), base)
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 def test_te_ops_backward_fusion_uses_eager_implementations():
     """A backward fusion prevents the group from using basic-op custom ops."""
@@ -2610,22 +2616,6 @@ def test_te_ops_backward_fusion_uses_eager_implementations():
         OperationFuser.backward_fusion_functions.remove(_fuse_backward_scale_pair)
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@pytest.mark.parametrize("compile_model", [False, True], ids=["eager", "compiled"])
-def test_te_ops_setup_context_saves_parameter(compile_model):
-    """Backward observes mutation of a tensor used by the forward."""
-    op = _ScaleOp()
-    model = te.ops.Sequential(op)
-    if compile_model:
-        model = torch.compile(model, fullgraph=True)
-    x = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True)
-    y = model(x)
-    with torch.no_grad():
-        op.scale.add_(1)
-    with pytest.raises(RuntimeError, match="modified by an inplace operation"):
-        y.sum().backward()
-
-
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 def test_te_ops_unsupported_group_still_compiles_eagerly():
     """An operation without the compute halves runs its eager implementation.
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index d35958109b..4fa0299f2d 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -830,14 +830,16 @@ def _pack_fwd_result(result: Any) -> List[torch.Tensor]:
     return flat
 
 
-def _pack_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List[torch.Tensor]:
+def _pack_bwd_result(
+    grads: Any, num_grad_inputs: Optional[int], op_qualname: str
+) -> List[torch.Tensor]:
     """Pack a backward-impl return tuple into the op's ``Tensor[]`` payload.
 
-    Each grad occupies exactly one slot (validated against ``num_grad_inputs``);
-    a :class:`TensorSpec` grad is materialized into a single tensor.
+    Each grad occupies exactly one slot (validated against ``num_grad_inputs``
+    when given); a :class:`TensorSpec` grad is materialized into a single tensor.
     """
     grads = list(grads)
-    if len(grads) != num_grad_inputs:
+    if num_grad_inputs is not None and len(grads) != num_grad_inputs:
         raise RuntimeError(
             f"{op_qualname} expected bwd_impl to return {num_grad_inputs} grads "
             f"(one per input_tensors_for_grad entry), got {len(grads)}"
@@ -1260,7 +1262,7 @@ def _register_backward_op(
     arg_type: type,
     impl: Callable[[Any], Any],
     fake_impl: Callable[[Any], Any],
-    num_grad_inputs: int,
+    num_grad_inputs: Optional[int],
 ) -> _RegisteredOp:
     # Pass-through body: a subclass input reaches the base op through the
     # dispatch rule, never through the wrapper body.
@@ -1298,7 +1300,7 @@ def register_custom_op(
     bwd_arg_type: type,
     bwd_impl: Callable[[Any], Any],
     bwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
-    num_grad_inputs: int,
+    num_grad_inputs: Optional[int] = None,
 ) -> Optional[Tuple[Callable[[Any], Any], Callable[[Any], Any]]]:
     """Register an op's forward and backward as two independent custom ops.
 
@@ -1315,7 +1317,8 @@ def register_custom_op(
 
     * ``fwd_impl(fwd_args) -> (output, aux)`` -- ``aux`` is a tuple of fresh tensors
     * ``fwd_fake_impl`` -- its data-free twin over :class:`TensorSpec`
-    * ``bwd_impl(bwd_args) -> tuple`` of ``num_grad_inputs`` gradients
+    * ``bwd_impl(bwd_args) -> tuple`` of gradients (``num_grad_inputs`` of them,
+      if given)
     * ``bwd_fake_impl`` -- its data-free twin
 
     Returns ``(forward_fn, backward_fn)``:
@@ -1355,7 +1358,7 @@ def _register_custom_op_impl(
     bwd_arg_type: type,
     bwd_impl: Callable[[Any], Any],
     bwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
-    num_grad_inputs: int,
+    num_grad_inputs: Optional[int],
 ) -> Tuple[Callable[[Any], Any], Callable[[Any], Any]]:
     """Body of :func:`register_custom_op`; see it for semantics."""
 
diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py
index 87a3f18d96..e51160a15c 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -374,10 +374,7 @@ def backward(
                 )
             for idx, dparams in zip(basic_op_idxs, fused_op_grad_params):
                 grad_params[idx] = dparams
-                # Dropping the reference frees the activation early; on the
-                # compiled path the graph owns that lifetime instead.
-                if not torch.compiler.is_compiling():
-                    basic_op_ctxs[idx].saved_tensors = None
+                basic_op_ctxs[idx].saved_tensors = None
             for idx, dxs in zip(basic_op_idxs, fused_op_grad_extra_inputs):
                 grad_extra_inputs[idx] = dxs
                 for input_idx, grad in enumerate(dxs):
@@ -745,6 +742,8 @@ def _custom_ops_unsupported_reason(
                 for idx, (op, basic_op_idxs) in enumerate(ops)
             ):
                 return f"a {mode} fusion"
+        if self._num_basic_ops != 1:
+            return "a group of several operations"
         for op, kwargs in zip(self._basic_ops, basic_op_kwargs, strict=True):
             # A kwarg an operation declares is resolved into its args container
             # like any other config. Anything else -- notably the preallocated
diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py
index e353d860ee..f2b9377cea 100644
--- a/transformer_engine/pytorch/ops/op.py
+++ b/transformer_engine/pytorch/ops/op.py
@@ -8,6 +8,7 @@
 import abc
 from collections.abc import Iterable, Sequence
 import dataclasses
+import inspect
 import pickle
 from typing import Any, Callable, Optional
 
@@ -24,6 +25,10 @@
 from ..tensor import Quantizer
 from ..dynamo import is_value_opaque_quantizer, register_custom_op
 
+_FIXED_RESOLVE_FWD_KWARGS = frozenset(
+    ("requires_grad", "prev_op_grad_output_quantizer", "next_op_input_quantizer")
+)
+
 
 @dataclasses.dataclass
 class OperationContext:
@@ -193,10 +198,9 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta):
     # op_backward, so no operation writes that plumbing itself.
     fwd_args_type: Optional[type] = None
     bwd_args_type: Optional[type] = None
-    # Gradients returned by backward_compute: the input's, then any parameters'.
-    num_grad_inputs: int = 1
-    # Forward kwargs this operation accepts, resolved into fwd_args_type like
-    # any other config. A kwarg carries no gradient and must not be mutated.
+    # Forward kwargs this operation accepts: the keyword-only parameters of its
+    # resolve_fwd_args beyond the fixed ones. A kwarg carries no gradient and
+    # must not be mutated.
     fwd_kwarg_names: tuple[str, ...] = ()
     # (forward_fn, backward_fn) pair, or None if the operation cannot be compiled.
     compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None
@@ -215,6 +219,14 @@ def __init_subclass__(cls, **kwargs) -> None:
             # framework's actual requirement -- check it where it is declared.
             if not dataclasses.is_dataclass(arg_type):
                 raise TypeError(f"{cls.__name__}.{name} must be a dataclass")
+        params = inspect.signature(cls.resolve_fwd_args).parameters
+        if any(p.kind is p.VAR_KEYWORD for p in params.values()):
+            raise TypeError(f"{cls.__name__}.resolve_fwd_args must name its keyword arguments")
+        cls.fwd_kwarg_names = tuple(
+            name
+            for name, p in params.items()
+            if p.kind is p.KEYWORD_ONLY and name not in _FIXED_RESOLVE_FWD_KWARGS
+        )
         # One registration per class. The compute halves are bound here, so a
         # subclass that only swaps kernels (the activations) still gets its own
         # op without repeating any of this.
@@ -226,7 +238,6 @@ def __init_subclass__(cls, **kwargs) -> None:
             bwd_arg_type=cls.bwd_args_type,
             bwd_impl=cls.backward_compute,
             bwd_fake_impl=cls.backward_fake,
-            num_grad_inputs=cls.num_grad_inputs,
         )
 
     def __init__(self) -> None:
@@ -344,7 +355,7 @@ def forward_fake(cls, args: Any) -> tuple[Any, tuple]:
 
     @classmethod
     def backward_compute(cls, args: Any) -> tuple:
-        """Pure backward: ``num_grad_inputs`` gradients."""
+        """Pure backward: the input's gradient, then the parameters'."""
         raise NotImplementedError
 
     @classmethod

From e1b33b0deb74d486827c2489cc84701df5cef502 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 9 Sep 2026 15:51:37 +0200
Subject: [PATCH 24/27] [PyTorch] Tidy BasicOperation's compile plumbing

op_forward / compiled_op_forward and op_backward / compiled_op_backward
share one body each, parametrized by the compute callable. Registration
moves into _register_compile_ops. The base resolve_fwd_args no longer
takes **kwargs, which the registration forbids on subclasses. Redundant
checks and long comments removed.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/ops/op.py | 304 ++++++++++++---------------
 1 file changed, 130 insertions(+), 174 deletions(-)

diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py
index f2b9377cea..0bb3eb7ae2 100644
--- a/transformer_engine/pytorch/ops/op.py
+++ b/transformer_engine/pytorch/ops/op.py
@@ -25,7 +25,8 @@
 from ..tensor import Quantizer
 from ..dynamo import is_value_opaque_quantizer, register_custom_op
 
-_FIXED_RESOLVE_FWD_KWARGS = frozenset(
+# Keyword-only parameters every resolve_fwd_args takes; the rest are forward kwargs.
+_RESOLVE_FWD_ARGS_PARAMS = frozenset(
     ("requires_grad", "prev_op_grad_output_quantizer", "next_op_input_quantizer")
 )
 
@@ -192,32 +193,25 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta):
     # Number of extra tensor outputs
     num_extra_outputs: int = 0
 
-    # torch.compile support. An operation opts in by declaring the two arg
-    # containers and implementing the four compute classmethods below; the base
-    # class then registers its custom ops and drives them from op_forward /
-    # op_backward, so no operation writes that plumbing itself.
+    # torch.compile support: an operation declares its argument containers and
+    # implements the compute halves (see op_forward); custom ops are registered
+    # once per class.
     fwd_args_type: Optional[type] = None
     bwd_args_type: Optional[type] = None
-    # Forward kwargs this operation accepts: the keyword-only parameters of its
-    # resolve_fwd_args beyond the fixed ones. A kwarg carries no gradient and
-    # must not be mutated.
+    # Forward kwargs accepted by resolve_fwd_args. Read-only, no gradient.
     fwd_kwarg_names: tuple[str, ...] = ()
-    # (forward_fn, backward_fn) pair, or None if the operation cannot be compiled.
+    # (forward_fn, backward_fn), or None if the operation cannot be compiled.
     compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None
 
     def __init_subclass__(cls, **kwargs) -> None:
         super().__init_subclass__(**kwargs)
-        if cls.fwd_args_type is None or cls.bwd_args_type is None:
-            return
-        if getattr(cls.forward_compute, "__isabstractmethod__", False):
-            return
-        for name, arg_type in (
-            ("fwd_args_type", cls.fwd_args_type),
-            ("bwd_args_type", cls.bwd_args_type),
-        ):
-            # The op schema is built from the container's fields, so this is the
-            # framework's actual requirement -- check it where it is declared.
-            if not dataclasses.is_dataclass(arg_type):
+        if cls.fwd_args_type is not None and cls.bwd_args_type is not None:
+            cls._register_compile_ops()
+
+    @classmethod
+    def _register_compile_ops(cls) -> None:
+        for name in ("fwd_args_type", "bwd_args_type"):
+            if not dataclasses.is_dataclass(getattr(cls, name)):
                 raise TypeError(f"{cls.__name__}.{name} must be a dataclass")
         params = inspect.signature(cls.resolve_fwd_args).parameters
         if any(p.kind is p.VAR_KEYWORD for p in params.values()):
@@ -225,11 +219,8 @@ def __init_subclass__(cls, **kwargs) -> None:
         cls.fwd_kwarg_names = tuple(
             name
             for name, p in params.items()
-            if p.kind is p.KEYWORD_ONLY and name not in _FIXED_RESOLVE_FWD_KWARGS
+            if p.kind is p.KEYWORD_ONLY and name not in _RESOLVE_FWD_ARGS_PARAMS
         )
-        # One registration per class. The compute halves are bound here, so a
-        # subclass that only swaps kernels (the activations) still gets its own
-        # op without repeating any of this.
         cls.compile_ops = register_custom_op(
             op_name=cls.__name__.lower(),
             fwd_arg_type=cls.fwd_args_type,
@@ -328,91 +319,6 @@ def set_extra_output_channel(
         self._extra_output_to_caller[index] = output_to_caller
         return self
 
-    # ------------------------------------------------------------------ #
-    # Compute halves. Classmethods, not free functions: they belong to the
-    # operation, and binding to the class is what lets a family of operations
-    # share one implementation while dispatching to per-class kernels.
-    # ------------------------------------------------------------------ #
-
-    @classmethod
-    def forward_compute(cls, args: Any) -> tuple[Any, tuple]:
-        """Forward computation over explicit arguments.
-
-        Takes everything through ``args``; must not read ``self`` or global
-        state, both of which are invisible to the compiler at this point.
-        """
-        raise NotImplementedError
-
-    @classmethod
-    def forward_fake(cls, args: Any) -> tuple[Any, tuple]:
-        """Allocation-free twin of :meth:`forward_compute` over ``TensorSpec``.
-
-        Runs as a meta kernel, outside the traced frame, and more than once per
-        compile, so it must be a pure function of ``args`` -- a read of global
-        state here is unguarded and can silently disagree with the real impl.
-        """
-        raise NotImplementedError
-
-    @classmethod
-    def backward_compute(cls, args: Any) -> tuple:
-        """Pure backward: the input's gradient, then the parameters'."""
-        raise NotImplementedError
-
-    @classmethod
-    def backward_fake(cls, args: Any) -> tuple:
-        """Allocation-free twin of :meth:`backward_compute`."""
-        raise NotImplementedError
-
-    def compile_unsupported_reason(self) -> Optional[str]:
-        """Why this operation cannot go through its custom op, or ``None``.
-
-        Asked per operation, but acted on per fuser group: a pipeline compiles
-        as a whole, so one unsupported operation sends the whole group to eager.
-        Recipe-level limits are not checked here -- they belong to whoever reads
-        the recipe, which is the fuser.
-        """
-        if self.compile_ops is None:
-            return f"{self.__class__.__name__} without compute halves"
-        for mode in ("forward", "backward"):
-            for index in range(self.num_quantizers(mode)):
-                quantizer = self.get_quantizer(mode, index)
-                if quantizer is not None and not is_value_opaque_quantizer(quantizer):
-                    # Delayed scaling holds live scale/amax tensors, so its
-                    # quantizer cannot be specialized on and would be baked into
-                    # the graph as a stale constant.
-                    return (
-                        f"{type(quantizer).__name__} (not a torch.compile value-opaque quantizer)"
-                    )
-        return None
-
-    def resolve_fwd_args(
-        self,
-        input_: torch.Tensor,
-        *,
-        requires_grad: bool,
-        prev_op_grad_output_quantizer: Optional[Quantizer] = None,
-        next_op_input_quantizer: Optional[Quantizer] = None,
-        **kwargs: Any,
-    ) -> Any:
-        """Gather the forward's inputs into a flat, ``self``-free container.
-
-        This is where module config and global state are read, so it belongs in
-        the traced region where Dynamo guards those reads -- never inside the
-        custom op. ``kwargs`` are the caller's forward kwargs, restricted to
-        ``fwd_kwarg_names``; an operation declaring them supplies their defaults
-        here, since a kwarg may be absent.
-        """
-        raise NotImplementedError
-
-    def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> Any:
-        """Rebuild the backward's inputs from the forward's saved state."""
-        raise NotImplementedError
-
-    def setup_context(self, ctx: OperationContext, args: Any, aux: tuple) -> None:
-        """Prepare backward state from the original arguments and fresh auxiliary tensors."""
-        del args
-        ctx.save_for_backward(*aux)
-
     @property
     def is_fused_op(self) -> bool:
         return False
@@ -658,9 +564,8 @@ def op_forward(
     ) -> torch.Tensor:
         """Forward pass
 
-        Operations that declare the compute halves inherit this: it resolves the
-        arguments, runs the forward, and records what the backward will need. The
-        rest override it.
+        Operations that implement the compute halves inherit this;
+        the rest override it.
 
         Parameters
         ----------
@@ -679,102 +584,153 @@ def op_forward(
             Output tensor
 
         """
-        if self.fwd_args_type is None:
-            raise NotImplementedError(
-                f"{self.__class__.__name__} implements neither op_forward nor the compute halves"
-            )
-        unsupported = sorted(name for name in kwargs if name not in self.fwd_kwarg_names)
-        if unsupported:
-            raise ValueError(
-                f"{self.__class__.__name__} forward does not accept keyword arguments {unsupported}"
-            )
-        args = self.resolve_fwd_args(
+        return self._forward(
+            self.forward_compute,
+            ctx,
             input_,
-            requires_grad=ctx.requires_grad,
             prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
             next_op_input_quantizer=next_op_input_quantizer,
             **kwargs,
         )
-        output, aux = self.forward_compute(args)
-        if ctx.requires_grad:
-            self.setup_context(ctx, args, aux)
-        return output
+
+    def op_backward(
+        self,
+        ctx: OperationContext,
+        grad_output: torch.Tensor,
+    ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]:
+        """Backward pass
+
+        Operations that implement the compute halves inherit this;
+        the rest override it.
+
+        Parameters
+        ----------
+        ctx: OperationContext
+            Context to coordinate between forward and backward passes
+        grad_output: torch.Tensor
+            Loss gradient w.r.t. operation output
+
+        Returns
+        -------
+        torch.Tensor
+            Loss gradient w.r.t. operation input
+        Iterable of torch.Tensor:
+            Loss gradients w.r.t. parameters
+
+        """
+        return self._backward(self.backward_compute, ctx, grad_output)
 
     def compiled_op_forward(
         self,
         ctx: OperationContext,
         input_: torch.Tensor,
         *,
-        prev_op_grad_output_quantizer: Optional[Quantizer],
-        next_op_input_quantizer: Optional[Quantizer],
+        prev_op_grad_output_quantizer: Optional[Quantizer] = None,
+        next_op_input_quantizer: Optional[Quantizer] = None,
         **kwargs: Any,
     ) -> torch.Tensor:
-        """:meth:`op_forward` routed through this operation's custom op.
-
-        Same bookkeeping, but the computation crosses an op boundary so Dynamo
-        sees one graph node instead of tracing into the kernels. ``kwargs`` are
-        not validated here -- the fuser's gate already rejected a group whose
-        kwargs an operation does not declare.
-        """
-        args = self.resolve_fwd_args(
+        """:meth:`op_forward` through this operation's custom op."""
+        return self._forward(
+            self.compile_ops[0],
+            ctx,
             input_,
-            requires_grad=ctx.requires_grad,
             prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
             next_op_input_quantizer=next_op_input_quantizer,
             **kwargs,
         )
-        output, aux = self.compile_ops[0](args)
-        if ctx.requires_grad:
-            self.setup_context(ctx, args, aux)
-        return output
 
     def compiled_op_backward(
         self,
         ctx: OperationContext,
         grad_output: torch.Tensor,
     ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]:
-        """:meth:`op_backward` routed through this operation's custom op."""
-        grads = self.compile_ops[1](self.resolve_bwd_args(ctx, grad_output))
-        grad_input = grads[0]
-        if grad_input is None:
-            grad_input = grad_output
-        return grad_input, tuple(grads[1:])
+        """:meth:`op_backward` through this operation's custom op."""
+        return self._backward(self.compile_ops[1], ctx, grad_output)
 
-    def op_backward(
+    def _forward(
+        self,
+        compute: Callable[[Any], tuple[torch.Tensor, tuple]],
+        ctx: OperationContext,
+        input_: torch.Tensor,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        args = self.resolve_fwd_args(input_, requires_grad=ctx.requires_grad, **kwargs)
+        output, aux = compute(args)
+        if ctx.requires_grad:
+            self.setup_context(ctx, args, aux)
+        return output
+
+    def _backward(
         self,
+        compute: Callable[[Any], tuple],
         ctx: OperationContext,
         grad_output: torch.Tensor,
     ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]:
-        """Backward pass
+        grads = compute(self.resolve_bwd_args(ctx, grad_output))
+        grad_input = grad_output if grads[0] is None else grads[0]
+        return grad_input, tuple(grads[1:])
 
-        Counterpart to the inherited :meth:`op_forward`.
+    # Compute halves: forward_compute / backward_compute are the eager kernels
+    # and the custom-op bodies; the *_fake twins run on TensorSpec. Neither may
+    # read self, global state, or mutate its arguments. A None grad_input
+    # means "grad_output, unchanged".
 
-        Parameters
-        ----------
-        ctx: OperationContext
-            Context to coordinate between forward and backward passes
-        grad_output: torch.Tensor
-            Loss gradient w.r.t. operation output
+    @classmethod
+    def forward_compute(cls, args: Any) -> tuple[torch.Tensor, tuple]:
+        """Forward over ``args``; returns ``(output, aux)``, aux being fresh tensors."""
+        raise NotImplementedError
 
-        Returns
-        -------
-        torch.Tensor
-            Loss gradient w.r.t. operation input
-        Iterable of torch.Tensor:
-            Loss gradients w.r.t. parameters
+    @classmethod
+    def forward_fake(cls, args: Any) -> tuple[Any, tuple]:
+        """Shape-only twin of :meth:`forward_compute`."""
+        raise NotImplementedError
+
+    @classmethod
+    def backward_compute(cls, args: Any) -> tuple:
+        """Backward over ``args``; returns grad_input, then parameter grads."""
+        raise NotImplementedError
 
+    @classmethod
+    def backward_fake(cls, args: Any) -> tuple:
+        """Shape-only twin of :meth:`backward_compute`."""
+        raise NotImplementedError
+
+    def resolve_fwd_args(
+        self,
+        input_: torch.Tensor,
+        *,
+        requires_grad: bool,
+        prev_op_grad_output_quantizer: Optional[Quantizer] = None,
+        next_op_input_quantizer: Optional[Quantizer] = None,
+    ) -> Any:
+        """Build ``fwd_args_type`` from the input, module state and forward kwargs.
+
+        Forward kwargs are declared as additional keyword-only parameters
+        with defaults.
         """
-        if self.bwd_args_type is None:
-            raise NotImplementedError(
-                f"{self.__class__.__name__} implements neither op_backward nor the compute halves"
-            )
-        grads = self.backward_compute(self.resolve_bwd_args(ctx, grad_output))
-        grad_input = grads[0]
-        if grad_input is None:
-            # "The incoming gradient, unchanged": a custom op may not return one
-            # of its own inputs, so the compute half hands back None instead.
-            grad_input = grad_output
-        return grad_input, tuple(grads[1:])
+        raise NotImplementedError
+
+    def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> Any:
+        """Build ``bwd_args_type`` from the saved context."""
+        raise NotImplementedError
+
+    def setup_context(self, ctx: OperationContext, args: Any, aux: tuple) -> None:
+        """Save backward state from the forward args and the fresh aux tensors."""
+        del args
+        ctx.save_for_backward(*aux)
+
+    def compile_unsupported_reason(self) -> Optional[str]:
+        """Why this operation cannot run through its custom op, or ``None``."""
+        if self.compile_ops is None:
+            return f"{self.__class__.__name__} without compute halves"
+        for mode in ("forward", "backward"):
+            for index in range(self.num_quantizers(mode)):
+                quantizer = self.get_quantizer(mode, index)
+                if quantizer is not None and not is_value_opaque_quantizer(quantizer):
+                    return (
+                        f"{type(quantizer).__name__} (not a torch.compile value-opaque quantizer)"
+                    )
+        return None
 
     def fuser_forward(
         self,

From 048162bf084808ee083ce1853c5acd93ee1a8baf Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 9 Sep 2026 16:03:40 +0200
Subject: [PATCH 25/27] [PyTorch] Name an op's custom-op methods after Linear's

forward_compute / backward_compute become forward_impl / backward_impl,
matching module/linear.py and register_custom_op's parameters. The fixed
resolve_fwd_args parameters are read off the base signature instead of a
separate constant.

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py  | 12 ++++----
 transformer_engine/pytorch/ops/op.py | 43 ++++++++++++----------------
 2 files changed, 24 insertions(+), 31 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 9ded6d9b6a..9252051387 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2375,7 +2375,7 @@ class _ScaleOp(BasicOperation):
     """Test-only operation: multiply by a learnable scalar.
 
     Exists so the fuser's compiled path can be exercised without depending on
-    which real operations happen to declare their compute halves. It is the
+    which real operations happen to have a custom op. It is the
     smallest operation that still has a parameter gradient and a saved tensor.
     """
 
@@ -2387,7 +2387,7 @@ def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16)
         self.scale = torch.nn.Parameter(torch.full((), 2.0, device=device, dtype=dtype))
 
     @classmethod
-    def forward_compute(cls, args):
+    def forward_impl(cls, args):
         return args.input_ * args.scale, ()
 
     @classmethod
@@ -2396,7 +2396,7 @@ def forward_fake(cls, args):
         return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), ()
 
     @classmethod
-    def backward_compute(cls, args):
+    def backward_impl(cls, args):
         dy = args.grad_output
         return dy * args.scale, (dy * args.saved_input).sum()
 
@@ -2479,7 +2479,7 @@ def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16)
         self.scale = torch.nn.Parameter(torch.full((), 2.0, device=device, dtype=dtype))
 
     @classmethod
-    def forward_compute(cls, args):
+    def forward_impl(cls, args):
         offset = args.offset
         if isinstance(offset, QuantizedTensor):
             offset = offset.dequantize()
@@ -2492,7 +2492,7 @@ def forward_fake(cls, args):
         return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), ()
 
     @classmethod
-    def backward_compute(cls, args):
+    def backward_impl(cls, args):
         dy = args.grad_output
         return (
             dy * args.scale * args.extra_scale,
@@ -2618,7 +2618,7 @@ def test_te_ops_backward_fusion_uses_eager_implementations():
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 def test_te_ops_unsupported_group_still_compiles_eagerly():
-    """An operation without the compute halves runs its eager implementation.
+    """An operation without a custom op runs its eager implementation.
 
     Note that this is not a fallback: under ``fullgraph=True`` there is no
     leaving the graph, so the pipeline is traced either way and only the choice
diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py
index 0bb3eb7ae2..d50cd86637 100644
--- a/transformer_engine/pytorch/ops/op.py
+++ b/transformer_engine/pytorch/ops/op.py
@@ -25,11 +25,6 @@
 from ..tensor import Quantizer
 from ..dynamo import is_value_opaque_quantizer, register_custom_op
 
-# Keyword-only parameters every resolve_fwd_args takes; the rest are forward kwargs.
-_RESOLVE_FWD_ARGS_PARAMS = frozenset(
-    ("requires_grad", "prev_op_grad_output_quantizer", "next_op_input_quantizer")
-)
-
 
 @dataclasses.dataclass
 class OperationContext:
@@ -194,8 +189,8 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta):
     num_extra_outputs: int = 0
 
     # torch.compile support: an operation declares its argument containers and
-    # implements the compute halves (see op_forward); custom ops are registered
-    # once per class.
+    # implements forward_impl / backward_impl and their fakes (see op_forward);
+    # its custom op is registered once per class.
     fwd_args_type: Optional[type] = None
     bwd_args_type: Optional[type] = None
     # Forward kwargs accepted by resolve_fwd_args. Read-only, no gradient.
@@ -216,18 +211,19 @@ def _register_compile_ops(cls) -> None:
         params = inspect.signature(cls.resolve_fwd_args).parameters
         if any(p.kind is p.VAR_KEYWORD for p in params.values()):
             raise TypeError(f"{cls.__name__}.resolve_fwd_args must name its keyword arguments")
+        base_params = inspect.signature(BasicOperation.resolve_fwd_args).parameters
         cls.fwd_kwarg_names = tuple(
             name
             for name, p in params.items()
-            if p.kind is p.KEYWORD_ONLY and name not in _RESOLVE_FWD_ARGS_PARAMS
+            if p.kind is p.KEYWORD_ONLY and name not in base_params
         )
         cls.compile_ops = register_custom_op(
             op_name=cls.__name__.lower(),
             fwd_arg_type=cls.fwd_args_type,
-            fwd_impl=cls.forward_compute,
+            fwd_impl=cls.forward_impl,
             fwd_fake_impl=cls.forward_fake,
             bwd_arg_type=cls.bwd_args_type,
-            bwd_impl=cls.backward_compute,
+            bwd_impl=cls.backward_impl,
             bwd_fake_impl=cls.backward_fake,
         )
 
@@ -564,8 +560,7 @@ def op_forward(
     ) -> torch.Tensor:
         """Forward pass
 
-        Operations that implement the compute halves inherit this;
-        the rest override it.
+        Operations with a custom op inherit this; the rest override it.
 
         Parameters
         ----------
@@ -585,7 +580,7 @@ def op_forward(
 
         """
         return self._forward(
-            self.forward_compute,
+            self.forward_impl,
             ctx,
             input_,
             prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
@@ -600,8 +595,7 @@ def op_backward(
     ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]:
         """Backward pass
 
-        Operations that implement the compute halves inherit this;
-        the rest override it.
+        Operations with a custom op inherit this; the rest override it.
 
         Parameters
         ----------
@@ -618,7 +612,7 @@ def op_backward(
             Loss gradients w.r.t. parameters
 
         """
-        return self._backward(self.backward_compute, ctx, grad_output)
+        return self._backward(self.backward_impl, ctx, grad_output)
 
     def compiled_op_forward(
         self,
@@ -670,29 +664,28 @@ def _backward(
         grad_input = grad_output if grads[0] is None else grads[0]
         return grad_input, tuple(grads[1:])
 
-    # Compute halves: forward_compute / backward_compute are the eager kernels
-    # and the custom-op bodies; the *_fake twins run on TensorSpec. Neither may
-    # read self, global state, or mutate its arguments. A None grad_input
-    # means "grad_output, unchanged".
+    # Custom op implementation: *_impl runs eagerly and as the op body, *_fake
+    # on TensorSpec. Neither may read self, global state, or mutate its
+    # arguments. A None grad_input means "grad_output, unchanged".
 
     @classmethod
-    def forward_compute(cls, args: Any) -> tuple[torch.Tensor, tuple]:
+    def forward_impl(cls, args: Any) -> tuple[torch.Tensor, tuple]:
         """Forward over ``args``; returns ``(output, aux)``, aux being fresh tensors."""
         raise NotImplementedError
 
     @classmethod
     def forward_fake(cls, args: Any) -> tuple[Any, tuple]:
-        """Shape-only twin of :meth:`forward_compute`."""
+        """Shape-only twin of :meth:`forward_impl`."""
         raise NotImplementedError
 
     @classmethod
-    def backward_compute(cls, args: Any) -> tuple:
+    def backward_impl(cls, args: Any) -> tuple:
         """Backward over ``args``; returns grad_input, then parameter grads."""
         raise NotImplementedError
 
     @classmethod
     def backward_fake(cls, args: Any) -> tuple:
-        """Shape-only twin of :meth:`backward_compute`."""
+        """Shape-only twin of :meth:`backward_impl`."""
         raise NotImplementedError
 
     def resolve_fwd_args(
@@ -722,7 +715,7 @@ def setup_context(self, ctx: OperationContext, args: Any, aux: tuple) -> None:
     def compile_unsupported_reason(self) -> Optional[str]:
         """Why this operation cannot run through its custom op, or ``None``."""
         if self.compile_ops is None:
-            return f"{self.__class__.__name__} without compute halves"
+            return f"{self.__class__.__name__} without a custom op"
         for mode in ("forward", "backward"):
             for index in range(self.num_quantizers(mode)):
                 quantizer = self.get_quantizer(mode, index)

From 2f8582f8667a16ecccc0c5d29be95a0309c58a7f Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 9 Sep 2026 18:08:15 +0200
Subject: [PATCH 26/27] [PyTorch] Guard compiled fuser cleanup and delayed
 scaling

Skip tensor storage cleanup while tracing eager operation bodies. Reject delayed-scaling state in the fuser, including CustomRecipe, before compiled execution can omit its backward scale update.

Add regression coverage for input preservation and built-in/custom delayed-scaling rejection after eager warmup.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py     | 77 ++++++++++++++++++++++++-
 transformer_engine/pytorch/ops/fuser.py | 11 +++-
 transformer_engine/pytorch/utils.py     |  2 +
 3 files changed, 88 insertions(+), 2 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 9252051387..620d765ad4 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -37,7 +37,12 @@
 from transformer_engine.common import recipe
 from transformer_engine.pytorch.constants import FP8FwdTensorIdx, FP8BwdTensorIdx
 from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
-from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole
+from transformer_engine.pytorch.quantization import (
+    FP8GlobalStateManager,
+    QuantizerRole,
+    DelayedScalingRequest,
+)
+from transformer_engine.pytorch.utils import clear_tensor_data
 from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear
 from transformer_engine.pytorch.ops.fuser import OperationFuser
 from transformer_engine.pytorch.ops.op import BasicOperation
@@ -2632,6 +2637,76 @@ def test_te_ops_unsupported_group_still_compiles_eagerly():
     _assert_sequential_matches_eager(lambda: te.ops.Sequential(te.ops.Identity()), base)
 
 
+class _CleanupOp(BasicOperation):
+    def op_forward(self, ctx, input_, **kwargs):
+        ctx.save_for_backward(input_)
+        return input_ * 2
+
+    def op_backward(self, ctx, grad_output):
+        dx = grad_output * 2
+        clear_tensor_data(ctx.saved_tensors[0])
+        return dx, ()
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("backend", ["eager", "inductor"])
+def test_te_ops_eager_implementation_preserves_saved_input(backend):
+    torch._dynamo.reset()
+    model = te.ops.Sequential(_CleanupOp())
+    compiled = torch.compile(model, backend=backend, fullgraph=True)
+    base = torch.randn(16, 64, device="cuda")
+    for fn in (model, compiled):
+        inp = base.clone().requires_grad_(True)
+        out = fn(inp)
+        out.sum().backward()
+        torch.testing.assert_close(out, base * 2)
+        torch.testing.assert_close(inp, base)
+        torch.testing.assert_close(inp.grad, torch.full_like(base, 2))
+
+
+class _DelayedScalingOp(BasicOperation):
+    def num_quantizers(self, mode):
+        return 1
+
+    def get_quantizer_roles(self, mode):
+        tensor_type = "input" if mode == "forward" else "grad_output"
+        return [QuantizerRole(module_type="test", tensor_type=tensor_type)]
+
+    def op_forward(self, ctx, input_, **kwargs):
+        ctx.amax = self.get_quantizer("backward", 0).amax
+        return input_ * 2
+
+    def op_backward(self, ctx, grad_output):
+        ctx.amax.copy_(grad_output.abs().max())
+        return grad_output * 2, ()
+
+
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("custom_recipe", [False, True])
+def test_te_ops_delayed_scaling_state_rejected(custom_recipe):
+    torch._dynamo.reset()
+    fp8_recipe = (
+        recipe.CustomRecipe(
+            qfactory=lambda role: DelayedScalingRequest(amax_history_len=4, reduce_amax=False)
+        )
+        if custom_recipe
+        else recipe.DelayedScaling(amax_history_len=4, reduce_amax=False)
+    )
+    op = _DelayedScalingOp()
+    model = te.ops.Sequential(op)
+    inp = torch.randn(16, 64, device="cuda", requires_grad=True)
+    with te.autocast(recipe=fp8_recipe):
+        out = model(inp)
+    out.sum().backward()
+
+    state = op._fp8_metas["backward"]["scaling_bwd"]
+    torch.testing.assert_close(state.amax_history[-1], torch.ones_like(state.scale))
+    compiled = torch.compile(model, fullgraph=True)
+    with te.autocast(recipe=fp8_recipe):
+        with pytest.raises(Exception, match="Delayed scaling is not supported under torch.compile"):
+            compiled(inp)
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_te_ops_forward_kwargs_compile():
diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py
index 3e1e85bae2..86174ca27a 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -12,7 +12,7 @@
 
 import torch
 
-from ..quantization import FP8GlobalStateManager, Recipe
+from ..quantization import FP8GlobalStateManager, Recipe, _has_delayed_scaling_state
 from ..quantized_tensor import prepare_for_saving, restore_from_func_ctx
 from ..utils import warn_compile_eager_fallback
 from .op import (
@@ -834,6 +834,15 @@ def __call__(
 
         # Initialization before forward
         for idx, op in enumerate(self._basic_ops):
+            if torch.compiler.is_compiling() and op._fp8_metas is not None:
+                if any(
+                    meta is not None and _has_delayed_scaling_state(meta)
+                    for meta in op._fp8_metas.values()
+                ):
+                    raise RuntimeError(
+                        "Delayed scaling is not supported under torch.compile in OperationFuser, "
+                        "including CustomRecipe with DelayedScalingRequest."
+                    )
             op.pre_fuser_forward(requires_grad=idx >= self.first_op_requiring_backward)
 
         # Fuser forward pass
diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py
index c4c6024f47..00bc627031 100644
--- a/transformer_engine/pytorch/utils.py
+++ b/transformer_engine/pytorch/utils.py
@@ -131,6 +131,8 @@ def clear_tensor_data(*tensors: Tuple[Optional[torch.Tensor], ...]) -> None:
 
     Must be used carefully.
     """
+    if torch.compiler.is_compiling():
+        return
 
     for t in tensors:
         if t is not None:

From 4c6070344d219ddddbbf02077ef8ba37fa2fe26f Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Wed, 9 Sep 2026 18:42:56 +0200
Subject: [PATCH 27/27] [PyTorch] Remove added fuser regression tests

Remove the test additions from 2f8582f8 at the requested scope. Keep both production fixes.

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

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 620d765ad4..9252051387 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -37,12 +37,7 @@
 from transformer_engine.common import recipe
 from transformer_engine.pytorch.constants import FP8FwdTensorIdx, FP8BwdTensorIdx
 from transformer_engine.pytorch.module.base import TransformerEngineBaseModule
-from transformer_engine.pytorch.quantization import (
-    FP8GlobalStateManager,
-    QuantizerRole,
-    DelayedScalingRequest,
-)
-from transformer_engine.pytorch.utils import clear_tensor_data
+from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole
 from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear
 from transformer_engine.pytorch.ops.fuser import OperationFuser
 from transformer_engine.pytorch.ops.op import BasicOperation
@@ -2637,76 +2632,6 @@ def test_te_ops_unsupported_group_still_compiles_eagerly():
     _assert_sequential_matches_eager(lambda: te.ops.Sequential(te.ops.Identity()), base)
 
 
-class _CleanupOp(BasicOperation):
-    def op_forward(self, ctx, input_, **kwargs):
-        ctx.save_for_backward(input_)
-        return input_ * 2
-
-    def op_backward(self, ctx, grad_output):
-        dx = grad_output * 2
-        clear_tensor_data(ctx.saved_tensors[0])
-        return dx, ()
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@pytest.mark.parametrize("backend", ["eager", "inductor"])
-def test_te_ops_eager_implementation_preserves_saved_input(backend):
-    torch._dynamo.reset()
-    model = te.ops.Sequential(_CleanupOp())
-    compiled = torch.compile(model, backend=backend, fullgraph=True)
-    base = torch.randn(16, 64, device="cuda")
-    for fn in (model, compiled):
-        inp = base.clone().requires_grad_(True)
-        out = fn(inp)
-        out.sum().backward()
-        torch.testing.assert_close(out, base * 2)
-        torch.testing.assert_close(inp, base)
-        torch.testing.assert_close(inp.grad, torch.full_like(base, 2))
-
-
-class _DelayedScalingOp(BasicOperation):
-    def num_quantizers(self, mode):
-        return 1
-
-    def get_quantizer_roles(self, mode):
-        tensor_type = "input" if mode == "forward" else "grad_output"
-        return [QuantizerRole(module_type="test", tensor_type=tensor_type)]
-
-    def op_forward(self, ctx, input_, **kwargs):
-        ctx.amax = self.get_quantizer("backward", 0).amax
-        return input_ * 2
-
-    def op_backward(self, ctx, grad_output):
-        ctx.amax.copy_(grad_output.abs().max())
-        return grad_output * 2, ()
-
-
-@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
-@pytest.mark.parametrize("custom_recipe", [False, True])
-def test_te_ops_delayed_scaling_state_rejected(custom_recipe):
-    torch._dynamo.reset()
-    fp8_recipe = (
-        recipe.CustomRecipe(
-            qfactory=lambda role: DelayedScalingRequest(amax_history_len=4, reduce_amax=False)
-        )
-        if custom_recipe
-        else recipe.DelayedScaling(amax_history_len=4, reduce_amax=False)
-    )
-    op = _DelayedScalingOp()
-    model = te.ops.Sequential(op)
-    inp = torch.randn(16, 64, device="cuda", requires_grad=True)
-    with te.autocast(recipe=fp8_recipe):
-        out = model(inp)
-    out.sum().backward()
-
-    state = op._fp8_metas["backward"]["scaling_bwd"]
-    torch.testing.assert_close(state.amax_history[-1], torch.ones_like(state.scale))
-    compiled = torch.compile(model, fullgraph=True)
-    with te.autocast(recipe=fp8_recipe):
-        with pytest.raises(Exception, match="Delayed scaling is not supported under torch.compile"):
-            compiled(inp)
-
-
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
 @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
 def test_te_ops_forward_kwargs_compile():