From 2c2d06d9208ac9a4bece75288c47989c32a2381d Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 28 Jan 2026 15:43:27 +0800 Subject: [PATCH] [Enhancement] Update GPUARWorker and GPUGenerationWorker to enforce v1 model runner - Added a warning in GPUARWorker and GPUGenerationWorker to indicate that the v2 model runner does not include omni hooks, forcing the use of the v1 model runner. - Updated import statements in gpu_ar_model_runner.py and gpu_generation_model_runner.py to include necessary functions from parallel_state. - Removed unused imports from both model runner files to streamline the code. Signed-off-by: [Your Name] Signed-off-by: tzhouam --- vllm_omni/worker/gpu_ar_model_runner.py | 6 ++---- vllm_omni/worker/gpu_ar_worker.py | 5 +++++ vllm_omni/worker/gpu_generation_model_runner.py | 4 ++-- vllm_omni/worker/gpu_generation_worker.py | 5 +++++ 4 files changed, 14 insertions(+), 6 deletions(-) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index ecc839fd3..078217fcc 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -13,7 +13,8 @@ import numpy as np import torch from vllm.config import CUDAGraphMode from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer -from vllm.distributed.kv_transfer import get_kv_transfer_group +from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group +from vllm.distributed.parallel_state import get_pp_group, get_tp_group from vllm.forward_context import set_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( @@ -28,9 +29,6 @@ from vllm.v1.worker.gpu_model_runner import ( EMPTY_MODEL_RUNNER_OUTPUT, AsyncGPUModelRunnerOutput, IntermediateTensors, - get_pp_group, - get_tp_group, - has_kv_transfer_group, ) from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import is_residual_scattered_for_sp diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index d2dafce68..946e52516 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -90,6 +90,11 @@ class GPUARWorker(GPUWorker): num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 init_workspace_manager(self.device, num_ubatches) + if self.use_v2_model_runner: + # OMNI: v2 model runner does not yet include omni hooks. + logger.warning("OMNI GPUARWorker forces v1 model runner for omni hooks.") + self.use_v2_model_runner = False + # Construct the model runner self.model_runner = GPUARModelRunner(self.vllm_config, self.device) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 87e1419c8..ee9fe01de 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -15,6 +15,8 @@ import torch from vllm.config import CUDAGraphMode from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group +from vllm.distributed.parallel_state import get_pp_group +from vllm.forward_context import set_forward_context from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( RoutedExpertsCapturer, ) @@ -29,8 +31,6 @@ from vllm.v1.worker.gpu_model_runner import ( AsyncGPUModelRunnerOutput, IntermediateTensors, PerLayerAttnMetadata, - get_pp_group, - set_forward_context, ) from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 19f8ab84b..d2f254b3c 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -90,6 +90,11 @@ class GPUGenerationWorker(GPUWorker): num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 init_workspace_manager(self.device, num_ubatches) + if self.use_v2_model_runner: + # OMNI: v2 model runner does not yet include omni hooks. + logger.warning("OMNI GPUGenerationWorker forces v1 model runner for omni hooks.") + self.use_v2_model_runner = False + self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device) if self.rank == 0: