[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."""