[Feature] Add precomputed custom voices and Qwen3-TTS ref-context cache (#3492)

Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Dan250124 <416947747@qq.com>
Co-authored-by: Yueqian Lin <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>

tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py

@@ -25,6 +25,7 @@ class _FakeDecoder(nn.Module):
         self.total_upsample = total_upsample
         self.decode_calls: list[dict[str, int]] = []
         self.batched_decode_calls: list[dict[str, int]] = []
+        self.decode_codes: list[torch.Tensor] = []
         self.cudagraph_calls: list[dict[str, int | torch.device]] = []
 
     def to(self, *args, **kwargs):
@@ -37,6 +38,7 @@ class _FakeDecoder(nn.Module):
         chunk_size: int = 300,
         left_context_size: int = 25,
     ) -> torch.Tensor:
+        self.decode_codes.append(codes.detach().cpu().clone())
         self.decode_calls.append(
             {
                 "chunk_size": chunk_size,
@@ -181,6 +183,136 @@ def test_forward_trims_trailing_padding_without_context():
     torch.testing.assert_close(audio, expected)
 
 
+def test_forward_reuses_cached_ref_context_for_followup_chunk():
+    model = _make_model()
+
+    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": {
+                    "left_context_size": 2,
+                    "ref_context_size": 2,
+                    "ref_context_request_id": "rid",
+                    "ref_context_included": True,
+                }
+            }
+        ],
+    )
+
+    followup_codes = torch.tensor([5, 6, 7, 15, 16, 17], dtype=torch.long)
+    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),
+                }
+            }
+        ],
+    )
+
+    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
+
+
 def test_connector_codec_chunking_does_not_override_decode_chunking():
     model = _make_model(
         async_chunk=True,

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -2,6 +2,7 @@ from collections import OrderedDict
 from types import SimpleNamespace
 
 import numpy as np
+import pytest
 import torch
 
 from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
@@ -9,6 +10,7 @@ from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
     PRECOMPUTED_REF_CODE_KEY,
     PRECOMPUTED_REF_IDS_KEY,
     PRECOMPUTED_TEXT_IDS_KEY,
+    REF_AUDIO_CACHE_KEY,
     Qwen3TTSPromptEmbedsBuilder,
 )
 from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
@@ -493,6 +495,41 @@ def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exi
     assert NORMALIZED_REF_AUDIO_KEY not in buf["r1"]
 
 
+def test_base_voice_clone_batch_preprocess_uses_serving_artifact_cache_key_without_normalize():
+    builder = _make_minimal_builder()
+    ref_code = torch.full((2, 2), 5, dtype=torch.long)
+    builder.put_ref_audio_artifacts("same-ref", ref_code=ref_code)
+    builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
+    builder._encode_ref_audio_batch_fn = lambda *a, **kw: (_ for _ in ()).throw(
+        AssertionError("speech tokenizer not expected")
+    )
+
+    class FakeTextTokenizer:
+        def __call__(self, texts, *, padding=False):
+            return {"input_ids": [[7, 8, 9] for _ in texts]}
+
+    builder._text_tokenizer = FakeTextTokenizer()
+    buf = {
+        "r1": {
+            "task_type": ["Base"],
+            "text": ["one"],
+            "ref_audio": ["a.wav"],
+            "ref_text": ["hello"],
+            "x_vector_only_mode": [False],
+            REF_AUDIO_CACHE_KEY: ["same-ref"],
+        }
+    }
+
+    builder.preprocess_batch(
+        req_ids=["r1"],
+        model_intermediate_buffer=buf,
+        device=torch.device("cpu"),
+    )
+
+    assert torch.equal(buf["r1"]["codes"][PRECOMPUTED_REF_CODE_KEY], ref_code)
+    assert NORMALIZED_REF_AUDIO_KEY not in buf["r1"]
+
+
 def test_base_voice_clone_uses_batched_ref_code_without_serial_encode():
     builder = _make_minimal_builder()
     device_param = torch.nn.Parameter(torch.empty(0))
@@ -536,3 +573,112 @@ def test_base_voice_clone_uses_batched_ref_code_without_serial_encode():
     assert speaker_wav_ids == [id(ref_audio)]
     assert ref_code_len == 2
     assert torch.equal(out_ref_code, ref_code)
+
+
+def test_voice_clone_prompt_artifacts_do_not_resolve_ref_audio_for_cache_fill():
+    builder = _make_minimal_builder()
+    device_param = torch.nn.Parameter(torch.empty(0))
+    builder._text_embedding = _stub_text_embedding(device_param)
+    builder._text_projection = lambda embeds: embeds
+    builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+
+    class FakeTokenizer:
+        def __call__(self, *_args, **_kwargs):
+            return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
+
+    builder._text_tokenizer = FakeTokenizer()
+    builder._generate_icl_prompt = lambda **kwargs: (
+        torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
+        torch.ones((1, 4), device=kwargs["ref_code"].device),
+    )
+    builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
+    builder.encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw(AssertionError("encode not expected"))
+    builder.extract_speaker_embedding = lambda _wav, _sr: (_ for _ in ()).throw(
+        AssertionError("speaker embedding not expected")
+    )
+    ref_code = torch.arange(4, dtype=torch.long).reshape(2, 2)
+
+    _prompt, _trailing, ref_code_len, out_ref_code = builder.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],
+            "voice_clone_prompt": {
+                "ref_code": ref_code,
+                "ref_spk_embedding": torch.ones(4, dtype=torch.bfloat16),
+                "icl_mode": True,
+            },
+        },
+    )
+
+    assert ref_code_len == 2
+    assert torch.equal(out_ref_code, ref_code)
+
+
+def test_ref_audio_artifact_only_uses_cache_without_ref_audio_payload():
+    builder = _make_minimal_builder()
+    device_param = torch.nn.Parameter(torch.empty(0))
+    builder._text_embedding = _stub_text_embedding(device_param)
+    builder._text_projection = lambda embeds: embeds
+    builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+
+    class FakeTokenizer:
+        def __call__(self, *_args, **_kwargs):
+            return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
+
+    builder._text_tokenizer = FakeTokenizer()
+    builder._generate_icl_prompt = lambda **kwargs: (
+        torch.ones((1, 2, 4), device=kwargs["ref_code"].device),
+        torch.ones((1, 4), device=kwargs["ref_code"].device),
+    )
+    builder.normalize_ref_audio = lambda _raw: (_ for _ in ()).throw(AssertionError("normalize not expected"))
+    builder.encode_ref_audio_to_code = lambda _wav, _sr: (_ for _ in ()).throw(AssertionError("encode not expected"))
+    builder.extract_speaker_embedding = lambda _wav, _sr: (_ for _ in ()).throw(
+        AssertionError("speaker embedding not expected")
+    )
+    ref_code = torch.arange(4, dtype=torch.long).reshape(2, 2)
+    builder.put_ref_audio_artifacts(
+        "same-ref",
+        ref_code=ref_code,
+        ref_spk_embedding=torch.ones(4, dtype=torch.bfloat16),
+    )
+
+    _prompt, _trailing, ref_code_len, out_ref_code = builder.build_prompt_embeds(
+        task_type="Base",
+        info_dict={
+            "text": ["hello"],
+            "ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
+            "non_streaming_mode": [False],
+            REF_AUDIO_CACHE_KEY: ["same-ref"],
+        },
+    )
+
+    assert ref_code_len == 2
+    assert torch.equal(out_ref_code, ref_code)
+
+
+def test_ref_audio_artifact_only_cache_miss_fails_fast():
+    builder = _make_minimal_builder()
+    device_param = torch.nn.Parameter(torch.empty(0))
+    builder._text_embedding = _stub_text_embedding(device_param)
+    builder._text_projection = lambda embeds: embeds
+    builder._codec_embed = lambda ids: torch.zeros((*ids.shape, 4), device=ids.device)
+
+    class FakeTokenizer:
+        def __call__(self, *_args, **_kwargs):
+            return {"input_ids": torch.arange(8, dtype=torch.long).reshape(1, -1)}
+
+    builder._text_tokenizer = FakeTokenizer()
+
+    with pytest.raises(RuntimeError, match="artifact cache miss"):
+        builder.build_prompt_embeds(
+            task_type="Base",
+            info_dict={
+                "text": ["hello"],
+                "ref_ids": torch.arange(8, dtype=torch.long).reshape(1, -1),
+                "non_streaming_mode": [False],
+                REF_AUDIO_CACHE_KEY: ["missing-ref"],
+            },
+        )

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_prompt_len.py

@@ -79,6 +79,27 @@ def test_estimate_prompt_len_handles_1d_ref_code() -> None:
     assert est > 50, f"got {est}; 1D ref_code must contribute its own length"
 
 
+def test_estimate_prompt_len_uses_ref_code_length_without_ref_audio() -> None:
+    est = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+        additional_information={
+            "task_type": ["Base"],
+            "text": ["hello"],
+            "ref_text": ["reference transcript"],
+            "x_vector_only_mode": [False],
+            "ref_code_length": [6],
+            "_qwen3_tts_ref_audio_cache_key": ["same-ref"],
+            "language": ["English"],
+        },
+        task_type="Base",
+        tokenize_prompt=lambda _text: list(range(10)),
+        codec_language_id=None,
+        spk_is_dialect=None,
+        estimate_ref_code_len=lambda _ref_audio: None,
+    )
+
+    assert est == 15
+
+
 def test_estimate_prompt_len_uses_list_of_int_ref_ids_from_voice_clone_prompt() -> None:
     """`voice_clone_prompt.ref_ids` as a list-of-int (the natural pre-tokenized shape)
     must be read as a sequence, not unwrapped to its first element. Without the fix