[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>

.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:

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,
     }
 

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

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]

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",

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

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

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

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

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

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)

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

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

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")
                 ):

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}")

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

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"

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

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

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