From 3f238fe6ccb4d747be97b351fa8a6f3883241489 Mon Sep 17 00:00:00 2001 From: Deep Shah Date: Wed, 9 Sep 2026 01:27:34 -0700 Subject: [PATCH] [Bugfix] Fix Qwen3-Omni streaming crops in mixed-length batches A shorter Code2Wav row inferred its right-edge trim from the padded batch output length. That made the trim zero, shifted the overlap boundary and retained padded samples. Measure the trim from the decoded batch window and use each request's valid end. Validation run c2w-0909a on two A100-80GB GPUs reproduced a 7680-sample first chunk instead of 7125. The candidate passed 48 fixed-shape pretrained prefix checks, preserved all eight singleton PCM outputs, and completed the 16-request serving check with correct chunk counts. The regression fails 22 cases on the old crop and passes all 48 with the correction. Related model and graph tests pass 126/126. GPU evidence was collected before the content-identical update to 6fb7b0a0. Fresh extracted-method red/green and changed-file hooks pass after the update. No general speech-quality or cross-shape numerical parity claim is made. AI assistance: Codex developed the fix, tests, and validation material. Signed-off-by: Deep Shah --- .../test_qwen3_omni_code2wav_streaming.py | 66 +++++++++++++------ .../models/qwen3_omni/qwen3_omni_code2wav.py | 21 +++--- 2 files changed, 53 insertions(+), 34 deletions(-) diff --git a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py index 7be4f6afad7..69dd4e9ceb1 100644 --- a/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py +++ b/tests/model_executor/models/qwen3_omni/test_qwen3_omni_code2wav_streaming.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Regression tests for ``Qwen3OmniMoeCode2Wav.chunked_decode_streaming``. Background @@ -210,26 +210,50 @@ def test_no_shift_when_decoder_returns_nominal_length(): torch.testing.assert_close(streamed, gt) -def test_batched_streaming_only_shifts_the_boundary_row(): - """Per-row tail: in a shared window the shorter row had right context from - the longer row's frames, so it lost nothing (tail==0, no shift); only the - row at the decoded boundary (code_seq_len == window length) is shifted.""" - model = _ShortTrimDecoder() - # One window of W=30 frames, two rows. codes[:, 0, :] are absolute indices. - idx = torch.arange(30, dtype=torch.long) - window = idx.view(1, 1, 30).expand(2, _Q, 30).contiguous() +@pytest.mark.parametrize("cudagraph", [False, True]) +@pytest.mark.parametrize("trim", [_TRIM, 0]) +@pytest.mark.parametrize("short_frames,left_context", [(1, 0), (4, 0), (29, 4), (26, 25), (37, 25)]) +@pytest.mark.parametrize("reverse", [False, True]) +def test_streaming_batch_matches_independent_requests(cudagraph, trim, short_frames, left_context, reverse): + """Padding a request beside a longer one cannot supply valid future codec frames. - long_frames, short_frames = 30, 20 # row0 spans the full window, row1 is shorter - lc = 5 - out = Qwen3OmniMoeCode2Wav.chunked_decode_streaming( + Compare values as well as lengths: an incorrect start and end can shift + the whole chunk while preserving its duration. The independent request + is the oracle, rather than another copy of the slicing arithmetic. + """ + model = _ShortTrimDecoder(trim=trim, cudagraph=cudagraph) + short = _make_codes(short_frames) + long = _make_codes(50) + 100 + windows = [(short, left_context), (long, 25)] + if reverse: + windows.reverse() + batch = torch.zeros(2, _Q, 50, dtype=torch.long) + for i, (window, _) in enumerate(windows): + batch[i, :, : window.shape[-1]] = window[0] + actual = Qwen3OmniMoeCode2Wav.chunked_decode_streaming( model, - window, - left_context_size=[lc, lc], - seq_token_counts=[long_frames * _Q, short_frames * _Q], + batch, + left_context_size=[lc for _, lc in windows], + seq_token_counts=[window.numel() for window, _ in windows], ) - - batch_len = 30 * _UP - _TRIM # actual decoder output width - # Row 0 (boundary row): tail == TRIM -> start = lc*up - TRIM - assert out[0].shape[-1] == batch_len - (lc * _UP - _TRIM) - # Row 1 (shorter): code_seq_len*up < batch_len -> tail == 0 -> start = lc*up - assert out[1].shape[-1] == short_frames * _UP - lc * _UP + assert len(actual) == len(windows) + for output, (window, lc) in zip(actual, windows): + expected = _decode_chunk(model, window, lc) + torch.testing.assert_close(output.reshape(-1), expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("cudagraph", [False, True]) +def test_streaming_preserves_samples_when_batch_composition_changes(cudagraph): + """A singleton -> mixed batch -> singleton stream cannot skip then repeat samples.""" + model = _ShortTrimDecoder(cudagraph=cudagraph) + codes = _make_codes(54) + first = _decode_chunk(model, codes[:, :, :4], 0) + middle_batch = torch.zeros(2, _Q, 50, dtype=torch.long) + middle_batch[0, :, :29] = codes[0, :, :29] + middle_batch[1] = _make_codes(50)[0] + 100 + middle = Qwen3OmniMoeCode2Wav.chunked_decode_streaming( + model, middle_batch, left_context_size=[4, 25], seq_token_counts=[29 * _Q, 50 * _Q] + )[0].reshape(-1) + last = _decode_chunk(model, codes[:, :, 4:54], 25) + streamed = torch.cat((first, middle, last)) + torch.testing.assert_close(streamed, _single_shot(model, codes), rtol=0, atol=0) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py index 67352eecb6c..bc8f6b3b49e 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -316,21 +316,16 @@ def chunked_decode_streaming( else: # Fallback: assume all batch elements share the same sequence length. code_seq_lens = [codes.shape[-1]] * codes.shape[0] + # The decoder trims the same right-edge tail from every batch row. + # Infer it from the decoded window, not a shorter request's length: + # padding that request with codec zeros does not provide valid context. + tail = max(0, int(codes.shape[-1] * self.total_upsample) - batch_wav.shape[-1]) for idx, code_seq_len in enumerate(code_seq_lens): - # The eager decoder (self(codes)) emits a fixed `tail` fewer samples than - # code_seq_len*total_upsample: the causal conv stack trims the sequence end - # because it lacks right context. The original slice assumed full length, so - # under async_chunk every streaming chunk silently dropped its last `tail` - # samples -> a ~23ms gap at every chunk boundary + cumulative ~1.2% time - # compression. Shift the start back by `tail` so this chunk refills the gap the - # previous chunk left (those frames are re-decoded here with proper context). - # `tail` is measured per call, so this is a no-op only when the decoder returns - # the nominal full length (tail==0, e.g. a C==0 decoder like Qwen3-TTS); for - # Qwen3-Omni Code2Wav both the eager forward AND the CUDA-graph wrapper - # (_trim_replay_output -> actual*upsample - 555) are short, so it applies to both. - tail = max(0, int(code_seq_len * self.total_upsample) - batch_wav.shape[-1]) + # Refill the previous chunk's tail using the re-decoded left context, + # and exclude the current row's padded tail just as singleton decode does. start = max(0, left_context_size[idx] * self.total_upsample - tail) - wav_chunk = batch_wav[idx, :, start : code_seq_len * self.total_upsample] + end = max(0, code_seq_len * self.total_upsample - tail) + wav_chunk = batch_wav[idx, :, start:end] wavs.append(wav_chunk) return wavs