From a46b9b083c3ae6e1bb04aef9652f80b9304a0a17 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 4 Sep 2026 13:05:54 +0200
Subject: [PATCH 1/2] [PyTorch][torch.compile] Differentiable fp8_output and
 fp8_grad for Linear

Quantized op outputs carry a grad-handle slot; _QuantizedOutputFn rebuilds
the wrapper so autograd reaches the op through the handle. Quantized grads
are flattened across the backward op boundary and rebuilt from the bwd fake
specs. QuantizedTensor gets the AOTAutograd tangent coercion hooks, FP8
dequantize becomes a traceable custom op, and QuantizedTensor dispatch
re-dispatches instead of disabling torch dispatch.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 241 ++++++++++++++++--
 .../pytorch/dynamo/custom_op.py               | 141 +++++++---
 transformer_engine/pytorch/module/linear.py   |  41 +--
 .../pytorch/quantized_tensor.py               |  63 ++++-
 .../pytorch/tensor/float8_tensor.py           |  10 +
 .../tensor/storage/float8_tensor_storage.py   |  29 ++-
 6 files changed, 430 insertions(+), 95 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index f8e09d5ce1..fc127b4a97 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -1969,8 +1969,8 @@ def fn(inp):
 def test_te_linear_compile_with_fp8_output(compile_mode):
     """torch.compile of ``te.Linear(..., fp8_output=True)`` under no_grad:
     forward must return a working :class:`Float8Tensor` (exercises the output
-    rewrap path). The differentiable case falls back to eager, so it is not
-    covered here."""
+    rewrap path). The differentiable case is covered by
+    ``test_te_linear_compile_fp8_output_differentiable``."""
     dtype = torch.bfloat16
     device = "cuda"
     fp8_recipe = recipe.Float8CurrentScaling()
@@ -2009,16 +2009,235 @@ def fn(inp):
             )
 
 
+def _run_fp8_io_pair(fn, model, compile_mode, *, consume=None, backward=True):
+    """Run ``fn`` eagerly and compiled on the same input; ``consume`` (eager,
+    outside the graph) maps the output to the loss. Returns ``(eager, compiled)``
+    tuples of ``(out, inp.grad, weight.grad, bias.grad)``."""
+    dtype, device = torch.bfloat16, "cuda"
+    torch.manual_seed(0)
+    base = torch.randn(32, 64, dtype=dtype, device=device)
+
+    def run(f):
+        model.zero_grad(set_to_none=True)
+        inp = base.clone().requires_grad_(True)
+        out = f(inp)
+        if backward:
+            loss = consume(out) if consume is not None else out
+            loss.sum().backward()
+        return out, inp.grad, model.weight.grad, model.bias.grad
+
+    torch._dynamo.reset()
+    if compile_mode == "reduce-overhead":
+        _cudagraph_warmup(
+            lambda x: (consume or (lambda o: o))(fn(x)),
+            base.clone().requires_grad_(True),
+            backward=backward,
+        )
+    ref = run(fn)
+    compiled = torch.compile(fn, fullgraph=True, mode=compile_mode)
+    n_iters = 3 if compile_mode == "reduce-overhead" else 1
+    with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"):
+        for _ in range(n_iters):
+            got = run(compiled)
+    return ref, got
+
+
+def _assert_fp8_io_close(ref, got, *, atol=_EAGER_ATOL, rtol=_EAGER_RTOL):
+    for name, r, g in zip(("out", "inp.grad", "wgrad", "bgrad"), ref, got):
+        if isinstance(r, QuantizedTensor):
+            r = r.dequantize()
+        if isinstance(g, QuantizedTensor):
+            g = g.dequantize()
+        torch.testing.assert_close(g, r, atol=atol, rtol=rtol, msg=lambda m, n=name: f"{n}: {m}")
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("compile_mode", _compile_modes)
+@pytest.mark.parametrize("producer", ["te_linear", "torch_mul"])
+def test_te_linear_compile_fp8_grad(compile_mode, producer):
+    """``fp8_grad=True`` under torch.compile: the quantized dgrad is consumed
+    inside the graph -- by another TE Linear's backward directly, or by a torch
+    op's backward through the traceable FP8 dequantize -- matching eager exactly."""
+    fp8_recipe = recipe.Float8CurrentScaling()
+    model = te.Linear(64, 32, params_dtype=torch.bfloat16, device="cuda")
+    upstream = te.Linear(64, 64, params_dtype=torch.bfloat16, device="cuda")
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            hidden = upstream(inp) if producer == "te_linear" else inp * 2
+            return model(hidden, fp8_grad=True)
+
+    ref, got = _run_fp8_io_pair(fn, model, compile_mode)
+    _assert_fp8_io_close(ref, got)
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("compile_mode", _compile_modes)
+def test_te_linear_compile_fp8_grad_crosses_boundary(compile_mode):
+    """A quantized dgrad for a graph *input* is re-wrapped by AOTAutograd only
+    when the graph already carries a tensor subclass (here the ``fp8_output``);
+    then ``inp.grad`` leaves the graph as a :class:`Float8Tensor`, as in eager."""
+    fp8_recipe = recipe.Float8CurrentScaling()
+    model = te.Linear(64, 32, params_dtype=torch.bfloat16, device="cuda")
+    consumer = te.Linear(32, 16, params_dtype=torch.bfloat16, device="cuda")
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            return model(inp, fp8_output=True, fp8_grad=True)
+
+    def consume(out):
+        with te.autocast(recipe=fp8_recipe):
+            return consumer(out, fp8_grad=True)
+
+    ref, got = _run_fp8_io_pair(fn, model, compile_mode, consume=consume)
+    assert isinstance(ref[1], te.Float8Tensor)
+    assert isinstance(got[1], te.Float8Tensor), type(got[1]).__name__
+    _assert_fp8_io_close(ref, got)
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.xfail(
+    strict=True,
+    reason=(
+        "PyTorch limitation: AOTAutograd only wraps tensor-subclass grads when the graph "
+        "has a subclass input or output, so a quantized dgrad for a plain graph input is "
+        "lifted as a fake constant (assert_no_fake_params_or_buffers)."
+    ),
+)
+def test_te_linear_compile_fp8_grad_plain_graph_boundary_unsupported():
+    fp8_recipe = recipe.Float8CurrentScaling()
+    model = te.Linear(64, 32, params_dtype=torch.bfloat16, device="cuda")
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            return model(inp, fp8_grad=True)
+
+    ref, got = _run_fp8_io_pair(fn, model, "default")
+    _assert_fp8_io_close(ref, got)
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("compile_mode", _compile_modes)
+def test_te_linear_compile_fp8_output_dequantize_in_graph(compile_mode):
+    """Differentiable ``fp8_output=True`` consumed inside the graph by
+    ``.dequantize()`` (traceable FP8 dequantize op); fwd + bwd match eager."""
+    fp8_recipe = recipe.Float8CurrentScaling()
+    model = te.Linear(64, 32, params_dtype=torch.bfloat16, device="cuda")
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            return model(inp, fp8_output=True).dequantize()
+
+    ref, got = _run_fp8_io_pair(fn, model, compile_mode)
+    _assert_fp8_io_close(ref, got)
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("compile_mode", _compile_modes)
+def test_te_linear_compile_fp8_output_differentiable(compile_mode):
+    """Differentiable ``fp8_output=True``: the compiled forward returns a
+    :class:`Float8Tensor` with a grad_fn. Its gradient arriving from outside the
+    graph must be quantized (here from an eager ``fp8_grad=True`` consumer) --
+    AOTAutograd traces the backward with a quantized tangent guess
+    (``__coerce_tangent_metadata__``), so the E5M2 dgrad of the consumer is
+    accepted as-is and numerics match eager exactly."""
+    fp8_recipe = recipe.Float8CurrentScaling()
+    model = te.Linear(64, 32, params_dtype=torch.bfloat16, device="cuda")
+    consumer = te.Linear(32, 16, params_dtype=torch.bfloat16, device="cuda")
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            return model(inp, fp8_output=True)
+
+    def consume(out):
+        with te.autocast(recipe=fp8_recipe):
+            return consumer(out, fp8_grad=True)
+
+    ref, got = _run_fp8_io_pair(fn, model, compile_mode, consume=consume)
+    assert isinstance(got[0], te.Float8Tensor), type(got[0]).__name__
+    assert got[0].requires_grad
+    _assert_fp8_io_close(ref, got)
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_te_linear_compile_fp8_output_tangent_requantized():
+    """A quantized tangent whose FP8 dtype differs from the traced guess (E4M3
+    grads from an ``Format.E4M3`` consumer vs the E5M2 guess) is requantized by
+    ``__coerce_same_metadata_as_tangent__`` instead of failing; numerics then
+    only match to FP8 precision."""
+    fp8_recipe = recipe.Float8CurrentScaling()
+    e4m3_recipe = recipe.Float8CurrentScaling(fp8_format=recipe.Format.E4M3)
+    model = te.Linear(64, 32, params_dtype=torch.bfloat16, device="cuda")
+    consumer = te.Linear(32, 16, params_dtype=torch.bfloat16, device="cuda")
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            return model(inp, fp8_output=True)
+
+    def consume(out):
+        with te.autocast(recipe=e4m3_recipe):
+            return consumer(out, fp8_grad=True)
+
+    ref, got = _run_fp8_io_pair(fn, model, "default", consume=consume)
+    _assert_fp8_io_close(ref, got, atol=0.1, rtol=0.25)
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_te_linear_compile_fp8_output_plain_tangent_unsupported():
+    """PyTorch limitation: a plain-tensor gradient for a quantized graph output
+    (e.g. ``.dequantize()`` outside the graph) has no coercion hook and
+    AOTAutograd raises a tangent-metadata error. Pinned so the failure stays
+    loud and explicit rather than silently wrong."""
+    fp8_recipe = recipe.Float8CurrentScaling()
+    model = te.Linear(64, 32, params_dtype=torch.bfloat16, device="cuda")
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            return model(inp, fp8_output=True)
+
+    torch._dynamo.reset()
+    compiled = torch.compile(fn, fullgraph=True)
+    inp = torch.randn(32, 64, dtype=torch.bfloat16, device="cuda", requires_grad=True)
+    out = compiled(inp)
+    with pytest.raises(RuntimeError, match="tangent"):
+        out.dequantize().sum().backward()
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_te_linear_compile_plain_output_fp8_tangent():
+    """A quantized gradient (eager ``fp8_grad=True`` consumer) arriving for a
+    plain-tensor graph output is dequantized by
+    ``__coerce_same_metadata_as_tangent__`` to match the traced plain tangent."""
+    fp8_recipe = recipe.Float8CurrentScaling()
+    model = te.Linear(64, 32, params_dtype=torch.bfloat16, device="cuda")
+    consumer = te.Linear(32, 16, params_dtype=torch.bfloat16, device="cuda")
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            return model(inp)
+
+    def consume(out):
+        with te.autocast(recipe=fp8_recipe):
+            return consumer(out, fp8_grad=True)
+
+    ref, got = _run_fp8_io_pair(fn, model, "default", consume=consume)
+    _assert_fp8_io_close(ref, got)
+
+
 # Configs rejected by LinearFwdArgs.compile_unsupported_reason() that a
 # single-GPU unit test can construct. Distributed-only reasons (fsdp_group,
 # DistributedWeight) and CPU offloading need machinery this file doesn't have;
 # delayed scaling is a hard error (check_recipe_support), tested separately.
-# Modes: "bwd" = fwd+bwd vs eager; "fwd_grad" = grad-enabled forward only
-# (differentiable fp8_output backward hits a PyTorch limitation: the Float8
-# output crossing the graph-break boundary gets a plain-tensor tangent);
-# "no_grad" = forward under no_grad.
+# Modes: "bwd" = fwd+bwd vs eager; "no_grad" = forward under no_grad.
 _FALLBACK_CASES = [
-    "fp8_output_differentiable",
     "fuse_wgrad_accumulation",
     "delayed_wgrad",
     "quantized_input",
@@ -2034,14 +2253,6 @@ def _fallback_case(case, dtype, device):
         model_kwargs["delay_wgrad_compute"] = True
     model = te.Linear(64, 32, params_dtype=dtype, device=device, **model_kwargs)
 
-    if case == "fp8_output_differentiable":
-        fp8_recipe = recipe.Float8CurrentScaling()
-
-        def fn(inp):
-            with te.autocast(recipe=fp8_recipe):
-                return model(inp, fp8_output=True).dequantize()
-
-        return model, fn, "fwd_grad", None, "differentiable fp8_output=True"
     if case == "fuse_wgrad_accumulation":
         model.weight.main_grad = torch.zeros_like(model.weight, dtype=torch.float32)
         return model, model, "bwd", None, "fuse_wgrad_accumulation"
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 00846d615a..aa0ffe3223 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -749,15 +749,35 @@ def _spec_view(obj: Any, tensor_field_names: Sequence[str]) -> Any:
 # --------------------------------------------------------------------------- #
 
 
-def _spec_slot_count(spec: Optional[TensorSpec]) -> int:
-    """Flat ``Tensor[]`` slots the value for ``spec`` occupies."""
+def _spec_slot_count(spec: Optional[TensorSpec], *, user_output: bool = False) -> int:
+    """Flat ``Tensor[]`` slots the value for ``spec`` occupies.
+
+    A quantized *user output* carries one extra slot, its grad handle (see
+    :func:`_grad_handle`).
+    """
     if spec is None:
         return 1
-    return len(spec.inner_names())
+    count = len(spec.inner_names())
+    if user_output and spec.is_quantized:
+        count += 1
+    return count
+
+
+def _grad_handle(shape: Sequence[int], dtype: torch.dtype, device: Any) -> torch.Tensor:
+    """Autograd stand-in for a quantized op output.
+
+    The inner buffers of a quantized output (uint8 data, scales) cannot carry
+    a gradient, so the op also returns this stride-0 high-precision tensor of
+    the output's logical shape. :class:`_QuantizedOutputFn` attaches it to
+    the rebuilt wrapper, and its gradient is the wrapper's ``grad_output``.
+    """
+    return torch.empty((1,), dtype=dtype, device=device).expand(tuple(shape))
 
 
 def _flatten_value(
     value: Optional[Union[torch.Tensor, QuantizedTensorStorage, TensorSpec]],
+    *,
+    user_output: bool = False,
 ) -> List[torch.Tensor]:
     """Return the flat ``Tensor[]`` slots that represent one op output ``value``.
 
@@ -767,10 +787,17 @@ def _flatten_value(
     if value is None:
         return [_encode_none(None)]
     if isinstance(value, TensorSpec):
-        return [_encode_none(t) for t in value.create_inner_tensors()]
+        flat = [_encode_none(t) for t in value.create_inner_tensors()]
+        if user_output and value.is_quantized:
+            flat.append(_grad_handle(value.shape, value.dtype, flat[0].device))
+        return flat
     if hasattr(value, "__tensor_flatten__"):
         inner_names, _ = value.__tensor_flatten__()
-        return [_encode_none(getattr(value, n)) for n in inner_names]
+        flat = [_encode_none(getattr(value, n)) for n in inner_names]
+        if user_output:
+            dtype = getattr(value, "dtype", None) or value._dtype
+            flat.append(_grad_handle(value.shape, dtype, flat[0].device))
+        return flat
     if isinstance(value, torch.Tensor):
         return [_encode_none(value)]
     raise TypeError(
@@ -779,6 +806,25 @@ def _flatten_value(
     )
 
 
+class _QuantizedOutputFn(torch.autograd.Function):
+    """Rebuild a quantized op output as a differentiable wrapper.
+
+    ``handle`` is the op's grad handle for this output; the gradient of the
+    wrapper flows back through it as-is (plain or quantized).
+    """
+
+    @staticmethod
+    def forward(ctx, handle, spec, *inner):
+        # pylint: disable=missing-function-docstring,unused-argument
+        ctx.num_inner = len(inner)
+        return dataclasses.replace(spec, requires_grad=False).assemble(list(inner))
+
+    @staticmethod
+    def backward(ctx, grad):
+        # pylint: disable=missing-function-docstring
+        return (grad, None) + (None,) * ctx.num_inner
+
+
 # Trailing slots in every fwd-impl return: ``tensors_to_save, ctx_attrs``.
 # User-output count is ``len(result) - this``.
 _FWD_TRAILING_SLOTS = 2
@@ -816,7 +862,7 @@ def _pack_fwd_result(result: Any) -> List[torch.Tensor]:
     num_outputs = len(result) - _FWD_TRAILING_SLOTS
     flat: List[torch.Tensor] = []
     for value in result[:num_outputs]:
-        flat.extend(_flatten_value(value))
+        flat.extend(_flatten_value(value, user_output=True))
     saved = result[num_outputs]
     if saved is not None:
         for value in saved:
@@ -827,8 +873,9 @@ def _pack_fwd_result(result: Any) -> List[torch.Tensor]:
 def _pack_bwd_result(grads: Any, num_grad_inputs: 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.
+    One grad per ``input_tensors_for_grad`` entry (validated); like forward
+    outputs, a quantized grad is flattened to its inner buffers and rebuilt by
+    ``_autograd_backward`` from the bwd fake impl's specs.
     """
     grads = list(grads)
     if len(grads) != num_grad_inputs:
@@ -838,13 +885,27 @@ def _pack_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> List
         )
     out: List[torch.Tensor] = []
     for g in grads:
-        if isinstance(g, TensorSpec):
-            out.append(_encode_none(g.create_tensor()))
-        else:
-            out.append(_encode_none(g))
+        out.extend(_flatten_value(g))
     return out
 
 
+def _unpack_bwd_result(
+    grad_specs: Sequence[Optional[TensorSpec]], flat: Sequence[Optional[torch.Tensor]]
+) -> List[Any]:
+    """Rebuild the per-input grads from the backward op's flat return."""
+    grads: List[Any] = []
+    cursor = 0
+    for spec in grad_specs:
+        n = _spec_slot_count(spec)
+        chunk = [_decode_none(t) for t in flat[cursor : cursor + n]]
+        cursor += n
+        if spec is None or not spec.is_quantized:
+            grads.append(chunk[0])
+        else:
+            grads.append(spec.assemble(chunk))
+    return grads
+
+
 @dataclasses.dataclass(frozen=True)
 class _OutputPlan:
     """Per-trace layout of an op's flat ``Tensor[]`` return.
@@ -874,7 +935,7 @@ def parse(cls, result: Tuple[Any, ...]) -> "_OutputPlan":
         cursor = 0
         user_ranges: List[Tuple[int, int]] = []
         for spec in user_specs:
-            n = _spec_slot_count(spec)
+            n = _spec_slot_count(spec, user_output=True)
             user_ranges.append((cursor, cursor + n))
             cursor += n
         return cls(
@@ -893,12 +954,27 @@ def _assemble(
         # ``spec is None`` is the op-boundary sentinel for an absent output.
         return spec.assemble(chunk) if spec is not None else None
 
-    def user_outputs(self, flat: Sequence[Optional[torch.Tensor]]) -> List[Any]:
-        """Rebuild the structured user outputs from the op's flat return."""
-        return [
-            self._assemble(spec, flat, start, stop)
-            for spec, (start, stop) in zip(self.user_specs, self.user_ranges)
-        ]
+    def user_outputs(
+        self, flat: Sequence[Optional[torch.Tensor]], *, differentiable: bool = True
+    ) -> List[Any]:
+        """Rebuild the structured user outputs from the op's flat return.
+
+        A quantized output whose grad handle requires grad is rebuilt through
+        :class:`_QuantizedOutputFn` (unless ``differentiable=False``, e.g. from
+        ``setup_context``) so autograd reaches the op through the handle.
+        """
+        outputs: List[Any] = []
+        for spec, (start, stop) in zip(self.user_specs, self.user_ranges):
+            if spec is None or not spec.is_quantized:
+                outputs.append(self._assemble(spec, flat, start, stop))
+                continue
+            handle = flat[stop - 1]
+            if differentiable and handle.requires_grad:
+                inner = [_decode_none(t) for t in flat[start : stop - 1]]
+                outputs.append(_QuantizedOutputFn.apply(handle, spec, *inner))
+            else:
+                outputs.append(self._assemble(spec, flat, start, stop - 1))
+        return outputs
 
     def saved_tensors(self, flat: Sequence[Optional[torch.Tensor]]) -> List[Any]:
         """Rebuild the saved-for-backward tensors from the op's flat return."""
@@ -916,17 +992,13 @@ def _slice_user_grads(
 ) -> List[Any]:
     """Gradient of each user output, sliced from the op's flat grad list.
 
-    A single-slot output yields its tensor grad; a flattened quantized output
-    yields the tuple of its inner-buffer grads. Takes the bare ranges (not the
-    whole :class:`_OutputPlan`) so ``setup_context`` only has to stash those on
-    ``ctx`` -- keeping the specs (and the quantizers they reference) off the
-    autograd tape.
+    A plain output's grad is its single slot; a quantized output's grad is the
+    grad of its trailing grad handle (plain or quantized tensor). Takes the bare
+    ranges (not the whole :class:`_OutputPlan`) so ``setup_context`` only has to
+    stash those on ``ctx`` -- keeping the specs (and the quantizers they
+    reference) off the autograd tape.
     """
-    grads: List[Any] = []
-    for start, stop in user_ranges:
-        chunk = [_decode_none(g) for g in flat_grads[start:stop]]
-        grads.append(chunk[0] if stop - start == 1 else tuple(chunk))
-    return grads
+    return [_decode_none(flat_grads[stop - 1]) for _, stop in user_ranges]
 
 
 # --------------------------------------------------------------------------- #
@@ -1001,13 +1073,16 @@ def _register_autograd_for_op(
     grad_targets: List[int],
     setup_context_user: Callable[..., Any],
     fwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
+    bwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
 ) -> None:
     """Wire ``register_autograd`` on a forward op so its backward calls ``bwd_op``.
 
     ``setup_context`` re-runs the spec fwd fake impl to parse the
     :class:`_OutputPlan`, reassembles the outputs / saved tensors from it, hands
     the saved tuple + ``ctx_attrs`` to the module's ``setup_context`` and stashes
-    the plan on ``ctx`` so backward can slice its grads per user output.
+    the plan on ``ctx`` so backward can slice its grads per user output. The
+    backward runs the spec bwd fake impl to learn the grads' layout and rebuilds
+    quantized grads from the op's flat return.
     """
     bwd_takes_grad_tuple = any(f.name == "grad_outputs" for f in bwd_plan.fields)
 
@@ -1019,7 +1094,7 @@ def _setup_context(ctx, inputs, output):
         spec_obj = _spec_view(fwd_obj, fwd_plan.tensor_field_names())
 
         out_plan = _OutputPlan.parse(fwd_fake_impl(spec_obj))
-        user_outputs = out_plan.user_outputs(output)
+        user_outputs = out_plan.user_outputs(output, differentiable=False)
         saved_list = out_plan.saved_tensors(output)
 
         bwd_obj = bwd_plan.arg_type()
@@ -1053,9 +1128,10 @@ def _autograd_backward(ctx, *grad_outputs):
             bwd_obj.grad_outputs = tuple(user_grads)
         else:
             bwd_obj.grad_output = user_grads[0]
+        grad_specs = bwd_fake_impl(_spec_view(bwd_obj, bwd_plan.tensor_field_names()))
         kwargs = bwd_plan.pack(bwd_obj)
         bwd_args_flat = [kwargs[name] for name in bwd_plan.slot_names]
-        grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)]
+        grads = _unpack_bwd_result(grad_specs, bwd_op(*bwd_args_flat))
         ctx.backward_objects = None
         # One grad per input schema slot: default None, but a ``Tensor[]`` slot
         # (always recorded in ``fwd_tensor_list_lengths``) needs a
@@ -1330,6 +1406,7 @@ def _register_custom_op_impl(
         "grad_targets": grad_targets,
         "setup_context_user": setup_context,
         "fwd_fake_impl": fwd_fake_impl,
+        "bwd_fake_impl": bwd_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)
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 55fc69ef7f..eb71827336 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -190,26 +190,11 @@ def compile_unsupported_reason(self) -> Optional[str]:
             return "a quantized input tensor"
         if self.fsdp_group is not None:
             return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP"
-        if (
-            self.fp8_output
-            and self.is_grad_enabled
-            and (self.input_requires_grad or self.weight_requires_grad or self.bias_requires_grad)
-        ):
-            return "differentiable fp8_output=True"
         if self.cpu_offloading:
             return "CPU activation offloading"
         if self.wgrad_store is not None:
             # Non-None only when delayed wgrad compute is on (see Linear.forward).
             return "delayed wgrad compute (wgrad_store)"
-        if (
-            self.grad_input_quantizer is not None
-            and self.is_grad_enabled
-            and self.input_requires_grad
-            and not (self.ub_overlap_rs_dgrad or self.ub_bulk_wgrad)
-        ):
-            # A quantized dgrad can't cross the op boundary: grads are packed
-            # one plain Tensor[] slot each (_pack_bwd_result).
-            return "a quantized input grad (fp8_grad=True)"
         if self.cache_weight and self.fp8:
             # The cached workspace is updated in place on the first microbatch,
             # which the functional op (mutates_args=()) can't express. Without
@@ -237,7 +222,7 @@ class LinearBwdArgs:
     """Single-argument bag for the backward path of :class:`_Linear`."""
 
     # --- Saved / restored tensors (populated at backward entry) ---
-    grad_output: Optional[torch.Tensor] = None
+    grad_output: Optional[TensorOrQuantized] = None
     inputmat: Optional[TensorOrQuantized] = None
     weight_fp8: Optional[TensorOrQuantized] = None
     saved_weight: Optional[TensorOrQuantized] = None
@@ -2382,9 +2367,7 @@ def forward(
                 fp8_grad = True
 
         if torch.compiler.is_compiling() and _linear_op is not None:
-            reason = self._compile_eager_fallback_reason(
-                inp, is_first_microbatch, fp8_output, fp8_grad, is_grad_enabled, debug
-            )
+            reason = self._compile_eager_fallback_reason(inp, is_first_microbatch, debug)
             if reason is not None:
                 # A break inside the try/finally below would skip the whole frame.
                 warn_compile_eager_fallback(reason)
@@ -2628,29 +2611,19 @@ def _compile_eager_fallback_reason(
         self,
         inp: torch.Tensor,
         is_first_microbatch: Optional[bool],
-        fp8_output: bool,
-        fp8_grad: bool,
-        is_grad_enabled: bool,
         debug: bool,
     ) -> Optional[str]:
         """Why this call can't use the compiled op (else None), decided before
         prepare_forward. Quantizer checks stay in compile_unsupported_reason."""
         if debug:
             return "debug instrumentation (nvidia-dlfw-inspect)"
-        weight_tensor, bias_tensor = self._get_weight_and_bias_tensors()
+        weight_tensor, _ = self._get_weight_and_bias_tensors()
         if is_distributed_weight(weight_tensor):
             return "a DistributedWeight (custom weight parallelism, e.g. GTP)"
         if isinstance(inp, (QuantizedTensor, QuantizedTensorStorage)):
             return "a quantized input tensor"
         if self.fsdp_group is not None:
             return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP"
-        any_requires_grad = (
-            inp.requires_grad
-            or weight_tensor.requires_grad
-            or (bias_tensor is not None and bias_tensor.requires_grad)
-        )
-        if fp8_output and is_grad_enabled and any_requires_grad:
-            return "differentiable fp8_output=True"
         if is_cpu_offload_enabled():
             return "CPU activation offloading"
         if self.wgrad_store is not None and self.wgrad_store.delay_wgrad_compute():
@@ -2658,14 +2631,6 @@ def _compile_eager_fallback_reason(
         if self.fuse_wgrad_accumulation:
             return "fuse_wgrad_accumulation (main_grad)"
         fp8 = FP8GlobalStateManager.is_fp8_enabled()
-        needs_dgrad = is_grad_enabled and inp.requires_grad
-        if (
-            fp8
-            and fp8_grad
-            and needs_dgrad
-            and not (self.ub_overlap_rs_dgrad or self.ub_bulk_wgrad)
-        ):
-            return "a quantized input grad (fp8_grad=True)"
         if fp8 and is_first_microbatch is not None and not self.is_fsdp2:
             return "FP8 weight caching (is_first_microbatch)"
         return None
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index 7149a5a163..b8fa4e748e 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -7,6 +7,7 @@
 from __future__ import annotations
 from typing import NamedTuple, Optional, Tuple, Iterable, Any, Dict, Union, get_type_hints
 import abc
+import copy
 import warnings
 import math
 
@@ -586,6 +587,16 @@ def set_usage(
         if columnwise is not None:
             self.columnwise_usage = columnwise
 
+    def tangent_quantizer(self) -> Quantizer:
+        """Quantizer describing the expected gradient of a tensor made by this quantizer.
+
+        Used by torch.compile to guess the layout of a quantized gradient before
+        it exists (rowwise data only, as produced by a quantized dgrad).
+        """
+        quantizer = self.copy() if hasattr(self, "copy") else copy.copy(self)
+        quantizer.set_usage(rowwise=True, columnwise=False)
+        return quantizer
+
     def onnx_quantize(self, tensor: torch.Tensor) -> QuantizedTensor:
         """Symbolic function for ONNX export"""
         raise NotImplementedError(
@@ -794,6 +805,46 @@ def quantize_(self, tensor: torch.Tensor) -> QuantizedTensor:
             f"{self.__class__.__name__} class does not implement quantize_ function"
         )
 
+    def __coerce_tangent_metadata__(self) -> QuantizedTensor:
+        """torch.compile hook: the gradient layout to trace the backward with."""
+        quantizer = getattr(self, "_quantizer", None)
+        if quantizer is None:
+            return self
+        return self._rewrap_with_quantizer(quantizer.tangent_quantizer())
+
+    def __coerce_same_metadata_as_tangent__(
+        self, expected_meta: Dict[str, Any], expected_type: Optional[type] = None
+    ) -> torch.Tensor:
+        """torch.compile hook: convert this runtime gradient to the traced layout."""
+        if expected_type is not None and not issubclass(expected_type, QuantizedTensorStorage):
+            return self.dequantize()
+        kwargs = expected_meta["nontensor_kwargs"]
+        quantizer = kwargs.get("quantizer")
+        if quantizer is None:
+            return self
+        same_type = expected_type is None or expected_type is type(self)
+        own = self._flatten_nontensor_kwargs()
+        if same_type and all(own.get(k) == v for k, v in kwargs.items() if k != "quantizer"):
+            self.update_usage(
+                rowwise_usage=quantizer.rowwise_usage, columnwise_usage=quantizer.columnwise_usage
+            )
+            return self
+        quantizer = quantizer.copy() if hasattr(quantizer, "copy") else copy.copy(quantizer)
+        quantizer.internal = False
+        return quantizer(self.dequantize())
+
+    def _rewrap_with_quantizer(self, quantizer: Quantizer) -> QuantizedTensor:
+        """Same buffers under ``quantizer``'s metadata; ``self`` if they don't cover its layout."""
+        names = list(quantizer.inner_tensor_specs(tuple(self.shape)).keys())
+        present, ctx = self.__tensor_flatten__()
+        if any(name not in present for name in names):
+            return self
+        ctx = dict(ctx)
+        ctx["requires_grad"] = False
+        ctx["nontensor_kwargs"] = quantizer.storage_metadata(self.dtype)["nontensor_kwargs"]
+        inner = {name: getattr(self, name) for name in names}
+        return type(self).__tensor_unflatten__(inner, ctx, tuple(self.shape), tuple(self.stride()))
+
     def detach(self) -> QuantizedTensor:
         """Create new quantized tensor with same data
 
@@ -1012,7 +1063,7 @@ def maybe_update_inplace(arg, new_arg, schema_arg):
             new_kwargs = tree_map(maybe_unwrap, kwargs)
             schema_args = func._schema.arguments
             args_len = len(args)
-            super().__torch_dispatch__(func, types, new_args, new_kwargs)
+            func(*new_args, **(new_kwargs or {}))
             for arg, new_arg, schema_arg in zip(args, new_args, schema_args):
                 maybe_update_inplace(arg, new_arg, schema_arg)
             for kwarg, new_kwarg, schema_arg in zip(kwargs, new_kwargs, schema_args[args_len:]):
@@ -1020,12 +1071,12 @@ def maybe_update_inplace(arg, new_arg, schema_arg):
                 maybe_update_inplace(kwargs[kwarg], new_kwargs[new_kwarg], schema_arg)
             return None
 
-        # Default op: dequantize and perform op
+        # Default op: dequantize and perform op. Re-dispatch normally (not via
+        # the dispatch-disabling ``super()`` path) so inner tensor modes such as
+        # FakeTensor/FunctionalTensor under torch.compile still see the call.
         args = tree_map(maybe_unwrap, args)
-        if kwargs is not None:
-            kwargs = tree_map(maybe_unwrap, kwargs)
-        out = super().__torch_dispatch__(func, types, args, kwargs)
-        return out
+        kwargs = tree_map(maybe_unwrap, kwargs) if kwargs is not None else {}
+        return func(*args, **kwargs)
 
     __torch_function__ = torch._C._disabled_torch_function_impl
 
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index cf37c36c59..5d28bfd64d 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -82,6 +82,11 @@ def __init__(
         self.amax = amax
         self.dtype = DType.cast(fp8_dtype)
 
+    def tangent_quantizer(self) -> Quantizer:
+        quantizer = super().tangent_quantizer()
+        quantizer.dtype = DType.kFloat8E5M2
+        return quantizer
+
     def copy(self) -> Float8Quantizer:
         """Create shallow copy"""
 
@@ -260,6 +265,11 @@ def __getstate__(self):
         state["amax_reduction_group"] = None
         return state
 
+    def tangent_quantizer(self) -> Quantizer:
+        quantizer = super().tangent_quantizer()
+        quantizer.dtype = DType.kFloat8E5M2
+        return quantizer
+
     def copy(self) -> Float8CurrentScalingQuantizer:
         """Create shallow copy"""
 
diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py
index 32c27eb583..64f9e51740 100644
--- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py
+++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py
@@ -18,6 +18,29 @@
 from ...utils import is_non_tn_fp8_gemm_supported, _empty_tensor
 
 
+def _dequantize_fp8_impl(
+    data: torch.Tensor, scale_inv: torch.Tensor, fp8_dtype: int, dtype: torch.dtype
+) -> torch.Tensor:
+    tensor = Float8TensorStorage(
+        data=data, fp8_scale_inv=scale_inv, fp8_dtype=DType(fp8_dtype), fake_dtype=dtype
+    )
+    return tex.dequantize(tensor, torch_to_transformer_engine_dtype[dtype])
+
+
+try:
+    _dequantize_fp8 = torch.library.custom_op(
+        "transformer_engine_compile::dequantize_fp8", mutates_args=()
+    )(_dequantize_fp8_impl)
+
+    @_dequantize_fp8.register_fake
+    def _(data: torch.Tensor, scale_inv: torch.Tensor, fp8_dtype: int, dtype: torch.dtype):
+        del scale_inv, fp8_dtype
+        return torch.empty(data.shape, dtype=dtype, device=data.device)
+
+except Exception:  # pylint: disable=broad-exception-caught
+    _dequantize_fp8 = _dequantize_fp8_impl
+
+
 class _FromFloat8Func(torch.autograd.Function):
     """Cast from FP8 to other dtype"""
 
@@ -28,8 +51,6 @@ def forward(
         dtype: torch.dtype,
     ) -> torch.Tensor:
         # pylint: disable=missing-function-docstring
-        te_dtype = torch_to_transformer_engine_dtype[dtype]
-
         # Make sure FP8 data is in expected format
         if tensor._data is not None:
             if tensor._data.numel() == 0:
@@ -41,8 +62,8 @@ def forward(
                     tensor._data.view(fp8_torch_dtype).float()
                     * tensor._scale_inv.to(tensor._data.device)
                 ).to(dtype)
-            # Cast from FP8
-            return tex.dequantize(tensor, te_dtype)
+            # Cast from FP8 (custom op so torch.compile can trace it)
+            return _dequantize_fp8(tensor._data, tensor._scale_inv, int(tensor._fp8_dtype), dtype)
 
         if tensor._transpose is not None and not tensor._transpose_invalid:
             # A columnwise-only tensor stores the logical last dimension first

From 19a1039b54a2610373725879c9776777abe82d3c Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Fri, 4 Sep 2026 13:54:57 +0200
Subject: [PATCH 2/2] [PyTorch][torch.compile] Quantized input for the Linear
 custom op

Declare inp as TensorOrQuantized so a Float8Tensor input is flattened by the
wrapper op like the other quantized operands, and drop the quantized-input
fallback. This lets Linear(fp8_output=True) feed another Linear inside one
graph; the consumer's dgrad (plain or quantized) reaches the producer through
the grad handle.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py         | 57 ++++++++++++++++-----
 transformer_engine/pytorch/module/linear.py | 13 ++---
 2 files changed, 46 insertions(+), 24 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index fc127b4a97..49ca1d3ada 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -2232,6 +2232,49 @@ def consume(out):
     _assert_fp8_io_close(ref, got)
 
 
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("compile_mode", _compile_modes)
+@pytest.mark.parametrize("consumer_grad", ["bf16", "fp8"])
+def test_te_linear_compile_fp8_output_chain(compile_mode, consumer_grad):
+    """``Linear(fp8_output=True)`` feeding a second ``Linear`` inside one graph:
+    the quantized input crosses the op boundary as a subclass, and the consumer's
+    dgrad (plain or quantized) reaches the producer through the grad handle."""
+    fp8_recipe = recipe.Float8CurrentScaling()
+    model = te.Linear(64, 32, params_dtype=torch.bfloat16, device="cuda")
+    consumer = te.Linear(32, 16, params_dtype=torch.bfloat16, device="cuda")
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            return consumer(model(inp, fp8_output=True), fp8_grad=consumer_grad == "fp8")
+
+    ref, got = _run_fp8_io_pair(fn, model, compile_mode)
+    _assert_fp8_io_close(ref, got)
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+def test_te_linear_compile_quantized_input():
+    """An externally quantized :class:`Float8Tensor` input (no grad) goes through
+    the compiled op via the wrapper op's subclass flattening."""
+    fp8_recipe = recipe.Float8CurrentScaling()
+    model = te.Linear(64, 32, params_dtype=torch.bfloat16, device="cuda")
+    quantizer = Float8CurrentScalingQuantizer(fp8_dtype=tex.DType.kFloat8E4M3, device="cuda")
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            return model(inp)
+
+    torch.manual_seed(0)
+    inp = quantizer(torch.randn(32, 64, dtype=torch.bfloat16, device="cuda"))
+    torch._dynamo.reset()
+    compiled = torch.compile(fn, fullgraph=True)
+    with torch.no_grad():
+        ref = fn(inp)
+        out = compiled(inp)
+    torch.testing.assert_close(out, ref, atol=_EAGER_ATOL, rtol=_EAGER_RTOL)
+
+
 # Configs rejected by LinearFwdArgs.compile_unsupported_reason() that a
 # single-GPU unit test can construct. Distributed-only reasons (fsdp_group,
 # DistributedWeight) and CPU offloading need machinery this file doesn't have;
@@ -2240,7 +2283,6 @@ def consume(out):
 _FALLBACK_CASES = [
     "fuse_wgrad_accumulation",
     "delayed_wgrad",
-    "quantized_input",
 ]
 
 
@@ -2258,14 +2300,6 @@ def _fallback_case(case, dtype, device):
         return model, model, "bwd", None, "fuse_wgrad_accumulation"
     if case == "delayed_wgrad":
         return model, model, "bwd", model.backward_dw, "delayed wgrad compute"
-    if case == "quantized_input":
-        fp8_recipe = recipe.Float8CurrentScaling()
-
-        def fn(inp):
-            with te.autocast(recipe=fp8_recipe):
-                return model(inp)
-
-        return model, fn, "no_grad", None, "a quantized input tensor"
     raise ValueError(case)
 
 
@@ -2285,11 +2319,6 @@ def test_te_linear_compile_eager_fallback(case):
     def make_inp():
         torch.manual_seed(1)
         x = torch.randn(32, 64, dtype=dtype, device=device)
-        if case == "quantized_input":
-            quantizer = Float8CurrentScalingQuantizer(
-                fp8_dtype=tex.DType.kFloat8E4M3, device=device
-            )
-            return quantizer(x)
         return x.requires_grad_(mode != "no_grad")
 
     torch._dynamo.reset()
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index eb71827336..02a60ce89c 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -112,7 +112,7 @@ class LinearFwdArgs:
 
     # --- Differentiable tensors (also passed positionally to autograd) ---
     weight: TensorOrQuantized
-    inp: torch.Tensor
+    inp: TensorOrQuantized
     bias: Optional[torch.Tensor]
 
     # --- Non-differentiable cached tensors ---
@@ -186,8 +186,6 @@ def compile_unsupported_reason(self) -> Optional[str]:
             return "debug instrumentation (nvidia-dlfw-inspect)"
         if is_distributed_weight(self.weight):
             return "a DistributedWeight (custom weight parallelism, e.g. GTP)"
-        if isinstance(self.inp, (QuantizedTensor, QuantizedTensorStorage)):
-            return "a quantized input tensor"
         if self.fsdp_group is not None:
             return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP"
         if self.cpu_offloading:
@@ -2367,7 +2365,7 @@ def forward(
                 fp8_grad = True
 
         if torch.compiler.is_compiling() and _linear_op is not None:
-            reason = self._compile_eager_fallback_reason(inp, is_first_microbatch, debug)
+            reason = self._compile_eager_fallback_reason(is_first_microbatch, debug)
             if reason is not None:
                 # A break inside the try/finally below would skip the whole frame.
                 warn_compile_eager_fallback(reason)
@@ -2608,10 +2606,7 @@ def _get_weight_tensors(self) -> List[Union[torch.Tensor, QuantizedTensorStorage
         return unfused_weights
 
     def _compile_eager_fallback_reason(
-        self,
-        inp: torch.Tensor,
-        is_first_microbatch: Optional[bool],
-        debug: bool,
+        self, is_first_microbatch: Optional[bool], debug: bool
     ) -> Optional[str]:
         """Why this call can't use the compiled op (else None), decided before
         prepare_forward. Quantizer checks stay in compile_unsupported_reason."""
@@ -2620,8 +2615,6 @@ def _compile_eager_fallback_reason(
         weight_tensor, _ = self._get_weight_and_bias_tensors()
         if is_distributed_weight(weight_tensor):
             return "a DistributedWeight (custom weight parallelism, e.g. GTP)"
-        if isinstance(inp, (QuantizedTensor, QuantizedTensorStorage)):
-            return "a quantized input tensor"
         if self.fsdp_group is not None:
             return "manual TE FSDP (fsdp_group); use FSDP2 or MCore FSDP"
         if is_cpu_offload_enabled():