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

tests/compile/fullgraph/test_basic_correctness.py

@@ -56,8 +56,10 @@ class TestSetting:
         ),
     ],
 )
+@pytest.mark.parametrize("use_v2_model_runner", [False, True])
 def test_compile_correctness(
     test_setting: TestSetting,
+    use_v2_model_runner: bool,
 ):
     model = test_setting.model
     model_args = test_setting.model_args
@@ -75,6 +77,10 @@ def test_compile_correctness(
         f"--attention-backend={attn_backend}",
     ]
 
+    # Pin every compared setting to the same model runner so that only the
+    # compilation mode differs within a comparison.
+    runner_env = {"VLLM_USE_V2_MODEL_RUNNER": str(int(use_v2_model_runner))}
+
     all_args: list[list[str]] = []
     all_envs: list[dict[str, str] | None] = []
 
@@ -86,7 +92,7 @@ def test_compile_correctness(
         CompilationMode.VLLM_COMPILE,
     ]:
         all_args.append(final_args + [f"-cc.mode={mode.name}", "-cc.backend=inductor"])
-        all_envs.append({})
+        all_envs.append(dict(runner_env))
     # inductor will change the output, so we only compare if the output
     # is close, not exactly the same.
     compare_all_settings(
@@ -94,7 +100,6 @@ def test_compile_correctness(
         all_args,
         all_envs,
         method=method if method != "generate" else "generate_close",
-        force_v1_runner=True,
     )
 
     all_envs.clear()
@@ -108,5 +113,5 @@ def test_compile_correctness(
         CompilationMode.VLLM_COMPILE,
     ]:
         all_args.append(final_args + [f"-cc.mode={mode.name}", "-cc.backend=eager"])
-        all_envs.append({})
-    compare_all_settings(model, all_args, all_envs, method=method, force_v1_runner=True)
+        all_envs.append(dict(runner_env))
+    compare_all_settings(model, all_args, all_envs, method=method)

tests/compile/test_config.py

@@ -215,7 +215,9 @@ def test_use_cudagraphs(
 
 # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
 @pytest.mark.forked
-def test_stock_torch_compile(vllm_runner, monkeypatch):
+@pytest.mark.parametrize("use_v2_model_runner", [False, True])
+def test_stock_torch_compile(vllm_runner, monkeypatch, use_v2_model_runner):
+    monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", str(int(use_v2_model_runner)))
     # Disable multiprocessing so that the counter is in the same process
     monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
 
@@ -226,9 +228,10 @@ def test_stock_torch_compile(vllm_runner, monkeypatch):
             "facebook/opt-125m",
             compilation_config={"mode": CompilationMode.STOCK_TORCH_COMPILE},
             gpu_memory_utilization=0.4,
-        ) as _,
+        ) as runner,
     ):
-        pass
+        outputs = runner.generate_greedy(["Hello, my name is"], max_tokens=5)
+        assert outputs[0][0]
 
 
 # forked needed to workaround https://github.com/vllm-project/vllm/issues/21073

tests/compile/test_wrapper.py

@@ -7,7 +7,11 @@ import os
 import pytest
 import torch
 
-from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper
+from vllm.compilation.counter import compilation_counter
+from vllm.compilation.wrapper import (
+    TorchCompileWithNoGuardsWrapper,
+    compile_model_with_stock_torch,
+)
 from vllm.config import (
     CompilationConfig,
     CompilationMode,
@@ -117,6 +121,44 @@ def test_torch_compile_wrapper(use_bytecode_hook, monkeypatch):
         raise AssertionError("expected an exception to be raised")
 
 
+def test_compile_model_with_stock_torch(monkeypatch):
+    """The whole model is compiled in place with the configured backend."""
+    graphs: list[torch.fx.GraphModule] = []
+
+    def recording_backend(gm: torch.fx.GraphModule, example_inputs):
+        graphs.append(gm)
+        return gm.forward
+
+    monkeypatch.setattr(
+        CompilationConfig,
+        "init_backend",
+        lambda self, vllm_config, *args, **kwargs: recording_backend,
+    )
+
+    class Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = torch.nn.Linear(4, 2)
+
+        def forward(self, x: torch.Tensor):
+            return self.linear(x).relu()
+
+    vllm_config = VllmConfig()
+    vllm_config.compilation_config = CompilationConfig(
+        mode=CompilationMode.STOCK_TORCH_COMPILE
+    )
+    model = Model()
+    x = torch.randn(3, 4)
+    expected = model(x)
+
+    torch._dynamo.reset()
+    with compilation_counter.expect(stock_torch_compile_count=1):
+        compile_model_with_stock_torch(model, vllm_config)
+    assert not graphs  # compilation is deferred to the first call
+    torch.testing.assert_close(model(x), expected)
+    assert len(graphs) == 1
+
+
 if __name__ == "__main__":
     # Run with both parameter values
 

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -7,6 +7,7 @@ from typing import Any
 
 import torch
 
+from vllm.config import CompilationConfig, CompilationMode
 from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry
 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
 from vllm.v1.worker.gpu import eplb_utils as eplb
@@ -68,7 +69,9 @@ def _make_runner(**overrides: Any) -> Any:
         enable_batch_sharded_sampling=False,
         eplb_config=SimpleNamespace(log_balancedness=True),
     )
+    runner.compilation_config = CompilationConfig(mode=CompilationMode.NONE)
     runner.vllm_config = SimpleNamespace(
+        compilation_config=runner.compilation_config,
         load_config=runner.load_config,
         model_config=runner.model_config,
     )

vllm/compilation/wrapper.py

@@ -12,7 +12,13 @@ from typing import Any, ParamSpec, TypeVar
 import torch
 
 import vllm.envs as envs
-from vllm.config import CompilationMode, CUDAGraphMode, get_current_vllm_config
+from vllm.compilation.counter import compilation_counter
+from vllm.config import (
+    CompilationMode,
+    CUDAGraphMode,
+    VllmConfig,
+    get_current_vllm_config,
+)
 from vllm.config.compilation import DynamicShapesType
 from vllm.logger import init_logger
 from vllm.utils.nvtx_pytorch_hooks import layerwise_nvtx_marker_context
@@ -44,6 +50,17 @@ def _compilation_context() -> Generator[None, None, None]:
         torch._dynamo.config.accumulated_cache_size_limit = original_accumulated_cache
 
 
+def compile_model_with_stock_torch(
+    model: torch.nn.Module, vllm_config: VllmConfig
+) -> None:
+    from vllm.env_override import _apply_constrain_to_fx_strides_patch
+
+    _apply_constrain_to_fx_strides_patch()
+    backend = vllm_config.compilation_config.init_backend(vllm_config)
+    compilation_counter.stock_torch_compile_count += 1
+    model.compile(fullgraph=True, backend=backend)
+
+
 class TorchCompileWithNoGuardsWrapper:
     """A wrapper class for torch.compile, it ensures that all guards are dropped
     when CompilationMode is not CompilationMode.STOCK_TORCH_COMPILE.

vllm/config/vllm.py

@@ -2041,6 +2041,7 @@ class VllmConfig:
         self._validate_mm_processor_device()
 
         if self.use_v2_model_runner:
+            self._disable_cudagraphs_for_v2_stock_torch_compile()
             self._validate_v2_model_runner()
         else:
             self._validate_v1_model_runner()
@@ -3045,9 +3046,6 @@ class VllmConfig:
         unsupported: list[str] = []
         speculative_config = self.speculative_config
 
-        if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE:
-            unsupported.append("stock torch.compile")
-
         if (
             self.compilation_config.pass_config.enable_sp
             and self.parallel_config.tensor_parallel_size > 1
@@ -3301,6 +3299,28 @@ class VllmConfig:
                 "Enable attribution or use --enforce-eager."
             )
 
+    def _disable_cudagraphs_for_v2_stock_torch_compile(self) -> None:
+        """Run stock torch.compile without CUDA graphs in Model Runner V2.
+
+        V1 never wraps a stock-compiled model in CUDAGraphWrapper, so it runs
+        without CUDA graphs. V2's CUDA graph manager would otherwise capture
+        FULL graphs around the stock-compiled model, so disable them to match.
+        """
+        compilation_config = self.compilation_config
+        if (
+            compilation_config.mode != CompilationMode.STOCK_TORCH_COMPILE
+            or compilation_config.cudagraph_mode == CUDAGraphMode.NONE
+        ):
+            return
+        logger.info_once(
+            "CUDA graphs are not supported with stock torch.compile in Model "
+            "Runner V2. Overriding cudagraph_mode %s to NONE.",
+            compilation_config.cudagraph_mode.name,
+        )
+        compilation_config.cudagraph_mode = CUDAGraphMode.NONE
+        compilation_config.max_cudagraph_capture_size = 0
+        compilation_config.cudagraph_capture_sizes = []
+
     def _validate_v2_model_runner(self) -> None:
         """Check for features not yet supported by the V2 model runner."""
         if not HAS_TRITON:

vllm/v1/worker/gpu/model_runner.py

@@ -31,8 +31,9 @@ import torch.nn as nn
 import vllm.envs as envs
 from vllm.compilation.counter import compilation_counter
 from vllm.compilation.cuda_graph import CUDAGraphStat
+from vllm.compilation.wrapper import compile_model_with_stock_torch
 from vllm.config import VllmConfig
-from vllm.config.compilation import CUDAGraphMode
+from vllm.config.compilation import CompilationMode, CUDAGraphMode
 from vllm.distributed.aux_output_connector.worker import (
     AuxOutputWorkerConnector,
     get_aux_output_connector,
@@ -522,6 +523,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
 
         get_offloader().post_init()
 
+        if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE:
+            compile_model_with_stock_torch(self.model, self.vllm_config)
+
     @property
     def dp_size(self) -> int:
         # Not cached: elastic EP rewrites parallel_config in place on reconfigure.