[Refactor] Share workspace and model runner init between GPU and XPU workers (#59200)
Signed-off-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
aoshen02 committed
Sep 30, 2026 at 22:29 UTC
643bbef55c37edb5c3fd341aae1e20e026d0e730
2 files changed
+16
-26
vllm/v1/worker/gpu_worker.py
+13
-13
@@ -522,6 +522,10 @@ class Worker(WorkerBase):
522
else:
523
raise RuntimeError(f"Unsupported device type: {self.device_config.device}")
524
525
+ self._init_workspace_and_model_runner()
526
+
527
+ def _init_workspace_and_model_runner(self) -> None:
528
+ """Set up the workspace manager, build the model runner, report usage."""
529
# DSpark target and draft CUDA graphs retain workspace views concurrently.
530
num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
531
init_workspace_manager(
@@ -529,8 +533,13 @@ class Worker(WorkerBase):
533
num_ubatches,
534
_num_workspace_lanes(self.vllm_config, self.use_v2_model_runner),
535
)
536
+ self.model_runner: GPUModelRunner = self._make_model_runner()
537
+ if self.rank == 0:
538
+ # If usage stat is enabled, collect relevant info.
539
+ report_usage_stats(self.vllm_config)
540
533
- # Construct the model runner
541
+ def _make_model_runner(self) -> "GPUModelRunner":
542
+ """Construct the platform's model runner; subclasses override this."""
543
if self.use_v2_model_runner:
544
if self.vllm_config.is_mm_encoder_only:
545
from vllm.v1.worker.mm_encoder_model_runner import (
@@ -542,19 +551,10 @@ class Worker(WorkerBase):
551
)
552
553
# HACK(woosuk): This is a temporary fix to avoid type errors.
545
- self.model_runner: GPUModelRunner = GPUModelRunnerV2( # type: ignore
546
- self.vllm_config, self.device
547
- )
548
- else:
549
- from vllm.v1.worker.gpu_model_runner import (
550
- GPUModelRunner as GPUModelRunnerV1,
551
- )
552
-
553
- self.model_runner = GPUModelRunnerV1(self.vllm_config, self.device)
554
+ return GPUModelRunnerV2(self.vllm_config, self.device) # type: ignore
555
+ from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1
556
555
- if self.rank == 0:
556
- # If usage stat is enabled, collect relevant info.
557
- report_usage_stats(self.vllm_config)
557
+ return GPUModelRunnerV1(self.vllm_config, self.device)
558
559
def handle_ft_command(self, ft_request):
560
assert self.worker_sentinel is not None
vllm/v1/worker/xpu_worker.py
+3
-13
@@ -10,9 +10,7 @@ from vllm.logger import init_logger
10
from vllm.platforms import current_platform
11
from vllm.utils.mem_utils import MemorySnapshot, format_gib
12
from vllm.utils.torch_utils import set_random_seed
13
-from vllm.v1.utils import report_usage_stats
13
from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
15
-from vllm.v1.worker.workspace import init_workspace_manager
14
from vllm.v1.worker.xpu_model_runner import (
15
XPUMMEncoderModelRunner,
16
XPUModelRunner,
@@ -214,11 +212,9 @@ class XPUWorker(Worker):
212
"worker requested memory: %sGiB", format_gib(self.requested_memory)
213
)
214
217
- # Initialize workspace manager
218
- num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
219
- init_workspace_manager(self.device, num_ubatches)
215
+ self._init_workspace_and_model_runner()
216
221
- # Construct the model runner
217
+ def _make_model_runner(self):
218
if self.use_v2_model_runner:
219
model_runner = (
220
XPUMMEncoderModelRunner
@@ -227,13 +223,7 @@ class XPUWorker(Worker):
223
)
224
else:
225
model_runner = XPUModelRunner # type: ignore[assignment]
230
- self.model_runner = model_runner( # type: ignore
231
- self.vllm_config, self.device
232
- )
233
-
234
- if self.rank == 0:
235
- # If usage stat is enabled, collect relevant info.
236
- report_usage_stats(self.vllm_config)
226
+ return model_runner(self.vllm_config, self.device) # type: ignore
227
228
def shutdown(self) -> None:
229
logger.info(