[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] <your.email@example.com>

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

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

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)
 

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

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: