Avoid bytecode hook and simplify TorchCompileWrapperWithCustomDipatch by laithsakka · Pull Request #25110 · vllm-project/vllm

laithsakka · 2025-09-17T23:14:49Z

Two main changes:

Add an option to not use the bytecode hook! torch2.8 have a way to drop guard without the hack.
We will switch totally to not use in a different PR after internal validation of perf post landing this.
Since @zou3519 have concerns about the perf. which i highly doubt that there will be issues.
TorchCompileWrapperWithCustomDipatch has complexities that are not used in the code base.
the current usage is to by pass guard evaluation so I am just introducing a much simpler TorchCompileWithNoGuardsWrapper.

Next PR I will add a debug mode option to keep DS gaurds and fail if get violated.

performance:

How does this effect run-time?
I ran the following benchmarks:

vllm bench latency
--model Qwen/Qwen2-1.5B-Instruct
--input-len 128
--output-len 256
--num-iters 50
--dtype float16

after:

10% percentile latency: 0.9403901444951771 seconds
25% percentile latency: 0.9434001832487411 seconds
50% percentile latency: 0.9502700129960431 seconds
75% percentile latency: 0.9560497467537061 seconds
90% percentile latency: 0.9588484280975536 seconds
99% percentile latency: 0.9626266451006813 seconds

before:

10% percentile latency: 0.9399352639040444 seconds
25% percentile latency: 0.9432627985042927 seconds
50% percentile latency: 0.9482277854986023 seconds
75% percentile latency: 0.9578490345011232 seconds
90% percentile latency: 0.9600405636985669 seconds
99% percentile latency: 0.9649963746999856 seconds

vllm bench throughput --model Qwen/Qwen2-1.5B-Instruct --input-len 512 --output-len 128 --num-prompts 1000

after

Throughput: 130.14 requests/s, 83113.04 total tokens/s, 16658.08 output tokens/s
Total num prompt tokens:  510637
Total num output tokens:  128000

before

Throughput: 129.54 requests/s, 82726.85 total tokens/s, 16580.68 output tokens/s
Total num prompt tokens:  510637
Total num output tokens:  128000

gemini-code-assist

Code Review

This pull request refactors the torch.compile wrapper to simplify its implementation. It removes the complex bytecode hooking mechanism in favor of the guard_filter_fn option available in newer PyTorch versions to drop guards. The new TorchCompileGuardsStripWrapper is much cleaner. The changes also include updating the support_torch_compile decorator and related tests to use the new wrapper. The overall change improves code clarity and maintainability. I've found one critical issue in the test suite that needs to be addressed.

tests/compile/test_wrapper.py

vllm/compilation/decorators.py

mergify · 2025-09-21T01:07:44Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @laithsakka.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

ProExpertProg · 2025-09-23T14:30:08Z

@bigPYJ1151 is planning to fix the CPU torch issue in vLLM so once that's done, we can upgrade to torch==2.8 everywhere and merge this PR

vllm/compilation/wrapper.py

ProExpertProg

Looks good, can you reformat?

mergify · 2025-11-12T05:25:34Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @laithsakka.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

vllm/compilation/wrapper.py

…patcher with TorchCompileGuardsStripWrapper Signed-off-by: Laith Sakka <lsakka@meta.com>

laithsakka · 2025-11-13T16:43:31Z

rebase again

laithsakka · 2025-11-14T21:34:15Z

@ProExpertProg can you help me land it before another rebase is needed :)

hmellor · 2025-11-15T13:08:02Z

FYI the docs build in this PR was failing so now it's failing on main too.

We're working on a fix.

ProExpertProg · 2025-11-15T14:16:57Z

Sorry about that, I looked at the error message and it looked unrelated, I should have checked nightly.

DarkLight1337 · 2025-11-15T14:30:03Z

Fixed by #28772

hmellor · 2025-11-15T16:16:32Z

It was an issue with mkdocstrings which I've reported to them.

For future reference, the docs are built for every commit on main. So if the docs are green on main, it's a real failure.

dtrifiro · 2025-12-05T17:43:39Z

vllm/v1/worker/tpu_model_runner.py

+            compiled_model.compiled = False
+            TorchCompileWithNoGuardsWrapper.__init__(compiled_model)


This change broke tpu on v0.11.2, calling __init__ here results in an attempt to use VllmConfig when it's not defined, resulting in the following traceback:

(Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] WorkerProc hit an exception. (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] Traceback (most recent call last): (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] File "/home/runner/actions-runner/_work/nm-cicd/nm-cicd/.venv/lib/python3.12/site-packages/vllm/v1/executor/multiproc_executor.py", line 810, in worker_busy_loop (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] output = func(*args, **kwargs) (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] ^^^^^^^^^^^^^^^^^^^^^ (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] File "/home/runner/actions-runner/_work/nm-cicd/nm-cicd/.venv/lib/python3.12/site-packages/vllm/v1/worker/tpu_worker.py", line 214, in determine_available_memory (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] self.model_runner.reset_dynamo_cache() (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] File "/home/runner/actions-runner/_work/nm-cicd/nm-cicd/.venv/lib/python3.12/site-packages/vllm/v1/worker/tpu_model_runner.py", line 1907, in reset_dynamo_cache (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] TorchCompileWithNoGuardsWrapper.__init__(compiled_model) (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] File "/home/runner/actions-runner/_work/nm-cicd/nm-cicd/.venv/lib/python3.12/site-packages/vllm/compilation/wrapper.py", line 91, in __init__ (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] vllm_config = get_current_vllm_config() (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] ^^^^^^^^^^^^^^^^^^^^^^^^^ (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] File "/home/runner/actions-runner/_work/nm-cicd/nm-cicd/.venv/lib/python3.12/site-packages/vllm/config/vllm.py", line 1136, in get_current_vllm_config (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] return VllmConfig() (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] ^^^^^^^^^^^^ (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] File "/home/runner/actions-runner/_work/nm-cicd/nm-cicd/.venv/lib/python3.12/site-packages/pydantic/_internal/_dataclasses.py", line 121, in __init__ (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] s.__pydantic_validator__.validate_python(ArgsKwargs(args, kwargs), self_instance=s) (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] File "/home/runner/actions-runner/_work/nm-cicd/nm-cicd/.venv/lib/python3.12/site-packages/vllm/config/vllm.py", line 609, in __post_init__ (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] current_platform.check_and_update_config(self) (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] File "/home/runner/actions-runner/_work/nm-cicd/nm-cicd/.venv/lib/python3.12/site-packages/vllm/platforms/tpu.py", line 165, in check_and_update_config (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] cache_config.block_size = PallasAttentionBackend.get_page_size(vllm_config) # type: ignore[assignment] (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] File "/home/runner/actions-runner/_work/nm-cicd/nm-cicd/.venv/lib/python3.12/site-packages/vllm/v1/attention/backends/pallas.py", line 162, in get_page_size (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] if vllm_config.model_config.max_model_len > 8192: (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] AttributeError: 'NoneType' object has no attribute 'max_model_len' (Worker_TP2 pid=3714832) ERROR 12-04 10:29:38 [multiproc_executor.py:815] Traceback (most recent call last):

cc @mgoin

Fix is here: #30331

laithsakka requested review from NickLucche, ProExpertProg, WoosukKwon, hmellor, houseroad, mgoin, robertgshaw2-redhat, simon-mo, tlrmchlsmth, yewentao256, youkaichao and zou3519 as code owners September 17, 2025 23:14

mergify bot added v1 tpu Related to Google TPUs labels Sep 17, 2025

gemini-code-assist bot reviewed Sep 17, 2025

View reviewed changes

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

laithsakka marked this pull request as draft September 17, 2025 23:18

laithsakka force-pushed the something branch from 4c055dc to 42e4dd4 Compare September 17, 2025 23:44

mergify bot removed the tpu Related to Google TPUs label Sep 17, 2025

laithsakka force-pushed the something branch 3 times, most recently from 6f72573 to bcc0f99 Compare September 18, 2025 00:05

mergify bot added the tpu Related to Google TPUs label Sep 18, 2025

laithsakka marked this pull request as ready for review September 18, 2025 00:21

laithsakka commented Sep 19, 2025

View reviewed changes

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

mergify bot added the needs-rebase label Sep 21, 2025

laithsakka mentioned this pull request Oct 3, 2025

Add option to use unbacked, and backed size obl dynamic shapes for more sounds compilation. #26199

Merged

ProExpertProg reviewed Oct 7, 2025

View reviewed changes

vllm/compilation/wrapper.py Show resolved Hide resolved

ProExpertProg approved these changes Oct 7, 2025

View reviewed changes

mergify bot added the llama Related to Llama models label Nov 11, 2025

laithsakka force-pushed the something branch 3 times, most recently from 2b04ba5 to 1e420a3 Compare November 12, 2025 05:24

mergify bot added the needs-rebase label Nov 12, 2025

laithsakka force-pushed the something branch from 1e420a3 to 3500329 Compare November 12, 2025 08:07

zou3519 self-requested a review November 12, 2025 14:39

zou3519 reviewed Nov 12, 2025

View reviewed changes

vllm/compilation/wrapper.py Show resolved Hide resolved

zou3519 reviewed Nov 12, 2025

View reviewed changes

vllm/compilation/wrapper.py Show resolved Hide resolved

laithsakka force-pushed the something branch from 3500329 to 84fbf7d Compare November 12, 2025 21:47

laithsakka requested review from DarkLight1337 and ywang96 as code owners November 12, 2025 21:47

mergify bot added multi-modality Related to multi-modality (#4194) qwen Related to Qwen models labels Nov 12, 2025

remove the bytecode hook and replace TorchCompileWrapperWithCustomDis…

84b93a5

…patcher with TorchCompileGuardsStripWrapper Signed-off-by: Laith Sakka <lsakka@meta.com>

laithsakka force-pushed the something branch from 84fbf7d to 84b93a5 Compare November 13, 2025 16:42

zou3519 approved these changes Nov 14, 2025

View reviewed changes

ProExpertProg approved these changes Nov 14, 2025

View reviewed changes

mgoin approved these changes Nov 14, 2025

View reviewed changes

This was referenced Nov 25, 2025

upgrade torch npu version vllm-project/vllm-ascend#4433

Merged

upgrade to vllm 0.11.2 vllm-project/vllm-ascend#4400

Merged

dtrifiro reviewed Dec 5, 2025

View reviewed changes

dtrifiro mentioned this pull request Dec 9, 2025

[Bugfix] tpu_model_runner: set vllm config context when calling reset_dynamo_cache() #30331

Merged

		compiled_model.compiled = False
		TorchCompileWithNoGuardsWrapper.__init__(compiled_model)

Uh oh!

Conversation

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

Uh oh!

performance:

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

mergify bot commented Sep 21, 2025

Uh oh!

ProExpertProg commented Sep 23, 2025

Uh oh!

Uh oh!

ProExpertProg left a comment

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Nov 12, 2025

Uh oh!

Uh oh!

Uh oh!

laithsakka commented Nov 13, 2025

Uh oh!

laithsakka commented Nov 14, 2025

Uh oh!

hmellor commented Nov 15, 2025

Uh oh!

ProExpertProg commented Nov 15, 2025

Uh oh!

DarkLight1337 commented Nov 15, 2025

Uh oh!

hmellor commented Nov 15, 2025

Uh oh!

dtrifiro Dec 5, 2025

Choose a reason for hiding this comment

Uh oh!

dtrifiro Dec 9, 2025

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

8 participants

laithsakka commented Sep 17, 2025 •

edited by github-actions bot

Loading