[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"] == []

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -316,10 +316,6 @@ class OmniARScheduler(VLLMScheduler):
                 finished = self._handle_stopped_request(request)
                 if finished:
                     kv_transfer_params = self._free_request(request)
-                    if self.chunk_transfer_adapter is not None:
-                        self.chunk_transfer_adapter.cleanup_receiver(
-                            request.request_id,
-                        )
                 if status_before_stop == RequestStatus.RUNNING:
                     stopped_running_reqs.add(request)
                 elif status_before_stop == RequestStatus.WAITING_FOR_CHUNK:

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -160,11 +160,15 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
                 new_ids = payload_data.get("code_predictor_codes", [])
                 request.prompt_token_ids = new_ids
-                # Pass additional fields (like left_context_size) to the request
-                # Only pass chunk context metadata in additional_information
-                request.additional_information = {}
-                if "left_context_size" in payload_data:
-                    request.additional_information["left_context_size"] = payload_data["left_context_size"]
+                # Preserve previously attached request metadata (e.g. prompt
+                # conditioning tensors) and update only per-chunk fields.
+                prev_info = getattr(request, "additional_information", None)
+                info = dict(prev_info) if isinstance(prev_info, dict) else {}
+                for key, value in payload_data.items():
+                    if key in {"code_predictor_codes", "finished"}:
+                        continue
+                    info[key] = value
+                request.additional_information = info
                 request.num_computed_tokens = 0
 
                 # Empty chunk with more data expected: keep polling.
@@ -240,9 +244,23 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         if success:
             self.put_req_chunk[external_req_id] += 1
             logger.debug(f"[Stage-{stage_id}] Sent {connector_put_key}")
+            finished_flag = payload_data.get("finished")
+            is_payload_finished = False
+            if isinstance(finished_flag, torch.Tensor):
+                is_payload_finished = finished_flag.numel() == 1 and bool(finished_flag.item())
+            elif finished_flag is not None:
+                is_payload_finished = bool(finished_flag)
+
+            # Reclaim per-request async state only after the terminal payload
+            # has been sent successfully. This avoids cleanup->save races.
+            if is_payload_finished:
+                self.cleanup(request.request_id, external_req_id)
 
         if is_finished:
-            self.cleanup_sender(external_req_id)
+            self.code_prompt_token_ids.pop(external_req_id, None)
+            cached_ic = getattr(self, "_cached_ic", None)
+            if cached_ic is not None:
+                cached_ic.pop(external_req_id, None)
 
     ########################################################################
     # Cleanup

vllm_omni/entrypoints/openai/serving_speech.py

@@ -1041,9 +1041,20 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         streaming needs per-chunk delta slicing; non-streaming needs full concatenation.
         """
         mm = getattr(res, "multimodal_output", None)
+        ro = None
         if not mm:
             ro = getattr(res, "request_output", None)
             mm = getattr(ro, "multimodal_output", None) if ro else None
+        if not mm:
+            if ro is None:
+                ro = getattr(res, "request_output", None)
+            outputs = getattr(ro, "outputs", None) if ro else None
+            if outputs:
+                for completion_output in outputs:
+                    completion_mm = getattr(completion_output, "multimodal_output", None)
+                    if completion_mm:
+                        mm = completion_mm
+                        break
         if not mm:
             return None, None
         key = "audio" if "audio" in mm else ("model_outputs" if "model_outputs" in mm else None)
@@ -1332,6 +1343,29 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
 
         sampling_params_list = self.engine_client.default_sampling_params_list
 
+        # CosyVoice3: set dynamic min/max tokens based on text length.
+        # The official model requires min_token_text_ratio to prevent early
+        # EOS and max_token_text_ratio to cap generation length.
+        if self._tts_model_type == "cosyvoice3" and sampling_params_list:
+            import copy
+
+            sampling_params_list = copy.deepcopy(sampling_params_list)
+            text_len = len(request.input)  # rough char-level estimate
+            # Use the model's configured ratios (defaults: min=2, max=20)
+            hf_cfg = self.model_config.hf_config
+            min_ratio = getattr(hf_cfg, "min_token_text_ratio", 2)
+            max_ratio = getattr(hf_cfg, "max_token_text_ratio", 20)
+            min_tokens = max(1, int(text_len * min_ratio))
+            max_tokens = min(2048, int(text_len * max_ratio))
+            sampling_params_list[0].min_tokens = min_tokens
+            sampling_params_list[0].max_tokens = max_tokens
+            logger.info(
+                "CosyVoice3 dynamic tokens: text_len=%d, min_tokens=%d, max_tokens=%d",
+                text_len,
+                min_tokens,
+                max_tokens,
+            )
+
         # Fish defaults come from stage_configs YAML. Only override when the caller
         # explicitly requests a different generation length.
         if self._is_fish_speech and request.max_new_tokens is not None and sampling_params_list:

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py

@@ -174,7 +174,7 @@ class CausalConditionalCFM(ConditionalCFM):
         super().__init__(in_channels, cfm_params, n_spks, spk_emb_dim, estimator)
 
     @torch.inference_mode()
-    def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None):
+    def forward(self, mu, mask, n_timesteps, temperature=1.0, spks=None, cond=None, streaming: bool = False):
         """Forward diffusion
 
         Args:
@@ -277,7 +277,9 @@ class CausalMaskedDiffWithDiT(torch.nn.Module):
         prompt_feat,
         prompt_feat_len,
         embedding,
-        finalize,
+        streaming: bool = True,
+        finalize: bool = False,
+        n_timesteps: int = 10,
     ):
         assert token.shape[0] == 1
         # xvec projection
@@ -314,7 +316,8 @@ class CausalMaskedDiffWithDiT(torch.nn.Module):
             mask=mask.unsqueeze(1),
             spks=embedding,
             cond=conds,
-            n_timesteps=10,
+            n_timesteps=max(1, int(n_timesteps)),
+            streaming=streaming,
         )
 
         feat = feat[:, :, mel_len1:]

vllm_omni/model_executor/models/cosyvoice3/config.py

@@ -7,7 +7,9 @@ class CosyVoice3Config(PretrainedConfig):
     model_type = "cosyvoice3"
 
     def __init__(self, **kwargs):
-        # Set speech EOS so vLLM stops generation at the right token
+        # Set primary speech EOS so vLLM stops generation at the right token.
+        # The official CosyVoice3 treats ALL tokens >= speech_token_size
+        # (6561-6760) as stop signals; see stop_token_ids in the YAML configs.
         kwargs.setdefault("eos_token_id", 6562)
         super().__init__(**kwargs)
         self.sample_rate = 24000

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -2,14 +2,16 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import os
 from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import replace
 from functools import partial
+from threading import Lock
 
-import numpy as np
 import torch
 import torch.nn as nn
 from transformers.feature_extraction_utils import BatchFeature
 from vllm.config import VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
+from vllm.forward_context import get_forward_context, is_forward_context_available
 from vllm.inputs import MultiModalDataDict
 from vllm.logger import init_logger
 from vllm.model_executor.models.interfaces import SupportsMultiModal
@@ -26,6 +28,9 @@ from vllm.multimodal.processing import (
     PromptUpdate,
 )
 from vllm.sequence import IntermediateTensors
+from vllm.v1.outputs import SamplerOutput
+from vllm.v1.sample.metadata import SamplingMetadata
+from vllm.v1.sample.sampler import Sampler
 
 from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
 from vllm_omni.model_executor.models.cosyvoice3.utils import (
@@ -267,6 +272,8 @@ class CosyVoice3Model(
     supports_multimodal_raw_input_only = True
     supports_multimodal = True
     requires_raw_input_tokens = True
+    prefer_model_sampler = True
+    _sampling_eps = 1e-5
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -305,6 +312,8 @@ class CosyVoice3Model(
             self.code2wav = CosyVoice3Code2Wav(self.config)
             self.model = self.code2wav.flow_model
             self.hift = self.code2wav.hift
+            # Keep additional information synchronized for async_chunk updates.
+            self.enable_update_additional_information = True
 
             # Expose streaming parameters
             self.token_overlap_len = self.code2wav.token_overlap_len
@@ -313,6 +322,9 @@ class CosyVoice3Model(
             self.mel_cache_len = self.code2wav.mel_cache_len
             self.source_cache_len = self.code2wav.source_cache_len
             self.speech_window = self.code2wav.speech_window
+            self._stream_audio_cache_by_req: dict[str, torch.Tensor] = {}
+            self._stream_audio_cache_lock = Lock()
+            self._stream_vocoder_cache_by_req: dict[str, dict[str, torch.Tensor]] = {}
         else:
             raise ValueError(f"Model stage not supported {self.model_stage}")
 
@@ -331,19 +343,277 @@ class CosyVoice3Model(
         # Use parent's cache config - critical for PagedAttention to work correctly
         return parent_config.with_hf_config(qwen_hf_config, architectures=["Qwen2Model"])
 
+    @staticmethod
+    def _as_tensor(value: object) -> torch.Tensor | None:
+        """Extract tensor payload from runtime info fields."""
+        if isinstance(value, list):
+            if not value:
+                return None
+            value = value[0]
+        if isinstance(value, torch.Tensor):
+            return value
+        return None
+
+    @staticmethod
+    def _as_str(value: object) -> str | None:
+        """Extract string payload from runtime info fields."""
+        if isinstance(value, list):
+            if not value:
+                return None
+            value = value[0]
+        if value is None:
+            return None
+        return str(value)
+
+    @staticmethod
+    def _as_bool(value: object) -> bool:
+        """Extract boolean payload from runtime info fields."""
+        if isinstance(value, list):
+            if not value:
+                return False
+            value = value[0]
+        if isinstance(value, torch.Tensor):
+            if value.numel() == 0:
+                return False
+            return bool(value.reshape(-1)[0].item())
+        if value is None:
+            return False
+        return bool(value)
+
+    @staticmethod
+    def _cross_fade_audio(audio: torch.Tensor, prev_tail: torch.Tensor) -> torch.Tensor:
+        """Blend previous chunk tail into current chunk head using a Hamming window.
+
+        This mirrors upstream CosyVoice's `fade_in_out(...)` semantics:
+        update the current head in-place using a 2*overlap window, then
+        concatenate the unchanged remainder.
+        """
+        if audio.numel() == 0 or prev_tail.numel() == 0:
+            return audio
+        overlap = min(int(audio.numel()), int(prev_tail.numel()))
+        if overlap <= 0:
+            return audio
+        window = torch.hamming_window(2 * overlap, periodic=False, dtype=audio.dtype, device=audio.device)
+        fade_in = window[:overlap]
+        fade_out = window[overlap:]
+        blended = audio[:overlap] * fade_in + prev_tail[-overlap:].to(device=audio.device, dtype=audio.dtype) * fade_out
+        if overlap == int(audio.numel()):
+            return blended
+        return torch.cat([blended, audio[overlap:]], dim=0)
+
+    def _stitch_stream_audio(self, req_id: str | None, audio: torch.Tensor, stream_finished: bool) -> torch.Tensor:
+        """Pass-through stitching for async_chunk.
+
+        Chunk overlap is already removed in mel domain via token_offset_tokens.
+        Applying an additional waveform-domain fade/cache step introduces either
+        duplicated overlap (if no tail trim) or duration shrink (if tail trim).
+        """
+        if req_id is not None and stream_finished and hasattr(self, "_stream_audio_cache_by_req"):
+            with self._stream_audio_cache_lock:
+                self._stream_audio_cache_by_req.pop(req_id, None)
+                if hasattr(self, "_stream_vocoder_cache_by_req"):
+                    self._stream_vocoder_cache_by_req.pop(req_id, None)
+        return audio
+
+    @staticmethod
+    def _split_request_ids(ids: torch.Tensor, seq_token_counts: list[int] | None = None) -> list[torch.Tensor]:
+        """Split concatenated input_ids into per-request segments."""
+        if seq_token_counts is not None:
+            boundaries = [0]
+            for count in seq_token_counts:
+                boundaries.append(boundaries[-1] + int(count))
+            total = ids.numel()
+            return [ids[boundaries[i] : min(boundaries[i + 1], total)] for i in range(len(seq_token_counts))]
+
+        if is_forward_context_available():
+            slices = get_forward_context().ubatch_slices
+            if slices is not None and len(slices) > 1 and not any(hasattr(s, "token_slice") for s in slices):
+                boundaries = [0]
+                for s in slices:
+                    boundaries.append(boundaries[-1] + int(s))
+                return [ids[boundaries[i] : boundaries[i + 1]] for i in range(len(boundaries) - 1)]
+
+        return [ids]
+
+    def _sanitize_codec_tokens(self, req_ids: torch.Tensor) -> torch.Tensor:
+        """Filter non-code tokens before feeding flow token embedding."""
+        vocab_size = int(self.code2wav.input_embedding.num_embeddings)
+        valid_mask = (req_ids >= 0) & (req_ids < vocab_size)
+        return req_ids[valid_mask]
+
+    @staticmethod
+    def _req_scalar(param: torch.Tensor | None, req_idx: int, default: float | int) -> float | int:
+        if param is None or param.numel() == 0:
+            return default
+        index = min(req_idx, int(param.numel()) - 1)
+        value = param.reshape(-1)[index].item()
+        if isinstance(default, int):
+            return int(value)
+        return float(value)
+
+    @staticmethod
+    def _multinomial_sample(probs: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor:
+        return torch.multinomial(probs, 1, replacement=True, generator=generator).reshape(())
+
+    @classmethod
+    def _nucleus_sample_one(
+        cls,
+        weighted_scores: torch.Tensor,
+        *,
+        top_p: float,
+        top_k: int,
+        generator: torch.Generator | None,
+    ) -> int:
+        probs = weighted_scores.softmax(dim=0)
+        sorted_prob, sorted_idx = probs.sort(descending=True, stable=True)
+        kept_probs: list[torch.Tensor] = []
+        kept_indices: list[torch.Tensor] = []
+        cum_prob = 0.0
+        max_keep = len(sorted_idx) if top_k <= 0 else min(int(top_k), len(sorted_idx))
+        for i in range(len(sorted_idx)):
+            if cum_prob < top_p and len(kept_probs) < max_keep:
+                cum_prob += float(sorted_prob[i].item())
+                kept_probs.append(sorted_prob[i])
+                kept_indices.append(sorted_idx[i])
+            else:
+                break
+
+        if not kept_probs:
+            return int(sorted_idx[0].item())
+
+        sample_probs = torch.stack(kept_probs)
+        sample_idx = cls._multinomial_sample(sample_probs, generator=generator)
+        return int(torch.stack(kept_indices)[int(sample_idx.item())].item())
+
+    @classmethod
+    def _ras_sample_one(
+        cls,
+        weighted_scores: torch.Tensor,
+        decoded_tokens: Sequence[int],
+        *,
+        top_p: float,
+        top_k: int,
+        win_size: int,
+        tau_r: float,
+        generator: torch.Generator | None,
+    ) -> int:
+        top_id = cls._nucleus_sample_one(
+            weighted_scores,
+            top_p=top_p,
+            top_k=top_k,
+            generator=generator,
+        )
+        if win_size > 0 and decoded_tokens:
+            recent = torch.as_tensor(
+                list(decoded_tokens[-win_size:]),
+                device=weighted_scores.device,
+                dtype=torch.long,
+            )
+            rep_num = int((recent == top_id).sum().item())
+            if rep_num >= win_size * tau_r:
+                weighted_scores = weighted_scores.clone()
+                weighted_scores[top_id] = float("-inf")
+                fallback_probs = weighted_scores.softmax(dim=0)
+                top_id = int(cls._multinomial_sample(fallback_probs, generator=generator).item())
+        return top_id
+
+    def _cosyvoice3_ras_enabled(self, sampling_metadata: SamplingMetadata) -> bool:
+        if self.model_stage != "cosyvoice3_talker":
+            return False
+        if sampling_metadata.max_num_logprobs is not None:
+            return False
+        if sampling_metadata.temperature is None:
+            return False
+        if bool(sampling_metadata.bad_words_token_ids):
+            return False
+        if torch.any(sampling_metadata.frequency_penalties != 0):
+            return False
+        if torch.any(sampling_metadata.presence_penalties != 0):
+            return False
+        return True
+
+    def sample(
+        self,
+        logits: torch.Tensor,
+        sampling_metadata: SamplingMetadata,
+    ) -> SamplerOutput | None:
+        if logits is None or logits.numel() == 0:
+            return None
+        if self.model_stage != "cosyvoice3_talker":
+            return None
+
+        sampler = getattr(self, "_talker_sampler", None)
+        if sampler is None:
+            sampler = Sampler()
+            self._talker_sampler = sampler
+
+        if not self._cosyvoice3_ras_enabled(sampling_metadata):
+            return sampler(logits=logits, sampling_metadata=sampling_metadata)
+
+        logits = logits.to(torch.float32)
+        sampling_for_processors = replace(sampling_metadata, no_penalties=True)
+        logits = sampler.apply_logits_processors(logits, sampling_for_processors, predict_bonus_token=False)
+
+        sampling_cfg = dict(self.config.llm.get("sampling", {}))
+        default_top_p = float(sampling_cfg.get("top_p", 0.8))
+        default_top_k = int(sampling_cfg.get("top_k", 25))
+        win_size = int(sampling_cfg.get("win_size", 10))
+        tau_r = float(sampling_cfg.get("tau_r", 0.1))
+
+        sampled_ids: list[int] = []
+        for req_idx in range(int(logits.shape[0])):
+            row_logits = logits[req_idx]
+
+            temperature = float(self._req_scalar(sampling_metadata.temperature, req_idx, 1.0))
+            if temperature < self._sampling_eps:
+                sampled_ids.append(int(torch.argmax(row_logits).item()))
+                continue
+
+            top_p = float(self._req_scalar(sampling_metadata.top_p, req_idx, default_top_p))
+            top_k = int(self._req_scalar(sampling_metadata.top_k, req_idx, default_top_k))
+            generator = sampling_metadata.generators.get(req_idx)
+            weighted_scores = torch.log_softmax(row_logits / max(temperature, self._sampling_eps), dim=0)
+            decoded_tokens = (
+                sampling_metadata.output_token_ids[req_idx] if req_idx < len(sampling_metadata.output_token_ids) else []
+            )
+            sampled_ids.append(
+                self._ras_sample_one(
+                    weighted_scores,
+                    decoded_tokens,
+                    top_p=top_p,
+                    top_k=top_k,
+                    win_size=win_size,
+                    tau_r=tau_r,
+                    generator=generator,
+                )
+            )
+
+        sampled = torch.tensor(sampled_ids, device=logits.device, dtype=torch.int32)
+        return SamplerOutput(sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=None)
+
     def compute_logits(self, hidden_states: torch.Tensor | OmniOutput) -> torch.Tensor | None:
         if isinstance(hidden_states, OmniOutput):
             hidden_states = hidden_states.text_hidden_states
         if self.model_stage == "cosyvoice3_talker":
             logits = self.model.llm_decoder(hidden_states)
+            # The decoder outputs speech_token_size + 200 logits.  The official
+            # CosyVoice3 treats ALL tokens >= speech_token_size (the last 200)
+            # as stop signals.  Merge their probabilities into a single EOS
+            # token (6562) via logsumexp so that vLLM's stop_token_ids=[6562]
+            # fires with the correct aggregate stop probability.
+            speech_token_size = self.config.llm["speech_token_size"]
+            eos_idx = self.config.llm["eos_token_id"]
+            stop_logits = logits[..., speech_token_size:]  # last 200
+            merged_stop = torch.logsumexp(stop_logits, dim=-1, keepdim=True)
+            logits[..., speech_token_size:] = float("-inf")  # mask all
+            logits[..., eos_idx] = merged_stop.squeeze(-1)  # restore merged
+            # Pad to full vocab_size for vLLM token handling.
             vocab_size = self.config.vocab_size
             pad_size = vocab_size - logits.size(-1)
-            pad_shape = logits.shape[:-1] + (pad_size,)
-            pad = logits.new_full(pad_shape, float("-inf"))
-            eos_token_val = logits[..., self.config.llm["eos_token_id"]].clone()
-            logits[..., -200:] = float("-inf")
-            logits[..., self.config.llm["eos_token_id"]] = eos_token_val
-            logits = torch.cat([logits, pad], dim=-1)
+            if pad_size > 0:
+                pad_shape = logits.shape[:-1] + (pad_size,)
+                pad = logits.new_full(pad_shape, float("-inf"))
+                logits = torch.cat([logits, pad], dim=-1)
             return logits
         else:
             raise RuntimeError(f"compute_logits is only valid for {self.model_stage}.")
@@ -380,6 +650,7 @@ class CosyVoice3Model(
             hidden = int(self.config.hidden_size)
             return torch.zeros(
                 (input_ids.shape[0], hidden),
+                device=input_ids.device,
             )
         else:
             raise RuntimeError(f"embed_input_ids is not valid for {self.model_stage}.")
@@ -412,28 +683,116 @@ class CosyVoice3Model(
 
             return OmniOutput(text_hidden_states=hidden_states, multimodal_outputs=multimodal_outputs)
         elif self.model_stage == "cosyvoice3_code2wav":
-            runtime_info = kwargs.get("runtime_additional_information", [])
-            if not runtime_info:
-                length = 30 * 24000
-                audio = np.zeros((length,))
-                return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audio})
+            runtime_info = kwargs.get("model_intermediate_buffer")
+            if runtime_info is None:
+                runtime_info = kwargs.get("runtime_additional_information", [])
+            if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs:
+                logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
 
-            # Remove the last eos token and add batch dimension
-            token = input_ids[..., :-1].unsqueeze(0)
+            seq_token_counts = kwargs.get("seq_token_counts")
+            flat_ids = input_ids.reshape(-1).to(dtype=torch.long)
+            request_ids_list = self._split_request_ids(flat_ids, seq_token_counts)
 
-            # Generate audio using code2wav
-            tts_speech = self.code2wav(
-                token=token,
-                prompt_token=runtime_info[0]["speech_token"][:1],
-                prompt_feat=runtime_info[0]["speech_feat"][:1],
-                embedding=runtime_info[0]["embedding"][:1],
-                n_timesteps=10,
-            )
+            num_reqs = max(1, len(request_ids_list))
+            sample_rate = torch.tensor(int(self.config.sample_rate), dtype=torch.int32)
+            empty_audio = torch.zeros((0,), dtype=torch.float32, device=input_ids.device)
+            audios: list[torch.Tensor] = [empty_audio] * num_reqs
+            srs: list[torch.Tensor] = [sample_rate] * num_reqs
+            if not isinstance(runtime_info, list):
+                runtime_info = []
 
-            return OmniOutput(
-                text_hidden_states=None,
-                multimodal_outputs={"audio": tts_speech, "sr": torch.tensor(22050)},
-            )
+            for idx, req_ids in enumerate(request_ids_list):
+                info = runtime_info[idx] if idx < len(runtime_info) and isinstance(runtime_info[idx], dict) else {}
+                req_id = self._as_str(info.get("req_id")) if info else None
+                stream_finished = self._as_bool(info.get("stream_finished")) if info else False
+                speech_token = self._as_tensor(info.get("speech_token")) if info else None
+                speech_feat = self._as_tensor(info.get("speech_feat")) if info else None
+                embedding = self._as_tensor(info.get("embedding")) if info else None
+                if speech_token is None or speech_feat is None or embedding is None:
+                    if stream_finished and req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"):
+                        with self._stream_audio_cache_lock:
+                            self._stream_vocoder_cache_by_req.pop(req_id, None)
+                    audios[idx] = self._stitch_stream_audio(req_id, empty_audio, stream_finished)
+                    if (
+                        req_ids.numel() > 0
+                        and info
+                        and ("token_offset" in info or "left_context_size" in info or "generated_len" in info)
+                    ):
+                        info_keys = ",".join(sorted(info.keys())) if info else ""
+                        logger.warning_once(
+                            "CosyVoice3 code2wav missing prompt conditioning for non-empty codec tokens: "
+                            "raw_len=%d info_keys=%s",
+                            int(req_ids.numel()),
+                            info_keys,
+                        )
+                    continue
+
+                token = self._sanitize_codec_tokens(req_ids)
+                if token.numel() == 0:
+                    audios[idx] = self._stitch_stream_audio(req_id, empty_audio, stream_finished)
+                    if req_ids.numel() > 0:
+                        logger.warning_once(
+                            "CosyVoice3 code2wav received no valid codec tokens after filtering: "
+                            "raw_len=%d raw_range=[%d,%d] vocab_size=%d",
+                            req_ids.numel(),
+                            int(req_ids.min().item()),
+                            int(req_ids.max().item()),
+                            int(self.code2wav.input_embedding.num_embeddings),
+                        )
+                    continue
+
+                # `generated_len` is injected for many models by the generic
+                # runner, so only explicit chunk-routing fields should switch
+                # code2wav into the streaming path.
+                uses_streaming_decode = bool(info) and (
+                    "stream_finished" in info or "token_offset" in info or "left_context_size" in info
+                )
+                if uses_streaming_decode:
+                    token_offset = 0
+                    try:
+                        if info and "token_offset" in info:
+                            token_offset = max(0, int(info.get("token_offset", 0)))
+                        elif info:
+                            token_offset = max(0, int(info.get("left_context_size", 0)))
+                    except (TypeError, ValueError):
+                        token_offset = 0
+
+                    cache_state = None
+                    if req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"):
+                        with self._stream_audio_cache_lock:
+                            cache_state = self._stream_vocoder_cache_by_req.get(req_id)
+
+                    tts_speech, new_cache_state = self.code2wav.forward_streaming(
+                        token=token.unsqueeze(0),
+                        prompt_token=speech_token[:1],
+                        prompt_feat=speech_feat[:1],
+                        embedding=embedding[:1],
+                        cache_state=cache_state,
+                        n_timesteps=10,
+                        token_offset_tokens=token_offset,
+                        finalize=stream_finished,
+                    )
+
+                    if req_id is not None and hasattr(self, "_stream_vocoder_cache_by_req"):
+                        with self._stream_audio_cache_lock:
+                            if new_cache_state is None or stream_finished:
+                                self._stream_vocoder_cache_by_req.pop(req_id, None)
+                            else:
+                                self._stream_vocoder_cache_by_req[req_id] = new_cache_state
+                else:
+                    tts_speech = self.code2wav.forward(
+                        token=token.unsqueeze(0),
+                        prompt_token=speech_token[:1],
+                        prompt_feat=speech_feat[:1],
+                        embedding=embedding[:1],
+                        n_timesteps=10,
+                    )
+
+                audio = tts_speech.reshape(-1).to(dtype=torch.float32)
+
+                audios[idx] = self._stitch_stream_audio(req_id, audio, stream_finished)
+
+            return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audios, "sr": srs})
         else:
             raise ValueError(f"Unsupported model_stage: {self.model_stage}")
 

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py

@@ -11,11 +11,12 @@ This module contains the code2wav (token-to-waveform) stage which uses:
 
 from __future__ import annotations
 
+from contextlib import nullcontext
+
 import numpy as np
 import torch
 import torch.nn as nn
 from omegaconf import DictConfig
-from torch.nn import functional as F
 from vllm.logger import init_logger
 
 from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT
@@ -29,7 +30,6 @@ from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
 )
 from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer
 from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
-from vllm_omni.model_executor.models.cosyvoice3.utils import make_pad_mask
 
 logger = init_logger(__name__)
 
@@ -150,6 +150,140 @@ class CosyVoice3Code2Wav(nn.Module):
         """Speaker embedding affine layer."""
         return self.flow_model.spk_embed_affine_layer
 
+    @torch.inference_mode()
+    def _forward_mel(
+        self,
+        token: torch.Tensor,
+        prompt_token: torch.Tensor,
+        prompt_feat: torch.Tensor,
+        embedding: torch.Tensor,
+        n_timesteps: int = 10,
+        token_offset_tokens: int = 0,
+        streaming: bool = True,
+        finalize: bool = False,
+    ) -> torch.Tensor:
+        """Generate mel features via the upstream flow-model inference path."""
+        flow_weight = next(self.flow_model.parameters())
+        device = flow_weight.device
+        dtype = flow_weight.dtype
+
+        token = token.to(device=device, dtype=torch.int32)
+        prompt_token = prompt_token.to(device=device, dtype=torch.int32)
+        prompt_feat = prompt_feat.to(device=device, dtype=dtype)
+        embedding = embedding.to(device=device, dtype=dtype)
+        token_len = torch.tensor([token.shape[1]], device=device, dtype=torch.int32)
+        prompt_token_len = torch.tensor([prompt_token.shape[1]], device=device, dtype=torch.int32)
+        prompt_feat_len = torch.tensor([prompt_feat.shape[1]], device=device, dtype=torch.int32)
+
+        with nullcontext():
+            feat, _ = self.flow_model.inference(
+                token=token,
+                token_len=token_len,
+                prompt_token=prompt_token,
+                prompt_token_len=prompt_token_len,
+                prompt_feat=prompt_feat,
+                prompt_feat_len=prompt_feat_len,
+                embedding=embedding,
+                streaming=streaming,
+                finalize=finalize,
+                n_timesteps=n_timesteps,
+            )
+
+        trim_mel = max(0, int(token_offset_tokens)) * int(self.token_mel_ratio)
+        if trim_mel > 0:
+            feat = feat[:, :, trim_mel:]
+
+        return feat
+
+    @staticmethod
+    def _fade_speech(
+        speech: torch.Tensor,
+        prev_speech: torch.Tensor,
+    ) -> torch.Tensor:
+        """Blend previous speech tail into current speech head."""
+        if speech.numel() == 0 or prev_speech.numel() == 0:
+            return speech
+        overlap = min(int(speech.shape[-1]), int(prev_speech.shape[-1]))
+        if overlap <= 0:
+            return speech
+        window = torch.hamming_window(2 * overlap, periodic=False, dtype=speech.dtype, device=speech.device)
+        fade_in = window[:overlap].view(1, -1)
+        fade_out = window[overlap:].view(1, -1)
+        blended_head = (
+            speech[:, :overlap] * fade_in
+            + prev_speech[:, -overlap:].to(device=speech.device, dtype=speech.dtype) * fade_out
+        )
+        if overlap == int(speech.shape[-1]):
+            return blended_head
+        return torch.cat([blended_head, speech[:, overlap:]], dim=-1)
+
+    @torch.inference_mode()
+    def forward_streaming(
+        self,
+        token: torch.Tensor,
+        prompt_token: torch.Tensor,
+        prompt_feat: torch.Tensor,
+        embedding: torch.Tensor,
+        *,
+        cache_state: dict[str, torch.Tensor] | None = None,
+        n_timesteps: int = 10,
+        token_offset_tokens: int = 0,
+        finalize: bool = False,
+    ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]:
+        """Decode streaming audio using cumulative mel + emitted-speech offset.
+
+        This mirrors upstream CosyVoice3 streaming semantics more closely than
+        waveform-domain overlap-add: keep a cumulative mel history per request,
+        re-run causal HiFT on the history, and emit only the newly grown speech
+        suffix. That preserves causal look-right handling without double
+        trimming or duplicated overlap at chunk boundaries.
+        """
+        with nullcontext():
+            feat = self._forward_mel(
+                token=token,
+                prompt_token=prompt_token,
+                prompt_feat=prompt_feat,
+                embedding=embedding,
+                n_timesteps=n_timesteps,
+                token_offset_tokens=token_offset_tokens,
+                streaming=True,
+                finalize=finalize,
+            )
+        hift_weight = self.hift.m_source.l_linear.weight
+        chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype)
+
+        cached_mel = None if not cache_state else cache_state.get("mel")
+        speech_offset_obj = None if not cache_state else cache_state.get("speech_offset")
+        try:
+            speech_offset = int(speech_offset_obj) if speech_offset_obj is not None else 0
+        except (TypeError, ValueError):
+            speech_offset = 0
+
+        if isinstance(cached_mel, torch.Tensor) and cached_mel.numel() > 0:
+            cached_mel = cached_mel.to(device=chunk_mel.device, dtype=chunk_mel.dtype)
+            tts_mel = torch.cat([cached_mel, chunk_mel], dim=-1) if chunk_mel.numel() > 0 else cached_mel
+        else:
+            tts_mel = chunk_mel
+
+        if tts_mel.shape[-1] == 0:
+            tts_speech = torch.zeros((chunk_mel.shape[0], 1, 0), device=chunk_mel.device, dtype=chunk_mel.dtype)
+        else:
+            with nullcontext():
+                tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=finalize)
+
+        tts_speech = tts_speech.reshape(tts_speech.shape[0], -1)
+        speech_offset = max(0, min(speech_offset, int(tts_speech.shape[-1])))
+        emitted_speech = tts_speech[:, speech_offset:]
+
+        if finalize:
+            return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), None
+
+        new_state = {
+            "mel": tts_mel.detach().cpu().contiguous(),
+            "speech_offset": int(tts_speech.shape[-1]),
+        }
+        return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), new_state
+
     @torch.inference_mode()
     def forward(
         self,
@@ -159,76 +293,18 @@ class CosyVoice3Code2Wav(nn.Module):
         embedding: torch.Tensor,
         n_timesteps: int = 10,
     ) -> torch.Tensor:
-        """Generate audio waveform from speech tokens.
-
-        Args:
-            token: Speech tokens from talker stage [batch, seq_len]
-            prompt_token: Prompt speech tokens [batch, prompt_len]
-            prompt_feat: Prompt mel features [batch, feat_len, mel_dim]
-            embedding: Speaker embedding [batch, spk_dim]
-            n_timesteps: Number of diffusion steps
-
-        Returns:
-            Audio waveform [batch, 1, audio_len]
-        """
-        device = token.device
-        dtype = next(self.flow_model.parameters()).dtype
-
-        # Normalize and project speaker embedding
-        embedding = embedding.to(device=device, dtype=dtype)
-        embedding = F.normalize(embedding, dim=1)
-        embedding = self.spk_embed_affine_layer(embedding)
-
-        # Prepare tokens
-        prompt_token = prompt_token.to(device=device)
-        token_len1, token_len2 = prompt_token.shape[1], token.shape[1]
-        prompt_token_len = torch.tensor([token_len1], device=device, dtype=torch.int32)
-        token_len = torch.tensor([token_len2], device=device, dtype=torch.int32)
-
-        # Concatenate prompt and target tokens
-        full_token = torch.cat([prompt_token, token], dim=1)
-        full_token_len = prompt_token_len + token_len
-
-        # Create mask
-        mask = (~make_pad_mask(full_token_len)).unsqueeze(-1).to(embedding)
-
-        # Token embedding (clamp to valid codebook range; EOS/padding tokens may exceed vocab_size)
-        token_emb = (
-            self.input_embedding(torch.clamp(full_token, min=0, max=self.input_embedding.num_embeddings - 1)) * mask
-        )
-
-        # Pre-lookahead processing
-        h = self.pre_lookahead_layer(token_emb)
-        h = h.repeat_interleave(self.token_mel_ratio, dim=1)
-
-        # Calculate mel lengths
-        mel_len1 = prompt_feat.shape[1]
-        mel_len2 = h.shape[1] - mel_len1
-
-        # Build conditioning
-        conds = torch.zeros(
-            [1, mel_len1 + mel_len2, self.output_size],
-            device=device,
-            dtype=h.dtype,
-        )
-        conds[:, :mel_len1] = prompt_feat
-        conds = conds.transpose(1, 2)
-
-        # Create mel mask
-        mel_mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h)
-
-        # Run flow matching decoder
-        feat, _ = self.decoder(
-            mu=h.transpose(1, 2).contiguous(),
-            mask=mel_mask.unsqueeze(1),
-            spks=embedding,
-            cond=conds,
+        """Generate audio waveform from speech tokens."""
+        feat = self._forward_mel(
+            token=token,
+            prompt_token=prompt_token,
+            prompt_feat=prompt_feat,
+            embedding=embedding,
             n_timesteps=n_timesteps,
+            token_offset_tokens=0,
+            streaming=False,
+            finalize=True,
         )
 
-        # Extract generated portion (after prompt)
-        feat = feat[:, :, mel_len1:]
-
         # Run vocoder
         hift_weight = self.hift.m_source.l_linear.weight
         tts_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype)
@@ -240,7 +316,7 @@ class CosyVoice3Code2Wav(nn.Module):
                 dtype=tts_mel.dtype,
             )
         else:
-            tts_speech, _ = self.hift.inference(speech_feat=tts_mel)
+            tts_speech, _ = self.hift.inference(speech_feat=tts_mel, finalize=True)
 
         return tts_speech
 

vllm_omni/model_executor/stage_configs/cosyvoice3.yaml

@@ -24,7 +24,10 @@ stage_args:
       dtype: "float32"
     default_sampling_params:
       max_tokens: 2048
-      stop_token_ids: [6562]  # speech EOS token
+      stop_token_ids: [6562]  # merged speech stop token (logsumexp of all 200 stop logits)
+      top_k: 25
+      top_p: 0.8
+      repetition_penalty: 1.0001  # near-identity; forces vLLM to track output_token_ids for RAS
 
   - stage_id: 1
     runtime:

vllm_omni/model_executor/stage_configs/cosyvoice3_async_chunk.yaml

@@ -0,0 +1,85 @@
+# Stage config for running CosyVoice3 with async_chunk architecture
+# Stage 0: Talker (text prompt -> speech tokens streamed by chunks)
+# Stage 1: Code2Wav (flow matching -> acoustic features -> waveform)
+async_chunk: true
+
+stage_args:
+  - stage_id: 0
+    is_comprehension: true
+    runtime:
+      devices: 0
+      max_batch_size: 1
+    engine_args:
+      model_stage: cosyvoice3_talker
+      model_arch: CosyVoice3Model
+      worker_type: ar
+      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
+      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.cosyvoice3.talker2code2wav_async_chunk
+      trust_remote_code: true
+      gpu_memory_utilization: 0.4
+      engine_output_type: latent
+      disable_hybrid_kv_cache_manager: true
+      enable_prefix_caching: false
+      enforce_eager: true
+      mm_processor_cache_gb: 0
+      skip_mm_profiling: true
+      dtype: "float32"
+    default_sampling_params:
+      max_tokens: 2048
+      stop_token_ids: [6562]  # merged speech stop token (logsumexp of all 200 stop logits)
+      top_k: 25
+      top_p: 0.8
+      repetition_penalty: 1.0001  # near-identity; forces vLLM to track output_token_ids for RAS
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+
+  - stage_id: 1
+    runtime:
+      devices: 0
+      max_batch_size: 1
+    engine_args:
+      model_stage: cosyvoice3_code2wav
+      model_arch: CosyVoice3Model
+      worker_type: generation
+      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
+      trust_remote_code: true
+      engine_output_type: latent
+      gpu_memory_utilization: 0.2
+      enforce_eager: true
+      disable_hybrid_kv_cache_manager: true
+      enable_prefix_caching: false
+      skip_mm_profiling: true
+      max_model_len: 32768
+      dtype: "float32"
+    default_sampling_params:
+      max_tokens: 2048
+    engine_input_source: [0]
+    final_output: true
+    final_output_type: audio
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+
+runtime:
+  enabled: true
+  defaults:
+    window_size: -1
+    max_inflight: 1
+
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        codec_streaming: true
+        connector_get_sleep_s: 0.01
+        connector_get_max_wait_first_chunk: 3000
+        connector_get_max_wait: 300
+        codec_chunk_frames: 25
+        codec_left_context_frames: 25
+        codec_vocab_size: 6561
+
+  edges:
+    - from: 0
+      to: 1
+      window_size: -1

vllm_omni/model_executor/stage_input_processors/cosyvoice3.py

@@ -1,10 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from collections import defaultdict
+from contextlib import nullcontext
 from typing import Any
 
+import numpy as np
+import torch
 from vllm.inputs import TextPrompt
 
 from vllm_omni.inputs.data import OmniTokensPrompt
 
 
+def _ensure_list(x: Any) -> list[Any]:
+    if hasattr(x, "_x"):
+        return list(x._x)
+    if isinstance(x, list):
+        return list(x)
+    if isinstance(x, tuple):
+        return list(x)
+    if x is None:
+        return []
+    try:
+        return list(x)
+    except TypeError:
+        return [x]
+
+
+def _to_cpu_tensor(x: Any) -> torch.Tensor | None:
+    if isinstance(x, list):
+        if not x:
+            return None
+        x = x[0]
+    if isinstance(x, torch.Tensor):
+        return x.detach().cpu()
+    return None
+
+
+def _decode_additional_information(raw_info: Any) -> dict[str, Any]:
+    if raw_info is None:
+        return {}
+    if isinstance(raw_info, dict):
+        return raw_info
+
+    entries = getattr(raw_info, "entries", None)
+    if not isinstance(entries, dict):
+        return {}
+
+    decoded: dict[str, Any] = {}
+    for key, entry in entries.items():
+        tensor_data = getattr(entry, "tensor_data", None)
+        if tensor_data is not None:
+            dtype_name = getattr(entry, "tensor_dtype", "float32")
+            tensor_shape = getattr(entry, "tensor_shape", None)
+            if tensor_shape is None:
+                continue
+            dt = np.dtype(dtype_name)
+            arr = np.frombuffer(tensor_data, dtype=dt).reshape(tensor_shape)
+            decoded[key] = torch.from_numpy(arr.copy())
+        else:
+            decoded[key] = getattr(entry, "list_data", None)
+    return decoded
+
+
 def text2flow(
     stage_list: list[Any],
     engine_input_source: list[int],
@@ -15,18 +72,178 @@ def text2flow(
     source_stage_id = engine_input_source[0]
     source_outputs = stage_list[source_stage_id].engine_outputs
 
-    if not isinstance(prompt, list):
-        prompt = [prompt]
+    engine_inputs: list[OmniTokensPrompt] = []
+    for source_output in source_outputs:
+        output = source_output.outputs[0]
+        multi_modal_data = output.multimodal_output
+        if multi_modal_data is None:
+            raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}")
 
-    source_output = source_outputs[0]
-    output = source_output.outputs[0]
+        output_ids = _ensure_list(output.token_ids)
+        prefix_ids = _ensure_list(source_output.prompt_token_ids)
+        additional_info = dict(multi_modal_data)
+        additional_info["prefix_ids"] = prefix_ids
+        engine_inputs.append(OmniTokensPrompt(prompt_token_ids=output_ids, additional_information=additional_info))
+    return engine_inputs
 
-    multi_modal_data = output.multimodal_output
-    if multi_modal_data is None:
-        raise RuntimeError(f"Missing multimodal_output for request {source_output.request_id}")
 
-    output_ids = output.token_ids
-    prefix_ids = source_output.prompt_token_ids
-    multi_modal_data["prefix_ids"] = prefix_ids
-    engine_input = OmniTokensPrompt(prompt_token_ids=output_ids, additional_information=multi_modal_data)
-    return [engine_input]
+def talker2code2wav_async_chunk(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any] | None,
+    request: Any,
+    is_finished: bool = False,
+) -> dict[str, Any] | None:
+    """CosyVoice3 async_chunk processor: talker token stream -> code2wav chunks."""
+    with nullcontext():
+        request_id = request.external_req_id
+        finished = bool(is_finished or request.is_finished())
+
+        connector = getattr(transfer_manager, "connector", None)
+        raw_cfg = getattr(connector, "config", {}) or {}
+        cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {}
+        chunk_size = int(cfg.get("codec_chunk_frames", 25))
+        code_vocab_size = int(cfg.get("codec_vocab_size", 6561))
+        pre_lookahead_len = int(cfg.get("codec_pre_lookahead_frames", 3))
+        max_chunk_size = int(cfg.get("codec_max_chunk_frames", 4 * chunk_size))
+        stream_scale_factor = int(cfg.get("codec_stream_scale_factor", 2))
+        if chunk_size <= 0 or pre_lookahead_len < 0 or max_chunk_size <= 0 or stream_scale_factor <= 0:
+            raise ValueError(
+                f"Invalid codec chunk config: codec_chunk_frames={chunk_size}, "
+                f"codec_pre_lookahead_frames={pre_lookahead_len}, "
+                f"codec_max_chunk_frames={max_chunk_size}, "
+                f"codec_stream_scale_factor={stream_scale_factor}"
+            )
+
+        request_state = transfer_manager.request_payload.get(request_id)
+        if not isinstance(request_state, dict) or "_cosyvoice3_async_state" not in request_state:
+            with nullcontext():
+                info = _decode_additional_information(getattr(request, "additional_information", None))
+                prompt_payload = {}
+                for key in ("speech_token", "speech_feat", "embedding"):
+                    value = _to_cpu_tensor(info.get(key))
+                    if value is not None:
+                        prompt_payload[key] = value
+                if isinstance(pooling_output, dict):
+                    for key in ("speech_token", "speech_feat", "embedding"):
+                        if key in prompt_payload:
+                            continue
+                        value = _to_cpu_tensor(pooling_output.get(key))
+                        if value is not None:
+                            prompt_payload[key] = value
+                prompt_token = prompt_payload.get("speech_token")
+                prompt_token_len = (
+                    int(prompt_token.shape[1])
+                    if isinstance(prompt_token, torch.Tensor) and prompt_token.ndim >= 2
+                    else 0
+                )
+                prompt_token_pad = (
+                    ((prompt_token_len + chunk_size - 1) // chunk_size) * chunk_size - prompt_token_len
+                    if prompt_token_len > 0
+                    else 0
+                )
+            request_state = {
+                "_cosyvoice3_async_state": {
+                    "seen_len": 0,
+                    "sent_prompt": False,
+                    "emitted_chunks": 0,
+                    "emitted_token_len": 0,
+                    "token_hop_len": chunk_size,
+                    "prompt_token_pad": prompt_token_pad,
+                    "pre_lookahead_len": pre_lookahead_len,
+                    "token_max_hop_len": max(chunk_size, max_chunk_size),
+                    "stream_scale_factor": stream_scale_factor,
+                    "terminal_sent": False,
+                    "prompt_payload": prompt_payload,
+                }
+            }
+            transfer_manager.request_payload[request_id] = request_state
+
+        state = request_state["_cosyvoice3_async_state"]
+        if bool(state.get("terminal_sent", False)):
+            return None
+
+        with nullcontext():
+            output_token_ids = _ensure_list(getattr(request, "output_token_ids", []))
+            seen_len = int(state.get("seen_len", 0))
+            new_tokens = output_token_ids[seen_len:] if seen_len < len(output_token_ids) else []
+            state["seen_len"] = len(output_token_ids)
+
+        if not hasattr(transfer_manager, "code_prompt_token_ids"):
+            transfer_manager.code_prompt_token_ids = defaultdict(list)
+        token_frames = transfer_manager.code_prompt_token_ids[request_id]
+        for tok in new_tokens:
+            tok_int = int(tok)
+            if 0 <= tok_int < code_vocab_size:
+                token_frames.append([tok_int])
+
+        length = len(token_frames)
+        if length <= 0:
+            if not finished:
+                return None
+            payload: dict[str, Any] = {
+                "code_predictor_codes": [],
+                "finished": torch.tensor(True, dtype=torch.bool),
+            }
+            if not state.get("sent_prompt", False):
+                payload.update(state.get("prompt_payload", {}))
+                state["sent_prompt"] = True
+            state["terminal_sent"] = True
+            return payload
+
+        emitted_token_len = int(state.get("emitted_token_len", 0))
+        if finished and length <= emitted_token_len:
+            payload = {
+                "code_predictor_codes": [],
+                "finished": torch.tensor(True, dtype=torch.bool),
+            }
+            if not state.get("sent_prompt", False):
+                payload.update(state.get("prompt_payload", {}))
+                state["sent_prompt"] = True
+            state["terminal_sent"] = True
+            return payload
+
+        with nullcontext():
+            token_hop_len = max(1, int(state.get("token_hop_len", chunk_size)))
+            prompt_token_pad = max(0, int(state.get("prompt_token_pad", 0)))
+            pre_lookahead_len = max(0, int(state.get("pre_lookahead_len", pre_lookahead_len)))
+            available = max(0, length - emitted_token_len)
+            this_token_hop_len = token_hop_len + prompt_token_pad if emitted_token_len == 0 else token_hop_len
+            required = this_token_hop_len + pre_lookahead_len
+
+            if not finished:
+                if available < required:
+                    return None
+                prefix_len = emitted_token_len + required
+                token_offset = emitted_token_len
+            else:
+                if available <= 0:
+                    return None
+                prefix_len = length
+                token_offset = emitted_token_len
+
+        with nullcontext():
+            code_predictor_codes = [int(frame[0]) for frame in token_frames[:prefix_len]]
+
+        payload = {
+            "code_predictor_codes": code_predictor_codes,
+            "token_offset": token_offset,
+            "left_context_size": token_offset,
+            "req_id": [request_id],
+            "stream_finished": torch.tensor(finished, dtype=torch.bool),
+            "finished": torch.tensor(finished, dtype=torch.bool),
+        }
+        if not state.get("sent_prompt", False):
+            payload.update(state.get("prompt_payload", {}))
+            state["sent_prompt"] = True
+
+        if not finished:
+            state["emitted_token_len"] = emitted_token_len + this_token_hop_len
+            state["token_hop_len"] = min(
+                int(state.get("token_max_hop_len", chunk_size)),
+                max(chunk_size, token_hop_len * int(state.get("stream_scale_factor", 1))),
+            )
+        else:
+            state["terminal_sent"] = True
+
+        state["emitted_chunks"] = int(state.get("emitted_chunks", 0)) + 1
+        return payload

vllm_omni/worker/gpu_ar_model_runner.py

@@ -6,7 +6,9 @@ and also outputs sampled tokens.
 
 from __future__ import annotations
 
+from contextlib import nullcontext
 from copy import copy
+from dataclasses import replace
 from typing import Any, NamedTuple
 
 import numpy as np
@@ -89,6 +91,53 @@ class GPUARModelRunner(OmniGPUModelRunner):
         with maybe_disable_pin_memory_for_ray(self, total_bytes):
             return super()._make_buffer(*size, dtype=dtype, numpy=numpy)
 
+    def _build_model_sampler_output_token_ids(self) -> list[list[int]]:
+        """Build decoded-token history for custom model samplers.
+
+        vLLM only populates sampling_metadata.output_token_ids when penalties or
+        logits processors require it. CosyVoice3's custom RAS sampler also
+        depends on this history, so we reconstruct it directly from the input
+        batch for prefer_model_sampler models.
+        """
+        req_output_token_ids = getattr(self.input_batch, "req_output_token_ids", [])
+        req_ids = list(getattr(self.input_batch, "req_ids", []))
+        output_token_ids = [list(req_output_token_ids[idx] or []) for idx in range(len(req_ids))]
+
+        sampled_token_ids_cpu = getattr(self.input_batch, "sampled_token_ids_cpu", None)
+        async_copy_ready_event = getattr(self.input_batch, "async_copy_ready_event", None)
+        prev_req_id_to_index = getattr(self.input_batch, "prev_req_id_to_index", None)
+        if sampled_token_ids_cpu is None or not output_token_ids or prev_req_id_to_index is None:
+            return output_token_ids
+
+        sampled_token_ids: list[list[int]] | None = None
+        for index, req_id in enumerate(req_ids):
+            prev_index = prev_req_id_to_index.get(req_id)
+            if prev_index is None:
+                continue
+            req_history = output_token_ids[index]
+            if not req_history or req_history[-1] != -1:
+                continue
+            if sampled_token_ids is None:
+                assert async_copy_ready_event is not None
+                async_copy_ready_event.synchronize()
+                sampled_token_ids = sampled_token_ids_cpu.tolist()
+            new_ids = list(sampled_token_ids[prev_index])
+            if not new_ids:
+                continue
+            num_sampled_ids = len(new_ids) if new_ids[-1] != -1 else new_ids.index(-1)
+            first_placeholder = req_history.index(-1)
+            num_placeholders = len(req_history) - first_placeholder
+            num_to_replace = min(num_sampled_ids, num_placeholders)
+            req_history[first_placeholder : first_placeholder + num_to_replace] = new_ids[:num_to_replace]
+
+        return output_token_ids
+
+    def _sampling_metadata_for_model_sampler(self, sampling_metadata):
+        output_token_ids = self._build_model_sampler_output_token_ids()
+        if output_token_ids == sampling_metadata.output_token_ids:
+            return sampling_metadata
+        return replace(sampling_metadata, output_token_ids=output_token_ids)
+
     @torch.inference_mode()
     def execute_model(
         self,
@@ -302,6 +351,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
         # (wait_for_save + clear metadata) until after draft model runs.
         defer_kv_connector_finalize = self.speculative_config is not None
         with (
+            nullcontext(),
             set_forward_context(
                 attn_metadata,
                 self.vllm_config,
@@ -424,6 +474,39 @@ class GPUARModelRunner(OmniGPUModelRunner):
 
         return None
 
+    def _sample(
+        self,
+        logits: torch.Tensor | None,
+        spec_decode_metadata: Any,
+    ):
+        sampling_metadata = self.input_batch.sampling_metadata
+        if spec_decode_metadata is None:
+            model_sample = getattr(self.model, "sample", None)
+            if logits is not None and callable(model_sample) and getattr(self.model, "prefer_model_sampler", False):
+                # Apply logit bias (min_tokens, allowed_token_ids) before
+                # the custom model sampler — the standard GPU sampler does
+                # this internally, but prefer_model_sampler bypasses it.
+                if hasattr(self.sampler, "logit_bias_state"):
+                    self.sampler.logit_bias_state.apply_logit_bias(
+                        logits,
+                        self.input_batch.expanded_idx_mapping,
+                        self.input_batch.idx_mapping_np,
+                        self.input_batch.positions[self.input_batch.logits_indices],
+                    )
+                sampler_output = model_sample(
+                    logits,
+                    self._sampling_metadata_for_model_sampler(sampling_metadata),
+                )
+                if sampler_output is not None:
+                    return sampler_output
+            self.input_batch.update_async_output_token_ids()
+            return self.sampler(
+                logits=logits,
+                sampling_metadata=sampling_metadata,
+            )
+
+        return super()._sample(logits, spec_decode_metadata)
+
     @torch.inference_mode()
     def sample_tokens(
         self,