Skip to content

Support conditional torch.compile per module #22269

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 4 commits into
base: main
Choose a base branch
from

Conversation

sarckk
Copy link
Collaborator

@sarckk sarckk commented Aug 5, 2025

Essential Elements of an Effective PR Description Checklist

  • The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
  • The test plan, such as providing test command.
  • The test results, such as pasting the results comparison before and after, or e2e results
  • (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

Required fro #22628

#21044 shows an example of compiling multiple submodules within a model. We also want to provide the user the flexibility of triggering this on/off.

This PR adds a new arg enable_if to the support_torch_compile decorator, which is a function that takes in VllmConfig and returns a bool of whether to enable compile or not.

Test Plan

added new unit tests

pytest tests/compile/test_decorator.py
pytest tests/compile/piecewise/test_multiple_graphs.py

Test Result

Tests pass

cc: @zou3519 @BoyuanFeng

Copy link

github-actions bot commented Aug 5, 2025

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

@sarckk sarckk changed the title Support conditional torch.compile Support conditional torch.compile per module Aug 5, 2025
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

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

Code Review

This pull request introduces a compile_cond argument to the support_torch_compile decorator, allowing for conditional compilation of models. The implementation is clean and well-tested. I've found one minor issue in the new test file that should be addressed. Overall, this is a great addition that increases the flexibility of the compilation framework.

@sarckk sarckk force-pushed the conditional-compile branch 2 times, most recently from a51f22b to 8e302bb Compare August 6, 2025 18:33
@mergify mergify bot added the ci/build label Aug 6, 2025
@ProExpertProg ProExpertProg added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 11, 2025
Copy link
Collaborator

@zou3519 zou3519 left a comment

Choose a reason for hiding this comment

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

minor nit around the name but otherwise lgtm

@sarckk sarckk force-pushed the conditional-compile branch from 8bf5031 to d48647a Compare August 11, 2025 23:20
@DarkLight1337
Copy link
Member

Can you merge from main to fix the CI failures?

@sarckk sarckk force-pushed the conditional-compile branch from d48647a to 58b56f2 Compare August 12, 2025 16:22
sarckk added 4 commits August 13, 2025 20:49
Signed-off-by: Yong Hoon Shin <[email protected]>
Signed-off-by: Yong Hoon Shin <[email protected]>
Signed-off-by: Yong Hoon Shin <[email protected]>
@sarckk sarckk force-pushed the conditional-compile branch from 58b56f2 to bd849d6 Compare August 14, 2025 03:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ci/build ready ONLY add when PR is ready to merge/full CI is needed
Projects
None yet
Development

Successfully merging this pull request may close these issues.

5 participants