Keep baseline TTS ASR evaluation unchanged and retain failed-audio evidence

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

.buildkite/cuda/test-nightly.yml

@@ -156,6 +156,11 @@ steps:
       - label: ":full_moon: TTS · Function Test with L4"
         source_file_dependencies: tts_qwen3_tts_function
         timeout_in_minutes: 120
+        env:
+          VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR: .
+        artifact_paths:
+          - failed_speech_*.wav
+          - failed_speech_*.json
         commands:
           - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
           - export VLLM_USE_DEEP_GEMM="0"

docs/user_guide/examples/offline_inference/text_to_speech.md

@@ -130,7 +130,7 @@ python examples/offline_inference/text_to_speech/cosyvoice3/end2end.py \
 ### Notes
 
 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
-- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
+- H200 automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when the CUDA MPS control tool is available; other devices and environments use `cosyvoice3.yaml`. See the [serving defaults and tradeoffs](../online_serving/text_to_speech.md#cosyvoice3). Pass `--deploy-config cosyvoice3.yaml` for the previous profile or `--deploy-config <path>` to override. The example uses the resolved profile's sampling parameters.
 - `async_chunk: true` is the default; pass `--no-async-chunk` to switch to the legacy synchronous path.
 
 ---

docs/user_guide/examples/online_serving/text_to_speech.md

@@ -155,16 +155,22 @@ vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remot
 ./cosyvoice3/run_server.sh
 ```
 
-Streaming is on by default via `async_chunk: true` in `vllm_omni/deploy/cosyvoice3.yaml`. Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path.
+Streaming is on by default. On CUDA Hopper GPUs with at least 140 GiB of device memory (H200), serving automatically selects `cosyvoice3_packed_streaming_optimized_standard.yaml` when `nvidia-cuda-mps-control` is available on `PATH`: AR32 / codec32, standard sampling, packed Flow, cached ISTFT, HiFT decode graphs, and local CUDA MPS. No `--hf-overrides` is needed. Other devices and environments without the MPS control tool use `cosyvoice3.yaml` with RAS sampling and AR8 / codec8.
 
-Cross-request Stage-1 flow batching is opt-in. Enable it in the server environment when concurrent requests should share a flow-estimator call:
+The optimized profile is not a universal hardware configuration. Its packed kernels are Hopper-only; the talker's Model Runner V2 is unsupported on NPU/XPU, and CUDA MPS requires a local NVIDIA GPU and its control tool. Larger batches, 25-frame chunks, and graph allocations also require per-device validation on smaller or slower GPUs. Automatic selection retains the generic profile outside the validated device class; an explicit deploy configuration takes precedence.
+
+The H200 default targets streaming throughput at C32/C64. It changes sampling from RAS to standard (temperature 0.7, top-p 0.8, top-k 20, repetition penalty 1.21); output lengths and quality can differ. HiFT captures a recurring exact shape on its third use, which adds latency to that request, and its graph pool needs memory headroom beyond the engine reservations. This profile is not a claim of optimal first-audio latency at low concurrency. Pass `--deploy-config cosyvoice3.yaml` to retain the previous configuration, or `--deploy-config <path>` to select another profile.
+
+Pass `--no-async-chunk` (or `NO_ASYNC_CHUNK=1 ./cosyvoice3/run_server.sh`) for the legacy synchronous path; use `--deploy-config cosyvoice3.yaml` alongside it when selecting that path on H200.
+
+Cross-request Stage-1 flow batching is enabled in the H200 profile. For the generic configuration, enable it in the server environment when concurrent requests should share a flow-estimator call:
 
 ```bash
 export COSYVOICE3_BATCH_FLOW=1
 vllm serve FunAudioLLM/Fun-CosyVoice3-0.5B-2512 --omni --port 8091 --trust-remote-code
 ```
 
-Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
+Batching preserves output lengths and streaming cache alignment, but the different GEMM shapes can produce small waveform differences compared with processing each request separately. Use `--deploy-config cosyvoice3.yaml` and leave `COSYVOICE3_BATCH_FLOW` unset (or set it to `0`) when request-independent numerical behavior is required. Set `COSYVOICE3_BATCH_FLOW_DEBUG=1` to log the observed group-size distribution and enable detailed Stage-1 profiler scopes; diagnostics are disabled by default to avoid per-step profiling overhead.
 
 ### CLI client
 
@@ -194,7 +200,7 @@ The client supports `--api-base`, `--model`, `--text`, `--ref-audio`, `--ref-tex
 ### Notes
 
 - Stage 0 (`talker`) emits speech tokens; stage 1 (`code2wav`) runs flow matching + HiFiGAN to synthesize waveform.
-- Deploy config auto-loads from `vllm_omni/deploy/cosyvoice3.yaml` based on HF `model_type`. Pass `--deploy-config <path>` to override.
+- Deploy config auto-loads the device-specific CosyVoice3 default described above. Pass `--deploy-config <path>` to override.
 - For offline inference and the end-to-end script, see the [offline CosyVoice3 section](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_speech/README.md#cosyvoice3).
 
 ---

examples/offline_inference/text_to_speech/cosyvoice3/end2end.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 import argparse
+import copy
 import functools
 import json
 import os
@@ -10,13 +11,11 @@ from typing import Any
 
 import numpy as np
 import soundfile as sf
-from vllm import SamplingParams
 from vllm.multimodal.media.audio import load_audio
 
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
 from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
-from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
 
 # Upstream zero-shot reference clip
 ZERO_SHOT_PROMPT_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
@@ -59,8 +58,7 @@ def run_e2e():
         "--deploy-config",
         type=str,
         default=None,
-        help="Override the deploy config path. If unset, auto-loads "
-        "vllm_omni/deploy/cosyvoice3.yaml based on the HF model_type.",
+        help="Override the deploy config path. If unset, selects the bundled CosyVoice3 profile for the device.",
     )
     parser.add_argument("--text", type=str, default="Hello, this is a test of the CosyVoice system capability.")
     parser.add_argument(
@@ -105,8 +103,6 @@ def run_e2e():
         profiler_config=args.profiler_config,
     )
 
-    sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
-
     print("Model initialized. Preparing inputs...")
     ref_audio_path = args.ref_audio or _default_ref_audio()
     if not os.path.exists(ref_audio_path):
@@ -138,7 +134,7 @@ def run_e2e():
 
     print(f"Generating for prompt: {args.text}")
 
-    config = CosyVoice3Config()
+    config = omni.engine.stage_vllm_configs[0].model_config.hf_config
     tokenizer = get_qwen_tokenizer(
         token_path=args.tokenizer,
         skip_special_tokens=config.skip_special_tokens,
@@ -149,29 +145,10 @@ def run_e2e():
     min_len = int(base_len * config.min_token_text_ratio)
     max_len = int(base_len * config.max_token_text_ratio)
 
-    # Build SamplingParams for each stage (GPT, S2Mel, Vocoder)
-    gpt_sampling = SamplingParams(
-        temperature=1.0,
-        top_p=sampling_cfg["top_p"],
-        top_k=sampling_cfg["top_k"],
-        repetition_penalty=2.0,
-        min_tokens=min_len,
-        max_tokens=max_len,
-        stop_token_ids=[sampling_cfg["eos_token_id"]],
-        # allowed_token_ids=list(range(6561+3)),
-        detokenize=False,
-    )
-    # Not used
-    s2mel_sampling = SamplingParams(
-        temperature=1.0,
-        top_p=1.0,
-        top_k=-1,
-        repetition_penalty=2.0,
-        max_tokens=256,
-        detokenize=False,
-    )
-
-    sampling_params_list = [gpt_sampling, s2mel_sampling]
+    # Keep the deploy profile's sampling policy, penalties and required stops.
+    sampling_params_list = copy.deepcopy(omni.default_sampling_params_list)
+    sampling_params_list[0].max_tokens = max(1, min(2048, max_len))
+    sampling_params_list[0].min_tokens = min(max(1, min_len), sampling_params_list[0].max_tokens)
 
     profiler_enabled = args.profiler_config is not None
     if profiler_enabled:

examples/online_serving/text_to_speech/cosyvoice3/run_server.sh

@@ -1,11 +1,15 @@
 #!/bin/bash
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Launch vLLM-Omni server for CosyVoice3 TTS
 #
 # Usage:
 #   ./run_server.sh
 #   CUDA_VISIBLE_DEVICES=0 ./run_server.sh
 #
-# Streaming (async-chunk) is on by default via vllm_omni/deploy/cosyvoice3.yaml.
+# Streaming (async-chunk) is on by default. H200 selects the optimized profile;
+# other devices use vllm_omni/deploy/cosyvoice3.yaml.
 # Set NO_ASYNC_CHUNK=1 to use the legacy synchronous path.
 
 set -e

tests/config/test_config_factory.py

@@ -571,6 +571,56 @@ class TestPipelineDiscovery:
         assert "qwen3_omni_moe_thinker_only" in OMNI_PIPELINES
         assert "qwen3_tts" in OMNI_PIPELINES
 
+    @pytest.mark.parametrize(
+        "cuda,major,memory_gib,mps,optimized",
+        [
+            (False, None, 0, True, False),
+            (True, None, 141, True, False),
+            (True, 8, 141, True, False),
+            (True, 9, 80, True, False),
+            (True, 10, 180, True, False),
+            (True, 9, 141, False, False),
+            (True, 9, 141, True, True),
+        ],
+    )
+    def test_cosyvoice3_device_default_and_explicit_overrides(
+        self, monkeypatch, cuda, major, memory_gib, mps, optimized
+    ):
+        from vllm.platforms import current_platform
+        from vllm.platforms.interface import DeviceCapability
+
+        from vllm_omni.platforms import current_omni_platform
+
+        monkeypatch.setattr(current_platform, "is_cuda", lambda: cuda)
+        monkeypatch.setattr(current_omni_platform, "device_name", "cuda" if cuda else "cpu")
+        capability = DeviceCapability(major, 0) if major is not None else None
+        monkeypatch.setattr(current_platform, "get_device_capability", lambda: capability)
+        monkeypatch.setattr(current_platform, "get_device_total_memory", lambda: memory_gib * 1024**3)
+        monkeypatch.setattr(
+            "vllm_omni.model_executor.models.cosyvoice3.pipeline.shutil.which",
+            lambda command: "/usr/bin/nvidia-cuda-mps-control" if mps else None,
+        )
+        pipeline = resolve_pipeline_config("cosyvoice3")
+        assert pipeline is not None
+        config = VllmOmniConfig.from_pipeline_config(pipeline)
+        talker, codec = config.stage_configs
+        assert talker.scheduler_config.max_num_seqs == codec.scheduler_config.max_num_seqs == (32 if optimized else 8)
+        assert talker.runtime_config.cuda_mps == codec.runtime_config.cuda_mps == optimized
+        assert talker.model_config.hf_overrides == ({"cosyvoice3_sampling_mode": "standard"} if optimized else None)
+        assert pipeline.stages[0].sampling_constraints["stop_token_ids"] == list(range(6561, 6761))
+        if optimized:
+            assert codec.runtime_config.env["COSYVOICE3_CACHED_ISTFT"] == "1"
+            assert codec.runtime_config.env["COSYVOICE3_HIFT_GRAPH"] == "1"
+            override = VllmOmniConfig.from_pipeline_config(
+                pipeline, cli_overrides={"hf_overrides": {"cosyvoice3_sampling_mode": "ras", "seed": 7}}
+            )
+            assert override.stage_by_id(0).model_config.hf_overrides == {"cosyvoice3_sampling_mode": "ras", "seed": 7}
+        explicit = VllmOmniConfig.from_pipeline_config(
+            pipeline, deploy_config_path=get_deploy_config_path("cosyvoice3.yaml")
+        )
+        assert explicit.stage_by_id(0).scheduler_config.max_num_seqs == 8
+        assert explicit.stage_by_id(0).model_config.hf_overrides is None
+
     def test_registry_resolver_qwen3_omni_all_stages(self):
         """Test that providing the HF config for qwen3 omni with audio enabled uses all stages."""
         pipeline = resolve_pipeline_config(

tests/helpers/assertions.py

@@ -4,9 +4,11 @@
 """Assertion and response validation helpers for tests."""
 
 import base64
+import hashlib
 import io
 import json
 import math
+import os
 import tempfile
 import threading
 import wave
@@ -999,6 +1001,31 @@ def _assert_transcript_matches(
     )
 
 
+def _retain_failed_speech_audio(
+    audio_bytes: bytes | None, transcript: str, expected_text: Any, error: AssertionError
+) -> None:
+    """Retain the exact ASR input when a CI speech-content assertion fails."""
+    directory = os.environ.get("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR")
+    if not directory or not audio_bytes:
+        return
+    stem = f"failed_speech_{hashlib.sha256(audio_bytes).hexdigest()[:16]}"
+    root = Path(directory)
+    try:
+        root.mkdir(parents=True, exist_ok=True)
+        (root / f"{stem}.wav").write_bytes(audio_bytes)
+        (root / f"{stem}.json").write_text(
+            json.dumps(
+                {"transcript": transcript, "expected_text": str(expected_text), "assertion": str(error)},
+                ensure_ascii=False,
+                indent=2,
+            )
+        )
+        print(f"Retained failed speech ASR input: {root / f'{stem}.wav'}")
+    except OSError as exc:
+        # Diagnostic I/O must never replace or suppress the content failure.
+        print(f"Could not retain failed speech ASR input: {exc}")
+
+
 def assert_audio_speech_response(response: Any, request_config: dict[str, Any], run_level: str | None = None) -> None:
     """Validate speech API results from :class:`~tests.helpers.client.OmniResponse`.
 
@@ -1063,14 +1090,21 @@ def assert_audio_speech_response(response: Any, request_config: dict[str, Any],
             if expected_text:
                 print(f"audio content is: {transcript}")
                 print(f"input text is: {expected_text}")
-                _assert_transcript_matches(
-                    transcript,
-                    _speech_audio_for_transcription(getattr(response, "audio_bytes", None), request_config),
-                    expected_text,
-                    threshold=0.9,
-                    escalation_model=request_config.get("transcript_escalation_model"),
-                    language=request_config.get("transcript_language"),
+                transcription_audio = _speech_audio_for_transcription(
+                    getattr(response, "audio_bytes", None), request_config
                 )
+                try:
+                    _assert_transcript_matches(
+                        transcript,
+                        transcription_audio,
+                        expected_text,
+                        threshold=0.9,
+                        escalation_model=request_config.get("transcript_escalation_model"),
+                        language=request_config.get("transcript_language"),
+                    )
+                except AssertionError as exc:
+                    _retain_failed_speech_audio(transcription_audio, transcript, expected_text, exc)
+                    raise
         _assert_preset_voice_gender_from_audio(
             response.audio_bytes,
             request_config.get("voice"),

tests/helpers/tests/test_assertions.py

@@ -2,7 +2,9 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import io
+import json
 import wave
+from contextlib import nullcontext
 from types import SimpleNamespace
 
 import numpy as np
@@ -249,6 +251,85 @@ def test_pcm_transcript_escalation_uses_same_valid_wav_and_rejects_two_wrong_tra
         _assert_wav_wraps_pcm(wav_bytes, pcm_bytes, 24_000)
 
 
+@pytest.mark.parametrize("response_format", ["wav", "pcm"])
+@pytest.mark.parametrize("escalation_model", [None, "large-v3"])
+def test_failed_speech_capture_retains_exact_asr_input_and_original_failure(
+    monkeypatch, tmp_path, response_format, escalation_model
+):
+    pcm_bytes = _pcm_sine()
+    wav = io.BytesIO()
+    with wave.open(wav, "wb") as writer:
+        writer.setnchannels(1)
+        writer.setsampwidth(2)
+        writer.setframerate(24_000)
+        writer.writeframes(pcm_bytes)
+    audio_bytes = wav.getvalue() if response_format == "wav" else pcm_bytes
+    transcribed = []
+
+    def transcribe(audio_bytes, **kwargs):
+        transcribed.append(audio_bytes)
+        return "this is unrelated speech"
+
+    directory = tmp_path / "failed_audio"
+    monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory))
+    monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", transcribe)
+    config = {
+        "input": "please read this requested sentence accurately",
+        "response_format": response_format,
+        "transcript_pcm_sample_rate": 24_000,
+        "transcript_escalation_model": escalation_model,
+    }
+    response = OmniResponse(success=True, audio_bytes=audio_bytes, audio_format=f"audio/{response_format}")
+    with pytest.raises(AssertionError, match="Transcript doesn't match input") as error:
+        assert_audio_speech_response(response, config, "full_model")
+
+    clips = list(directory.glob("*.wav"))
+    assert len(clips) == 1
+    assert all(clip == clips[0].read_bytes() for clip in transcribed)
+    _assert_wav_wraps_pcm(clips[0].read_bytes(), pcm_bytes, 24_000)
+    metadata = json.loads(clips[0].with_suffix(".json").read_text())
+    assert metadata["expected_text"] == config["input"]
+    assert metadata["transcript"] == "this is unrelated speech"
+    assert metadata["assertion"] == str(error.value)
+
+
+@pytest.mark.parametrize("capture_enabled", [False, True])
+@pytest.mark.parametrize("content_failure", [False, True])
+def test_speech_capture_is_opt_in_and_does_not_write_successes(monkeypatch, tmp_path, capture_enabled, content_failure):
+    directory = tmp_path / "failed_audio"
+    if capture_enabled:
+        monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory))
+    else:
+        monkeypatch.delenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", raising=False)
+    expected = "please read this requested sentence accurately"
+    monkeypatch.setattr(
+        assertions,
+        "convert_audio_bytes_to_text",
+        lambda *_args, **_kwargs: "this is unrelated speech" if content_failure else expected,
+    )
+    with pytest.raises(AssertionError, match="Transcript doesn't match input") if content_failure else nullcontext():
+        assert_audio_speech_response(
+            OmniResponse(success=True, audio_bytes=b"audio", audio_format="audio/wav"),
+            {"input": expected, "response_format": "wav"},
+            "full_model",
+        )
+    assert directory.exists() is (capture_enabled and content_failure)
+    assert len(list(tmp_path.rglob("*.wav"))) == int(capture_enabled and content_failure)
+
+
+def test_speech_capture_io_failure_does_not_mask_content_failure(monkeypatch, tmp_path):
+    directory = tmp_path / "blocked"
+    directory.write_text("a file cannot be an output directory")
+    monkeypatch.setenv("VLLM_OMNI_FAILED_SPEECH_AUDIO_DIR", str(directory))
+    monkeypatch.setattr(assertions, "convert_audio_bytes_to_text", lambda *_args, **_kwargs: "this is unrelated speech")
+    with pytest.raises(AssertionError, match="Transcript doesn't match input"):
+        assert_audio_speech_response(
+            OmniResponse(success=True, audio_bytes=b"audio", audio_format="audio/wav"),
+            {"input": "please read this requested sentence accurately", "response_format": "wav"},
+            "full_model",
+        )
+
+
 @pytest.mark.parametrize(
     "pcm_bytes",
     [

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -928,4 +928,51 @@ def test_talker_output_contract(monkeypatch, mrv2, packed):
         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)
+    assert output is hidden
+
+
+def test_request_conditioning_normalizes_singletons_without_changing_schema():
+    from vllm_omni.data_entry_keys import to_struct
+    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning
+
+    token = torch.tensor([[7, 8]], dtype=torch.int32)
+    feat = torch.randn(1, 4, 80)
+    speaker = torch.randn(1, 192)
+    length = torch.tensor([[2]], dtype=torch.int32)
+    raw = {"embed": {"speech_token": token, "speech_feat": feat, "embedding": speaker, "speech_token_len": [length]}}
+    result = to_struct(_normalize_request_conditioning(raw))
+    assert result.embed.speech_token is token
+    assert result.embed.speech_feat is feat
+    assert result.embed.embedding is speaker
+    assert result.embed.speech_token_len is length
+    assert raw["embed"]["speech_token_len"] == [length]
+
+
+def test_request_conditioning_rejects_unsplit_voices():
+    from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3 import _normalize_request_conditioning
+
+    with pytest.raises(ValueError, match="unsplit batch"):
+        _normalize_request_conditioning({"embed": {"embedding": [torch.ones(1, 192), torch.zeros(1, 192)]}})
+
+
+def test_request_conditioning_reaches_code2wav_with_true_prompt_length():
+    model = _make_code2wav_model()
+    output = model.forward(
+        torch.tensor([1, 2]),
+        torch.arange(2),
+        model_intermediate_buffer=[
+            {
+                "embed": {
+                    "speech_token": torch.tensor([[3, 2, 0, 0]], dtype=torch.int32),
+                    "speech_feat": torch.randn(1, 8, 80),
+                    "embedding": torch.randn(1, 192),
+                    "speech_token_len": [torch.tensor([[2]], dtype=torch.int32)],
+                },
+                "meta": {"req_id": ["voice"]},
+            }
+        ],
+    )
+    call = model.code2wav.forward_calls[0]
+    assert call["prompt_token"].shape == (1, 2)
+    assert call["prompt_feat"].shape == (1, 4, 80)
+    assert output.multimodal_outputs

tests/model_executor/models/qwen3_tts/test_qwen3_tts_prompt_text_batch.py

@@ -0,0 +1,140 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+    Qwen3TTSConfig,
+    Qwen3TTSTalkerConfig,
+)
+from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+    PRECOMPUTED_TEXT_IDS_KEY,
+    Qwen3TTSPromptEmbedsBuilder,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_batch_preprocess_projects_new_non_streaming_texts_once():
+    torch.manual_seed(0)
+    text_embedding = torch.nn.Embedding(64, 4)
+    projection = torch.nn.Linear(4, 4)
+    projected_shapes: list[tuple[int, ...]] = []
+
+    def text_projection(x: torch.Tensor) -> torch.Tensor:
+        projected_shapes.append(tuple(x.shape))
+        return projection(x)
+
+    builder = Qwen3TTSPromptEmbedsBuilder.__new__(Qwen3TTSPromptEmbedsBuilder)
+    builder._text_embedding = text_embedding
+    builder._text_projection = text_projection
+    builder._batched_text_embeds = {}
+
+    ids_a = list(range(1, 13))
+    ids_b = list(range(20, 31))
+    # Serving stores the assistant-template ids wrapped in a one-element list.
+    serving_ids = [ids_a]
+    buf: dict[str, dict[str, Any]] = {
+        "a": {"req_id": "a", "task_type": ["CustomVoice"], PRECOMPUTED_TEXT_IDS_KEY: [ids_a]},
+        "b": {"req_id": "b", "task_type": ["VoiceDesign"], PRECOMPUTED_TEXT_IDS_KEY: [torch.tensor(ids_b)]},
+        "streaming": {
+            "req_id": "streaming",
+            "task_type": ["CustomVoice"],
+            "non_streaming_mode": [False],
+            PRECOMPUTED_TEXT_IDS_KEY: serving_ids,
+        },
+        "base": {"req_id": "base", "task_type": ["Base"], PRECOMPUTED_TEXT_IDS_KEY: serving_ids},
+        "built": {
+            "req_id": "built",
+            "task_type": ["CustomVoice"],
+            "embed": {"prefill": torch.zeros(1)},
+            PRECOMPUTED_TEXT_IDS_KEY: serving_ids,
+        },
+    }
+
+    builder.preprocess_infos_batch(req_infos=list(buf.values()), device=torch.device("cpu"))
+
+    # One embedding + projection over both requests' text tokens (template stripped).
+    assert projected_shapes == [(1, (len(ids_a) - 8) + (len(ids_b) - 8), 4)]
+    assert set(builder._batched_text_embeds) == {"a", "b"}
+    with torch.no_grad():
+        for req_id, ids in (("a", ids_a), ("b", ids_b)):
+            expected = projection(text_embedding(torch.tensor([ids[3:-5]])))
+            torch.testing.assert_close(builder._batched_text_embeds[req_id], expected)
+            assert buf[req_id][PRECOMPUTED_TEXT_IDS_KEY].tolist() == [ids]
+    for skipped in ("streaming", "base", "built"):
+        assert buf[skipped][PRECOMPUTED_TEXT_IDS_KEY] is serving_ids
+
+
+class _Tokenizer:
+    def __call__(self, text: str, **kwargs: Any) -> dict[str, torch.Tensor]:
+        return {"input_ids": torch.tensor([[2, 3, 4, 5]])}
+
+
+def _prompt_builder() -> Qwen3TTSPromptEmbedsBuilder:
+    """Use the real constructor and prompt assembly with small embedding tables."""
+    torch.manual_seed(42)
+    config = Qwen3TTSConfig(tts_bos_token_id=50, tts_eos_token_id=51, tts_pad_token_id=52)
+    talker_config = Qwen3TTSTalkerConfig(
+        codec_nothink_id=10,
+        codec_think_id=11,
+        codec_think_bos_id=12,
+        codec_think_eos_id=13,
+        codec_pad_id=14,
+        codec_bos_id=15,
+        codec_language_id={"english": 16},
+        spk_id={"vivian": 17, "serena": 18},
+    )
+    builder = Qwen3TTSPromptEmbedsBuilder(
+        config=config,
+        talker_config=talker_config,
+        model_path="",
+        text_embedding=torch.nn.Embedding(64, 4),
+        text_projection=torch.nn.Linear(4, 4),
+        codec_embed=torch.nn.Embedding(64, 4),
+        residual_code_embeddings=lambda: [],
+        speaker_encoder=torch.nn.Identity(),
+        tts_pad_embed=torch.zeros(1, 4),
+        encode_ref_audio_batch=lambda *args, **kwargs: [],
+    )
+    builder._text_tokenizer = _Tokenizer()
+    return builder
+
+
+@pytest.mark.parametrize("task_type", ["CustomVoice", "VoiceDesign"])
+@pytest.mark.parametrize("non_streaming", [False, True])
+def test_batched_prompt_consumption_matches_serial_and_clears_cache(task_type: str, non_streaming: bool):
+    batched, serial = _prompt_builder(), _prompt_builder()
+
+    def request(req_id: str, ids: list[int], speaker: str) -> dict[str, Any]:
+        return {
+            "req_id": req_id,
+            "text": ["hello"],
+            "task_type": [task_type],
+            "language": ["English"],
+            "speaker": [speaker],
+            "instruct": ["Speak calmly"],
+            "non_streaming_mode": [non_streaming],
+            PRECOMPUTED_TEXT_IDS_KEY: [ids],
+        }
+
+    with torch.inference_mode():
+        # The same request id can be used again after its earlier prompt was consumed.
+        for ids_a, ids_b in [(list(range(12)), list(range(20, 31))), (list(range(10, 24)), list(range(30, 46)))]:
+            pending = [request("a", ids_a, "Vivian"), request("b", ids_b, "Serena")]
+            expected = [request("a", ids_a, "Vivian"), request("b", ids_b, "Serena")]
+            batched.preprocess_infos_batch(req_infos=pending, device=torch.device("cpu"))
+            assert set(batched._batched_text_embeds) == ({"a", "b"} if non_streaming else set())
+            for actual_info, expected_info in zip(pending, expected, strict=True):
+                actual = batched.build_prompt_embeds(task_type=task_type, info_dict=actual_info)
+                reference = serial.build_prompt_embeds(task_type=task_type, info_dict=expected_info)
+                torch.testing.assert_close(actual[0], reference[0])
+                torch.testing.assert_close(actual[1], reference[1])
+                assert actual[2:] == reference[2:]
+                assert PRECOMPUTED_TEXT_IDS_KEY not in actual_info
+            assert not batched._batched_text_embeds
+        batched.preprocess_infos_batch(req_infos=[], device=torch.device("cpu"))
+        assert not batched._batched_text_embeds

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -13,6 +13,7 @@ from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
     text2flow_full_payload,
     text2flow_token_only,
 )
+from vllm_omni.outputs.mm_outputs import MultimodalPayload
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -121,6 +122,24 @@ def test_text2flow_full_payload_does_not_send_codec_ids():
     assert torch.equal(payload["embed"]["speech_token"], torch.tensor([[1, 2]], dtype=torch.long))
 
 
+def test_text2flow_full_payload_accepts_structured_multimodal_payload():
+    speech_token = torch.tensor([[1, 2]], dtype=torch.long)
+    payload = text2flow_full_payload(
+        None,
+        MultimodalPayload(
+            tensors={
+                "embed.speech_token": speech_token,
+                "codes.audio": torch.tensor([7, 8, 9], dtype=torch.long),
+            }
+        ),
+        SimpleNamespace(external_req_id="req-structured"),
+    )
+
+    assert payload is not None
+    assert "codes" not in payload
+    assert torch.equal(payload["embed"]["speech_token"], speech_token)
+
+
 def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
     transfer_manager = _transfer_manager()
     request = SimpleNamespace(
@@ -355,3 +374,36 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio(
     assert payload_final is not None
     assert payload_final.meta.finished.item() is True
     assert payload_final.codes.audio.tolist() == []
+
+
+def test_full_payload_retains_true_length_for_padded_conditioning():
+    from vllm_omni.data_entry_keys import to_struct
+
+    token = torch.tensor([[7, 8, 0, 0]], dtype=torch.int32)
+    feat = torch.randn(1, 8, 80)
+    speaker = torch.randn(1, 192)
+    length = torch.tensor([[2]], dtype=torch.int32)
+    result = text2flow_full_payload(
+        None,
+        {
+            "embed.speech_token": token,
+            "embed.speech_feat": feat,
+            "embed.embedding": speaker,
+            "embed.speech_token_len": length,
+        },
+        SimpleNamespace(request_id="padded"),
+    )
+    payload = to_struct(result)
+    assert payload.embed.speech_token_len is length
+
+
+def test_reference_features_are_snapshots_under_output_accumulation():
+    from vllm_omni.outputs.mm_outputs import MultimodalPayload
+    from vllm_omni.outputs.output_modality import OutputModality
+
+    first = torch.ones(1, 174, 80)
+    latest = torch.zeros(1, 1500, 80)
+    payload = MultimodalPayload.from_dict({"embed.speech_feat": first})
+    payload = payload.merged_with(MultimodalPayload.from_dict({"embed.speech_feat": latest}))
+    payload.consolidate_tensors(OutputModality.LATENT)
+    assert payload["embed.speech_feat"] is latest

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -971,12 +971,14 @@ def test_cosyvoice3_text2flow_full_payload_nested_fallback() -> None:
 
 
 def test_cosyvoice3_full_payload_replace_keys_present() -> None:
-    """Confirm _FULL_PAYLOAD_REPLACE_KEYS lists the three embed.* keys."""
+    """Reference tensors and their length travel as complete snapshots."""
     from vllm_omni.model_executor.stage_input_processors.cosyvoice3 import (
         _FULL_PAYLOAD_REPLACE_KEYS,
     )
 
-    assert _FULL_PAYLOAD_REPLACE_KEYS == frozenset({"embed.speech_token", "embed.speech_feat", "embed.embedding"})
+    assert _FULL_PAYLOAD_REPLACE_KEYS == frozenset(
+        {"embed.speech_token", "embed.speech_token_len", "embed.speech_feat", "embed.embedding"}
+    )
 
 
 def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None:

tests/model_executor/test_seeded_exponential_platform.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.utils import seeded_exponential
+
+pytestmark = [pytest.mark.core_model]
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("is_nvidia", [False, True], ids=["rocm", "cuda"])
+def test_cuda_tensor_requires_nvidia_distribution_contract(monkeypatch, is_nvidia: bool) -> None:
+    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: is_nvidia)
+    # HIP tensors expose the same Torch device predicate as NVIDIA tensors.
+    q = Mock(spec=torch.Tensor)
+    q.is_cuda = True
+    q.dtype = torch.float32
+    q.dim.return_value = 2
+    q.is_contiguous.return_value = True
+    generators = {0: object(), 1: object()}
+
+    assert seeded_exponential.batched_seeded_exponential_supported(q, generators) is is_nvidia
+
+
+@pytest.mark.cpu
+def test_shared_generator_keeps_ordered_torch_draws(monkeypatch) -> None:
+    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: True)
+    q = Mock(spec=torch.Tensor)
+    q.is_cuda = True
+    q.dtype = torch.float32
+    q.dim.return_value = 2
+    q.is_contiguous.return_value = True
+    shared = object()
+
+    assert not seeded_exponential.batched_seeded_exponential_supported(q, {0: shared, 1: shared})
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.parametrize("seeded_rows", [(0, 1, 2), (0, 2)], ids=["all_seeded", "mixed"])
+@pytest.mark.parametrize("use_fp64_gumbel", [False, True])
+def test_torch_fallback_preserves_samples_and_generator_state(monkeypatch, seeded_rows, use_fp64_gumbel) -> None:
+    from vllm.v1.sample.ops import topk_topp_sampler as sampler_ops
+
+    import vllm_omni.patch  # noqa: F401
+
+    original = sampler_ops.random_sample.__wrapped__
+    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: False)
+
+    def reject_kernel(*args, **kwargs):
+        raise AssertionError("the CUDA distribution kernel must not run on the fallback path")
+
+    monkeypatch.setattr(seeded_exponential, "fill_exponential_rows", reject_kernel)
+    probs = torch.arange(1, 258, device="cuda", dtype=torch.float32).expand(3, -1).contiguous()
+    probs /= probs.sum(dim=-1, keepdim=True)
+    expected_gens = {row: torch.Generator(device="cuda").manual_seed(42 + row) for row in seeded_rows}
+    actual_gens = {row: torch.Generator(device="cuda").manual_seed(42 + row) for row in seeded_rows}
+    expected_default = torch.cuda.get_rng_state()
+    actual_default = expected_default.clone()
+    try:
+        for _ in range(5):
+            torch.cuda.set_rng_state(expected_default)
+            # FP32 sampling divides probabilities in place. Give both paths
+            # the same untouched input on every step.
+            expected = original(probs.clone(), expected_gens, use_fp64_gumbel)
+            expected_default = torch.cuda.get_rng_state()
+            torch.cuda.set_rng_state(actual_default)
+            actual = sampler_ops.random_sample(probs.clone(), actual_gens, use_fp64_gumbel)
+            actual_default = torch.cuda.get_rng_state()
+            assert torch.equal(actual, expected)
+            assert torch.equal(actual_default, expected_default)
+            for row in seeded_rows:
+                assert torch.equal(actual_gens[row].get_state(), expected_gens[row].get_state())
+    finally:
+        torch.cuda.set_rng_state(expected_default)
+
+
+@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
+@pytest.mark.parametrize("rows", [None, [0, 2, 4]], ids=["grouped_codebooks", "selected_rows"])
+def test_direct_fill_fallback_preserves_ordered_draws(monkeypatch, rows) -> None:
+    monkeypatch.setattr(seeded_exponential.current_platform, "is_cuda", lambda: False)
+
+    def reject_kernel(*args, **kwargs):
+        raise AssertionError("the NVIDIA kernel must not run for direct fallback callers")
+
+    monkeypatch.setattr(seeded_exponential, "torch_exponential_policy", reject_kernel)
+    expected = torch.ones(6, 257, device="cuda")
+    actual = expected.clone()
+    expected_gen = torch.Generator(device="cuda").manual_seed(71)
+    actual_gen = torch.Generator(device="cuda").manual_seed(71)
+    # A shared seeded generator and an unseeded request make loop order and
+    # default-state advancement observable, including grouped codebooks.
+    expected_gens = [expected_gen, None, expected_gen]
+    actual_gens = [actual_gen, None, actual_gen]
+    original_default = torch.cuda.get_rng_state()
+    expected_default = original_default.clone()
+    actual_default = original_default.clone()
+    try:
+        for _ in range(5):
+            torch.cuda.set_rng_state(expected_default)
+            targets = expected.reshape(3, -1) if rows is None else [expected[row] for row in rows]
+            for target, generator in zip(targets, expected_gens):
+                target.exponential_(generator=generator)
+            expected_default = torch.cuda.get_rng_state()
+            torch.cuda.set_rng_state(actual_default)
+            assert seeded_exponential.fill_exponential_rows(actual, actual_gens, rows) is actual
+            actual_default = torch.cuda.get_rng_state()
+            assert torch.equal(actual, expected)
+            assert torch.equal(actual_gen.get_state(), expected_gen.get_state())
+            assert torch.equal(actual_default, expected_default)
+    finally:
+        torch.cuda.set_rng_state(original_default)

tests/worker_v2/test_omni_model_state.py

@@ -24,6 +24,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 class _DummyInputBatch:
     is_prefilling_np: np.ndarray
+    req_ids: list[str]
 
     input_ids: SimpleNamespace
     query_start_loc: torch.Tensor
@@ -701,6 +702,58 @@ def test_publish_sampled_embeddings_is_opt_in() -> None:
     assert state.publish_sampled_embeddings(SimpleNamespace(num_reqs=1), torch.tensor([[1]])) is None
 
 
+def test_split_settled_rows_matches_per_row_check():
+    batch = _DummyInputBatch(np.array([3, 0, 2, 1, 4]))
+    batch.req_ids = ["c", "a", "x", "b", "e"]
+    batch.num_scheduled_tokens = [1, 1, 1, 4, 1]
+    batch.query_start_loc_np = [0, 1, 2, 3, 7]
+    req_indices = batch.idx_mapping_np.tolist()
+    # Slot 2 still names a finished request, slot 1 schedules a prefill
+    # chunk, slot 4 is not settled and slot 7 is not scheduled.
+    settled = {3: "c", 0: "a", 2: "old", 1: "b", 7: "z"}
+    split = OmniModelState._split_settled_rows(batch, req_indices, settled)
+    assert split is not None
+    settled_rows, remaining = split
+    assert settled_rows == [(0, 3, 0, "c"), (1, 0, 1, "a")]
+    assert remaining == [(2, 2), (3, 1), (4, 4)]
+    assert OmniModelState._split_settled_rows(batch, req_indices, {2: "old"}) is None
+
+
+def test_settled_row_split_preserves_mixed_batch_and_pending_replay(monkeypatch):
+    """Exercise run_preprocess, including the scalar fallback during replay."""
+
+    def prepare():
+        state = _make_eager_state()
+        _fill_buffers(state, "prefill", "decode")
+        state._eager_fastpath = True
+        state._eager_settled = {1: "decode"}
+        state._eager_ready = {1: "decode"}
+        state._eager_embeds[1].fill_(7)
+        state.model.eager_settled_text_step = lambda: torch.ones(_EAGER_DIM)
+        state.model.eager_decode_settled = lambda info: False
+        state.model.preprocess = lambda input_ids, input_embeds, **info: (input_ids, input_embeds, {})
+        batch = _EagerBatch([3, 1])
+        batch.req_ids = ["prefill", "decode"]
+        inputs = {"input_ids": torch.zeros(4, dtype=torch.long), "inputs_embeds": torch.zeros(4, _EAGER_DIM)}
+        return state, batch, inputs
+
+    vector, batch, inputs = prepare()
+    vector.run_preprocess(batch, inputs)
+    reference, ref_batch, ref_inputs = prepare()
+    monkeypatch.setattr(reference, "_split_settled_rows", lambda *args: None)
+    reference.run_preprocess(ref_batch, ref_inputs)
+    torch.testing.assert_close(inputs["inputs_embeds"], ref_inputs["inputs_embeds"], rtol=0, atol=0)
+    assert vector._eager_rows[1] == reference._eager_rows[1]
+    assert inputs["inputs_embeds"][3].tolist() == [8.0] * _EAGER_DIM
+
+    replay, replay_batch, replay_inputs = prepare()
+    replay._eager_state._restore_audio["prefill"] = [torch.ones(1)]
+    assert replay._eager_state.has_pending_replay()
+    monkeypatch.setattr(replay, "_split_settled_rows", MagicMock(side_effect=AssertionError("replay must skip split")))
+    replay.run_preprocess(replay_batch, replay_inputs)
+    torch.testing.assert_close(replay_inputs["inputs_embeds"], ref_inputs["inputs_embeds"], rtol=0, atol=0)
+
+
 @pytest.mark.parametrize("prefilling", [False, True])
 def test_identity_preprocess_skips_only_decode_rows(prefilling):
     state = _make_state(has_preprocess=True)

vllm_omni/config/pipeline_registry.py

@@ -54,7 +54,7 @@ from vllm_omni.model_executor.models.cosmos3.pipeline import (
     COSMOS3_OMNI_DEPLOY_PIPELINE,
     COSMOS3_POLICY_PIPELINE,
 )
-from vllm_omni.model_executor.models.cosyvoice3.pipeline import COSYVOICE3_PIPELINE
+from vllm_omni.model_executor.models.cosyvoice3.pipeline import resolve_cosyvoice3_pipeline
 from vllm_omni.model_executor.models.covo_audio.pipeline import COVO_AUDIO_PIPELINE
 from vllm_omni.model_executor.models.dreamzero.pipeline import DREAMZERO_PIPELINE
 from vllm_omni.model_executor.models.fish_speech.pipeline import FISH_SPEECH_PIPELINE
@@ -150,7 +150,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     # key (see deploy/cosmos3_policy_droid.yaml and deploy/cosmos3_omni.yaml).
     "cosmos3_omni_deploy": COSMOS3_OMNI_DEPLOY_PIPELINE,
     "cosmos3_policy": COSMOS3_POLICY_PIPELINE,
-    "cosyvoice3": COSYVOICE3_PIPELINE,
+    "cosyvoice3": resolve_cosyvoice3_pipeline,
     "covo_audio": COVO_AUDIO_PIPELINE,
     "dreamzero": DREAMZERO_PIPELINE,
     "fish_qwen3_omni": FISH_SPEECH_PIPELINE,

vllm_omni/deploy/cosyvoice3_packed_streaming_optimized_standard.yaml

@@ -1,5 +1,7 @@
-# Opt-in single-Hopper streaming throughput profile: AR32 / codec32, Flow10.
-# Use --hf-overrides '{"cosyvoice3_sampling_mode": "standard"}' and stream=true.
+# Single-Hopper streaming throughput profile: AR32 / codec32, Flow10.
+# Automatically selected on Hopper devices with at least 140 GiB (H200)
+# when nvidia-cuda-mps-control is available on PATH.
+# Standard sampling is configured below; request stream=true for first audio.
 # F0 and HiFT stay FP32; only deterministic HiFT decode is graphed.
 # HiFT graphs capture on the third use of each exact shape, adding latency then.
 # The codec graph pool is not capped by gpu_memory_utilization; leave memory headroom.
@@ -7,6 +9,8 @@ base_config: cosyvoice3_packed_streaming.yaml
 stages:
   - stage_id: 0
     max_num_seqs: 32
+    hf_overrides:
+      cosyvoice3_sampling_mode: standard
     default_sampling_params:
       max_tokens: 2048
       repetition_penalty: 1.21

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py

@@ -74,6 +74,27 @@ logger = init_logger(__name__)
 _RUNTIME_COMPONENTS_CACHE: dict[str, dict] = {}
 
 
+def _normalize_request_conditioning(payload: dict) -> dict:
+    """Unwrap singleton conditioning containers at the per-request boundary.
+
+    A processor/prefix-cache passthrough can retain a one-item tensor list.
+    The code2wav schema consumes one request, so that wrapper carries no
+    batch dimension. A multi-item list here is an unsplit batch and must
+    fail instead of silently choosing another request's voice.
+    """
+    embed = payload.get("embed")
+    if not isinstance(embed, Mapping):
+        return payload
+    normalized = dict(embed)
+    for key in ("speech_token", "speech_feat", "embedding", "speech_token_len"):
+        value = normalized.get(key)
+        if isinstance(value, (list, tuple)):
+            if len(value) != 1:
+                raise ValueError(f"CosyVoice3 per-request {key} contains an unsplit batch of {len(value)} items")
+            normalized[key] = value[0]
+    return {**payload, "embed": normalized}
+
+
 def _cosyvoice3_trt_enabled() -> bool:
     """COSYVOICE3_TRT env toggle (default on) for the optional TensorRT paths.
 
@@ -1501,9 +1522,10 @@ class CosyVoice3Model(
             # [total_tokens, hidden]
             hidden_states = self.model.llm(inputs_embeds, positions)
 
-            if getattr(self, "_mrv2_tensor_output", False) or cosyvoice3_packed_inference_enabled():
-                return hidden_states
-            return self.make_omni_output(hidden_states, **kwargs)
+            # Both runners attach the live step's conditioning after replay.
+            # Capturing an OmniOutput here would retain the dummy prefill lists
+            # and re-emit them during decode, replacing the real reference.
+            return hidden_states
         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
@@ -1539,7 +1561,7 @@ class CosyVoice3Model(
 
             for idx, req_ids in enumerate(request_ids_list):
                 raw = runtime_info[idx] if idx < len(runtime_info) and isinstance(runtime_info[idx], dict) else {}
-                payload = to_struct(raw)
+                payload = to_struct(_normalize_request_conditioning(raw))
                 meta = payload.meta
                 embed = payload.embed
 

vllm_omni/model_executor/models/cosyvoice3/pipeline.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """CosyVoice3 pipeline topology (frozen).
 
 Stage 0: Talker   — text prompt → speech tokens (LLM autoregressive).
@@ -12,11 +12,17 @@ Stage 1: Code2Wav — flow-matching decoder → acoustic features → waveform.
     through the shared-memory connector.
 """
 
+import shutil
+from dataclasses import replace
+
+from transformers import PretrainedConfig
+
 from vllm_omni.config.stage_config import (
     PipelineConfig,
     StageExecutionType,
     StagePipelineConfig,
 )
+from vllm_omni.transformers_utils.configs.cosyvoice3 import CosyVoice3Config
 
 _PROC = "vllm_omni.model_executor.stage_input_processors.cosyvoice3"
 
@@ -35,8 +41,9 @@ COSYVOICE3_PIPELINE = PipelineConfig(
             async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
             custom_process_next_stage_input_func=f"{_PROC}.text2flow_full_payload",
             sampling_constraints={
-                # merged speech stop token (logsumexp of all 200 stop logits)
-                "stop_token_ids": [6562],
+                # Standard sampling can emit any of the 200 control tokens.
+                # RAS merges their logits into 6562, which is also in this range.
+                "stop_token_ids": list(range(6561, 6761)),
             },
         ),
         StagePipelineConfig(
@@ -52,3 +59,28 @@ COSYVOICE3_PIPELINE = PipelineConfig(
         ),
     ),
 )
+
+
+def resolve_cosyvoice3_pipeline(hf_config: PretrainedConfig | None = None) -> PipelineConfig | None:
+    """Select a streaming default compatible with the device and runtime.
+
+    Query the platform through NVML without initializing CUDA in the parent.
+    Explicit deploy configs still take precedence over this pipeline default.
+    """
+    from vllm.platforms import current_platform
+
+    if hf_config is not None and not isinstance(hf_config, CosyVoice3Config):
+        return None
+    if current_platform.is_cuda():
+        capability = current_platform.get_device_capability()
+        if (
+            capability is not None
+            and capability.major == 9
+            and current_platform.get_device_total_memory() >= 140 * 1024**3
+            and shutil.which("nvidia-cuda-mps-control") is not None
+        ):
+            return replace(
+                COSYVOICE3_PIPELINE,
+                default_deploy_config_name="cosyvoice3_packed_streaming_optimized_standard.yaml",
+            )
+    return COSYVOICE3_PIPELINE

vllm_omni/model_executor/models/moss_tts/reference_encoder_graphs.py

@@ -65,6 +65,24 @@ def _cached_reference_rope(self, q: torch.Tensor, k: torch.Tensor, cache: dict):
     return qo.view(batch, heads, seqlen, dim), ko.view(batch, heads, seqlen, dim)
 
 
+def _reference_flash_attn_varlen_func():
+    """Resolve the platform's variable-length FlashAttention entrypoint.
+
+    vLLM's bundled extension is NVIDIA-only. ROCm exposes the compatible
+    unpacked-QKV API through AITER, with upstream flash-attn as its fallback.
+    """
+    if torch.version.hip is not None:
+        try:
+            from aiter import flash_attn_varlen_func
+
+            return flash_attn_varlen_func, True
+        except ImportError:
+            from flash_attn import flash_attn_varlen_func
+    else:
+        from vllm.vllm_flash_attn import flash_attn_varlen_func
+    return flash_attn_varlen_func, False
+
+
 def _windowed_attention(
     self,
     x: torch.Tensor,
@@ -86,7 +104,7 @@ def _windowed_attention(
     q, k = self._apply_dense_rope(q, k) if rope_cache is None else _cached_reference_rope(self, q, k, rope_cache)
     if q.dtype not in (torch.bfloat16, torch.float16):
         return sdpa(x, input_lengths)
-    from vllm.vllm_flash_attn import flash_attn_varlen_func
+    flash_attn_varlen_func, is_aiter = _reference_flash_attn_varlen_func()
 
     heads, dim = q.shape[1], q.shape[-1]
 
@@ -95,7 +113,16 @@ def _windowed_attention(
 
     starts = torch.arange(0, (batch + 1) * seqlen, seqlen, device=x.device, dtype=torch.int32)
     window = [self.context - 1, 0] if self.causal and self.context is not None else [-1, -1]
-    out = flash_attn_varlen_func(
+    kwargs = {}
+    if torch.version.hip is None:
+        # Only vLLM's NVIDIA wrapper accepts its backend-selection argument.
+        kwargs["fa_version"] = fa_version
+    elif is_aiter and torch.is_grad_enabled() and any(t.requires_grad for t in (q, k, v)):
+        # AITER's autograd wrapper requires LSE to be retained for backward.
+        # The inference path avoids this allocation, and upstream flash-attn
+        # manages its own autograd state without this AITER-specific keyword.
+        kwargs["return_lse"] = True
+    result = flash_attn_varlen_func(
         packed(q),
         packed(k),
         packed(v),
@@ -105,8 +132,10 @@ def _windowed_attention(
         cu_seqlens_k=starts,
         causal=self.causal,
         window_size=window,
-        fa_version=fa_version,
-    ).view(batch, seqlen, heads, dim)
+        **kwargs,
+    )
+    out = result[0] if isinstance(result, tuple) else result
+    out = out.view(batch, seqlen, heads, dim)
     if not skip_padded_query_mask:
         valid = (torch.arange(seqlen, device=x.device).view(1, seqlen) < input_lengths.view(-1, 1)).view(
             batch, seqlen, 1, 1

vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py

@@ -370,6 +370,11 @@ class Qwen3TTSPromptEmbedsBuilder:
         # CustomVoice prompt pieces that do not depend on the text, per
         # (device, speaker, language): bounded by the model's voices.
         self._custom_voice_prompt_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
+        # Speaker codec id and embedding per (device, speaker).
+        self._speaker_embed_cache: dict[tuple[str, str], tuple[int, torch.Tensor]] = {}
+        # Text-conditioning embeddings projected by ``preprocess_infos_batch``
+        # for the current step's new requests, keyed by request id.
+        self._batched_text_embeds: dict[str, torch.Tensor] = {}
 
         self._ref_audio_artifact_cache_max_entries = int(ref_audio_artifact_cache_max_entries)
         self._ref_audio_artifact_cache: OrderedDict[str, dict[str, torch.Tensor | bool]] = OrderedDict()
@@ -752,6 +757,18 @@ class Qwen3TTSPromptEmbedsBuilder:
             return bool(value)
         return False
 
+    @staticmethod
+    def _non_streaming_mode(info_dict: dict[str, Any], task_type: str) -> bool:
+        value = info_dict.get("non_streaming_mode")
+        if isinstance(value, list):
+            value = value[0] if value else None
+        if isinstance(value, bool):
+            return value
+        # Match official inference defaults:
+        # - CustomVoice/VoiceDesign: non_streaming_mode=True
+        # - Base: non_streaming_mode=False
+        return task_type in ("CustomVoice", "VoiceDesign")
+
     @staticmethod
     def _needs_initial_prompt_preprocess(info_dict: dict[str, Any]) -> bool:
         embed = info_dict.get("embed")
@@ -806,15 +823,16 @@ class Qwen3TTSPromptEmbedsBuilder:
         req_infos: list[dict[str, Any]],
         device: torch.device,
     ) -> None:
-        """Batch Base voice-clone ref-audio codec extraction for current prefill requests.
+        """Batch per-request prompt work for current prefill requests.
 
-        Tokenizes assistant / ref text and runs the talker-supplied
-        ``encode_ref_audio_batch`` callable once per sample-rate group,
-        stashing the precomputed artifacts under
+        For Base voice cloning, tokenizes assistant / ref text and runs the
+        talker-supplied ``encode_ref_audio_batch`` callable once per
+        sample-rate group, stashing the precomputed artifacts under
         :data:`PRECOMPUTED_TEXT_IDS_KEY`, :data:`PRECOMPUTED_REF_IDS_KEY`,
-        :data:`PRECOMPUTED_REF_CODE_KEY`, and :data:`NORMALIZED_REF_AUDIO_KEY`
-        so the per-request :meth:`build_prompt_embeds` call can skip the
-        equivalent serial work.
+        :data:`PRECOMPUTED_REF_CODE_KEY`, and :data:`NORMALIZED_REF_AUDIO_KEY`.
+        For non-streaming CustomVoice / VoiceDesign, projects every request's
+        text at once (:meth:`_project_texts_batch`). The per-request
+        :meth:`build_prompt_embeds` call then skips the equivalent serial work.
         """
         pending_text: list[tuple[dict[str, Any], str]] = []
         pending_ref_text: list[tuple[dict[str, Any], str]] = []
@@ -919,6 +937,54 @@ class Qwen3TTSPromptEmbedsBuilder:
                 if isinstance(cache_key, str) and cache_key:
                     self.put_ref_audio_artifacts(cache_key, ref_code=ref_code_t)
 
+        self._batched_text_embeds = self._project_texts_batch(req_infos)
+
+    def _project_texts_batch(self, req_infos: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
+        """Embed and project the text of new non-streaming CustomVoice / VoiceDesign requests at once.
+
+        Each request otherwise uploads its text ids and runs the text
+        embedding + projection MLP on its own. Here every request's ids go up
+        in one pinned host-to-device copy and the MLP runs once over the
+        concatenated text tokens; the rows are independent, so each request's
+        slice equals its own per-request result. The device ids replace
+        :data:`PRECOMPUTED_TEXT_IDS_KEY` so ``build_prompt_embeds`` skips the
+        per-request copy too.
+        """
+        items: list[tuple[dict[str, Any], str, np.ndarray]] = []
+        for info_dict in req_infos:
+            req_id = info_dict.get("req_id")
+            if not isinstance(req_id, str) or not self._needs_initial_prompt_preprocess(info_dict):
+                continue
+            task_type = first_value(info_dict.get("task_type"), "CustomVoice")
+            if task_type not in ("CustomVoice", "VoiceDesign") or not self._non_streaming_mode(info_dict, task_type):
+                continue
+            ids = first_value(info_dict.get(PRECOMPUTED_TEXT_IDS_KEY))
+            if isinstance(ids, torch.Tensor) and ids.device.type == "cpu":
+                ids = ids.reshape(-1).numpy()
+            elif isinstance(ids, list) and ids and all(isinstance(v, (int, np.integer)) for v in ids):
+                ids = np.asarray(ids)
+            elif not (isinstance(ids, np.ndarray) and ids.ndim == 1):
+                continue
+            # The assistant template wraps the text in 3 leading and 5 trailing tokens.
+            if ids.size <= 8:
+                continue
+            items.append((info_dict, req_id, ids.astype(np.int64, copy=False)))
+        if not items:
+            return {}
+        full_ids = [ids for _info, _req_id, ids in items]
+        text_ids = [ids[3:-5] for ids in full_ids]
+        packed = to_device_nonblocking(torch.from_numpy(np.concatenate(full_ids + text_ids)), self._device())
+        num_full = sum(ids.size for ids in full_ids)
+        projected = self._text_projection(self._text_embedding(packed[num_full:].view(1, -1)))
+        embeds: dict[str, torch.Tensor] = {}
+        full_offset = text_offset = 0
+        for (info_dict, req_id, ids), text in zip(items, text_ids, strict=True):
+            info_dict[PRECOMPUTED_TEXT_IDS_KEY] = packed[full_offset : full_offset + ids.size].view(1, -1)
+            embeds[req_id] = projected[:, text_offset : text_offset + text.size]
+            full_offset += ids.size
+            text_offset += text.size
+        return embeds
+
     # -------------------- ICL prompt assembly --------------------
 
     def _generate_icl_prompt(
@@ -1010,18 +1076,7 @@ class Qwen3TTSPromptEmbedsBuilder:
 
         text = (info_dict.get("text") or [""])[0]
         language = (info_dict.get("language") or ["Auto"])[0]
-        non_streaming_mode_val = info_dict.get("non_streaming_mode")
-        if isinstance(non_streaming_mode_val, list):
-            non_streaming_mode_raw = non_streaming_mode_val[0] if non_streaming_mode_val else None
-        else:
-            non_streaming_mode_raw = non_streaming_mode_val
-        if isinstance(non_streaming_mode_raw, bool):
-            non_streaming_mode = non_streaming_mode_raw
-        else:
-            # Match official inference defaults:
-            # - CustomVoice/VoiceDesign: non_streaming_mode=True
-            # - Base: non_streaming_mode=False
-            non_streaming_mode = task_type in ("CustomVoice", "VoiceDesign")
+        non_streaming_mode = self._non_streaming_mode(info_dict, task_type)
 
         # Text ids for assistant template (always).
         tok = self.get_text_tokenizer()
@@ -1412,15 +1467,19 @@ class Qwen3TTSPromptEmbedsBuilder:
             )
             if not speaker:
                 raise ValueError("CustomVoice requires additional_information.speaker.")
-            spk_id_map = {k.lower(): v for k, v in (getattr(talker_config, "spk_id", None) or {}).items()}
-            if speaker not in spk_id_map:
-                raise ValueError(f"Unsupported speaker: {speaker}")
-            spk_id = spk_id_map[speaker]
-            # Keep it at least 1D; embedding on a 0-d tensor can return 1D.
-            spk_tensor = index_to_device([spk_id], input_ids.device)
-            spk_embed = codec_embed(spk_tensor)
-            if spk_embed.ndim in (1, 2):
-                spk_embed = spk_embed.view(1, 1, -1)
+            speaker_key = (str(input_ids.device), speaker)
+            cached_speaker = self._speaker_embed_cache.get(speaker_key)
+            if cached_speaker is None:
+                spk_id_map = {k.lower(): v for k, v in (getattr(talker_config, "spk_id", None) or {}).items()}
+                if speaker not in spk_id_map:
+                    raise ValueError(f"Unsupported speaker: {speaker}")
+                spk_id = spk_id_map[speaker]
+                # Keep it at least 1D; embedding on a 0-d tensor can return 1D.
+                spk_embed = codec_embed(index_to_device([spk_id], input_ids.device))
+                if spk_embed.ndim in (1, 2):
+                    spk_embed = spk_embed.view(1, 1, -1)
+                cached_speaker = self._speaker_embed_cache[speaker_key] = (spk_id, spk_embed)
+            spk_id, spk_embed = cached_speaker
             speaker_embed = spk_embed
             codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1)
 
@@ -1452,7 +1511,9 @@ class Qwen3TTSPromptEmbedsBuilder:
             talker_prompt = prefix
 
             if non_streaming_mode:
-                text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
+                text_all = self._batched_text_embeds.pop(str(info_dict.get("req_id")), None)
+                if text_all is None:
+                    text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
                 text_all = torch.cat([text_all, tts_eos_embed], dim=1)
                 talker_prompt = torch.cat([talker_prompt, text_all + pad_row, tail], dim=1)
                 trailing_text_hidden = tts_pad_embed
@@ -1477,7 +1538,9 @@ class Qwen3TTSPromptEmbedsBuilder:
             talker_prompt = torch.cat((role_embed, codec_prefix), dim=1)
 
             if non_streaming_mode:
-                text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
+                text_all = self._batched_text_embeds.pop(str(info_dict.get("req_id")), None)
+                if text_all is None:
+                    text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
                 text_all = torch.cat([text_all, tts_eos_embed], dim=1)
                 pad_ids = torch.full(
                     (1, int(text_all.shape[1])),

vllm_omni/model_executor/stage_input_processors/cosyvoice3.py

@@ -292,7 +292,9 @@ def talker2code2wav_async_chunk(
 # CONCAT across the (already trivial) per-request accumulator history so a
 # regression where decode unexpectedly re-emits them does not silently
 # duplicate the prefill tensor.  See mixin._FULL_PAYLOAD_REPLACE_KEYS.
-_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset({"embed.speech_token", "embed.speech_feat", "embed.embedding"})
+_FULL_PAYLOAD_REPLACE_KEYS: frozenset[str] = frozenset(
+    {"embed.speech_token", "embed.speech_feat", "embed.embedding", "embed.speech_token_len"}
+)
 
 
 def text2flow_token_only(
@@ -347,20 +349,20 @@ def text2flow_full_payload(
     """
     del transfer_manager
     rid = getattr(request, "external_req_id", None) or getattr(request, "request_id", "?")
-    if not isinstance(pooling_output, dict):
+    if not isinstance(pooling_output, Mapping):
         logger.warning(
-            "cosyvoice3.text2flow_full_payload: pooling_output not a dict "
+            "cosyvoice3.text2flow_full_payload: pooling_output not a mapping "
             "(type=%s) for req=%s; consumer wait gate may hang.",
             type(pooling_output).__name__,
             rid,
         )
         return None
     embed_out: dict[str, Any] = {}
-    for key in ("speech_token", "speech_feat", "embedding"):
+    for key in ("speech_token", "speech_feat", "embedding", "speech_token_len"):
         v = pooling_output.get(f"embed.{key}")
         if v is None:
             nested = pooling_output.get("embed")
-            if isinstance(nested, dict):
+            if isinstance(nested, Mapping):
                 v = nested.get(key)
         if isinstance(v, torch.Tensor) and v.numel() > 0:
             embed_out[key] = v

vllm_omni/outputs/output_modality.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Output modality types for vLLM-Omni.
 
 This module defines the OutputModality enum and TensorAccumulationStrategy
@@ -87,7 +90,7 @@ class OutputModality(Flag):
             try:
                 result |= cls[p.upper()]
             except KeyError:
-                raise ValueError(f"Unknown modality: {p!r}. Supported: {[m.name.lower() for m in cls]}")
+                raise ValueError(f"Unknown modality: {p!r}. Supported: {[(m.name or '').lower() for m in cls]}")
         return result
 
     @property
@@ -123,7 +126,14 @@ class TensorAccumulationStrategy(Enum):
 # re-emitted unchanged at every step. Pipelines that produce such keys
 # register the correct strategy here via ``register_key_accumulation_strategy``
 # instead of forcing every tensor of the modality through one default.
-_KEY_ACCUMULATION_OVERRIDES: dict[str, TensorAccumulationStrategy] = {}
+# Schema-defined prompt conditioning is a complete reference snapshot. Keep
+# these defaults in the output process too, which may never import the model's
+# worker-side stage input processor. Feature lengths can change between
+# snapshots and cannot be concatenated as generated content deltas.
+_KEY_ACCUMULATION_OVERRIDES: dict[str, TensorAccumulationStrategy] = {
+    key: TensorAccumulationStrategy.REPLACE
+    for key in ("embed.speech_token", "embed.speech_feat", "embed.embedding", "embed.speech_token_len")
+}
 
 
 def register_key_accumulation_strategy(key: str, strategy: TensorAccumulationStrategy) -> None:

vllm_omni/utils/seeded_exponential.py

@@ -12,6 +12,7 @@ every value and every generator state afterwards matches the per-request loop.
 """
 
 import torch
+from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 
 
@@ -73,10 +74,18 @@ def torch_exponential_policy(numel: int, device: torch.device) -> tuple[int, int
 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.
+    ``out`` is a contiguous float32 matrix. Without ``rows``, each generator
+    owns an equal contiguous slice of the flattened output (which may span
+    several codebook rows). ``None`` generators use the default generator in
+    order. NVIDIA uses the batched CUDA kernel; other platforms retain
+    ordered Torch draws, including direct Higgs Audio callers.
     """
+    if not current_platform.is_cuda():
+        targets = out.reshape(len(generators), -1) if rows is None else [out[row] for row in rows]
+        for target, generator in zip(targets, generators):
+            target.exponential_(generator=generator)
+        return out
+
     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)
@@ -109,9 +118,15 @@ def fill_exponential_rows(out: torch.Tensor, generators: list, rows: list[int] |
 
 
 def batched_seeded_exponential_supported(q: torch.Tensor, generators: dict) -> bool:
-    """Whether ``fill_exponential_rows`` reproduces the per-row loop for ``q``."""
+    """Whether the NVIDIA CUDA kernel reproduces the per-row loop for ``q``.
+
+    ROCm tensors also report ``is_cuda``, but HIP uses a different distribution
+    implementation and cannot compile the CUDA libdevice transform above.
+    Keep those draws on Torch's original per-row path, including RNG state.
+    """
     return (
-        len(generators) > 1
+        current_platform.is_cuda()
+        and len(generators) > 1
         and q.is_cuda
         and q.dtype == torch.float32
         and q.dim() == 2

vllm_omni/worker_v2/model_states/eager_mtp.py

@@ -111,6 +111,10 @@ class EagerMTPState:
         if self._audio_buffer is not None:
             self._audio_buffer.finish(req_ids)
 
+    def has_pending_replay(self) -> bool:
+        """Whether any preempted request still has saved inputs to replay."""
+        return bool(self._restore_audio)
+
     def replay_inputs(self, req_id: str, req_idx: int, offset: int, ids: torch.Tensor, embeds: torch.Tensor) -> bool:
         """Rebuild freed Talker KV from the exact conditioned inputs, without rerunning MTP."""
         if req_id not in self._restore_audio:
@@ -139,12 +143,15 @@ class EagerMTPState:
             return
         history_slices: list[torch.Tensor] = []
         input_slices: list[torch.Tensor] = []
+        histories = self._talker_inputs
+        on_cpu = embeds.is_cpu
         for i, req_id in enumerate(input_batch.req_ids):
             start, end = input_batch.query_start_loc_np[i : i + 2]
             offset = int(input_batch.num_computed_tokens_np[i])
             length = offset + int(end - start)
-            history = self._talker_inputs.get(req_id)
-            if history is None or history.embeds.device != embeds.device:
+            history = histories.get(req_id)
+            # A suspended request's inputs wait on the host; bring them back.
+            if history is None or history.embeds.is_cpu != on_cpu:
                 storage = torch.empty(
                     self.owner.vllm_config.model_config.max_model_len,
                     embeds.shape[-1],
@@ -405,13 +412,15 @@ class EagerMTPState:
                         p = int(ref.shape[0])
                 pos_list.append(p)
                 positions[req_id] = p + 1
-        meta_list = (
-            [i for i, _r, _q, _p in entries]
-            + [r for _i, r, _q, _p in entries]
-            + [int(p) for _i, _r, _q, p in entries]
-            + [int(q in fa_requests) for _i, _r, q, _p in entries]
-            + pos_list
-        )
+        # One transpose of the entry tuples instead of a Python pass per field.
+        rows_i, rows_req_idx, rows_req_id, rows_prefill = zip(*entries) if entries else ((), (), (), ())
+        meta_list = [
+            *rows_i,
+            *rows_req_idx,
+            *map(int, rows_prefill),
+            *[int(req_id in fa_requests) for req_id in rows_req_id],
+            *pos_list,
+        ]
         meta = index_to_device(meta_list, device)
         rows = meta[bsz : 2 * bsz]
         sampled = sampled_token_ids.reshape(input_batch.num_reqs, -1)
@@ -430,7 +439,7 @@ class EagerMTPState:
             req_indices, batch_ids, batch_emb, batch_hidden, batch_step, mtp_batch_descriptor_dispatcher
         )
         assert codes is not None
-        has_prefill = any(p for _i, _r, _q, p in entries)
+        has_prefill = any(rows_prefill)
         fa_in_kernel = isinstance(first_audio, torch.Tensor) and not has_prefill
         valid = eager_post(
             meta, bsz, last_tokens, layer0, frame_embeds.reshape(frame_embeds.shape[0], -1).contiguous(),
@@ -438,8 +447,7 @@ class EagerMTPState:
             first_audio if fa_in_kernel else None, self._first_audio_valid if fa_in_kernel else None,
             owner.model._codebook_vocab_size,
         )  # fmt: skip
-        for _i, req_idx, req_id, _prefill in entries:
-            owner._eager_ready[req_idx] = req_id
+        owner._eager_ready.update(zip(rows_req_idx, rows_req_id))
         if pos_list:
             assert stream is not None
             assert isinstance(stream_out, torch.Tensor)
@@ -452,10 +460,11 @@ class EagerMTPState:
             for t in (frame_codes, meta, last_tokens, valid, stream_out):
                 t.record_stream(side)
             with torch.cuda.stream(side):
-                for _i, req_idx, req_id, _p in entries:
-                    saved = self._restore_audio.pop(req_id, None)
-                    if saved is not None:
-                        stream.restore_slot(req_idx, saved)
+                if self._restore_audio:
+                    for _i, req_idx, req_id, _p in entries:
+                        saved = self._restore_audio.pop(req_id, None)
+                        if saved is not None:
+                            stream.restore_slot(req_idx, saved)
                 model.prime_stream_decoder(primes)
                 pcm = decode(frame_codes, rows, meta[4 * bsz : 5 * bsz])
                 stream_out.index_copy_(0, last_tokens, pcm.reshape(bsz, -1).to(stream_out.dtype))

vllm_omni/worker_v2/model_states/omni_model_state.py

@@ -16,7 +16,7 @@ from __future__ import annotations
 import inspect
 import threading
 import types
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from dataclasses import replace
 from typing import Any, cast
 
@@ -657,6 +657,44 @@ class OmniModelState(DefaultModelState):
                 else:
                     del updates[key]
 
+    @staticmethod
+    def _split_settled_rows(
+        input_batch: InputBatch,
+        req_indices: list[int],
+        settled: dict[int, str],
+    ) -> tuple[list[tuple[int, int, int, str]], list[tuple[int, int]]] | None:
+        """Vectorized settled-row test of ``run_preprocess``.
+
+        In steady decode almost every row is settled: it schedules one token
+        and ``settled`` maps its slot to its request id. ``settled`` entries
+        are dropped together with the slot buffer, so comparing against the
+        row's request id equals the per-row ``buf["req_id"]`` test. Returns the
+        settled rows and the remaining ``(row, slot)`` pairs, both in row
+        order, or ``None`` when no row is settled.
+        """
+        n = input_batch.num_reqs
+        slots = np.asarray(input_batch.idx_mapping_np[:n])
+        owners = np.full(max(int(slots.max()), max(settled)) + 1, None, dtype=object)
+        owners[np.fromiter(settled.keys(), dtype=np.int64, count=len(settled))] = np.array(
+            list(settled.values()), dtype=object
+        )
+        req_ids = np.empty(n, dtype=object)
+        req_ids[:] = input_batch.req_ids[:n]
+        is_settled = (np.asarray(input_batch.num_scheduled_tokens[:n]) == 1) & (owners[slots] == req_ids)
+        if not is_settled.any():
+            return None
+        rows = np.flatnonzero(is_settled)
+        settled_rows = list(
+            zip(
+                rows.tolist(),
+                slots[rows].tolist(),
+                np.asarray(input_batch.query_start_loc_np[:n])[rows].tolist(),
+                req_ids[rows].tolist(),
+            )
+        )
+        remaining = [(i, req_indices[i]) for i in np.flatnonzero(~is_settled).tolist()]
+        return settled_rows, remaining
+
     def run_preprocess(
         self,
         input_batch: InputBatch,
@@ -694,7 +732,7 @@ class OmniModelState(DefaultModelState):
         mtp_batches: list[tuple[int, int, tuple[torch.Tensor, torch.Tensor]]] = []
         prepacked_mtp_inputs: tuple[torch.Tensor, torch.Tensor] | None = None
 
-        req_indices = [int(input_batch.idx_mapping_np[i]) for i in range(input_batch.num_reqs)]
+        req_indices = np.asarray(input_batch.idx_mapping_np[: input_batch.num_reqs]).tolist()
         self._stage_batched_preprocess_inputs(req_indices, embeds.device)
 
         preprocess_entries: list[tuple[int, int, int, int, dict[str, Any], bool]] = []
@@ -704,7 +742,12 @@ class OmniModelState(DefaultModelState):
         # Declared identity decode has no per-row updates or hook work.
         skip_decode_rows = self._decode_preprocess_is_identity and not self._eager_mtp
         is_prefilling_np = getattr(input_batch, "is_prefilling_np", None) if skip_decode_rows else None
-        for i, req_idx in enumerate(req_indices):
+        rows: Iterable[tuple[int, int]] = enumerate(req_indices)
+        if settled and is_prefilling_np is None and not self._eager_state.has_pending_replay():
+            split = self._split_settled_rows(input_batch, req_indices, settled)
+            if split is not None:
+                settled_rows, rows = split
+        for i, req_idx in rows:
             if is_prefilling_np is not None and not is_prefilling_np[i]:
                 continue
             buf = self.intermediate_buffer.buffers[req_idx]