From 643bbef55c37edb5c3fd341aae1e20e026d0e730 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Wed, 30 Sep 2026 22:29:12 +0800 Subject: [PATCH] [Refactor] Share workspace and model runner init between GPU and XPU workers (#59200) Signed-off-by: aoshen02 Co-authored-by: Claude Opus 5.5 --- vllm/v1/worker/gpu_worker.py | 26 +++++++++++++------------- vllm/v1/worker/xpu_worker.py | 16 +++------------- 2 files changed, 16 insertions(+), 26 deletions(-) diff --git a/vllm/v1/worker/gpu_worker.py b/vllm/v1/worker/gpu_worker.py index 8b7a0260a7..247265b659 100644 --- a/vllm/v1/worker/gpu_worker.py +++ b/vllm/v1/worker/gpu_worker.py @@ -522,6 +522,10 @@ class Worker(WorkerBase): else: raise RuntimeError(f"Unsupported device type: {self.device_config.device}") + self._init_workspace_and_model_runner() + + def _init_workspace_and_model_runner(self) -> None: + """Set up the workspace manager, build the model runner, report usage.""" # DSpark target and draft CUDA graphs retain workspace views concurrently. num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 init_workspace_manager( @@ -529,8 +533,13 @@ class Worker(WorkerBase): num_ubatches, _num_workspace_lanes(self.vllm_config, self.use_v2_model_runner), ) + self.model_runner: GPUModelRunner = self._make_model_runner() + if self.rank == 0: + # If usage stat is enabled, collect relevant info. + report_usage_stats(self.vllm_config) - # Construct the model runner + def _make_model_runner(self) -> "GPUModelRunner": + """Construct the platform's model runner; subclasses override this.""" if self.use_v2_model_runner: if self.vllm_config.is_mm_encoder_only: from vllm.v1.worker.mm_encoder_model_runner import ( @@ -542,19 +551,10 @@ class Worker(WorkerBase): ) # HACK(woosuk): This is a temporary fix to avoid type errors. - self.model_runner: GPUModelRunner = GPUModelRunnerV2( # type: ignore - self.vllm_config, self.device - ) - else: - from vllm.v1.worker.gpu_model_runner import ( - GPUModelRunner as GPUModelRunnerV1, - ) + return GPUModelRunnerV2(self.vllm_config, self.device) # type: ignore + from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1 - self.model_runner = GPUModelRunnerV1(self.vllm_config, self.device) - - if self.rank == 0: - # If usage stat is enabled, collect relevant info. - report_usage_stats(self.vllm_config) + return GPUModelRunnerV1(self.vllm_config, self.device) def handle_ft_command(self, ft_request): assert self.worker_sentinel is not None diff --git a/vllm/v1/worker/xpu_worker.py b/vllm/v1/worker/xpu_worker.py index 8a8166c435..f7539d2c20 100644 --- a/vllm/v1/worker/xpu_worker.py +++ b/vllm/v1/worker/xpu_worker.py @@ -10,9 +10,7 @@ from vllm.logger import init_logger from vllm.platforms import current_platform from vllm.utils.mem_utils import MemorySnapshot, format_gib from vllm.utils.torch_utils import set_random_seed -from vllm.v1.utils import report_usage_stats from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment -from vllm.v1.worker.workspace import init_workspace_manager from vllm.v1.worker.xpu_model_runner import ( XPUMMEncoderModelRunner, XPUModelRunner, @@ -214,11 +212,9 @@ class XPUWorker(Worker): "worker requested memory: %sGiB", format_gib(self.requested_memory) ) - # Initialize workspace manager - num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 - init_workspace_manager(self.device, num_ubatches) + self._init_workspace_and_model_runner() - # Construct the model runner + def _make_model_runner(self): if self.use_v2_model_runner: model_runner = ( XPUMMEncoderModelRunner @@ -227,13 +223,7 @@ class XPUWorker(Worker): ) else: model_runner = XPUModelRunner # type: ignore[assignment] - self.model_runner = model_runner( # type: ignore - self.vllm_config, self.device - ) - - if self.rank == 0: - # If usage stat is enabled, collect relevant info. - report_usage_stats(self.vllm_config) + return model_runner(self.vllm_config, self.device) # type: ignore def shutdown(self) -> None: logger.info(