[TTS][Perf] Optimize Qwen3-TTS high-concurrency serving (#3662)

Signed-off-by: Sy03 <1370724210@qq.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: SYLAR <125541396+lishunyang12@users.noreply.github.com>

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -403,3 +403,57 @@ class TestCodePredictorWrapperConfig:
         assert config.use_parallel_embedding is False
         assert config.return_proj_buf is False
         assert config.sampling_mode == "per_call"
+
+    def test_prefix_graph_config_helpers(self, loaded_target_classes) -> None:
+        """Prefix graph helpers parse deploy config values and keep valid seq lens only."""
+        _ = loaded_target_classes
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        wrapper_cls = common_mod.CodePredictorWrapper
+
+        assert wrapper_cls._parse_positive_int_set("64; 128,0,-1") == {
+            64,
+            128,
+        }
+        assert wrapper_cls._parse_positive_int_set([2, "4", 0]) == {2, 4}
+        with pytest.raises(ValueError, match="Invalid positive int config value 'bad'"):
+            wrapper_cls._parse_positive_int_set("2,bad")
+
+        wrapper = object.__new__(wrapper_cls)
+        wrapper._prefix_graph_seq_lens = {1, 2, 4, 8, 99}
+        assert wrapper._prefix_seq_lens(6) == [2, 4]
+
+    def test_prefix_graph_env_requires_cuda_graphs(
+        self,
+        mocker: MockerFixture,
+        loaded_target_classes,
+    ) -> None:
+        """Avoid prefix warmup on shared code-predictor users that disable CUDA graphs."""
+        _ = loaded_target_classes
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+
+        cp_config, _ = _make_tiny_config(loaded_target_classes)
+        vllm_config = _make_vllm_config(mocker, max_num_seqs=2)
+        vllm_config.model_config.stage_connector_config = {
+            "extra": {
+                "code_predictor_prefix_graphs": True,
+                "code_predictor_prefix_graph_buckets": [2],
+                "code_predictor_prefix_graph_seq_lens": "2,3",
+            }
+        }
+
+        no_graph_wrapper = common_mod.CodePredictorWrapper(
+            vllm_config=vllm_config,
+            cp_config=cp_config,
+            wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=False),
+        )
+        assert no_graph_wrapper._prefix_graphs_enabled is False
+        assert no_graph_wrapper._prefix_graph_buckets == {2}
+        assert no_graph_wrapper._prefix_graph_seq_lens == {2, 3}
+
+        graph_wrapper = common_mod.CodePredictorWrapper(
+            vllm_config=vllm_config,
+            cp_config=cp_config,
+            wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=True),
+        )
+        assert graph_wrapper._prefix_graphs_enabled is True

tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py

@@ -87,6 +87,7 @@ def wrapper(decoder):
     w = CUDAGraphDecoderWrapper(
         decoder=decoder,
         capture_sizes=[25, 50, 100],
+        capture_batch_sizes=[1, 2],
         num_quantizers=NUM_QUANTIZERS,
         enabled=True,
     )
@@ -94,8 +95,8 @@ def wrapper(decoder):
     return w
 
 
-def _random_codes(seq_len, device=DEVICE):
-    return torch.randint(0, 100, (1, NUM_QUANTIZERS, seq_len), dtype=torch.long, device=device)
+def _random_codes(seq_len, batch_size=1, device=DEVICE):
+    return torch.randint(0, 100, (batch_size, NUM_QUANTIZERS, seq_len), dtype=torch.long, device=device)
 
 
 # ──────────────────────────────────────────────────────────────────
@@ -213,6 +214,44 @@ def test_chunked_decode_exact_size_equivalence(decoder, wrapper, total_len):
     torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
 
 
+def test_chunked_decode_output_survives_later_replay(wrapper):
+    """Chunked output must not alias graph static buffers overwritten by later replays."""
+    codes = _random_codes(100)
+    overwrite_codes = _random_codes(100)
+
+    with torch.no_grad():
+        graph_out = wrapper.chunked_decode_with_cudagraph(codes, chunk_size=50, left_context_size=0)
+        expected = graph_out.clone()
+        _ = wrapper.decode(overwrite_codes[..., :50])
+        _ = wrapper.decode(overwrite_codes)
+
+    torch.testing.assert_close(graph_out, expected, atol=0, rtol=0)
+
+
+def test_batched_chunked_decode_variable_lengths_matches_per_request_eager(decoder, wrapper):
+    """Variable-length chunk batching should match independent chunked decodes."""
+    long_codes = _random_codes(100)
+    short_codes = _random_codes(50)
+    padded_codes = torch.zeros(2, NUM_QUANTIZERS, 100, dtype=torch.long, device=DEVICE)
+    padded_codes[0, :, :] = long_codes[0]
+    padded_codes[1, :, :50] = short_codes[0]
+
+    with torch.no_grad():
+        eager_long = _eager_chunked(decoder, long_codes, chunk_size=50, left_context_size=0)
+        eager_short = _eager_chunked(decoder, short_codes, chunk_size=50, left_context_size=0)
+        graph_out = wrapper.batched_chunked_decode_with_cudagraph(
+            padded_codes,
+            [100, 50],
+            chunk_size=50,
+            left_context_size=0,
+            max_batch_size=2,
+        )
+
+    assert graph_out.shape == (2, 1, 100 * TOTAL_UPSAMPLE)
+    torch.testing.assert_close(graph_out[0:1], eager_long, atol=1e-6, rtol=1e-6)
+    torch.testing.assert_close(graph_out[1:2, :, : 50 * TOTAL_UPSAMPLE], eager_short, atol=1e-6, rtol=1e-6)
+
+
 def _eager_chunked(decoder, codes, chunk_size, left_context_size):
     """Eager chunked decode matching the real decoder's chunked_decode logic."""
     wavs = []
@@ -254,15 +293,95 @@ def test_disabled_wrapper_matches_eager(decoder, wrapper):
     torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
 
 
-def test_batch_size_gt1_falls_back(decoder, wrapper):
-    """Batch size > 1 should fall back to eager (bit-identical)."""
-    codes = torch.randint(0, 100, (2, NUM_QUANTIZERS, 25), dtype=torch.long, device=DEVICE)
+def test_batch_size_gt1_uses_matching_graph(decoder, wrapper):
+    """Captured batch size > 1 should replay a matching graph."""
+    assert (2, 25) in wrapper.graphs
+    codes = _random_codes(25, batch_size=2)
     with torch.no_grad():
         eager_out = decoder(codes)
         graph_out = wrapper.decode(codes)
     torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
 
 
+def test_uncaptured_batch_size_falls_back(decoder, wrapper):
+    """Uncaptured batch sizes should fall back to eager."""
+    assert (3, 25) not in wrapper.graphs
+    codes = _random_codes(25, batch_size=3)
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = wrapper.decode(codes)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+def test_extra_capture_shape_uses_sparse_graph(decoder):
+    """Extra capture shapes should not expand to a full batch x size product."""
+    sparse_wrapper = CUDAGraphDecoderWrapper(
+        decoder=decoder,
+        capture_sizes=[25],
+        capture_batch_sizes=[1],
+        extra_capture_shapes=[(2, 50)],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    sparse_wrapper.warmup(DEVICE)
+
+    assert (1, 25) in sparse_wrapper.graphs
+    assert (2, 50) in sparse_wrapper.graphs
+    assert (2, 25) not in sparse_wrapper.graphs
+
+    codes = _random_codes(50, batch_size=2)
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = sparse_wrapper.decode(codes)
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
+def test_compile_shape_supports_exact_and_padded_buckets(decoder, monkeypatch):
+    """Configured torch.compile shapes should replay exact and padded CUDA Graph buckets."""
+
+    compile_kwargs = {}
+
+    def _fake_compile(model, **_kwargs):
+        compile_kwargs.update(_kwargs)
+
+        def _compiled(codes):
+            return model(codes) + 0.125
+
+        return _compiled
+
+    monkeypatch.setattr(torch, "compile", _fake_compile)
+
+    compiled_wrapper = CUDAGraphDecoderWrapper(
+        decoder=decoder,
+        capture_sizes=[25, 50],
+        capture_batch_sizes=[1],
+        compile_shapes=[(1, 25), (1, 50)],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    compiled_wrapper.warmup(DEVICE)
+
+    exact_codes = _random_codes(25)
+    padded_codes = _random_codes(30)
+    uncaptured_codes = _random_codes(60)
+    padded_static = torch.zeros(1, NUM_QUANTIZERS, 50, dtype=torch.long, device=DEVICE)
+    padded_static[:, :, :30] = padded_codes
+    with torch.no_grad():
+        exact_eager = decoder(exact_codes)
+        exact_out = compiled_wrapper.decode(exact_codes)
+        padded_graph_expected = decoder(padded_static)[..., : 30 * TOTAL_UPSAMPLE]
+        padded_out = compiled_wrapper.decode(padded_codes)
+        uncaptured_eager = decoder(uncaptured_codes)
+        uncaptured_out = compiled_wrapper.decode(uncaptured_codes)
+
+    torch.testing.assert_close(exact_out, exact_eager + 0.125, atol=0, rtol=0)
+    torch.testing.assert_close(padded_out, padded_graph_expected + 0.125, atol=0, rtol=0)
+    torch.testing.assert_close(uncaptured_out, uncaptured_eager, atol=0, rtol=0)
+    assert compile_kwargs["mode"] == "default"
+    assert compile_kwargs["fullgraph"] is False
+    assert compile_kwargs["dynamic"] is False
+
+
 def test_deterministic_across_calls(decoder, wrapper):
     """Same input should produce identical CUDA graph output across calls."""
     codes = _random_codes(30)

tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py

@@ -24,6 +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.cudagraph_calls: list[dict[str, int | torch.device]] = []
 
     def to(self, *args, **kwargs):
@@ -40,12 +41,40 @@ class _FakeDecoder(nn.Module):
             {
                 "chunk_size": chunk_size,
                 "left_context_size": left_context_size,
+                "codes_shape": tuple(codes.shape),
             }
         )
+        batch = codes.shape[0]
         frames = codes.shape[-1]
         wav_len = frames * self.total_upsample + 6
-        wav = torch.arange(wav_len, dtype=torch.float32)
-        return wav.view(1, 1, -1)
+        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
+
+    def batched_chunked_decode(
+        self,
+        codes: torch.Tensor,
+        lengths: list[int],
+        *,
+        chunk_size: int = 300,
+        left_context_size: int = 25,
+        max_batch_size: int = 0,
+    ) -> torch.Tensor:
+        self.batched_decode_calls.append(
+            {
+                "chunk_size": chunk_size,
+                "left_context_size": left_context_size,
+                "max_batch_size": max_batch_size,
+                "codes_shape": tuple(codes.shape),
+                "lengths": tuple(lengths),
+            }
+        )
+        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
 
     def enable_cudagraph(self, **kwargs):
         self.cudagraph_calls.append(kwargs)
@@ -177,6 +206,7 @@ def test_connector_codec_chunking_does_not_override_decode_chunking():
     assert model.decoder.decode_calls[-1] == {
         "chunk_size": 300,
         "left_context_size": 25,
+        "codes_shape": (1, _NUM_QUANTIZERS, 6),
     }
 
 
@@ -199,6 +229,175 @@ def test_decode_chunking_can_be_overridden_separately():
     assert model._decode_left_context_frames == 17
 
 
+def test_malformed_codec_length_warning_is_rate_limited():
+    model = _make_model()
+
+    with patch("vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code2wav.logger.warning") as warning:
+        out1 = model.forward(input_ids=torch.arange(1, dtype=torch.long))
+        out2 = model.forward(input_ids=torch.arange(1, dtype=torch.long))
+
+    assert warning.call_count == 1
+    assert "not divisible by num_quantizers" in warning.call_args[0][0]
+    assert "suppressing repeats" in warning.call_args[0][0]
+    assert out1.multimodal_outputs["model_outputs"][0].numel() == 0
+    assert out2.multimodal_outputs["model_outputs"][0].numel() == 0
+
+
+def test_forward_batches_equal_length_requests_in_one_decoder_call():
+    model = _make_model()
+
+    out = model.forward(
+        input_ids=torch.arange(12, dtype=torch.long),
+        seq_token_counts=[6, 6],
+        runtime_additional_information=[
+            {"meta": {"left_context_size": 0}},
+            {"meta": {"left_context_size": 1}},
+        ],
+    )
+
+    assert model.decoder.decode_calls == [
+        {
+            "chunk_size": 300,
+            "left_context_size": 25,
+            "codes_shape": (2, _NUM_QUANTIZERS, 3),
+        }
+    ]
+    audios = out.multimodal_outputs["model_outputs"]
+    torch.testing.assert_close(audios[0], torch.arange(12, dtype=torch.float32))
+    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():
+    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),
+        seq_token_counts=[4, 8],
+        runtime_additional_information=[
+            {"meta": {"left_context_size": 0}},
+            {"meta": {"left_context_size": 2}},
+        ],
+    )
+
+    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, 4),
+            "lengths": (2, 4),
+        }
+    ]
+    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(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,
@@ -216,6 +415,10 @@ 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,
         "device": torch.device("cuda"),
         "codec_chunk_frames": 25,
         "codec_left_context_frames": 72,
@@ -224,6 +427,106 @@ def test_decode_chunking_override_is_passed_to_cudagraph():
     }
 
 
+def test_cudagraph_capture_shapes_can_be_configured():
+    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]],
+            }
+        },
+    )
+
+    _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():
+    model = _make_model(
+        async_chunk=True,
+        device=torch.device("cuda"),
+        stage_connector_config={
+            "extra": {
+                "decode_compile_shapes": ["1:325", [1, 73]],
+            }
+        },
+    )
+
+    _load_weights_noop(model)
+
+    call = model.decoder.cudagraph_calls[-1]
+    assert call["compile_shapes"] == [(1, 73), (1, 325)]
+
+
+def test_decode_tf32_can_be_configured():
+    old_matmul_tf32 = torch.backends.cuda.matmul.allow_tf32
+    old_cudnn_tf32 = torch.backends.cudnn.allow_tf32
+    old_matmul_precision = torch.get_float32_matmul_precision()
+    try:
+        torch.backends.cuda.matmul.allow_tf32 = False
+        torch.backends.cudnn.allow_tf32 = False
+        torch.set_float32_matmul_precision("highest")
+        model = _make_model(
+            async_chunk=True,
+            device=torch.device("cuda"),
+            stage_connector_config={
+                "extra": {
+                    "decode_enable_tf32": "true",
+                }
+            },
+        )
+
+        _load_weights_noop(model)
+
+        assert torch.backends.cuda.matmul.allow_tf32 is True
+        assert torch.backends.cudnn.allow_tf32 is True
+        assert torch.get_float32_matmul_precision() == "high"
+    finally:
+        torch.backends.cuda.matmul.allow_tf32 = old_matmul_tf32
+        torch.backends.cudnn.allow_tf32 = old_cudnn_tf32
+        torch.set_float32_matmul_precision(old_matmul_precision)
+
+
+def test_decode_batch_bucket_frames_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():
+    model = _make_model(
+        async_chunk=True,
+        stage_connector_config={
+            "extra": {
+                "decode_batch_max_size": -1,
+            }
+        },
+    )
+
+    with pytest.raises(ValueError, match="decode_batch_max_size"):
+        _load_weights_noop(model)
+
+
 def test_invalid_decode_chunking_is_rejected():
     model = _make_model(
         async_chunk=True,

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -0,0 +1,491 @@
+from types import SimpleNamespace
+
+import numpy as np
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+    _NORMALIZED_REF_AUDIO_KEY,
+    _PRECOMPUTED_REF_CODE_KEY,
+    _PRECOMPUTED_REF_IDS_KEY,
+    _PRECOMPUTED_TEXT_IDS_KEY,
+    Qwen3TTSTalkerForConditionalGeneration,
+)
+
+
+def _make_minimal_talker():
+    model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+    model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16)
+    return model
+
+
+def test_single_token_prefill_uses_prefill_path():
+    model = _make_minimal_talker()
+    full_prompt_embeds = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    trailing_text = torch.ones((2, 4), dtype=torch.float32)
+    tts_pad = torch.full((1, 4), 0.5, dtype=torch.float32)
+    ref_code = torch.arange(32, dtype=torch.long).reshape(2, 16)
+
+    def fake_build_prompt_embeds(*, task_type, info_dict):
+        return full_prompt_embeds, trailing_text, tts_pad, 2, ref_code
+
+    model._build_prompt_embeds = fake_build_prompt_embeds
+
+    input_ids = torch.tensor([123], dtype=torch.long)
+    out_ids, out_embeds, update = model.preprocess(
+        input_ids=input_ids,
+        input_embeds=None,
+        text=["hello"],
+        task_type=["CustomVoice"],
+        _omni_is_prefill=True,
+        _omni_num_computed_tokens=0,
+        _omni_prompt_len=3,
+    )
+
+    assert out_ids.tolist() == [7]
+    assert torch.equal(out_embeds.cpu(), full_prompt_embeds[:1].to(torch.bfloat16))
+    assert update["meta"]["talker_prefill_offset"] == 1
+    assert update["meta"]["talker_text_offset"] == 0
+    assert update["meta"]["ref_code_len"] == 2
+    assert torch.equal(update["embed"]["prefill"], full_prompt_embeds)
+    assert torch.equal(update["embed"]["tts_pad"], tts_pad)
+    assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text)
+    assert torch.equal(update["codes"]["ref"], ref_code)
+    assert update["codes"]["audio"].shape == (1, 16)
+
+
+def test_single_token_prefill_can_be_inferred_from_token_progress():
+    model = _make_minimal_talker()
+    full_prompt_embeds = torch.arange(8, dtype=torch.float32).reshape(2, 4)
+    trailing_text = torch.ones((1, 4), dtype=torch.float32)
+    tts_pad = torch.zeros((1, 4), dtype=torch.float32)
+
+    def fake_build_prompt_embeds(*, task_type, info_dict):
+        return full_prompt_embeds, trailing_text, tts_pad, None, None
+
+    model._build_prompt_embeds = fake_build_prompt_embeds
+
+    out_ids, out_embeds, update = model.preprocess(
+        input_ids=torch.tensor([123], dtype=torch.long),
+        input_embeds=None,
+        text=["hello"],
+        task_type=["CustomVoice"],
+        _omni_num_computed_tokens=0,
+        _omni_prompt_len=2,
+    )
+
+    assert out_ids.tolist() == [7]
+    assert torch.equal(out_embeds.cpu(), full_prompt_embeds[:1].to(torch.bfloat16))
+    assert update["meta"]["talker_prefill_offset"] == 1
+
+
+def test_decode_advances_trailing_text_by_offset_without_rewriting_tail():
+    model = _make_minimal_talker()
+
+    def fake_embed_input_ids(input_ids):
+        return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
+
+    model.embed_input_ids = fake_embed_input_ids
+    trailing_text = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
+
+    out_ids, out_embeds, update = model.preprocess(
+        input_ids=torch.tensor([123], dtype=torch.long),
+        input_embeds=None,
+        text=["hello"],
+        task_type=["CustomVoice"],
+        hidden_states={"trailing_text": trailing_text, "last": last_hidden},
+        embed={"tts_pad": tts_pad},
+        meta={"talker_text_offset": 1},
+        _omni_is_prefill=False,
+        _omni_num_computed_tokens=2,
+        _omni_prompt_len=2,
+    )
+
+    assert out_ids.tolist() == [123]
+    assert torch.equal(out_embeds.cpu(), torch.full((1, 4), 123.0, dtype=torch.bfloat16))
+    assert "hidden_states" not in update
+    assert update["meta"]["talker_text_offset"] == 2
+    past_hidden, text_step = update["mtp_inputs"]
+    assert torch.equal(past_hidden.cpu(), last_hidden.reshape(1, -1).to(torch.bfloat16))
+    assert torch.equal(text_step.cpu(), trailing_text[1:2].to(torch.bfloat16))
+
+
+def test_decode_advances_trailing_text_offset_across_multiple_steps():
+    model = _make_minimal_talker()
+
+    def fake_embed_input_ids(input_ids):
+        return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
+
+    model.embed_input_ids = fake_embed_input_ids
+    trailing_text = torch.arange(8, dtype=torch.float32).reshape(2, 4)
+    state_tail = trailing_text
+    last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
+    meta = {"talker_text_offset": 0}
+    seen_steps = []
+
+    for _ in range(3):
+        _, _, update = model.preprocess(
+            input_ids=torch.tensor([123], dtype=torch.long),
+            input_embeds=None,
+            text=["hello"],
+            task_type=["CustomVoice"],
+            hidden_states={"trailing_text": state_tail, "last": last_hidden},
+            embed={"tts_pad": tts_pad},
+            meta=meta,
+            _omni_is_prefill=False,
+            _omni_num_computed_tokens=2,
+            _omni_prompt_len=2,
+        )
+        seen_steps.append(update["mtp_inputs"][1].cpu())
+        if "hidden_states" in update and "trailing_text" in update["hidden_states"]:
+            state_tail = update["hidden_states"]["trailing_text"]
+        meta = update["meta"]
+
+    assert torch.equal(seen_steps[0], trailing_text[0:1].to(torch.bfloat16))
+    assert torch.equal(seen_steps[1], trailing_text[1:2].to(torch.bfloat16))
+    assert torch.equal(seen_steps[2], tts_pad.to(torch.bfloat16))
+    assert meta["talker_text_offset"] == 0
+    assert state_tail.numel() == 0
+
+
+def test_decode_compacts_long_trailing_text_after_large_offset():
+    model = _make_minimal_talker()
+
+    def fake_embed_input_ids(input_ids):
+        return input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
+
+    model.embed_input_ids = fake_embed_input_ids
+    trailing_text = torch.arange(130 * 4, dtype=torch.float32).reshape(130, 4)
+    last_hidden = torch.full((4,), 2.0, dtype=torch.float32)
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
+
+    _, _, update = model.preprocess(
+        input_ids=torch.tensor([123], dtype=torch.long),
+        input_embeds=None,
+        text=["hello"],
+        task_type=["CustomVoice"],
+        hidden_states={"trailing_text": trailing_text, "last": last_hidden},
+        embed={"tts_pad": tts_pad},
+        meta={"talker_text_offset": 64},
+        _omni_is_prefill=False,
+        _omni_num_computed_tokens=2,
+        _omni_prompt_len=2,
+    )
+
+    assert torch.equal(update["mtp_inputs"][1].cpu(), trailing_text[64:65].to(torch.bfloat16))
+    assert update["meta"]["talker_text_offset"] == 0
+    assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text[65:])
+
+
+def test_decode_batch_preprocess_matches_decode_state_updates():
+    model = _make_minimal_talker()
+
+    def fake_embed_input_ids(input_ids):
+        return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4)
+
+    model.embed_input_ids = fake_embed_input_ids
+    trailing_a = torch.arange(12, dtype=torch.float32).reshape(3, 4)
+    trailing_b = torch.arange(8, dtype=torch.float32).reshape(2, 4) + 100
+    last_a = torch.full((4,), 2.0, dtype=torch.float32)
+    last_b = torch.full((4,), 3.0, dtype=torch.float32)
+    tts_pad = torch.full((1, 4), -1.0, dtype=torch.float32)
+
+    out_ids, out_embeds, past_hidden, text_step, updates = model.preprocess_decode_batch(
+        input_ids=torch.tensor([101, 202], dtype=torch.long),
+        req_infos=[
+            {
+                "text": ["hello"],
+                "task_type": ["Base"],
+                "hidden_states": {"trailing_text": trailing_a, "last": last_a},
+                "embed": {"tts_pad": tts_pad},
+                "meta": {"talker_text_offset": 1},
+            },
+            {
+                "text": ["world"],
+                "task_type": ["CustomVoice"],
+                "hidden_states": {"trailing_text": trailing_b, "last": last_b},
+                "embed": {"tts_pad": tts_pad},
+                "meta": {"talker_text_offset": 2},
+            },
+        ],
+    )
+
+    assert out_ids.tolist() == [101, 202]
+    assert torch.equal(out_embeds.cpu(), torch.tensor([[101.0] * 4, [202.0] * 4], dtype=torch.bfloat16))
+    assert torch.equal(past_hidden.cpu(), torch.stack([last_a, last_b]).to(torch.bfloat16))
+    assert torch.equal(text_step[0].cpu(), trailing_a[1].to(torch.bfloat16))
+    assert torch.equal(text_step[1].cpu(), tts_pad.reshape(-1).to(torch.bfloat16))
+    assert updates[0]["meta"]["talker_text_offset"] == 2
+    assert updates[0]["meta"]["codec_streaming"] is True
+    assert "hidden_states" not in updates[0]
+    assert updates[1]["meta"]["talker_text_offset"] == 0
+    assert updates[1]["meta"]["codec_streaming"] is False
+    assert updates[1]["hidden_states"]["trailing_text"].numel() == 0
+
+
+def test_base_voice_clone_normalizes_ref_audio_once_for_ref_code_and_speaker():
+    model = _make_minimal_talker()
+    device_param = torch.nn.Parameter(torch.empty(0))
+    model.parameters = lambda: iter([device_param])
+    model.config = SimpleNamespace(
+        tts_bos_token_id=100,
+        tts_eos_token_id=101,
+        tts_pad_token_id=102,
+    )
+    model.talker_config = SimpleNamespace(
+        codec_nothink_id=10,
+        codec_think_bos_id=11,
+        codec_think_eos_id=12,
+        codec_think_id=13,
+        codec_language_id={},
+        codec_pad_id=7,
+        codec_bos_id=8,
+        num_code_groups=2,
+        spk_is_dialect={},
+    )
+
+    class FakeTokenizer:
+        def __call__(self, *_args, **_kwargs):
+            return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
+
+    model._get_tokenizer = lambda: FakeTokenizer()
+    model.text_embedding = lambda ids: torch.ones((*ids.shape, 4), device=ids.device)
+    model.text_projection = lambda embeds: embeds
+    model.embed_input_ids = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+    model._generate_icl_prompt = lambda **kwargs: (
+        torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
+        torch.ones((1, 4), device=kwargs["ref_code"].device),
+    )
+
+    normalize_calls = []
+    ref_audio = np.arange(1024, dtype=np.float32)
+    model._normalize_ref_audio = lambda raw: normalize_calls.append(raw) or (ref_audio, 16000)
+
+    ref_audio_ids = []
+    model._encode_ref_audio_to_code = lambda wav, _sr: (
+        ref_audio_ids.append(id(wav)) or torch.ones((2, 2), dtype=torch.long)
+    )
+    model._extract_speaker_embedding = lambda wav, _sr: (
+        ref_audio_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16)
+    )
+
+    _prompt, _trailing, _pad, ref_code_len, ref_code = model._build_prompt_embeds(
+        task_type="Base",
+        info_dict={
+            "text": ["hello"],
+            "ref_audio": ["ref.wav"],
+            "ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
+            "non_streaming_mode": [False],
+        },
+    )
+
+    assert normalize_calls == ["ref.wav"]
+    assert ref_audio_ids == [id(ref_audio), id(ref_audio)]
+    assert ref_code_len == 2
+    assert torch.equal(ref_code, torch.ones((2, 2), dtype=torch.long))
+
+
+def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
+    model = _make_minimal_talker()
+    wav1 = np.arange(2048, dtype=np.float32)
+    wav2 = np.arange(3072, dtype=np.float32)
+    normalize_calls = []
+    model._normalize_ref_audio = lambda raw: (
+        normalize_calls.append(raw)
+        or (
+            wav1 if raw == "a.wav" else wav2,
+            16000,
+        )
+    )
+
+    class FakeSpeechTokenizer:
+        def __init__(self):
+            self.calls = []
+
+        def encode(self, audios, *, sr, return_dict):
+            self.calls.append((audios, sr, return_dict))
+            return SimpleNamespace(
+                audio_codes=[
+                    torch.full((2, 2), 11, dtype=torch.long),
+                    torch.full((3, 2), 22, dtype=torch.long),
+                ]
+            )
+
+    tok = FakeSpeechTokenizer()
+    model._ensure_speech_tokenizer_loaded = lambda: tok
+
+    class FakeTextTokenizer:
+        def __init__(self):
+            self.calls = []
+
+        def __call__(self, texts, *, padding=False):
+            self.calls.append((texts, padding))
+            return {"input_ids": [[idx + 1, idx + 2, idx + 3] for idx, _ in enumerate(texts)]}
+
+    text_tok = FakeTextTokenizer()
+    model._get_tokenizer = lambda: text_tok
+    buf = {
+        "r1": {
+            "task_type": ["Base"],
+            "text": ["one"],
+            "ref_audio": ["a.wav"],
+            "ref_text": ["hello"],
+            "x_vector_only_mode": [False],
+        },
+        "r2": {
+            "task_type": ["Base"],
+            "text": ["two"],
+            "ref_audio": ["b.wav"],
+            "ref_text": ["world"],
+            "x_vector_only_mode": [False],
+        },
+    }
+
+    model.preprocess_batch(
+        req_ids=["r1", "r2"],
+        model_intermediate_buffer=buf,
+        device=torch.device("cpu"),
+    )
+
+    assert normalize_calls == ["a.wav", "b.wav"]
+    assert len(tok.calls) == 1
+    audios, sr, return_dict = tok.calls[0]
+    assert audios[0] is wav1
+    assert audios[1] is wav2
+    assert sr == 16000
+    assert return_dict is True
+    assert torch.equal(buf["r1"]["codes"][_PRECOMPUTED_REF_CODE_KEY], torch.full((2, 2), 11))
+    assert torch.equal(buf["r2"]["codes"][_PRECOMPUTED_REF_CODE_KEY], torch.full((3, 2), 22))
+    assert buf["r1"][_NORMALIZED_REF_AUDIO_KEY][0] is wav1
+    assert buf["r2"][_NORMALIZED_REF_AUDIO_KEY][0] is wav2
+    assert len(text_tok.calls) == 2
+    assert torch.equal(buf["r1"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([1, 2, 3]))
+    assert torch.equal(buf["r2"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4]))
+    assert torch.equal(buf["r1"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([1, 2, 3]))
+    assert torch.equal(buf["r2"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4]))
+
+
+def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_without_speech_tokenizer():
+    model = _make_minimal_talker()
+    wav = np.arange(2048, dtype=np.float32)
+    model._normalize_ref_audio = lambda raw: (wav, 16000)
+    model._ensure_speech_tokenizer_loaded = lambda: (_ for _ in ()).throw(
+        AssertionError("singleton should not load speech tokenizer")
+    )
+
+    class FakeTextTokenizer:
+        def __call__(self, texts, *, padding=False):
+            return {"input_ids": [[7, 8, 9] for _ in texts]}
+
+    model._get_tokenizer = lambda: FakeTextTokenizer()
+    buf = {
+        "r1": {
+            "task_type": ["Base"],
+            "text": ["one"],
+            "ref_audio": ["a.wav"],
+            "ref_text": ["hello"],
+            "x_vector_only_mode": [False],
+        }
+    }
+
+    model.preprocess_batch(
+        req_ids=["r1"],
+        model_intermediate_buffer=buf,
+        device=torch.device("cpu"),
+    )
+
+    assert buf["r1"][_NORMALIZED_REF_AUDIO_KEY][0] is wav
+    assert torch.equal(buf["r1"][_PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([7, 8, 9]))
+    assert torch.equal(buf["r1"][_PRECOMPUTED_REF_IDS_KEY], torch.tensor([7, 8, 9]))
+
+
+def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exists():
+    model = _make_minimal_talker()
+    model._normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
+    model._get_tokenizer = lambda: (_ for _ in ()).throw(AssertionError("tokenizer not expected"))
+    model._ensure_speech_tokenizer_loaded = lambda: (_ for _ in ()).throw(
+        AssertionError("speech tokenizer not expected")
+    )
+    buf = {
+        "r1": {
+            "task_type": ["Base"],
+            "text": ["one"],
+            "ref_audio": ["a.wav"],
+            "ref_text": ["hello"],
+            "x_vector_only_mode": [False],
+            "embed": {"prefill": torch.ones((1, 4))},
+        }
+    }
+
+    model.preprocess_batch(
+        req_ids=["r1"],
+        model_intermediate_buffer=buf,
+        device=torch.device("cpu"),
+    )
+
+    assert _PRECOMPUTED_TEXT_IDS_KEY not in buf["r1"]
+    assert _NORMALIZED_REF_AUDIO_KEY not in buf["r1"]
+
+
+def test_base_voice_clone_uses_batched_ref_code_without_serial_encode():
+    model = _make_minimal_talker()
+    device_param = torch.nn.Parameter(torch.empty(0))
+    model.parameters = lambda: iter([device_param])
+    model.config = SimpleNamespace(
+        tts_bos_token_id=100,
+        tts_eos_token_id=101,
+        tts_pad_token_id=102,
+    )
+    model.talker_config = SimpleNamespace(
+        codec_nothink_id=10,
+        codec_think_bos_id=11,
+        codec_think_eos_id=12,
+        codec_think_id=13,
+        codec_language_id={},
+        codec_pad_id=7,
+        codec_bos_id=8,
+        num_code_groups=2,
+        spk_is_dialect={},
+    )
+
+    class FakeTokenizer:
+        def __call__(self, *_args, **_kwargs):
+            return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
+
+    model._get_tokenizer = lambda: FakeTokenizer()
+    model.text_embedding = lambda ids: torch.ones((*ids.shape, 4), device=ids.device)
+    model.text_projection = lambda embeds: embeds
+    model.embed_input_ids = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+    model._generate_icl_prompt = lambda **kwargs: (
+        torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
+        torch.ones((1, 4), device=kwargs["ref_code"].device),
+    )
+
+    ref_audio = np.arange(2048, dtype=np.float32)
+    ref_code = torch.arange(4, dtype=torch.long).reshape(2, 2)
+    model._normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("serial normalize not expected"))
+    model._encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw(
+        AssertionError("serial encode not expected")
+    )
+    speaker_wav_ids = []
+    model._extract_speaker_embedding = lambda wav, _sr: (
+        speaker_wav_ids.append(id(wav)) or torch.ones(4, dtype=torch.bfloat16)
+    )
+
+    _prompt, _trailing, _pad, ref_code_len, out_ref_code = model._build_prompt_embeds(
+        task_type="Base",
+        info_dict={
+            "text": ["hello"],
+            "ref_audio": ["ref.wav"],
+            "ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
+            "non_streaming_mode": [False],
+            "codes": {_PRECOMPUTED_REF_CODE_KEY: ref_code},
+            _NORMALIZED_REF_AUDIO_KEY: (ref_audio, 16000),
+        },
+    )
+
+    assert speaker_wav_ids == [id(ref_audio)]
+    assert ref_code_len == 2
+    assert torch.equal(out_ref_code, ref_code)