[MRV2] Support stock torch.compile mode (#59079)
Signed-off-by: cheesecake <farzanaman99@gmail.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Cheese Cake committed
Sep 29, 2026 at 08:47 UTC
4e59d53ab534fbb4aaa751eb37ecfbcf438855b5
7 files changed
+107
-13
tests/compile/fullgraph/test_basic_correctness.py
+9
-4
@@ -56,8 +56,10 @@ class TestSetting:
56
),
57
],
58
)
59
+@pytest.mark.parametrize("use_v2_model_runner", [False, True])
60
def test_compile_correctness(
61
test_setting: TestSetting,
62
+ use_v2_model_runner: bool,
63
):
64
model = test_setting.model
65
model_args = test_setting.model_args
@@ -75,6 +77,10 @@ def test_compile_correctness(
77
f"--attention-backend={attn_backend}",
78
]
79
80
+ # Pin every compared setting to the same model runner so that only the
81
+ # compilation mode differs within a comparison.
82
+ runner_env = {"VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner))}
83
+
84
all_args: list[list[str]] = []
85
all_envs: list[dict[str, str] | None] = []
86
@@ -86,7 +92,7 @@ def test_compile_correctness(
92
CompilationMode.VLLM_COMPILE,
93
]:
94
all_args.append(final_args + [f"-cc.mode={mode.name}", "-cc.backend=inductor"])
89
- all_envs.append({})
95
+ all_envs.append(dict(runner_env))
96
# inductor will change the output, so we only compare if the output
97
# is close, not exactly the same.
98
compare_all_settings(
@@ -94,7 +100,6 @@ def test_compile_correctness(
100
all_args,
101
all_envs,
102
method=method if method != "generate" else "generate_close",
97
- force_v1_runner=True,
103
)
104
105
all_envs.clear()
@@ -108,5 +113,5 @@ def test_compile_correctness(
113
CompilationMode.VLLM_COMPILE,
114
]:
115
all_args.append(final_args + [f"-cc.mode={mode.name}", "-cc.backend=eager"])
111
- all_envs.append({})
112
- compare_all_settings(model, all_args, all_envs, method=method, force_v1_runner=True)
116
+ all_envs.append(dict(runner_env))
117
+ compare_all_settings(model, all_args, all_envs, method=method)
tests/compile/test_config.py
+6
-3
@@ -215,7 +215,9 @@ def test_use_cudagraphs(
215
216
# forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
217
@pytest.mark.forked
218
-def test_stock_torch_compile(vllm_runner, monkeypatch):
218
+@pytest.mark.parametrize("use_v2_model_runner", [False, True])
219
+def test_stock_torch_compile(vllm_runner, monkeypatch, use_v2_model_runner):
220
+ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner)))
221
# Disable multiprocessing so that the counter is in the same process
222
monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
223
@@ -226,9 +228,10 @@ def test_stock_torch_compile(vllm_runner, monkeypatch):
228
"facebook/opt-125m",
229
compilation_config={"mode": CompilationMode.STOCK_TORCH_COMPILE},
230
gpu_memory_utilization=0.4,
229
- ) as _,
231
+ ) as runner,
232
):
231
- pass
233
+ outputs = runner.generate_greedy(["Hello, my name is"], max_tokens=5)
234
+ assert outputs[0][0]
235
236
237
# forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
tests/compile/test_wrapper.py
+43
-1
@@ -7,7 +7,11 @@ import os
7
import pytest
8
import torch
9
10
-from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
10
+from vllm.compilation.counter import compilation_counter
11
+from vllm.compilation.wrapper import (
12
+ TorchCompileWithNoGuardsWrapper,
13
+ compile_model_with_stock_torch,
14
+)
15
from vllm.config import (
16
CompilationConfig,
17
CompilationMode,
@@ -117,6 +121,44 @@ def test_torch_compile_wrapper(use_bytecode_hook, monkeypatch):
121
raise AssertionError("expected an exception to be raised")
122
123
124
+def test_compile_model_with_stock_torch(monkeypatch):
125
+ """The whole model is compiled in place with the configured backend."""
126
+ graphs: list[torch.fx.GraphModule] = []
127
+
128
+ def recording_backend(gm: torch.fx.GraphModule, example_inputs):
129
+ graphs.append(gm)
130
+ return gm.forward
131
+
132
+ monkeypatch.setattr(
133
+ CompilationConfig,
134
+ "init_backend",
135
+ lambda self, vllm_config, *args, **kwargs: recording_backend,
136
+ )
137
+
138
+ class Model(torch.nn.Module):
139
+ def __init__(self):
140
+ super().__init__()
141
+ self.linear = torch.nn.Linear(4, 2)
142
+
143
+ def forward(self, x: torch.Tensor):
144
+ return self.linear(x).relu()
145
+
146
+ vllm_config = VllmConfig()
147
+ vllm_config.compilation_config = CompilationConfig(
148
+ mode=CompilationMode.STOCK_TORCH_COMPILE
149
+ )
150
+ model = Model()
151
+ x = torch.randn(3, 4)
152
+ expected = model(x)
153
+
154
+ torch._dynamo.reset()
155
+ with compilation_counter.expect(stock_torch_compile_count=1):
156
+ compile_model_with_stock_torch(model, vllm_config)
157
+ assert not graphs # compilation is deferred to the first call
158
+ torch.testing.assert_close(model(x), expected)
159
+ assert len(graphs) == 1
160
+
161
+
162
if __name__ == "__main__":
163
# Run with both parameter values
164
tests/v1/worker/test_gpu_model_runner_v2_eplb.py
+3
@@ -7,6 +7,7 @@ from typing import Any
7
8
import torch
9
10
+from vllm.config import CompilationConfig, CompilationMode
11
from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry
12
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
13
from vllm.v1.worker.gpu import eplb_utils as eplb
@@ -68,7 +69,9 @@ def _make_runner(**overrides: Any) -> Any:
69
enable_batch_sharded_sampling=False,
70
eplb_config=SimpleNamespace(log_balancedness=True),
71
)
72
+ runner.compilation_config = CompilationConfig(mode=CompilationMode.NONE)
73
runner.vllm_config = SimpleNamespace(
74
+ compilation_config=runner.compilation_config,
75
load_config=runner.load_config,
76
model_config=runner.model_config,
77
)
vllm/compilation/wrapper.py
+18
-1
@@ -12,7 +12,13 @@ from typing import Any, ParamSpec, TypeVar
12
import torch
13
14
import vllm.envs as envs
15
-from vllm.config import CompilationMode, CUDAGraphMode, get_current_vllm_config
15
+from vllm.compilation.counter import compilation_counter
16
+from vllm.config import (
17
+ CompilationMode,
18
+ CUDAGraphMode,
19
+ VllmConfig,
20
+ get_current_vllm_config,
21
+)
22
from vllm.config.compilation import DynamicShapesType
23
from vllm.logger import init_logger
24
from vllm.utils.nvtx_pytorch_hooks import layerwise_nvtx_marker_context
@@ -44,6 +50,17 @@ def _compilation_context() -> Generator[None, None, None]:
50
torch._dynamo.config.accumulated_cache_size_limit = original_accumulated_cache
51
52
53
+def compile_model_with_stock_torch(
54
+ model: torch.nn.Module, vllm_config: VllmConfig
55
+) -> None:
56
+ from vllm.env_override import _apply_constrain_to_fx_strides_patch
57
+
58
+ _apply_constrain_to_fx_strides_patch()
59
+ backend = vllm_config.compilation_config.init_backend(vllm_config)
60
+ compilation_counter.stock_torch_compile_count += 1
61
+ model.compile(fullgraph=True, backend=backend)
62
+
63
+
64
class TorchCompileWithNoGuardsWrapper:
65
"""A wrapper class for torch.compile, it ensures that all guards are dropped
66
when CompilationMode is not CompilationMode.STOCK_TORCH_COMPILE.
vllm/config/vllm.py
+23
-3
@@ -2041,6 +2041,7 @@ class VllmConfig:
2041
self._validate_mm_processor_device()
2042
2043
if self.use_v2_model_runner:
2044
+ self._disable_cudagraphs_for_v2_stock_torch_compile()
2045
self._validate_v2_model_runner()
2046
else:
2047
self._validate_v1_model_runner()
@@ -3045,9 +3046,6 @@ class VllmConfig:
3046
unsupported: list[str] = []
3047
speculative_config = self.speculative_config
3048
3048
- if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE:
3049
- unsupported.append("stock torch.compile")
3050
-
3049
if (
3050
self.compilation_config.pass_config.enable_sp
3051
and self.parallel_config.tensor_parallel_size > 1
@@ -3301,6 +3299,28 @@ class VllmConfig:
3299
"Enable attribution or use --enforce-eager."
3300
)
3301
3302
+ def _disable_cudagraphs_for_v2_stock_torch_compile(self) -> None:
3303
+ """Run stock torch.compile without CUDA graphs in Model Runner V2.
3304
+
3305
+ V1 never wraps a stock-compiled model in CUDAGraphWrapper, so it runs
3306
+ without CUDA graphs. V2's CUDA graph manager would otherwise capture
3307
+ FULL graphs around the stock-compiled model, so disable them to match.
3308
+ """
3309
+ compilation_config = self.compilation_config
3310
+ if (
3311
+ compilation_config.mode != CompilationMode.STOCK_TORCH_COMPILE
3312
+ or compilation_config.cudagraph_mode == CUDAGraphMode.NONE
3313
+ ):
3314
+ return
3315
+ logger.info_once(
3316
+ "CUDA graphs are not supported with stock torch.compile in Model "
3317
+ "Runner V2. Overriding cudagraph_mode %s to NONE.",
3318
+ compilation_config.cudagraph_mode.name,
3319
+ )
3320
+ compilation_config.cudagraph_mode = CUDAGraphMode.NONE
3321
+ compilation_config.max_cudagraph_capture_size = 0
3322
+ compilation_config.cudagraph_capture_sizes = []
3323
+
3324
def _validate_v2_model_runner(self) -> None:
3325
"""Check for features not yet supported by the V2 model runner."""
3326
if not HAS_TRITON:
vllm/v1/worker/gpu/model_runner.py
+5
-1
@@ -31,8 +31,9 @@ import torch.nn as nn
31
import vllm.envs as envs
32
from vllm.compilation.counter import compilation_counter
33
from vllm.compilation.cuda_graph import CUDAGraphStat
34
+from vllm.compilation.wrapper import compile_model_with_stock_torch
35
from vllm.config import VllmConfig
35
-from vllm.config.compilation import CUDAGraphMode
36
+from vllm.config.compilation import CompilationMode, CUDAGraphMode
37
from vllm.distributed.aux_output_connector.worker import (
38
AuxOutputWorkerConnector,
39
get_aux_output_connector,
@@ -522,6 +523,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
523
524
get_offloader().post_init()
525
526
+ if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE:
527
+ compile_model_with_stock_torch(self.model, self.vllm_config)
528
+
529
@property
530
def dp_size(self) -> int:
531
# Not cached: elastic EP rewrites parallel_config in place on reconfigure.