[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 torch.compile support for te.pytorch.Linear by registering its forward and backward as declarative custom operators.

Introduces dataclass-driven argument flattening, fake implementations, output reconstruction, and autograd wiring.
Splits Linear computation from context saving and adds compiled support for quantized outputs, dynamic shapes, microbatch weight caching, and distributed execution.
Adds compile and CUDA-graph coverage across supported quantization recipes and communication-overlap configurations.

Confidence Score: 5/5

The PR appears safe to merge based on the eligible follow-up findings.

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

Important Files Changed

Filename	Overview
transformer_engine/pytorch/dynamo/custom_op.py	Adds the declarative custom-op framework that derives schemas from dataclasses, handles quantized tensor flattening, supplies fake implementations, and registers autograd.
transformer_engine/pytorch/module/linear.py	Refactors Linear forward and backward into reusable computation, fake, and context-setup stages and routes supported compiled calls through the new custom operator.
transformer_engine/pytorch/tensor/float8_tensor.py	Makes Float8Tensor representation safe for fake, meta, and functional tensors used during compilation.
transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py	Applies corresponding safe representation behavior to Float8TensorStorage.
transformer_engine/pytorch/cpp_extensions/gemm.py	Exposes the cuBLAS workspace helper and prevents first allocation during CUDA-graph capture.
tests/pytorch/test_torch_compile.py	Adds comprehensive compiled Linear coverage for supported recipes, quantized weights and outputs, microbatch caching, dynamic shapes, fallback behavior, and CUDA-graph modes.
tests/pytorch/distributed/run_numerics.py	Extends distributed Linear numerics checks to compiled execution and validates input gradients.
tests/pytorch/distributed/run_layer_with_overlap.py	Adds compiled Userbuffers execution, reduce-overhead warmup, and CUDA-graph capture assertions.

Sequence Diagram

sequenceDiagram
  participant User
  participant Linear as te.Linear
  participant Adapter as register_custom_op adapter
  participant Fake as TensorSpec fake implementation
  participant Op as torch.library custom op
  participant Kernel as TE GEMM kernels
  User->>Linear: forward(input)
  Linear->>Adapter: packed LinearFwdArgs
  Adapter->>Fake: derive output and saved-tensor plan
  Fake-->>Adapter: TensorSpec structure
  Adapter->>Op: flattened tensor and opaque slots
  Op->>Kernel: execute Linear computation
  Kernel-->>Op: outputs and saved tensors
  Op-->>Adapter: flat Tensor[] payload
  Adapter-->>User: reconstructed output
  User->>Op: backward(output gradients)
  Op->>Kernel: execute Linear backward
  Kernel-->>User: input, weight, and bias gradients

_{Reviews (11): Last reviewed commit: "Add an output plan and lift the single-g..." | 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: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?

The benefit of such approach would be time saved. If torch.compile + TE CI time will be big we may do that.

So did you measure the time increase of the CI due to this PR?

1.5 min for L0(not just this PR, all torch.compile tests), ~5min for L1

- Backward fake now returns grad_bias whenever bias is used on the FP8 backward path (grad_output_preprocess computes bgrad independent of requires_wgrad); previously a frozen weight silently dropped the bias gradient under torch.compile. - Forward fake now mirrors quantize_weight's workspace invalidation: a cached workspace missing buffers for the quantizer's current usage is dropped and a fresh new_weight_workspace is declared, instead of always assuming a cache hit (previously crashed with an output size/stride mismatch when a rowwise-only cache met a training step). Both verified against eager on RTX Ada. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

copyreg.pickle is process-wide: with the reducer installed, torch.save of any object graph reaching a ProcessGroup silently succeeded and the checkpoint failed only at torch.load (the reconstruct stub raises). Restore the loud failure at save time; the cost is that inductor bypasses the FX disk cache for compiled distributed graphs (with its own warning) until the cache-key pickler handles real opaque objects upstream. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

DelayedScaling quantizers are not value-opaque, so the compiled path falls back to eager, which errors under fullgraph=True. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

… UB test to FP8 Under ub_overlap_rs_dgrad the impl returns the plain high-precision reduce-scatter output as dgrad (the grad_input_quantizer only feeds the communication buffer), while the fake declared a quantized dgrad spec -- an op output-contract mismatch. test_linear_with_overlap_compile now also runs fp8_current_scaling and mxfp8 for the column-parallel cases (bulk and DGRAD+RS); FP8 row-parallel stays skipped (forced differentiable fp8_output is unsupported under compile) and delayed scaling is excluded like elsewhere. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

The compiled path handles neither the external weight subclass at the op boundary nor the materialize/refresh logic in the fakes; gate it in compile_unsupported_reason like the other unsupported configs. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

- fsdp_group: fall back regardless of grad mode; the adapter rejects the field for any non-trivial value, so inference with manual TE FSDP could reach the op and fail there instead. - fp8_output: also fall back when only the bias requires grad; the backward tangent for the quantized output was mis-guessed by AOTAutograd (RuntimeError: Expected a Float8Tensor tangent but got a plain Tensor). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Registration failures now only record the reason; importing TE on a build without opaque-object support stays silent. The warning is emitted from Linear.forward when a compiled call finds the op unregistered; under fullgraph=True the resulting error names warn_if_compile_disabled. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…erEngine logger) Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

The impl disables save_original_input when the input quantizer cannot reconstruct the wgrad operand from the original input (e.g. NVFP4 with stochastic rounding); the fake kept it on and declared the saved-input slot as an alias while the impl saved a quantized storage. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

The inp field crosses the op boundary as a plain Tensor slot, so a quantized activation (e.g. the fp8_output of a previous layer) breaks fake propagation even under no_grad; gate it until the boundary supports quantized inputs. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…eal microbatch cache checks - Tolerances tightened to exact: all compute runs inside the op and the loss grad is ones, so eager and compiled are bit-identical (measured on both compile modes, bf16 and fp8). - torch._dynamo counters reads degrade with a warning instead of failing when the private API changes. - is_first_microbatch test: eager reference on a separate module (shared cache made it unable to catch corruption or rebuilds), structural asserts that the compiled step creates the cache and later steps reuse the same object, eager priming of FP8 state before tracing (in-graph quantizer creation breaks recompiles; upstream Dynamo bug). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

… under reduce-overhead run_numerics now checks dgrad (gathered per parallel mode) in every linear case; run_layer_with_overlap warms the compiled model up under reduce-overhead so the measured run replays captured graphs, and asserts inductor recorded no cudagraph skips. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…ve a redundant skip The a==b/hash/dict-key block and its other_kwargs parametrization were already removed once (6f66c3e) as covered by the __fx_repr__ round-trip; a rebase resurrected them. The fp8_available skip in test_te_linear_compiles is dead: _all_recipes is availability-gated at construction. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…tween iterations The warmup iterations kept warmup gradients alive in the cudagraph pool, tripping cudagraph_trees' check_memory_pool on the next capture (Detected N tensor(s) in the cudagraph pool not tracked as outputs). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Replace the reference-opaque ProcessGroup graph input with the pattern traceable functional collectives use: the adapter ships pg.group_name (a plain string in the value bundle) and re-resolves the live group from the c10d registry inside the op, in the same process -- from_slots(to_slots(pg)) is the identical object by construction. This removes the opaque PG from example_inputs entirely, so inductor's FX disk cache works for compiled distributed graphs again (verified: second process gets fxgraph_cache_hit=2, no pickle bypass) without any upstream change. The reference-opaque registration machinery and the PG_REFERENCE_OPAQUE fallback gate are no longer needed. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Blocked by an upstream Dynamo bug: FP8 state created inside the first compiled call comes back as FakeScriptObject/None in the graph outputs, so any later recompile crashes. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Each parametrized case resets dynamo and compiles once with a single shape, so automatic dynamic shapes never trigger; verified on 4xGB200 (run_numerics 126/126, comm-GEMM overlap compile suite green). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

The previous fetch in forward never executed for real: Dynamo ignores the lru_cache wrapper and traces the wrapped function, so the workspace was first allocated by the op impl at runtime - under reduce-overhead on capture-first torch builds that lands in the CUDA-graph pool and trips 'cudagraph pool not tracked as outputs'. Allocate both variants (plain and UB) in reset_parameters, which always runs eagerly, drop the dead cublas_workspace bundle field, and fail fast if a workspace would first be allocated during stream capture. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

ptrendx · 2026-08-18T21:22:27Z

    )


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


So did you measure the time increase of the CI due to this PR?

ptrendx · 2026-08-18T22:39:32Z

+class _SimpleBundleAdapter(_Adapter):
+    """Aggregates every simple-typed field into a single OpaqueValueBundle.
+
+    Unlike the per-field adapters, at most one of these exists per op (none if the


I don't understand why we cannot put everything into a single bundle.

We can put qunatizers and processgroup (it was not possible some multiple commits ago, but now we can do it).
We cannot put any tensors inside it (it is opaque value object which cannot contain tensors).

ptrendx · 2026-08-18T22:39:32Z

+    The check must run at call time (not in ``_get_adapters``): the annotation
+    alone -- e.g. ``Optional[Any]`` -- is valid when the value is ``None``, so
+    only the runtime value can decide.


Does that mean that this will run every time the compiled function is invoked?

- drop unused Float8Quantizer import; import all quantizers from their tensor modules - import the local tests utils.py by explicit sys.path so a cutedsl top-level utils package cannot shadow it - cache OpaqueValueBundle hash at construction - guard is_simple_value when the opaque-object API is unavailable - inline the one-line _unflatten_value helper - point the --compile/--use-cuda-graphs error at --compile-mode reduce-overhead Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

pytorch/pytorch#187041; #187057 fixes the cold-compile path (merged), #193190 fixes the FX-graph-cache-hit path (in review). Both verified against this test on nightly 2.15.0.dev20260815. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

The dedicated _QuantizerAdapter only existed because the bundle matches fields by annotation and quantizer fields are annotated with the abstract Quantizer base, which is not a registered opaque type itself. Match the base class in _SimpleBundleAdapter instead and drop the adapter; the per-quantizer schema slots carried no gradients. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Dynamo silently drops warnings.warn in traced code, and the graph break inside the forward's try/finally makes it skip the whole frame and re-run it with is_compiling() == False, so neither warn_compile_eager_fallback nor warn_if_compile_disabled ever emitted. Emit them at trace time via torch._dynamo.comptime instead (once per compilation) and warn before the explicit graph break, which ends the trace. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Parametrized test over the single-GPU-constructible reasons rejected by compile_unsupported_reason (differentiable fp8_output, wgrad fusion / delay, quantized input): fallback warning fires, numerics match eager, fullgraph=True fails with the explicit reason. Delayed scaling is a hard error under fullgraph (check_recipe_support) and is tested separately. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Both warn functions duplicated the is_compiling()/comptime dispatch; move it into _compile_safe_warn so callers just pass the message. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Replace the _FIELD_ADAPTERS registry + per-class try_build with a single _build_field_adapter factory dispatching on the annotation. Remove OpaqueValueBundle.__getattr__: nothing uses attribute access (consumers go through __getitem__/get/as_dict), and without it default copy/deepcopy/pickle work with no special-casing -- which matters since Dynamo's guard machinery deepcopies value-opaque objects. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Drop get(): the __kind__ tag is set at every construction site, so plain indexing (loud KeyError) is the right access. _storage_unflatten's only caller always passes a bundle, so drop the dead dict(meta) branch. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

A live group can't cross as a value, so the bundle stores its c10d registry name and the op re-resolves it (same scheme the dedicated adapter used). Drops _ProcessGroupAdapter and the tp_group__pg schema slot; per-field adapters are now only the two tensor kinds. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…e_final_attempt Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> # Conflicts: # tests/pytorch/test_torch_compile.py

Tests merged from main leave pending delayed-scaling amax reductions in FP8GlobalStateManager; a later autocast __exit__ then calls raw tex bindings, graph-breaking the fullgraph=True Linear tests. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Parse the args dataclass's annotations once, at registration, into an immutable _ArgPlan: per-field _FieldPlan records (_FieldKind + schema slots) plus the derived layout -- schema string, slot order, gradient placement, tensor-or-quantized offsets -- with duplicate-slot-name validation. pack/unpack interpret the plan on each call. Replaces the adapter classes and the four layout helpers that each re-walked them; the op schema and Linear semantics are unchanged. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Parse the fwd fake-impl result into a per-trace _OutputPlan (logical outputs / saved tensors with their flat Tensor[] ranges) and use it as the single structure behind forward_fn, setup_context and backward. Backward now slices grads per user output from the plan stashed on ctx: a grad_outputs field on the backward args receives the whole tuple, otherwise grad_output receives the first output's grad -- removing the flat_grads[0] single-output assumption. Also reject unions mixing tensor types with other members at registration. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

pggPL added 4 commits August 11, 2026 13:35

Skip use_compile numerics cases for DelayedScaling

23b093d

DelayedScaling quantizers are not value-opaque, so the compiled path falls back to eager, which errors under fullgraph=True. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

pggPL and others added 15 commits August 11, 2026 15:02

Log the compile-disabled reason at registration time (INFO, Transform…

1822afc

…erEngine logger) Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Release ctx.backward_objects after the compiled-op backward

dc5cecd

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

Add train/eval mode-switch test as xfail

0dcaadd

Blocked by an upstream Dynamo bug: FP8 state created inside the first compiled call comes back as FakeScriptObject/None in the graph outputs, so any later recompile crashes. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

pggPL mentioned this pull request Aug 13, 2026

[PyTorch] [torch.compile] torch.compile support for GroupedLinear pggPL/TransformerEngine#32

Draft

13 tasks

ptrendx reviewed Aug 18, 2026

View reviewed changes

pggPL and others added 13 commits August 19, 2026 12:33

Merge remote-tracking branch 'upstream/main' into linear_torch_compil…

1fe48cd

…e_final_attempt Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com> # Conflicts: # tests/pytorch/test_torch_compile.py

		)


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

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!

Uh oh!

Uh oh!

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!

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

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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!

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

pggPL Aug 19, 2026 •

edited

Loading