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

Merged

pggPL merged 70 commits into

NVIDIA:mainfrom

pggPL:linear_torch_compile_final_attempt

Sep 2, 2026

Commits on Aug 5, 2026

[PyTorch] [torch.compile] torch.compile support for Linear
pggPL
committed
c6544d0
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot]
committed
dfa79c2
[PyTorch] Keep the broad-except pylint disable on the anchored line
pggPL
committed
116d477
[PyTorch] [torch.compile] Style pass on the Linear custom-op path

pggPL
and
claude
committed
dc4ff77
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot]
committed
250bd71

Commits on Aug 10, 2026

Commits on Aug 11, 2026

Commits on Aug 12, 2026

Commits on Aug 19, 2026

[PyTorch] [torch.compile] Address trivial review comments

pggPL
and
claude
committed
31b385e
Reference the upstream PyTorch fixes in the train/eval xfail

pggPL
and
claude
committed
3bdce4e
Fold quantizers into the shared simple-value bundle

pggPL
and
claude
committed
47fee78
Make the eager-fallback warnings actually fire under torch.compile

pggPL
and
claude
committed
6815201
Test the eager-fallback paths of the compiled Linear op

pggPL
and
claude
committed
db8c13c
Fold the compile-aware branching into the warn helper

pggPL
and
claude
committed
43a0783
Simplify adapter selection and drop the bundle's __getattr__

pggPL
and
claude
committed
0dc980a
Trim OpaqueValueBundle to its actual consumers

pggPL
and
claude
committed
355a2cd
Fold ProcessGroup fields into the shared simple-value bundle
pggPL
committed
d3f6849
Merge remote-tracking branch 'upstream/main' into linear_torch_compile_final_attempt
pggPL
committed
1fe48cd
Reset FP8 global state between torch.compile tests

pggPL
and
claude
committed
727def7

Commits on Aug 21, 2026

Rewrite the custom-op arg boundary as a parsed plan

pggPL
and
claude
committed
97825f3
Add an output plan and lift the single-grad-output limit

pggPL
and
claude
committed
97bb4be
Trim derived state from the plans

pggPL
and
claude
committed
8da8dea

Commits on Sep 1, 2026

Fall back to eager for FP8 weight caching under torch.compile
pggPL
committed
9b90839
Strengthen compile tests: bias grads in the main matrix, stateful side effects in fallback cases
pggPL
committed
8e9b95e
Narrow NVTE_TORCH_COMPILE doc to its actual scope (internal jit fusions)
pggPL
committed
f03cd95
Fix rank-1 input under torch.compile
pggPL
committed
51f20f3
Drop is_first_microbatch from the train/eval switch test
pggPL
committed
4722dfd
Bwd fake: dgrad is a plain tensor under ub_bulk_wgrad too
pggPL
committed
d69dbf8
Limit the weight-caching compile fallback to FP8
pggPL
committed
3d74303
Fall back to eager for fp8_grad=True under torch.compile
pggPL
committed
780d39f
Protect the custom ops from dead-code elimination
pggPL
committed
865c7f6
Preserve _transpose_invalid across the flatten/unflatten round trip
pggPL
committed
cbf3077
Rework cuBLAS workspace preallocation
pggPL
committed
91c9512
Decide the compile fallback before the try/finally in Linear.forward
pggPL
committed
54883ca

Commits on Sep 2, 2026