diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py
index dff94cef9d..de62e56e53 100644
--- a/tests/pytorch/test_sanity.py
+++ b/tests/pytorch/test_sanity.py
@@ -1154,6 +1154,41 @@ def test_quantized_model_init_high_precision_init_val():
     ), "clear_high_precision_init_val() not work"
 
 
+@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8)
+@pytest.mark.parametrize("move", ["cuda", "cpu", "half"])
+def test_quantized_param_attrs_survive_apply(move):
+    """Attributes attached to a quantized parameter survive nn.Module._apply.
+
+    Quantized parameters implement the flatten protocol, so ``_apply`` moves them
+    with ``swap_tensors``, which exchanges the parameter's whole ``__dict__``.
+    Anything attached from the outside rides out on the discarded tensor unless
+    the module re-attaches it.
+    """
+    with quantized_model_init(preserve_high_precision_init_val=True):
+        model = Linear(64, 64)
+
+    weight = model.weight
+    expected = weight.get_high_precision_init_val()
+    weight.probe_attr = "attached-from-outside"
+
+    if move == "cuda":
+        model = model.cuda()
+    elif move == "cpu":
+        model = model.cpu()
+    else:
+        model = model.half()
+
+    weight = model.weight
+    assert hasattr(weight, "get_high_precision_init_val"), f"accessor lost by .{move}()"
+    assert hasattr(weight, "clear_high_precision_init_val"), f"accessor lost by .{move}()"
+    torch.testing.assert_close(weight.get_high_precision_init_val(), expected, rtol=0, atol=0)
+    assert weight.probe_attr == "attached-from-outside", f"custom attr lost by .{move}()"
+
+    # The accessor must read the surviving parameter, not the discarded one.
+    weight.clear_high_precision_init_val()
+    assert weight.get_high_precision_init_val() is None
+
+
 @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8)
 def test_grouped_linear_single_param_preserves_high_precision_init(monkeypatch):
     """Grouped MXFP8 and discrete weights produce identical FP32 master initialization."""
diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py
index d720ca7a4a..6ab29d5f53 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
@@ -27,6 +28,9 @@
 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.quantized_tensor import QuantizedTensor, Quantizer
+from transformer_engine.pytorch.dynamo import TensorProto, to_tensor_proto
 from transformer_engine.pytorch import (
     is_fp8_available,
     is_mxfp8_available,
@@ -711,14 +715,7 @@ def _hw_available(quantizer):
     pytest.param(_mxfp8, id="mxfp8"),
     pytest.param(_blockwise, id="float8_blockwise"),
     pytest.param(_current_scaling, id="float8_current_scaling"),
-    pytest.param(
-        _nvfp4,
-        id="nvfp4",
-        marks=pytest.mark.skipif(
-            not torch.cuda.is_available(),
-            reason="NVFP4Quantizer requires CUDA to construct",
-        ),
-    ),
+    pytest.param(_nvfp4, id="nvfp4"),
 ]
 
 
@@ -732,13 +729,15 @@ def test_quantizer_value_object(factory):
     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
     # 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):
+    if _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)
 
@@ -807,7 +806,7 @@ def test_quantizer_value_object_fullgraph(factory):
     unlike merely passing the quantizer through.
     """
     q = factory()
-    if not (torch.cuda.is_available() and _hw_available(q)):
+    if not _hw_available(q):
         pytest.skip("format not supported on this HW")
 
     op = _QDQ_OPS[type(q)]
@@ -820,3 +819,379 @@ 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.
+# Format support is gated at runtime, in the tests that run a kernel; the rest is
+# pure Python and works on any HW.
+_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"),
+]
+
+
+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.inner_tensor_specs(shape))
+    ctx = quantizer.create_metadata(shape, dtype=dtype)
+    buffers = quantizer.alloc_tensors(shape, device=device)
+    inner = {name: buffers[name] for name in names}
+    storage_cls = ctx["cls"]
+    # 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):
+    """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.inner_tensor_specs(shape))
+
+    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.inner_tensor_specs(shape)
+    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.inner_tensor_specs(shape))
+    # 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)
+
+
+_USAGE_COMBOS = [
+    pytest.param(True, True, id="rowwise_columnwise"),
+    pytest.param(True, False, id="rowwise_only"),
+    pytest.param(False, True, id="columnwise_only"),
+]
+
+
+@pytest.mark.parametrize("factory, shape", _PROTO_QUANTIZERS)
+@pytest.mark.parametrize("rowwise, columnwise", _USAGE_COMBOS)
+@pytest.mark.parametrize("internal", [False, True], ids=["wrapper", "internal"])
+def test_python_alloc_matches_cpp_make_empty(factory, shape, rowwise, columnwise, internal):
+    """Pure-Python allocation is interchangeable with the C++ allocation.
+
+    Builds the same quantized tensor twice: via ``Quantizer.make_empty``
+    (``tex.create_empty_quantized_tensor``, the C++ path) and via the Python
+    primitives ``inner_tensor_specs`` + ``create_metadata`` + ``alloc_tensors``
+    + ``__tensor_unflatten__`` (exactly what ``TensorProto.create_tensor``
+    does). Checks:
+
+    * structural parity -- same concrete class, buffer set, per-buffer
+      shape/dtype/device, logical shape/dtype and flatten context;
+    * functional parity -- the real C++ quantize kernel writes bit-identical
+      results into the Python-allocated buffers as into the C++-allocated
+      ones, proving the Python buffer description matches the layout
+      (padding/alignment) the kernels expect.
+    """
+
+    # Two independent, identically-configured quantizers so no state can leak
+    # between the two allocation paths.
+    def make_quantizer():
+        q = factory()
+        q.set_usage(rowwise=rowwise, columnwise=columnwise)
+        q.internal = internal
+        return q
+
+    q_ref = make_quantizer()
+    if not _hw_available(q_ref):
+        pytest.skip("format not supported on this HW")
+    q_py = make_quantizer()
+
+    ref = q_ref.make_empty(shape, dtype=torch.bfloat16, device="cuda")
+    py = _build_from_primitives(q_py, shape, torch.bfloat16, device="cuda")
+
+    # --- Structural parity ---
+    assert type(py) is type(ref)
+    ref_names, ref_ctx = ref.__tensor_flatten__()
+    py_names, py_ctx = py.__tensor_flatten__()
+    assert set(py_names) == set(ref_names)
+    for name in ref_names:
+        rbuf, pbuf = getattr(ref, name), getattr(py, name)
+        assert tuple(pbuf.shape) == tuple(rbuf.shape), name
+        assert pbuf.dtype == rbuf.dtype, name
+        assert pbuf.device == rbuf.device, name
+
+    # Logical shape / dtype (bare storages are not torch.Tensors: they expose
+    # size() and _dtype instead of .shape / .dtype).
+    if isinstance(ref, QuantizedTensor):
+        assert tuple(py.shape) == tuple(ref.shape) == tuple(shape)
+        assert py.dtype == ref.dtype == torch.bfloat16
+    else:
+        assert tuple(py.size()) == tuple(ref.size()) == tuple(shape)
+        # pylint: disable=protected-access
+        assert py._dtype == ref._dtype == torch.bfloat16
+
+    # Flatten context. The quantizer entry needs special handling: production
+    # quantizers get a value-based __eq__ from register_value_opaque_quantizer,
+    # but fall back to field-wise comparison for classes that don't define one
+    # (plain object.__eq__ is identity, which would spuriously fail).
+    assert set(py_ctx) == set(ref_ctx)
+    for key in ("cls", "is_tensor", "requires_grad"):
+        assert py_ctx[key] == ref_ctx[key], key
+    ref_kwargs, py_kwargs = ref_ctx["nontensor_kwargs"], py_ctx["nontensor_kwargs"]
+    assert set(py_kwargs) == set(ref_kwargs)
+    for key in ref_kwargs:
+        rv, pv = ref_kwargs[key], py_kwargs[key]
+        if isinstance(rv, Quantizer) or isinstance(pv, Quantizer):
+            assert type(pv) is type(rv), key
+            assert (pv.rowwise_usage, pv.columnwise_usage, pv.internal) == (
+                rv.rowwise_usage,
+                rv.columnwise_usage,
+                rv.internal,
+            ), key
+            if type(rv).__eq__ is not object.__eq__:
+                assert pv == rv, key
+        else:
+            assert pv == rv, key
+
+    # --- Functional parity: run the real C++ quantize kernel into both ---
+    x = torch.randn(*shape, dtype=torch.bfloat16, device="cuda")
+
+    def _quantize_into(quantizer, dst):
+        if internal:
+            # update_quantized() only accepts the wrapper classes; internal
+            # (bare storage) tensors are filled through the same underlying
+            # kernel binding directly.
+            tex.quantize(x, quantizer, dst, None)
+        else:
+            quantizer.update_quantized(x, dst)
+
+    # Scale-inv padding is never written by the kernel and both paths allocate it
+    # uninitialized; zero it so the comparison below covers only kernel output.
+    for name in ref_names:
+        getattr(ref, name).zero_()
+        getattr(py, name).zero_()
+
+    # Some combos are rejected by the quantize kernel itself regardless of who
+    # allocated the tensor (e.g. FP8 current-scaling columnwise-only on
+    # TN-capable archs: there is no rowwise data buffer and
+    # nvte_compute_scale_from_amax asserts on it). Parity then means the
+    # Python-allocated tensor is rejected the same way -- not a silent skip.
+    try:
+        _quantize_into(q_ref, ref)
+    except RuntimeError:
+        with pytest.raises(RuntimeError):
+            _quantize_into(q_py, py)
+        return
+    _quantize_into(q_py, py)
+    for name in ref_names:
+        torch.testing.assert_close(
+            getattr(py, name), getattr(ref, name), rtol=0.0, atol=0.0, equal_nan=True
+        )
+
+    # Value check through dequantize(). Some layouts cannot dequantize at all
+    # (e.g. FP8 columnwise-only raises NotImplementedError) -- the C++-allocated
+    # reference defines what is supported, and when it raises, the bitwise
+    # buffer equality above already proves value parity.
+    try:
+        expected = ref.dequantize()
+    except NotImplementedError:
+        expected = None
+    if expected is not None:
+        torch.testing.assert_close(py.dequantize(), expected, rtol=0.0, atol=0.0)
+
+
+# ----- 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 inner_tensor_specs.
+    bufs = q.inner_tensor_specs(shape)
+    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.inner_tensor_specs(shape))
+    # 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..cab7da48e6
--- /dev/null
+++ b/transformer_engine/pytorch/dynamo/tensor_proto.py
@@ -0,0 +1,164 @@
+# 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
+from torch._prims_common import make_contiguous_strides_for
+
+
+@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. Raises on plain (non-quantized) protos --
+        a real plain ``torch.Tensor`` has no ``update_usage`` either.
+        """
+        if self.quantizer is None:
+            raise ValueError("update_usage called on a non-quantized TensorProto")
+        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. ``_INNER_TENSORS`` order, keeping
+        only the present buffers. ``inner_tensor_specs`` 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.inner_tensor_specs(tuple(self.shape)).keys())
+        storage_cls = self.quantizer.storage_metadata(self.dtype)["cls"]
+        flatten_order = [attr for attr, _ in storage_cls._INNER_TENSORS]
+        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"_INNER_TENSORS {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 assemble(self, inner_tensors: List[torch.Tensor]) -> torch.Tensor:
+        """Rebuild the tensor from ready-made ``inner_tensors`` (in :meth:`inner_names`
+        order). Shared by :meth:`create_tensor` (fresh buffers) and the custom-op
+        boundary (buffers arriving from an op's flat ``Tensor[]`` payload).
+
+        Non-quantized protos are the single inner tensor as-is; quantized protos
+        are reassembled into the storage/wrapper via ``__tensor_unflatten__``.
+        """
+        if self.quantizer is None:
+            return inner_tensors[0]
+        shape = tuple(self.shape)
+        ctx = self.create_metadata()
+        inner = dict(zip(self.inner_names(), inner_tensors))
+        storage_cls = ctx["cls"]
+        return storage_cls.__tensor_unflatten__(
+            inner, ctx, shape, make_contiguous_strides_for(shape)
+        )
+
+    def create_tensor(self) -> torch.Tensor:
+        """Materialize an (uninitialized) tensor matching this proto (traceable).
+
+        Quantized protos reassemble freshly-allocated :meth:`create_inner_tensors`
+        buffers via :meth:`assemble`.
+        """
+        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,
+            )
+        return self.assemble(self.create_inner_tensors())
+
+
+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 (fake) dtype via
+    ``_dtype`` rather than ``.dtype``.
+    """
+    requires_grad = bool(getattr(tensor, "requires_grad", False))
+    dtype = getattr(tensor, "dtype", None)
+    if dtype is None:
+        dtype = getattr(tensor, "_dtype", None)
+    return TensorProto(
+        shape=tuple(tensor.shape),
+        dtype=dtype,
+        quantizer=getattr(tensor, "_quantizer", None),
+        requires_grad=requires_grad,
+        device=tensor.device,
+    )
diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py
index a0cdfa1378..0e922fa072 100644
--- a/transformer_engine/pytorch/module/base.py
+++ b/transformer_engine/pytorch/module/base.py
@@ -927,6 +927,35 @@ def module_setattr(self, name: str, value: Any) -> None:
         """
         super().__setattr__(name, value)
 
+    def _apply(self, *args, **kwargs):
+        """Re-attach attributes that ``swap_tensors`` moves off a quantized parameter.
+
+        ``_apply`` moves wrapper subclasses by exchanging the parameter's whole
+        ``__dict__``, which carries the inner buffers over but takes externally
+        attached state (``_high_precision_init_val``, ``main_grad``, ...) with it.
+        """
+        snapshots = {
+            name: (param, dict(param.__dict__))
+            for name, param in self._parameters.items()
+            if isinstance(param, QuantizedTensorStorage)
+        }
+        out = super()._apply(*args, **kwargs)
+        for name, (old_param, attrs) in snapshots.items():
+            new_param = self._parameters.get(name)
+            if new_param is None:
+                raise RuntimeError(
+                    f"{type(self).__name__}.{name} disappeared during _apply; the state"
+                    " attached to it cannot be restored"
+                )
+            for key, value in attrs.items():
+                # Still present -> tensor state; the post-swap value is the right one.
+                if key in new_param.__dict__:
+                    continue
+                if isinstance(value, MethodType) and value.__self__ is old_param:
+                    value = MethodType(value.__func__, new_param)
+                setattr(new_param, key, value)
+        return out
+
     @property
     def output_quantizer_role(self) -> Optional[QuantizerRole]:
         """Caller-configurable :class:`QuantizerRole` for the forward output quantizer.
diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py
index 654f84159f..022af2c4c0 100644
--- a/transformer_engine/pytorch/quantized_tensor.py
+++ b/transformer_engine/pytorch/quantized_tensor.py
@@ -5,7 +5,7 @@
 """Pure Python base classes for quantization."""
 
 from __future__ import annotations
-from typing import Optional, Tuple, Iterable, Any, Dict, Union
+from typing import NamedTuple, Optional, Tuple, Iterable, Any, Dict, Union, get_type_hints
 import abc
 import warnings
 import math
@@ -23,13 +23,33 @@
     _stride_from_shape,
 )
 
-
 # Custom ops that should pass through __torch_dispatch__ without unwrapping
 # QuantizedTensor subclasses (e.g. Float8Tensor). Register ops here that
 # handle quantized tensors internally.
 _quantized_tensor_passthrough_ops: set = set()
 
 
+class InnerTensor(NamedTuple):
+    """Marks a storage field as a flat tensor buffer.
+
+    Annotate the field with it -- ``_scale_inv: Annotated[torch.Tensor,
+    InnerTensor("fp8_scale_inv")]`` -- and ``__init_subclass__`` collects the
+    declarations into ``_INNER_TENSORS``, in field order.
+    """
+
+    ctor_kwarg: str
+
+
+def _collect_inner_tensor_fields(cls: type) -> Tuple[Tuple[str, str], ...]:
+    """Buffers a storage class declares, as ``(attribute, constructor kwarg)``."""
+    fields = []
+    for attr, hint in get_type_hints(cls, include_extras=True).items():
+        for meta in getattr(hint, "__metadata__", ()):
+            if isinstance(meta, InnerTensor):
+                fields.append((attr, meta.ctor_kwarg))
+    return tuple(fields)
+
+
 class QuantizedTensorStorage:
     r"""Base class for all TensorStorage classes.
 
@@ -49,6 +69,20 @@ class QuantizedTensorStorage:
     _dtype: torch.dtype
     _quantizer: Optional[Quantizer]
 
+    @property
+    def shape(self) -> torch.Size:
+        """Logical tensor shape, valid on bare storages and wrapper tensors alike.
+
+        Wrapper subclasses (also ``torch.Tensor``) defer to the native tensor
+        shape (``size()`` on a bare storage may reconstruct the shape from the
+        columnwise buffer, which is not necessarily the outer shape); bare
+        storages derive it from ``size()``.
+        """
+        if isinstance(self, torch.Tensor):
+            # pylint: disable=unnecessary-dunder-call
+            return torch._C.TensorBase.shape.__get__(self, type(self))
+        return torch.Size(self.size())
+
     def update_usage(
         self,
         rowwise_usage: Optional[bool] = None,
@@ -133,6 +167,61 @@ 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) -----
+
+    # Collected from the subclasses' :class:`InnerTensor` field annotations; everything
+    # else returned by :meth:`get_metadata` is treated as non-tensor context.
+    _INNER_TENSORS: Tuple[Tuple[str, str], ...] = ()
+
+    def __init_subclass__(cls, **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
+        cls._INNER_TENSORS = _collect_inner_tensor_fields(cls)
+
+    def _flatten_nontensor_kwargs(self) -> Dict[str, Any]:
+        """Non-tensor constructor kwargs (scalars, dtype, quantizer)."""
+        tensor_kwargs = {kwarg for _, kwarg in self._INNER_TENSORS}
+        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._INNER_TENSORS if getattr(self, attr) is not None]
+        ctx = {
+            "cls": type(self),
+            "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 = 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._INNER_TENSORS:
+            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],
@@ -350,6 +439,71 @@ def make_empty(
             result.requires_grad_(True)
         return result
 
+    # ----- Data-free buffer/metadata primitives backing TensorProto -----
+
+    def inner_tensor_specs(
+        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
+        ``_INNER_TENSORS`` and respect the quantizer's usage flags.
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not implement inner_tensor_specs; "
+            "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.inner_tensor_specs(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"],
+            "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 6699cb1ea1..2beb487cd9 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:
 
         return quantizer
 
+    # ----- 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 inner_tensor_specs(
+        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 2e0491d49a..0bdbd8fefc 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
@@ -387,6 +387,36 @@ def supports_only_rowwise_all_gather(self) -> bool:
         """
         return True
 
+    # ----- 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 inner_tensor_specs(
+        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 3045662216..8adcdb5c0c 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
@@ -58,6 +58,38 @@ def copy(self) -> MXFP8Quantizer:
 
         return quantizer
 
+    # ----- 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 inner_tensor_specs(
+        self, shape: Tuple[int, ...]
+    ) -> Dict[str, Tuple[Tuple[int, ...], torch.dtype]]:
+        shape = tuple(shape)
+        buffers: Dict[str, Tuple[Tuple[int, ...], torch.dtype]] = {}
+        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 4614ea53a6..ae40830901 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
@@ -345,6 +345,55 @@ def _canonicalized_amax_reduction_group(self) -> dist_group_type:
     def _get_compatible_recipe(self) -> Union[type[Recipe], None]:
         return NVFP4BlockScaling
 
+    # ----- 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 inner_tensor_specs(
+        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).
+        # Order matches NVFP4TensorStorage._INNER_TENSORS (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"] = (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"] = (
+                type(self).convert_shape_for_fp4(type(self).get_columnwise_shape(shape)),
+                torch.uint8,
+            )
+            buffers["_columnwise_scale_inv"] = (
+                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
+
 
 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 06516d0b9e..f2bf573484 100644
--- a/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py
+++ b/transformer_engine/pytorch/tensor/storage/float8_blockwise_tensor_storage.py
@@ -7,12 +7,12 @@
 from __future__ import annotations
 from collections.abc import Iterable
 import math
-from typing import Optional, Dict, Any, Tuple, Union
+from typing import Annotated, Optional, Dict, Any, Tuple, Union
 import torch
 
 import transformer_engine_torch as tex
 
-from ...quantized_tensor import QuantizedTensorStorage, Quantizer
+from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer
 from .._quantization_helpers import safe_quantized_repr
 
 from ...constants import TE_DType_To_Torch, DType
@@ -119,12 +119,12 @@ class Float8BlockwiseQTensorStorage(QuantizedTensorStorage):
     be instantiated directly for performance-critical internal usage.
     """
 
-    _rowwise_data: Optional[torch.Tensor]
-    _columnwise_data: Optional[torch.Tensor]
+    _rowwise_data: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_data")]
+    _rowwise_scale_inv: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_scale_inv")]
+    _columnwise_data: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_data")]
+    _columnwise_scale_inv: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_scale_inv")]
     _quantizer: Quantizer
     _fp8_dtype: DType
-    _rowwise_scale_inv: Optional[torch.Tensor]
-    _columnwise_scale_inv: Optional[torch.Tensor]
     _is_2D_scaled: bool
 
     def __new__(
diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py
index 429ddde97f..c82821e88f 100644
--- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py
+++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py
@@ -5,12 +5,12 @@
 """Mixin class holding data specific for Float8Tensor"""
 
 from __future__ import annotations
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Annotated, Any, Dict, Optional, Tuple, Union
 import torch
 
 import transformer_engine_torch as tex
 
-from ...quantized_tensor import QuantizedTensorStorage, Quantizer
+from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer
 from .._quantization_helpers import safe_quantized_repr
 
 from ...constants import TE_DType as torch_to_transformer_engine_dtype, TE_DType_To_Torch, DType
@@ -66,15 +66,16 @@ class Float8TensorStorage(QuantizedTensorStorage):
 
     """
 
-    _data: Optional[torch.Tensor]
+    _data: Annotated[Optional[torch.Tensor], InnerTensor("data")]
     _quantizer: Optional[Quantizer]
     _fp8_dtype: DType
-    _scale_inv: torch.Tensor
 
     # FP8 transpose cache
-    _transpose: Optional[torch.Tensor]
+    _transpose: Annotated[Optional[torch.Tensor], InnerTensor("data_transpose")]
     _transpose_invalid: bool
 
+    _scale_inv: Annotated[torch.Tensor, InnerTensor("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 606ac9e74b..e36f7baa37 100644
--- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py
+++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py
@@ -5,14 +5,14 @@
 """Mixin class holding data specific for MXFP8Tensor"""
 
 from __future__ import annotations
-from typing import Optional, Dict, Any, Tuple, Union
+from typing import Annotated, Optional, Dict, Any, Tuple, Union
 from collections.abc import Iterable
 import math
 import torch
 
 import transformer_engine_torch as tex
 
-from ...quantized_tensor import QuantizedTensorStorage, Quantizer
+from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer
 from .._quantization_helpers import safe_quantized_repr
 
 from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType
@@ -66,14 +66,12 @@ class MXFP8TensorStorage(QuantizedTensorStorage):
 
     """
 
-    # Row-scaled FP8 data
-    _rowwise_data: Optional[torch.Tensor]
-    # Column-scaled FP8 data
-    _columnwise_data: Optional[torch.Tensor]
-    # Scaling factors for row-scaled FP8 data
-    _rowwise_scale_inv: torch.Tensor
-    # Scaling factors for column-scaled FP8 data
-    _columnwise_scale_inv: torch.Tensor
+    # Row-scaled FP8 data and its scaling factors
+    _rowwise_data: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_data")]
+    _rowwise_scale_inv: Annotated[torch.Tensor, InnerTensor("rowwise_scale_inv")]
+    # Column-scaled FP8 data and its scaling factors
+    _columnwise_data: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_data")]
+    _columnwise_scale_inv: Annotated[torch.Tensor, InnerTensor("columnwise_scale_inv")]
 
     # Builder class for casting to MXFP8
     _quantizer: Optional[Quantizer]
diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py
index 09f040ba67..7e0861c967 100644
--- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py
+++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py
@@ -8,14 +8,14 @@
 from collections.abc import Iterable
 import functools
 import math
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Annotated, Any, Dict, Optional, Tuple, Union
 import warnings
 
 import torch
 
 import transformer_engine_torch as tex
 
-from ...quantized_tensor import QuantizedTensorStorage, Quantizer
+from ...quantized_tensor import InnerTensor, QuantizedTensorStorage, Quantizer
 from .._quantization_helpers import safe_quantized_repr
 
 from ...constants import TE_DType as torch_to_transformer_engine_dtype, DType
@@ -80,20 +80,15 @@ class NVFP4TensorStorage(QuantizedTensorStorage):
 
     """
 
-    # Row-scaled FP4 data
-    _rowwise_data: Optional[torch.Tensor]
-    # Column-scaled FP4 data
-    _columnwise_data: Optional[torch.Tensor]
-    # Block scaling factors for row-scaled FP4 data
-    _rowwise_scale_inv: torch.Tensor
-    # Block scaling factors for column-scaled FP4 data
-    _columnwise_scale_inv: torch.Tensor
-    # Input absolute maximum value (used to compute tensor scale for
-    # row-scaled FP4 data)
-    _amax_rowwise: torch.Tensor
-    # Input absolute maximum value (used to compute tensor scale for
-    # column-scaled FP4 data)
-    _amax_columnwise: torch.Tensor
+    # Row-scaled FP4 data and its block scaling factors
+    _rowwise_data: Annotated[Optional[torch.Tensor], InnerTensor("rowwise_data")]
+    _rowwise_scale_inv: Annotated[torch.Tensor, InnerTensor("rowwise_scale_inv")]
+    # Column-scaled FP4 data and its block scaling factors
+    _columnwise_data: Annotated[Optional[torch.Tensor], InnerTensor("columnwise_data")]
+    _columnwise_scale_inv: Annotated[torch.Tensor, InnerTensor("columnwise_scale_inv")]
+    # Input absolute maximum values, used to compute the tensor scale
+    _amax_rowwise: Annotated[torch.Tensor, InnerTensor("amax_rowwise")]
+    _amax_columnwise: Annotated[torch.Tensor, InnerTensor("amax_columnwise")]
 
     # Builder class for casting to MXFP8
     _quantizer: Optional[Quantizer]