From 1ae0c64fe01993cfd07e5fde0c21ccc03a4ac9c1 Mon Sep 17 00:00:00 2001 From: Ziming Huang Date: Tue, 30 Dec 2025 18:46:13 +0800 Subject: [PATCH] Support Qwen3 Omni thinker cuda graph (#523) Signed-off-by: ZeldaHuang Co-authored-by: Hongsheng Liu --- .../stage_configs/qwen3_omni_ci.yaml | 2 +- .../stage_configs/qwen3_omni_ci.yaml | 2 +- .../models/qwen3_omni/qwen3_omni.py | 53 ++++++++++--------- vllm_omni/worker/gpu_ar_model_runner.py | 3 ++ vllm_omni/worker/gpu_model_runner.py | 8 ++- 5 files changed, 38 insertions(+), 30 deletions(-) diff --git a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml index fec6169d8..5e3be865e 100644 --- a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml @@ -15,7 +15,7 @@ stage_args: worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 - enforce_eager: true + enforce_eager: false trust_remote_code: true engine_output_type: latent # Output hidden states for talker distributed_executor_backend: "mp" diff --git a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml index ff889043f..7420e2745 100644 --- a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml @@ -15,7 +15,7 @@ stage_args: worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 - enforce_eager: true + enforce_eager: false trust_remote_code: true engine_output_type: latent # Output hidden states for talker distributed_executor_backend: "mp" diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 7695adec7..504ca4486 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -188,6 +188,12 @@ class Qwen3OmniMoeForConditionalGeneration( self.make_empty_intermediate_tensors = ( self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None ) + if self.model_stage == "thinker": + self.tts_tokens = torch.tensor( + [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]], + device=self._module_device(self.thinker), + dtype=torch.long, + ) # ==================== Device utilities ==================== @@ -302,28 +308,28 @@ class Qwen3OmniMoeForConditionalGeneration( # ========== Stage 1: Thinker ========== if self.model_stage == "thinker": - # Normalize to batched inputs if needed - _added_batch_dim = False - if input_ids is not None and input_ids.ndim == 1: - input_ids = input_ids.unsqueeze(0) - _added_batch_dim = True - if positions is not None and positions.ndim == 1: - positions = positions.unsqueeze(0) - _added_batch_dim = True - if inputs_embeds is not None and inputs_embeds.ndim == 2: - inputs_embeds = inputs_embeds.unsqueeze(0) - _added_batch_dim = True - thinker_dev = self._module_device(self.thinker) + if is_npu(): + # Normalize to batched inputs if needed + _added_batch_dim = False + if input_ids is not None and input_ids.ndim == 1: + input_ids = input_ids.unsqueeze(0) + _added_batch_dim = True + if positions is not None and positions.ndim == 1: + positions = positions.unsqueeze(0) + _added_batch_dim = True + if inputs_embeds is not None and inputs_embeds.ndim == 2: + inputs_embeds = inputs_embeds.unsqueeze(0) + _added_batch_dim = True - # Handle None input_ids - if input_ids is None: - input_ids = torch.zeros( - inputs_embeds.shape[1], - dtype=torch.long, - device=thinker_dev, - ).unsqueeze(0) - _added_batch_dim = True + # Handle None input_ids + if input_ids is None: + input_ids = torch.zeros( + inputs_embeds.shape[1], + dtype=torch.long, + device=thinker_dev, + ).unsqueeze(0) + _added_batch_dim = True # Move to thinker device if input_ids is not None and input_ids.device != thinker_dev: @@ -365,12 +371,7 @@ class Qwen3OmniMoeForConditionalGeneration( # These will later be projected into talker text space by the talker stage. multimodal_outputs = captured_layer_dict if captured_layer_dict is not None else {} try: - tts_tokens = torch.tensor( - [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]], - device=self._module_device(self.thinker), - dtype=torch.long, - ) - thinker_tts_embeds = self.thinker.embed_input_ids(tts_tokens) # [1,3,thinker_hidden] + thinker_tts_embeds = self.thinker.embed_input_ids(self.tts_tokens) # [1,3,thinker_hidden] if ( isinstance(thinker_tts_embeds, torch.Tensor) and thinker_tts_embeds.ndim == 3 diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 77b32c681..859027a8f 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -29,6 +29,7 @@ from vllm.v1.worker.gpu_model_runner import ( ) from vllm.v1.worker.utils import is_residual_scattered_for_sp +from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner @@ -184,6 +185,8 @@ class GPUARModelRunner(OmniGPUModelRunner): logits_index=logits_indices, sampler=self.sampler, ) + if isinstance(model_output, tuple): + model_output = OmniOutput(*model_output) with record_function_or_nullcontext("gpu_model_runner: postprocess"): if self.use_aux_hidden_state_outputs: diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 2cb76815a..9e155f4ec 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -312,14 +312,18 @@ class OmniGPUModelRunner(GPUModelRunner): @torch.inference_mode() def extract_multimodal_outputs(self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput) -> dict: - if hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs: + if ( + hasattr(self.model, "have_multimodal_outputs") + and self.model.have_multimodal_outputs + and isinstance(hidden_states, OmniOutput) + ): text_hidden_states = hidden_states.text_hidden_states multimodal_outputs = hidden_states.multimodal_outputs elif isinstance(hidden_states, torch.Tensor): text_hidden_states = hidden_states multimodal_outputs = {} - elif isinstance(hidden_states, list): + elif isinstance(hidden_states, list) or isinstance(hidden_states, tuple): text_hidden_states = hidden_states[0] multimodal_outputs = {} else: