Skip to content

Apply Torch-TRT decomps during dynamo.trace so compile skips a second AOT retrace. - #4596

Open
micwill755 wants to merge 1 commit into
early-lowering-cachefrom
one-shot-decomp-in-trace
Open

Apply Torch-TRT decomps during dynamo.trace so compile skips a second AOT retrace.#4596
micwill755 wants to merge 1 commit into
early-lowering-cachefrom
one-shot-decomp-in-trace

Conversation

@micwill755

@micwill755 micwill755 commented Aug 26, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Apply Torch-TRT's decomp table during dynamo.trace() (the path torch_tensorrt.compile(nn.Module) already uses) so AOT Autograd runs once, matching Inductor / torch.compile(backend="torch_tensorrt").
  • dynamo.compile() skips ExportedProgram.run_decompositions when the exported program stamp matches the current settings. Stock torch.export EPs, or a settings mismatch, still decompose.
  • No new user-facing API (export_for_tensorrt, export_and_compile, skip_decompositions).

PR 5 of the compile-startup stack; based on #4540.

Test plan

  • tests/py/dynamo/models/test_one_shot_decomp.py (stamp match, skip when stamped, stock export still decomps, settings mismatch re-decomps, inplace add_ decomposed during export, trace() stamps)
  • Flux.1-dev transformer export+decomp only: second AOT 14.4s → 0s
  • Flux2 Klein 4B NVFP4 cold compile (Victor harness): run_decompositions 2.59s → skip; compile phase 33.7s → 28.8s; steady-state unchanged

Stack created with GitHub Stacks CLIGive Feedback 💬

@meta-cla meta-cla Bot added the cla signed label Aug 26, 2026
@github-actions github-actions Bot added component: tests Issues re: Tests component: lowering Issues re: The lowering / preprocessing passes component: core Issues re: The core compiler component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Aug 26, 2026
@github-actions
github-actions Bot requested a review from zewenli98 August 26, 2026 21:45
@micwill755
micwill755 force-pushed the one-shot-decomp-in-trace branch from a41ed2f to 20b8bcd Compare September 1, 2026 18:30
@micwill755
micwill755 force-pushed the one-shot-decomp-in-trace branch from 20b8bcd to d48152c Compare September 2, 2026 07:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed component: api [Python] Issues re: Python API component: core Issues re: The core compiler component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths component: lowering Issues re: The lowering / preprocessing passes component: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant