[Bugfix][Qwen3-Omni] Fix Code2Wav CUDA-graph output length surplus (#4466)

tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py

@@ -175,6 +175,77 @@ def test_padded_output_bounded(decoder, wrapper, seq_len):
     assert max_diff < 1.0, f"Max diff {max_diff} exceeds bound"
 
 
+@pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99])
+def test_padded_short_output_decode_length_matches_eager(seq_len):
+    """Streaming ``decode()`` of a shorter-than-nominal decoder must match eager length.
+
+    Regression for #4466: Qwen3-Omni Code2Wav returns fewer samples than
+    ``seq_len * total_upsample``. The graph ``decode()`` path used to trim to the
+    nominal length, so for padded (non-capture-size) chunks it leaked a fixed
+    surplus of stale buffer-tail samples. The streaming concat turned that
+    per-chunk surplus into a ~46 ms "lag in the middle". The trim must instead be
+    relative to the captured output length, matching the eager forward exactly.
+    """
+    short_decoder = ShortOutputDecoder().to(DEVICE).eval()
+    short_wrapper = CUDAGraphDecoderWrapper(
+        decoder=short_decoder,
+        capture_sizes=[25, 50, 100],
+        capture_batch_sizes=[1],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    short_wrapper.warmup(DEVICE)
+    codes = _random_codes(seq_len)
+    with torch.no_grad():
+        eager_out = short_decoder(codes)
+        graph_out = short_wrapper.decode(codes)
+    surplus = graph_out.shape[-1] - eager_out.shape[-1]
+    assert surplus == 0, (
+        f"graph length {graph_out.shape[-1]} != eager length {eager_out.shape[-1]} "
+        f"(surplus {surplus}); the padded graph output is leaking stale tail samples"
+    )
+    # Interior content (away from the conv receptive-field boundary) must still match.
+    boundary = 3 * TOTAL_UPSAMPLE
+    if eager_out.shape[-1] > boundary:
+        torch.testing.assert_close(graph_out[..., :-boundary], eager_out[..., :-boundary], atol=1e-5, rtol=1e-5)
+
+
+def test_compiled_padded_short_output_length_matches_eager(monkeypatch):
+    """Compiled replay path must also trim short-output decoders to eager length.
+
+    Same regression as the graph path (#4466) but exercising the torch.compile +
+    CUDA Graph branch in ``_decode``. ``torch.compile`` is faked (as in the other
+    compiled-path tests) so the test stays fast and deterministic.
+    """
+
+    def _fake_compile(model, **_kwargs):
+        def _compiled(codes):
+            return model(codes)
+
+        return _compiled
+
+    monkeypatch.setattr(torch, "compile", _fake_compile)
+
+    short_decoder = ShortOutputDecoder().to(DEVICE).eval()
+    compiled_wrapper = CUDAGraphDecoderWrapper(
+        decoder=short_decoder,
+        capture_sizes=[25, 50],
+        capture_batch_sizes=[1],
+        compile_shapes=[(1, 25), (1, 50)],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    compiled_wrapper.warmup(DEVICE)
+    codes = _random_codes(30)  # not a capture size -> padded to bucket 50 on the compiled path
+    with torch.no_grad():
+        eager_out = short_decoder(codes)
+        graph_out = compiled_wrapper.decode(codes)
+    surplus = graph_out.shape[-1] - eager_out.shape[-1]
+    assert surplus == 0, (
+        f"compiled graph length {graph_out.shape[-1]} != eager length {eager_out.shape[-1]} (surplus {surplus})"
+    )
+
+
 # ──────────────────────────────────────────────────────────────────
 # 3. Fallback to eager (size exceeds all capture sizes) → bit-identical
 # ──────────────────────────────────────────────────────────────────

vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py

@@ -44,7 +44,13 @@ def _batched_chunked_decode(
     left_context_size: int = 25,
     max_batch_size: int = 0,
 ) -> torch.Tensor:
-    """Decode a padded batch by grouping same-round chunks across requests."""
+    """Decode a padded batch by grouping same-round chunks across requests.
+
+    Assumes ``decode_fn`` returns exactly ``chunk_len * total_upsample`` samples per
+    chunk (causal-conv trim constant C = 0, e.g. Qwen3-TTS) and raises otherwise. It
+    is not used by short-output decoders such as Qwen3-Omni Code2Wav (C = 555), which
+    go through the per-request ``chunked_decode`` / ``chunked_decode_streaming`` paths.
+    """
     if codes.dim() < 3:
         raise ValueError(f"Expected codes with shape [B, Q, F], got {tuple(codes.shape)}")
     if chunk_size <= 0:
@@ -511,8 +517,7 @@ def _decode(self, codes: torch.Tensor, *, clone_graph_output: bool) -> torch.Ten
                 static_input.zero_()
                 static_input[:, :, :actual_size] = codes
             self._compiled_graphs[compile_key].replay()
-            actual_out_len = actual_size * self.decoder.total_upsample
-            output = self._compiled_static_outputs[compile_key][..., :actual_out_len]
+            output = self._trim_replay_output(self._compiled_static_outputs[compile_key], actual_size, compiled_size)
             if clone_graph_output:
                 return output.clone()
             return output
@@ -542,12 +547,26 @@ def _decode(self, codes: torch.Tensor, *, clone_graph_output: bool) -> torch.Ten
             static_input[:, :, :actual_size] = codes
         self.graphs[graph_key].replay()
 
-        actual_out_len = actual_size * self.decoder.total_upsample
-        output = self.static_outputs[graph_key][..., :actual_out_len]
+        output = self._trim_replay_output(self.static_outputs[graph_key], actual_size, padded_size)
         if clone_graph_output:
             return output.clone()
         return output
 
+    def _trim_replay_output(self, static_output: torch.Tensor, actual_size: int, padded_size: int) -> torch.Tensor:
+        """Trim a graph/compiled replay output to the eager-equivalent length.
+
+        The captured ``static_output`` already reflects the decoder's TRUE output
+        length for ``padded_size`` frames, which for causal decoders (e.g.
+        Qwen3-Omni Code2Wav) is shorter than ``padded_size * total_upsample`` by a
+        fixed amount. Each zero-padded frame beyond ``actual_size`` contributes
+        ``total_upsample`` trailing samples that are stale buffer content, so trim
+        relative to the captured length instead of the nominal length. This is a
+        no-op for decoders whose output equals the nominal length (e.g. Qwen3-TTS).
+        """
+        drop = (padded_size - actual_size) * self.decoder.total_upsample
+        actual_out_len = max(0, static_output.shape[-1] - drop)
+        return static_output[..., :actual_out_len]
+
     def decode(self, codes: torch.Tensor) -> torch.Tensor:
         return self._decode(codes, clone_graph_output=True)