[Model][Core] Enable async_chunk streaming pipeline for CosyVoice3 (#1703)

Signed-off-by: linyueqian <linyueqian@outlook.com>
Signed-off-by: indevn <indevn@outlook.com>
Co-authored-by: linyueqian <linyueqian@outlook.com>

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -133,6 +133,22 @@ def test_save_async(build_adapter):
     assert task["is_finished"] is False
 
 
+def test_send_single_request_cleans_up_after_finished_payload(build_adapter, monkeypatch):
+    adapter, _ = build_adapter(stage_id=1)
+    request = _req("req-finished", RequestStatus.FINISHED_STOPPED, external_req_id="ext-finished")
+
+    adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": True}
+    cleanup_calls = []
+    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
+
+    adapter._send_single_request({"pooling_output": None, "request": request, "is_finished": True})
+
+    assert len(cleanup_calls) == 1
+    args, _ = cleanup_calls[0]
+    assert args[0] == "req-finished"
+    assert args[1] == "ext-finished"
+
+
 def test_update_request_payload(build_adapter):
     adapter, _ = build_adapter()
 
@@ -409,3 +425,86 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
     args, _ = cleanup_calls[0]
     assert args[0] == "req-s1"
     assert args[1] == "ext-s1"
+
+
+def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerFixture):
+    """OmniARScheduler should enqueue save; adapter cleanup is handled in save thread."""
+    cleanup_calls = []
+    save_calls = []
+
+    adapter_mock = mocker.MagicMock()
+    adapter_mock.cleanup = lambda *a, **kw: cleanup_calls.append((a, kw))
+    adapter_mock.save_async = lambda *a, **kw: save_calls.append((a, kw))
+
+    from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+
+    scheduler = mocker.MagicMock()
+    scheduler.chunk_transfer_adapter = adapter_mock
+    scheduler.connector = None
+    scheduler.perf_metrics = None
+    scheduler.log_stats = False
+    scheduler.recompute_kv_load_failures = False
+    scheduler.structured_output_manager = mocker.MagicMock()
+    scheduler.structured_output_manager.should_advance.return_value = False
+    scheduler.finished_req_ids_dict = {}
+    scheduler.kv_cache_manager = mocker.MagicMock()
+    scheduler.kv_cache_manager.take_events.return_value = None
+    scheduler.kv_event_publisher = mocker.MagicMock()
+    scheduler.waiting_for_transfer_free = set()
+    scheduler.transfer_triggered_requests = set()
+    scheduler.active_kv_transfers = set()
+
+    request = _HashableRequest(
+        request_id="req-ar",
+        external_req_id="ext-ar",
+        status=RequestStatus.RUNNING,
+        is_finished=lambda: False,
+        num_computed_tokens=1,
+        num_prompt_tokens=1,
+        prompt_token_ids=[1],
+        num_output_placeholders=0,
+        sampling_params=None,
+        pooling_params=None,
+        stop_reason=None,
+        client_index=0,
+        take_events=lambda: [],
+        trace_headers=None,
+        num_cached_tokens=0,
+        num_external_computed_tokens=0,
+        num_nans_in_logits=0,
+        get_finished_reason=lambda: "stop",
+    )
+    scheduler.requests = {"req-ar": request}
+
+    scheduler._update_request_with_output = mocker.MagicMock(return_value=([], True))
+    scheduler._process_kv_transfer_trigger = mocker.MagicMock(return_value=False)
+    scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
+    scheduler._free_request = mocker.MagicMock(return_value=None)
+    scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
+    scheduler.running = [request]
+    scheduler.waiting = mocker.MagicMock()
+    scheduler.waiting.remove_requests = mocker.MagicMock()
+    scheduler.make_spec_decoding_stats = mocker.MagicMock(return_value=None)
+    scheduler.make_stats = mocker.MagicMock(return_value=None)
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"req-ar": 1},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    model_runner_output = SimpleNamespace(
+        sampled_token_ids=[[123]],
+        logprobs=None,
+        prompt_logprobs_dict={},
+        pooler_output=None,
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        req_id_to_index={"req-ar": 0},
+        kv_extracted_req_ids=None,
+    )
+
+    OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
+
+    assert len(cleanup_calls) == 0
+    assert len(save_calls) == 1

tests/e2e/offline_inference/test_cosyvoice3.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Offline E2E smoke test for CosyVoice3 zero-shot reference inference.
+
+This test uses the official upstream zero-shot prompt text/audio pair and
+verifies a stable reference recipe:
+- config-derived top_p/top_k and token-length ratios
+- model EOS token as the stop token
+- a conservative repetition penalty to avoid degenerate loops
+"""
+
+from __future__ import annotations
+
+import functools
+import io
+import os
+import tempfile
+from pathlib import Path
+from urllib.request import urlopen
+
+import numpy as np
+import pytest
+import soundfile as sf
+import yaml
+from huggingface_hub import snapshot_download
+from vllm.sampling_params import SamplingParams
+
+from tests.conftest import OmniRunner
+from tests.utils import hardware_test
+from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
+from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
+
+MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
+MODEL_DIR_ENV = "VLLM_OMNI_COSYVOICE3_MODEL_DIR"
+
+REFERENCE_PROMPT_WAV_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
+REFERENCE_PROMPT_TEXT = "You are a helpful assistant.<|endofprompt|>希望你以后能够做的比我还好呦。"
+REFERENCE_SYNTH_TEXT = (
+    "CosyVoice is undergoing a comprehensive upgrade, providing more accurate, "
+    "stable, faster, and better voice generation capabilities."
+)
+REFERENCE_STAGE0_TEMPERATURE = 1.0
+REFERENCE_STAGE0_REPETITION_PENALTY = 2.0
+
+
+def _stage_config(name: str) -> str:
+    return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
+
+
+STAGE_CONFIGS = [
+    _stage_config("cosyvoice3.yaml"),
+    _stage_config("cosyvoice3_async_chunk.yaml"),
+]
+
+
+@functools.lru_cache(maxsize=1)
+def _load_reference_prompt_wav() -> tuple[np.ndarray, int]:
+    with urlopen(REFERENCE_PROMPT_WAV_URL, timeout=30) as resp:
+        data = resp.read()
+    audio, sr = sf.read(io.BytesIO(data), dtype="float32", always_2d=False)
+    if isinstance(audio, np.ndarray) and audio.ndim > 1:
+        audio = np.mean(audio, axis=-1)
+    return np.asarray(audio, dtype=np.float32), int(sr)
+
+
+@functools.lru_cache(maxsize=1)
+def _resolve_model_dir() -> Path:
+    override = os.environ.get(MODEL_DIR_ENV)
+    if override:
+        return Path(override).expanduser().resolve()
+    return Path(snapshot_download(MODEL, allow_patterns=["*"]))
+
+
+def _reference_zero_shot_stage0_sampling(*, text: str) -> SamplingParams:
+    config = CosyVoice3Config()
+    sampling_cfg = config.llm.get("sampling", {})
+    eos_token_id = int(config.llm["eos_token_id"])
+    model_dir = _resolve_model_dir()
+    tokenizer = get_qwen_tokenizer(
+        token_path=str(model_dir / config.qwen_pretrain_path),
+        skip_special_tokens=config.skip_special_tokens,
+        version=config.version,
+    )
+    text_len = max(1, len(tokenizer.encode(text, allowed_special=config.allowed_special)))
+    return SamplingParams(
+        temperature=REFERENCE_STAGE0_TEMPERATURE,
+        top_p=float(sampling_cfg.get("top_p", 0.8)),
+        top_k=int(sampling_cfg.get("top_k", 25)),
+        repetition_penalty=REFERENCE_STAGE0_REPETITION_PENALTY,
+        stop_token_ids=[eos_token_id],
+        min_tokens=int(text_len * config.min_token_text_ratio),
+        max_tokens=int(text_len * config.max_token_text_ratio),
+    )
+
+
+def _concat_audio(audio_val) -> np.ndarray:
+    import torch
+
+    if isinstance(audio_val, list):
+        tensors = []
+        for t in audio_val:
+            if t is None:
+                continue
+            if hasattr(t, "detach"):
+                t = t.detach()
+            if hasattr(t, "cpu"):
+                t = t.cpu()
+            if hasattr(t, "float"):
+                t = t.float()
+            if isinstance(t, torch.Tensor):
+                tensors.append(t.reshape(-1))
+        if not tensors:
+            return np.zeros((0,), dtype=np.float32)
+        return torch.cat(tensors, dim=-1).numpy().astype(np.float32, copy=False)
+
+    if hasattr(audio_val, "detach"):
+        audio_val = audio_val.detach()
+    if hasattr(audio_val, "cpu"):
+        audio_val = audio_val.cpu()
+    if hasattr(audio_val, "float"):
+        audio_val = audio_val.float()
+    if hasattr(audio_val, "numpy"):
+        audio_val = audio_val.numpy()
+    audio_np = np.asarray(audio_val, dtype=np.float32)
+    return audio_np.reshape(-1)
+
+
+def _get_stage_engine_outputs(omni_runner: OmniRunner, stage_id: int):
+    stage_list = getattr(omni_runner.omni, "stage_list", None)
+    if stage_list is not None:
+        return getattr(stage_list[stage_id], "engine_outputs", None) or []
+
+    stage_clients = getattr(getattr(omni_runner.omni, "engine", None), "stage_clients", None)
+    if stage_clients is not None:
+        return getattr(stage_clients[stage_id], "engine_outputs", None) or []
+
+    raise AttributeError("Unable to locate stage outputs on Omni runner")
+
+
+def _patched_stage_config(base_stage_config: str, model_dir: Path, tmp_dir: Path) -> str:
+    cfg = yaml.safe_load(Path(base_stage_config).read_text(encoding="utf-8"))
+    tokenizer_path = str(model_dir / "CosyVoice-BlankEN")
+    for stage in cfg.get("stage_args", []):
+        engine_args = stage.setdefault("engine_args", {})
+        engine_args["tokenizer"] = tokenizer_path
+        engine_args["enforce_eager"] = True
+        engine_args["hf_overrides"] = {"architectures": ["CosyVoice3Model"]}
+    out_path = tmp_dir / Path(base_stage_config).name
+    out_path.write_text(yaml.safe_dump(cfg, sort_keys=False), encoding="utf-8")
+    return str(out_path)
+
+
+def _build_reference_inputs(prompt_audio: tuple[np.ndarray, int]) -> list[dict[str, object]]:
+    return [
+        {
+            "prompt": REFERENCE_SYNTH_TEXT,
+            "multi_modal_data": {"audio": prompt_audio},
+            "modalities": ["audio"],
+            "mm_processor_kwargs": {"prompt_text": REFERENCE_PROMPT_TEXT},
+        }
+    ]
+
+
+@pytest.mark.core_model
+@pytest.mark.omni
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("base_stage_config", STAGE_CONFIGS)
+def test_cosyvoice3_offline_reference_zero_shot(base_stage_config: str) -> None:
+    """CosyVoice3 zero-shot reference inference should stop cleanly and produce sane audio."""
+    prompt_audio, prompt_sr = _load_reference_prompt_wav()
+    model_dir = _resolve_model_dir()
+    expected_stop_token = int(CosyVoice3Config().llm["eos_token_id"])
+
+    with tempfile.TemporaryDirectory(prefix="cv3-e2e-") as tmp:
+        stage_config = _patched_stage_config(base_stage_config, model_dir, Path(tmp))
+        with OmniRunner(
+            str(model_dir), seed=42, stage_configs_path=stage_config, stage_init_timeout=300
+        ) as omni_runner:
+            sampling_params_list = omni_runner.get_default_sampling_params_list()
+            sampling_params_list[0] = _reference_zero_shot_stage0_sampling(text=REFERENCE_SYNTH_TEXT)
+
+            outputs = omni_runner.omni.generate(
+                _build_reference_inputs((prompt_audio, prompt_sr)), sampling_params_list
+            )
+
+            assert outputs, "No outputs returned"
+            audio_mm = outputs[0].multimodal_output
+            assert "audio" in audio_mm, "No audio output found"
+
+            audio = _concat_audio(audio_mm["audio"])
+            assert audio.size > 0, "Generated audio is empty"
+
+            sr_val = audio_mm.get("sr", 24000)
+            if isinstance(sr_val, list) and sr_val:
+                sr_val = sr_val[-1]
+            if hasattr(sr_val, "item"):
+                sr_val = sr_val.item()
+            sr = int(sr_val)
+            assert sr == 24000, f"Unexpected sample_rate={sr}"
+
+            duration_s = audio.size / sr
+            assert 2.8 <= duration_s <= 8.8, f"Unexpected duration={duration_s:.3f}s (samples={audio.size}, sr={sr})"
+
+            stage0_outputs = _get_stage_engine_outputs(omni_runner, 0)
+            if stage0_outputs:
+                completion = stage0_outputs[0].outputs[0]
+                finish_reason = getattr(completion, "finish_reason", None)
+                stop_reason = getattr(completion, "stop_reason", None)
+                num_tokens = len(getattr(completion, "token_ids", []) or [])
+
+                assert finish_reason == "stop", f"Stage-0 finish_reason={finish_reason}, expected 'stop'"
+                assert int(stop_reason) == expected_stop_token, (
+                    f"Stage-0 stop_reason={stop_reason}, expected {expected_stop_token}"
+                )
+                assert 80 <= num_tokens <= 220, f"Stage-0 num_tokens={num_tokens}, expected sane stop-bound range"
+            else:
+                assert "async_chunk" in Path(base_stage_config).name, "Stage-0 produced no engine outputs"

tests/e2e/online_serving/test_cosyvoice3_tts.py

@@ -50,8 +50,18 @@ tts_server_params = [
     )
 ]
 
+tts_async_chunk_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_stage_config("cosyvoice3_async_chunk.yaml"),
+            server_args=["--trust-remote-code", "--disable-log-stats"],
+        ),
+        id="cosyvoice3_async_chunk",
+    )
+]
+
 
-@pytest.mark.advanced_model
 @pytest.mark.core_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
@@ -76,17 +86,16 @@ def test_voice_clone_zh_001(omni_server, openai_client) -> None:
     openai_client.send_audio_speech_request(request_config)
 
 
-@pytest.mark.advanced_model
+@pytest.mark.core_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
-@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
-@pytest.mark.skip(reason="CosyVoice3 does not support async_chunk streaming yet")
+@pytest.mark.parametrize("omni_server", tts_async_chunk_server_params, indirect=True)
 def test_voice_clone_zh_002(omni_server, openai_client) -> None:
     """
-    Test voice cloning TTS with Chinese text via OpenAI API.
-    Deploy Setting: default yaml
+    Test voice cloning TTS with Chinese text via async_chunk streaming.
+    Deploy Setting: cosyvoice3_async_chunk.yaml
     Input Modal: text + ref_audio + ref_text
-    Output Modal: audio
+    Output Modal: audio (streamed)
     Input Setting: stream=True
     Datasets: single request
     """
@@ -101,7 +110,7 @@ def test_voice_clone_zh_002(omni_server, openai_client) -> None:
     openai_client.send_audio_speech_request(request_config)
 
 
-@pytest.mark.advanced_model
+@pytest.mark.core_model
 @pytest.mark.omni
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)

tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Unit tests for CosyVoice3 components."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 import torch.nn as nn
@@ -247,3 +249,32 @@ class TestSDPAFallback:
 
         assert out.shape == (batch, seq_len, heads, dim)
         assert out.dtype == torch.float32
+
+
+def test_code2wav_forward_finalizes_hift_tail():
+    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav
+
+    class DummyHiFT(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.m_source = SimpleNamespace(l_linear=SimpleNamespace(weight=torch.ones(1, dtype=torch.float32)))
+            self.finalize_calls: list[bool] = []
+
+        def inference(self, speech_feat, finalize=True):
+            self.finalize_calls.append(bool(finalize))
+            return torch.zeros((speech_feat.shape[0], 1, speech_feat.shape[-1]), dtype=speech_feat.dtype), None
+
+    model = object.__new__(CosyVoice3Code2Wav)
+    nn.Module.__init__(model)
+    model.hift = DummyHiFT()
+    model._forward_mel = lambda **_: torch.ones((1, 80, 8), dtype=torch.float32)
+
+    out = model.forward(
+        token=torch.tensor([[1, 2, 3]], dtype=torch.int32),
+        prompt_token=torch.tensor([[4, 5]], dtype=torch.int32),
+        prompt_feat=torch.ones((1, 4, 80), dtype=torch.float32),
+        embedding=torch.ones((1, 192), dtype=torch.float32),
+    )
+
+    assert out.shape == (1, 1, 8)
+    assert model.hift.finalize_calls == [True]

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -0,0 +1,463 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from threading import Lock
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.nn as nn
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.logits_processor.state import LogitsProcessors
+from vllm.v1.sample.metadata import SamplingMetadata
+
+from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import CosyVoice3Model
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyCode2Wav:
+    def __init__(
+        self,
+        vocab_size: int,
+        num_samples: int = 32,
+        outputs: list[tuple[torch.Tensor, dict[str, object] | None]] | None = None,
+    ):
+        self.input_embedding = SimpleNamespace(num_embeddings=vocab_size)
+        self.num_samples = num_samples
+        self.outputs = list(outputs or [])
+        self.forward_calls: list[dict[str, object]] = []
+        self.forward_streaming_calls: list[dict[str, object]] = []
+
+    def forward(self, **kwargs):
+        self.forward_calls.append(kwargs)
+        token = kwargs["token"]
+        num_samples = int(token.shape[-1])
+        return torch.linspace(-1.0, 1.0, max(num_samples, 1), dtype=torch.float32).reshape(1, 1, -1)
+
+    def forward_streaming(self, **kwargs):
+        self.forward_streaming_calls.append(kwargs)
+        if self.outputs:
+            return self.outputs.pop(0)
+
+        token = kwargs["token"]
+        num_samples = int(token.shape[-1])
+        audio = torch.linspace(-1.0, 1.0, max(num_samples, 1), dtype=torch.float32).reshape(1, 1, -1)
+        new_state = None
+        if not kwargs.get("finalize", False):
+            new_state = {
+                "mel": torch.ones((1, 80, max(num_samples, 1)), dtype=torch.float32),
+                "speech_offset": audio.shape[-1],
+            }
+        return audio, new_state
+
+
+def _make_code2wav_model(
+    *,
+    with_stride_cfg: bool = False,
+    num_samples: int = 32,
+    outputs: list[tuple[torch.Tensor, dict[str, object] | None]] | None = None,
+) -> CosyVoice3Model:
+    model = object.__new__(CosyVoice3Model)
+    nn.Module.__init__(model)
+    model.model_stage = "cosyvoice3_code2wav"
+    hift_cfg = {} if not with_stride_cfg else {"upsample_rates": [8, 5, 3], "istft_params": {"hop_len": 4}}
+    model.config = SimpleNamespace(
+        sample_rate=24000,
+        hift=hift_cfg,
+        token_frame_rate=25 if with_stride_cfg else 0,
+        token_mel_ratio=2 if with_stride_cfg else 0,
+    )
+    model.code2wav = _DummyCode2Wav(vocab_size=4, num_samples=num_samples, outputs=outputs)
+    model.source_cache_len = 4
+    model.speech_window = torch.hamming_window(8, periodic=False)
+    model._stream_audio_cache_by_req = {}
+    model._stream_audio_cache_lock = Lock()
+    model._stream_vocoder_cache_by_req = {}
+    return model
+
+
+def _make_talker_model() -> CosyVoice3Model:
+    model = object.__new__(CosyVoice3Model)
+    nn.Module.__init__(model)
+    model.model_stage = "cosyvoice3_talker"
+    model.config = SimpleNamespace(
+        llm={
+            "speech_token_size": 6561,
+            "eos_token_id": 6562,
+            "sampling": {
+                "top_p": 0.8,
+                "top_k": 25,
+                "win_size": 10,
+                "tau_r": 0.1,
+            },
+        },
+        vocab_size=151923,
+    )
+    return model
+
+
+def _make_sampling_metadata(
+    *,
+    output_token_ids: list[list[int]],
+    repetition_penalty: float = 2.0,
+) -> SamplingMetadata:
+    return SamplingMetadata(
+        temperature=torch.tensor([1.0], dtype=torch.float32),
+        all_greedy=False,
+        all_random=True,
+        top_p=torch.tensor([0.8], dtype=torch.float32),
+        top_k=torch.tensor([25], dtype=torch.int32),
+        generators={},
+        max_num_logprobs=None,
+        no_penalties=False,
+        prompt_token_ids=None,
+        frequency_penalties=torch.zeros(1, dtype=torch.float32),
+        presence_penalties=torch.zeros(1, dtype=torch.float32),
+        repetition_penalties=torch.tensor([repetition_penalty], dtype=torch.float32),
+        output_token_ids=output_token_ids,
+        allowed_token_ids_mask=None,
+        bad_words_token_ids={},
+        logitsprocs=LogitsProcessors(),
+    )
+
+
+def test_split_request_ids_uses_seq_token_counts():
+    ids = torch.tensor([10, 11, 12, 13, 14], dtype=torch.long)
+    chunks = CosyVoice3Model._split_request_ids(ids, [2, 2, 2])
+    assert [c.tolist() for c in chunks] == [[10, 11], [12, 13], [14]]
+
+
+def test_split_request_ids_honors_single_request_seq_token_counts():
+    ids = torch.tensor([10, 11, 12, 13, 14], dtype=torch.long)
+    chunks = CosyVoice3Model._split_request_ids(ids, [3])
+    assert [c.tolist() for c in chunks] == [[10, 11, 12]]
+
+
+def test_sanitize_codec_tokens_filters_out_of_range():
+    model = _make_code2wav_model()
+    raw = torch.tensor([-1, 0, 3, 4, 99], dtype=torch.long)
+    clean = model._sanitize_codec_tokens(raw)
+    assert clean.tolist() == [0, 3]
+
+
+def test_forward_prefers_token_offset_when_present():
+    model = _make_code2wav_model()
+
+    runtime_info = [
+        {
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "token_offset": 2,
+            "left_context_size": 1,
+        }
+    ]
+
+    out = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    assert len(out.multimodal_outputs["audio"]) == 1
+    assert out.multimodal_outputs["audio"][0].numel() > 0
+    assert len(model.code2wav.forward_streaming_calls) == 1
+    call = model.code2wav.forward_streaming_calls[0]
+    assert call["token"].shape == (1, 3)
+    assert call["token_offset_tokens"] == 2
+    assert call["finalize"] is False
+
+
+def test_forward_falls_back_to_left_context_size_for_backward_compat():
+    model = _make_code2wav_model()
+
+    runtime_info = [
+        {
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "left_context_size": 2,
+        }
+    ]
+
+    model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    assert model.code2wav.forward_streaming_calls[0]["token_offset_tokens"] == 2
+
+
+def test_forward_ignores_single_request_padded_tail_tokens():
+    model = _make_code2wav_model(with_stride_cfg=True)
+    runtime_info = [
+        {
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "token_offset": 0,
+        }
+    ]
+
+    out = model.forward(
+        input_ids=torch.tensor([0, 1, 2, 3, 3], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2, 3, 4], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    # The padded tail must not contribute to code2wav length.
+    assert out.multimodal_outputs["audio"][0].numel() == 3
+    assert model.code2wav.forward_streaming_calls[0]["token"].tolist() == [[0, 1, 2]]
+
+
+def test_forward_uses_non_stream_decode_without_chunk_metadata():
+    model = _make_code2wav_model()
+
+    runtime_info = [
+        {
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "prefix_ids": [101, 102],
+            "generated_len": 3,
+        }
+    ]
+
+    out = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+
+    assert out.multimodal_outputs["audio"][0].numel() == 3
+    assert len(model.code2wav.forward_calls) == 1
+    assert len(model.code2wav.forward_streaming_calls) == 0
+    call = model.code2wav.forward_calls[0]
+    assert call["token"].tolist() == [[0, 1, 2]]
+
+
+def test_forward_reuses_streaming_cache_state_between_chunks():
+    model = _make_code2wav_model(
+        outputs=[
+            (
+                torch.arange(4, dtype=torch.float32).reshape(1, 1, -1),
+                {"mel": torch.ones((1, 80, 3), dtype=torch.float32), "speech_offset": 4},
+            ),
+            (
+                torch.full((1, 1, 2), 9.0, dtype=torch.float32),
+                {"mel": torch.ones((1, 80, 5), dtype=torch.float32), "speech_offset": 6},
+            ),
+        ]
+    )
+    runtime_info = [
+        {
+            "req_id": ["rid-stream"],
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "token_offset": 0,
+            "stream_finished": torch.tensor(False),
+        }
+    ]
+
+    out1 = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+    assert out1.multimodal_outputs["audio"][0].tolist() == [0.0, 1.0, 2.0, 3.0]
+    assert model.code2wav.forward_streaming_calls[0]["cache_state"] is None
+
+    out2 = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+    assert out2.multimodal_outputs["audio"][0].tolist() == [9.0, 9.0]
+    cache_state = model.code2wav.forward_streaming_calls[1]["cache_state"]
+    assert cache_state is not None
+    assert cache_state["speech_offset"] == 4
+    assert "rid-stream" in model._stream_vocoder_cache_by_req
+
+
+def test_forward_clears_streaming_cache_on_terminal_chunk():
+    model = _make_code2wav_model(
+        outputs=[
+            (
+                torch.arange(4, dtype=torch.float32).reshape(1, 1, -1),
+                {"mel": torch.ones((1, 80, 3), dtype=torch.float32), "speech_offset": 4},
+            ),
+            (
+                torch.full((1, 1, 1), 7.0, dtype=torch.float32),
+                None,
+            ),
+        ]
+    )
+    runtime_info = [
+        {
+            "req_id": ["rid-stream"],
+            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            "token_offset": 0,
+            "stream_finished": torch.tensor(False),
+        }
+    ]
+
+    model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+    assert "rid-stream" in model._stream_vocoder_cache_by_req
+
+    runtime_info[0]["stream_finished"] = torch.tensor(True)
+    out = model.forward(
+        input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
+        positions=torch.tensor([0, 1, 2], dtype=torch.long),
+        model_intermediate_buffer=runtime_info,
+        seq_token_counts=[3],
+    )
+    assert out.multimodal_outputs["audio"][0].tolist() == [7.0]
+    assert "rid-stream" not in model._stream_vocoder_cache_by_req
+
+
+def test_sample_uses_ras_rejection_for_recent_repetition():
+    model = _make_talker_model()
+    metadata = _make_sampling_metadata(output_token_ids=[[1] * 10])
+    logits = torch.tensor([[-1e9, 10.0, 0.0]], dtype=torch.float32)
+
+    out = model.sample(logits, metadata)
+
+    assert out is not None
+    assert out.sampled_token_ids.tolist() == [[2]]
+
+
+def test_sample_tolerates_padded_rows_without_history():
+    model = _make_talker_model()
+    metadata = _make_sampling_metadata(output_token_ids=[[1] * 10])
+    logits = torch.tensor(
+        [
+            [-1e9, 10.0, 0.0],
+            [-1e9, 0.0, 10.0],
+        ],
+        dtype=torch.float32,
+    )
+
+    out = model.sample(logits, metadata)
+
+    assert out is not None
+    assert out.sampled_token_ids.shape == (2, 1)
+
+
+def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in():
+    metadata = _make_sampling_metadata(output_token_ids=[[1, 2, 3]])
+    expected = SamplerOutput(
+        sampled_token_ids=torch.tensor([[7]], dtype=torch.int32),
+        logprobs_tensors=None,
+    )
+    calls: list[torch.Tensor] = []
+
+    class _DummyInputBatch:
+        def __init__(self):
+            self.sampling_metadata = metadata
+            self.updated = False
+
+        def update_async_output_token_ids(self):
+            self.updated = True
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.input_batch = _DummyInputBatch()
+    runner.model = SimpleNamespace(
+        prefer_model_sampler=True,
+        sample=lambda logits, sampling_metadata: calls.append(logits.clone()) or expected,
+    )
+    runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
+
+    out = runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
+
+    assert out is expected
+    assert runner.input_batch.updated is False
+    assert len(calls) == 1
+
+
+def test_gpu_ar_model_runner_supplies_req_output_history_to_model_sampler():
+    metadata = _make_sampling_metadata(output_token_ids=[])
+    seen_histories: list[list[list[int]]] = []
+
+    class _DummyInputBatch:
+        def __init__(self):
+            self.sampling_metadata = metadata
+            self.req_output_token_ids = [[1, 2, 3]]
+            self.req_ids = ["rid-1"]
+            self.sampled_token_ids_cpu = None
+            self.async_copy_ready_event = None
+            self.prev_req_id_to_index = None
+
+        def update_async_output_token_ids(self):
+            raise AssertionError("fallback async repair should not run for model sampler path")
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.input_batch = _DummyInputBatch()
+    runner.model = SimpleNamespace(
+        prefer_model_sampler=True,
+        sample=lambda logits, sampling_metadata: seen_histories.append(
+            [list(x) for x in sampling_metadata.output_token_ids]
+        )
+        or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None),
+    )
+    runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
+
+    runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
+
+    assert seen_histories == [[[1, 2, 3]]]
+
+
+def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler():
+    metadata = _make_sampling_metadata(output_token_ids=[])
+    seen_histories: list[list[list[int]]] = []
+
+    class _ReadyEvent:
+        def __init__(self):
+            self.synced = False
+
+        def synchronize(self):
+            self.synced = True
+
+    class _DummyInputBatch:
+        def __init__(self):
+            self.sampling_metadata = metadata
+            self.req_output_token_ids = [[11, -1]]
+            self.req_ids = ["rid-1"]
+            self.sampled_token_ids_cpu = torch.tensor([[29]], dtype=torch.int32)
+            self.async_copy_ready_event = _ReadyEvent()
+            self.prev_req_id_to_index = {"rid-1": 0}
+
+        def update_async_output_token_ids(self):
+            raise AssertionError("fallback async repair should not run for model sampler path")
+
+    runner = object.__new__(GPUARModelRunner)
+    runner.input_batch = _DummyInputBatch()
+    runner.model = SimpleNamespace(
+        prefer_model_sampler=True,
+        sample=lambda logits, sampling_metadata: seen_histories.append(
+            [list(x) for x in sampling_metadata.output_token_ids]
+        )
+        or SamplerOutput(sampled_token_ids=torch.tensor([[7]], dtype=torch.int32), logprobs_tensors=None),
+    )
+    runner.sampler = lambda **_: (_ for _ in ()).throw(AssertionError("fallback sampler should not be used"))
+
+    runner._sample(torch.tensor([[0.1, 0.2]], dtype=torch.float32), spec_decode_metadata=None)
+
+    assert runner.input_batch.async_copy_ready_event.synced is True
+    assert seen_histories == [[[11, 29]]]

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -0,0 +1,267 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import torch
+
+from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import talker2code2wav_async_chunk, text2flow
+
+
+def _source_output(request_id: str, prompt_ids: list[int], out_ids: list[int], mm: dict):
+    return SimpleNamespace(
+        request_id=request_id,
+        prompt_token_ids=prompt_ids,
+        outputs=[SimpleNamespace(token_ids=out_ids, multimodal_output=mm)],
+    )
+
+
+def _transfer_manager(
+    *,
+    chunk_frames: int = 2,
+    pre_lookahead_frames: int = 0,
+    stream_scale_factor: int = 1,
+    max_chunk_frames: int | None = None,
+):
+    if max_chunk_frames is None:
+        max_chunk_frames = chunk_frames
+    return SimpleNamespace(
+        code_prompt_token_ids=defaultdict(list),
+        request_payload={},
+        connector=SimpleNamespace(
+            config={
+                "extra": {
+                    "codec_chunk_frames": chunk_frames,
+                    "codec_pre_lookahead_frames": pre_lookahead_frames,
+                    "codec_max_chunk_frames": max_chunk_frames,
+                    "codec_stream_scale_factor": stream_scale_factor,
+                    "codec_vocab_size": 6561,
+                }
+            }
+        ),
+    )
+
+
+def test_text2flow_supports_batched_source_outputs():
+    stage_list = [
+        SimpleNamespace(
+            engine_outputs=[
+                _source_output("req-0", [10, 11], [1, 2, 3], {"speech_token": torch.tensor([[1, 2]])}),
+                _source_output("req-1", [20, 21], [4, 5], {"speech_token": torch.tensor([[3, 4]])}),
+            ]
+        )
+    ]
+
+    outputs = text2flow(stage_list=stage_list, engine_input_source=[0], prompt=None)
+
+    assert len(outputs) == 2
+    assert outputs[0]["prompt_token_ids"] == [1, 2, 3]
+    assert outputs[1]["prompt_token_ids"] == [4, 5]
+    assert outputs[0]["additional_information"]["prefix_ids"] == [10, 11]
+    assert outputs[1]["additional_information"]["prefix_ids"] == [20, 21]
+
+
+def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
+    transfer_manager = _transfer_manager()
+    request = SimpleNamespace(
+        external_req_id="rid-0",
+        output_token_ids=[1, 2, 6562, 3],
+        additional_information={
+            "speech_token": [torch.tensor([[11, 12, 13]])],
+            "speech_feat": [torch.tensor([[[0.1, 0.2], [0.3, 0.4]]])],
+            "embedding": [torch.tensor([[0.5, 0.6]])],
+        },
+        is_finished=lambda: True,
+    )
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=True,
+    )
+
+    assert payload is not None
+    assert payload["finished"].item() is True
+    assert payload["code_predictor_codes"] == [1, 2, 3]
+    assert payload["token_offset"] == 0
+    assert payload["left_context_size"] == 0
+    assert payload["req_id"] == ["rid-0"]
+    assert payload["stream_finished"].item() is True
+    assert "speech_token" in payload
+    assert "speech_feat" in payload
+    assert "embedding" in payload
+
+
+def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes():
+    transfer_manager = _transfer_manager(chunk_frames=25)
+    request = SimpleNamespace(
+        external_req_id="rid-eof",
+        output_token_ids=[6561, 6562],  # all filtered out
+        additional_information={},
+        is_finished=lambda: True,
+    )
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=True,
+    )
+
+    assert payload is not None
+    assert payload["code_predictor_codes"] == []
+    assert payload["finished"].item() is True
+
+
+def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():
+    transfer_manager = _transfer_manager()
+    request = SimpleNamespace(
+        external_req_id="rid-stable",
+        output_token_ids=[1, 2],
+        additional_information={},
+        is_finished=lambda: False,
+    )
+
+    payload1 = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+    payload2 = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+
+    assert payload1 is not None
+    assert payload1["code_predictor_codes"] == [1, 2]
+    assert payload1["token_offset"] == 0
+    assert payload2 is None
+
+
+def test_talker2code2wav_async_chunk_waits_for_prelookahead_and_emits_cumulative_prefix():
+    transfer_manager = _transfer_manager(pre_lookahead_frames=1)
+    request = SimpleNamespace(
+        external_req_id="rid-pre",
+        output_token_ids=[1, 2],
+        additional_information={},
+        is_finished=lambda: False,
+    )
+
+    payload_pending = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+    request.output_token_ids = [1, 2, 3]
+    payload_ready = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+
+    assert payload_pending is None
+    assert payload_ready is not None
+    assert payload_ready["code_predictor_codes"] == [1, 2, 3]
+    assert payload_ready["token_offset"] == 0
+    assert payload_ready["finished"].item() is False
+
+
+def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():
+    transfer_manager = _transfer_manager(pre_lookahead_frames=1)
+    request = SimpleNamespace(
+        external_req_id="rid-tail",
+        output_token_ids=[3, 4, 5],
+        additional_information={},
+        is_finished=lambda: False,
+    )
+
+    payload_stream = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+    request.output_token_ids = [3, 4, 5, 6]
+    payload_final = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=True,
+    )
+
+    assert payload_stream is not None
+    assert payload_stream["finished"].item() is False
+    assert payload_stream["code_predictor_codes"] == [3, 4, 5]
+    assert payload_stream["token_offset"] == 0
+    assert payload_final is not None
+    assert payload_final["finished"].item() is True
+    assert payload_final["code_predictor_codes"] == [3, 4, 5, 6]
+    assert payload_final["token_offset"] == 2
+
+
+def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():
+    transfer_manager = _transfer_manager(pre_lookahead_frames=1)
+    request = SimpleNamespace(
+        external_req_id="rid-pad",
+        output_token_ids=[8, 9, 10],
+        additional_information={
+            "speech_token": [torch.tensor([[1, 2, 3]])],
+        },
+        is_finished=lambda: False,
+    )
+
+    payload_pending = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+    request.output_token_ids = [8, 9, 10, 11]
+    payload_ready = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+
+    assert payload_pending is None
+    assert payload_ready is not None
+    assert payload_ready["code_predictor_codes"] == [8, 9, 10, 11]
+    assert payload_ready["token_offset"] == 0
+
+
+def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio():
+    transfer_manager = _transfer_manager()
+    request = SimpleNamespace(
+        external_req_id="rid-eof-tail",
+        output_token_ids=[3, 4],
+        additional_information={},
+        is_finished=lambda: False,
+    )
+
+    payload_stream = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=False,
+    )
+    payload_final = talker2code2wav_async_chunk(
+        transfer_manager=transfer_manager,
+        pooling_output=None,
+        request=request,
+        is_finished=True,
+    )
+
+    assert payload_stream is not None
+    assert payload_stream["finished"].item() is False
+    assert payload_stream["code_predictor_codes"] == [3, 4]
+    assert payload_final is not None
+    assert payload_final["finished"].item() is True
+    assert payload_final["code_predictor_codes"] == []