diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py index 4c94eb520..45b40d026 100644 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py @@ -24,7 +24,7 @@ class _FakeDecoder(nn.Module): super().__init__() self.total_upsample = total_upsample self.decode_calls: list[dict[str, int]] = [] - self.batched_decode_calls: list[dict[str, int]] = [] + self.batched_decode_calls: list[dict[str, object]] = [] self.decode_codes: list[torch.Tensor] = [] self.cudagraph_calls: list[dict[str, int | torch.device]] = [] @@ -35,6 +35,7 @@ class _FakeDecoder(nn.Module): self, codes: torch.Tensor, *, + caches=None, chunk_size: int = 300, left_context_size: int = 25, ) -> torch.Tensor: @@ -58,6 +59,7 @@ class _FakeDecoder(nn.Module): codes: torch.Tensor, lengths: list[int], *, + caches=None, chunk_size: int = 300, left_context_size: int = 25, max_batch_size: int = 0, @@ -69,14 +71,28 @@ class _FakeDecoder(nn.Module): "max_batch_size": max_batch_size, "codes_shape": tuple(codes.shape), "lengths": tuple(lengths), + "caches": caches, } ) + if caches is not None: + outputs = [] + for row, (length, cache) in enumerate(zip(lengths, caches, strict=True)): + output = self.chunked_decode( + codes[row : row + 1, :, :length], + caches=cache, + chunk_size=chunk_size, + left_context_size=left_context_size, + ) + outputs.append(output[0]) + return outputs + batch = codes.shape[0] frames = codes.shape[-1] wav_len = frames * self.total_upsample + 6 wav = torch.arange(wav_len, dtype=torch.float32).view(1, 1, -1) offsets = torch.arange(batch, dtype=torch.float32).view(batch, 1, 1) * 1000 - return wav.expand(batch, 1, wav_len) + offsets + outputs = wav.expand(batch, 1, wav_len) + offsets + return [outputs[row, :, : lengths[row] * self.total_upsample] for row in range(batch)] def enable_cudagraph(self, **kwargs): self.cudagraph_calls.append(kwargs) @@ -157,7 +173,7 @@ def _load_weights_noop(model: Qwen3TTSCode2Wav) -> set[str]: return model.load_weights(iter(())) -def test_forward_trims_context_on_exact_frame_boundaries(): +def test_forward_uses_decoder_audio_contract_on_exact_frame_boundaries(): model = _make_model() out = model.forward( @@ -170,7 +186,7 @@ def test_forward_trims_context_on_exact_frame_boundaries(): torch.testing.assert_close(audio, expected) -def test_forward_trims_trailing_padding_without_context(): +def test_forward_uses_decoder_audio_contract_without_context(): model = _make_model() out = model.forward( @@ -183,24 +199,40 @@ def test_forward_trims_trailing_padding_without_context(): torch.testing.assert_close(audio, expected) -def test_forward_reuses_cached_ref_context_for_followup_chunk(): +@pytest.mark.parametrize("skipped_length", [0, 1], ids=["empty", "malformed"]) +def test_stateless_batch_uses_original_request_index_for_left_context(skipped_length): model = _make_model() - first_codes = torch.tensor( - [ - 9, - 8, - 1, - 2, - 9, - 8, - 3, - 4, + skipped_request = torch.zeros(skipped_length, dtype=torch.long) + valid_request = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long) + out = model.forward( + input_ids=torch.cat([skipped_request, valid_request]), + seq_token_counts=[skipped_length, valid_request.numel()], + runtime_additional_information=[ + { + "codes": {"audio": torch.zeros(skipped_length, dtype=torch.float32)}, + "meta": {"left_context_size": 0}, + }, + { + "codes": {"audio": torch.zeros(8, dtype=torch.float32)}, + "meta": {"left_context_size": 2}, + }, ], - dtype=torch.long, ) - model.forward( - input_ids=first_codes, + wav_output = out.multimodal_outputs["model_outputs"] + torch.testing.assert_close(wav_output[0], torch.empty(0)) + torch.testing.assert_close(wav_output[1], torch.arange(8, 16, dtype=torch.float32)) + + +def test_stateless_forward_trims_reference_prefix_from_decoder_output(): + model = _make_model() + + def _incremental_decode(codes, lengths, *, caches, **_kwargs): + return torch.arange(16, dtype=torch.float32).view(1, 1, -1) + + model.decoder.batched_chunked_decode = _incremental_decode + out = model.forward( + input_ids=torch.arange(8, dtype=torch.long), runtime_additional_information=[ { "meta": { @@ -213,104 +245,48 @@ def test_forward_reuses_cached_ref_context_for_followup_chunk(): ], ) - followup_codes = torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long) + audio = out.multimodal_outputs["model_outputs"][0] + torch.testing.assert_close(audio, torch.arange(8, 16, dtype=torch.float32)) + + +def test_request_aware_forward_passes_only_delta_frames_to_decoder(): + model = _make_model(async_chunk=True) + model.forward( - input_ids=followup_codes, - runtime_additional_information=[ - { - "meta": { - "left_context_size": 3, - "ref_context_size": 2, - "ref_context_request_id": "rid", - "ref_context_included": False, - "finished": torch.tensor(True), - } + input_ids=torch.arange(6, dtype=torch.long), + runtime_additional_information=[{"meta": {"request_id": "rid", "left_context_size": 0}}], + ) + model.forward( + input_ids=torch.arange(4, dtype=torch.long), + runtime_additional_information=[{"meta": {"request_id": "rid", "left_context_size": 0}}], + ) + + assert [tuple(codes.shape) for codes in model.decoder.decode_codes[-2:]] == [ + (1, _NUM_QUANTIZERS, 3), + (1, _NUM_QUANTIZERS, 2), + ] + + +def test_four_frame_first_chunk_uses_configured_initial_size_without_ramp(): + model = _make_model( + async_chunk=True, + stage_connector_config={ + "extra": { + "codec_chunk_frames": 25, + "initial_codec_chunk_frames": 4, } - ], + }, + ) + _load_weights_noop(model) + + model.forward( + input_ids=torch.arange(8, dtype=torch.long), + runtime_additional_information=[{"meta": {"request_id": "four-frame-request", "left_context_size": 0}}], ) - torch.testing.assert_close( - model.decoder.decode_codes[-1], - torch.tensor( - [ - [ - [9, 8, 5, 6, 7], - [9, 8, 15, 16, 17], - ] - ], - dtype=torch.long, - ), - ) - assert "rid" not in model._ref_context_cache - - -def test_forward_fails_fast_when_ref_context_cache_is_missing(): - model = _make_model() - - followup_codes = torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long) - with pytest.raises(ValueError, match="Missing Qwen3-TTS ref context cache"): - model.forward( - input_ids=followup_codes, - runtime_additional_information=[ - { - "meta": { - "left_context_size": 3, - "ref_context_size": 2, - "ref_context_request_id": "rid", - "ref_context_included": False, - } - } - ], - ) - - -def test_ref_context_cache_evicts_lru_entries_when_requests_abort(): - model = _make_model() - model._ref_context_cache_max_entries = 2 - - first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long) - for rid in ("a", "b", "c"): - model.forward( - input_ids=first_codes, - runtime_additional_information=[ - { - "meta": { - "left_context_size": 2, - "ref_context_size": 2, - "ref_context_request_id": rid, - "ref_context_included": True, - } - } - ], - ) - - assert list(model._ref_context_cache) == ["b", "c"] - assert model._ref_context_cache_bytes == sum(model._tensor_nbytes(t) for t in model._ref_context_cache.values()) - - -def test_ref_context_cache_evicts_to_byte_cap(): - model = _make_model() - model._ref_context_cache_max_entries = 100 - model._ref_context_cache_max_bytes = 33 - - first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long) - for rid in ("a", "b"): - model.forward( - input_ids=first_codes, - runtime_additional_information=[ - { - "meta": { - "left_context_size": 2, - "ref_context_size": 2, - "ref_context_request_id": rid, - "ref_context_included": True, - } - } - ], - ) - - assert list(model._ref_context_cache) == ["b"] - assert model._ref_context_cache_bytes == 32 + assert model.decoder._initial_codec_chunk_frames == 4 + assert model.decoder._incremental_chunk_ramp == [] + assert tuple(model.decoder.decode_codes[-1].shape) == (1, _NUM_QUANTIZERS, 4) def test_connector_codec_chunking_does_not_override_decode_chunking(): @@ -332,7 +308,7 @@ def test_connector_codec_chunking_does_not_override_decode_chunking(): model.forward( input_ids=torch.arange(12, dtype=torch.long), - runtime_additional_information=[{"meta": {"left_context_size": 0}}], + runtime_additional_information=[{"meta": {"request_id": "rid", "left_context_size": 0}}], ) assert model.decoder.decode_calls[-1] == { @@ -391,6 +367,66 @@ def test_forward_emits_zero_samples_for_empty_codec_payload(): assert out.multimodal_outputs["model_outputs"][0].numel() == 0 +def test_dummy_runtime_information_provides_finished_request_state(): + model = _make_model(async_chunk=True) + + runtime_info = model.get_dummy_runtime_additional_information(2) + + assert len(runtime_info) == 2 + assert [info["meta"]["request_id"] for info in runtime_info] == [ + "__qwen3_tts_dummy_run__0", + "__qwen3_tts_dummy_run__1", + ] + assert all(info["meta"]["finished"] for info in runtime_info) + + +def test_dummy_request_state_is_marked_for_graph_stats_suppression(): + model = _make_model(async_chunk=True) + runtime_info = model.get_dummy_runtime_additional_information(1) + + model.forward( + input_ids=torch.arange(4, dtype=torch.long), + runtime_additional_information=runtime_info, + ) + + cache = model.decoder.batched_decode_calls[0]["caches"][0] + assert cache["_is_dummy_run"] is True + assert model._decoder_state_cache == {} + + +def test_decoder_state_cache_capacity_warns_without_evicting_or_failing(): + model = _make_model(async_chunk=True) + assert model._decoder_state_cache_warn_entries == 512 + model._decoder_state_cache.update({f"active-{index}": {"prefix_frames": 0} for index in range(512)}) + + with patch("vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav.logger.warning_once") as warning: + model.forward( + input_ids=torch.arange(4, dtype=torch.long), + runtime_additional_information=[{"meta": {"request_id": "new-request"}}], + ) + + warning.assert_called_once() + assert len(model._decoder_state_cache) == 513 + assert "new-request" in model._decoder_state_cache + + +def test_decoder_state_cache_uses_scheduler_request_id_for_finished_cleanup(): + model = _make_model(async_chunk=True) + + model.forward( + input_ids=torch.arange(4, dtype=torch.long), + request_ids=["internal-request-id"], + runtime_additional_information=[{"meta": {"request_id": "external-request-id"}}], + ) + + assert "internal-request-id" in model._decoder_state_cache + assert "external-request-id" not in model._decoder_state_cache + + model.on_requests_finished(["internal-request-id"]) + + assert model._decoder_state_cache == {} + + def test_forward_batches_equal_length_requests_in_one_decoder_call(): model = _make_model() @@ -403,11 +439,15 @@ def test_forward_batches_equal_length_requests_in_one_decoder_call(): ], ) - assert model.decoder.decode_calls == [ + assert model.decoder.decode_calls == [] + assert model.decoder.batched_decode_calls == [ { "chunk_size": 300, "left_context_size": 25, + "max_batch_size": 0, "codes_shape": (2, _NUM_QUANTIZERS, 3), + "lengths": (3, 3), + "caches": None, } ] audios = out.multimodal_outputs["model_outputs"] @@ -415,10 +455,81 @@ def test_forward_batches_equal_length_requests_in_one_decoder_call(): torch.testing.assert_close(audios[1], torch.arange(1004, 1012, dtype=torch.float32)) -def test_forward_uses_variable_length_chunk_batching_for_long_bucket_group(): +def test_forward_batches_request_aware_decoder_states(): + model = _make_model(async_chunk=True) + calls = [] + + def _batched_chunked_decode(codes, lengths, caches, **kwargs): + calls.append((codes, lengths, caches)) + return [torch.arange(16, dtype=torch.float32).view(1, -1) + row * 100 for row in range(2)] + + model.decoder.batched_chunked_decode = _batched_chunked_decode + per_request = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long) + out = model.forward( + input_ids=torch.cat([per_request, per_request]), + seq_token_counts=[8, 8], + runtime_additional_information=[ + { + "meta": { + "left_context_size": 2, + "ref_context_size": 2, + "ref_context_request_id": request_id, + "request_id": request_id, + "ref_context_included": True, + } + } + for request_id in ("a", "b") + ], + ) + + assert len(calls) == 1 + assert tuple(calls[0][0].shape) == (2, 2, 4) + assert calls[0][1] == [4, 4] + assert [cache["prefix_frames"] for cache in calls[0][2]] == [2, 2] + audios = out.multimodal_outputs["model_outputs"] + torch.testing.assert_close(audios[0], torch.arange(16, dtype=torch.float32)) + torch.testing.assert_close(audios[1], torch.arange(100, 116, dtype=torch.float32)) + + +def test_forward_allows_followup_delta_without_reference_prefix(): + model = _make_model(async_chunk=True) + first_codes = torch.tensor([9, 8, 1, 2, 9, 8, 3, 4], dtype=torch.long) + model.forward( + input_ids=first_codes, + runtime_additional_information=[ + { + "meta": { + "request_id": "rid", + "ref_context_request_id": "rid", + "ref_context_size": 2, + "ref_context_included": True, + } + } + ], + ) + + model.forward( + input_ids=torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long), + runtime_additional_information=[ + { + "meta": { + "request_id": "rid", + "ref_context_request_id": "rid", + "ref_context_size": 2, + "ref_context_included": False, + } + } + ], + ) + + assert [tuple(codes.shape) for codes in model.decoder.decode_codes[-2:]] == [ + (1, _NUM_QUANTIZERS, 4), + (1, _NUM_QUANTIZERS, 3), + ] + + +def test_forward_passes_variable_length_padded_batch_to_decoder(): model = _make_model() - model._decode_batch_bucket_frames = [4] - model._decode_variable_chunk_batch_min_frames = 1 out = model.forward( input_ids=torch.arange(12, dtype=torch.long), @@ -437,6 +548,7 @@ def test_forward_uses_variable_length_chunk_batching_for_long_bucket_group(): "max_batch_size": 0, "codes_shape": (2, _NUM_QUANTIZERS, 4), "lengths": (2, 4), + "caches": None, } ] audios = out.multimodal_outputs["model_outputs"] @@ -444,108 +556,6 @@ def test_forward_uses_variable_length_chunk_batching_for_long_bucket_group(): torch.testing.assert_close(audios[1], torch.arange(1008, 1016, dtype=torch.float32)) -def test_forward_bucket_batches_different_length_requests_and_trims_rows(): - model = _make_model() - model._decode_batch_bucket_frames = [4] - - out = model.forward( - input_ids=torch.arange(18, dtype=torch.long), - seq_token_counts=[4, 6, 8], - runtime_additional_information=[ - {"meta": {"left_context_size": 0}}, - {"meta": {"left_context_size": 1}}, - {"meta": {"left_context_size": 2}}, - ], - ) - - assert model.decoder.decode_calls == [ - { - "chunk_size": 300, - "left_context_size": 25, - "codes_shape": (3, _NUM_QUANTIZERS, 4), - } - ] - assert model.decoder.batched_decode_calls == [] - audios = out.multimodal_outputs["model_outputs"] - torch.testing.assert_close(audios[0], torch.arange(8, dtype=torch.float32)) - torch.testing.assert_close(audios[1], torch.arange(1004, 1012, dtype=torch.float32)) - torch.testing.assert_close(audios[2], torch.arange(2008, 2016, dtype=torch.float32)) - - -def test_forward_bucket_pads_only_to_group_max_frame_length(): - model = _make_model() - model._decode_batch_bucket_frames = [8] - - out = model.forward( - input_ids=torch.arange(10, dtype=torch.long), - seq_token_counts=[4, 6], - runtime_additional_information=[ - {"meta": {"left_context_size": 0}}, - {"meta": {"left_context_size": 0}}, - ], - ) - - assert model.decoder.decode_calls == [ - { - "chunk_size": 300, - "left_context_size": 25, - "codes_shape": (2, _NUM_QUANTIZERS, 3), - } - ] - assert model.decoder.batched_decode_calls == [] - audios = out.multimodal_outputs["model_outputs"] - torch.testing.assert_close(audios[0], torch.arange(8, dtype=torch.float32)) - torch.testing.assert_close(audios[1], torch.arange(1000, 1012, dtype=torch.float32)) - - -def test_forward_splits_bucket_groups_by_configured_max_batch_size(): - model = _make_model() - model._decode_batch_bucket_frames = [4] - model._decode_batch_max_size = 2 - - model.forward( - input_ids=torch.arange(18, dtype=torch.long), - seq_token_counts=[4, 6, 8], - runtime_additional_information=[ - {"meta": {"left_context_size": 0}}, - {"meta": {"left_context_size": 0}}, - {"meta": {"left_context_size": 0}}, - ], - ) - - assert model.decoder.decode_calls == [ - { - "chunk_size": 300, - "left_context_size": 25, - "codes_shape": (2, _NUM_QUANTIZERS, 3), - }, - { - "chunk_size": 300, - "left_context_size": 25, - "codes_shape": (1, _NUM_QUANTIZERS, 4), - }, - ] - assert model.decoder.batched_decode_calls == [] - - -def test_forward_does_not_pad_singleton_bucket_group(): - model = _make_model() - model._decode_batch_bucket_frames = [4] - - model.forward( - input_ids=torch.arange(4, dtype=torch.long), - runtime_additional_information=[{"meta": {"left_context_size": 0}}], - ) - - assert model.decoder.decode_calls == [ - { - "chunk_size": 300, - "left_context_size": 25, - "codes_shape": (1, _NUM_QUANTIZERS, 2), - } - ] - - def test_decode_chunking_override_is_passed_to_cudagraph(): model = _make_model( async_chunk=True, @@ -563,27 +573,26 @@ def test_decode_chunking_override_is_passed_to_cudagraph(): _load_weights_noop(model) assert model.decoder.cudagraph_calls[-1] == { - "capture_sizes": None, "capture_batch_sizes": None, - "extra_capture_shapes": None, - "compile_shapes": None, + "stateless_capture_sizes": None, "device": torch.device("cuda"), "codec_chunk_frames": 25, "codec_left_context_frames": 72, + "initial_codec_chunk_frames": 1, + "codec_chunk_ramp": None, + "async_chunk": True, "decode_chunk_size": 400, "decode_left_context": 17, } -def test_cudagraph_capture_shapes_can_be_configured(): +def test_cudagraph_batch_config_is_preserved(): model = _make_model( async_chunk=True, device=torch.device("cuda"), stage_connector_config={ "extra": { - "decode_cudagraph_capture_sizes": "97,325", "decode_cudagraph_batch_sizes": [1, 2, 4, 8], - "decode_cudagraph_extra_capture_shapes": ["3:325", [5, 325]], } }, ) @@ -591,26 +600,39 @@ def test_cudagraph_capture_shapes_can_be_configured(): _load_weights_noop(model) call = model.decoder.cudagraph_calls[-1] - assert call["capture_sizes"] == [97, 325] assert call["capture_batch_sizes"] == [1, 2, 4, 8] - assert call["extra_capture_shapes"] == [(3, 325), (5, 325)] -def test_decode_compile_shapes_can_be_configured(): +def test_stateless_cudagraph_capture_sizes_are_passed_from_config(): + model = _make_model( + async_chunk=False, + device=torch.device("cuda"), + stage_connector_config={ + "extra": { + "decode_cudagraph_capture_sizes": [97, 400], + } + }, + ) + + _load_weights_noop(model) + + assert model.decoder.cudagraph_calls[-1]["stateless_capture_sizes"] == [97, 400] + + +def test_async_cudagraph_ignores_stateless_capture_sizes(): model = _make_model( async_chunk=True, device=torch.device("cuda"), stage_connector_config={ "extra": { - "decode_compile_shapes": ["1:325", [1, 73]], + "decode_cudagraph_capture_sizes": [97, 400], } }, ) _load_weights_noop(model) - call = model.decoder.cudagraph_calls[-1] - assert call["compile_shapes"] == [(1, 73), (1, 325)] + assert model.decoder.cudagraph_calls[-1]["stateless_capture_sizes"] is None def test_decode_tf32_can_be_configured(): @@ -642,23 +664,19 @@ def test_decode_tf32_can_be_configured(): torch.set_float32_matmul_precision(old_matmul_precision) -def test_decode_batch_bucket_frames_can_be_configured(): +def test_decode_batch_max_size_can_be_configured(): model = _make_model( async_chunk=True, stage_connector_config={ "extra": { - "decode_batch_bucket_frames": "73,169", "decode_batch_max_size": 10, - "decode_variable_chunk_batch_min_frames": 512, } }, ) _load_weights_noop(model) - assert model._decode_batch_bucket_frames == [73, 169] assert model._decode_batch_max_size == 10 - assert model._decode_variable_chunk_batch_min_frames == 512 def test_invalid_decode_batch_max_size_is_rejected(): diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_incremental_decode.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_incremental_decode.py new file mode 100644 index 000000000..58cb04884 --- /dev/null +++ b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_incremental_decode.py @@ -0,0 +1,1253 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import copy +from collections import Counter +from types import SimpleNamespace + +import pytest +import torch +import torch.nn as nn +from transformers.cache_utils import DynamicCache + +from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav import Qwen3TTSCode2Wav +from vllm_omni.model_executor.models.qwen3_tts.segmented_graph_wrapper import CUDAGraphDecoderWrapper +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import ( + Qwen3TTSTokenizerV2DecoderConfig, +) +from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( + _CONV_CONTEXT_FRAME, + _DOWNSTREAM_CONTEXT_FRAME, + Qwen3TTSTokenizerV2Decoder, + Qwen3TTSTokenizerV2DecoderTransformerModel, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Shared test helpers + + +def _decoder_stub(**kwargs): + decoder = SimpleNamespace(**kwargs) + decoder.batched_request_decode = Qwen3TTSTokenizerV2Decoder.batched_request_decode.__get__(decoder) + return decoder + + +def _make_transformer_model() -> Qwen3TTSTokenizerV2DecoderTransformerModel: + config = Qwen3TTSTokenizerV2DecoderConfig( + hidden_size=32, + latent_dim=16, + max_position_embeddings=512, + num_attention_heads=4, + num_key_value_heads=4, + intermediate_size=64, + num_hidden_layers=1, + num_quantizers=2, + decoder_dim=32, + upsample_rates=(2,), + upsampling_ratios=(2,), + ) + config._attn_implementation = "sdpa" + return Qwen3TTSTokenizerV2DecoderTransformerModel._from_config(config).eval() + + +def _make_small_decoder() -> Qwen3TTSTokenizerV2Decoder: + config = Qwen3TTSTokenizerV2DecoderConfig( + codebook_size=32, + hidden_size=16, + latent_dim=16, + codebook_dim=16, + num_attention_heads=2, + num_key_value_heads=2, + intermediate_size=32, + num_hidden_layers=1, + num_quantizers=2, + decoder_dim=32, + upsample_rates=(8, 5, 4, 3), + upsampling_ratios=(2, 2), + sliding_window=72, + ) + return Qwen3TTSTokenizerV2Decoder(config).eval() + + +def _run_downstream(decoder: Qwen3TTSTokenizerV2Decoder, hidden: torch.Tensor) -> torch.Tensor: + for blocks in decoder.upsample: + for block in blocks: + hidden = block(hidden) + for block in decoder.decoder: + hidden = block(hidden) + return hidden + + +class _RecordingDecoder(nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(sliding_window=72) + self.prefix_frames = None + + def decode_suffix(self, codes, caches, prefix_frames): + self.prefix_frames = prefix_frames + return codes + + +def _make_stateful_eager_decoder_pair() -> tuple[Qwen3TTSTokenizerV2Decoder, Qwen3TTSTokenizerV2Decoder]: + batched_decoder = _make_small_decoder() + per_request_decoder = copy.deepcopy(batched_decoder) + batched_decoder._incremental_chunk_frames = 25 + per_request_decoder._incremental_chunk_frames = 25 + return batched_decoder, per_request_decoder + + +def test_dummy_forward_uses_exact_decode_during_outer_graph_capture(monkeypatch): + decoder = _make_small_decoder() + codes = torch.zeros(1, decoder.config.num_quantizers, 4, dtype=torch.long) + cache = {"prefix_frames": 0, "_is_dummy_run": True} + expected = torch.ones(1, 1, 7) + + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + monkeypatch.setattr(decoder, "_forward_exact", lambda _codes: expected) + + def _unexpected_stateful_decode(*_args, **_kwargs): + raise AssertionError("capturing dummy run must not initialize decoder state") + + monkeypatch.setattr(decoder, "_decode_xvec_first_chunk", _unexpected_stateful_decode) + + output = decoder(codes, cache) + + assert output is expected + assert cache == {"prefix_frames": 0, "_is_dummy_run": True} + + +def test_code2wav_full_graph_dummy_forward_uses_exact_batched_decode(monkeypatch): + decoder = _make_small_decoder() + expected = torch.arange(7, dtype=torch.float32).view(1, 1, -1) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True) + monkeypatch.setattr(decoder, "_forward_exact", lambda _codes: expected) + + def _unexpected_stateful_decode(*_args, **_kwargs): + raise AssertionError("FULL graph dummy capture must not initialize decoder state") + + monkeypatch.setattr(decoder, "_decode_xvec_first_chunk", _unexpected_stateful_decode) + + model = Qwen3TTSCode2Wav.__new__(Qwen3TTSCode2Wav) + nn.Module.__init__(model) + model.decoder = decoder + model._async_chunk = True + model._num_quantizers = decoder.config.num_quantizers + model._total_upsample = decoder.total_upsample + model._output_sample_rate = 24000 + model._decode_chunk_frames = 300 + model._decode_left_context_frames = 25 + model._decode_batch_max_size = 0 + model._decoder_state_cache = {} + model._decoder_state_cache_warn_entries = 512 + model._logged_codec_stats = True + model._logged_malformed_codec_lengths = set() + model._batch_stats_enabled = False + model._batch_stats_log_every = 0 + model._batch_stats_forwards = 0 + + runtime_info = model.get_dummy_runtime_additional_information(1) + output = model.forward( + input_ids=torch.arange(8, dtype=torch.long), + seq_token_counts=[8], + runtime_additional_information=runtime_info, + ) + + torch.testing.assert_close(output.multimodal_outputs["model_outputs"][0], expected[0, 0]) + assert model._decoder_state_cache == {} + + +def _make_stateful_first_chunk_inputs( + decoder: Qwen3TTSTokenizerV2Decoder, + mode: str, +) -> tuple[torch.Tensor, list[int], list[dict], list[dict]]: + if mode == "xvec": + lengths = [1, 1] + prefix_frames = [0, 0] + else: + lengths = [49, 73] + prefix_frames = [48, 72] + + codes = torch.zeros(2, decoder.config.num_quantizers, max(lengths), dtype=torch.long) + for row, length in enumerate(lengths): + codes[row, :, :length] = torch.randint( + 0, + decoder.config.codebook_size, + (decoder.config.num_quantizers, length), + ) + batched_caches = [{"prefix_frames": frames} for frames in prefix_frames] + per_request_caches = copy.deepcopy(batched_caches) + return codes, lengths, batched_caches, per_request_caches + + +def _decode_per_request( + decoder: Qwen3TTSTokenizerV2Decoder, + codes: torch.Tensor, + lengths: list[int], + caches: list[dict], +) -> torch.Tensor: + return torch.cat( + [ + decoder.chunked_decode(codes[row : row + 1, :, :length], caches=caches[row]) + for row, length in enumerate(lengths) + ], + dim=0, + ) + + +def _initialize_stateful_batch( + mode: str, +) -> tuple[Qwen3TTSTokenizerV2Decoder, Qwen3TTSTokenizerV2Decoder, list[dict], list[dict]]: + batched_decoder, per_request_decoder = _make_stateful_eager_decoder_pair() + codes, lengths, batched_caches, per_request_caches = _make_stateful_first_chunk_inputs(batched_decoder, mode) + batched_decoder.batched_chunked_decode(codes, lengths, caches=batched_caches) + _decode_per_request(per_request_decoder, codes, lengths, per_request_caches) + return batched_decoder, per_request_decoder, batched_caches, per_request_caches + + +# ----------------------------------------------------------------------------- +# Result correctness: components, end-to-end incremental decode, rolling, and batching +# ----------------------------------------------------------------------------- + + +def test_transformer_incremental_kv_matches_full_decode(): + torch.manual_seed(0) + model = _make_transformer_model() + prefix = torch.randn(1, 11, model.config.latent_dim) + suffix = torch.randn(1, 6, model.config.latent_dim) + + prefix_output = model(inputs_embeds=prefix, use_cache=True) + assert isinstance(prefix_output.past_key_values, DynamicCache) + assert prefix_output.past_key_values.get_seq_length() == prefix.shape[1] + suffix_output = model( + inputs_embeds=suffix, + past_key_values=prefix_output.past_key_values, + use_cache=True, + ) + full_output = model(inputs_embeds=torch.cat([prefix, suffix], dim=1)) + + torch.testing.assert_close( + suffix_output.last_hidden_state, + full_output.last_hidden_state[:, prefix.shape[1] :, :], + atol=1e-5, + rtol=1e-4, + ) + assert suffix_output.past_key_values.get_seq_length() == prefix.shape[1] + suffix.shape[1] + + +def test_transformer_left_padded_prefix_mask_matches_unpadded_incremental_decode(): + torch.manual_seed(0) + model = _make_transformer_model() + prefix = torch.randn(1, 11, model.config.latent_dim) + suffix = torch.randn(1, 6, model.config.latent_dim) + prefix_bucket = 16 + prefix_pad = prefix_bucket - prefix.shape[1] + + prefix_output = model(inputs_embeds=prefix, use_cache=True) + suffix_output = model( + inputs_embeds=suffix, + past_key_values=prefix_output.past_key_values, + use_cache=True, + ) + + padded_prefix = torch.cat([torch.zeros_like(prefix[:, :prefix_pad]), prefix], dim=1) + prefix_mask = torch.cat( + [ + torch.zeros(1, prefix_pad, dtype=torch.bool), + torch.ones(1, prefix.shape[1], dtype=torch.bool), + ], + dim=1, + ) + padded_prefix_output = model( + inputs_embeds=padded_prefix, + attention_mask=prefix_mask, + use_cache=True, + ) + assert padded_prefix_output.past_key_values.get_seq_length() == prefix_bucket + padded_suffix_output = model( + inputs_embeds=suffix, + attention_mask=torch.cat( + [prefix_mask, torch.ones(1, suffix.shape[1], dtype=torch.bool)], + dim=1, + ), + past_key_values=padded_prefix_output.past_key_values, + use_cache=True, + ) + + torch.testing.assert_close( + padded_suffix_output.last_hidden_state, + suffix_output.last_hidden_state, + atol=1e-5, + rtol=1e-4, + ) + assert padded_suffix_output.past_key_values.get_seq_length() == prefix_bucket + suffix.shape[1] + + +def test_pre_conv_two_frame_context_matches_full_decode_and_one_frame_does_not(): + torch.manual_seed(2) + decoder = _make_small_decoder() + previous_frames = 19 + new_frames = 7 + hidden = torch.randn(1, decoder.config.codebook_dim, previous_frames + new_frames) + + with torch.no_grad(): + full = decoder.pre_conv(hidden)[..., -new_frames:] + with_full_context = decoder.pre_conv(hidden[..., -(_CONV_CONTEXT_FRAME + new_frames) :])[..., -new_frames:] + with_short_context = decoder.pre_conv(hidden[..., -((_CONV_CONTEXT_FRAME - 1) + new_frames) :])[ + ..., -new_frames: + ] + + assert _CONV_CONTEXT_FRAME == 2 + torch.testing.assert_close(with_full_context, full, atol=1e-6, rtol=1e-5) + assert not torch.allclose(with_short_context, full, atol=1e-6, rtol=1e-5) + + +def test_downstream_twelve_frame_context_matches_full_decode_and_context_is_required(): + torch.manual_seed(3) + decoder = _make_small_decoder() + previous_frames = 20 + new_frames = 5 + hidden = torch.randn(1, decoder.config.latent_dim, previous_frames + new_frames) + + with torch.no_grad(): + full = _run_downstream(decoder, hidden)[..., -new_frames * decoder.total_upsample :] + with_full_context = _run_downstream( + decoder, + hidden[..., -(_DOWNSTREAM_CONTEXT_FRAME + new_frames) :], + )[..., -new_frames * decoder.total_upsample :] + without_context = _run_downstream( + decoder, + hidden[..., -new_frames:], + )[..., -new_frames * decoder.total_upsample :] + + assert _DOWNSTREAM_CONTEXT_FRAME == 12 + torch.testing.assert_close(with_full_context, full) + assert (without_context - full).abs().max().item() > 1e-10 + + +@pytest.mark.parametrize("mode", ["xvec", "icl"]) +def test_incremental_request_audio_matches_full_decode(mode): + torch.manual_seed(3) + decoder = _make_small_decoder() + decoder._incremental_chunk_frames = 25 + prefix_frames = 0 if mode == "xvec" else 10 + suffix_frames = 51 + codes = torch.randint( + 0, + decoder.config.codebook_size, + (1, decoder.config.num_quantizers, prefix_frames + suffix_frames), + ) + suffix_codes = codes[..., prefix_frames:] + caches = {"prefix_frames": prefix_frames} + + with torch.no_grad(): + first = decoder(codes[..., : prefix_frames + 1], caches=caches) + second = decoder(suffix_codes[..., 1:26], caches=caches) + third = decoder(suffix_codes[..., 26:51], caches=caches) + incremental = torch.cat( + [first[..., -decoder.total_upsample :], second, third], + dim=-1, + ) + full = decoder._forward_exact(codes) + full_suffix = full[..., prefix_frames * decoder.total_upsample :] + + assert incremental.shape[-1] == suffix_frames * decoder.total_upsample + assert full_suffix.shape == incremental.shape + assert caches["decoder_prefix_frames"] == prefix_frames + assert caches["suffix_frames"] == suffix_frames + assert caches["past_key_values"].get_seq_length() == prefix_frames + assert all(layer.is_initialized for layer in caches["past_key_values"].layers) + torch.testing.assert_close(incremental, full_suffix, atol=1e-5, rtol=1e-4) + + +def test_single_frame_xvec_prefix_keeps_all_next_chunk_conv_frames(): + torch.manual_seed(4) + config = Qwen3TTSTokenizerV2DecoderConfig( + codebook_size=32, + hidden_size=16, + latent_dim=16, + codebook_dim=16, + num_attention_heads=2, + num_key_value_heads=2, + intermediate_size=32, + num_hidden_layers=1, + num_quantizers=2, + decoder_dim=32, + upsample_rates=(8, 5, 4, 3), + upsampling_ratios=(2, 2), + sliding_window=72, + ) + decoder = Qwen3TTSTokenizerV2Decoder(config).eval() + decoder._incremental_chunk_frames = 25 + codes = torch.randint(0, config.codebook_size, (1, config.num_quantizers, 26)) + caches = {"prefix_frames": 0} + + with torch.no_grad(): + decoder(codes[..., :1], caches=caches) + incremental = decoder(codes[..., 1:], caches=caches) + full = decoder._forward_exact(codes) + + assert caches["suffix_frames"] == 26 + assert caches["suffix_quantized"].shape[-1] == 26 + assert caches["suffix_conv"].shape[1] == 26 + torch.testing.assert_close( + incremental, + full[..., -25 * decoder.total_upsample :], + atol=1e-5, + rtol=1e-4, + ) + + +def test_xvec_rolling_matches_truncated_suffix_window(): + torch.manual_seed(5) + config = Qwen3TTSTokenizerV2DecoderConfig( + codebook_size=32, + hidden_size=16, + latent_dim=16, + codebook_dim=16, + num_attention_heads=2, + num_key_value_heads=2, + intermediate_size=32, + num_hidden_layers=1, + num_quantizers=2, + decoder_dim=32, + upsample_rates=(8, 5, 4, 3), + upsampling_ratios=(2, 2), + sliding_window=72, + ) + decoder = Qwen3TTSTokenizerV2Decoder(config).eval() + decoder._incremental_chunk_frames = 25 + codes = torch.randint(0, config.codebook_size, (1, config.num_quantizers, 101)) + caches = {"prefix_frames": 0} + + with torch.no_grad(): + decoder(codes[..., :1], caches=caches) + decoder(codes[..., 1:26], caches=caches) + decoder(codes[..., 26:51], caches=caches) + decoder(codes[..., 51:76], caches=caches) + rolling = decoder(codes[..., 76:101], caches=caches) + truncated = decoder._forward_exact(codes[..., -(config.sliding_window + 25) :]) + + assert caches["ref_hidden"].shape[-1] == 0 + assert caches["suffix_quantized"].shape[-1] == config.sliding_window + assert caches["suffix_conv"].shape[1] == config.sliding_window - 2 + torch.testing.assert_close( + rolling, + truncated[..., -25 * decoder.total_upsample :], + atol=1e-5, + rtol=1e-4, + ) + + +def test_icl_rolling_matches_reference_plus_truncated_suffix_window(): + torch.manual_seed(6) + config = Qwen3TTSTokenizerV2DecoderConfig( + codebook_size=32, + hidden_size=16, + latent_dim=16, + codebook_dim=16, + num_attention_heads=2, + num_key_value_heads=2, + intermediate_size=32, + num_hidden_layers=1, + num_quantizers=2, + decoder_dim=32, + upsample_rates=(8, 5, 4, 3), + upsampling_ratios=(2, 2), + sliding_window=72, + ) + decoder = Qwen3TTSTokenizerV2Decoder(config).eval() + decoder._incremental_chunk_frames = 25 + prefix_frames = 10 + suffix_frames = 101 + codes = torch.randint( + 0, + config.codebook_size, + (1, config.num_quantizers, prefix_frames + suffix_frames), + ) + prefix_codes = codes[..., :prefix_frames] + suffix_codes = codes[..., prefix_frames:] + caches = {"prefix_frames": prefix_frames} + + with torch.no_grad(): + decoder(codes[..., : prefix_frames + 1], caches=caches) + decoder(suffix_codes[..., 1:26], caches=caches) + decoder(suffix_codes[..., 26:51], caches=caches) + decoder(suffix_codes[..., 51:76], caches=caches) + rolling = decoder(suffix_codes[..., 76:101], caches=caches) + truncated_codes = torch.cat( + [prefix_codes, suffix_codes[..., -(config.sliding_window + 25) :]], + dim=-1, + ) + truncated = decoder._forward_exact(truncated_codes) + + assert caches["decoder_prefix_frames"] == prefix_frames + assert caches["ref_hidden"].shape[-1] == prefix_frames + assert caches["suffix_quantized"].shape[-1] == config.sliding_window + assert caches["suffix_conv"].shape[1] == config.sliding_window - 2 + torch.testing.assert_close( + rolling, + truncated[..., -25 * decoder.total_upsample :], + atol=1e-5, + rtol=1e-4, + ) + + +@pytest.mark.parametrize("mode", ["xvec", "icl"]) +def test_batched_stateful_first_chunk_matches_per_request(mode): + torch.manual_seed(7) + batched_decoder, per_request_decoder = _make_stateful_eager_decoder_pair() + codes, lengths, batched_caches, per_request_caches = _make_stateful_first_chunk_inputs(batched_decoder, mode) + + with torch.no_grad(): + batched = batched_decoder.batched_chunked_decode(codes, lengths, caches=batched_caches) + per_request = _decode_per_request(per_request_decoder, codes, lengths, per_request_caches) + + torch.testing.assert_close(torch.stack(batched), per_request, atol=1e-5, rtol=1e-4) + for batched_cache, per_request_cache in zip(batched_caches, per_request_caches, strict=True): + batched_suffix_frames = batched_cache.get("suffix_frames", batched_cache["suffix_quantized"].shape[-1]) + per_request_suffix_frames = per_request_cache.get( + "suffix_frames", per_request_cache["suffix_quantized"].shape[-1] + ) + assert batched_suffix_frames == per_request_suffix_frames == 1 + torch.testing.assert_close(batched_cache["suffix_quantized"], per_request_cache["suffix_quantized"]) + torch.testing.assert_close(batched_cache["suffix_conv"], per_request_cache["suffix_conv"]) + + +@pytest.mark.parametrize("mode", ["xvec", "icl"]) +def test_batched_stateful_growing_suffix_matches_per_request(mode): + torch.manual_seed(8) + with torch.no_grad(): + batched_decoder, per_request_decoder, batched_caches, per_request_caches = _initialize_stateful_batch(mode) + codes = torch.randint(0, batched_decoder.config.codebook_size, (2, batched_decoder.config.num_quantizers, 25)) + batched = batched_decoder.batched_chunked_decode(codes, [25, 25], caches=batched_caches) + per_request = _decode_per_request(per_request_decoder, codes, [25, 25], per_request_caches) + + torch.testing.assert_close(torch.stack(batched), per_request, atol=1e-5, rtol=1e-4) + for batched_cache, per_request_cache in zip(batched_caches, per_request_caches, strict=True): + assert batched_cache["suffix_frames"] == per_request_cache["suffix_frames"] == 26 + torch.testing.assert_close(batched_cache["suffix_quantized"], per_request_cache["suffix_quantized"]) + torch.testing.assert_close(batched_cache["suffix_conv"], per_request_cache["suffix_conv"]) + + +@pytest.mark.parametrize("mode", ["xvec", "icl"]) +def test_batched_stateful_rolling_suffix_matches_per_request(mode): + torch.manual_seed(9) + with torch.no_grad(): + batched_decoder, per_request_decoder, batched_caches, per_request_caches = _initialize_stateful_batch(mode) + for _ in range(3): + codes = torch.randint( + 0, + batched_decoder.config.codebook_size, + (2, batched_decoder.config.num_quantizers, 25), + ) + batched_decoder.batched_chunked_decode(codes, [25, 25], caches=batched_caches) + _decode_per_request(per_request_decoder, codes, [25, 25], per_request_caches) + + rolling_codes = torch.randint( + 0, + batched_decoder.config.codebook_size, + (2, batched_decoder.config.num_quantizers, 25), + ) + batched = batched_decoder.batched_chunked_decode(rolling_codes, [25, 25], caches=batched_caches) + per_request = _decode_per_request(per_request_decoder, rolling_codes, [25, 25], per_request_caches) + + torch.testing.assert_close(torch.stack(batched), per_request, atol=1e-5, rtol=1e-4) + for batched_cache, per_request_cache in zip(batched_caches, per_request_caches, strict=True): + assert batched_cache["suffix_frames"] == per_request_cache["suffix_frames"] == 97 + assert batched_cache["suffix_quantized"].shape[-1] == batched_decoder.config.sliding_window + assert batched_cache["suffix_conv"].shape[1] == batched_decoder.config.sliding_window - 2 + torch.testing.assert_close(batched_cache["suffix_quantized"], per_request_cache["suffix_quantized"]) + torch.testing.assert_close(batched_cache["suffix_conv"], per_request_cache["suffix_conv"]) + + +# ----------------------------------------------------------------------------- +# Warm-up, capture, replay, and cache initialization regressions +# ----------------------------------------------------------------------------- + + +def test_segmented_capture_sizes_are_derived_from_decoder_and_chunk_config(): + decoder = _RecordingDecoder() + decoder.config.sliding_window = 10 + + wrapper = CUDAGraphDecoderWrapper( + decoder, + initial_chunk_frames=2, + codec_chunk_frames=4, + enabled=False, + ) + + assert wrapper.prefix_length == 10 + assert wrapper.icl_capture_sizes == [6, 10, 14] + assert wrapper._icl_previous_frames_by_target == {6: 2, 10: 6, 14: 10} + assert wrapper._xvec_previous_frames_by_target == {6: 2, 10: 6, 14: 10} + + +def test_stateless_capture_sizes_extend_defaults_with_configured_sizes(): + decoder = _RecordingDecoder() + wrapper = CUDAGraphDecoderWrapper( + decoder, + async_chunk=False, + decode_chunk_size=300, + decode_left_context=25, + stateless_capture_sizes=[97, 400], + enabled=False, + ) + + assert wrapper.stateless_capture_sizes == [97, 150, 325, 400] + + +def test_capture_uses_cached_conv_lengths(): + decoder = _RecordingDecoder() + decoder.config.sliding_window = 72 + wrapper = CUDAGraphDecoderWrapper( + decoder, + initial_chunk_frames=1, + codec_chunk_frames=25, + enabled=False, + ) + + assert [wrapper._cached_conv_frames(frames) for frames in (1, 26, 51, 72)] == [1, 26, 51, 70] + + +def test_icl_prefix_graph_cache_keeps_physical_prefix_length(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + batch_size = 1 + num_quantizers = 2 + prefix_length = 72 + initial_chunk_frames = 1 + physical_frames = prefix_length + initial_chunk_frames + total_upsample = 1920 + + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.icl_prefix_states = { + batch_size: { + "graph": type("_Graph", (), {"replay": lambda self: None})(), + "input": { + "codes": torch.zeros(batch_size, num_quantizers, physical_frames), + "hidden_mask": torch.ones(batch_size, 1, physical_frames), + "prefix_attention_mask": torch.ones(batch_size, prefix_length, dtype=torch.bool), + "suffix_attention_mask": torch.ones(batch_size, physical_frames, dtype=torch.bool), + }, + "output": torch.zeros(batch_size, 1, physical_frames * total_upsample), + "cache": { + "decoder_prefix_frames": prefix_length, + "ref_hidden": torch.zeros(batch_size, 2, prefix_length), + "ref_conv": torch.zeros(batch_size, prefix_length, 2), + "prefix_hidden": torch.zeros(batch_size, prefix_length, 2), + "past_key_values": object(), + "suffix_quantized": torch.zeros(batch_size, num_quantizers, initial_chunk_frames), + "suffix_conv": torch.zeros(batch_size, initial_chunk_frames, 2), + }, + } + } + wrapper.prefix_length = prefix_length + wrapper.initial_chunk_frames = initial_chunk_frames + wrapper.capture_batch_sizes = [1] + wrapper.decoder = type( + "_Decoder", + (), + { + "total_upsample": total_upsample, + "_slice_dynamic_cache": staticmethod(lambda cache, row: cache), + }, + )() + wrapper._ensure_suffix_buffers = lambda caches: None + + caches = {"prefix_frames": 48} + codes = torch.zeros(batch_size, num_quantizers, caches["prefix_frames"] + initial_chunk_frames) + outputs = wrapper._decode_icl_prefix_batch([codes], [caches]) + + assert outputs is not None + output = outputs[0] + assert caches["prefix_frames"] == 48 + assert caches["prefix_pad_frames"] == 24 + assert output.shape == (batch_size, 1, initial_chunk_frames * total_upsample) + assert caches["decoder_prefix_frames"] == prefix_length + + +def test_xvec_dummy_cache_is_initialized_before_graph_capture(): + config = Qwen3TTSTokenizerV2DecoderConfig( + codebook_size=32, + hidden_size=16, + latent_dim=16, + codebook_dim=16, + num_attention_heads=2, + num_key_value_heads=2, + intermediate_size=32, + num_hidden_layers=1, + num_quantizers=2, + decoder_dim=32, + upsample_rates=(8, 5, 4, 3), + upsampling_ratios=(2, 2), + sliding_window=72, + ) + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.prefix_length = 72 + wrapper.decoder = SimpleNamespace(config=config) + + cache = wrapper._make_dummy_xvec_cache(2, torch.device("cpu"), torch.float32)["past_key_values"] + + assert cache.get_seq_length() == 0 + assert all(layer.is_initialized for layer in cache.layers) + assert all(layer.keys.device.type == "cpu" for layer in cache.layers) + + +def test_batched_suffix_graph_updates_ping_pong_buffers(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.prefix_length = 72 + wrapper.codec_chunk_frames = 25 + wrapper._icl_previous_frames_by_target = {26: 1, 51: 26, 76: 51, 97: 72} + wrapper._xvec_previous_frames_by_target = {26: 1, 51: 26, 76: 51, 97: 72} + wrapper.capture_batch_sizes = [1, 2] + key = ("icl", 2, 26) + graph_output = torch.stack( + [ + torch.arange(50, dtype=torch.float32).view(1, -1), + torch.arange(50, dtype=torch.float32).view(1, -1) + 100, + ] + ) + graph_next_quantized = torch.stack([torch.full((2, 26), 1.0), torch.full((2, 26), 2.0)]) + graph_next_conv = torch.stack([torch.full((26, 3), 3.0), torch.full((26, 3), 4.0)]) + wrapper.combined_states = { + key: { + "graph": type("_Graph", (), {"replay": lambda self: None})(), + "input": { + "codes": torch.zeros(2, 2, 25), + "quantized": torch.zeros(2, 2, 1), + "conv": torch.zeros(2, 1, 3), + "mask": torch.ones(2, 98, dtype=torch.bool), + }, + "output": (graph_output, graph_next_quantized, graph_next_conv), + "cache": {}, + } + } + wrapper.decoder = type("_Decoder", (), {"total_upsample": 2})() + monkeypatch.setattr(wrapper, "_copy_request_caches_to_batch", lambda caches, static: None) + + caches = [ + { + "suffix_quantized": torch.zeros(1, 2, 1), + "suffix_conv": torch.zeros(1, 1, 3), + "suffix_frames": 1, + } + for _ in range(2) + ] + codes = [torch.zeros(1, 2, 26) for _ in range(2)] + + outputs = wrapper._decode_suffix_batch("icl", 26, codes, caches, [25, 25]) + + assert outputs is not None + assert [cache["_suffix_buffer_index"] for cache in caches] == [1, 1] + assert caches[0]["suffix_quantized"].data_ptr() == caches[0]["_suffix_quantized_buffers"][1].data_ptr() + assert caches[1]["suffix_conv"].data_ptr() == caches[1]["_suffix_conv_buffers"][1].data_ptr() + torch.testing.assert_close(caches[0]["suffix_quantized"], torch.full((1, 2, 26), 1.0)) + torch.testing.assert_close(caches[1]["suffix_quantized"], torch.full((1, 2, 26), 2.0)) + torch.testing.assert_close(caches[0]["suffix_conv"], torch.full((1, 26, 3), 3.0)) + torch.testing.assert_close(caches[1]["suffix_conv"], torch.full((1, 26, 3), 4.0)) + + +def test_dummy_request_does_not_record_pre_capture_shape_fallback(): + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.prefix_length = 72 + wrapper.initial_chunk_frames = 1 + wrapper.codec_chunk_frames = 25 + wrapper._icl_previous_frames_by_target = {26: 1} + wrapper._xvec_previous_frames_by_target = {26: 1} + wrapper.xvec_prefix_states = {} + wrapper._stats_enabled = True + wrapper._stats_total_requests = 0 + wrapper._stats_hit_requests = 0 + wrapper._stats_fallback_requests = 0 + wrapper._stats_replays = Counter() + wrapper._stats_fallbacks = Counter() + wrapper._maybe_log_stats = lambda: None + wrapper.decoder = _decoder_stub( + _decode_xvec_first_chunk=lambda codes, _cache: codes[:, :1, :], + ) + + wrapper._batched_request_decode( + [torch.zeros(1, 2, 4)], + [{"prefix_frames": 0, "_is_dummy_run": True}], + ) + + assert wrapper._stats_total_requests == 0 + assert wrapper._stats_fallback_requests == 0 + assert wrapper._stats_fallbacks == Counter() + assert wrapper._suppress_stats is False + + +def test_xvec_first_chunk_replays_prefix_graph(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.prefix_length = 72 + wrapper.initial_chunk_frames = 1 + wrapper.codec_chunk_frames = 25 + wrapper.capture_batch_sizes = [1, 2] + wrapper._icl_previous_frames_by_target = {26: 1} + wrapper._xvec_previous_frames_by_target = {26: 1} + replayed = [] + wrapper.xvec_prefix_states = { + 2: { + "graph": SimpleNamespace(replay=lambda: replayed.append(True)), + "input": {"codes": torch.zeros(2, 2, 1)}, + "output": torch.arange(4, dtype=torch.float32).view(2, 1, 2), + "cache": { + "ref_hidden": torch.zeros(2, 2, 0), + "ref_conv": torch.zeros(2, 0, 3), + "prefix_hidden": torch.zeros(2, 0, 3), + "ref_upsample": torch.zeros(2, 3, 0), + "ref_wav": torch.zeros(2, 1, 0), + "suffix_quantized": torch.ones(2, 2, 1), + "suffix_conv": torch.ones(2, 1, 3), + "past_key_values": SimpleNamespace(layers=[]), + }, + } + } + wrapper._record_graph_hit = lambda *_args: None + wrapper._record_graph_fallback = lambda *_args: None + wrapper._ensure_suffix_buffers = lambda cache: None + wrapper.decoder = _decoder_stub( + _is_suffix_cache_rolling=lambda previous, cached: False, + _decode_xvec_first_chunk=lambda *_args: pytest.fail("unexpected eager fallback"), + _slice_dynamic_cache=Qwen3TTSTokenizerV2Decoder._slice_dynamic_cache, + ) + + caches = [{"prefix_frames": 0}, {"prefix_frames": 0}] + outputs = wrapper._batched_request_decode( + [torch.full((1, 2, 1), 3), torch.full((1, 2, 1), 4)], + caches, + ) + + assert replayed == [True] + torch.testing.assert_close(wrapper.xvec_prefix_states[2]["input"]["codes"][0], torch.full((2, 1), 3.0)) + torch.testing.assert_close(wrapper.xvec_prefix_states[2]["input"]["codes"][1], torch.full((2, 1), 4.0)) + assert [cache["decoder_prefix_frames"] for cache in caches] == [0, 0] + assert [cache["suffix_frames"] for cache in caches] == [1, 1] + assert len(outputs) == 2 + + +# ----------------------------------------------------------------------------- +# Routing, phase grouping, fallback, and batch splitting +# ----------------------------------------------------------------------------- + + +def test_eager_fallback_uses_dynamic_decoder_prefix_length(): + decoder = _RecordingDecoder() + wrapper = CUDAGraphDecoderWrapper(decoder, enabled=False) + codes = torch.zeros(1, 2, 2) + + output = wrapper._decode_request_fallback( + codes, + { + "prefix_frames": 48, + "decoder_prefix_frames": 48, + "suffix_quantized": torch.zeros(1, 2, 1), + }, + ) + + assert output is codes + assert decoder.prefix_frames == 48 + + +def test_decode_modes_enforce_cache_contracts(): + decoder = _RecordingDecoder() + async_wrapper = CUDAGraphDecoderWrapper(decoder, async_chunk=True, enabled=False) + stateless_wrapper = CUDAGraphDecoderWrapper(decoder, async_chunk=False, enabled=False) + codes = torch.zeros(1, 2, 1) + + with pytest.raises(ValueError, match="caches are required"): + async_wrapper.chunked_decode_with_cudagraph(codes, caches=None) + with pytest.raises(ValueError, match="caches must be None"): + stateless_wrapper.chunked_decode_with_cudagraph(codes, caches={}) + + +def test_single_request_xvec_first_chunk_uses_batched_graph_router(monkeypatch): + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.async_chunk = True + wrapper.decoder = _decoder_stub(total_upsample=1) + codes = torch.zeros(1, 2, 1) + cache = {"prefix_frames": 0} + calls = [] + + def _batched_request_decode(codes_list, request_caches): + calls.append((codes_list, request_caches)) + request_caches[0]["suffix_quantized"] = codes_list[0] + return [codes_list[0][:, :1]] + + monkeypatch.setattr(wrapper, "_batched_request_decode", _batched_request_decode) + + output = wrapper.chunked_decode_with_cudagraph(codes, caches=cache) + + assert len(calls) == 1 + torch.testing.assert_close(calls[0][0][0], codes) + assert len(calls[0][1]) == 1 + assert calls[0][1][0] is cache + torch.testing.assert_close(output, codes[:, :1]) + + +def test_single_request_suffix_uses_batched_graph_router(monkeypatch): + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.async_chunk = True + wrapper.decoder = _decoder_stub(total_upsample=1) + codes = torch.arange(6, dtype=torch.float32).reshape(1, 2, 3) + cache = { + "prefix_frames": 0, + "suffix_frames": 1, + "suffix_quantized": torch.zeros(1, 2, 1), + } + calls = [] + + def _batched_request_decode(codes_list, request_caches): + calls.append((codes_list, request_caches)) + return [codes_list[0][:, :1]] + + monkeypatch.setattr(wrapper, "_batched_request_decode", _batched_request_decode) + + output = wrapper.chunked_decode_with_cudagraph(codes, caches=cache) + + assert len(calls) == 1 + torch.testing.assert_close(calls[0][0][0], codes) + assert calls[0][1][0] is cache + torch.testing.assert_close(output, codes[:, :1]) + + +def test_stateless_variable_lengths_share_capture_bucket(monkeypatch): + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.stateless_capture_sizes = [25] + wrapper.capture_batch_sizes = [1, 2] + wrapper.stateless_states = {(2, 25): {}} + wrapper.decoder = _decoder_stub(total_upsample=2) + calls = [] + + def _decode_stateless(codes_chunk, *, clone_graph_output): + calls.append((codes_chunk.clone(), clone_graph_output)) + return codes_chunk[:, :1].repeat_interleave(2, dim=-1).float() + + monkeypatch.setattr(wrapper, "_decode_stateless", _decode_stateless) + codes = torch.arange(100).reshape(2, 2, 25) + + output = wrapper._batched_stateless_chunked_decode( + codes, + [25, 12], + chunk_size=25, + left_context_size=0, + max_batch_size=0, + ) + + assert len(calls) == 1 + assert calls[0][0].shape == (2, 2, 25) + assert calls[0][1] is False + torch.testing.assert_close(output[0], codes[0, :1].repeat_interleave(2, dim=-1).float()) + torch.testing.assert_close(output[1], codes[1, :1, :12].repeat_interleave(2, dim=-1).float()) + + +def test_batched_chunked_decode_groups_exact_phases(monkeypatch): + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.async_chunk = True + wrapper.prefix_length = 72 + wrapper.initial_chunk_frames = 1 + wrapper.codec_chunk_frames = 25 + wrapper._icl_previous_frames_by_target = {26: 1, 51: 26, 76: 51, 97: 72} + wrapper._xvec_previous_frames_by_target = {26: 1, 51: 26, 76: 51, 97: 72} + wrapper.decoder = _decoder_stub(_is_suffix_cache_rolling=lambda previous, cached: previous >= 72 and cached == 72) + calls: list[tuple[str, int, int]] = [] + + def _icl_prefix(codes, caches): + calls.append(("icl_prefix", 0, len(codes))) + return [torch.full((1, 1, 1), 1.0) for _ in codes] + + def _suffix(mode, target, codes, caches, new_frames): + calls.append((f"suffix:{mode}", target, len(codes))) + return [torch.full((1, 1, 1), float(target)) for _ in codes] + + monkeypatch.setattr(wrapper, "_decode_icl_prefix_batch", _icl_prefix) + monkeypatch.setattr(wrapper, "_decode_suffix_batch", _suffix) + + class _KV: + @staticmethod + def get_seq_length(): + return 72 + + prefix_frames = 48 + codes_list = [torch.zeros(1, 2, prefix_frames + 1)] + caches = [{"prefix_frames": prefix_frames}] + for previous, _suffix_frames, cached_frames in ( + (1, 26, 1), + (26, 51, 26), + (51, 76, 51), + (76, 97, 72), + ): + codes_list.append(torch.zeros(1, 2, 25)) + caches.append( + { + "prefix_frames": prefix_frames, + "decoder_prefix_frames": 72, + "ref_wav": torch.zeros(1, 1, 1), + "past_key_values": _KV(), + "suffix_frames": previous, + "suffix_quantized": torch.zeros(1, 2, cached_frames), + } + ) + + lengths = [codes.shape[-1] for codes in codes_list] + padded_codes = torch.zeros(len(codes_list), 2, max(lengths)) + for row, codes in enumerate(codes_list): + padded_codes[row, :, : codes.shape[-1]] = codes[0] + outputs = wrapper.batched_chunked_decode_with_cudagraph(padded_codes, lengths, caches=caches) + + assert calls == [ + ("icl_prefix", 0, 1), + ("suffix:icl", 26, 1), + ("suffix:icl", 51, 1), + ("suffix:icl", 76, 1), + ("suffix:icl", 97, 1), + ] + assert [int(output.item()) for output in outputs] == [1, 26, 51, 76, 97] + + +def test_xvec_delta_chunks_route_through_all_reachable_graph_phases(monkeypatch): + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.async_chunk = True + wrapper.prefix_length = 72 + wrapper.initial_chunk_frames = 1 + wrapper.codec_chunk_frames = 25 + wrapper._icl_previous_frames_by_target = {26: 1, 51: 26, 76: 51, 97: 72} + wrapper._xvec_previous_frames_by_target = {26: 1, 51: 26, 76: 51, 97: 72} + wrapper.decoder = _decoder_stub(_is_suffix_cache_rolling=lambda previous, cached: previous >= 72 and cached == 72) + calls: list[tuple[str, int]] = [] + + def _suffix(mode, target, codes, caches, new_frames): + calls.append((mode, target)) + return [torch.zeros(1, 1, 1) for _ in codes] + + monkeypatch.setattr(wrapper, "_decode_suffix_batch", _suffix) + + class _EmptyKV: + @staticmethod + def get_seq_length(): + return 0 + + caches = [ + { + "prefix_frames": 0, + "decoder_prefix_frames": 0, + "ref_wav": torch.zeros(1, 1, 0), + "past_key_values": _EmptyKV(), + "suffix_frames": previous, + "suffix_quantized": torch.zeros(1, 2, cached_frames), + } + for previous, cached_frames in ((1, 1), (26, 26), (51, 51), (76, 72)) + ] + lengths = [25, 25, 25, 25] + codes = torch.zeros(4, 2, 25) + + wrapper.batched_chunked_decode_with_cudagraph(codes, lengths, caches=caches) + + assert calls == [("xvec", 26), ("xvec", 51), ("xvec", 76), ("xvec", 97)] + + +def test_graph_batch_overflow_splits_at_largest_bucket(): + assert CUDAGraphDecoderWrapper._split_for_graph_buckets(16, {1, 2, 4, 8}) == [(0, 8), (8, 16)] + assert CUDAGraphDecoderWrapper._split_for_graph_buckets(10, {1, 2, 4, 8}) == [(0, 8), (8, 10)] + + +def test_missing_graph_phase_falls_back_instead_of_returning_empty(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.icl_prefix_states = {} + wrapper.combined_states = {} + wrapper._record_graph_fallback = lambda *_args: None + + codes = [torch.zeros(1, 2, 25)] + caches = [{}] + + assert wrapper._decode_icl_prefix_batch(codes, caches) is None + assert wrapper._decode_suffix_batch("xvec", 50, codes, caches, [25]) is None + + +def test_xvec_first_chunk_uses_prefixless_initializer(monkeypatch): + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.prefix_length = 72 + wrapper.initial_chunk_frames = 1 + wrapper.codec_chunk_frames = 25 + wrapper._icl_previous_frames_by_target = {26: 1} + wrapper._xvec_previous_frames_by_target = {26: 1} + calls = [] + wrapper.decoder = _decoder_stub( + _decode_xvec_first_chunk=lambda codes, cache: calls.append("xvec") or codes[:, :1, :], + ) + + output = wrapper._batched_request_decode( + [torch.zeros(1, 2, 25)], + [{"prefix_frames": 0}], + ) + + assert calls == ["xvec"] + assert output[0].shape == (1, 1, 25) + + +def test_eager_backend_max_batch_size_splits_each_phase_group(): + decoder = _decoder_stub() + batch_sizes = [] + + def decode_xvec(codes, _caches): + batch_sizes.append(len(codes)) + return [code[:, :1, :].clone() for code in codes] + + outputs = decoder.batched_request_decode( + [torch.zeros(1, 2, 1) for _ in range(5)], + [{"prefix_frames": 0} for _ in range(5)], + prefix_length=72, + initial_chunk_frames=1, + codec_chunk_frames=25, + transitions_by_mode={"icl": {26: 1}, "xvec": {26: 1}}, + decode_icl_prefix_batch=None, + decode_xvec_prefix_batch=decode_xvec, + decode_suffix_batch=None, + decode_fallback=lambda *_args: pytest.fail("unexpected per-request fallback"), + backend_max_batch_size=2, + ) + + assert batch_sizes == [2, 2, 1] + assert len(outputs) == 5 + + +def test_mixed_phase_batch_restores_outputs_and_caches_to_original_slots(): + decoder = _decoder_stub( + _is_suffix_cache_rolling=lambda _previous, _cached: False, + ) + + class _KV: + def __init__(self, length): + self.length = length + + def get_seq_length(self): + return self.length + + codes_list = [ + torch.zeros(1, 2, 1), + torch.zeros(1, 2, 25), + torch.zeros(1, 2, 49), + torch.zeros(1, 2, 25), + ] + caches = [ + {"slot": 0, "prefix_frames": 0}, + { + "slot": 1, + "prefix_frames": 0, + "decoder_prefix_frames": 0, + "past_key_values": _KV(0), + "suffix_frames": 26, + "suffix_quantized": torch.zeros(1, 2, 26), + }, + {"slot": 2, "prefix_frames": 48}, + { + "slot": 3, + "prefix_frames": 48, + "decoder_prefix_frames": 72, + "past_key_values": _KV(72), + "suffix_frames": 1, + "suffix_quantized": torch.zeros(1, 2, 1), + }, + ] + + def _outputs_for_slots(group_caches, backend): + outputs = [] + for cache in group_caches: + cache["backend"] = backend + outputs.append(torch.full((1, 1, 1), float(cache["slot"]))) + return outputs + + outputs = decoder.batched_request_decode( + codes_list, + caches, + prefix_length=72, + initial_chunk_frames=1, + codec_chunk_frames=25, + transitions_by_mode={ + "icl": {26: 1, 51: 26, 76: 51, 97: 72}, + "xvec": {26: 1, 51: 26, 76: 51, 97: 72}, + }, + decode_icl_prefix_batch=lambda _codes, group_caches: _outputs_for_slots(group_caches, "icl_prefix"), + decode_xvec_prefix_batch=lambda _codes, group_caches: _outputs_for_slots(group_caches, "xvec_first"), + decode_suffix_batch=lambda mode, target, _codes, group_caches, _new_frames: _outputs_for_slots( + group_caches, f"suffix:{mode}:{target}" + ), + decode_fallback=lambda *_args: pytest.fail("unexpected eager fallback"), + ) + + assert [int(output.item()) for output in outputs] == [0, 1, 2, 3] + assert [cache["backend"] for cache in caches] == [ + "xvec_first", + "suffix:xvec:51", + "icl_prefix", + "suffix:icl:26", + ] + + +# ----------------------------------------------------------------------------- +# Padding and output trimming regressions +# ----------------------------------------------------------------------------- + + +def test_stateless_replay_pads_to_bucket_and_trims_output(monkeypatch): + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + class _Graph: + def replay(self): + return None + + wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper) + wrapper.async_chunk = False + wrapper.enabled = True + wrapper._warmed_up = True + wrapper.stateless_capture_sizes = [4] + wrapper.stateless_states = { + (1, 4): { + "graph": _Graph(), + "input": {"codes": torch.full((1, 2, 4), -1, dtype=torch.long)}, + "output": torch.arange(4, dtype=torch.float32).view(1, 1, 4), + "cache": None, + } + } + wrapper.decoder = SimpleNamespace(total_upsample=1) + codes = torch.tensor([[[1, 2], [3, 4]]]) + + output = wrapper._decode_stateless(codes, clone_graph_output=True) + + torch.testing.assert_close(output, torch.tensor([[[0.0, 1.0]]])) + torch.testing.assert_close( + wrapper.stateless_states[(1, 4)]["input"]["codes"], + torch.tensor([[[1, 2, 0, 0], [3, 4, 0, 0]]]), + ) + + +def test_stateful_eager_partial_final_chunk_is_right_trimmed(): + torch.manual_seed(10) + with torch.no_grad(): + batched_decoder, per_request_decoder, batched_caches, per_request_caches = _initialize_stateful_batch("xvec") + codes = torch.randint( + 0, + batched_decoder.config.codebook_size, + (2, batched_decoder.config.num_quantizers, 25), + ) + lengths = [25, 12] + batched = batched_decoder.batched_chunked_decode(codes, lengths, caches=batched_caches) + per_request = [ + per_request_decoder.chunked_decode( + codes[row : row + 1, :, :length], + caches=per_request_caches[row], + ) + for row, length in enumerate(lengths) + ] + + full_samples = lengths[0] * batched_decoder.total_upsample + partial_samples = lengths[1] * batched_decoder.total_upsample + assert [output.shape[-1] for output in batched] == [full_samples, partial_samples] + torch.testing.assert_close(batched[0], per_request[0][0], atol=1e-5, rtol=1e-4) + torch.testing.assert_close(batched[1], per_request[1][0], atol=1e-5, rtol=1e-4) diff --git a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_tokenizer_mask_cache.py b/tests/model_executor/models/qwen3_tts/test_qwen3_tts_tokenizer_mask_cache.py deleted file mode 100644 index aeb538adc..000000000 --- a/tests/model_executor/models/qwen3_tts/test_qwen3_tts_tokenizer_mask_cache.py +++ /dev/null @@ -1,125 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -import inspect - -import pytest -import torch -from transformers.masking_utils import create_sliding_window_causal_mask - -from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import ( - Qwen3TTSTokenizerV2DecoderConfig, -) -from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( - Qwen3TTSTokenizerV2DecoderTransformerModel, -) - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def _make_config(attention_implementation: str) -> Qwen3TTSTokenizerV2DecoderConfig: - config = Qwen3TTSTokenizerV2DecoderConfig( - hidden_size=32, - latent_dim=16, - max_position_embeddings=512, - num_attention_heads=4, - num_key_value_heads=4, - intermediate_size=64, - num_hidden_layers=1, - num_quantizers=2, - decoder_dim=32, - upsample_rates=(2,), - upsampling_ratios=(2,), - ) - config._attn_implementation = attention_implementation - return config - - -def _make_model( - attention_implementation: str, - dtype: torch.dtype = torch.float32, -) -> Qwen3TTSTokenizerV2DecoderTransformerModel: - model = Qwen3TTSTokenizerV2DecoderTransformerModel._from_config(_make_config(attention_implementation)) - return model.to(dtype=dtype).eval() - - -def _create_reference_mask( - model: Qwen3TTSTokenizerV2DecoderTransformerModel, - inputs_embeds: torch.Tensor, -) -> torch.Tensor | None: - sequence_length = inputs_embeds.shape[1] - mask_kwargs = { - "config": model.config, - "attention_mask": None, - "past_key_values": None, - "position_ids": None, - } - sig = inspect.signature(create_sliding_window_causal_mask) - if "input_embeds" in sig.parameters: - mask_kwargs["input_embeds"] = inputs_embeds[:1] - mask_kwargs["cache_position"] = torch.arange(sequence_length) - else: - mask_kwargs["inputs_embeds"] = inputs_embeds[:1] - return create_sliding_window_causal_mask(**mask_kwargs) - - -@pytest.mark.parametrize("sequence_length", [17, 72, 73, 169, 325]) -@pytest.mark.parametrize( - ("attention_implementation", "dtype"), - [("sdpa", torch.float32), ("eager", torch.float32), ("eager", torch.float16)], -) -def test_cached_sliding_mask_matches_transformers( - sequence_length: int, - attention_implementation: str, - dtype: torch.dtype, -): - model = _make_model(attention_implementation, dtype) - inputs_embeds = torch.empty(2, sequence_length, model.config.hidden_size, dtype=dtype) - - actual = model._get_sliding_attention_mask(inputs_embeds) - expected = _create_reference_mask(model, inputs_embeds) - - if expected is None: - assert actual is None - assert model._get_sliding_attention_mask(inputs_embeds) is None - return - - assert actual is not None - torch.testing.assert_close(actual, expected, atol=0, rtol=0) - assert actual.shape == (1, 1, sequence_length, sequence_length) - assert actual.is_contiguous() - assert model._get_sliding_attention_mask(inputs_embeds) is actual - - -def test_cached_eager_mask_tracks_compute_dtype(): - model = _make_model("eager") - - fp32_inputs = torch.empty(1, 73, model.config.hidden_size) - fp32_mask = model._get_sliding_attention_mask(fp32_inputs) - assert fp32_mask is not None - assert fp32_mask.dtype == torch.float32 - - model.half() - fp16_inputs = fp32_inputs.half() - fp16_mask = model._get_sliding_attention_mask(fp16_inputs) - assert fp16_mask is not None - assert fp16_mask.dtype == torch.float16 - assert torch.isfinite(fp16_mask).all() - assert fp16_mask.min() == torch.finfo(torch.float16).min - - -@pytest.mark.parametrize( - ("kwargs", "message"), - [ - ({"use_cache": True}, "use_cache=True"), - ({"past_key_values": object()}, "past_key_values"), - ({"cache_position": torch.arange(1, 18)}, "zero-based cache_position"), - ({"position_ids": torch.arange(1, 18).unsqueeze(0)}, "zero-based position_ids"), - ], -) -def test_cached_mask_rejects_unsupported_cache_and_positions(kwargs: dict, message: str): - model = _make_model("sdpa") - inputs_embeds = torch.randn(1, 17, model.config.latent_dim) - - with pytest.raises(ValueError, match=message): - model(inputs_embeds=inputs_embeds, **kwargs) diff --git a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py index 96a938513..c1a6041c7 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py @@ -788,10 +788,16 @@ def test_qwen3_tts_code2wav_forward_decodes_connector_payload() -> None: self.last_codes = codes.detach().clone() return codes.sum(dim=1).to(torch.float32) + def batched_chunked_decode(self, codes, lengths, caches, **kwargs): + assert lengths == [codes.shape[-1]] + assert caches is None + return self.chunked_decode(codes, **kwargs) + decoder = _Decoder() model = Qwen3TTSCode2Wav.__new__(Qwen3TTSCode2Wav) torch.nn.Module.__init__(model) model.decoder = decoder + model._async_chunk = False model._num_quantizers = 2 model._total_upsample = 1 model._output_sample_rate = 24000 diff --git a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py index c41541cc8..013c0b99e 100644 --- a/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py @@ -156,8 +156,9 @@ def test_streaming_phases(config, n_frames, finished, expected): else: exp_ctx, exp_window = expected assert payload is not None - assert payload.meta.left_context_size == exp_ctx - assert len(payload.codes.audio) == _Q * exp_window + assert payload.meta.left_context_size == 0 + expected_delta = exp_window if finished else exp_window - exp_ctx + assert len(payload.codes.audio) == _Q * expected_delta def test_dynamic_ic_adapts_to_load(): @@ -202,7 +203,8 @@ def test_ic_load_change_mid_request(): assert _call(tm, "r", n_frames=25) is None p3 = _call(tm, "r", n_frames=27) assert p3 is not None - assert p3.meta.left_context_size == 2 + assert p3.meta.left_context_size == 0 + assert len(p3.codes.audio) == _Q * 25 # A *new* request under high load gets IC=16 (not IC=2). # Frame 2 would emit under IC=2 but must hold under IC=16. @@ -227,7 +229,8 @@ def test_connector_initial_chunk_config_overrides_dynamic_ic(): assert _call(tm, "r", n_frames=25) is None p2 = _call(tm, "r", n_frames=29) assert p2 is not None - assert p2.meta.left_context_size == 4 + assert p2.meta.left_context_size == 0 + assert len(p2.codes.audio) == _Q * 25 @pytest.mark.parametrize( @@ -282,8 +285,8 @@ def test_first_streaming_chunk_prepends_ref_code_context(): assert len(payload.codes.audio) == _Q * 12 -def test_followup_ref_code_context_is_sent_as_metadata_handle(): - """Follow-up chunks keep full ref context semantically without resending it.""" +def test_followup_sends_only_codec_delta_without_ref_metadata(): + """Follow-up chunks rely on decoder state and resend neither reference codes nor metadata.""" tm = _tm() rid = "r-ref2" tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(35)] @@ -299,13 +302,11 @@ def test_followup_ref_code_context_is_sent_as_metadata_handle(): ) assert payload is not None - # ref_code (2 frames) is represented in metadata so Code2Wav can restore it - # from its request-local cache. It must not be resent in codes.audio. - assert payload.meta.left_context_size == 10 + 2 - assert payload.meta.ref_context_size == 2 - assert payload.meta.ref_context_request_id == rid - assert payload.meta.ref_context_included is False - assert len(payload.codes.audio) == _Q * 35 + assert payload.meta.left_context_size == 0 + assert payload.meta.ref_context_size is None + assert payload.meta.ref_context_request_id is None + assert payload.meta.ref_context_included is None + assert len(payload.codes.audio) == _Q * 25 def test_streaming_ref_code_context_is_bounded_for_batchable_shapes(): @@ -332,8 +333,8 @@ def test_streaming_ref_code_context_is_bounded_for_batchable_shapes(): ) assert payload is not None - assert payload.meta.left_context_size == 3 + 3 - assert len(payload.codes.audio) == _Q * (3 + 3 + 4) + assert payload.meta.left_context_size == 3 + assert len(payload.codes.audio) == _Q * (3 + 4) frames = payload.codes.audio.reshape(_Q, -1).transpose(0, 1) torch.testing.assert_close(frames[:3], ref_code[-3:]) @@ -700,36 +701,40 @@ class TestChunkRampEmission: assert self._emit(tm, rid, 4) is None p1 = self._emit(tm, rid, 5) assert p1 is not None - assert len(p1.codes.audio) == _Q * 5 - assert p1.meta.left_context_size == 1 + assert len(p1.codes.audio) == _Q * 4 + assert p1.meta.left_context_size == 0 tm.ramp_chunk_count[rid] = 2 for n in range(6, 13): assert self._emit(tm, rid, n) is None p2 = self._emit(tm, rid, 13) assert p2 is not None - assert p2.meta.left_context_size == 5 + assert p2.meta.left_context_size == 0 + assert len(p2.codes.audio) == _Q * 8 tm.ramp_chunk_count[rid] = 3 for n in range(14, 29): assert self._emit(tm, rid, n) is None p3 = self._emit(tm, rid, 29) assert p3 is not None - assert p3.meta.left_context_size == 13 + assert p3.meta.left_context_size == 0 + assert len(p3.codes.audio) == _Q * 16 tm.ramp_chunk_count[rid] = 4 for n in range(30, 54): assert self._emit(tm, rid, n) is None p4 = self._emit(tm, rid, 54) assert p4 is not None - assert p4.meta.left_context_size == 25 + assert p4.meta.left_context_size == 0 + assert len(p4.codes.audio) == _Q * 25 tm.ramp_chunk_count[rid] = 5 for n in range(55, 79): assert self._emit(tm, rid, n) is None p5 = self._emit(tm, rid, 79) assert p5 is not None - assert p5.meta.left_context_size == 25 + assert p5.meta.left_context_size == 0 + assert len(p5.codes.audio) == _Q * 25 def test_ramp_finished_flush_mid_chunk(self): tm = _tm(chunk_ramp=self.RAMP) @@ -742,8 +747,8 @@ class TestChunkRampEmission: p_fin = self._emit(tm, rid, 3, finished=True) assert p_fin is not None assert p_fin.meta.finished.item() is True - assert len(p_fin.codes.audio) == _Q * (1 + 2) - assert p_fin.meta.left_context_size == 1 + assert len(p_fin.codes.audio) == _Q * 2 + assert p_fin.meta.left_context_size == 0 def test_ramp_finished_no_new_frames(self): tm = _tm(chunk_ramp=self.RAMP) @@ -783,7 +788,8 @@ class TestChunkRampEmission: assert self._emit(tm, rid, n) is None p1 = self._emit(tm, rid, 26) assert p1 is not None - assert p1.meta.left_context_size == 1 + assert p1.meta.left_context_size == 0 + assert len(p1.codes.audio) == _Q * 25 def test_ramp_with_ref_code_first_chunk(self): tm = _tm(chunk_ramp=self.RAMP, left_context=25) @@ -820,14 +826,16 @@ class TestChunkRampEmission: assert self._emit(tm, rid, n) is None p2 = self._emit(tm, rid, 30) assert p2 is not None - assert p2.meta.left_context_size == 5 + assert p2.meta.left_context_size == 0 + assert len(p2.codes.audio) == _Q * 25 tm.ramp_chunk_count[rid] = 3 for n in range(31, 55): assert self._emit(tm, rid, n) is None p3 = self._emit(tm, rid, 55) assert p3 is not None - assert p3.meta.left_context_size == 25 + assert p3.meta.left_context_size == 0 + assert len(p3.codes.audio) == _Q * 25 def test_ramp_profile_4_4_8_16_25(self): """Ramp [4,4,8,16,25]: chunk 0=4 frames (320ms audio covers chunk 1 @@ -846,35 +854,40 @@ class TestChunkRampEmission: assert self._emit(tm, rid, n) is None p1 = self._emit(tm, rid, 8) assert p1 is not None - assert p1.meta.left_context_size == 4 + assert p1.meta.left_context_size == 0 + assert len(p1.codes.audio) == _Q * 4 tm.ramp_chunk_count[rid] = 2 for n in range(9, 16): assert self._emit(tm, rid, n) is None p2 = self._emit(tm, rid, 16) assert p2 is not None - assert p2.meta.left_context_size == 8 + assert p2.meta.left_context_size == 0 + assert len(p2.codes.audio) == _Q * 8 tm.ramp_chunk_count[rid] = 3 for n in range(17, 32): assert self._emit(tm, rid, n) is None p3 = self._emit(tm, rid, 32) assert p3 is not None - assert p3.meta.left_context_size == 16 + assert p3.meta.left_context_size == 0 + assert len(p3.codes.audio) == _Q * 16 tm.ramp_chunk_count[rid] = 4 for n in range(33, 57): assert self._emit(tm, rid, n) is None p4 = self._emit(tm, rid, 57) assert p4 is not None - assert p4.meta.left_context_size == 25 + assert p4.meta.left_context_size == 0 + assert len(p4.codes.audio) == _Q * 25 tm.ramp_chunk_count[rid] = 5 for n in range(58, 82): assert self._emit(tm, rid, n) is None p5 = self._emit(tm, rid, 82) assert p5 is not None - assert p5.meta.left_context_size == 25 + assert p5.meta.left_context_size == 0 + assert len(p5.codes.audio) == _Q * 25 def test_ramp_resets_on_segment_boundary(self): """After segment 1 emits chunks, a segment boundary clears diff --git a/vllm_omni/deploy/qwen3_tts.yaml b/vllm_omni/deploy/qwen3_tts.yaml index ad154bed2..08531869a 100644 --- a/vllm_omni/deploy/qwen3_tts.yaml +++ b/vllm_omni/deploy/qwen3_tts.yaml @@ -33,11 +33,19 @@ connectors: # initial_codec_chunk_frames (including the per-request API field). # Example: # codec_chunk_ramp: [4, 4, 8, 16, 25] - # Common Stage1 decode buckets: - # no-ref first/steady chunks: 25 / 97 frames - # Base ref-context first/steady chunks: 73 / 169 frames - # decoder internal non-streaming chunks: 325 frames - decode_cudagraph_capture_sizes: [25, 73, 97, 169, 325] + # Decoder CUDA Graph sizing: + # * async_chunk=true derives all stateful ICL/xvec shapes from + # codec_chunk_ramp when enabled, otherwise from + # initial_codec_chunk_frames, codec_chunk_frames, and sliding_window; + # decode_cudagraph_capture_sizes is ignored in this mode. + # * async_chunk=false captures stateless shapes 150 and + # decode_chunk_frames + decode_left_context_frames (325 by default). + # Extra stateless shapes may be added (union, not intersection) with: + # decode_cudagraph_capture_sizes: [97, 400] + # Graph batch size defaults to 1. Add larger buckets only when needed: + # decode_cudagraph_batch_sizes: [1, 2, 4] + # Limit eager batches to reduce temporary stateful cache allocations; CUDA Graph batching uses the buckets above. + decode_batch_max_size: 4 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml index 25bac72c0..ff086beb5 100644 --- a/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml +++ b/vllm_omni/deploy/qwen3_tts_high_concurrency.yaml @@ -38,16 +38,12 @@ connectors: # initial_codec_chunk_frames (including the per-request API field). # Example: # codec_chunk_ramp: [4, 4, 8, 16, 25] - # Common Stage1 decode buckets: - # no-ref first/steady chunks: 25 / 97 frames - # Base ref-context first/steady chunks: 73 / 169 frames - # decoder internal non-streaming chunks: 325 frames - decode_cudagraph_capture_sizes: [25, 73, 97, 169, 325] - # Keep B>1 captures opt-in; c64 e2e validation did not show a stable win. + # async_chunk derives stateful ICL/xvec graph shapes from the ramp when + # enabled, otherwise from the fixed chunk settings above, and ignores + # decode_cudagraph_capture_sizes. Keep B>1 + # captures opt-in; c64 e2e validation did not show a stable win. decode_cudagraph_batch_sizes: [1] - decode_compile_shapes: [] decode_batch_max_size: 1 - decode_batch_bucket_frames: [] decode_enable_tf32: false stages: diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py index 6a5866c7e..87df87889 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py @@ -1,7 +1,7 @@ from __future__ import annotations import os -from collections import Counter, OrderedDict +from collections import Counter from collections.abc import Iterable from typing import Any @@ -14,6 +14,7 @@ from vllm.model_executor.model_loader import DefaultModelLoader from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import ( Qwen3TTSTokenizerV2Config, @@ -24,8 +25,7 @@ from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import ( logger = init_logger(__name__) -_REF_CONTEXT_CACHE_MAX_ENTRIES = 4096 -_REF_CONTEXT_CACHE_MAX_BYTES = 64 * 1024 * 1024 +_DUMMY_REQUEST_ID = "__qwen3_tts_dummy_run__" def _codec_ids_from_payload_or_input( @@ -56,10 +56,16 @@ class Qwen3TTSCode2Wav(nn.Module): input_modalities = "audio" + # Ask the model runner for the scheduler-side request IDs. Stateful + # decoder caches must use the same IDs delivered by on_requests_finished; + # payload metadata carries an external ID which may differ. + requires_request_ids = True + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.vllm_config = vllm_config self.model_path = vllm_config.model_config.model + self._async_chunk = bool(getattr(vllm_config.model_config, "async_chunk", False)) self.have_multimodal_outputs = True self.has_preprocess = False @@ -69,9 +75,7 @@ class Qwen3TTSCode2Wav(nn.Module): self._decode_chunk_frames = 300 self._decode_left_context_frames = 25 - self._decode_batch_bucket_frames: list[int] = [] self._decode_batch_max_size = 0 - self._decode_variable_chunk_batch_min_frames = self._decode_chunk_frames + self._decode_left_context_frames + 1 self._logged_codec_stats = False self._logged_malformed_codec_lengths: set[tuple[int, int]] = set() self._batch_stats_enabled = os.environ.get("VLLM_OMNI_QWEN3_CODE2WAV_BATCH_STATS", "").lower() in ( @@ -103,10 +107,8 @@ class Qwen3TTSCode2Wav(nn.Module): self._output_sample_rate = int(tok_config.output_sample_rate) self._total_upsample = int(self.decoder.total_upsample) self._decoder_sliding_window = int(getattr(dec_config, "sliding_window", 0) or 0) - self._ref_context_cache: OrderedDict[str, torch.Tensor] = OrderedDict() - self._ref_context_cache_bytes = 0 - self._ref_context_cache_max_entries = _REF_CONTEXT_CACHE_MAX_ENTRIES - self._ref_context_cache_max_bytes = _REF_CONTEXT_CACHE_MAX_BYTES + self._decoder_state_cache: dict[str, dict[str, Any]] = {} + self._decoder_state_cache_warn_entries = 512 def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: # This stage ignores token embeddings. Keep a stable dummy embedding for vLLM runner. @@ -114,6 +116,13 @@ class Qwen3TTSCode2Wav(nn.Module): return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32) return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) + def get_dummy_runtime_additional_information(self, num_reqs: int) -> list[dict[str, Any]]: + """Provide request state metadata for vLLM's external graph dummy run.""" + return [ + {"meta": {"request_id": f"{_DUMMY_REQUEST_ID}{index}", "finished": True}} + for index in range(max(0, int(num_reqs))) + ] + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> None: return None @@ -145,10 +154,10 @@ class Qwen3TTSCode2Wav(nn.Module): device: torch.device, codec_chunk_frames: int, codec_left_context_frames: int, - decode_cudagraph_capture_sizes: list[int] | None, + initial_codec_chunk_frames: int, + codec_chunk_ramp: list[int] | None, decode_cudagraph_batch_sizes: list[int] | None, - decode_cudagraph_extra_capture_shapes: list[tuple[int, int]] | None, - decode_compile_shapes: list[tuple[int, int]] | None, + decode_cudagraph_capture_sizes: list[int] | None, ) -> None: """Enable inner Code2Wav CUDA graph unless stage is enforce_eager.""" if not hasattr(self.decoder, "enable_cudagraph") or device.type != "cuda": @@ -160,7 +169,8 @@ class Qwen3TTSCode2Wav(nn.Module): return if ( - codec_chunk_frames > 0 + self._async_chunk + and codec_chunk_frames > 0 and codec_left_context_frames > 0 and self._decoder_sliding_window and codec_left_context_frames < self._decoder_sliding_window @@ -177,24 +187,19 @@ class Qwen3TTSCode2Wav(nn.Module): ) self.decoder.enable_cudagraph( - capture_sizes=decode_cudagraph_capture_sizes, capture_batch_sizes=decode_cudagraph_batch_sizes, - extra_capture_shapes=decode_cudagraph_extra_capture_shapes, - compile_shapes=decode_compile_shapes, + stateless_capture_sizes=decode_cudagraph_capture_sizes, device=device, codec_chunk_frames=codec_chunk_frames, codec_left_context_frames=codec_left_context_frames, + initial_codec_chunk_frames=initial_codec_chunk_frames, + codec_chunk_ramp=codec_chunk_ramp, + async_chunk=self._async_chunk, decode_chunk_size=self._decode_chunk_frames, decode_left_context=self._decode_left_context_frames, ) logger.info("Code2Wav decoder CUDA Graph enabled") - def _get_decode_batch_bucket_frames(self, actual_frames: int) -> int: - for bucket_frames in self._decode_batch_bucket_frames: - if actual_frames <= bucket_frames: - return bucket_frames - return actual_frames - def _record_decode_batch_stats( self, *, @@ -212,48 +217,9 @@ class Qwen3TTSCode2Wav(nn.Module): self._batch_stats_actual_frames.update(actual_frames) self._batch_stats_bucket_groups[(group_size, bucket_frames)] += 1 - @staticmethod - def _tensor_nbytes(tensor: torch.Tensor) -> int: - return int(tensor.numel() * tensor.element_size()) - - def _evict_ref_context_cache_if_needed(self) -> None: - evicted = 0 - while len(self._ref_context_cache) > self._ref_context_cache_max_entries: - _, cached = self._ref_context_cache.popitem(last=False) - self._ref_context_cache_bytes -= self._tensor_nbytes(cached) - evicted += 1 - while self._ref_context_cache_bytes > self._ref_context_cache_max_bytes and len(self._ref_context_cache) > 1: - _, cached = self._ref_context_cache.popitem(last=False) - self._ref_context_cache_bytes -= self._tensor_nbytes(cached) - evicted += 1 - if evicted: - logger.debug( - "Evicted %d Qwen3-TTS ref context cache entries; entries=%d bytes=%d", - evicted, - len(self._ref_context_cache), - self._ref_context_cache_bytes, - ) - - def _cache_ref_context(self, request_id: str, tensor: torch.Tensor) -> None: - previous = self._ref_context_cache.pop(request_id, None) - if previous is not None: - self._ref_context_cache_bytes -= self._tensor_nbytes(previous) - cached = tensor.detach().contiguous() - self._ref_context_cache[request_id] = cached - self._ref_context_cache.move_to_end(request_id) - self._ref_context_cache_bytes += self._tensor_nbytes(cached) - self._evict_ref_context_cache_if_needed() - - def _get_ref_context(self, request_id: str) -> torch.Tensor | None: - cached = self._ref_context_cache.get(request_id) - if cached is not None: - self._ref_context_cache.move_to_end(request_id) - return cached - - def _pop_ref_context(self, request_id: str) -> None: - cached = self._ref_context_cache.pop(request_id, None) - if cached is not None: - self._ref_context_cache_bytes -= self._tensor_nbytes(cached) + def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: + for req_id in finished_req_ids: + self._decoder_state_cache.pop(req_id, None) def log_decode_batch_stats(self) -> None: if not self._batch_stats_enabled or self._batch_stats_requests == 0: @@ -299,7 +265,6 @@ class Qwen3TTSCode2Wav(nn.Module): self._batch_stats_forwards += 1 decoder = self.decoder q = int(self._num_quantizers) - upsample = int(self._total_upsample) sr_val = int(self._output_sample_rate) sr_tensor = torch.tensor(sr_val, dtype=torch.int32) empty = torch.zeros((0,), dtype=torch.float32) @@ -314,11 +279,13 @@ class Qwen3TTSCode2Wav(nn.Module): ids = input_ids.reshape(-1).to(dtype=torch.long) request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts")) - parsed: list[tuple[int, int]] = [] - valid_codes_qf: list[torch.Tensor] = [] + internal_request_ids = kwargs.get("request_ids") + valid_codes_qf: list[tuple[str | None, torch.Tensor]] = [] valid_indices: list[int] = [] left_context_size = [0] * len(request_ids_list) ref_context_size = [0] * len(request_ids_list) + segment_finished_flags = [False] * len(request_ids_list) + request_state_ids: list[str | None] = [None] * len(request_ids_list) ref_context_request_ids: list[str | None] = [None] * len(request_ids_list) ref_context_included = [False] * len(request_ids_list) finished_flags = [False] * len(request_ids_list) @@ -346,28 +313,42 @@ class Qwen3TTSCode2Wav(nn.Module): if runtime_infos: for i, info in enumerate(runtime_infos): - if i >= len(left_context_size): + if i >= len(ref_context_size): break if not isinstance(info, dict): continue meta = info.get("meta", {}) + if "is_segment_finished" in meta: + segment_finished_flags[i] = _meta_bool(meta["is_segment_finished"]) if "left_context_size" in meta: left_context_size[i] = _meta_int(meta["left_context_size"]) if "ref_context_size" in meta: ref_context_size[i] = _meta_int(meta["ref_context_size"]) if "ref_context_request_id" in meta: ref_context_request_ids[i] = _meta_str(meta["ref_context_request_id"]) + if "request_id" in meta: + request_state_ids[i] = _meta_str(meta["request_id"]) if "ref_context_included" in meta: ref_context_included[i] = _meta_bool(meta["ref_context_included"]) if "finished" in meta: finished_flags[i] = _meta_bool(meta["finished"]) + + # Normal runner calls provide scheduler-side IDs, which are also used + # by scheduler_output.finished_req_ids. Direct forward calls and CUDA + # Graph dummy runs have no runner IDs, so retain the payload ID as a + # fallback for those paths. + cache_request_ids = request_state_ids.copy() + if internal_request_ids is not None: + for i, request_id in enumerate(internal_request_ids): + if i >= len(cache_request_ids): + break + cache_request_ids[i] = str(request_id) + for i, req_ids in enumerate(request_ids_list): runtime_info = runtime_infos[i] if i < len(runtime_infos) else None req_ids = _codec_ids_from_payload_or_input(req_ids, runtime_info) if req_ids.numel() < 1: - parsed.append((0, 0)) continue - ctx_frames = left_context_size[i] ref_ctx_frames = ref_context_size[i] flat = req_ids n = flat.numel() @@ -382,36 +363,54 @@ class Qwen3TTSCode2Wav(nn.Module): n, q, ) - parsed.append((0, 0)) continue frames = n // q # [q*F] -> [Q, F] for direct decoder call (decoder expects [B, Q, F]) codes_qf = flat.reshape(q, frames) ref_req_id = ref_context_request_ids[i] - if ref_req_id is not None and ref_ctx_frames > 0: - if ref_context_included[i]: - if frames < ref_ctx_frames: - raise ValueError( - "Qwen3-TTS ref context metadata says ref prefix is included, " - f"but frames={frames} < ref_context_size={ref_ctx_frames}" + state_req_id = cache_request_ids[i] if self._async_chunk else None + is_new_state = state_req_id is not None and state_req_id not in self._decoder_state_cache + if is_new_state and ref_req_id is not None and ref_ctx_frames > 0: + if not ref_context_included[i] or frames < ref_ctx_frames: + raise ValueError("Qwen3-TTS async_chunk first ICL chunk must include its declared reference prefix") + valid_codes_qf.append((state_req_id, codes_qf)) + if state_req_id is not None: + state = self._decoder_state_cache.get(state_req_id) + if state is None: + if len(self._decoder_state_cache) >= self._decoder_state_cache_warn_entries: + logger.warning_once( + "Qwen3-TTS decoder state cache exceeded the expected active-request envelope: " + "entries=%d threshold=%d. Keeping all active states; check request cleanup paths.", + len(self._decoder_state_cache), + self._decoder_state_cache_warn_entries, ) - self._cache_ref_context(ref_req_id, codes_qf[:, :ref_ctx_frames]) - else: - cached_ref = self._get_ref_context(ref_req_id) - if cached_ref is None: + state = {} + self._decoder_state_cache[state_req_id] = state + state.setdefault("prefix_frames", 0) + if state_req_id.startswith(_DUMMY_REQUEST_ID): + state["_is_dummy_run"] = True + if ref_ctx_frames > 0: + cached_prefix_frames = state.setdefault("prefix_frames", ref_ctx_frames) + if cached_prefix_frames == 0: + state["prefix_frames"] = ref_ctx_frames + cached_prefix_frames = ref_ctx_frames + if cached_prefix_frames != ref_ctx_frames: raise ValueError( - "Missing Qwen3-TTS ref context cache for " - f"request {ref_req_id!r}; first chunk must include ref_code" + "Qwen3-TTS ref context size changed within request " + f"{ref_req_id!r}: cached={cached_prefix_frames}, current={ref_ctx_frames}" ) - cached_ref = cached_ref.to(device=codes_qf.device, dtype=codes_qf.dtype) - codes_qf = torch.cat((cached_ref, codes_qf), dim=1) - frames = int(codes_qf.shape[1]) - parsed.append((ctx_frames, frames)) - valid_codes_qf.append(codes_qf) valid_indices.append(i) num_req = len(request_ids_list) if not valid_codes_qf: + for req_id, finished, segment_finished in zip( + cache_request_ids, + finished_flags, + segment_finished_flags, + strict=False, + ): + if req_id is not None and (finished or segment_finished): + self._decoder_state_cache.pop(req_id, None) return OmniOutput( text_hidden_states=None, multimodal_outputs={ @@ -423,7 +422,7 @@ class Qwen3TTSCode2Wav(nn.Module): if not self._logged_codec_stats: self._logged_codec_stats = True try: - c = valid_codes_qf[0] + _, c = valid_codes_qf[0] logger.info( "Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] batch=%d", c.shape[1], @@ -436,92 +435,54 @@ class Qwen3TTSCode2Wav(nn.Module): except Exception: pass - wav_tensors: list[torch.Tensor | None] = [None] * len(valid_codes_qf) - - def _decode_group_chunks(group_chunks: list[list[tuple[int, torch.Tensor]]]) -> None: - for group_chunk in group_chunks: - actual_frames = [int(codes_qf.shape[1]) for _, codes_qf in group_chunk] - target_frames = max(actual_frames) - is_equal_length_batch = all(frames == target_frames for frames in actual_frames) - use_variable_length_batch = ( - len(group_chunk) > 1 - and not is_equal_length_batch - and target_frames >= self._decode_variable_chunk_batch_min_frames - and hasattr(decoder, "batched_chunked_decode") + request_states: list[dict[str, Any]] | None = None + if self._async_chunk: + missing_state = [index for index, (state_req_id, _) in enumerate(valid_codes_qf) if state_req_id is None] + if missing_state: + raise ValueError( + "Qwen3-TTS async_chunk Code2Wav inputs require a request_id; " + f"missing state for batch indices {missing_state}" ) - if len(group_chunk) == 1: - codes_bqf = group_chunk[0][1].unsqueeze(0) - elif is_equal_length_batch: - codes_bqf = torch.stack([codes_qf for _, codes_qf in group_chunk], dim=0) - else: - first = group_chunk[0][1] - codes_bqf = first.new_zeros((len(group_chunk), q, target_frames)) - for row, (_, codes_qf) in enumerate(group_chunk): - codes_bqf[row, :, : codes_qf.shape[1]] = codes_qf - self._record_decode_batch_stats( - group_size=len(group_chunk), - bucket_frames=target_frames, - actual_frames=actual_frames, - ) - try: - if use_variable_length_batch: - wav_batch = decoder.batched_chunked_decode( - codes_bqf, - actual_frames, - chunk_size=self._decode_chunk_frames, - left_context_size=self._decode_left_context_frames, - max_batch_size=self._decode_batch_max_size, - ) # [B, 1, wav_len] - else: - wav_batch = decoder.chunked_decode( - codes_bqf, - chunk_size=self._decode_chunk_frames, - left_context_size=self._decode_left_context_frames, - ) # [B, 1, wav_len] - except TypeError: - # Unit-test fakes and older decoder shims may not accept the - # explicit chunk kwargs; production Qwen3-TTS decoders do. - wav_batch = decoder.chunked_decode(codes_bqf) # [B, 1, wav_len] + request_states = [ + self._decoder_state_cache[state_req_id] + for state_req_id, _ in valid_codes_qf + if state_req_id is not None + ] - if wav_batch.dim() == 3 and wav_batch.shape[1] == 1: - wav_rows = wav_batch[:, 0, :] - elif wav_batch.dim() == 2: - wav_rows = wav_batch - else: - raise ValueError( - "Code2Wav decoder returned unexpected shape " - f"{tuple(wav_batch.shape)} for batch size {len(group_chunk)}" - ) - if wav_rows.shape[0] != len(group_chunk): - raise ValueError( - f"Code2Wav decoder returned batch size {wav_rows.shape[0]} " - f"for input batch size {len(group_chunk)}" - ) - for row, (j, _) in enumerate(group_chunk): - wav_tensors[j] = wav_rows[row] + request_lengths = [int(codes_qf.shape[-1]) for _, codes_qf in valid_codes_qf] + max_request_length = max(request_lengths) + request_codes = valid_codes_qf[0][1].new_zeros((len(valid_codes_qf), q, max_request_length)) + for row, (_, codes_qf) in enumerate(valid_codes_qf): + request_codes[row, :, : codes_qf.shape[-1]].copy_(codes_qf) - # Group by configured frame buckets instead of only exact lengths. - # For ordinary async streaming windows this is the real batching - # opportunity; decoder-internal variable chunk batching is gated to - # longer inputs where repeated full chunks can amortize its overhead. - grouped_codes: dict[int, list[tuple[int, torch.Tensor]]] = {} - for j, codes_qf in enumerate(valid_codes_qf): - frames = int(codes_qf.shape[1]) - grouped_codes.setdefault(self._get_decode_batch_bucket_frames(frames), []).append((j, codes_qf)) - - for _bucket_frames, group in grouped_codes.items(): - if self._decode_batch_max_size > 0 and len(group) > self._decode_batch_max_size: - # Keep each decoder call inside the configured CUDA graph batch - # envelope. Sorting by length lowers right-padding within each - # split while outputs are restored by original request index. - group = sorted(group, key=lambda item: int(item[1].shape[1])) - group_chunks = [ - group[start : start + self._decode_batch_max_size] - for start in range(0, len(group), self._decode_batch_max_size) - ] - else: - group_chunks = [group] - _decode_group_chunks(group_chunks) + self._record_decode_batch_stats( + group_size=len(valid_codes_qf), + bucket_frames=max_request_length, + actual_frames=request_lengths, + ) + request_wavs = decoder.batched_chunked_decode( + request_codes, + request_lengths, + caches=request_states, + chunk_size=self._decode_chunk_frames, + left_context_size=self._decode_left_context_frames, + max_batch_size=self._decode_batch_max_size, + ) + if len(request_wavs) != len(valid_codes_qf): + raise ValueError( + f"Qwen3-TTS batched decoder returned {len(request_wavs)} outputs for {len(valid_codes_qf)} requests" + ) + wav_tensors: list[torch.Tensor] = [] + for row in range(len(valid_codes_qf)): + wav = request_wavs[row] + if wav.dim() == 2 and wav.shape[0] == 1: + wav = wav[0] + elif wav.dim() != 1: + raise ValueError(f"Qwen3-TTS batched decoder returned unexpected row shape {tuple(wav.shape)}") + if request_states is None: + start = left_context_size[valid_indices[row]] * self._total_upsample + wav = wav[start:] + wav_tensors.append(wav) if self._batch_stats_log_every > 0 and self._batch_stats_forwards % self._batch_stats_log_every == 0: self.log_decode_batch_stats() @@ -530,27 +491,22 @@ class Qwen3TTSCode2Wav(nn.Module): srs = [sr_tensor] * num_req for j, idx in enumerate(valid_indices): - ctx_frames, actual_frames = parsed[idx] wav = wav_tensors[j] assert wav is not None - # Slice on exact codec-frame boundaries instead of proportionally. - start = max(0, ctx_frames * upsample) - end = max(start, actual_frames * upsample) - if start >= wav.shape[0]: - logger.warning( - "Context trim start %d >= decoded length %d; returning empty audio.", - start, - wav.shape[0], - ) + if wav.numel() == 0: continue - wav = wav[start : min(end, wav.shape[0])] if wav.shape[0] > 0: # Decoder already runs in fp32, so the .to(float32) is a redundant dispatch. audios[idx] = (wav if wav.dtype == torch.float32 else wav.to(torch.float32)).reshape(-1) - for req_id, finished in zip(ref_context_request_ids, finished_flags, strict=False): - if req_id is not None and finished: - self._pop_ref_context(req_id) + for req_id, finished, segment_finished in zip( + cache_request_ids, + finished_flags, + segment_finished_flags, + strict=False, + ): + if req_id is not None and (finished or segment_finished): + self._decoder_state_cache.pop(req_id, None) return OmniOutput( text_hidden_states=None, @@ -707,6 +663,8 @@ class Qwen3TTSCode2Wav(nn.Module): if isinstance(extra_cfg, dict): codec_chunk_frames = int(extra_cfg.get("codec_chunk_frames") or 0) codec_left_context_frames = int(extra_cfg.get("codec_left_context_frames") or 0) + initial_codec_chunk_frames = int(extra_cfg.get("initial_codec_chunk_frames") or 1) + codec_chunk_ramp = parse_chunk_ramp(extra_cfg, steady=codec_chunk_frames) if self._async_chunk else None decode_chunk_frames = _get_int_config("decode_chunk_frames", self._decode_chunk_frames) decode_left_context_frames = _get_int_config( "decode_left_context_frames", @@ -720,33 +678,22 @@ class Qwen3TTSCode2Wav(nn.Module): ) self._decode_chunk_frames = decode_chunk_frames self._decode_left_context_frames = decode_left_context_frames - decode_cudagraph_capture_sizes = _get_int_list_config("decode_cudagraph_capture_sizes") decode_cudagraph_batch_sizes = _get_int_list_config("decode_cudagraph_batch_sizes") - decode_cudagraph_extra_capture_shapes = _get_int_pair_list_config("decode_cudagraph_extra_capture_shapes") - decode_compile_shapes = _get_int_pair_list_config("decode_compile_shapes") - decode_batch_bucket_frames = _get_int_list_config("decode_batch_bucket_frames") - if decode_batch_bucket_frames is not None: - self._decode_batch_bucket_frames = decode_batch_bucket_frames + decode_cudagraph_capture_sizes = ( + None if self._async_chunk else _get_int_list_config("decode_cudagraph_capture_sizes") + ) decode_batch_max_size = _get_int_config("decode_batch_max_size", self._decode_batch_max_size) if decode_batch_max_size < 0: raise ValueError(f"Invalid Qwen3-TTS Code2Wav config decode_batch_max_size={decode_batch_max_size}") self._decode_batch_max_size = decode_batch_max_size - decode_variable_chunk_batch_min_frames = _get_int_config( - "decode_variable_chunk_batch_min_frames", - self._decode_variable_chunk_batch_min_frames, - ) - if decode_variable_chunk_batch_min_frames < 0: - raise ValueError( - "Invalid Qwen3-TTS Code2Wav config " - f"decode_variable_chunk_batch_min_frames={decode_variable_chunk_batch_min_frames}" - ) - self._decode_variable_chunk_batch_min_frames = decode_variable_chunk_batch_min_frames decode_enable_tf32 = _get_bool_config("decode_enable_tf32", False) else: - decode_cudagraph_capture_sizes = None + codec_chunk_frames = 0 + codec_left_context_frames = 0 + initial_codec_chunk_frames = 1 + codec_chunk_ramp = None decode_cudagraph_batch_sizes = None - decode_cudagraph_extra_capture_shapes = None - decode_compile_shapes = None + decode_cudagraph_capture_sizes = None decode_enable_tf32 = False if decode_enable_tf32 and device.type == "cuda": @@ -764,16 +711,20 @@ class Qwen3TTSCode2Wav(nn.Module): torch.get_float32_matmul_precision(), ) + self.decoder._initial_codec_chunk_frames = initial_codec_chunk_frames + self.decoder._incremental_chunk_frames = codec_chunk_frames or 25 + self.decoder._incremental_chunk_ramp = list(codec_chunk_ramp or ()) + if hasattr(self.decoder, "enable_cudagraph") and device.type == "cuda": try: self._maybe_enable_decoder_cudagraph( device=device, codec_chunk_frames=codec_chunk_frames, codec_left_context_frames=codec_left_context_frames, - decode_cudagraph_capture_sizes=decode_cudagraph_capture_sizes, + initial_codec_chunk_frames=initial_codec_chunk_frames, + codec_chunk_ramp=codec_chunk_ramp, decode_cudagraph_batch_sizes=decode_cudagraph_batch_sizes, - decode_cudagraph_extra_capture_shapes=decode_cudagraph_extra_capture_shapes, - decode_compile_shapes=decode_compile_shapes, + decode_cudagraph_capture_sizes=decode_cudagraph_capture_sizes, ) except Exception: logger.warning( diff --git a/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py b/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py new file mode 100644 index 000000000..a1eb071ee --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py @@ -0,0 +1,1097 @@ +# Copyright 2026 The Alibaba Qwen team. +# SPDX-License-Identifier: Apache-2.0 +""" +CUDA Graph wrapper for Qwen3TTSTokenizerV2Decoder. + +This module provides CUDA Graph acceleration for the speech tokenizer decoder, +reducing kernel launch overhead during inference. +""" + +import bisect +import os +from collections import Counter + +import torch +from torch.cuda import CUDAGraph +from transformers.cache_utils import DynamicCache +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + + +class CUDAGraphDecoderWrapper: + """ + CUDA Graph wrapper for Qwen3TTSTokenizerV2Decoder. + + This wrapper captures the decoder forward pass for fixed input sizes + and replays them during inference to reduce kernel launch overhead. + + Usage: + wrapper = CUDAGraphDecoderWrapper(decoder) + wrapper.warmup(device) + + # During inference: + output = wrapper.chunked_decode_with_cudagraph(codes, caches) # Uses CUDA graph if possible + """ + + def __init__( + self, + decoder: torch.nn.Module, + capture_batch_sizes: list[int] | None = None, + stateless_capture_sizes: list[int] | None = None, + num_quantizers: int = 8, + enabled: bool = True, + async_chunk: bool = True, + initial_chunk_frames: int = 1, + codec_chunk_frames: int = 25, + codec_chunk_ramp: list[int] | None = None, + decode_chunk_size: int = 300, + decode_left_context: int = 25, + ): + self.decoder = decoder + self.capture_batch_sizes = sorted(set(capture_batch_sizes or [1])) + self.num_quantizers = num_quantizers + self.enabled = enabled + self.async_chunk = bool(async_chunk) + self._warmed_up = False + + self.combined_states: dict[tuple[str, int, int], dict] = {} + self.icl_prefix_states: dict[int, dict] = {} + self.xvec_prefix_states: dict[int, dict] = {} + self.stateless_states: dict[tuple[int, int], dict] = {} + + self._device = None + self.prefix_length = int(getattr(self.decoder.config, "sliding_window", 0) or 0) + self.codec_chunk_frames = int(codec_chunk_frames) + self.codec_chunk_ramp = [int(size) for size in codec_chunk_ramp or ()] + self.initial_chunk_frames = self.codec_chunk_ramp[0] if self.codec_chunk_ramp else int(initial_chunk_frames) + self.decode_chunk_size = int(decode_chunk_size) + self.decode_left_context = int(decode_left_context) + if self.prefix_length <= 2: + raise ValueError(f"decoder sliding_window must be greater than 2, got {self.prefix_length}") + if ( + self.initial_chunk_frames <= 0 + or self.codec_chunk_frames <= 0 + or any(size <= 0 for size in self.codec_chunk_ramp) + ): + raise ValueError( + "initial_chunk_frames and codec_chunk_frames must be positive, " + f"got {self.initial_chunk_frames} and {self.codec_chunk_frames}" + ) + self._icl_previous_frames_by_target = self._derive_icl_suffix_transitions() + self._xvec_previous_frames_by_target = self._derive_xvec_suffix_transitions() + self.icl_capture_sizes = sorted(self._icl_previous_frames_by_target) + self.stateless_capture_sizes = self._derive_stateless_capture_sizes(stateless_capture_sizes) + self._stats_enabled = os.environ.get( + "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS", + "", + ).lower() in ("1", "true", "yes", "on") + self._stats_log_every = int( + os.environ.get( + "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS_LOG_EVERY", + "100", + ) + or 0 + ) + self._stats_file = os.environ.get( + "VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS_FILE", + "", + ) + self._stats_total_requests = 0 + self._stats_hit_requests = 0 + self._stats_fallback_requests = 0 + self._stats_replays: Counter[tuple[str, int, int]] = Counter() + self._stats_fallbacks: Counter[tuple[str, int]] = Counter() + + def _derive_icl_suffix_transitions(self) -> dict[int, int]: + return self._derive_suffix_transitions() + + def _derive_xvec_suffix_transitions(self) -> dict[int, int]: + return self._derive_suffix_transitions() + + def _derive_suffix_transitions(self) -> dict[int, int]: + transitions: dict[int, int] = {} + previous = self.initial_chunk_frames + for new_frames in self.codec_chunk_ramp[1:]: + if previous >= self.prefix_length: + break + target = previous + new_frames + transitions[target] = previous + previous = target + while previous < self.prefix_length: + target = previous + self.codec_chunk_frames + transitions[target] = previous + previous = target + transitions[self.prefix_length + self.codec_chunk_frames] = self.prefix_length + return transitions + + def _derive_stateless_capture_sizes(self, configured_sizes: list[int] | None) -> list[int]: + max_frames = self.decode_chunk_size + self.decode_left_context + sizes = {150, max_frames} + sizes.update(int(size) for size in configured_sizes or () if int(size) > 0) + return sorted(sizes) + + def _transitions_for_mode(self, mode: str) -> dict[int, int]: + return self._xvec_previous_frames_by_target if mode == "xvec" else self._icl_previous_frames_by_target + + def _cached_conv_frames(self, previous_frames: int) -> int: + return min(previous_frames, self.prefix_length - 2) + + def _record_graph_hit(self, phase: str, batch_size: int, request_count: int) -> None: + if ( + not getattr(self, "_stats_enabled", False) + or getattr(self, "_suppress_stats", False) + or torch.cuda.is_current_stream_capturing() + ): + return + self._stats_total_requests += request_count + self._stats_hit_requests += request_count + self._stats_replays[(phase, batch_size, request_count)] += 1 + self._maybe_log_stats() + + def _record_graph_fallback(self, phase: str, request_count: int) -> None: + if ( + not getattr(self, "_stats_enabled", False) + or getattr(self, "_suppress_stats", False) + or torch.cuda.is_current_stream_capturing() + ): + return + self._stats_total_requests += request_count + self._stats_fallback_requests += request_count + self._stats_fallbacks[(phase, request_count)] += 1 + self._maybe_log_stats() + + def _maybe_log_stats(self) -> None: + if self._stats_log_every > 0 and self._stats_total_requests % self._stats_log_every == 0: + self.log_decode_stats() + + def log_decode_stats(self) -> None: + if not getattr(self, "_stats_enabled", False) or self._stats_total_requests == 0: + return + hit_rate = 100.0 * self._stats_hit_requests / self._stats_total_requests + message = ( + "Segmented Code2Wav CUDA Graph stats: " + f"requests={self._stats_total_requests} hits={self._stats_hit_requests} " + f"fallbacks={self._stats_fallback_requests} hit_rate={hit_rate:.2f}% " + f"replays={self._stats_replays.most_common(20)} " + f"fallback_groups={self._stats_fallbacks.most_common(20)}" + ) + logger.warning("%s", message) + stats_file = getattr(self, "_stats_file", "") + if stats_file: + fd = os.open(stats_file, os.O_APPEND | os.O_CREAT | os.O_WRONLY, 0o644) + try: + os.write(fd, f"pid={os.getpid()} {message}\n".encode()) + finally: + os.close(fd) + + def _get_icl_capture_shapes(self) -> list[tuple[int, int]]: + shapes = {(batch_size, size) for batch_size in self.capture_batch_sizes for size in self.icl_capture_sizes} + return sorted(shapes) + + def _make_dummy_icl_cache( + self, + batch_size: int, + device: torch.device, + model_dtype: torch.dtype, + ) -> dict: + config = self.decoder.config + dummy_hidden = torch.zeros( + batch_size, + config.codebook_dim, + self.prefix_length, + device=device, + dtype=model_dtype, + ) + dummy_conv = torch.zeros( + batch_size, + self.prefix_length, + config.latent_dim, + device=device, + dtype=model_dtype, + ) + dummy_kv = DynamicCache(config=config) + assert self.prefix_length == config.sliding_window + head_dim = getattr( + config, + "head_dim", + config.hidden_size // config.num_attention_heads, + ) + prefix_kv_shape = (batch_size, config.num_key_value_heads, self.prefix_length, head_dim) + + prefix_keys = torch.zeros(prefix_kv_shape, device=device, dtype=model_dtype) + prefix_values = torch.zeros_like(prefix_keys) + + for layer_idx in range(config.num_hidden_layers): + dummy_kv.update(prefix_keys, prefix_values, layer_idx) + + expected_kv_length = self.prefix_length - 1 + if dummy_kv.get_seq_length() != self.prefix_length or any( + layer.keys.shape[-2] != expected_kv_length for layer in dummy_kv.layers + ): + raise RuntimeError( + "Dummy sliding-window cache does not match the captured prefix " + f"state: expected logical length {self.prefix_length} and " + f"physical KV length {expected_kv_length}" + ) + + dummy_prefix_hidden = torch.zeros( + batch_size, + self.prefix_length, + config.latent_dim, + device=device, + dtype=model_dtype, + ) + return { + "ref_hidden": dummy_hidden, + "ref_conv": dummy_conv, + "past_key_values": dummy_kv, + "prefix_hidden": dummy_prefix_hidden, + } + + def _make_dummy_xvec_cache( + self, + batch_size: int, + device: torch.device, + model_dtype: torch.dtype, + ) -> dict: + config = self.decoder.config + head_dim = getattr( + config, + "head_dim", + config.hidden_size // config.num_attention_heads, + ) + empty_prefix_cache = DynamicCache(config=config) + empty_prefix_cache.early_initialization( + batch_size=batch_size, + num_heads=config.num_key_value_heads, + head_dim=head_dim, + dtype=model_dtype, + device=device, + ) + return { + "ref_hidden": torch.zeros(batch_size, config.codebook_dim, 0, device=device, dtype=model_dtype), + "ref_conv": torch.zeros(batch_size, 0, config.latent_dim, device=device, dtype=model_dtype), + "past_key_values": empty_prefix_cache, + "prefix_hidden": torch.zeros(batch_size, 0, config.latent_dim, device=device, dtype=model_dtype), + } + + def warmup( + self, + device: torch.device, + dtype: torch.dtype = torch.long, + ): + if device.type != "cuda" or not self.enabled or self._warmed_up: + return + + self._device = device + self.decoder.eval() + + self.capture_batch_sizes = [bs for bs in self.capture_batch_sizes if bs > 0] + if not self.capture_batch_sizes: + self.capture_batch_sizes = [1] + + if not self.async_chunk: + capture_shapes = [ + (batch_size, size) for batch_size in self.capture_batch_sizes for size in self.stateless_capture_sizes + ] + logger.info( + "Starting stateless CUDA Graph warmup for %d shapes: batch_sizes=%s seq_lens=%s", + len(capture_shapes), + self.capture_batch_sizes, + self.stateless_capture_sizes, + ) + for batch_size, size in capture_shapes: + try: + self._capture_stateless(batch_size, size, device, dtype) + logger.info("Captured stateless CUDA Graph for batch=%d size=%d", batch_size, size) + except Exception: + logger.warning( + "Failed to capture stateless CUDA Graph for batch=%d size=%d", + batch_size, + size, + exc_info=True, + ) + self._warmed_up = bool(self.stateless_states) + return + + if not self.icl_capture_sizes: + raise ValueError("No reachable segmented capture sizes were configured") + icl_capture_shapes = self._get_icl_capture_shapes() + + for batch_size in self.capture_batch_sizes: + try: + self._capture_icl_prefix(batch_size, device, dtype) + logger.info("Captured ICL prefix CUDA Graph for batch=%d", batch_size) + except Exception: + logger.warning("Failed to capture ICL prefix CUDA Graph for batch=%d", batch_size, exc_info=True) + try: + self._capture_xvec_prefix(batch_size, device, dtype) + logger.info("Captured xvec prefix CUDA Graph for batch=%d", batch_size) + except Exception: + logger.warning("Failed to capture xvec prefix CUDA Graph for batch=%d", batch_size, exc_info=True) + + logger.info( + "Starting CUDA Graph warmup for %d shapes: batch_sizes=%s seq_lens=%s", + len(icl_capture_shapes), + self.capture_batch_sizes, + self.icl_capture_sizes, + ) + + for batch_size, size in icl_capture_shapes: + try: + caches = self._make_dummy_icl_cache(batch_size, device, next(self.decoder.parameters()).dtype) + self._capture_combined_suffix("icl", batch_size, size, caches, device, dtype) + logger.info("Captured combined suffix CUDA Graph for batch=%d size=%d", batch_size, size) + except Exception: + logger.warning( + "Failed to capture combined suffix graph for batch=%d size=%d", + batch_size, + size, + exc_info=True, + ) + + model_dtype = next(self.decoder.parameters()).dtype + for batch_size in self.capture_batch_sizes: + for size, previous_frames in self._xvec_previous_frames_by_target.items(): + try: + caches = self._make_dummy_xvec_cache(batch_size, device, model_dtype) + self._capture_combined_suffix("xvec", batch_size, size, caches, device, dtype) + logger.info("Captured xvec suffix CUDA Graph for batch=%d size=%d", batch_size, size) + except Exception: + logger.warning( + "Failed to capture xvec suffix graph for batch=%d size=%d", + batch_size, + size, + exc_info=True, + ) + + self._warmed_up = bool(self.icl_prefix_states) or bool(self.xvec_prefix_states) or bool(self.combined_states) + + def _capture_stateless( + self, + batch_size: int, + size: int, + device: torch.device, + dtype: torch.dtype, + ) -> None: + key = (batch_size, size) + static_input = torch.zeros(batch_size, self.num_quantizers, size, dtype=dtype, device=device) + with torch.no_grad(): + _ = self.decoder._forward_exact(static_input) + torch.accelerator.synchronize(device) + + graph = CUDAGraph() + with torch.no_grad(): + with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + static_output = self.decoder._forward_exact(static_input) + + self.stateless_states[key] = { + "graph": graph, + "input": {"codes": static_input}, + "output": static_output, + "cache": None, + } + + def _run_combined_suffix( + self, + mode: str, + codes: torch.Tensor, + old_quantized: torch.Tensor, + old_conv: torch.Tensor, + caches: dict, + target_frames: int, + attention_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + previous_frames = self._transitions_for_mode(mode)[target_frames] + new_frames = target_frames - previous_frames + rolling = self.decoder._is_suffix_cache_rolling(previous_frames, int(old_quantized.shape[-1])) + return self.decoder._decode_suffix( + codes, + old_quantized, + old_conv, + caches, + self.prefix_length if mode == "icl" else 0, + new_frames, + rolling, + attention_mask=attention_mask, + ) + + def _capture_combined_suffix( + self, + mode: str, + batch_size: int, + target_frames: int, + caches: dict, + device: torch.device, + code_dtype: torch.dtype, + ) -> None: + previous_frames = self._transitions_for_mode(mode)[target_frames] + new_frames = target_frames - previous_frames + model_dtype = next(self.decoder.parameters()).dtype + key = (mode, batch_size, target_frames) + static_codes = torch.zeros(batch_size, self.num_quantizers, new_frames, dtype=code_dtype, device=device) + static_quantized = torch.zeros( + batch_size, + self.decoder.config.codebook_dim, + min(previous_frames, self.prefix_length), + dtype=model_dtype, + device=device, + ) + static_conv = torch.zeros( + batch_size, + self._cached_conv_frames(previous_frames), + self.decoder.config.latent_dim, + dtype=model_dtype, + device=device, + ) + prefix_frames = self.prefix_length if mode == "icl" else 0 + static_mask = torch.ones(batch_size, prefix_frames + target_frames, dtype=torch.bool, device=device) + + with torch.no_grad(): + _ = self._run_combined_suffix( + mode, + static_codes, + static_quantized, + static_conv, + caches, + target_frames, + static_mask, + ) + torch.accelerator.synchronize(device) + + graph = CUDAGraph() + with torch.no_grad(): + with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + outputs = self._run_combined_suffix( + mode, + static_codes, + static_quantized, + static_conv, + caches, + target_frames, + static_mask, + ) + + self.combined_states[key] = { + "graph": graph, + "input": { + "codes": static_codes, + "quantized": static_quantized, + "conv": static_conv, + "mask": static_mask, + }, + "output": outputs, + "cache": caches, + } + + def _capture_icl_prefix(self, batch_size: int, device: torch.device, dtype: torch.dtype): + size = self.prefix_length + self.initial_chunk_frames + caches: dict = {} + static_input = torch.zeros(batch_size, self.num_quantizers, size, dtype=dtype, device=device) + with torch.no_grad(): + _ = self.decoder._decode_icl_first_chunk(static_input, caches, self.prefix_length) + torch.accelerator.synchronize(device) + + config = self.decoder.config + model_dtype = next(self.decoder.parameters()).dtype + head_dim = getattr( + config, + "head_dim", + config.hidden_size // config.num_attention_heads, + ) + prefix_cache = DynamicCache(config=config) + prefix_cache.early_initialization( + batch_size=batch_size, + num_heads=config.num_key_value_heads, + head_dim=head_dim, + dtype=model_dtype, + device=device, + ) + prefix_hidden_mask = torch.ones(batch_size, 1, size, dtype=model_dtype, device=device) + prefix_attention_mask = torch.ones(batch_size, self.prefix_length, dtype=torch.bool, device=device) + suffix_attention_mask = torch.ones(batch_size, size, dtype=torch.bool, device=device) + + graph = CUDAGraph() + with torch.no_grad(): + with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + static_output = self.decoder._decode_icl_first_chunk( + static_input, + caches, + self.prefix_length, + prefix_cache=prefix_cache, + prefix_hidden_mask=prefix_hidden_mask, + prefix_attention_mask=prefix_attention_mask, + suffix_attention_mask=suffix_attention_mask, + ) + + self.icl_prefix_states[batch_size] = { + "graph": graph, + "input": { + "codes": static_input, + "hidden_mask": prefix_hidden_mask, + "prefix_attention_mask": prefix_attention_mask, + "suffix_attention_mask": suffix_attention_mask, + }, + "output": static_output, + "cache": caches, + } + + def _capture_xvec_prefix(self, batch_size: int, device: torch.device, dtype: torch.dtype) -> None: + static_input = torch.zeros( + batch_size, + self.num_quantizers, + self.initial_chunk_frames, + dtype=dtype, + device=device, + ) + with torch.no_grad(): + _ = self.decoder._decode_xvec_first_chunk(static_input, {}) + torch.accelerator.synchronize(device) + + config = self.decoder.config + model_dtype = next(self.decoder.parameters()).dtype + head_dim = getattr( + config, + "head_dim", + config.hidden_size // config.num_attention_heads, + ) + prefix_cache = DynamicCache(config=config) + prefix_cache.early_initialization( + batch_size=batch_size, + num_heads=config.num_key_value_heads, + head_dim=head_dim, + dtype=model_dtype, + device=device, + ) + caches: dict = {"past_key_values": prefix_cache} + graph = CUDAGraph() + with torch.no_grad(): + with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + static_output = self.decoder._decode_xvec_first_chunk(static_input, caches) + + self.xvec_prefix_states[batch_size] = { + "graph": graph, + "input": {"codes": static_input}, + "output": static_output, + "cache": caches, + } + + def _ensure_suffix_buffers(self, caches: dict) -> None: + if "_suffix_quantized_buffers" in caches: + return + + suffix_quantized = caches["suffix_quantized"] + suffix_conv = caches["suffix_conv"] + quantized_capacity = self.prefix_length + conv_capacity = self.prefix_length - 2 + quantized_buffers = [ + suffix_quantized.new_zeros((*suffix_quantized.shape[:-1], quantized_capacity)) for _ in range(2) + ] + conv_buffers = [ + suffix_conv.new_zeros((suffix_conv.shape[0], conv_capacity, suffix_conv.shape[-1])) for _ in range(2) + ] + quantized_valid = min(int(suffix_quantized.shape[-1]), quantized_capacity) + conv_valid = min(int(suffix_conv.shape[1]), conv_capacity) + quantized_buffers[0][:, :, :quantized_valid].copy_(suffix_quantized[:, :, -quantized_valid:]) + conv_buffers[0][:, :conv_valid, :].copy_(suffix_conv[:, -conv_valid:, :]) + caches["_suffix_quantized_buffers"] = quantized_buffers + caches["_suffix_conv_buffers"] = conv_buffers + caches["_suffix_buffer_index"] = 0 + caches.setdefault("suffix_frames", int(suffix_quantized.shape[-1])) + caches["suffix_quantized"] = quantized_buffers[0][:, :, :quantized_valid] + caches["suffix_conv"] = conv_buffers[0][:, :conv_valid, :] + + def _copy_request_caches_to_batch(self, request_caches: list[dict], static_caches: dict) -> None: + tensor_cache_keys = ("ref_hidden", "ref_conv", "prefix_hidden") + for row, request_cache in enumerate(request_caches): + for key in tensor_cache_keys: + static_caches[key][row : row + 1].copy_(request_cache[key]) + + for static_layer, request_layer in zip( + static_caches["past_key_values"].layers, + request_cache["past_key_values"].layers, + strict=True, + ): + if request_layer.keys is None or request_layer.values is None: + if static_layer.keys.numel() == 0 and static_layer.values.numel() == 0: + continue + raise RuntimeError("Uninitialized request KV cache cannot populate a non-empty graph cache") + static_layer.keys[row : row + 1].copy_(request_layer.keys) + static_layer.values[row : row + 1].copy_(request_layer.values) + + def _select_batch_bucket(self, request_count: int, available: set[int]) -> int | None: + return next((size for size in self.capture_batch_sizes if size >= request_count and size in available), None) + + @staticmethod + def _split_for_graph_buckets(request_count: int, available: set[int]) -> list[tuple[int, int]]: + if request_count <= 0 or not available: + return [] + largest = max(available) + ranges: list[tuple[int, int]] = [] + start = 0 + while request_count - start > largest: + ranges.append((start, start + largest)) + start += largest + ranges.append((start, request_count)) + return ranges + + def _decode_icl_prefix_batch( + self, + codes_list: list[torch.Tensor], + request_caches: list[dict], + ) -> list[torch.Tensor] | None: + available = set(self.icl_prefix_states) + if torch.cuda.is_current_stream_capturing(): + self._record_graph_fallback("icl_prefix:capture", len(codes_list)) + return None + if not available: + self._record_graph_fallback("icl_prefix:no_graph", len(codes_list)) + return None + if len(codes_list) > max(available, default=0): + outputs: list[torch.Tensor] = [] + for start, end in self._split_for_graph_buckets(len(codes_list), available): + chunk_outputs = self._decode_icl_prefix_batch( + codes_list[start:end], + request_caches[start:end], + ) + if chunk_outputs is None: + return None + outputs.extend(chunk_outputs) + return outputs + + batch_size = self._select_batch_bucket(len(codes_list), available) + if batch_size is None: + self._record_graph_fallback("icl_prefix:no_graph", len(codes_list)) + return None + + state = self.icl_prefix_states[batch_size] + static_inputs = state["input"] + static_input = static_inputs["codes"] + hidden_mask = static_inputs["hidden_mask"] + prefix_mask = static_inputs["prefix_attention_mask"] + suffix_mask = static_inputs["suffix_attention_mask"] + static_input.zero_() + hidden_mask.fill_(1) + prefix_mask.fill_(1) + suffix_mask.fill_(1) + + prefix_pads: list[int] = [] + for row, (codes, cache) in enumerate(zip(codes_list, request_caches, strict=True)): + prefix_frames = int(cache["prefix_frames"]) + suffix_frames = int(codes.shape[-1]) - prefix_frames + if not 0 < prefix_frames <= self.prefix_length or suffix_frames != self.initial_chunk_frames: + return None + prefix_pad = self.prefix_length - prefix_frames + prefix_pads.append(prefix_pad) + static_input[row, :, prefix_pad : self.prefix_length].copy_(codes[0, :, :prefix_frames]) + static_input[row, :, self.prefix_length :].copy_(codes[0, :, prefix_frames:]) + hidden_mask[row, :, :prefix_pad] = 0 + prefix_mask[row, :prefix_pad] = 0 + suffix_mask[row, :prefix_pad] = 0 + + state["graph"].replay() + self._record_graph_hit("icl_prefix", batch_size, len(codes_list)) + static_caches = state["cache"] + outputs: list[torch.Tensor] = [] + tensor_keys = ("ref_hidden", "ref_conv", "prefix_hidden") + for row, (cache, prefix_pad) in enumerate(zip(request_caches, prefix_pads, strict=True)): + for key in tensor_keys: + cache[key] = static_caches[key][row : row + 1].clone() + cache["past_key_values"] = self.decoder._slice_dynamic_cache(static_caches["past_key_values"], row) + cache["decoder_prefix_frames"] = self.prefix_length + cache["prefix_pad_frames"] = prefix_pad + cache["suffix_quantized"] = static_caches["suffix_quantized"][row : row + 1].clone() + cache["suffix_conv"] = static_caches["suffix_conv"][row : row + 1].clone() + cache["suffix_frames"] = self.initial_chunk_frames + outputs.append( + state["output"][row : row + 1, :, -self.initial_chunk_frames * self.decoder.total_upsample :].clone() + ) + return outputs + + def _decode_xvec_prefix_batch( + self, + codes_list: list[torch.Tensor], + request_caches: list[dict], + ) -> list[torch.Tensor] | None: + available = set(getattr(self, "xvec_prefix_states", {})) + if not available: + self._record_graph_fallback("xvec_prefix:no_graph", len(codes_list)) + return None + if torch.cuda.is_current_stream_capturing(): + self._record_graph_fallback("xvec_prefix:capture", len(codes_list)) + return None + if any(int(codes.shape[-1]) != self.initial_chunk_frames for codes in codes_list): + self._record_graph_fallback("xvec_prefix:shape", len(codes_list)) + return None + if len(codes_list) > max(available, default=0): + outputs: list[torch.Tensor] = [] + for start, end in self._split_for_graph_buckets(len(codes_list), available): + chunk_outputs = self._decode_xvec_prefix_batch( + codes_list[start:end], + request_caches[start:end], + ) + if chunk_outputs is None: + return None + outputs.extend(chunk_outputs) + return outputs + + batch_size = self._select_batch_bucket(len(codes_list), available) + if batch_size is None: + self._record_graph_fallback("xvec_prefix:no_graph", len(codes_list)) + return None + + state = self.xvec_prefix_states[batch_size] + static_input = state["input"]["codes"] + static_input.zero_() + for row, codes in enumerate(codes_list): + static_input[row].copy_(codes[0]) + + state["graph"].replay() + self._record_graph_hit("xvec_prefix", batch_size, len(codes_list)) + static_caches = state["cache"] + outputs: list[torch.Tensor] = [] + tensor_keys = ("ref_hidden", "ref_conv", "prefix_hidden") + for row, cache in enumerate(request_caches): + for key in tensor_keys: + cache[key] = static_caches[key][row : row + 1].clone() + cache["past_key_values"] = self.decoder._slice_dynamic_cache(static_caches["past_key_values"], row) + cache["decoder_prefix_frames"] = 0 + cache["suffix_quantized"] = static_caches["suffix_quantized"][row : row + 1].clone() + cache["suffix_conv"] = static_caches["suffix_conv"][row : row + 1].clone() + cache["suffix_frames"] = self.initial_chunk_frames + self._ensure_suffix_buffers(cache) + outputs.append(state["output"][row : row + 1].clone()) + return outputs + + def _decode_suffix_batch( + self, + mode: str, + target_frames: int, + codes_list: list[torch.Tensor], + request_caches: list[dict], + new_frames_list: list[int], + ) -> list[torch.Tensor] | None: + available = { + batch + for graph_mode, batch, target in self.combined_states + if graph_mode == mode and target == target_frames + } + if torch.cuda.is_current_stream_capturing(): + self._record_graph_fallback(f"suffix_{target_frames}:capture", len(codes_list)) + return None + if not available: + self._record_graph_fallback(f"suffix_{target_frames}:no_graph", len(codes_list)) + return None + previous_frames = self._transitions_for_mode(mode).get(target_frames) + if previous_frames is None: + return None + expected_new_frames = target_frames - previous_frames + if any(not 0 < new_frames <= expected_new_frames for new_frames in new_frames_list): + self._record_graph_fallback(f"suffix_{target_frames}:shape", len(codes_list)) + return None + if len(codes_list) > max(available, default=0): + outputs: list[torch.Tensor] = [] + for start, end in self._split_for_graph_buckets(len(codes_list), available): + chunk_outputs = self._decode_suffix_batch( + mode, + target_frames, + codes_list[start:end], + request_caches[start:end], + new_frames_list[start:end], + ) + if chunk_outputs is None: + return None + outputs.extend(chunk_outputs) + return outputs + + batch_size = self._select_batch_bucket(len(codes_list), available) + if batch_size is None: + self._record_graph_fallback(f"suffix_{target_frames}:no_graph", len(codes_list)) + return None + key = (mode, batch_size, target_frames) + state = self.combined_states[key] + static_inputs = state["input"] + static_codes = static_inputs["codes"] + static_quantized = static_inputs["quantized"] + static_conv = static_inputs["conv"] + static_mask = static_inputs["mask"] + static_caches = state["cache"] + static_codes.zero_() + static_mask.fill_(1) + + for row, (codes, cache, new_frames) in enumerate(zip(codes_list, request_caches, new_frames_list, strict=True)): + self._ensure_suffix_buffers(cache) + static_codes[row, :, :new_frames].copy_(codes[0, :, -new_frames:]) + static_quantized[row : row + 1].copy_(cache["suffix_quantized"]) + static_conv[row : row + 1].copy_(cache["suffix_conv"]) + if mode == "icl": + static_mask[row, : int(cache.get("prefix_pad_frames", 0))] = 0 + + self._copy_request_caches_to_batch(request_caches, static_caches) + state["graph"].replay() + self._record_graph_hit(f"suffix_{target_frames}", batch_size, len(codes_list)) + graph_output, graph_next_quantized, graph_next_conv = state["output"] + outputs: list[torch.Tensor] = [] + for row, (cache, new_frames) in enumerate(zip(request_caches, new_frames_list, strict=True)): + actual_suffix_frames = previous_frames + new_frames + output = self._trim_replay_output( + graph_output[row : row + 1], + actual_suffix_frames, + target_frames, + ).clone() + outputs.append(output) + if new_frames == expected_new_frames: + current_index = int(cache["_suffix_buffer_index"]) + next_index = 1 - current_index + next_quantized = cache["_suffix_quantized_buffers"][next_index] + next_conv = cache["_suffix_conv_buffers"][next_index] + quantized_valid = int(graph_next_quantized.shape[-1]) + conv_valid = int(graph_next_conv.shape[1]) + next_quantized[:, :, :quantized_valid].copy_(graph_next_quantized[row : row + 1]) + next_conv[:, :conv_valid, :].copy_(graph_next_conv[row : row + 1]) + cache["_suffix_buffer_index"] = next_index + cache["suffix_frames"] = target_frames + cache["suffix_quantized"] = next_quantized[:, :, :quantized_valid] + cache["suffix_conv"] = next_conv[:, :conv_valid, :] + return outputs + + def _batched_request_decode( + self, + codes_list: list[torch.Tensor], + request_caches: list[dict], + ) -> list[torch.Tensor]: + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + return [self.decoder._forward_exact(codes) for codes in codes_list] + previous_suppression = getattr(self, "_suppress_stats", False) + self._suppress_stats = previous_suppression or ( + bool(request_caches) and all(cache.get("_is_dummy_run", False) for cache in request_caches) + ) + try: + return self.decoder.batched_request_decode( + codes_list, + request_caches, + prefix_length=self.prefix_length, + initial_chunk_frames=self.initial_chunk_frames, + codec_chunk_frames=self.codec_chunk_frames, + transitions_by_mode={ + "icl": self._icl_previous_frames_by_target, + "xvec": self._xvec_previous_frames_by_target, + }, + decode_icl_prefix_batch=self._decode_icl_prefix_batch, + decode_xvec_prefix_batch=self._decode_xvec_prefix_batch, + decode_suffix_batch=self._decode_suffix_batch, + decode_fallback=self._decode_request_fallback, + ) + finally: + self._suppress_stats = previous_suppression + + def _decode_request_fallback(self, codes: torch.Tensor, cache: dict) -> torch.Tensor: + if "suffix_quantized" not in cache: + if int(cache["prefix_frames"]) == 0: + return self.decoder._decode_xvec_first_chunk(codes, cache) + prefix_frames = int(cache["prefix_frames"]) + output = self.decoder._decode_icl_first_chunk(codes, cache, prefix_frames) + suffix_frames = max(0, int(codes.shape[-1]) - prefix_frames) + return output[..., -suffix_frames * self.decoder.total_upsample :] + return self._decode_suffix_eager_fallback(codes, cache) + + def batched_chunked_decode_with_cudagraph( + self, + codes: torch.Tensor, + lengths: list[int], + caches: list[dict] | None = None, + chunk_size: int = 300, + left_context_size: int = 25, + max_batch_size: int = 0, + ) -> list[torch.Tensor]: + if not self.async_chunk: + if caches is not None: + raise ValueError("caches must be None when async_chunk=False") + return self._batched_stateless_chunked_decode(codes, lengths, chunk_size, left_context_size, max_batch_size) + + if caches is None: + raise ValueError("caches are required when async_chunk=True") + if len(caches) != codes.shape[0] or len(lengths) != codes.shape[0]: + raise ValueError("codes, lengths, and caches must have the same batch size") + codes_list = [codes[row : row + 1, :, :length] for row, length in enumerate(lengths)] + outputs = self._batched_request_decode(codes_list, caches) + return [output[0] for output in outputs] + + def _batched_stateless_chunked_decode( + self, + codes: torch.Tensor, + lengths: list[int], + chunk_size: int, + left_context_size: int, + max_batch_size: int, + ) -> list[torch.Tensor]: + batch_size = int(codes.shape[0]) + if len(lengths) != batch_size: + raise ValueError("codes and lengths must have the same batch size") + if any(length < 0 or length > codes.shape[-1] for length in lengths): + raise ValueError(f"Decode lengths must be within [0, {codes.shape[-1]}], got {lengths}") + max_length = max(lengths, default=0) + if max_length == 0: + return [codes.new_empty((1, 0), dtype=torch.float32) for _ in range(batch_size)] + + total_upsample = self.decoder.total_upsample + wav_out: list[torch.Tensor] | None = None + num_rounds = (max_length + chunk_size - 1) // chunk_size + for round_index in range(num_rounds): + start = round_index * chunk_size + groups: dict[int, list[tuple[int, int, int, int, int]]] = {} + for request_index, request_length in enumerate(lengths): + if start >= request_length: + continue + end = min(start + chunk_size, request_length) + context = left_context_size if start - left_context_size > 0 else start + input_start = start - context + actual_frames = end - input_start + bucket_index = bisect.bisect_left(self.stateless_capture_sizes, actual_frames) + padded_frames = ( + self.stateless_capture_sizes[bucket_index] + if bucket_index < len(self.stateless_capture_sizes) + else actual_frames + ) + groups.setdefault(padded_frames, []).append((request_index, input_start, end, start, context)) + + for padded_frames, jobs in groups.items(): + available_batches = { + graph_batch for graph_batch, graph_frames in self.stateless_states if graph_frames == padded_frames + } + offset = 0 + while offset < len(jobs): + remaining = len(jobs) - offset + graph_batch = ( + max(available_batches) + if available_batches and remaining > max(available_batches) + else self._select_batch_bucket(remaining, available_batches) + ) + if graph_batch is not None: + request_count = min(remaining, graph_batch) + decode_batch = graph_batch + else: + request_count = min(remaining, max_batch_size) if max_batch_size > 0 else remaining + decode_batch = request_count + job_batch = jobs[offset : offset + request_count] + codes_chunk = codes.new_zeros((decode_batch, codes.shape[1], padded_frames)) + for row, (request_index, input_start, input_end, _, _) in enumerate(job_batch): + chunk = codes[request_index, :, input_start:input_end] + codes_chunk[row, :, : chunk.shape[-1]].copy_(chunk) + wav_chunk = self._decode_stateless(codes_chunk, clone_graph_output=False) + if wav_out is None: + wav_out = [ + wav_chunk.new_empty((*wav_chunk.shape[1:-1], length * total_upsample)) for length in lengths + ] + for row, (request_index, _, input_end, chunk_start, context) in enumerate(job_batch): + src_start = context * total_upsample + dst_start = chunk_start * total_upsample + dst_end = input_end * total_upsample + src_end = src_start + dst_end - dst_start + wav_out[request_index][..., dst_start:dst_end].copy_(wav_chunk[row, ..., src_start:src_end]) + offset += request_count + + if wav_out is None: + return [codes.new_empty((1, 0), dtype=torch.float32) for _ in range(batch_size)] + return wav_out + + def _decode_suffix_eager_fallback(self, codes: torch.Tensor, caches: dict) -> torch.Tensor: + for key in ( + "_suffix_quantized_buffers", + "_suffix_conv_buffers", + "_suffix_buffer_index", + ): + caches.pop(key, None) + decoder_prefix_frames = int(caches.get("decoder_prefix_frames", caches["prefix_frames"])) + return self.decoder.decode_suffix(codes, caches, decoder_prefix_frames) + + def _decode_stateless(self, codes: torch.Tensor, *, clone_graph_output: bool) -> torch.Tensor: + if self.async_chunk: + raise RuntimeError("Stateless decode is unavailable when async_chunk=True") + if not self.enabled or not self._warmed_up or torch.cuda.is_current_stream_capturing(): + return self.decoder._forward_exact(codes) + + batch_size = int(codes.shape[0]) + actual_frames = int(codes.shape[-1]) + bucket_index = bisect.bisect_left(self.stateless_capture_sizes, actual_frames) + padded_frames = ( + self.stateless_capture_sizes[bucket_index] if bucket_index < len(self.stateless_capture_sizes) else None + ) + key = (batch_size, padded_frames) if padded_frames is not None else None + if key is None or key not in self.stateless_states: + return self.decoder._forward_exact(codes) + + state = self.stateless_states[key] + static_input = state["input"]["codes"] + if actual_frames == padded_frames: + static_input.copy_(codes) + else: + static_input.zero_() + static_input[:, :, :actual_frames].copy_(codes) + state["graph"].replay() + output = self._trim_replay_output( + state["output"], + actual_frames, + padded_frames, + ) + return output.clone() if clone_graph_output else 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 chunked_decode_with_cudagraph( + self, + codes: torch.Tensor, + caches: dict | None = None, + chunk_size: int = 300, + left_context_size: int = 25, + ) -> torch.Tensor: + if self.async_chunk and caches is None: + raise ValueError("caches are required when async_chunk=True") + if not self.async_chunk and caches is not None: + raise ValueError("caches must be None when async_chunk=False") + + wavs = [] + start_index = 0 + total_len = codes.shape[-1] + total_upsample = self.decoder.total_upsample + while start_index < total_len: + end_index = min(start_index + chunk_size, total_len) + context_size = left_context_size if start_index - left_context_size > 0 else start_index + + codes_chunk = codes[..., start_index - context_size : end_index] + if not self.async_chunk: + wav_chunk = self._decode_stateless(codes_chunk, clone_graph_output=False) + else: + wav_chunk = self._batched_request_decode([codes_chunk], [caches])[0] + + # 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. + if self.async_chunk: + wavs.append(wav_chunk.clone()) + else: + wavs.append(wav_chunk[..., context_size * total_upsample :].clone()) + start_index = end_index + + if not wavs: + return self.decoder._forward_exact(codes) + output = torch.cat(wavs, dim=-1) + if output.is_cuda and not torch.cuda.is_current_stream_capturing(): + torch.cuda.current_stream(output.device).synchronize() + return output diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py index 9d6eb7aaa..ea11e64a7 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py @@ -13,6 +13,7 @@ # limitations under the License. """PyTorch Qwen3TTSTokenizerV2 model.""" +import copy import inspect import math from collections.abc import Callable @@ -25,9 +26,9 @@ from torch import nn from torch.nn import functional as F from transformers import MimiConfig, MimiModel from transformers.activations import ACT2FN -from transformers.cache_utils import Cache +from transformers.cache_utils import Cache, DynamicCache from transformers.integrations import use_kernel_forward_from_hub -from transformers.masking_utils import create_sliding_window_causal_mask +from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask from transformers.modeling_flash_attention_utils import FlashAttentionKwargs from transformers.modeling_layers import GradientCheckpointingLayer from transformers.modeling_outputs import BaseModelOutputWithPast @@ -520,64 +521,15 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr self.norm = Qwen3TTSTokenizerV2DecoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.rotary_emb = Qwen3TTSTokenizerV2DecoderRotatoryEmbedding(config=config) self.gradient_checkpointing = False + self.has_sliding_layers = "sliding_attention" in self.config.layer_types self.window_size = config.sliding_window self.input_proj = nn.Linear(config.latent_dim, config.hidden_size) self.output_proj = nn.Linear(config.hidden_size, config.latent_dim) - # Code2Wav is cache-free and every decoder layer uses sliding attention. - # Cache an exact, compact mask per warmup shape instead of retaining a - # max_position_embeddings square buffer. CUDA Graph warmup populates - # these entries before capture, and dtype/device are part of the key. - self._sliding_attention_mask_cache: dict[tuple[int, torch.dtype, torch.device, str], torch.Tensor | None] = {} - # Initialize weights and apply final processing self.post_init() - def _apply(self, fn, recurse=True): - # Cached masks are keyed by dtype/device and should not retain storage - # from the model's previous placement. - self._sliding_attention_mask_cache.clear() - return super()._apply(fn, recurse=recurse) - - def _get_sliding_attention_mask(self, inputs_embeds: torch.Tensor) -> torch.Tensor | None: - sequence_length = int(inputs_embeds.shape[1]) - max_position_embeddings = int(self.config.max_position_embeddings) - if sequence_length > max_position_embeddings: - raise ValueError( - f"Input length {sequence_length} exceeds max_position_embeddings {max_position_embeddings}" - ) - - attention_implementation = str(self.config._attn_implementation) - cache_key = ( - sequence_length, - inputs_embeds.dtype, - inputs_embeds.device, - attention_implementation, - ) - if cache_key not in self._sliding_attention_mask_cache: - # A batch-one mask broadcasts across requests and avoids caching a - # duplicate tensor for every batch size. - mask_inputs = inputs_embeds[:1] - mask_kwargs = { - "config": self.config, - "attention_mask": None, - "past_key_values": None, - "position_ids": None, - } - sig = inspect.signature(create_sliding_window_causal_mask) - if "input_embeds" in sig.parameters: - mask_kwargs["input_embeds"] = mask_inputs - mask_kwargs["cache_position"] = torch.arange(sequence_length, device=inputs_embeds.device) - else: - mask_kwargs["inputs_embeds"] = mask_inputs - sliding_attention_mask = create_sliding_window_causal_mask(**mask_kwargs) - if sliding_attention_mask is not None: - sliding_attention_mask = sliding_attention_mask.contiguous() - self._sliding_attention_mask_cache[cache_key] = sliding_attention_mask - - return self._sliding_attention_mask_cache[cache_key] - # Note: @check_model_inputs decorator removed for vLLM compatibility # The decorator causes "unexpected keyword argument 'inputs_embeds'" error @auto_docstring @@ -600,48 +552,44 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr raise ValueError("input_ids is not expected") if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds") - if use_cache: - raise ValueError("Code2Wav sliding-mask caching does not support use_cache=True") - if past_key_values is not None: - raise ValueError("Code2Wav sliding-mask caching does not support past_key_values") - if attention_mask is not None and not isinstance(attention_mask, dict): - raise ValueError("Code2Wav cached masks require attention_mask to be None or a prepared mask mapping") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) inputs_embeds = self.input_proj(inputs_embeds) - sequence_length = inputs_embeds.shape[1] - expected_cache_position = torch.arange(sequence_length, device=inputs_embeds.device) + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) if cache_position is None: - cache_position = expected_cache_position - elif attention_mask is None and ( - cache_position.shape != expected_cache_position.shape - or cache_position.device != expected_cache_position.device - or not torch.equal(cache_position, expected_cache_position) - ): - raise ValueError("Cached Code2Wav masks require contiguous zero-based cache_position") + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, + past_seen_tokens + inputs_embeds.shape[1], + device=inputs_embeds.device, + ) if position_ids is None: position_ids = cache_position.unsqueeze(0) - elif attention_mask is None: - expected_position_ids = expected_cache_position.unsqueeze(0) - if ( - position_ids.ndim != 2 - or position_ids.shape[-1] != sequence_length - or position_ids.device != expected_position_ids.device - or not torch.equal(position_ids, expected_position_ids.expand_as(position_ids)) - ): - raise ValueError("Cached Code2Wav masks require contiguous zero-based position_ids") hidden_states = inputs_embeds - if isinstance(attention_mask, dict): - sliding_attention_mask = attention_mask["sliding_attention"] - else: - sliding_attention_mask = self._get_sliding_attention_mask(inputs_embeds) + if not isinstance(causal_mask_mapping := attention_mask, dict): + mask_kwargs = { + "config": self.config, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + sig = inspect.signature(create_causal_mask) + if "input_embeds" in sig.parameters: + mask_kwargs["input_embeds"] = inputs_embeds + mask_kwargs["cache_position"] = cache_position + else: + mask_kwargs["inputs_embeds"] = inputs_embeds + causal_mask_mapping = {"full_attention": create_causal_mask(**mask_kwargs)} + if self.has_sliding_layers: + causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs) # create position embeddings to be shared across the decoder layers position_embeddings = self.rotary_emb(hidden_states, position_ids) @@ -649,7 +597,7 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr for decoder_layer in self.layers[: self.config.num_hidden_layers]: hidden_states = decoder_layer( hidden_states, - attention_mask=sliding_attention_mask, + attention_mask=causal_mask_mapping[decoder_layer.attention_type], position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, @@ -662,7 +610,7 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr hidden_states = self.output_proj(hidden_states) return BaseModelOutputWithPast( last_hidden_state=hidden_states, - past_key_values=None, + past_key_values=past_key_values if use_cache else None, ) @@ -852,10 +800,15 @@ class SplitResidualVectorQuantizer(nn.Module): return quantized +_CONV_CONTEXT_FRAME = 2 +_DOWNSTREAM_CONTEXT_FRAME = 12 + + class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): def __init__(self, config: Qwen3TTSTokenizerV2DecoderConfig): super().__init__(config) self.total_upsample = np.prod(config.upsample_rates + config.upsampling_ratios) + self.upsample_factor = np.prod(config.upsampling_ratios) self.pre_transformer = Qwen3TTSTokenizerV2DecoderTransformerModel._from_config(config) self.quantizer = SplitResidualVectorQuantizer( @@ -913,17 +866,19 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): def enable_cudagraph( self, - capture_sizes: list[int] | None = None, capture_batch_sizes: list[int] | None = None, - extra_capture_shapes: list[tuple[int, int]] | None = None, - compile_shapes: list[tuple[int, int]] | None = None, + stateless_capture_sizes: list[int] | None = None, device: torch.device | None = None, codec_chunk_frames: int = 0, codec_left_context_frames: int = 0, + initial_codec_chunk_frames: int = 1, + codec_chunk_ramp: list[int] | None = None, + async_chunk: bool = True, decode_chunk_size: int = 300, decode_left_context: int = 25, ): - from ..cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper + # from ..cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper + from ..segmented_graph_wrapper import CUDAGraphDecoderWrapper if device is None: device = next(self.parameters()).device @@ -933,34 +888,35 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): self._cudagraph_wrapper = CUDAGraphDecoderWrapper( decoder=self, - capture_sizes=capture_sizes, capture_batch_sizes=capture_batch_sizes, - extra_capture_shapes=extra_capture_shapes, - compile_shapes=compile_shapes, + stateless_capture_sizes=stateless_capture_sizes, num_quantizers=self.config.num_quantizers, enabled=True, - ) - self._cudagraph_wrapper.warmup( - device, - dtype=torch.long, - codec_chunk_frames=codec_chunk_frames, - codec_left_context_frames=codec_left_context_frames, + async_chunk=async_chunk, + initial_chunk_frames=initial_codec_chunk_frames, + codec_chunk_frames=codec_chunk_frames or 25, + codec_chunk_ramp=codec_chunk_ramp, decode_chunk_size=decode_chunk_size, decode_left_context=decode_left_context, ) + self._incremental_chunk_frames = codec_chunk_frames or 25 + self._incremental_chunk_ramp = list(codec_chunk_ramp or ()) + self._cudagraph_wrapper.warmup( + device, + dtype=torch.long, + ) self._cudagraph_enabled = True logger.info( - "CUDA Graph enabled for decoder: batch_sizes=%s seq_lens=%s extra_shapes=%s compile_shapes=%s", + "CUDA Graph enabled for decoder: batch_sizes=%s seq_lens=%s", self._cudagraph_wrapper.capture_batch_sizes, - self._cudagraph_wrapper.capture_sizes, - self._cudagraph_wrapper.extra_capture_shapes, - self._cudagraph_wrapper.compile_shapes, + ( + self._cudagraph_wrapper.icl_capture_sizes + if async_chunk + else self._cudagraph_wrapper.stateless_capture_sizes + ), ) - def forward(self, codes): - if codes.shape[1] != self.config.num_quantizers: - raise ValueError(f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}") - + def _forward_exact(self, codes): hidden = self.quantizer.decode(codes) hidden = self.pre_conv(hidden).transpose(1, 2) @@ -969,15 +925,267 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): for blocks in self.upsample: for block in blocks: hidden = block(hidden) + wav = hidden for block in self.decoder: wav = block(wav) return wav.clamp(min=-1, max=1) - def chunked_decode(self, codes, chunk_size=300, left_context_size=25): + def _decode_icl_first_chunk( + self, + codes, + caches, + prefix_frames, + prefix_cache=None, + prefix_hidden_mask=None, + prefix_attention_mask=None, + suffix_attention_mask=None, + ): + # The CUDA graph path may left-pad a shorter request prefix to the + # captured sliding-window length. Keep the decoder's physical layout + # separate from the request metadata so eager fallbacks slice the + # cached tensors at the correct boundary. + caches["decoder_prefix_frames"] = prefix_frames + hidden = self.quantizer.decode(codes) + if prefix_hidden_mask is not None: + hidden = hidden * prefix_hidden_mask + caches.update({"ref_hidden": hidden[:, :, :prefix_frames]}) + caches.update({"suffix_quantized": hidden[:, :, prefix_frames:]}) + + hidden = self.pre_conv(hidden).transpose(1, 2) + caches.update({"ref_conv": hidden[:, :prefix_frames, :]}) + caches.update({"suffix_conv": hidden[:, prefix_frames:, :]}) + + if hidden.shape[1] < prefix_frames: + raise ValueError( + "Qwen3-TTS ICL prefix cache requires at least prefix_frames inputs, " + f"got prefix_frames={prefix_frames}, total_frames={hidden.shape[1]}" + ) + prefix_result = self.pre_transformer( + inputs_embeds=hidden[:, :prefix_frames, :], + attention_mask=prefix_attention_mask, + past_key_values=prefix_cache, + use_cache=True, + ) + + prefix_cache = prefix_result.past_key_values + working_cache = copy.deepcopy(prefix_cache) + + suffix_result = self.pre_transformer( + inputs_embeds=hidden[:, prefix_frames:, :], + attention_mask=suffix_attention_mask, + past_key_values=working_cache, + use_cache=True, + ) + + prefix_hidden = prefix_result.last_hidden_state + if prefix_attention_mask is not None: + prefix_hidden = prefix_hidden * prefix_attention_mask.unsqueeze(-1) + hidden = torch.cat( + [ + prefix_hidden, + suffix_result.last_hidden_state, + ], + dim=1, + ) + caches.update({"past_key_values": prefix_cache}) + caches.update({"prefix_hidden": prefix_hidden}) + + hidden = hidden.permute(0, 2, 1) + + for blocks in self.upsample: + for block in blocks: + hidden = block(hidden) + wav = hidden + + for block in self.decoder: + wav = block(wav) + wav = wav.clamp(min=-1, max=1) + return wav + + def _decode_xvec_first_chunk(self, codes: torch.Tensor, caches: dict) -> torch.Tensor: + """Initialize prefixless incremental state from the first codec chunk.""" + hidden = self.quantizer.decode(codes) + caches["decoder_prefix_frames"] = 0 + caches["ref_hidden"] = hidden[:, :, :0] + caches["suffix_quantized"] = hidden + + hidden = self.pre_conv(hidden).transpose(1, 2) + caches["ref_conv"] = hidden[:, :0, :] + caches["suffix_conv"] = hidden + + empty_prefix_cache = caches.get("past_key_values") + if empty_prefix_cache is None: + empty_prefix_cache = DynamicCache(config=self.config) + head_dim = getattr( + self.config, + "head_dim", + self.config.hidden_size // self.config.num_attention_heads, + ) + empty_prefix_cache.early_initialization( + batch_size=int(codes.shape[0]), + num_heads=self.config.num_key_value_heads, + head_dim=head_dim, + dtype=hidden.dtype, + device=hidden.device, + ) + working_cache = copy.deepcopy(empty_prefix_cache) + suffix_hidden = self.pre_transformer( + inputs_embeds=hidden, + past_key_values=working_cache, + use_cache=True, + ).last_hidden_state + caches["past_key_values"] = empty_prefix_cache + caches["prefix_hidden"] = suffix_hidden[:, :0, :] + + hidden = suffix_hidden.permute(0, 2, 1) + for blocks in self.upsample: + for block in blocks: + hidden = block(hidden) + wav = hidden + for block in self.decoder: + wav = block(wav) + caches["suffix_frames"] = int(codes.shape[-1]) + return wav.clamp(min=-1, max=1) + + def _decode_suffix( + self, + new_codes: torch.Tensor, + old_quantized: torch.Tensor, + old_conv: torch.Tensor, + caches: dict, + prefix_frames: int, + new_frames: int, + rolling: bool, + attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ref_hidden = caches["ref_hidden"] + conv_context = torch.cat([ref_hidden, old_quantized], dim=-1)[:, :, -_CONV_CONTEXT_FRAME:] + suffix_cache_length = int(self.config.sliding_window) + + new_quantized = self.quantizer.decode(new_codes) + new_conv_input = torch.cat([conv_context, new_quantized], dim=-1) + new_conv = self.pre_conv(new_conv_input) + new_conv = new_conv[:, :, -new_codes.shape[-1] :].transpose(1, 2) + + if rolling: + boundary_input = torch.cat( + [ref_hidden[:, :, -_CONV_CONTEXT_FRAME:], old_quantized[:, :, :_CONV_CONTEXT_FRAME]], + dim=-1, + ) + boundary_conv = self.pre_conv(boundary_input) + boundary_conv = boundary_conv[:, :, -_CONV_CONTEXT_FRAME:].transpose(1, 2) + suffix_conv = torch.cat([boundary_conv, old_conv, new_conv], dim=1) + else: + suffix_conv = torch.cat([old_conv, new_conv], dim=1) + + next_quantized = torch.cat([old_quantized, new_quantized], dim=-1)[:, :, -suffix_cache_length:] + next_conv = suffix_conv[:, -(suffix_cache_length - _CONV_CONTEXT_FRAME) :, :] + + hidden = torch.cat([caches["ref_conv"], suffix_conv], dim=1) + hidden = hidden.transpose(1, 2).contiguous().transpose(1, 2) + working_cache = copy.deepcopy(caches["past_key_values"]) + suffix_hidden = self.pre_transformer( + inputs_embeds=hidden[:, prefix_frames:, :], + attention_mask=attention_mask, + past_key_values=working_cache, + use_cache=True, + ).last_hidden_state + hidden = torch.cat([caches["prefix_hidden"], suffix_hidden], dim=1) + + required_frames = min(hidden.shape[1], new_frames + _DOWNSTREAM_CONTEXT_FRAME) + hidden = hidden[:, -required_frames:, :].permute(0, 2, 1) + for blocks in self.upsample: + for block in blocks: + hidden = block(hidden) + + wav = hidden + for block in self.decoder: + wav = block(wav) + wav = wav[..., -new_frames * self.total_upsample :].clamp(min=-1, max=1) + return wav, next_quantized, next_conv + + def _is_suffix_cache_rolling(self, previous_suffix_frames: int, cached_frames: int) -> bool: + suffix_cache_length = int(self.config.sliding_window) + return previous_suffix_frames >= suffix_cache_length and cached_frames == suffix_cache_length + + def decode_suffix(self, new_codes, caches, prefix_frames): + previous_suffix_frames = int(caches.get("suffix_frames", caches["suffix_quantized"].shape[-1])) + cached_frames = int(caches["suffix_quantized"].shape[-1]) + suffix_cache_length = int(self.config.sliding_window) + rolling = self._is_suffix_cache_rolling(previous_suffix_frames, cached_frames) + retained_frames = suffix_cache_length if rolling else previous_suffix_frames + new_frames = int(new_codes.shape[-1]) + max_new_frames = int(getattr(self, "_incremental_chunk_frames", 25)) + if not 0 < new_frames <= max_new_frames: + raise ValueError( + f"Qwen3-TTS incremental decode expected 1..{max_new_frames} new frames, got new_frames={new_frames}" + ) + + attention_mask = torch.ones( + new_codes.shape[0], + prefix_frames + retained_frames + new_frames, + dtype=torch.bool, + device=new_codes.device, + ) + attention_mask[:, : int(caches.get("prefix_pad_frames", 0))] = 0 + output, next_quantized, next_conv = self._decode_suffix( + new_codes, + caches["suffix_quantized"], + caches["suffix_conv"], + caches, + prefix_frames, + new_frames, + rolling, + attention_mask=attention_mask, + ) + caches["suffix_quantized"] = next_quantized + caches["suffix_conv"] = next_conv + caches["suffix_frames"] = retained_frames + new_frames + return output + + def forward(self, codes, caches=None): + if codes.shape[1] != self.config.num_quantizers: + raise ValueError(f"Expected {self.config.num_quantizers} layer of codes, got {codes.shape[1]}") + + if caches is None or caches.get("_is_dummy_run", False): + return self._forward_exact(codes) + + prefix_frames = int(caches["prefix_frames"]) + if prefix_frames < 0: + raise ValueError( + "Qwen3-TTS ICL prefix cache received invalid frame counts: " + f"prefix_frames={prefix_frames}, total_frames={codes.shape[-1]}" + ) + + if "suffix_quantized" not in caches: + if codes.shape[-1] < prefix_frames: + raise ValueError( + "Qwen3-TTS ICL first chunk is shorter than its prefix: " + f"prefix_frames={prefix_frames}, total_frames={codes.shape[-1]}" + ) + if prefix_frames == 0: + return self._decode_xvec_first_chunk(codes, caches) + return self._decode_icl_first_chunk(codes, caches, prefix_frames) + else: + decoder_prefix_frames = int(caches.get("decoder_prefix_frames", prefix_frames)) + return self.decode_suffix(codes, caches, decoder_prefix_frames) + + def chunked_decode( + self, + codes, + caches=None, + chunk_size=300, + left_context_size=25, + ): # Use CUDA graph if enabled if self._cudagraph_enabled and self._cudagraph_wrapper is not None: - return self._cudagraph_wrapper.chunked_decode_with_cudagraph(codes, chunk_size, left_context_size) + return self._cudagraph_wrapper.chunked_decode_with_cudagraph( + codes, + caches=caches, + chunk_size=chunk_size, + left_context_size=left_context_size, + ) # Original implementation (eager mode) wavs = [] @@ -986,31 +1194,385 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): end_index = min(start_index + chunk_size, codes.shape[-1]) context_size = left_context_size if start_index - left_context_size > 0 else start_index codes_chunk = codes[..., start_index - context_size : end_index] - wav_chunk = self(codes_chunk) - wavs.append(wav_chunk[..., context_size * self.total_upsample :]) + is_first_stateful_chunk = caches is not None and "suffix_quantized" not in caches + wav_chunk = self(codes_chunk, caches) + if caches is not None: + if is_first_stateful_chunk: + logical_prefix_frames = int(caches["prefix_frames"]) + suffix_frames = codes_chunk.shape[-1] - logical_prefix_frames + wav_chunk = wav_chunk[..., -suffix_frames * self.total_upsample :] + wavs.append(wav_chunk) + else: + wavs.append(wav_chunk[..., context_size * self.total_upsample :]) start_index = end_index - return torch.cat(wavs, dim=-1) + output = torch.cat(wavs, dim=-1) + return output + + def batched_request_decode( + self, + codes_list: list[torch.Tensor], + request_caches: list[dict[str, Any]], + *, + prefix_length: int, + initial_chunk_frames: int, + codec_chunk_frames: int, + transitions_by_mode: dict[str, dict[int, int]], + decode_icl_prefix_batch: Callable[[list[torch.Tensor], list[dict[str, Any]]], list[torch.Tensor] | None] | None, + decode_xvec_prefix_batch: Callable[[list[torch.Tensor], list[dict[str, Any]]], list[torch.Tensor] | None] + | None, + decode_suffix_batch: Callable[ + [str, int, list[torch.Tensor], list[dict[str, Any]], list[int]], list[torch.Tensor] | None + ] + | None, + decode_fallback: Callable[[torch.Tensor, dict[str, Any]], torch.Tensor], + backend_max_batch_size: int = 0, + ) -> list[torch.Tensor]: + """Group stateful requests once, then execute them with a supplied backend. + + The grouping policy is shared by CUDA Graph and eager execution. A + backend may decline a grouped decode by returning ``None``; requests in + that group are then decoded individually through ``decode_fallback``. + """ + if len(codes_list) != len(request_caches): + raise ValueError("codes_list and request_caches must have the same length") + + if torch.cuda.is_available() and torch.cuda.is_current_stream_capturing(): + return [self._forward_exact(codes) for codes in codes_list] + + outputs: list[torch.Tensor | None] = [None] * len(codes_list) + groups: dict[tuple[str, int], list[int]] = {} + for index, (codes, cache) in enumerate(zip(codes_list, request_caches, strict=True)): + if "suffix_quantized" not in cache: + prefix_frames = int(cache["prefix_frames"]) + suffix_frames = int(codes.shape[-1]) - prefix_frames + if prefix_frames == 0: + phase = "xvec_first" + elif 0 < prefix_frames <= prefix_length and suffix_frames == initial_chunk_frames: + phase = "icl_prefix" + else: + phase = "eager" + groups.setdefault((phase, 0), []).append(index) + continue + + request_kv = cache["past_key_values"] + decoder_prefix = int(cache.get("decoder_prefix_frames", cache["prefix_frames"])) + mode = "xvec" if decoder_prefix == 0 else "icl" + expected_kv_length = 0 if mode == "xvec" else prefix_length + if decoder_prefix != expected_kv_length or request_kv.get_seq_length() != expected_kv_length: + groups.setdefault(("eager", 0), []).append(index) + continue + + previous = int(cache.get("suffix_frames", cache["suffix_quantized"].shape[-1])) + rolling = self._is_suffix_cache_rolling(previous, int(cache["suffix_quantized"].shape[-1])) + target = ( + prefix_length + codec_chunk_frames + if rolling + else next( + (target for target, source in transitions_by_mode[mode].items() if source == previous), + None, + ) + ) + new_frames = int(codes.shape[-1]) + if target is None or not 0 < new_frames <= codec_chunk_frames: + groups.setdefault(("eager", 0), []).append(index) + else: + groups.setdefault((f"suffix:{mode}", target), []).append(index) + + for (phase, target), indices in groups.items(): + group_splits = ( + [ + indices[start : start + backend_max_batch_size] + for start in range(0, len(indices), backend_max_batch_size) + ] + if backend_max_batch_size > 0 + else [indices] + ) + for split_indices in group_splits: + group_codes = [codes_list[index] for index in split_indices] + group_caches = [request_caches[index] for index in split_indices] + group_outputs: list[torch.Tensor] | None = None + if phase == "icl_prefix" and decode_icl_prefix_batch is not None: + group_outputs = decode_icl_prefix_batch(group_codes, group_caches) + elif phase == "xvec_first" and decode_xvec_prefix_batch is not None: + group_outputs = decode_xvec_prefix_batch(group_codes, group_caches) + elif phase.startswith("suffix:") and decode_suffix_batch is not None: + mode = phase.removeprefix("suffix:") + new_frames = [int(codes.shape[-1]) for codes in group_codes] + group_outputs = decode_suffix_batch(mode, target, group_codes, group_caches, new_frames) + + if group_outputs is None: + group_outputs = [ + decode_fallback(codes_list[index], request_caches[index]) for index in split_indices + ] + for index, output in zip(split_indices, group_outputs, strict=True): + outputs[index] = output + + if any(output is None for output in outputs): + raise RuntimeError("stateful batched decode did not produce an output for every request") + return [output for output in outputs if output is not None] + + @staticmethod + def _batch_dynamic_caches(request_caches: list[DynamicCache]) -> DynamicCache: + batched_cache = copy.deepcopy(request_caches[0]) + for batched_layer, request_layers in zip( + batched_cache.layers, + zip(*(cache.layers for cache in request_caches), strict=True), + strict=True, + ): + if any(layer.keys is None or layer.values is None for layer in request_layers): + if not all(layer.keys is None and layer.values is None for layer in request_layers): + raise ValueError("Cannot batch partially initialized request KV caches") + batched_layer.keys = None + batched_layer.values = None + continue + batched_layer.keys = torch.cat([layer.keys for layer in request_layers], dim=0) + batched_layer.values = torch.cat([layer.values for layer in request_layers], dim=0) + return batched_cache + + @staticmethod + def _slice_dynamic_cache(cache: DynamicCache, row: int) -> DynamicCache: + request_cache = copy.deepcopy(cache) + for layer in request_cache.layers: + if layer.keys is not None: + layer.keys = layer.keys[row : row + 1].clone() + if layer.values is not None: + layer.values = layer.values[row : row + 1].clone() + return request_cache + + @staticmethod + def _cache_tensors_are_batchable(request_caches: list[dict[str, Any]], keys: tuple[str, ...]) -> bool: + return all( + key in cache + and cache[key].shape[1:] == request_caches[0][key].shape[1:] + and cache[key].dtype == request_caches[0][key].dtype + and cache[key].device == request_caches[0][key].device + for cache in request_caches + for key in keys + ) + + def _decode_icl_prefix_eager_batch( + self, + codes_list: list[torch.Tensor], + request_caches: list[dict[str, Any]], + *, + prefix_length: int, + initial_chunk_frames: int, + ) -> list[torch.Tensor] | None: + batch_size = len(codes_list) + total_frames = prefix_length + initial_chunk_frames + batched_codes = codes_list[0].new_zeros((batch_size, codes_list[0].shape[1], total_frames)) + hidden_mask = torch.ones((batch_size, 1, total_frames), dtype=self.dtype, device=batched_codes.device) + prefix_mask = torch.ones((batch_size, prefix_length), dtype=torch.bool, device=batched_codes.device) + suffix_mask = torch.ones((batch_size, total_frames), dtype=torch.bool, device=batched_codes.device) + prefix_pads: list[int] = [] + for row, (codes, cache) in enumerate(zip(codes_list, request_caches, strict=True)): + prefix_frames = int(cache["prefix_frames"]) + suffix_frames = int(codes.shape[-1]) - prefix_frames + if not 0 < prefix_frames <= prefix_length or suffix_frames != initial_chunk_frames: + return None + prefix_pad = prefix_length - prefix_frames + prefix_pads.append(prefix_pad) + batched_codes[row, :, prefix_pad:prefix_length].copy_(codes[0, :, :prefix_frames]) + batched_codes[row, :, prefix_length:].copy_(codes[0, :, prefix_frames:]) + hidden_mask[row, :, :prefix_pad] = 0 + prefix_mask[row, :prefix_pad] = 0 + suffix_mask[row, :prefix_pad] = 0 + + batched_cache: dict[str, Any] = {} + output = self._decode_icl_first_chunk( + batched_codes, + batched_cache, + prefix_length, + prefix_hidden_mask=hidden_mask, + prefix_attention_mask=prefix_mask, + suffix_attention_mask=suffix_mask, + ) + outputs: list[torch.Tensor] = [] + for row, (cache, prefix_pad) in enumerate(zip(request_caches, prefix_pads, strict=True)): + for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"): + cache[key] = batched_cache[key][row : row + 1].clone() + cache["past_key_values"] = self._slice_dynamic_cache(batched_cache["past_key_values"], row) + cache["decoder_prefix_frames"] = prefix_length + cache["prefix_pad_frames"] = prefix_pad + cache["suffix_frames"] = initial_chunk_frames + outputs.append(output[row : row + 1, :, -initial_chunk_frames * self.total_upsample :].clone()) + return outputs + + def _decode_xvec_prefix_eager_batch( + self, + codes_list: list[torch.Tensor], + request_caches: list[dict[str, Any]], + *, + initial_chunk_frames: int, + ) -> list[torch.Tensor] | None: + if any(int(codes.shape[-1]) != initial_chunk_frames for codes in codes_list): + return None + batched_codes = torch.cat(codes_list, dim=0) + batched_cache: dict[str, Any] = {} + output = self._decode_xvec_first_chunk(batched_codes, batched_cache) + outputs: list[torch.Tensor] = [] + for row, cache in enumerate(request_caches): + for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"): + cache[key] = batched_cache[key][row : row + 1].clone() + cache["past_key_values"] = self._slice_dynamic_cache(batched_cache["past_key_values"], row) + cache["decoder_prefix_frames"] = 0 + cache["suffix_frames"] = initial_chunk_frames + outputs.append(output[row : row + 1].clone()) + return outputs + + def _decode_suffix_eager_batch( + self, + mode: str, + target_frames: int, + codes_list: list[torch.Tensor], + request_caches: list[dict[str, Any]], + new_frames_list: list[int], + *, + prefix_length: int, + codec_chunk_frames: int, + transitions_by_mode: dict[str, dict[int, int]], + ) -> list[torch.Tensor] | None: + tensor_keys = ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv") + if not self._cache_tensors_are_batchable(request_caches, tensor_keys): + return None + previous_frames = transitions_by_mode[mode].get(target_frames) + if previous_frames is None: + return None + expected_new_frames = target_frames - previous_frames + if any(not 0 < new_frames <= expected_new_frames for new_frames in new_frames_list): + return None + + batched_codes = codes_list[0].new_zeros((len(codes_list), codes_list[0].shape[1], expected_new_frames)) + for row, (codes, new_frames) in enumerate(zip(codes_list, new_frames_list, strict=True)): + batched_codes[row, :, :new_frames].copy_(codes[0, :, -new_frames:]) + batched_cache = { + key: torch.cat([cache[key] for cache in request_caches], dim=0) + for key in ("ref_hidden", "ref_conv", "prefix_hidden") + } + try: + batched_cache["past_key_values"] = self._batch_dynamic_caches( + [cache["past_key_values"] for cache in request_caches] + ) + except ValueError: + return None + attention_mask = torch.ones( + len(codes_list), + (prefix_length if mode == "icl" else 0) + previous_frames + expected_new_frames, + dtype=torch.bool, + device=batched_codes.device, + ) + if mode == "icl": + for row, cache in enumerate(request_caches): + attention_mask[row, : int(cache.get("prefix_pad_frames", 0))] = 0 + old_quantized = torch.cat([cache["suffix_quantized"] for cache in request_caches], dim=0) + old_conv = torch.cat([cache["suffix_conv"] for cache in request_caches], dim=0) + rolling = self._is_suffix_cache_rolling(previous_frames, int(old_quantized.shape[-1])) + output, next_quantized, next_conv = self._decode_suffix( + batched_codes, + old_quantized, + old_conv, + batched_cache, + prefix_length if mode == "icl" else 0, + expected_new_frames, + rolling, + attention_mask=attention_mask, + ) + + outputs: list[torch.Tensor] = [] + for row, (cache, new_frames) in enumerate(zip(request_caches, new_frames_list, strict=True)): + outputs.append(output[row : row + 1, :, : new_frames * self.total_upsample].clone()) + if new_frames == expected_new_frames: + cache["suffix_quantized"] = next_quantized[row : row + 1].clone() + cache["suffix_conv"] = next_conv[row : row + 1].clone() + cache["suffix_frames"] = target_frames + return outputs def batched_chunked_decode( self, codes, lengths, + caches=None, chunk_size=300, left_context_size=25, max_batch_size=0, ): - if self._cudagraph_enabled and self._cudagraph_wrapper is not None: + if ( + self._cudagraph_enabled + and self._cudagraph_wrapper is not None + and hasattr(self._cudagraph_wrapper, "batched_chunked_decode_with_cudagraph") + ): return self._cudagraph_wrapper.batched_chunked_decode_with_cudagraph( codes, lengths, + caches=caches, chunk_size=chunk_size, left_context_size=left_context_size, max_batch_size=max_batch_size, ) + if caches is not None: + prefix_length = int(getattr(self.config, "sliding_window", 0) or 0) + initial_codec_chunk_frames = int(getattr(self, "_initial_codec_chunk_frames", 1)) + codec_chunk_frames = int(getattr(self, "_incremental_chunk_frames", 25)) + chunk_ramp = list(getattr(self, "_incremental_chunk_ramp", ()) or ()) + initial_chunk_frames = int(chunk_ramp[0]) if chunk_ramp else initial_codec_chunk_frames + transitions: dict[int, int] = {} + previous = initial_chunk_frames + for new_frames in chunk_ramp[1:]: + if previous >= prefix_length: + break + target = previous + int(new_frames) + transitions[target] = previous + previous = target + while previous < prefix_length: + target = previous + codec_chunk_frames + transitions[target] = previous + previous = target + transitions[prefix_length + codec_chunk_frames] = prefix_length + codes_list = [codes[row : row + 1, :, :length] for row, length in enumerate(lengths)] + outputs = self.batched_request_decode( + codes_list, + caches, + prefix_length=prefix_length, + initial_chunk_frames=initial_chunk_frames, + codec_chunk_frames=codec_chunk_frames, + transitions_by_mode={"icl": transitions, "xvec": transitions}, + decode_icl_prefix_batch=lambda group_codes, group_caches: self._decode_icl_prefix_eager_batch( + group_codes, + group_caches, + prefix_length=prefix_length, + initial_chunk_frames=initial_chunk_frames, + ), + decode_xvec_prefix_batch=lambda group_codes, group_caches: self._decode_xvec_prefix_eager_batch( + group_codes, + group_caches, + initial_chunk_frames=initial_chunk_frames, + ), + decode_suffix_batch=lambda mode, target, group_codes, group_caches, new_frames: ( + self._decode_suffix_eager_batch( + mode, + target, + group_codes, + group_caches, + new_frames, + prefix_length=prefix_length, + codec_chunk_frames=codec_chunk_frames, + transitions_by_mode={"icl": transitions, "xvec": transitions}, + ) + ), + decode_fallback=lambda request_codes, cache: self.chunked_decode( + request_codes, + caches=cache, + chunk_size=chunk_size, + left_context_size=left_context_size, + ), + backend_max_batch_size=max_batch_size, + ) + return [output[0] for output in outputs] + from ..cuda_graph_decoder_wrapper import _batched_chunked_decode - return _batched_chunked_decode( + padded = _batched_chunked_decode( codes, lengths, decode_fn=self, @@ -1019,6 +1581,7 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel): left_context_size=left_context_size, max_batch_size=max_batch_size, ) + return [padded[row, :, : length * self.total_upsample] for row, length in enumerate(lengths)] class Qwen3TTSTokenizerV2Encoder(MimiModel): diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py index 9b5e5ea38..da5fa840d 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_tts.py @@ -164,21 +164,31 @@ def talker2code2wav_async_chunk( if finished: return OmniPayloadStruct( codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), - meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + meta=MetaStruct( + request_id=request_id, + left_context_size=0, + finished=torch.tensor(True, dtype=torch.bool), + ), ) return None if ramp is not None: chunk_index = transfer_manager.ramp_chunk_count.get(request_id, 0) + first_chunk = chunk_index <= 0 emit, context_length = compute_ramp_emit(length, chunk_index, ramp, chunk_size, finished) if not emit: return None if context_length == 0: return OmniPayloadStruct( codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), - meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + meta=MetaStruct( + request_id=request_id, + left_context_size=0, + finished=torch.tensor(True, dtype=torch.bool), + ), ) else: + first_chunk = int(transfer_manager.put_req_chunk.get(request_id, 0)) <= 0 use_first_chunk = initial_chunk_size > 0 and initial_chunk_size < chunk_size if use_first_chunk and length <= initial_chunk_size: @@ -193,15 +203,18 @@ def talker2code2wav_async_chunk( chunk_length = adjusted % chunk_size context_length = chunk_length if chunk_length != 0 else chunk_size - end_index = min(length, left_context_size_config + context_length) - left_context_size = max(0, end_index - context_length) - window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] + if finished and first_chunk: + context_length = length - # Prepend the bounded ref_code tail to the first emitted chunk so Code2Wav - # can cache the same reference context that mainline would otherwise send on - # every chunk. Follow-up chunks send only a metadata handle; Code2Wav restores - # this context locally before decode. + # Code2Wav keeps the quantizer/conv/Transformer context per request. Ship + # only the newly completed codec frames after the first chunk. + window_frames = transfer_manager.code_prompt_token_ids[request_id][-context_length:] + left_context_size = 0 + + # ICL reference codes are part of the first decoder initialization only. + # Follow-up chunks rely entirely on the request-local decoder cache. ref_code = request_payload.get(request_id) + emitted_chunks = int(transfer_manager.put_req_chunk.get(request_id, 0)) ref_context_size = 0 ref_context_request_id: str | None = None ref_context_included = False @@ -231,13 +244,12 @@ def talker2code2wav_async_chunk( ref_context = ref_context[-ref_code_context_frames:] ref_context_size = int(ref_context.shape[0]) if ref_context.ndim > 1 else 0 if ref_context_size > 0: - ref_context_request_id = request_id - emitted_chunks = int(transfer_manager.put_req_chunk.get(request_id, 0)) if emitted_chunks <= 0: + ref_context_request_id = request_id ref_frames = ref_context.tolist() window_frames = ref_frames + window_frames ref_context_included = True - left_context_size += ref_context_size + left_context_size = ref_context_size num_quantizers = len(window_frames[0]) num_frames = len(window_frames) @@ -247,6 +259,7 @@ def talker2code2wav_async_chunk( ) meta = MetaStruct( + request_id=request_id, left_context_size=left_context_size, finished=torch.tensor(finished, dtype=torch.bool), )