[Model] Add opt-in packed Flow and batched HiFT inference for CosyVoice3 (#8224)

Signed-off-by: Sy03 <1370724210@qq.com>
Signed-off-by: zack <huixindaddy@yahoo.com>
Co-authored-by: zack <huixindaddy@yahoo.com>

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())

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

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)

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())

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)

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

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)

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

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]