[PyTorch][torch.compile] Add TensorProto mechanism

[pggPL]

Description

This PR introduces TensorSpec — a data-free description of a tensor (or quantized tensor) that captures everything needed to rebuild it without holding any storage: its logical shape/dtype and, for quantized tensors, the value-opaque quantizer that defines the buffer layout.

The key property is that TensorSpec.create_tensor() materializes a quantized tensor purely in Python — via Quantizer.alloc_tensors plus the storage's __tensor_unflatten__ — so it traces under torch.compile(fullgraph=True) with no graph break, unlike make_empty, which goes through the opaque C++ tex.create_empty_quantized_tensor. This is the foundation for writing torch.library custom-op fake implementations of quantized ops; the consumers land in the follow-up Linear custom-op PR.

This builds on the value-opaque quantizer work, so a TensorSpec is itself safe to treat as a compile-time constant.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

dynamo/tensor_spec.py (new) — TensorSpec dataclass (shape, dtype, quantizer, requires_grad, device) with is_quantized, update_usage(), inner_names(), create_metadata(), create_inner_tensors(), assemble() and create_tensor(), plus a to_tensor_spec() helper that builds a spec from a plain torch.Tensor, a QuantizedTensorStorage or a QuantizedTensor. Exported from transformer_engine.pytorch.dynamo.

quantized_tensor.py

• Add the PyTorch wrapper-subclass flatten protocol (__tensor_flatten__ / __tensor_unflatten__) to QuantizedTensorStorage.
• Declare the flat buffers on the field itself: _scale_inv: Annotated[torch.Tensor, InnerTensor("fp8_scale_inv")]. __init_subclass__ collects these into _INNER_TENSORS in field order, so the attribute-to-constructor-kwarg mapping lives next to the attribute instead of in a parallel tuple.
• The flatten context carries the storage class itself, so __tensor_unflatten__ needs no registry lookup.
• Add pure-Python, traceable allocation primitives to Quantizer: inner_tensor_specs (buffer geometry), storage_metadata (concrete class + non-tensor constructor kwargs), and the alloc_tensors / create_metadata built on top of them. The base implementations raise NotImplementedError, so a quantizer that does not implement them simply cannot be used with TensorSpec.
• Add a shape property that is valid on bare storages as well as wrapper tensors.

Quantizers — implement inner_tensor_specs and storage_metadata for Float8CurrentScalingQuantizer, MXFP8Quantizer, Float8BlockQuantizer and NVFP4Quantizer. The FP8 description mirrors the C++ allocation in csrc/quantizer.cpp, including the non-TN-capable-arch case where a single _data buffer backs both directions.

Storage classes — declare InnerTensor fields for Float8TensorStorage, MXFP8TensorStorage, Float8BlockwiseQTensorStorage and NVFP4TensorStorage.

module/base.py — override nn.Module._apply in TransformerEngineBaseModule. This is a consequence of the flatten protocol, not part of the new API: once a parameter implements it, _apply moves the parameter with torch.utils.swap_tensors, which exchanges its whole __dict__. Inner buffers ride across correctly, but state attached from the outside (_high_precision_init_val and its accessors, main_grad, user attributes) would be left behind on the discarded tensor. The override snapshots those attributes and re-attaches the ones the swap did not carry over, restoring the pre-PR behaviour of .to() / .cuda() / .half().

Tests

• tests/pytorch/test_torch_compile.py: quantizer primitives under FakeTensorMode, storage flatten/unflatten round-trip, TensorSpec behaviour in eager and fake mode, fullgraph=True tracing, and to_tensor_spec round-trips — across FP8 current scaling, MXFP8, FP8 blockwise and NVFP4.
• test_python_alloc_matches_cpp_make_empty builds the same tensor twice, via make_empty (C++) and via the Python primitives, then checks structural parity (class, buffer set, per-buffer shape/dtype/device, logical shape/dtype, flatten context) and functional parity — the real quantize kernel writes bit-identical results into both — across quantizer families x rowwise/columnwise x wrapper/internal.
• tests/pytorch/test_sanity.py: attributes attached to a quantized parameter survive nn.Module._apply for .cuda(), .cpu() and .half().

Known limitations

• The flatten protocol covers the four storage classes listed above. HybridQuantizedTensorStorage and IdentityTensorStorage declare no InnerTensor fields, so flattening them raises rather than silently passing their buffers through the context; hybrid storage holds nested storages rather than flat buffers, which the current model does not express.
• HybridQuantizer and IdentityQuantizer are not registered as value-opaque quantizers.
• Single-device only; tensor/sequence-parallel shape effects are not modelled.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR introduces TensorSpec — a data-free description of a tensor (plain or quantized) that captures shape, dtype, and quantizer, enabling torch.compile(fullgraph=True)-traceable allocation of quantized tensors via pure Python (alloc_tensors + storage __tensor_flatten__/__tensor_unflatten__) rather than the opaque C++ create_empty_quantized_tensor. It also fixes a swap_tensors-driven attribute loss on TransformerEngineBaseModule._apply.

• TensorSpec / to_tensor_spec (dynamo/tensor_spec.py): data-free spec with create_tensor() that traces under fullgraph=True; to_tensor_spec is documented to work on bare QuantizedTensorStorage but fails with AttributeError because tensor.device is not defined on bare storages.
• Flatten protocol (quantized_tensor.py): InnerTensor annotation marker drives __init_subclass__-collected _INNER_TENSORS; __tensor_flatten__/__tensor_unflatten__ on QuantizedTensorStorage; Quantizer allocation hooks implemented for all four quantizer types.
• _apply fix (module/base.py): snapshots parameter __dict__ before _apply and re-attaches missing keys after swap_tensors; handles only _parameters, not _buffers.

Confidence Score: 4/5

Safe to merge with a targeted fix; to_tensor_spec on a bare storage will raise AttributeError at runtime.

The core TensorSpec machinery, flatten protocol, and _apply fix are well-designed and well-tested. The one concrete defect is in to_tensor_spec: it accesses tensor.device unconditionally, but bare QuantizedTensorStorage objects have no device attribute. The PR added a shape property and a _dtype fallback for bare storages, making the missing device an oversight. The current test suite avoids this path, so there is no regression guard.

Files Needing Attention: transformer_engine/pytorch/dynamo/tensor_spec.py needs a device guard for bare storages. transformer_engine/pytorch/quantized_tensor.py — the get_type_hints ordering assumption in _collect_inner_tensor_fields should be documented for future subclass authors.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/dynamo/tensor_spec.py	New TensorSpec dataclass and to_tensor_spec helper. to_tensor_spec fails with AttributeError when called on a bare QuantizedTensorStorage because it accesses tensor.device which is not defined on bare storages; shape was covered but device was not.
transformer_engine/pytorch/quantized_tensor.py	Adds InnerTensor annotation marker, _collect_inner_tensor_fields, QuantizedTensorStorage.shape property, __init_subclass__-driven _INNER_TENSORS, __tensor_flatten__/__tensor_unflatten__, and pure-Python allocation primitives on Quantizer. Implementation is correct but relies on get_type_hints MRO order for _INNER_TENSORS which is implicitly fragile for future subclasses.
transformer_engine/pytorch/module/base.py	Adds _apply override to re-attach externally-bound attributes lost during swap_tensors. Handles parameters only; registered QuantizedTensorStorage buffers are not snapshotted.
transformer_engine/pytorch/tensor/float8_tensor.py	Adds storage_metadata and inner_tensor_specs to Float8CurrentScalingQuantizer, correctly handling the non-TN (Blackwell+) path where _data backs both usages.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Adds storage_metadata and inner_tensor_specs to NVFP4Quantizer. Uses type(self) to call @staticmethods (correct workaround for PyTorch guard issue #182741).
tests/pytorch/test_torch_compile.py	Adds comprehensive TensorSpec tests covering all four quantizer types and three usage combos with fullgraph=True compile validation.

Sequence Diagram

sequenceDiagram
    participant User
    participant TensorSpec
    participant Quantizer
    participant Storage as QuantizedTensorStorage

    User->>TensorSpec: TensorSpec(shape, dtype, quantizer)
    TensorSpec->>Quantizer: copy() [isolate usage mutations]

    User->>TensorSpec: create_tensor()
    TensorSpec->>TensorSpec: create_inner_tensors()
    TensorSpec->>Quantizer: alloc_tensors(shape, device)
    Quantizer->>Quantizer: inner_tensor_specs(shape)
    Quantizer-->>TensorSpec: "{attr: Tensor} inner tensors"

    TensorSpec->>TensorSpec: assemble(inner_tensors)
    TensorSpec->>Quantizer: create_metadata(shape, dtype)
    Quantizer->>Quantizer: storage_metadata(dtype)
    Quantizer-->>TensorSpec: "ctx {cls, is_tensor, nontensor_kwargs}"

    TensorSpec->>Storage: cls.__tensor_unflatten__(inner, ctx, shape, stride)
    Storage-->>TensorSpec: QuantizedTensor / QuantizedTensorStorage
    TensorSpec-->>User: materialized tensor (FakeTensor under FakeTensorMode)

Loading

_{Reviews (22): Last reviewed commit: "Merge branch 'main' into tensor_proto_me..." | Re-trigger Greptile}

[pggPL]

/te-ci pytorch

[pggPL]

/te-ci pytorch

[pggPL]

/te-ci pytorch L1

[pggPL]

/te-ci pytorch L1

[ptrendx]

A more general question - considering that PyTorch went with Tensor/FakeTensor naming, shouldn't we
follow suit with QuantizedTensor/FakeQuantizedTensor rather than introducing a completely new name?

[pggPL]

Not really, because TensorProto represents both Tensor and QuantizedTensor. I changed the name to TensorSpec.

[ptrendx]

/te-ci pytorch