diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index 54fed272c..3b5213d9e 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -16,6 +16,15 @@ steps: - "pytest -sv tests/model_executor -m 'core_model and cpu'" mirror_hardwares: l4_1 + - label: "Simple · CosyVoice3 Packed Flow Test" + timeout_in_minutes: 15 + source_file_dependencies: + - tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py + - vllm_omni/model_executor/models/cosyvoice3/ + commands: + - "pytest -sv tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py -k packed -m 'core_model and cuda' --run-level=core_model" + mirror_hardwares: h100_1 + - label: "Simple · Diffusion Test" timeout_in_minutes: 30 commands: diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 14613b1cb..b66c930fc 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PUBLIC_OMNI: 32, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29, - EnvironmentVariableCategory.MODEL_SPECIFIC: 66, + EnvironmentVariableCategory.MODEL_SPECIFIC: 74, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, EnvironmentVariableCategory.INTERNAL: 6, } @@ -178,10 +178,10 @@ def test_inventory_matches_reviewed_snapshot_counts(): if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { - ModelEnvironmentVariableDisposition.PROMOTE: 38, + ModelEnvironmentVariableDisposition.PROMOTE: 45, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5, ModelEnvironmentVariableDisposition.EXTERNAL: 0, - ModelEnvironmentVariableDisposition.INTERNALIZE: 16, + ModelEnvironmentVariableDisposition.INTERNALIZE: 17, ModelEnvironmentVariableDisposition.DEPRECATE_REMOVE: 7, } diff --git a/tests/entrypoints/openai_api/test_reference_waveform_cache.py b/tests/entrypoints/openai_api/test_reference_waveform_cache.py index 963e782bc..e171d7fd8 100644 --- a/tests/entrypoints/openai_api/test_reference_waveform_cache.py +++ b/tests/entrypoints/openai_api/test_reference_waveform_cache.py @@ -106,3 +106,86 @@ def test_numeric_waveform_cache_limits(server, entries, budget): server._put_resolved_ref_audio("b", np.ones(2, dtype=np.float32), 24000, "artifact-b") assert list(server._ref_audio_resolve_cache) == ["b"] assert server._ref_audio_resolve_cache_bytes == 8 + + +@pytest.mark.asyncio +async def test_cosyvoice_array_prompt_owns_buffer_without_list_roundtrip(mocker, monkeypatch): + from types import SimpleNamespace + + from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext + from vllm_omni.entrypoints.openai.tts_adapters.cosyvoice3 import CosyVoice3Adapter + + monkeypatch.setenv("COSYVOICE3_REFERENCE_PREFETCH", "0") + waveform = np.linspace(-1, 1, 24000, dtype=np.float32) + server = mocker.Mock() + server._resolve_ref_audio_array = mocker.AsyncMock(return_value=(waveform, 24000, "key")) + server._resolve_ref_audio = mocker.AsyncMock(side_effect=AssertionError("unnecessary list conversion")) + adapter = CosyVoice3Adapter(SpeechServingContext(server=server)) + request = SimpleNamespace(input="Target text.", ref_audio="reference", ref_text="Reference text.", voice=None) + prompt = await adapter._build_prompt(request) + actual, rate = prompt["multi_modal_data"]["audio"] + np.testing.assert_array_equal(actual, waveform) + assert rate == 24000 and actual.dtype == np.float32 + assert not np.shares_memory(actual, waveform) + actual[0] = 100 + assert waveform[0] == -1 + assert prompt["prompt"] == request.input + assert prompt["mm_processor_kwargs"]["prompt_text"].endswith("<|endofprompt|>Reference text.") + server._resolve_ref_audio.assert_not_awaited() + server._resolve_ref_audio_array.assert_awaited_once_with("reference") + + +@pytest.mark.asyncio +async def test_cosyvoice_reference_prefetch_resamples_and_warms_conditioning(mocker): + from types import SimpleNamespace + + import vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 as cosyvoice3 + from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext + from vllm_omni.entrypoints.openai.tts_adapters.cosyvoice3 import CosyVoice3Adapter + + waveform = np.linspace(-1, 1, 24000, dtype=np.float32) + server = mocker.Mock() + server._resolve_ref_audio_array = mocker.AsyncMock(return_value=(waveform, 24000, "key")) + server.model_config.hf_config = SimpleNamespace(target_sr=16000) + server.model_config.model = "model-dir" + warm = mocker.patch.object(cosyvoice3, "prefetch_reference_conditioning") + adapter = CosyVoice3Adapter(SpeechServingContext(server=server)) + request = SimpleNamespace(input="Target text.", ref_audio="reference", ref_text="Reference text.", voice=None) + prompt = await adapter._build_prompt(request) + audio, rate = prompt["multi_modal_data"]["audio"] + # The processor receives samples at its own rate, so its resampling is a + # no-op and its content-addressed cache lookup hits the warmed entry. + assert rate == 16000 and prompt["mm_processor_kwargs"]["sample_rate"] == 16000 + assert audio.shape == (16000,) and audio.dtype == np.float32 + warm.assert_called_once() + model_dir, config, (warmed, warmed_rate) = warm.call_args.args + assert model_dir == "model-dir" and config is server.model_config.hf_config + assert warmed is audio and warmed_rate == 16000 + assert waveform[0] == -1 + + +@pytest.mark.parametrize("mode", ["ras", "standard"]) +def test_cosyvoice_sampling_mode_sets_control_stops_without_mutating_defaults(mocker, mode): + from vllm.sampling_params import SamplingParams + + from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest + from vllm_omni.entrypoints.openai.tts_adapters.base import SpeechServingContext + from vllm_omni.entrypoints.openai.tts_adapters.cosyvoice3 import CosyVoice3Adapter + from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config + + config = CosyVoice3Config() + config.cosyvoice3_sampling_mode = mode + serving = mocker.Mock(spec=OmniOpenAIServingSpeech) + serving.model_config = mocker.Mock(hf_config=config) + adapter = CosyVoice3Adapter(SpeechServingContext(server=serving)) + adapter._tokenizer = mocker.Mock() + mocker.patch("vllm_omni.model_executor.models.cosyvoice3.utils.extract_text_token", return_value=(None, 9)) + defaults = [SamplingParams(stop_token_ids=[6562])] + request = OpenAICreateSpeechRequest(input="Example text.", max_new_tokens=2048) + result = adapter.apply_sampling_overrides(defaults, request) + assert defaults[0].stop_token_ids == [6562] + assert result[0].min_tokens == 18 + assert result[0].max_tokens == 2048 + expected = set(range(6561, 6761)) if mode == "standard" else {6562} + assert set(result[0].stop_token_ids) == expected + assert expected <= result[0].all_stop_token_ids diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_f0_weight_norm.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_f0_weight_norm.py new file mode 100644 index 000000000..2e149a815 --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_f0_weight_norm.py @@ -0,0 +1,283 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""F0 folding and joint C4/C5 loading regressions with the existing CPU F0 policy.""" + +import pytest +import torch +from torch import nn +from torch.nn.utils import parametrize + +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core import hifigan + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.filterwarnings("ignore:.*weight_norm.*:FutureWarning"), +] + + +@pytest.fixture(autouse=True) +def fixed_seed(): + with torch.random.fork_rng(devices=[]): + torch.random.default_generator.manual_seed(5) + yield + + +@pytest.fixture(params=["parametrized", "legacy"]) +def weight_norm_api(request, monkeypatch): + if request.param == "legacy": + monkeypatch.setattr(hifigan, "weight_norm", nn.utils.weight_norm) + return request.param + + +@pytest.fixture( + params=[torch.enable_grad, torch.no_grad, torch.inference_mode], ids=["grad", "no_grad", "inference_mode"] +) +def fold_context(request): + return request.param + + +@pytest.fixture( + params=[ + pytest.param("cpu", marks=pytest.mark.cpu), + pytest.param("cuda", marks=pytest.mark.cuda), + ] +) +def runtime(request): + """Only the generator may run on CUDA; production F0 still runs on CPU.""" + on_cuda = request.param == "cuda" + if on_cuda and not torch.cuda.is_available(): + pytest.skip("CUDA is required for the generator/device-placement regressions") + with torch.random.fork_rng(devices=[0] if on_cuda else []): + torch.manual_seed(5) + with torch.backends.cudnn.flags(benchmark=False, deterministic=True, allow_tf32=False): + yield torch.device("cuda:0" if on_cuda else "cpu") + + +def _has_weight_norm(module): + return parametrize.is_parametrized(module, "weight") or ( + hasattr(module, "weight_g") and hasattr(module, "weight_v") + ) + + +def _f0_layers(f0): + return [layer for layer in f0.condnet if isinstance(layer, hifigan.CausalConv1d)] + + +def _make_hift(): + # Source noise/phase tensors are not in state_dict. Give independent + # instances identical construction RNG without changing the caller's RNG. + with torch.random.fork_rng(devices=[]): + torch.random.default_generator.manual_seed(5) + return hifigan.CausalHiFTGenerator( + base_channels=32, + sampling_rate=24000, + upsample_rates=[8, 5, 3], + upsample_kernel_sizes=[16, 11, 7], + resblock_kernel_sizes=[3], + resblock_dilation_sizes=[[1, 3, 5]], + source_resblock_kernel_sizes=[3, 3, 3], + source_resblock_dilation_sizes=[[1, 3, 5]] * 3, + f0_predictor=hifigan.CausalConvRNNF0Predictor(cond_channels=16), + ) + + +def _make_loader(): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + # Real load_weights and vocoder; only the unrelated Flow is reduced. + model = CosyVoice3Code2Wav.__new__(CosyVoice3Code2Wav) + nn.Module.__init__(model) + model.flow_model = nn.Linear(2, 2) + model.hift = _make_hift().train() + model._hift_window_len = 64 + return model + + +def _write_checkpoint(model, directory, legacy_checkpoint=False): + state = {name: value.detach().clone() for name, value in model.hift.state_dict().items()} + # Change conv_pre and F0 gains to detect stale constructor weights without + # amplifying every generator layer and overflowing its magnitude exp(). + for name, value in state.items(): + if name.startswith(("conv_pre.", "f0_predictor.")) and name.endswith( + ("weight_g", "parametrizations.weight.original0") + ): + value.mul_(2) + if legacy_checkpoint: + state = { + name.replace("parametrizations.weight.original0", "weight_g").replace( + "parametrizations.weight.original1", "weight_v" + ): value + for name, value in state.items() + } + torch.save(model.flow_model.state_dict(), directory / "flow.pt") + torch.save({"generator." + name: value for name, value in state.items()}, directory / "hift.pt") + return state + + +@pytest.mark.cpu +@pytest.mark.parametrize("finalize", [False, True]) +def test_f0_fold_is_exact_and_idempotent(weight_norm_api, finalize): + f0 = hifigan.CausalConvRNNF0Predictor(cond_channels=16).eval() + layers = _f0_layers(f0) + assert len(layers) == 5 and all(_has_weight_norm(layer) for layer in layers) + assert hasattr(layers[0], "weight_g") == (weight_norm_api == "legacy") + classifier = {name: value.detach().clone() for name, value in f0.classifier.state_dict().items()} + context = f0.left_context_frames + mel = torch.randn(1, 80, 24) + with torch.no_grad(): + before = f0(mel, finalize=finalize) + assert f0.remove_weight_norm() == 5 + weights = [layer.weight for layer in layers] + assert f0.remove_weight_norm() == 0 + assert all(layer.weight is weight for layer, weight in zip(layers, weights)) + assert not any(_has_weight_norm(module) for module in f0.modules()) + assert f0.left_context_frames == context + with torch.no_grad(): + after = f0(mel, finalize=finalize) + assert before.numel() > 0 and torch.isfinite(before).all() + torch.testing.assert_close(after, before, rtol=0, atol=0) + for name, value in f0.classifier.state_dict().items(): + torch.testing.assert_close(value, classifier[name], rtol=0, atol=0) + + +@pytest.mark.cpu +def test_f0_fold_preserves_autograd_usability(weight_norm_api, fold_context): + f0 = hifigan.CausalConvRNNF0Predictor(cond_channels=16).eval() + mel = torch.randn(1, 80, 24, requires_grad=True) + with torch.no_grad(): + before = f0(mel) + with fold_context(): + assert f0.remove_weight_norm() == 5 + assert f0.remove_weight_norm() == 0 + for layer in _f0_layers(f0): + assert isinstance(layer.weight, nn.Parameter) + assert not layer.weight.is_inference() + out = f0(mel) + torch.testing.assert_close(out, before, rtol=0, atol=0) + out.sum().backward() + assert mel.grad is not None and torch.isfinite(mel.grad).all() + + +@pytest.mark.cpu +def test_f0_fold_leaves_generator_norms_unchanged(weight_norm_api): + hift = _make_hift() + generator_norms = { + name + for name, module in hift.named_modules() + if not name.startswith("f0_predictor.") and _has_weight_norm(module) + } + assert len(generator_norms) == 41 + assert sum(_has_weight_norm(module) for module in hift.f0_predictor.modules()) == 5 + assert hift.f0_predictor.remove_weight_norm() == 5 + assert {name for name, module in hift.named_modules() if _has_weight_norm(module)} == generator_norms + + +@pytest.mark.parametrize( + "weight_norm_api,legacy_checkpoint", + [("parametrized", False), ("parametrized", True), ("legacy", True)], + indirect=["weight_norm_api"], +) +def test_load_folds_loaded_generator_and_f0_weights( + tmp_path, weight_norm_api, legacy_checkpoint, runtime, fold_context +): + device = runtime + model = _make_loader() + assert model.hift.training and model.hift.f0_predictor.training and model.flow_model.training + assert hasattr(model.hift.f0_predictor.condnet[0], "weight_g") == (weight_norm_api == "legacy") + state = _write_checkpoint(model, tmp_path, legacy_checkpoint) + # Do not forward the loader instance before loading. A separate unfolded + # reference refreshes legacy cached weights from the changed checkpoint. + reference = _make_hift() + reference.load_state_dict(state, strict=True) + reference.to(device).eval() + reference.f0_predictor.to(device="cpu", dtype=torch.float32) + normalized = {name: module for name, module in reference.named_modules() if _has_weight_norm(module)} + assert len(normalized) == 46 # 41 reduced-fixture generator layers + 5 F0 + mel = torch.randn(1, 80, 24, device=device) + expected = reference.inference(mel, finalize=True)[0] + + with fold_context(): + model.load_weights(str(tmp_path), device) + + assert not model.hift.training and not model.hift.f0_predictor.training and not model.flow_model.training + assert not any(_has_weight_norm(module) for module in model.hift.modules()) + for name, ref_layer in normalized.items(): + layer = model.hift.get_submodule(name) + assert isinstance(layer.weight, nn.Parameter) and not layer.weight.is_inference() + assert layer.weight.device == device + assert layer.weight.dtype == torch.float32 + torch.testing.assert_close(layer.weight, ref_layer.weight, rtol=0, atol=0) + assert {p.device for p in model.hift.f0_predictor.parameters()} == {device} + assert {p.dtype for p in model.hift.f0_predictor.parameters()} == {torch.float32} + assert model.hift.remove_weight_norm() == 0 + assert model.hift.f0_predictor.remove_weight_norm() == 0 + actual = model.hift.inference(mel, finalize=True)[0] + assert expected.numel() > 0 and torch.isfinite(expected).all() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.cpu +def test_invalid_checkpoint_does_not_fold_either_component(tmp_path, weight_norm_api, monkeypatch): + model = _make_loader() + state = _write_checkpoint(model, tmp_path) + state.pop("f0_predictor.classifier.bias") + torch.save({"generator." + name: value for name, value in state.items()}, tmp_path / "hift.pt") + + def unexpected_fold(*args, **kwargs): + pytest.fail("Folding must not run before strict checkpoint validation") + + monkeypatch.setattr(model.hift, "remove_weight_norm", unexpected_fold) + monkeypatch.setattr(model.hift.f0_predictor, "remove_weight_norm", unexpected_fold) + with pytest.raises(RuntimeError, match="Missing key"): + model.load_weights(str(tmp_path), torch.device("cpu")) + assert sum(_has_weight_norm(module) for module in model.hift.modules()) == 46 + + +@pytest.mark.parametrize("reference_fold_generator", [False, True], ids=["unfolded", "c4_only"]) +def test_joint_fold_preserves_streaming_audio_and_caches(tmp_path, weight_norm_api, runtime, reference_fold_generator): + device = runtime + model = _make_loader() + state = _write_checkpoint(model, tmp_path) + reference = _make_loader() + reference.hift.load_state_dict(state, strict=True) + reference.hift.to(device).eval() + reference.hift.f0_predictor.to(device="cpu", dtype=torch.float32) + # Compare both against no folding and against C4 alone, isolating C5's + # incremental effect while exercising the real loader/streaming path. + if reference_fold_generator: + assert reference.hift.remove_weight_norm() == 41 + model.load_weights(str(tmp_path), device) + assert not any(_has_weight_norm(module) for module in model.hift.modules()) + expected_norms = 5 if reference_fold_generator else 46 + assert sum(_has_weight_norm(module) for module in reference.hift.modules()) == expected_norms + chunks = torch.randn(1, 80, 144, device=device).split(24, dim=-1) + cache = ref_cache = None + shifted = False + for i, chunk in enumerate(chunks): + finalize = i == len(chunks) - 1 + expected, ref_cache = reference._stream_hift_from_feat(chunk, cache_state=ref_cache, finalize=finalize) + actual, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=finalize) + assert expected.numel() > 0 and torch.isfinite(expected).all() + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + if finalize: + assert cache is None and ref_cache is None + else: + assert cache.keys() == ref_cache.keys() + for key, value in ref_cache.items(): + if isinstance(value, torch.Tensor): + torch.testing.assert_close(cache[key], value, rtol=0, atol=0) + else: + assert cache[key] == value + shifted |= cache["mel_offset"] > 0 + assert shifted + + +@pytest.mark.cpu +def test_shipped_f0_config_folds_exactly_five_layers(weight_norm_api): + from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config + + f0 = hifigan.CausalConvRNNF0Predictor(**CosyVoice3Config().hift["f0_predictor"]) + assert sum(_has_weight_norm(module) for module in f0.modules()) == 5 + assert f0.remove_weight_norm() == 5 + assert f0.remove_weight_norm() == 0 + assert not any(_has_weight_norm(module) for module in f0.modules()) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_flow_graph.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_flow_graph.py new file mode 100644 index 000000000..1d850232d --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_flow_graph.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +import pytest +import torch + +from tests.helpers.mark import hardware_test +from tests.model_executor.models.cosyvoice3.test_cosyvoice3_full_batch import MelOutput, _items +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="requires CUDA")] + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("streaming", [False, True]) +@torch.inference_mode() +def test_flow_graph_replay_matches_eager_solve(monkeypatch, streaming): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("the opt-in packed backend requires Hopper FA3") + import vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav as module + from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config + + monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA") + monkeypatch.setenv("COSYVOICE3_FLOW_GRAPH", "1") + monkeypatch.setenv("COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS", "0" if streaming else "1") + monkeypatch.setenv("COSYVOICE3_PACKED_STREAMING", "1" if streaming else "0") + monkeypatch.setattr(module, "CausalHiFTGenerator", MelOutput) + config = CosyVoice3Config() + config.flow["pre_lookahead_layer"]["channels"] = 32 + # 256 channels in 16 conv groups keep the packed convolution (and graphs). + config.flow["decoder"]["estimator"].update(dim=256, depth=2, heads=4, dim_head=64) + model = module.CosyVoice3Code2Wav(config).cuda().bfloat16().eval() + if streaming: + monkeypatch.setattr( + model, "_stream_hift_from_feat", lambda mel, cache_state, finalize: (mel, None if finalize else cache_state) + ) + else: + model.hift.upsample_rates = [1] + model.hift.istft_params = {"hop_len": 1} + + def run(items, graph): + estimator = getattr(model, "_packed_estimator", None) + runner = None if estimator is None else estimator.graph_runner + if estimator is not None and not graph: + estimator.graph_runner = None + torch.manual_seed(3) + try: + if streaming: + return [audio for audio, _ in model.forward_streaming_batch(items, n_timesteps=3)] + return model.forward_batch(items, n_timesteps=3) + finally: + if estimator is not None: + estimator.graph_runner = runner + + extra = dict(token_offset_tokens=2, finalize=False, cache_state=None) if streaming else {} + # The second batch reuses the first one's padded layout with new lengths. + batches = [_items([(7, 60), (13, 41), (3, 97)], **extra), _items([(5, 71), (9, 30), (11, 88)], **extra)] + run(batches[0], graph=False) # compile + for items in batches: + actual = run(items, graph=True) + expected = run(items, graph=False) + for output, reference in zip(actual, expected): + torch.testing.assert_close(output, reference, rtol=1e-2, atol=1e-2) + runner = model._packed_estimator.graph_runner + assert len(runner.graphs) == 1 diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py new file mode 100644 index 000000000..3722bbcc7 --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from tests.helpers.mark import hardware_test +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="requires CUDA")] + + +class MelOutput(nn.Module): + def __init__(self, **kwargs): + super().__init__() + self.anchor = nn.Parameter(torch.zeros(1)) + self.m_source = SimpleNamespace(l_linear=SimpleNamespace(weight=self.anchor)) + + def inference(self, speech_feat, finalize): + assert finalize + return speech_feat, None, None + + +def _items(shapes, **extra): + return [ + dict( + token=torch.randint(0, 64, (1, length), device="cuda"), + prompt_token=torch.randint(0, 64, (1, prompt)), + prompt_feat=torch.randn(1, prompt * 2, 80, device="cuda"), + embedding=torch.randn(1, 192, device="cuda"), + **extra, + ) + for prompt, length in shapes + ] + + +def tiny_flow(monkeypatch): + import vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav as module + from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config + + monkeypatch.setenv("DIFFUSION_ATTENTION_BACKEND", "TORCH_SDPA") + monkeypatch.setattr(module, "CausalHiFTGenerator", MelOutput) + config = CosyVoice3Config() + config.flow["pre_lookahead_layer"]["channels"] = 32 + config.flow["decoder"]["estimator"].update(dim=32, depth=2, heads=4, dim_head=8) + return module.CosyVoice3Code2Wav(config).cuda().eval() + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("offset", [0, 3]) +@pytest.mark.parametrize("count", [3, 11]) +@torch.inference_mode() +def test_full_batch_ragged_prompt_matches_independent_flow(monkeypatch, offset, count): + model = tiny_flow(monkeypatch) + items = _items(([(7, 17), (13, 9), (3, 31)] * 4)[:count], token_offset_tokens=offset) + # Fix Flow noise so scheduling changes do not change the comparison input. + monkeypatch.setattr(torch, "randn", lambda shape, **kw: torch.zeros(shape, **kw)) + expected = [model.forward(**item, n_timesteps=3) for item in items] + actual = model.forward_batch(items, n_timesteps=3) + for result, reference in zip(actual, expected): + torch.testing.assert_close(result, reference, rtol=3e-4, atol=3e-5) + # Reordering unrelated requests must not move their conditioning or crop points. + reverse = model.forward_batch(items[::-1], n_timesteps=3) + for result, reference in zip(reverse, expected[::-1]): + torch.testing.assert_close(result, reference, rtol=3e-4, atol=3e-5) + assert model.forward_batch([]) == [] + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@torch.inference_mode() +def test_packed_full_response_preserves_request_isolation(monkeypatch): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("the opt-in packed backend requires Hopper FA3") + monkeypatch.setenv("COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS", "1") + model = tiny_flow(monkeypatch).bfloat16() + model.hift.upsample_rates = [1] + model.hift.istft_params = {"hop_len": 1} + items = _items([(7, 17), (13, 9), (3, 31)]) + monkeypatch.setattr(torch, "randn", lambda shape, **kw: torch.zeros(shape, **kw)) + expected = [model.forward(**item, n_timesteps=3) for item in items] + actual = model.forward_batch(items, n_timesteps=3) + for output, reference in zip(actual, expected): + torch.testing.assert_close(output, reference, rtol=0.04, atol=0.04) + # A different batch order must not leak the other requests' conditioning. + actual = model.forward_batch(items[::-1], n_timesteps=3) + for output, reference in zip(actual, expected[::-1]): + torch.testing.assert_close(output, reference, rtol=0.04, atol=0.04) + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_packed_stream_noise_growth_preserves_prefix_and_global_rng(): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + model = CosyVoice3Code2Wav.__new__(CosyVoice3Code2Wav) + nn.Module.__init__(model) + like = torch.empty(1, 80, 1, dtype=torch.bfloat16) + rng = torch.random.get_rng_state().clone() + first = model._stream_position_noise(97, like).clone() + grown = model._stream_position_noise(30001, like).clone() + torch.testing.assert_close(grown[..., :97], first, rtol=0, atol=0) + torch.testing.assert_close(torch.random.get_rng_state(), rng, rtol=0, atol=0) + direct = CosyVoice3Code2Wav.__new__(CosyVoice3Code2Wav) + nn.Module.__init__(direct) + torch.testing.assert_close(direct._stream_position_noise(30001, like), grown, rtol=0, atol=0) + assert model._stream_position_noise(0, like).shape == (1, 80, 0) + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@torch.inference_mode() +def test_packed_stream_attention_matches_chunk_causal_reference(): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("the opt-in packed backend requires Hopper FA3") + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.packed_dit import ( + RaggedRowAttention, + pack_rows, + packed_fa3, + ) + + lengths = [67, 103] + rows = pack_rows(lengths, torch.device("cuda")) + attention = RaggedRowAttention(rows, heads=2, head_dim=64, chunk_size=50) + generator = torch.Generator(device="cuda").manual_seed(7) + q, k, v = [ + torch.randn(sum(lengths), 2, 64, device="cuda", dtype=torch.bfloat16, generator=generator) for _ in range(3) + ] + actual = packed_fa3( + q, + k[:, None], + v[:, None], + attention.cache_seqlens, + attention.page_table, + attention.cu_seqlens_q, + attention.max_seqlen_q, + ) + start = 0 + for length in lengths: + position = torch.arange(length, device="cuda") + mask = position[None, :] < ((position[:, None] // 50 + 1) * 50) + expected = torch.nn.functional.scaled_dot_product_attention( + q[start : start + length].transpose(0, 1).float(), + k[start : start + length].transpose(0, 1).float(), + v[start : start + length].transpose(0, 1).float(), + attn_mask=mask, + ).transpose(0, 1) + torch.testing.assert_close(actual[start : start + length].float(), expected, rtol=0.03, atol=0.01) + start += length + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@torch.inference_mode() +def test_packed_stream_mixed_finalization_and_ragged_requests_stay_aligned(monkeypatch): + if torch.cuda.get_device_capability()[0] != 9: + pytest.skip("the opt-in packed backend requires Hopper FA3") + monkeypatch.setenv("COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS", "0") + monkeypatch.setenv("COSYVOICE3_PACKED_STREAMING", "1") + model = tiny_flow(monkeypatch).bfloat16() + monkeypatch.setattr( + model, "_stream_hift_from_feat", lambda mel, cache_state, finalize: (mel, None if finalize else cache_state) + ) + items = [] + for index, (prompt, generated) in enumerate(([(7, 17), (13, 9), (3, 31)] * 6)[:17]): + items.append( + dict( + token=torch.randint(0, 64, (1, generated), device="cuda"), + prompt_token=torch.randint(0, 64, (1, prompt)), + prompt_feat=torch.randn(1, prompt * 2, 80, device="cuda"), + embedding=torch.randn(1, 192, device="cuda"), + token_offset_tokens=3, + finalize=index % 3 == 0, + cache_state={"request_marker": torch.tensor(index)}, + ) + ) + expected = [model.forward_streaming_batch([item], n_timesteps=3)[0] for item in items] + for ordered, reference in [(items, expected), (items[::-1], expected[::-1])]: + actual = model.forward_streaming_batch(ordered, n_timesteps=3) + for (audio, state), (ref_audio, ref_state), item in zip(actual, reference, ordered): + torch.testing.assert_close(audio, ref_audio, rtol=0.04, atol=0.04) + assert state is ref_state + assert state is (None if item["finalize"] else item["cache_state"]) + assert model.forward_streaming_batch([]) == [] + with pytest.raises(ValueError, match="mixed finalization"): + model.forward_batch(items, n_timesteps=3, stream_items=True) + monkeypatch.setenv("COSYVOICE3_PACKED_STREAMING", "0") + with pytest.raises(ValueError, match="COSYVOICE3_PACKED_STREAMING"): + model.forward_batch(items[:1], n_timesteps=3, stream_items=True) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py new file mode 100644 index 000000000..713d437ca --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_weight_norm.py @@ -0,0 +1,264 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""CPU regressions for folding frozen HiFT generator weight normalization.""" + +import pytest +import torch +from torch import nn +from torch.nn.utils import parametrize + +from vllm_omni.model_executor.models.cosyvoice3.code2wav_core import hifigan + +pytestmark = [ + pytest.mark.core_model, + pytest.mark.cpu, + # The legacy weight_norm API warns on construction; the "legacy" params + # below exercise it on purpose. + pytest.mark.filterwarnings("ignore:.*weight_norm.*:FutureWarning"), +] + + +@pytest.fixture(autouse=True) +def fixed_seed(): + with torch.random.fork_rng(devices=[]): + torch.manual_seed(0) + yield + + +@pytest.fixture(params=["parametrized", "legacy"]) +def weight_norm_api(request, monkeypatch): + if request.param == "legacy": + monkeypatch.setattr(hifigan, "weight_norm", nn.utils.weight_norm) + return request.param + + +def _has_weight_norm(module): + return parametrize.is_parametrized(module, "weight") or ( + hasattr(module, "weight_g") and hasattr(module, "weight_v") + ) + + +def _make_hift(causal=True, sampling_rate=24000): + # Note: pairing the non-causal generator with CausalConvRNNF0Predictor is + # not a shipped configuration; it exists only to exercise the shared + # ResBlock/conv fold code on both generator classes. + cls = hifigan.CausalHiFTGenerator if causal else hifigan.HiFTGenerator + return cls( + base_channels=32, + sampling_rate=sampling_rate, + # Keep the real CosyVoice3 upsample ratio and SineGen2 (24 kHz), + # while shrinking channels and residual blocks for CPU tests. + upsample_rates=[8, 5, 3], + upsample_kernel_sizes=[16, 11, 7], + resblock_kernel_sizes=[3], + resblock_dilation_sizes=[[1, 3, 5]], + source_resblock_kernel_sizes=[3, 3, 3], + source_resblock_dilation_sizes=[[1, 3, 5]] * 3, + f0_predictor=hifigan.CausalConvRNNF0Predictor(cond_channels=16), + ).eval() + + +@pytest.mark.parametrize("causal", [False, True]) +def test_resblock_fold_is_exact_and_idempotent(weight_norm_api, causal): + block = hifigan.ResBlock(channels=4, causal=causal).eval() + x = torch.randn(1, 4, 12) + conv = block.convs1[0] + # Guard: the fixture really selected the intended API, so the legacy + # cases cannot silently re-test the parametrizations path. + if weight_norm_api == "legacy": + assert hasattr(conv, "weight_g") and not parametrize.is_parametrized(conv, "weight") + else: + assert parametrize.is_parametrized(conv, "weight") + + with torch.no_grad(): + before = block(x) + # Fold in normal grad mode, like the production load path (vLLM's loader + # does not wrap load_weights in no_grad/inference_mode): the materialized + # weight must be a real Parameter. + first = block.remove_weight_norm() + assert isinstance(conv.weight, nn.Parameter) + second = block.remove_weight_norm() + with torch.no_grad(): + once = block(x) + twice = block(x) + + assert (first, second) == (6, 0) # 3 dilations x {convs1, convs2}; no-op after + assert not any(_has_weight_norm(module) for module in block.modules()) + torch.testing.assert_close(once, before, rtol=0, atol=0) + torch.testing.assert_close(twice, before, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "causal,sampling_rate,finalize", + [(False, 22050, True), (True, 22050, True), (True, 24000, True), (True, 24000, False)], +) +def test_hift_fold_preserves_audio_and_leaves_f0_unchanged(weight_norm_api, causal, sampling_rate, finalize): + hift = _make_hift(causal, sampling_rate) + mel = torch.randn(1, 80, 20) + if weight_norm_api == "legacy": + assert hasattr(hift.conv_pre, "weight_g") and not parametrize.is_parametrized(hift.conv_pre, "weight") + normalized = {name: module for name, module in hift.named_modules() if _has_weight_norm(module)} + f0_state = {name: value.clone() for name, value in hift.f0_predictor.state_dict().items()} + assert any(name.startswith("source_resblocks.") for name in normalized) + assert any(name.startswith("resblocks.") for name in normalized) + # The C5-scope layers must start out weight-normed, or the "f0 untouched" + # asserts below would pass vacuously. + assert any(name.startswith("f0_predictor.") for name in normalized) + + # Non-causal HiFT draws fresh source noise on each invocation. + rng = torch.get_rng_state() + kwargs = {"finalize": finalize} if causal else {} + before = hift.inference(mel, **kwargs)[0] + # 3 ups + conv_pre + conv_post + (3 resblocks + 3 source resblocks) x 6 convs + assert hift.remove_weight_norm() == 41 + assert hift.remove_weight_norm() == 0 + torch.set_rng_state(rng) + after = hift.inference(mel, **kwargs)[0] + + assert before.numel() > 0 + assert torch.isfinite(before).all() + torch.testing.assert_close(after, before, rtol=0, atol=0) + for name, module in normalized.items(): + # The generator entry point must leave F0 to its separate fold call. + assert _has_weight_norm(module) == name.startswith("f0_predictor.") + for name, value in hift.f0_predictor.state_dict().items(): + torch.testing.assert_close(value, f0_state[name], rtol=0, atol=0) + + +@pytest.mark.parametrize( + "weight_norm_api,legacy_checkpoint", + [("parametrized", False), ("parametrized", True), ("legacy", True)], + indirect=["weight_norm_api"], +) +def test_load_weights_folds_loaded_not_initial_weights(tmp_path, weight_norm_api, legacy_checkpoint): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav + + # Exercise the real loader with a small vocoder and synthetic checkpoints. + # __new__ skips CosyVoice3Code2Wav.__init__ (which would build the DiT and + # flow decoder); load_weights only touches .flow_model and .hift. + model = CosyVoice3Code2Wav.__new__(CosyVoice3Code2Wav) + nn.Module.__init__(model) + model.flow_model = nn.Linear(2, 2) + # Start in training mode to check the loader's transition to eval. + model.hift = _make_hift().train() + assert model.hift.training + conv = model.hift.conv_pre + assert hasattr(conv, "weight_g") == (weight_norm_api == "legacy") + state = {name: value.clone() for name, value in model.hift.state_dict().items()} + gain = "weight_g" if weight_norm_api == "legacy" else "parametrizations.weight.original0" + state[f"conv_pre.{gain}"].mul_(2) + expected_weight = conv.weight.detach().clone() * 2 + assert not torch.equal(conv.weight, expected_weight) + # No forward before loading: a legacy fold must recompute from the loaded + # g/v parameters, not reuse the cached effective weight from construction. + flow_state = {name: torch.ones_like(value) for name, value in model.flow_model.state_dict().items()} + if legacy_checkpoint: + state = { + name.replace("parametrizations.weight.original0", "weight_g").replace( + "parametrizations.weight.original1", "weight_v" + ): value + for name, value in state.items() + } + torch.save(flow_state, tmp_path / "flow.pt") + torch.save({"generator." + name: value for name, value in state.items()}, tmp_path / "hift.pt") + + model.load_weights(str(tmp_path), torch.device("cpu")) + + torch.testing.assert_close(model.hift.conv_pre.weight, expected_weight, rtol=0, atol=0) + # Production loads in normal grad mode, so the folded weight is a Parameter. + assert isinstance(model.hift.conv_pre.weight, nn.Parameter) + assert not model.hift.conv_pre.weight.is_inference() + assert not any( + _has_weight_norm(module) for name, module in model.hift.named_modules() if not name.startswith("f0_predictor.") + ) + assert not any(_has_weight_norm(module) for module in model.hift.f0_predictor.modules()) + assert not model.hift.training + assert not model.flow_model.training + for name, value in model.flow_model.state_dict().items(): + torch.testing.assert_close(value, flow_state[name], rtol=0, atol=0) + + +@pytest.mark.parametrize( + "fold_context", + [torch.enable_grad, torch.no_grad, torch.inference_mode], + ids=["grad", "no_grad", "inference_mode"], +) +def test_fold_stays_usable_outside_grad_context(weight_norm_api, fold_context): + """Both APIs must leave non-inference Parameters usable by autograd.""" + block = hifigan.ResBlock(channels=4).eval() + x = torch.randn(1, 4, 12, requires_grad=True) + assert hasattr(block.convs1[0], "weight_g") == (weight_norm_api == "legacy") + with torch.no_grad(): + before = block(x) + with fold_context(): + assert block.remove_weight_norm() == 6 + assert block.remove_weight_norm() == 0 + for conv in (*block.convs1, *block.convs2): + assert isinstance(conv.weight, nn.Parameter) + assert not conv.weight.is_inference() + # no_grad would hide the failure when autograd tries to save a weight. + out = block(x) + torch.testing.assert_close(out, before, rtol=0, atol=0) + out.sum().backward() + assert x.grad is not None + assert torch.isfinite(x.grad).all() + + +@pytest.mark.parametrize("sampling_rate", [22050, 24000]) +def test_fold_preserves_multichunk_audio_and_finalize(weight_norm_api, sampling_rate): + from tests.model_executor.models.cosyvoice3.test_cosyvoice3_incremental_hift import _make_model + + hift = _make_hift(sampling_rate=sampling_rate) + model = _make_model(hift, window_len=64) + mel = torch.randn(1, 80, 192) + chunks = mel.split(24, dim=-1) + + def stream(): + cache = None + audio = [] + window_shifted = False + for i, chunk in enumerate(chunks): + speech, cache = model._stream_hift_from_feat(chunk, cache_state=cache, finalize=i == len(chunks) - 1) + audio.append(speech) + if cache is not None: + window_shifted |= cache["mel_offset"] > 0 + assert window_shifted + assert cache is None + return torch.cat(audio, dim=-1) + + before = stream() + assert hift.remove_weight_norm() == 41 + after = stream() + assert before.numel() > 0 + assert torch.isfinite(before).all() + torch.testing.assert_close(after, before, rtol=0, atol=0) + + +def test_shipped_config_folds_77_generator_and_five_f0_layers(weight_norm_api): + """Pin separate C4/C5 counts for both APIs on the shipped configuration. + + Generator: conv_pre 1 + ups 3 + 9 resblocks x 6 + 3 source resblocks x 6 + + conv_post 1 = 77. F0 adds five condnet convolutions. Construction only + (no forward), so the full 512-channel model stays cheap on CPU. + """ + from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config + + hift_cfg = dict(CosyVoice3Config().hift) + f0_cfg = hift_cfg.pop("f0_predictor") + f0 = hifigan.CausalConvRNNF0Predictor( + num_class=f0_cfg["num_class"], + in_channels=f0_cfg["in_channels"], + cond_channels=f0_cfg["cond_channels"], + ) + generator = hifigan.CausalHiFTGenerator(f0_predictor=f0, **hift_cfg) + + assert sum(_has_weight_norm(module) for module in generator.modules()) == 82 + assert generator.remove_weight_norm() == 77 + assert generator.remove_weight_norm() == 0 + remaining = [name for name, m in generator.named_modules() if _has_weight_norm(m)] + assert len(remaining) == 5 + assert all(name.startswith("f0_predictor.") for name in remaining) + assert f0.remove_weight_norm() == 5 + assert f0.remove_weight_norm() == 0 + assert generator.remove_weight_norm() == 0 + assert not any(_has_weight_norm(module) for module in generator.modules()) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py index 1d583de6c..6853d0ff4 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py @@ -511,6 +511,37 @@ def test_sample_keeps_only_finite_token_after_ras_rejection(): assert out.sampled_token_ids.tolist() == [[1]] +def test_batched_ras_matches_serial_seeded_requests_and_rng_state(): + model = _make_talker_model() + logits = torch.tensor([[2.0, 1.0, 0.0, -1.0], [0.0, 3.0, 1.0, -2.0], [1.0, 0.0, 2.0, -1.0]]) + metadata = _make_sampling_metadata(output_token_ids=[[0] * 10, [], [2] * 10]) + metadata.temperature = torch.tensor([0.5, 1.0, 1.5]) + metadata.top_k = torch.tensor([1, 2, 4]) + metadata.top_p = torch.tensor([0.7, 0.8, 0.9]) + metadata.generators = {i: torch.Generator().manual_seed(400 + i) for i in range(3)} + reference_generators = {i: torch.Generator().manual_seed(400 + i) for i in range(3)} + for _ in range(6): + expected = [] + for i in range(3): + scores = torch.log_softmax(logits[i] / metadata.temperature[i], dim=0) + expected.append( + model._ras_sample_one( + scores, + metadata.output_token_ids[i], + top_p=float(metadata.top_p[i]), + top_k=int(metadata.top_k[i]), + win_size=10, + tau_r=0.1, + generator=reference_generators[i], + ) + ) + actual = model._ras_sample_batch(logits, metadata, default_top_p=0.8, default_top_k=25, win_size=10, tau_r=0.1) + assert actual.tolist() == expected + for i, token in enumerate(expected): + assert torch.equal(metadata.generators[i].get_state(), reference_generators[i].get_state()) + metadata.output_token_ids[i].append(token) + + def test_gpu_ar_model_runner_prefers_model_sampler_when_opted_in(): metadata = _make_sampling_metadata(output_token_ids=[[1, 2, 3]]) expected = SamplerOutput( @@ -637,3 +668,264 @@ def test_gpu_ar_model_runner_repairs_async_placeholders_for_model_sampler(): assert runner.input_batch.async_copy_ready_event.synced is True assert runner.input_batch.update_async_called is True assert seen_histories == [[[11, 29]]] + + +@pytest.mark.parametrize( + "order", + [("decode", "short", "decode", "long"), ("short", "decode", "long", "decode"), ("long", "short", "decode")], +) +def test_embed_input_ids_preserves_interleaved_request_boundaries(order): + model = _make_talker_model() + speech = nn.Embedding.from_pretrained(torch.arange(100 * 4, dtype=torch.float32).reshape(100, 4)) + text = nn.Embedding.from_pretrained(torch.arange(100 * 4, dtype=torch.float32).reshape(100, 4) + 1000) + model.model = SimpleNamespace( + speech_embedding=speech, sos=0, task_id=1, llm=SimpleNamespace(model=SimpleNamespace(embed_tokens=text)) + ) + requests = { + "decode": (torch.tensor([7]), torch.tensor([False]), None, speech.weight[7:8]), + "short": ( + torch.tensor([1, 1, 1, 1, 21]), + torch.tensor([True] * 4 + [False]), + speech.weight[2:4], + torch.cat((speech.weight[0:1], text.weight[21:22], speech.weight[1:2], speech.weight[2:4])), + ), + "long": ( + torch.tensor([1, 1, 1, 1, 1, 30, 31]), + torch.tensor([True] * 5 + [False] * 2), + speech.weight[4:7], + torch.cat((speech.weight[0:1], text.weight[30:32], speech.weight[1:2], speech.weight[4:7])), + ), + } + rows = [requests[name] for name in order] + boundaries = [0] + for ids, *_ in rows: + boundaries.append(boundaries[-1] + ids.numel()) + actual = model.embed_input_ids( + torch.cat([r[0] for r in rows]), + multimodal_embeddings=[r[2] for r in rows if r[2] is not None], + is_multimodal=torch.cat([r[1] for r in rows]), + query_start_loc=boundaries, + ) + torch.testing.assert_close(actual, torch.cat([r[3] for r in rows]), rtol=0, atol=0) + conditioning = [torch.full((1, 192), i + 1.0) for i in range(sum(name != "decode" for name in order))] + aligned = model._align_prompt_conditioning(conditioning) + assert len(aligned) == len(order) + index = 0 + for name, value in zip(order, aligned): + if name == "decode": + assert value is None + else: + assert value is conditioning[index] + index += 1 + with pytest.raises(ValueError, match="conditioning must match"): + model._align_prompt_conditioning(conditioning[:-1]) + + +def test_mrv2_conditioning_follows_encoded_prompts_to_batch_rows(): + from collections import OrderedDict + + model = _make_talker_model() + speech = nn.Embedding.from_pretrained(torch.arange(100 * 4, dtype=torch.float32).reshape(100, 4)) + text = nn.Embedding.from_pretrained(torch.arange(100 * 4, dtype=torch.float32).reshape(100, 4) + 1000) + model.model = SimpleNamespace( + speech_embedding=speech, sos=0, task_id=1, llm=SimpleNamespace(model=SimpleNamespace(embed_tokens=text)) + ) + model._mrv2_encoded_conditioning = OrderedDict() + tokens = [torch.tensor([2, 3]), torch.tensor([4, 5, 6])] + feats = [torch.full((4, 2), 1.0), torch.full((6, 2), 2.0)] + speakers = [torch.full((1, 3), 1.0), torch.full((1, 3), 2.0)] + encoded = model.embed_multimodal(speech_token=tokens, speech_feat=feats, embedding=speakers) + # Batch: decode, short prompt, decode, long prompt, short again (same + # reference: the encoder cache hands back the same tensor). + ids = torch.tensor([7, 1, 1, 1, 1, 21, 8, 1, 1, 1, 1, 1, 30, 1, 1, 1, 1, 22]) + is_mm = torch.tensor([False] + [True] * 4 + [False, False] + [True] * 5 + [False] + [True] * 4 + [False]) + model.embed_input_ids( + ids, + multimodal_embeddings=[encoded[0], encoded[1], encoded[0]], + is_multimodal=is_mm, + query_start_loc=[0, 1, 6, 7, 13, 18], + ) + output = model.make_omni_output(torch.zeros(18, 4), model_intermediate_buffer=[None] * 5) + embed = output.multimodal_outputs["embed"] + assert [value is None for value in embed["speech_token"]] == [True, False, True, False, False] + for row, item in ((1, 0), (3, 1), (4, 0)): + torch.testing.assert_close(embed["speech_token"][row], tokens[item][None]) + torch.testing.assert_close(embed["speech_feat"][row], feats[item][None]) + torch.testing.assert_close(embed["embedding"][row], speakers[item]) + # A following decode step carries no conditioning. + model.embed_input_ids( + torch.tensor([9, 10]), multimodal_embeddings=[], is_multimodal=torch.zeros(2, dtype=torch.bool) + ) + assert model.make_omni_output(torch.zeros(2, 4)).multimodal_outputs == {} + with pytest.raises(ValueError, match="lost its prompt conditioning"): + model.embed_input_ids( + torch.tensor([1, 1, 1, 1, 21]), + multimodal_embeddings=[speech.weight[2:4].clone()], + is_multimodal=torch.tensor([True] * 4 + [False]), + query_start_loc=[0, 5], + ) + + +def test_talker_declares_its_narrow_logits_head(): + model = _make_talker_model() + assert model.logits_vocab_size == 6561 + 200 + model.model_stage = "cosyvoice3_code2wav" + assert model.logits_vocab_size is None + + +def test_embed_input_ids_rejects_cross_request_placeholder_block(): + model = _make_talker_model() + speech = nn.Embedding(100, 4) + model.model = SimpleNamespace( + speech_embedding=speech, + sos=0, + task_id=1, + llm=SimpleNamespace(model=SimpleNamespace(embed_tokens=nn.Embedding(100, 4))), + ) + with pytest.raises(ValueError, match="complete prefill block"): + model.embed_input_ids( + torch.tensor([7, 1, 1, 1, 1, 20]), + multimodal_embeddings=[speech.weight[2:4]], + is_multimodal=torch.tensor([False, True, True, True, True, False]), + ) + + +def test_full_response_singleton_uses_the_optimized_batch_path(monkeypatch): + import vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 as module + + monkeypatch.setenv("COSYVOICE3_BATCH_FLOW", "1") + monkeypatch.setattr(module, "cosyvoice3_packed_inference_enabled", lambda: True) + model = _make_code2wav_model() + model.forward( + input_ids=torch.tensor([0, 1, 2]), + positions=torch.arange(3), + seq_token_counts=[3], + model_intermediate_buffer=[ + { + "embed": { + "speech_token": torch.tensor([[1, 2, 3]]), + "speech_feat": torch.ones(1, 6, 2), + "embedding": torch.ones(1, 2), + }, + "meta": {"req_id": ["single"], "stream_finished": torch.tensor(True), "left_context_size": 0}, + } + ], + ) + assert len(model.code2wav.forward_streaming_batch_calls) == 1 + assert model.code2wav.forward_streaming_batch_calls[0][0]["finalize"] + + +def test_full_response_adapter_uses_live_prefill_then_empty_decode(monkeypatch): + from vllm.v1.worker.gpu_model_runner import GPUModelRunner + + from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner + + model = _make_talker_model() + hidden = torch.ones(2, 4) + runner = object.__new__(OmniGPUModelRunner) + runner.model = model + runner._build_model_kwargs_extra = lambda: {} + # Replay the same graph tensor while the live conditioning changes. + monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *_args, **_kwargs: hidden) + model._conditioning_request_rows = [1] + model._conditioning_request_count = 2 + prompt = torch.tensor([[7, 8]]) + output = runner._model_forward(speech_token=prompt) + values = output.multimodal_outputs["embed"]["speech_token"] + assert values[0] is None + torch.testing.assert_close(values[1], prompt) + decode = runner._model_forward() + assert decode.text_hidden_states is hidden + assert decode.multimodal_outputs == {} + + +def test_packed_flow_preserves_torch_estimator_when_speaker_trt_is_enabled(monkeypatch): + import vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 as module + + model = _make_code2wav_model() + model._code2wav_trt_done = False + monkeypatch.setattr(module.torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(module, "cosyvoice3_packed_inference_enabled", lambda: True) + monkeypatch.setattr(module, "_cosyvoice3_trt_enabled", lambda: True) + + def unexpected_estimator_resolution(): + raise AssertionError("Packed Flow must not resolve or replace its torch estimator") + + monkeypatch.setattr(model, "_resolve_flow_estimator_onnx", unexpected_estimator_resolution) + model._maybe_enable_code2wav_trt() + assert model._code2wav_trt_done + + +def test_standard_sampling_preserves_parameters_and_excludes_prompt_penalties(monkeypatch): + model = _make_talker_model() + model.config.cosyvoice3_sampling_mode = "standard" + metadata = _make_sampling_metadata(output_token_ids=[[2, 3]], repetition_penalty=1.21) + metadata.prompt_token_ids = torch.tensor([[2, 4]], dtype=torch.long) + metadata.temperature = torch.tensor([0.7]) + metadata.top_k = torch.tensor([20]) + captured = {} + expected = SamplerOutput(sampled_token_ids=torch.tensor([[5]], dtype=torch.int32), logprobs_tensors=None) + + def standard_sampler(*, logits, sampling_metadata): + captured["metadata"] = sampling_metadata + return expected + + model._talker_sampler = standard_sampler + monkeypatch.setattr(model, "_ras_sample_batch", lambda *a, **k: pytest.fail("RAS must not run in standard mode")) + assert model.sample(torch.zeros(1, 8), metadata) is expected + actual = captured["metadata"] + assert actual is not metadata + assert actual.output_token_ids == [[2, 3]] + torch.testing.assert_close(actual.prompt_token_ids, torch.tensor([[8, 8]])) + torch.testing.assert_close(metadata.prompt_token_ids, torch.tensor([[2, 4]])) + assert actual.temperature is metadata.temperature + assert actual.top_k is metadata.top_k + assert actual.top_p is metadata.top_p + assert actual.repetition_penalties is metadata.repetition_penalties + + +@pytest.mark.parametrize("mode", ["ras", "standard"]) +def test_sampling_mode_preserves_or_merges_control_logits(mode): + model = _make_talker_model() + model.config.cosyvoice3_sampling_mode = mode + model.config.vocab_size = 6770 + model.model = nn.Module() + model.model.llm_decoder = lambda hidden: hidden.clone() + hidden = torch.linspace(-2, 2, 6761).reshape(1, -1) + logits = model.compute_logits(hidden) + torch.testing.assert_close(logits[:, :6561], hidden[:, :6561]) + assert torch.isneginf(logits[:, 6761:]).all() + if mode == "standard": + torch.testing.assert_close(logits[:, 6561:6761], hidden[:, 6561:]) + else: + torch.testing.assert_close(logits[:, 6562], torch.logsumexp(hidden[:, 6561:], dim=-1)) + assert torch.isneginf(logits[:, 6561]).all() + assert torch.isneginf(logits[:, 6563:6761]).all() + + +def test_cancelled_stream_releases_only_its_vocoder_state(): + model_cls, _ = _cosyvoice3_model_and_runner() + model = SimpleNamespace( + _stream_audio_cache_lock=Lock(), + _stream_vocoder_cache_by_req={"cancelled": {"mel": torch.ones(2)}, "live": {"mel": torch.zeros(2)}}, + ) + model_cls.on_requests_finished(model, {"cancelled", "unknown"}) + model_cls.on_requests_finished(model, {"cancelled"}) + assert set(model._stream_vocoder_cache_by_req) == {"live"} + + +@pytest.mark.parametrize("mrv2,packed", [(False, False), (False, True), (True, False)]) +def test_talker_output_contract(monkeypatch, mrv2, packed): + import vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 as mod + + monkeypatch.setattr(mod, "cosyvoice3_packed_inference_enabled", lambda: packed) + model = _make_talker_model() + hidden = torch.zeros(2, 4) + model.model = SimpleNamespace(llm=lambda embeddings, positions: embeddings) + sentinel = object() + model.make_omni_output = lambda *args, **kwargs: sentinel + if mrv2: + model._sampling_eps = 1e-6 + model.mrv2_custom_sampler(SimpleNamespace(penalties_state=None)) + output = model.forward(torch.ones(2, dtype=torch.long), torch.arange(2), inputs_embeds=hidden) + assert output is (hidden if mrv2 or packed else sentinel) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_mrv2.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_mrv2.py new file mode 100644 index 000000000..7b965a107 --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_mrv2.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Model Runner V2 penalties must count generated speech tokens only.""" + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from tests.helpers.mark import hardware_test +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="requires CUDA")] + + +class _Uva: + def copy_to_uva(self) -> None: + pass + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_generated_only_penalty_writes_skip_prompt_and_keep_resumed_outputs(): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _generated_only_penalty_writes + + vocab, rows, width = 6761, 3, 16 + tokens = torch.zeros(rows, width, dtype=torch.int32, device="cuda") + # Row 0: fresh request whose prompt holds text ids beyond the speech head. + tokens[0, :4] = torch.tensor([151645, 12, 7, 99999]) + # Row 1: resumed request, prompt [5, 6] then generated [7, 7, 9]. + tokens[1, :5] = torch.tensor([5, 6, 7, 7, 9]) + prompt_len = torch.tensor([4, 2, 0], dtype=torch.int32, device="cuda") + prefill_len = torch.tensor([4, 5, 0], dtype=torch.int32) + state = SimpleNamespace( + _new_penalties_reqs=[0, 1], + device=torch.device("cuda"), + prompt_bin_mask=torch.full((rows, (vocab + 31) // 32), -1, dtype=torch.int32, device="cuda"), + output_bin_counts=torch.full((rows, vocab), 3, dtype=torch.int32, device="cuda"), + req_states=SimpleNamespace( + all_token_ids=SimpleNamespace(gpu=tokens), + prompt_len=SimpleNamespace(gpu=prompt_len), + prefill_len=SimpleNamespace(gpu=prefill_len.cuda(), np=np.asarray(prefill_len)), + ), + repetition_penalty=_Uva(), + frequency_penalty=_Uva(), + presence_penalty=_Uva(), + ) + _generated_only_penalty_writes(state) + torch.accelerator.synchronize() + assert state._new_penalties_reqs == [] + # No prompt token is penalized for the new requests. + assert int(state.prompt_bin_mask[:2].abs().sum()) == 0 + counts = state.output_bin_counts + assert int(counts[0].sum()) == 0 + assert counts[1, 7].item() == 2 and counts[1, 9].item() == 1 and int(counts[1].sum()) == 3 + # Rows that were not added keep their state. + assert int(counts[2, 0]) == 3 and int(state.prompt_bin_mask[2, 0]) == -1 + + +class _States: + def __init__(self, temperature, top_k, top_p, seeds, vocab): + self.vocab_size = vocab + self.temperature = SimpleNamespace(gpu=temperature) + self.seeds = SimpleNamespace(gpu=seeds) + self.top_k = SimpleNamespace(gpu=top_k, np=top_k.cpu().numpy()) + self.top_p = SimpleNamespace(gpu=top_p, np=top_p.cpu().numpy()) + + def get_top_k_top_p(self, expanded_idx_mapping, idx_mapping_np): + rows = expanded_idx_mapping.long() + return self.top_k.gpu[rows], self.top_p.gpu[rows] + + +def _ras_sampler(tokens, prompt_len, total_len, temperature, top_k, top_p, seeds, vocab): + return SimpleNamespace( + logit_bias_state=SimpleNamespace(apply_logit_bias=lambda *args: None), + sampling_states=_States(temperature, top_k, top_p, seeds, vocab), + req_states=SimpleNamespace( + all_token_ids=SimpleNamespace(gpu=tokens), + prompt_len=SimpleNamespace(gpu=prompt_len), + total_len=SimpleNamespace(gpu=total_len), + ), + ) + + +def _sample(sampler, logits, pos, fused): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _ras_mrv2_sample + + rows = logits.shape[0] + idx = torch.arange(rows, dtype=torch.int32, device=logits.device) + sampled, _ = _ras_mrv2_sample( + sampler, + logits, + idx, + idx, + np.arange(rows), + pos, + None, + None, + not fused, # logprobs keep the tensor path + default_top_p=0.8, + default_top_k=25, + win_size=10, + tau_r=0.1, + eps=1e-5, + ) + return sampled + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("fused", [True, False]) +def test_ras_mrv2_sample_rejects_recent_repeats_on_device(fused): + dev, vocab, rows = "cuda", 64, 4 + logits = torch.full((rows, vocab), -4.0, device=dev) + logits[:, 7] = 8.0 # token 7 dominates every row + tokens = torch.zeros(rows, 32, dtype=torch.int32, device=dev) + # Row 0: token 7 generated recently -> rejected. Row 1: 7 only in the prompt. + # Row 2: no output yet. Row 3: greedy, never rejected. + tokens[0, :6] = torch.tensor([1, 2, 3, 7, 4, 5]) + tokens[1, :3] = torch.tensor([7, 7, 9]) + tokens[3, :4] = torch.tensor([1, 7, 7, 7]) + prompt_len = torch.tensor([3, 2, 5, 1], dtype=torch.int32, device=dev) + total_len = torch.tensor([6, 3, 5, 4], dtype=torch.int32, device=dev) + sampler = _ras_sampler( + tokens, + prompt_len, + total_len, + torch.tensor([1.0, 1.0, 1.0, 0.0], device=dev), + torch.full((rows,), 1, dtype=torch.int32, device=dev), + torch.full((rows,), 0.8, device=dev), + torch.arange(rows, dtype=torch.int64, device=dev), + vocab, + ) + sampled = _sample(sampler, logits, torch.full((rows,), 11, dtype=torch.int64, device=dev), fused).tolist() + assert sampled[0] != 7 and 0 <= sampled[0] < vocab + assert sampled[1:] == [7, 7, 7] + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("fused", [True, False]) +def test_ras_mrv2_sample_draws_from_top_p_capped_by_top_k(fused): + dev, vocab, rows = "cuda", 32, 4096 + probs = torch.tensor([0.4, 0.3, 0.2, 0.1] + [0.0] * (vocab - 4), device=dev) + 1e-9 + logits = probs.log().expand(rows, vocab).contiguous() + no_history = torch.zeros(rows, dtype=torch.int32, device=dev) + # Row halves: top-p 0.6 keeps {0, 1}; top-p 0.95 with top-k 3 keeps {0, 1, 2}. + top_p = torch.where(torch.arange(rows, device=dev) < rows // 2, 0.6, 0.95).float() + sampler = _ras_sampler( + torch.zeros(rows, 4, dtype=torch.int32, device=dev), + no_history, + no_history, + torch.ones(rows, device=dev), + torch.full((rows,), 3, dtype=torch.int32, device=dev), + top_p, + torch.arange(rows, dtype=torch.int64, device=dev), + vocab, + ) + sampled = _sample(sampler, logits, torch.zeros(rows, dtype=torch.int64, device=dev), fused) + low, high = sampled[: rows // 2], sampled[rows // 2 :] + assert set(low.tolist()) == {0, 1} and set(high.tolist()) == {0, 1, 2} + # Kept tokens are renormalized: 0.4 / 0.7 of the top-p 0.6 rows pick token 0. + assert abs((low == 0).float().mean().item() - 4 / 7) < 0.04 + assert abs((high == 0).float().mean().item() - 4 / 9) < 0.04 + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_fused_ras_replacement_follows_the_remaining_distribution(): + dev, vocab, rows = "cuda", 16, 8192 + probs = torch.tensor([0.5, 0.25, 0.125, 0.125] + [0.0] * (vocab - 4), device=dev) + 1e-12 + logits = probs.log().expand(rows, vocab).contiguous() + tokens = torch.zeros(rows, 4, dtype=torch.int32, device=dev) # every row just generated token 0 + sampler = _ras_sampler( + tokens, + torch.zeros(rows, dtype=torch.int32, device=dev), + torch.ones(rows, dtype=torch.int32, device=dev), + torch.ones(rows, device=dev), + torch.full((rows,), 1, dtype=torch.int32, device=dev), # nucleus is always token 0 + torch.full((rows,), 0.8, device=dev), + torch.arange(rows, dtype=torch.int64, device=dev), + vocab, + ) + sampled = _sample(sampler, logits, torch.zeros(rows, dtype=torch.int64, device=dev), True) + freq = torch.bincount(sampled, minlength=vocab).float() / rows + assert freq[0] == 0 + for token, want in ((1, 0.5), (2, 0.25), (3, 0.25)): + assert abs(freq[token].item() - want) < 0.03 diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_packed_conv.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_packed_conv.py new file mode 100644 index 000000000..c9c180863 --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_packed_conv.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The packed causal conv kernel must match the padded cuDNN position embedding.""" + +import pytest +import torch +import torch.nn.functional as F + +from tests.helpers.mark import hardware_test +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="requires CUDA")] + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("lengths", [[1], [31, 30, 32], [500, 7, 129, 64]]) +@torch.inference_mode() +def test_packed_causal_conv_mish_matches_padded_cudnn(lengths): + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.packed_conv import ( + pack_conv_weight, + packed_causal_conv_mish, + ) + from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.packed_dit import ( + gather_rows, + pack_rows, + scatter_rows, + ) + + torch.manual_seed(0) + channels, groups, taps = 256, 16, 31 + conv = torch.nn.Conv1d(channels, channels, taps, groups=groups).cuda().to(torch.bfloat16) + rows = pack_rows(lengths, torch.device("cuda")) + x = torch.randn(rows.total, channels, device="cuda", dtype=torch.bfloat16) + + padded = scatter_rows(x.unsqueeze(0), rows, rows.width).permute(0, 2, 1) + reference = F.mish(conv(F.pad(padded, (taps - 1, 0)))) + reference = gather_rows(reference.permute(0, 2, 1), rows)[0] + output = packed_causal_conv_mish(x, pack_conv_weight(conv), conv.bias, rows.positions) + + # Rows never mix and match the left-zero-padded convolution up to bf16 rounding. + torch.testing.assert_close(output.float(), reference.float(), atol=0.04, rtol=0.02) diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_reference_cache.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_reference_cache.py new file mode 100644 index 000000000..12952e96c --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_reference_cache.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def processor(monkeypatch): + from vllm_omni.model_executor.models.cosyvoice3 import cosyvoice3 as module + from vllm_omni.utils.speaker_cache import SpeakerEmbeddingCache + + calls = [] + config = SimpleNamespace(allowed_special=[], sample_rate=24000) + ctx = SimpleNamespace(get_hf_config=lambda: config, model_config=SimpleNamespace(model="model-a")) + obj = SimpleNamespace( + info=SimpleNamespace(ctx=ctx), + tokenizer=None, + feat_extractor=None, + campplus_trt=None, + campplus_session=None, + _speaker_cache=SpeakerEmbeddingCache(max_bytes=1024), + _ensure_cached_runtime_components=lambda *args: None, + ) + + def speech(audio, device): + calls.append(audio) + return torch.tensor([[1, 2]], dtype=torch.int32), torch.tensor([2], dtype=torch.int32) + + obj._extract_speech_token_via_s3 = speech + monkeypatch.setattr( + module, "extract_text_token", lambda text, *args: (torch.tensor([[ord(c) for c in text]]), len(text)) + ) + monkeypatch.setattr(module, "extract_speech_feat", lambda *args: (torch.ones(1, 4, 2), torch.tensor([4]))) + monkeypatch.setattr(module, "extract_spk_embedding", lambda *args: torch.ones(1, 2)) + + def run(audio, text="target", reference="reference"): + return module.CosyVoice3MultiModalProcessor._call_hf_processor( + obj, text, {"audio": audio}, {"prompt_text": reference}, {} + ) + + return obj, run, calls + + +def test_reference_cache_reuses_audio_but_not_text_and_owns_tensors(processor): + obj, run, calls = processor + audio = (np.arange(8, dtype=np.float32), 16000) + first = run(audio) + first["speech_feat"].zero_() + first["speech_token"].zero_() + second = run((audio[0].copy(), 16000), text="different", reference="changed") + assert len(calls) == 1 + assert torch.all(second["speech_feat"] == 1) + assert second["speech_token"].tolist() == [[1, 2]] + assert not torch.equal(first["input_ids"], second["input_ids"]) + second["embedding"].zero_() + assert torch.all(run(audio)["embedding"] == 1) + assert obj._speaker_cache.stats()["hits"] == 2 + + +@pytest.mark.parametrize("change", ["samples", "rate", "dtype", "model"]) +def test_reference_cache_distinguishes_inputs(processor, change): + obj, run, calls = processor + wave = np.arange(8, dtype=np.float32) + run((wave, 16000)) + rate = 16000 + if change == "samples": + wave[0] = 20 + elif change == "rate": + rate = 24000 + elif change == "dtype": + wave = wave.astype(np.float64) + else: + obj.info.ctx.model_config.model = "model-b" + run((wave, rate)) + assert len(calls) == 2 diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_speech_token_batcher.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_speech_token_batcher.py new file mode 100644 index 000000000..a43951593 --- /dev/null +++ b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_speech_token_batcher.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Batched reference speech-token extraction returns each caller its own tokens.""" + +import threading +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _s3(calls): + def padding(mels): + lens = torch.tensor([m.shape[-1] for m in mels]) + out = torch.zeros(len(mels), 1, int(lens.max())) + for i, m in enumerate(mels): + out[i, :, : m.shape[-1]] = m + return out, lens + + def quantize(mels, lens): + calls.append(int(mels.shape[0])) + if bool((mels < 0).any()): + raise RuntimeError("bad mel") + # Token t of a row is its first mel value plus t; padding shows up as garbage. + codes = mels[:, 0, :1].long() + torch.arange(mels.shape[-1]).unsqueeze(0) + return codes, lens + + return SimpleNamespace(padding=padding), SimpleNamespace(quantize=quantize) + + +def test_speech_token_batcher_returns_each_rows_tokens_under_concurrency(): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _SpeechTokenBatcher + + calls: list[int] = [] + s3, model = _s3(calls) + batcher = _SpeechTokenBatcher(model, s3, "cpu", max_batch=8) + results: dict[int, torch.Tensor] = {} + start = threading.Barrier(24) + + def worker(i): + start.wait() + results[i] = batcher.tokens(torch.full((1, 3 + i % 5), float(100 * i))) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(24)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + for i in range(24): + assert results[i].tolist() == [100 * i + t for t in range(3 + i % 5)] + assert sum(calls) == 24 and max(calls) <= 8 + + +def test_speech_token_batcher_propagates_errors_and_keeps_serving(): + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _SpeechTokenBatcher + + s3, model = _s3([]) + batcher = _SpeechTokenBatcher(model, s3, "cpu") + with pytest.raises(RuntimeError, match="bad mel"): + batcher.tokens(torch.full((1, 4), -1.0)) + assert batcher.tokens(torch.full((1, 2), 7.0)).tolist() == [7, 8] diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 7a25034a1..aa94a6879 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -149,6 +149,13 @@ _MODEL_PROMOTE = ( "COSMOS3_SOUND_TOKENIZER_PATH", "COSYVOICE3_BATCH_FLOW", "COSYVOICE3_ESTIMATOR_ONNX", + "COSYVOICE3_FLOW_CUDAGRAPH", + "COSYVOICE3_FLOW_GRAPH", + "COSYVOICE3_FLOW_GRAPH_MAX_TOTAL", + "COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS", + "COSYVOICE3_PACKED_STREAMING", + "COSYVOICE3_REFERENCE_PREFETCH", + "COSYVOICE3_REFERENCE_WORKERS", "COSYVOICE3_TRT", "COSYVOICE3_TRT_CACHE", "FASTVIDEO_VSA_SM100A", @@ -200,6 +207,7 @@ _MODEL_EXTERNAL: tuple[str, ...] = () _MODEL_INTERNALIZE = ( "COSYVOICE3_BATCH_FLOW_DEBUG", "COSYVOICE3_F0_ON_CPU", + "COSYVOICE3_REFERENCE_CACHE_DEBUG", "DZ_PHASE_TIMING", "MAGI2_ALLOW_UNSUPPORTED_TOPOLOGY", "MAGI2_ROUTER_BIAS_SOURCE", diff --git a/vllm_omni/deploy/cosyvoice3_full_response.yaml b/vllm_omni/deploy/cosyvoice3_full_response.yaml new file mode 100644 index 000000000..bb0fd216d --- /dev/null +++ b/vllm_omni/deploy/cosyvoice3_full_response.yaml @@ -0,0 +1,62 @@ +# Opt-in single-Hopper full-response throughput profile (not streaming TTFP). +# Stage 1 selects the packed backend through its environment below. +# The 4096-frame internal boundary exceeds this profile's 2048-token output cap. +async_chunk: true +connectors: + connector_of_shared_memory: + extra: + codec_chunk_frames: 4096 + codec_left_context_frames: 25 + codec_streaming: true + codec_vocab_size: 6561 + connector_get_max_wait: 300 + connector_get_max_wait_first_chunk: 3000 + connector_get_sleep_s: 0.01 + name: SharedMemoryConnector +cuda_mps: true +pipeline: cosyvoice3 +stages: +- default_sampling_params: + max_tokens: 2048 + repetition_penalty: 1.0001 + temperature: 1.0 + top_k: 25 + top_p: 0.8 + devices: '0' + disable_hybrid_kv_cache_manager: true + # Prompt embedding rearrangement requires the complete prefill block. + enable_chunked_prefill: false + enable_prefix_caching: false + enforce_eager: false + gpu_memory_utilization: 0.4 + max_num_batched_tokens: 32768 + max_num_seqs: 64 + mm_processor_cache_gb: 0 + # The talker runs on Model Runner V2 (both RAS and standard sampling); + # set model_runner: v1 to fall back. The codec stage stays on V1. + model_runner: v2 + output_connectors: + to_stage_1: connector_of_shared_memory + skip_mm_profiling: true + stage_id: 0 + trust_remote_code: true +- default_sampling_params: + max_tokens: 2048 + devices: '0' + disable_hybrid_kv_cache_manager: true + enable_prefix_caching: false + enforce_eager: true + gpu_memory_utilization: 0.25 + input_connectors: + from_stage_0: connector_of_shared_memory + max_model_len: 32768 + max_num_batched_tokens: 32768 + max_num_seqs: 8 + skip_mm_profiling: true + stage_id: 1 + env: + COSYVOICE3_TRT: "0" + COSYVOICE3_BATCH_FLOW: "1" + COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS: "1" + COSYVOICE3_PACKED_STREAMING: "0" + trust_remote_code: true diff --git a/vllm_omni/deploy/cosyvoice3_full_response_high_concurrency.yaml b/vllm_omni/deploy/cosyvoice3_full_response_high_concurrency.yaml new file mode 100644 index 000000000..e0c1e9f2e --- /dev/null +++ b/vllm_omni/deploy/cosyvoice3_full_response_high_concurrency.yaml @@ -0,0 +1,19 @@ +# Opt-in single-Hopper C64 throughput profile: one AR and two codec replicas. +# Inherit sampling, connector, and runtime settings from the full-response profile. +# Each stage has aggregate capacity 64; this is not a streaming TTFP profile. +# One Model Runner V2 talker serves all 64 requests: a decode step costs about +# the same at 32 and 64 rows, so a second talker only adds processes. The +# codec is the throughput limit; two replicas overlap its launch-bound steps. +base_config: cosyvoice3_full_response.yaml + +stages: + - stage_id: 0 + devices: "0" + num_replicas: 1 + gpu_memory_utilization: 0.35 + max_num_seqs: 64 + - stage_id: 1 + devices: "0,0" + num_replicas: 2 + gpu_memory_utilization: 0.2 + max_num_seqs: 32 diff --git a/vllm_omni/deploy/cosyvoice3_full_response_high_concurrency_standard.yaml b/vllm_omni/deploy/cosyvoice3_full_response_high_concurrency_standard.yaml new file mode 100644 index 000000000..80b08c3d7 --- /dev/null +++ b/vllm_omni/deploy/cosyvoice3_full_response_high_concurrency_standard.yaml @@ -0,0 +1,15 @@ +# Standard (non-RAS) sampling variant of the C64 throughput profile, used for +# SGLang-parity validation. Launch with +# --hf-overrides '{"cosyvoice3_sampling_mode": "standard"}'. +# Standard sampling needs a real repetition penalty (RAS handles repetition +# itself); these defaults match the SGLang reference. +base_config: cosyvoice3_full_response_high_concurrency.yaml + +stages: + - stage_id: 0 + default_sampling_params: + max_tokens: 2048 + repetition_penalty: 1.21 + temperature: 0.7 + top_k: 20 + top_p: 0.8 diff --git a/vllm_omni/deploy/cosyvoice3_packed_streaming.yaml b/vllm_omni/deploy/cosyvoice3_packed_streaming.yaml new file mode 100644 index 000000000..1a8449ca8 --- /dev/null +++ b/vllm_omni/deploy/cosyvoice3_packed_streaming.yaml @@ -0,0 +1,18 @@ +# Opt-in single-Hopper streaming profile: AR64 / codec16, Flow10. +# Request stream=true to measure first audible audio; non-stream latency is not TTFP. +# Keeps each request's HiFT mel/phase state on GPU and uses mainline GPU FP32 F0. +# Shared serving settings come from the full-response profile; only chunk +# boundaries, codec admission capacity, and stage environment differ. +base_config: cosyvoice3_full_response.yaml +connectors: + connector_of_shared_memory: + extra: + codec_chunk_frames: 25 +stages: + - stage_id: 1 + env: + COSYVOICE3_TRT: "0" + COSYVOICE3_BATCH_FLOW: "1" + COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS: "0" + COSYVOICE3_PACKED_STREAMING: "1" + max_num_seqs: 16 diff --git a/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py b/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py index 1196c5c70..34287f583 100644 --- a/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py +++ b/vllm_omni/entrypoints/openai/tts_adapters/cosyvoice3.py @@ -2,7 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """CosyVoice3 TTS serving adapter.""" +import asyncio import os +from concurrent.futures import ThreadPoolExecutor from typing import TYPE_CHECKING, Any import numpy as np @@ -10,6 +12,7 @@ from vllm.logger import init_logger from vllm_omni.entrypoints.openai.tts_adapters import register_tts_adapter from vllm_omni.entrypoints.openai.tts_adapters.base import ARTTSAdapter, PreparedRequest +from vllm_omni.model_executor.models.cosyvoice3.runtime import cosyvoice3_standard_sampling if TYPE_CHECKING: from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest @@ -34,13 +37,49 @@ class CosyVoice3Adapter(ARTTSAdapter): # Rebuilding it would repeat snapshot resolution and add ~100 ms to the # TTFP-critical sampling override on every request. self._tokenizer = None + # Cold references are resampled and conditioned here, in parallel + # workers, instead of serially inside the input processor. + self._reference_executor: ThreadPoolExecutor | None = None + self._reference_resampler = None + + async def _prefetch_reference(self, wav: np.ndarray, sr: int) -> tuple[np.ndarray, int]: + """Resample to the processor rate and warm its content-addressed cache. + + Uses the processor's own resampler, so the processor receives exactly + these samples (its resampling becomes a no-op) and its cache lookup for + the same bytes hits instead of recomputing on the serving event loop. + """ + from vllm.multimodal.parse import MultiModalDataParser + + from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import prefetch_reference_conditioning + + server = self.ctx.server + hf_cfg = server.model_config.hf_config + target_sr = getattr(hf_cfg, "target_sr", None) + if target_sr is None: + return wav, sr + if self._reference_executor is None: + workers = int(os.environ.get("COSYVOICE3_REFERENCE_WORKERS", "16")) + self._reference_executor = ThreadPoolExecutor(max_workers=workers, thread_name_prefix="cosyvoice3-ref") + self._reference_resampler = MultiModalDataParser(target_sr=target_sr).audio_resampler + model_dir = server.model_config.model + resampler = self._reference_resampler + assert resampler is not None + + def prepare() -> np.ndarray: + audio = resampler.resample(wav, orig_sr=sr) + prefetch_reference_conditioning(model_dir, hf_cfg, (audio, target_sr)) + return audio + + audio = await asyncio.get_running_loop().run_in_executor(self._reference_executor, prepare) + return audio, int(target_sr) async def _build_prompt( self, request: "OpenAICreateSpeechRequest", *, has_inline_ref_audio: bool = False ) -> dict[str, Any]: """Build the multimodal CosyVoice3 voice-cloning prompt.""" server = self.ctx.server - wav_samples, sr, _ = await server._resolve_ref_audio(request.ref_audio) + wav_samples, sr, _ = await server._resolve_ref_audio_array(request.ref_audio) ref_text = request.ref_text or "" if _PROMPT_DELIMITER not in ref_text: ref_text = f"{_PROMPT_PREFIX}{ref_text}" @@ -50,9 +89,13 @@ class CosyVoice3Adapter(ARTTSAdapter): if voice_lower in server.uploaded_speakers and not has_inline_ref_audio: mm_kwargs["voice_name"] = voice_lower mm_kwargs["voice_created_at"] = server._voice_created_at(voice_lower) + wav = np.array(wav_samples, dtype=np.float32, copy=True) + if "voice_name" not in mm_kwargs and os.environ.get("COSYVOICE3_REFERENCE_PREFETCH", "1") != "0": + wav, sr = await self._prefetch_reference(wav, sr) + mm_kwargs["sample_rate"] = sr return { "prompt": request.input, - "multi_modal_data": {"audio": (np.asarray(wav_samples, dtype=np.float32), sr)}, + "multi_modal_data": {"audio": (wav, sr)}, "mm_processor_kwargs": mm_kwargs, } @@ -110,6 +153,11 @@ class CosyVoice3Adapter(ARTTSAdapter): server = self.ctx.server sampling_params_list = copy.deepcopy(sampling_params_list) hf_cfg = server.model_config.hf_config + if cosyvoice3_standard_sampling(hf_cfg): + first_control = int(hf_cfg.llm["speech_token_size"]) + controls = list(range(first_control, first_control + 200)) + sampling_params_list[0].stop_token_ids = controls + sampling_params_list[0].all_stop_token_ids.update(controls) # Build the Qwen tokenizer once per process (resolving the model dir via # snapshot_download at most once) and reuse it across requests. tokenizer = self._tokenizer diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py index f3ac19d94..ed81d618d 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py @@ -14,7 +14,9 @@ import torch import torch.nn as nn import torch.nn.functional as F from torch.nn import Conv1d, ConvTranspose1d -from torch.nn.utils import remove_weight_norm +from torch.nn.utils import parametrize +from torch.nn.utils import remove_weight_norm as remove_legacy_weight_norm +from torch.nn.utils.weight_norm import WeightNorm try: from torch.nn.utils.parametrizations import weight_norm @@ -43,6 +45,44 @@ def get_padding(kernel_size, dilation=1): return int((kernel_size * dilation - dilation) / 2) +def _normalize_folded_weight(module: nn.Module, name: str = "weight") -> None: + """Keep ordinary Parameters; replace plain or inference tensors. + + In inference_mode, the legacy API can leave an inference Parameter, + while the parametrization API can leave a plain inference tensor. + Rebuild those outside inference_mode so they remain usable afterward. + """ + weight = getattr(module, name) + if isinstance(weight, nn.Parameter) and not weight.is_inference(): + return + with torch.inference_mode(False): + frozen = nn.Parameter(weight.detach().clone(), requires_grad=False) + delattr(module, name) + module.register_parameter(name, frozen) + + +def _fold_weight_norm(module: nn.Module) -> int: + """Fold one weight-normalized layer into a plain ``weight``, in place. + + Supports both APIs: the modern ``torch.nn.utils.parametrizations`` API + (what this module is built with when torch provides it) and the legacy + ``WeightNorm`` forward-pre-hook. Returns 1 when a weight norm was + folded and 0 when there was nothing to fold (e.g. on a second call). + """ + if parametrize.is_parametrized(module, "weight"): + # Removes every parametrization registered on "weight"; only + # weight_norm is applied to the layers in this module. + parametrize.remove_parametrizations(module, "weight", leave_parametrized=True) + _normalize_folded_weight(module) + return 1 + for hook in list(module._forward_pre_hooks.values()): + if isinstance(hook, WeightNorm): + remove_legacy_weight_norm(module, name=hook.name) + _normalize_folded_weight(module, hook.name) + return 1 + return 0 + + class ResBlock(torch.nn.Module): """Residual block module in HiFiGAN/BigVGAN.""" @@ -94,10 +134,13 @@ class ResBlock(torch.nn.Module): x = xt + x return x - def remove_weight_norm(self): + def remove_weight_norm(self) -> int: + """Fold this block's conv weight norms; returns the folded count.""" + folded = 0 for idx in range(len(self.convs1)): - remove_weight_norm(self.convs1[idx]) - remove_weight_norm(self.convs2[idx]) + folded += _fold_weight_norm(self.convs1[idx]) + folded += _fold_weight_norm(self.convs2[idx]) + return folded def _carry_phase_at_boundary( @@ -577,18 +620,24 @@ class HiFTGenerator(nn.Module): ) self.f0_predictor = f0_predictor - def remove_weight_norm(self): + def remove_weight_norm(self) -> int: + """Fold the generator's frozen weight norms into plain weights. + + Returns how many convolutions were folded. ``source_downs`` and + ``m_source`` carry no weight norm. ``f0_predictor`` is excluded: + causal inference runs it on CPU for precision, and folding it is + tracked separately in RFC #6870 (C5). + """ + folded = 0 for layer in self.ups: - remove_weight_norm(layer) + folded += _fold_weight_norm(layer) for block in self.resblocks: - block.remove_weight_norm() - remove_weight_norm(self.conv_pre) - remove_weight_norm(self.conv_post) - self.m_source.remove_weight_norm() - for layer in self.source_downs: - remove_weight_norm(layer) + folded += block.remove_weight_norm() + folded += _fold_weight_norm(self.conv_pre) + folded += _fold_weight_norm(self.conv_post) for block in self.source_resblocks: - block.remove_weight_norm() + folded += block.remove_weight_norm() + return folded def _stft(self, x): if x.device.type == "npu": @@ -1101,6 +1150,14 @@ class CausalConvRNNF0Predictor(nn.Module): if isinstance(layer, CausalConv1d) and layer.causal_type == "left" ) + def remove_weight_norm(self) -> int: + """Fold the five F0 convolutions after loading on the execution device. + + Repeated calls are no-ops. Load original checkpoints into a fresh + predictor, before folding changes the state-dict keys. + """ + return sum(_fold_weight_norm(layer) for layer in self.condnet) + def forward(self, x: torch.Tensor, finalize: bool = True) -> torch.Tensor: if finalize is True: x = self.condnet[0](x) diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_conv.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_conv.py new file mode 100644 index 000000000..f58d43b7b --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_conv.py @@ -0,0 +1,86 @@ +# ruff: noqa: N803 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Grouped causal Conv1d + Mish directly on packed DiT rows. + +The DiT position embedding is two grouped (16 groups), 31-tap left-padded +convolutions. cuDNN runs them as one small kernel per group and needs a padded +``(rows, width)`` scatter, transposes and a gather around every call. This +implicit-GEMM kernel reads the packed ``[tokens, channels]`` sequence directly: +taps before a row's first frame are masked, which is exactly the left zero +padding, so rows never mix and no padded layout is materialized. +""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _packed_causal_conv_mish_kernel( + x_ptr, + w_ptr, + bias_ptr, + pos_ptr, + out_ptr, + total, + channels, + TAPS: tl.constexpr, + GROUP: tl.constexpr, + BLOCK_M: tl.constexpr, +): + pid_m = tl.program_id(0) + group = tl.program_id(1) + rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + row_mask = rows < total + positions = tl.load(pos_ptr + rows, mask=row_mask, other=0) + cin = tl.arange(0, GROUP) + cout = tl.arange(0, GROUP) + acc = tl.zeros((BLOCK_M, GROUP), dtype=tl.float32) + for tap in range(TAPS): + shift = TAPS - 1 - tap + valid = row_mask & (positions >= shift) + src = (rows - shift).to(tl.int64) + a = tl.load( + x_ptr + src[:, None] * channels + group * GROUP + cin[None, :], + mask=valid[:, None], + other=0.0, + ) + b = tl.load(w_ptr + ((group * TAPS + tap) * GROUP + cin[:, None]) * GROUP + cout[None, :]) + acc = tl.dot(a, b, acc) + acc += tl.load(bias_ptr + group * GROUP + cout).to(tl.float32)[None, :] + # Match eager rounding: the convolution output is stored in the input + # dtype before Mish computes in float32 and rounds again. + conv = acc.to(out_ptr.dtype.element_ty).to(tl.float32) + softplus = tl.where(conv > 20.0, conv, tl.log(1.0 + tl.exp(conv))) + e2 = tl.exp(-2.0 * softplus) + mish = conv * (1.0 - e2) / (1.0 + e2) + tl.store( + out_ptr + rows.to(tl.int64)[:, None] * channels + group * GROUP + cout[None, :], + mish.to(out_ptr.dtype.element_ty), + mask=row_mask[:, None], + ) + + +def pack_conv_weight(conv: torch.nn.Conv1d) -> torch.Tensor: + """[Cout, Cin/G, K] grouped weight -> [G, K, Cin/G, Cout/G] contiguous.""" + cout, cin_g, taps = conv.weight.shape + groups = conv.groups + return conv.weight.detach().view(groups, cout // groups, cin_g, taps).permute(0, 3, 2, 1).contiguous() + + +def packed_causal_conv_mish( + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + positions: torch.Tensor, +) -> torch.Tensor: + """x: [T, C] packed rows; positions: [T] int frame index within each row.""" + total, channels = x.shape + groups, taps, group, _ = weight.shape + out = torch.empty_like(x) + if total: + block_m = 64 + _packed_causal_conv_mish_kernel[(triton.cdiv(total, block_m), groups)]( + x, weight, bias, positions, out, total, channels, TAPS=taps, GROUP=group, BLOCK_M=block_m + ) + return out diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py new file mode 100644 index 000000000..ff6535b87 --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py @@ -0,0 +1,832 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Opt-in packed Flow for full-response and streaming inference on Hopper GPUs. + +Adapted from SGLang-Omni packed_dit.py, commit +127f34b57446a3cb5588ec987da0771a3be675be (Apache-2.0): +https://github.com/sgl-project/sglang-omni +Full-context and chunk-causal attention use vLLM's FA3 binding. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass + +import torch +import torch._dynamo as dynamo +import torch.nn.functional as F +from vllm.vllm_flash_attn import flash_attn_varlen_func + +from .packed_conv import pack_conv_weight, packed_causal_conv_mish + +FA3_PAGE_SIZE = 1 + + +class RaggedRowAttention: + """FA3 page metadata isolating each request and each causal query chunk.""" + + def __init__( + self, + rows: PackedRows, + *, + heads: int, + head_dim: int, + chunk_size: int | None = None, + width: int | None = None, + ) -> None: + self.heads, self.head_dim = heads, head_dim + device = rows.row_ids.device + segment_rows, segment_ends, offsets = [], [], [0] + for row, length in enumerate(rows.lengths): + if length == 0 and width is not None: + # Bucketed layouts keep one (empty) segment per row slot so + # metadata shapes are identical across batches in a bucket. + segment_rows.append(row) + segment_ends.append(0) + offsets.append(offsets[-1]) + continue + span = chunk_size or length + for start in range(0, length, span): + end = min(start + span, length) + segment_rows.append(row) + segment_ends.append(end) + offsets.append(offsets[-1] + end - start) + self.cache_seqlens = torch.tensor(segment_ends, dtype=torch.int32, device=device) + self.cu_seqlens_q = torch.tensor(offsets, dtype=torch.int32, device=device) + self.max_seqlen_q = max(b - a for a, b in zip(offsets, offsets[1:])) + starts = rows.starts_host[segment_rows].to(device) + page = torch.arange(width or max(segment_ends), dtype=torch.int32, device=device) + self.page_table = torch.where(page[None] < self.cache_seqlens[:, None], starts[:, None] + page[None], 0) + + +class PackedDiT: + """Adapt a loaded DiT to compiled, variable-length inference with FA3.""" + + def __init__(self, dit: torch.nn.Module) -> None: + self.dit = dit + self.compiled_full_forward = None + self.graph_buckets = False + self.graph_runner: FlowGraphRunner | None = None + conv_pos = dit.input_embed.conv_pos_embed + group = conv_pos.conv1[0].in_channels // conv_pos.conv1[0].groups + # The packed kernel's tensor-core tiles need a power-of-two group of at + # least 16 channels in half precision; other shapes keep cuDNN. + self.conv_weights = ( + (pack_conv_weight(conv_pos.conv1[0]), pack_conv_weight(conv_pos.conv2[0])) + if group >= 16 + and group & (group - 1) == 0 + and conv_pos.conv1[0].weight.dtype in (torch.float16, torch.bfloat16) + else None + ) + self._modulations: dict = {} + + @torch.inference_mode() + def step_modulations(self, time_span: torch.Tensor, key: tuple) -> tuple[list[torch.Tensor], list[torch.Tensor]]: + """Per-step adaLN modulations, identical for every request and batch. + + They depend only on the Euler time, which follows a fixed schedule. + Replay the solver's own time recurrence and compute each step's + modulation exactly as the forward would (same M=1 shapes and autocast), + once per schedule, instead of 23 GEMVs and their activations per call. + """ + cached = self._modulations.get(key) + if cached is not None: + return cached + dit = self.dit + blocks, finals = [], [] + flow_time = torch.zeros(1, device=time_span.device, dtype=time_span.dtype) + t, dt = time_span[0], time_span[1] - time_span[0] + with torch.autocast("cuda", dtype=time_span.dtype): + for step in range(1, len(time_span)): + flow_time[:] = t + emb = dit.time_embed(flow_time) + blocks.append( + torch.cat([block.attn_norm.linear(block.attn_norm.silu(emb)) for block in dit.transformer_blocks]) + ) + finals.append(dit.norm_out.linear(dit.norm_out.silu(emb))) + t = t + dt + if step < len(time_span) - 1: + dt = time_span[step + 1] - t + cached = self._modulations[key] = (blocks, finals) + return cached + + def compile(self, dtype: torch.dtype) -> None: + if dtype not in (torch.float16, torch.bfloat16): + raise ValueError("Packed Flow requires FP16 or BF16") + import os + + options = {"emulate_precision_casts": True} + if os.environ.get("COSYVOICE3_FLOW_CUDAGRAPH") == "1": + # Inductor CUDA graphs replay one recording per concrete shape; + # full-context batches are padded to recurring buckets. + options["triton.cudagraphs"] = True + self.graph_buckets = True + self.compiled_full_forward = torch.compile( + self.forward_full, + backend="inductor", + dynamic=True, + fullgraph=True, + options=options, + ) + if ( + os.environ.get("COSYVOICE3_FLOW_GRAPH", "0") == "1" + and not self.graph_buckets + and self.conv_weights is not None + ): + self.graph_runner = FlowGraphRunner(self) + + def row_attention(self, rows: PackedRows, *, streaming: bool, width: int | None = None) -> RaggedRowAttention: + attn = self.dit.transformer_blocks[0].attn + result = RaggedRowAttention( + rows, + heads=attn.heads, + head_dim=attn.inner_dim // attn.heads, + chunk_size=self.dit.static_chunk_size if streaming else None, + width=width, + ) + mark_packed_compile_metadata(rows, result) + return result + + def forward_full( + self, + x: torch.Tensor, + mu: torch.Tensor, + spks: torch.Tensor, + cond: torch.Tensor, + t: torch.Tensor, + rows: PackedRows, + attention: RaggedRowAttention, + block_modulation: torch.Tensor | None = None, + final_modulation: torch.Tensor | None = None, + rope: tuple[torch.Tensor, torch.Tensor] | None = None, + ) -> torch.Tensor: + return forward_packed_tensor_geometry( + self, + x, + mu, + spks, + cond, + t, + rows, + attention, + max_seqlen_q=attention.page_table.shape[1], + block_modulation=block_modulation, + final_modulation=final_modulation, + rope=rope, + ) + + +@torch.library.custom_op( + "vllm_omni_cosyvoice3_perf::packed_fa3", + mutates_args=(), + device_types="cuda", +) +def packed_fa3( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_seqlens: torch.Tensor, + page_table: torch.Tensor, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, +) -> torch.Tensor: + """Alias-free FA3 boundary for the compiled PackedDiT path.""" + return flash_attn_varlen_func( + q, + k_cache, + v_cache, + max_seqlen_q=max_seqlen_q, + cu_seqlens_q=cu_seqlens_q, + max_seqlen_k=page_table.shape[1], + seqused_k=cache_seqlens, + block_table=page_table, + causal=False, + fa_version=3, + ) + + +@packed_fa3.register_fake +def fake_packed_fa3( + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + cache_seqlens: torch.Tensor, + page_table: torch.Tensor, + cu_seqlens_q: torch.Tensor, + max_seqlen_q: int, +) -> torch.Tensor: + return torch.empty_like(q) + + +@torch.library.custom_op( + "vllm_omni_cosyvoice3_perf::native_mish", + mutates_args=(), + device_types="cuda", +) +def native_mish(x: torch.Tensor) -> torch.Tensor: + """Preserve eager CUDA Mish arithmetic across the Inductor boundary.""" + return F.mish(x) + + +@native_mish.register_fake +def fake_native_mish(x: torch.Tensor) -> torch.Tensor: + return torch.empty_like(x) + + +@torch.library.custom_op( + "vllm_omni_cosyvoice3_perf::packed_conv_mish", + mutates_args=(), + device_types="cuda", +) +def packed_conv_mish( + x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, positions: torch.Tensor +) -> torch.Tensor: + """Grouped causal conv + Mish on packed rows (see packed_conv.py).""" + return packed_causal_conv_mish(x, weight, bias, positions) + + +@packed_conv_mish.register_fake +def fake_packed_conv_mish( + x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, positions: torch.Tensor +) -> torch.Tensor: + return torch.empty_like(x) + + +@torch.library.custom_op( + "vllm_omni_cosyvoice3_perf::native_layer_norm", + mutates_args=(), + device_types="cuda", +) +def native_layer_norm( + x: torch.Tensor, + normalized_size: int, + eps: float, +) -> torch.Tensor: + """Keep CUDA autocast's eager FP32 LayerNorm contract.""" + return F.layer_norm(x.float(), (normalized_size,), None, None, eps) + + +@native_layer_norm.register_fake +def fake_native_layer_norm( + x: torch.Tensor, + normalized_size: int, + eps: float, +) -> torch.Tensor: + return torch.empty_like(x, dtype=torch.float32) + + +@dataclass(frozen=True) +class PackedRows: + lengths: tuple[int, ...] + starts_host: torch.Tensor + row_ids: torch.Tensor + positions: torch.Tensor + + @property + def total(self) -> int: + return sum(self.lengths) + + @property + def width(self) -> int: + return max(self.lengths) + + +def pack_rows(lengths: Sequence[int], device: torch.device) -> PackedRows: + lengths = tuple(int(length) for length in lengths) + starts_host = F.pad(torch.tensor(lengths, dtype=torch.int64).cumsum(0), (1, 0)) + starts = starts_host.to(device) + total = int(starts_host[-1]) + row_ids = torch.repeat_interleave( + torch.arange(len(lengths), device=device), + torch.tensor(lengths, dtype=torch.int64, device=device), + output_size=total, + ) + positions = torch.arange(total, device=device) - starts[row_ids] + return PackedRows( + lengths=lengths, + starts_host=starts_host.to(torch.int32), + row_ids=row_ids, + positions=positions, + ) + + +def gather_rows(padded: torch.Tensor, rows: PackedRows) -> torch.Tensor: + """(rows, width, channels) -> (1, total, channels), each row's first + length frames in row order.""" + width = padded.shape[1] + flat = padded.reshape(padded.shape[0] * width, padded.shape[2]) + return flat[rows.row_ids * width + rows.positions].unsqueeze(0) + + +def scatter_rows(packed: torch.Tensor, rows: PackedRows, width: int) -> torch.Tensor: + """(1, total, channels) -> (rows, width, channels), zero past each row's + length.""" + channels = packed.shape[2] + flat = packed.new_zeros(len(rows.lengths) * width, channels) + flat[rows.row_ids * width + rows.positions] = packed[0] + return flat.view(len(rows.lengths), width, channels) + + +def mark_packed_compile_metadata(rows: PackedRows, attention: RaggedRowAttention) -> None: + dynamo.mark_dynamic(attention.page_table, (0, 1)) + dynamo.mark_dynamic(attention.cu_seqlens_q, 0) + dynamo.mark_dynamic(attention.cache_seqlens, 0) + dynamo.mark_dynamic(rows.starts_host, 0) + dynamo.mark_dynamic(rows.row_ids, 0) + dynamo.mark_dynamic(rows.positions, 0) + + +def gather_rows_tensor_geometry(padded: torch.Tensor, rows: PackedRows) -> torch.Tensor: + width = padded.shape[1] + flat = padded.reshape(padded.shape[0] * width, padded.shape[2]) + return flat[rows.row_ids * width + rows.positions].unsqueeze(0) + + +def scatter_rows_tensor_geometry(packed: torch.Tensor, rows: PackedRows, row_count: int, width: int) -> torch.Tensor: + channels = packed.shape[2] + flat = packed.new_zeros(row_count * width, channels) + flat[rows.row_ids * width + rows.positions] = packed[0] + return flat.view(row_count, width, channels) + + +def bucketed_width(width): + """Round a frame count up so convolution shapes recur across batches. + + cuDNN and cuFFT build a plan for every new shape; live batches have unique + lengths. Right padding is exact for the causal convolutions it feeds. + """ + step = 32 if width <= 512 else (64 if width <= 2048 else 128) + return (width + step - 1) // step * step + + +def conv_pos_embed_tensor_geometry( + estimator: PackedDiT, + h: torch.Tensor, + rows: PackedRows, + attention: RaggedRowAttention, +) -> torch.Tensor: + module = estimator.dit.input_embed.conv_pos_embed + if estimator.conv_weights is not None: + # Both causal convolutions run on the packed sequence: taps before a + # row start are masked (the left zero padding), so no padded layout, + # transposes or per-group cuDNN launches are needed. + first, second = estimator.conv_weights + embedded = packed_conv_mish(h[0], first, module.conv1[0].bias, rows.positions) + embedded = packed_conv_mish(embedded, second, module.conv2[0].bias, rows.positions) + return embedded.unsqueeze(0) + row_count = rows.starts_host.shape[0] - 1 + # Branch-free rounding keeps this symbolic under dynamic compilation; a + # size comparison here would add guards and recompile mid-serving. + width = (attention.page_table.shape[1] + 63) // 64 * 64 + padded = scatter_rows_tensor_geometry(h, rows, row_count, width) + embedded = padded.permute(0, 2, 1) + embedded = F.pad(embedded, (module.kernel_size - 1, 0, 0, 0)) + embedded = module.conv1[0](embedded) + embedded = native_mish(embedded) + embedded = F.pad(embedded, (module.kernel_size - 1, 0, 0, 0)) + embedded = module.conv2[0](embedded) + embedded = native_mish(embedded) + embedded = embedded.permute(0, 2, 1) + return gather_rows_tensor_geometry(embedded, rows) + + +def rope_tensor_geometry( + estimator: PackedDiT, + rows: PackedRows, + attention: RaggedRowAttention, +) -> tuple[torch.Tensor, torch.Tensor]: + width = attention.page_table.shape[1] + freqs, scale = estimator.dit.rotary_embed.forward_from_seq_len(width) + assert not isinstance(scale, torch.Tensor), "the DiT's RoPE has no xpos scale" + freqs = freqs[:, rows.positions] + return freqs.cos(), freqs.sin() + + +def ragged_attention_tensor_geometry( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention: RaggedRowAttention, + max_seqlen_q: int, +) -> torch.Tensor: + page_shape = (-1, FA3_PAGE_SIZE, attention.heads, attention.head_dim) + output = packed_fa3( + query[0].reshape(-1, attention.heads, attention.head_dim), + key[0].reshape(page_shape), + value[0].reshape(page_shape), + attention.cache_seqlens, + attention.page_table, + attention.cu_seqlens_q, + max_seqlen_q, + ) + return output.reshape(1, -1, attention.heads * attention.head_dim) + + +def attend_tensor_geometry( + attn: torch.nn.Module, + x: torch.Tensor, + rope: tuple[torch.Tensor, torch.Tensor], + attention: RaggedRowAttention, + max_seqlen_q: int, +) -> torch.Tensor: + # note (ratish): under autocast to_q, to_k and to_v would each cast the + # float32 norm output again. + x = x.to(attn.to_q.weight.dtype) + query = attn.to_q(x) + key = attn.to_k(x) + value = attn.to_v(x) + rotate_in_place(query, *rope) + rotate_in_place(key, *rope) + output = ragged_attention_tensor_geometry(query, key, value, attention, max_seqlen_q).to(query.dtype) + return attn.to_out[1](attn.to_out[0](output)) + + +_FUSED_LAYER_NORM = __import__("os").environ.get("COSYVOICE3_FUSED_LAYERNORM", "1") != "0" + + +def layer_norm_preserving_eager(layer_norm: torch.nn.LayerNorm, x: torch.Tensor) -> torch.Tensor: + normalized_size = int(layer_norm.normalized_shape[0]) + if _FUSED_LAYER_NORM: + # Same FP32 LayerNorm contract as autocast, but visible to Inductor so + # the upcast, normalization, modulation and downcast fuse into one + # kernel instead of a LayerNorm, two copies and elementwise passes. + return F.layer_norm(x.float(), (normalized_size,), None, None, float(layer_norm.eps)) + return native_layer_norm(x, normalized_size, float(layer_norm.eps)) + + +def attn_norm_preserving_eager( + block: torch.nn.Module, + h: torch.Tensor, + time_embedding: torch.Tensor | None, + modulation: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + if modulation is None: + modulation = block.attn_norm.linear(block.attn_norm.silu(time_embedding)) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(modulation, 6, dim=1) + normalized = layer_norm_preserving_eager(block.attn_norm.norm, h) + normalized = normalized * (1 + scale_msa[:, None]) + shift_msa[:, None] + return normalized, gate_msa, shift_mlp, scale_mlp, gate_mlp + + +def final_norm_preserving_eager( + norm_out: torch.nn.Module, + h: torch.Tensor, + time_embedding: torch.Tensor | None, + modulation: torch.Tensor | None = None, +) -> torch.Tensor: + if modulation is None: + modulation = norm_out.linear(norm_out.silu(time_embedding)) + scale, shift = torch.chunk(modulation, 2, dim=1) + normalized = layer_norm_preserving_eager(norm_out.norm, h) + return normalized * (1 + scale)[:, None, :] + shift[:, None, :] + + +def forward_packed_tensor_geometry( + estimator: PackedDiT, + x: torch.Tensor, + mu: torch.Tensor, + spks: torch.Tensor, + cond: torch.Tensor, + t: torch.Tensor, + rows: PackedRows, + attention: RaggedRowAttention, + *, + max_seqlen_q: int, + block_modulation: torch.Tensor | None = None, + final_modulation: torch.Tensor | None = None, + rope: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + dit = estimator.dit + t = dit.time_embed(t) if block_modulation is None else None + h = dit.input_embed.proj(torch.cat((x, cond, mu, spks), dim=-1)) + h = conv_pos_embed_tensor_geometry(estimator, h, rows, attention) + h + if rope is None: + rope = rope_tensor_geometry(estimator, rows, attention) + residual = h + for index, block in enumerate(dit.transformer_blocks): + modulation = None if block_modulation is None else block_modulation[index : index + 1] + norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = attn_norm_preserving_eager(block, h, t, modulation) + h = h + gate_msa.unsqueeze(1) * attend_tensor_geometry(block.attn, norm, rope, attention, max_seqlen_q) + ff_norm = layer_norm_preserving_eager(block.ff_norm, h) + ff_norm = ff_norm * (1 + scale_mlp[:, None]) + shift_mlp[:, None] + h = h + gate_mlp.unsqueeze(1) * block.ff(ff_norm) + if dit.long_skip_connection is not None: + h = dit.long_skip_connection(torch.cat((h, residual), dim=-1)) + h = final_norm_preserving_eager(dit.norm_out, h, t, final_modulation) + return dit.proj_out(h) + + +def rotate_in_place(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> None: + """x: (1, total, heads * head_dim). Interleaved RoPE in float32 on the + rotary dims, rounded back into x.""" + # note (ratish): the DiT rotates only the first rotary dims of the + # flattened heads, so the rest of x is never copied. + rotary = x[..., : cos.shape[-1]] + half = torch.stack((-rotary[..., 1::2], rotary[..., ::2]), dim=-1).flatten(-2) + rotary.copy_(rotary * cos + half * sin) + + +def solve_flow_euler_packed( + estimator: PackedDiT, + noise: torch.Tensor, + time_span: torch.Tensor, + mu: torch.Tensor, + spks: torch.Tensor, + cond: torch.Tensor, + rows: PackedRows, + *, + cfg_rate: float, + streaming: bool, + modulation_key: tuple | None = None, +) -> torch.Tensor: + """Euler steps over a packed sequence with classifier free guidance: the + conditional rows and their unconditional twins share one DiT call.""" + if estimator.graph_runner is not None and modulation_key is not None: + solved = estimator.graph_runner.run( + noise, time_span, mu, spks, cond, rows, cfg_rate=cfg_rate, streaming=streaming, key=modulation_key + ) + if solved is not None: + return solved + real_total = noise.shape[1] + width = None + if estimator.graph_buckets and not streaming: + rows, noise, mu, spks, cond, width = pad_to_graph_bucket(rows, noise, mu, spks, cond) + twin_rows = pack_rows(rows.lengths * 2, noise.device) + attention = estimator.row_attention(twin_rows, streaming=streaming, width=width) + mu_cfg = torch.cat((mu, torch.zeros_like(mu)), dim=1) + cond_cfg = torch.cat((cond, torch.zeros_like(cond)), dim=1) + spks_cfg = torch.cat((spks, torch.zeros_like(spks)), dim=0) + spks_cfg = spks_cfg[twin_rows.row_ids].unsqueeze(0) + flow_time = torch.zeros(1, device=noise.device, dtype=spks.dtype) + forward = estimator.compiled_full_forward + if forward is None: + raise ValueError("Packed Flow requires a compiled packed path") + modulations = None if modulation_key is None else estimator.step_modulations(time_span, modulation_key) + x = euler_steps( + forward, noise, time_span, mu_cfg, spks_cfg, cond_cfg, flow_time, twin_rows, attention, modulations, cfg_rate + ) + return x[:, :real_total].float() + + +def euler_steps( + forward, + noise: torch.Tensor, + time_span: torch.Tensor, + mu_cfg: torch.Tensor, + spks_cfg: torch.Tensor, + cond_cfg: torch.Tensor, + flow_time: torch.Tensor, + twin_rows: PackedRows, + attention: RaggedRowAttention, + modulations: tuple[list[torch.Tensor], list[torch.Tensor]] | None, + cfg_rate: float, +) -> torch.Tensor: + total = noise.shape[1] + x = noise + t, dt = time_span[0], time_span[1] - time_span[0] + for step in range(1, len(time_span)): + flow_time[:] = t + vector_field = forward( + torch.cat((x, x), dim=1), + mu_cfg, + spks_cfg, + cond_cfg, + flow_time, + twin_rows, + attention, + None if modulations is None else modulations[0][step - 1], + None if modulations is None else modulations[1][step - 1], + ) + conditional = vector_field[:, :total] + unconditional = vector_field[:, total:] + x = x + dt * ((1.0 + cfg_rate) * conditional - cfg_rate * unconditional) + t = t + dt + if step < len(time_span) - 1: + dt = time_span[step + 1] - t + return x + + +_GRAPH_ROW_BUCKETS = (1, 2, 4, 8, 16) +_GRAPH_TOTAL_STEP = 512 +_GRAPH_WIDTH_STEP = 256 + + +def pad_to_graph_bucket(rows, noise, mu, spks, cond): + """Pad a full-context batch so CUDA-graph shapes recur across batches. + + Rows are padded to a bucketed count with empty rows, and the packed length + to a bucketed total with one trailing padding row; attention isolates every + row, so padding never reaches real frames, which stay first in the packing. + """ + n = len(rows.lengths) + slots = next((b for b in _GRAPH_ROW_BUCKETS if b >= n), n) + total = rows.total + pad = -total % _GRAPH_TOTAL_STEP + width = max(max(rows.lengths), pad, 1) + width = -(-width // _GRAPH_WIDTH_STEP) * _GRAPH_WIDTH_STEP + lengths = tuple(rows.lengths) + (0,) * (slots - n) + (pad,) + padded_rows = pack_rows(lengths, noise.device) + extend = (0, 0, 0, pad) + spks = torch.cat((spks, spks.new_zeros((slots - n + 1, spks.shape[1])))) + return ( + padded_rows, + F.pad(noise, extend), + F.pad(mu, extend), + spks, + F.pad(cond, extend), + width, + ) + + +_FLOW_GRAPH_SLOTS = (1, 2, 4, 8) +_FLOW_GRAPH_MAX_TOTAL = int(__import__("os").environ.get("COSYVOICE3_FLOW_GRAPH_MAX_TOTAL", "16384")) +_FLOW_GRAPH_SEGMENT_STEP = 16 + + +def flow_graph_total(total: int) -> int: + """Packed-length bucket: at most ~12% padding, few recurring layouts.""" + step = 128 if total <= 1024 else (256 if total <= 4096 else 512) + return -(-total // step) * step + + +@dataclass +class _FlowGraph: + graph: torch.cuda.CUDAGraph + noise: torch.Tensor + mu: torch.Tensor + cond: torch.Tensor + spks: torch.Tensor + time_span: torch.Tensor + table: torch.Tensor + row_ids: torch.Tensor + positions: torch.Tensor + cache_seqlens: torch.Tensor + cu_seqlens_q: torch.Tensor + page_table: torch.Tensor + output: torch.Tensor | None = None + + +class FlowGraphRunner: + """Replay the whole CFG Euler solve as one CUDA graph per padded layout. + + Small codec batches are launch bound: every step walks Dynamo guards, the + Inductor wrapper and one FA3 wrapper per block, several hundred host + microseconds per step before the GPU sees work. The layout (row slots, + packed length, attention segments) is padded to recurring buckets exactly + like the bucketed layouts above: padding rows are isolated by attention + and by the positional masks of the packed convolution, so real frames only + see their own row. Metadata lives in static buffers that are rewritten + before each replay. + """ + + def __init__(self, estimator: PackedDiT) -> None: + self.estimator = estimator + self.graphs: dict[tuple, _FlowGraph] = {} + self.pool = None + attn = estimator.dit.transformer_blocks[0].attn + self.heads, self.head_dim = attn.heads, attn.inner_dim // attn.heads + + def _layout(self, lengths: tuple[int, ...], streaming: bool): + n, total = len(lengths), sum(lengths) + slots = next((b for b in _FLOW_GRAPH_SLOTS if b >= n), None) + padded_total = flow_graph_total(total) + if slots is None or padded_total > _FLOW_GRAPH_MAX_TOTAL: + return None + twin = (tuple(lengths) + (0,) * (slots - n) + (padded_total - total,)) * 2 + chunk = self.estimator.dit.static_chunk_size if streaming else None + seg_rows, seg_ends, offsets = [], [], [0] + for row, length in enumerate(twin): + if length == 0: + seg_rows.append(row) + seg_ends.append(0) + offsets.append(offsets[-1]) + continue + span = chunk or length + for start in range(0, length, span): + end = min(start + span, length) + seg_rows.append(row) + seg_ends.append(end) + offsets.append(offsets[-1] + end - start) + segments = len(seg_ends) + if streaming: + segments = -(-segments // _FLOW_GRAPH_SEGMENT_STEP) * _FLOW_GRAPH_SEGMENT_STEP + extra = segments - len(seg_ends) + seg_rows += [0] * extra + seg_ends += [0] * extra + offsets += [offsets[-1]] * extra + return (streaming, slots, padded_total, segments), twin, seg_rows, seg_ends, offsets + + def _allocate(self, key, noise, spks, time_span) -> _FlowGraph: + _, slots, padded_total, segments = key[:4] + device, dtype = noise.device, noise.dtype + channels = noise.shape[-1] + rows = 2 * (slots + 1) + return _FlowGraph( + graph=torch.cuda.CUDAGraph(), + noise=torch.zeros((1, padded_total, channels), device=device, dtype=dtype), + mu=torch.zeros((1, padded_total, channels), device=device, dtype=dtype), + cond=torch.zeros((1, padded_total, channels), device=device, dtype=dtype), + spks=torch.zeros((slots + 1, spks.shape[-1]), device=device, dtype=spks.dtype), + time_span=torch.zeros_like(time_span), + table=torch.zeros(rows + 3 * segments + 1, device=device, dtype=torch.int32), + row_ids=torch.zeros(2 * padded_total, device=device, dtype=torch.int64), + positions=torch.zeros(2 * padded_total, device=device, dtype=torch.int64), + cache_seqlens=torch.zeros(segments, device=device, dtype=torch.int32), + cu_seqlens_q=torch.zeros(segments + 1, device=device, dtype=torch.int32), + page_table=torch.zeros((segments, padded_total), device=device, dtype=torch.int32), + ) + + @staticmethod + def _load(state: _FlowGraph, twin, seg_rows, seg_ends, offsets, noise, mu, cond, spks, time_span) -> None: + total = noise.shape[1] + host = torch.tensor(list(twin) + seg_rows + seg_ends + offsets, dtype=torch.int32, pin_memory=True) + state.table.copy_(host, non_blocking=True) + rows, segments = len(twin), len(seg_ends) + lengths = state.table[:rows].long() + starts = F.pad(lengths.cumsum(0), (1, 0))[:-1] + width = state.row_ids.shape[0] + row_ids = torch.repeat_interleave(torch.arange(rows, device=noise.device), lengths, output_size=width) + state.row_ids.copy_(row_ids) + state.positions.copy_(torch.arange(width, device=noise.device) - starts[row_ids]) + seg_rows_t = state.table[rows : rows + segments] + ends = state.table[rows + segments : rows + 2 * segments] + state.cache_seqlens.copy_(ends) + state.cu_seqlens_q.copy_(state.table[rows + 2 * segments :]) + page = torch.arange(state.page_table.shape[1], device=noise.device, dtype=torch.int32) + state.page_table.copy_( + torch.where(page[None] < ends[:, None], starts[seg_rows_t.long()].int()[:, None] + page[None], 0) + ) + state.noise[:, :total].copy_(noise) + state.mu[:, :total].copy_(mu) + state.cond[:, :total].copy_(cond) + state.spks.zero_() + state.spks[: spks.shape[0]].copy_(spks) + state.time_span.copy_(time_span) + + def _solve(self, state: _FlowGraph, cfg_rate: float, key: tuple) -> torch.Tensor: + estimator = self.estimator + rows = PackedRows(lengths=(), starts_host=state.table, row_ids=state.row_ids, positions=state.positions) + attention = object.__new__(RaggedRowAttention) + attention.heads, attention.head_dim = self.heads, self.head_dim + attention.cache_seqlens, attention.cu_seqlens_q = state.cache_seqlens, state.cu_seqlens_q + attention.page_table = state.page_table + attention.max_seqlen_q = state.page_table.shape[1] + mark_packed_compile_metadata(rows, attention) + mu_cfg = torch.cat((state.mu, torch.zeros_like(state.mu)), dim=1) + cond_cfg = torch.cat((state.cond, torch.zeros_like(state.cond)), dim=1) + spks_cfg = torch.cat((state.spks, torch.zeros_like(state.spks)), dim=0)[state.row_ids].unsqueeze(0) + flow_time = torch.zeros(1, device=state.noise.device, dtype=state.spks.dtype) + modulations = estimator.step_modulations(state.time_span, key[4]) + return euler_steps( + estimator.compiled_full_forward, + state.noise, + state.time_span, + mu_cfg, + spks_cfg, + cond_cfg, + flow_time, + rows, + attention, + modulations, + cfg_rate, + ) + + def run( + self, + noise: torch.Tensor, + time_span: torch.Tensor, + mu: torch.Tensor, + spks: torch.Tensor, + cond: torch.Tensor, + rows: PackedRows, + *, + cfg_rate: float, + streaming: bool, + key: tuple, + ) -> torch.Tensor | None: + layout = self._layout(rows.lengths, streaming) + if layout is None or torch.cuda.is_current_stream_capturing(): + return None + graph_key, twin, seg_rows, seg_ends, offsets = layout + graph_key = graph_key + (key, float(cfg_rate)) + state = self.graphs.get(graph_key) + fresh = state is None + if fresh: + state = self._allocate(graph_key, noise, spks, time_span) + self._load(state, twin, seg_rows, seg_ends, offsets, noise, mu, cond, spks, time_span) + if fresh: + # Warm up (and compile new dynamic shapes) off the capture, then + # record the solve into the shared pool. + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + self._solve(state, cfg_rate, graph_key) + torch.cuda.current_stream().wait_stream(stream) + with torch.cuda.graph(state.graph, pool=self.pool): + state.output = self._solve(state, cfg_rate, graph_key) + if self.pool is None: + self.pool = state.graph.pool() + self.graphs[graph_key] = state + state.graph.replay() + return state.output[:, : noise.shape[1]].float() diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 0ddc6dafa..654675474 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -1,11 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import hashlib import os +from collections import OrderedDict from collections.abc import Iterable, Mapping, Sequence from dataclasses import replace from functools import partial from math import gcd from threading import Lock +from types import MethodType +from typing import Any import numpy as np import onnxruntime @@ -39,9 +43,12 @@ from vllm.v1.sample.sampler import Sampler from vllm_omni.data_entry_keys import EmbeddingsStruct, OmniPayloadStruct, to_dict, to_struct from vllm_omni.inputs.mm_processor import OmniMultiModalProcessor +from vllm_omni.model_executor.models.cosyvoice3.ras_sampler import MAX_FUSED_TOP_K, fused_ras_sample from vllm_omni.model_executor.models.cosyvoice3.runtime import ( cosyvoice3_batch_flow_debug, cosyvoice3_batch_flow_enabled, + cosyvoice3_packed_inference_enabled, + cosyvoice3_standard_sampling, ) from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer from vllm_omni.model_executor.models.cosyvoice3.utils import ( @@ -92,6 +99,267 @@ def _campplus_onnx_providers() -> list[str]: return ["CPUExecutionProvider"] +def _audio_conditioning(proc, audio, model_dir: str, config, mm_kwargs: Mapping[str, object]) -> dict: + """Reference-audio conditioning (speech tokens, mel, speaker embedding). + + Reuses bounded, process-local artifacts keyed by voice name or by the + waveform content, and returns independent tensors so downstream mutation + cannot poison a cache hit. + """ + device = "cpu" + voice_name = mm_kwargs.get("voice_name") + cache_key = None + if voice_name and isinstance(voice_name, str): + cache_key = proc._speaker_cache.make_cache_key( + voice_name, + model_type="cosyvoice3", + created_at=int(mm_kwargs.get("voice_created_at") or 0), + ) + else: + # Cache audio artifacts only: reference/target text is tokenized + # separately. Keep dtype and sample rate in the identity to avoid + # aliases between numerically distinct preprocessing inputs. + waveform, sample_rate = audio + waveform = np.ascontiguousarray(waveform) + digest = hashlib.sha256() + digest.update(str((int(sample_rate), waveform.shape, waveform.dtype.str)).encode()) + digest.update(memoryview(waveform).cast("B")) + backend = "trt" if proc.campplus_trt is not None else "onnx" + cache_key = proc._speaker_cache.make_cache_key( + digest.hexdigest(), model_type=f"cosyvoice3-reference:{model_dir}:{backend}" + ) + cached = proc._speaker_cache.get(cache_key) + if os.environ.get("COSYVOICE3_REFERENCE_CACHE_DEBUG") == "1": + _REFERENCE_CACHE_DEBUG[cached is not None] += 1 + if sum(_REFERENCE_CACHE_DEBUG.values()) % 128 == 0: + logger.info( + "CosyVoice3 reference cache lookups: hits=%d misses=%d model_dir=%s", + _REFERENCE_CACHE_DEBUG[True], + _REFERENCE_CACHE_DEBUG[False], + model_dir, + ) + if cached is None: + # Speech-token extraction via the S3Tokenizer PyTorch model on GPU + # (~30x faster than the bundled ``speech_tokenizer_v3.onnx`` CPU path). + speech_token, speech_token_len = proc._extract_speech_token_via_s3(audio, device) + speech_feat, speech_feat_len = extract_speech_feat(audio, proc.feat_extractor, device) + if config.sample_rate == 24000: + token_len = min(int(speech_feat.shape[1] / 2), speech_token.shape[1]) + speech_feat, speech_feat_len[:] = speech_feat[:, : 2 * token_len], 2 * token_len + speech_token, speech_token_len[:] = speech_token[:, :token_len], token_len + if proc.campplus_trt is not None: + embedding = extract_spk_embedding_trt(audio, proc.campplus_trt, device) + else: + embedding = extract_spk_embedding(audio, proc.campplus_session, device) + cached = { + "speech_feat": speech_feat.detach().cpu().clone(), + "speech_token": speech_token.detach().cpu().clone(), + "speech_token_len": speech_token_len.detach().cpu().clone(), + "embedding": embedding.detach().cpu().clone(), + } + proc._speaker_cache.put(cache_key, cached) + return {key: value.clone() for key, value in cached.items()} + + +_PROMPT_SCRATCH_TOKENS = 1 << 20 +_MRV2_CONDITIONING_ENTRIES = 4096 + + +def _generated_only_penalty_writes(state) -> None: + """``PenaltiesState.apply_staged_writes`` that leaves the prompt unpenalized.""" + if state._new_penalties_reqs: + from vllm.utils.torch_utils import async_tensor_h2d + from vllm.v1.worker.gpu.sample.penalties import bincount + + scratch = getattr(state, "_prompt_scratch_mask", None) + if scratch is None: + columns = max(state.prompt_bin_mask.shape[1], _PROMPT_SCRATCH_TOKENS // 32) + scratch = state._prompt_scratch_mask = torch.zeros( + state.prompt_bin_mask.shape[0], columns, dtype=torch.int32, device=state.device + ) + rows = async_tensor_h2d(state._new_penalties_reqs, dtype=torch.int32, device=state.device) + prefill_lens = state.req_states.prefill_len.np[state._new_penalties_reqs] + # Prompt bits land in the scratch mask; resumed output tokens (between + # prompt_len and prefill_len) still seed the output counts. + bincount( + rows, + state.req_states.all_token_ids.gpu, + state.req_states.prompt_len.gpu, + state.req_states.prefill_len.gpu, + scratch, + state.output_bin_counts, + int(prefill_lens.max()), + ) + state.prompt_bin_mask.index_fill_(0, rows.long(), 0) + state._new_penalties_reqs.clear() + state.repetition_penalty.copy_to_uva() + state.frequency_penalty.copy_to_uva() + state.presence_penalty.copy_to_uva() + + +def _ras_mrv2_sample( + sampler, + logits: torch.Tensor, + expanded_idx_mapping: torch.Tensor, + idx_mapping: torch.Tensor, + idx_mapping_np: np.ndarray, + pos: torch.Tensor, + input_ids: torch.Tensor, + expanded_local_pos: torch.Tensor, + return_logprobs: bool = False, + *, + default_top_p: float, + default_top_k: int, + win_size: int, + tau_r: float, + eps: float, +) -> tuple[torch.Tensor, torch.Tensor]: + """``Sampler.sample`` for MRv2: CosyVoice3 repetition-aware sampling on the GPU. + + Same distribution as ``CosyVoice3Model._ras_sample_batch`` (V1): top-p over + the full distribution capped at top-k; a draw repeating within the last + ``win_size`` generated tokens is replaced by a draw from the full remaining + distribution. History is read from the runner's token table and rejection + is resolved per row on the device, so no step waits on the host. Penalties + are not applied, as in V1 RAS. Draws use the request seed and position; the + replacement draw uses a disjoint noise stream. + """ + from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample + + logits = torch.empty_like(logits, dtype=torch.float32).copy_(logits) + sampler.logit_bias_state.apply_logit_bias(logits, expanded_idx_mapping, idx_mapping_np, pos) + states = sampler.sampling_states + temperature = states.temperature.gpu + req_states = sampler.req_states + top_k_np = states.top_k.np[idx_mapping_np] + use_top_k = bool((top_k_np != states.vocab_size).any()) + if logits.is_cuda and not return_logprobs and (top_k_np.max() if use_top_k else default_top_k) <= MAX_FUSED_TOP_K: + # One kernel per step instead of ~70 (see ras_sampler.py). + use_top_p = bool((states.top_p.np[idx_mapping_np] != 1.0).any()) + sampled = fused_ras_sample( + logits, + expanded_idx_mapping, + temperature, + states.top_k.gpu if use_top_k else None, + states.top_p.gpu if use_top_p else None, + states.seeds.gpu, + pos, + req_states.all_token_ids.gpu, + req_states.total_len.gpu, + req_states.prompt_len.gpu, + default_top_k=default_top_k, + default_top_p=default_top_p, + win_size=win_size, + tau_r=tau_r, + eps=eps, + ) + return sampled, logits + rows = expanded_idx_mapping.long() + row_temperature = temperature[rows] + scores = torch.log_softmax(logits / row_temperature.clamp_min(eps).unsqueeze(1), dim=1) + num_rows, vocab = scores.shape + top_k, top_p = states.get_top_k_top_p(expanded_idx_mapping, idx_mapping_np) + if top_p is None: + top_p = scores.new_full((num_rows,), default_top_p) + if top_k is None: + top_k = torch.full((num_rows,), default_top_k, dtype=torch.int32, device=scores.device) + sorted_scores, sorted_ids = scores.sort(dim=1, descending=True, stable=True) + sorted_probs = sorted_scores.exp() + # Top-p is decided on the full distribution, before the top-k cap. + keep = (sorted_probs.cumsum(dim=1) - sorted_probs) < top_p.unsqueeze(1) + ranks = torch.arange(vocab, device=scores.device) + keep &= (top_k.unsqueeze(1) <= 0) | (ranks.unsqueeze(0) < top_k.unsqueeze(1)) + nucleus = sorted_scores.masked_fill(~keep, float("-inf")) + seeds = states.seeds.gpu + draws = gumbel_sample(nucleus, expanded_idx_mapping, temperature, seeds, pos, False, False) + sampled = sorted_ids.gather(1, draws.unsqueeze(1)) + if win_size <= 0: + return sampled.squeeze(1), scores + + total_len = req_states.total_len.gpu[rows].long() + prompt_len = req_states.prompt_len.gpu[rows].long() + offsets = total_len.unsqueeze(1) - 1 - torch.arange(win_size, device=scores.device).unsqueeze(0) + in_output = offsets >= prompt_len.unsqueeze(1) + history = req_states.all_token_ids.gpu[rows.unsqueeze(1), offsets.clamp_min(0)] + repeats = ((history == sampled) & in_output).sum(dim=1) + rejected = (repeats >= win_size * tau_r) & in_output.any(dim=1) & (row_temperature > 0) + remaining = scores.scatter(1, sampled, float("-inf")) + remaining = torch.where(torch.isfinite(remaining).any(dim=1, keepdim=True), remaining, scores) + # Without top-k/top-p, from a noise stream independent of the first draw. + replacement = gumbel_sample(remaining, expanded_idx_mapping, temperature, seeds, pos, False, True) + return torch.where(rejected, replacement, sampled.squeeze(1)), scores + + +class _SpeechTokenBatcher: + """Run concurrent reference speech-token extractions as one S3 batch. + + Cold references are conditioned by several worker threads. One S3 call per + reference runs batch-1 FP32 GEMMs, which cost ~3x more GPU time per + reference than a batch of 8 and compete with the talker and codec stages. + A single thread drains every request already queued into one padded + ``quantize`` call; it never waits for more, so an idle server adds no + latency. Padding changes ~0.1% of prompt tokens versus batch-1 extraction. + """ + + def __init__(self, model, s3, device, max_batch: int = 16) -> None: + import queue + import threading + + self._model, self._s3, self._device, self._max_batch = model, s3, device, max_batch + self._queue: queue.SimpleQueue = queue.SimpleQueue() + self._empty = queue.Empty + threading.Thread(target=self._run, name="cosyvoice3-s3-batch", daemon=True).start() + + def tokens(self, mel: torch.Tensor) -> torch.Tensor: + from concurrent.futures import Future + + future: Future = Future() + self._queue.put((mel, future)) + return future.result() + + def _run(self) -> None: + while True: + items = [self._queue.get()] + while len(items) < self._max_batch: + try: + items.append(self._queue.get_nowait()) + except self._empty: + break + try: + mels, lens = self._s3.padding([mel for mel, _ in items]) + with torch.inference_mode(): + codes, codes_lens = self._model.quantize(mels.to(self._device), lens.to(self._device)) + codes = codes.cpu() + for (_, future), row, n in zip(items, codes, codes_lens.tolist()): + future.set_result(row[:n].clone()) + except BaseException as exc: # noqa: BLE001 - surfaced to every waiting caller + for _, future in items: + if not future.done(): + future.set_exception(exc) + + +_REFERENCE_CONDITIONERS: dict[str, "CosyVoice3MultiModalProcessor"] = {} +_REFERENCE_CACHE_DEBUG = {True: 0, False: 0} +_REFERENCE_CONDITIONERS_LOCK = Lock() + + +def prefetch_reference_conditioning(model_dir: str, config, audio) -> None: + """Warm the reference cache for ``audio`` (already at ``target_sr``). + + Runs in a worker thread before the request enters the synchronous input + processor, so concurrent cold references are prepared in parallel (the + ONNX/TensorRT/torch work releases the GIL) and the processor hits the + cache instead of blocking the serving event loop. + """ + with _REFERENCE_CONDITIONERS_LOCK: + conditioner = _REFERENCE_CONDITIONERS.get(model_dir) + if conditioner is None: + conditioner = object.__new__(CosyVoice3MultiModalProcessor) + conditioner._ensure_cached_runtime_components(model_dir, config) + _REFERENCE_CONDITIONERS[model_dir] = conditioner + _audio_conditioning(conditioner, audio, model_dir, config, {}) + + class CosyVoice3MultiModalProcessingInfo(BaseProcessingInfo): def get_hf_config(self): """If the config is not already present pass it @@ -244,7 +512,13 @@ class CosyVoice3MultiModalProcessor(OmniMultiModalProcessor[CosyVoice3MultiModal def _ensure_s3_model(cls): if cls._s3_model is not None: return cls._s3_model + with _REFERENCE_CONDITIONERS_LOCK: + if cls._s3_model is None: + cls._s3_model = cls._load_s3_model() + return cls._s3_model + @classmethod + def _load_s3_model(cls): # s3tokenizer is imported lazily (kept off the module top level) so # callers that don't use the CosyVoice3 talker need not install it. try: @@ -258,8 +532,18 @@ class CosyVoice3MultiModalProcessor(OmniMultiModalProcessor[CosyVoice3MultiModal model = _s3.load_model("speech_tokenizer_v3_25hz") device = current_omni_platform.get_torch_device() model = model.to(device).eval() - cls._s3_model = (model, _s3, device) - return cls._s3_model + return (model, _s3, device) + + _s3_batcher: "_SpeechTokenBatcher | None" = None + + @classmethod + def _ensure_s3_batcher(cls) -> "_SpeechTokenBatcher": + if cls._s3_batcher is None: + model, s3, device = cls._ensure_s3_model() + with _REFERENCE_CONDITIONERS_LOCK: + if cls._s3_batcher is None: + cls._s3_batcher = _SpeechTokenBatcher(model, s3, device) + return cls._s3_batcher def _extract_speech_token_via_s3(self, audio, return_device): """Drop-in replacement for ``extract_speech_token`` that uses the @@ -280,12 +564,15 @@ class CosyVoice3MultiModalProcessor(OmniMultiModalProcessor[CosyVoice3MultiModal audio_t = torch.from_numpy(wav) mel = _s3.log_mel_spectrogram(audio_t) - mels_p, mels_lens = _s3.padding([mel]) - with torch.inference_mode(): - codes, codes_lens = model.quantize(mels_p.to(dev), mels_lens.to(dev)) - n = int(codes_lens[0].item()) - speech_token = codes[:1, :n].to(dtype=torch.int32, device=return_device) - speech_token_len = torch.tensor([n], dtype=torch.int32, device=return_device) + if torch.device(dev).type == "cuda": + codes = self._ensure_s3_batcher().tokens(mel) + else: + mels_p, mels_lens = _s3.padding([mel]) + with torch.inference_mode(): + codes, codes_lens = model.quantize(mels_p.to(dev), mels_lens.to(dev)) + codes = codes[0, : int(codes_lens[0].item())] + speech_token = codes.reshape(1, -1).to(dtype=torch.int32, device=return_device) + speech_token_len = torch.tensor([speech_token.shape[1]], dtype=torch.int32, device=return_device) return speech_token, speech_token_len def _call_hf_processor( @@ -342,71 +629,17 @@ class CosyVoice3MultiModalProcessor(OmniMultiModalProcessor[CosyVoice3MultiModal int(text_token_len), int(input_len), ) - device = "cpu" - - # Speaker cache: skip 3 ONNX sessions on cache hit - voice_name = mm_kwargs.get("voice_name") - cache_key = None - if voice_name and isinstance(voice_name, str): - cache_key = self._speaker_cache.make_cache_key( - voice_name, - model_type="cosyvoice3", - created_at=int(mm_kwargs.get("voice_created_at") or 0), - ) - cached = self._speaker_cache.get(cache_key) - if cached is not None: - ft = BatchFeature( - { - "input_ids": input_ids, - "speech_feat": cached["speech_feat"].clone(), - "speech_token": cached["speech_token"].clone(), - "speech_token_len": [cached["speech_token_len"].clone()], - "embedding": cached["embedding"].clone(), - } - ) - return ft - - # Speech-token extraction via the S3Tokenizer PyTorch model on GPU - # (~30x faster than the bundled ``speech_tokenizer_v3.onnx`` CPU ONNX - # path in this venv). - speech_token, speech_token_len = self._extract_speech_token_via_s3(audio, device) - - speech_feat, speech_feat_len = extract_speech_feat(audio, self.feat_extractor, device) - - if config.sample_rate == 24000: - token_len = min(int(speech_feat.shape[1] / 2), speech_token.shape[1]) - speech_feat, speech_feat_len[:] = speech_feat[:, : 2 * token_len], 2 * token_len - speech_token, speech_token_len[:] = speech_token[:, :token_len], token_len - - if self.campplus_trt is not None: - embedding = extract_spk_embedding_trt(audio, self.campplus_trt, device) - else: - embedding = extract_spk_embedding(audio, self.campplus_session, device) - - # Cache the extracted artifacts for named speakers - if cache_key is not None: - self._speaker_cache.put( - cache_key, - { - "speech_feat": speech_feat.detach().cpu(), - "speech_token": speech_token.detach().cpu(), - "speech_token_len": speech_token_len.detach().cpu(), - "embedding": embedding.detach().cpu(), - }, - ) - - ft = BatchFeature( + conditioning = _audio_conditioning(self, audio, model_dir, config, mm_kwargs) + return BatchFeature( { "input_ids": input_ids, - "speech_feat": speech_feat, - "speech_token": speech_token, - "speech_token_len": [speech_token_len], - "embedding": embedding, + "speech_feat": conditioning["speech_feat"], + "speech_token": conditioning["speech_token"], + "speech_token_len": [conditioning["speech_token_len"]], + "embedding": conditioning["embedding"], } ) - return ft - def _get_mm_fields_config( self, hf_inputs: "BatchFeature", @@ -487,13 +720,31 @@ class CosyVoice3Model( supports_multimodal_raw_input_only = True supports_multimodal = True requires_raw_input_tokens = True + supports_embed_input_ids_query_start_loc = True prefer_model_sampler = True _sampling_eps = 1e-5 + @property + def logits_vocab_size(self) -> int | None: + """Width of the logits ``compute_logits`` returns (the speech head). + + The runners drop stop ids beyond it from min-tokens masking: vLLM folds + the text tokenizer's EOS into every request's stop set, which this + head can never emit (and MRv2's unchecked mask kernel would write past + the logits row). + """ + if self.model_stage != "cosyvoice3_talker": + return None + return int(self.config.llm["speech_token_size"]) + 200 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() self.config = vllm_config.model_config.hf_config + standard_sampling = cosyvoice3_standard_sampling(self.config) + logger.info("CosyVoice3 sampling policy: %s", "standard" if standard_sampling else "ras") self.have_multimodal_outputs = True + # Full-response codec consumes tokens and conditioning, not LLM hidden states. + self.omni_pooler_payload_include_hidden = not cosyvoice3_packed_inference_enabled() self.model_stage = vllm_config.model_config.model_stage model_dir = vllm_config.model_config.model if not os.path.isdir(model_dir): @@ -552,6 +803,13 @@ class CosyVoice3Model( # Use parent's cache config - critical for PagedAttention to work correctly return parent_config.with_hf_config(qwen_hf_config, architectures=["Qwen2Model"]) + def on_requests_finished(self, request_ids: set[str]) -> None: + """Release per-request vocoder state on completion, cancellation or error.""" + if hasattr(self, "_stream_vocoder_cache_by_req"): + with self._stream_audio_cache_lock: + for request_id in request_ids: + self._stream_vocoder_cache_by_req.pop(request_id, None) + def _stitch_stream_audio(self, req_id: str | None, audio: torch.Tensor, stream_finished: bool) -> torch.Tensor: """Pass-through stitching for async_chunk. @@ -692,6 +950,74 @@ class CosyVoice3Model( return False return True + def _ras_sample_batch( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + *, + default_top_p: float, + default_top_k: int, + win_size: int, + tau_r: float, + ) -> torch.Tensor: + """Batch random RAS without per-row scalar parameter/token transfers. + + Rejection flags cross to the host once per batch so only requests that + actually reject a token consume a second RNG draw. Per-request seeded + generators retain their ownership when rejection compacts the batch. + """ + batch_size = logits.shape[0] + + def parameter(value: torch.Tensor | None, default: float | int) -> torch.Tensor: + if value is None or value.numel() == 0: + return torch.full((batch_size,), default, device=logits.device) + value = value.reshape(-1).to(device=logits.device) + if value.numel() < batch_size: + value = torch.cat((value, value[-1:].expand(batch_size - value.numel()))) + return value[:batch_size] + + temperature = parameter(sampling_metadata.temperature, 1.0) + top_p = parameter(sampling_metadata.top_p, default_top_p) + top_k = parameter(sampling_metadata.top_k, default_top_k) + weighted_scores = torch.log_softmax(logits / temperature.clamp_min(self._sampling_eps).unsqueeze(1), dim=1) + sorted_probs, sorted_ids = weighted_scores.softmax(dim=1).sort(dim=1, descending=True, stable=True) + # Keep top-p based on the full distribution, before top-k masking. + keep = (sorted_probs.cumsum(dim=1) - sorted_probs) < top_p.unsqueeze(1) + ranks = torch.arange(logits.shape[1], device=logits.device) + keep &= (top_k.unsqueeze(1) <= 0) | (ranks.unsqueeze(0) < top_k.unsqueeze(1)) + weights = sorted_probs * keep.to(sorted_probs.dtype) + generators = {i: g for i, g in sampling_metadata.generators.items() if i < batch_size} + draws = random_sample(weights, generators).reshape(-1, 1).long() + sampled = sorted_ids.gather(1, draws).squeeze(1) + + histories = sampling_metadata.output_token_ids + if win_size > 0 and any(histories[:batch_size]): + recent = [] + for i in range(batch_size): + row = list(histories[i][-win_size:]) if i < len(histories) else [] + recent.append([-1] * (win_size - len(row)) + row) + history = torch.tensor(recent, dtype=torch.long, device=logits.device) + repeated = ((history == sampled.unsqueeze(1)).sum(dim=1) >= win_size * tau_r) & (history >= 0).any(dim=1) + rejected_rows = [i for i, reject in enumerate(repeated.tolist()) if reject] + if rejected_rows: + rows = torch.tensor(rejected_rows, device=logits.device, dtype=torch.long) + scores = weighted_scores.index_select(0, rows) + rejected_ids = sampled.index_select(0, rows).unsqueeze(1) + original = scores.gather(1, rejected_ids) + scores.scatter_(1, rejected_ids, float("-inf")) + replacement = torch.where( + torch.isfinite(scores).any(dim=1, keepdim=True), + torch.full_like(original, float("-inf")), + original, + ) + scores.scatter_(1, rejected_ids, replacement) + compact_generators = {i: generators[row] for i, row in enumerate(rejected_rows) if row in generators} + # RAS rejection samples the complete remaining distribution, + # without applying top-k/top-p a second time. + replacement_ids = random_sample(scores.softmax(dim=1), compact_generators).reshape(-1).long() + sampled.index_copy_(0, rows, replacement_ids) + return sampled.to(torch.int32) + def sample( self, logits: torch.Tensor, @@ -702,14 +1028,31 @@ class CosyVoice3Model( if self.model_stage != "cosyvoice3_talker": return None + if cosyvoice3_standard_sampling(self.config): + sampler = getattr(self, "_talker_sampler", None) + if sampler is None: + sampler = Sampler() + self._talker_sampler = sampler + # Penalize generated tokens only, never text or reference + # speech in the multimodal prompt. Padding is ignored by vLLM. + if sampling_metadata.prompt_token_ids is not None: + sampling_metadata = replace( + sampling_metadata, + prompt_token_ids=torch.full_like(sampling_metadata.prompt_token_ids, logits.shape[-1]), + ) + return sampler(logits=logits, sampling_metadata=sampling_metadata) + if not self._cosyvoice3_ras_enabled(sampling_metadata): sampler = getattr(self, "_talker_sampler", None) if sampler is None: sampler = Sampler() self._talker_sampler = sampler - return sampler(logits=logits, sampling_metadata=sampling_metadata) + return sampler(logits=self._full_vocab_logits(logits), sampling_metadata=sampling_metadata) logits = logits.to(torch.float32) + mask = sampling_metadata.allowed_token_ids_mask + if mask is not None and mask.shape[-1] != logits.shape[-1]: + sampling_metadata = replace(sampling_metadata, allowed_token_ids_mask=mask[..., : logits.shape[-1]]) # Apply logits processors directly — RAS handles its own repetition # logic. We avoid instantiating Sampler() here because its import # chain pulls in flashinfer / GPU deps that fail in CPU-only tests. @@ -717,6 +1060,14 @@ class CosyVoice3Model( logits.masked_fill_(sampling_metadata.allowed_token_ids_mask, float("-inf")) for processor in sampling_metadata.logitsprocs.non_argmax_invariant: logits = processor.apply(logits) + # The text tokenizer vocabulary is much larger than the speech head. + # compute_logits pads its output for the generic sampler/processors; + # RAS only needs speech codes and the merged stop token. Keep the + # original token indices, and trim only after full-vocabulary masks + # and processors have run. The generic sampler fallback above retains + # its full-vocabulary contract (including logprobs and bad words). + speech_head_size = int(self.config.llm["speech_token_size"]) + 200 + logits = logits[..., :speech_head_size] finite_logits = torch.isfinite(logits) if not finite_logits.any(dim=-1).all().item(): raise ValueError("CosyVoice3 sampling received a row with no finite logits") @@ -728,6 +1079,17 @@ class CosyVoice3Model( win_size = int(sampling_cfg.get("win_size", 10)) tau_r = float(sampling_cfg.get("tau_r", 0.1)) + if sampling_metadata.all_random: + sampled = self._ras_sample_batch( + logits, + sampling_metadata, + default_top_p=default_top_p, + default_top_k=default_top_k, + win_size=win_size, + tau_r=tau_r, + ) + return SamplerOutput(sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=None) + sampled_ids: list[int] = [] for req_idx in range(int(logits.shape[0])): row_logits = logits[req_idx] @@ -759,6 +1121,48 @@ class CosyVoice3Model( sampled = torch.tensor(sampled_ids, device=logits.device, dtype=torch.int32) return SamplerOutput(sampled_token_ids=sampled.unsqueeze(-1), logprobs_tensors=None) + def mrv2_custom_sampler(self, sampler: Any) -> None: + """Model Runner V2: the V1 sampling policies on the MRv2 sampler. + + RAS replaces the sampling step (``_ras_mrv2_sample``). Standard + sampling penalizes generated speech tokens only: the V1 path hides the + prompt from penalties (SGLang parity: text and reference speech are + never penalized), while MRv2 keeps penalty statistics on the GPU and + bins every prompt token of a new request, whose text ids also lie + outside the narrow speech head. Bin the prompt into a scratch mask + instead, keeping resumed output tokens counted. + """ + if self.model_stage != "cosyvoice3_talker": + return None + if not cosyvoice3_standard_sampling(self.config): + sampling_cfg = dict(self.config.llm.get("sampling", {})) + ras = partial( + _ras_mrv2_sample, + 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)), + eps=self._sampling_eps, + ) + sampler.sample = MethodType(ras, sampler) + state = getattr(sampler, "penalties_state", None) + if state is not None: + state.apply_staged_writes = MethodType(_generated_only_penalty_writes, state) + # MRv2 captures tensor outputs and constructs OmniOutput afterwards. + # This contract also applies when only the codec stage enables packed Flow. + self._mrv2_tensor_output = True + # MRv2 hands make_omni_output no multimodal kwargs; carry each prompt's + # conditioning from the encoder call to its batch row instead. + self._mrv2_encoded_conditioning = OrderedDict() + return None + + def _full_vocab_logits(self, logits: torch.Tensor) -> torch.Tensor: + """Pad the speech head to the text vocabulary for the generic sampler contract.""" + pad_size = int(self.config.vocab_size) - logits.size(-1) + if pad_size <= 0: + return logits + return torch.cat([logits, logits.new_full(logits.shape[:-1] + (pad_size,), float("-inf"))], dim=-1) + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput) -> torch.Tensor | None: if isinstance(hidden_states, OmniOutput): hidden_states = hidden_states.text_hidden_states @@ -771,17 +1175,15 @@ class CosyVoice3Model( # 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) - 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) + if not cosyvoice3_standard_sampling(self.config): + 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 + # The model owns sampling and the runner accepts a narrow codec + # head, so keep only the speech head: padding every decode step to + # the ~152k text vocabulary multiplied sampler and penalty work. + # The generic-sampler fallback pads on demand (_full_vocab_logits). return logits else: raise RuntimeError(f"compute_logits is only valid for {self.model_stage}.") @@ -799,8 +1201,25 @@ class CosyVoice3Model( # [1, T, emb] tensor (the caller's extend() iterates dim 0). if isinstance(speech_token, (list, tuple)): emb_dim = self.model.speech_embedding.weight.shape[1] - return [self.model.speech_embedding(t).reshape(-1, emb_dim) for t in speech_token] - return self.model.speech_embedding(speech_token) + embeddings = [self.model.speech_embedding(t).reshape(-1, emb_dim) for t in speech_token] + else: + embeddings = self.model.speech_embedding(speech_token) + encoded = getattr(self, "_mrv2_encoded_conditioning", None) + if encoded is not None: + conditioning = self._split_prompt_conditioning( + speech_token, kwargs.get("speech_feat"), kwargs.get("embedding"), kwargs.get("speech_token_len") + ) + for item, embedding in enumerate(embeddings): + # The encoder cache hands this item's tensor back (whole) to + # embed_input_ids, where its storage identifies the prompt. + # Identical references share one cached tensor, so entries + # stay until a newer encoding reuses the storage. + key = embedding.data_ptr() + encoded[key] = tuple(None if values is None else values[item] for values in conditioning) + encoded.move_to_end(key) + while len(encoded) > _MRV2_CONDITIONING_ENTRIES: + encoded.popitem(last=False) + return embeddings else: raise RuntimeError(f"embed_multimodal is only valid for {self.model_stage}.") @@ -809,86 +1228,80 @@ class CosyVoice3Model( input_ids: torch.Tensor, multimodal_embeddings=None, is_multimodal=None, + query_start_loc: Sequence[int] | None = None, ) -> torch.Tensor: if self.model_stage == "cosyvoice3_talker": - if is_multimodal is not None and any(is_multimodal): - # Per-request rearrange for the talker prompt layout - # [SOS_ph, TASK_ID_ph, AUDIO_ph * pstoken_len, ...text_tokens...] - # → [SOS_emb, ...text_embs..., TASK_ID_emb, AUDIO_embs]. - # - # In batched prefill ``input_ids`` is the flat concatenation of - # N sequences and ``multimodal_embeddings`` is a list with one - # tensor per multimodal request. Each request's audio - # placeholders form one contiguous ``True`` group in - # ``is_multimodal``; the 2 positions immediately before each - # group are the SOS / TASK_ID placeholders. Walking the groups - # lets us locate every request's segment without needing a - # ``seq_token_counts`` kwarg (which only the generation runner - # injects, not the AR runner used here). - embed_tokens = self.model.llm.model.embed_tokens(input_ids) - sos = self.model.speech_embedding.weight[self.model.sos].reshape(1, -1) - task_id = self.model.speech_embedding.weight[self.model.task_id].reshape(1, -1) + # Decode-only steps carry no speech embeddings; checking the + # (device) placeholder mask would synchronize every step. + if not multimodal_embeddings or is_multimodal is None or not torch.any(is_multimodal): + return self.model.speech_embedding.weight[input_ids] - is_mm = is_multimodal.to(torch.bool).reshape(-1) - # vLLM v1's multimodal placeholder range covers the FULL - # ``[SOS_ph, TASK_ID_ph, AUDIO_ph * pstoken_len]`` block - # (length ``2 + pstoken_len``, see ``_get_prompt_updates``'s - # ``insertion_end`` which inserts ``[1] * (1 + 1 + token_len)`` - # and ``mm_position.length`` matches that). So each contiguous - # ``True`` group in ``is_mm`` starts at a request's SOS - # placeholder, NOT at its first audio placeholder. - prev_false = torch.cat([torch.ones(1, dtype=torch.bool, device=is_mm.device), ~is_mm[:-1]]) - group_starts = (is_mm & prev_false).nonzero(as_tuple=True)[0].tolist() - if len(group_starts) != len(multimodal_embeddings): - raise RuntimeError( - f"cosyvoice3 talker: found {len(group_starts)} placeholder " - f"blocks in is_multimodal but {len(multimodal_embeddings)} " - f"multimodal_embeddings tensors — these must match 1:1." - ) + # Requests can interleave new prefills and ongoing decodes after + # scheduler slot reuse. A placeholder group identifies the start + # of a prompt, but cannot identify its end: use runner boundaries. + total_tokens = input_ids.numel() + boundaries = list(query_start_loc) if query_start_loc is not None else [0, total_tokens] + if not boundaries or boundaries[0] != 0 or boundaries[-1] != total_tokens: + raise ValueError("CosyVoice3 input embedding boundaries must cover all scheduled tokens") + if any(end <= start for start, end in zip(boundaries, boundaries[1:])): + raise ValueError("CosyVoice3 input embedding boundaries must be strictly increasing") + mm_mask = is_multimodal.reshape(-1).tolist() + # Conditioning contains prefill rows only, whereas the downstream + # payload splitter indexes the full (prefill + decode) request batch. + self._conditioning_request_rows = [ + i for i, (start, end) in enumerate(zip(boundaries, boundaries[1:])) if any(mm_mask[start:end]) + ] + self._conditioning_request_count = len(boundaries) - 1 + self._mrv2_step_conditioning = [] - total_tokens = int(embed_tokens.shape[0]) - # vLLM v1 packs cached (decode) requests first, then new - # (prefill) requests. The decode tokens at positions - # ``[0:group_starts[0]]`` get a speech_embedding lookup (the - # talker is generating codec tokens here, so input ids are - # codec ids, not text ids). Positions starting at - # ``group_starts[0]`` are the first new prefill request's - # SOS placeholder. - segments: list[torch.Tensor] = [] - first_prefill = group_starts[0] - if first_prefill > 0: - decode_ids = input_ids[:first_prefill].to(dtype=torch.long) - segments.append(self.model.speech_embedding.weight[decode_ids]) - - for i, req_start in enumerate(group_starts): - pstoken_len_i = int(multimodal_embeddings[i].shape[0]) - # Real text starts after the 2 + pstoken_len_i leading - # placeholders of this request's prompt. - text_start = req_start + 2 + pstoken_len_i - # Next request's segment starts at its own SOS placeholder, - # which is exactly the next True group's first index. - if i + 1 < len(group_starts): - req_end = group_starts[i + 1] - else: - req_end = total_tokens - segments.append(sos) - segments.append(embed_tokens[text_start:req_end]) - segments.append(task_id) - segments.append(multimodal_embeddings[i]) - embed_tokens = torch.cat(segments, dim=0) - else: - embed_tokens = self.model.speech_embedding.weight[input_ids] - return embed_tokens + text_embeds = self.model.llm.model.embed_tokens(input_ids) + sos = self.model.speech_embedding.weight[self.model.sos].reshape(1, -1) + task_id = self.model.speech_embedding.weight[self.model.task_id].reshape(1, -1) + segments = [] + mm_index = 0 + for start, end in zip(boundaries, boundaries[1:]): + request_mask = mm_mask[start:end] + if not any(request_mask): + segments.append(self.model.speech_embedding.weight[input_ids[start:end]]) + continue + if multimodal_embeddings is None or mm_index >= len(multimodal_embeddings): + raise ValueError("CosyVoice3 prefill is missing its speech embedding") + speech = multimodal_embeddings[mm_index] + mm_index += 1 + encoded = getattr(self, "_mrv2_encoded_conditioning", None) + if encoded is not None: + key = speech.data_ptr() + conditioning = encoded.get(key) + if conditioning is None: + raise ValueError("CosyVoice3 MRv2 prefill lost its prompt conditioning") + encoded.move_to_end(key) + self._mrv2_step_conditioning.append(conditioning) + prefix_len = 2 + speech.shape[0] + # The processor marks SOS, TASK_ID and all speech placeholders. + # Rearrangement needs the complete prompt block in this call. + if request_mask != [True] * prefix_len + [False] * (end - start - prefix_len): + raise ValueError("CosyVoice3 requires a complete prefill block within each request boundary") + segments.extend((sos, text_embeds[start + prefix_len : end], task_id, speech)) + if mm_index != len(multimodal_embeddings): + raise ValueError("CosyVoice3 speech embeddings must match prefill requests one to one") + return torch.cat(segments, dim=0) elif self.model_stage == "cosyvoice3_code2wav": assert input_ids.dim() == 1 - hidden = int(self.config.hidden_size) - return torch.zeros( - (input_ids.shape[0], hidden), - device=input_ids.device, - ) + return torch.zeros((input_ids.shape[0], int(self.config.hidden_size)), device=input_ids.device) else: raise RuntimeError(f"embed_input_ids is not valid for {self.model_stage}.") + def _align_prompt_conditioning(self, values): + rows = getattr(self, "_conditioning_request_rows", None) + if values is None or rows is None: + return values + if len(rows) != len(values): + raise ValueError("CosyVoice3 conditioning must match the prefill request rows") + aligned = [None] * self._conditioning_request_count + for row, value in zip(rows, values): + aligned[row] = value + return aligned + @staticmethod def _split_prompt_conditioning(speech_token, speech_feat, embedding, speech_token_len): """Split collated prompt conditioning into per-request lists. @@ -980,6 +1393,9 @@ class CosyVoice3Model( return self._code2wav_trt_done = True + # Packed Flow owns the torch estimator; speaker TRT remains independent. + if cosyvoice3_packed_inference_enabled(): + return if not (_cosyvoice3_trt_enabled() and torch.cuda.is_available()): return onnx_path = self._resolve_flow_estimator_onnx() @@ -1005,6 +1421,70 @@ class CosyVoice3Model( exc, ) + def make_omni_output(self, hidden_states: torch.Tensor, **kwargs) -> OmniOutput: + """Attach live prefill conditioning outside the full-response CUDA graph. + + Graph replay reuses capture-time output containers. Returning conditioning + from the captured forward would replay dummy prompt tensors on every + decode step. The runner calls this adapter with the current step's kwargs. + """ + multimodal_outputs = {} + + step_conditioning = getattr(self, "_mrv2_step_conditioning", None) + if "speech_token" not in kwargs and step_conditioning: + self._mrv2_step_conditioning = None + speech_token, speech_feat, embedding, speech_token_len = ( + list(values) for values in zip(*step_conditioning) + ) + multimodal_outputs = to_dict( + OmniPayloadStruct( + embed=EmbeddingsStruct( + speech_token=self._align_prompt_conditioning(speech_token), + speech_feat=self._align_prompt_conditioning(speech_feat), + speech_token_len=self._align_prompt_conditioning(speech_token_len), + embedding=self._align_prompt_conditioning(embedding), + ), + ) + ) + elif "speech_token" in kwargs: + # Prompt conditioning tensors for code2wav: live under + # ``embed.*`` per OmniPayloadStruct schema. + # + # vLLM hands these mm fields to forward() as collated, padded + # batch tensors (speech_token [B, maxT], speech_feat + # [B, 2*maxT, F], embedding [B, D]). Emitting them raw makes the + # downstream per-request payload split fragile at batch>1: it + # intermittently de-batches speech_token to 1-D and leaks the + # whole [B, D] embedding to every request, which corrupts voice + # conditioning and crashes code2wav (`prompt_token.shape[1]`). + # Instead, split into an explicit per-request list of unpadded, + # correctly-ranked tensors so ``to_payload_element`` splits them + # deterministically by request index (list[idx]). + speech_token_list, speech_feat_list, embedding_list, speech_token_len_list = ( + self._split_prompt_conditioning( + kwargs.get("speech_token"), + kwargs.get("speech_feat"), + kwargs.get("embedding"), + kwargs.get("speech_token_len"), + ) + ) + speech_token_list = self._align_prompt_conditioning(speech_token_list) + speech_feat_list = self._align_prompt_conditioning(speech_feat_list) + embedding_list = self._align_prompt_conditioning(embedding_list) + speech_token_len_list = self._align_prompt_conditioning(speech_token_len_list) + multimodal_outputs = to_dict( + OmniPayloadStruct( + embed=EmbeddingsStruct( + speech_token=speech_token_list, + speech_feat=speech_feat_list, + speech_token_len=speech_token_len_list, + embedding=embedding_list, + ), + ) + ) + + return OmniOutput(text_hidden_states=hidden_states, multimodal_outputs=multimodal_outputs) + def forward( self, input_ids: torch.Tensor, @@ -1013,7 +1493,7 @@ class CosyVoice3Model( inputs_embeds: torch.Tensor | None = None, additional_information: dict[str, object] | None = None, **kwargs: object, - ) -> OmniOutput: + ) -> torch.Tensor | OmniOutput: if self.model_stage == "cosyvoice3_talker": if inputs_embeds is None: inputs_embeds = self.embed_input_ids(input_ids) @@ -1021,42 +1501,9 @@ class CosyVoice3Model( # [total_tokens, hidden] hidden_states = self.model.llm(inputs_embeds, positions) - multimodal_outputs = {} - - if "speech_token" in kwargs: - # Prompt conditioning tensors for code2wav: live under - # ``embed.*`` per OmniPayloadStruct schema. - # - # vLLM hands these mm fields to forward() as collated, padded - # batch tensors (speech_token [B, maxT], speech_feat - # [B, 2*maxT, F], embedding [B, D]). Emitting them raw makes the - # downstream per-request payload split fragile at batch>1: it - # intermittently de-batches speech_token to 1-D and leaks the - # whole [B, D] embedding to every request, which corrupts voice - # conditioning and crashes code2wav (`prompt_token.shape[1]`). - # Instead, split into an explicit per-request list of unpadded, - # correctly-ranked tensors so ``to_payload_element`` splits them - # deterministically by request index (list[idx]). - speech_token_list, speech_feat_list, embedding_list, speech_token_len_list = ( - self._split_prompt_conditioning( - kwargs.get("speech_token"), - kwargs.get("speech_feat"), - kwargs.get("embedding"), - kwargs.get("speech_token_len"), - ) - ) - multimodal_outputs = to_dict( - OmniPayloadStruct( - embed=EmbeddingsStruct( - speech_token=speech_token_list, - speech_feat=speech_feat_list, - speech_token_len=speech_token_len_list, - embedding=embedding_list, - ), - ) - ) - - return OmniOutput(text_hidden_states=hidden_states, multimodal_outputs=multimodal_outputs) + if getattr(self, "_mrv2_tensor_output", False) or cosyvoice3_packed_inference_enabled(): + return hidden_states + return self.make_omni_output(hidden_states, **kwargs) elif self.model_stage == "cosyvoice3_code2wav": # Lazily swap the flow-decoder estimator to a TensorRT engine on the # first code2wav step (after weights are loaded), gated by the same @@ -1202,7 +1649,7 @@ class CosyVoice3Model( item_shapes, ) if ( - len(streaming_flow_items) > 1 + (len(streaming_flow_items) > 1 or cosyvoice3_packed_inference_enabled()) and cosyvoice3_batch_flow_enabled() and hasattr(self.code2wav, "forward_streaming_batch") ): diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index b35aefe52..6fd13c173 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py @@ -9,9 +9,10 @@ This module contains the code2wav (token-to-waveform) stage which uses: 3. HiFiGAN vocoder for waveform synthesis """ +import os from collections import Counter from collections.abc import Mapping -from typing import Any, TypedDict, cast +from typing import Any, Literal, TypedDict, cast, overload import numpy as np import torch @@ -34,6 +35,9 @@ from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreL from vllm_omni.model_executor.models.cosyvoice3.runtime import ( cosyvoice3_batch_flow_debug, cosyvoice3_batch_flow_profile, + cosyvoice3_full_response_enabled, + cosyvoice3_packed_inference_enabled, + cosyvoice3_packed_streaming_enabled, ) from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config @@ -57,6 +61,14 @@ def _build_dit_estimator(estimator_config: Mapping[str, Any]) -> DiT: return DiT(**estimator_config) +class StreamingHiFTState(TypedDict, total=False): + """Request-owned mel history, its absolute offset, and oscillator phase.""" + + mel: torch.Tensor + mel_offset: int + phase_acc: torch.Tensor | None + + class StreamingFlowItem(TypedDict, total=False): """One entry of the batched-streaming item list passed to forward_streaming_batch.""" @@ -67,7 +79,7 @@ class StreamingFlowItem(TypedDict, total=False): prompt_token: torch.Tensor prompt_feat: torch.Tensor embedding: torch.Tensor - cache_state: dict[str, torch.Tensor] | None + cache_state: StreamingHiFTState | None token_offset_tokens: int finalize: bool @@ -114,7 +126,6 @@ class CosyVoice3Code2Wav(nn.Module): pre_lookahead_layer=pre_lookahead_layer, decoder=decoder, ) - # Build HiFiGAN vocoder f0_predictor = CausalConvRNNF0Predictor( num_class=config.hift["f0_predictor"]["num_class"], @@ -256,9 +267,9 @@ class CosyVoice3Code2Wav(nn.Module): self, feat: torch.Tensor, *, - cache_state: dict[str, torch.Tensor] | None = None, + cache_state: StreamingHiFTState | None = None, finalize: bool = False, - ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: + ) -> tuple[torch.Tensor, StreamingHiFTState | None]: hift_weight = self.hift.m_source.l_linear.weight chunk_mel = feat.to(device=hift_weight.device, dtype=hift_weight.dtype) @@ -325,10 +336,24 @@ class CosyVoice3Code2Wav(nn.Module): if finalize: return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), None + # Clone GPU views so each request owns its cache and does not retain + # the full batch allocation or alias a later inference output. new_state = { - "mel": f0_input_mel.detach().cpu().contiguous(), + "mel": ( + f0_input_mel.detach().clone() + if cosyvoice3_packed_streaming_enabled() + else f0_input_mel.detach().cpu().contiguous() + ), "mel_offset": window_offset - f0_margin, - "phase_acc": new_phase_acc.detach().cpu().contiguous() if new_phase_acc is not None else None, + "phase_acc": ( + ( + new_phase_acc.detach().clone() + if cosyvoice3_packed_streaming_enabled() + else new_phase_acc.detach().cpu().contiguous() + ) + if new_phase_acc is not None + else None + ), } return emitted_speech.reshape(emitted_speech.shape[0], 1, -1), new_state @@ -338,14 +363,34 @@ class CosyVoice3Code2Wav(nn.Module): items: list[StreamingFlowItem], *, n_timesteps: int = 10, - ) -> list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]]: + ) -> list[tuple[torch.Tensor, StreamingHiFTState | None]]: """Batch the flow-matching mel path, then run HiFT per request. Items are grouped by prompt condition shape and finalization state. Codec tokens may have different lengths; those are padded within the group and passed to the flow as per-row token lengths. """ - results: list[tuple[torch.Tensor, dict[str, torch.Tensor] | None] | None] = [None] * len(items) + if items and cosyvoice3_packed_streaming_enabled(): + results = [None] * len(items) + for finalize in (False, True): + group = [(i, x) for i, x in enumerate(items) if bool(x.get("finalize", False)) == finalize] + if group: + values = self.forward_batch([x for _, x in group], n_timesteps=n_timesteps, stream_items=True) + for (i, _), value in zip(group, values): + results[i] = value + return results + if ( + items + and cosyvoice3_packed_inference_enabled() + and all( + item.get("finalize", False) + and item.get("cache_state") is None + and int(item.get("token_offset_tokens", 0)) == 0 + for item in items + ) + ): + return [(speech, None) for speech in self.forward_batch(items, n_timesteps=n_timesteps)] + results: list[tuple[torch.Tensor, StreamingHiFTState | None] | None] = [None] * len(items) groups: dict[tuple[int, int, int, bool], list[tuple[int, StreamingFlowItem]]] = {} for index, item in enumerate(items): assert isinstance(item["token"], torch.Tensor) @@ -445,7 +490,7 @@ class CosyVoice3Code2Wav(nn.Module): ) assert all(result is not None for result in results), "every streaming item must produce exactly one result" - return cast(list[tuple[torch.Tensor, dict[str, torch.Tensor] | None]], results) + return cast(list[tuple[torch.Tensor, StreamingHiFTState | None]], results) @torch.inference_mode() def forward_streaming( @@ -455,11 +500,11 @@ class CosyVoice3Code2Wav(nn.Module): prompt_feat: torch.Tensor, embedding: torch.Tensor, *, - cache_state: dict[str, torch.Tensor] | None = None, + cache_state: StreamingHiFTState | None = None, n_timesteps: int = 10, token_offset_tokens: int = 0, finalize: bool = False, - ) -> tuple[torch.Tensor, dict[str, torch.Tensor] | None]: + ) -> tuple[torch.Tensor, StreamingHiFTState | None]: """Decode streaming audio using cumulative mel + emitted-speech offset. This mirrors upstream CosyVoice3 streaming semantics more closely than @@ -480,6 +525,176 @@ class CosyVoice3Code2Wav(nn.Module): ) return self._stream_hift_from_feat(feat, cache_state=cache_state, finalize=finalize) + @torch.inference_mode() + def _stream_position_noise(self, width: int, like: torch.Tensor) -> torch.Tensor: + """Extend fixed-position noise without changing previously emitted prefixes. + + Fixed-size CPU blocks preserve the initial experiment's seed-0 noise, + isolate the global RNG, and make growth independent of request ordering. + """ + block_frames = 15000 + cached = getattr(self, "_position_noise", None) + present = 0 if cached is None else cached.shape[-1] + blocks = [] if cached is None else [cached.to(like)] + for index in range(present // block_frames, (width + block_frames - 1) // block_frames): + generator = torch.Generator(device="cpu").manual_seed(index) + blocks.append(torch.randn((1, 80, block_frames), generator=generator).to(like)) + if not blocks: + return like.new_empty((1, 80, 0)) + self._position_noise = torch.cat(blocks, dim=-1) if len(blocks) > 1 else blocks[0] + return self._position_noise[..., :width] + + @overload + def forward_batch( + self, items: list[StreamingFlowItem], *, n_timesteps: int = 10, stream_items: Literal[False] = False + ) -> list[torch.Tensor]: ... + + @overload + def forward_batch( + self, items: list[StreamingFlowItem], *, n_timesteps: int = 10, stream_items: Literal[True] + ) -> list[tuple[torch.Tensor, StreamingHiFTState | None]]: ... + + @torch.inference_mode() + def forward_batch( + self, items: list[StreamingFlowItem], *, n_timesteps: int = 10, stream_items: bool = False + ) -> list[torch.Tensor] | list[tuple[torch.Tensor, StreamingHiFTState | None]]: + """Batch Flow across different reference lengths and restore request order. + + Build each prompt/continuation contiguously before right padding. Padding + the prompt and continuation separately would insert a gap in conditioning. + The default keeps exact-length CPU-F0 HiFT. The opt-in Hopper profile + uses packed Flow, batched HiFT, and FP64 GPU F0. + """ + if not items: + return [] + if stream_items: + if not cosyvoice3_packed_streaming_enabled(): + raise ValueError("Packed streaming requires COSYVOICE3_PACKED_STREAMING on Hopper") + if any(bool(item.get("finalize", False)) != bool(items[0].get("finalize", False)) for item in items): + raise ValueError("Use forward_streaming_batch to group mixed finalization states") + if len(items) > 8: + ordered = sorted( + enumerate(items), key=lambda pair: pair[1]["token"].shape[1] + pair[1]["prompt_token"].shape[1] + ) + outputs: list[torch.Tensor | tuple[torch.Tensor, StreamingHiFTState | None] | None] = [None] * len(items) + for start in range(0, len(ordered), 8): + group = ordered[start : start + 8] + for (index, _), speech in zip( + group, + self.forward_batch([item for _, item in group], n_timesteps=n_timesteps, stream_items=stream_items), + ): + outputs[index] = speech + return cast(list[torch.Tensor] | list[tuple[torch.Tensor, StreamingHiFTState | None]], outputs) + flow = self.flow_model + weight = next(flow.parameters()) + mus, conds, speakers, lengths, prompt_lengths = [], [], [], [], [] + for item in items: + tokens = torch.cat( + tuple(item[key].to(device=weight.device, dtype=torch.long) for key in ("prompt_token", "token")), dim=1 + ) + embedded = flow.input_embedding(tokens.clamp(min=0)) + if stream_items and not item.get("finalize", False) and int(flow.pre_lookahead_len) > 0: + lookahead = int(flow.pre_lookahead_len) + hidden = flow.pre_lookahead_layer(embedded[:, :-lookahead], context=embedded[:, -lookahead:]) + else: + hidden = flow.pre_lookahead_layer(embedded) + mu = hidden.repeat_interleave(self.token_mel_ratio, dim=1) + prompt = item["prompt_feat"].to(device=weight.device, dtype=mu.dtype) + cond = torch.zeros_like(mu) + cond[:, : prompt.shape[1]] = prompt + speaker = torch.nn.functional.normalize(item["embedding"].to(weight), dim=1) + speakers.append(flow.spk_embed_affine_layer(speaker)) + mus.append(mu.transpose(1, 2)) + conds.append(cond.transpose(1, 2)) + lengths.append(mu.shape[1]) + prompt_lengths.append(prompt.shape[1]) + width = max(lengths) + mu = torch.cat([torch.nn.functional.pad(x, (0, width - x.shape[-1])) for x in mus]) + cond = torch.cat([torch.nn.functional.pad(x, (0, width - x.shape[-1])) for x in conds]) + mask = ( + (torch.arange(width, device=weight.device)[None, :] < torch.tensor(lengths, device=weight.device)[:, None]) + .unsqueeze(1) + .to(mu) + ) + if cosyvoice3_packed_inference_enabled(): + from .code2wav_core.packed_dit import ( + PackedDiT, + gather_rows, + pack_rows, + scatter_rows, + solve_flow_euler_packed, + ) + + if not hasattr(self, "_packed_estimator"): + self._packed_estimator = PackedDiT(flow.decoder.estimator) + self._packed_estimator.compile(mu.dtype) + rows = pack_rows(lengths, mu.device) + if stream_items: + noise = self._stream_position_noise(width, mu).expand(len(items), -1, -1) + else: + noise = torch.randn(mu.shape, device=mu.device, dtype=mu.dtype) + time_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype) + if flow.decoder.t_scheduler == "cosine": + time_span = 1 - torch.cos(time_span * 0.5 * torch.pi) + with torch.autocast("cuda", dtype=mu.dtype): + packed = solve_flow_euler_packed( + self._packed_estimator, + gather_rows(noise.transpose(1, 2), rows), + time_span, + gather_rows(mu.transpose(1, 2), rows), + torch.cat(speakers), + gather_rows(cond.transpose(1, 2), rows), + rows, + cfg_rate=flow.decoder.inference_cfg_rate, + streaming=stream_items and not items[0].get("finalize", False), + modulation_key=(n_timesteps, flow.decoder.t_scheduler, mu.dtype, mu.device), + ) + feat = scatter_rows(packed, rows, width).transpose(1, 2) + + else: + feat, _ = flow.decoder( + mu=mu, mask=mask, spks=torch.cat(speakers), cond=cond, n_timesteps=n_timesteps, streaming=False + ) + if stream_items: + results = [] + for row, item in enumerate(items): + start = prompt_lengths[row] + max(0, int(item.get("token_offset_tokens", 0))) * self.token_mel_ratio + mel = feat[row : row + 1, :, start : lengths[row]].float() + results.append( + self._stream_hift_from_feat( + mel, cache_state=item.get("cache_state"), finalize=bool(item.get("finalize", False)) + ) + ) + return results + hift_weight = self.hift.m_source.l_linear.weight + if cosyvoice3_full_response_enabled(): + mels = [] + for row, item in enumerate(items): + start = prompt_lengths[row] + max(0, int(item.get("token_offset_tokens", 0))) * self.token_mel_ratio + mels.append(feat[row : row + 1, :, start : lengths[row]].float().to(hift_weight)) + width = max(mel.shape[-1] for mel in mels) + if width == 0: + return [mel.new_zeros((1, 1, 0)) for mel in mels] + from .code2wav_core.packed_dit import bucketed_width + + # Recurring shapes let cuDNN/cuFFT reuse plans; rows are already + # right padded to the batch maximum and trimmed below. + width = bucketed_width(width) + padded = torch.cat([torch.nn.functional.pad(mel, (0, width - mel.shape[-1])) for mel in mels]) + speech, _, _ = self.hift.inference(speech_feat=padded, finalize=True) + stride = int(np.prod(self.hift.upsample_rates) * self.hift.istft_params["hop_len"]) + return [speech[row : row + 1, ..., : mel.shape[-1] * stride] for row, mel in enumerate(mels)] + results = [] + for row, item in enumerate(items): + start = prompt_lengths[row] + max(0, int(item.get("token_offset_tokens", 0))) * self.token_mel_ratio + mel = feat[row : row + 1, :, start : lengths[row]].float().to(hift_weight) + if mel.shape[-1] == 0: + speech = mel.new_zeros((1, 1, 0)) + else: + speech, _, _ = self.hift.inference(speech_feat=mel, finalize=True) + results.append(speech) + return results + @torch.inference_mode() def forward( self, @@ -524,7 +739,6 @@ class CosyVoice3Code2Wav(nn.Module): model_dir: Model directory containing flow.pt and hift.pt device: Device to load weights to """ - import os # Load flow weights flow_path = os.path.join(model_dir, "flow.pt") @@ -538,5 +752,21 @@ class CosyVoice3Code2Wav(nn.Module): k.replace("generator.", ""): v for k, v in torch.load(hift_path, map_location=device).items() } self.hift.load_state_dict(hift_state_dict, strict=True) - self.hift.to(device).eval() + self.hift.to(device) + # Fold after loading and device placement to avoid recomputing weights + # for every chunk. Folding changes state_dict keys, so reloading the + # original checkpoint into this instance is unsupported. + folded = self.hift.remove_weight_norm() + logger.info("Folded %d weight-norm layers in HiFT generator", folded) + if folded == 0: + logger.warning("HiFT generator had no weight-norm layers to fold; check config drift") + # Fold on the same device as mainline F0 inference. Folding on CPU + # and moving to GPU later changes the normalized FP32 weights. + f0_device = device if device.type == "cuda" and os.getenv("COSYVOICE3_F0_ON_CPU", "0") != "1" else "cpu" + self.hift.f0_predictor.to(device=f0_device, dtype=torch.float32) + f0_folded = self.hift.f0_predictor.remove_weight_norm() + logger.info("Folded %d weight-norm layers in HiFT F0 predictor", f0_folded) + if f0_folded == 0: + logger.warning("HiFT F0 predictor had no weight-norm layers to fold; check config drift") + self.hift.eval() logger.info(f"Loaded hift weights from {hift_path}") diff --git a/vllm_omni/model_executor/models/cosyvoice3/ras_sampler.py b/vllm_omni/model_executor/models/cosyvoice3/ras_sampler.py new file mode 100644 index 000000000..a373d6f37 --- /dev/null +++ b/vllm_omni/model_executor/models/cosyvoice3/ras_sampler.py @@ -0,0 +1,155 @@ +# ruff: noqa: N803 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused CosyVoice3 repetition-aware sampling for Model Runner V2. + +One program per row replaces the ~70 small kernels (softmax, full-vocabulary +sort, cumulative sum, masks, two Gumbel draws, history gather, rejection) of +the tensor implementation. The nucleus is the top-p prefix of the full +distribution capped at top-k, found by repeated argmax instead of a sort; a +draw that repeats within the last ``win_size`` generated tokens is replaced by +a draw from the full remaining distribution, from a disjoint noise stream. +""" + +import torch +from vllm.triton_utils import tl, triton + +MAX_FUSED_TOP_K = 64 + + +@triton.jit +def _gumbel(seed, offset): + u = tl.maximum(tl.rand(seed, offset), 4.6566127342e-10) + return -tl.log(-tl.log(u)) + + +@triton.jit +def _ras_sample_kernel( + out_ptr, + logits_ptr, + logits_stride, + vocab, + idx_mapping_ptr, + temperature_ptr, + top_k_ptr, + top_p_ptr, + seeds_ptr, + pos_ptr, + tokens_ptr, + tokens_stride, + total_len_ptr, + prompt_len_ptr, + default_top_k, + default_top_p, + eps, + threshold, + win_size, + HAS_TOP_K: tl.constexpr, + HAS_TOP_P: tl.constexpr, + WIN: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + req = tl.load(idx_mapping_ptr + row).to(tl.int64) + cols = tl.arange(0, BLOCK) + mask = cols < vocab + x = tl.load(logits_ptr + row * logits_stride + cols, mask=mask, other=float("-inf")).to(tl.float32) + temperature = tl.load(temperature_ptr + req).to(tl.float32) + if temperature == 0.0: + tl.store(out_ptr + row, tl.argmax(x, axis=0).to(tl.int64)) + return + x = x / tl.maximum(temperature, eps) + peak = tl.max(x, axis=0) + weights = tl.where(mask, tl.exp(x - peak), 0.0) + total = tl.sum(weights, axis=0) + log_probs = x - peak - tl.log(total) + top_k = tl.load(top_k_ptr + req) if HAS_TOP_K else default_top_k + top_p = tl.load(top_p_ptr + req).to(tl.float32) if HAS_TOP_P else default_top_p + seed = tl.load(seeds_ptr + req) + pos = tl.load(pos_ptr + row) + nucleus_seed = tl.randint(seed, pos) + + # Walk the distribution in descending order (ties: lowest id first, like a + # stable descending sort) while the top-p prefix and the top-k cap allow. + remaining = weights + before = 0.0 + best = float("-inf") + sampled = tl.argmax(x, axis=0).to(tl.int64) + for _ in range(top_k): + if before < top_p: + value = tl.max(remaining, axis=0) + token = tl.argmax(remaining, axis=0) + prob = value / total + score = tl.log(prob) + _gumbel(nucleus_seed, token) + if score > best: + best = score + sampled = token.to(tl.int64) + before += prob + remaining = tl.where(cols == token, -1.0, remaining) + + # Repetition-aware rejection over the last WIN generated tokens. + total_len = tl.load(total_len_ptr + req).to(tl.int64) + prompt_len = tl.load(prompt_len_ptr + req).to(tl.int64) + offsets = total_len - 1 - tl.arange(0, WIN) + in_output = (offsets >= prompt_len) & (tl.arange(0, WIN) < win_size) + history = tl.load(tokens_ptr + req * tokens_stride + tl.maximum(offsets, 0), mask=in_output, other=-1) + repeats = tl.sum((in_output & (history == sampled)).to(tl.float32), axis=0) + if (repeats >= threshold) & (tl.sum(in_output.to(tl.int32), axis=0) > 0): + # Offset keeps the replacement noise disjoint from the nucleus draw's. + replacement_seed = tl.randint(seed, pos + 1073741824) + scores = tl.where(mask & (cols != sampled), log_probs + _gumbel(replacement_seed, cols), float("-inf")) + if tl.max(scores, axis=0) > float("-inf"): + sampled = tl.argmax(scores, axis=0).to(tl.int64) + tl.store(out_ptr + row, sampled) + + +def fused_ras_sample( + logits: torch.Tensor, + idx_mapping: torch.Tensor, + temperature: torch.Tensor, + top_k: torch.Tensor | None, + top_p: torch.Tensor | None, + seeds: torch.Tensor, + pos: torch.Tensor, + all_token_ids: torch.Tensor, + total_len: torch.Tensor, + prompt_len: torch.Tensor, + *, + default_top_k: int, + default_top_p: float, + win_size: int, + tau_r: float, + eps: float, +) -> torch.Tensor: + """Sample one token per logits row; per-request tensors are indexed by ``idx_mapping``.""" + rows, vocab = logits.shape + out = torch.empty(rows, dtype=torch.int64, device=logits.device) + if rows == 0: + return out + _ras_sample_kernel[(rows,)]( + out, + logits, + logits.stride(0), + vocab, + idx_mapping, + temperature, + temperature if top_k is None else top_k, + temperature if top_p is None else top_p, + seeds, + pos, + all_token_ids, + all_token_ids.stride(0), + total_len, + prompt_len, + int(default_top_k), + float(default_top_p), + float(eps), + float(win_size * tau_r), + int(win_size), + HAS_TOP_K=top_k is not None, + HAS_TOP_P=top_p is not None, + WIN=triton.next_power_of_2(max(int(win_size), 1)), + BLOCK=triton.next_power_of_2(vocab), + num_warps=8, + ) + return out diff --git a/vllm_omni/model_executor/models/cosyvoice3/runtime.py b/vllm_omni/model_executor/models/cosyvoice3/runtime.py index d94725148..0b8231509 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/runtime.py +++ b/vllm_omni/model_executor/models/cosyvoice3/runtime.py @@ -29,3 +29,38 @@ def cosyvoice3_batch_flow_profile(name: str): if cosyvoice3_batch_flow_debug(): return torch.profiler.record_function(name) return nullcontext() + + +def cosyvoice3_full_response_enabled() -> bool: + """Opt-in Hopper full-response path; other devices retain the standard path.""" + return ( + _env_flag("COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS", "0") + and torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] == 9 + ) + + +def cosyvoice3_packed_streaming_enabled() -> bool: + """Opt-in Hopper packed, chunk-causal Flow with request-owned GPU HiFT state.""" + return ( + _env_flag("COSYVOICE3_PACKED_STREAMING", "0") + and torch.cuda.is_available() + and torch.cuda.get_device_capability()[0] == 9 + ) + + +def cosyvoice3_packed_inference_enabled() -> bool: + """Either optimized profile needs live conditioning outside CUDA graphs.""" + return cosyvoice3_full_response_enabled() or cosyvoice3_packed_streaming_enabled() + + +def cosyvoice3_standard_sampling(config: object) -> bool: + """Select ordinary sampling; RAS remains the default checkpoint behavior. + + The optional HF override is shared by API stop configuration and the model + sampler. Reject misspellings instead of silently selecting another policy. + """ + mode = getattr(config, "cosyvoice3_sampling_mode", "ras") + if mode not in ("ras", "standard"): + raise ValueError("cosyvoice3_sampling_mode must be 'ras' or 'standard'") + return mode == "standard" diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 777374f91..85b583b99 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -633,3 +633,77 @@ def _patch_cumem_free_callback_cuda() -> None: _patch_cumem_free_callback_cuda() + + +def _patch_batched_seeded_random_sample() -> None: + """Draw per-request seeded sampling noise in one launch. + + vLLM's ``random_sample`` calls ``exponential_(generator=...)`` once per + seeded request and step. The replacement reproduces those draws bit for + bit (values and generator states) with one kernel; unsupported cases keep + the original loop. + """ + try: + import inspect + + from vllm.v1.sample.ops import topk_topp_sampler as sampler_ops + except ImportError: + return + original = getattr(sampler_ops, "random_sample", None) + if original is None or getattr(original, "_omni_batched_seeded", False): + return + if list(inspect.signature(original).parameters) != ["probs", "generators", "use_fp64_gumbel"] or not all( + hasattr(sampler_ops, name) for name in ("empty_exponential_noise_like", "sample_with_exponential_noise") + ): + _PATCH_LOGGER.debug("[seeded-sampling] random_sample signature changed; batching disabled") + return + + def random_sample( + probs: torch.Tensor, + generators: dict[int, torch.Generator], + use_fp64_gumbel: bool = False, + ) -> torch.Tensor: + q = sampler_ops.empty_exponential_noise_like(probs, use_fp64_gumbel) + if len(generators) != probs.shape[0]: + q.exponential_() + if generators: + from vllm_omni.utils.seeded_exponential import ( + batched_seeded_exponential_supported, + fill_exponential_rows, + ) + + if batched_seeded_exponential_supported(q, generators): + fill_exponential_rows(q, list(generators.values()), list(generators)) + else: + for i, generator in generators.items(): + q[i].exponential_(generator=generator) + return sampler_ops.sample_with_exponential_noise(probs, q) + + random_sample._omni_batched_seeded = True # type: ignore[attr-defined] + random_sample.__wrapped__ = original # type: ignore[attr-defined] + sampler_ops.random_sample = random_sample + + +_patch_batched_seeded_random_sample() + + +def _patch_mrv2_stop_token_capacity() -> None: + """Let Model Runner V2 enforce ``min_tokens`` over large stop sets. + + MRV2 keeps each request's stop token ids for ``min_tokens`` masking in a + fixed-width GPU table capped at 128 ids and rejects larger sets. Speech + LMs stop on every special id past the codebook (CosyVoice3: 200). The + masking kernel sizes its block from the table widths, which other tables + already raise to 1024, so widening this one only grows a small buffer. + """ + try: + from vllm.v1.worker.gpu.sample import logit_bias + except ImportError: + return + logit_bias.MAX_NUM_STOP_TOKEN_IDS = max( + logit_bias.MAX_NUM_STOP_TOKEN_IDS, + logit_bias.MAX_NUM_ALLOWED_TOKEN_IDS, + ) + + +_patch_mrv2_stop_token_capacity() diff --git a/vllm_omni/utils/seeded_exponential.py b/vllm_omni/utils/seeded_exponential.py new file mode 100644 index 000000000..390c75747 --- /dev/null +++ b/vllm_omni/utils/seeded_exponential.py @@ -0,0 +1,121 @@ +# ruff: noqa: N803 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Batched per-request ``exponential_(generator=...)`` that is bit-identical to torch. + +Seeded sampling draws exponential noise one request at a time, which costs a +kernel launch and host overhead per request and step. This kernel reproduces +torch's CUDA distribution kernel exactly for float32 (the grid-stride +Philox4x32-10 ``curand_uniform4`` stream, the fast-math ``__logf`` transform +and the per-call Philox offset increment), so a batch draws in one launch while +every value and every generator state afterwards matches the per-request loop. +""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _generator_exponential_kernel( + out_ptr, + rows_ptr, + seed_ptr, + offset_ptr, + numel, + threads, + HAS_ROWS: tl.constexpr, + BLOCK: tl.constexpr, +): + request = tl.program_id(0) + elements = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = elements < numel + seed = tl.load(seed_ptr + request).to(tl.uint64) + offset = tl.load(offset_ptr + request).to(tl.uint64) + # torch's grid-stride kernel: thread t draws curand_uniform4 once per grid + # stride and element j = t + threads * (4 * call + component). + thread = (elements % threads).to(tl.uint64) + stride = (elements // threads).to(tl.uint64) + counter = offset // 4 + stride // 4 + component = stride % 4 + c0 = (counter & 0xFFFFFFFF).to(tl.uint32) + c1 = (counter >> 32).to(tl.uint32) + c2 = (thread & 0xFFFFFFFF).to(tl.uint32) + c3 = (thread >> 32).to(tl.uint32) + r0, r1, r2, r3 = tl.philox(seed, c0, c1, c2, c3, 10) + bits = tl.where(component == 0, r0, tl.where(component == 1, r1, tl.where(component == 2, r2, r3))) + # curand_uniform: (0, 1] from 32 bits. + uniform = bits.to(tl.float32) * 2.3283064365386963e-10 + 1.1641532182693481e-10 + # at::exponential with lambda 1 under fast math (at::log is __logf). + log = tl.where( + uniform >= 1.0 - 5.960464477539063e-08, + -5.960464477539063e-08, + tl.extra.cuda.libdevice.fast_logf(uniform), + ) + row = tl.load(rows_ptr + request).to(tl.int64) if HAS_ROWS else request.to(tl.int64) + tl.store(out_ptr + row * numel + elements, -log, mask=mask) + + +_POLICY_CACHE: dict[tuple[int, int], tuple[int, int]] = {} + + +def torch_exponential_policy(numel: int, device: torch.device) -> tuple[int, int]: + """(grid threads, Philox offset increment) of torch's CUDA ``exponential_``.""" + key = (numel, device.index or 0) + policy = _POLICY_CACHE.get(key) + if policy is None: + props = torch.cuda.get_device_properties(device) + blocks = min(props.multi_processor_count * (props.max_threads_per_multi_processor // 256), -(-numel // 256)) + threads = blocks * 256 + policy = _POLICY_CACHE[key] = (threads, ((numel - 1) // (threads * 4) + 1) * 4) + return policy + + +def fill_exponential_rows(out: torch.Tensor, generators: list, rows: list[int] | None = None) -> torch.Tensor: + """Draw ``out[rows[i]].exponential_(generator=generators[i])`` for all ``i`` in one launch. + + ``out`` is a contiguous float32 CUDA matrix; ``rows`` defaults to + ``range(len(generators))``. ``None`` generators use the default CUDA + generator in order. Each generator is advanced exactly as torch would. + """ + numel = int(out.shape[-1]) if rows is not None else int(out.numel()) // len(generators) + device = out.device + threads, increment = torch_exponential_policy(numel, device) + default = None + seeds, offsets = [], [] + for generator in generators: + if generator is None: + if default is None: + default = torch.cuda.default_generators[device.index or 0] + generator = default + offset = generator.get_offset() + generator.set_offset(offset + increment) + seed = generator.initial_seed() + seeds.append(seed - (1 << 64) if seed >= 1 << 63 else seed) + offsets.append(offset) + table = [seeds, offsets] if rows is None else [seeds, offsets, list(rows)] + state = torch.tensor(table, dtype=torch.int64, pin_memory=True).to(device, non_blocking=True) + block = 1024 + _generator_exponential_kernel[(len(generators), triton.cdiv(numel, block))]( + out, + state[2] if rows is not None else state[0], + state[0], + state[1], + numel, + threads, + HAS_ROWS=rows is not None, + BLOCK=block, + ) + return out + + +def batched_seeded_exponential_supported(q: torch.Tensor, generators: dict) -> bool: + """Whether ``fill_exponential_rows`` reproduces the per-row loop for ``q``.""" + return ( + len(generators) > 1 + and q.is_cuda + and q.dtype == torch.float32 + and q.dim() == 2 + and q.is_contiguous() + # A generator shared by two rows makes the loop order observable. + and len({id(g) for g in generators.values()}) == len(generators) + ) diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py index 1e5963cc2..08b7acdd0 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -413,6 +413,32 @@ class OmniModelState(DefaultModelState): # Input preparation # ------------------------------------------------------------------ + def prepare_inputs_embeds( + self, + scheduled_encoder_inputs: dict[str, list[int]], + input_batch: InputBatch, + req_states: RequestState, + ) -> torch.Tensor: + # Models that rearrange multimodal prompts per request (CosyVoice3) + # need the request boundaries the V1 runner passes as query_start_loc. + if ( + not self.supports_mm_inputs + or self.mm_pruner is not None + or not getattr(self.model, "supports_embed_input_ids_query_start_loc", False) + ): + return super().prepare_inputs_embeds(scheduled_encoder_inputs, input_batch, req_states) + self.execute_mm_encoder(scheduled_encoder_inputs) + mm_embeds, is_mm_embed = self.gather_mm_embeddings(input_batch) + kwargs: dict[str, Any] = {"multimodal_embeddings": mm_embeds, "is_multimodal": is_mm_embed} + if mm_embeds: + kwargs["query_start_loc"] = input_batch.query_start_loc_np.tolist() + embeds = self.model.embed_input_ids(input_batch.input_ids[: input_batch.num_tokens], **kwargs) + inputs_embeds = self.encoder_runner.inputs_embeds + inputs_embeds[: embeds.shape[0]] = embeds + if self.prompt_embeds_state is not None: + self.prompt_embeds_state.apply(input_batch, req_states.num_computed_tokens.gpu, inputs_embeds) + return inputs_embeds[: input_batch.num_tokens_after_padding] + def prepare_inputs(self, input_batch: InputBatch, req_states: RequestState) -> dict[str, Any]: # Forward-only stages have no preprocess hook. Honor their declared # GPU inputs before gathering the model's runtime view, so payload