[Pref] Support Qwen3 Omni code2wav batch infernce with async chunk (#1246)

Signed-off-by: ZeldaHuang <hzm414167@alibaba-inc.com>
Signed-off-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Signed-off-by: Ziming Huang <1520787127@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -15,6 +15,7 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (
     Qwen3OmniMoeThinkerConfig,
 )
 from vllm.config import VllmConfig
+from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding
 from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP
@@ -355,9 +356,18 @@ class Qwen3OmniMoeForConditionalGeneration(
         # ========== Stage 3: Code2Wav ==========
         elif self.model_stage == "code2wav":
             # Extract codec codes from input
-            codes = []
             if input_ids.shape[0] % 16 == 0:
-                codes.append(input_ids.reshape(1, 16, -1))
+                ubatch_slices = get_forward_context().ubatch_slices
+                if ubatch_slices is not None:
+                    max_seq_len = max(ubatch_slices) // 16
+                    batch_size = len(ubatch_slices)
+                    split_codes = torch.split(input_ids, ubatch_slices, dim=0)
+                    codes = torch.zeros((batch_size, 16, max_seq_len), device=input_ids.device, dtype=input_ids.dtype)
+                    for idx, code in enumerate(split_codes):
+                        seq_len = code.shape[0] // 16
+                        codes[idx, :, :seq_len] = code.reshape(16, seq_len)
+                else:
+                    codes = input_ids.reshape(1, 16, -1)
             else:
                 logger.warning(
                     (
@@ -373,17 +383,10 @@ class Qwen3OmniMoeForConditionalGeneration(
                         torch.zeros(16 - input_ids.shape[0] % 16, dtype=torch.long, device=input_ids.device),
                     ]
                 )
-                codes.append(input_ids_flatten.reshape(1, 16, -1))
+                codes = input_ids_flatten.reshape(1, 16, -1)
 
             # Generate audio from codec codes
-            audio_tensors = []
-            for code in codes:
-                audio_tensor = self.generate_audio(code, voice_type)
-                audio_tensors.append(audio_tensor)
-            if len(audio_tensors) > 1:
-                logger.warning(
-                    "Batched input for code2wav is not supported yet, only the first audio tensor will be returned"
-                )
+            audio_tensors = self.generate_audio(codes, voice_type)
 
             return audio_tensors
 
@@ -448,14 +451,14 @@ class Qwen3OmniMoeForConditionalGeneration(
             audio_tensors = model_outputs
             return OmniOutput(
                 text_hidden_states=None,
-                multimodal_outputs={"model_outputs": audio_tensors[0].reshape(1, -1)},
+                multimodal_outputs={"model_outputs": [audio_tensor.reshape(1, -1) for audio_tensor in audio_tensors]},
             )
 
         return model_outputs
 
     # ==================== Audio Generation ====================
 
-    def generate_audio(self, code: torch.Tensor, voice_type: str) -> torch.Tensor:
+    def generate_audio(self, code: torch.Tensor, voice_type: str) -> list[torch.Tensor]:
         """
         Generate audio waveform from codec codes.
 
@@ -484,20 +487,20 @@ class Qwen3OmniMoeForConditionalGeneration(
             talker_codes = talker_codes.expand(1, 16, -1)
 
         if self.vllm_config.model_config.async_chunk:
-            audio_tensor = self.code2wav.chunked_decode_streaming(
+            audio_tensors = self.code2wav.chunked_decode_streaming(
                 talker_codes,
                 chunk_size=25,
                 left_context_size=25,
             )
         else:
             # Use chunked decode for memory efficiency
-            audio_tensor = self.code2wav.chunked_decode(
+            audio_tensors = self.code2wav.chunked_decode(
                 talker_codes,
                 chunk_size=300,
                 left_context_size=25,
             )
 
-        return audio_tensor
+        return audio_tensors
 
     # ==================== Thinker-Talker Projection ====================
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py

@@ -22,6 +22,7 @@ from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
     SnakeBeta,
 )
 from vllm.config import VllmConfig  # type: ignore
+from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger  # type: ignore
 from vllm.model_executor.models.utils import (  # type: ignore
     AutoWeightsLoader,
@@ -162,7 +163,7 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
         codes: torch.Tensor,
         chunk_size: int = 300,
         left_context_size: int = 25,
-    ) -> torch.Tensor:
+    ) -> list[torch.Tensor]:
         """
         Decode long sequences in chunks to avoid OOM.
 
@@ -174,7 +175,9 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
             left_context_size: Number of overlapping frames for context
 
         Returns:
-            waveform: [batch, 1, waveform_len] - Complete waveform
+            list[torch.Tensor]: Complete waveform decoded from the input
+                codes. For ``batch_size == 1``, this is a list containing a
+                single tensor with shape ``[1, waveform_len]``.
         """
         wavs = []
         start_index = 0
@@ -194,14 +197,26 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
 
             start_index = end_index
 
-        return torch.cat(wavs, dim=-1)
+        ubatch_slices = get_forward_context().ubatch_slices
+        if ubatch_slices is not None:
+            code_seq_lens = [seq_len // self.config.num_quantizers for seq_len in ubatch_slices]
+        else:
+            # Fallback: assume all batch elements share the same sequence length.
+            # Create one entry per batch so that each element is processed.
+            code_seq_lens = [codes.shape[-1]] * codes.shape[0]
+        batch_wav = torch.cat(wavs, dim=-1)
+        wavs = []
+        for idx, code_seq_len in enumerate(code_seq_lens):
+            wav_chunk = batch_wav[idx, :, : code_seq_len * self.total_upsample]
+            wavs.append(wav_chunk)
+        return wavs
 
     def chunked_decode_streaming(
         self,
         codes: torch.Tensor,
         chunk_size: int = 25,
         left_context_size: int = 25,
-    ) -> torch.Tensor:
+    ) -> list[torch.Tensor]:
         """
         Decode long sequences in chunks to avoid OOM.
 
@@ -213,21 +228,31 @@ class Qwen3OmniMoeCode2Wav(nn.Module):
             left_context_size: Number of overlapping frames for context
 
         Returns:
-            waveform: [batch, 1, waveform_len] - Complete waveform
+            list[torch.Tensor]: Complete waveform decoded from the input
+                codes. For ``batch_size == 1``, this is a list containing a
+                single tensor with shape ``[1, waveform_len]``.
         """
-        wavs = []
-        end_index = codes.shape[-1]
-        # TODO: need to optimize algorithms, current only support
-        # chunk_size = left_context_size = 25
-        if end_index <= chunk_size:
-            context_size = 0
-        else:
-            context_size = left_context_size
         # Decode chunk
-        wav_chunk = self(codes)
-        # Remove context from output (context_size * total_upsample samples)
-        wavs.append(wav_chunk[..., context_size * self.total_upsample :])
-        return torch.cat(wavs, dim=-1)
+        wavs = []
+        batch_wav = self(codes)
+        ubatch_slices = get_forward_context().ubatch_slices
+        if ubatch_slices is not None:
+            code_seq_lens = [seq_len // self.config.num_quantizers for seq_len in ubatch_slices]
+        else:
+            # Fallback: assume all batch elements share the same sequence length.
+            # Create one entry per batch so that each element is processed.
+            code_seq_lens = [codes.shape[-1]] * codes.shape[0]
+        for idx, code_seq_len in enumerate(code_seq_lens):
+            # TODO: need to optimize algorithms, current only support
+            # chunk_size = left_context_size = 25
+            if code_seq_len <= chunk_size:
+                context_size = 0
+            else:
+                context_size = left_context_size
+            # Remove context from output (context_size * total_upsample samples)
+            wav_chunk = batch_wav[idx, :, context_size * self.total_upsample : code_seq_len * self.total_upsample]
+            wavs.append(wav_chunk)
+        return wavs
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights from HuggingFace checkpoint."""

vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml

@@ -72,7 +72,7 @@ stage_args:
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
       devices: "1"
-      max_batch_size: 1
+      max_batch_size: 64
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration

vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml

@@ -73,7 +73,7 @@ stage_args:
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
       devices: "1"
-      max_batch_size: 1
+      max_batch_size: 64
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration

vllm_omni/model_executor/stage_configs/qwen3_omni_moe_multiconnector.yaml

@@ -78,7 +78,7 @@ stage_args:
     stage_type: llm  # Use llm stage type to launch OmniLLM
     runtime:
       devices: "1"
-      max_batch_size: 1
+      max_batch_size: 64
     engine_args:
       model_stage: code2wav
       model_arch: Qwen3OmniMoeForConditionalGeneration

vllm_omni/worker/gpu_generation_model_runner.py

@@ -50,12 +50,26 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
     """
 
     def _update_request_states(self, scheduler_output: SchedulerOutput):
+        # remove requests
+        for req_id in scheduler_output.finished_req_ids:
+            self.input_batch.remove_request(req_id)
+        scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
+        cached_req_ids = self.input_batch.req_id_to_index.keys()
+        resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids
+        unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids)
+        for req_id in unscheduled_req_ids:
+            self.input_batch.remove_request(req_id)
         cached_reqs = scheduler_output.scheduled_cached_reqs
-        for _, req_id in enumerate(cached_reqs.req_ids):
+        req_states = []
+        for req_id in cached_reqs.req_ids:
             req_state = self.requests.get(req_id)
             assert req_state is not None
             req_state.prompt_token_ids = cached_reqs.prompt_token_ids.get(req_id)
-            self.input_batch.remove_request(req_id)
+            req_states.append(req_state)
+            # Remove the request from the current input batch only if it is still present.
+            if req_id in self.input_batch.req_id_to_index:
+                self.input_batch.remove_request(req_id)
+        for req_state in req_states:
             # update the request state in self.input_batch
             self.input_batch.add_request(req_state)
             # Only relevant for models using M-RoPE (e.g, Qwen2-VL)
@@ -244,6 +258,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
             # Mark KV scales as calculated after the first forward pass
             self.calculate_kv_scales = False
 
+        if ubatch_slices_padded is None:
+            # reuse ubatch_slices_padded for code2wav batching
+            ubatch_slices_padded = tokens
+
         # Run the model.
         # Use persistent buffers for CUDA graphs.
         with (
@@ -342,11 +360,22 @@ class GPUGenerationModelRunner(OmniGPUModelRunner):
                     {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None}
                 )
         elif isinstance(multimodal_outputs, dict):
-            mm_payload = {}
-            for key, out in multimodal_outputs.items():
-                if out is not None and isinstance(out, torch.Tensor):
-                    mm_payload[key] = out.detach().to("cpu").contiguous()
-            pooler_output.append(mm_payload)
+            num_reqs = self.input_batch.num_reqs
+            for i in range(num_reqs):
+                mm_payload = {}
+                for key, out in multimodal_outputs.items():
+                    if isinstance(out, list):
+                        if len(out) != num_reqs:
+                            raise ValueError(
+                                f"Multimodal output list for key '{key}' has length {len(out)} "
+                                f"but expected {num_reqs} (one entry per request)."
+                            )
+                        mm_payload[key] = out[i].detach().to("cpu").contiguous()
+                    elif isinstance(out, torch.Tensor):
+                        mm_payload[key] = out.detach().to("cpu").contiguous()
+                    else:
+                        logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}")
+                pooler_output.append(mm_payload)
         else:
             raise RuntimeError("Unsupported diffusion output type")
         output = OmniModelRunnerOutput(