From cead615288028ce0bdf5f61c4bdb1dd65f762fb4 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Fri, 19 Jun 2026 23:58:45 -0700 Subject: [PATCH] [Bugfix][Qwen3-Omni] Fix Code2Wav CUDA-graph output length surplus (#4466) (#4525) Signed-off-by: Yueqian Lin Co-authored-by: Sy03 <1370724210@qq.com> Co-authored-by: Hongsheng Liu --- .../qwen3_tts/test_cuda_graph_decoder.py | 71 +++++++++++++++++++ .../qwen3_tts/cuda_graph_decoder_wrapper.py | 29 ++++++-- 2 files changed, 95 insertions(+), 5 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 9a7460bf1..0f77455d7 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 @@ -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 # ────────────────────────────────────────────────────────────────── 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 604946a35..abeaf243d 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 @@ -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 @@ class CUDAGraphDecoderWrapper: 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 @@ class CUDAGraphDecoderWrapper: 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)