Merge branch 'main' into wentao-model-runner-v2-support-stock-torch-compile

Signed-off-by: yewentao256 <zhyanwentao@126.com>

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -7,6 +7,7 @@ from typing import Any
 
 import torch
 
+from vllm.config.compilation import CompilationMode
 from vllm.v1.worker.gpu import eplb_utils as eplb
 from vllm.v1.worker.gpu import model_runner as mrv2
 
@@ -56,6 +57,7 @@ def _make_runner(**overrides: Any) -> Any:
     runner.device = torch.device("cpu")
     runner.model_config = SimpleNamespace(model="test-model")
     runner.load_config = SimpleNamespace(load_format="hf")
+    runner.compilation_config = SimpleNamespace(mode=CompilationMode.NONE)
     runner.parallel_config = SimpleNamespace(
         enable_eplb=True,
         enable_elastic_ep=False,
@@ -64,6 +66,7 @@ def _make_runner(**overrides: Any) -> Any:
     runner.vllm_config = SimpleNamespace(
         load_config=runner.load_config,
         model_config=runner.model_config,
+        compilation_config=runner.compilation_config,
     )
     runner.lora_config = None
     runner.use_aux_hidden_state_outputs = False

vllm/v1/worker/gpu/model_runner.py

@@ -29,7 +29,7 @@ import torch.nn as nn
 
 from vllm.compilation.counter import compilation_counter
 from vllm.config import VllmConfig
-from vllm.config.compilation import CUDAGraphMode
+from vllm.config.compilation import CompilationMode, CUDAGraphMode
 from vllm.distributed.parallel_state import (
     get_dcp_group,
     get_pp_group,
@@ -325,6 +325,14 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                 device=self.device,
             )
 
+        if self.compilation_config.mode == CompilationMode.STOCK_TORCH_COMPILE:
+            from vllm.env_override import _apply_constrain_to_fx_strides_patch
+
+            _apply_constrain_to_fx_strides_patch()
+            backend = self.compilation_config.init_backend(self.vllm_config)
+            compilation_counter.stock_torch_compile_count += 1
+            self.model.compile(fullgraph=True, backend=backend)
+
     def get_model(self) -> nn.Module:
         return self.model