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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py     | 199 ++++++++++++++++++++----
 transformer_engine/pytorch/ops/fuser.py |  11 +-
 transformer_engine/pytorch/ops/op.py    |  22 ++-
 3 files changed, 200 insertions(+), 32 deletions(-)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Drop test_te_ops_setup_context_saves_parameter.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

From 9302392fde601ae1d864e14d85249f67467f7c5c Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 11 Sep 2026 12:51:40 +0200
Subject: [PATCH 12/20] [PyTorch] Share compute interfaces between basic and
 fused operations

Move argument packing, compute dispatch, and context saving into FusibleOperation. Preserve nested custom-op results for extra outputs and gradients, while retaining existing pipeline compile gates.

Validate shared fused interfaces and custom-op graph nodes. Workstation tests: 1620 passed, 1039 skipped, 1 xpassed; full Python lint reports only existing fused_mla_q_uproj diagnostics.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 257 ++++++++++++---
 .../pytorch/dynamo/custom_op.py               |  60 ++--
 transformer_engine/pytorch/ops/fuser.py       |  45 +--
 transformer_engine/pytorch/ops/op.py          | 308 ++++++++----------
 4 files changed, 404 insertions(+), 266 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 9252051387..c1b3666d8d 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -40,7 +40,7 @@
 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.ops.op import BasicOperation, OperationContext
 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
@@ -2387,44 +2387,47 @@ def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16)
         self.scale = torch.nn.Parameter(torch.full((), 2.0, device=device, dtype=dtype))
 
     @classmethod
-    def forward_impl(cls, args):
-        return args.input_ * args.scale, ()
+    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), ()
+        return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), [()], ()
 
     @classmethod
-    def backward_impl(cls, args):
+    def backward_compute(cls, args):
         dy = args.grad_output
-        return dy * args.scale, (dy * args.saved_input).sum()
+        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),
+            [(TensorSpec(shape=(), dtype=dy.dtype, device=dy.device),)],
+            [()],
         )
 
-    def setup_context(self, ctx, args, aux):
+    def forward_setup_context(self, basic_op_ctxs, args, aux):
         del aux
+        ctx = basic_op_ctxs[0]
         ctx.save_for_backward(args.input_, args.scale)
 
-    def resolve_fwd_args(
+    def pack_forward_args(
         self,
+        basic_op_ctxs,
         input_,
         *,
-        requires_grad,
-        prev_op_grad_output_quantizer=None,
-        next_op_input_quantizer=None,
+        basic_op_extra_inputs,
+        prev_op_grad_output_quantizer,
+        next_op_input_quantizer,
+        basic_op_kwargs,
     ):
-        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, scale = ctx.saved_tensors
+    def pack_backward_args(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs):
+        x, scale = basic_op_ctxs[0].saved_tensors
         return _ScaleBwdArgs(grad_output=grad_output, saved_input=x, scale=scale)
 
 
@@ -2432,9 +2435,13 @@ 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], [(), ()]
+        dx, grad_params_1, _ = self.basic_ops[1].fuser_backward(
+            [basic_op_ctxs[1]], grad_output, basic_op_grad_extra_outputs=[()]
+        )
+        dx, grad_params_0, _ = self.basic_ops[0].fuser_backward(
+            [basic_op_ctxs[0]], dx, basic_op_grad_extra_outputs=[()]
+        )
+        return dx, grad_params_0 + grad_params_1, [(), ()]
 
 
 def _fuse_backward_scale_pair(ops, **unused):
@@ -2473,30 +2480,32 @@ class _ScaleWithKwargsOp(BasicOperation):
 
     fwd_args_type = _ScaleKwargsFwdArgs
     bwd_args_type = _ScaleKwargsBwdArgs
+    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_impl(cls, args):
+    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, ()
+        return out, [()], ()
 
     @classmethod
     def forward_fake(cls, args):
         x = args.input_
-        return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), ()
+        return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), [()], ()
 
     @classmethod
-    def backward_impl(cls, args):
+    def backward_compute(cls, args):
         dy = args.grad_output
         return (
             dy * args.scale * args.extra_scale,
-            (dy * args.saved_input).sum() * args.extra_scale,
+            [((dy * args.saved_input).sum() * args.extra_scale,)],
+            [()],
         )
 
     @classmethod
@@ -2504,35 +2513,39 @@ 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),
+            [(TensorSpec(shape=(), dtype=dy.dtype, device=dy.device),)],
+            [()],
         )
 
-    def setup_context(self, ctx, args, aux):
+    def forward_setup_context(self, basic_op_ctxs, args, aux):
         del aux
+        ctx = basic_op_ctxs[0]
         ctx.save_for_backward(args.input_, args.scale)
         ctx.extra_scale = args.extra_scale
 
-    def resolve_fwd_args(
+    def pack_forward_args(
         self,
+        basic_op_ctxs,
         input_,
         *,
-        requires_grad,
-        prev_op_grad_output_quantizer=None,
-        next_op_input_quantizer=None,
-        extra_scale=1.0,
-        offset=None,
+        basic_op_extra_inputs,
+        prev_op_grad_output_quantizer,
+        next_op_input_quantizer,
+        basic_op_kwargs,
     ):
-        del requires_grad, prev_op_grad_output_quantizer, next_op_input_quantizer
+        kwargs = basic_op_kwargs[0]
+        offset = kwargs.get("offset")
         if offset is None:
             offset = torch.zeros((), device=input_.device, dtype=input_.dtype)
         return _ScaleKwargsFwdArgs(
             input_=input_,
             scale=self.scale,
-            extra_scale=extra_scale,
+            extra_scale=kwargs.get("extra_scale", 1.0),
             offset=offset,
         )
 
-    def resolve_bwd_args(self, ctx, grad_output):
+    def pack_backward_args(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs):
+        ctx = basic_op_ctxs[0]
         x, scale = ctx.saved_tensors
         return _ScaleKwargsBwdArgs(
             grad_output=grad_output,
@@ -2542,7 +2555,156 @@ def resolve_bwd_args(self, ctx, grad_output):
         )
 
 
-def _assert_sequential_matches_eager(make_model, base, op_kwargs_seq=(None,)):
+@dataclasses.dataclass(slots=True)
+class _ScalePairFwdArgs:
+    input_: torch.Tensor
+    scale0: torch.Tensor
+    scale1: torch.Tensor
+    extra_input: torch.Tensor = None
+
+
+@dataclasses.dataclass(slots=True)
+class _ScalePairBwdArgs:
+    grad_output: torch.Tensor
+    input_: torch.Tensor
+    intermediate: torch.Tensor
+    scale0: torch.Tensor
+    scale1: torch.Tensor
+    grad_extra_output: torch.Tensor = None
+
+
+class _ScalePair(te.ops.FusedOperation):
+    """Two scales, optionally with a residual input and squared intermediate output."""
+
+    fwd_args_type = _ScalePairFwdArgs
+    bwd_args_type = _ScalePairBwdArgs
+
+    @classmethod
+    def forward_compute(cls, args):
+        intermediate = args.input_ * args.scale0
+        output = intermediate * args.scale1
+        extras = ()
+        if args.extra_input is not None:
+            output = output + args.extra_input
+            extras = (intermediate.square(), None)
+        return output, [(), extras], (intermediate,)
+
+    @classmethod
+    def forward_fake(cls, args):
+        x = args.input_
+        spec = TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device)
+        extras = (spec, None) if args.extra_input is not None else ()
+        return spec, [(), extras], (spec,)
+
+    @classmethod
+    def backward_compute(cls, args):
+        dy = args.grad_output
+        du = dy * args.scale1
+        extras = ()
+        if args.grad_extra_output is not None:
+            du = du + 2 * args.intermediate * args.grad_extra_output
+            extras = (dy.clone(),)
+        return (
+            du * args.scale0,
+            [((du * args.input_).sum(),), ((dy * args.intermediate).sum(),)],
+            [(), extras],
+        )
+
+    @classmethod
+    def backward_fake(cls, args):
+        dy = args.grad_output
+        spec = TensorSpec(shape=tuple(dy.shape), dtype=dy.dtype, device=dy.device)
+        scalar = TensorSpec(shape=(), dtype=dy.dtype, device=dy.device)
+        extras = (spec,) if args.grad_extra_output is not None else ()
+        return spec, [(scalar,), (scalar,)], [(), extras]
+
+    def pack_forward_args(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused):
+        extras = basic_op_extra_inputs[1]
+        return _ScalePairFwdArgs(
+            input_,
+            self.basic_ops[0].scale,
+            self.basic_ops[1].scale,
+            extras[0] if extras else None,
+        )
+
+    def forward_setup_context(self, basic_op_ctxs, args, aux):
+        basic_op_ctxs[0].save_for_backward(args.input_, args.scale0)
+        basic_op_ctxs[1].save_for_backward(aux[0], args.scale1)
+
+    def pack_backward_args(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs):
+        x, scale0 = basic_op_ctxs[0].saved_tensors
+        intermediate, scale1 = basic_op_ctxs[1].saved_tensors
+        extras = basic_op_grad_extra_outputs[1]
+        return _ScalePairBwdArgs(
+            grad_output, x, intermediate, scale0, scale1, extras[0] if extras else None
+        )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("with_extras", [False, True])
+@pytest.mark.parametrize("use_custom_ops", [False, True])
+def test_te_ops_fused_compute_contract(with_extras, use_custom_ops):
+    """Exercise the shared interface directly; pipeline fusion stays gated under compile."""
+    torch._dynamo.reset()
+    ops = [_ScaleOp(dtype=torch.float32), _ScaleOp(dtype=torch.float32)]
+    with torch.no_grad():
+        ops[1].scale.fill_(3.0)
+    fused = _ScalePair(ops)
+    x = torch.randn(8, 16, device="cuda", requires_grad=True)
+    residual = torch.randn_like(x, requires_grad=True)
+    dy = torch.randn_like(x)
+    dextra = torch.randn_like(x)
+
+    def run(input_, extra_input, grad_output, grad_extra_output):
+        ctxs = [OperationContext(), OperationContext()]
+        output, extra_outputs = fused.fuser_forward(
+            ctxs,
+            input_,
+            basic_op_extra_inputs=[(), (extra_input,) if with_extras else ()],
+            prev_op_grad_output_quantizer=None,
+            next_op_input_quantizer=None,
+            basic_op_kwargs=[{}, {}],
+            use_custom_ops=use_custom_ops,
+        )
+        for ctx in ctxs:
+            ctx.saved_tensors = ctx.to_save
+        grads = fused.fuser_backward(
+            ctxs,
+            grad_output,
+            basic_op_grad_extra_outputs=[(), (grad_extra_output, None) if with_extras else ()],
+            use_custom_ops=use_custom_ops,
+        )
+        return output, extra_outputs, grads
+
+    if use_custom_ops:
+        assert fused.compile_ops is not None
+        run = torch.compile(run, fullgraph=True)
+    with torch.no_grad():
+        output, extras, (dx, dparams, dextras) = run(x, residual, dy, dextra)
+
+    intermediate = x * ops[0].scale
+    reference = intermediate * ops[1].scale
+    loss = (reference * dy).sum()
+    if with_extras:
+        reference = reference + residual
+        loss = (reference * dy).sum() + (intermediate.square() * dextra).sum()
+    inputs = [x, ops[0].scale, ops[1].scale] + ([residual] if with_extras else [])
+    expected = torch.autograd.grad(loss, inputs)
+    torch.testing.assert_close(output, reference)
+    torch.testing.assert_close(dx, expected[0])
+    assert len(dparams) == 2 and all(len(group) == 1 for group in dparams)
+    torch.testing.assert_close(dparams[0][0], expected[1])
+    torch.testing.assert_close(dparams[1][0], expected[2])
+    assert extras[0] == () and dextras[0] == ()
+    if with_extras:
+        assert len(extras[1]) == 2 and extras[1][1] is None
+        torch.testing.assert_close(extras[1][0], intermediate.square())
+        torch.testing.assert_close(dextras[1][0], expected[3])
+    else:
+        assert extras[1] == () and dextras[1] == ()
+
+
+def _assert_sequential_matches_eager(make_model, base, op_kwargs_seq=(None,), *, graphs=None):
     """Run a Sequential eagerly and compiled on identical inputs; compare both
     the output and every parameter gradient.
 
@@ -2556,7 +2718,14 @@ def _assert_sequential_matches_eager(make_model, base, op_kwargs_seq=(None,)):
     """
     eager_model = make_model()
     compiled_model = make_model()
-    compiled = torch.compile(compiled_model, fullgraph=True)
+    backend = "inductor"
+    if graphs is not None:
+
+        def backend(graph, inputs):
+            graphs.append(graph)
+            return torch._dynamo.lookup_backend("inductor")(graph, inputs)
+
+    compiled = torch.compile(compiled_model, fullgraph=True, backend=backend)
 
     for op_kwargs in op_kwargs_seq:
         call_kwargs = {} if op_kwargs is None else {"op_kwargs": op_kwargs}
@@ -2590,7 +2759,21 @@ def test_te_ops_single_op_group_compiles():
     """
     torch._dynamo.reset()
     base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
-    _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp()), base)
+    graphs = []
+    _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp()), base, graphs=graphs)
+    targets = {
+        str(node.target).removesuffix(".default")
+        for graph in graphs
+        for module in graph.modules()
+        if isinstance(module, torch.fx.GraphModule)
+        for node in module.graph.nodes
+        if node.op == "call_function"
+    }
+    for name in ("_scaleop", "_scaleop_backward"):
+        assert targets & {
+            f"transformer_engine_compile.{name}",
+            f"transformer_engine_compile.{name}_base",
+        }, targets
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 8d80eebf61..db14a22bef 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -103,6 +103,7 @@
 import torch
 
 from torch._prims_common import make_contiguous_strides_for
+from torch.utils._pytree import tree_flatten, tree_unflatten
 
 from .tensor_spec import TensorSpec, to_tensor_spec
 from ..quantized_tensor import (
@@ -1313,20 +1314,13 @@ def register_custom_op(
     Both ops are two-tier, so ``QuantizedTensor`` subclass inputs pass through
     without dequantization.
 
-    Callable contracts:
-
-    * ``fwd_impl(fwd_args) -> (output, aux)`` -- ``aux`` is a tuple of fresh tensors
-    * ``fwd_fake_impl`` -- its data-free twin over :class:`TensorSpec`
-    * ``bwd_impl(bwd_args) -> tuple`` of gradients (``num_grad_inputs`` of them,
-      if given)
-    * ``bwd_fake_impl`` -- its data-free twin
-
-    Returns ``(forward_fn, backward_fn)``:
+    Implementations return nested tuples/lists of tensors or None. Fake
+    implementations return the same structure with TensorSpec leaves. Forward
+    outputs must be fresh tensors; context saving remains the caller's job.
+    num_grad_inputs, if given, counts flattened backward result leaves.
 
-    * ``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 (forward_fn, backward_fn), preserving each implementation's result
+    structure, including per-basic-op extra outputs and gradients.
 
     Returns ``None`` if registration fails (recorded once), so callers can fall
     back to eager rather than breaking import.
@@ -1364,43 +1358,43 @@ def _register_custom_op_impl(
 
     def adapt_forward(impl):
         def wrapped(args):
-            result = impl(args)
-            if not isinstance(result, tuple) or len(result) != 2:
-                raise TypeError(
-                    "autograd-free fwd impl must return an (output, aux) tuple, got"
-                    f" {type(result).__name__}"
-                )
-            output, aux = result
-            return output, tuple(aux), None
+            values, _ = tree_flatten(impl(args))
+            return (*values, (), None)
+
+        return wrapped
+
+    def adapt_backward(impl):
+        def wrapped(args):
+            values, _ = tree_flatten(impl(args))
+            return tuple(values)
 
         return wrapped
 
-    adapted_fwd_fake_impl = adapt_forward(fwd_fake_impl)
     fwd_op = _register_forward_op(
         name=op_name,
         arg_type=fwd_arg_type,
         impl=adapt_forward(fwd_impl),
-        fake_impl=adapted_fwd_fake_impl,
+        fake_impl=adapt_forward(fwd_fake_impl),
     )
     bwd_op = _register_backward_op(
         name=f"{op_name}_backward",
         arg_type=bwd_arg_type,
-        impl=bwd_impl,
-        fake_impl=bwd_fake_impl,
+        impl=adapt_backward(bwd_impl),
+        fake_impl=adapt_backward(bwd_fake_impl),
         num_grad_inputs=num_grad_inputs,
     )
 
     def forward_fn(fwd_args):
-        out_plan, payload = _run_forward(fwd_op, adapted_fwd_fake_impl, fwd_args)
-        outputs = out_plan.user_outputs(payload)
-        aux = out_plan.saved_tensors(payload)
-        return outputs[0], tuple(aux)
+        spec_args = _spec_view(fwd_args, fwd_op.plan.tensor_field_names())
+        specs, structure = tree_flatten(fwd_fake_impl(spec_args))
+        out_plan = _OutputPlan.parse((*specs, (), None))
+        return tree_unflatten(out_plan.user_outputs(fwd_op(fwd_args)), structure)
 
     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.
-        return tuple(_decode_none(t) for t in bwd_op(bwd_args))
+        spec_args = _spec_view(bwd_args, bwd_op.plan.tensor_field_names())
+        _, structure = tree_flatten(bwd_fake_impl(spec_args))
+        grads = [_decode_none(t) for t in bwd_op(bwd_args)]
+        return tree_unflatten(grads, structure)
 
     return forward_fn, backward_fn
 
diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py
index 86174ca27a..10cbe6493f 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -166,24 +166,16 @@ def forward(
             if next_op is not None:
                 next_op_input_quantizer = next_op.get_input_quantizer()
 
-            if use_custom_ops:
-                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,
-                    **basic_op_kwargs[basic_op_idxs[0]],
-                )
-                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],
-                )
+            compile_kwargs = {"use_custom_ops": True} if use_custom_ops 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],
+                **compile_kwargs,
+            )
             if len(fused_op_extra_outputs) != len(basic_op_idxs):
                 raise RuntimeError(
                     f"Expected {type(op).__name__} to generate extra outputs for "
@@ -364,16 +356,13 @@ 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_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 = [()]
-            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,
-                )
+            compile_kwargs = {"use_custom_ops": True} if func_ctx.use_custom_ops 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,
+                **compile_kwargs,
+            )
             for idx, dparams in zip(basic_op_idxs, fused_op_grad_params):
                 grad_params[idx] = dparams
                 basic_op_ctxs[idx].saved_tensors = None
diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py
index d50cd86637..53b6d83fbb 100644
--- a/transformer_engine/pytorch/ops/op.py
+++ b/transformer_engine/pytorch/ops/op.py
@@ -8,7 +8,6 @@
 import abc
 from collections.abc import Iterable, Sequence
 import dataclasses
-import inspect
 import pickle
 from typing import Any, Callable, Optional
 
@@ -61,6 +60,47 @@ def save_for_backward(self, *tensors: Optional[torch.Tensor]) -> None:
 class FusibleOperation(torch.nn.Module, metaclass=abc.ABCMeta):
     """Tensor operation supported by the operation fuser"""
 
+    # One custom-op registration per operation class.
+    fwd_args_type: Optional[type] = None
+    bwd_args_type: Optional[type] = None
+    # Supported read-only forward kwargs; no gradients.
+    fwd_kwarg_names: tuple[str, ...] = ()
+    # (forward_fn, backward_fn), or None if the operation cannot be compiled.
+    compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None
+
+    def __init_subclass__(cls, **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
+        if cls.fwd_args_type is not None and cls.bwd_args_type is not None:
+            cls._register_compile_ops()
+
+    @classmethod
+    def _register_compile_ops(cls) -> None:
+        for name in ("fwd_args_type", "bwd_args_type"):
+            if not dataclasses.is_dataclass(getattr(cls, name)):
+                raise TypeError(f"{cls.__name__}.{name} must be a dataclass")
+        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,
+        )
+
+    def compile_unsupported_reason(self) -> Optional[str]:
+        """Why this operation cannot run through its custom op, or ``None``."""
+        if self.compile_ops is None:
+            return f"{self.__class__.__name__} without a custom op"
+        basic_ops = self.basic_ops if self.is_fused_op else (self,)
+        for op in basic_ops:
+            for mode in ("forward", "backward"):
+                for index in range(op.num_quantizers(mode)):
+                    quantizer = op.get_quantizer(mode, index)
+                    if quantizer is not None and not is_value_opaque_quantizer(quantizer):
+                        return f"{type(quantizer).__name__} (not a torch.compile value-opaque quantizer)"
+        return None
+
     @property
     @abc.abstractmethod
     def is_fused_op(self) -> bool:
@@ -91,6 +131,7 @@ def fuser_forward(
         prev_op_grad_output_quantizer: Optional[Quantizer],
         next_op_input_quantizer: Optional[Quantizer],
         basic_op_kwargs: list[dict[str, Any]],
+        use_custom_ops: bool = False,
     ) -> tuple[torch.Tensor, Sequence[Sequence[Optional[torch.Tensor]]]]:
         """Forward pass
 
@@ -126,9 +167,19 @@ def fuser_forward(
             channel owned by this fused operation may be ``None``.
 
         """
-        raise NotImplementedError(
-            f"Forward pass is not implemented for operation ({self.__class__.__name__})"
+        args = self.pack_forward_args(
+            basic_op_ctxs,
+            input_,
+            basic_op_extra_inputs=basic_op_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,
         )
+        compute = self.compile_ops[0] if use_custom_ops else self.forward_compute
+        output, extra_outputs, aux = compute(args)
+        if any(ctx.requires_grad for ctx in basic_op_ctxs):
+            self.forward_setup_context(basic_op_ctxs, args, aux)
+        return output, extra_outputs
 
     def fuser_backward(
         self,
@@ -136,6 +187,7 @@ def fuser_backward(
         grad_output: torch.Tensor,
         *,
         basic_op_grad_extra_outputs: Sequence[Sequence[Optional[torch.Tensor]]],
+        use_custom_ops: bool = False,
     ) -> tuple[
         torch.Tensor,
         Sequence[Sequence[Optional[torch.Tensor]]],
@@ -170,9 +222,67 @@ def fuser_backward(
             operations
 
         """
-        raise NotImplementedError(
-            f"Backward pass is not implemented for operation ({self.__class__.__name__})"
+        args = self.pack_backward_args(
+            basic_op_ctxs,
+            grad_output,
+            basic_op_grad_extra_outputs=basic_op_grad_extra_outputs,
         )
+        compute = self.compile_ops[1] if use_custom_ops else self.backward_compute
+        grad_input, grad_params, grad_extra_inputs = compute(args)
+        if grad_input is None:
+            grad_input = grad_output
+        return grad_input, grad_params, grad_extra_inputs
+
+    @classmethod
+    def forward_compute(cls, args: Any) -> tuple:
+        """Return (output, extra_outputs per basic op, fresh aux tensors)."""
+        raise NotImplementedError
+
+    @classmethod
+    def forward_fake(cls, args: Any) -> tuple:
+        """Shape-only twin of forward_compute, using TensorSpec."""
+        raise NotImplementedError
+
+    @classmethod
+    def backward_compute(cls, args: Any) -> tuple:
+        """Return (grad_input, grad_params per basic op, grad_extra_inputs per basic op).
+
+        A None grad_input passes grad_output through unchanged.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    def backward_fake(cls, args: Any) -> tuple:
+        """Shape-only twin of backward_compute, using TensorSpec."""
+        raise NotImplementedError
+
+    def pack_forward_args(
+        self,
+        basic_op_ctxs: list[OperationContext],
+        input_: torch.Tensor,
+        *,
+        basic_op_extra_inputs: Sequence[Sequence[Optional[torch.Tensor]]],
+        prev_op_grad_output_quantizer: Optional[Quantizer],
+        next_op_input_quantizer: Optional[Quantizer],
+        basic_op_kwargs: list[dict[str, Any]],
+    ) -> Any:
+        """Gather inputs and module state into fwd_args_type."""
+        raise NotImplementedError
+
+    def pack_backward_args(
+        self,
+        basic_op_ctxs: list[OperationContext],
+        grad_output: torch.Tensor,
+        *,
+        basic_op_grad_extra_outputs: Sequence[Sequence[Optional[torch.Tensor]]],
+    ) -> Any:
+        """Gather gradients and saved context into bwd_args_type."""
+        raise NotImplementedError
+
+    def forward_setup_context(
+        self, basic_op_ctxs: list[OperationContext], args: Any, aux: tuple
+    ) -> None:
+        """Save state needed by the basic operations' backward passes."""
 
 
 class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta):
@@ -188,45 +298,6 @@ class BasicOperation(FusibleOperation, metaclass=abc.ABCMeta):
     # Number of extra tensor outputs
     num_extra_outputs: int = 0
 
-    # torch.compile support: an operation declares its argument containers and
-    # implements forward_impl / backward_impl and their fakes (see op_forward);
-    # its custom op is registered once per class.
-    fwd_args_type: Optional[type] = None
-    bwd_args_type: Optional[type] = None
-    # Forward kwargs accepted by resolve_fwd_args. Read-only, no gradient.
-    fwd_kwarg_names: tuple[str, ...] = ()
-    # (forward_fn, backward_fn), or None if the operation cannot be compiled.
-    compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None
-
-    def __init_subclass__(cls, **kwargs) -> None:
-        super().__init_subclass__(**kwargs)
-        if cls.fwd_args_type is not None and cls.bwd_args_type is not None:
-            cls._register_compile_ops()
-
-    @classmethod
-    def _register_compile_ops(cls) -> None:
-        for name in ("fwd_args_type", "bwd_args_type"):
-            if not dataclasses.is_dataclass(getattr(cls, name)):
-                raise TypeError(f"{cls.__name__}.{name} must be a dataclass")
-        params = inspect.signature(cls.resolve_fwd_args).parameters
-        if any(p.kind is p.VAR_KEYWORD for p in params.values()):
-            raise TypeError(f"{cls.__name__}.resolve_fwd_args must name its keyword arguments")
-        base_params = inspect.signature(BasicOperation.resolve_fwd_args).parameters
-        cls.fwd_kwarg_names = tuple(
-            name
-            for name, p in params.items()
-            if p.kind is p.KEYWORD_ONLY and name not in base_params
-        )
-        cls.compile_ops = register_custom_op(
-            op_name=cls.__name__.lower(),
-            fwd_arg_type=cls.fwd_args_type,
-            fwd_impl=cls.forward_impl,
-            fwd_fake_impl=cls.forward_fake,
-            bwd_arg_type=cls.bwd_args_type,
-            bwd_impl=cls.backward_impl,
-            bwd_fake_impl=cls.backward_fake,
-        )
-
     def __init__(self) -> None:
         super().__init__()
 
@@ -560,7 +631,7 @@ def op_forward(
     ) -> torch.Tensor:
         """Forward pass
 
-        Operations with a custom op inherit this; the rest override it.
+        Convenience interface for operations without extra tensor inputs or outputs.
 
         Parameters
         ----------
@@ -579,14 +650,7 @@ def op_forward(
             Output tensor
 
         """
-        return self._forward(
-            self.forward_impl,
-            ctx,
-            input_,
-            prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
-            next_op_input_quantizer=next_op_input_quantizer,
-            **kwargs,
-        )
+        raise NotImplementedError
 
     def op_backward(
         self,
@@ -595,7 +659,7 @@ def op_backward(
     ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]:
         """Backward pass
 
-        Operations with a custom op inherit this; the rest override it.
+        Convenience interface for operations without extra tensor inputs or outputs.
 
         Parameters
         ----------
@@ -612,119 +676,8 @@ def op_backward(
             Loss gradients w.r.t. parameters
 
         """
-        return self._backward(self.backward_impl, ctx, grad_output)
-
-    def compiled_op_forward(
-        self,
-        ctx: OperationContext,
-        input_: torch.Tensor,
-        *,
-        prev_op_grad_output_quantizer: Optional[Quantizer] = None,
-        next_op_input_quantizer: Optional[Quantizer] = None,
-        **kwargs: Any,
-    ) -> torch.Tensor:
-        """:meth:`op_forward` through this operation's custom op."""
-        return self._forward(
-            self.compile_ops[0],
-            ctx,
-            input_,
-            prev_op_grad_output_quantizer=prev_op_grad_output_quantizer,
-            next_op_input_quantizer=next_op_input_quantizer,
-            **kwargs,
-        )
-
-    def compiled_op_backward(
-        self,
-        ctx: OperationContext,
-        grad_output: torch.Tensor,
-    ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]:
-        """:meth:`op_backward` through this operation's custom op."""
-        return self._backward(self.compile_ops[1], ctx, grad_output)
-
-    def _forward(
-        self,
-        compute: Callable[[Any], tuple[torch.Tensor, tuple]],
-        ctx: OperationContext,
-        input_: torch.Tensor,
-        **kwargs: Any,
-    ) -> torch.Tensor:
-        args = self.resolve_fwd_args(input_, requires_grad=ctx.requires_grad, **kwargs)
-        output, aux = compute(args)
-        if ctx.requires_grad:
-            self.setup_context(ctx, args, aux)
-        return output
-
-    def _backward(
-        self,
-        compute: Callable[[Any], tuple],
-        ctx: OperationContext,
-        grad_output: torch.Tensor,
-    ) -> tuple[torch.Tensor, Iterable[Optional[torch.Tensor]]]:
-        grads = compute(self.resolve_bwd_args(ctx, grad_output))
-        grad_input = grad_output if grads[0] is None else grads[0]
-        return grad_input, tuple(grads[1:])
-
-    # Custom op implementation: *_impl runs eagerly and as the op body, *_fake
-    # on TensorSpec. Neither may read self, global state, or mutate its
-    # arguments. A None grad_input means "grad_output, unchanged".
-
-    @classmethod
-    def forward_impl(cls, args: Any) -> tuple[torch.Tensor, tuple]:
-        """Forward over ``args``; returns ``(output, aux)``, aux being fresh tensors."""
-        raise NotImplementedError
-
-    @classmethod
-    def forward_fake(cls, args: Any) -> tuple[Any, tuple]:
-        """Shape-only twin of :meth:`forward_impl`."""
-        raise NotImplementedError
-
-    @classmethod
-    def backward_impl(cls, args: Any) -> tuple:
-        """Backward over ``args``; returns grad_input, then parameter grads."""
-        raise NotImplementedError
-
-    @classmethod
-    def backward_fake(cls, args: Any) -> tuple:
-        """Shape-only twin of :meth:`backward_impl`."""
-        raise NotImplementedError
-
-    def resolve_fwd_args(
-        self,
-        input_: torch.Tensor,
-        *,
-        requires_grad: bool,
-        prev_op_grad_output_quantizer: Optional[Quantizer] = None,
-        next_op_input_quantizer: Optional[Quantizer] = None,
-    ) -> Any:
-        """Build ``fwd_args_type`` from the input, module state and forward kwargs.
-
-        Forward kwargs are declared as additional keyword-only parameters
-        with defaults.
-        """
-        raise NotImplementedError
-
-    def resolve_bwd_args(self, ctx: OperationContext, grad_output: torch.Tensor) -> Any:
-        """Build ``bwd_args_type`` from the saved context."""
         raise NotImplementedError
 
-    def setup_context(self, ctx: OperationContext, args: Any, aux: tuple) -> None:
-        """Save backward state from the forward args and the fresh aux tensors."""
-        del args
-        ctx.save_for_backward(*aux)
-
-    def compile_unsupported_reason(self) -> Optional[str]:
-        """Why this operation cannot run through its custom op, or ``None``."""
-        if self.compile_ops is None:
-            return f"{self.__class__.__name__} without a custom op"
-        for mode in ("forward", "backward"):
-            for index in range(self.num_quantizers(mode)):
-                quantizer = self.get_quantizer(mode, index)
-                if quantizer is not None and not is_value_opaque_quantizer(quantizer):
-                    return (
-                        f"{type(quantizer).__name__} (not a torch.compile value-opaque quantizer)"
-                    )
-        return None
-
     def fuser_forward(
         self,
         basic_op_ctxs: list[OperationContext],
@@ -734,7 +687,18 @@ def fuser_forward(
         prev_op_grad_output_quantizer: Optional[Quantizer],
         next_op_input_quantizer: Optional[Quantizer],
         basic_op_kwargs: list[dict[str, Any]],
-    ) -> tuple[torch.Tensor, list[tuple[()]]]:
+        use_custom_ops: bool = False,
+    ) -> tuple[torch.Tensor, Sequence[Sequence[Optional[torch.Tensor]]]]:
+        if use_custom_ops or type(self).op_forward is BasicOperation.op_forward:
+            return super().fuser_forward(
+                basic_op_ctxs,
+                input_,
+                basic_op_extra_inputs=basic_op_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,
+                use_custom_ops=use_custom_ops,
+            )
         if self.num_extra_inputs > 0 or self.num_extra_outputs > 0:
             raise RuntimeError(
                 "{self.__class__.__name__} operation has "
@@ -757,11 +721,19 @@ def fuser_backward(
         grad_output: torch.Tensor,
         *,
         basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]],
+        use_custom_ops: bool = False,
     ) -> tuple[
         torch.Tensor,
-        list[Iterable[Optional[torch.Tensor]]],
-        list[tuple[()]],
+        Sequence[Sequence[Optional[torch.Tensor]]],
+        Sequence[Sequence[Optional[torch.Tensor]]],
     ]:
+        if use_custom_ops or type(self).op_backward is BasicOperation.op_backward:
+            return super().fuser_backward(
+                basic_op_ctxs,
+                grad_output,
+                basic_op_grad_extra_outputs=basic_op_grad_extra_outputs,
+                use_custom_ops=use_custom_ops,
+            )
         if self.num_extra_inputs > 0 or self.num_extra_outputs > 0:
             raise RuntimeError(
                 "{self.__class__.__name__} operation has "

From 8c7643122d827962bdecdfe0a1e7ecf79ffcc5df Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Fri, 11 Sep 2026 10:53:00 +0000
Subject: [PATCH 13/20] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 transformer_engine/pytorch/ops/op.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py
index 53b6d83fbb..69df67200d 100644
--- a/transformer_engine/pytorch/ops/op.py
+++ b/transformer_engine/pytorch/ops/op.py
@@ -98,7 +98,10 @@ def compile_unsupported_reason(self) -> Optional[str]:
                 for index in range(op.num_quantizers(mode)):
                     quantizer = op.get_quantizer(mode, index)
                     if quantizer is not None and not is_value_opaque_quantizer(quantizer):
-                        return f"{type(quantizer).__name__} (not a torch.compile value-opaque quantizer)"
+                        return (
+                            f"{type(quantizer).__name__} (not a torch.compile value-opaque"
+                            " quantizer)"
+                        )
         return None
 
     @property

From 4c14b36f07d017be7ad8f8de3f546df1c49d67f3 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 11 Sep 2026 13:03:13 +0200
Subject: [PATCH 14/20] [PyTorch] Register one custom operation at a time

Make register_custom_op accept one argument type, implementation, and fake implementation. Let FusibleOperation call it independently for forward and backward. Preserve nested tensor results through the same adapter in either direction.

Validation: 13 focused compile tests passed; standalone eager/fullgraph op ran without a backward registration; changed-file pylint passed.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/dynamo/custom_op.py               | 126 +++++-------------
 transformer_engine/pytorch/ops/op.py          |  25 ++--
 2 files changed, 47 insertions(+), 104 deletions(-)

diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index db14a22bef..c067d2b9b8 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -8,14 +8,14 @@
 ``torch.compile(fullgraph=True)`` traces them as single graph nodes.
 ``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``).
+``register_custom_op`` registers one op without autograd; ``ops/fuser.py`` uses
+separate registrations for forward and backward and drives autograd itself.
 
 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 an ``(output, aux)`` tuple; the autograd-wired
-API keeps its saved-tensor and context-metadata contract.
+The autograd-free API preserves nested tensor results; 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[]``.
@@ -1288,115 +1288,51 @@ def _run_forward(
 
 
 # --------------------------------------------------------------------------- #
-# Op registration: the autograd-free pair, and the autograd-wired variant
+# Op registration: a single autograd-free op, 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: Optional[int] = None,
-) -> 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.
+    arg_type: type,
+    impl: Callable[[Any], Any],
+    fake_impl: Callable[[Any], Any],
+) -> Optional[Callable[[Any], Any]]:
+    """Register one custom op without autograd wiring.
 
-    Implementations return nested tuples/lists of tensors or None. Fake
-    implementations return the same structure with TensorSpec leaves. Forward
-    outputs must be fresh tensors; context saving remains the caller's job.
-    num_grad_inputs, if given, counts flattened backward result leaves.
+    arg_type is a dataclass defining the input schema. Results may be nested
+    tuples/lists of fresh tensors or None; fake_impl mirrors them with TensorSpec.
+    The returned callable takes an args instance and preserves the result structure.
+    Returns None if registration fails, so callers can fall back to eager.
+    """
 
-    Returns (forward_fn, backward_fn), preserving each implementation's result
-    structure, including per-basic-op extra outputs and gradients.
+    def pack_result(result):
+        values, _ = tree_flatten(result)
+        return [tensor for value in values for tensor in _flatten_value(value)]
 
-    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,
+        op = _register_op(
+            name=op_name,
+            arg_type=arg_type,
+            impl=impl,
+            fake_impl=fake_impl,
+            pack_result=pack_result,
+            flatten_in_body=True,
         )
     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})"
+            f"could not register custom op '{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: Optional[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):
-            values, _ = tree_flatten(impl(args))
-            return (*values, (), None)
-
-        return wrapped
-
-    def adapt_backward(impl):
-        def wrapped(args):
-            values, _ = tree_flatten(impl(args))
-            return tuple(values)
-
-        return wrapped
-
-    fwd_op = _register_forward_op(
-        name=op_name,
-        arg_type=fwd_arg_type,
-        impl=adapt_forward(fwd_impl),
-        fake_impl=adapt_forward(fwd_fake_impl),
-    )
-    bwd_op = _register_backward_op(
-        name=f"{op_name}_backward",
-        arg_type=bwd_arg_type,
-        impl=adapt_backward(bwd_impl),
-        fake_impl=adapt_backward(bwd_fake_impl),
-        num_grad_inputs=num_grad_inputs,
-    )
-
-    def forward_fn(fwd_args):
-        spec_args = _spec_view(fwd_args, fwd_op.plan.tensor_field_names())
-        specs, structure = tree_flatten(fwd_fake_impl(spec_args))
+    def call(args):
+        spec_args = _spec_view(args, op.plan.tensor_field_names())
+        specs, structure = tree_flatten(fake_impl(spec_args))
         out_plan = _OutputPlan.parse((*specs, (), None))
-        return tree_unflatten(out_plan.user_outputs(fwd_op(fwd_args)), structure)
+        return tree_unflatten(out_plan.user_outputs(op(args)), structure)
 
-    def backward_fn(bwd_args):
-        spec_args = _spec_view(bwd_args, bwd_op.plan.tensor_field_names())
-        _, structure = tree_flatten(bwd_fake_impl(spec_args))
-        grads = [_decode_none(t) for t in bwd_op(bwd_args)]
-        return tree_unflatten(grads, structure)
-
-    return forward_fn, backward_fn
+    return call
 
 
 def register_custom_op_with_autograd(
diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py
index 69df67200d..5c60f3db9a 100644
--- a/transformer_engine/pytorch/ops/op.py
+++ b/transformer_engine/pytorch/ops/op.py
@@ -60,7 +60,7 @@ def save_for_backward(self, *tensors: Optional[torch.Tensor]) -> None:
 class FusibleOperation(torch.nn.Module, metaclass=abc.ABCMeta):
     """Tensor operation supported by the operation fuser"""
 
-    # One custom-op registration per operation class.
+    # Custom ops are registered once per operation class.
     fwd_args_type: Optional[type] = None
     bwd_args_type: Optional[type] = None
     # Supported read-only forward kwargs; no gradients.
@@ -78,14 +78,21 @@ def _register_compile_ops(cls) -> None:
         for name in ("fwd_args_type", "bwd_args_type"):
             if not dataclasses.is_dataclass(getattr(cls, name)):
                 raise TypeError(f"{cls.__name__}.{name} must be a dataclass")
-        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,
+        name = cls.__name__.lower()
+        forward = register_custom_op(
+            op_name=name,
+            arg_type=cls.fwd_args_type,
+            impl=cls.forward_compute,
+            fake_impl=cls.forward_fake,
+        )
+        backward = register_custom_op(
+            op_name=f"{name}_backward",
+            arg_type=cls.bwd_args_type,
+            impl=cls.backward_compute,
+            fake_impl=cls.backward_fake,
+        )
+        cls.compile_ops = (
+            (forward, backward) if forward is not None and backward is not None else None
         )
 
     def compile_unsupported_reason(self) -> Optional[str]:

From b648a57588ae8f440ce83d6f4ed18774cf28ac67 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 11 Sep 2026 13:10:49 +0200
Subject: [PATCH 15/20] [PyTorch] Inline single-use custom op registration
 helpers

Inline forward/backward registration and forward result reconstruction in the autograd-wired adapter. Preserve registration options and gradient validation while removing three single-use helpers.

Validation: 13 focused ops/Linear compile tests passed; changed-file pylint and formatting checks passed.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/dynamo/custom_op.py               | 64 +++++--------------
 1 file changed, 15 insertions(+), 49 deletions(-)

diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index c067d2b9b8..d8efbaf4fa 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -1244,49 +1244,6 @@ def _register_op(
     )
 
 
-def _register_forward_op(
-    *, name: str, arg_type: type, impl: Callable[[Any], Any], fake_impl: Callable[[Any], Any]
-) -> _RegisteredOp:
-    return _register_op(
-        name=name,
-        arg_type=arg_type,
-        impl=impl,
-        fake_impl=fake_impl,
-        pack_result=_pack_fwd_result,
-        flatten_in_body=True,
-    )
-
-
-def _register_backward_op(
-    *,
-    name: str,
-    arg_type: type,
-    impl: Callable[[Any], Any],
-    fake_impl: Callable[[Any], Any],
-    num_grad_inputs: Optional[int],
-) -> _RegisteredOp:
-    # Pass-through body: a subclass input reaches the base op through the
-    # dispatch rule, never through the wrapper body.
-    qualname = f"{_TE_OP_NAMESPACE}::{name}_base"
-    return _register_op(
-        name=name,
-        arg_type=arg_type,
-        impl=impl,
-        fake_impl=fake_impl,
-        pack_result=lambda g: _pack_bwd_result(g, num_grad_inputs, qualname),
-        flatten_in_body=False,
-    )
-
-
-def _run_forward(
-    fwd_op: _RegisteredOp, fwd_fake_impl: Callable[[Any], Tuple[Any, ...]], fwd_args: Any
-) -> Tuple[_OutputPlan, List[torch.Tensor]]:
-    """Run the forward op on ``fwd_args``: its output plan and flat payload."""
-    spec_obj = _spec_view(fwd_args, fwd_op.plan.tensor_field_names())
-    out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj))
-    return out_plan, fwd_op(fwd_args)
-
-
 # --------------------------------------------------------------------------- #
 # Op registration: a single autograd-free op, and the autograd-wired variant
 # --------------------------------------------------------------------------- #
@@ -1440,15 +1397,23 @@ def _register_custom_op_with_autograd_impl(
     if missing:
         raise ValueError(f"input_tensors_for_grad names not in {fwd_arg_type.__name__}: {missing}")
 
-    fwd_op = _register_forward_op(
-        name=op_name, arg_type=fwd_arg_type, impl=fwd_impl, fake_impl=fwd_fake_impl
+    fwd_op = _register_op(
+        name=op_name,
+        arg_type=fwd_arg_type,
+        impl=fwd_impl,
+        fake_impl=fwd_fake_impl,
+        pack_result=_pack_fwd_result,
+        flatten_in_body=True,
     )
-    bwd_op = _register_backward_op(
+    bwd_qualname = f"{_TE_OP_NAMESPACE}::{op_name}_backward_base"
+    num_grad_inputs = len(input_tensors_for_grad)
+    bwd_op = _register_op(
         name=f"{op_name}_backward",
         arg_type=bwd_arg_type,
         impl=bwd_impl,
         fake_impl=bwd_fake_impl,
-        num_grad_inputs=len(input_tensors_for_grad),
+        pack_result=lambda grads: _pack_bwd_result(grads, num_grad_inputs, bwd_qualname),
+        flatten_in_body=False,
     )
 
     autograd_common = {
@@ -1464,8 +1429,9 @@ def _register_custom_op_with_autograd_impl(
     )
 
     def forward_fn(fwd_args):
-        out_plan, payload = _run_forward(fwd_op, fwd_fake_impl, fwd_args)
-        outputs = out_plan.user_outputs(payload)
+        spec_args = _spec_view(fwd_args, fwd_op.plan.tensor_field_names())
+        out_plan = _OutputPlan.parse(fwd_fake_impl(spec_args))
+        outputs = out_plan.user_outputs(fwd_op(fwd_args))
         if len(outputs) == 1:
             return outputs[0]
         return tuple(outputs)

From dd41bbf3bc4b399223522b7a5dcec7d262c7cc2d Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 11 Sep 2026 13:27:19 +0200
Subject: [PATCH 16/20] [PyTorch] Use uniform quantized-input handling in
 custom op wrappers

Remove flatten_in_body and always give wrapper bodies the same quantized-input slot adapter used by their dispatch rules. Already packed storage and ordinary tensors pass through unchanged.

Validation: test_torch_compile.py 126 passed, 46 skipped, 1 xpassed. Six standalone autograd cases passed for ordinary tensors, FP8 wrappers, and bare FP8 storage in eager and Inductor. Pylint and repository Black checks passed.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/dynamo/custom_op.py | 11 ++---------
 1 file changed, 2 insertions(+), 9 deletions(-)

diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index d8efbaf4fa..a56b763db8 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -1198,13 +1198,9 @@ def _register_op(
     impl: Callable[[Any], Any],
     fake_impl: Callable[[Any], Any],
     pack_result: Callable[[Any], List[torch.Tensor]],
-    flatten_in_body: bool,
 ) -> _RegisteredOp:
     """Define one two-tier custom op: the base kernel, the wrapper op that lets
     ``QuantizedTensor`` subclasses be inputs, and the passthrough registrations.
-
-    ``flatten_in_body`` also flattens subclass inputs inside the wrapper body,
-    not only through the ``register_torch_dispatch`` rules.
     """
     plan = _parse_arg_type(arg_type)
     schema = f"{plan.schema_str} -> Tensor[]"
@@ -1225,8 +1221,8 @@ def _register_op(
         wrapper_op_name=name,
         schema_str=schema,
         base_op=base_op,
-        slot_offsets=slot_offsets if flatten_in_body else (),
-        subclasses=subclasses if flatten_in_body else (),
+        slot_offsets=slot_offsets,
+        subclasses=subclasses,
     )
     wrapper_op = getattr(namespace, name)
 
@@ -1275,7 +1271,6 @@ def pack_result(result):
             impl=impl,
             fake_impl=fake_impl,
             pack_result=pack_result,
-            flatten_in_body=True,
         )
     except (ImportError, AttributeError, RuntimeError, TypeError) as e:
         record_compile_disabled(
@@ -1403,7 +1398,6 @@ def _register_custom_op_with_autograd_impl(
         impl=fwd_impl,
         fake_impl=fwd_fake_impl,
         pack_result=_pack_fwd_result,
-        flatten_in_body=True,
     )
     bwd_qualname = f"{_TE_OP_NAMESPACE}::{op_name}_backward_base"
     num_grad_inputs = len(input_tensors_for_grad)
@@ -1413,7 +1407,6 @@ def _register_custom_op_with_autograd_impl(
         impl=bwd_impl,
         fake_impl=bwd_fake_impl,
         pack_result=lambda grads: _pack_bwd_result(grads, num_grad_inputs, bwd_qualname),
-        flatten_in_body=False,
     )
 
     autograd_common = {

From ff17443090274e56db8b13b21d971743197ce480 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 11 Sep 2026 13:58:19 +0200
Subject: [PATCH 17/20] [PyTorch] Consolidate ops compile tests around shared
 affine fixtures

Replace duplicated scale operations with one affine op and use a fixed fused contract. Compare outputs and random-gradient backward results against native PyTorch autograd. Parameterize FP32/BF16 pipeline cases and assert custom-op graph selection and reuse across tensor kwargs.

Reduce the ops test section from 497 to 297 lines. Use exactly representable BF16 products to avoid eager versus fused-reduction rounding differences. Validation: test_torch_compile.py 128 passed, 46 skipped, 1 xpassed; repository Black and diff checks passed.

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

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index c1b3666d8d..9c7fe9062a 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2355,312 +2355,244 @@ def fn(inp):
 
 
 @dataclasses.dataclass(slots=True)
-class _ScaleFwdArgs:
-    """Flat, ``self``-free inputs to the test operation's forward."""
-
+class _AffineFwdArgs:
     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 have a custom op. It is the
-    smallest operation that still has a parameter gradient and a saved tensor.
-    """
-
-    fwd_args_type = _ScaleFwdArgs
-    bwd_args_type = _ScaleBwdArgs
-
-    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 forward_setup_context(self, basic_op_ctxs, args, aux):
-        del aux
-        ctx = basic_op_ctxs[0]
-        ctx.save_for_backward(args.input_, args.scale)
-
-    def pack_forward_args(
-        self,
-        basic_op_ctxs,
-        input_,
-        *,
-        basic_op_extra_inputs,
-        prev_op_grad_output_quantizer,
-        next_op_input_quantizer,
-        basic_op_kwargs,
-    ):
-        return _ScaleFwdArgs(input_=input_, scale=self.scale)
-
-    def pack_backward_args(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs):
-        x, scale = basic_op_ctxs[0].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].fuser_backward(
-            [basic_op_ctxs[1]], grad_output, basic_op_grad_extra_outputs=[()]
-        )
-        dx, grad_params_0, _ = self.basic_ops[0].fuser_backward(
-            [basic_op_ctxs[0]], dx, basic_op_grad_extra_outputs=[()]
-        )
-        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
+    weight: torch.Tensor
+    gain: float = 1.0
+    offset: Union[torch.Tensor, QuantizedTensorStorage] = None
 
 
 @dataclasses.dataclass(slots=True)
-class _ScaleKwargsFwdArgs:
-    """Flat inputs to the kwarg-taking test operation's forward."""
-
+class _AffineBwdArgs:
+    grad_output: torch.Tensor
     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
+    weight: torch.Tensor
+    gain: float
 
 
-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.
-    """
+class _AffineOp(BasicOperation):
+    """Learnable scale with read-only gain and offset kwargs."""
 
-    fwd_args_type = _ScaleKwargsFwdArgs
-    bwd_args_type = _ScaleKwargsBwdArgs
-    fwd_kwarg_names = ("extra_scale", "offset")
+    fwd_args_type = _AffineFwdArgs
+    bwd_args_type = _AffineBwdArgs
+    fwd_kwarg_names = ("gain", "offset")
 
-    def __init__(self, *, device: str = "cuda", dtype: torch.dtype = torch.bfloat16) -> None:
+    def __init__(self, weight=2.0, dtype=torch.float32):
         super().__init__()
-        self.scale = torch.nn.Parameter(torch.full((), 2.0, device=device, dtype=dtype))
+        self.weight = torch.nn.Parameter(torch.tensor(weight, dtype=dtype, device="cuda"))
 
     @classmethod
     def forward_compute(cls, args):
+        output = args.input_ * args.weight * args.gain
         offset = args.offset
-        if isinstance(offset, QuantizedTensor):
+        if isinstance(offset, QuantizedTensorStorage):
             offset = offset.dequantize()
-        out = args.input_ * args.scale * args.extra_scale + offset
-        return out, [()], ()
+        if offset is not None:
+            output = output + offset
+        return output, [()], ()
 
     @classmethod
     def forward_fake(cls, args):
-        x = args.input_
-        return TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device), [()], ()
+        return args.input_, [()], ()
 
     @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,)],
-            [()],
-        )
+        return dy * args.weight * args.gain, [((dy * args.input_).sum() * args.gain,)], [()]
 
     @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),)],
-            [()],
-        )
+        return dy, [(TensorSpec(shape=(), dtype=dy.dtype, device=dy.device),)], [()]
+
+    def pack_forward_args(self, basic_op_ctxs, input_, *, basic_op_kwargs, **unused):
+        return _AffineFwdArgs(input_, self.weight, **basic_op_kwargs[0])
 
     def forward_setup_context(self, basic_op_ctxs, args, aux):
-        del aux
         ctx = basic_op_ctxs[0]
-        ctx.save_for_backward(args.input_, args.scale)
-        ctx.extra_scale = args.extra_scale
-
-    def pack_forward_args(
-        self,
-        basic_op_ctxs,
-        input_,
-        *,
-        basic_op_extra_inputs,
-        prev_op_grad_output_quantizer,
-        next_op_input_quantizer,
-        basic_op_kwargs,
-    ):
-        kwargs = basic_op_kwargs[0]
-        offset = kwargs.get("offset")
-        if offset is None:
-            offset = torch.zeros((), device=input_.device, dtype=input_.dtype)
-        return _ScaleKwargsFwdArgs(
-            input_=input_,
-            scale=self.scale,
-            extra_scale=kwargs.get("extra_scale", 1.0),
-            offset=offset,
-        )
+        ctx.save_for_backward(args.input_, args.weight)
+        ctx.gain = args.gain
 
-    def pack_backward_args(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs):
+    def pack_backward_args(self, basic_op_ctxs, grad_output, **unused):
         ctx = basic_op_ctxs[0]
-        x, scale = ctx.saved_tensors
-        return _ScaleKwargsBwdArgs(
-            grad_output=grad_output,
-            saved_input=x,
-            scale=scale,
-            extra_scale=ctx.extra_scale,
-        )
+        return _AffineBwdArgs(grad_output, *ctx.saved_tensors, ctx.gain)
+
+
+class _BackwardAffinePair(te.ops.FusedOperation):
+    def fuser_backward(self, basic_op_ctxs, grad_output, **unused):
+        grads = []
+        for op, ctx in reversed(list(zip(self.basic_ops, basic_op_ctxs))):
+            grad_output, params, _ = op.fuser_backward(
+                [ctx], grad_output, basic_op_grad_extra_outputs=[()]
+            )
+            grads.insert(0, params[0])
+        return grad_output, grads, [(), ()]
 
 
 @dataclasses.dataclass(slots=True)
-class _ScalePairFwdArgs:
+class _AffinePairFwdArgs:
     input_: torch.Tensor
-    scale0: torch.Tensor
-    scale1: torch.Tensor
-    extra_input: torch.Tensor = None
+    weight0: torch.Tensor
+    weight1: torch.Tensor
+    residual: torch.Tensor
 
 
 @dataclasses.dataclass(slots=True)
-class _ScalePairBwdArgs:
+class _AffinePairBwdArgs:
     grad_output: torch.Tensor
     input_: torch.Tensor
     intermediate: torch.Tensor
-    scale0: torch.Tensor
-    scale1: torch.Tensor
-    grad_extra_output: torch.Tensor = None
+    weight0: torch.Tensor
+    weight1: torch.Tensor
+    grad_extra_output: torch.Tensor
 
 
-class _ScalePair(te.ops.FusedOperation):
-    """Two scales, optionally with a residual input and squared intermediate output."""
+class _AffinePair(te.ops.FusedOperation):
+    """Two scales with a residual input and a squared intermediate output."""
 
-    fwd_args_type = _ScalePairFwdArgs
-    bwd_args_type = _ScalePairBwdArgs
+    fwd_args_type = _AffinePairFwdArgs
+    bwd_args_type = _AffinePairBwdArgs
 
     @classmethod
     def forward_compute(cls, args):
-        intermediate = args.input_ * args.scale0
-        output = intermediate * args.scale1
-        extras = ()
-        if args.extra_input is not None:
-            output = output + args.extra_input
-            extras = (intermediate.square(), None)
-        return output, [(), extras], (intermediate,)
+        intermediate = args.input_ * args.weight0
+        output = intermediate * args.weight1 + args.residual
+        return output, [(), (intermediate.square(), None)], (intermediate,)
 
     @classmethod
     def forward_fake(cls, args):
-        x = args.input_
-        spec = TensorSpec(shape=tuple(x.shape), dtype=x.dtype, device=x.device)
-        extras = (spec, None) if args.extra_input is not None else ()
-        return spec, [(), extras], (spec,)
+        return args.input_, [(), (args.input_, None)], (args.input_,)
 
     @classmethod
     def backward_compute(cls, args):
         dy = args.grad_output
-        du = dy * args.scale1
-        extras = ()
-        if args.grad_extra_output is not None:
-            du = du + 2 * args.intermediate * args.grad_extra_output
-            extras = (dy.clone(),)
+        du = dy * args.weight1 + 2 * args.intermediate * args.grad_extra_output
         return (
-            du * args.scale0,
+            du * args.weight0,
             [((du * args.input_).sum(),), ((dy * args.intermediate).sum(),)],
-            [(), extras],
+            [(), (dy.clone(),)],
         )
 
     @classmethod
     def backward_fake(cls, args):
         dy = args.grad_output
-        spec = TensorSpec(shape=tuple(dy.shape), dtype=dy.dtype, device=dy.device)
         scalar = TensorSpec(shape=(), dtype=dy.dtype, device=dy.device)
-        extras = (spec,) if args.grad_extra_output is not None else ()
-        return spec, [(scalar,), (scalar,)], [(), extras]
+        return dy, [(scalar,), (scalar,)], [(), (dy,)]
 
     def pack_forward_args(self, basic_op_ctxs, input_, *, basic_op_extra_inputs, **unused):
-        extras = basic_op_extra_inputs[1]
-        return _ScalePairFwdArgs(
+        return _AffinePairFwdArgs(
             input_,
-            self.basic_ops[0].scale,
-            self.basic_ops[1].scale,
-            extras[0] if extras else None,
+            self.basic_ops[0].weight,
+            self.basic_ops[1].weight,
+            basic_op_extra_inputs[1][0],
         )
 
     def forward_setup_context(self, basic_op_ctxs, args, aux):
-        basic_op_ctxs[0].save_for_backward(args.input_, args.scale0)
-        basic_op_ctxs[1].save_for_backward(aux[0], args.scale1)
+        basic_op_ctxs[0].save_for_backward(args.input_, args.weight0)
+        basic_op_ctxs[1].save_for_backward(aux[0], args.weight1)
 
     def pack_backward_args(self, basic_op_ctxs, grad_output, *, basic_op_grad_extra_outputs):
-        x, scale0 = basic_op_ctxs[0].saved_tensors
-        intermediate, scale1 = basic_op_ctxs[1].saved_tensors
-        extras = basic_op_grad_extra_outputs[1]
-        return _ScalePairBwdArgs(
-            grad_output, x, intermediate, scale0, scale1, extras[0] if extras else None
+        x, weight0 = basic_op_ctxs[0].saved_tensors
+        intermediate, weight1 = basic_op_ctxs[1].saved_tensors
+        return _AffinePairBwdArgs(
+            grad_output, x, intermediate, weight0, weight1, basic_op_grad_extra_outputs[1][0]
         )
 
 
+def _compile_with_graphs(fn):
+    graphs = []
+
+    def backend(graph, inputs):
+        graphs.append(graph)
+        return torch._dynamo.lookup_backend("inductor")(graph, inputs)
+
+    return torch.compile(fn, fullgraph=True, backend=backend), graphs
+
+
+def _assert_custom_ops(graphs, name, present=True):
+    targets = {
+        str(node.target).removesuffix(".default").removesuffix("_base")
+        for graph in graphs
+        for module in graph.modules()
+        if isinstance(module, torch.fx.GraphModule)
+        for node in module.graph.nodes
+        if node.op == "call_function"
+    }
+    for suffix in ("", "_backward"):
+        assert (f"transformer_engine_compile.{name}{suffix}" in targets) == present, targets
+
+
+def _check_ops(fn, model, x, dy, kwargs=None):
+    """Compare the pipeline with native PyTorch forward and autograd."""
+    params = tuple(model.parameters())
+    kwargs = kwargs or {}
+    offset = kwargs.get("offset", 0)
+    if isinstance(offset, QuantizedTensorStorage):
+        offset = offset.dequantize()
+    reference = x
+    for i, weight in enumerate(params):
+        reference = reference * weight
+        if i == 0:
+            reference = reference * kwargs.get("gain", 1.0) + offset
+    output = fn(x, op_kwargs={0: kwargs})
+    actual = output, torch.autograd.grad(output, (x, *params), dy)
+    expected = reference, torch.autograd.grad(reference, (x, *params), dy)
+    torch.testing.assert_close(actual, expected)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize(
+    "case,reason",
+    [
+        ("single", None),
+        ("multi", "several operations"),
+        ("backward_fusion", "backward fusion"),
+        ("legacy", "without a custom op"),
+    ],
+)
+def test_te_ops_pipeline(case, reason, dtype, monkeypatch):
+    torch._dynamo.reset()
+    if case == "backward_fusion":
+
+        def fuse(ops, **unused):
+            if len(ops) == 2 and all(isinstance(op, _AffineOp) for op in ops):
+                return [_BackwardAffinePair(ops)]
+            return ops
+
+        monkeypatch.setattr(OperationFuser, "backward_fusion_functions", [fuse])
+    ops = (
+        [te.ops.Identity()]
+        if case == "legacy"
+        else [_AffineOp(float(i + 2), dtype) for i in range(1 if case == "single" else 2)]
+    )
+    model = te.ops.Sequential(*ops)
+    compiled, graphs = _compile_with_graphs(model)
+    # Keep products exact in BF16 across eager and fused reductions.
+    x = (torch.randint(-8, 9, (8, 16), device="cuda").to(dtype) / 8).requires_grad_()
+    dy = torch.randint_like(x, -8, 9) / 8
+    with pytest.warns(UserWarning, match=reason) if reason else contextlib.nullcontext():
+        _check_ops(compiled, model, x, dy)
+    _check_ops(model, model, x, dy)
+    _assert_custom_ops(graphs, "_affineop", present=reason is None)
+
+
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-@pytest.mark.parametrize("with_extras", [False, True])
 @pytest.mark.parametrize("use_custom_ops", [False, True])
-def test_te_ops_fused_compute_contract(with_extras, use_custom_ops):
+def test_te_ops_fused_compute_contract(use_custom_ops):
     """Exercise the shared interface directly; pipeline fusion stays gated under compile."""
     torch._dynamo.reset()
-    ops = [_ScaleOp(dtype=torch.float32), _ScaleOp(dtype=torch.float32)]
-    with torch.no_grad():
-        ops[1].scale.fill_(3.0)
-    fused = _ScalePair(ops)
+    fused = _AffinePair([_AffineOp(2.0), _AffineOp(3.0)])
     x = torch.randn(8, 16, device="cuda", requires_grad=True)
     residual = torch.randn_like(x, requires_grad=True)
-    dy = torch.randn_like(x)
-    dextra = torch.randn_like(x)
+    dy, dextra = torch.randn_like(x), torch.randn_like(x)
 
     def run(input_, extra_input, grad_output, grad_extra_output):
         ctxs = [OperationContext(), OperationContext()]
-        output, extra_outputs = fused.fuser_forward(
+        output, extras = fused.fuser_forward(
             ctxs,
             input_,
-            basic_op_extra_inputs=[(), (extra_input,) if with_extras else ()],
+            basic_op_extra_inputs=[(), (extra_input,)],
             prev_op_grad_output_quantizer=None,
             next_op_input_quantizer=None,
             basic_op_kwargs=[{}, {}],
@@ -2671,179 +2603,47 @@ def run(input_, extra_input, grad_output, grad_extra_output):
         grads = fused.fuser_backward(
             ctxs,
             grad_output,
-            basic_op_grad_extra_outputs=[(), (grad_extra_output, None) if with_extras else ()],
+            basic_op_grad_extra_outputs=[(), (grad_extra_output, None)],
             use_custom_ops=use_custom_ops,
         )
-        return output, extra_outputs, grads
+        return output, extras, grads
 
+    graphs = []
     if use_custom_ops:
-        assert fused.compile_ops is not None
-        run = torch.compile(run, fullgraph=True)
+        run, graphs = _compile_with_graphs(run)
     with torch.no_grad():
-        output, extras, (dx, dparams, dextras) = run(x, residual, dy, dextra)
-
-    intermediate = x * ops[0].scale
-    reference = intermediate * ops[1].scale
-    loss = (reference * dy).sum()
-    if with_extras:
-        reference = reference + residual
-        loss = (reference * dy).sum() + (intermediate.square() * dextra).sum()
-    inputs = [x, ops[0].scale, ops[1].scale] + ([residual] if with_extras else [])
-    expected = torch.autograd.grad(loss, inputs)
-    torch.testing.assert_close(output, reference)
-    torch.testing.assert_close(dx, expected[0])
-    assert len(dparams) == 2 and all(len(group) == 1 for group in dparams)
-    torch.testing.assert_close(dparams[0][0], expected[1])
-    torch.testing.assert_close(dparams[1][0], expected[2])
-    assert extras[0] == () and dextras[0] == ()
-    if with_extras:
-        assert len(extras[1]) == 2 and extras[1][1] is None
-        torch.testing.assert_close(extras[1][0], intermediate.square())
-        torch.testing.assert_close(dextras[1][0], expected[3])
-    else:
-        assert extras[1] == () and dextras[1] == ()
-
-
-def _assert_sequential_matches_eager(make_model, base, op_kwargs_seq=(None,), *, graphs=None):
-    """Run a Sequential eagerly and compiled on identical inputs; compare both
-    the output and every parameter gradient.
-
-    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()
-    backend = "inductor"
-    if graphs is not None:
-
-        def backend(graph, inputs):
-            graphs.append(graph)
-            return torch._dynamo.lookup_backend("inductor")(graph, inputs)
-
-    compiled = torch.compile(compiled_model, fullgraph=True, backend=backend)
-
-    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")
-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()
-    base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
-    graphs = []
-    _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp()), base, graphs=graphs)
-    targets = {
-        str(node.target).removesuffix(".default")
-        for graph in graphs
-        for module in graph.modules()
-        if isinstance(module, torch.fx.GraphModule)
-        for node in module.graph.nodes
-        if node.op == "call_function"
-    }
-    for name in ("_scaleop", "_scaleop_backward"):
-        assert targets & {
-            f"transformer_engine_compile.{name}",
-            f"transformer_engine_compile.{name}_base",
-        }, targets
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-def test_te_ops_multi_op_group_uses_eager_implementations():
-    """A group of several operations is gated onto the eager implementations."""
-    torch._dynamo.reset()
-    base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
-    with pytest.warns(UserWarning, match="several operations"):
-        _assert_sequential_matches_eager(lambda: te.ops.Sequential(_ScaleOp(), _ScaleOp()), base)
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
-def test_te_ops_backward_fusion_uses_eager_implementations():
-    """A backward fusion prevents the group from using basic-op custom ops."""
-    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")
-def test_te_ops_unsupported_group_still_compiles_eagerly():
-    """An operation without a custom op runs its eager implementation.
-
-    Note that this is not a fallback: under ``fullgraph=True`` there is no
-    leaving the graph, so the pipeline is traced either way and only the choice
-    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()
-    assert te.ops.Identity().compile_unsupported_reason() is not None
-
-    base = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda")
-    _assert_sequential_matches_eager(lambda: te.ops.Sequential(te.ops.Identity()), base)
+        actual = run(x, residual, dy, dextra)
+    weight0, weight1 = tuple(fused.parameters())
+    intermediate = x * weight0
+    output, extra = intermediate * weight1 + residual, intermediate.square()
+    dx, dw0, dw1, dr = torch.autograd.grad(
+        (output, extra), (x, weight0, weight1, residual), (dy, dextra)
+    )
+    expected = output, [(), (extra, None)], (dx, [(dw0,), (dw1,)], [(), (dr,)])
+    torch.testing.assert_close(actual, expected)
+    _assert_custom_ops(graphs, "_affinepair", present=use_custom_ops)
 
 
 @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():
-    """A tensor forward kwarg reaches the operation through its custom op.
-
-    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()
+    model = te.ops.Sequential(_AffineOp())
+    compiled, graphs = _compile_with_graphs(model)
     quantizer = Float8CurrentScalingQuantizer(
-        fp8_dtype=tex.DType.kFloat8E4M3,
-        device=torch.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: {"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}},
-        ),
+        fp8_dtype=tex.DType.kFloat8E4M3, device=torch.device("cuda")
     )
+    x = torch.randn(8, 16, device="cuda", requires_grad=True)
+    dy = torch.randn_like(x)
+    for i, value in enumerate((0.5, 0.5, 1.5, -0.5)):
+        offset = quantizer(torch.full((16,), value, device="cuda"))
+        _check_ops(compiled, model, x, dy, {"offset": offset})
+        if i == 1:
+            warmed_graph_count = len(graphs)
+        elif i > 1:
+            assert len(graphs) == warmed_graph_count, "Tensor kwarg triggered recompilation"
+    _assert_custom_ops(graphs, "_affineop")
+    with pytest.warns(UserWarning, match="non-tensor keyword arguments"):
+        for gain in (3.0, 5.0):
+            _check_ops(compiled, model, x, dy, {"gain": gain})
+    _assert_custom_ops(graphs[-1:], "_affineop", present=False)

From 465ffaa691d47d368f7459c996c1bb16f9ecb695 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 11 Sep 2026 14:41:16 +0200
Subject: [PATCH 18/20] Remove redundant forward kwarg name declarations

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py     | 1 -
 transformer_engine/pytorch/ops/fuser.py | 7 -------
 transformer_engine/pytorch/ops/op.py    | 2 --
 3 files changed, 10 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 9c7fe9062a..f8f5c4de08 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2375,7 +2375,6 @@ class _AffineOp(BasicOperation):
 
     fwd_args_type = _AffineFwdArgs
     bwd_args_type = _AffineBwdArgs
-    fwd_kwarg_names = ("gain", "offset")
 
     def __init__(self, weight=2.0, dtype=torch.float32):
         super().__init__()
diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py
index 10cbe6493f..e5e40e2945 100644
--- a/transformer_engine/pytorch/ops/fuser.py
+++ b/transformer_engine/pytorch/ops/fuser.py
@@ -750,13 +750,6 @@ def _custom_ops_unsupported_reason(
         if self._num_basic_ops != 1:
             return "a group of several operations"
         for op, kwargs in zip(self._basic_ops, basic_op_kwargs, strict=True):
-            # A kwarg an operation declares is resolved into its args container
-            # like any other config. Anything else -- notably the preallocated
-            # buffers of the grouped operations -- is written to by the op, and a
-            # custom op may not mutate a tensor from an enclosing scope.
-            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
diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py
index 5c60f3db9a..5968d66430 100644
--- a/transformer_engine/pytorch/ops/op.py
+++ b/transformer_engine/pytorch/ops/op.py
@@ -63,8 +63,6 @@ class FusibleOperation(torch.nn.Module, metaclass=abc.ABCMeta):
     # Custom ops are registered once per operation class.
     fwd_args_type: Optional[type] = None
     bwd_args_type: Optional[type] = None
-    # Supported read-only forward kwargs; no gradients.
-    fwd_kwarg_names: tuple[str, ...] = ()
     # (forward_fn, backward_fn), or None if the operation cannot be compiled.
     compile_ops: Optional[tuple[Callable[..., Any], Callable[..., Any]]] = None
 

From 464907c5d9935d10501d86411af9f040448bc079 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 11 Sep 2026 15:06:37 +0200
Subject: [PATCH 19/20] [PyTorch] Reject custom-op collisions and restore
 compiled storage safely

Raise on duplicate custom-op names instead of replacing registered kernels. Save bare quantized storage as buffers and metadata while compiling, then reconstruct a fresh storage object in backward.

Validation: 8 focused regression tests passed; changed-file Black and pylint passed.
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/dynamo/custom_op.py               |  4 ++++
 .../pytorch/quantized_tensor.py               | 20 +++++++++++++++++--
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index a56b763db8..24b3991c57 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -1208,6 +1208,10 @@ def _register_op(
     slot_offsets = plan.tensor_or_quantized_offsets()
     namespace = getattr(torch.ops, _TE_OP_NAMESPACE)
 
+    for registered_name in (name, f"{name}_base"):
+        if hasattr(namespace, registered_name):
+            raise ValueError(f"Custom op '{_TE_OP_NAMESPACE}::{registered_name}' already exists")
+
     base_def = _register_base_op(
         op_name=f"{name}_base",
         schema_str=schema,
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index 7149a5a163..4d64bc43c7 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -282,11 +282,16 @@ def __tensor_unflatten__(
         )
 
 
+class _SavedQuantizedTensor(NamedTuple):
+    inner_names: tuple[str, ...]
+    metadata: Dict[str, Any]
+
+
 def prepare_for_saving(
     *tensors: Union[torch.Tensor, QuantizedTensorStorage],
 ) -> Tuple[
     list[Optional[Union[torch.Tensor, torch.nn.Parameter]]],
-    list[Optional[QuantizedTensorStorage]],
+    list[Optional[Union[QuantizedTensorStorage, _SavedQuantizedTensor]]],
 ]:
     """Prepare tensors for saving. Needed because save_for_backward accepts only
     torch.Tensor/torch.nn.Parameter types, while we want to be able to save
@@ -297,6 +302,10 @@ def prepare_for_saving(
         if tensor is None or isinstance(tensor, torch.Tensor):
             tensor_list.append(tensor)
             tensor_objects_list.append(None)
+        elif torch.compiler.is_compiling():
+            inner_names, metadata = tensor.__tensor_flatten__()
+            tensor_list.extend(getattr(tensor, name) for name in inner_names)
+            tensor_objects_list.append(_SavedQuantizedTensor(tuple(inner_names), metadata))
         else:
             t, t_obj = tensor.prepare_for_saving()
             tensor_list.extend(t)
@@ -306,7 +315,7 @@ def prepare_for_saving(
 
 
 def restore_from_saved(
-    tensors: list[Optional[Union[torch.Tensor, QuantizedTensorStorage]]],
+    tensors: list[Optional[Union[torch.Tensor, QuantizedTensorStorage, _SavedQuantizedTensor]]],
     saved_tensors: list[Optional[Union[torch.Tensor, torch.nn.Parameter]]],
     return_saved_tensors: bool = False,
 ) -> (
@@ -324,6 +333,13 @@ def restore_from_saved(
         if tensor is None or isinstance(tensor, torch.Tensor):
             tensor_objects.append(saved_tensors[0])
             saved_tensors = saved_tensors[1:]
+        elif isinstance(tensor, _SavedQuantizedTensor):
+            count = len(tensor.inner_names)
+            inner = dict(zip(tensor.inner_names, saved_tensors[:count]))
+            tensor_objects.append(
+                QuantizedTensorStorage.__tensor_unflatten__(inner, tensor.metadata, None, None)
+            )
+            saved_tensors = saved_tensors[count:]
         else:
             saved_tensors = tensor.restore_from_saved(saved_tensors)
             tensor_objects.append(tensor)

From 01762bee8b9fbc6e3a0088bfc99660fc2d91fea7 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 11 Sep 2026 15:13:05 +0200
Subject: [PATCH 20/20] [PyTorch] Document custom-op name collision errors

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

diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 24b3991c57..a87374095a 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -1261,7 +1261,8 @@ def register_custom_op(
     arg_type is a dataclass defining the input schema. Results may be nested
     tuples/lists of fresh tensors or None; fake_impl mirrors them with TensorSpec.
     The returned callable takes an args instance and preserves the result structure.
-    Returns None if registration fails, so callers can fall back to eager.
+    Returns None for unsupported registration APIs so callers can fall back to eager.
+    Duplicate operator names raise ValueError.
     """
 
     def pack_result(result):
@@ -1351,9 +1352,9 @@ def register_custom_op_with_autograd(
     ``bwd_arg_type``.
 
     Registration touches experimental ``torch.library`` / opaque-object APIs
-    that may be missing on older PyTorch. If it fails, this warns once and
-    returns ``None`` instead of raising, so callers can fall back to eager under
-    ``torch.compile`` (a graph break) rather than breaking import.
+    that may be missing on older PyTorch. Unavailable or incompatible APIs cause
+    a warning and return ``None``, allowing eager execution under ``torch.compile``
+    (a graph break). Duplicate operator names raise ``ValueError``.
     """
     try:
         return _register_custom_op_with_autograd_impl(