[PyTorch] [torch.compile] torch.compile support for Linear

[pggPL]

Description

This PR adds torch.compile support for te.pytorch.Linear, building on the TensorSpec mechanism already in main.

_Linear's forward and backward are registered as torch.library custom ops, so a module containing te.Linear traces under torch.compile(fullgraph=True) without graph breaks. The fake (meta) implementations describe the produced tensors through TensorSpec instead of allocating them, which is what makes the quantized outputs traceable — the compiler sees the full quantized-tensor structure (data, scales, transposes) without any device allocation at trace time.

The bulk of the diff is transformer_engine/pytorch/dynamo/custom_op.py: a declarative register_custom_op helper. Custom ops require flat lists of tensors, while the TE forward/backward take dataclass "argument bundles" holding tensors, quantized tensors, quantizers, process groups and plain Python values. The helper derives the op schema from the dataclass field annotations, flattens each field to op slots via a per-kind adapter, and rebuilds the bundle on the other side, so ops are declared by writing a dataclass rather than by hand-maintaining a schema string.

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/custom_op.py (new): register_custom_op — declarative registration of forward/backward custom ops from dataclass argument bundles. Handles per-field adapters for plain tensors, quantized tensors, quantizers, opaque value bundles and reference-opaque types (e.g. process groups), schema generation, TensorSpec-based fake outputs and autograd wiring. Falls back to eager with a single warning if registration fails.
• dynamo/__init__.py: export register_custom_op.
• module/linear.py: split the forward into pure computation and context saving, add allocation-free fake forward/backward on TensorSpec, and register _Linear through register_custom_op. Eager behavior is unchanged.
• dynamo/quantizer_opaque.py, dynamo/tensor_spec.py, tensor/_quantization_helpers.py, tensor/float8_tensor.py, tensor/storage/float8_tensor_storage.py, utils.py: small supporting changes (idempotent spec conversion, weight-workspace quantizer preservation, keeping attributes attached to quantized parameters across _apply).
• tests/pytorch/test_torch_compile.py: coverage for the compiled Linear — fullgraph compilation, quantized FP8 weights, FP8 output, is_first_microbatch, dynamic shapes, parametrized over the supported recipes (FP8 per-tensor/current scaling, MXFP8, NVFP4).
• tests/pytorch/distributed/*: exercise the compiled path in the distributed numerics and comm-GEMM-overlap runs.

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 adds full-graph torch.compile support for PyTorch Linear through a declarative custom-op framework and allocation-free fake implementations.

• Registers compiled forward and backward operations with structured argument flattening and autograd integration.
• Preserves quantized tensor, quantizer, weight-workspace, and distributed process-group state across compiled execution.
• Adds single-GPU, dynamic-shape, quantized-output, CUDA-graph, and distributed overlap coverage.

Confidence Score: 5/5

The PR appears safe to merge based on the eligible follow-up findings available for this review.

No blocking failure remains in the available follow-up review scope.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/dynamo/custom_op.py	Introduces the generic custom-op registration, argument adapters, fake execution, output reconstruction, and autograd plumbing used by compiled TE modules.
transformer_engine/pytorch/module/linear.py	Refactors Linear computation and context management into real and TensorSpec-based fake implementations and selects the registered custom op during compilation.
transformer_engine/pytorch/dynamo/tensor_spec.py	Extends allocation-free tensor specifications used to represent ordinary and quantized outputs during fake tracing.
transformer_engine/pytorch/dynamo/quantizer_opaque.py	Supports value-opaque quantizer representation required for stable compiled graph specialization.
transformer_engine/pytorch/tensor/float8_tensor.py	Updates Float8 tensor metadata and quantizer handling to support custom-op flattening and reconstruction.
transformer_engine/pytorch/utils.py	Preserves externally attached parameter attributes when module transformations replace or mutate parameter objects.
tests/pytorch/test_torch_compile.py	Adds compiled Linear coverage across quantization recipes, CUDA-graph mode, quantized weights and outputs, microbatch caching, gradients, and dynamic shapes.
tests/pytorch/distributed/run_numerics.py	Adds full-graph compilation to distributed Linear numerical comparisons using static shapes.
tests/pytorch/distributed/test_comm_gemm_overlap.py	Exercises compiled Linear with Userbuffers communication and GEMM overlap in default and reduced-overhead modes.

Sequence Diagram

sequenceDiagram
    participant User
    participant Linear
    participant Wrapper as Custom-op wrapper
    participant Base as Base custom op
    participant Kernel as Linear implementation
    participant Autograd

    User->>Linear: forward(input)
    Linear->>Wrapper: structured LinearFwdArgs
    Wrapper->>Wrapper: flatten tensors, quantized storage, and opaque values
    Wrapper->>Base: flat custom-op inputs
    Base->>Kernel: reconstructed argument bundle
    Kernel-->>Base: output, workspace, and saved tensors
    Base-->>Wrapper: flat Tensor[] payload
    Wrapper-->>Linear: reconstructed outputs
    Linear-->>User: output
    User->>Autograd: backward(grad_output)
    Autograd->>Base: saved state and flat gradient inputs
    Base->>Kernel: Linear backward implementation
    Kernel-->>Autograd: input, weight, and bias gradients

Loading

_{Reviews (2): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}