[PyTorch][torch.compile] Add TensorProto mechanism

[pggPL]

Note

This PR is stacked on top of #3152 ([PyTorch][torch.compile] Make quantizers opaque value objects).
Until #3152 is merged, the diff below also includes its changes. Review/merge #3152 first.

Description

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

The key property is that TensorProto.create_tensor() materializes a quantized tensor purely in Python (via Quantizer.alloc_tensors + 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.

This builds on the value-opaque quantizer work (so a TensorProto 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.py: Add TensorProto dataclass (shape, dtype, quantizer, requires_grad, device) with is_quantized, inner_names(), create_metadata() and create_tensor(), plus a to_tensor_proto() helper that builds a proto from a plain torch.Tensor or a QuantizedTensorStorage/QuantizedTensor.
• quantized_tensor.py:
  • Add the PyTorch wrapper-subclass flatten protocol (__tensor_flatten__ / __tensor_unflatten__) to QuantizedTensorStorage, driven by a per-class _FLATTEN_TENSOR_BUFFERS declaration of (attribute_name, constructor_kwarg) pairs.
  • Add a _STORAGE_REGISTRY (populated via __init_subclass__) so __tensor_unflatten__ can resolve a concrete storage/wrapper class from its qualname inside an FX graph.
  • Add pure-Python, traceable allocation hooks to Quantizer: alloc_tensors, create_metadata, and the opt-in overrides _describe_buffers, _storage_scalars, _resolve_storage_cls.
• Quantizers: Implement the allocation hooks for Float8CurrentScalingQuantizer, MXFP8Quantizer and Float8BlockQuantizer.
• Storage classes: Declare _FLATTEN_TENSOR_BUFFERS for Float8TensorStorage, MXFP8TensorStorage and Float8BlockwiseQTensorStorage.
• ops/basic/basic_linear.py: Add allocation-free _functional_forward_fake / _functional_backward_fake that operate on TensorProto and return output/gradient protos, as a basis for custom-op fake impls (single-device only; TP/SP shape effects not yet modeled).
• Tests: Add tests/pytorch/test_tensor_proto.py (CPU smoke tests for _describe_buffers/alloc_tensors/create_metadata, flatten round-trip, and to_tensor_proto) and torch.compile fullgraph tests in test_torch_compile.py.

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 TensorProto — a data-free descriptor of a tensor or quantized tensor — enabling torch.compile(fullgraph=True) fake implementations of quantized ops by materialising tensors purely in Python via Quantizer.alloc_tensors + QuantizedTensorStorage.__tensor_unflatten__, without touching the opaque C++ create_empty_quantized_tensor.

• dynamo/tensor_proto.py: New TensorProto dataclass with inner_names(), create_metadata(), create_inner_tensors(), and create_tensor(); re-orders _describe_buffers output to the canonical _FLATTEN_TENSOR_BUFFERS slot order so the fake layout aligns with the real op's flatten protocol.
• quantized_tensor.py: Adds __tensor_flatten__/__tensor_unflatten__ to QuantizedTensorStorage (driven by per-class _FLATTEN_TENSOR_BUFFERS), a _STORAGE_REGISTRY for class lookup by qualname inside FX graphs, and _describe_buffers/alloc_tensors/create_metadata allocation hooks to Quantizer.
• module/linear.py: Adds _linear_forward_impl_fake/_linear_backward_impl_fake (skeleton fake impls not yet wired to dispatch), and changes backward_needs_input/columnwise_usage to use captured args.*_requires_grad flags to avoid divergence during CUDA-graph-tree capture where live requires_grad is detached to False.

Confidence Score: 5/5

Safe to merge. The new TensorProto mechanism is purely additive; existing eager paths are unchanged except for the requires_grad flag capture fix in linear.py, which is a correctness improvement for CUDA-graph-tree compilation.

All new code is on non-default code paths: TensorProto is only constructed explicitly, the flatten protocol is only invoked by torch.compile, and the fake forward/backward in linear.py are not yet registered to any dispatch mechanism. The requires_grad flag change in the real _linear_forward_impl correctly follows the intent captured at op-call time. The _STORAGE_REGISTRY registration is idempotent and qualname-stable for top-level classes. Test coverage spans eager, FakeTensorMode, and fullgraph=True compilation for all four quantizer families.

transformer_engine/pytorch/quantized_tensor.py (the __tensor_unflatten__ tensor-path device=None fallback) and transformer_engine/pytorch/module/linear.py (the inp.is_quantized duck-typing in _linear_forward_impl_fake before it is wired to dispatch).

Important Files Changed

Filename	Overview
transformer_engine/pytorch/dynamo/tensor_proto.py	New file introducing TensorProto dataclass and to_tensor_proto helper; correctly copies quantizer in __post_init__, reorders buffers to canonical flatten order in inner_names(), and assembles tensors via __tensor_unflatten__ without storage; one minor docstring inaccuracy about NVFP4 buffer ordering.
transformer_engine/pytorch/quantized_tensor.py	Adds _STORAGE_REGISTRY (populated via __init_subclass__), __tensor_flatten__/__tensor_unflatten__ protocol to QuantizedTensorStorage, and _describe_buffers/alloc_tensors/create_metadata hooks to Quantizer; in the tensor-path of __tensor_unflatten__, device can be None if all inner tensors are None.
transformer_engine/pytorch/module/linear.py	Adds _linear_forward_impl_fake and _linear_backward_impl_fake (not yet wired to dispatch), changes backward_needs_input/columnwise_usage to use captured args.*_requires_grad flags, and expands weight-workspace alias dedup to new_workspace/weight_workspace; inp.is_quantized in the fake forward has a semantic mismatch with real torch.Tensor.is_quantized.
transformer_engine/pytorch/tensor/float8_tensor.py	Adds _storage_metadata and _describe_buffers to Float8CurrentScalingQuantizer; correctly mirrors C++ buffer layout (single _data on non-TN arches) with _scale_inv always present.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Adds _storage_metadata and _describe_buffers to NVFP4Quantizer; calls static methods via type(self) to avoid torch.compile guard issues; amax buffers correctly placed at end matching _FLATTEN_TENSOR_BUFFERS order.
transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py	Adds _FLATTEN_TENSOR_BUFFERS; attribute-to-kwarg mapping consistent with constructor signature.
transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py	Adds _FLATTEN_TENSOR_BUFFERS with all six buffer slots; ordering matches NVFP4Quantizer._describe_buffers.
tests/pytorch/test_torch_compile.py	Comprehensive tests covering primitives, flatten round-trip, and torch.compile(fullgraph=True) traceability for all four quantizer families.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant C as Caller (fake impl)
    participant TP as TensorProto
    participant Q as Quantizer
    participant SR as _STORAGE_REGISTRY
    participant SC as StorageClass

    C->>TP: TensorProto(shape, dtype, quantizer)
    TP->>Q: copy()
    C->>TP: inner_names()
    TP->>Q: _describe_buffers(shape)
    Q-->>TP: "attr -> (shape, dtype)"
    TP->>Q: _storage_metadata(dtype)
    Q-->>TP: cls + nontensor_kwargs
    TP-->>C: canonical buffer names
    C->>TP: create_tensor()
    TP->>Q: create_metadata(shape, dtype)
    Q-->>TP: ctx with cls qualname
    TP->>Q: alloc_tensors(shape, device)
    Q-->>TP: "attr_name -> torch.empty"
    TP->>SR: lookup by qualname
    SR-->>TP: StorageClass
    TP->>SC: __tensor_unflatten__(inner, ctx, shape, stride)
    SC-->>TP: QuantizedTensorStorage / QuantizedTensor
    TP-->>C: assembled tensor

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant C as Caller (fake impl)
    participant TP as TensorProto
    participant Q as Quantizer
    participant SR as _STORAGE_REGISTRY
    participant SC as StorageClass

    C->>TP: TensorProto(shape, dtype, quantizer)
    TP->>Q: copy()
    C->>TP: inner_names()
    TP->>Q: _describe_buffers(shape)
    Q-->>TP: "attr -> (shape, dtype)"
    TP->>Q: _storage_metadata(dtype)
    Q-->>TP: cls + nontensor_kwargs
    TP-->>C: canonical buffer names
    C->>TP: create_tensor()
    TP->>Q: create_metadata(shape, dtype)
    Q-->>TP: ctx with cls qualname
    TP->>Q: alloc_tensors(shape, device)
    Q-->>TP: "attr_name -> torch.empty"
    TP->>SR: lookup by qualname
    SR-->>TP: StorageClass
    TP->>SC: __tensor_unflatten__(inner, ctx, shape, stride)
    SC-->>TP: QuantizedTensorStorage / QuantizedTensor
    TP-->>C: assembled tensor

Loading

_{Reviews (8): Last reviewed commit: "[PyTorch] Workaround torch.compile stati..." | Re-trigger Greptile}

[ptrendx]

I do not see any tests that would ensure that this and the actual C++ allocation produces compatible
results.