Support Qwen3 Omni thinker cuda graph (#523)

Signed-off-by: ZeldaHuang <hzm414167@alibaba-inc.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

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"

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"

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

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:

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: