Refactor OmniEngineArgs and AsyncOmniEngineArgs to improve compatibility with partial payloads by using getattr for various attributes. Update GPUGenerationModelRunner and OmniGPUModelRunner to replace the activate_lora parameter with num_active_loras for better clarity and backward compatibility. This change enhances the handling of LoRA activation during model runs.

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

vllm_omni/engine/arg_utils.py

@@ -2,10 +2,9 @@ from dataclasses import dataclass, field
 from typing import Any
 
 import vllm.envs as envs
-from vllm.engine.arg_utils import EngineArgs
+from vllm.engine.arg_utils import AsyncEngineArgs, EngineArgs
 from vllm.logger import init_logger
 from vllm.transformers_utils.gguf_utils import is_gguf
-from vllm.v1.engine.async_llm import AsyncEngineArgs
 
 from vllm_omni.config import OmniModelConfig
 from vllm_omni.plugins import load_omni_general_plugins
@@ -107,6 +106,22 @@ class OmniEngineArgs(EngineArgs):
         # register omni models to avoid model not found error
         self._ensure_omni_models_registered()
 
+        # Keep compatibility when async args are constructed from partial payloads.
+        language_model_only = getattr(self, "language_model_only", False)
+        limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {})
+        enable_mm_embeds = getattr(self, "enable_mm_embeds", False)
+        interleave_mm_strings = getattr(self, "interleave_mm_strings", False)
+        media_io_kwargs = getattr(self, "media_io_kwargs", {})
+        skip_mm_profiling = getattr(self, "skip_mm_profiling", False)
+        mm_processor_kwargs = getattr(self, "mm_processor_kwargs", None)
+        mm_processor_cache_gb = getattr(self, "mm_processor_cache_gb", 4)
+        mm_processor_cache_type = getattr(self, "mm_processor_cache_type", None)
+        mm_shm_cache_max_object_size_mb = getattr(self, "mm_shm_cache_max_object_size_mb", 128)
+        mm_encoder_only = getattr(self, "mm_encoder_only", False)
+        mm_encoder_tp_mode = getattr(self, "mm_encoder_tp_mode", "weights")
+        mm_encoder_attn_backend = getattr(self, "mm_encoder_attn_backend", None)
+        video_pruning_rate = getattr(self, "video_pruning_rate", 0.0)
+
         # Build stage_connector_config from stage_connector_spec
         stage_connector_config = {
             "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"),
@@ -147,20 +162,20 @@ class OmniEngineArgs(EngineArgs):
             skip_tokenizer_init=self.skip_tokenizer_init,
             enable_prompt_embeds=self.enable_prompt_embeds,
             served_model_name=self.served_model_name,
-            language_model_only=self.language_model_only,
-            limit_mm_per_prompt=self.limit_mm_per_prompt,
-            enable_mm_embeds=self.enable_mm_embeds,
-            interleave_mm_strings=self.interleave_mm_strings,
-            media_io_kwargs=self.media_io_kwargs,
-            skip_mm_profiling=self.skip_mm_profiling,
+            language_model_only=language_model_only,
+            limit_mm_per_prompt=limit_mm_per_prompt,
+            enable_mm_embeds=enable_mm_embeds,
+            interleave_mm_strings=interleave_mm_strings,
+            media_io_kwargs=media_io_kwargs,
+            skip_mm_profiling=skip_mm_profiling,
             config_format=self.config_format,
-            mm_processor_kwargs=self.mm_processor_kwargs,
-            mm_processor_cache_gb=self.mm_processor_cache_gb,
-            mm_processor_cache_type=self.mm_processor_cache_type,
-            mm_shm_cache_max_object_size_mb=self.mm_shm_cache_max_object_size_mb,
-            mm_encoder_only=self.mm_encoder_only,
-            mm_encoder_tp_mode=self.mm_encoder_tp_mode,
-            mm_encoder_attn_backend=self.mm_encoder_attn_backend,
+            mm_processor_kwargs=mm_processor_kwargs,
+            mm_processor_cache_gb=mm_processor_cache_gb,
+            mm_processor_cache_type=mm_processor_cache_type,
+            mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
+            mm_encoder_only=mm_encoder_only,
+            mm_encoder_tp_mode=mm_encoder_tp_mode,
+            mm_encoder_attn_backend=mm_encoder_attn_backend,
             pooler_config=self.pooler_config,
             logits_processor_pattern=self.logits_processor_pattern,
             generation_config=self.generation_config,
@@ -169,7 +184,7 @@ class OmniEngineArgs(EngineArgs):
             model_impl=self.model_impl,
             override_attention_dtype=self.override_attention_dtype,
             logits_processors=self.logits_processors,
-            video_pruning_rate=self.video_pruning_rate,
+            video_pruning_rate=video_pruning_rate,
             io_processor_plugin=self.io_processor_plugin,
             # Omni-specific fields
             stage_id=self.stage_id,
@@ -246,6 +261,22 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
         # register omni models to avoid model not found error
         self._ensure_omni_models_registered()
 
+        # Keep compatibility when async args are constructed from partial payloads.
+        language_model_only = getattr(self, "language_model_only", False)
+        limit_mm_per_prompt = getattr(self, "limit_mm_per_prompt", {})
+        enable_mm_embeds = getattr(self, "enable_mm_embeds", False)
+        interleave_mm_strings = getattr(self, "interleave_mm_strings", False)
+        media_io_kwargs = getattr(self, "media_io_kwargs", {})
+        skip_mm_profiling = getattr(self, "skip_mm_profiling", False)
+        mm_processor_kwargs = getattr(self, "mm_processor_kwargs", None)
+        mm_processor_cache_gb = getattr(self, "mm_processor_cache_gb", 4)
+        mm_processor_cache_type = getattr(self, "mm_processor_cache_type", None)
+        mm_shm_cache_max_object_size_mb = getattr(self, "mm_shm_cache_max_object_size_mb", 128)
+        mm_encoder_only = getattr(self, "mm_encoder_only", False)
+        mm_encoder_tp_mode = getattr(self, "mm_encoder_tp_mode", "weights")
+        mm_encoder_attn_backend = getattr(self, "mm_encoder_attn_backend", None)
+        video_pruning_rate = getattr(self, "video_pruning_rate", 0.0)
+
         # Build stage_connector_config from stage_connector_spec
         stage_connector_config = {
             "name": self.stage_connector_spec.get("name", "SharedMemoryConnector"),
@@ -286,20 +317,20 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
             skip_tokenizer_init=self.skip_tokenizer_init,
             enable_prompt_embeds=self.enable_prompt_embeds,
             served_model_name=self.served_model_name,
-            language_model_only=self.language_model_only,
-            limit_mm_per_prompt=self.limit_mm_per_prompt,
-            enable_mm_embeds=self.enable_mm_embeds,
-            interleave_mm_strings=self.interleave_mm_strings,
-            media_io_kwargs=self.media_io_kwargs,
-            skip_mm_profiling=self.skip_mm_profiling,
+            language_model_only=language_model_only,
+            limit_mm_per_prompt=limit_mm_per_prompt,
+            enable_mm_embeds=enable_mm_embeds,
+            interleave_mm_strings=interleave_mm_strings,
+            media_io_kwargs=media_io_kwargs,
+            skip_mm_profiling=skip_mm_profiling,
             config_format=self.config_format,
-            mm_processor_kwargs=self.mm_processor_kwargs,
-            mm_processor_cache_gb=self.mm_processor_cache_gb,
-            mm_processor_cache_type=self.mm_processor_cache_type,
-            mm_shm_cache_max_object_size_mb=self.mm_shm_cache_max_object_size_mb,
-            mm_encoder_only=self.mm_encoder_only,
-            mm_encoder_tp_mode=self.mm_encoder_tp_mode,
-            mm_encoder_attn_backend=self.mm_encoder_attn_backend,
+            mm_processor_kwargs=mm_processor_kwargs,
+            mm_processor_cache_gb=mm_processor_cache_gb,
+            mm_processor_cache_type=mm_processor_cache_type,
+            mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb,
+            mm_encoder_only=mm_encoder_only,
+            mm_encoder_tp_mode=mm_encoder_tp_mode,
+            mm_encoder_attn_backend=mm_encoder_attn_backend,
             pooler_config=self.pooler_config,
             logits_processor_pattern=self.logits_processor_pattern,
             generation_config=self.generation_config,
@@ -308,7 +339,7 @@ class AsyncOmniEngineArgs(AsyncEngineArgs):
             model_impl=self.model_impl,
             override_attention_dtype=self.override_attention_dtype,
             logits_processors=self.logits_processors,
-            video_pruning_rate=self.video_pruning_rate,
+            video_pruning_rate=video_pruning_rate,
             io_processor_plugin=self.io_processor_plugin,
             # Omni-specific fields
             stage_id=self.stage_id,

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py

@@ -144,7 +144,6 @@ class Qwen3MoeLLMModel(_Qwen3MoeLLMModel):
         positions: torch.Tensor,
         intermediate_tensors: IntermediateTensors | None = None,
         inputs_embeds: torch.Tensor | None = None,
-        *,
         capture_layer_indices: Sequence[int] | None = None,
         return_hidden_states: bool = False,
         deepstack_input_embeds: IntermediateTensors | None = None,

vllm_omni/worker/gpu_generation_model_runner.py

@@ -432,8 +432,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
         is_profile: bool = False,
         create_mixed_batch: bool = False,
         remove_lora: bool = True,
-        activate_lora: bool = False,
         is_graph_capturing: bool = False,
+        num_active_loras: int = 0,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Run a dummy forward pass to warm up/profile run or capture the
@@ -456,7 +456,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             create_mixed_batch: If True, create a mixed batch with both decode
                 (1 token) and prefill (multiple tokens) requests.
             remove_lora: If False, dummy LoRAs are not destroyed after the run
-            activate_lora: If False, dummy_run is performed without LoRAs.
+            num_active_loras: Number of active LoRAs to capture for.
         """
         mm_config = self.vllm_config.model_config.multimodal_config
         if mm_config and mm_config.mm_encoder_only:
@@ -534,7 +534,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                 # `force_has_lora` is used for cudagraph capture; because LoRA is
                 # activated later in the context manager, but we need to know the
                 # LoRA state when determining the batch descriptor for capture
-                force_has_lora=activate_lora,
+                force_has_lora=num_active_loras > 0,
+                # Capture shape specialization for specific active LoRA counts.
+                force_num_active_loras=num_active_loras,
             )
         )
 
@@ -603,8 +605,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             self.lora_config,
             num_scheduled_tokens,
             num_sampled_tokens,
-            activate_lora,
             remove_lora,
+            num_active_loras,
         ):
             # Make sure padding doesn't exceed max_num_tokens
             assert num_tokens_padded <= self.max_num_tokens
@@ -691,7 +693,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                 # lora cases when cudagraph_specialize_lora is enabled. This is a
                 # short term mitigation for issue mentioned in
                 # https://github.com/vllm-project/vllm/issues/28334
-                if self.compilation_config.cudagraph_specialize_lora and activate_lora:
+                if self.compilation_config.cudagraph_specialize_lora and num_active_loras > 0:
                     use_cudagraphs = False
 
                 self.drafter.dummy_run(

vllm_omni/worker/gpu_model_runner.py

@@ -424,8 +424,9 @@ class OmniGPUModelRunner(GPUModelRunner):
         is_profile: bool = False,
         create_mixed_batch: bool = False,
         remove_lora: bool = True,
-        activate_lora: bool = False,
         is_graph_capturing: bool = False,
+        num_active_loras: int = 0,
+        activate_lora: bool | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """
         Run a dummy forward pass to warm up/profile run or capture the
@@ -448,8 +449,12 @@ class OmniGPUModelRunner(GPUModelRunner):
             create_mixed_batch: If True, create a mixed batch with both decode
                 (1 token) and prefill (multiple tokens) requests.
             remove_lora: If False, dummy LoRAs are not destroyed after the run
-            activate_lora: If False, dummy_run is performed without LoRAs.
+            num_active_loras: Number of active LoRAs to capture for.
+            activate_lora: Backward-compatible override for LoRA activation.
         """
+        if activate_lora is None:
+            activate_lora = num_active_loras > 0
+
         mm_config = self.vllm_config.model_config.multimodal_config
         if mm_config and mm_config.mm_encoder_only:
             # The current dummy run only covers LM execution, so we can skip it.
@@ -527,6 +532,8 @@ class OmniGPUModelRunner(GPUModelRunner):
                 # activated later in the context manager, but we need to know the
                 # LoRA state when determining the batch descriptor for capture
                 force_has_lora=activate_lora,
+                # Capture shape specialization for specific active LoRA counts.
+                force_num_active_loras=num_active_loras,
             )
         )