[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

This PR adds torch.compile support for te.pytorch.Linear through data-free fake implementations and registered forward/backward custom operators.

Introduces a declarative framework for flattening structured TE arguments and outputs across custom-op boundaries.
Preserves quantized tensor metadata, process-group references, and autograd state across compiled execution.
Adds single-device and distributed coverage for compilation, quantization recipes, dynamic shapes, and CUDA-graph modes.

Confidence Score: 5/5

The PR appears safe to merge because no blocking failure remains in the eligible follow-up review scope.

No blocking failure remains.

Important Files Changed

Filename	Overview
transformer_engine/pytorch/dynamo/custom_op.py	Adds the custom-op registration, structured argument adaptation, fake-output planning, quantized tensor reconstruction, and autograd wiring used by compiled TE modules.
transformer_engine/pytorch/module/linear.py	Splits Linear computation from context handling, implements TensorSpec-based fake execution, registers compiled forward and backward operators, and retains eager fallback for unsupported configurations.
transformer_engine/pytorch/csrc/extensions/swizzle.cpp	Allocates swizzled scale storage as byte-backed buffers to match the native swizzle representation.
transformer_engine/pytorch/tensor/float8_tensor.py	Updates Float8 tensor and quantizer behavior needed to preserve representation metadata across compiled execution.
tests/pytorch/test_torch_compile.py	Adds broad compiled Linear coverage for forward and backward numerics, quantized weights and outputs, dynamic shapes, fallbacks, and CUDA-graph execution.
tests/pytorch/distributed/run_numerics.py	Extends distributed Linear numerical checks to compiled execution and validates input gradients across tensor- and sequence-parallel layouts.
tests/pytorch/distributed/run_layer_with_overlap.py	Adds compiled Userbuffers execution, including reduce-overhead warmup and CUDA-graph skip validation.

Sequence Diagram

sequenceDiagram
    participant User as PyTorch model
    participant Linear as te.Linear
    participant Fake as TensorSpec fake implementation
    participant Fwd as Registered forward op
    participant Native as TE kernels
    participant Bwd as Registered backward op

    User->>Linear: torch.compile forward(input)
    Linear->>Fake: Describe outputs and saved state
    Fake-->>Linear: Output plan and TensorSpecs
    Linear->>Fwd: Flattened argument slots
    Fwd->>Native: Execute Linear forward
    Native-->>Fwd: Outputs and tensors to save
    Fwd-->>User: Reconstructed logical outputs
    User->>Bwd: Output gradients
    Bwd->>Native: Restored state and flattened gradients
    Native-->>Bwd: Input, weight, and bias gradients
    Bwd-->>User: Reconstructed gradients

_{Reviews (17): Last reviewed commit: "Fix stale amax groups in compiled Linear" | 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>

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

The cached weight workspace is updated in place on the first microbatch, which the functional custom op (mutates_args=()) can't express. Covers skip_fp8_weight_update too (it implies caching). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…e effects in fallback cases Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

The fwd fake declared (inp.shape[0], out_features) while the real impl views the output to (1, out_features) for 1D inputs. The bwd impl rederives the input shape from grad_output, which cannot recover rank-1, so the autograd glue now stashes the true input shapes on ctx (SymInt-safe) and views the returned grads back to them. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Weight caching now falls back to eager under torch.compile, so with the argument the xfail test never reached the train/eval switching it covers. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

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

A quantized dgrad can't cross the op boundary (grads are packed one plain Tensor[] slot each), so AOT tracing crashed instead of falling back. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

An op with an unused output is legally DCE'd (mutates_args=()), silently dropping its collectives and state updates. Register the ops in FX's side-effect registry by default; NVTE_COMPILE_OP_SIDE_EFFECTS selects token (ordered effect tokens, blocks reordering too, but incompatible with cudagraph trees today) or 0 (off). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

An allocated transpose may be stale (invalidated by _reset_caches); reconstruction derived validity from presence and silently revalidated it, so the compiled path could consume a pre-update transpose. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

- skip under fake tensors/modes (a fake workspace poisoned the process-global cache) - also preallocate in _apply, covering .to()/.cuda()/to_empty() flows (meta-device init never ran the reset_parameters path) - preallocate the UB workspace only when comm overlap is actually enabled, not on ub_name alone (MHA sets it unconditionally) - drop the stream-capture assert: it broke previously-working eager CUDA-graph captures, while the cudagraph-trees hazard it aimed at (warmup-pool allocation) never triggers it anyway Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

A graph break inside try/finally cannot build a resume function, so Dynamo marked the shared Linear.forward code object SKIP: one unsupported config silently reverted every te.Linear in the process to eager. Hoist the config checks into _compile_eager_fallback_reason and exit through a dynamo-disabled eager re-entry; only quantizer-dependent conditions remain in the late check. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

ptrendx · 2026-09-01T22:57:30Z

/te-ci pytorch L0 L1

…e_final_attempt

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

…nd LayerNormMLP Register the LayerNormLinear and LayerNormMLP forward/backward as torch.library custom ops on top of the framework from NVIDIA#3053, so both trace under fullgraph compile with the FP8/MXFP8/NVFP4 recipes. Builds on the restructuring of the two modules into FwdArgs/BwdArgs + impl functions. Framework (dynamo/custom_op.py): - A bwd args dataclass may declare GRAD_OUTPUT_FIELDS naming one grad field per user output; the LN modules have two differentiable outputs (out and the returned norm output). - Dict[str, <simple>] annotations are bundle-simple (activation_params). module/_common.py: sp_out_leading / sp_inp_leading / fake_workspace_valid shared by Linear and the LN modules (moved out of linear.py). layernorm_linear.py / layernorm_mlp.py: - compile_unsupported_reason gate, data-free forward/backward fakes on TensorSpec, register_custom_op, @no_torch_dynamo eager wrapper, compiled dispatch and eager fallback in Module.forward, mirroring Linear. - inp_shape is rederived from grad_output in backward instead of being stored (SymInt dims are not hashable in the value bundle). - LayerNormMLP: the recipe object no longer rides on the backward args; the properties the backward needs are bools computed in Module.forward. The activation tables are split into per-activation (act, dact) pairs and the fused dbias kernels, so the backward looks them up without a recipe. - Eager fallbacks: debug, quantized input, fsdp_group, differentiable fp8_output, CPU offload, delayed wgrad, FP8 weight caching, fuse_wgrad_accumulation, non-opaque quantizers; MLP also activation checkpointing and gemm_gelu_fusion under FP8. tests/pytorch/test_torch_compile.py: LN module coverage (recipes x compile modes, norm/activation/return_layernorm_output variants, FP8 primary weight, frozen weights / no bias, dynamic shapes, eager fallbacks). Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

…nd LayerNormMLP Register the LayerNormLinear and LayerNormMLP forward/backward as torch.library custom ops on top of the framework from NVIDIA#3053, so both trace under fullgraph compile with the FP8/MXFP8/NVFP4 recipes. Builds on the module restructuring and the framework groundwork in the two preceding PRs. layernorm_linear.py / layernorm_mlp.py: - compile_unsupported_reason gate, data-free forward/backward fakes on TensorSpec, register_custom_op, @no_torch_dynamo eager wrapper, compiled dispatch and eager fallback in Module.forward, mirroring Linear. - Eager fallbacks: debug, quantized input, fsdp_group, differentiable fp8_output, CPU offload, delayed wgrad, FP8 weight caching, fuse_wgrad_accumulation, non-opaque quantizers; MLP also activation checkpointing and gemm_gelu_fusion under FP8. tests/pytorch/test_torch_compile.py: LN module coverage (recipes x compile modes, norm/activation/return_layernorm_output variants, FP8 primary weight, frozen weights / no bias, dynamic shapes, eager fallbacks). 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 added 14 commits August 31, 2026 19:14

Trim swizzle.cpp comments

8bd8bfb

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

Derive tensor_field_names from fields instead of storing it on _ArgPlan

9d7e71d

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

Strengthen compile tests: bias grads in the main matrix, stateful sid…

8e9b95e

…e effects in fallback cases Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Narrow NVTE_TORCH_COMPILE doc to its actual scope (internal jit fusions)

f03cd95

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

Drop is_first_microbatch from the train/eval switch test

4722dfd

Weight caching now falls back to eager under torch.compile, so with the argument the xfail test never reached the train/eval switching it covers. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

Bwd fake: dgrad is a plain tensor under ub_bulk_wgrad too

d69dbf8

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

Limit the weight-caching compile fallback to FP8

3d74303

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

Fall back to eager for fp8_grad=True under torch.compile

780d39f

A quantized dgrad can't cross the op boundary (grads are packed one plain Tensor[] slot each), so AOT tracing crashed instead of falling back. Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>

ptrendx previously approved these changes Sep 1, 2026

View reviewed changes

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

9b5c068

…e_final_attempt

pggPL dismissed ptrendx’s stale review via 9b5c068 September 2, 2026 07:55

Fix stale amax groups in compiled Linear

dcfda40

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

ptrendx approved these changes Sep 2, 2026

View reviewed changes

pggPL merged commit 15f882d into NVIDIA:main Sep 2, 2026
12 of 16 checks passed

pggPL mentioned this pull request Sep 3, 2026

[PyTorch] [torch.compile] Prepare LayerNormLinear and LayerNormMLP for torch.compile pggPL/TransformerEngine#42

Open

8 tasks

This was referenced Sep 3, 2026

[PyTorch] [torch.compile] torch.compile support for LayerNormLinear and LayerNormMLP pggPL/TransformerEngine#43

Open

[PyTorch] [torch.compile] Prepare the custom-op framework and LayerNorm modules for compile wiring pggPL/TransformerEngine#44

Open

This was referenced Sep 3, 2026

[PyTorch] torch.compile support for FusedAttention #3472

Draft

[PyTorch] torch.compile for an OperationFuser group holding one operation pggPL/TransformerEngine#45

Open

[PyTorch] Split FusedAttnFunc into single-argument forward/backward helpers #3480

Merged

Navigation Menu

Uh oh!

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

[PyTorch] [torch.compile] torch.compile support for Linear#3053
pggPL merged 70 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!

Uh oh!

Uh oh!

Uh oh!

ptrendx Aug 10, 2026

Uh oh!

pggPL Aug 12, 2026

Uh oh!

ptrendx Aug 18, 2026

Uh oh!

pggPL Aug 19, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ptrendx commented Sep 1, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		)


		@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!

ptrendx Aug 10, 2026

Choose a reason for hiding this comment

Uh oh!

pggPL Aug 12, 2026

Choose a reason for hiding this comment

Uh oh!

ptrendx Aug 18, 2026

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!

ptrendx commented Sep 1, 2026

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