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