[Feature] Support Qwen3 Omni talker cudagraph (#669)

Signed-off-by: ZeldaHuang <hzm414167@alibaba-inc.com>

tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml

@@ -45,7 +45,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.3
-       enforce_eager: true
+       enforce_eager: false
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
       #  tensor_parallel_size: 2

tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml

@@ -45,7 +45,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.3
-       enforce_eager: true
+       enforce_eager: false
        trust_remote_code: true
        engine_output_type: latent  # Output codec codes for code2wav
       #  tensor_parallel_size: 2

vllm_omni/entrypoints/omni_stage.py

@@ -1226,6 +1226,7 @@ async def _stage_worker_async(
                     gen_output = res
                     _gen_t1 = _time.time()
                     _gen_ms = (_gen_t1 - _gen_t0) * 1000.0
+                    _gen_t0 = _gen_t1
                     await generation_out_q.put((rid, gen_output, _gen_ms))
         except Exception as e:
             logger.exception("Failed on request %s: %s", rid, e)

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -8,6 +8,13 @@ from functools import cached_property
 
 import torch
 import torch.nn as nn
+from transformers.generation.logits_process import (
+    LogitsProcessorList,
+    RepetitionPenaltyLogitsProcessor,
+    SuppressTokensLogitsProcessor,
+    TemperatureLogitsWarper,
+    TopKLogitsWarper,
+)
 from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
     Qwen3OmniMoeCode2WavConfig,
     Qwen3OmniMoeConfig,
@@ -146,13 +153,6 @@ class Qwen3OmniMoeForConditionalGeneration(
             self.talker.init_multi_modal(thinker_config)
             self.model = self.talker
             self.code2wav = None
-            from transformers.generation.logits_process import (
-                LogitsProcessorList,
-                RepetitionPenaltyLogitsProcessor,
-                SuppressTokensLogitsProcessor,
-                TemperatureLogitsWarper,
-                TopKLogitsWarper,
-            )
 
             self.logits_processors = LogitsProcessorList(
                 [
@@ -367,46 +367,17 @@ class Qwen3OmniMoeForConditionalGeneration(
                 **capture_kwargs,
                 **kwargs,
             )
-            # Compute thinker-side TTS token embeddings for BOS/EOS/PAD and expose via multimodal outputs.
-            # 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:
-                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
-                    and thinker_tts_embeds.shape[1] == 3
-                ):
-                    bos_eos_pad = thinker_tts_embeds.to(text_hidden_states.device).chunk(3, dim=1)  # 3 * [1,1,H]
-                    multimodal_outputs["tts_bos_embed"] = [bos_eos_pad[0]]
-                    multimodal_outputs["tts_eos_embed"] = [bos_eos_pad[1]]
-                    multimodal_outputs["tts_pad_embed"] = [bos_eos_pad[2]]
-            except Exception:
-                # Best-effort; absence will be handled by talker with fallbacks
-                pass
-
-            # Return text-only output (with multimodal sidecar)
-            return OmniOutput(
-                text_hidden_states=(text_hidden_states.reshape(-1, text_hidden_states.shape[-1])),
-                multimodal_outputs=multimodal_outputs,
-            )
+            return text_hidden_states, captured_layer_dict
 
         # ========== Stage 2.1: Talker ==========
         elif self.model_stage == "talker":
             if input_ids is None:
                 input_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device)
-                self.code_predictor_hidden_pool = torch.zeros_like(inputs_embeds)
-                is_profile = True
-            else:
-                is_profile = False
 
             # Ensure we have base embeddings when only ids are provided
             if inputs_embeds is None and input_ids is not None:
                 inputs_embeds = self.talker.embed_input_ids(input_ids)
 
-            if not is_profile and input_ids.shape[0] > 1:
-                self.input_ids_list = input_ids.tolist()
-
             # TODO(Peiqi): temporal hack here to support voice_type.
             if not hasattr(self, "voice_type"):
                 self.voice_type = voice_type
@@ -418,19 +389,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                     positions=positions,
                     inputs_embeds=inputs_embeds,
                 )
-
-            # merge the code_predictor_codes from the info_dict list into a single tensor
-            multimodal_outputs: dict = None
-            if not is_profile:
-                # Here is the only place to use runtime_additional_information. After MTP in the
-                # preprocess function, the code_predictor_codes are stored in the info_dict list.
-                # We need to merge the tensors from different requests into a single tensor.
-                # In the future, we may allow user to custom an aggregated function.
-                info_dicts = kwargs.get("runtime_additional_information")
-                code_predictor_codes = [info.get("code_predictor_codes") for info in info_dicts]
-                multimodal_outputs = {"code_predictor_codes": torch.cat(code_predictor_codes, dim=0)}
-
-            return OmniOutput(text_hidden_states=talker_hidden, multimodal_outputs=multimodal_outputs)
+            return talker_hidden
 
         # ========== Stage 3: Code2Wav ==========
         elif self.model_stage == "code2wav":
@@ -459,10 +418,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                     "Batched input for code2wav is not supported yet, only the first audio tensor will be returned"
                 )
 
-            return OmniOutput(
-                text_hidden_states=None,
-                multimodal_outputs={"model_outputs": audio_tensors[0].reshape(1, -1)},
-            )
+            return audio_tensors
 
         # Fallback (shouldn't reach here)
         return OmniOutput(
@@ -473,6 +429,63 @@ class Qwen3OmniMoeForConditionalGeneration(
             multimodal_outputs=None,
         )
 
+    def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs) -> OmniOutput:
+        """
+        Make an OmniOutput object from model outputs.
+        Args:
+            model_outputs: Model outputs
+        """
+        if isinstance(model_outputs, OmniOutput):
+            return model_outputs
+
+        if self.model_stage == "thinker":
+            text_hidden_states, captured_layer_dict = model_outputs
+            # Compute thinker-side TTS token embeddings for BOS/EOS/PAD and expose via multimodal outputs.
+            # 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:
+                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
+                    and thinker_tts_embeds.shape[1] == 3
+                ):
+                    bos_eos_pad = thinker_tts_embeds.to(text_hidden_states.device).chunk(3, dim=1)  # 3 * [1,1,H]
+                    multimodal_outputs["tts_bos_embed"] = [bos_eos_pad[0]]
+                    multimodal_outputs["tts_eos_embed"] = [bos_eos_pad[1]]
+                    multimodal_outputs["tts_pad_embed"] = [bos_eos_pad[2]]
+            except Exception:
+                # Best-effort; absence will be handled by talker with fallbacks
+                pass
+
+            # Return text-only output (with multimodal sidecar)
+            return OmniOutput(
+                text_hidden_states=(text_hidden_states.reshape(-1, text_hidden_states.shape[-1])),
+                multimodal_outputs=multimodal_outputs,
+            )
+        elif self.model_stage == "talker":
+            talker_hidden = model_outputs
+            # merge the code_predictor_codes from the info_dict list into a single tensor
+            multimodal_outputs: dict = None
+            # Here is the only place to use runtime_additional_information. After MTP in the
+            # preprocess function, the code_predictor_codes are stored in the info_dict list.
+            # We need to merge the tensors from different requests into a single tensor.
+            # In the future, we may allow user to custom an aggregated function.
+            info_dicts = kwargs.get("runtime_additional_information")
+            code_predictor_codes = [info.get("code_predictor_codes") for info in info_dicts]
+            multimodal_outputs = {"code_predictor_codes": torch.cat(code_predictor_codes, dim=0)}
+            span_len = multimodal_outputs["code_predictor_codes"].shape[0]
+            talker_hidden = talker_hidden[:span_len]
+            return OmniOutput(text_hidden_states=talker_hidden, multimodal_outputs=multimodal_outputs)
+        elif self.model_stage == "code2wav":
+            audio_tensors = model_outputs
+            return OmniOutput(
+                text_hidden_states=None,
+                multimodal_outputs={"model_outputs": audio_tensors[0].reshape(1, -1)},
+            )
+
+        return model_outputs
+
     # ==================== Audio Generation ====================
 
     def generate_audio(self, code: torch.Tensor, voice_type: str) -> torch.Tensor:
@@ -624,9 +637,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 input_ids, input_embeds, **info_dict
             )
 
-        # execute talker MTP
-        input_embeds, code_predictor_codes = self.talker_mtp(input_ids, input_embeds, last_talker_hidden, text_step)
-        update_dict["code_predictor_codes"] = code_predictor_codes
+        update_dict["mtp_inputs"] = last_talker_hidden, text_step
 
         return input_ids, input_embeds, update_dict
 
@@ -657,7 +668,7 @@ class Qwen3OmniMoeForConditionalGeneration(
                 dtype=torch.long,
             )
         inputs_embeds = (inputs_embeds + text_step).reshape(-1, self.talker_config.text_config.hidden_size)
-        return inputs_embeds, code_predictor_codes.squeeze(-1).detach().to("cpu").contiguous()
+        return inputs_embeds, code_predictor_codes.squeeze(-1)
 
     def talker_preprocess_prefill(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
         # Containers to return per-request updates (e.g., code_predictor_hidden_per_request)
@@ -872,7 +883,6 @@ class Qwen3OmniMoeForConditionalGeneration(
 
     def talker_preprocess_decode(self, input_ids: torch.Tensor, input_embeds: torch.Tensor, **info_dict: dict):
         update_dict: dict[str, dict] = {}
-
         try:
             q_tail = info_dict.get("tailing_text_hidden")
             if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0:

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py

@@ -14,6 +14,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from transformers import Cache, PretrainedConfig
 from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeRotaryEmbedding
+from vllm.compilation.decorators import support_torch_compile
 from vllm.config import CacheConfig, ModelConfig, VllmConfig
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -463,6 +464,7 @@ class Qwen3OmniCodePredictorBaseModel(nn.Module):
         return self.codec_embedding
 
 
+@support_torch_compile
 class Qwen3OmniMoeTalkerCodePredictor(nn.Module):
     """
     Code predictor wrapper matching HF structure.

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py

@@ -4,6 +4,11 @@ from typing import Any
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from transformers.generation.logits_process import (
+    LogitsProcessorList,
+    TopKLogitsWarper,
+    TopPLogitsWarper,
+)
 from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
     Qwen3OmniMoeTalkerConfig,
 )
@@ -134,6 +139,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
         self.code_predictor = Qwen3OmniMoeTalkerCodePredictor(
             vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor")
         )
+        self.empty_code = torch.empty((1, 0), dtype=torch.long)
 
     def code_predictor_forward(
         self,
@@ -185,6 +191,12 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
         all_codes_per_position = []
         middle_hidden_states = []  # Collect hidden states for each position
 
+        logits_processors = LogitsProcessorList(
+            [
+                TopKLogitsWarper(top_k=top_k),
+                TopPLogitsWarper(top_p=top_p),
+            ]
+        )
         # Generate residual layers for each position
         for pos in range(seq_len):
             layer0_code = input_ids[:, pos : pos + 1]  # [batch, 1]
@@ -220,19 +232,13 @@ class Qwen3OmniMoeTalkerForConditionalGeneration(
 
                 # Use the corresponding lm_head for this layer
                 logits = self.code_predictor.lm_head[layer_idx](hidden_state[:, -1:, :])  # [batch, 1, vocab_size]
-                from transformers.generation.logits_process import (
-                    LogitsProcessorList,
-                    TopKLogitsWarper,
-                    TopPLogitsWarper,
-                )
 
-                logits_processors = LogitsProcessorList(
-                    [
-                        TopKLogitsWarper(top_k=top_k),
-                        TopPLogitsWarper(top_p=top_p),
-                    ]
-                )
-                input_ids_for_logits_processors = torch.tensor([pos_codes[1:]]).to(logits.device, dtype=torch.long)
+                if len(pos_codes) > 1:
+                    input_ids_for_logits_processors = torch.cat(pos_codes[1:], dim=1).to(
+                        device=logits.device, dtype=torch.long
+                    )
+                else:
+                    input_ids_for_logits_processors = self.empty_code
                 logits = logits_processors(input_ids_for_logits_processors, logits.squeeze(0)).unsqueeze(0)
 
                 # Sample from the filtered distribution

vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml

@@ -16,7 +16,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"
@@ -47,7 +47,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.3
-      enforce_eager: true
+      enforce_eager: false
       trust_remote_code: true
       engine_output_type: latent  # Output codec codes for code2wav
       # tensor_parallel_size: 2

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -160,11 +160,11 @@ def talker2code2wav(
     # Process each talker output
     for i, talker_output in enumerate(talker_outputs):
         output = talker_output.outputs[0]
-
+        seq_len = len(output.token_ids)
         # Extract codec codes from talker output
         # Expected shape: [8, seq_len] (8-layer RVQ codes)
         codec_codes = (
-            output.multimodal_output["code_predictor_codes"]
+            output.multimodal_output["code_predictor_codes"][-seq_len:]
             .to(torch.long)
             .transpose(0, 1)
             .cpu()

vllm_omni/worker/gpu_ar_model_runner.py

@@ -29,7 +29,6 @@ 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
 
@@ -185,8 +184,6 @@ 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

@@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
 import torch
+from vllm.compilation.cuda_graph import CUDAGraphWrapper
 from vllm.config import CUDAGraphMode
 from vllm.distributed.parallel_state import get_pp_group
 from vllm.forward_context import set_forward_context
@@ -38,6 +39,25 @@ class OmniGPUModelRunner(GPUModelRunner):
         self._omni_num_scheduled_tokens_np: np.ndarray | None = None
         self._omni_last_model_output: object | None = None
 
+    def load_model(self, *args, **kwargs) -> None:
+        super().load_model(*args, **kwargs)
+        # TODO move this model specific logic to a separate class
+        if hasattr(self.model, "talker_mtp") and self.model.talker is not None:
+            self.talker_mtp = self.model.talker_mtp
+            cudagraph_mode = self.compilation_config.cudagraph_mode
+            assert cudagraph_mode is not None
+            if cudagraph_mode.has_full_cudagraphs():
+                self.talker_mtp = CUDAGraphWrapper(
+                    self.model.talker_mtp, self.vllm_config, runtime_mode=CUDAGraphMode.FULL
+                )
+            hidden_size = self.model_config.hf_config.talker_config.text_config.hidden_size
+            self.talker_mtp_input_ids = self._make_buffer(self.max_num_reqs, dtype=torch.int32)
+            self.talker_mtp_inputs_embeds = self._make_buffer(
+                self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False
+            )
+            self.last_talker_hidden = self._make_buffer(self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False)
+            self.text_step = self._make_buffer(self.max_num_reqs, hidden_size, dtype=self.dtype, numpy=False)
+
     def _init_mrope_positions(self, req_state: CachedRequestState):
         """Initialize M-RoPE positions for multimodal inputs.
 
@@ -543,6 +563,13 @@ class OmniGPUModelRunner(GPUModelRunner):
                     ubatch_slices=ubatch_slices,
                 ),
             ):
+                if self.model.talker is not None and hasattr(self.model, "talker_mtp") and num_tokens_padded == 1:
+                    outputs = self.talker_mtp(
+                        self.talker_mtp_input_ids.gpu[:num_tokens_padded],
+                        self.talker_mtp_inputs_embeds.gpu[:num_tokens_padded],
+                        self.last_talker_hidden.gpu[:num_tokens_padded],
+                        self.text_step.gpu[:num_tokens_padded],
+                    )
                 outputs = self.model(
                     input_ids=input_ids,
                     positions=positions,
@@ -866,6 +893,34 @@ class OmniGPUModelRunner(GPUModelRunner):
                 req_input_ids, req_embeds, update_dict = self.model.preprocess(
                     input_ids=input_ids[s:e], input_embeds=inputs_embeds[s:e], **req_infos
                 )
+                # run talker mtp decode
+                if hasattr(self.model, "talker_mtp"):
+                    _cudagraph_mode, batch_desc, _, _ = self._determine_batch_execution_and_padding(
+                        num_tokens=span_len,
+                        num_reqs=1,
+                        num_scheduled_tokens_np=num_scheduled_tokens_np[req_index],
+                        max_num_scheduled_tokens=1,
+                        force_eager=span_len > 1,
+                        use_cascade_attn=False,
+                    )
+                    last_talker_hidden, text_step = update_dict.pop("mtp_inputs")
+                    if _cudagraph_mode != CUDAGraphMode.NONE:
+                        self.talker_mtp_input_ids.gpu[:span_len].copy_(req_input_ids)
+                        self.talker_mtp_inputs_embeds.gpu[:span_len].copy_(req_embeds)
+                        self.last_talker_hidden.gpu[:span_len].copy_(last_talker_hidden)
+                        self.text_step.gpu[:span_len].copy_(text_step)
+                        req_input_ids = self.talker_mtp_input_ids.gpu[:span_len]
+                        req_embeds = self.talker_mtp_inputs_embeds.gpu[:span_len]
+                        last_talker_hidden = self.last_talker_hidden.gpu[:span_len]
+                        text_step = self.text_step.gpu[:span_len]
+                    with set_forward_context(
+                        None, self.vllm_config, cudagraph_runtime_mode=_cudagraph_mode, batch_descriptor=batch_desc
+                    ):
+                        req_embeds, code_predictor_codes = self.talker_mtp(
+                            req_input_ids, req_embeds, last_talker_hidden, text_step
+                        )
+                        update_dict["code_predictor_codes"] = code_predictor_codes
+
                 # TODO(Peiqi): the merge stage could move out from the critical path
                 self._merge_additional_information_update(req_id, update_dict)
 
@@ -909,6 +964,8 @@ class OmniGPUModelRunner(GPUModelRunner):
             **model_kwargs,
             **model_kwargs_extra,
         )
+        if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"):
+            model_output = self.model.make_omni_output(model_output, **model_kwargs_extra)
         # Cache model output so later sample_tokens can consume multimodal results.
         self._omni_last_model_output = model_output
         return model_output