[PyTorch][torch.compile] Add TensorProto mechanism

[pggPL]

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 prototype of a tensor or quantized tensor — as infrastructure for torch.library custom-op fake implementations that trace under torch.compile(fullgraph=True). It also adds the PyTorch subclass flatten/unflatten protocol to QuantizedTensorStorage and pure-Python traceable allocation hooks to Quantizer.

• TensorProto (dynamo/tensor_proto.py): new dataclass capturing shape, dtype, and an optional value-opaque quantizer; create_tensor() materializes a quantized tensor via alloc_tensors + __tensor_unflatten__ entirely in Python, enabling fullgraph tracing.
• QuantizedTensorStorage flatten protocol (quantized_tensor.py): adds _STORAGE_REGISTRY (via __init_subclass__), _FLATTEN_TENSOR_BUFFERS on all four storage classes, and a shared __tensor_flatten__ / __tensor_unflatten__ implementation.
• Fake linear impls (module/linear.py): allocation-free shape-only twins returning TensorProto descriptors; not yet registered as custom-op fake impls (follow-up PR), plus correctness fix capturing requires_grad into args rather than reading from live tensors under compile.

Confidence Score: 5/5

Safe to merge. Core allocation and flatten/unflatten paths are exercised by comprehensive CPU smoke tests including fullgraph compile checks.

Structural parity tests validate that Python-allocated tensors match C++ make_empty in buffer layout and kernel output. The only outstanding items are a private PyTorch API import and implicit stride forwarding via **kwargs, both non-blocking style issues. The fake forward/backward impls are intentionally incomplete and cannot be called incorrectly in the current state.

transformer_engine/pytorch/dynamo/tensor_proto.py (private API import) and transformer_engine/pytorch/quantized_tensor.py (tensor_unflatten stride forwarding) are worth a second look before the follow-up custom-op registration PR lands.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/dynamo/tensor_proto.py	New TensorProto dataclass with create_tensor/inner_names/create_metadata; uses private torch._prims_common.make_contiguous_strides_for where the tests use the public meta-device alternative.
transformer_engine/pytorch/quantized_tensor.py	Adds _STORAGE_REGISTRY, init_subclass registration, tensor_flatten/tensor_unflatten protocol, and Quantizer allocation primitives; stride forwarding via **kwargs is implicit.
transformer_engine/pytorch/module/linear.py	Adds _linear_forward_impl_fake/_linear_backward_impl_fake, fixes requires_grad capture, and extends weight alias dedup; fake impls not yet registered.
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Adds _storage_metadata and _describe_buffers for NVFP4Quantizer; uses type(self).convert_shape_for_fp4 to avoid torch.compile guard issues.
transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py	Declares _FLATTEN_TENSOR_BUFFERS for all six NVFP4 buffers including amax buffers.
tests/pytorch/test_torch_compile.py	Comprehensive tests for allocation primitives, flatten/unflatten roundtrip, structural and functional parity with C++ make_empty, and TensorProto API.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant User as User / torch.compile
    participant TP as TensorProto
    participant Q as Quantizer
    participant SR as _STORAGE_REGISTRY
    participant S as QuantizedTensorStorage

    User->>TP: TensorProto(shape, dtype, quantizer)
    TP->>Q: copy() quantizer
    User->>TP: create_tensor()
    TP->>Q: create_metadata(shape, dtype)
    Q-->>TP: "ctx {cls: qualname, nontensor_kwargs}"
    TP->>Q: alloc_tensors(shape, device)
    Q->>Q: _describe_buffers(shape)
    Q-->>TP: "{attr: torch.empty(...)}"
    TP->>TP: inner_names() ordered attr list
    TP->>SR: _STORAGE_REGISTRY[ctx[cls]]
    SR-->>TP: storage_cls
    TP->>S: storage_cls.__tensor_unflatten__(inner, ctx, shape, stride)
    S-->>User: QuantizedTensor / QuantizedTensorStorage
    User->>S: __tensor_flatten__()
    S->>S: _FLATTEN_TENSOR_BUFFERS present attrs
    S->>S: _flatten_nontensor_kwargs() via get_metadata()
    S-->>User: (names, ctx)

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 User as User / torch.compile
    participant TP as TensorProto
    participant Q as Quantizer
    participant SR as _STORAGE_REGISTRY
    participant S as QuantizedTensorStorage

    User->>TP: TensorProto(shape, dtype, quantizer)
    TP->>Q: copy() quantizer
    User->>TP: create_tensor()
    TP->>Q: create_metadata(shape, dtype)
    Q-->>TP: "ctx {cls: qualname, nontensor_kwargs}"
    TP->>Q: alloc_tensors(shape, device)
    Q->>Q: _describe_buffers(shape)
    Q-->>TP: "{attr: torch.empty(...)}"
    TP->>TP: inner_names() ordered attr list
    TP->>SR: _STORAGE_REGISTRY[ctx[cls]]
    SR-->>TP: storage_cls
    TP->>S: storage_cls.__tensor_unflatten__(inner, ctx, shape, stride)
    S-->>User: QuantizedTensor / QuantizedTensorStorage
    User->>S: __tensor_flatten__()
    S->>S: _FLATTEN_TENSOR_BUFFERS present attrs
    S->>S: _flatten_nontensor_kwargs() via get_metadata()
    S-->>User: (names, ctx)

Loading

_{Reviews (13): Last reviewed commit: "Merge branch 'main' of https://github.co..." | Re-trigger Greptile}

[shino16]

Note: We will update this comment as part of the custom op PR.