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)

@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

@narendasan narendasan left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Here can we use the same caching interface as engine caching? In theory users could provide their own implementation for things like remote cache or mem cache. We should use the same cache object for both engine and graph cache if possible

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.
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.

3 participants