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 1/5] [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 2/5] [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 3/5] [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 4/5] [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 5/5] [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]]],