From 0099eebb3f8e2e5cd08b8d988ab2719f34cba4ca Mon Sep 17 00:00:00 2001 From: tzhouam Date: Thu, 12 Feb 2026 03:33:50 +0000 Subject: [PATCH] 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 --- vllm_omni/engine/arg_utils.py | 91 +++++++++++++------ .../qwen3_omni/qwen3_omni_moe_thinker.py | 1 - .../worker/gpu_generation_model_runner.py | 12 ++- vllm_omni/worker/gpu_model_runner.py | 11 ++- 4 files changed, 77 insertions(+), 38 deletions(-) diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index f83998baf..4e1cb7c8c 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/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, diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index ee1944a77..3dda39423 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/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, diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 80999e9b7..b8b770dd5 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/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( diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 5735ef1e9..6d4423fc3 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/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, ) )