[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 prototype of a tensor that captures shape, dtype, and quantizer, enabling pure-Python quantized-tensor allocation via __tensor_flatten__/__tensor_unflatten__. This is foundational infrastructure for writing torch.library custom-op fake implementations that trace under torch.compile(fullgraph=True) without graph breaks.

• TensorProto (dynamo/tensor_proto.py): new dataclass with create_tensor() / inner_names() / create_metadata() that materialises quantized tensors purely in Python through Quantizer.alloc_tensors + storage __tensor_unflatten__.
• QuantizedTensorStorage (quantized_tensor.py): adds __tensor_flatten__/__tensor_unflatten__ protocol, _FLATTEN_TENSOR_BUFFERS class-level declarations, and _STORAGE_REGISTRY auto-populated via __init_subclass__; Quantizer gains _describe_buffers, alloc_tensors, and create_metadata pure-Python allocation hooks.
• Quantizers (float8_tensor.py, mxfp8_tensor.py, float8_blockwise_tensor.py, nvfp4_tensor.py): implement the allocation hooks and register as value-opaque types; NVFP4Quantizer adds _rebuild_derived_state to restore rht_matrix after FX reconstruction; NVFP4Quantizer also adds _with_random_sign_mask field.
• BasicLinear (module/linear.py): preparatory _linear_forward_impl_fake / _linear_backward_impl_fake shape-propagation stubs for future custom-op fake registrations.

Confidence Score: 4/5

Safe to merge after fixing the wrong flag in _rebuild_derived_state; the copy() omission of with_random_sign_mask is a related follow-up. The core TensorProto and flatten/unflatten machinery is correct and well-tested.

The NVFP4Quantizer._rebuild_derived_state method calls get_rht_matrix(self._with_random_sign_mask, ...), but the original __init__ calls get_rht_matrix(with_rht, ...). These are distinct parameters: with_rht gates whether the transform is used at all; _with_random_sign_mask only controls the sign pattern within it. When they differ — e.g. with_rht=True, with_random_sign_mask=False — the rebuilt quantizer silently receives a different matrix and produces wrong quantization results. The default path (both True) passes unchanged, which is why the existing tests don't catch it.

transformer_engine/pytorch/tensor/nvfp4_tensor.py — both _rebuild_derived_state (wrong flag) and copy() (missing with_random_sign_mask forwarding) need attention.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/tensor/nvfp4_tensor.py	Adds _value_fields, _storage_metadata, _describe_buffers, _rebuild_derived_state, and registration. _rebuild_derived_state uses self._with_random_sign_mask as the flag to get_rht_matrix, but init uses with_rht — when they differ these produce different matrices. copy() also omits with_random_sign_mask, leaving the copy with the wrong value for that field.
transformer_engine/pytorch/dynamo/tensor_proto.py	New file implementing TensorProto dataclass; core logic is sound. Minor: storage_cls lookup in create_tensor is redundant since tensor_unflatten re-resolves from the context dict.
transformer_engine/pytorch/quantized_tensor.py	Adds flatten/unflatten protocol and _STORAGE_REGISTRY to QuantizedTensorStorage; the init_subclass registration and tensor_unflatten static dispatch are correct. Registry uses bare qualname without module path, which could silently collide for two classes with identical names across modules.
transformer_engine/pytorch/dynamo/quantizer_opaque.py	Adds _rebuild_quantizer with _rebuild_derived_state hook mechanism. The process-group guard and LRU-cache-based device-local reconstruction are well thought out.
transformer_engine/pytorch/module/linear.py	Adds _linear_forward_impl_fake and _linear_backward_impl_fake shape-propagation stubs (not yet wired to a custom op). The requires_grad fix (args.*_requires_grad instead of live tensor flags) is correct and important. New alias tracking for the cached FP8 weight is correctly handled in both real and fake paths.
tests/pytorch/test_torch_compile.py	Comprehensive tests for value-opaque quantizers, TensorProto primitives, flatten/unflatten roundtrips, and fullgraph compile. Tests cover the default NVFP4 path (both with_rht=True and with_random_sign_mask=True) but do not exercise the with_rht=True/with_random_sign_mask=False case that exposes the _rebuild_derived_state bug.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant C as Caller / register_fake
    participant TP as TensorProto
    participant Q as Quantizer
    participant S as QuantizedTensorStorage
    participant R as _STORAGE_REGISTRY

    Note over C,R: TensorProto.create_tensor() flow
    C->>TP: TensorProto(shape, dtype, quantizer)
    TP->>Q: copy() [__post_init__]

    C->>TP: create_tensor()
    TP->>TP: create_metadata()
    TP->>Q: create_metadata(shape, dtype)
    Q->>Q: _storage_metadata(dtype)
    Q-->>TP: "ctx {cls: qualname, nontensor_kwargs, ...}"

    TP->>TP: create_inner_tensors()
    TP->>Q: alloc_tensors(shape, device)
    Q->>Q: _describe_buffers(shape)
    Q-->>TP: "{attr: torch.empty(buf_shape, buf_dtype)}"

    TP->>TP: inner_names() [reorder to _FLATTEN_TENSOR_BUFFERS order]
    TP->>R: _STORAGE_REGISTRY[ctx["cls"]]
    R-->>TP: storage_cls

    TP->>S: __tensor_unflatten__(inner, ctx, shape, stride)
    S->>R: _STORAGE_REGISTRY[ctx["cls"]]
    R-->>S: cls
    S-->>C: QuantizedTensorStorage instance

    Note over C,R: FX graph rebuild via _rebuild_quantizer
    C->>C: eval(quantizer.__fx_repr__())
    C->>C: _rebuild_quantizer(cls, items)
    C->>C: object.__setattr__(obj, field, value) for each field
    C->>Q: _rebuild_derived_state() [NVFP4 only]
    Q->>Q: get_rht_matrix(_with_random_sign_mask, device)

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 / register_fake
    participant TP as TensorProto
    participant Q as Quantizer
    participant S as QuantizedTensorStorage
    participant R as _STORAGE_REGISTRY

    Note over C,R: TensorProto.create_tensor() flow
    C->>TP: TensorProto(shape, dtype, quantizer)
    TP->>Q: copy() [__post_init__]

    C->>TP: create_tensor()
    TP->>TP: create_metadata()
    TP->>Q: create_metadata(shape, dtype)
    Q->>Q: _storage_metadata(dtype)
    Q-->>TP: "ctx {cls: qualname, nontensor_kwargs, ...}"

    TP->>TP: create_inner_tensors()
    TP->>Q: alloc_tensors(shape, device)
    Q->>Q: _describe_buffers(shape)
    Q-->>TP: "{attr: torch.empty(buf_shape, buf_dtype)}"

    TP->>TP: inner_names() [reorder to _FLATTEN_TENSOR_BUFFERS order]
    TP->>R: _STORAGE_REGISTRY[ctx["cls"]]
    R-->>TP: storage_cls

    TP->>S: __tensor_unflatten__(inner, ctx, shape, stride)
    S->>R: _STORAGE_REGISTRY[ctx["cls"]]
    R-->>S: cls
    S-->>C: QuantizedTensorStorage instance

    Note over C,R: FX graph rebuild via _rebuild_quantizer
    C->>C: eval(quantizer.__fx_repr__())
    C->>C: _rebuild_quantizer(cls, items)
    C->>C: object.__setattr__(obj, field, value) for each field
    C->>Q: _rebuild_derived_state() [NVFP4 only]
    Q->>Q: get_rht_matrix(_with_random_sign_mask, device)

Loading

Comments Outside Diff (1)

1. transformer_engine/pytorch/tensor/nvfp4_tensor.py, line 231-245 (link)

   copy() does not forward with_random_sign_mask to the constructor, so the copy always gets _with_random_sign_mask = True (the default) regardless of the original's value. If _rebuild_derived_state is fixed to use self.with_rht, this inconsistency becomes harmless for reconstruction; but a copy used directly (without going through FX reconstruction) would still carry the wrong _with_random_sign_mask and rht_matrix_random_sign_mask_t, which can affect future behavior if those fields are ever inspected independently.

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

[greptile-apps]

_rebuild_derived_state passes self._with_random_sign_mask to get_rht_matrix, but __init__ passes with_rht. These are distinct parameters: with_rht controls whether the Hadamard transform is applied at all; _with_random_sign_mask controls whether random signs are used within it. When they differ (e.g. with_rht=True, with_random_sign_mask=False) the rebuilt quantizer gets a different matrix than the original. Since get_rht_matrix is LRU-cached by (flag, device), get_rht_matrix(True, d) and get_rht_matrix(False, d) return different objects, so the kernel would receive the wrong transform matrix and silently produce incorrect quantization results.

Suggested change

	self.rht_matrix = get_rht_matrix(self._with_random_sign_mask, torch.cuda.current_device())
	self.rht_matrix = get_rht_matrix(self.with_rht, torch.cuda.current_device())