From ab59673f21d804729557ac53d4c839e6d7353afb Mon Sep 17 00:00:00 2001
From: Sy03 <1370724210@qq.com>
Date: Tue, 19 May 2026 02:59:23 +0800
Subject: [PATCH] [Bugfix][Qwen3-Omni] Handle short Code2Wav chunk outputs
 (#3687)

Signed-off-by: Sy03 <1370724210@qq.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
---
 .../qwen3_tts/test_cuda_graph_decoder.py      | 37 +++++++++++++++++++
 .../qwen3_tts/cuda_graph_decoder_wrapper.py   | 21 ++++-------
 2 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py
index dbd217246..e1ee68991 100644
--- a/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py
+++ b/tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py
@@ -74,6 +74,14 @@ class SyntheticDecoder(nn.Module):
         return self.out(x).clamp(min=-1, max=1)
 
 
+class ShortOutputDecoder(SyntheticDecoder):
+    """Decoder variant that returns fewer samples than seq_len * total_upsample."""
+
+    def forward(self, codes):
+        out = super().forward(codes)
+        return out[..., :-5] if out.shape[-1] > 5 else out
+
+
 @pytest.fixture(scope="module")
 def decoder():
     """Create a synthetic decoder on CUDA with fixed weights."""
@@ -228,6 +236,35 @@ def test_chunked_decode_output_survives_later_replay(wrapper):
     torch.testing.assert_close(graph_out, expected, atol=0, rtol=0)
 
 
+def test_chunked_decode_preserves_short_chunk_concat_semantics():
+    """Chunked decode must allow shorter-than-nominal chunk outputs.
+
+    Qwen3-Omni non-async startup can hit a decoder path where a chunk returns
+    fewer waveform samples than ``code_len * total_upsample``. The wrapper must
+    preserve the original eager concat behavior instead of copying into exact
+    nominal slices.
+    """
+    torch.manual_seed(123)
+    short_decoder = ShortOutputDecoder().to(DEVICE).eval()
+    short_wrapper = CUDAGraphDecoderWrapper(
+        decoder=short_decoder,
+        capture_sizes=[50],
+        capture_batch_sizes=[1],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    short_wrapper.warmup(DEVICE)
+    codes = _random_codes(100)
+
+    with torch.no_grad():
+        eager_out = _eager_chunked(short_decoder, codes, chunk_size=50, left_context_size=0)
+        graph_out = short_wrapper.chunked_decode_with_cudagraph(codes, chunk_size=50, left_context_size=0)
+
+    assert graph_out.shape == eager_out.shape
+    assert graph_out.shape[-1] < 100 * TOTAL_UPSAMPLE
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
 def test_batched_chunked_decode_variable_lengths_matches_per_request_eager(decoder, wrapper):
     """Variable-length chunk batching should match independent chunked decodes."""
     long_codes = _random_codes(100)
diff --git a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
index a9cdc8948..1c25e0842 100644
--- a/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
+++ b/vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py
@@ -554,10 +554,10 @@ class CUDAGraphDecoderWrapper:
         chunk_size: int = 300,
         left_context_size: int = 25,
     ) -> torch.Tensor:
+        wavs = []
         start_index = 0
         total_len = codes.shape[-1]
         total_upsample = self.decoder.total_upsample
-        wav_out = None
 
         while start_index < total_len:
             end_index = min(start_index + chunk_size, total_len)
@@ -566,21 +566,16 @@ class CUDAGraphDecoderWrapper:
             codes_chunk = codes[..., start_index - context_size : end_index]
             wav_chunk = self._decode(codes_chunk, clone_graph_output=False)
 
-            if wav_out is None:
-                wav_out = torch.empty(
-                    (*wav_chunk.shape[:-1], total_len * total_upsample),
-                    dtype=wav_chunk.dtype,
-                    device=wav_chunk.device,
-                )
-            src_start = context_size * total_upsample
-            dst_start = start_index * total_upsample
-            dst_end = end_index * total_upsample
-            wav_out[..., dst_start:dst_end].copy_(wav_chunk[..., src_start:])
+            # Keep origin/main's concat semantics: Qwen3-Omni can return a chunk
+            # that is shorter than the nominal code_len * total_upsample length.
+            # Clone each slice because graph outputs are static buffers that later
+            # replays may overwrite.
+            wavs.append(wav_chunk[..., context_size * total_upsample :].clone())
             start_index = end_index
 
-        if wav_out is None:
+        if not wavs:
             return self.decoder(codes)
-        return wav_out
+        return torch.cat(wavs, dim=-1)
 
     def batched_chunked_decode_with_cudagraph(
         self,