[Model] Optimize Qwen3-TTS streaming throughput and first audio on CUDA (#8163)

Signed-off-by: Sy03 <1370724210@qq.com>

tests/engine/test_async_omni_engine_stage_init.py

@@ -10,6 +10,7 @@ import threading
 import time
 import types
 from dataclasses import dataclass, field
+from subprocess import CompletedProcess
 
 import pytest
 from omegaconf import OmegaConf
@@ -70,6 +71,49 @@ def test_stage_runtime_env_accepts_typed_runtime_config(monkeypatch):
     assert env_key not in os.environ
 
 
+@pytest.mark.parametrize(
+    "invalid_key,invalid_value",
+    [("INVALID=ENV", "value"), ("INVALID\0ENV", "value"), ("INVALID_ENV", "value\0")],
+    ids=["equals-in-key", "nul-in-key", "nul-in-value"],
+)
+def test_stage_runtime_env_restores_partial_application(monkeypatch, invalid_key, invalid_value):
+    existing_key = "VLLM_OMNI_TEST_EXISTING_STAGE_ENV"
+    new_key = "VLLM_OMNI_TEST_NEW_STAGE_ENV"
+    monkeypatch.setenv(existing_key, "original")
+    monkeypatch.delenv(new_key, raising=False)
+    runtime_config = OmniStageRuntimeConfig(
+        env={existing_key: "overridden", new_key: "temporary", invalid_key: invalid_value}
+    )
+
+    with pytest.raises(ValueError):
+        with stage_runtime_env(0, runtime_config):
+            pytest.fail("Invalid environment must prevent stage launch")
+
+    assert os.environ[existing_key] == "original"
+    assert new_key not in os.environ
+
+
+def test_stage_runtime_env_restores_after_launch_error(monkeypatch):
+    env_key = "VLLM_OMNI_TEST_STAGE_ENV_LAUNCH_ERROR"
+    monkeypatch.setenv(env_key, "original")
+
+    with pytest.raises(RuntimeError, match="stage launch failed"):
+        with stage_runtime_env(0, OmniStageRuntimeConfig(env={env_key: "temporary"})):
+            assert os.environ[env_key] == "temporary"
+            raise RuntimeError("stage launch failed")
+
+    assert os.environ[env_key] == "original"
+
+
+def test_stage_runtime_env_restores_first_value_for_normalized_keys(monkeypatch):
+    monkeypatch.setenv("123", "original")
+
+    with stage_runtime_env(0, {"env": {123: "first", "123": "second"}}):
+        assert os.environ["123"] == "second"
+
+    assert os.environ["123"] == "original"
+
+
 def test_orchestrator_startup_timeout_warns_how_to_raise_limits(monkeypatch):
     engine = object.__new__(AsyncOmniEngine)
     engine.orchestrator_thread = types.SimpleNamespace(is_alive=lambda: True)
@@ -1019,6 +1063,63 @@ def test_stage_runtime_launches_shared_engines_with_per_client_addresses(monkeyp
     assert os.environ.get("VLLM_OMNI_TEST_STAGE_RUNTIME_ENV") is None
 
 
+@pytest.mark.parametrize("client_count", [1, 2])
+def test_stage_launch_uses_configured_mps_pipe_and_rejects_conflicts(monkeypatch, client_count):
+    import vllm_omni.engine.stage_runtime as runtime_mod
+    from vllm_omni.engine import cuda_mps
+
+    runtime = _make_stage_runtime()
+    plans = [_make_llm_plan(stage_id, stage_id=stage_id, vllm_config=_FakeVllmConfig()) for stage_id in (0, 1)]
+    for stage_id, plan in enumerate(plans):
+        plan.replicas[0].engine_args_dict = {}
+        plan.replicas[0].metadata.runtime_cfg = OmniStageRuntimeConfig(
+            cuda_mps=True,
+            devices="0",
+            env={"CUDA_MPS_PIPE_DIRECTORY": f"/operator/stage-{stage_id}"},
+        )
+    monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/parent")
+    monkeypatch.setattr(runtime, "_prepare_stage_plans", lambda: plans)
+    monkeypatch.setattr(runtime, "_resolve_replica_physical_devices", lambda *_: "0")
+    monkeypatch.setattr(runtime_mod, "acquire_device_locks", lambda *_args, **_kwargs: [])
+    monkeypatch.setattr(runtime_mod, "physical_gpu_uuid", lambda _: "GPU-example")
+    monkeypatch.setattr(runtime_mod.current_omni_platform, "is_cuda", lambda: True)
+    monkeypatch.setattr(cuda_mps.shutil, "which", lambda _: "/bin/mps-control")
+    controls = []
+
+    def run_control(args, **kwargs):
+        controls.append(kwargs)
+        assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent"
+        assert kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/stage-0"
+        return CompletedProcess(args, 0, stdout="")
+
+    monkeypatch.setattr(cuda_mps.subprocess, "run", run_control)
+    launched = []
+
+    @contextlib.contextmanager
+    def launch(**kwargs):
+        launched.append(kwargs["stage_id"])
+        assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/stage-0"
+        assert kwargs["stage_visible_devices"] == "GPU-example"
+        yield StageReplicaResources(
+            addresses=EngineZmqAddresses(
+                inputs=[f"ipc://input-{i}" for i in range(client_count)],
+                outputs=[f"ipc://output-{i}" for i in range(client_count)],
+            )
+        )
+
+    monkeypatch.setattr(runtime_mod, "launch_stage_replica", launch)
+    try:
+        with pytest.raises(ValueError, match="Conflicting CUDA_MPS_PIPE_DIRECTORY"):
+            with runtime.launch_stage_engines(client_count):
+                pytest.fail("Conflicting colocated MPS policies must prevent launch")
+        assert launched == [0]
+        assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent"
+    finally:
+        runtime.shutdown()
+    assert len(controls) == 1
+    assert controls[0]["input"] == "get_server_list\n"
+
+
 @pytest.mark.parametrize("client_count", [1, 2])
 def test_stage_runtime_overlapping_devices_acquire_real_locks_once(monkeypatch, tmp_path, client_count):
     import fcntl

tests/engine/test_cuda_mps.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import os
+from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from vllm_omni.config.omni_config import OmniStageRuntimeConfig
+from vllm_omni.engine import cuda_mps
+from vllm_omni.engine.stage_runtime import StageRuntime
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def controls(monkeypatch):
+    monkeypatch.delenv("CUDA_MPS_PIPE_DIRECTORY", raising=False)
+    monkeypatch.setattr(cuda_mps.shutil, "which", lambda name: "/bin/mps-control")
+    run = Mock(return_value=SimpleNamespace(stdout=""))
+    monkeypatch.setattr(cuda_mps.subprocess, "run", run)
+    return run
+
+
+def test_private_server_has_scoped_environment_and_idempotent_close(controls):
+    original = dict(os.environ)
+    server = cuda_mps.CudaMPSServer("GPU-example")
+    pipe = Path(server.env["CUDA_MPS_PIPE_DIRECTORY"])
+    assert pipe.is_dir()
+    assert controls.call_args_list[0].args[0] == ["/bin/mps-control", "-d"]
+    assert controls.call_args_list[0].kwargs["env"]["CUDA_VISIBLE_DEVICES"] == "GPU-example"
+    assert dict(os.environ) == original
+    server.close()
+    count = controls.call_count
+    server.close()
+    assert controls.call_count == count
+    assert controls.call_args.kwargs["input"] == "quit\n"
+    assert not pipe.exists()
+
+
+def test_operator_daemon_is_checked_but_never_stopped(controls, monkeypatch):
+    monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/mps")
+    server = cuda_mps.CudaMPSServer("GPU-example")
+    server.close()
+    assert controls.call_count == 1
+    assert controls.call_args.kwargs["input"] == "get_server_list\n"
+
+
+def test_explicit_pipe_overrides_parent_without_taking_ownership(controls, monkeypatch):
+    monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/parent")
+    server = cuda_mps.CudaMPSServer("GPU-example", pipe_directory="/operator/stage")
+
+    assert controls.call_args.kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/stage"
+    assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent"
+    server.close()
+    assert controls.call_count == 1
+
+
+def test_explicit_empty_pipe_selects_private_daemon(controls, monkeypatch):
+    monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/parent")
+    server = cuda_mps.CudaMPSServer("GPU-example", pipe_directory="")
+    pipe = server.env["CUDA_MPS_PIPE_DIRECTORY"]
+    try:
+        assert pipe != "/operator/parent"
+        assert controls.call_args_list[0].args[0] == ["/bin/mps-control", "-d"]
+        assert controls.call_args_list[0].kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == pipe
+        assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent"
+    finally:
+        server.close()
+    assert controls.call_args.kwargs["input"] == "quit\n"
+    assert controls.call_args.kwargs["env"]["CUDA_MPS_PIPE_DIRECTORY"] == pipe
+    assert not Path(pipe).exists()
+
+
+@pytest.fixture
+def mps_runtime(monkeypatch):
+    from vllm_omni.engine import stage_runtime
+
+    runtime = StageRuntime.__new__(StageRuntime)
+    runtime._mps_servers = {}
+    monkeypatch.setattr(stage_runtime, "physical_gpu_uuid", lambda device: "GPU-example")
+    monkeypatch.setattr(stage_runtime.current_omni_platform, "is_cuda", lambda: True)
+    yield runtime
+    runtime._close_mps_servers()
+
+
+@pytest.mark.parametrize("pipe_directory", ["/operator/stage", ""])
+def test_runtime_shares_first_stage_pipe_policy(controls, monkeypatch, mps_runtime, pipe_directory):
+    monkeypatch.setenv("CUDA_MPS_PIPE_DIRECTORY", "/operator/parent")
+    config = OmniStageRuntimeConfig(cuda_mps=True, env={"CUDA_MPS_PIPE_DIRECTORY": pipe_directory})
+    first_env = mps_runtime._mps_environment("0", config)
+    selected_pipe = first_env["CUDA_MPS_PIPE_DIRECTORY"]
+    if pipe_directory:
+        assert selected_pipe == pipe_directory
+    else:
+        assert selected_pipe != "/operator/parent"
+        assert Path(selected_pipe).is_dir()
+    call_count = controls.call_count
+
+    assert mps_runtime._mps_environment("0", {"cuda_mps": True}) is first_env
+    assert mps_runtime._mps_environment("0", config) is first_env
+    assert controls.call_count == call_count
+    assert os.environ["CUDA_MPS_PIPE_DIRECTORY"] == "/operator/parent"
+
+
+@pytest.mark.parametrize(
+    "first_pipe,conflicting_pipe",
+    [("/operator/first", "/operator/other"), ("/operator/first", ""), ("", "/operator/other")],
+)
+def test_runtime_rejects_conflicting_explicit_stage_pipes(controls, mps_runtime, first_pipe, conflicting_pipe):
+    mps_runtime._mps_environment("0", {"cuda_mps": True, "env": {"CUDA_MPS_PIPE_DIRECTORY": first_pipe}})
+    call_count = controls.call_count
+
+    with pytest.raises(ValueError, match="Conflicting CUDA_MPS_PIPE_DIRECTORY"):
+        mps_runtime._mps_environment("0", {"cuda_mps": True, "env": {"CUDA_MPS_PIPE_DIRECTORY": conflicting_pipe}})
+
+    assert controls.call_count == call_count
+
+
+def test_startup_failure_closes_only_the_private_daemon(controls):
+    controls.side_effect = [SimpleNamespace(stdout=""), RuntimeError("health failed"), SimpleNamespace(stdout="")]
+    with pytest.raises(RuntimeError, match="health failed"):
+        cuda_mps.CudaMPSServer("GPU-example")
+    assert controls.call_args.kwargs["input"] == "quit\n"
+
+
+def test_runtime_opt_in_and_shared_gpu_ownership(monkeypatch):
+    from vllm_omni.engine import stage_runtime
+
+    runtime = StageRuntime.__new__(StageRuntime)
+    runtime._mps_servers = {}
+    create = Mock(return_value=SimpleNamespace(env={"CUDA_VISIBLE_DEVICES": "GPU-example"}, close=Mock()))
+    monkeypatch.setattr(stage_runtime, "CudaMPSServer", create)
+    monkeypatch.setattr(stage_runtime, "physical_gpu_uuid", lambda device: "GPU-example")
+    monkeypatch.setattr(stage_runtime, "current_omni_platform", SimpleNamespace(is_cuda=lambda: True))
+    assert runtime._mps_environment("0", {}) == {}
+    assert create.call_count == 0
+    for _ in range(2):
+        assert runtime._mps_environment("0", {"cuda_mps": True}) == {"CUDA_VISIBLE_DEVICES": "GPU-example"}
+    assert create.call_count == 1
+    runtime._close_mps_servers()
+    runtime._close_mps_servers()
+    create.return_value.close.assert_called_once()
+
+
+@pytest.mark.parametrize("devices", [None, "0,1"])
+def test_runtime_rejects_ambiguous_placement(devices, monkeypatch):
+    from vllm_omni.engine import stage_runtime
+
+    runtime = StageRuntime.__new__(StageRuntime)
+    monkeypatch.setattr(stage_runtime, "current_omni_platform", SimpleNamespace(is_cuda=lambda: True))
+    with pytest.raises(ValueError, match="exactly one explicit GPU"):
+        runtime._mps_environment(devices, {"cuda_mps": True})
+
+
+def test_failed_shutdown_preserves_socket_for_retry(controls):
+    server = cuda_mps.CudaMPSServer("GPU-example")
+    pipe = Path(server.env["CUDA_MPS_PIPE_DIRECTORY"])
+    controls.side_effect = RuntimeError("quit failed")
+    with pytest.raises(RuntimeError, match="quit failed"):
+        server.close()
+    assert pipe.is_dir()
+    controls.side_effect = None
+    server.close()
+    assert not pipe.exists()
+
+
+@pytest.mark.parametrize("launch_mode,stage_type", [("remote", "llm"), ("local", "diffusion")])
+def test_mps_rejects_unsupported_stages(launch_mode, stage_type):
+    runtime = StageRuntime.__new__(StageRuntime)
+    metadata = SimpleNamespace(runtime_cfg={"cuda_mps": True}, stage_type=stage_type)
+    stage_cfg = SimpleNamespace(runtime=metadata.runtime_cfg)
+    plans = [
+        SimpleNamespace(replicas=[SimpleNamespace(metadata=metadata, stage_cfg=stage_cfg, launch_mode=launch_mode)])
+    ]
+    with pytest.raises(ValueError, match="local EngineCore"):
+        runtime._validate_mps_topology(plans)
+    metadata.runtime_cfg["cuda_mps"] = False
+    runtime._validate_mps_topology(plans)
+
+
+def test_mps_rejects_parallel_initialization():
+    runtime = StageRuntime.__new__(StageRuntime)
+    runtime._parallel_stage_init = True
+    replica = SimpleNamespace(
+        metadata=SimpleNamespace(stage_type="llm"),
+        stage_cfg=SimpleNamespace(runtime={"cuda_mps": True}),
+        launch_mode="local",
+    )
+    with pytest.raises(ValueError, match="parallel_stage_init=false"):
+        runtime._validate_mps_topology([SimpleNamespace(replicas=[replica])])
+
+
+def test_failed_launch_command_still_checks_private_daemon_cleanup(controls):
+    controls.side_effect = [RuntimeError("launch failed after fork"), SimpleNamespace(stdout="")]
+    with pytest.raises(RuntimeError, match="launch failed after fork"):
+        cuda_mps.CudaMPSServer("GPU-example")
+    assert controls.call_args.kwargs["input"] == "quit\n"

tests/engine/test_data_entry_keys.py

@@ -425,3 +425,11 @@ class TestSerializeDeserializePayload:
         restored = deserialize_payload(wire)
 
         assert restored["meta"]["minimax_h3_prepared_reference_videos"] == descriptor
+
+
+def test_first_audio_metadata_round_trips_as_a_tensor():
+    from vllm_omni.data_entry_keys import to_dict, to_struct
+
+    payload = {"meta": {"first_audio": torch.tensor(True)}}
+    restored = to_dict(to_struct(payload))
+    assert restored["meta"]["first_audio"].item() is True

tests/engine/test_orchestrator_first_audio.py

@@ -0,0 +1,345 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+import torch
+from vllm.sampling_params import RequestOutputKind
+from vllm.v1.engine import FinishReason
+
+from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs
+from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker
+from vllm_omni.engine.messages import ErrorMessage, OutputMessage
+from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
+from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _CodecPool:
+    final_output = True
+
+    def __init__(self, output_kind, registered, replica_id):
+        self.output_processor = MultimodalOutputProcessor(
+            tokenizer=None, log_stats=False, engine_core_output_type="audio"
+        )
+        self.output_kind = output_kind
+        self.replica_id = replica_id
+        self.bound_replica = None
+        if registered:
+            self.register()
+
+    def register(self):
+        self.bound_replica = self.replica_id
+        self.output_processor.request_states["r"] = OmniRequestState(
+            request_id="r",
+            external_req_id="r",
+            parent_req=None,
+            request_index=0,
+            lora_request=None,
+            prompt=None,
+            prompt_token_ids=[0],
+            prompt_embeds=None,
+            logprobs_processor=None,
+            detokenizer=None,
+            max_tokens_param=None,
+            arrival_time=0.0,
+            queue=None,
+            log_stats=False,
+            stream_interval=1,
+            output_kind=self.output_kind,
+        )
+        self.output_processor.external_req_ids["r"].append("r")
+
+    def get_bound_replica_id(self, request_id):
+        return self.bound_replica
+
+    async def process_llm_raw_outputs(self, replica_id, raw, **kwargs):
+        assert replica_id == self.replica_id
+        return self.output_processor.process_outputs(raw.outputs, raw.timestamp).request_outputs
+
+
+def _orchestrator(output_kind=RequestOutputKind.DELTA, *, registered=True, replica_id=2, final_stage_id=1):
+    obj = Orchestrator.__new__(Orchestrator)
+    obj.stage_pools = [SimpleNamespace(final_output=False), _CodecPool(output_kind, registered, replica_id)]
+    obj.request_states = {"r": OrchestratorRequestState(request_id="r", final_stage_id=final_stage_id)}
+    obj.output_async_queue = asyncio.Queue()
+    obj._cfg_tracker = CfgCompanionTracker()
+    obj._pd_kv_params = {}
+    obj._running_counter = None
+    obj._abort_request_ids = AsyncMock(return_value=[])
+    obj._release_request_bindings = Mock()
+    obj._finish_raw_terminal_requests = AsyncMock()
+
+    async def process(stage, replica, raw, terminals):
+        await obj._route_upstream_first_audio(stage, replica, raw)
+        return await obj.stage_pools[stage].process_llm_raw_outputs(replica, raw)
+
+    async def route(stage, replica, outputs):
+        for output in outputs:
+            await obj.output_async_queue.put(
+                OutputMessage(
+                    request_id=output.request_id,
+                    stage_id=stage,
+                    replica_id=replica,
+                    engine_outputs=output,
+                    metrics=None,
+                    finished=output.finished,
+                )
+            )
+
+    obj._process_llm_stage_outputs = process
+    obj._handle_processed_outputs = route
+    return obj
+
+
+def _raw(*, first=False, required=False, terminal=False, samples=(1, 1, 1, 1)):
+    return OmniEngineCoreOutputs(
+        outputs=[
+            OmniEngineCoreOutput(
+                request_id="r",
+                new_token_ids=[],
+                finish_reason=FinishReason.STOP if terminal else None,
+                multimodal_output={
+                    "model_outputs": torch.tensor(samples, dtype=torch.float32),
+                    "sr": torch.tensor(24000),
+                    "_omni_first_audio": torch.tensor(first),
+                    "_omni_first_audio_required": torch.tensor(required),
+                },
+            )
+        ]
+    )
+
+
+async def _codec_output(obj, raw):
+    replica = obj.stage_pools[1].replica_id
+    processed = await obj._process_llm_stage_outputs(1, replica, raw, set())
+    await obj._handle_processed_outputs(1, replica, processed)
+
+
+def _messages(obj):
+    messages = []
+    while not obj.output_async_queue.empty():
+        messages.append(obj.output_async_queue.get_nowait())
+    return messages
+
+
+def _audio(output):
+    audio = output.outputs[0].multimodal_output["audio"]
+    return torch.cat(audio) if isinstance(audio, list) else audio
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("kind", list(RequestOutputKind))
+@pytest.mark.parametrize("codec_overtook", [False, True])
+async def test_first_audio_obeys_codec_output_kind(kind, codec_overtook):
+    obj = _orchestrator(kind)
+    suffix = _raw(required=True, samples=(20, 21))
+    terminal = _raw(terminal=True, samples=(30, 31))
+    if codec_overtook:
+        await _codec_output(obj, suffix)
+        await _codec_output(obj, terminal)
+        assert obj.output_async_queue.empty()
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11)))
+    if kind == RequestOutputKind.FINAL_ONLY and not codec_overtook:
+        assert obj.output_async_queue.empty()
+    if not codec_overtook:
+        await _codec_output(obj, suffix)
+        await _codec_output(obj, terminal)
+    messages = _messages(obj)
+    # The source replica is zero; output ownership belongs to codec replica two.
+    assert all((message.stage_id, message.replica_id) == (1, 2) for message in messages)
+    actual = [_audio(message.engine_outputs).tolist() for message in messages]
+    expected = {
+        RequestOutputKind.DELTA: [[10, 11], [20, 21], [30, 31]],
+        RequestOutputKind.CUMULATIVE: [[10, 11], [10, 11, 20, 21], [10, 11, 20, 21, 30, 31]],
+        RequestOutputKind.FINAL_ONLY: [[10, 11, 20, 21, 30, 31]],
+    }
+    assert actual == expected[kind]
+    assert messages[-1].finished
+    assert obj.request_states["r"].pending_first_audio_outputs == []
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True))
+    assert obj.output_async_queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_first_audio_waits_for_codec_registration_and_ignores_duplicates():
+    obj = _orchestrator(registered=False)
+    req_state = obj.request_states["r"]
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11)))
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(90, 91)))
+    assert obj.output_async_queue.empty()
+    assert not req_state.upstream_first_audio
+    assert req_state.pending_upstream_first_audio is not None
+    obj.stage_pools[1].register()
+    await obj._flush_upstream_first_audio(req_state)
+    messages = _messages(obj)
+    assert len(messages) == 1
+    assert _audio(messages[0].engine_outputs).tolist() == [10, 11]
+    assert req_state.upstream_first_audio
+    assert req_state.pending_upstream_first_audio is None
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("processing_failed", [False, True])
+async def test_prewarm_flushes_first_audio_after_successful_submission(mocker, processing_failed):
+    obj = _orchestrator(registered=False, final_stage_id=2 if processing_failed else 1)
+    req_state = obj.request_states["r"]
+    req_state.sampling_params_list = [None] * (req_state.final_stage_id + 1)
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11)))
+    pool = obj.stage_pools[1]
+    pool.stage_type = "llm"
+    pool.stage_vllm_config = SimpleNamespace(model_config=SimpleNamespace(hf_config=None))
+    obj.stage_pools[0].get_bound_replica_id = lambda request_id: 0
+    obj._stage_receives_async_chunks = lambda stage: True
+    obj._build_payload_sender_info = Mock(return_value={})
+    obj._on_stage_submitted = Mock()
+    obj._emit_tx_edge = Mock()
+    mocker.patch(
+        "vllm_omni.engine.orchestrator.build_engine_core_request_from_tokens",
+        return_value=SimpleNamespace(request_id="r"),
+    )
+
+    async def submit(*args, **kwargs):
+        pool.register()
+
+    pool.submit_initial = submit
+    if processing_failed:
+        pool.process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed"))
+        later_pool = SimpleNamespace(submit_initial=AsyncMock())
+        obj.stage_pools.append(later_pool)
+    submitted = await obj._prewarm_async_chunk_stages("r", SimpleNamespace(prompt_token_ids=[0]), req_state)
+    assert submitted is not processing_failed
+    messages = _messages(obj)
+    assert len(messages) == 1
+    if processing_failed:
+        assert isinstance(messages[0], ErrorMessage)
+        later_pool.submit_initial.assert_not_called()
+        assert "r" not in obj.request_states
+    else:
+        assert _audio(messages[0].engine_outputs).tolist() == [10, 11]
+
+
+@pytest.mark.asyncio
+async def test_pending_first_audio_is_discarded_on_cancellation():
+    obj = _orchestrator(registered=False)
+    req_state = obj.request_states["r"]
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True))
+    await obj._cleanup_request_ids(["r"], abort=True)
+    obj.stage_pools[1].register()
+    await obj._flush_upstream_first_audio(req_state)
+    assert obj.output_async_queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_cancellation_while_routing_first_audio_does_not_release_codec_outputs():
+    obj = _orchestrator()
+    await obj._route_upstream_first_audio(1, 2, _raw(required=True, terminal=True))
+
+    async def cancel(stage, replica, outputs):
+        await obj._cleanup_request_ids(["r"], abort=True)
+
+    obj._handle_processed_outputs = cancel
+    obj._process_llm_stage_outputs = AsyncMock()
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True))
+    obj._process_llm_stage_outputs.assert_not_called()
+    assert obj.output_async_queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_request_ending_at_source_does_not_emit_first_audio():
+    obj = _orchestrator(final_stage_id=0)
+    raw = _raw(first=True)
+    await obj._route_upstream_first_audio(0, 0, raw)
+    assert not raw.outputs and obj.output_async_queue.empty()
+    assert obj.request_states["r"].pending_upstream_first_audio is None
+
+
+@pytest.mark.asyncio
+async def test_later_aligner_receives_full_codec_waveform():
+    from vllm_omni.model_executor.stage_input_processors.forced_aligner import code2wav2aligner
+
+    obj = _orchestrator(RequestOutputKind.CUMULATIVE, final_stage_id=2)
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11)))
+    await _codec_output(obj, _raw(required=True, terminal=True, samples=(20, 21)))
+    messages = _messages(obj)
+    assert all(message.stage_id == 1 for message in messages)
+    inputs = code2wav2aligner(
+        [messages[-1].engine_outputs], prompt={"additional_information": {"text": ["Hello world"]}}
+    )
+    assert inputs[0]["multi_modal_data"]["audio"][0].tolist() == [10, 11, 20, 21]
+
+
+@pytest.mark.asyncio
+async def test_first_audio_processing_failure_fails_only_its_request():
+    obj = _orchestrator()
+    obj.request_states["healthy"] = OrchestratorRequestState(request_id="healthy", final_stage_id=1)
+    obj.stage_pools[1].process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed"))
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True))
+    error = obj.output_async_queue.get_nowait()
+    assert isinstance(error, ErrorMessage)
+    assert error.request_id == "r" and error.stage_id == 1
+    assert "processing failed" in error.error
+    assert set(obj.request_states) == {"healthy"}
+
+
+@pytest.mark.asyncio
+async def test_unmarked_upstream_audio_and_regular_codec_outputs_are_untouched():
+    obj = _orchestrator()
+    for stage in (0, 1):
+        raw = _raw()
+        await obj._route_upstream_first_audio(stage, 0, raw)
+        assert len(raw.outputs) == 1
+        assert torch.equal(raw.outputs[0].multimodal_output["model_outputs"], torch.ones(4))
+    assert obj.output_async_queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_cancelled_request_does_not_receive_late_first_audio():
+    obj = _orchestrator()
+    obj.request_states.clear()
+    raw = _raw(first=True)
+    await obj._route_upstream_first_audio(0, 0, raw)
+    assert not raw.outputs and obj.output_async_queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_first_audio_supports_both_speech_and_chat_consumers():
+    from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
+
+    obj = _orchestrator()
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True, samples=(10, 11)))
+    result = obj.output_async_queue.get_nowait().engine_outputs
+    assert not result.finished
+    assert _audio(result).tolist() == [10, 11]
+    payload, key = OmniOpenAIServingSpeech._extract_audio_output(result)
+    assert payload is not None and key is not None
+    audio = payload[key]
+    audio = torch.cat(audio) if isinstance(audio, list) else audio
+    assert audio.tolist() == [10, 11]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("audio", [None, [], torch.empty(0)])
+@pytest.mark.parametrize("codec_overtook", [False, True])
+async def test_invalid_first_audio_reports_error_and_releases_request(audio, codec_overtook):
+    obj = _orchestrator()
+    obj.request_states["healthy"] = OrchestratorRequestState(request_id="healthy", final_stage_id=1)
+    if codec_overtook:
+        await obj._route_upstream_first_audio(1, 2, _raw(required=True, terminal=True))
+        assert obj.request_states["r"].pending_first_audio_outputs
+    first = _raw(first=True)
+    first.outputs[0].multimodal_output["model_outputs"] = audio
+    await obj._route_upstream_first_audio(0, 0, first)
+
+    error = obj.output_async_queue.get_nowait()
+    assert isinstance(error, ErrorMessage)
+    assert error.request_id == "r" and "first-audio" in error.error
+    assert not first.outputs and set(obj.request_states) == {"healthy"}
+    obj._abort_request_ids.assert_awaited_once_with(["r"])
+    obj._release_request_bindings.assert_called_once_with(["r"])
+    await obj._route_upstream_first_audio(0, 0, _raw(first=True))
+    assert obj.output_async_queue.empty()