Support conditional torch.compile per module #22269

sarckk · 2025-08-05T18:06:04Z

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

github-actions · 2025-08-05T18:06:15Z

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

🚀

gemini-code-assist

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.

tests/compile/test_decorator.py

vllm/compilation/decorators.py

tests/compile/test_decorator.py

.buildkite/test-pipeline.yaml

tests/compile/test_decorator.py

zou3519

minor nit around the name but otherwise lgtm

DarkLight1337 · 2025-08-12T12:52:13Z

Can you merge from main to fix the CI failures?

ProExpertProg · 2025-08-15T19:40:19Z

@sarckk failures appear related, could you take a look?

sarckk · 2025-08-18T17:17:52Z

Fixed tests that were broken due to changes in #20059.

Signed-off-by: Yong Hoon Shin <[email protected]>

sarckk · 2025-08-19T20:09:54Z

@ProExpertProg thanks for the rebases. CI seems quite noisy, should we keep rebasing until tests pass?

ProExpertProg · 2025-08-19T22:03:35Z

Yeah check if these failures are happening on main and restart ci

…oject#22269) Signed-off-by: Yong Hoon Shin <[email protected]> Signed-off-by: Duncan Moss <[email protected]>

…oject#22269) Signed-off-by: Yong Hoon Shin <[email protected]>

…oject#22269) Signed-off-by: Yong Hoon Shin <[email protected]> Signed-off-by: Xiao Yu <[email protected]>

…oject#22269) Signed-off-by: Yong Hoon Shin <[email protected]>

sarckk requested review from ProExpertProg, youkaichao and zou3519 as code owners August 5, 2025 18:06

sarckk changed the title ~~Support conditional torch.compile~~ Support conditional torch.compile per module Aug 5, 2025

gemini-code-assist bot reviewed Aug 5, 2025

View reviewed changes

tests/compile/test_decorator.py Outdated Show resolved Hide resolved

sarckk force-pushed the conditional-compile branch 2 times, most recently from a51f22b to 8e302bb Compare August 6, 2025 18:33

mgoin reviewed Aug 6, 2025

View reviewed changes

vllm/compilation/decorators.py Outdated Show resolved Hide resolved

tests/compile/test_decorator.py Outdated Show resolved Hide resolved

mergify bot added the ci/build label Aug 6, 2025

zou3519 reviewed Aug 11, 2025

View reviewed changes

.buildkite/test-pipeline.yaml Outdated Show resolved Hide resolved

zou3519 reviewed Aug 11, 2025

View reviewed changes

tests/compile/test_decorator.py Outdated Show resolved Hide resolved

ProExpertProg approved these changes Aug 11, 2025

View reviewed changes

ProExpertProg added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 11, 2025

zou3519 reviewed Aug 11, 2025

View reviewed changes

tests/compile/test_decorator.py Outdated Show resolved Hide resolved

zou3519 approved these changes Aug 11, 2025

View reviewed changes

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

sarckk force-pushed the conditional-compile branch 2 times, most recently from 58b56f2 to bd849d6 Compare August 14, 2025 03:49

ProExpertProg enabled auto-merge (squash) August 15, 2025 17:28

auto-merge was automatically disabled August 18, 2025 17:17
Head branch was pushed to by a user without write access

sarckk force-pushed the conditional-compile branch from e11fe22 to 5d00f1e Compare August 18, 2025 17:17

ProExpertProg force-pushed the conditional-compile branch from 5d00f1e to 30df952 Compare August 18, 2025 21:27

sarckk added 4 commits August 19, 2025 12:21

Support conditional torch.compile

95d0767

Signed-off-by: Yong Hoon Shin <[email protected]>

Fix typo

949f953

Signed-off-by: Yong Hoon Shin <[email protected]>

Address comments

da91b26

Signed-off-by: Yong Hoon Shin <[email protected]>

Rename compile_cond -> enable_if

c196704

Signed-off-by: Yong Hoon Shin <[email protected]>

Fix compile tests

7c416e5

Signed-off-by: Yong Hoon Shin <[email protected]>

ProExpertProg force-pushed the conditional-compile branch from 30df952 to 7c416e5 Compare August 19, 2025 16:21

ProExpertProg enabled auto-merge (squash) August 19, 2025 16:24

ProExpertProg disabled auto-merge August 19, 2025 16:24

ProExpertProg enabled auto-merge (squash) August 19, 2025 16:25

ProExpertProg merged commit dfd2382 into vllm-project:main Aug 20, 2025
71 of 72 checks passed

djmmoss pushed a commit to djmmoss/vllm that referenced this pull request Aug 21, 2025

[torch.compile] Support conditional torch.compile per module (vllm-pr…

285cd2b

…oject#22269) Signed-off-by: Yong Hoon Shin <[email protected]> Signed-off-by: Duncan Moss <[email protected]>

shanes-cerebras pushed a commit to smsegal/vllm that referenced this pull request Aug 24, 2025

[torch.compile] Support conditional torch.compile per module (vllm-pr…

b861219

…oject#22269) Signed-off-by: Yong Hoon Shin <[email protected]>

epwalsh pushed a commit to epwalsh/vllm that referenced this pull request Aug 28, 2025

[torch.compile] Support conditional torch.compile per module (vllm-pr…

70c90cd

…oject#22269) Signed-off-by: Yong Hoon Shin <[email protected]>

xiao-llm pushed a commit to xiao-llm/vllm that referenced this pull request Aug 28, 2025

[torch.compile] Support conditional torch.compile per module (vllm-pr…

b1ac3e8

…oject#22269) Signed-off-by: Yong Hoon Shin <[email protected]> Signed-off-by: Xiao Yu <[email protected]>

zhewenl pushed a commit to zhewenl/vllm that referenced this pull request Aug 28, 2025

[torch.compile] Support conditional torch.compile per module (vllm-pr…

3c73dca

…oject#22269) Signed-off-by: Yong Hoon Shin <[email protected]>

mengxingkongzhouhan pushed a commit to mengxingkongzhouhan/vllm that referenced this pull request Aug 30, 2025

[torch.compile] Support conditional torch.compile per module (vllm-pr…

b828d56

…oject#22269) Signed-off-by: Yong Hoon Shin <[email protected]>

zhewenl pushed a commit to zhewenl/vllm that referenced this pull request Sep 3, 2025

[torch.compile] Support conditional torch.compile per module (vllm-pr…

e146c9e

…oject#22269) Signed-off-by: Yong Hoon Shin <[email protected]>

FeiDaLI pushed a commit to FeiDaLI/vllm that referenced this pull request Sep 25, 2025

[torch.compile] Support conditional torch.compile per module (vllm-pr…

82cf9eb

…oject#22269) Signed-off-by: Yong Hoon Shin <[email protected]>

Uh oh!

Support conditional torch.compile per module #22269

Support conditional torch.compile per module #22269

Uh oh!

Conversation

sarckk commented Aug 5, 2025 • edited by github-actions bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Test Plan

Test Result

Uh oh!

github-actions bot commented Aug 5, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zou3519 left a comment

Choose a reason for hiding this comment

Uh oh!

DarkLight1337 commented Aug 12, 2025

Uh oh!

ProExpertProg commented Aug 15, 2025

Uh oh!

sarckk commented Aug 18, 2025

Uh oh!

sarckk commented Aug 19, 2025

Uh oh!

ProExpertProg commented Aug 19, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

sarckk commented Aug 5, 2025 •

edited by github-actions bot

Loading