From f3401df703909f211dc8617e8832ba2f2d3592a0 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Sat, 6 Jun 2026 14:11:29 +0200
Subject: [PATCH 01/28] [PyTorch] Make tensorless quantizers opaque value
 objects for torch.compile

Give tensorless quantizers (MXFP8, FP8 blockwise, FP8 current-scaling,
NVFP4) value-object semantics so torch.compile can treat them as baked-in
constants:

- Add opt-in value identity to the base Quantizer (_value_fields /
  _value_key / __eq__ / __hash__). Quantizers holding live tensors
  (delayed-scaling Float8Quantizer) and custom quantizers keep identity
  semantics.
- New transformer_engine/pytorch/dynamo.py houses the torch.compile glue:
  __fx_repr__, value-key reconstruction and register_value_opaque_quantizer
  (gracefully a no-op on PyTorch builds without the opaque-object API).
- Register the four tensorless quantizers as value opaque types.

Also fix CustomRecipe state caching in TransformerEngineBaseModule:
set_meta_tensor now rebuilds quantizers when the CustomRecipe instance
changes (e.g. nested te.autocast regions) instead of reusing the first
recipe's state, since every CustomRecipe shares the CustomRecipeState type
but carries its own qfactory.

Move the quantizer value-object tests into tests/pytorch/test_torch_compile.py
and add that file to the L0 pytorch unittest QA suite.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 154 ++++++++++++++++++
 transformer_engine/pytorch/dynamo.py          | 120 ++++++++++++++
 .../pytorch/quantized_tensor.py               |  67 ++++++++
 .../pytorch/tensor/float8_blockwise_tensor.py |   7 +
 .../pytorch/tensor/float8_tensor.py           |   9 +
 .../pytorch/tensor/mxfp8_tensor.py            |   7 +
 .../pytorch/tensor/nvfp4_tensor.py            |  25 +++
 7 files changed, 389 insertions(+)
 create mode 100644 transformer_engine/pytorch/dynamo.py

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 1286492a6e..8adcc88e61 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -24,6 +24,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 QuantizerRole
 from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear
 from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer
 from transformer_engine.pytorch.quantization import QuantizerRole
@@ -32,6 +33,14 @@
     is_mxfp8_available,
     is_fp8_block_scaling_available,
     is_nvfp4_available,
+    Float8Quantizer,
+    Float8BlockQuantizer,
+    MXFP8Quantizer,
+    NVFP4Quantizer,
+)
+from transformer_engine.pytorch.dynamo import (
+    register_value_opaque_quantizer,
+    _quantizer_from_value_key,
 )
 from utils import recipe_id
 
@@ -384,3 +393,148 @@ def fn(inp):
 
     out = compiled(inp)
     out.sum().backward()
+
+
+# ---------------------------------------------------------------------------
+# Value-opaque quantizers: eager value semantics + FX reconstruction
+#
+# The tensorless quantizers (current-scaling FP8, FP8 blockwise, MXFP8, NVFP4)
+# are torch.compile *value* opaque types: they provide value-based
+# ``__eq__`` / ``__hash__`` and an evaluable ``__fx_repr__`` (see
+# ``torch._library.opaque_object`` Note [Opaque Objects]). These tests exercise
+# the eager value semantics and the FX reconstruction round-trip. They are
+# CPU-friendly except for NVFP4 (whose constructor touches the current CUDA
+# device).
+# ---------------------------------------------------------------------------
+
+
+def _mxfp8(dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True):
+    return MXFP8Quantizer(fp8_dtype=dtype, rowwise=rowwise, columnwise=columnwise)
+
+
+def _blockwise(dtype=tex.DType.kFloat8E4M3, force_pow_2_scales=True, block_scaling_dim=2):
+    return Float8BlockQuantizer(
+        fp8_dtype=dtype,
+        rowwise=True,
+        columnwise=True,
+        force_pow_2_scales=force_pow_2_scales,
+        block_scaling_dim=block_scaling_dim,
+    )
+
+
+def _current_scaling(dtype=tex.DType.kFloat8E4M3, force_pow_2_scales=False, amax_epsilon=0.0):
+    return Float8CurrentScalingQuantizer(
+        fp8_dtype=dtype,
+        device=torch.device("cpu"),
+        force_pow_2_scales=force_pow_2_scales,
+        amax_epsilon=amax_epsilon,
+    )
+
+
+# (factory, kwargs_for_a_different_but_valid_config)
+_CPU_VALUE_QUANTIZERS = [
+    pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"),
+    pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"),
+    pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"),
+]
+
+
+@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS)
+def test_quantizer_value_equality(factory, other_kwargs):
+    """Same config -> equal & same hash; different config -> not equal."""
+    a = factory()
+    b = factory()
+    assert a is not b
+    assert a == b
+    assert hash(a) == hash(b)
+
+    c = factory(**other_kwargs)
+    assert a != c
+    # Usage flags participate in the value.
+    d = factory()
+    d.set_usage(rowwise=False, columnwise=False)
+    assert a != d
+
+
+@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS)
+def test_quantizer_usable_in_set_and_dict(factory, other_kwargs):
+    a = factory()
+    b = factory()
+    c = factory(**other_kwargs)
+    assert len({a, b, c}) == 2
+    mapping = {a: "x"}
+    assert mapping[b] == "x"
+
+
+@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS)
+def test_quantizer_cross_type_inequality(factory, other_kwargs):
+    a = factory()
+    other = _current_scaling() if not isinstance(a, Float8CurrentScalingQuantizer) else _mxfp8()
+    assert a != other
+
+
+@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS)
+def test_quantizer_fx_repr_roundtrip(factory, other_kwargs):
+    """``__fx_repr__`` returns an evaluable expression rebuilding an equal object."""
+    a = factory()
+    repr_str, globals_ = a.__fx_repr__()
+    assert isinstance(repr_str, str)
+    assert isinstance(globals_, dict)
+    rebuilt = eval(repr_str, dict(globals_))  # pylint: disable=eval-used
+    assert rebuilt == a
+    assert rebuilt is not a
+    assert hash(rebuilt) == hash(a)
+
+
+@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS)
+def test_quantizer_value_key_reconstruction(factory, other_kwargs):
+    a = factory()
+    rebuilt = _quantizer_from_value_key(a._value_key())
+    assert type(rebuilt) is type(a)
+    assert rebuilt == a
+    # The deprecated amax-reduction process group is never carried in the value.
+    assert getattr(rebuilt, "amax_reduction_group", None) is None
+
+
+def test_quantizer_delayed_scaling_keeps_identity_semantics():
+    """Float8Quantizer holds live tensors -> identity (not value) semantics."""
+    scale = torch.ones(1)
+    amax = torch.zeros(1)
+    a = Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3)
+    b = Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3)
+    assert a._value_fields() is None
+    assert a == a
+    assert a != b  # distinct instances are not equal despite identical config
+    assert hash(a) == object.__hash__(a)
+
+
+@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS)
+def test_quantizer_registration_is_idempotent_and_tolerant(factory, other_kwargs):
+    """Re-registering must not raise, regardless of PyTorch opaque-object support."""
+    cls = type(factory())
+    register_value_opaque_quantizer(cls)
+    register_value_opaque_quantizer(cls)
+
+    if not _opaque_available:
+        pytest.skip("PyTorch build without opaque-object API")
+    from torch._library.opaque_object import is_opaque_value_type
+
+    assert is_opaque_value_type(cls)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4Quantizer requires CUDA")
+def test_quantizer_nvfp4_value_semantics():
+    a = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1)
+    b = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1)
+    assert a == b
+    assert hash(a) == hash(b)
+
+    c = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1, with_rht=not a.with_rht)
+    assert a != c
+
+    rebuilt = _quantizer_from_value_key(a._value_key())
+    assert rebuilt == a
+    assert rebuilt.amax_reduction_group is None
+
+    repr_str, globals_ = a.__fx_repr__()
+    assert eval(repr_str, dict(globals_)) == a  # pylint: disable=eval-used
diff --git a/transformer_engine/pytorch/dynamo.py b/transformer_engine/pytorch/dynamo.py
new file mode 100644
index 0000000000..26aff6f59c
--- /dev/null
+++ b/transformer_engine/pytorch/dynamo.py
@@ -0,0 +1,120 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""torch.compile glue for Transformer Engine quantizers.
+
+This module isolates the torch.compile-specific plumbing that turns a
+*tensorless* quantizer into a torch.compile **value** opaque type:
+
+  * :func:`register_value_opaque_quantizer` -- attaches the ``__fx_repr__`` used
+    by FX codegen and registers the quantizer class with
+    ``torch._library.opaque_object``. It is a no-op (other than populating the
+    local registry) on PyTorch builds without the opaque-object API, so
+    importing Transformer Engine never fails on older PyTorch -- only
+    torch.compile specialization on the quantizer is unavailable there.
+  * :func:`_quantizer_from_value_key` -- rebuilds a quantizer constant from its
+    value key inside the generated FX graph.
+
+The eager value semantics (``__eq__`` / ``__hash__`` / ``_value_key`` /
+``_value_fields``) live on the quantizer itself; see
+:class:`transformer_engine.pytorch.quantized_tensor.Quantizer`.
+
+See ``torch._library.opaque_object`` Note [Opaque Objects] for the contract a
+value-typed opaque object must satisfy (``__eq__`` / ``__hash__`` /
+``__fx_repr__``).
+"""
+
+from __future__ import annotations
+from typing import Any, Dict, Tuple
+
+from .constants import DType
+
+
+# Maps a quantizer class qualname to the class object. A value key stores only
+# the qualname, so reconstruction looks the class up here. Populated by
+# ``register_value_opaque_quantizer`` at import time of each tensor module; this
+# avoids importing the tensor modules into this module (which would create an
+# import cycle).
+_QUANTIZER_VALUE_REGISTRY: Dict[str, type] = {}
+
+
+def _quantizer_from_value_key(key: Tuple[Any, ...]) -> Any:
+    """Rebuild a tensorless quantizer from its value key.
+
+    Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the
+    generated FX code calls this to materialize the quantizer constant. The
+    deprecated amax-reduction process group is never part of the value, so a
+    reconstructed quantizer always starts with no stored group.
+    """
+    qualname, items = key[0], key[1]
+    cls = _QUANTIZER_VALUE_REGISTRY[qualname]
+    # Bypass ``__init__`` and restore the value attributes directly: the value
+    # key already captures every value-defining field (including derived ones),
+    # and the constructors have heterogeneous signatures / side effects.
+    obj = cls.__new__(cls)
+    field_names = set()
+    for name, value in items:
+        if name == "dtype":
+            value = DType.cast(value)
+        object.__setattr__(obj, name, value)
+        field_names.add(name)
+    # The deprecated amax-reduction process group is excluded from the value;
+    # restore it as ``None`` for quantizers that still carry the fallback so
+    # attribute access keeps working.
+    if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"):
+        object.__setattr__(obj, "amax_reduction_group", None)
+    return obj
+
+
+def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]:
+    """``__fx_repr__`` for value-opaque quantizers (attached at registration).
+
+    Returns an evaluable expression that rebuilds the quantizer via
+    :func:`_quantizer_from_value_key`, together with the globals needed to
+    evaluate it.
+    """
+    return (
+        f"_quantizer_from_value_key({self._value_key()!r})",
+        {"_quantizer_from_value_key": _quantizer_from_value_key},
+    )
+
+
+def register_value_opaque_quantizer(cls: type) -> None:
+    """Register a tensorless quantizer class as a torch.compile value opaque type.
+
+    Attaches ``__fx_repr__`` and registers the class with
+    ``torch._library.opaque_object``. Safe to call on any PyTorch build: on
+    versions without the opaque-object API it only records the class in the
+    local registry and attaches ``__fx_repr__`` (both harmless), so Transformer
+    Engine keeps importing and running in eager mode.
+
+    The quantizer class must already provide value ``__eq__`` / ``__hash__`` and
+    a non-``None`` ``_value_fields`` (see
+    :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`).
+    """
+    _QUANTIZER_VALUE_REGISTRY[cls.__qualname__] = cls
+
+    # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the
+    # class, so attach it before registering.
+    if "__fx_repr__" not in cls.__dict__:
+        cls.__fx_repr__ = _quantizer_fx_repr
+
+    try:
+        from torch._library.opaque_object import (  # pylint: disable=import-outside-toplevel
+            register_opaque_type,
+            is_opaque_value_type,
+        )
+    except (ImportError, AttributeError):
+        # Older PyTorch without the opaque-object API: eager value semantics
+        # still work; torch.compile specialization on the quantizer does not.
+        return
+
+    if is_opaque_value_type(cls):
+        return
+
+    try:
+        register_opaque_type(cls, typ="value")
+    except (ImportError, AttributeError, RuntimeError, TypeError):
+        # Tolerate partial / experimental opaque-object support.
+        pass
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index cfe488aae5..f37b3cc63d 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -408,6 +408,73 @@ def get_usages(self) -> Dict[str, bool]:
             "columnwise": self.columnwise_usage,
         }
 
+    # ----- Value-object identity (torch.compile opaque value support) -----
+    # A *tensorless* quantizer (one whose entire state is a handful of plain,
+    # reproducible scalars -- no live tensors, no process groups) behaves like a
+    # value: two instances with the same configuration are interchangeable. Such
+    # quantizers opt into value-based ``__eq__`` / ``__hash__`` by overriding
+    # ``_value_fields``. Quantizers that keep the default (e.g. delayed-scaling
+    # ``Float8Quantizer``, which holds live scale/amax tensors, and any custom
+    # quantizer) retain the default identity semantics.
+    #
+    # This is the eager-side half of registering the quantizer as a torch.compile
+    # *value* opaque type; the torch.compile glue (``__fx_repr__``, FX
+    # reconstruction and ``register_opaque_type``) lives in
+    # ``transformer_engine.pytorch.dynamo``.
+
+    #: Attributes shared by every quantizer that take part in value identity.
+    _BASE_VALUE_FIELDS: Tuple[str, ...] = (
+        "rowwise_usage",
+        "columnwise_usage",
+        "internal",
+        "optimize_for_gemm",
+    )
+
+    def _value_fields(self) -> Optional[Tuple[str, ...]]:
+        """Subclass-specific value-defining attribute names, or ``None``.
+
+        Returning ``None`` (the default) means the quantizer is *not* a value
+        object and keeps identity-based equality/hashing. Tensorless quantizers
+        override this to return the tuple of attribute names that, together with
+        :attr:`_BASE_VALUE_FIELDS`, fully determine their value (excluding
+        non-value state such as a deprecated amax-reduction process group).
+        """
+        return None
+
+    def _value_key(self) -> Tuple[Any, ...]:
+        """Hashable, reproducible key identifying this quantizer's value.
+
+        Only valid for value quantizers (``_value_fields()`` is not ``None``).
+        """
+        fields = self._value_fields()  # pylint: disable=assignment-from-none
+        assert fields is not None, f"{type(self).__name__} is not a value quantizer"
+        items = []
+        for name in self._BASE_VALUE_FIELDS + tuple(fields):
+            value = getattr(self, name)
+            if name == "dtype":
+                # ``DType`` is an ``IntEnum``; store the int so the key stays
+                # plain: hashable and ``repr``-reproducible for FX codegen.
+                value = int(value)
+            items.append((name, value))
+        return (type(self).__qualname__, tuple(items))
+
+    def __eq__(self, other: object) -> Any:
+        # Value quantizers compare by configuration; everything else keeps the
+        # default identity semantics (returning ``NotImplemented`` makes Python
+        # fall back to identity).
+        if self is other:
+            return True
+        if self._value_fields() is None or type(self) is not type(other):
+            return NotImplemented
+        if other._value_fields() is None:
+            return NotImplemented
+        return self._value_key() == other._value_key()
+
+    def __hash__(self) -> int:
+        if self._value_fields() is None:
+            return object.__hash__(self)
+        return hash(self._value_key())
+
 
 class QuantizedTensor(torch.Tensor):
     """Abstract base class for tensor with quantized data
diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
index ba46508d74..92c22acd0b 100644
--- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
@@ -14,6 +14,7 @@
 from transformer_engine.common.recipe import Float8BlockScaling, Recipe
 from .storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage
 from ..quantized_tensor import QuantizedTensor, Quantizer
+from ..dynamo import register_value_opaque_quantizer
 from ._quantization_helpers import _IdentityFunc
 from ..constants import DType
 from ..utils import devices_match, round_up_to_nearest_multiple
@@ -69,6 +70,9 @@ def copy(self) -> Float8BlockQuantizer:
 
         return quantizer
 
+    def _value_fields(self) -> Tuple[str, ...]:
+        return ("dtype", "block_len", "amax_epsilon", "force_pow_2_scales", "block_scaling_dim")
+
     def update_quantized(
         self,
         src: torch.Tensor,
@@ -211,6 +215,9 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]:
         return Float8BlockScaling
 
 
+register_value_opaque_quantizer(Float8BlockQuantizer)
+
+
 class Float8BlockwiseQTensor(Float8BlockwiseQTensorStorage, QuantizedTensor):
     """Tensor class with FP8 data quantized via NxN blocks or 1xN blocks.
 
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index e26abf7df0..56b1ecfb09 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -18,6 +18,7 @@
 from ..utils import canonicalize_process_group, devices_match
 from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func
 from ..quantized_tensor import QuantizedTensor, Quantizer
+from ..dynamo import register_value_opaque_quantizer
 from ._quantization_helpers import _IdentityFunc
 from ..constants import dist_group_type, DType
 
@@ -386,6 +387,14 @@ def supports_only_rowwise_all_gather(self) -> bool:
         """
         return True
 
+    def _value_fields(self) -> Tuple[str, ...]:
+        # ``amax_reduction_group`` is intentionally excluded: it is a deprecated
+        # process group (not a value) and is restored as ``None`` on rebuild.
+        return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction")
+
+
+register_value_opaque_quantizer(Float8CurrentScalingQuantizer)
+
 
 class Float8Tensor(Float8TensorStorage, QuantizedTensor):
     """Experimental tensor class with FP8 data
diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
index d759aaf5c4..a3746b3088 100644
--- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py
+++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
@@ -18,6 +18,7 @@
 from ..utils import devices_match, round_up_to_nearest_multiple
 from .storage.mxfp8_tensor_storage import MXFP8TensorStorage, _FromMXFP8Func
 from ..quantized_tensor import QuantizedTensor, Quantizer
+from ..dynamo import register_value_opaque_quantizer
 from ._quantization_helpers import _IdentityFunc
 
 aten = torch.ops.aten
@@ -57,6 +58,9 @@ def copy(self) -> MXFP8Quantizer:
 
         return quantizer
 
+    def _value_fields(self) -> Tuple[str, ...]:
+        return ("dtype",)
+
     def update_quantized(
         self,
         src: torch.Tensor,
@@ -1058,3 +1062,6 @@ def backward(
             )
             return dgrad, None
         return grad.view(ctx.shape), None
+
+
+register_value_opaque_quantizer(MXFP8Quantizer)
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index aa92be004f..573c78907c 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -23,6 +23,7 @@
 
 from .storage.nvfp4_tensor_storage import NVFP4TensorStorage, _FromNVFP4Func
 from ..quantized_tensor import QuantizedTensor, Quantizer
+from ..dynamo import register_value_opaque_quantizer
 from ._quantization_helpers import _IdentityFunc
 
 aten = torch.ops.aten
@@ -333,6 +334,30 @@ def _canonicalized_amax_reduction_group(self) -> dist_group_type:
     def _get_compatible_recipe(self) -> Union[type[Recipe], None]:
         return NVFP4BlockScaling
 
+    def _value_fields(self) -> Tuple[str, ...]:
+        # ``amax_reduction_group`` is intentionally excluded: it is a deprecated
+        # process group (not a value) and is restored as ``None`` on rebuild.
+        # ``rht_matrix_random_sign_mask_t`` is derived (from
+        # ``_with_random_sign_mask`` and the device) but is stored verbatim so
+        # reconstruction does not need to touch the device.
+        return (
+            "dtype",
+            "with_rht",
+            "with_post_rht_amax",
+            "with_2d_quantization",
+            "stochastic_rounding",
+            "row_scaled_nvfp4",
+            "nvfp4_use_4over6",
+            "nvfp4_e4m3_max",
+            "nvfp4_4over6_err_mode",
+            "_with_random_sign_mask",
+            "rht_matrix_random_sign_mask_t",
+            "with_amax_reduction",
+        )
+
+
+register_value_opaque_quantizer(NVFP4Quantizer)
+
 
 class NVFP4Tensor(NVFP4TensorStorage, QuantizedTensor):
     """Quantized tensor class with FP4 data

From c4ad54c1b1f28d9d3402cff92ce244fb9f1a2bc4 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Sat, 6 Jun 2026 14:51:06 +0200
Subject: [PATCH 02/28] [PyTorch] Drop quantizer value registry; reconstruct
 via __fx_repr__ globals

Follow-up to the value-opaque quantizer support:

- Remove the module-level _QUANTIZER_VALUE_REGISTRY (qualname -> class) and
  _quantizer_from_value_key. __fx_repr__ now captures the quantizer class
  directly in the FX globals and reconstructs via _rebuild_quantizer(cls, items),
  matching how PyTorch's own value opaque types (e.g. DTensor placements)
  reconstruct themselves. This removes global mutable state and the qualname
  collision risk.
- Consolidate the quantizer value-object tests in test_torch_compile.py down to
  two functions and exercise reconstruction through the public __fx_repr__ path
  instead of internal helpers.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 138 +++---------------
 transformer_engine/pytorch/dynamo.py          |  53 +++----
 .../pytorch/quantized_tensor.py               |  22 +--
 3 files changed, 51 insertions(+), 162 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 8adcc88e61..4491e23881 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -36,11 +36,6 @@
     Float8Quantizer,
     Float8BlockQuantizer,
     MXFP8Quantizer,
-    NVFP4Quantizer,
-)
-from transformer_engine.pytorch.dynamo import (
-    register_value_opaque_quantizer,
-    _quantizer_from_value_key,
 )
 from utils import recipe_id
 
@@ -396,145 +391,60 @@ def fn(inp):
 
 
 # ---------------------------------------------------------------------------
-# Value-opaque quantizers: eager value semantics + FX reconstruction
+# Value-opaque quantizers
 #
-# The tensorless quantizers (current-scaling FP8, FP8 blockwise, MXFP8, NVFP4)
-# are torch.compile *value* opaque types: they provide value-based
-# ``__eq__`` / ``__hash__`` and an evaluable ``__fx_repr__`` (see
-# ``torch._library.opaque_object`` Note [Opaque Objects]). These tests exercise
-# the eager value semantics and the FX reconstruction round-trip. They are
-# CPU-friendly except for NVFP4 (whose constructor touches the current CUDA
-# device).
+# Tensorless quantizers (MXFP8, FP8 blockwise, FP8 current-scaling) are
+# torch.compile *value* opaque types: value-based ``__eq__`` / ``__hash__`` plus
+# an evaluable ``__fx_repr__`` that rebuilds an equal object (see
+# ``torch._library.opaque_object`` Note [Opaque Objects]).
 # ---------------------------------------------------------------------------
 
 
-def _mxfp8(dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True):
-    return MXFP8Quantizer(fp8_dtype=dtype, rowwise=rowwise, columnwise=columnwise)
+def _mxfp8(dtype=tex.DType.kFloat8E4M3):
+    return MXFP8Quantizer(fp8_dtype=dtype)
 
 
-def _blockwise(dtype=tex.DType.kFloat8E4M3, force_pow_2_scales=True, block_scaling_dim=2):
+def _blockwise(force_pow_2_scales=True):
     return Float8BlockQuantizer(
-        fp8_dtype=dtype,
+        fp8_dtype=tex.DType.kFloat8E4M3,
         rowwise=True,
         columnwise=True,
         force_pow_2_scales=force_pow_2_scales,
-        block_scaling_dim=block_scaling_dim,
     )
 
 
-def _current_scaling(dtype=tex.DType.kFloat8E4M3, force_pow_2_scales=False, amax_epsilon=0.0):
+def _current_scaling(amax_epsilon=0.0):
     return Float8CurrentScalingQuantizer(
-        fp8_dtype=dtype,
+        fp8_dtype=tex.DType.kFloat8E4M3,
         device=torch.device("cpu"),
-        force_pow_2_scales=force_pow_2_scales,
         amax_epsilon=amax_epsilon,
     )
 
 
-# (factory, kwargs_for_a_different_but_valid_config)
-_CPU_VALUE_QUANTIZERS = [
+# (factory, kwargs producing a different-but-valid config)
+_VALUE_QUANTIZERS = [
     pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"),
     pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"),
     pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"),
 ]
 
 
-@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS)
-def test_quantizer_value_equality(factory, other_kwargs):
-    """Same config -> equal & same hash; different config -> not equal."""
-    a = factory()
-    b = factory()
+@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS)
+def test_quantizer_value_object(factory, other_kwargs):
+    """Value semantics + ``__fx_repr__`` round-trip via the production FX path."""
+    a, b = factory(), factory()
+    # Same config -> equal, same hash, interchangeable as a dict/set key.
     assert a is not b
     assert a == b
     assert hash(a) == hash(b)
+    assert {a: "x"}[b] == "x"
+    # Different config -> not equal.
+    assert a != factory(**other_kwargs)
 
-    c = factory(**other_kwargs)
-    assert a != c
-    # Usage flags participate in the value.
-    d = factory()
-    d.set_usage(rowwise=False, columnwise=False)
-    assert a != d
-
-
-@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS)
-def test_quantizer_usable_in_set_and_dict(factory, other_kwargs):
-    a = factory()
-    b = factory()
-    c = factory(**other_kwargs)
-    assert len({a, b, c}) == 2
-    mapping = {a: "x"}
-    assert mapping[b] == "x"
-
-
-@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS)
-def test_quantizer_cross_type_inequality(factory, other_kwargs):
-    a = factory()
-    other = _current_scaling() if not isinstance(a, Float8CurrentScalingQuantizer) else _mxfp8()
-    assert a != other
-
-
-@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS)
-def test_quantizer_fx_repr_roundtrip(factory, other_kwargs):
-    """``__fx_repr__`` returns an evaluable expression rebuilding an equal object."""
-    a = factory()
+    # ``__fx_repr__`` (used by torch.compile codegen) rebuilds an equal object.
     repr_str, globals_ = a.__fx_repr__()
-    assert isinstance(repr_str, str)
-    assert isinstance(globals_, dict)
     rebuilt = eval(repr_str, dict(globals_))  # pylint: disable=eval-used
-    assert rebuilt == a
-    assert rebuilt is not a
+    assert rebuilt == a and rebuilt is not a
     assert hash(rebuilt) == hash(a)
-
-
-@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS)
-def test_quantizer_value_key_reconstruction(factory, other_kwargs):
-    a = factory()
-    rebuilt = _quantizer_from_value_key(a._value_key())
-    assert type(rebuilt) is type(a)
-    assert rebuilt == a
-    # The deprecated amax-reduction process group is never carried in the value.
+    # The deprecated amax-reduction group is never part of the value.
     assert getattr(rebuilt, "amax_reduction_group", None) is None
-
-
-def test_quantizer_delayed_scaling_keeps_identity_semantics():
-    """Float8Quantizer holds live tensors -> identity (not value) semantics."""
-    scale = torch.ones(1)
-    amax = torch.zeros(1)
-    a = Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3)
-    b = Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3)
-    assert a._value_fields() is None
-    assert a == a
-    assert a != b  # distinct instances are not equal despite identical config
-    assert hash(a) == object.__hash__(a)
-
-
-@pytest.mark.parametrize("factory, other_kwargs", _CPU_VALUE_QUANTIZERS)
-def test_quantizer_registration_is_idempotent_and_tolerant(factory, other_kwargs):
-    """Re-registering must not raise, regardless of PyTorch opaque-object support."""
-    cls = type(factory())
-    register_value_opaque_quantizer(cls)
-    register_value_opaque_quantizer(cls)
-
-    if not _opaque_available:
-        pytest.skip("PyTorch build without opaque-object API")
-    from torch._library.opaque_object import is_opaque_value_type
-
-    assert is_opaque_value_type(cls)
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4Quantizer requires CUDA")
-def test_quantizer_nvfp4_value_semantics():
-    a = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1)
-    b = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1)
-    assert a == b
-    assert hash(a) == hash(b)
-
-    c = NVFP4Quantizer(fp4_dtype=tex.DType.kFloat4E2M1, with_rht=not a.with_rht)
-    assert a != c
-
-    rebuilt = _quantizer_from_value_key(a._value_key())
-    assert rebuilt == a
-    assert rebuilt.amax_reduction_group is None
-
-    repr_str, globals_ = a.__fx_repr__()
-    assert eval(repr_str, dict(globals_)) == a  # pylint: disable=eval-used
diff --git a/transformer_engine/pytorch/dynamo.py b/transformer_engine/pytorch/dynamo.py
index 26aff6f59c..16f73920ae 100644
--- a/transformer_engine/pytorch/dynamo.py
+++ b/transformer_engine/pytorch/dynamo.py
@@ -9,12 +9,14 @@
 
   * :func:`register_value_opaque_quantizer` -- attaches the ``__fx_repr__`` used
     by FX codegen and registers the quantizer class with
-    ``torch._library.opaque_object``. It is a no-op (other than populating the
-    local registry) on PyTorch builds without the opaque-object API, so
-    importing Transformer Engine never fails on older PyTorch -- only
-    torch.compile specialization on the quantizer is unavailable there.
-  * :func:`_quantizer_from_value_key` -- rebuilds a quantizer constant from its
-    value key inside the generated FX graph.
+    ``torch._library.opaque_object``. It is a no-op on PyTorch builds without
+    the opaque-object API, so importing Transformer Engine never fails on older
+    PyTorch -- only torch.compile specialization on the quantizer is
+    unavailable there.
+  * :func:`_rebuild_quantizer` -- rebuilds a quantizer constant from its value
+    items inside the generated FX graph. The quantizer class is captured
+    directly in the FX globals (see :func:`_quantizer_fx_repr`), so no global
+    class registry is needed.
 
 The eager value semantics (``__eq__`` / ``__hash__`` / ``_value_key`` /
 ``_value_fields``) live on the quantizer itself; see
@@ -22,7 +24,10 @@
 
 See ``torch._library.opaque_object`` Note [Opaque Objects] for the contract a
 value-typed opaque object must satisfy (``__eq__`` / ``__hash__`` /
-``__fx_repr__``).
+``__fx_repr__``). The ``__fx_repr__`` contract -- ``(repr_string, {name: type})``
+where ``repr_string`` references the names in the dict -- is exactly how
+PyTorch's own value opaque types (e.g. DTensor placements) reconstruct
+themselves, including across the on-disk compile cache.
 """
 
 from __future__ import annotations
@@ -31,26 +36,16 @@
 from .constants import DType
 
 
-# Maps a quantizer class qualname to the class object. A value key stores only
-# the qualname, so reconstruction looks the class up here. Populated by
-# ``register_value_opaque_quantizer`` at import time of each tensor module; this
-# avoids importing the tensor modules into this module (which would create an
-# import cycle).
-_QUANTIZER_VALUE_REGISTRY: Dict[str, type] = {}
-
-
-def _quantizer_from_value_key(key: Tuple[Any, ...]) -> Any:
-    """Rebuild a tensorless quantizer from its value key.
+def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any:
+    """Rebuild a tensorless quantizer of type *cls* from its value items.
 
     Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the
     generated FX code calls this to materialize the quantizer constant. The
     deprecated amax-reduction process group is never part of the value, so a
     reconstructed quantizer always starts with no stored group.
     """
-    qualname, items = key[0], key[1]
-    cls = _QUANTIZER_VALUE_REGISTRY[qualname]
     # Bypass ``__init__`` and restore the value attributes directly: the value
-    # key already captures every value-defining field (including derived ones),
+    # items already capture every value-defining field (including derived ones),
     # and the constructors have heterogeneous signatures / side effects.
     obj = cls.__new__(cls)
     field_names = set()
@@ -71,12 +66,15 @@ def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]:
     """``__fx_repr__`` for value-opaque quantizers (attached at registration).
 
     Returns an evaluable expression that rebuilds the quantizer via
-    :func:`_quantizer_from_value_key`, together with the globals needed to
-    evaluate it.
+    :func:`_rebuild_quantizer`, capturing both the helper and the quantizer
+    class itself in the FX globals so codegen can resolve them with no global
+    registry and no qualname collisions.
     """
+    cls = type(self)
+    items = self._value_key()[1]
     return (
-        f"_quantizer_from_value_key({self._value_key()!r})",
-        {"_quantizer_from_value_key": _quantizer_from_value_key},
+        f"_rebuild_quantizer({cls.__name__}, {items!r})",
+        {"_rebuild_quantizer": _rebuild_quantizer, cls.__name__: cls},
     )
 
 
@@ -85,16 +83,13 @@ def register_value_opaque_quantizer(cls: type) -> None:
 
     Attaches ``__fx_repr__`` and registers the class with
     ``torch._library.opaque_object``. Safe to call on any PyTorch build: on
-    versions without the opaque-object API it only records the class in the
-    local registry and attaches ``__fx_repr__`` (both harmless), so Transformer
-    Engine keeps importing and running in eager mode.
+    versions without the opaque-object API it only attaches ``__fx_repr__``
+    (harmless), so Transformer Engine keeps importing and running in eager mode.
 
     The quantizer class must already provide value ``__eq__`` / ``__hash__`` and
     a non-``None`` ``_value_fields`` (see
     :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`).
     """
-    _QUANTIZER_VALUE_REGISTRY[cls.__qualname__] = cls
-
     # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the
     # class, so attach it before registering.
     if "__fx_repr__" not in cls.__dict__:
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index f37b3cc63d..6809893a40 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -408,20 +408,6 @@ def get_usages(self) -> Dict[str, bool]:
             "columnwise": self.columnwise_usage,
         }
 
-    # ----- Value-object identity (torch.compile opaque value support) -----
-    # A *tensorless* quantizer (one whose entire state is a handful of plain,
-    # reproducible scalars -- no live tensors, no process groups) behaves like a
-    # value: two instances with the same configuration are interchangeable. Such
-    # quantizers opt into value-based ``__eq__`` / ``__hash__`` by overriding
-    # ``_value_fields``. Quantizers that keep the default (e.g. delayed-scaling
-    # ``Float8Quantizer``, which holds live scale/amax tensors, and any custom
-    # quantizer) retain the default identity semantics.
-    #
-    # This is the eager-side half of registering the quantizer as a torch.compile
-    # *value* opaque type; the torch.compile glue (``__fx_repr__``, FX
-    # reconstruction and ``register_opaque_type``) lives in
-    # ``transformer_engine.pytorch.dynamo``.
-
     #: Attributes shared by every quantizer that take part in value identity.
     _BASE_VALUE_FIELDS: Tuple[str, ...] = (
         "rowwise_usage",
@@ -433,11 +419,9 @@ def get_usages(self) -> Dict[str, bool]:
     def _value_fields(self) -> Optional[Tuple[str, ...]]:
         """Subclass-specific value-defining attribute names, or ``None``.
 
-        Returning ``None`` (the default) means the quantizer is *not* a value
-        object and keeps identity-based equality/hashing. Tensorless quantizers
-        override this to return the tuple of attribute names that, together with
-        :attr:`_BASE_VALUE_FIELDS`, fully determine their value (excluding
-        non-value state such as a deprecated amax-reduction process group).
+        Returning ``None`` (the default) means the quantizer cannot be represented as
+        a value opaque object and keeps identity-based equality/hashing.
+        This also means, that torch.compile will not be able to optimize the quantizer.
         """
         return None
 

From a06324bd4f10354d76b5fbec133f07709bd6a1da Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Sun, 7 Jun 2026 15:42:31 +0200
Subject: [PATCH 03/28] [PyTorch] Split dynamo.py into a dynamo/ package

Replace the single dynamo.py module with a dynamo/ package so the
torch.compile glue can grow with a clear responsibility split across the
stacked branches. This branch owns the value-opaque quantizer layer.

  * dynamo/quantizer_opaque.py -- register_value_opaque_quantizer and helpers
  * dynamo/__init__.py -- re-exports the public API so callers keep importing
    from transformer_engine.pytorch.dynamo unchanged

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/dynamo/__init__.py  | 18 ++++++++++++++++++
 .../{dynamo.py => dynamo/quantizer_opaque.py}  |  7 +++----
 2 files changed, 21 insertions(+), 4 deletions(-)
 create mode 100644 transformer_engine/pytorch/dynamo/__init__.py
 rename transformer_engine/pytorch/{dynamo.py => dynamo/quantizer_opaque.py} (95%)

diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py
new file mode 100644
index 0000000000..44ca61d470
--- /dev/null
+++ b/transformer_engine/pytorch/dynamo/__init__.py
@@ -0,0 +1,18 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""torch.compile glue for Transformer Engine.
+
+Public API is re-exported here so callers keep importing from
+``transformer_engine.pytorch.dynamo`` regardless of the internal module layout:
+
+  * :mod:`.quantizer_opaque` -- make a tensorless quantizer a torch.compile
+    *value* opaque type (:func:`register_value_opaque_quantizer`).
+"""
+
+from .quantizer_opaque import register_value_opaque_quantizer
+
+__all__ = [
+    "register_value_opaque_quantizer",
+]
diff --git a/transformer_engine/pytorch/dynamo.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
similarity index 95%
rename from transformer_engine/pytorch/dynamo.py
rename to transformer_engine/pytorch/dynamo/quantizer_opaque.py
index 16f73920ae..c7455846ee 100644
--- a/transformer_engine/pytorch/dynamo.py
+++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
@@ -2,10 +2,9 @@
 #
 # See LICENSE for license information.
 
-"""torch.compile glue for Transformer Engine quantizers.
+"""Value-opaque quantizers for torch.compile.
 
-This module isolates the torch.compile-specific plumbing that turns a
-*tensorless* quantizer into a torch.compile **value** opaque type:
+Turns a *tensorless* quantizer into a torch.compile **value** opaque type:
 
   * :func:`register_value_opaque_quantizer` -- attaches the ``__fx_repr__`` used
     by FX codegen and registers the quantizer class with
@@ -33,7 +32,7 @@ class registry is needed.
 from __future__ import annotations
 from typing import Any, Dict, Tuple
 
-from .constants import DType
+from ..constants import DType
 
 
 def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any:

From ea5b396b9e1e3ee67905c58b50e30989f5ab095d Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 8 Jun 2026 12:39:40 +0200
Subject: [PATCH 04/28] [PyTorch] Raise in quantizer __fx_repr__ when a process
 group is stored

A value-opaque quantizer must not carry live distributed state. Scan the
quantizer attributes in __fx_repr__ and raise TypeError if any holds a
torch.distributed.ProcessGroup (e.g. a non-None deprecated amax_reduction_group),
so it cannot be silently baked into a torch.compile FX graph. Clarify the related
comments accordingly.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           |  7 --
 transformer_engine/pytorch/dynamo/__init__.py |  9 +--
 .../pytorch/dynamo/quantizer_opaque.py        | 70 ++++++++++---------
 .../pytorch/quantized_tensor.py               |  4 +-
 .../pytorch/tensor/float8_tensor.py           |  3 +-
 .../pytorch/tensor/nvfp4_tensor.py            |  3 +-
 6 files changed, 45 insertions(+), 51 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 4491e23881..9a7a4a356a 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -392,11 +392,6 @@ def fn(inp):
 
 # ---------------------------------------------------------------------------
 # Value-opaque quantizers
-#
-# Tensorless quantizers (MXFP8, FP8 blockwise, FP8 current-scaling) are
-# torch.compile *value* opaque types: value-based ``__eq__`` / ``__hash__`` plus
-# an evaluable ``__fx_repr__`` that rebuilds an equal object (see
-# ``torch._library.opaque_object`` Note [Opaque Objects]).
 # ---------------------------------------------------------------------------
 
 
@@ -446,5 +441,3 @@ def test_quantizer_value_object(factory, other_kwargs):
     rebuilt = eval(repr_str, dict(globals_))  # pylint: disable=eval-used
     assert rebuilt == a and rebuilt is not a
     assert hash(rebuilt) == hash(a)
-    # The deprecated amax-reduction group is never part of the value.
-    assert getattr(rebuilt, "amax_reduction_group", None) is None
diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py
index 44ca61d470..aae8b9cff6 100644
--- a/transformer_engine/pytorch/dynamo/__init__.py
+++ b/transformer_engine/pytorch/dynamo/__init__.py
@@ -2,14 +2,7 @@
 #
 # See LICENSE for license information.
 
-"""torch.compile glue for Transformer Engine.
-
-Public API is re-exported here so callers keep importing from
-``transformer_engine.pytorch.dynamo`` regardless of the internal module layout:
-
-  * :mod:`.quantizer_opaque` -- make a tensorless quantizer a torch.compile
-    *value* opaque type (:func:`register_value_opaque_quantizer`).
-"""
+"""torch.compile glue for Transformer Engine."""
 
 from .quantizer_opaque import register_value_opaque_quantizer
 
diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
index c7455846ee..55bc8326e6 100644
--- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py
+++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
@@ -2,46 +2,35 @@
 #
 # See LICENSE for license information.
 
-"""Value-opaque quantizers for torch.compile.
-
-Turns a *tensorless* quantizer into a torch.compile **value** opaque type:
-
-  * :func:`register_value_opaque_quantizer` -- attaches the ``__fx_repr__`` used
-    by FX codegen and registers the quantizer class with
-    ``torch._library.opaque_object``. It is a no-op on PyTorch builds without
-    the opaque-object API, so importing Transformer Engine never fails on older
-    PyTorch -- only torch.compile specialization on the quantizer is
-    unavailable there.
-  * :func:`_rebuild_quantizer` -- rebuilds a quantizer constant from its value
-    items inside the generated FX graph. The quantizer class is captured
-    directly in the FX globals (see :func:`_quantizer_fx_repr`), so no global
-    class registry is needed.
-
-The eager value semantics (``__eq__`` / ``__hash__`` / ``_value_key`` /
-``_value_fields``) live on the quantizer itself; see
-:class:`transformer_engine.pytorch.quantized_tensor.Quantizer`.
-
-See ``torch._library.opaque_object`` Note [Opaque Objects] for the contract a
-value-typed opaque object must satisfy (``__eq__`` / ``__hash__`` /
-``__fx_repr__``). The ``__fx_repr__`` contract -- ``(repr_string, {name: type})``
-where ``repr_string`` references the names in the dict -- is exactly how
-PyTorch's own value opaque types (e.g. DTensor placements) reconstruct
-themselves, including across the on-disk compile cache.
-"""
+"""Value-opaque quantizers for torch.compile."""
 
 from __future__ import annotations
 from typing import Any, Dict, Tuple
 
-from ..constants import DType
+from ..constants import DType, dist_group_type
+
+
+def _contains_process_group(value: Any) -> bool:
+    """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``.
+
+    Checks the value directly and one level of ``tuple``/``list`` nesting, which
+    covers the shapes a quantizer value field could plausibly take.
+    """
+    if isinstance(value, dist_group_type):
+        return True
+    if isinstance(value, (tuple, list)):
+        return any(_contains_process_group(item) for item in value)
+    return False
 
 
 def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any:
     """Rebuild a tensorless quantizer of type *cls* from its value items.
 
     Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the
-    generated FX code calls this to materialize the quantizer constant. The
-    deprecated amax-reduction process group is never part of the value, so a
-    reconstructed quantizer always starts with no stored group.
+    generated FX code calls this to materialize the quantizer constant. A
+    quantizer that actually stores a process group never reaches this path:
+    ``__fx_repr__`` raises for it. The deprecated amax-reduction group is not a
+    value field, so the rebuilt quantizer simply has no group attribute.
     """
     # Bypass ``__init__`` and restore the value attributes directly: the value
     # items already capture every value-defining field (including derived ones),
@@ -53,9 +42,9 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any:
             value = DType.cast(value)
         object.__setattr__(obj, name, value)
         field_names.add(name)
-    # The deprecated amax-reduction process group is excluded from the value;
-    # restore it as ``None`` for quantizers that still carry the fallback so
-    # attribute access keeps working.
+    # The deprecated amax-reduction group is not a value field. Quantizers that
+    # actually hold a group error out in ``__fx_repr__`` before reaching here, so
+    # this only initializes the (groupless) attribute to keep access working.
     if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"):
         object.__setattr__(obj, "amax_reduction_group", None)
     return obj
@@ -68,8 +57,23 @@ def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]:
     :func:`_rebuild_quantizer`, capturing both the helper and the quantizer
     class itself in the FX globals so codegen can resolve them with no global
     registry and no qualname collisions.
+
+    Raises ``TypeError`` if the quantizer stores a process group (e.g. a
+    non-``None`` deprecated ``amax_reduction_group``): live distributed state
+    must never be baked into the graph as a constant, so such a quantizer cannot
+    be used with ``torch.compile``. Pass the reduction group per quantize call
+    instead of storing it on the quantizer.
     """
     cls = type(self)
+    for name, value in vars(self).items():
+        if _contains_process_group(value):
+            raise TypeError(
+                f"{cls.__name__} cannot be used with torch.compile: attribute "
+                f"{name!r} holds a torch.distributed.ProcessGroup, which is live "
+                "distributed state and must not be baked into an FX graph as a "
+                "constant. Pass the amax reduction group per quantize call instead "
+                "of storing it on the quantizer."
+            )
     items = self._value_key()[1]
     return (
         f"_rebuild_quantizer({cls.__name__}, {items!r})",
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index 6809893a40..3612c1080d 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -421,7 +421,9 @@ def _value_fields(self) -> Optional[Tuple[str, ...]]:
 
         Returning ``None`` (the default) means the quantizer cannot be represented as
         a value opaque object and keeps identity-based equality/hashing.
-        This also means, that torch.compile will not be able to optimize the quantizer.
+        This also means that passing such a quantizer as an argument to a custom op
+        causes a graph break under torch.compile, since it cannot be baked into the
+        FX graph as a constant.
         """
         return None
 
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index 56b1ecfb09..0310c3855c 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -389,7 +389,8 @@ def supports_only_rowwise_all_gather(self) -> bool:
 
     def _value_fields(self) -> Tuple[str, ...]:
         # ``amax_reduction_group`` is intentionally excluded: it is a deprecated
-        # process group (not a value) and is restored as ``None`` on rebuild.
+        # process group (not a value). If one is actually stored, ``__fx_repr__``
+        # raises so it can never be baked into a torch.compile graph.
         return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction")
 
 
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 573c78907c..4bca783922 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -336,7 +336,8 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]:
 
     def _value_fields(self) -> Tuple[str, ...]:
         # ``amax_reduction_group`` is intentionally excluded: it is a deprecated
-        # process group (not a value) and is restored as ``None`` on rebuild.
+        # process group (not a value). If one is actually stored, ``__fx_repr__``
+        # raises so it can never be baked into a torch.compile graph.
         # ``rht_matrix_random_sign_mask_t`` is derived (from
         # ``_with_random_sign_mask`` and the device) but is stored verbatim so
         # reconstruction does not need to touch the device.

From aa65e34e2ed8ba784543caa703ef7962062d3546 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 8 Jun 2026 16:17:36 +0200
Subject: [PATCH 05/28] [PyTorch] Cover NVFP4 in quantizer value-object test

NVFP4Quantizer is registered as a value-opaque quantizer but was missing
from the value-semantics / __fx_repr__ round-trip test. Add it to
_VALUE_QUANTIZERS (skipped without CUDA, which it needs to construct).

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

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 9a7a4a356a..3405001e04 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -27,7 +27,7 @@
 from transformer_engine.pytorch.quantization import QuantizerRole
 from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear
 from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer
-from transformer_engine.pytorch.quantization import QuantizerRole
+from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
 from transformer_engine.pytorch import (
     is_fp8_available,
     is_mxfp8_available,
@@ -416,11 +416,29 @@ def _current_scaling(amax_epsilon=0.0):
     )
 
 
+def _nvfp4(with_rht=False):
+    return NVFP4Quantizer(
+        fp4_dtype=tex.DType.kFloat4E2M1,
+        rowwise=True,
+        columnwise=True,
+        with_rht=with_rht,
+    )
+
+
 # (factory, kwargs producing a different-but-valid config)
 _VALUE_QUANTIZERS = [
     pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"),
     pytest.param(_blockwise, {"force_pow_2_scales": False}, id="float8_blockwise"),
     pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"),
+    pytest.param(
+        _nvfp4,
+        {"with_rht": True},
+        id="nvfp4",
+        marks=pytest.mark.skipif(
+            not torch.cuda.is_available(),
+            reason="NVFP4Quantizer requires CUDA to construct",
+        ),
+    ),
 ]
 
 

From e1b1db6b1a9de58a6868f1a706659733792476cc Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 16 Jun 2026 17:13:15 +0200
Subject: [PATCH 06/28] Reject a value quantizer that carries an amax reduction
 group in __eq__/__hash__

The amax reduction group is excluded from the value key, so a value quantizer
that stored one would compare/hash equal to a groupless one and let torch.compile
reuse a graph that skips the reduction. __eq__/__hash__ now raise (mirroring
__fx_repr__, which already rejects any process-group-bearing quantizer). The
group should be passed per quantize call, not stored on the quantizer.

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

diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index 3612c1080d..b7357b0e7b 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -444,6 +444,18 @@ def _value_key(self) -> Tuple[Any, ...]:
             items.append((name, value))
         return (type(self).__qualname__, tuple(items))
 
+    def _check_value_has_no_amax_reduction_group(self) -> None:
+        # The amax reduction group is not part of the value key, so a value
+        # quantizer that stores one would compare/hash equal to a groupless one
+        # and let torch.compile reuse a graph that skips the reduction. Reject it
+        # (mirrors ``__fx_repr__``); pass the group per quantize call instead.
+        if getattr(self, "amax_reduction_group", None) is not None:
+            raise TypeError(
+                f"{type(self).__name__} with a non-None amax_reduction_group cannot be "
+                "used as a value object; pass the amax reduction group per quantize call "
+                "instead of storing it on the quantizer."
+            )
+
     def __eq__(self, other: object) -> Any:
         # Value quantizers compare by configuration; everything else keeps the
         # default identity semantics (returning ``NotImplemented`` makes Python
@@ -454,11 +466,14 @@ def __eq__(self, other: object) -> Any:
             return NotImplemented
         if other._value_fields() is None:
             return NotImplemented
+        self._check_value_has_no_amax_reduction_group()
+        other._check_value_has_no_amax_reduction_group()
         return self._value_key() == other._value_key()
 
     def __hash__(self) -> int:
         if self._value_fields() is None:
             return object.__hash__(self)
+        self._check_value_has_no_amax_reduction_group()
         return hash(self._value_key())
 
 

From 8c33d0ec4dd8ff255a75827fd643d3619e6ae9af Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 16 Jun 2026 18:03:29 +0200
Subject: [PATCH 07/28] Recognize value-opaque quantizers via a class flag

Add is_value_opaque_quantizer() + the _te_compile_value_opaque flag stamped at
registration, so dynamo-traced code can detect registered quantizers (and fall
back to eager for unregistered ones).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/dynamo/__init__.py     |  3 ++-
 .../pytorch/dynamo/quantizer_opaque.py            | 15 +++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py
index aae8b9cff6..ee860c78e3 100644
--- a/transformer_engine/pytorch/dynamo/__init__.py
+++ b/transformer_engine/pytorch/dynamo/__init__.py
@@ -4,8 +4,9 @@
 
 """torch.compile glue for Transformer Engine."""
 
-from .quantizer_opaque import register_value_opaque_quantizer
+from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer
 
 __all__ = [
     "register_value_opaque_quantizer",
+    "is_value_opaque_quantizer",
 ]
diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
index 55bc8326e6..6cb9552b71 100644
--- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py
+++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
@@ -10,6 +10,17 @@
 from ..constants import DType, dist_group_type
 
 
+# Class attribute stamped on quantizers registered as torch.compile value-opaque
+# types.
+_VALUE_OPAQUE_FLAG = "_te_compile_value_opaque"
+
+
+def is_value_opaque_quantizer(quantizer: Any) -> bool:
+    """Whether *quantizer*'s class is registered as a torch.compile value-opaque
+    type."""
+    return getattr(quantizer, _VALUE_OPAQUE_FLAG, False)
+
+
 def _contains_process_group(value: Any) -> bool:
     """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``.
 
@@ -93,6 +104,10 @@ def register_value_opaque_quantizer(cls: type) -> None:
     a non-``None`` ``_value_fields`` (see
     :class:`transformer_engine.pytorch.quantized_tensor.Quantizer`).
     """
+    # Stamp the class so it can be recognized as value-opaque in dynamo-traced
+    # code (used to fall back to eager for unregistered quantizers).
+    setattr(cls, _VALUE_OPAQUE_FLAG, True)
+
     # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the
     # class, so attach it before registering.
     if "__fx_repr__" not in cls.__dict__:

From 945f62dadd18d5d0bb7b68c54f4bfb9767699521 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 29 Jun 2026 11:34:37 +0200
Subject: [PATCH 08/28] Address review: narrow opaque-type except, add
 fullgraph test, fix nvfp4 value key

- Narrow register_opaque_type except to (RuntimeError, TypeError): the API is
  already imported above, so ImportError/AttributeError there only mask real errors.
- Add test_quantizer_value_object_fullgraph exercising torch.compile(fullgraph=True)
  end-to-end to verify opaque-type registration took effect.
- Restore missing NVFP4Quantizer._with_random_sign_mask assignment required by
  _value_fields()/_value_key().

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py            | 18 ++++++++++++++++++
 .../pytorch/dynamo/quantizer_opaque.py         |  5 +++--
 .../pytorch/tensor/nvfp4_tensor.py             |  1 +
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 3405001e04..5e1f753f08 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -459,3 +459,21 @@ def test_quantizer_value_object(factory, other_kwargs):
     rebuilt = eval(repr_str, dict(globals_))  # pylint: disable=eval-used
     assert rebuilt == a and rebuilt is not a
     assert hash(rebuilt) == hash(a)
+
+
+@pytest.mark.skipif(
+    not _opaque_available,
+    reason="torch.compile opaque-object support requires PyTorch >= 2.11",
+)
+@pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS)
+def test_quantizer_value_object_fullgraph(factory, other_kwargs):
+    """Quantizer survives torch.compile(fullgraph=True) - verifies registration took effect."""
+
+    def fn(quantizer):
+        return quantizer
+
+    torch._dynamo.reset()
+    compiled = torch.compile(fn, fullgraph=True)
+
+    quantizer = factory()
+    assert compiled(quantizer) is quantizer
diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
index 6cb9552b71..6d630d665c 100644
--- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py
+++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
@@ -128,6 +128,7 @@ def register_value_opaque_quantizer(cls: type) -> None:
 
     try:
         register_opaque_type(cls, typ="value")
-    except (ImportError, AttributeError, RuntimeError, TypeError):
-        # Tolerate partial / experimental opaque-object support.
+    except (RuntimeError, TypeError):
+        # Keep TE importable: registration must never crash the import, e.g. on
+        # PyTorch versions with only partial / experimental opaque-object support.
         pass
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 4bca783922..ffc5f97eca 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -174,6 +174,7 @@ def __init__(
         self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper()
         if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"):
             raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.")
+        self._with_random_sign_mask = with_random_sign_mask
         self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht(
             with_random_sign_mask, torch.cuda.current_device()
         )

From e3c8f430883762ac7be289616bf79f40eb522a0d Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 29 Jun 2026 12:05:01 +0200
Subject: [PATCH 09/28] Restore NVFP4 rht_matrix on value-key rebuild; assert
 quantize round-trip

_rebuild_quantizer only restores value-key fields, so a reconstructed
NVFP4Quantizer was missing the derived rht_matrix tensor (not hashable, so not
in the value key) and failed at copy()/quantize time. Add a _rebuild_derived_state
hook (called by _rebuild_quantizer) that NVFP4Quantizer uses to rebuild rht_matrix
from _with_random_sign_mask (lru_cache -> cheap).

Extend test_quantizer_value_object to also quantize with the original and the
rebuilt quantizer and require bit-exact results (gated on HW support), so a
field the kernel needs but the value key omits can no longer slip through.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 27 +++++++++++++++++--
 .../pytorch/dynamo/quantizer_opaque.py        |  5 ++++
 .../pytorch/tensor/nvfp4_tensor.py            | 10 +++++++
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 5e1f753f08..da4f96d2f3 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -416,7 +416,10 @@ def _current_scaling(amax_epsilon=0.0):
     )
 
 
-def _nvfp4(with_rht=False):
+def _nvfp4(with_rht=True):
+    # Default with_rht=True so the quantize round-trip below exercises the
+    # derived ``rht_matrix`` tensor (the field most likely to be dropped on
+    # value-key reconstruction).
     return NVFP4Quantizer(
         fp4_dtype=tex.DType.kFloat4E2M1,
         rowwise=True,
@@ -425,6 +428,17 @@ def _nvfp4(with_rht=False):
     )
 
 
+def _hw_available(quantizer):
+    """Whether this HW can actually run the quantize kernel for *quantizer*."""
+    if isinstance(quantizer, MXFP8Quantizer):
+        return mxfp8_available
+    if isinstance(quantizer, NVFP4Quantizer):
+        return nvfp4_available
+    if isinstance(quantizer, Float8BlockQuantizer):
+        return fp8_block_scaling_available
+    return fp8_available  # Float8CurrentScalingQuantizer
+
+
 # (factory, kwargs producing a different-but-valid config)
 _VALUE_QUANTIZERS = [
     pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"),
@@ -432,7 +446,7 @@ def _nvfp4(with_rht=False):
     pytest.param(_current_scaling, {"amax_epsilon": 1e-4}, id="float8_current_scaling"),
     pytest.param(
         _nvfp4,
-        {"with_rht": True},
+        {"with_rht": False},
         id="nvfp4",
         marks=pytest.mark.skipif(
             not torch.cuda.is_available(),
@@ -460,6 +474,15 @@ def test_quantizer_value_object(factory, other_kwargs):
     assert rebuilt == a and rebuilt is not a
     assert hash(rebuilt) == hash(a)
 
+    # The rebuilt quantizer must also *behave* identically, not just compare
+    # equal: equality only looks at the value key, so a field the kernel needs
+    # but that is absent from the key (e.g. NVFP4's derived ``rht_matrix``) would
+    # slip through the checks above and only blow up at quantize time. Run the
+    # real quantize kernel on both and require bit-exact results.
+    if torch.cuda.is_available() and _hw_available(a):
+        x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")
+        torch.testing.assert_close(rebuilt(x).dequantize(), a(x).dequantize(), rtol=0.0, atol=0.0)
+
 
 @pytest.mark.skipif(
     not _opaque_available,
diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
index 6d630d665c..97532b2790 100644
--- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py
+++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
@@ -58,6 +58,11 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any:
     # this only initializes the (groupless) attribute to keep access working.
     if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"):
         object.__setattr__(obj, "amax_reduction_group", None)
+    # Restore non-value derived state that ``__init__`` would normally build but
+    # that cannot live in the value key (e.g. NVFP4's ``rht_matrix`` tensor).
+    finalize = getattr(obj, "_rebuild_derived_state", None)
+    if finalize is not None:
+        finalize()
     return obj
 
 
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index ffc5f97eca..f2f30cdcc5 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -186,6 +186,16 @@ def __getstate__(self):
         state["amax_reduction_group"] = None
         return state
 
+    def _rebuild_derived_state(self) -> None:
+        """Restore the derived ``rht_matrix`` after value-key reconstruction.
+
+        ``rht_matrix`` is a ``torch.Tensor`` built from ``_with_random_sign_mask``
+        and the device, so it cannot be part of the (hashable) value key.
+        ``_rebuild_quantizer`` calls this hook to rebuild it; the ``lru_cache`` on
+        :func:`get_rht_matrix` makes an already-seen (flag, device) a cheap hit.
+        """
+        self.rht_matrix = get_rht_matrix(self._with_random_sign_mask, torch.cuda.current_device())
+
     def update_quantized(
         self,
         src: torch.Tensor,

From 3f6862137b30aa30c2980cc1cb9cc750599b02fa Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 29 Jun 2026 14:46:49 +0200
Subject: [PATCH 10/28] Enforce process-group rejection in _value_key, not
 __fx_repr__; add test

Move the ProcessGroup guard out of the (overridable) __fx_repr__ into
Quantizer._value_key -- the single point every value-materialization path
(__eq__/__hash__/__fx_repr__) goes through -- so a custom __fx_repr__ can no
longer bypass it. Generalizes the old amax-only check to any field holding a
ProcessGroup. Add a test that a value quantizer carrying a live group raises.

Addresses review on NVIDIA#3152.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 21 ++++++++
 .../pytorch/dynamo/quantizer_opaque.py        | 34 +++----------
 .../pytorch/quantized_tensor.py               | 50 +++++++++++++------
 3 files changed, 61 insertions(+), 44 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index da4f96d2f3..dc98e13708 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -484,6 +484,27 @@ def test_quantizer_value_object(factory, other_kwargs):
         torch.testing.assert_close(rebuilt(x).dequantize(), a(x).dequantize(), rtol=0.0, atol=0.0)
 
 
+def test_value_quantizer_rejects_process_group():
+    """A value quantizer holding a live ProcessGroup must refuse to be turned
+    into a value key / FX constant (raise), not silently drop the group."""
+    import torch.distributed as dist  # pylint: disable=import-outside-toplevel
+
+    created = not dist.is_initialized()
+    if created:
+        dist.init_process_group(backend="gloo", store=dist.HashStore(), rank=0, world_size=1)
+    try:
+        q = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)
+        q.amax_reduction_group = dist.group.WORLD
+        # Every value-materialization path must reject it (hash, eq, __fx_repr__).
+        with pytest.raises(TypeError):
+            hash(q)
+        with pytest.raises(TypeError):
+            q.__fx_repr__()
+    finally:
+        if created:
+            dist.destroy_process_group()
+
+
 @pytest.mark.skipif(
     not _opaque_available,
     reason="torch.compile opaque-object support requires PyTorch >= 2.11",
diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
index 97532b2790..37e718689d 100644
--- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py
+++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
@@ -7,7 +7,7 @@
 from __future__ import annotations
 from typing import Any, Dict, Tuple
 
-from ..constants import DType, dist_group_type
+from ..constants import DType
 
 
 # Class attribute stamped on quantizers registered as torch.compile value-opaque
@@ -21,19 +21,6 @@ def is_value_opaque_quantizer(quantizer: Any) -> bool:
     return getattr(quantizer, _VALUE_OPAQUE_FLAG, False)
 
 
-def _contains_process_group(value: Any) -> bool:
-    """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``.
-
-    Checks the value directly and one level of ``tuple``/``list`` nesting, which
-    covers the shapes a quantizer value field could plausibly take.
-    """
-    if isinstance(value, dist_group_type):
-        return True
-    if isinstance(value, (tuple, list)):
-        return any(_contains_process_group(item) for item in value)
-    return False
-
-
 def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any:
     """Rebuild a tensorless quantizer of type *cls* from its value items.
 
@@ -74,22 +61,13 @@ def _quantizer_fx_repr(self: Any) -> Tuple[str, Dict[str, Any]]:
     class itself in the FX globals so codegen can resolve them with no global
     registry and no qualname collisions.
 
-    Raises ``TypeError`` if the quantizer stores a process group (e.g. a
-    non-``None`` deprecated ``amax_reduction_group``): live distributed state
-    must never be baked into the graph as a constant, so such a quantizer cannot
-    be used with ``torch.compile``. Pass the reduction group per quantize call
-    instead of storing it on the quantizer.
+    Raises ``TypeError`` (via :meth:`Quantizer._value_key`) if the quantizer
+    stores a process group (e.g. a non-``None`` deprecated
+    ``amax_reduction_group``): live distributed state must never be baked into
+    the graph as a constant. Pass the reduction group per quantize call instead
+    of storing it on the quantizer.
     """
     cls = type(self)
-    for name, value in vars(self).items():
-        if _contains_process_group(value):
-            raise TypeError(
-                f"{cls.__name__} cannot be used with torch.compile: attribute "
-                f"{name!r} holds a torch.distributed.ProcessGroup, which is live "
-                "distributed state and must not be baked into an FX graph as a "
-                "constant. Pass the amax reduction group per quantize call instead "
-                "of storing it on the quantizer."
-            )
     items = self._value_key()[1]
     return (
         f"_rebuild_quantizer({cls.__name__}, {items!r})",
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index b7357b0e7b..033a35f8e1 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -16,6 +16,7 @@
 import transformer_engine_torch as tex
 
 from transformer_engine.common.recipe import Recipe
+from transformer_engine.pytorch.constants import dist_group_type
 from transformer_engine.pytorch.tensor._quantization_helpers import (
     _QuantizeFunc,
     _IdentityFunc,
@@ -23,6 +24,19 @@
 )
 
 
+def _contains_process_group(value: Any) -> bool:
+    """Whether *value* is (or nests) a ``torch.distributed.ProcessGroup``.
+
+    Checks the value directly and one level of ``tuple``/``list`` nesting, which
+    covers the shapes a quantizer value field could plausibly take.
+    """
+    if isinstance(value, dist_group_type):
+        return True
+    if isinstance(value, (tuple, list)):
+        return any(_contains_process_group(item) for item in value)
+    return False
+
+
 # Custom ops that should pass through __torch_dispatch__ without unwrapping
 # QuantizedTensor subclasses (e.g. Float8Tensor). Register ops here that
 # handle quantized tensors internally.
@@ -427,6 +441,24 @@ def _value_fields(self) -> Optional[Tuple[str, ...]]:
         """
         return None
 
+    def _check_value_has_no_process_group(self) -> None:
+        # A value quantizer is baked into the FX graph as a constant via its
+        # value key, which cannot carry live distributed state. Enforced here --
+        # the single point every value-materialization path (``__eq__`` /
+        # ``__hash__`` / ``__fx_repr__``) goes through -- so a custom
+        # ``__fx_repr__`` cannot bypass it. Reject any field holding a
+        # ProcessGroup (e.g. the deprecated ``amax_reduction_group``) rather than
+        # silently dropping it; pass the reduction group per quantize call.
+        for name, value in vars(self).items():
+            if _contains_process_group(value):
+                raise TypeError(
+                    f"{type(self).__name__} cannot be used as a torch.compile value "
+                    f"object: attribute {name!r} holds a torch.distributed.ProcessGroup, "
+                    "which is live distributed state and must not be baked into an FX "
+                    "graph. Pass the amax reduction group per quantize call instead of "
+                    "storing it on the quantizer."
+                )
+
     def _value_key(self) -> Tuple[Any, ...]:
         """Hashable, reproducible key identifying this quantizer's value.
 
@@ -434,6 +466,7 @@ def _value_key(self) -> Tuple[Any, ...]:
         """
         fields = self._value_fields()  # pylint: disable=assignment-from-none
         assert fields is not None, f"{type(self).__name__} is not a value quantizer"
+        self._check_value_has_no_process_group()
         items = []
         for name in self._BASE_VALUE_FIELDS + tuple(fields):
             value = getattr(self, name)
@@ -444,36 +477,21 @@ def _value_key(self) -> Tuple[Any, ...]:
             items.append((name, value))
         return (type(self).__qualname__, tuple(items))
 
-    def _check_value_has_no_amax_reduction_group(self) -> None:
-        # The amax reduction group is not part of the value key, so a value
-        # quantizer that stores one would compare/hash equal to a groupless one
-        # and let torch.compile reuse a graph that skips the reduction. Reject it
-        # (mirrors ``__fx_repr__``); pass the group per quantize call instead.
-        if getattr(self, "amax_reduction_group", None) is not None:
-            raise TypeError(
-                f"{type(self).__name__} with a non-None amax_reduction_group cannot be "
-                "used as a value object; pass the amax reduction group per quantize call "
-                "instead of storing it on the quantizer."
-            )
-
     def __eq__(self, other: object) -> Any:
         # Value quantizers compare by configuration; everything else keeps the
         # default identity semantics (returning ``NotImplemented`` makes Python
-        # fall back to identity).
+        # fall back to identity). ``_value_key`` rejects a stored ProcessGroup.
         if self is other:
             return True
         if self._value_fields() is None or type(self) is not type(other):
             return NotImplemented
         if other._value_fields() is None:
             return NotImplemented
-        self._check_value_has_no_amax_reduction_group()
-        other._check_value_has_no_amax_reduction_group()
         return self._value_key() == other._value_key()
 
     def __hash__(self) -> int:
         if self._value_fields() is None:
             return object.__hash__(self)
-        self._check_value_has_no_amax_reduction_group()
         return hash(self._value_key())
 
 

From 32d17683f036390a6825902ebacba0e7bf592050 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 29 Jun 2026 15:08:39 +0200
Subject: [PATCH 11/28] Strengthen fullgraph test: quantize/dequantize via a
 custom op, not passthrough

Replace the trivial pass-through fullgraph test with one that drives each
production quantizer through a minimal custom op (quantize + dequantize) under
torch.compile(fullgraph=True) and compares to eager -- so the opaque-type
registration is actually exercised inside the graph (a graph break would make
fullgraph=True raise). Op registration sits right before the test. Also drop
stale comments referencing the old __fx_repr__-side process-group guard.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 54 ++++++++++++++++---
 .../pytorch/dynamo/quantizer_opaque.py        | 10 ++--
 .../pytorch/tensor/nvfp4_tensor.py            |  3 +-
 3 files changed, 51 insertions(+), 16 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index dc98e13708..63cb82eca8 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -505,19 +505,59 @@ def test_value_quantizer_rejects_process_group():
             dist.destroy_process_group()
 
 
+if _opaque_available:
+    # A minimal custom op taking a tensor and a value-opaque quantizer that
+    # quantizes + dequantizes inside it, one per production quantizer class.
+    # ``test_quantizer_value_object_fullgraph`` drives this under
+    # ``torch.compile(fullgraph=True)`` so the quantizer is used *inside* the
+    # graph -- proving the opaque-type registration took effect (a graph break
+    # would make ``fullgraph=True`` raise).
+    _qdq_lib = torch.library.Library("test_te_qdq", "DEF")
+    _QDQ_OPS = {}
+    for _qcls in (
+        MXFP8Quantizer,
+        Float8BlockQuantizer,
+        Float8CurrentScalingQuantizer,
+        NVFP4Quantizer,
+    ):
+        _op = f"qdq_{_qcls.__name__}"
+        _qdq_lib.define(f"{_op}(Tensor x, {get_opaque_type_name(_qcls)} q) -> Tensor")
+
+        @torch.library.impl(f"test_te_qdq::{_op}", "CompositeExplicitAutograd", lib=_qdq_lib)
+        def _qdq_impl(x, q):
+            return q(x).dequantize()
+
+        @torch.library.register_fake(f"test_te_qdq::{_op}", lib=_qdq_lib)
+        def _qdq_fake(x, q):
+            return torch.empty_like(x)
+
+        _QDQ_OPS[_qcls] = getattr(torch.ops.test_te_qdq, _op)
+
+
 @pytest.mark.skipif(
     not _opaque_available,
     reason="torch.compile opaque-object support requires PyTorch >= 2.11",
 )
 @pytest.mark.parametrize("factory, other_kwargs", _VALUE_QUANTIZERS)
 def test_quantizer_value_object_fullgraph(factory, other_kwargs):
-    """Quantizer survives torch.compile(fullgraph=True) - verifies registration took effect."""
+    """Quantizer is usable *inside* a torch.compile(fullgraph=True) graph.
 
-    def fn(quantizer):
-        return quantizer
+    A custom op quantizes+dequantizes with the (opaque value) quantizer; the
+    compiled result must match eager. ``fullgraph=True`` raises on any graph
+    break, so this proves the opaque-type registration actually took effect --
+    unlike merely passing the quantizer through.
+    """
+    q = factory()
+    if not (torch.cuda.is_available() and _hw_available(q)):
+        pytest.skip("format not supported on this HW")
 
-    torch._dynamo.reset()
-    compiled = torch.compile(fn, fullgraph=True)
+    op = _QDQ_OPS[type(q)]
+    x = torch.randn(128, 256, dtype=torch.bfloat16, device="cuda")
+
+    def fn(inp):
+        return op(inp, q)
 
-    quantizer = factory()
-    assert compiled(quantizer) is quantizer
+    ref = fn(x)
+    torch._dynamo.reset()
+    out = torch.compile(fn, fullgraph=True)(x)
+    torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0)
diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
index 37e718689d..409cb979d1 100644
--- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py
+++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
@@ -25,10 +25,7 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any:
     """Rebuild a tensorless quantizer of type *cls* from its value items.
 
     Referenced by the ``__fx_repr__`` emitted for value-opaque quantizers; the
-    generated FX code calls this to materialize the quantizer constant. A
-    quantizer that actually stores a process group never reaches this path:
-    ``__fx_repr__`` raises for it. The deprecated amax-reduction group is not a
-    value field, so the rebuilt quantizer simply has no group attribute.
+    generated FX code calls this to materialize the quantizer constant.
     """
     # Bypass ``__init__`` and restore the value attributes directly: the value
     # items already capture every value-defining field (including derived ones),
@@ -40,9 +37,8 @@ def _rebuild_quantizer(cls: type, items: Tuple[Tuple[str, Any], ...]) -> Any:
             value = DType.cast(value)
         object.__setattr__(obj, name, value)
         field_names.add(name)
-    # The deprecated amax-reduction group is not a value field. Quantizers that
-    # actually hold a group error out in ``__fx_repr__`` before reaching here, so
-    # this only initializes the (groupless) attribute to keep access working.
+    # The deprecated amax-reduction group is not a value field; initialize it to
+    # None so attribute access keeps working on the rebuilt quantizer.
     if "with_amax_reduction" in field_names and not hasattr(obj, "amax_reduction_group"):
         object.__setattr__(obj, "amax_reduction_group", None)
     # Restore non-value derived state that ``__init__`` would normally build but
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index f2f30cdcc5..d9cd534606 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -347,8 +347,7 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]:
 
     def _value_fields(self) -> Tuple[str, ...]:
         # ``amax_reduction_group`` is intentionally excluded: it is a deprecated
-        # process group (not a value). If one is actually stored, ``__fx_repr__``
-        # raises so it can never be baked into a torch.compile graph.
+        # process group, not a value (``_value_key`` rejects a stored group).
         # ``rht_matrix_random_sign_mask_t`` is derived (from
         # ``_with_random_sign_mask`` and the device) but is stored verbatim so
         # reconstruction does not need to touch the device.

From 28bde9e7040a8b4c87b992cd5717b7518d19e000 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 29 Jun 2026 15:29:11 +0200
Subject: [PATCH 12/28] Clarify comments: rht_matrix_random_sign_mask_t
 derivation; why the opaque flag

- rht_matrix_random_sign_mask_t is a device-independent int derived from
  _with_random_sign_mask (the device only places a throwaway tensor); fix the
  misleading comment.
- Explain why registration uses a class attribute, not a registry set:
  is_value_opaque_quantizer is traced inside the compile graph and dynamo can
  bake a getattr constant but cannot do 'type(q) in set' on the opaque class.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/dynamo/quantizer_opaque.py | 8 ++++++--
 transformer_engine/pytorch/tensor/nvfp4_tensor.py     | 6 +++---
 2 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
index 409cb979d1..33e831ef38 100644
--- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py
+++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
@@ -10,8 +10,12 @@
 from ..constants import DType
 
 
-# Class attribute stamped on quantizers registered as torch.compile value-opaque
-# types.
+# Registration marks the class with this attribute instead of recording it in a
+# module-level set. ``is_value_opaque_quantizer`` runs *inside* the torch.compile
+# graph (``Linear.forward`` consults it): Dynamo can trace a ``getattr`` on the
+# opaque quantizer and bake the result as a constant, but cannot evaluate
+# ``type(q) in some_set`` -- it has no equality/hash rules for the opaque class
+# object, so a set/dict lookup graph-breaks under ``fullgraph=True``.
 _VALUE_OPAQUE_FLAG = "_te_compile_value_opaque"
 
 
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index d9cd534606..c8c0a7b854 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -348,9 +348,9 @@ def _get_compatible_recipe(self) -> Union[type[Recipe], None]:
     def _value_fields(self) -> Tuple[str, ...]:
         # ``amax_reduction_group`` is intentionally excluded: it is a deprecated
         # process group, not a value (``_value_key`` rejects a stored group).
-        # ``rht_matrix_random_sign_mask_t`` is derived (from
-        # ``_with_random_sign_mask`` and the device) but is stored verbatim so
-        # reconstruction does not need to touch the device.
+        # ``rht_matrix_random_sign_mask_t`` is a device-independent int derived
+        # from ``_with_random_sign_mask``; kept in the key so the rebuilt
+        # quantizer carries it without recomputation.
         return (
             "dtype",
             "with_rht",

From 2c3c5df0fb3488b8e95670066ac53f958640a6de Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 29 Jun 2026 15:30:44 +0200
Subject: [PATCH 13/28] Reword opaque-flag comment: self-contained, no Linear
 reference

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 .../pytorch/dynamo/quantizer_opaque.py               | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
index 33e831ef38..98349e12ba 100644
--- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py
+++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
@@ -10,12 +10,12 @@
 from ..constants import DType
 
 
-# Registration marks the class with this attribute instead of recording it in a
-# module-level set. ``is_value_opaque_quantizer`` runs *inside* the torch.compile
-# graph (``Linear.forward`` consults it): Dynamo can trace a ``getattr`` on the
-# opaque quantizer and bake the result as a constant, but cannot evaluate
-# ``type(q) in some_set`` -- it has no equality/hash rules for the opaque class
-# object, so a set/dict lookup graph-breaks under ``fullgraph=True``.
+# Registration marks the class with this attribute rather than recording it in a
+# module-level set. It looks odd but is a deliberate workaround: the check must
+# stay traceable when it runs inside a torch.compile graph -- Dynamo can bake a
+# ``getattr`` on the opaque quantizer into a constant, but cannot evaluate
+# ``type(q) in some_set`` (no equality/hash rules for the opaque class object),
+# which would graph-break under ``fullgraph=True``.
 _VALUE_OPAQUE_FLAG = "_te_compile_value_opaque"
 
 

From 826f271ebb28466cdd94b9ee5cc6875eb61d4a97 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 29 Jun 2026 15:38:53 +0200
Subject: [PATCH 14/28] Cover is_opaque_value_type with the import-safety guard
 too

is_opaque_value_type(cls) sat between the import guard and the
register_opaque_type guard, so on a partial/experimental opaque-object build it
could raise RuntimeError/TypeError and crash TE import. Move it inside the same
except so the 'registration never crashes import' promise holds for both calls.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/dynamo/quantizer_opaque.py | 11 +++++------
 1 file changed, 5 insertions(+), 6 deletions(-)

diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
index 98349e12ba..8b8b3caa69 100644
--- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py
+++ b/transformer_engine/pytorch/dynamo/quantizer_opaque.py
@@ -106,12 +106,11 @@ def register_value_opaque_quantizer(cls: type) -> None:
         # still work; torch.compile specialization on the quantizer does not.
         return
 
-    if is_opaque_value_type(cls):
-        return
-
     try:
-        register_opaque_type(cls, typ="value")
+        if not is_opaque_value_type(cls):
+            register_opaque_type(cls, typ="value")
     except (RuntimeError, TypeError):
-        # Keep TE importable: registration must never crash the import, e.g. on
-        # PyTorch versions with only partial / experimental opaque-object support.
+        # Keep TE importable: neither the opaque-type query nor the registration
+        # must crash the import, e.g. on PyTorch versions with only partial /
+        # experimental opaque-object support.
         pass

From ad1ccce0e53e01833c4342d3ed8eb342a0064929 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 16 Jun 2026 17:29:49 +0200
Subject: [PATCH 15/28] Add TensorProto mechanism for data-free quantized
 tensor allocation

Squashed PR #8 (tensor_proto_mechanism) onto the rebased base. Adds TensorProto
(pure-Python, torch.compile-traceable quantized-tensor allocation via
Quantizer.alloc_tensors + storage __tensor_flatten__/__tensor_unflatten__),
Linear fake fwd/bwd impls for the custom-op path, and tests.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py           | 251 ++++++++++++++++
 transformer_engine/pytorch/dynamo/__init__.py |   3 +
 .../pytorch/dynamo/tensor_proto.py            | 172 +++++++++++
 transformer_engine/pytorch/module/linear.py   | 275 +++++++++++++++++-
 .../pytorch/quantized_tensor.py               | 127 ++++++++
 .../pytorch/tensor/float8_blockwise_tensor.py |  35 ++-
 .../pytorch/tensor/float8_tensor.py           |  34 ++-
 .../pytorch/tensor/mxfp8_tensor.py            |  36 ++-
 .../pytorch/tensor/nvfp4_tensor.py            |  45 ++-
 .../float8_blockwise_tensor_storage.py        |   9 +
 .../tensor/storage/float8_tensor_storage.py   |   8 +
 .../tensor/storage/mxfp8_tensor_storage.py    |   9 +
 .../tensor/storage/nvfp4_tensor_storage.py    |  11 +
 13 files changed, 1007 insertions(+), 8 deletions(-)
 create mode 100644 transformer_engine/pytorch/dynamo/tensor_proto.py

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 63cb82eca8..4a3f22a291 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -6,6 +6,7 @@
 
 import pytest
 import torch
+from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
 
 try:
     from torch._opaque_base import OpaqueBaseMeta
@@ -28,6 +29,9 @@
 from transformer_engine.pytorch.ops.basic.basic_linear import BasicLinear
 from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer
 from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
+from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, _STORAGE_REGISTRY
+from transformer_engine.pytorch.dynamo import TensorProto, to_tensor_proto
+from transformer_engine.pytorch.dynamo.tensor_proto import _contiguous_stride
 from transformer_engine.pytorch import (
     is_fp8_available,
     is_mxfp8_available,
@@ -473,6 +477,8 @@ def test_quantizer_value_object(factory, other_kwargs):
     rebuilt = eval(repr_str, dict(globals_))  # pylint: disable=eval-used
     assert rebuilt == a and rebuilt is not a
     assert hash(rebuilt) == hash(a)
+    # The deprecated amax-reduction group is never part of the value.
+    assert getattr(rebuilt, "amax_reduction_group", None) is None
 
     # The rebuilt quantizer must also *behave* identically, not just compare
     # equal: equality only looks at the value key, so a field the kernel needs
@@ -561,3 +567,248 @@ def fn(inp):
     torch._dynamo.reset()
     out = torch.compile(fn, fullgraph=True)(x)
     torch.testing.assert_close(out, ref, rtol=0.0, atol=0.0)
+
+
+# ---------------------------------------------------------------------------
+# torch.compile-traceable allocation primitives + TensorProto
+# ---------------------------------------------------------------------------
+
+
+# (factory, logical shape) -- shapes respect MXFP8 (mult. of 32) / blockwise (128)
+# / NVFP4 (mult. of 16) constraints.
+_PROTO_QUANTIZERS = [
+    pytest.param(_current_scaling, (4, 8), id="fp8_current_scaling"),
+    pytest.param(_mxfp8, (64, 128), id="mxfp8"),
+    pytest.param(_blockwise, (128, 256), id="fp8_blockwise"),
+    pytest.param(
+        _nvfp4,
+        (64, 128),
+        id="nvfp4",
+        marks=pytest.mark.skipif(
+            not torch.cuda.is_available(),
+            reason="NVFP4Quantizer requires CUDA to construct",
+        ),
+    ),
+]
+
+
+def _build_from_primitives(quantizer, shape, dtype, device="cpu"):
+    """Assemble a quantized tensor straight from the quantizer primitives:
+    ``alloc_tensors`` (buffers) + ``create_metadata`` (ctx) + the storage's
+    ``__tensor_unflatten__`` -- i.e. exactly what ``TensorProto.create_tensor``
+    does, but without going through :class:`TensorProto`.
+    """
+    names = tuple(quantizer._describe_buffers(shape))  # pylint: disable=protected-access
+    ctx = quantizer.create_metadata(shape, dtype=dtype)
+    buffers = quantizer.alloc_tensors(shape, device=device)
+    inner = {name: buffers[name] for name in names}
+    storage_cls = _STORAGE_REGISTRY[ctx["cls"]]
+    return storage_cls.__tensor_unflatten__(inner, ctx, tuple(shape), _contiguous_stride(shape))
+
+
+def _signature(tensor, names):
+    """Comparable shape/dtype fingerprint of a tensor and its inner buffers."""
+    sig = {"__shape__": tuple(tensor.shape), "__dtype__": tensor.dtype}
+    for name in names:
+        buf = getattr(tensor, name)
+        sig[name] = (tuple(buf.shape), buf.dtype)
+    return sig
+
+
+def _skip_if_dequantize_unsupported(q):
+    """Skip when this HW can't run ``dequantize()`` for the quantizer's format.
+
+    ``dequantize()`` runs the real kernel on CUDA, so each format has its own
+    availability gate (mirrors the ``is_*_available`` checks in test_numerics).
+    """
+    if isinstance(q, MXFP8Quantizer):
+        if not mxfp8_available:
+            pytest.skip(reason_for_no_mxfp8)
+    elif isinstance(q, NVFP4Quantizer):
+        if not nvfp4_available:
+            pytest.skip("NVFP4 is not available")
+    elif isinstance(q, Float8BlockQuantizer):
+        if not fp8_block_scaling_available:
+            pytest.skip("FP8 block scaling is not available")
+    elif not fp8_available:  # Float8 current scaling
+        pytest.skip(reason_for_no_fp8)
+
+
+# ----- Quantizer primitives -----
+
+
+@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
+def test_primitives_unflatten_compiles(factory, shape):
+    """create_metadata + alloc_tensors + __tensor_unflatten__ compose and trace
+    under ``fullgraph=True`` (CPU), without TensorProto."""
+    q = factory()
+    names = tuple(q._describe_buffers(shape))  # pylint: disable=protected-access
+
+    def fn(x):
+        t = _build_from_primitives(q, shape, x.dtype, device=x.device)
+        # Read every buffer into the result so the alloc + unflatten can't be
+        # eliminated as dead code -- forces the whole build path into the graph.
+        acc = x.new_zeros(())
+        for name in names:
+            acc = acc + getattr(t, name).float().sum()
+        return acc
+
+    x = torch.zeros(*shape, dtype=torch.bfloat16)
+    torch._dynamo.reset()
+    out = torch.compile(fn, fullgraph=True)(x)
+    assert out.shape == ()
+
+
+@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
+def test_alloc_tensors_fake(factory, shape):
+    """``alloc_tensors`` produces FakeTensors with the described shapes/dtypes."""
+    q = factory()
+    bufs = q._describe_buffers(shape)  # pylint: disable=protected-access
+    with FakeTensorMode():
+        alloc = q.alloc_tensors(shape, device="cpu")
+    assert set(alloc) == set(bufs)
+    for name, (buf_shape, buf_dtype) in bufs.items():
+        assert isinstance(alloc[name], FakeTensor)
+        assert tuple(alloc[name].shape) == tuple(buf_shape)
+        assert alloc[name].dtype == buf_dtype
+
+
+@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
+def test_storage_flatten_unflatten_roundtrip(factory, shape):
+    """Storage ``__tensor_flatten__`` / ``__tensor_unflatten__`` round-trips.
+
+    Build a tensor from ``alloc_tensors`` + ``create_metadata``, flatten it, then
+    unflatten and verify shape/dtype and every inner buffer match before vs after.
+    """
+    q = factory()
+    _skip_if_dequantize_unsupported(q)
+
+    tensor = _build_from_primitives(q, shape, torch.bfloat16)
+    names = tuple(q._describe_buffers(shape))  # pylint: disable=protected-access
+    # Fill buffers with deterministic data (empty() may contain NaNs) so the
+    # round-trip can be checked by value via dequantize().
+    for name in names:
+        buf = getattr(tensor, name)
+        buf.copy_(torch.arange(buf.numel(), device=buf.device).reshape(buf.shape))
+    before = _signature(tensor, names)
+    expected = tensor.dequantize()
+
+    flat_names, flat_ctx = tensor.__tensor_flatten__()
+    assert set(flat_names) == set(names)
+    inner = {name: getattr(tensor, name) for name in flat_names}
+    rebuilt = type(tensor).__tensor_unflatten__(
+        inner, flat_ctx, tuple(tensor.shape), tensor.stride()
+    )
+
+    assert isinstance(rebuilt, QuantizedTensor)
+    assert _signature(rebuilt, flat_names) == before
+    # The reconstructed tensor dequantizes to the same values.
+    torch.testing.assert_close(rebuilt.dequantize(), expected, atol=0, rtol=0, equal_nan=True)
+
+
+# ----- TensorProto -----
+
+
+@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
+def test_tensor_proto_matches_primitives(factory, shape):
+    """TensorProto is a thin wrapper: its ``create_metadata`` /
+    ``create_inner_tensors`` / ``create_tensor`` match building everything
+    directly from the quantizer primitives."""
+    q = factory()
+    proto = TensorProto(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu"))
+    assert proto.is_quantized
+
+    # Metadata matches the quantizer's.
+    assert proto.create_metadata() == q.create_metadata(shape, dtype=torch.bfloat16)
+
+    # inner_names + create_inner_tensors match _describe_buffers.
+    bufs = q._describe_buffers(shape)  # pylint: disable=protected-access
+    names = tuple(bufs)
+    assert proto.inner_names() == names
+    inner = proto.create_inner_tensors()
+    assert len(inner) == len(names)
+    for name, buf in zip(names, inner):
+        exp_shape, exp_dtype = bufs[name]
+        assert tuple(buf.shape) == tuple(exp_shape)
+        assert buf.dtype == exp_dtype
+
+    # The assembled tensor matches one built directly from the primitives.
+    direct = _build_from_primitives(q, shape, torch.bfloat16)
+    assert _signature(proto.create_tensor(), names) == _signature(direct, names)
+
+
+@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
+def test_tensor_proto_create_tensor_eager(factory, shape):
+    """``create_tensor`` (no fake) yields a real quantized tensor."""
+    q = factory()
+    proto = TensorProto(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu"))
+    out = proto.create_tensor()
+    assert isinstance(out, QuantizedTensor)
+    assert tuple(out.shape) == tuple(shape)
+    assert out.dtype == torch.bfloat16
+    for name in proto.inner_names():
+        assert not isinstance(getattr(out, name), FakeTensor)
+
+
+@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
+def test_tensor_proto_create_tensor_fake(factory, shape):
+    """``create_tensor`` under ``FakeTensorMode`` yields a fake-backed quantized
+    tensor with the right shape/dtype and fake inner buffers."""
+    q = factory()
+    proto = TensorProto(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu"))
+    with FakeTensorMode():
+        out = proto.create_tensor()
+    assert isinstance(out, QuantizedTensor)
+    assert tuple(out.shape) == tuple(shape)
+    assert out.dtype == torch.bfloat16
+    for name in proto.inner_names():
+        assert isinstance(getattr(out, name), FakeTensor)
+
+
+@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
+def test_tensor_proto_create_tensor_compiles(factory, shape):
+    """``TensorProto.create_tensor`` traces under ``fullgraph=True`` (CPU)."""
+    q = factory()
+
+    def fn(x):
+        proto = TensorProto(shape=tuple(x.shape), dtype=x.dtype, quantizer=q, device=x.device)
+        t = proto.create_tensor()
+        acc = x.new_zeros(())
+        for name in proto.inner_names():
+            acc = acc + getattr(t, name).float().sum()
+        return acc
+
+    x = torch.zeros(*shape, dtype=torch.bfloat16)
+    torch._dynamo.reset()
+    out = torch.compile(fn, fullgraph=True)(x)
+    assert out.shape == ()
+
+
+def test_to_tensor_proto_plain():
+    """``to_tensor_proto`` describes a plain tensor."""
+    t = torch.empty(2, 3, dtype=torch.float32)
+    proto = to_tensor_proto(t)
+    assert not proto.is_quantized
+    assert proto.shape == (2, 3)
+    assert proto.dtype == torch.float32
+    assert proto.inner_names() == ("data",)
+
+
+@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
+def test_to_tensor_proto_quantized(factory, shape):
+    """``to_tensor_proto`` round-trips a quantized tensor back into a proto."""
+    q = factory()
+    tensor = TensorProto(
+        shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu")
+    ).create_tensor()
+
+    proto = to_tensor_proto(tensor)
+    assert proto.is_quantized
+    assert proto.shape == tuple(shape)
+    assert proto.dtype == torch.bfloat16
+    # Same buffer layout as the original tensor.
+    assert proto.inner_names() == tuple(q._describe_buffers(shape))  # pylint: disable=protected-access
+    # Rebuilding from the derived proto matches the original tensor's structure.
+    assert _signature(proto.create_tensor(), proto.inner_names()) == _signature(
+        tensor, proto.inner_names()
+    )
diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py
index ee860c78e3..f932a7d9c3 100644
--- a/transformer_engine/pytorch/dynamo/__init__.py
+++ b/transformer_engine/pytorch/dynamo/__init__.py
@@ -5,8 +5,11 @@
 """torch.compile glue for Transformer Engine."""
 
 from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer
+from .tensor_proto import TensorProto, to_tensor_proto
 
 __all__ = [
     "register_value_opaque_quantizer",
     "is_value_opaque_quantizer",
+    "TensorProto",
+    "to_tensor_proto",
 ]
diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py
new file mode 100644
index 0000000000..b5248e3ee4
--- /dev/null
+++ b/transformer_engine/pytorch/dynamo/tensor_proto.py
@@ -0,0 +1,172 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""TensorProto: a data-free description of a tensor / quantized tensor."""
+
+from __future__ import annotations
+import copy as _copy
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+
+
+def _contiguous_stride(shape: Tuple[int, ...]) -> Tuple[int, ...]:
+    """Row-major (contiguous) stride for ``shape``."""
+    stride: list = []
+    acc = 1
+    for dim in reversed(shape):
+        stride.append(acc)
+        acc *= dim
+    return tuple(reversed(stride))
+
+
+@dataclass
+class TensorProto:
+    """A data-free *prototype* of a tensor or quantized tensor.
+
+    Captures ``shape`` / ``dtype`` and, for quantized tensors, the
+    (value-opaque) ``quantizer`` -- enough to rebuild a tensor without holding
+    storage. The common abstraction over plain ``torch.Tensor``,
+    ``QuantizedTensorStorage`` and ``QuantizedTensor``, used for custom-op fake
+    impls and for reassembling a quantized tensor from bare buffers.
+    """
+
+    shape: Tuple[int, ...]
+    dtype: torch.dtype
+    quantizer: Optional[Any] = None
+    requires_grad: bool = False
+    device: Optional[torch.device] = field(default=None)
+
+    def __post_init__(self) -> None:
+        # Own a private copy of the quantizer so usage changes (update_usage)
+        # never touch the shared, value-opaque quantizer. The copy inherits the
+        # quantizer's current row-/column-wise usage as this proto's layout.
+        if self.quantizer is not None:
+            q = self.quantizer
+            self.quantizer = q.copy() if hasattr(q, "copy") else _copy.copy(q)
+
+    @property
+    def is_quantized(self) -> bool:
+        """Whether this proto describes a quantized tensor."""
+        return self.quantizer is not None
+
+    def update_usage(
+        self,
+        *,
+        rowwise_usage: Optional[bool] = None,
+        columnwise_usage: Optional[bool] = None,
+    ) -> None:
+        """Mirror ``QuantizedTensor.update_usage`` on the proto's buffer layout.
+
+        Applied to the proto's own quantizer copy, so the shared (value-opaque)
+        quantizer is never mutated. No-op for plain (non-quantized) protos.
+        """
+        if self.quantizer is None:
+            return
+        self.quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage)
+
+    def inner_names(self) -> Tuple[str, ...]:
+        """Names of the flat tensor buffers backing this proto, in order.
+
+        The real op flattens a quantized output via the storage's
+        ``__tensor_flatten__`` -- i.e. ``_FLATTEN_TENSOR_BUFFERS`` order, keeping
+        only the present buffers. ``_describe_buffers`` may emit the same buffers
+        in a different (per-usage) order (e.g. NVFP4 groups each amax right after
+        its scale), so reorder to the canonical flatten order here to keep the
+        fake layout aligned with the real one slot-for-slot.
+        """
+        if self.quantizer is None:
+            return ("data",)
+        # pylint: disable=protected-access
+        described = list(self.quantizer._describe_buffers(tuple(self.shape)).keys())
+        storage_cls = self.quantizer._storage_metadata(self.dtype)["cls"]
+        flatten_order = [attr for attr, _ in storage_cls._FLATTEN_TENSOR_BUFFERS]
+        ordered = [name for name in flatten_order if name in described]
+        ordered += [name for name in described if name not in flatten_order]
+        return tuple(ordered)
+
+    def create_metadata(self) -> Dict[str, Any]:
+        """Data-free ``__tensor_unflatten__`` context describing this tensor."""
+        if self.quantizer is None:
+            return {
+                "is_tensor": True,
+                "is_quantized": False,
+                "dtype": self.dtype,
+                "requires_grad": self.requires_grad,
+            }
+        return self.quantizer.create_metadata(
+            tuple(self.shape), dtype=self.dtype, requires_grad=self.requires_grad
+        )
+
+    def create_inner_tensors(self) -> List[torch.Tensor]:
+        """Materialize the flat inner buffers (in :meth:`inner_names` order).
+
+        Under ``register_fake`` the ``torch.empty`` calls produce ``FakeTensor``s;
+        ``requires_grad`` is left default (managed by ``register_autograd``).
+        """
+        device = self.device if self.device is not None else torch.device("cuda")
+        if self.quantizer is None:
+            return [torch.empty(tuple(self.shape), dtype=self.dtype, device=device)]
+        inner = self.quantizer.alloc_tensors(tuple(self.shape), device=device)
+        return [inner[name] for name in self.inner_names()]
+
+    def create_tensor(self) -> torch.Tensor:
+        """Materialize an (uninitialized) tensor matching this proto (traceable).
+
+        Quantized protos reassemble the :meth:`create_inner_tensors` buffers via
+        the storage's ``__tensor_unflatten__``.
+        """
+        if self.quantizer is None:
+            device = self.device if self.device is not None else torch.device("cuda")
+            return torch.empty(
+                tuple(self.shape),
+                dtype=self.dtype,
+                device=device,
+                requires_grad=self.requires_grad,
+            )
+        from ..quantized_tensor import (  # pylint: disable=import-outside-toplevel
+            _STORAGE_REGISTRY,
+        )
+
+        shape = tuple(self.shape)
+        ctx = self.create_metadata()
+        inner = dict(zip(self.inner_names(), self.create_inner_tensors()))
+        storage_cls = _STORAGE_REGISTRY[ctx["cls"]]
+        return storage_cls.__tensor_unflatten__(inner, ctx, shape, _contiguous_stride(shape))
+
+
+def to_tensor_proto(tensor: Any) -> TensorProto:
+    """Build a :class:`TensorProto` describing ``tensor``.
+
+    Works for plain ``torch.Tensor`` and for ``QuantizedTensorStorage`` /
+    ``QuantizedTensor``. A *bare* storage exposes its shape via ``.size()`` and
+    its (fake) dtype via ``_dtype`` rather than ``.shape`` / ``.dtype``.
+    """
+    from ..quantized_tensor import (  # pylint: disable=import-outside-toplevel
+        QuantizedTensorStorage,
+    )
+
+    requires_grad = bool(getattr(tensor, "requires_grad", False))
+    if isinstance(tensor, QuantizedTensorStorage):
+        shape = getattr(tensor, "shape", None)
+        if shape is None:
+            shape = tensor.size()
+        dtype = getattr(tensor, "dtype", None)
+        if dtype is None:
+            dtype = getattr(tensor, "_dtype", None)
+        return TensorProto(
+            shape=tuple(shape),
+            dtype=dtype,
+            quantizer=getattr(tensor, "_quantizer", None),
+            requires_grad=requires_grad,
+            device=tensor.device,
+        )
+    return TensorProto(
+        shape=tuple(tensor.shape),
+        dtype=tensor.dtype,
+        quantizer=None,
+        requires_grad=requires_grad,
+        device=tensor.device,
+    )
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index fed367bce5..d479d2f67b 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -68,6 +68,7 @@
     prepare_for_saving,
     restore_from_func_ctx,
 )
+from ..dynamo import TensorProto
 from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer
 from ..tensor.mxfp8_tensor import MXFP8Quantizer
 from ..tensor.utils import clear_columnwise_cache, is_custom
@@ -92,7 +93,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 ---
@@ -301,7 +302,15 @@ def _linear_forward_impl(
 
     # Configure tensor-parallel communication
     tp_world_size = get_distributed_world_size(tp_group)
-    backward_needs_input = is_grad_enabled and weight.requires_grad
+    # Use the requires-grad flags captured into ``args`` at op-call time rather
+    # than the live tensors': the fake impl (``_linear_forward_impl_fake``) keys
+    # the number of FP8 inner buffers it emits off ``args.*_requires_grad``, so
+    # the real impl must agree to keep the custom-op output arity stable. Under
+    # ``torch.compile`` with CUDA-graph trees (``mode="reduce-overhead"``) the
+    # static graph inputs are detached during capture, so live
+    # ``weight.requires_grad`` / ``inp.requires_grad`` flip to False mid-capture
+    # and would otherwise diverge from the fake (schema/arity mismatch).
+    backward_needs_input = is_grad_enabled and args.weight_requires_grad
     with_input_all_gather_nccl = (
         parallel_mode == "column" and sequence_parallel and not ub_overlap_ag_fprop
     )
@@ -418,7 +427,7 @@ def _linear_forward_impl(
         # No need to set the quantizer states if weight is already quantized
         # for debug mode we create quantizer every iteration, thus we need to set the quantizer states
         if weight_quantizer is not None and (not isinstance(weight, QuantizedTensor) or debug):
-            columnwise_usage = is_grad_enabled and inp.requires_grad and not is_fsdp2
+            columnwise_usage = is_grad_enabled and args.input_requires_grad and not is_fsdp2
             if backward_override is not None:
                 columnwise_usage = False
             if not columnwise_usage:
@@ -622,6 +631,204 @@ def _linear_forward_impl(
     return out, new_weight_workspace, tensors_to_save_from_forward, None, ctx_attrs
 
 
+def _linear_forward_impl_fake(
+    args: LinearFwdArgs,
+) -> Tuple[TensorProto, Optional[TensorProto], Optional[Tuple[Any, ...]], None, Optional[Dict]]:
+    """Shape/metadata-only twin of :func:`_linear_forward_impl` for torch.compile,
+    returning ``TensorProto`` descriptors for the outputs and saved tensors instead
+    of allocating real data."""
+    if args.fsdp_group is not None and args.is_grad_enabled:
+        raise NotImplementedError(
+            "Compile-time Linear forward does not support manual TE FSDP "
+            "(fsdp_group is not None); use FSDP2 or MCore FSDP."
+        )
+
+    weight = args.weight
+    inp = args.inp
+    bias = args.bias
+    input_quantizer = args.input_quantizer
+    weight_quantizer = args.weight_quantizer
+    output_quantizer = args.output_quantizer
+    fp8 = args.fp8
+    debug = args.debug
+    fp8_or_debug = fp8 or debug
+    is_grad_enabled = args.is_grad_enabled
+    activation_dtype = args.activation_dtype
+    save_original_input = args.save_original_input
+    if args.backward_override == "high_precision":
+        save_original_input = True
+
+    out_features, _ = weight.shape
+    backward_needs_input = is_grad_enabled and args.weight_requires_grad
+
+    own_quantized_input = False
+    inputmat_is_storage = False
+    inputmat_aliases_inp = False
+    if fp8_or_debug:
+        if inp.is_quantized:
+            # Primary-quantized input reused as-is.
+            inputmat_is_storage = True
+            inputmat_aliases_inp = True
+        else:
+            if input_quantizer is None:
+                raise ValueError("Missing quantizer for input tensor")
+            input_quantizer.set_usage(
+                rowwise=True,
+                columnwise=(
+                    backward_needs_input
+                    and not save_original_input
+                    and args.backward_override is None
+                ),
+            )
+            own_quantized_input = True
+            inputmat_is_storage = True
+    else:
+        inputmat_aliases_inp = inp.dtype == activation_dtype
+
+    if save_original_input:
+        inputmat_aliases_inp = True
+        inputmat_is_storage = False
+
+    # ------------------------------------------------------
+    # Weight pipeline -- mirror ``quantize_weight`` / ``cast_if_needed``.
+    # ``new_weight_workspace`` is a fresh fake storage only on the
+    # cache-miss + ``cache_weight`` path, else ``None``.
+    # ------------------------------------------------------
+    new_weight_workspace = None
+    weightmat = None
+    weightmat_is_storage = False
+    weightmat_aliases_weight = False
+    if fp8_or_debug:
+        if weight_quantizer is not None and (not weight.is_quantized or debug):
+            columnwise_usage = is_grad_enabled and args.input_requires_grad and not args.is_fsdp2
+            if args.backward_override is not None:
+                columnwise_usage = False
+            if not columnwise_usage:
+                columnwise_usage = (
+                    is_fp8_activation_recompute_enabled()
+                    and not in_fp8_activation_recompute_phase()
+                )
+            weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)
+        elif weight.is_quantized:
+            weight_quantizer = weight.quantizer
+
+        if weight.is_quantized:
+            # Primary-quantized weight: the impl reuses it as ``weightmat``.
+            weightmat = weight
+            weightmat_is_storage = True
+            weightmat_aliases_weight = True
+        else:
+            weightmat = TensorProto(
+                shape=tuple(weight.shape),
+                dtype=activation_dtype,
+                quantizer=weight_quantizer,
+                device=weight.device,
+            )
+            weightmat_is_storage = True
+            update_ws = args.is_first_microbatch is None or args.is_first_microbatch
+            if args.cache_weight and update_ws and args.weight_workspace is None:
+                new_weight_workspace = TensorProto(
+                    shape=tuple(weight.shape),
+                    dtype=activation_dtype,
+                    quantizer=weight_quantizer,
+                    device=weight.device,
+                )
+    else:
+        weightmat_aliases_weight = weight.dtype == activation_dtype
+        weightmat = TensorProto(
+            shape=tuple(weight.shape), dtype=activation_dtype, device=weight.device
+        )
+
+    if output_quantizer is not None:
+        output_quantizer.set_usage(rowwise=True, columnwise=False)
+
+    # ------------------------------------------------------
+    # Output tensor: y = x @ w^T (quantized iff an output quantizer is set).
+    # ------------------------------------------------------
+    out_leading = inp.shape[0]
+    if args.parallel_mode == "column" and args.sequence_parallel:
+        out_leading = out_leading * args.tp_size
+    elif args.parallel_mode == "row" and args.sequence_parallel:
+        out_leading = out_leading // args.tp_size
+    out = TensorProto(
+        shape=(out_leading, *tuple(inp.shape[1:-1]), out_features),
+        dtype=activation_dtype,
+        quantizer=output_quantizer,
+        requires_grad=is_grad_enabled
+        and (args.input_requires_grad or args.weight_requires_grad),
+        device=inp.device,
+    )
+
+    # ------------------------------------------------------
+    # Backward state -- saved-tensor layout
+    # (saved_inputmat, wt_save, saved_weight, bias) with name-based aliasing.
+    # ------------------------------------------------------
+    tensors_to_save_from_forward = None
+    ctx_attrs = None
+    if is_grad_enabled:
+        # Slot 0 -- ``saved_inputmat``.
+        inputmat_alias = None
+        saved_inputmat = None
+        if backward_needs_input:
+            if inputmat_aliases_inp:
+                inputmat_alias = "inp"
+            elif inputmat_is_storage:
+                saved_inputmat = TensorProto(
+                    shape=tuple(inp.shape),
+                    dtype=activation_dtype,
+                    quantizer=input_quantizer,
+                    device=inp.device,
+                )
+                # Mirror ``_linear_forward_impl``'s post-quantization
+                # ``inputmat.update_usage(...)`` so the saved input's buffer layout
+                # matches -- driven by the same conditions as the real impl.
+                if own_quantized_input and not save_original_input:
+                    if args.backward_override is not None:
+                        saved_inputmat.update_usage(rowwise_usage=True, columnwise_usage=False)
+                    elif (
+                        args.backward_input_needs_gather
+                        and weight_quantizer is not None
+                        and weight_quantizer.supports_only_rowwise_all_gather()
+                    ):
+                        saved_inputmat.update_usage(rowwise_usage=True, columnwise_usage=False)
+                    else:
+                        saved_inputmat.update_usage(rowwise_usage=False, columnwise_usage=True)
+            else:
+                saved_inputmat = TensorProto(
+                    shape=tuple(inp.shape), dtype=activation_dtype, device=inp.device
+                )
+
+        # Slot 1 -- ``wt_save``.
+        wt_alias = None
+        wt_save = None
+        if weightmat_aliases_weight:
+            wt_alias = "weight"
+        elif args.is_fsdp2:
+            pass  # FSDP2 re-quantizes from the gathered weight in backward.
+        elif weightmat_is_storage:
+            wt_save = weightmat
+        else:
+            wt_save = TensorProto(
+                shape=tuple(weight.shape), dtype=activation_dtype, device=weight.device
+            )
+
+        # Slot 2 -- ``saved_weight`` (always aliased to ``weight``).
+        # Slot 3 -- ``bias`` (aliased to ``bias`` when present, else absent).
+        saved_tensor_aliases = (
+            inputmat_alias,
+            wt_alias,
+            "weight",
+            "bias" if bias is not None else None,
+        )
+        tensors_to_save_from_forward = (saved_inputmat, wt_save, None, None)
+        ctx_attrs = {
+            "fsdp_shapes": [],
+            "saved_tensor_aliases": saved_tensor_aliases,
+        }
+
+    return out, new_weight_workspace, tensors_to_save_from_forward, None, ctx_attrs
+
+
 def _linear_setup_ctx(
     bwd_args: LinearBwdArgs,
     fwd_args: LinearFwdArgs,
@@ -1311,6 +1518,68 @@ def wgrad_gemm(
     )
 
 
+def _linear_backward_impl_fake(
+    args: LinearBwdArgs,
+) -> Tuple[Optional[TensorProto], Optional[TensorProto], Optional[TensorProto]]:
+    """Allocation-free fake of :func:`_linear_backward` on ``TensorProto``.
+
+    The saved-tensor fields of ``args`` carry
+    :class:`~transformer_engine.pytorch.dynamo.TensorProto` instances. Returns
+    ``(wgrad, dgrad, grad_bias)`` protos describing the nature of the gradients,
+    mirroring the real backward's return contract without allocating storage.
+
+    Tensor-/sequence-parallel gather/scatter happens inside the eager backward
+    custom op and is opaque to ``torch.compile``: ``dgrad`` always carries the
+    rank-local input shape and ``wgrad`` the local weight shape, so no extra
+    shape modeling is needed here.
+    """
+    if args.fsdp_group is not None:
+        raise NotImplementedError(
+            "Fake Linear backward does not support manual TE FSDP "
+            "(fsdp_group is not None); use FSDP2 or MCore FSDP."
+        )
+
+    weight = args.saved_weight if args.saved_weight is not None else args.weight_fp8
+    out_dtype = args.activation_dtype
+    out_features, in_features = weight.shape
+
+    # Mirror ``_linear_backward``: ``set_usage`` on ``grad_input_quantizer``
+    # influences ``dgrad``'s buffer layout.
+    if args.grad_input_quantizer is not None:
+        args.grad_input_quantizer.set_usage(rowwise=True, columnwise=False)
+
+    dgrad = None
+    if args.requires_dgrad:
+        # dgrad has the logical input shape and may be quantized for the next op.
+        dgrad = TensorProto(
+            shape=tuple(args.inp_shape),
+            dtype=out_dtype,
+            quantizer=args.grad_input_quantizer,
+            device=args.grad_output.device,
+        )
+
+    wgrad = None
+    if args.requires_wgrad and not args.fuse_wgrad_accumulation:
+        # wgrad has the weight's shape; quantized iff an fp8 wgrad output is
+        # requested (mirrors ``quantization_params=grad_weight_quantizer``),
+        # otherwise high precision. Under fuse_wgrad_accumulation the grad is
+        # written into ``main_grad`` in place and no wgrad tensor is returned.
+        wgrad = TensorProto(
+            shape=(out_features, in_features),
+            dtype=out_dtype,
+            quantizer=args.grad_weight_quantizer,
+            device=weight.device,
+        )
+
+    grad_bias = None
+    if args.use_bias and args.requires_wgrad:
+        grad_bias = TensorProto(
+            shape=(out_features,), dtype=out_dtype, device=args.grad_output.device
+        )
+
+    return wgrad, dgrad, grad_bias
+
+
 class _Linear(torch.autograd.Function):
     """Linear semi-top level module
     Calls custom cuda extensions.
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index 033a35f8e1..ce358aaaf4 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -43,6 +43,10 @@ def _contains_process_group(value: Any) -> bool:
 _quantized_tensor_passthrough_ops: set = set()
 
 
+#: Maps storage / wrapper class qualname -> class object, for ``__tensor_unflatten__``.
+_STORAGE_REGISTRY: Dict[str, type] = {}
+
+
 class QuantizedTensorStorage:
     r"""Base class for all TensorStorage classes.
 
@@ -146,6 +150,64 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None:
             f"{self.__class__.__name__} class does not implement copy_from_storage function"
         )
 
+    # ----- PyTorch subclass flatten protocol (torch.compile / TensorProto) -----
+
+    # Subclasses declare their tensor buffers once, as ``(attribute_name,
+    # constructor_kwarg)`` pairs in flatten order; everything else returned by
+    # :meth:`get_metadata` is treated as non-tensor context.
+    _FLATTEN_TENSOR_BUFFERS: Tuple[Tuple[str, str], ...] = ()
+
+    def __init_subclass__(cls, **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
+        # Register every storage / wrapper class so ``__tensor_unflatten__`` can
+        # resolve the concrete class from its qualname inside an FX graph.
+        _STORAGE_REGISTRY[cls.__qualname__] = cls
+
+    def _flatten_nontensor_kwargs(self) -> Dict[str, Any]:
+        """Non-tensor constructor kwargs (scalars, dtype, quantizer)."""
+        tensor_kwargs = {kwarg for _, kwarg in self._FLATTEN_TENSOR_BUFFERS}
+        return {k: v for k, v in self.get_metadata().items() if k not in tensor_kwargs}
+
+    def __tensor_flatten__(self) -> Tuple[list, Dict[str, Any]]:
+        """Return ``(inner_tensor_attr_names, context)``; see class comment."""
+        present = [
+            attr for attr, _ in self._FLATTEN_TENSOR_BUFFERS if getattr(self, attr) is not None
+        ]
+        ctx = {
+            "cls": type(self).__qualname__,
+            "is_tensor": isinstance(self, QuantizedTensor),
+            "requires_grad": bool(self.requires_grad) if isinstance(self, QuantizedTensor) else False,
+            "nontensor_kwargs": self._flatten_nontensor_kwargs(),
+        }
+        return present, ctx
+
+    @staticmethod
+    def __tensor_unflatten__(
+        inner_tensors: Dict[str, torch.Tensor],
+        ctx: Dict[str, Any],
+        outer_size: Iterable[int],
+        outer_stride: Optional[Iterable[int]],
+    ) -> QuantizedTensorStorage:
+        """Rebuild a storage / wrapper from flat tensors + context."""
+        cls = _STORAGE_REGISTRY[ctx["cls"]]
+        kwargs: Dict[str, Any] = dict(ctx["nontensor_kwargs"])
+        # Map each declared buffer back to its constructor kwarg (absent -> None).
+        for attr, kwarg in cls._FLATTEN_TENSOR_BUFFERS:
+            kwargs[kwarg] = inner_tensors.get(attr)
+        if not ctx["is_tensor"]:
+            return cls(**kwargs)
+        # Wrapper subclass: it also needs outer shape / dtype / device / stride.
+        fake_dtype = kwargs.get("fake_dtype")
+        device = next((t.device for t in inner_tensors.values() if t is not None), None)
+        return cls(
+            shape=tuple(outer_size),
+            dtype=fake_dtype,
+            requires_grad=ctx["requires_grad"],
+            device=device,
+            stride=tuple(outer_stride) if outer_stride is not None else None,
+            **kwargs,
+        )
+
 
 def prepare_for_saving(
     *tensors: Union[torch.Tensor, QuantizedTensorStorage],
@@ -363,6 +425,71 @@ def make_empty(
             result.requires_grad_(True)
         return result
 
+    # ----- Data-free buffer/metadata primitives backing TensorProto -----
+
+    def _describe_buffers(
+        self, shape: Tuple[int, ...]
+    ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]:
+        """Return ``{attr_name: (buffer_shape, buffer_dtype)}`` for the buffers
+        this quantizer would allocate for a logical tensor of ``shape``.
+
+        Keys must match the buffer attribute names declared in the storage's
+        ``_FLATTEN_TENSOR_BUFFERS`` and respect the quantizer's usage flags.
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not implement _describe_buffers; "
+            "it cannot be used with TensorProto / pure-Python allocation"
+        )
+
+    def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]:
+        """Non-tensor context for the produced storage.
+
+        Returns ``{"cls": <type>, "nontensor_kwargs": {...}}`` where ``cls`` is
+        the concrete class to instantiate (wrapper subclass for user-visible
+        tensors, bare storage class for ``internal`` quantizers) and
+        ``nontensor_kwargs`` are its non-tensor constructor kwargs (e.g.
+        ``fp8_dtype``, ``quantizer``, ``fake_dtype``).
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not implement _storage_metadata; "
+            "it cannot be used with TensorProto / pure-Python allocation"
+        )
+
+    def alloc_tensors(
+        self,
+        shape: Iterable[int],
+        *,
+        device: Optional[Union[torch.device, str]] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """Allocate (uninitialized) the flat buffers for ``shape``.
+
+        Returns ``{attr_name: torch.Tensor}`` suitable as the ``inner_tensors``
+        argument of the storage's ``__tensor_unflatten__``.
+        """
+        device = torch.device(device if device is not None else "cuda")
+        return {
+            attr: torch.empty(buf_shape, dtype=buf_dtype, device=device)
+            for attr, (buf_shape, buf_dtype) in self._describe_buffers(tuple(shape)).items()
+        }
+
+    def create_metadata(
+        self,
+        _shape: Iterable[int],
+        *,
+        dtype: torch.dtype,
+        requires_grad: bool = False,
+    ) -> Dict[str, Any]:
+        """Build the data-free ``__tensor_unflatten__`` context describing the
+        quantized tensor this quantizer would produce for ``shape`` / ``dtype``.
+        """
+        meta = self._storage_metadata(dtype)
+        return {
+            "cls": meta["cls"].__qualname__,
+            "is_tensor": not self.internal,
+            "requires_grad": requires_grad,
+            "nontensor_kwargs": meta["nontensor_kwargs"],
+        }
+
     def calibrate(self, tensor: torch.Tensor) -> None:
         """Calibrate quantizer state
 
diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
index 92c22acd0b..c816b4fb04 100644
--- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
@@ -7,7 +7,7 @@
 from collections.abc import Iterable
 import math
 import warnings
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 import torch
 import transformer_engine_torch as tex
@@ -73,6 +73,39 @@ def copy(self) -> Float8BlockQuantizer:
     def _value_fields(self) -> Tuple[str, ...]:
         return ("dtype", "block_len", "amax_epsilon", "force_pow_2_scales", "block_scaling_dim")
 
+    # ----- TensorProto / pure-Python allocation -----
+
+    def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]:
+        return {
+            "cls": Float8BlockwiseQTensorStorage if self.internal else Float8BlockwiseQTensor,
+            "nontensor_kwargs": {
+                "fp8_dtype": self.dtype,
+                "quantizer": self,
+                "is_2D_scaled": self.block_scaling_dim == 2,
+                "fake_dtype": fake_dtype,
+            },
+        }
+
+    def _describe_buffers(
+        self, shape: Tuple[int, ...]
+    ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]:
+        shape = tuple(shape)
+        buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {}
+        # Blockwise FP8 scales are FP32; columnwise data is stored transposed.
+        if self.rowwise_usage:
+            buffers["_rowwise_data"] = (shape, torch.uint8)
+            buffers["_rowwise_scale_inv"] = (
+                tuple(self.get_scale_shape(shape, columnwise=False)),
+                torch.float32,
+            )
+        if self.columnwise_usage:
+            buffers["_columnwise_data"] = (tuple(self.get_columnwise_shape(shape)), torch.uint8)
+            buffers["_columnwise_scale_inv"] = (
+                tuple(self.get_scale_shape(shape, columnwise=True)),
+                torch.float32,
+            )
+        return buffers
+
     def update_quantized(
         self,
         src: torch.Tensor,
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index 0310c3855c..6d3a53b3d7 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -4,7 +4,7 @@
 
 """Tensor class with FP8 data"""
 from __future__ import annotations
-from typing import Any, Optional, Tuple, Iterable, Union
+from typing import Any, Dict, Optional, Tuple, Iterable, Union
 import warnings
 import torch
 from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState
@@ -15,7 +15,7 @@
     Float8CurrentScaling,
     Recipe,
 )
-from ..utils import canonicalize_process_group, devices_match
+from ..utils import canonicalize_process_group, devices_match, is_non_tn_fp8_gemm_supported
 from .storage.float8_tensor_storage import Float8TensorStorage, _FromFloat8Func
 from ..quantized_tensor import QuantizedTensor, Quantizer
 from ..dynamo import register_value_opaque_quantizer
@@ -393,6 +393,36 @@ def _value_fields(self) -> Tuple[str, ...]:
         # raises so it can never be baked into a torch.compile graph.
         return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction")
 
+    # ----- TensorProto / pure-Python allocation -----
+
+    def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]:
+        return {
+            "cls": Float8TensorStorage if self.internal else Float8Tensor,
+            "nontensor_kwargs": {
+                "fp8_dtype": self.dtype,
+                "quantizer": self,
+                "fake_dtype": fake_dtype,
+            },
+        }
+
+    def _describe_buffers(
+        self, shape: Tuple[int, ...]
+    ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]:
+        shape = tuple(shape)
+        buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {}
+        # Mirror the C++ quantizer allocation (csrc/quantizer.cpp): on non-TN-capable
+        # archs (Blackwell+) a single ``_data`` buffer backs both row- and column-wise
+        # usage and no separate transpose is materialized. This must match what the
+        # real kernel produces so the torch.compile fake layout lines up slot-for-slot.
+        non_tn = is_non_tn_fp8_gemm_supported()
+        if self.rowwise_usage or non_tn:
+            buffers["_data"] = (shape, torch.uint8)
+        if self.columnwise_usage and not non_tn:
+            buffers["_transpose"] = ((shape[-1], *shape[:-1]), torch.uint8)
+        # Per-tensor scale-inv is always present for current scaling.
+        buffers["_scale_inv"] = ((1,), torch.float32)
+        return buffers
+
 
 register_value_opaque_quantizer(Float8CurrentScalingQuantizer)
 
diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
index a3746b3088..3e41bf7bb1 100644
--- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py
+++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
@@ -6,7 +6,7 @@
 from __future__ import annotations
 from collections.abc import Iterable
 import math
-from typing import Optional, Tuple, Union, Any
+from typing import Optional, Tuple, Union, Any, Dict
 import warnings
 
 import torch
@@ -61,6 +61,40 @@ def copy(self) -> MXFP8Quantizer:
     def _value_fields(self) -> Tuple[str, ...]:
         return ("dtype",)
 
+    # ----- TensorProto / pure-Python allocation -----
+
+    def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]:
+        return {
+            "cls": MXFP8TensorStorage if self.internal else MXFP8Tensor,
+            "nontensor_kwargs": {
+                "fp8_dtype": self.dtype,
+                "quantizer": self,
+                "with_gemm_swizzled_scales": self.optimize_for_gemm,
+                "fake_dtype": fake_dtype,
+            },
+        }
+
+    def _describe_buffers(
+        self, shape: Tuple[int, ...]
+    ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]:
+        shape = tuple(shape)
+        buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {}
+        # MXFP8 block scales are stored as uint8 (E8M0); data buffers keep the
+        # logical shape (rowwise) and its transpose (columnwise).
+        if self.rowwise_usage:
+            buffers["_rowwise_data"] = (shape, torch.uint8)
+            buffers["_rowwise_scale_inv"] = (
+                tuple(self.get_scale_shape(shape, columnwise=False)),
+                torch.uint8,
+            )
+        if self.columnwise_usage:
+            buffers["_columnwise_data"] = (shape, torch.uint8)
+            buffers["_columnwise_scale_inv"] = (
+                tuple(self.get_scale_shape(shape, columnwise=True)),
+                torch.uint8,
+            )
+        return buffers
+
     def update_quantized(
         self,
         src: torch.Tensor,
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index c8c0a7b854..8827b61467 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -7,7 +7,7 @@
 from collections.abc import Iterable
 import math
 import warnings
-from typing import Dict, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Tuple, Union
 import functools
 
 import torch
@@ -366,6 +366,49 @@ def _value_fields(self) -> Tuple[str, ...]:
             "with_amax_reduction",
         )
 
+    # ----- TensorProto / pure-Python allocation -----
+
+    def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]:
+        return {
+            "cls": NVFP4TensorStorage if self.internal else NVFP4Tensor,
+            "nontensor_kwargs": {
+                "fp4_dtype": self.dtype,
+                "quantizer": self,
+                "with_gemm_swizzled_scales": self.optimize_for_gemm,
+                "row_scaled_nvfp4": self.row_scaled_nvfp4,
+                "nvfp4_use_4over6": self.nvfp4_use_4over6,
+                "nvfp4_e4m3_max": self.nvfp4_e4m3_max,
+                "fake_dtype": fake_dtype,
+            },
+        }
+
+    def _describe_buffers(
+        self, shape: Tuple[int, ...]
+    ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]:
+        shape = tuple(shape)
+        buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {}
+        # FP4 data packs 2 values per byte (uint8); block scales are E4M3 stored
+        # as uint8; amax buffers are FP32 (per-row when row-scaled, else scalar).
+        if self.rowwise_usage:
+            buffers["_rowwise_data"] = (self.convert_shape_for_fp4(shape), torch.uint8)
+            buffers["_rowwise_scale_inv"] = (
+                tuple(self.get_scale_shape(shape, columnwise=False)),
+                torch.uint8,
+            )
+            amax_rowwise_shape = (math.prod(shape[:-1]),) if self.row_scaled_nvfp4 else (1,)
+            buffers["_amax_rowwise"] = (amax_rowwise_shape, torch.float32)
+        if self.columnwise_usage:
+            buffers["_columnwise_data"] = (
+                self.convert_shape_for_fp4(self.get_columnwise_shape(shape)),
+                torch.uint8,
+            )
+            buffers["_columnwise_scale_inv"] = (
+                tuple(self.get_scale_shape(shape, columnwise=True)),
+                torch.uint8,
+            )
+            buffers["_amax_columnwise"] = ((1,), torch.float32)
+        return buffers
+
 
 register_value_opaque_quantizer(NVFP4Quantizer)
 
diff --git a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py
index f7a3dae70b..ec2c40ef9a 100644
--- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py
+++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py
@@ -35,6 +35,15 @@ class Float8BlockwiseQTensorStorage(QuantizedTensorStorage):
     _columnwise_scale_inv: Optional[torch.Tensor]
     _is_2D_scaled: bool
 
+    # (attribute_name, constructor_kwarg) for each tensor buffer; drives
+    # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage).
+    _FLATTEN_TENSOR_BUFFERS = (
+        ("_rowwise_data", "rowwise_data"),
+        ("_rowwise_scale_inv", "rowwise_scale_inv"),
+        ("_columnwise_data", "columnwise_data"),
+        ("_columnwise_scale_inv", "columnwise_scale_inv"),
+    )
+
     def __new__(
         cls,
         rowwise_data: Optional[torch.Tensor],
diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py
index a97162f91c..1419a02559 100644
--- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py
+++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py
@@ -75,6 +75,14 @@ class Float8TensorStorage(QuantizedTensorStorage):
     _transpose: Optional[torch.Tensor]
     _transpose_invalid: bool
 
+    # (attribute_name, constructor_kwarg) for each tensor buffer; drives
+    # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage).
+    _FLATTEN_TENSOR_BUFFERS = (
+        ("_data", "data"),
+        ("_transpose", "data_transpose"),
+        ("_scale_inv", "fp8_scale_inv"),
+    )
+
     def __new__(
         cls,
         *args,
diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py
index ea592cd989..a2a3bf2f4c 100644
--- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py
+++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py
@@ -82,6 +82,15 @@ class MXFP8TensorStorage(QuantizedTensorStorage):
     # GEMM
     _with_gemm_swizzled_scales: bool
 
+    # (attribute_name, constructor_kwarg) for each tensor buffer; drives
+    # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage).
+    _FLATTEN_TENSOR_BUFFERS = (
+        ("_rowwise_data", "rowwise_data"),
+        ("_rowwise_scale_inv", "rowwise_scale_inv"),
+        ("_columnwise_data", "columnwise_data"),
+        ("_columnwise_scale_inv", "columnwise_scale_inv"),
+    )
+
     def __new__(
         cls,
         rowwise_data: Optional[torch.Tensor],
diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py
index 53bb5e7c11..5ed6d1d641 100644
--- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py
+++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py
@@ -108,6 +108,17 @@ class NVFP4TensorStorage(QuantizedTensorStorage):
     # Global E4M3 scale bound used by this NVFP4 tensor
     _nvfp4_e4m3_max: int
 
+    # (attribute_name, constructor_kwarg) for each tensor buffer; drives
+    # __tensor_flatten__ / __tensor_unflatten__ (see QuantizedTensorStorage).
+    _FLATTEN_TENSOR_BUFFERS = (
+        ("_rowwise_data", "rowwise_data"),
+        ("_rowwise_scale_inv", "rowwise_scale_inv"),
+        ("_columnwise_data", "columnwise_data"),
+        ("_columnwise_scale_inv", "columnwise_scale_inv"),
+        ("_amax_rowwise", "amax_rowwise"),
+        ("_amax_columnwise", "amax_columnwise"),
+    )
+
     def __new__(
         cls,
         rowwise_data: Optional[torch.Tensor],

From ea3df7a22dfb9c803dfe315d40fb22c99af76b89 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 22 Jun 2026 12:45:34 +0200
Subject: [PATCH 16/28] [PyTorch] torch.compile: dedup cached FP8 weight from
 saved-for-backward

The cached FP8 weight is the same tensor returned as new_weight_workspace (cache miss) or passed in as weight_workspace (cache hit). A custom op may not return a tensor that aliases an input or another return, so mark those slots and reconstruct wt_save in _linear_setup_ctx instead of saving it twice. Mirrored in the fake impl so the saved-slot layout matches.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/module/linear.py | 40 ++++++++++++++++-----
 1 file changed, 32 insertions(+), 8 deletions(-)

diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index d479d2f67b..2b2ee05a3e 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -607,12 +607,24 @@ def _linear_forward_impl(
         if is_fsdp2 and weightmat is not weight:
             wt_save = None
 
-        # Dedup save slots that alias forward inputs; ``_linear_setup_ctx``
-        # rebuilds the refs from ``inp`` / ``weight`` / ``bias``.
-        # Needed for torch.compile to work correctly.
+        # Dedup save slots that alias forward inputs or other op returns;
+        # ``_linear_setup_ctx`` rebuilds the refs. Needed because a custom op may
+        # not return a tensor that aliases an input or another return: the cached
+        # FP8 weight is the same tensor as ``new_weight_workspace`` (a return, on a
+        # cache miss) or ``weight_workspace`` (an input, on a cache hit).
+        if wt_save is None:
+            wt_alias = None
+        elif wt_save is weight:
+            wt_alias = "weight"
+        elif new_weight_workspace is not None and wt_save is new_weight_workspace:
+            wt_alias = "new_workspace"
+        elif args.weight_workspace is not None and wt_save is args.weight_workspace:
+            wt_alias = "weight_workspace"
+        else:
+            wt_alias = None
         saved_tensor_aliases = (
             "inp" if saved_inputmat is inp else None,
-            "weight" if wt_save is weight else None,
+            wt_alias,
             "weight",  # ``saved_weight`` slot is always the weight parameter
             "bias" if bias is not None else None,
         )
@@ -798,13 +810,20 @@ def _linear_forward_impl_fake(
                     shape=tuple(inp.shape), dtype=activation_dtype, device=inp.device
                 )
 
-        # Slot 1 -- ``wt_save``.
+        # Slot 1 -- ``wt_save``. Mirror the real impl's alias dedup: the cached
+        # FP8 weight is shared with ``new_weight_workspace`` (a return, on a cache
+        # miss) or the ``weight_workspace`` input (on a cache hit), so it is
+        # reconstructed in ``_linear_setup_ctx`` rather than saved twice.
         wt_alias = None
         wt_save = None
         if weightmat_aliases_weight:
             wt_alias = "weight"
         elif args.is_fsdp2:
             pass  # FSDP2 re-quantizes from the gathered weight in backward.
+        elif weightmat_is_storage and new_weight_workspace is not None:
+            wt_alias = "new_workspace"
+        elif weightmat_is_storage and args.weight_workspace is not None:
+            wt_alias = "weight_workspace"
         elif weightmat_is_storage:
             wt_save = weightmat
         else:
@@ -832,7 +851,7 @@ def _linear_forward_impl_fake(
 def _linear_setup_ctx(
     bwd_args: LinearBwdArgs,
     fwd_args: LinearFwdArgs,
-    out: torch.Tensor,
+    fwd_outputs: Tuple[Any, ...],
     ctx_attrs: Dict,
     tensors_to_save_from_forward: Tuple[Any, ...],
 ) -> Tuple[Any, ...]:
@@ -845,7 +864,8 @@ def _linear_setup_ctx(
     for FSDP2 re-quantization) without having to mutate the structured
     metadata returned by ``prepare_for_saving``.
     """
-    del out  # No-op; kept for symmetry with the compile-time helper signature.
+    # ``fwd_outputs`` are the op's user outputs ``(out, new_weight_workspace)``;
+    # only ``new_weight_workspace`` is needed here, to rebuild the deduped weight.
 
     inp = fwd_args.inp
     weight = fwd_args.weight
@@ -935,6 +955,10 @@ def _linear_setup_ctx(
         saved_inputmat = inp
     if wt_save_alias == "weight":
         wt_save = weight
+    elif wt_save_alias == "new_workspace":
+        wt_save = fwd_outputs[1]
+    elif wt_save_alias == "weight_workspace":
+        wt_save = fwd_args.weight_workspace
     if saved_weight_alias == "weight":
         saved_weight = weight
     if bias_alias == "bias":
@@ -1619,7 +1643,7 @@ def forward(
             tensors_to_save_from_setup = _linear_setup_ctx(
                 bwd_args,
                 fwd_args,
-                out,
+                (out, new_weight_workspace),
                 ctx_attrs,
                 tensors_to_save_from_forward,
             )

From 4997929645510e3157e9cbe7fe5b260d438273a5 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 22 Jun 2026 13:17:19 +0200
Subject: [PATCH 17/28] [PyTorch] nvfp4: emit _describe_buffers in canonical
 flatten order

NVFP4Quantizer._describe_buffers grouped each amax right after its scale (per-usage), diverging from NVFP4TensorStorage._FLATTEN_TENSOR_BUFFERS (amax buffers last). The order is functionally irrelevant (buffers are consumed by name in alloc_tensors and reordered in TensorProto.inner_names), but aligning it makes describe/flatten agree and fixes test_to_tensor_proto_quantized[nvfp4].

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/tensor/nvfp4_tensor.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 8827b61467..0cc436c399 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -389,14 +389,14 @@ def _describe_buffers(
         buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {}
         # FP4 data packs 2 values per byte (uint8); block scales are E4M3 stored
         # as uint8; amax buffers are FP32 (per-row when row-scaled, else scalar).
+        # Order matches NVFP4TensorStorage._FLATTEN_TENSOR_BUFFERS (the canonical
+        # __tensor_flatten__ order): data + scale_inv per usage first, amax last.
         if self.rowwise_usage:
             buffers["_rowwise_data"] = (self.convert_shape_for_fp4(shape), torch.uint8)
             buffers["_rowwise_scale_inv"] = (
                 tuple(self.get_scale_shape(shape, columnwise=False)),
                 torch.uint8,
             )
-            amax_rowwise_shape = (math.prod(shape[:-1]),) if self.row_scaled_nvfp4 else (1,)
-            buffers["_amax_rowwise"] = (amax_rowwise_shape, torch.float32)
         if self.columnwise_usage:
             buffers["_columnwise_data"] = (
                 self.convert_shape_for_fp4(self.get_columnwise_shape(shape)),
@@ -406,6 +406,10 @@ def _describe_buffers(
                 tuple(self.get_scale_shape(shape, columnwise=True)),
                 torch.uint8,
             )
+        if self.rowwise_usage:
+            amax_rowwise_shape = (math.prod(shape[:-1]),) if self.row_scaled_nvfp4 else (1,)
+            buffers["_amax_rowwise"] = (amax_rowwise_shape, torch.float32)
+        if self.columnwise_usage:
             buffers["_amax_columnwise"] = ((1,), torch.float32)
         return buffers
 

From 50c11cd415466267893f351e276ff8a432a63320 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 29 Jun 2026 08:00:58 +0200
Subject: [PATCH 18/28] Address review: error on undescribed buffers, gate
 nvfp4 test on HW support

- TensorProto.inner_names now raises if the quantizer describes buffer(s) absent
  from the storage's _FLATTEN_TENSOR_BUFFERS, instead of silently appending them.
- Gate the nvfp4 proto-quantizer param on nvfp4_available so it skips on hardware
  without NVFP4 support rather than failing.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 tests/pytorch/test_torch_compile.py               | 14 +++++++++-----
 transformer_engine/pytorch/dynamo/tensor_proto.py | 11 ++++++++---
 transformer_engine/pytorch/module/linear.py       |  3 +--
 transformer_engine/pytorch/quantized_tensor.py    |  4 +++-
 transformer_engine/pytorch/tensor/mxfp8_tensor.py |  2 --
 5 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 4a3f22a291..0c8f34e5e2 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -31,7 +31,6 @@
 from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
 from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, _STORAGE_REGISTRY
 from transformer_engine.pytorch.dynamo import TensorProto, to_tensor_proto
-from transformer_engine.pytorch.dynamo.tensor_proto import _contiguous_stride
 from transformer_engine.pytorch import (
     is_fp8_available,
     is_mxfp8_available,
@@ -585,8 +584,8 @@ def fn(inp):
         (64, 128),
         id="nvfp4",
         marks=pytest.mark.skipif(
-            not torch.cuda.is_available(),
-            reason="NVFP4Quantizer requires CUDA to construct",
+            not nvfp4_available,
+            reason="NVFP4 is not available",
         ),
     ),
 ]
@@ -603,7 +602,10 @@ def _build_from_primitives(quantizer, shape, dtype, device="cpu"):
     buffers = quantizer.alloc_tensors(shape, device=device)
     inner = {name: buffers[name] for name in names}
     storage_cls = _STORAGE_REGISTRY[ctx["cls"]]
-    return storage_cls.__tensor_unflatten__(inner, ctx, tuple(shape), _contiguous_stride(shape))
+    # Row-major (contiguous) outer stride for ``__tensor_unflatten__``; ``meta``
+    # device computes it without allocating storage.
+    outer_stride = torch.empty(tuple(shape), device="meta").stride()
+    return storage_cls.__tensor_unflatten__(inner, ctx, tuple(shape), outer_stride)
 
 
 def _signature(tensor, names):
@@ -807,7 +809,9 @@ def test_to_tensor_proto_quantized(factory, shape):
     assert proto.shape == tuple(shape)
     assert proto.dtype == torch.bfloat16
     # Same buffer layout as the original tensor.
-    assert proto.inner_names() == tuple(q._describe_buffers(shape))  # pylint: disable=protected-access
+    assert proto.inner_names() == tuple(
+        q._describe_buffers(shape)
+    )  # pylint: disable=protected-access
     # Rebuilding from the derived proto matches the original tensor's structure.
     assert _signature(proto.create_tensor(), proto.inner_names()) == _signature(
         tensor, proto.inner_names()
diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py
index b5248e3ee4..911b4151bf 100644
--- a/transformer_engine/pytorch/dynamo/tensor_proto.py
+++ b/transformer_engine/pytorch/dynamo/tensor_proto.py
@@ -83,9 +83,14 @@ def inner_names(self) -> Tuple[str, ...]:
         described = list(self.quantizer._describe_buffers(tuple(self.shape)).keys())
         storage_cls = self.quantizer._storage_metadata(self.dtype)["cls"]
         flatten_order = [attr for attr, _ in storage_cls._FLATTEN_TENSOR_BUFFERS]
-        ordered = [name for name in flatten_order if name in described]
-        ordered += [name for name in described if name not in flatten_order]
-        return tuple(ordered)
+        extra = [name for name in described if name not in flatten_order]
+        if extra:
+            raise RuntimeError(
+                f"{storage_cls.__name__} describes buffer(s) {extra} absent from its "
+                f"_FLATTEN_TENSOR_BUFFERS {flatten_order}; the fake layout cannot be "
+                "aligned with the real one slot-for-slot."
+            )
+        return tuple(name for name in flatten_order if name in described)
 
     def create_metadata(self) -> Dict[str, Any]:
         """Data-free ``__tensor_unflatten__`` context describing this tensor."""
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 2b2ee05a3e..473483faff 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -766,8 +766,7 @@ def _linear_forward_impl_fake(
         shape=(out_leading, *tuple(inp.shape[1:-1]), out_features),
         dtype=activation_dtype,
         quantizer=output_quantizer,
-        requires_grad=is_grad_enabled
-        and (args.input_requires_grad or args.weight_requires_grad),
+        requires_grad=is_grad_enabled and (args.input_requires_grad or args.weight_requires_grad),
         device=inp.device,
     )
 
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index ce358aaaf4..86c4c94f7c 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -176,7 +176,9 @@ def __tensor_flatten__(self) -> Tuple[list, Dict[str, Any]]:
         ctx = {
             "cls": type(self).__qualname__,
             "is_tensor": isinstance(self, QuantizedTensor),
-            "requires_grad": bool(self.requires_grad) if isinstance(self, QuantizedTensor) else False,
+            "requires_grad": (
+                bool(self.requires_grad) if isinstance(self, QuantizedTensor) else False
+            ),
             "nontensor_kwargs": self._flatten_nontensor_kwargs(),
         }
         return present, ctx
diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
index 3e41bf7bb1..f804a96f24 100644
--- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py
+++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
@@ -79,8 +79,6 @@ def _describe_buffers(
     ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]:
         shape = tuple(shape)
         buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {}
-        # MXFP8 block scales are stored as uint8 (E8M0); data buffers keep the
-        # logical shape (rowwise) and its transpose (columnwise).
         if self.rowwise_usage:
             buffers["_rowwise_data"] = (shape, torch.uint8)
             buffers["_rowwise_scale_inv"] = (

From ff48e52a0d794acc17717575329f142e4cb09272 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 29 Jun 2026 16:25:44 +0200
Subject: [PATCH 19/28] [PyTorch] Workaround torch.compile staticmethod guard
 bug in NVFP4 _describe_buffers

Access NVFP4Quantizer @staticmethods (convert_shape_for_fp4, get_columnwise_shape)
via the class instead of the instance. Under torch.compile, instance access of a
@staticmethod on a value-opaque object crashes Dynamo guard generation with
"'function' object has no attribute '__func__'" (pytorch/pytorch#182741).
Temporary workaround until the PyTorch-side fix lands.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/tensor/nvfp4_tensor.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 0cc436c399..8d51480f51 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -391,15 +391,17 @@ def _describe_buffers(
         # as uint8; amax buffers are FP32 (per-row when row-scaled, else scalar).
         # Order matches NVFP4TensorStorage._FLATTEN_TENSOR_BUFFERS (the canonical
         # __tensor_flatten__ order): data + scale_inv per usage first, amax last.
+        # Workaround: call @staticmethods via the class, not the instance --
+        # instance access breaks torch.compile guard generation (pytorch #182741).
         if self.rowwise_usage:
-            buffers["_rowwise_data"] = (self.convert_shape_for_fp4(shape), torch.uint8)
+            buffers["_rowwise_data"] = (type(self).convert_shape_for_fp4(shape), torch.uint8)
             buffers["_rowwise_scale_inv"] = (
                 tuple(self.get_scale_shape(shape, columnwise=False)),
                 torch.uint8,
             )
         if self.columnwise_usage:
             buffers["_columnwise_data"] = (
-                self.convert_shape_for_fp4(self.get_columnwise_shape(shape)),
+                type(self).convert_shape_for_fp4(type(self).get_columnwise_shape(shape)),
                 torch.uint8,
             )
             buffers["_columnwise_scale_inv"] = (

From e1e271cc5f51835a6e28b10c4cbb96efefd74da9 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 15 Jun 2026 13:28:53 +0200
Subject: [PATCH 20/28] [PyTorch] torch.compile: wrap pybind11 UB methods as
 compile-time constants; fix SP memory leak; test suite hook-up
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Wrap CommOverlapCore pybind11 methods that return compile-time constants
so torch.compile(fullgraph=True) can trace through them without graph
breaks:
- `is_fp8_ubuf()` → `ub_is_fp8()` / `get_ub_is_fp8()` in base.py;
  `_ub_is_fp8()` in gemm.py
- `with_cublasmp()` → `ub_is_cublasmp()` in base.py

All callers in linear.py, layernorm_linear.py, layernorm_mlp.py,
base.py, gemm.py, userbuffers_backward_linear.py and
userbuffers_forward_linear.py updated.

Fix quantized grad_output not being freed early for column-parallel SP
backward. Row-parallel SP already called clear_tensor_data(grad_output)
to release the gathered tensor; column-parallel SP quantizes grad_output
to Float8TensorStorage but never freed it before returning.  Under
torch.compile reduce-overhead this leaves 3 live pool tensors at
recording end and triggers "Detected 3 tensor(s) in the cudagraph pool
not tracked as outputs".  Extend the existing clear_tensor_data guard to
cover both parallel modes.

Fix custom-recipe quantizer state being re-initialised on every forward
call even when the recipe object has not changed. The existing early-exit
for CustomRecipeState was missing an identity check on the recipe object,
so any repeated call with the same recipe would bypass the early-return
and rebuild quantizers unnecessarily.  Add `if recipe_state.recipe is
recipe: return` to restore the intended caching behaviour.

Add test_torch_compile.py to L0_pytorch_unittest so the autocast and
existing compile tests run in CI.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
(cherry picked from commit bfce3a7d02a0c87e0c3472bd40f2fadf68a0e6a4)
---
 transformer_engine/pytorch/module/layernorm_linear.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 43799b003c..1214edb937 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -424,7 +424,11 @@ def forward(
         if ub_overlap_rs_fprop:
             # cuBLASMp writes the reduce-scattered output directly into the
             # GEMM output tensor; Userbuffers writes it into the extra-output buffer.
-            out = gemm_out if ub_obj is not None and ub_obj.with_cublasmp() else reduce_scatter_out
+            out = (
+                gemm_out
+                if ub_obj is not None and ub_obj.with_cublasmp()
+                else reduce_scatter_out
+            )
         elif parallel_mode == "row" and tp_size > 1:
             nvtx_range_push(f"{nvtx_label}.row_parallel_comm")
             out = gemm_out

From 598a07cdc546fd6c9a90c43d640bfc696eef726f Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Mon, 15 Jun 2026 14:42:37 +0000
Subject: [PATCH 21/28] [pre-commit.ci] auto fixes from pre-commit.com hooks

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

(cherry picked from commit afe364bb6ef8004a03129b9e39047df6871317a6)
---
 transformer_engine/pytorch/module/layernorm_linear.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py
index 1214edb937..43799b003c 100644
--- a/transformer_engine/pytorch/module/layernorm_linear.py
+++ b/transformer_engine/pytorch/module/layernorm_linear.py
@@ -424,11 +424,7 @@ def forward(
         if ub_overlap_rs_fprop:
             # cuBLASMp writes the reduce-scattered output directly into the
             # GEMM output tensor; Userbuffers writes it into the extra-output buffer.
-            out = (
-                gemm_out
-                if ub_obj is not None and ub_obj.with_cublasmp()
-                else reduce_scatter_out
-            )
+            out = gemm_out if ub_obj is not None and ub_obj.with_cublasmp() else reduce_scatter_out
         elif parallel_mode == "row" and tp_size > 1:
             nvtx_range_push(f"{nvtx_label}.row_parallel_comm")
             out = gemm_out

From 05af2a0a15e0fc694f30b84f0d637771fe3b1a25 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 16 Jun 2026 16:32:11 +0200
Subject: [PATCH 22/28] Provide explicit QuantizerRoles in torch.compile
 custom-recipe test

ToyLinear now overrides get_quantizer_roles so CustomRecipeState doesn't hit
the no-roles warning, which graph-breaks under fullgraph=True. qfactory
dispatches on role.tensor_type instead of a pre-baked string key.

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
(cherry picked from commit 22f80e40846365bd46bb97a5febb10ca02889136)
---
 tests/pytorch/test_torch_compile.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 0c8f34e5e2..f88bf9ae2c 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -31,6 +31,7 @@
 from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
 from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, _STORAGE_REGISTRY
 from transformer_engine.pytorch.dynamo import TensorProto, to_tensor_proto
+from transformer_engine.pytorch.quantization import QuantizerRole
 from transformer_engine.pytorch import (
     is_fp8_available,
     is_mxfp8_available,

From 9bd16fd2c69b3ae0c0fa2f5f2adcc056a120d82a Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 16 Jun 2026 17:31:37 +0200
Subject: [PATCH 23/28] Add torch.compile custom-op path for Linear

Squashed PR #9 (linear_compile) onto the rebased base.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
(cherry picked from commit 84dbc6b2b704a1ca82dd0da37361e62228d9935a)
---
 .../distributed/run_layer_with_overlap.py     |   28 +
 tests/pytorch/distributed/run_numerics.py     |   51 +-
 .../distributed/test_comm_gemm_overlap.py     |   40 +
 tests/pytorch/test_torch_compile.py           |  343 ++++-
 transformer_engine/pytorch/dynamo/__init__.py |    2 +
 .../pytorch/dynamo/custom_op.py               | 1322 +++++++++++++++++
 transformer_engine/pytorch/module/linear.py   |  203 ++-
 transformer_engine/pytorch/utils.py           |   14 +
 8 files changed, 1968 insertions(+), 35 deletions(-)
 create mode 100644 transformer_engine/pytorch/dynamo/custom_op.py

diff --git a/tests/pytorch/distributed/run_layer_with_overlap.py b/tests/pytorch/distributed/run_layer_with_overlap.py
index 46795415e5..e65824ce85 100644
--- a/tests/pytorch/distributed/run_layer_with_overlap.py
+++ b/tests/pytorch/distributed/run_layer_with_overlap.py
@@ -200,6 +200,19 @@ def _parse_args(argv=None, namespace=None):
     parser.add_argument(
         "--use-cuda-graphs", action="store_true", default=False, help="Use CUDA Graphs."
     )
+    parser.add_argument(
+        "--compile",
+        action="store_true",
+        default=False,
+        help="Wrap each layer in torch.compile (tests Userbuffers on the compiled path).",
+    )
+    parser.add_argument(
+        "--compile-mode",
+        type=str,
+        default="default",
+        choices=["default", "reduce-overhead"],
+        help="torch.compile mode used when --compile is set.",
+    )
     parser.add_argument(
         "--ub-cfg", type=str, default=None, help="Optional TP config yaml file input."
     )
@@ -485,6 +498,9 @@ def dist_print(msg, src=None, end="\n", debug=False, error=False):
         torch.testing.assert_close(test_param, ref_param, rtol=0.0, atol=0.0)
     dist_print("Copied parameters from test model to reference model...", debug=True)
 
+    if opts.compile and opts.use_cuda_graphs:
+        raise ValueError("--compile and --use-cuda-graphs are mutually exclusive.")
+
     # Fp8 recipe setup
     fp8_format = Format.HYBRID
     fp8_recipe = None
@@ -535,6 +551,18 @@ def run_fwd_bwd(model, x):
             loss.backward()
         return out
 
+    if opts.compile:
+        for i, layer in enumerate(test_model.layers):
+            # dynamic=False for now: symbolic shapes would land in an OpaqueValueBundle
+            # op arg whose hash chokes on non-nested SymInt (see run_numerics).
+            test_model.layers[i] = torch.compile(
+                layer, fullgraph=True, mode=opts.compile_mode, dynamic=False
+            )
+        dist_print(
+            f"Compiled test model layers with torch.compile (mode={opts.compile_mode})...",
+            debug=True,
+        )
+
     torch_rng_state = torch.get_rng_state()
     cuda_rng_state = torch.cuda.get_rng_state(torch.device(f"cuda:{LOCAL_RANK}"))
     if opts.use_cuda_graphs:
diff --git a/tests/pytorch/distributed/run_numerics.py b/tests/pytorch/distributed/run_numerics.py
index fe02f990b4..1590979ba3 100644
--- a/tests/pytorch/distributed/run_numerics.py
+++ b/tests/pytorch/distributed/run_numerics.py
@@ -310,22 +310,45 @@ def _copy_params(model_distributed, model_single):
 
 
 def _apply_models(
-    model_single_node, model_distributed, input_single_node, input_distributed, **kwargs
+    model_single_node,
+    model_distributed,
+    input_single_node,
+    input_distributed,
+    use_compile=False,
+    compile_mode="default",
+    **kwargs,
 ):
     _alloc_main_grad(model_single_node, model_distributed)  # for fuse_wgrad_accumulation=True
     input_single_node.requires_grad_()
     input_distributed.requires_grad_()
+    forward_single_node = model_single_node
+    forward_distributed = model_distributed
+    if use_compile:
+        # Each parametrized case compiles the same module.forward code object with
+        # a different shape/recipe; with dynamic=False those guards accumulate and
+        # eventually trip Dynamo's recompile_limit. Reset so every case starts from
+        # a clean compile cache (mirrors the single-GPU torch.compile tests).
+        torch._dynamo.reset()
+        # dynamic=False for now: a symbolic shape would land in an OpaqueValueBundle
+        # (value-opaque op arg) whose hash chokes on non-nested SymInt. Force static
+        # shapes (recompile per shape) until the bundle handles symbolic shapes.
+        forward_single_node = torch.compile(
+            model_single_node, fullgraph=True, mode=compile_mode, dynamic=False
+        )
+        forward_distributed = torch.compile(
+            model_distributed, fullgraph=True, mode=compile_mode, dynamic=False
+        )
     with te.autocast(
         enabled=QUANTIZATION is not None,
         recipe=quantization_recipe(),
     ):
-        output_single_node = model_single_node(input_single_node, **kwargs)
+        output_single_node = forward_single_node(input_single_node, **kwargs)
     with te.autocast(
         enabled=QUANTIZATION is not None,
         recipe=quantization_recipe(),
         amax_reduction_group=NCCL_WORLD,
     ):
-        output_distributed = model_distributed(input_distributed, **kwargs)
+        output_distributed = forward_distributed(input_distributed, **kwargs)
     return output_single_node, output_distributed
 
 
@@ -641,12 +664,20 @@ def test_quantized_all_gather():
 #                   Linear                 #
 ############################################
 @run_distributed_test()
-def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs):
+def _test_linear(
+    parallel_mode=None,
+    sequence_parallel=False,
+    use_compile=False,
+    compile_mode="default",
+    **kwargs,
+):
     """Test the linear layer with specified parallel mode and sequence parallelization.
 
     Args:
         parallel_mode (str): 'row' or 'column' parallelism.
         sequence_parallel (bool): Enable sequence parallelism if True.
+        use_compile (bool): Wrap the modules in ``torch.compile`` before running.
+        compile_mode (str): ``torch.compile`` mode ("default" or "reduce-overhead").
         kwargs (dict): Additional arguments for the linear layer.
     """
     # Set parameter data type
@@ -696,7 +727,12 @@ def _test_linear(parallel_mode=None, sequence_parallel=False, **kwargs):
 
     # Apply models
     output_single_node, output_distributed = _apply_models(
-        model_single_node, model_distributed, input_single_node, input_distributed
+        model_single_node,
+        model_distributed,
+        input_single_node,
+        input_distributed,
+        use_compile=use_compile,
+        compile_mode=compile_mode,
     )
 
     if "return_bias" in kwargs:
@@ -740,6 +776,8 @@ def test_linear():
         {"params_dtype": torch.float16 if QUANTIZATION != "nvfp4" else torch.bfloat16},
         {"delay_wgrad_compute": True},
         {"save_original_input": True},
+        {"use_compile": True},
+        {"use_compile": True, "compile_mode": "reduce-overhead"},
     ]
 
     for kwargs in kwargs_list:
@@ -747,6 +785,9 @@ def test_linear():
             continue
         if kwargs.get("delay_wgrad_compute", False) and NVTE_TEST_NVINSPECT_ENABLED:
             continue
+        # debug instrumentation forces the eager fallback, so compile is a no-op there.
+        if kwargs.get("use_compile", False) and NVTE_TEST_NVINSPECT_ENABLED:
+            continue
         for parallel_mode in ["column", "row"]:
             for sequence_parallel in [False, True]:
                 _test_linear(parallel_mode, sequence_parallel, **kwargs)
diff --git a/tests/pytorch/distributed/test_comm_gemm_overlap.py b/tests/pytorch/distributed/test_comm_gemm_overlap.py
index 6b1ad870e9..6521ed7dbc 100644
--- a/tests/pytorch/distributed/test_comm_gemm_overlap.py
+++ b/tests/pytorch/distributed/test_comm_gemm_overlap.py
@@ -111,6 +111,8 @@ def _run_layer_with_overlap(
     quantization,
     num_layers=1,
     use_cublasmp=False,
+    compile=False,
+    compile_mode="default",
 ):
     test_path = TEST_ROOT / "run_layer_with_overlap.py"
     test_cmd = LAUNCH_CMD + [
@@ -129,6 +131,10 @@ def _run_layer_with_overlap(
     if overlap_rs_dgrad:
         test_cmd.append("--overlap-rs-dgrad")
 
+    if compile:
+        test_cmd.append("--compile")
+        test_cmd.append(f"--compile-mode={compile_mode}")
+
     if fp8:
         if quantization in ("fp8_delayed_scaling", "fp8_current_scaling") and not fp8_available:
             pytest.skip(reason_for_no_fp8)
@@ -281,6 +287,40 @@ def test_layers_with_overlap_bf16(
     )
 
 
+@pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"])
+@pytest.mark.parametrize(
+    "linear_parallel_mode,overlap_rs_dgrad",
+    [
+        ("row", False),
+        ("column", False),
+        ("column", True),
+    ],
+    ids=[
+        "ROW-PARALLEL",
+        "COL-PARALLEL - BULK DGRAD/WGRAD",
+        "COL-PARALLEL - DGRAD+RS",
+    ],
+)
+def test_linear_with_overlap_compile(linear_parallel_mode, overlap_rs_dgrad, compile_mode):
+    """te.Linear comm+GEMM overlap (Userbuffers) under torch.compile (BF16).
+
+    Userbuffers is expected to stay on Linear's compiled custom-op path (the
+    collective lives inside the opaque op), so this checks that torch.compile +
+    Userbuffers stays numerically correct against the eager, non-overlap reference.
+    ``compile_mode="reduce-overhead"`` additionally exercises CUDA-graph trees on
+    top of the Userbuffers collectives.
+    """
+    _run_layer_with_overlap(
+        te.Linear.__name__,
+        linear_parallel_mode,
+        overlap_rs_dgrad,
+        False,
+        None,
+        compile=True,
+        compile_mode=compile_mode,
+    )
+
+
 @pytest.mark.parametrize("use_cublasmp", (False, True))
 @pytest.mark.parametrize(
     "quantization",
diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index f88bf9ae2c..a0501cca8e 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -3,9 +3,11 @@
 # See LICENSE for license information.
 
 import abc
+import contextlib
 
 import pytest
 import torch
+from torch._dynamo.utils import counters
 from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode
 
 try:
@@ -31,7 +33,6 @@
 from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
 from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, _STORAGE_REGISTRY
 from transformer_engine.pytorch.dynamo import TensorProto, to_tensor_proto
-from transformer_engine.pytorch.quantization import QuantizerRole
 from transformer_engine.pytorch import (
     is_fp8_available,
     is_mxfp8_available,
@@ -88,6 +89,75 @@ def nvfp4_4over6():
     _all_recipes.append(nvfp4_row_scaled())
 
 
+# torch.compile modes exercised by the te.Linear tests: the default backend and
+# "reduce-overhead" (CUDA-graph trees), to ensure the custom-op path is
+# CUDA-graph capturable.
+_compile_modes = ["default", "reduce-overhead"]
+
+
+def _cudagraph_warmup(fn, inp, *, backward: bool) -> None:
+    """
+    Force TE's lazily-created global scratch to be allocated before capture.
+    """
+    out = fn(inp)
+    if backward:
+        out.sum().backward()
+
+
+@contextlib.contextmanager
+def _assert_no_cudagraph_skips(enabled: bool):
+    """Assert ``torch.compile(mode="reduce-overhead")`` actually captured CUDA
+    graphs for every graph instead of silently running it eagerly.
+
+    Inductor bumps ``counters["inductor"]["cudagraph_skips"]`` whenever it
+    declines to capture a cudagraph (input mutation, CPU scalars, cudagraph-unsafe
+    ops, ...) and falls back to eager for that graph. ``fullgraph=True`` only rules
+    out *dynamo* graph breaks, not these *inductor*-level skips, so this guards that
+    the reduce-overhead path didn't degrade to eager. No-op when ``enabled`` is
+    False (e.g. the default backend, where cudagraphs don't apply).
+    """
+    before = counters["inductor"]["cudagraph_skips"]
+    yield
+    if enabled:
+        skipped = counters["inductor"]["cudagraph_skips"] - before
+        assert skipped == 0, (
+            f"reduce-overhead fell back to eager: {skipped} cudagraph skip(s); "
+            "see the 'skipping cudagraphs due to ...' log for the reason"
+        )
+
+
+# bf16 output tolerance: eager and compiled run the same kernels, so they should
+# agree closely; the slack only absorbs reduction-order / cuda-graph differences.
+_EAGER_ATOL, _EAGER_RTOL = 1e-2, 1.6e-2
+
+
+def _assert_close_eager_compiled(fn, compiled, model, base):
+    """Run ``fn`` eagerly and ``compiled`` on identical inputs; assert the
+    forward output and the input / weight gradients match.
+
+    Guards the compiled custom-op path against silently diverging from eager
+    execution -- a wrong-but-same-shape result would slip past shape / grad
+    presence checks alone.
+    """
+    inp_eager = base.detach().clone().requires_grad_(True)
+    model.zero_grad(set_to_none=True)
+    out_eager = fn(inp_eager)
+    out_eager.sum().backward()
+    ref_out = out_eager.detach().clone()
+    ref_wgrad = model.weight.grad.detach().clone()
+    ref_igrad = inp_eager.grad.detach().clone()
+
+    inp_compiled = base.detach().clone().requires_grad_(True)
+    model.zero_grad(set_to_none=True)
+    # Clone before a later cuda-graph replay overwrites the static output buffer.
+    out_compiled = compiled(inp_compiled).clone()
+    out_compiled.sum().backward()
+
+    torch.testing.assert_close(out_compiled, ref_out, atol=_EAGER_ATOL, rtol=_EAGER_RTOL)
+    torch.testing.assert_close(inp_compiled.grad, ref_igrad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL)
+    torch.testing.assert_close(model.weight.grad, ref_wgrad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL)
+
+
 # ---------------------------------------------------------------------------
 # ToyQuantizer – opaque value-type quantizer for torch.compile
 # (requires torch opaque object support, not available in older PyTorch)
@@ -724,9 +794,13 @@ def test_tensor_proto_matches_primitives(factory, shape):
     # Metadata matches the quantizer's.
     assert proto.create_metadata() == q.create_metadata(shape, dtype=torch.bfloat16)
 
-    # inner_names + create_inner_tensors match _describe_buffers.
+    # inner_names follows the storage's canonical __tensor_flatten__ order (the
+    # order the real op flattens its outputs to), while create_inner_tensors
+    # matches the _describe_buffers geometry (a name->shape/dtype mapping).
     bufs = q._describe_buffers(shape)  # pylint: disable=protected-access
-    names = tuple(bufs)
+    direct = _build_from_primitives(q, shape, torch.bfloat16)
+    names = tuple(direct.__tensor_flatten__()[0])
+    assert set(names) == set(bufs)
     assert proto.inner_names() == names
     inner = proto.create_inner_tensors()
     assert len(inner) == len(names)
@@ -736,7 +810,6 @@ def test_tensor_proto_matches_primitives(factory, shape):
         assert buf.dtype == exp_dtype
 
     # The assembled tensor matches one built directly from the primitives.
-    direct = _build_from_primitives(q, shape, torch.bfloat16)
     assert _signature(proto.create_tensor(), names) == _signature(direct, names)
 
 
@@ -817,3 +890,265 @@ def test_to_tensor_proto_quantized(factory, shape):
     assert _signature(proto.create_tensor(), proto.inner_names()) == _signature(
         tensor, proto.inner_names()
     )
+
+
+# ---------------------------------------------------------------------------
+# te.Linear
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+@pytest.mark.parametrize("compile_mode", _compile_modes)
+@pytest.mark.parametrize(
+    "fp8_recipe",
+    [None, *_all_recipes],
+    ids=lambda r: "bf16" if r is None else type(r).__name__,
+)
+def test_te_linear_compiles(fp8_recipe, compile_mode):
+    """
+    torch.compile(fullgraph=True) of ``te.Linear`` under every built-in
+    recipe (plus the bf16-only baseline with no autocast), for both the default
+    backend and ``mode="reduce-overhead"`` (CUDA-graph trees).
+    """
+    if fp8_recipe is not None and not fp8_available:
+        pytest.skip(reason_for_no_fp8)
+
+    dtype = torch.bfloat16
+    device = "cuda"
+
+    # FP8 GEMMs require leading dimensions divisible by 16.
+    model = te.Linear(64, 32, params_dtype=dtype, device=device)
+
+    def fn(inp):
+        if fp8_recipe is None:
+            return model(inp)
+        with te.autocast(recipe=fp8_recipe):
+            return model(inp)
+
+    torch._dynamo.reset()
+    if compile_mode == "reduce-overhead":
+        _cudagraph_warmup(
+            fn,
+            torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True),
+            backward=True,
+        )
+        model.zero_grad(set_to_none=True)
+    compiled = torch.compile(fn, fullgraph=True, mode=compile_mode)
+
+    # ``reduce-overhead`` warms up on the first call(s) and replays a captured
+    # CUDA graph afterwards, so iterate a few times to actually exercise replay.
+    n_iters = 3 if compile_mode == "reduce-overhead" else 1
+    with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"):
+        for _ in range(n_iters):
+            base = torch.randn(32, 64, dtype=dtype, device=device)
+            _assert_close_eager_compiled(fn, compiled, model, base)
+
+
+@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_with_quantized_fp8_weight(compile_mode):
+    """torch.compile should handle Linear weights initialized as FP8 tensors,
+    for both the default backend and ``mode="reduce-overhead"``.
+
+    Exercises the two-tier op + ``register_torch_dispatch`` flattening of a
+    ``Float8Tensor`` weight *input* in
+    :mod:`transformer_engine.pytorch.dynamo`.
+    """
+    dtype = torch.bfloat16
+    device = "cuda"
+    fp8_recipe = recipe.Float8CurrentScaling()
+
+    with te.quantized_model_init(enabled=True, recipe=fp8_recipe):
+        model = te.Linear(64, 32, params_dtype=dtype, device=device)
+
+    assert isinstance(model.weight, te.Float8Tensor)
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            return model(inp)
+
+    torch._dynamo.reset()
+    if compile_mode == "reduce-overhead":
+        _cudagraph_warmup(
+            fn,
+            torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True),
+            backward=True,
+        )
+        model.zero_grad(set_to_none=True)
+    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):
+            base = torch.randn(32, 64, dtype=dtype, device=device)
+            _assert_close_eager_compiled(fn, compiled, model, base)
+
+
+@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_with_fp8_output(compile_mode):
+    """torch.compile of ``te.Linear(..., fp8_output=True)`` without gradient:
+    forward returns a :class:`Float8Tensor`. Covers the default backend and
+    ``mode="reduce-overhead"``.
+
+    Exercises the output-rewrap path in
+    :mod:`transformer_engine.pytorch.dynamo`: when an output quantizer is
+    active, the op returns the flat inner data tensors and the framework
+    rewraps them into a ``Float8Tensor`` via ``__tensor_unflatten__``. A
+    differentiable FP8 output is unsupported under compile (``Linear.forward``
+    falls back to eager), so this test covers the supported case: an FP8 output
+    that does not require grad (inference / ``torch.no_grad``).
+    """
+    dtype = torch.bfloat16
+    device = "cuda"
+    fp8_recipe = recipe.Float8CurrentScaling()
+
+    model = te.Linear(64, 32, params_dtype=dtype, device=device)
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            return model(inp, fp8_output=True)
+
+    torch._dynamo.reset()
+    if compile_mode == "reduce-overhead":
+        with torch.no_grad():
+            _cudagraph_warmup(fn, torch.randn(32, 64, dtype=dtype, device=device), backward=False)
+    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):
+            inp = torch.randn(32, 64, dtype=dtype, device=device)
+            with torch.no_grad():
+                out_eager = fn(inp)
+                out = compiled(inp)
+            assert isinstance(
+                out, te.Float8Tensor
+            ), f"expected Float8Tensor output, got {type(out).__name__}"
+            assert out.shape == (32, 32)
+            assert (
+                out._quantizer is not None
+            ), "FP8 output lost its quantizer on the torch.compile path"
+            # The rewrap rebuilt a fully-functional Float8Tensor: dequantizing it
+            # outside the compiled region exercises scale + data + dtype wiring.
+            deq = out.dequantize()
+            assert deq.shape == (32, 32)
+            assert deq.dtype == dtype
+            # Compiled FP8 output must match the eager FP8 output value-wise.
+            torch.testing.assert_close(
+                deq, out_eager.dequantize(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL
+            )
+
+
+@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_is_first_microbatch(compile_mode):
+    """torch.compile of ``te.Linear`` across a multi-step microbatch schedule that
+    drives FP8 weight caching via ``is_first_microbatch``, for the default backend
+    and ``mode="reduce-overhead"`` (CUDA-graph trees).
+
+    ``is_first_microbatch=True`` quantizes and caches the FP8 weight; subsequent
+    ``False`` steps must reuse the cached FP8 weight instead of re-quantizing. This
+    exercises that cache path under compile and checks it stays numerically aligned
+    with eager. ``is_first_microbatch`` is a Python bool, so each distinct value is
+    its own dynamo guard/graph.
+    """
+    dtype = torch.bfloat16
+    device = "cuda"
+    fp8_recipe = recipe.Float8CurrentScaling()
+    model = te.Linear(64, 32, params_dtype=dtype, device=device)
+
+    # First microbatch caches the FP8 weight, the rest reuse the cache.
+    schedule = [True, False, False]
+    is_first = schedule[0]  # rebound each step; closed over by ``fn``.
+
+    def fn(inp):
+        with te.autocast(recipe=fp8_recipe):
+            return model(inp, is_first_microbatch=is_first)
+
+    torch._dynamo.reset()
+    if compile_mode == "reduce-overhead":
+        _cudagraph_warmup(
+            fn,
+            torch.randn(32, 64, dtype=dtype, device=device, requires_grad=True),
+            backward=True,
+        )
+        model.zero_grad(set_to_none=True)
+    compiled = torch.compile(fn, fullgraph=True, mode=compile_mode)
+
+    with _assert_no_cudagraph_skips(compile_mode == "reduce-overhead"):
+        for is_first in schedule:
+            base = torch.randn(32, 64, dtype=dtype, device=device)
+            _assert_close_eager_compiled(fn, compiled, model, base)
+
+
+@pytest.mark.skipif(not _opaque_available, reason="torch opaque object API not available")
+def test_te_linear_dynamic_shapes():
+    """torch.compile(dynamic=True) of ``te.Linear`` with varying batch sizes.
+
+    Verifies that the compiled graph handles symbolic (dynamic) leading
+    dimensions without graph breaks or recompilations after the initial trace.
+    Key correctness property: a graph compiled for batch=16 must produce
+    numerically correct results for batch=32 without triggering a recompile.
+
+    This exercises two fixes for dynamic shapes:
+    1. ``_linear_setup_ctx`` no longer stores ``inp_shape`` in the value bundle
+       (torch.Size with SymInt dims is not hashable in OpaqueValueBundle).
+    2. ``_linear_backward_impl_fake`` derives dgrad shape from grad_output +
+       weight + SP config instead of relying on the stored ``inp_shape``.
+    3. ``_linear_backward`` reconstructs ``inp_shape`` on-the-fly from the same
+       tensor sources when it is None (compiled mode).
+
+    FP8 + dynamic=True is tracked separately (requires resolving
+    ``UnsafeScriptObjectError`` for TorchScript quantizer objects with Dynamo).
+    """
+    dtype = torch.bfloat16
+    device = "cuda"
+    in_features, out_features = 64, 32
+    model = te.Linear(in_features, out_features, params_dtype=dtype, device=device)
+
+    def fn(inp):
+        return model(inp)
+
+    torch._dynamo.reset()
+    compiled = torch.compile(fn, fullgraph=True)
+
+    batch_sizes = [16, 32, 48]
+
+    for i, batch in enumerate(batch_sizes):
+        inp = torch.randn(batch, in_features, dtype=dtype, device=device, requires_grad=True)
+        # Mark batch dim as dynamic so Dynamo traces once and reuses across batch sizes.
+        torch._dynamo.mark_dynamic(inp, 0)
+        out = compiled(inp)
+        assert out.shape == (batch, out_features), f"wrong output shape for batch={batch}"
+        out.sum().backward()
+        assert inp.grad is not None, f"no input gradient for batch={batch}"
+        assert inp.grad.shape == inp.shape, f"wrong grad shape for batch={batch}"
+
+        # Verify numerics against eager on each distinct batch size.
+        inp_eager = inp.detach().clone().requires_grad_(True)
+        model.zero_grad(set_to_none=True)
+        out_eager = model(inp_eager)
+        out_eager.sum().backward()
+        torch.testing.assert_close(
+            out.detach(), out_eager.detach(), atol=_EAGER_ATOL, rtol=_EAGER_RTOL,
+            msg=f"forward mismatch at batch={batch}",
+        )
+        torch.testing.assert_close(
+            inp.grad, inp_eager.grad, atol=_EAGER_ATOL, rtol=_EAGER_RTOL,
+            msg=f"dgrad mismatch at batch={batch}",
+        )
+
+        if i == 0:
+            # After the first (tracing) call, record the recompile counter
+            # baseline -- subsequent batch sizes must not trigger recompiles.
+            recompile_count_baseline = counters["stats"].get("recompile_reasons", 0)
+
+    recompile_count_after = counters["stats"].get("recompile_reasons", 0)
+    assert recompile_count_after == recompile_count_baseline, (
+        f"Unexpected recompilation(s) across different batch sizes: "
+        f"{recompile_count_after - recompile_count_baseline} recompile(s) detected"
+    )
diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py
index f932a7d9c3..66e0525a9e 100644
--- a/transformer_engine/pytorch/dynamo/__init__.py
+++ b/transformer_engine/pytorch/dynamo/__init__.py
@@ -6,10 +6,12 @@
 
 from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer
 from .tensor_proto import TensorProto, to_tensor_proto
+from .custom_op import register_custom_op
 
 __all__ = [
     "register_value_opaque_quantizer",
     "is_value_opaque_quantizer",
     "TensorProto",
     "to_tensor_proto",
+    "register_custom_op",
 ]
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
new file mode 100644
index 0000000000..ef943d4e8e
--- /dev/null
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -0,0 +1,1322 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""torch.compile custom-op framework for Transformer Engine."""
+
+from __future__ import annotations
+import dataclasses
+import warnings
+from enum import Enum
+from typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Optional,
+    Sequence,
+    Tuple,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+import torch
+
+from .tensor_proto import TensorProto, to_tensor_proto, _contiguous_stride
+from ..quantized_tensor import (
+    QuantizedTensor,
+    QuantizedTensorStorage,
+    Quantizer,
+    _STORAGE_REGISTRY,
+    _quantized_tensor_passthrough_ops,
+    prepare_for_saving,
+)
+
+_TE_OP_NAMESPACE = "transformer_engine_compile"
+_TE_LIB = torch.library.Library(_TE_OP_NAMESPACE, "FRAGMENT")  # noqa: TOR901
+
+
+# ``None`` entries in an op's flat ``Tensor[]`` return are smuggled through a
+# 0-element uint8 tensor: a non-nullable ``Tensor[]`` schema is required for
+# ``register_autograd`` to attach a ``grad_fn`` to the outputs.
+_NONE_SENTINEL_DTYPE = torch.uint8
+
+
+def _encode_none(t: Optional[torch.Tensor]) -> torch.Tensor:
+    """Replace ``None`` with a 0-element uint8 sentinel tensor."""
+    if t is None:
+        return torch.empty(0, dtype=_NONE_SENTINEL_DTYPE)
+    return t
+
+
+def _decode_none(t: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+    """Inverse of :func:`_encode_none`."""
+    if t is None:
+        return None
+    if t.numel() == 0 and t.dtype == _NONE_SENTINEL_DTYPE:
+        return None
+    return t
+
+
+# --------------------------------------------------------------------------- #
+# OpaqueValueBundle: bundle of simple / value-opaque Python values
+# --------------------------------------------------------------------------- #
+
+
+class OpaqueValueBundle:
+    """Opaque value-type bundle of simple Python values.
+
+    Wraps a ``{name: value}`` dict so many small non-Tensor args pass through a
+    single custom-op input; registered as a torch.compile *value* opaque type
+    (Dynamo specializes the graph on its contents). Allowed values: primitives
+    in :attr:`PRIMITIVE_TYPES` (incl. ``torch.Size``), ``enum.Enum``, any
+    registered value-opaque type (e.g. TE quantizers), plus nested tuples /
+    lists / dicts thereof (so a bundle can carry a ``__tensor_flatten__``
+    context verbatim).
+    """
+
+    PRIMITIVE_TYPES: Tuple[type, ...] = (
+        type(None),
+        bool,
+        int,
+        float,
+        str,
+        torch.dtype,
+        torch.device,
+        torch.Size,
+    )
+
+    @classmethod
+    def is_simple_value(cls, value: Any) -> bool:
+        """Whether ``value`` may be stored inside an instance (recursive)."""
+        if isinstance(value, cls.PRIMITIVE_TYPES):
+            return True
+        if isinstance(value, Enum):
+            return True
+        if _is_opaque_value_type(type(value)):
+            return True
+        if isinstance(value, dict):
+            return all(
+                isinstance(k, str) and cls.is_simple_value(v) for k, v in value.items()
+            )
+        if isinstance(value, (list, tuple)):
+            return all(cls.is_simple_value(v) for v in value)
+        return False
+
+    @classmethod
+    def _to_hashable(cls, value: Any) -> Any:
+        if isinstance(value, dict):
+            return tuple(sorted((k, cls._to_hashable(v)) for k, v in value.items()))
+        if isinstance(value, (list, tuple, torch.Size)):
+            return tuple(cls._to_hashable(v) for v in value)
+        return value
+
+    @classmethod
+    def _fmt_simple(cls, value: Any) -> str:
+        """Repr for a value, evaluable in a context with ``torch`` globals."""
+        if isinstance(value, torch.dtype):
+            return f"__import__('torch').{str(value).split('.')[-1]}"
+        if isinstance(value, torch.device):
+            return f"__import__('torch').device({str(value)!r})"
+        if isinstance(value, torch.Size):
+            return f"__import__('torch').Size({list(value)!r})"
+        # Enum before primitives: IntEnum is also ``int`` but must render as
+        # ``EnumName.MEMBER`` (the Enum class is added to globals by ``_collect``).
+        if isinstance(value, Enum):
+            return f"{type(value).__name__}.{value.name}"
+        if isinstance(value, dict):
+            body = ", ".join(f"{k!r}: {cls._fmt_simple(v)}" for k, v in value.items())
+            return f"{{{body}}}"
+        if isinstance(value, list):
+            return "[" + ", ".join(cls._fmt_simple(v) for v in value) + "]"
+        if isinstance(value, tuple):
+            body = ", ".join(cls._fmt_simple(v) for v in value)
+            return f"({body},)" if len(value) == 1 else f"({body})"
+        if _is_opaque_value_type(type(value)):
+            return value.__fx_repr__()[0]
+        return repr(value)
+
+    def __init__(self, data: Optional[Dict[str, Any]] = None) -> None:
+        data = dict(data) if data else {}
+        for k, v in data.items():
+            if not OpaqueValueBundle.is_simple_value(v):
+                raise TypeError(
+                    f"OpaqueValueBundle field '{k}' has unsupported type "
+                    f"{type(v).__name__}; only simple primitives, Enum, "
+                    "torch.Size, registered value-opaque types and nested "
+                    "tuples / lists / dicts thereof are allowed."
+                )
+        self._data: Dict[str, Any] = data
+        self._frozen: Tuple[Tuple[str, Any], ...] = tuple(
+            (k, OpaqueValueBundle._to_hashable(v)) for k, v in sorted(data.items())
+        )
+
+    def __getitem__(self, key: str) -> Any:
+        return self._data[key]
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self._data[name]
+        except KeyError as e:
+            raise AttributeError(name) from e
+
+    def get(self, key: str, default: Any = None) -> Any:
+        """Return ``self._data.get(key, default)``."""
+        return self._data.get(key, default)
+
+    def as_dict(self) -> Dict[str, Any]:
+        """Return a shallow copy of the stored mapping."""
+        return dict(self._data)
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, OpaqueValueBundle):
+            return NotImplemented
+        return self._frozen == other._frozen
+
+    def __hash__(self) -> int:
+        return hash(self._frozen)
+
+    def __fx_repr__(self) -> Tuple[str, Dict[str, Any]]:
+        items = ", ".join(
+            f"{k!r}: {OpaqueValueBundle._fmt_simple(v)}" for k, v in self._data.items()
+        )
+        globals_: Dict[str, Any] = {"OpaqueValueBundle": OpaqueValueBundle}
+
+        def _collect(value: Any) -> None:
+            if isinstance(value, dict):
+                for v in value.values():
+                    _collect(v)
+                return
+            if isinstance(value, (list, tuple)):
+                for v in value:
+                    _collect(v)
+                return
+            if isinstance(value, Enum):
+                globals_[type(value).__name__] = type(value)
+                return
+            if isinstance(value, OpaqueValueBundle.PRIMITIVE_TYPES):
+                return
+            if _is_opaque_value_type(type(value)):
+                _, extra = value.__fx_repr__()
+                globals_.update(extra)
+
+        for v in self._data.values():
+            _collect(v)
+        return (f"OpaqueValueBundle({{{items}}})", globals_)
+
+
+try:
+    from torch._library.opaque_object import (  # pylint: disable=import-outside-toplevel
+        get_opaque_type_name,
+        is_opaque_value_type as _is_opaque_value_type,
+        is_opaque_reference_type as _is_opaque_reference_type,
+        register_opaque_type,
+    )
+
+    register_opaque_type(OpaqueValueBundle, typ="value")
+    _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle)
+except Exception:  # pylint: disable=broad-exception-caught  # pragma: no cover - older torch without opaque_object
+    _is_opaque_value_type = None
+    _is_opaque_reference_type = None
+    _OPAQUE_VALUE_BUNDLE_TYPE_NAME = None
+
+
+def _pg_pickle_stub(*args: Any) -> None:  # pragma: no cover
+    raise RuntimeError("ProcessGroup cannot be unpickled — cache-key use only")
+
+
+def _ensure_distributed_opaque_types() -> None:
+    """Register ``torch.distributed.ProcessGroup`` as a *reference* opaque type.
+
+    A process group is live distributed state: unlike a value-opaque quantizer
+    (which Dynamo bakes into the graph as a constant), it must be carried through
+    the custom op as a graph *input*. PyTorch supports this via
+    ``register_opaque_type(ProcessGroup, typ="reference")`` but only auto-runs it
+    when ``torch.distributed.tensor`` (DTensor) is imported; TE may not import
+    that, so trigger the same idempotent registration here. Best-effort: on
+    builds without the opaque-object / distributed APIs this is a no-op and the
+    process-group field simply falls back to eager under torch.compile.
+
+    Also registers a ``copyreg`` reducer that lets ``FxGraphCachePickler`` hash
+    graphs containing a ``ProcessGroup`` input without crashing.  Without this,
+    inductor logs "Failed to pickle cache key" warnings and bypasses the FX
+    graph disk cache for every distributed compiled call.  The reducer encodes
+    the group as (world_size, rank, backend) — enough to distinguish configs —
+    and raises on reconstruct since deserialization is never needed for hashing.
+    """
+    if _is_opaque_reference_type is None:
+        return
+    try:  # pylint: disable=import-outside-toplevel
+        from torch.distributed.device_mesh import _register_distributed_opaque_types
+
+        _register_distributed_opaque_types()
+    except Exception:  # pylint: disable=broad-exception-caught
+        pass
+
+    # Workaround for PyTorch issue: FxGraphCachePickler handles FakeScriptObject
+    # but not the real ProcessGroup that appears in example_inputs at inductor
+    # compile time.  Register a copyreg reducer so the pickler can hash the key.
+    try:  # pylint: disable=import-outside-toplevel
+        import copyreg
+        import torch.distributed as dist
+        from torch._C._distributed_c10d import ProcessGroup
+
+        if ProcessGroup not in copyreg.dispatch_table:
+
+            def _pg_reduce(pg: ProcessGroup) -> tuple:  # type: ignore[valid-type]
+                try:
+                    return _pg_pickle_stub, (
+                        dist.get_world_size(pg),
+                        dist.get_rank(pg),
+                        dist.get_backend(pg),
+                    )
+                except Exception:  # pylint: disable=broad-exception-caught
+                    return _pg_pickle_stub, (id(pg),)
+
+            copyreg.pickle(ProcessGroup, _pg_reduce)
+    except Exception:  # pylint: disable=broad-exception-caught
+        pass
+
+
+_ensure_distributed_opaque_types()
+
+
+# --------------------------------------------------------------------------- #
+# Storage flatten / unflatten (value-opaque quantizer; no ProcessGroup)
+# --------------------------------------------------------------------------- #
+
+
+def _storage_flatten(value: Any) -> Tuple["OpaqueValueBundle", List[torch.Tensor]]:
+    """Split a ``QuantizedTensor`` / bare storage into ``(meta, Tensor[])``.
+
+    The flatten context (embedding the value-opaque quantizer) plus inner names
+    and -- for a wrapper subclass -- the outer geometry are stashed in the bundle
+    so :func:`_storage_unflatten` can rebuild without PyTorch's ``outer_size``.
+    """
+    inner_names, ctx = value.__tensor_flatten__()
+    meta = dict(ctx)
+    meta["_inner_names"] = list(inner_names)
+    if isinstance(value, torch.Tensor):
+        meta["_outer_shape"] = torch.Size(value.shape)
+    tensors = [getattr(value, name) for name in inner_names]
+    return OpaqueValueBundle(meta), tensors
+
+
+def _storage_unflatten(meta: Any, tensors: List[torch.Tensor]) -> Any:
+    """Inverse of :func:`_storage_flatten`."""
+    meta_dict = meta.as_dict() if isinstance(meta, OpaqueValueBundle) else dict(meta)
+    inner_names = meta_dict["_inner_names"]
+    inner = dict(zip(inner_names, tensors))
+    outer_shape = meta_dict.get("_outer_shape")
+    stride = _contiguous_stride(tuple(outer_shape)) if outer_shape is not None else None
+    return QuantizedTensorStorage.__tensor_unflatten__(inner, meta_dict, outer_shape, stride)
+
+
+# --------------------------------------------------------------------------- #
+# Field buckets: dataclass field <-> flat torch.library slot(s)
+# --------------------------------------------------------------------------- #
+
+
+def _strip_optional(annot: Any) -> Tuple[Any, bool]:
+    """If ``annot`` is ``Optional[X]`` return ``(X, True)``; else ``(annot, False)``."""
+    if get_origin(annot) is Union:
+        args = get_args(annot)
+        if type(None) in args:
+            non_none = [a for a in args if a is not type(None)]
+            if len(non_none) == 1:
+                return non_none[0], True
+    return annot, False
+
+
+class _Bucket:
+    """Maps one (or, for the aggregating bucket, several) dataclass field(s)
+    to/from a contiguous run of custom-op schema *slots*.
+
+    A custom op only takes flat, simply-typed arguments, but a TE op takes a
+    single ``@dataclass`` of mixed fields. Each bucket knows how to translate
+    its kind of field both ways. ``try_build`` and ``schema_slots`` run once at
+    registration (to build the op's schema); ``pack`` and ``unpack`` run on each
+    call and must agree on the slot layout that ``schema_slots`` declares.
+    """
+
+    @classmethod
+    def try_build(cls, name: str, annot: Any) -> Optional["_Bucket"]:
+        """Decide whether this bucket type handles the field ``name`` given its
+        type annotation ``annot``; return a configured bucket if so, else
+        ``None`` so the next candidate is tried.
+
+        Called once per field at registration, in :data:`_FIELD_BUCKETS`
+        priority order.
+        """
+        raise NotImplementedError
+
+    def schema_slots(self) -> List[Tuple[str, str]]:
+        """Declare the schema slots this field occupies, each as a
+        ``(slot_name, schema_type)`` pair (e.g. ``("bias", "Tensor?")``).
+
+        Concatenated across all buckets to form the op's schema string.
+        """
+        raise NotImplementedError
+
+    def pack(self, owner: Any) -> List[Tuple[str, Any]]:
+        """Read this field from the dataclass ``owner`` and produce the concrete
+        value for each of its schema slots, as ``(slot_name, value)`` pairs.
+
+        Composite values are flattened to fit the (tensor-only) slots: e.g. a
+        quantized tensor is split into its plain inner buffers plus a metadata
+        bundle. Inverse of :meth:`unpack`.
+        """
+        raise NotImplementedError
+
+    def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None:
+        """Read this field's slots back from the op arguments ``args`` and write
+        the reconstructed field value into ``kwargs`` (rebuilding any flattened
+        composite). The filled ``kwargs`` are then used to rebuild the original
+        dataclass for the eager implementation. Inverse of :meth:`pack`.
+        """
+        raise NotImplementedError
+
+    def grad_slot(self) -> Optional[int]:
+        """Index (within this bucket's :meth:`schema_slots`) of the slot that
+        carries a gradient, or ``None`` if the field is not differentiable.
+
+        Used to map ``input_tensors_for_grad`` names onto backward grad-output
+        positions. Non-tensor buckets (quantizers, metadata) return ``None``.
+        """
+        return None
+
+
+class _UniversalKind(Enum):
+    """What a universal-tensor slot group carries, tagged in its ``__meta``."""
+
+    NONE = "none"
+    TENSOR = "tensor"
+    STORAGE = "storage"
+
+
+class _UniversalTensorBucket(_Bucket):
+    """``Tensor | QuantizedTensorStorage`` (also subclass tensor) field.
+
+    Three slots regardless of value: ``<name>`` (``Tensor?`` -- plain / subclass
+    tensor passes through, ``None`` for bare storage), ``<name>__tensors``
+    (``Tensor[]`` flat inner tensors when flattened), ``<name>__meta``
+    (``OpaqueValueBundle`` flatten metadata + a ``__kind__`` marker).
+    """
+
+    KIND_KEY = "__kind__"
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+
+    def slot_name(self) -> str:
+        """Primary slot name for a plain / subclass tensor."""
+        return self.name
+
+    def slot_tensors(self) -> str:
+        """Flat inner-tensor slot name."""
+        return self.name + "__tensors"
+
+    def slot_meta(self) -> str:
+        """Flatten-metadata slot name."""
+        return self.name + "__meta"
+
+    def schema_slots(self) -> List[Tuple[str, str]]:
+        return [
+            (self.slot_name(), "Tensor?"),
+            (self.slot_tensors(), "Tensor[]"),
+            (self.slot_meta(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME),
+        ]
+
+    @staticmethod
+    def _is_tensor_storage_union(annot: Any) -> bool:
+        if get_origin(annot) is not Union:
+            return False
+        members = [a for a in get_args(annot) if a is not type(None)]
+        if torch.Tensor not in members:
+            return False
+        return any(
+            isinstance(m, type) and issubclass(m, QuantizedTensorStorage) for m in members
+        )
+
+    @classmethod
+    def try_build(cls, name: str, annot: Any) -> Optional["_UniversalTensorBucket"]:
+        if cls._is_tensor_storage_union(annot):
+            return cls(name)
+        return None
+
+    def pack(self, owner: Any) -> List[Tuple[str, Any]]:
+        value = getattr(owner, self.name)
+        if value is None:
+            return [
+                (self.slot_name(), None),
+                (self.slot_tensors(), []),
+                (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _UniversalKind.NONE})),
+            ]
+        if isinstance(value, torch.Tensor):
+            # Plain tensor *and* subclass (e.g. Float8Tensor) pass through the
+            # ``Tensor?`` slot; subclass flattening (if any) is done by the
+            # outer op's ``register_torch_dispatch`` rule.
+            return [
+                (self.slot_name(), value),
+                (self.slot_tensors(), []),
+                (self.slot_meta(), OpaqueValueBundle({self.KIND_KEY: _UniversalKind.TENSOR})),
+            ]
+        if isinstance(value, QuantizedTensorStorage):
+            meta, tensors = _storage_flatten(value)
+            meta._data[self.KIND_KEY] = _UniversalKind.STORAGE
+            return [
+                (self.slot_name(), None),
+                (self.slot_tensors(), list(tensors)),
+                (self.slot_meta(), meta),
+            ]
+        raise TypeError(
+            f"field {self.name!r} expected None, torch.Tensor, or "
+            f"QuantizedTensorStorage, got {type(value).__name__}"
+        )
+
+    def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None:
+        meta = args[self.slot_meta()]
+        kind = meta.get(self.KIND_KEY)
+        if kind == _UniversalKind.NONE:
+            kwargs[self.name] = None
+        elif kind == _UniversalKind.TENSOR:
+            kwargs[self.name] = args[self.slot_name()]
+        else:
+            kwargs[self.name] = _storage_unflatten(meta, args[self.slot_tensors()])
+
+    def grad_slot(self) -> Optional[int]:
+        # Gradient flows to the plain / subclass tensor slot (``slot_name``,
+        # the first of the three).
+        return 0
+
+
+class _TensorBucket(_Bucket):
+    """``Tensor`` / ``Optional[Tensor]`` -> single ``Tensor`` / ``Tensor?`` slot."""
+
+    def __init__(self, name: str, is_optional: bool) -> None:
+        self.name = name
+        self.type_str = "Tensor?" if is_optional else "Tensor"
+
+    @classmethod
+    def try_build(cls, name: str, annot: Any) -> Optional["_TensorBucket"]:
+        stripped, is_optional = _strip_optional(annot)
+        if stripped is torch.Tensor:
+            return cls(name, is_optional)
+        return None
+
+    def schema_slots(self) -> List[Tuple[str, str]]:
+        return [(self.name, self.type_str)]
+
+    def pack(self, owner: Any) -> List[Tuple[str, Any]]:
+        return [(self.name, getattr(owner, self.name))]
+
+    def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None:
+        kwargs[self.name] = args[self.name]
+
+    def grad_slot(self) -> Optional[int]:
+        return 0
+
+
+class _QuantizerBucket(_Bucket):
+    """``Quantizer`` / ``Optional[Quantizer]`` -> one own ``OpaqueValueBundle`` slot.
+
+    Each quantizer gets its own dedicated slot. The field is annotated with the
+    base ``Quantizer`` (not itself a registered opaque type), so the simple
+    bundle would not claim it.
+    """
+
+    KEY = "q"
+
+    def __init__(self, name: str) -> None:
+        self.name = name
+
+    def slot(self) -> str:
+        """Opaque quantizer metadata slot name."""
+        return self.name + "__q"
+
+    @classmethod
+    def try_build(cls, name: str, annot: Any) -> Optional["_QuantizerBucket"]:
+        stripped, _ = _strip_optional(annot)
+        if not isinstance(stripped, type):
+            return None
+        if issubclass(stripped, Quantizer):
+            return cls(name)
+        return None
+
+    def schema_slots(self) -> List[Tuple[str, str]]:
+        return [(self.slot(), _OPAQUE_VALUE_BUNDLE_TYPE_NAME)]
+
+    def pack(self, owner: Any) -> List[Tuple[str, Any]]:
+        return [(self.slot(), OpaqueValueBundle({self.KEY: getattr(owner, self.name)}))]
+
+    def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None:
+        kwargs[self.name] = args[self.slot()][self.KEY]
+
+
+class _ReferenceOpaqueBucket(_Bucket):
+    """``ProcessGroup`` (or any reference-opaque type) -> one own opaque slot.
+
+    A reference-opaque object is live, stateful black-box data (e.g. a
+    ``torch.distributed.ProcessGroup``): it cannot be specialized on or baked
+    into the graph as a constant the way a value-opaque quantizer is. torch.compile
+    instead carries it through as a graph *input*, so it passes straight through
+    its own schema slot (no ``OpaqueValueBundle`` wrapper). The field is annotated
+    with a concrete type registered via ``register_opaque_type(..., typ="reference")``.
+
+    On the fake / setup-context path the slot holds a ``FakeScriptObject`` (or
+    ``None``); it is assigned to the field verbatim, so the fake impl must never
+    read the object's contents.
+    """
+
+    def __init__(self, name: str, type_name: str, is_optional: bool) -> None:
+        self.name = name
+        self.type_str = f"{type_name}?" if is_optional else type_name
+
+    @classmethod
+    def try_build(cls, name: str, annot: Any) -> Optional["_ReferenceOpaqueBucket"]:
+        if _is_opaque_reference_type is None:
+            return None
+        stripped, is_optional = _strip_optional(annot)
+        if not isinstance(stripped, type):
+            return None
+        if _is_opaque_reference_type(stripped):
+            return cls(name, get_opaque_type_name(stripped), is_optional)
+        return None
+
+    def schema_slots(self) -> List[Tuple[str, str]]:
+        return [(self.name, self.type_str)]
+
+    def pack(self, owner: Any) -> List[Tuple[str, Any]]:
+        return [(self.name, getattr(owner, self.name))]
+
+    def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None:
+        kwargs[self.name] = args[self.name]
+
+
+class _SimpleBundleBucket(_Bucket):
+    """Aggregates every simple-typed field into a single OpaqueValueBundle."""
+
+    SLOT = "_simple_meta"
+
+    def __init__(self, names: List[str]) -> None:
+        self.names = list(names)
+
+    @classmethod
+    def matches_field(cls, annot: Any) -> bool:
+        """Whether ``annot`` (Optional-aware, recursive) is bundle-simple."""
+        annot, _ = _strip_optional(annot)
+        if annot in OpaqueValueBundle.PRIMITIVE_TYPES:
+            return True
+        if isinstance(annot, type) and issubclass(annot, Enum):
+            return True
+        if (
+            isinstance(annot, type)
+            and _is_opaque_value_type is not None
+            and _is_opaque_value_type(annot)
+        ):
+            return True
+        if get_origin(annot) in (tuple, list):
+            inner = [a for a in get_args(annot) if a is not Ellipsis]
+            return bool(inner) and all(cls.matches_field(a) for a in inner)
+        return False
+
+    def schema_slots(self) -> List[Tuple[str, str]]:
+        return [(self.SLOT, _OPAQUE_VALUE_BUNDLE_TYPE_NAME)]
+
+    def pack(self, owner: Any) -> List[Tuple[str, Any]]:
+        return [(self.SLOT, OpaqueValueBundle({n: getattr(owner, n) for n in self.names}))]
+
+    def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None:
+        if self.SLOT not in args:
+            return
+        meta = args[self.SLOT]
+        for n in self.names:
+            kwargs[n] = meta[n]
+
+
+class _UnknownBucket(_Bucket):
+    """Fallback for fields no other bucket claims.
+
+    Emits no slot; pack rejects non-trivial values (anything other than
+    ``None`` / all-``None`` sequence); unpack restores the field as ``None``.
+    """
+
+    def __init__(self, name: str, owner_cls_name: str) -> None:
+        self.name = name
+        self.owner_cls_name = owner_cls_name
+
+    @staticmethod
+    def _is_trivial(value: Any) -> bool:
+        if value is None:
+            return True
+        if isinstance(value, (list, tuple)):
+            return all(v is None for v in value)
+        return False
+
+    def schema_slots(self) -> List[Tuple[str, str]]:
+        return []
+
+    def pack(self, owner: Any) -> List[Tuple[str, Any]]:
+        value = getattr(owner, self.name, None)
+        if not self._is_trivial(value):
+            raise TypeError(
+                f"{self.owner_cls_name} field {self.name!r} has a type not "
+                "supported by torch.compile (not Tensor, simple, or Quantizer) "
+                "and carries a non-trivial value; add a matching bucket in "
+                "dynamo.py to handle it."
+            )
+        return []
+
+    def unpack(self, args: Dict[str, Any], kwargs: Dict[str, Any]) -> None:
+        kwargs[self.name] = None
+
+
+# Buckets, in priority order, owning ``try_build`` for a single field.
+_FIELD_BUCKETS: Tuple[type, ...] = (
+    _UniversalTensorBucket,
+    _TensorBucket,
+    _ReferenceOpaqueBucket,
+    _QuantizerBucket,
+)
+
+
+def _resolved_field_annotations(cls: type) -> List[Tuple[str, Any]]:
+    """Return ``[(field_name, resolved_type), ...]`` for a dataclass."""
+    if not dataclasses.is_dataclass(cls):
+        raise TypeError(f"{cls.__name__} must be a @dataclass to be a TE op arg container.")
+    try:
+        hints = get_type_hints(cls)
+    except Exception:  # pylint: disable=broad-exception-caught
+        hints = {}
+    return [(f.name, hints.get(f.name, f.type)) for f in dataclasses.fields(cls)]
+
+
+def _get_buckets(cls: type) -> List[_Bucket]:
+    """Build the bucket list for a dataclass from its field annotations."""
+    if _OPAQUE_VALUE_BUNDLE_TYPE_NAME is None:
+        raise RuntimeError(
+            f"{cls.__name__} cannot be turned into a TE custom op: OpaqueValueBundle "
+            "is not registered as a torch._library value-opaque type (PyTorch build "
+            "without opaque-object support)."
+        )
+    buckets: List[_Bucket] = []
+    simple_names: List[str] = []
+    for name, annot in _resolved_field_annotations(cls):
+        built: Optional[_Bucket] = None
+        for bucket_cls in _FIELD_BUCKETS:
+            built = bucket_cls.try_build(name, annot)
+            if built is not None:
+                break
+        if built is not None:
+            buckets.append(built)
+        elif _SimpleBundleBucket.matches_field(annot):
+            simple_names.append(name)
+        else:
+            buckets.append(_UnknownBucket(name, cls.__name__))
+    if simple_names:
+        buckets.append(_SimpleBundleBucket(simple_names))
+    return buckets
+
+
+def _tensor_field_names(buckets: List[_Bucket]) -> List[str]:
+    """Names of fields carrying tensors (for building the proto view)."""
+    return [b.name for b in buckets if isinstance(b, (_TensorBucket, _UniversalTensorBucket))]
+
+
+def _build_schema(buckets: List[_Bucket]) -> Tuple[str, List[str]]:
+    """Return ``(schema_arg_str, slot_names)`` for a bucket list."""
+    spec = [slot for b in buckets for slot in b.schema_slots()]
+    names = [name for name, _ in spec]
+    schema_str = "(" + ", ".join(f"{type_str} {name}" for name, type_str in spec) + ")"
+    return schema_str, names
+
+
+def _pack(obj: Any, buckets: List[_Bucket]) -> Dict[str, Any]:
+    """Build the op's flat ``{slot_name: value}`` argument dict from an args
+    dataclass ``obj`` (e.g. ``LinearFwdArgs``), by collecting every bucket's
+    packed slot(s). Inverse of :func:`_unpack`.
+    """
+    out: Dict[str, Any] = {}
+    for bucket in buckets:
+        for name, value in bucket.pack(obj):
+            out[name] = value
+    return out
+
+
+def _unpack(cls: type, args: Dict[str, Any], buckets: List[_Bucket]) -> Any:
+    """Rebuild a fresh args dataclass ``cls`` (e.g. ``LinearFwdArgs``) from the
+    op's flat slot ``args`` dict, by letting every bucket restore its field(s).
+    Inverse of :func:`_pack`.
+    """
+    kwargs: Dict[str, Any] = {}
+    for bucket in buckets:
+        bucket.unpack(args, kwargs)
+    obj = cls.__new__(cls)
+    for k, v in kwargs.items():
+        object.__setattr__(obj, k, v)
+    return obj
+
+
+def _proto_view(obj: Any, tensor_field_names: Sequence[str]) -> Any:
+    """Copy of dataclass ``obj`` with each tensor field replaced by a :class:`TensorProto`.
+
+    Only tensor fields have a ``TensorProto`` equivalent, so quantizer / scalar
+    fields are simply carried over unchanged; the fake impl works purely on
+    geometry. Built with :func:`dataclasses.replace` (the only such construction
+    Dynamo can trace).
+    """
+    overrides: Dict[str, Any] = {}
+    for name in tensor_field_names:
+        value = getattr(obj, name, None)
+        if value is not None and not isinstance(value, TensorProto):
+            overrides[name] = to_tensor_proto(value)
+    if not overrides:
+        return obj
+    return dataclasses.replace(obj, **overrides)
+
+
+# --------------------------------------------------------------------------- #
+# Op outputs <-> flat ``Tensor[]`` payload: this is how an op returns / saves
+# quantized tensors (and wrapper subclasses). Outputs are flattened to their
+# inner buffers on the way out and rebuilt via ``__tensor_unflatten__`` on the
+# way back; on the fake side a TensorProto supplies the geometry.
+# --------------------------------------------------------------------------- #
+
+
+def _proto_slot_count(proto: Optional[TensorProto]) -> int:
+    """Flat ``Tensor[]`` slots the value for ``proto`` occupies."""
+    if proto is None:
+        return 1
+    return len(proto.inner_names())
+
+
+def _proto_reassemble(
+    proto: Optional[TensorProto],
+    chunk: List[Optional[torch.Tensor]],
+) -> Optional[Union[torch.Tensor, QuantizedTensorStorage]]:
+    """Rebuild the value described by ``proto`` from its flat tensors ``chunk``.
+
+    ``proto`` describes one output: ``None`` (-> ``None``), a plain tensor
+    (``chunk`` is the single tensor, returned as-is), or a quantized tensor
+    (``chunk`` are its inner tensors, reassembled into the wrapper subclass via
+    ``__tensor_unflatten__``).
+    """
+    if proto is None:
+        return None
+    if proto.quantizer is None:
+        return chunk[0]
+    inner_names = proto.inner_names()
+    meta = proto.create_metadata()
+    shape = tuple(proto.shape)
+    stride = _contiguous_stride(shape)
+    storage_cls = _STORAGE_REGISTRY[meta["cls"]]
+    inner_dict = dict(zip(inner_names, chunk))
+    return storage_cls.__tensor_unflatten__(inner_dict, meta, shape, stride)
+
+
+def _value_to_flat_tensors(
+    value: Optional[Union[torch.Tensor, QuantizedTensorStorage, TensorProto]],
+) -> List[torch.Tensor]:
+    """Return the flat ``Tensor[]`` slots that represent one op output ``value``.
+
+    Inverse of :func:`_proto_reassemble`; the slot count matches
+    :func:`_proto_slot_count`.
+    """
+    if value is None:
+        return [_encode_none(None)]
+    if isinstance(value, TensorProto):
+        return [_encode_none(t) for t in value.create_inner_tensors()]
+    if isinstance(value, torch.Tensor):
+        if type(value) is not torch.Tensor and hasattr(  # pylint: disable=unidiomatic-typecheck
+            value, "__tensor_flatten__"
+        ):
+            inner_names, _ = value.__tensor_flatten__()
+            return [_encode_none(getattr(value, n)) for n in inner_names]
+        return [_encode_none(value)]
+    if hasattr(value, "__tensor_flatten__"):
+        inner_names, _ = value.__tensor_flatten__()
+        return [_encode_none(getattr(value, n)) for n in inner_names]
+    raise TypeError(
+        f"unsupported value type {type(value).__name__}; expected None / "
+        "torch.Tensor / tensor subclass / bare storage / TensorProto."
+    )
+
+
+# Trailing slots in every fwd-impl return: ``tensors_to_save, tensor_objects,
+# ctx_attrs``. User-output count is ``len(result) - this``.
+_FWD_TRAILING_SLOTS = 3
+
+
+def _format_fwd_result(result: Any) -> List[torch.Tensor]:
+    """Pack a fwd-impl return tuple into the op's ``Tensor[]`` payload.
+
+    User outputs first, then saved-for-backward tensors in declaration order.
+    """
+    num_outputs = len(result) - _FWD_TRAILING_SLOTS
+    flat: List[torch.Tensor] = []
+    for value in result[:num_outputs]:
+        flat.extend(_value_to_flat_tensors(value))
+    saved = result[num_outputs] or ()
+    for value in saved:
+        flat.extend(_value_to_flat_tensors(value))
+    return flat
+
+
+def _format_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:`TensorProto` grad is materialized into a single tensor.
+    """
+    grads = list(grads)
+    if len(grads) != num_grad_inputs:
+        raise RuntimeError(
+            f"{op_qualname} expected backward_impl to return {num_grad_inputs} grads "
+            f"(one per input_tensors_for_grad entry), got {len(grads)}"
+        )
+    out: List[torch.Tensor] = []
+    for g in grads:
+        if isinstance(g, TensorProto):
+            out.append(_encode_none(g.create_tensor()))
+        else:
+            out.append(_encode_none(g))
+    return out
+
+
+def _split_fwd_fake_result(
+    result: Tuple[Any, ...],
+) -> Tuple[List[Any], List[Any], Dict[str, Any]]:
+    """Slice a fwd fake-impl return into ``(user_fakes, saved_fakes, ctx_attrs)``."""
+    num_outputs = len(result) - _FWD_TRAILING_SLOTS
+    saved = result[num_outputs]
+    ctx_attrs = result[num_outputs + 2]
+    user_fakes = list(result[:num_outputs])
+    saved_fakes = list(saved) if saved is not None else []
+    ctx_attrs = dict(ctx_attrs) if ctx_attrs else {}
+    return user_fakes, saved_fakes, ctx_attrs
+
+
+# --------------------------------------------------------------------------- #
+# Op registration
+# --------------------------------------------------------------------------- #
+
+
+def _resolve_grad_targets(
+    fwd_buckets: List[_Bucket],
+    fwd_arg_type: type,
+    input_tensors_for_grad: List[str],
+) -> Tuple[List[Any], List[int]]:
+    """Validate ``input_tensors_for_grad`` and resolve the grad-output layout.
+
+    Returns ``(fwd_slot_defaults, grad_targets)``: the per-slot no-grad template
+    (``[]`` for ``Tensor[]`` slots, ``None`` otherwise) and, for each requested
+    input name, the schema-slot index its gradient maps to.
+    """
+    fwd_slot_defaults: List[Any] = []
+    name_to_slot: Dict[str, int] = {}
+    slot_offset = 0
+    for bucket in fwd_buckets:
+        slots = bucket.schema_slots()
+        for _, type_str in slots:
+            fwd_slot_defaults.append([] if type_str.endswith("[]") else None)
+        grad_slot = bucket.grad_slot()
+        if grad_slot is not None:
+            name_to_slot[bucket.name] = slot_offset + grad_slot
+        slot_offset += len(slots)
+
+    unknown = [n for n in input_tensors_for_grad if n not in name_to_slot]
+    if unknown:
+        raise ValueError(
+            f"input_tensors_for_grad contains names not in {fwd_arg_type.__name__} "
+            f"schema: {unknown}"
+        )
+    grad_targets = [name_to_slot[n] for n in input_tensors_for_grad]
+    return fwd_slot_defaults, grad_targets
+
+
+def _register_kernel(
+    *,
+    op_name: str,
+    op_qualname: str,
+    arg_type: type,
+    arg_names: List[str],
+    buckets: List[_Bucket],
+    tensor_field_names: List[str],
+    impl: Callable[[Any], Any],
+    fake_impl: Callable[[Any], Any],
+    format_result: Callable[[Any], List[torch.Tensor]],
+) -> None:
+    """Wire the real ``impl`` + the ``fake_impl`` (proto) into the library.
+
+    The real kernel rebuilds the dataclass and runs ``impl``; the fake kernel
+    runs the proto fake impl on the :func:`_proto_view`. Both go through
+    ``format_result``.
+    """
+
+    def _impl(*flat: Any) -> List[torch.Tensor]:
+        kwargs = dict(zip(arg_names, flat))
+        obj = _unpack(arg_type, kwargs, buckets)
+        return format_result(impl(obj))
+
+    def _fake(*flat: Any) -> List[torch.Tensor]:
+        kwargs = dict(zip(arg_names, flat))
+        obj = _unpack(arg_type, kwargs, buckets)
+        proto_obj = _proto_view(obj, tensor_field_names)
+        return format_result(fake_impl(proto_obj))
+
+    _TE_LIB.impl(op_name, _impl, "CompositeExplicitAutograd")
+    torch.library.register_fake(op_qualname, _fake, lib=_TE_LIB)
+
+
+def _register_autograd_for_op(
+    *,
+    fwd_op_name: str,
+    bwd_op_name: str,
+    fwd_arg_type: type,
+    fwd_arg_names: List[str],
+    fwd_buckets: List[_Bucket],
+    fwd_tensor_field_names: List[str],
+    bwd_arg_names: List[str],
+    bwd_buckets: List[_Bucket],
+    fwd_slot_defaults: List[Any],
+    grad_targets: List[int],
+    setup_context_user: Callable[..., Any],
+    backward_obj_type: type,
+    fwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
+) -> None:
+    """Wire ``register_autograd`` on a forward op so its backward calls ``bwd_op_name``.
+
+    ``setup_context`` re-runs the proto fwd fake impl to recover output / saved
+    templates, reassembles each flat output chunk, and hands the saved tuple +
+    ``ctx_attrs`` to the module's ``setup_context``.
+    """
+    fwd_qualname = f"{_TE_OP_NAMESPACE}::{fwd_op_name}"
+
+    def _setup_context(ctx, inputs, output):
+        ctx._te_fwd_tensor_list_lengths = {
+            i: len(value) for i, value in enumerate(inputs) if isinstance(value, list)
+        }
+        kwargs = dict(zip(fwd_arg_names, inputs))
+        fwd_obj = _unpack(fwd_arg_type, kwargs, fwd_buckets)
+        proto_obj = _proto_view(fwd_obj, fwd_tensor_field_names)
+
+        user_fakes, saved_fakes, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(proto_obj))
+
+        cursor = 0
+        user_outputs: List[Any] = []
+        for proto in user_fakes:
+            n = _proto_slot_count(proto)
+            chunk = [_decode_none(t) for t in output[cursor : cursor + n]]
+            cursor += n
+            user_outputs.append(_proto_reassemble(proto, chunk))
+
+        saved_list: List[Any] = []
+        for proto in saved_fakes:
+            n = _proto_slot_count(proto)
+            chunk = [_decode_none(t) for t in output[cursor : cursor + n]]
+            cursor += n
+            saved_list.append(_proto_reassemble(proto, chunk))
+
+        bwd_obj = backward_obj_type()
+        tensors_to_save_from_setup = setup_context_user(
+            bwd_obj,
+            fwd_obj,
+            user_outputs[0] if len(user_fakes) == 1 else tuple(user_outputs),
+            ctx_attrs,
+            tuple(saved_list),
+        )
+        tensors_to_save, tensor_objects = prepare_for_saving(
+            *(tensors_to_save_from_setup or ())
+        )
+        ctx.tensor_objects = tensor_objects
+        ctx.save_for_backward(*tensors_to_save)
+        ctx.bwd_obj = bwd_obj
+
+    def _autograd_backward(ctx, *grad_outputs):
+        bwd_obj = ctx.bwd_obj
+        if hasattr(bwd_obj, "setup_saved_tensors"):
+            bwd_obj.setup_saved_tensors(ctx)
+        ctx.tensor_objects = None
+        per_output_grads = grad_outputs[0]
+        bwd_obj.grad_output = _decode_none(per_output_grads[0])
+        kwargs = _pack(bwd_obj, bwd_buckets)
+        bwd_args_flat = [kwargs[name] for name in bwd_arg_names]
+        bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), bwd_op_name)
+        grads = [_decode_none(g) for g in bwd_op(*bwd_args_flat)]
+        out: List[Any] = list(fwd_slot_defaults)
+        tensor_list_lengths = getattr(ctx, "_te_fwd_tensor_list_lengths", {})
+        for pos, length in tensor_list_lengths.items():
+            if isinstance(out[pos], list):
+                out[pos] = [None] * length
+        for pos, g in zip(grad_targets, grads):
+            out[pos] = g
+        return tuple(out)
+
+    torch.library.register_autograd(
+        fwd_qualname,
+        _autograd_backward,
+        setup_context=_setup_context,
+        lib=_TE_LIB,
+    )
+
+
+def _collect_universal_slot_offsets(buckets: List[_Bucket]) -> List[int]:
+    """Start index of each ``_UniversalTensorBucket`` group in the flat args."""
+    offsets: List[int] = []
+    pos = 0
+    for bucket in buckets:
+        if isinstance(bucket, _UniversalTensorBucket):
+            offsets.append(pos)
+        pos += len(bucket.schema_slots())
+    return offsets
+
+
+def _flatten_subclass_into_slots(
+    new_args: List[Any], slot_offsets: List[int], subclass: type
+) -> None:
+    """Rewrite each universal-bucket group whose ``Tensor?`` slot holds an
+    instance of ``subclass`` into the storage layout (3 slots: name / tensors / meta).
+    """
+    for offset in slot_offsets:
+        val = new_args[offset]
+        if val is None or not isinstance(val, subclass):
+            continue
+        meta, tensors = _storage_flatten(val)
+        meta._data[_UniversalTensorBucket.KIND_KEY] = _UniversalKind.STORAGE
+        new_args[offset] = None
+        new_args[offset + 1] = list(tensors)
+        new_args[offset + 2] = meta
+
+
+def _register_outer_forwarder(
+    *,
+    outer_op_name: str,
+    inner_op_name: str,
+    buckets: Optional[List[_Bucket]] = None,
+    subclass_list: Optional[List[type]] = None,
+) -> None:
+    """Register the outer op's default kernel + fake: forward to the inner op,
+    optionally flattening registered subclass inputs in place first.
+    """
+    inner_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_op_name)
+    input_flatten_enabled = bool(subclass_list) and buckets is not None
+    slot_offsets = _collect_universal_slot_offsets(buckets) if input_flatten_enabled else []
+
+    def _forward(*flat: Any) -> List[torch.Tensor]:
+        if not input_flatten_enabled:
+            return inner_op(*flat)
+        new_args = list(flat)
+        for sub in subclass_list:
+            _flatten_subclass_into_slots(new_args, slot_offsets, sub)
+        return inner_op(*new_args)
+
+    _TE_LIB.impl(outer_op_name, _forward, "CompositeExplicitAutograd")
+    torch.library.register_fake(f"{_TE_OP_NAMESPACE}::{outer_op_name}", _forward, lib=_TE_LIB)
+
+
+def _all_quantized_tensor_subclasses() -> List[type]:
+    """Return every imported ``QuantizedTensor`` wrapper subclass."""
+    import transformer_engine.pytorch.tensor  # noqa: F401  pylint: disable=import-outside-toplevel,unused-import
+    return [cls for cls in _STORAGE_REGISTRY.values() if issubclass(cls, QuantizedTensor)]
+
+
+def register_custom_op(
+    *,
+    op_name: str,
+    input_tensors_for_grad: List[str],
+    fwd_arg_type: type,
+    fwd_impl: Callable[[Any], Any],
+    setup_context: Callable[..., Any],
+    backward_arg_type: type,
+    backward_obj: type,
+    backward_impl: Callable[[Any], Any],
+    fwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
+    bwd_fake_impl: Callable[[Any], Tuple[Any, ...]],
+) -> Optional[Callable[..., Any]]:
+    """Register a TE module's forward + backward as torch custom ops.
+
+    Always two-tier: an inner ``<op>_base`` op carries the real schema /
+    autograd, and an outer ``<op>`` op forwards to it, flattening any
+    quantized-tensor wrapper inputs first via ``register_torch_dispatch`` (an
+    empty subclass list simply makes the outer op a pass-through, so a pure
+    plain-tensor / bf16 call goes straight through).
+
+    Returns ``forward_fn(fwd_arg_type_instance)`` -- a drop-in for
+    ``Function.apply`` under ``torch.compiler.is_compiling()`` that dispatches
+    through the outer op and returns the user-facing outputs.
+
+    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.
+    """
+    try:
+        return _register_custom_op_impl(
+            op_name=op_name,
+            input_tensors_for_grad=input_tensors_for_grad,
+            fwd_arg_type=fwd_arg_type,
+            fwd_impl=fwd_impl,
+            setup_context=setup_context,
+            backward_arg_type=backward_arg_type,
+            backward_obj=backward_obj,
+            backward_impl=backward_impl,
+            fwd_fake_impl=fwd_fake_impl,
+            bwd_fake_impl=bwd_fake_impl,
+        )
+    except (ImportError, AttributeError, RuntimeError, TypeError) as e:
+        warnings.warn(
+            f"Could not register the torch.compile custom op '{op_name}' "
+            f"({type(e).__name__}: {e}); modules using it will fall back to eager "
+            "execution under torch.compile (a graph break, incompatible with "
+            "fullgraph=True)."
+        )
+        return None
+
+
+def _register_custom_op_impl(
+    *,
+    op_name: str,
+    input_tensors_for_grad: List[str],
+    fwd_arg_type: type,
+    fwd_impl: Callable[[Any], Any],
+    setup_context: Callable[..., Any],
+    backward_arg_type: type,
+    backward_obj: type,
+    backward_impl: Callable[[Any], Any],
+    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."""
+    outer_fwd_name = op_name
+    outer_bwd_name = f"{op_name}_backward"
+    inner_fwd_name = f"{op_name}_base"
+    inner_bwd_name = f"{outer_bwd_name}_base"
+    subclass_list = _all_quantized_tensor_subclasses()
+
+    fwd_buckets = _get_buckets(fwd_arg_type)
+    bwd_buckets = _get_buckets(backward_arg_type)
+    fwd_tensor_field_names = _tensor_field_names(fwd_buckets)
+    bwd_tensor_field_names = _tensor_field_names(bwd_buckets)
+
+    fwd_schema_args, fwd_arg_names = _build_schema(fwd_buckets)
+    bwd_schema_args, bwd_arg_names = _build_schema(bwd_buckets)
+
+    num_grad_inputs = len(input_tensors_for_grad)
+    fwd_slot_defaults, grad_targets = _resolve_grad_targets(
+        fwd_buckets, fwd_arg_type, input_tensors_for_grad
+    )
+
+    _TE_LIB.define(f"{inner_fwd_name}{fwd_schema_args} -> Tensor[]")
+    _TE_LIB.define(f"{inner_bwd_name}{bwd_schema_args} -> Tensor[]")
+    _TE_LIB.define(f"{outer_fwd_name}{fwd_schema_args} -> Tensor[]")
+    _TE_LIB.define(f"{outer_bwd_name}{bwd_schema_args} -> Tensor[]")
+
+    inner_fwd_qualname = f"{_TE_OP_NAMESPACE}::{inner_fwd_name}"
+    inner_bwd_qualname = f"{_TE_OP_NAMESPACE}::{inner_bwd_name}"
+    outer_fwd_qualname = f"{_TE_OP_NAMESPACE}::{outer_fwd_name}"
+    outer_bwd_qualname = f"{_TE_OP_NAMESPACE}::{outer_bwd_name}"
+
+    _register_kernel(
+        op_name=inner_fwd_name,
+        op_qualname=inner_fwd_qualname,
+        arg_type=fwd_arg_type,
+        arg_names=fwd_arg_names,
+        buckets=fwd_buckets,
+        tensor_field_names=fwd_tensor_field_names,
+        impl=fwd_impl,
+        fake_impl=fwd_fake_impl,
+        format_result=_format_fwd_result,
+    )
+    _register_kernel(
+        op_name=inner_bwd_name,
+        op_qualname=inner_bwd_qualname,
+        arg_type=backward_arg_type,
+        arg_names=bwd_arg_names,
+        buckets=bwd_buckets,
+        tensor_field_names=bwd_tensor_field_names,
+        impl=backward_impl,
+        fake_impl=bwd_fake_impl,
+        format_result=lambda g: _format_bwd_result(g, num_grad_inputs, inner_bwd_qualname),
+    )
+
+    autograd_common = {
+        "fwd_arg_type": fwd_arg_type,
+        "fwd_arg_names": fwd_arg_names,
+        "fwd_buckets": fwd_buckets,
+        "fwd_tensor_field_names": fwd_tensor_field_names,
+        "bwd_arg_names": bwd_arg_names,
+        "bwd_buckets": bwd_buckets,
+        "fwd_slot_defaults": fwd_slot_defaults,
+        "grad_targets": grad_targets,
+        "setup_context_user": setup_context,
+        "backward_obj_type": backward_obj,
+        "fwd_fake_impl": fwd_fake_impl,
+    }
+    _register_autograd_for_op(
+        fwd_op_name=inner_fwd_name, bwd_op_name=inner_bwd_name, **autograd_common
+    )
+    _register_autograd_for_op(
+        fwd_op_name=outer_fwd_name, bwd_op_name=outer_bwd_name, **autograd_common
+    )
+
+    _register_outer_forwarder(
+        outer_op_name=outer_fwd_name,
+        inner_op_name=inner_fwd_name,
+        buckets=fwd_buckets,
+        subclass_list=list(subclass_list),
+    )
+    _register_outer_forwarder(outer_op_name=outer_bwd_name, inner_op_name=inner_bwd_name)
+
+    inner_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_fwd_name)
+    inner_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_bwd_name)
+    outer_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), outer_fwd_name)
+    outer_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), outer_bwd_name)
+
+    fwd_slot_offsets = _collect_universal_slot_offsets(fwd_buckets)
+    bwd_slot_offsets = _collect_universal_slot_offsets(bwd_buckets)
+
+    def _fwd_rule(mode, func, types, args, kwargs):
+        del mode, func, types, kwargs
+        new_args = list(args)
+        for sub in subclass_list:
+            _flatten_subclass_into_slots(new_args, fwd_slot_offsets, sub)
+        return inner_fwd_op(*new_args)
+
+    def _bwd_rule(mode, func, types, args, kwargs):
+        del mode, func, types, kwargs
+        new_args = list(args)
+        for sub in subclass_list:
+            _flatten_subclass_into_slots(new_args, bwd_slot_offsets, sub)
+        return inner_bwd_op(*new_args)
+
+    for sub in subclass_list:
+        torch.library.register_torch_dispatch(outer_fwd_qualname, sub, _fwd_rule, lib=_TE_LIB)
+        torch.library.register_torch_dispatch(outer_bwd_qualname, sub, _bwd_rule, lib=_TE_LIB)
+
+    _quantized_tensor_passthrough_ops.add(outer_fwd_op.default)
+    _quantized_tensor_passthrough_ops.add(outer_bwd_op.default)
+    _quantized_tensor_passthrough_ops.add(inner_fwd_op.default)
+    _quantized_tensor_passthrough_ops.add(inner_bwd_op.default)
+
+    def forward_fn(fwd_args):
+        proto_obj = _proto_view(fwd_args, fwd_tensor_field_names)
+        user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(proto_obj))
+        kwargs = _pack(fwd_args, fwd_buckets)
+        flat_in = [kwargs[name] for name in fwd_arg_names]
+        result = outer_fwd_op(*flat_in)
+
+        cursor = 0
+        outputs: List[Any] = []
+        for proto in user_fakes:
+            n = _proto_slot_count(proto)
+            chunk = [_decode_none(t) for t in result[cursor : cursor + n]]
+            cursor += n
+            outputs.append(_proto_reassemble(proto, chunk))
+
+        if len(outputs) == 1:
+            return outputs[0]
+        return tuple(outputs)
+
+    return forward_fn
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index 473483faff..d1ea805077 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -3,10 +3,12 @@
 # See LICENSE for license information.
 
 """Linear API"""
+
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, Optional, Tuple, Union, List
 from functools import reduce
 from operator import mul as multiply_op
+import math
 import warnings
 import weakref
 
@@ -38,10 +40,10 @@
     divide,
     init_method_constant,
     needs_quantized_gemm,
-    assert_dim_for_fp8_exec,
     nvtx_range_pop,
     nvtx_range_push,
     get_nvtx_range_context,
+    warn_compile_eager_fallback,
 )
 from ..distributed import (
     set_tensor_model_parallel_attributes,
@@ -58,9 +60,10 @@
 from ..cpp_extensions import (
     general_gemm,
 )
+from ..cpp_extensions.gemm import get_cublas_workspace
 from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type
-from ..jit import no_torch_dynamo
 from ..graph import is_graph_capturing
+from ..jit import no_torch_dynamo
 from ..quantized_tensor import (
     QuantizedTensor,
     QuantizedTensorStorage,
@@ -68,7 +71,7 @@
     prepare_for_saving,
     restore_from_func_ctx,
 )
-from ..dynamo import TensorProto
+from ..dynamo import TensorProto, register_custom_op, is_value_opaque_quantizer
 from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer
 from ..tensor.mxfp8_tensor import MXFP8Quantizer
 from ..tensor.utils import clear_columnwise_cache, is_custom
@@ -97,7 +100,24 @@ class LinearFwdArgs:
     bias: Optional[torch.Tensor]
 
     # --- Non-differentiable cached tensors ---
-    weight_workspace: Optional[torch.Tensor]
+    # Same union as ``weight`` so a cached quantized workspace is flattened to its
+    # inner tensors on the way into the op (symmetric with ``new_weight_workspace``
+    # on the way out); a plain ``Tensor?`` slot can't carry a quantized subclass
+    # across the torch.compile custom-op boundary.
+    weight_workspace: Optional[TensorOrQuantized]
+
+    # --- CUDA-graph workspace pinning (torch.compile / reduce-overhead) ---
+    # Fetched in the *traced* module forward and threaded in as op inputs purely so
+    # the process-global, lru_cached cuBLAS / NVFP4-RHT workspaces become graph
+    # inputs: they are then allocated at trace time in the normal allocator (external
+    # to the cudagraph private pool) instead of being created inside the op during
+    # capture, where a persistent allocation trips check_memory_pool ("tensor not
+    # tracked as outputs"). The op body never reads these; general_gemm / the
+    # quantizer fetch the same lru_cached globals by address. Pinning the workspace
+    # in the forward also covers the backward GEMM, which reuses the same cached
+    # global. None on eager / non-compiled paths.
+    cublas_workspace: Optional[torch.Tensor]
+    rht_matrix: Optional[torch.Tensor]
 
     # --- requires_grad flags (cached so backward does not re-query) ---
     input_requires_grad: bool
@@ -131,7 +151,9 @@ class LinearFwdArgs:
 
     # --- Tensor / sequence parallelism ---
     parallel_mode: Optional[str]
-    tp_group: Optional[Any]
+    # ProcessGroup is a *reference*-opaque type: carried through the torch.compile
+    # custom op as a graph input (never baked into the graph as a constant).
+    tp_group: Optional[dist_group_type]
     tp_size: int
     tensor_parallel: bool
     sequence_parallel: bool
@@ -159,6 +181,39 @@ class LinearFwdArgs:
     cpu_offloading: bool
     is_grad_enabled: bool
 
+    def compile_unsupported_reason(self) -> Optional[str]:
+        """Reason this config can't use the torch.compile custom-op path (else None)."""
+        if self.debug:
+            return "debug instrumentation (nvidia-dlfw-inspect)"
+        if self.fsdp_group is not None and self.is_grad_enabled:
+            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)
+        ):
+            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.fuse_wgrad_accumulation:
+            return "fuse_wgrad_accumulation (main_grad)"
+        for quantizer in (
+            self.input_quantizer,
+            self.weight_quantizer,
+            self.output_quantizer,
+            self.grad_input_quantizer,
+            self.grad_weight_quantizer,
+            self.grad_output_quantizer,
+        ):
+            # e.g. delayed-scaling Float8Quantizer and unregistered custom-recipe
+            # quantizers are not value-opaque and can't cross the custom-op boundary.
+            if quantizer is not None and not is_value_opaque_quantizer(quantizer):
+                return "a quantizer not registered as a torch.compile value-opaque type"
+        return None
+
 
 @dataclass(slots=True)
 class LinearBwdArgs:
@@ -196,7 +251,8 @@ class LinearBwdArgs:
 
     # --- Tensor / sequence parallelism ---
     parallel_mode: Optional[str] = None
-    tp_group: Optional[Any] = None
+    # Reference-opaque ProcessGroup (graph input), see LinearFwdArgs.tp_group.
+    tp_group: Optional[dist_group_type] = None
     tp_size: int = 1
     tensor_parallel: bool = False
     sequence_parallel: bool = False
@@ -296,9 +352,7 @@ def _linear_forward_impl(
     if ub_name is not None:
         nvtx_label = f"{nvtx_label}.{ub_name}"
 
-    # Make sure input dimensions are compatible
-    out_features, in_features = weight.shape
-    assert inp.shape[-1] == in_features, "GEMM not possible"
+    out_features = weight.shape[0]
 
     # Configure tensor-parallel communication
     tp_world_size = get_distributed_world_size(tp_group)
@@ -340,7 +394,6 @@ def _linear_forward_impl(
     inputmat_total = None  # Input tensor to pass to GEMM (gathered)
     own_quantized_input = False
     if fp8:
-        assert_dim_for_fp8_exec(inputmat, weight)
         if save_original_input:
             assert not isinstance(
                 input_quantizer, Float8Quantizer
@@ -887,7 +940,10 @@ def _linear_setup_ctx(
     bwd_args.use_bias = bias is not None
     bwd_args.requires_dgrad = fwd_args.input_requires_grad
     bwd_args.requires_wgrad = fwd_args.weight_requires_grad
-    bwd_args.inp_shape = inp.shape
+    # Don't store inp_shape in the value bundle: under torch.compile(dynamic=True)
+    # inp.shape contains SymInt dims which are not hashable in OpaqueValueBundle.
+    # The backward reconstructs inp_shape from grad_output + weight + SP config.
+    bwd_args.inp_shape = None
 
     # Numerical / dtype config
     bwd_args.activation_dtype = fwd_args.activation_dtype
@@ -1035,6 +1091,19 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ..
         )
         nvtx_range_pop(f"{nvtx_label}.fsdp_gather")
 
+        # Reconstruct inp_shape when not stored (compiled mode with dynamic shapes).
+        if bwd_args.inp_shape is None:
+            _w = saved_weight if saved_weight is not None else weight_fp8
+            in_features = _w.shape[-1]
+            go_leading = grad_output.shape[0]
+            if bwd_args.parallel_mode == "column" and bwd_args.sequence_parallel:
+                inp_leading = go_leading // bwd_args.tp_size
+            elif bwd_args.parallel_mode == "row" and bwd_args.sequence_parallel:
+                inp_leading = go_leading * bwd_args.tp_size
+            else:
+                inp_leading = go_leading
+            bwd_args.inp_shape = torch.Size([inp_leading, *grad_output.shape[1:-1], in_features])
+
         # Configure Userbuffers communication (comm+GEMM overlap)
         bwd_args.ub_obj_gradout = None
         ub_obj_dgrad = None
@@ -1574,8 +1643,19 @@ def _linear_backward_impl_fake(
     dgrad = None
     if args.requires_dgrad:
         # dgrad has the logical input shape and may be quantized for the next op.
+        # Derive shape from grad_output + weight + SP config instead of args.inp_shape:
+        # inp_shape is not stored in the value bundle under dynamic shapes (SymInt is
+        # not hashable in OpaqueValueBundle), so we reconstruct it here.
+        _in_features = weight.shape[-1]
+        _go_leading = args.grad_output.shape[0]
+        if args.parallel_mode == "column" and args.sequence_parallel:
+            _dgrad_leading = _go_leading // args.tp_size
+        elif args.parallel_mode == "row" and args.sequence_parallel:
+            _dgrad_leading = _go_leading * args.tp_size
+        else:
+            _dgrad_leading = _go_leading
         dgrad = TensorProto(
-            shape=tuple(args.inp_shape),
+            shape=(_dgrad_leading, *args.grad_output.shape[1:-1], _in_features),
             dtype=out_dtype,
             quantizer=args.grad_input_quantizer,
             device=args.grad_output.device,
@@ -1603,6 +1683,21 @@ def _linear_backward_impl_fake(
     return wgrad, dgrad, grad_bias
 
 
+# Custom op used under ``torch.compile``.
+_linear_op = register_custom_op(
+    op_name="linear",
+    input_tensors_for_grad=["weight", "inp", "bias"],
+    fwd_arg_type=LinearFwdArgs,
+    fwd_impl=_linear_forward_impl,
+    fwd_fake_impl=_linear_forward_impl_fake,
+    setup_context=_linear_setup_ctx,
+    backward_arg_type=LinearBwdArgs,
+    backward_obj=LinearBwdArgs,
+    backward_impl=_linear_backward,
+    bwd_fake_impl=_linear_backward_impl_fake,
+)
+
+
 class _Linear(torch.autograd.Function):
     """Linear semi-top level module
     Calls custom cuda extensions.
@@ -1687,6 +1782,20 @@ def backward(
         return result
 
 
+@no_torch_dynamo()
+def _linear_eager(
+    weight_tensor: torch.Tensor,
+    inp: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    fwd_args: LinearFwdArgs,
+    is_grad_enabled: bool,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+    """Run ``_Linear`` eagerly, bypassing Dynamo."""
+    if is_grad_enabled:
+        return _Linear.apply(weight_tensor, inp, bias, fwd_args)
+    return _Linear.forward(None, weight_tensor, inp, bias, fwd_args)
+
+
 class Linear(TransformerEngineBaseModule):
     """Applies a linear transformation to the incoming data :math:`y = xA^T + b`
 
@@ -2084,7 +2193,6 @@ def reset_parameters(self, defer_init=False):
                     elif self.parallel_mode == "column":
                         set_tensor_model_parallel_attributes(getattr(self, bias), True, 0, 1)
 
-    @no_torch_dynamo()
     def forward(
         self,
         inp: torch.Tensor,
@@ -2159,12 +2267,7 @@ def forward(
                 grad_output_quantizer,
             ) = quantizers
 
-            if is_grad_enabled:
-                linear_fn = _Linear.apply
-                autograd_ctx = []
-            else:
-                linear_fn = _Linear.forward
-                autograd_ctx = [None]
+            use_compiled_op = torch.compiler.is_compiling() and _linear_op is not None
 
             cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight"
             weight_workspace = (
@@ -2204,16 +2307,58 @@ def forward(
                 ub_bulk_dgrad = self.ub_bulk_dgrad
                 ub_bulk_wgrad = self.ub_bulk_wgrad
 
+            torch._check(
+                inp.shape[-1] == weight_tensor.shape[-1],
+                lambda: "GEMM not possible: input last dim must equal in_features",
+            )
+            if self.fp8:
+                torch._check(
+                    math.prod(inp.shape[:-1]) % 8 == 0,
+                    lambda: (
+                        "FP8 execution requires the product of all input dimensions except"
+                        " the last to be divisible by 8"
+                    ),
+                )
+                torch._check(
+                    inp.shape[-1] % 16 == 0,
+                    lambda: "FP8 execution requires the input last dimension to be divisible by 16",
+                )
+                torch._check(
+                    weight_tensor.shape[0] % 16 == 0,
+                    lambda: "FP8 execution requires out_features to be divisible by 16",
+                )
+                torch._check(
+                    weight_tensor.shape[1] % 16 == 0,
+                    lambda: "FP8 execution requires in_features to be divisible by 16",
+                )
+
             linear_bias_tensor = (
                 bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None
             )
             wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None
+
+            # Pin the lazily-cached cuBLAS (and NVFP4-RHT) workspaces as op inputs so
+            # they are materialized at trace time (external to the cudagraph pool)
+            # rather than inside the op during capture. See LinearFwdArgs for details.
+            cublas_workspace = None
+            rht_matrix = None
+            if use_compiled_op:
+                cublas_workspace = get_cublas_workspace(inp.device.index, False, False)
+                from ..tensor.nvfp4_tensor import NVFP4Quantizer, get_rht_matrix
+
+                if isinstance(input_quantizer, NVFP4Quantizer):
+                    rht_matrix = get_rht_matrix(
+                        input_quantizer._with_random_sign_mask, inp.device.index
+                    )
+
             fwd_args = LinearFwdArgs(
                 # tensors
                 weight=weight_tensor,
                 inp=inp,
                 bias=linear_bias_tensor,
                 weight_workspace=weight_workspace,
+                cublas_workspace=cublas_workspace,
+                rht_matrix=rht_matrix,
                 # requires_grad flags
                 input_requires_grad=inp.requires_grad,
                 weight_requires_grad=weight_tensor.requires_grad,
@@ -2268,13 +2413,19 @@ def forward(
                 cpu_offloading=is_cpu_offload_enabled(),
                 is_grad_enabled=is_grad_enabled,
             )
-            out, new_weight_workspace = linear_fn(
-                *autograd_ctx,
-                weight_tensor,
-                inp,
-                linear_bias_tensor,
-                fwd_args,
-            )
+
+            if use_compiled_op:
+                fallback_reason = fwd_args.compile_unsupported_reason()
+                if fallback_reason is not None:
+                    warn_compile_eager_fallback(fallback_reason)
+                    use_compiled_op = False
+
+            if use_compiled_op:
+                out, new_weight_workspace = _linear_op(fwd_args)
+            else:
+                out, new_weight_workspace = _linear_eager(
+                    weight_tensor, inp, linear_bias_tensor, fwd_args, is_grad_enabled
+                )
 
             if new_weight_workspace is not None and cache_name is not None:
                 if isinstance(new_weight_workspace, torch.Tensor):
diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py
index ffbfbc1fdd..8eeaf2417d 100644
--- a/transformer_engine/pytorch/utils.py
+++ b/transformer_engine/pytorch/utils.py
@@ -26,6 +26,20 @@
 ]
 
 
+def warn_compile_eager_fallback(reason: str) -> None:
+    """Warn that a TE module is running eagerly under ``torch.compile``.
+
+    Emitted when ``reason`` is unsupported on the module's compiled custom-op
+    path. Python's default warning filter dedups identical messages, so each
+    distinct ``reason`` is surfaced once.
+    """
+    warnings.warn(
+        f"Falling back to eager execution under torch.compile: {reason} is "
+        "unsupported on the compiled path (graph-breaks under fullgraph=True).",
+        stacklevel=2,
+    )
+
+
 @functools.lru_cache(maxsize=None)
 def get_cached_ones_tensor(
     num_elements: int,

From fdac6597edbca38fdcd8375af2d308596cafe809 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Mon, 22 Jun 2026 12:48:26 +0200
Subject: [PATCH 24/28] [PyTorch] torch.compile: register TE custom ops via
 torch.library.custom_op

Replace the low-level torch.library.Library (_TE_LIB.define/.impl + functional register_fake/register_autograd/register_torch_dispatch with lib=) with the standard torch.library.custom_op API, passing the dynamically built schema explicitly via schema=. register_fake/register_autograd/register_torch_dispatch are now methods on the returned CustomOpDef. Drops the TOR901 Library usage and is robust to re-registration (get_library_allowing_overwrite).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
(cherry picked from commit d590560c9a2ed26072ea48b0c8e3934ad62b9a29)
---
 .../pytorch/dynamo/custom_op.py               | 82 +++++++++----------
 1 file changed, 41 insertions(+), 41 deletions(-)

diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index ef943d4e8e..2b1b34643c 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -35,7 +35,6 @@
 )
 
 _TE_OP_NAMESPACE = "transformer_engine_compile"
-_TE_LIB = torch.library.Library(_TE_OP_NAMESPACE, "FRAGMENT")  # noqa: TOR901
 
 
 # ``None`` entries in an op's flat ``Tensor[]`` return are smuggled through a
@@ -939,7 +938,7 @@ def _resolve_grad_targets(
 def _register_kernel(
     *,
     op_name: str,
-    op_qualname: str,
+    schema_str: str,
     arg_type: type,
     arg_names: List[str],
     buckets: List[_Bucket],
@@ -947,8 +946,9 @@ def _register_kernel(
     impl: Callable[[Any], Any],
     fake_impl: Callable[[Any], Any],
     format_result: Callable[[Any], List[torch.Tensor]],
-) -> None:
-    """Wire the real ``impl`` + the ``fake_impl`` (proto) into the library.
+) -> Any:
+    """Define the op via ``torch.library.custom_op`` with the real ``impl`` + the
+    ``fake_impl`` (proto), returning the ``CustomOpDef``.
 
     The real kernel rebuilds the dataclass and runs ``impl``; the fake kernel
     runs the proto fake impl on the :func:`_proto_view`. Both go through
@@ -966,13 +966,16 @@ def _fake(*flat: Any) -> List[torch.Tensor]:
         proto_obj = _proto_view(obj, tensor_field_names)
         return format_result(fake_impl(proto_obj))
 
-    _TE_LIB.impl(op_name, _impl, "CompositeExplicitAutograd")
-    torch.library.register_fake(op_qualname, _fake, lib=_TE_LIB)
+    op = torch.library.custom_op(
+        f"{_TE_OP_NAMESPACE}::{op_name}", _impl, mutates_args=(), schema=schema_str
+    )
+    op.register_fake(_fake)
+    return op
 
 
 def _register_autograd_for_op(
     *,
-    fwd_op_name: str,
+    fwd_op: Any,
     bwd_op_name: str,
     fwd_arg_type: type,
     fwd_arg_names: List[str],
@@ -992,7 +995,6 @@ def _register_autograd_for_op(
     templates, reassembles each flat output chunk, and hands the saved tuple +
     ``ctx_attrs`` to the module's ``setup_context``.
     """
-    fwd_qualname = f"{_TE_OP_NAMESPACE}::{fwd_op_name}"
 
     def _setup_context(ctx, inputs, output):
         ctx._te_fwd_tensor_list_lengths = {
@@ -1054,12 +1056,7 @@ def _autograd_backward(ctx, *grad_outputs):
             out[pos] = g
         return tuple(out)
 
-    torch.library.register_autograd(
-        fwd_qualname,
-        _autograd_backward,
-        setup_context=_setup_context,
-        lib=_TE_LIB,
-    )
+    fwd_op.register_autograd(_autograd_backward, setup_context=_setup_context)
 
 
 def _collect_universal_slot_offsets(buckets: List[_Bucket]) -> List[int]:
@@ -1093,12 +1090,14 @@ def _flatten_subclass_into_slots(
 def _register_outer_forwarder(
     *,
     outer_op_name: str,
+    schema_str: str,
     inner_op_name: str,
     buckets: Optional[List[_Bucket]] = None,
     subclass_list: Optional[List[type]] = None,
-) -> None:
-    """Register the outer op's default kernel + fake: forward to the inner op,
-    optionally flattening registered subclass inputs in place first.
+) -> Any:
+    """Define the outer op via ``torch.library.custom_op``: forward to the inner
+    op, optionally flattening registered subclass inputs in place first. Returns
+    the ``CustomOpDef``.
     """
     inner_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_op_name)
     input_flatten_enabled = bool(subclass_list) and buckets is not None
@@ -1112,8 +1111,11 @@ def _forward(*flat: Any) -> List[torch.Tensor]:
             _flatten_subclass_into_slots(new_args, slot_offsets, sub)
         return inner_op(*new_args)
 
-    _TE_LIB.impl(outer_op_name, _forward, "CompositeExplicitAutograd")
-    torch.library.register_fake(f"{_TE_OP_NAMESPACE}::{outer_op_name}", _forward, lib=_TE_LIB)
+    op = torch.library.custom_op(
+        f"{_TE_OP_NAMESPACE}::{outer_op_name}", _forward, mutates_args=(), schema=schema_str
+    )
+    op.register_fake(_forward)
+    return op
 
 
 def _all_quantized_tensor_subclasses() -> List[type]:
@@ -1208,19 +1210,14 @@ def _register_custom_op_impl(
         fwd_buckets, fwd_arg_type, input_tensors_for_grad
     )
 
-    _TE_LIB.define(f"{inner_fwd_name}{fwd_schema_args} -> Tensor[]")
-    _TE_LIB.define(f"{inner_bwd_name}{bwd_schema_args} -> Tensor[]")
-    _TE_LIB.define(f"{outer_fwd_name}{fwd_schema_args} -> Tensor[]")
-    _TE_LIB.define(f"{outer_bwd_name}{bwd_schema_args} -> Tensor[]")
+    fwd_schema = f"{fwd_schema_args} -> Tensor[]"
+    bwd_schema = f"{bwd_schema_args} -> Tensor[]"
 
-    inner_fwd_qualname = f"{_TE_OP_NAMESPACE}::{inner_fwd_name}"
     inner_bwd_qualname = f"{_TE_OP_NAMESPACE}::{inner_bwd_name}"
-    outer_fwd_qualname = f"{_TE_OP_NAMESPACE}::{outer_fwd_name}"
-    outer_bwd_qualname = f"{_TE_OP_NAMESPACE}::{outer_bwd_name}"
 
-    _register_kernel(
+    inner_fwd_def = _register_kernel(
         op_name=inner_fwd_name,
-        op_qualname=inner_fwd_qualname,
+        schema_str=fwd_schema,
         arg_type=fwd_arg_type,
         arg_names=fwd_arg_names,
         buckets=fwd_buckets,
@@ -1231,7 +1228,7 @@ def _register_custom_op_impl(
     )
     _register_kernel(
         op_name=inner_bwd_name,
-        op_qualname=inner_bwd_qualname,
+        schema_str=bwd_schema,
         arg_type=backward_arg_type,
         arg_names=bwd_arg_names,
         buckets=bwd_buckets,
@@ -1241,6 +1238,17 @@ def _register_custom_op_impl(
         format_result=lambda g: _format_bwd_result(g, num_grad_inputs, inner_bwd_qualname),
     )
 
+    outer_fwd_def = _register_outer_forwarder(
+        outer_op_name=outer_fwd_name,
+        schema_str=fwd_schema,
+        inner_op_name=inner_fwd_name,
+        buckets=fwd_buckets,
+        subclass_list=list(subclass_list),
+    )
+    outer_bwd_def = _register_outer_forwarder(
+        outer_op_name=outer_bwd_name, schema_str=bwd_schema, inner_op_name=inner_bwd_name
+    )
+
     autograd_common = {
         "fwd_arg_type": fwd_arg_type,
         "fwd_arg_names": fwd_arg_names,
@@ -1255,19 +1263,11 @@ def _register_custom_op_impl(
         "fwd_fake_impl": fwd_fake_impl,
     }
     _register_autograd_for_op(
-        fwd_op_name=inner_fwd_name, bwd_op_name=inner_bwd_name, **autograd_common
+        fwd_op=inner_fwd_def, bwd_op_name=inner_bwd_name, **autograd_common
     )
     _register_autograd_for_op(
-        fwd_op_name=outer_fwd_name, bwd_op_name=outer_bwd_name, **autograd_common
-    )
-
-    _register_outer_forwarder(
-        outer_op_name=outer_fwd_name,
-        inner_op_name=inner_fwd_name,
-        buckets=fwd_buckets,
-        subclass_list=list(subclass_list),
+        fwd_op=outer_fwd_def, bwd_op_name=outer_bwd_name, **autograd_common
     )
-    _register_outer_forwarder(outer_op_name=outer_bwd_name, inner_op_name=inner_bwd_name)
 
     inner_fwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_fwd_name)
     inner_bwd_op = getattr(getattr(torch.ops, _TE_OP_NAMESPACE), inner_bwd_name)
@@ -1292,8 +1292,8 @@ def _bwd_rule(mode, func, types, args, kwargs):
         return inner_bwd_op(*new_args)
 
     for sub in subclass_list:
-        torch.library.register_torch_dispatch(outer_fwd_qualname, sub, _fwd_rule, lib=_TE_LIB)
-        torch.library.register_torch_dispatch(outer_bwd_qualname, sub, _bwd_rule, lib=_TE_LIB)
+        outer_fwd_def.register_torch_dispatch(sub, _fwd_rule)
+        outer_bwd_def.register_torch_dispatch(sub, _bwd_rule)
 
     _quantized_tensor_passthrough_ops.add(outer_fwd_op.default)
     _quantized_tensor_passthrough_ops.add(outer_bwd_op.default)

From 554f5a8fcd0f8316e5243fde4449307f87e70ef6 Mon Sep 17 00:00:00 2001
From: Pawel Gadzinski <pgadzinski@nvidia.com>
Date: Tue, 30 Jun 2026 14:06:46 +0200
Subject: [PATCH 25/28] [PyTorch] custom_op: note pytorch/pytorch#187434
 enables dropping None sentinel

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
---
 transformer_engine/pytorch/dynamo/custom_op.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 2b1b34643c..68b2aae59e 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -40,6 +40,10 @@
 # ``None`` entries in an op's flat ``Tensor[]`` return are smuggled through a
 # 0-element uint8 tensor: a non-nullable ``Tensor[]`` schema is required for
 # ``register_autograd`` to attach a ``grad_fn`` to the outputs.
+#
+# TODO: once https://github.com/pytorch/pytorch/pull/187434 lands, a nullable
+# ``Tensor?[]`` return schema lets ``None`` pass through directly and this
+# sentinel encoding (``_encode_none`` / ``_decode_none``) can be removed.
 _NONE_SENTINEL_DTYPE = torch.uint8
 
 

From b24d259b0cd9a68ca29b9c1a881b71d028a8aa15 Mon Sep 17 00:00:00 2001
From: kshitij12345 <kshitij12345@users.noreply.github.com>
Date: Wed, 1 Jul 2026 13:47:40 -0700
Subject: [PATCH 26/28] [PyTorch][torch.compile] Replace TensorProto with
 make_empty_traceable

Replace the 172-line TensorProto dataclass with a single function
make_empty_traceable(quantizer, shape, dtype, device) that directly allocates
traceable quantized tensors. The fake impls now return actual tensors (which
become FakeTensors under register_fake) instead of intermediate descriptors.

The key insight: make_empty_traceable stashes _te_flat_names and _te_flat_ctx
on the resulting tensor. Dynamo treats non-callable attributes on traceable
wrapper subclasses as constant metadata, so forward_fn can read slot counts
and reassembly info from these attributes without calling __tensor_flatten__
(which would cause a graph break since it returns non-Tensor Python objects).

This eliminates:
- TensorProto class and to_tensor_proto helper (tensor_proto.py deleted)
- _proto_view (converted tensor fields to TensorProto before fake impls)
- _tensor_field_names (identified fields for _proto_view)
- _proto_slot_count / _proto_reassemble (operated on TensorProto objects)
- TensorProto branch in _value_to_flat_tensors and _format_bwd_result

The fake impls in linear.py now use:
- isinstance(inp, QuantizedTensorStorage) instead of inp.is_quantized
- weight._quantizer instead of weight.quantizer (TensorProto field)
- make_empty_traceable(...) instead of TensorProto(...)
- Direct set_usage on quantizer instead of proto.update_usage()

Test Plan:

```
python -m pytest tests/pytorch/test_torch_compile.py -v -k 'not nvfp4'
```

Authored with Claude.
---
 tests/pytorch/test_torch_compile.py           | 120 +++++-------
 transformer_engine/pytorch/dynamo/__init__.py |   5 +-
 .../pytorch/dynamo/custom_op.py               | 122 +++---------
 .../pytorch/dynamo/tensor_proto.py            | 177 ------------------
 .../pytorch/dynamo/traceable_utils.py         | 133 +++++++++++++
 transformer_engine/pytorch/module/linear.py   | 100 +++++-----
 .../pytorch/quantized_tensor.py               |   8 +-
 .../pytorch/tensor/float8_blockwise_tensor.py |   2 +-
 .../pytorch/tensor/float8_tensor.py           |   2 +-
 .../pytorch/tensor/mxfp8_tensor.py            |   2 +-
 .../pytorch/tensor/nvfp4_tensor.py            |   2 +-
 11 files changed, 264 insertions(+), 409 deletions(-)
 delete mode 100644 transformer_engine/pytorch/dynamo/tensor_proto.py
 create mode 100644 transformer_engine/pytorch/dynamo/traceable_utils.py

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index a0501cca8e..716947026f 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -32,7 +32,7 @@
 from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer
 from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer
 from transformer_engine.pytorch.quantized_tensor import QuantizedTensor, _STORAGE_REGISTRY
-from transformer_engine.pytorch.dynamo import TensorProto, to_tensor_proto
+from transformer_engine.pytorch.dynamo.traceable_utils import make_empty_traceable, _contiguous_stride
 from transformer_engine.pytorch import (
     is_fp8_available,
     is_mxfp8_available,
@@ -640,7 +640,7 @@ def fn(inp):
 
 
 # ---------------------------------------------------------------------------
-# torch.compile-traceable allocation primitives + TensorProto
+# torch.compile-traceable allocation primitives
 # ---------------------------------------------------------------------------
 
 
@@ -665,9 +665,7 @@ def fn(inp):
 def _build_from_primitives(quantizer, shape, dtype, device="cpu"):
     """Assemble a quantized tensor straight from the quantizer primitives:
     ``alloc_tensors`` (buffers) + ``create_metadata`` (ctx) + the storage's
-    ``__tensor_unflatten__`` -- i.e. exactly what ``TensorProto.create_tensor``
-    does, but without going through :class:`TensorProto`.
-    """
+    ``__tensor_unflatten__``."""
     names = tuple(quantizer._describe_buffers(shape))  # pylint: disable=protected-access
     ctx = quantizer.create_metadata(shape, dtype=dtype)
     buffers = quantizer.alloc_tensors(shape, device=device)
@@ -713,7 +711,7 @@ def _skip_if_dequantize_unsupported(q):
 @pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
 def test_primitives_unflatten_compiles(factory, shape):
     """create_metadata + alloc_tensors + __tensor_unflatten__ compose and trace
-    under ``fullgraph=True`` (CPU), without TensorProto."""
+    under ``fullgraph=True`` (CPU)."""
     q = factory()
     names = tuple(q._describe_buffers(shape))  # pylint: disable=protected-access
 
@@ -779,78 +777,64 @@ def test_storage_flatten_unflatten_roundtrip(factory, shape):
     torch.testing.assert_close(rebuilt.dequantize(), expected, atol=0, rtol=0, equal_nan=True)
 
 
-# ----- TensorProto -----
+# ----- make_empty_traceable -----
 
 
 @pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
-def test_tensor_proto_matches_primitives(factory, shape):
-    """TensorProto is a thin wrapper: its ``create_metadata`` /
-    ``create_inner_tensors`` / ``create_tensor`` match building everything
-    directly from the quantizer primitives."""
+def test_make_empty_traceable_matches_primitives(factory, shape):
+    """make_empty_traceable is equivalent to building from quantizer primitives
+    directly (alloc_tensors + create_metadata + __tensor_unflatten__)."""
     q = factory()
-    proto = TensorProto(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu"))
-    assert proto.is_quantized
 
-    # Metadata matches the quantizer's.
-    assert proto.create_metadata() == q.create_metadata(shape, dtype=torch.bfloat16)
+    # Build via make_empty_traceable.
+    tensor = make_empty_traceable(q, shape, dtype=torch.bfloat16, device="cpu")
 
-    # inner_names follows the storage's canonical __tensor_flatten__ order (the
-    # order the real op flattens its outputs to), while create_inner_tensors
-    # matches the _describe_buffers geometry (a name->shape/dtype mapping).
-    bufs = q._describe_buffers(shape)  # pylint: disable=protected-access
+    # Build directly from primitives.
     direct = _build_from_primitives(q, shape, torch.bfloat16)
-    names = tuple(direct.__tensor_flatten__()[0])
-    assert set(names) == set(bufs)
-    assert proto.inner_names() == names
-    inner = proto.create_inner_tensors()
-    assert len(inner) == len(names)
-    for name, buf in zip(names, inner):
-        exp_shape, exp_dtype = bufs[name]
-        assert tuple(buf.shape) == tuple(exp_shape)
-        assert buf.dtype == exp_dtype
 
-    # The assembled tensor matches one built directly from the primitives.
-    assert _signature(proto.create_tensor(), names) == _signature(direct, names)
+    # Both should produce the same buffer layout.
+    names = tuple(q._describe_buffers(shape))  # pylint: disable=protected-access
+    assert _signature(tensor, names) == _signature(direct, names)
 
 
 @pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
-def test_tensor_proto_create_tensor_eager(factory, shape):
-    """``create_tensor`` (no fake) yields a real quantized tensor."""
+def test_make_empty_traceable_eager(factory, shape):
+    """make_empty_traceable (no fake) yields a real quantized tensor."""
     q = factory()
-    proto = TensorProto(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu"))
-    out = proto.create_tensor()
+    out = make_empty_traceable(q, shape, dtype=torch.bfloat16, device="cpu")
     assert isinstance(out, QuantizedTensor)
     assert tuple(out.shape) == tuple(shape)
     assert out.dtype == torch.bfloat16
-    for name in proto.inner_names():
+    names = tuple(q._describe_buffers(shape))  # pylint: disable=protected-access
+    for name in names:
         assert not isinstance(getattr(out, name), FakeTensor)
 
 
 @pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
-def test_tensor_proto_create_tensor_fake(factory, shape):
-    """``create_tensor`` under ``FakeTensorMode`` yields a fake-backed quantized
+def test_make_empty_traceable_fake(factory, shape):
+    """make_empty_traceable under FakeTensorMode yields a fake-backed quantized
     tensor with the right shape/dtype and fake inner buffers."""
     q = factory()
-    proto = TensorProto(shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu"))
     with FakeTensorMode():
-        out = proto.create_tensor()
+        out = make_empty_traceable(q, shape, dtype=torch.bfloat16, device="cpu")
     assert isinstance(out, QuantizedTensor)
     assert tuple(out.shape) == tuple(shape)
     assert out.dtype == torch.bfloat16
-    for name in proto.inner_names():
+    names = tuple(q._describe_buffers(shape))  # pylint: disable=protected-access
+    for name in names:
         assert isinstance(getattr(out, name), FakeTensor)
 
 
 @pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
-def test_tensor_proto_create_tensor_compiles(factory, shape):
-    """``TensorProto.create_tensor`` traces under ``fullgraph=True`` (CPU)."""
+def test_make_empty_traceable_compiles(factory, shape):
+    """make_empty_traceable traces under torch.compile(fullgraph=True) (CPU)."""
     q = factory()
+    names = tuple(q._describe_buffers(shape))  # pylint: disable=protected-access
 
     def fn(x):
-        proto = TensorProto(shape=tuple(x.shape), dtype=x.dtype, quantizer=q, device=x.device)
-        t = proto.create_tensor()
+        t = make_empty_traceable(q, tuple(x.shape), dtype=x.dtype, device=x.device)
         acc = x.new_zeros(())
-        for name in proto.inner_names():
+        for name in names:
             acc = acc + getattr(t, name).float().sum()
         return acc
 
@@ -860,36 +844,32 @@ def fn(x):
     assert out.shape == ()
 
 
-def test_to_tensor_proto_plain():
-    """``to_tensor_proto`` describes a plain tensor."""
-    t = torch.empty(2, 3, dtype=torch.float32)
-    proto = to_tensor_proto(t)
-    assert not proto.is_quantized
-    assert proto.shape == (2, 3)
-    assert proto.dtype == torch.float32
-    assert proto.inner_names() == ("data",)
+def test_make_empty_traceable_plain_tensor():
+    """For non-quantized tensors, make_empty_traceable produces a plain tensor."""
+    t = make_empty_traceable(None, (2, 3), dtype=torch.float32, device="cpu")
+    assert not isinstance(t, QuantizedTensor)
+    assert t.shape == (2, 3)
+    assert t.dtype == torch.float32
 
 
 @pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
-def test_to_tensor_proto_quantized(factory, shape):
-    """``to_tensor_proto`` round-trips a quantized tensor back into a proto."""
+def test_make_empty_traceable_roundtrip(factory, shape):
+    """A tensor built via make_empty_traceable can be flattened and unflattened."""
     q = factory()
-    tensor = TensorProto(
-        shape=shape, dtype=torch.bfloat16, quantizer=q, device=torch.device("cpu")
-    ).create_tensor()
-
-    proto = to_tensor_proto(tensor)
-    assert proto.is_quantized
-    assert proto.shape == tuple(shape)
-    assert proto.dtype == torch.bfloat16
-    # Same buffer layout as the original tensor.
-    assert proto.inner_names() == tuple(
-        q._describe_buffers(shape)
-    )  # pylint: disable=protected-access
-    # Rebuilding from the derived proto matches the original tensor's structure.
-    assert _signature(proto.create_tensor(), proto.inner_names()) == _signature(
-        tensor, proto.inner_names()
+    tensor = make_empty_traceable(q, shape, dtype=torch.bfloat16, device="cpu")
+    names = tuple(q._describe_buffers(shape))  # pylint: disable=protected-access
+
+    # Flatten.
+    flat_names, flat_ctx = tensor.__tensor_flatten__()
+    assert set(flat_names) == set(names)
+    inner = {name: getattr(tensor, name) for name in flat_names}
+
+    # Unflatten.
+    rebuilt = type(tensor).__tensor_unflatten__(
+        inner, flat_ctx, tuple(tensor.shape), tensor.stride()
     )
+    assert isinstance(rebuilt, QuantizedTensor)
+    assert _signature(rebuilt, flat_names) == _signature(tensor, names)
 
 
 # ---------------------------------------------------------------------------
diff --git a/transformer_engine/pytorch/dynamo/__init__.py b/transformer_engine/pytorch/dynamo/__init__.py
index 66e0525a9e..f4864cda5b 100644
--- a/transformer_engine/pytorch/dynamo/__init__.py
+++ b/transformer_engine/pytorch/dynamo/__init__.py
@@ -5,13 +5,12 @@
 """torch.compile glue for Transformer Engine."""
 
 from .quantizer_opaque import register_value_opaque_quantizer, is_value_opaque_quantizer
-from .tensor_proto import TensorProto, to_tensor_proto
+from .traceable_utils import make_empty_traceable
 from .custom_op import register_custom_op
 
 __all__ = [
     "register_value_opaque_quantizer",
     "is_value_opaque_quantizer",
-    "TensorProto",
-    "to_tensor_proto",
+    "make_empty_traceable",
     "register_custom_op",
 ]
diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py
index 68b2aae59e..9aca5b866f 100644
--- a/transformer_engine/pytorch/dynamo/custom_op.py
+++ b/transformer_engine/pytorch/dynamo/custom_op.py
@@ -14,7 +14,6 @@
     Dict,
     List,
     Optional,
-    Sequence,
     Tuple,
     Union,
     get_args,
@@ -24,7 +23,7 @@
 
 import torch
 
-from .tensor_proto import TensorProto, to_tensor_proto, _contiguous_stride
+from .traceable_utils import _contiguous_stride, _slot_count, _maybe_reassemble_tensor_subclass
 from ..quantized_tensor import (
     QuantizedTensor,
     QuantizedTensorStorage,
@@ -723,11 +722,6 @@ def _get_buckets(cls: type) -> List[_Bucket]:
     return buckets
 
 
-def _tensor_field_names(buckets: List[_Bucket]) -> List[str]:
-    """Names of fields carrying tensors (for building the proto view)."""
-    return [b.name for b in buckets if isinstance(b, (_TensorBucket, _UniversalTensorBucket))]
-
-
 def _build_schema(buckets: List[_Bucket]) -> Tuple[str, List[str]]:
     """Return ``(schema_arg_str, slot_names)`` for a bucket list."""
     spec = [slot for b in buckets for slot in b.schema_slots()]
@@ -761,76 +755,25 @@ def _unpack(cls: type, args: Dict[str, Any], buckets: List[_Bucket]) -> Any:
         object.__setattr__(obj, k, v)
     return obj
 
-
-def _proto_view(obj: Any, tensor_field_names: Sequence[str]) -> Any:
-    """Copy of dataclass ``obj`` with each tensor field replaced by a :class:`TensorProto`.
-
-    Only tensor fields have a ``TensorProto`` equivalent, so quantizer / scalar
-    fields are simply carried over unchanged; the fake impl works purely on
-    geometry. Built with :func:`dataclasses.replace` (the only such construction
-    Dynamo can trace).
-    """
-    overrides: Dict[str, Any] = {}
-    for name in tensor_field_names:
-        value = getattr(obj, name, None)
-        if value is not None and not isinstance(value, TensorProto):
-            overrides[name] = to_tensor_proto(value)
-    if not overrides:
-        return obj
-    return dataclasses.replace(obj, **overrides)
-
-
 # --------------------------------------------------------------------------- #
 # Op outputs <-> flat ``Tensor[]`` payload: this is how an op returns / saves
 # quantized tensors (and wrapper subclasses). Outputs are flattened to their
 # inner buffers on the way out and rebuilt via ``__tensor_unflatten__`` on the
-# way back; on the fake side a TensorProto supplies the geometry.
+# way back; the fake impl returns actual tensors whose __tensor_flatten__
+# provides the template for reassembly.
 # --------------------------------------------------------------------------- #
 
 
-def _proto_slot_count(proto: Optional[TensorProto]) -> int:
-    """Flat ``Tensor[]`` slots the value for ``proto`` occupies."""
-    if proto is None:
-        return 1
-    return len(proto.inner_names())
-
-
-def _proto_reassemble(
-    proto: Optional[TensorProto],
-    chunk: List[Optional[torch.Tensor]],
-) -> Optional[Union[torch.Tensor, QuantizedTensorStorage]]:
-    """Rebuild the value described by ``proto`` from its flat tensors ``chunk``.
-
-    ``proto`` describes one output: ``None`` (-> ``None``), a plain tensor
-    (``chunk`` is the single tensor, returned as-is), or a quantized tensor
-    (``chunk`` are its inner tensors, reassembled into the wrapper subclass via
-    ``__tensor_unflatten__``).
-    """
-    if proto is None:
-        return None
-    if proto.quantizer is None:
-        return chunk[0]
-    inner_names = proto.inner_names()
-    meta = proto.create_metadata()
-    shape = tuple(proto.shape)
-    stride = _contiguous_stride(shape)
-    storage_cls = _STORAGE_REGISTRY[meta["cls"]]
-    inner_dict = dict(zip(inner_names, chunk))
-    return storage_cls.__tensor_unflatten__(inner_dict, meta, shape, stride)
-
-
 def _value_to_flat_tensors(
-    value: Optional[Union[torch.Tensor, QuantizedTensorStorage, TensorProto]],
+    value: Optional[Union[torch.Tensor, QuantizedTensorStorage]],
 ) -> List[torch.Tensor]:
     """Return the flat ``Tensor[]`` slots that represent one op output ``value``.
 
-    Inverse of :func:`_proto_reassemble`; the slot count matches
-    :func:`_proto_slot_count`.
+    Inverse of :func:`_maybe_reassemble_tensor_subclass`; the slot count matches
+    :func:`_slot_count`.
     """
     if value is None:
         return [_encode_none(None)]
-    if isinstance(value, TensorProto):
-        return [_encode_none(t) for t in value.create_inner_tensors()]
     if isinstance(value, torch.Tensor):
         if type(value) is not torch.Tensor and hasattr(  # pylint: disable=unidiomatic-typecheck
             value, "__tensor_flatten__"
@@ -843,7 +786,7 @@ def _value_to_flat_tensors(
         return [_encode_none(getattr(value, n)) for n in inner_names]
     raise TypeError(
         f"unsupported value type {type(value).__name__}; expected None / "
-        "torch.Tensor / tensor subclass / bare storage / TensorProto."
+        "torch.Tensor / tensor subclass / bare storage."
     )
 
 
@@ -870,8 +813,7 @@ def _format_fwd_result(result: Any) -> List[torch.Tensor]:
 def _format_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:`TensorProto` grad is materialized into a single tensor.
+    Each grad occupies exactly one slot (validated against ``num_grad_inputs``).
     """
     grads = list(grads)
     if len(grads) != num_grad_inputs:
@@ -879,13 +821,7 @@ def _format_bwd_result(grads: Any, num_grad_inputs: int, op_qualname: str) -> Li
             f"{op_qualname} expected backward_impl to return {num_grad_inputs} grads "
             f"(one per input_tensors_for_grad entry), got {len(grads)}"
         )
-    out: List[torch.Tensor] = []
-    for g in grads:
-        if isinstance(g, TensorProto):
-            out.append(_encode_none(g.create_tensor()))
-        else:
-            out.append(_encode_none(g))
-    return out
+    return [_encode_none(g) for g in grads]
 
 
 def _split_fwd_fake_result(
@@ -946,16 +882,15 @@ def _register_kernel(
     arg_type: type,
     arg_names: List[str],
     buckets: List[_Bucket],
-    tensor_field_names: List[str],
     impl: Callable[[Any], Any],
     fake_impl: Callable[[Any], Any],
     format_result: Callable[[Any], List[torch.Tensor]],
 ) -> Any:
     """Define the op via ``torch.library.custom_op`` with the real ``impl`` + the
-    ``fake_impl`` (proto), returning the ``CustomOpDef``.
+    ``fake_impl``, returning the ``CustomOpDef``.
 
     The real kernel rebuilds the dataclass and runs ``impl``; the fake kernel
-    runs the proto fake impl on the :func:`_proto_view`. Both go through
+    runs the fake impl directly on the unpacked object. Both go through
     ``format_result``.
     """
 
@@ -967,8 +902,7 @@ def _impl(*flat: Any) -> List[torch.Tensor]:
     def _fake(*flat: Any) -> List[torch.Tensor]:
         kwargs = dict(zip(arg_names, flat))
         obj = _unpack(arg_type, kwargs, buckets)
-        proto_obj = _proto_view(obj, tensor_field_names)
-        return format_result(fake_impl(proto_obj))
+        return format_result(fake_impl(obj))
 
     op = torch.library.custom_op(
         f"{_TE_OP_NAMESPACE}::{op_name}", _impl, mutates_args=(), schema=schema_str
@@ -984,7 +918,6 @@ def _register_autograd_for_op(
     fwd_arg_type: type,
     fwd_arg_names: List[str],
     fwd_buckets: List[_Bucket],
-    fwd_tensor_field_names: List[str],
     bwd_arg_names: List[str],
     bwd_buckets: List[_Bucket],
     fwd_slot_defaults: List[Any],
@@ -995,7 +928,7 @@ def _register_autograd_for_op(
 ) -> None:
     """Wire ``register_autograd`` on a forward op so its backward calls ``bwd_op_name``.
 
-    ``setup_context`` re-runs the proto fwd fake impl to recover output / saved
+    ``setup_context`` re-runs the fwd fake impl to recover output / saved
     templates, reassembles each flat output chunk, and hands the saved tuple +
     ``ctx_attrs`` to the module's ``setup_context``.
     """
@@ -1006,24 +939,23 @@ def _setup_context(ctx, inputs, output):
         }
         kwargs = dict(zip(fwd_arg_names, inputs))
         fwd_obj = _unpack(fwd_arg_type, kwargs, fwd_buckets)
-        proto_obj = _proto_view(fwd_obj, fwd_tensor_field_names)
 
-        user_fakes, saved_fakes, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(proto_obj))
+        user_fakes, saved_fakes, ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(fwd_obj))
 
         cursor = 0
         user_outputs: List[Any] = []
-        for proto in user_fakes:
-            n = _proto_slot_count(proto)
+        for template in user_fakes:
+            n = _slot_count(template)
             chunk = [_decode_none(t) for t in output[cursor : cursor + n]]
             cursor += n
-            user_outputs.append(_proto_reassemble(proto, chunk))
+            user_outputs.append(_maybe_reassemble_tensor_subclass(template, chunk))
 
         saved_list: List[Any] = []
-        for proto in saved_fakes:
-            n = _proto_slot_count(proto)
+        for template in saved_fakes:
+            n = _slot_count(template)
             chunk = [_decode_none(t) for t in output[cursor : cursor + n]]
             cursor += n
-            saved_list.append(_proto_reassemble(proto, chunk))
+            saved_list.append(_maybe_reassemble_tensor_subclass(template, chunk))
 
         bwd_obj = backward_obj_type()
         tensors_to_save_from_setup = setup_context_user(
@@ -1203,8 +1135,6 @@ def _register_custom_op_impl(
 
     fwd_buckets = _get_buckets(fwd_arg_type)
     bwd_buckets = _get_buckets(backward_arg_type)
-    fwd_tensor_field_names = _tensor_field_names(fwd_buckets)
-    bwd_tensor_field_names = _tensor_field_names(bwd_buckets)
 
     fwd_schema_args, fwd_arg_names = _build_schema(fwd_buckets)
     bwd_schema_args, bwd_arg_names = _build_schema(bwd_buckets)
@@ -1225,7 +1155,6 @@ def _register_custom_op_impl(
         arg_type=fwd_arg_type,
         arg_names=fwd_arg_names,
         buckets=fwd_buckets,
-        tensor_field_names=fwd_tensor_field_names,
         impl=fwd_impl,
         fake_impl=fwd_fake_impl,
         format_result=_format_fwd_result,
@@ -1236,7 +1165,6 @@ def _register_custom_op_impl(
         arg_type=backward_arg_type,
         arg_names=bwd_arg_names,
         buckets=bwd_buckets,
-        tensor_field_names=bwd_tensor_field_names,
         impl=backward_impl,
         fake_impl=bwd_fake_impl,
         format_result=lambda g: _format_bwd_result(g, num_grad_inputs, inner_bwd_qualname),
@@ -1257,7 +1185,6 @@ def _register_custom_op_impl(
         "fwd_arg_type": fwd_arg_type,
         "fwd_arg_names": fwd_arg_names,
         "fwd_buckets": fwd_buckets,
-        "fwd_tensor_field_names": fwd_tensor_field_names,
         "bwd_arg_names": bwd_arg_names,
         "bwd_buckets": bwd_buckets,
         "fwd_slot_defaults": fwd_slot_defaults,
@@ -1305,19 +1232,18 @@ def _bwd_rule(mode, func, types, args, kwargs):
     _quantized_tensor_passthrough_ops.add(inner_bwd_op.default)
 
     def forward_fn(fwd_args):
-        proto_obj = _proto_view(fwd_args, fwd_tensor_field_names)
-        user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(proto_obj))
+        user_fakes, _saved_fakes, _ctx_attrs = _split_fwd_fake_result(fwd_fake_impl(fwd_args))
         kwargs = _pack(fwd_args, fwd_buckets)
         flat_in = [kwargs[name] for name in fwd_arg_names]
         result = outer_fwd_op(*flat_in)
 
         cursor = 0
         outputs: List[Any] = []
-        for proto in user_fakes:
-            n = _proto_slot_count(proto)
+        for template in user_fakes:
+            n = _slot_count(template)
             chunk = [_decode_none(t) for t in result[cursor : cursor + n]]
             cursor += n
-            outputs.append(_proto_reassemble(proto, chunk))
+            outputs.append(_maybe_reassemble_tensor_subclass(template, chunk))
 
         if len(outputs) == 1:
             return outputs[0]
diff --git a/transformer_engine/pytorch/dynamo/tensor_proto.py b/transformer_engine/pytorch/dynamo/tensor_proto.py
deleted file mode 100644
index 911b4151bf..0000000000
--- a/transformer_engine/pytorch/dynamo/tensor_proto.py
+++ /dev/null
@@ -1,177 +0,0 @@
-# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
-#
-# See LICENSE for license information.
-
-"""TensorProto: a data-free description of a tensor / quantized tensor."""
-
-from __future__ import annotations
-import copy as _copy
-from dataclasses import dataclass, field
-from typing import Any, Dict, List, Optional, Tuple
-
-import torch
-
-
-def _contiguous_stride(shape: Tuple[int, ...]) -> Tuple[int, ...]:
-    """Row-major (contiguous) stride for ``shape``."""
-    stride: list = []
-    acc = 1
-    for dim in reversed(shape):
-        stride.append(acc)
-        acc *= dim
-    return tuple(reversed(stride))
-
-
-@dataclass
-class TensorProto:
-    """A data-free *prototype* of a tensor or quantized tensor.
-
-    Captures ``shape`` / ``dtype`` and, for quantized tensors, the
-    (value-opaque) ``quantizer`` -- enough to rebuild a tensor without holding
-    storage. The common abstraction over plain ``torch.Tensor``,
-    ``QuantizedTensorStorage`` and ``QuantizedTensor``, used for custom-op fake
-    impls and for reassembling a quantized tensor from bare buffers.
-    """
-
-    shape: Tuple[int, ...]
-    dtype: torch.dtype
-    quantizer: Optional[Any] = None
-    requires_grad: bool = False
-    device: Optional[torch.device] = field(default=None)
-
-    def __post_init__(self) -> None:
-        # Own a private copy of the quantizer so usage changes (update_usage)
-        # never touch the shared, value-opaque quantizer. The copy inherits the
-        # quantizer's current row-/column-wise usage as this proto's layout.
-        if self.quantizer is not None:
-            q = self.quantizer
-            self.quantizer = q.copy() if hasattr(q, "copy") else _copy.copy(q)
-
-    @property
-    def is_quantized(self) -> bool:
-        """Whether this proto describes a quantized tensor."""
-        return self.quantizer is not None
-
-    def update_usage(
-        self,
-        *,
-        rowwise_usage: Optional[bool] = None,
-        columnwise_usage: Optional[bool] = None,
-    ) -> None:
-        """Mirror ``QuantizedTensor.update_usage`` on the proto's buffer layout.
-
-        Applied to the proto's own quantizer copy, so the shared (value-opaque)
-        quantizer is never mutated. No-op for plain (non-quantized) protos.
-        """
-        if self.quantizer is None:
-            return
-        self.quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage)
-
-    def inner_names(self) -> Tuple[str, ...]:
-        """Names of the flat tensor buffers backing this proto, in order.
-
-        The real op flattens a quantized output via the storage's
-        ``__tensor_flatten__`` -- i.e. ``_FLATTEN_TENSOR_BUFFERS`` order, keeping
-        only the present buffers. ``_describe_buffers`` may emit the same buffers
-        in a different (per-usage) order (e.g. NVFP4 groups each amax right after
-        its scale), so reorder to the canonical flatten order here to keep the
-        fake layout aligned with the real one slot-for-slot.
-        """
-        if self.quantizer is None:
-            return ("data",)
-        # pylint: disable=protected-access
-        described = list(self.quantizer._describe_buffers(tuple(self.shape)).keys())
-        storage_cls = self.quantizer._storage_metadata(self.dtype)["cls"]
-        flatten_order = [attr for attr, _ in storage_cls._FLATTEN_TENSOR_BUFFERS]
-        extra = [name for name in described if name not in flatten_order]
-        if extra:
-            raise RuntimeError(
-                f"{storage_cls.__name__} describes buffer(s) {extra} absent from its "
-                f"_FLATTEN_TENSOR_BUFFERS {flatten_order}; the fake layout cannot be "
-                "aligned with the real one slot-for-slot."
-            )
-        return tuple(name for name in flatten_order if name in described)
-
-    def create_metadata(self) -> Dict[str, Any]:
-        """Data-free ``__tensor_unflatten__`` context describing this tensor."""
-        if self.quantizer is None:
-            return {
-                "is_tensor": True,
-                "is_quantized": False,
-                "dtype": self.dtype,
-                "requires_grad": self.requires_grad,
-            }
-        return self.quantizer.create_metadata(
-            tuple(self.shape), dtype=self.dtype, requires_grad=self.requires_grad
-        )
-
-    def create_inner_tensors(self) -> List[torch.Tensor]:
-        """Materialize the flat inner buffers (in :meth:`inner_names` order).
-
-        Under ``register_fake`` the ``torch.empty`` calls produce ``FakeTensor``s;
-        ``requires_grad`` is left default (managed by ``register_autograd``).
-        """
-        device = self.device if self.device is not None else torch.device("cuda")
-        if self.quantizer is None:
-            return [torch.empty(tuple(self.shape), dtype=self.dtype, device=device)]
-        inner = self.quantizer.alloc_tensors(tuple(self.shape), device=device)
-        return [inner[name] for name in self.inner_names()]
-
-    def create_tensor(self) -> torch.Tensor:
-        """Materialize an (uninitialized) tensor matching this proto (traceable).
-
-        Quantized protos reassemble the :meth:`create_inner_tensors` buffers via
-        the storage's ``__tensor_unflatten__``.
-        """
-        if self.quantizer is None:
-            device = self.device if self.device is not None else torch.device("cuda")
-            return torch.empty(
-                tuple(self.shape),
-                dtype=self.dtype,
-                device=device,
-                requires_grad=self.requires_grad,
-            )
-        from ..quantized_tensor import (  # pylint: disable=import-outside-toplevel
-            _STORAGE_REGISTRY,
-        )
-
-        shape = tuple(self.shape)
-        ctx = self.create_metadata()
-        inner = dict(zip(self.inner_names(), self.create_inner_tensors()))
-        storage_cls = _STORAGE_REGISTRY[ctx["cls"]]
-        return storage_cls.__tensor_unflatten__(inner, ctx, shape, _contiguous_stride(shape))
-
-
-def to_tensor_proto(tensor: Any) -> TensorProto:
-    """Build a :class:`TensorProto` describing ``tensor``.
-
-    Works for plain ``torch.Tensor`` and for ``QuantizedTensorStorage`` /
-    ``QuantizedTensor``. A *bare* storage exposes its shape via ``.size()`` and
-    its (fake) dtype via ``_dtype`` rather than ``.shape`` / ``.dtype``.
-    """
-    from ..quantized_tensor import (  # pylint: disable=import-outside-toplevel
-        QuantizedTensorStorage,
-    )
-
-    requires_grad = bool(getattr(tensor, "requires_grad", False))
-    if isinstance(tensor, QuantizedTensorStorage):
-        shape = getattr(tensor, "shape", None)
-        if shape is None:
-            shape = tensor.size()
-        dtype = getattr(tensor, "dtype", None)
-        if dtype is None:
-            dtype = getattr(tensor, "_dtype", None)
-        return TensorProto(
-            shape=tuple(shape),
-            dtype=dtype,
-            quantizer=getattr(tensor, "_quantizer", None),
-            requires_grad=requires_grad,
-            device=tensor.device,
-        )
-    return TensorProto(
-        shape=tuple(tensor.shape),
-        dtype=tensor.dtype,
-        quantizer=None,
-        requires_grad=requires_grad,
-        device=tensor.device,
-    )
diff --git a/transformer_engine/pytorch/dynamo/traceable_utils.py b/transformer_engine/pytorch/dynamo/traceable_utils.py
new file mode 100644
index 0000000000..560599518a
--- /dev/null
+++ b/transformer_engine/pytorch/dynamo/traceable_utils.py
@@ -0,0 +1,133 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+"""Pure-Python, torch.compile-traceable quantized tensor allocation and reassembly."""
+
+from __future__ import annotations
+import copy as _copy
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+
+
+def _contiguous_stride(shape: Tuple[int, ...]) -> Tuple[int, ...]:
+    """Row-major (contiguous) stride for ``shape``."""
+    stride: list = []
+    acc = 1
+    for dim in reversed(shape):
+        stride.append(acc)
+        acc *= dim
+    return tuple(reversed(stride))
+
+
+def make_empty_traceable(
+    quantizer,
+    shape: Tuple[int, ...],
+    *,
+    dtype: torch.dtype = torch.float32,
+    device: Optional[Union[torch.device, str]] = None,
+    requires_grad: bool = False,
+) -> Any:
+    """Allocate a tensor purely in Python (traceable under torch.compile).
+
+    When ``quantizer`` is not None, produces a quantized tensor via
+    ``alloc_tensors`` + ``__tensor_unflatten__`` (the compile-friendly
+    equivalent of ``Quantizer.make_empty``).  The quantizer is copied first
+    so the caller's instance is never mutated.
+
+    When ``quantizer`` is None, falls back to ``torch.empty`` for a plain tensor.
+
+    Stashed metadata (``_te_flat_names``, ``_te_flat_ctx``)
+    ---------------------------------------------------------
+    The resulting quantized tensor has these attributes stashed so that
+    ``forward_fn`` (in custom_op.py) can read them at Dynamo trace time.
+
+    Why: ``forward_fn`` runs inside torch.compile's trace.  It needs the flat
+    buffer names and unflatten context to decode the custom op's flat Tensor[]
+    return back into a structured QuantizedTensor.  Calling
+    ``__tensor_flatten__()`` for this would cause a graph break (it returns
+    non-Tensor Python objects -- List[str] and Dict -- that Dynamo cannot
+    represent as graph nodes).  Accessing ``t._quantizer`` and calling
+    ``_describe_buffers()`` on it also fails: while Dynamo treats the retrieved
+    quantizer as constant metadata, it wraps it in a generic VariableTracker
+    that does not support method calls (unlike a closure-captured quantizer
+    which is recognized as a value-opaque constant).  Stashing the buffer names
+    and context as plain attributes sidesteps both issues -- Dynamo reads them
+    as constants without needing to call any methods.
+
+    Allocation cost: when ``forward_fn`` calls the fake impl to obtain these
+    templates, the ``torch.empty`` calls appear as nodes in the initial Dynamo
+    FX graph.  However, because the tensors themselves are never used (only
+    the stashed metadata is read), AOT autograd's dead-code elimination removes
+    them before any kernel code is generated.  They do not appear in the final
+    compiled graph.
+    """
+    device = torch.device(device if device is not None else "cuda")
+    shape = tuple(shape)
+    if quantizer is None:
+        return torch.empty(shape, dtype=dtype, device=device, requires_grad=requires_grad)
+    from ..quantized_tensor import _STORAGE_REGISTRY  # pylint: disable=import-outside-toplevel
+
+    # Copy so the caller's quantizer is not mutated by alloc_tensors internals.
+    # The caller is expected to have already called set_usage() on the quantizer
+    # before passing it here -- Dynamo tracks those mutations as explicit setattr
+    # nodes in the graph, so the copy captures the post-mutation state and the
+    # stashed _te_flat_names reflects the correct buffer layout.
+    q = quantizer.copy() if hasattr(quantizer, "copy") else _copy.copy(quantizer)
+    ctx = q.create_metadata(shape, dtype=dtype, requires_grad=requires_grad)
+    inner = q.alloc_tensors(shape, device=device)
+    storage_cls = _STORAGE_REGISTRY[ctx["cls"]]
+    result = storage_cls.__tensor_unflatten__(inner, ctx, shape, _contiguous_stride(shape))
+    if requires_grad and hasattr(result, "requires_grad_"):
+        result.requires_grad_(True)
+    # TODO: understand why Dynamo does not recognize the quantizer retrieved via
+    # t._quantizer as the same value-opaque type it would if captured from a
+    # closure. If that is fixed upstream, the stashed attributes become
+    # unnecessary and we could compute slot counts directly from the quantizer.
+    result._te_flat_names = tuple(inner.keys())
+    result._te_flat_ctx = ctx
+    return result
+
+
+# --------------------------------------------------------------------------- #
+# Slot counting and reassembly for the custom-op flat Tensor[] protocol.
+# --------------------------------------------------------------------------- #
+
+
+def _slot_count(value: Any) -> int:
+    """Number of flat tensor slots a value occupies in the op's Tensor[] return.
+
+    Reads ``_te_flat_names`` stashed by :func:`make_empty_traceable`, which is
+    safe to access at Dynamo trace time (treated as constant metadata on a
+    traceable wrapper subclass).  Plain tensors (no stashed names) occupy 1 slot.
+    """
+    if value is None:
+        return 1
+    names = getattr(value, "_te_flat_names", None)
+    if names is not None:
+        return len(names)
+    return 1
+
+
+def _maybe_reassemble_tensor_subclass(
+    template: Any,
+    chunk: List[Optional[torch.Tensor]],
+) -> Optional[Union[torch.Tensor, Any]]:
+    """Rebuild a value from its flat tensors using ``template`` for geometry.
+
+    ``template`` is a tensor produced by the fake impl via
+    :func:`make_empty_traceable`.  Uses the stashed ``_te_flat_names`` /
+    ``_te_flat_ctx`` attributes for reassembly (trace-safe).  For plain tensors
+    (no stashed metadata), returns the single chunk element directly.
+    """
+    if template is None:
+        return None
+    names = getattr(template, "_te_flat_names", None)
+    ctx = getattr(template, "_te_flat_ctx", None)
+    if names is None or ctx is None:
+        return chunk[0]
+    inner_dict = dict(zip(names, chunk))
+    shape = tuple(template.shape) if hasattr(template, "shape") else tuple(template.size())
+    stride = _contiguous_stride(shape)
+    return type(template).__tensor_unflatten__(inner_dict, ctx, shape, stride)
diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py
index d1ea805077..d1dd5f362b 100644
--- a/transformer_engine/pytorch/module/linear.py
+++ b/transformer_engine/pytorch/module/linear.py
@@ -71,7 +71,8 @@
     prepare_for_saving,
     restore_from_func_ctx,
 )
-from ..dynamo import TensorProto, register_custom_op, is_value_opaque_quantizer
+from ..dynamo import register_custom_op, is_value_opaque_quantizer
+from ..dynamo.traceable_utils import make_empty_traceable
 from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer
 from ..tensor.mxfp8_tensor import MXFP8Quantizer
 from ..tensor.utils import clear_columnwise_cache, is_custom
@@ -698,10 +699,10 @@ def _linear_forward_impl(
 
 def _linear_forward_impl_fake(
     args: LinearFwdArgs,
-) -> Tuple[TensorProto, Optional[TensorProto], Optional[Tuple[Any, ...]], None, Optional[Dict]]:
+) -> Tuple[Any, Optional[Any], Optional[Tuple[Any, ...]], None, Optional[Dict]]:
     """Shape/metadata-only twin of :func:`_linear_forward_impl` for torch.compile,
-    returning ``TensorProto`` descriptors for the outputs and saved tensors instead
-    of allocating real data."""
+    returning traceable tensors for the outputs and saved tensors instead of
+    allocating real data via C++ kernels."""
     if args.fsdp_group is not None and args.is_grad_enabled:
         raise NotImplementedError(
             "Compile-time Linear forward does not support manual TE FSDP "
@@ -730,7 +731,7 @@ def _linear_forward_impl_fake(
     inputmat_is_storage = False
     inputmat_aliases_inp = False
     if fp8_or_debug:
-        if inp.is_quantized:
+        if isinstance(inp, QuantizedTensorStorage):
             # Primary-quantized input reused as-is.
             inputmat_is_storage = True
             inputmat_aliases_inp = True
@@ -764,7 +765,7 @@ def _linear_forward_impl_fake(
     weightmat_is_storage = False
     weightmat_aliases_weight = False
     if fp8_or_debug:
-        if weight_quantizer is not None and (not weight.is_quantized or debug):
+        if weight_quantizer is not None and (not isinstance(weight, QuantizedTensorStorage) or debug):
             columnwise_usage = is_grad_enabled and args.input_requires_grad and not args.is_fsdp2
             if args.backward_override is not None:
                 columnwise_usage = False
@@ -774,34 +775,30 @@ def _linear_forward_impl_fake(
                     and not in_fp8_activation_recompute_phase()
                 )
             weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage)
-        elif weight.is_quantized:
-            weight_quantizer = weight.quantizer
+        elif isinstance(weight, QuantizedTensorStorage):
+            weight_quantizer = weight._quantizer
 
-        if weight.is_quantized:
+        if isinstance(weight, QuantizedTensorStorage):
             # Primary-quantized weight: the impl reuses it as ``weightmat``.
             weightmat = weight
             weightmat_is_storage = True
             weightmat_aliases_weight = True
         else:
-            weightmat = TensorProto(
-                shape=tuple(weight.shape),
-                dtype=activation_dtype,
-                quantizer=weight_quantizer,
-                device=weight.device,
+            weightmat = make_empty_traceable(
+                weight_quantizer, tuple(weight.shape),
+                dtype=activation_dtype, device=weight.device,
             )
             weightmat_is_storage = True
             update_ws = args.is_first_microbatch is None or args.is_first_microbatch
             if args.cache_weight and update_ws and args.weight_workspace is None:
-                new_weight_workspace = TensorProto(
-                    shape=tuple(weight.shape),
-                    dtype=activation_dtype,
-                    quantizer=weight_quantizer,
-                    device=weight.device,
+                new_weight_workspace = make_empty_traceable(
+                    weight_quantizer, tuple(weight.shape),
+                    dtype=activation_dtype, device=weight.device,
                 )
     else:
         weightmat_aliases_weight = weight.dtype == activation_dtype
-        weightmat = TensorProto(
-            shape=tuple(weight.shape), dtype=activation_dtype, device=weight.device
+        weightmat = make_empty_traceable(
+            None, tuple(weight.shape), dtype=activation_dtype, device=weight.device,
         )
 
     if output_quantizer is not None:
@@ -815,11 +812,10 @@ def _linear_forward_impl_fake(
         out_leading = out_leading * args.tp_size
     elif args.parallel_mode == "row" and args.sequence_parallel:
         out_leading = out_leading // args.tp_size
-    out = TensorProto(
-        shape=(out_leading, *tuple(inp.shape[1:-1]), out_features),
+    out = make_empty_traceable(
+        output_quantizer,
+        (out_leading, *tuple(inp.shape[1:-1]), out_features),
         dtype=activation_dtype,
-        quantizer=output_quantizer,
-        requires_grad=is_grad_enabled and (args.input_requires_grad or args.weight_requires_grad),
         device=inp.device,
     )
 
@@ -837,29 +833,29 @@ def _linear_forward_impl_fake(
             if inputmat_aliases_inp:
                 inputmat_alias = "inp"
             elif inputmat_is_storage:
-                saved_inputmat = TensorProto(
-                    shape=tuple(inp.shape),
-                    dtype=activation_dtype,
-                    quantizer=input_quantizer,
-                    device=inp.device,
-                )
                 # Mirror ``_linear_forward_impl``'s post-quantization
                 # ``inputmat.update_usage(...)`` so the saved input's buffer layout
                 # matches -- driven by the same conditions as the real impl.
+                # Copy the quantizer so the shared instance on fwd_args is not mutated.
+                save_q = input_quantizer.copy() if hasattr(input_quantizer, "copy") else input_quantizer
                 if own_quantized_input and not save_original_input:
                     if args.backward_override is not None:
-                        saved_inputmat.update_usage(rowwise_usage=True, columnwise_usage=False)
+                        save_q.set_usage(rowwise=True, columnwise=False)
                     elif (
                         args.backward_input_needs_gather
                         and weight_quantizer is not None
                         and weight_quantizer.supports_only_rowwise_all_gather()
                     ):
-                        saved_inputmat.update_usage(rowwise_usage=True, columnwise_usage=False)
+                        save_q.set_usage(rowwise=True, columnwise=False)
                     else:
-                        saved_inputmat.update_usage(rowwise_usage=False, columnwise_usage=True)
+                        save_q.set_usage(rowwise=False, columnwise=True)
+                saved_inputmat = make_empty_traceable(
+                    save_q, tuple(inp.shape),
+                    dtype=activation_dtype, device=inp.device,
+                )
             else:
-                saved_inputmat = TensorProto(
-                    shape=tuple(inp.shape), dtype=activation_dtype, device=inp.device
+                saved_inputmat = make_empty_traceable(
+                    None, tuple(inp.shape), dtype=activation_dtype, device=inp.device,
                 )
 
         # Slot 1 -- ``wt_save``. Mirror the real impl's alias dedup: the cached
@@ -879,8 +875,8 @@ def _linear_forward_impl_fake(
         elif weightmat_is_storage:
             wt_save = weightmat
         else:
-            wt_save = TensorProto(
-                shape=tuple(weight.shape), dtype=activation_dtype, device=weight.device
+            wt_save = make_empty_traceable(
+                None, tuple(weight.shape), dtype=activation_dtype, device=weight.device,
             )
 
         # Slot 2 -- ``saved_weight`` (always aliased to ``weight``).
@@ -1612,13 +1608,12 @@ def wgrad_gemm(
 
 def _linear_backward_impl_fake(
     args: LinearBwdArgs,
-) -> Tuple[Optional[TensorProto], Optional[TensorProto], Optional[TensorProto]]:
-    """Allocation-free fake of :func:`_linear_backward` on ``TensorProto``.
+) -> Tuple[Optional[Any], Optional[Any], Optional[Any]]:
+    """Traceable fake of :func:`_linear_backward`.
 
-    The saved-tensor fields of ``args`` carry
-    :class:`~transformer_engine.pytorch.dynamo.TensorProto` instances. Returns
-    ``(wgrad, dgrad, grad_bias)`` protos describing the nature of the gradients,
-    mirroring the real backward's return contract without allocating storage.
+    The saved-tensor fields of ``args`` carry real tensors (fake-backed under
+    tracing).  Returns ``(wgrad, dgrad, grad_bias)`` as traceable tensors,
+    mirroring the real backward's return contract without running C++ kernels.
 
     Tensor-/sequence-parallel gather/scatter happens inside the eager backward
     custom op and is opaque to ``torch.compile``: ``dgrad`` always carries the
@@ -1642,7 +1637,6 @@ def _linear_backward_impl_fake(
 
     dgrad = None
     if args.requires_dgrad:
-        # dgrad has the logical input shape and may be quantized for the next op.
         # Derive shape from grad_output + weight + SP config instead of args.inp_shape:
         # inp_shape is not stored in the value bundle under dynamic shapes (SymInt is
         # not hashable in OpaqueValueBundle), so we reconstruct it here.
@@ -1654,10 +1648,10 @@ def _linear_backward_impl_fake(
             _dgrad_leading = _go_leading * args.tp_size
         else:
             _dgrad_leading = _go_leading
-        dgrad = TensorProto(
-            shape=(_dgrad_leading, *args.grad_output.shape[1:-1], _in_features),
+        dgrad = make_empty_traceable(
+            args.grad_input_quantizer,
+            (_dgrad_leading, *args.grad_output.shape[1:-1], _in_features),
             dtype=out_dtype,
-            quantizer=args.grad_input_quantizer,
             device=args.grad_output.device,
         )
 
@@ -1667,17 +1661,17 @@ def _linear_backward_impl_fake(
         # requested (mirrors ``quantization_params=grad_weight_quantizer``),
         # otherwise high precision. Under fuse_wgrad_accumulation the grad is
         # written into ``main_grad`` in place and no wgrad tensor is returned.
-        wgrad = TensorProto(
-            shape=(out_features, in_features),
+        wgrad = make_empty_traceable(
+            args.grad_weight_quantizer,
+            (out_features, in_features),
             dtype=out_dtype,
-            quantizer=args.grad_weight_quantizer,
             device=weight.device,
         )
 
     grad_bias = None
     if args.use_bias and args.requires_wgrad:
-        grad_bias = TensorProto(
-            shape=(out_features,), dtype=out_dtype, device=args.grad_output.device
+        grad_bias = make_empty_traceable(
+            None, (out_features,), dtype=out_dtype, device=args.grad_output.device,
         )
 
     return wgrad, dgrad, grad_bias
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index 86c4c94f7c..2af15ca7b0 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -150,7 +150,7 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None:
             f"{self.__class__.__name__} class does not implement copy_from_storage function"
         )
 
-    # ----- PyTorch subclass flatten protocol (torch.compile / TensorProto) -----
+    # ----- PyTorch subclass flatten protocol (torch.compile / traceable allocation) -----
 
     # Subclasses declare their tensor buffers once, as ``(attribute_name,
     # constructor_kwarg)`` pairs in flatten order; everything else returned by
@@ -427,7 +427,7 @@ def make_empty(
             result.requires_grad_(True)
         return result
 
-    # ----- Data-free buffer/metadata primitives backing TensorProto -----
+    # ----- Data-free buffer/metadata primitives backing make_empty_traceable -----
 
     def _describe_buffers(
         self, shape: Tuple[int, ...]
@@ -440,7 +440,7 @@ def _describe_buffers(
         """
         raise NotImplementedError(
             f"{self.__class__.__name__} does not implement _describe_buffers; "
-            "it cannot be used with TensorProto / pure-Python allocation"
+            "it cannot be used with traceable allocation"
         )
 
     def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]:
@@ -454,7 +454,7 @@ def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]:
         """
         raise NotImplementedError(
             f"{self.__class__.__name__} does not implement _storage_metadata; "
-            "it cannot be used with TensorProto / pure-Python allocation"
+            "it cannot be used with traceable allocation"
         )
 
     def alloc_tensors(
diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
index c816b4fb04..04e5281f65 100644
--- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py
@@ -73,7 +73,7 @@ def copy(self) -> Float8BlockQuantizer:
     def _value_fields(self) -> Tuple[str, ...]:
         return ("dtype", "block_len", "amax_epsilon", "force_pow_2_scales", "block_scaling_dim")
 
-    # ----- TensorProto / pure-Python allocation -----
+    # ----- traceable allocation -----
 
     def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]:
         return {
diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py
index 6d3a53b3d7..423b055f15 100644
--- a/transformer_engine/pytorch/tensor/float8_tensor.py
+++ b/transformer_engine/pytorch/tensor/float8_tensor.py
@@ -393,7 +393,7 @@ def _value_fields(self) -> Tuple[str, ...]:
         # raises so it can never be baked into a torch.compile graph.
         return ("dtype", "force_pow_2_scales", "amax_epsilon", "with_amax_reduction")
 
-    # ----- TensorProto / pure-Python allocation -----
+    # ----- traceable allocation -----
 
     def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]:
         return {
diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
index f804a96f24..bd51a88aaf 100644
--- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py
+++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py
@@ -61,7 +61,7 @@ def copy(self) -> MXFP8Quantizer:
     def _value_fields(self) -> Tuple[str, ...]:
         return ("dtype",)
 
-    # ----- TensorProto / pure-Python allocation -----
+    # ----- traceable allocation -----
 
     def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]:
         return {
diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 8d51480f51..795eb499c0 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -366,7 +366,7 @@ def _value_fields(self) -> Tuple[str, ...]:
             "with_amax_reduction",
         )
 
-    # ----- TensorProto / pure-Python allocation -----
+    # ----- traceable allocation -----
 
     def _storage_metadata(self, fake_dtype: torch.dtype) -> Dict[str, Any]:
         return {

From c33cd00a83b127e45ce8d7846e6d99dc2b0b65eb Mon Sep 17 00:00:00 2001
From: kshitij12345 <kshitijkalambarkar@gmail.com>
Date: Thu, 2 Jul 2026 07:56:50 -0700
Subject: [PATCH 27/28] [PyTorch][torch.compile] Guard NVFP4Quantizer.copy()
 tensor access under is_compiling

Under Dynamo tracing, rht_matrix is a FakeTensor attached to an opaque
script object. Accessing it in copy() triggers SourcelessBuilder which
cannot wrap FakeTensor, causing an InternalTorchDynamoError.

The fake impl never runs real quantization, so rht_matrix is unnecessary
during tracing. Guard the tensor field copies with
torch.compiler.is_compiling() -- the matrix will be rebuilt lazily via
_rebuild_derived_state if the quantizer is later used outside tracing.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 transformer_engine/pytorch/tensor/nvfp4_tensor.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
index 795eb499c0..d22423353c 100644
--- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py
+++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py
@@ -245,8 +245,14 @@ def copy(self) -> NVFP4Quantizer:
         )
         quantizer.internal = self.internal
         quantizer.optimize_for_gemm = self.optimize_for_gemm
-        quantizer.rht_matrix = self.rht_matrix
         quantizer.rht_matrix_random_sign_mask_t = self.rht_matrix_random_sign_mask_t
+        if not torch.compiler.is_compiling():
+            # Under Dynamo tracing rht_matrix is a FakeTensor on an opaque script
+            # object; accessing it triggers SourcelessBuilder which cannot wrap
+            # FakeTensor.  The fake impl never runs real quantization so the matrix
+            # is unnecessary -- it will be rebuilt lazily via _rebuild_derived_state
+            # if the quantizer is later used outside tracing.
+            quantizer.rht_matrix = self.rht_matrix
 
         return quantizer
 

From a582eb53752afee4692ea5ec486451378e1a549e Mon Sep 17 00:00:00 2001
From: kshitij12345 <kshitijkalambarkar@gmail.com>
Date: Fri, 3 Jul 2026 05:22:50 -0700
Subject: [PATCH 28/28] [PyTorch][test] Fix NVFP4 value-object test: set
 with_post_rht_amax=True

The C++ quantize kernel requires with_post_rht_amax=True when with_rht
is enabled. The test factory was creating an NVFP4Quantizer with
with_rht=True but with_post_rht_amax defaulting to False, causing
'Pre-RHT amax is not supported yet' at quantize time.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 tests/pytorch/test_torch_compile.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index 716947026f..87dfd171f6 100644
--- a/tests/pytorch/test_torch_compile.py
+++ b/tests/pytorch/test_torch_compile.py
@@ -499,6 +499,7 @@ def _nvfp4(with_rht=True):
         rowwise=True,
         columnwise=True,
         with_rht=with_rht,
+        with_post_rht_amax=with_rht,
     )