From 6fc38e0467b2d8967cda0749e3d83a5e0561b4ff Mon Sep 17 00:00:00 2001 From: indevn Date: Sun, 5 Apr 2026 14:11:33 +0800 Subject: [PATCH] [Model][Core] Enable async_chunk streaming pipeline for CosyVoice3 (#1703) Signed-off-by: linyueqian Signed-off-by: indevn Co-authored-by: linyueqian --- .../test_chunk_transfer_adapter.py | 99 ++++ .../e2e/offline_inference/test_cosyvoice3.py | 218 +++++++++ .../e2e/online_serving/test_cosyvoice3_tts.py | 25 +- .../cosyvoice3/test_cosyvoice3_components.py | 31 ++ .../test_cosyvoice3_model_helpers.py | 463 ++++++++++++++++++ .../test_cosyvoice3_stage_input_processors.py | 267 ++++++++++ vllm_omni/core/sched/omni_ar_scheduler.py | 4 - .../chunk_transfer_adapter.py | 30 +- .../entrypoints/openai/serving_speech.py | 34 ++ .../models/cosyvoice3/assets/mel_filters.npz | Bin 0 -> 4271 bytes .../models/cosyvoice3/code2wav_core/cfm.py | 9 +- .../models/cosyvoice3/config.py | 4 +- .../models/cosyvoice3/cosyvoice3.py | 411 +++++++++++++++- .../models/cosyvoice3/cosyvoice3_code2wav.py | 216 +++++--- .../stage_configs/cosyvoice3.yaml | 5 +- .../stage_configs/cosyvoice3_async_chunk.yaml | 85 ++++ .../stage_input_processors/cosyvoice3.py | 241 ++++++++- vllm_omni/worker/gpu_ar_model_runner.py | 83 ++++ 18 files changed, 2094 insertions(+), 131 deletions(-) create mode 100644 tests/e2e/offline_inference/test_cosyvoice3.py create mode 100644 tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py create mode 100644 tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py create mode 100644 vllm_omni/model_executor/models/cosyvoice3/assets/mel_filters.npz create mode 100644 vllm_omni/model_executor/stage_configs/cosyvoice3_async_chunk.yaml diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index dddf49a05..7a3caba11 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/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 diff --git a/tests/e2e/offline_inference/test_cosyvoice3.py b/tests/e2e/offline_inference/test_cosyvoice3.py new file mode 100644 index 000000000..8c88d972d --- /dev/null +++ b/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" diff --git a/tests/e2e/online_serving/test_cosyvoice3_tts.py b/tests/e2e/online_serving/test_cosyvoice3_tts.py index 1845d7818..276b1782f 100644 --- a/tests/e2e/online_serving/test_cosyvoice3_tts.py +++ b/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) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py index 3b1471365..0f5202c3b 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_components.py +++ b/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] diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py new file mode 100644 index 000000000..9a78c54de --- /dev/null +++ b/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]]] diff --git a/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py b/tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py new file mode 100644 index 000000000..e26de3022 --- /dev/null +++ b/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"] == [] diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index af178d14d..0956d1856 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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: diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index e8e00eeca..393d0e801 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/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 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index f05126882..0a9e11b77 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/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: diff --git a/vllm_omni/model_executor/models/cosyvoice3/assets/mel_filters.npz b/vllm_omni/model_executor/models/cosyvoice3/assets/mel_filters.npz new file mode 100644 index 0000000000000000000000000000000000000000..28ea26909dbdfd608aef67afc4d74d7961ae4bb6 GIT binary patch literal 4271 zcmZ`-cQjmYw;lx1g6JcN7QKe3LG%_Oh!VX=^k~teM-XGQ(Mu4$_Y%?jkm$lFBkB+( z3yfKIgF zxGiAhze`A@t->QRNVV!%P+W=o}VHkB) z%g>qyRHfN1IQ4-=`Y@0T9qE#o+;4E3VQ!epW1Xt=ZG`I3U|62t?<>5h*W|9VvJc`KZ+)ghnA**Z~ET21Tjf_f8oe`vy zZQNtlOx?dDhS71hnOus5cqj)hfyF@H&4y?@9z{I#&cf>A+s2~~(I>TQF}SaR3_tqa z(7&ZdN^vR*t<~?{9DEoI>0PL@Sl?wa?Z{rGX`*eEx9Nh=z*J3HZL1*Py4z$TD#+;m zSSW(kcOTe(4hqgib_W6&xx+j~-u(p)Nn6?>a%wHk=h7Ay$%lcGoo;gAY zmVV7|!Nb;w(PlH@c24{ple2Y3<*9J@jE=sfLzwu_BiAFPE$0Axp`^Nq!H}eG0?r-X zFj@Pwp^al*p>K{@_Cz`q#(N0Y=OpZy^ z{P$KjLJuk_Y%I)$mh`b{uOW5C5Xcmxk!gt_Zg zw>}6fkD4zRK9!#ems~H%U$>V;_wK38Zf-baU$S!#i;7!HWsi}GuC>%@?lMdgkUGC& zh9gC?O-5BlS2#}?7x0?eP#bOL(cqE{M%LJD$CZnplD)CgQR#KCttD=dZK+Ck5R52; z*%5hZ+SXU7)8k%Y^_1U>yI*By(INn&+ir-_4$#dUwTlMNyR@iGQIaZ+eiYqucu)CB z#i{Ru1w+aU#}DHSyzjG_9c?ToB_YjU#f;N=qel98WBIjIc1!#ePwRR+(go&-by#}@ z+M+klVke5b@lWfZ+O&|c??YvRe)&W)qAgtc>t-IZtbRTG#X}49_Q$>P%-)=0W_QY-x%DPep2Vm9#ci zyQcCc4p2&dLtV1@rPe!%>Y^#9W8#ZH&}^@wJKT7N;R9A7cEq&;Y2CYvd@R+Mn&b5O zVyfS^*H#kD74=J5uhD)o`TXoX>>Si$!cT?TXRxj2pB)w_ljjhTby&Je;X|BESZZT= zC%G5!-$BJf&a~U78d_3zBjrvrkJ0CCl@Rfcf7I(`VTNPnI^B#B$zOfPW zG&mEd?R0+W<`l08O1dkcWKS8wB!Z*Cs%I1nMs-EeB-uu5?t@PuD3|z>je8DKi#X(B z{Z=Rz{4X%?-UnxnHQtkELIZ&=J;fK_t}yu8|IxG0(85e&K>H3!!~zlhyJrgti~o1i zzBS*jTgdG~Exp#B-T)6A+PB ztD-e`j^@XAx}|L&JSEFkRvS_%3b%m86z02#Hfn{Y+qIqQ_muywgt?roUA7oiS1xBD zFxmDMsj_cbBcn*^rn^KIMP{AlHM`NiVm*D&`z~7FH#hf<$L3HmJ+=NdiY5>W?nKD? z8Ox6{9dKyI1o8a-j9BtV-|=lm`<`v>tR^Cln&x1dMYzu{@wq5KW!#K14_QMnpH5K%Pavag+g6(i8i-#Eq zguc}rH3?BxH4SOqZW#7m*aT(U9-n#_Xn^Q19(}eH!xG`nI!GYziVQNcA0)`FDHD%~ zz2$HnxW4BQ{#*@u`dssbAa`|fESn$8i8FdxGZh48_Uf~_Q@tv?4in)6fwSed)k&ITqu|){^(WL~J z?Lb|0ro06J^>f>^2}^e-+$u5bU4IZNfO?75v8lstS15%XYw2ac^pkU34{QhDR(umt zPu~`w2?FP|nn3!RWZ3{?=77@teulahD9*S*k5KmY3*adlM)%{SR~bkZYlx1q@fkE= zI$7+kiw5!ha=dYlO>Z5KgxnZEJsaBm%v#nkX0MN-h%n&KA?N}xU3K3o-3Jpk?ANq2n9&Lh%K_CTvfiN ze>6w~NSSl8$#NEZ^t7h9YOxI=zcAG|a+m6AWei`3Jw7K;b;T${pJa^4RwRt%F>?>M zBmoQqm1`<_W7i!5P~THp-II)Ka^u;=z;}d{;SVj{G_4`9^HaEb!=@Pa;Dw)CH^DjsGxFqmb%o$Bkop$KnH8 zDYN)Bh)5=5!-*|f0Gh4)oZG=TEBr()g^DCtSQhmT3!ZN`Qd-E%@1cE}hm8&Vq5B+C zVF2_O)9IiZ(v(xzTwJIg5|}KVuE(;}|7dVIrT`$d=q_OG|3PY}x*URYkMXXJ6PT1$IFkNyvY_(9UglDi6TaeikPS(!Bnij z;Szn+)I_oxnRz7(WTYTp+IHSWQ?Xd~tQn(Q1r)kThM?NM< z?d6LaBG!H}R$zRy!Ij(}1?xe^+o+!;tqWJ3NgjHl1XNxzusxQ0I#6qzM(_00UPMw* zF*GWW_q&fqAN=uimSKgBu_@jD%MX3hpNY|*4r=e=k1lw2r**IyD(hcq?A+HtUgUy4Dqh5D7|G9q{)TsUj{g~c!xy>9wk^(LiXA4VKGz_zMvJMX#AgsR z34T3hhJ)#&sUaQ1+0PML(?YA~{5?=(MT}X^Vib%};uoI{qGW@wgJ&_M+8S8clsNz2 zPQkxMi`#3+Khwtl>>K>wxc{71{&!qGu&Zzz_wU(7TLTyG){PAu?!cXs?Dp-y0Ekcn AQvd(} literal 0 HcmV?d00001 diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py index 7281cd81f..36ff0d456 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/cfm.py +++ b/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:] diff --git a/vllm_omni/model_executor/models/cosyvoice3/config.py b/vllm_omni/model_executor/models/cosyvoice3/config.py index b4e44b7a8..518fe76b7 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/config.py +++ b/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 diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 18a16ba55..2fba8fb8a 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/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}") diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index 222d6d98a..3ad23cdb1 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/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 diff --git a/vllm_omni/model_executor/stage_configs/cosyvoice3.yaml b/vllm_omni/model_executor/stage_configs/cosyvoice3.yaml index bfb847f5e..8e0582723 100644 --- a/vllm_omni/model_executor/stage_configs/cosyvoice3.yaml +++ b/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: diff --git a/vllm_omni/model_executor/stage_configs/cosyvoice3_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/cosyvoice3_async_chunk.yaml new file mode 100644 index 000000000..ca7e9850a --- /dev/null +++ b/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 diff --git a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py b/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py index b7f21eca8..c722a125e 100644 --- a/vllm_omni/model_executor/stage_input_processors/cosyvoice3.py +++ b/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 diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index f1115ab4c..01ec23acb 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/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,