[PyTorch] [torch.compile] torch.compile support for Linear by pggPL · Pull Request #3053 · NVIDIA/TransformerEngine

pggPL · 2026-05-28T14:03:55Z

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

Register the Linear forward/backward as torch.library custom ops on top of the TensorSpec mechanism (NVIDIA#3153), so Linear traces under fullgraph compile with FP8/MXFP8/NVFP4 recipes. - transformer_engine/pytorch/dynamo/custom_op.py: custom-op registration framework (arg bundles, fake impls, autograd wiring) - module/linear.py: split forward into compute + ctx save, fake forward/backward - tests/pytorch/test_torch_compile.py: coverage for the compiled path Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

for more information, see https://pre-commit.ci

black wrapped the 122-char except clause, moving Exception onto its own line while the disable comment stayed on the closing paren, so pylint's W0718 no longer saw it. Shorten the line instead. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

greptile-apps · 2026-08-05T17:08:41Z

Greptile Summary

The PR adds full-graph torch.compile support for PyTorch Linear through a declarative custom-op framework and TensorSpec-based fake execution.

Registers flattened forward and backward custom ops with autograd integration.
Models plain and quantized tensor structures without trace-time device allocation.
Preserves quantized parameter metadata and cached weight workspaces across compiled execution.
Adds single-GPU and distributed compile coverage across supported quantization recipes, dynamic shapes, microbatch caching, and communication overlap.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains from the previously reported review findings.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/dynamo/custom_op.py	Introduces the declarative custom-op registration framework, adapters, TensorSpec fake execution, structured flattening, and autograd wiring.
transformer_engine/pytorch/module/linear.py	Splits Linear computation and context handling into eager and fake implementations and selects the registered custom op while compiling.
transformer_engine/pytorch/dynamo/quantizer_opaque.py	Supports reconstruction and value-opaque representation of quantizers used as compiled graph constants.
transformer_engine/pytorch/tensor/float8_tensor.py	Makes Float8 tensor representation safe for fake or non-materialized internal scale tensors.
transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py	Applies safe representation handling to bare Float8 storage objects.
transformer_engine/pytorch/utils.py	Adds compile fallback warnings and graph-compatible GEMM dimension validation.
tests/pytorch/test_torch_compile.py	Expands compiled Linear coverage across quantization, cached weights, FP8 output, dynamic shapes, and microbatch behavior.
tests/pytorch/distributed/run_numerics.py	Adds full-graph compilation modes to distributed Linear numerical comparisons.
tests/pytorch/distributed/run_layer_with_overlap.py	Adds compiled execution to the Userbuffers layer-overlap test harness.
tests/pytorch/distributed/test_comm_gemm_overlap.py	Adds BF16 compiled Linear coverage for row- and column-parallel communication-overlap configurations.

Sequence Diagram

sequenceDiagram
  participant User
  participant Linear
  participant Wrapper as Custom-op wrapper
  participant Base as Base custom op
  participant Fake as TensorSpec fake implementation
  participant CUDA as TE GEMM implementation
  participant Autograd

  User->>Linear: forward(input)
  Linear->>Wrapper: flattened argument bundle
  Wrapper->>Fake: describe outputs and saved tensors
  Fake-->>Wrapper: TensorSpec structures
  Wrapper->>Base: plain tensors, inner buffers, opaque values
  Base->>CUDA: execute Linear forward
  CUDA-->>Base: output, workspace, saved payload
  Base-->>Linear: reconstructed outputs
  Linear-->>User: output

  User->>Autograd: backward(grad_output)
  Autograd->>Base: reconstructed backward bundle
  Base->>CUDA: execute Linear backward
  CUDA-->>Autograd: weight, input, and bias gradients

_{Reviews (4): Last reviewed commit: "Restructure register_custom_op docstring..." | Re-trigger Greptile}

Naming consistency and de-duplication in the torch.compile custom-op framework and its Linear user. No functional change. Naming: - unify the register_custom_op API on fwd_*/bwd_* (backward_arg_type, backward_impl, backward_obj_type -> bwd_arg_type, bwd_impl) - _register_kernel -> _register_base_op, pairing with _register_wrapper_op - _format_*_result / _split_fwd_fake_result -> _pack_*_result / _unpack_fwd_fake_result - _value_to_flat_tensors / _spec_reassemble -> _flatten_value / _unflatten_value, matching _storage_flatten / _storage_unflatten - adapter slots: tensor_slot / inner_slot / meta_slot, META_SLOT, QUANTIZER_KEY - _linear_backward -> _linear_backward_impl and *_fake twins, so the real and fake implementations pair up by name - ctx attrs: drop the lone _te_ prefix, and use ctx.backward_objects as the eager path already does - move warn_compile_unsupported to utils as warn_compile_disabled, next to warn_compile_eager_fallback, so the two "unsupported" meanings are distinguishable - move the TensorOrQuantized alias next to the adapter that matches it De-duplication: - _unflatten_values() replaces three copies of the cursor/reassemble loop - _make_slot_forwarder() / _make_dispatch_rule() replace three copies of the subclass-flattening forward path - _sp_out_leading() / _sp_inp_leading() replace three copies of the sequence-parallel leading-dim arithmetic (two of them inverses) - check_gemm_dims() moves the fp8 dimension checks to utils - drop the duplicate backward_needs_input assignment in the forward impl Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

for more information, see https://pre-commit.ci

… dim checks and cleanups - check_gemm_dims: restore assert_dim_for_fp8_exec semantics (per-tensor leading%8 / last%16, out_features%8 not %16); rich error messages with dims on the eager path, constant torch._check messages under compile (Dynamo forbids tensor closures in _check message lambdas). - test_te_linear_dynamic_shapes: the recompile assertion compared a nonexistent counter (always 0==0); use stats/unique_graphs and absorb the one-time lazy is_fsdp2 hasattr-guard recompile with a warmup. - custom_op: None-sentinel dtype uint8 -> complex32; a genuinely empty FP8 uint8 buffer (batch=0) decoded as None and broke compilation. - OpaqueValueBundle: type-tag _to_hashable (list/tuple/Size no longer compare equal), guard __getattr__ against copy/pickle recursion on underscored probes, render non-finite floats evaluably in __fx_repr__. - Linear.forward: fetch the cuBLAS workspace only after the eager-fallback decision; explicit torch._dynamo.graph_break(msg=...) so fullgraph=True errors carry the fallback reason instead of breaking on warnings.warn. - warn_compile_disabled: move the 'use a newer PyTorch build' advice to the version-related call sites only. - Comment/docstring/typography/pylint-disable cleanups in custom_op; test cosmetics (use_compile arg name, argparse-time validation of --compile/--use-cuda-graphs, merged NVINSPECT skips, docstring fixes); export get_cublas_workspace from cpp_extensions. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

… eager dim asserts - check_gemm_dims is now a compile-only torch._check guard emitter, called from the compiled-op branch; eager dim validation returns to the op impl (assert + assert_dim_for_fp8_exec, as on main) so eager pays no overhead and keeps full error messages with dims. - Trim verbose test docstrings/comments (te.Linear section, warmup helper, cudagraph-skip helper); describe the dynamic-shape scope (leading dims) instead of the fix history. - Drop the stale 'FP8 with symbolic shapes unsupported' comments: FP8 with a mark_dynamic batch works on current nightly (verified: one graph reused across batch sizes, numerics match eager). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…wo float8 reprs Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

_sp_out_leading/_sp_inp_leading -> _out_leading_from_inp/_inp_leading_from_out; shorten the weight_workspace field comment; drop the to_tensor_spec caveat paragraph. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…navailable PG_REFERENCE_OPAQUE is computed once at import (Dynamo-friendly constant); compile_unsupported_reason reports a tp_group it cannot carry instead of the misleading _UnsupportedAdapter TypeError at trace time. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

… module-docstring duplication Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

ptrendx · 2026-08-10T18:44:56Z

+    if args.compile and args.use_cuda_graphs:
+        parser.error("--compile and --use-cuda-graphs are mutually exclusive.")


So how do we test the lower overhead mode? Shouldn't we just do that mode of the compilation if this
option is set?

ptrendx · 2026-08-10T18:48:35Z

+        # Each parametrized case compiles the same module.forward code object with
+        # a different shape/recipe; with dynamic=False those guards accumulate and
+        # eventually trip Dynamo's recompile_limit. Reset so every case starts from
+        # a clean compile cache (mirrors the single-GPU torch.compile tests).


One concern here is that when more things are going to try testing the compilation flow we would
spend a lot of time in the compilation during the CI. Maybe it would be good to have some thoughts
on which tests should actually be run in L0 and which could be done in the nightly CI for torch.compile.

ptrendx · 2026-08-10T18:50:29Z

+        # dynamic=False for now: a symbolic shape would land in an OpaqueValueBundle
+        # (value-opaque op arg) whose hash chokes on non-nested SymInt. Force static
+        # shapes (recompile per shape) until the bundle handles symbolic shapes.


Is the dynamic=False a limitation that we require users to set? If so, we need to have some documentation
about that somewhere.

ptrendx · 2026-08-10T18:53:04Z

    )


+@pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"])


Kind of a general comment, but do we expect to ever see a case that would work under reduce
overhead mode but not work under the default mode? If so then maybe we could just test the stricter
mode if things are supposed to work under both of them?

ptrendx · 2026-08-10T18:59:24Z

+
+
+@contextlib.contextmanager
+def _assert_no_cudagraph_skips(enabled: bool):


Generally speaking I think it makes sense, but since the counters map is technically an internal
thing in PyTorch, I think that we should guard ourselves a little bit in case the API changes.
E.g. our CI should not fail (but probably issue some kind of warning?) if the counters map is not
there or the specific keys are not found.

ptrendx · 2026-08-10T19:00:55Z

+
+
+# bf16 output tolerance: eager and compiled run the same kernels, so they should
+# agree closely; the slack only absorbs reduction-order / cuda-graph differences.


CUDA graphs and non-CUDA graph execution should actually be exactly identical. The compiled vs eager
could have some difference if you capture more that just TE, but should not have any difference in
the TE ops themselves.

ptrendx · 2026-08-10T21:29:17Z

-    pytest.param(_blockwise, id="float8_blockwise"),
-    pytest.param(_current_scaling, id="float8_current_scaling"),
-    pytest.param(_nvfp4, id="nvfp4"),
+    pytest.param(_mxfp8, {"dtype": tex.DType.kFloat8E5M2}, id="mxfp8"),


ptrendx · 2026-08-10T21:30:28Z

+    a, b = factory(), factory()
+    # Same config -> equal, same hash, interchangeable as a dict/set key.
+    assert a is not b
+    assert a == b
+    assert hash(a) == hash(b)
+    assert {a: "x"}[b] == "x"
+    # Different config -> not equal.
+    assert a != factory(**other_kwargs)


I think this is an artifact of the rebase or something, I vaguely recall we got rid of something
very similar from the previous PR.

ptrendx · 2026-08-10T23:21:10Z

+    if fp8_recipe is not None and not fp8_available:
+        pytest.skip(reason_for_no_fp8)


The fp8_recipe takes values from all_recipes which already skips unavailable recipes, so this
check is useless.

ptrendx · 2026-08-10T23:34:20Z

+        for is_first in schedule:
+            base = torch.randn(32, 64, dtype=dtype, device=device)
+            _assert_close_eager_compiled(fn, compiled, model, base)
+


A similar test that should be here is to run the Linear in training and eval modes back to back to
see whether changing from training -> inference (and more importantly from inference -> training)
works well.

pggPL mentioned this pull request Jun 3, 2026

[PyTorch] Torch.compile support #2590

Open

pggPL force-pushed the linear_torch_compile_final_attempt branch from 98cd401 to c6544d0 Compare August 5, 2026 16:07

pre-commit-ci Bot and others added 2 commits August 5, 2026 16:09

[pre-commit.ci] auto fixes from pre-commit.com hooks

dfa79c2

for more information, see https://pre-commit.ci

pggPL marked this pull request as ready for review August 5, 2026 17:05

pggPL requested a review from ksivaman as a code owner August 5, 2026 17:05

pggPL requested a review from ptrendx August 5, 2026 17:05

pggPL and others added 16 commits August 5, 2026 23:00

[pre-commit.ci] auto fixes from pre-commit.com hooks

250bd71

for more information, see https://pre-commit.ci

Shorten check_gemm_dims docstring

eb3f50f

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Drop tensor_can_be_materialized: inline an exact-class check in the t…

9695013

…wo float8 reprs Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Trim paraphrase comments in the Linear fake impls

c64a0b0

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Tighten custom_op module docstring intro

aeb3481

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Merge and simplify custom_op docstring paragraphs 2-3

53b811b

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Keep the impl-vs-op contrast as two paragraphs

948f94e

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Drop reference to a PyTorch PR that will not land

872387f

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Shorten _ensure_distributed_opaque_types docstring

dcc3c9a

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Fix leftover 'priority order' wording at _FIELD_ADAPTERS

1950585

Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Restructure register_custom_op docstring: caller contract first, drop…

9c4fadf

… module-docstring duplication Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

ptrendx reviewed Aug 10, 2026

View reviewed changes

Navigation Menu

Uh oh!

[PyTorch] [torch.compile] torch.compile support for Linear - #3053

[PyTorch] [torch.compile] torch.compile support for Linear#3053
pggPL wants to merge 19 commits into
NVIDIA:mainfrom
pggPL:linear_torch_compile_final_attempt

pggPL commented May 28, 2026 •

edited

Loading

Uh oh!

greptile-apps Bot commented Aug 5, 2026 •

edited

Loading

Uh oh!

ptrendx Aug 10, 2026

Uh oh!

ptrendx Aug 10, 2026

Uh oh!

ptrendx Aug 10, 2026

Uh oh!

ptrendx Aug 10, 2026

Uh oh!

ptrendx Aug 10, 2026

Uh oh!

ptrendx Aug 10, 2026

Uh oh!

ptrendx Aug 10, 2026

Uh oh!

ptrendx Aug 10, 2026

Uh oh!

ptrendx Aug 10, 2026

Uh oh!

ptrendx Aug 10, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		if args.compile and args.use_cuda_graphs:
		parser.error("--compile and --use-cuda-graphs are mutually exclusive.")

		)


		@pytest.mark.parametrize("compile_mode", ["default", "reduce-overhead"])



		@contextlib.contextmanager
		def _assert_no_cudagraph_skips(enabled: bool):



		# bf16 output tolerance: eager and compiled run the same kernels, so they should
		# agree closely; the slack only absorbs reduction-order / cuda-graph differences.

		if fp8_recipe is not None and not fp8_available:
		pytest.skip(reason_for_no_fp8)

Uh oh!

Conversation

pggPL commented May 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

Changes

Checklist:

Uh oh!

greptile-apps Bot commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Sequence Diagram

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pggPL commented May 28, 2026 •

edited

Loading

greptile-apps Bot commented Aug 5, 2026 •

edited

Loading