Skip to content

Cache post-lowering ATen graphs to skip decomp on warm compiles. - #4540

Draft
micwill755 wants to merge 1 commit into
batch-defer-fx-cleanupfrom
early-lowering-cache
Draft

Cache post-lowering ATen graphs to skip decomp on warm compiles.#4540
micwill755 wants to merge 1 commit into
batch-defer-fx-cleanupfrom
early-lowering-cache

Conversation

@micwill755

@micwill755 micwill755 commented Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add a disk cache for the post-lowering ATen GraphModule, keyed from the exported program before decomposition.
  • On a hit, skip run_decompositions and post_lowering and rebuild the FX graph from a node list + state_dict (torch.save), not torch.export.save or FX __reduce__ retrace.
  • Conservative: full compilation + fast partition only; bypass if the graph cannot be encoded.

Test plan

  • tests/py/dynamo/models/test_lowering_cache.py
  • Flux2 Klein 4B NVFP4 warm compile + inference (engine cache hits after the engine is built from the reconstructed graph)

Serialize the live FX node list and state_dict with torch.save so a hit can skip decomposition and post_lowering without retracing generated Python.
@meta-cla meta-cla Bot added the cla signed label Aug 20, 2026
@github-actions github-actions Bot added component: tests Issues re: Tests 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 20, 2026
@github-actions
github-actions Bot requested a review from zewenli98 August 20, 2026 22:40
@lanluo-nvidia lanluo-nvidia modified the milestones: v2.14.0, v2.15.0 Aug 20, 2026
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: tests Issues re: Tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants