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

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

.buildkite/common/ci_source_file_dependencies.yml

@@ -20,6 +20,7 @@ x-qwen3-tts: &qwen3_tts
   - vllm_omni/model_executor/models/common/alias_free_activation.py
   - vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
   - vllm_omni/deploy/qwen3_tts.yaml
+  - vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml
   - vllm_omni/entrypoints/openai/serving_speech.py
   - vllm_omni/entrypoints/openai/tts_adapters/
   - vllm_omni/engine/stage_pool.py
@@ -30,6 +31,9 @@ x-qwen3-tts: &qwen3_tts
   - vllm_omni/config/omni_config.py
   - vllm_omni/engine/orchestrator.py
   - vllm_omni/engine/stage_engine_core_proc.py
+  - vllm_omni/engine/cuda_mps.py
+  - vllm_omni/engine/stage_runtime.py
+  - vllm_omni/engine/stage_init_utils.py
   - vllm_omni/deploy/qwen3_tts_mrv2.yaml
   - vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml
   - vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_b4.yaml

docs/configuration/stage_configs.md

@@ -337,6 +337,41 @@ supply `mtp_sampling_params` and `get_mtp_seed(sampling_params)` for model-local
 (with `mtp_sample_steps` and `mtp_sample_vocab_size`) as needed. Qwen3-TTS retains
 its existing `talker_mtp` entry point for V1.
 
+### Qwen3-TTS first-frame delivery and rollback
+
+The standard `qwen3_tts.yaml`, `qwen3_tts_high_concurrency.yaml`,
+`qwen3_tts_mrv2.yaml` and `qwen3_tts_high_concurrency_mrv2.yaml` profiles
+enable the `talker_first_audio` connector option, as does the experimental
+single-GPU profile. This changes the default CUDA MRV2 streaming path:
+residual prediction runs eagerly after the Talker sample, the Talker loads
+an additional first-frame decoder with its weights and CUDA graphs, and the
+orchestrator orders its audio before subsequent Code2Wav chunks.
+
+Code2Wav retains full-audio prefix graphs alongside the optional state-only
+graphs. Each request's delivery marker selects the graph and audio trimming;
+enabling the option alone never suppresses audio or disables prefix batching
+for requests that retain regular codec delivery.
+
+The path requires asynchronous chunks, TP/PP 1, an in-process executor and
+disabled prefix caching. Requests with reference codes and unsupported
+runners or platforms retain regular Code2Wav delivery.
+
+To restore regular codec delivery and avoid loading the Talker's additional
+decoder, use a deploy overlay with the relevant base profile:
+
+```yaml
+base_config: qwen3_tts.yaml
+connectors:
+  connector_of_shared_memory:
+    extra:
+      talker_first_audio: false
+```
+
+First-packet latency measures when PCM starts arriving. Time to first audible
+audio (TTFA) also includes any leading silence in the generated audio. An
+earlier first packet therefore does not necessarily improve TTFA; measure
+both for the intended voice and workload.
+
 ### Included performance work
 
 - Request snapshots have a fast path for immutable scalar leaves, including
@@ -372,14 +407,34 @@ speaker similarity before adopting either batching preset for a production
 workload. Floating-point decoder outputs can differ across batch sizes; this PR
 does not claim bitwise or quality equivalence.
 
-### Optional MPS deployment
+### Experimental MPS deployment
 
-NVIDIA MPS is an optional operator setting for colocated CUDA processes, not a
-YAML option or a library default. This PR does not establish a throughput or
-first-packet latency benefit from MPS. Measure the exact deployment with and
-without MPS before enabling it.
+NVIDIA MPS lets colocated CUDA stage processes share GPU execution resources.
+It is disabled by default. The experimental
+`qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` profile sets `cuda_mps: true`,
+alongside cached residual prediction, fused sampling, time-major codec
+convolutions, first-frame delivery, and larger graph batches.
 
-Use only assigned GPUs and an independent MPS pipe directory. A private MPS
-server does not provide exclusive GPU ownership or MIG isolation. For a
-single-GPU deployment, explicitly place both stages on that GPU; the supplied
-high-concurrency profile places its two stages on different GPUs by default.
+The runtime requires `nvidia-cuda-mps-control` on `PATH` and one explicit CUDA
+GPU per local EngineCore stage, with `parallel_stage_init: false`.
+Use numeric GPU ordinals for stage placement and visibility so initialization
+locks identify the physical GPU before MPS remaps it. It starts a private MPS daemon for each selected
+GPU and stops its own daemon after the stages exit. If
+`CUDA_MPS_PIPE_DIRECTORY` already names an operator-managed daemon, the runtime
+reuses it without stopping it. Diffusion and remote stages are unsupported.
+Set `cuda_mps: false` in a deploy overlay to disable automatic MPS management.
+
+Stage `env.CUDA_MPS_PIPE_DIRECTORY` overrides the parent setting when selecting
+the daemon; an explicit empty string selects a private daemon. Stages sharing
+a GPU reuse the first stage's daemon. Later stages may omit the setting or
+explicitly match that policy; a conflicting explicit setting fails before
+that stage starts. The parent environment is unchanged.
+
+MPS can improve throughput under concurrent load while increasing first-packet
+or first-audible-audio latency, particularly at low request rates. For a
+latency-sensitive workload, compare the same profile with `cuda_mps: false`.
+An inherited operator-managed MPS daemon must also be disabled by its owner
+for that comparison; this switch only controls automatic MPS management.
+
+MPS does not reserve a GPU. Use only assigned GPUs, explicitly place stages on
+the intended GPU, and warm the complete pipeline before measuring performance.

docs/design/qwen3_omni_tts_performance_optimization.md

@@ -93,13 +93,13 @@ Each optimization stacks on the previous one. The summary plots below show the c
 
 **Benchmark environment:**
 
-| | Qwen3-Omni                  | Qwen3-TTS |
-| --- |-----------------------------| --- |
-| **GPU** | A100                        | H200 |
+| | Qwen3-Omni | Qwen3-TTS |
+| --- | ----------------------------- | --- |
+| **GPU** | A100 | H200 |
 | **Model** | Qwen3-Omni-30B-A3B-Instruct | Qwen3-TTS-12Hz-1.7B-CustomVoice |
-| **vLLM** | v0.17.0                     | v0.18.0 |
-| **vllm-omni** | commit 199f7832             | v0.18.0rc2 |
-| **CUDA** | 12.9                        | 12.8 |
+| **vLLM** | v0.17.0 | v0.18.0 |
+| **vllm-omni** | commit 199f7832 | v0.18.0rc2 |
+| **CUDA** | 12.9 | 12.8 |
 
 This post walks through each optimization in the same order they are typically enabled in practice, then ends with deployment playbooks for both models.
 
@@ -310,7 +310,68 @@ self._compiled_model_fwd = torch.compile(
 
 `mode="default"` is used instead of `mode="reduce-overhead"` to avoid conflicts with vLLM's own CUDA graph capture on the main Talker model. `dynamic=True` handles the growing sequence length without recompilation.
 
-These optimizations are always-on in the current codebase - all Qwen3-TTS benchmark results in this post include them.
+The historical benchmark results above use re-prefill. It remains the default
+for Qwen3-TTS and the shared Qwen3-Omni predictor.
+
+### Experimental Qwen3-TTS single-GPU throughput profile
+
+The opt-in `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` profile uses a
+Qwen3-TTS-specific execution path on CUDA with BF16 weights:
+
+- `code_predictor_kv_cache` reuses each layer's keys and values across the 15
+  residual-codebook predictions within one audio frame. The first step writes
+  two positions; subsequent steps write one. Scratch storage is bounded by the
+  configured batch capacity, overwritten for every frame, and never shared
+  between requests across frames. The worker's serialized CUDA stream owns it;
+  it is separate from the Talker's paged KV cache.
+- A short causal/GQA Triton kernel masks unwritten K/V loads. Dynamic-batch
+  `torch.compile` warms the predictor before the runner captures whole-MTP
+  CUDA graphs. This path creates no inner CUDA graphs. Unsupported devices and
+  dtypes retain re-prefill; FP16 retains its existing precision safeguards.
+- `code_predictor_fused_sampling` fuses temperature scaling, top-k filtering,
+  and Gumbel argmax when the runner supplies FP32 uniforms. It preserves top-k
+  ties and uses the supplied randomness. Explicit generator paths retain the
+  existing sampler. Qwen3-Omni does not read either TTS option.
+- `decode_cudnn_benchmark` selects convolution algorithms during Code2Wav
+  warmup and graph capture. The process's cuDNN flags are restored afterwards,
+  including on capture failure. Captured graphs retain the selected algorithms.
+
+- `decode_time_major_conv` keeps codec convolution activations in time-major
+  layout to reduce transposes and copies between layers.
+- `talker_first_audio` completes residual prediction immediately after the
+  first codebook sample and decodes the first frame in the Talker process.
+  The codec advances its state without repeating the delivered frame; the
+  orchestrator orders that frame before subsequent codec chunks. This path
+  requires CUDA, MRV2, asynchronous chunks, a single-process executor, TP/PP 1,
+  and disabled prefix caching. Reference-code requests retain codec delivery.
+  Supported Qwen3-TTS profiles enable it explicitly; other models do not opt in.
+
+The predictor, sampling, convolution-layout and autotuning options default to
+false outside this throughput profile. It uses fixed 1/25-frame chunks,
+72 frames of left context, Talker/codec capacities of 128/64, codec batches
+through 8, and synchronous codec scheduling. Client concurrency is independent
+of those capacities. This profile also sets `cuda_mps: true`; the runtime owns
+a private MPS daemon unless an operator supplies an existing MPS pipe directory.
+Ordinary profiles leave MPS disabled. See
+[experimental MPS deployment](../configuration/stage_configs.md#experimental-mps-deployment)
+for requirements and ownership. Startup includes compilation and graph capture.
+
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml \
+  --api-server-count 4 --stage-init-timeout 1800 --init-timeout 2400
+```
+
+KV reuse changes GEMM shapes and attention reduction order; cuDNN autotuning
+can also change waveform numerics. This path is not bitwise equivalent to
+re-prefill, and aggregate WER alone cannot establish perceptual or speaker
+similarity equivalence. Validate quality on the intended workload before
+adopting it. Qwen3-Omni keeps its existing predictor and audio-delivery behavior.
+
+For a controlled ablation, copy the profile and disable individual options while keeping capacities, chunking, hardware, client concurrency,
+and warmup identical. Report repeated runs, first-audio latency and failures
+alongside throughput. Warm requests measure a ready service; they do not measure
+cold-start latency.
 
 ---
 

tests/config/test_environment_variables.py

@@ -166,7 +166,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
     assert category_counts == {
         EnvironmentVariableCategory.PUBLIC_OMNI: 30,
         EnvironmentVariableCategory.INHERITED_VLLM: 20,
-        EnvironmentVariableCategory.PLATFORM_EXTERNAL: 27,
+        EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
         EnvironmentVariableCategory.MODEL_SPECIFIC: 58,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
         EnvironmentVariableCategory.INTERNAL: 2,

tests/config/test_omni_config.py

@@ -886,6 +886,7 @@ def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
 
 def test_runtime_config_fields_match_structured_runtime_scope():
     assert {f.name for f in fields(OmniStageRuntimeConfig)} == {
+        "cuda_mps",
         "additional_config",
         "distributed_executor_backend",
         "worker_cls",
@@ -2051,6 +2052,28 @@ def test_async_chunk_rejects_mismatched_connector_edge(disabled_stage, builder):
             builder(pipeline, user_deploy_config=deploy)
 
 
+@pytest.mark.parametrize("scope", ["pipeline", "stage"])
+@pytest.mark.parametrize("from_yaml", [False, True])
+@pytest.mark.parametrize("builder", [merge_pipeline_deploy, VllmOmniConfig.from_pipeline_config])
+def test_async_chunk_rejects_quoted_false_before_selecting_processors(tmp_path, scope, from_yaml, builder):
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    if from_yaml:
+        path = tmp_path / "quoted_false.yaml"
+        path.write_text(
+            'async_chunk: "false"\n' if scope == "pipeline" else 'stages:\n  - stage_id: 0\n    async_chunk: "false"\n'
+        )
+        deploy = load_deploy_config(path)
+    elif scope == "pipeline":
+        deploy = DeployConfig(async_chunk="false")
+    else:
+        deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0, async_chunk="false")])
+    with pytest.raises(ValueError, match="async_chunk must be a boolean"):
+        if builder is merge_pipeline_deploy:
+            builder(pipeline, deploy)
+        else:
+            builder(pipeline, user_deploy_config=deploy)
+
+
 @pytest.mark.parametrize("explicit", [False, True])
 def test_diffusion_quantization_origin_survives_projection_and_transport(monkeypatch, explicit):
     from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
@@ -2169,3 +2192,31 @@ def test_structured_diffusion_stage_keeps_shared_globals_outside_diffusion():
 def test_structured_diffusion_stage_rejects_explicit_shared_engine_field(field_name, config_kwargs):
     with pytest.raises(ValueError, match=rf"stage 0.*{field_name}"):
         _build_single_diffusion_config(**config_kwargs)
+
+
+def test_mps_is_explicit_only_in_experimental_single_gpu_profile():
+    assert not load_deploy_config(_DEPLOY_DIR / "qwen3_tts.yaml").cuda_mps
+    deploy = load_deploy_config(_DEPLOY_DIR / "qwen3_tts_high_concurrency_mrv2_single_gpu.yaml")
+    assert deploy.cuda_mps
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    stages = merge_pipeline_deploy(pipeline, deploy)
+    assert all(stage.yaml_runtime["cuda_mps"] is True for stage in stages)
+
+
+def test_mps_config_rejects_string_boolean(tmp_path):
+    path = tmp_path / "deploy.yaml"
+    path.write_text('cuda_mps: "false"\n')
+    with pytest.raises(ValueError, match="cuda_mps must be a boolean"):
+        load_deploy_config(path)
+
+
+def test_mps_stays_in_runtime_instead_of_engine_arguments():
+    from vllm_omni.engine.stage_init_utils import build_engine_args_dict_from_omni_stage_config
+
+    config = _from_pipeline_key(
+        "qwen3_tts", deploy_config_path=str(_DEPLOY_DIR / "qwen3_tts_high_concurrency_mrv2_single_gpu.yaml")
+    )
+    for stage in config.stage_configs:
+        assert stage.runtime_config.cuda_mps
+        args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
+        assert "cuda_mps" not in args

tests/e2e/online_serving/test_qwen3_tts_customvoice.py

@@ -66,6 +66,53 @@ default_tts_server_params = [
     )
 ]
 
+# Exercise the throughput path in the existing L4 lane with smaller capacities.
+# The production H200 profile retains its own admission and graph buckets.
+_FAST_PATH_STAGE_CONFIG = modify_stage_config(
+    get_deploy_config_path("qwen3_tts_high_concurrency_mrv2_single_gpu.yaml"),
+    updates={
+        # The temporary overlay lives outside deploy/, so anchor its parent.
+        "base_config": get_deploy_config_path("qwen3_tts_high_concurrency_mrv2.yaml"),
+        "cuda_mps": False,
+        "connectors.connector_of_shared_memory.extra.decode_batch_max_size": 2,
+        "connectors.connector_of_shared_memory.extra.decode_cudagraph_batch_sizes": [1, 2],
+        "stages": {
+            0: {
+                "max_num_seqs": 4,
+                "max_num_batched_tokens": 128,
+                "compilation_config.cudagraph_capture_sizes": [1, 2, 4, 8, 16, 32, 64, 128],
+                "compilation_config.max_cudagraph_capture_size": 128,
+            },
+            1: {"max_num_seqs": 4},
+        },
+    },
+)
+
+
+@pytest.mark.core_model
+@pytest.mark.advanced_model
+@pytest.mark.tts
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize(
+    "omni_server",
+    [OmniServerParams(model=MODEL, stage_config_path=_FAST_PATH_STAGE_CONFIG, server_args=["--trust-remote-code"])],
+    indirect=True,
+)
+def test_cached_predictor_streaming_audio(omni_server, online_client) -> None:
+    """Cover the optimized predictor, first audio and codec path with real weights at L3."""
+    online_client.send_audio_speech_request(
+        {
+            "model": omni_server.model,
+            "input": "The quick brown fox jumps over the lazy dog.",
+            "stream": True,
+            "stream_format": "audio",
+            "response_format": "wav",
+            "task_type": "CustomVoice",
+            "voice": "vivian",
+        },
+        request_num=4,
+    )
+
 
 @pytest.mark.core_model
 @pytest.mark.advanced_model

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

tests/model_executor/models/qwen3_tts/test_cached_code_predictor.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Cache isolation and graph replay for the opt-in TTS predictor."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.common.qwen3_code_predictor import CodePredictorBaseModel
+from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA")
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_cached_frames_match_causal_prefill_after_graph_replay(dtype):
+    from vllm_omni.model_executor.models.qwen3_tts.cached_code_predictor import FrameLocalKVCache
+
+    config = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64,
+        hidden_size=64,
+        intermediate_size=128,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=16,
+        num_code_groups=4,
+    )
+    torch.manual_seed(171)
+    model = CodePredictorBaseModel(config).to(device="cuda", dtype=dtype).eval()
+    inputs = torch.randn(3, 5, 64, device="cuda", dtype=dtype)
+    positions = torch.arange(5, device="cuda").expand(3, -1)
+    cache = FrameLocalKVCache(model, max_batch=3)
+    with torch.inference_mode():
+        for batch in (1, 3):
+            for _ in range(2):
+                for step in (1, 2, 3):
+                    cache(inputs, batch, step)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            output = torch.cat([cache(inputs, 3, step) for step in (1, 2, 3)], dim=1)
+        for _ in range(3):
+            inputs.normal_()
+            # No previous frame (including a smaller batch) may leak through.
+            cache(inputs, 1, 1)
+            cache.buffer.fill_(torch.nan)
+            graph.replay()
+            expected = model(inputs, positions)[:, :4]
+            tolerance = 1e-5 if dtype == torch.float32 else 0.04
+            torch.testing.assert_close(output, expected, atol=tolerance, rtol=tolerance)

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -681,6 +681,7 @@ class TestCodePredictorGraphReplay:
             sampling_mode="per_call",
             return_proj_buf=False,
         )
+        predictor._frame_cache = None
         predictor._prefix_graphs_enabled = False
         predictor._prefix_reprefill_enabled = False
         predictor._bucket_pos_ids = {1: torch.arange(4).unsqueeze(0)}
@@ -1398,3 +1399,36 @@ class TestMTPExecutionBuckets:
             )
             wrapper2._bucket_sizes = [1, 2, 4]  # stand-in for warmed-up state
             wrapper2.configure_mtp_execution_buckets([3])
+
+
+@pytest.mark.parametrize("enabled", [False, True])
+def test_tts_cache_falls_back_on_cpu(mocker, loaded_target_classes, enabled):
+    cp_config, talker_config = _make_tiny_config(loaded_target_classes)
+    cfg = _make_vllm_config(mocker)
+    cfg.model_config.stage_connector_config = {"extra": {"code_predictor_kv_cache": enabled}}
+    predictor = loaded_target_classes[2](vllm_config=cfg, config=cp_config, talker_config=talker_config)
+    common = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+    setup = mocker.patch.object(common.CodePredictorWrapper, "_setup_compile")
+    predictor._setup_compile()
+    setup.assert_called_once_with()
+    assert predictor._frame_cache is None
+
+
+def test_shared_predictor_ignores_tts_fast_path_options(mocker, loaded_target_classes):
+    cp_config, _ = _make_tiny_config(loaded_target_classes)
+    cfg = _make_vllm_config(mocker)
+    cfg.model_config.stage_connector_config = {
+        "extra": {"code_predictor_kv_cache": True, "code_predictor_fused_sampling": True}
+    }
+    common = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+    predictor = common.CodePredictorWrapper(
+        vllm_config=cfg,
+        cp_config=cp_config,
+        wrapper_config=common.CodePredictorWrapperConfig(),
+    )
+    assert not hasattr(predictor, "_frame_cache")
+    logits = torch.randn(3, cp_config.vocab_size)
+    uniforms = torch.rand_like(logits)
+    expected = (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
+    actual = predictor._sample_per_call(logits, 1.0, 0, None, uniforms)
+    torch.testing.assert_close(actual, expected)

tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py

@@ -0,0 +1,68 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Supplied-uniform samples must match the existing ATen path exactly."""
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
+]
+
+
+@pytest.mark.parametrize("batch", [1, 64])
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+@pytest.mark.parametrize("top_k", [-1, 1, 50, 2048])
+def test_topk_gumbel_exact_samples_and_graph(batch, dtype, top_k):
+    from vllm_omni.model_executor.models.qwen3_tts.code_predictor_sampling import sample_code_topk_gumbel
+
+    with torch.inference_mode():
+        torch.manual_seed(42)
+        logits = torch.randn(batch, 2048, device="cuda", dtype=dtype) * 5
+        uniforms = torch.rand(batch, 15, 2048, device="cuda")[:, 7, :]
+        uniforms.clamp_(1e-6, 1 - 1e-6)
+        for case in ("random", "ties", "all_masked", "part_masked", "uniform_zero", "uniform_one", "nan"):
+            if case == "ties":
+                logits.zero_()
+                uniforms.fill_(0.5)
+            elif case == "all_masked":
+                logits.fill_(-torch.inf)
+            elif case == "part_masked":
+                logits[:, 4:9] = 1
+            elif case == "uniform_zero":
+                logits.normal_()
+                uniforms.zero_()
+            elif case == "uniform_one":
+                uniforms.fill_(1)
+            elif case == "nan":
+                uniforms.fill_(0.5)
+                logits[:, 7] = torch.nan
+            scaled = logits * (1 / 0.9)
+            if top_k > 0:
+                threshold = scaled.topk(top_k, dim=-1).values[:, -1:]
+                scaled = scaled.masked_fill(scaled < threshold, -torch.inf)
+            expected = (scaled.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
+            actual = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9)
+            torch.testing.assert_close(actual, expected, atol=0, rtol=0)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9)
+        graph.replay()
+        torch.testing.assert_close(captured, expected, atol=0, rtol=0)
+
+
+@pytest.mark.parametrize("top_k", [0, 1, 50, 97])
+def test_topk_gumbel_accepts_strided_supplied_uniforms(top_k):
+    from vllm_omni.model_executor.models.qwen3_tts.code_predictor_sampling import sample_code_topk_gumbel
+
+    torch.manual_seed(81)
+    logits = torch.randn(3, 194, device="cuda", dtype=torch.bfloat16)[:, ::2]
+    uniforms = torch.rand(3, 15, 194, device="cuda")[:, 3, ::2].clamp_(1e-6, 1 - 1e-6)
+    scaled = logits * (1 / 0.9)
+    if top_k > 0:
+        scaled = scaled.masked_fill(scaled < scaled.topk(top_k).values[:, -1:], -torch.inf)
+    expected = (scaled.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
+    actual = sample_code_topk_gumbel(logits, uniforms, top_k, 1 / 0.9)
+    torch.testing.assert_close(actual, expected)

tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Prefix graphs must follow per-request first-audio delivery, not a capture hint."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from tests.model_executor.models.qwen3_tts.test_qwen3_tts_incremental_decode import _decoder_stub
+from vllm_omni.model_executor.models.qwen3_tts.segmented_graph_wrapper import CUDAGraphDecoderWrapper
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2Decoder,
+)
+
+pytestmark = [pytest.mark.core_model]
+DEVICE = torch.device("cuda:0")
+
+
+@pytest.mark.cpu
+@pytest.mark.parametrize("skip_flags", [(False, False), (True, True), (False, True), (True, False)])
+@pytest.mark.parametrize("initial_frames", [1, 3])
+@pytest.mark.parametrize("state_only_available", [False, True])
+def test_xvec_first_chunk_replays_prefix_graph(monkeypatch, skip_flags, initial_frames, state_only_available):
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
+    wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper)
+    wrapper.prefix_length = 72
+    wrapper.initial_chunk_frames = initial_frames
+    wrapper.codec_chunk_frames = 25
+    wrapper.capture_batch_sizes = [1, 2]
+    wrapper._icl_previous_frames_by_target = {26: 1}
+    wrapper._xvec_previous_frames_by_target = {initial_frames + 25: initial_frames}
+    replayed = []
+    wrapper.xvec_prefix_states = {
+        2: {
+            "graph": SimpleNamespace(replay=lambda: replayed.append("full")),
+            "input": {"codes": torch.zeros(2, 2, initial_frames)},
+            "output": torch.arange(4 * initial_frames, dtype=torch.float32).view(2, 1, 2 * initial_frames),
+            "cache": {
+                "ref_hidden": torch.zeros(2, 2, 0),
+                "ref_conv": torch.zeros(2, 0, 3),
+                "prefix_hidden": torch.zeros(2, 0, 3),
+                "ref_upsample": torch.zeros(2, 3, 0),
+                "ref_wav": torch.zeros(2, 1, 0),
+                "suffix_quantized": torch.ones(2, 2, initial_frames),
+                "suffix_conv": torch.ones(2, initial_frames, 3),
+                "past_key_values": SimpleNamespace(layers=[]),
+            },
+        }
+    }
+    wrapper.xvec_prefix_state_only_states = {}
+    if state_only_available and initial_frames == 1:
+        wrapper.xvec_prefix_state_only_states[2] = {
+            **wrapper.xvec_prefix_states[2],
+            "graph": SimpleNamespace(replay=lambda: replayed.append("state")),
+            "state_only": True,
+        }
+    wrapper._record_graph_hit = lambda *_args: None
+    wrapper._record_graph_fallback = lambda *_args: None
+    wrapper._ensure_suffix_buffers = lambda cache: None
+    wrapper.decoder = _decoder_stub(
+        capture_first_audio_state_only=True,
+        total_upsample=2,
+        _is_suffix_cache_rolling=lambda previous, cached: False,
+        _decode_xvec_first_chunk=lambda *_args: pytest.fail("unexpected eager fallback"),
+        _slice_dynamic_cache=Qwen3TTSTokenizerV2Decoder._slice_dynamic_cache,
+    )
+
+    caches = [{"prefix_frames": 0, "skip_first_audio": skip} for skip in skip_flags]
+    outputs = wrapper._batched_request_decode(
+        [torch.full((1, 2, initial_frames), 3), torch.full((1, 2, initial_frames), 4)],
+        caches,
+    )
+
+    assert replayed == ["state" if all(skip_flags) and state_only_available and initial_frames == 1 else "full"]
+    torch.testing.assert_close(wrapper.xvec_prefix_states[2]["input"]["codes"][0], torch.full((2, initial_frames), 3.0))
+    torch.testing.assert_close(wrapper.xvec_prefix_states[2]["input"]["codes"][1], torch.full((2, initial_frames), 4.0))
+    assert [cache["decoder_prefix_frames"] for cache in caches] == [0, 0]
+    assert [cache["suffix_frames"] for cache in caches] == [initial_frames, initial_frames]
+    assert len(outputs) == 2
+    for row, skip in enumerate(skip_flags):
+        expected = wrapper.xvec_prefix_states[2]["output"][row : row + 1, :, 2 if skip else 0 :]
+        torch.testing.assert_close(outputs[row], expected)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("initial_frames", [1, 3])
+@pytest.mark.parametrize("time_major", [False, True])
+@torch.inference_mode()
+def test_xvec_graphs_follow_request_first_audio_flags(initial_frames, time_major):
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.segmented_graph_wrapper import (
+        CUDAGraphDecoderWrapper as SegmentedWrapper,
+    )
+
+    decoder = _make_decoder().to(DEVICE)
+    decoder.precompute_snake_caches()
+    if time_major:
+        decoder.enable_time_major_conv()
+    decoder.capture_first_audio_state_only = True
+    wrapper = SegmentedWrapper(
+        decoder,
+        capture_modes=("xvec",),
+        capture_batch_sizes=[1, 2],
+        num_quantizers=2,
+        initial_chunk_frames=initial_frames,
+    )
+    # This regression covers prefix graphs. Avoid compiling unrelated suffix
+    # shapes, especially the time-major kernels for long continuation chunks.
+    wrapper._xvec_previous_frames_by_target = {}
+    wrapper.warmup(DEVICE)
+    assert set(wrapper.xvec_prefix_states) == {1, 2}
+    assert set(wrapper.xvec_prefix_state_only_states) == ({1, 2} if initial_frames == 1 else set())
+
+    for skip_flags in ((False, False), (True, True), (False, True), (True, False)):
+        codes = [torch.randint(0, 32, (1, 2, initial_frames), device=DEVICE) for _ in skip_flags]
+        caches = [{"prefix_frames": 0, "skip_first_audio": skip} for skip in skip_flags]
+        expected_caches = [{"prefix_frames": 0, "skip_first_audio": skip} for skip in skip_flags]
+        expected = [decoder._decode_stream_first_chunk(code, cache) for code, cache in zip(codes, expected_caches)]
+        actual = wrapper._decode_xvec_prefix_batch(codes, caches)
+        assert actual is not None  # Eligible batches must not silently become per-request eager calls.
+        for got, want, cache, expected_cache in zip(actual, expected, caches, expected_caches):
+            torch.testing.assert_close(got, want, atol=1e-4, rtol=1e-4)
+            for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"):
+                torch.testing.assert_close(cache[key], expected_cache[key], atol=1e-5, rtol=1e-5)

tests/model_executor/models/qwen3_tts/test_qwen3_tts_code2wav.py

@@ -752,3 +752,37 @@ def test_load_weights_uses_model_dtype_before_precomputing_caches(dtype):
     assert _load_weights_noop(model) == {"decoder.fake_weight"}
     assert model.decoder.weight.dtype is dtype
     assert cache_dtypes == [dtype]
+
+
+@pytest.mark.parametrize("capture_fails", [False, True])
+def test_decode_autotune_restores_process_flags(mocker, monkeypatch, capture_fails):
+    monkeypatch.setattr(torch.backends.cudnn, "benchmark", False)
+    monkeypatch.setattr(torch.backends.cudnn, "benchmark_limit", 5)
+    original = (torch.backends.cudnn.enabled, torch.backends.cudnn.deterministic, torch.backends.cudnn.allow_tf32)
+    model = _make_model(
+        async_chunk=True,
+        device=torch.device("cuda"),
+        stage_connector_config={"extra": {"decode_cudnn_benchmark": True}},
+    )
+
+    observed = []
+
+    def capture(**kwargs):
+        observed.append(
+            (
+                torch.backends.cudnn.benchmark,
+                torch.backends.cudnn.benchmark_limit,
+                torch.backends.cudnn.enabled,
+                torch.backends.cudnn.deterministic,
+                torch.backends.cudnn.allow_tf32,
+            )
+        )
+        if capture_fails:
+            raise RuntimeError("capture failed")
+
+    capture_mock = mocker.patch.object(model, "_maybe_enable_decoder_cudagraph", side_effect=capture)
+    _load_weights_noop(model)
+    capture_mock.assert_called_once()
+    assert observed == [(True, 10, *original)]
+    assert torch.backends.cudnn.benchmark is False
+    assert torch.backends.cudnn.benchmark_limit == 5

tests/model_executor/models/qwen3_tts/test_short_kv_attention.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Causal/GQA correctness and graph replay for the frame-local KV kernel."""
+
+import pytest
+import torch
+
+pytestmark = [
+    pytest.mark.core_model,
+    pytest.mark.cuda,
+    pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
+]
+
+
+@pytest.mark.parametrize("batch", [1, 3, 64])
+@pytest.mark.parametrize("heads,kv_heads,dim", [(16, 8, 128), (8, 8, 64)])
+@pytest.mark.parametrize("positions", [[0, 1], [2], [8], [15]])
+def test_short_attention_causal_gqa_and_graph(batch, heads, kv_heads, dim, positions):
+    from vllm_omni.model_executor.models.qwen3_tts.short_kv_attention import short_kv_attention
+
+    with torch.inference_mode():
+        torch.manual_seed(42)
+        pos = torch.tensor(positions, device="cuda")
+        q = torch.randn(batch, len(positions), heads, dim, device="cuda", dtype=torch.bfloat16).transpose(1, 2)
+        cache = torch.randn(2, batch + 3, kv_heads, 17, dim, device="cuda", dtype=torch.bfloat16)
+        k, v = cache[0, :batch], cache[1, :batch]
+        mask = torch.arange(17, device="cuda")[None, :] <= pos[:, None]
+        kr = k.float().repeat_interleave(heads // kv_heads, dim=1)
+        vr = v.float().repeat_interleave(heads // kv_heads, dim=1)
+        scores = (q.float() @ kr.transpose(-1, -2)) * dim**-0.5
+        expected = scores.masked_fill(~mask, -torch.inf).softmax(-1) @ vr
+        actual = short_kv_attention(q, k, v, pos, dim**-0.5)
+        torch.testing.assert_close(actual.float(), expected, atol=0.016, rtol=0.008)
+
+        # Unwritten cache slots must not affect a new frame, even if poisoned.
+        k[:, :, max(positions) + 1 :] = torch.nan
+        v[:, :, max(positions) + 1 :] = torch.nan
+        poisoned = short_kv_attention(q, k, v, pos, dim**-0.5)
+        torch.testing.assert_close(poisoned, actual, atol=0, rtol=0)
+        graph = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(graph):
+            captured = short_kv_attention(q, k, v, pos, dim**-0.5)
+        graph.replay()
+        torch.testing.assert_close(captured, actual, atol=0, rtol=0)

tests/model_executor/models/qwen3_tts/test_skip_delivered_first.py

@@ -0,0 +1,65 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Code2Wav drops only frames explicitly marked as delivered by the Talker."""
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2Decoder,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+UPSAMPLE = 4
+
+
+def _decoder(skip: bool):
+    decoder = Qwen3TTSTokenizerV2Decoder.__new__(Qwen3TTSTokenizerV2Decoder)
+    decoder.__dict__["capture_first_audio_state_only"] = skip
+    decoder.__dict__["total_upsample"] = UPSAMPLE
+    calls = []
+
+    def full(codes, caches):
+        calls.append("full")
+        frames = int(codes.shape[-1])
+        return torch.arange(frames * UPSAMPLE, dtype=torch.float32).view(1, 1, -1)
+
+    def state_only(codes, caches):
+        calls.append("state")
+        return torch.zeros(1)
+
+    decoder.__dict__["_decode_xvec_first_chunk"] = full
+    decoder.__dict__["_decode_xvec_first_chunk_state_only"] = state_only
+    return decoder, calls
+
+
+def test_one_frame_first_chunk_only_builds_state():
+    decoder, calls = _decoder(skip=True)
+    out = decoder._decode_stream_first_chunk(torch.zeros(1, 16, 1, dtype=torch.long), {"skip_first_audio": True})
+    assert out.shape[-1] == 0 and calls == ["state"]
+
+
+def test_longer_first_chunk_keeps_all_but_the_delivered_frame():
+    decoder, calls = _decoder(skip=True)
+    out = decoder._decode_stream_first_chunk(
+        torch.zeros(1, 16, 3, dtype=torch.long),
+        {"skip_first_audio": True},
+    )
+    assert calls == ["full"]
+    assert torch.equal(out, torch.arange(UPSAMPLE, 3 * UPSAMPLE, dtype=torch.float32).view(1, 1, -1))
+
+
+def test_without_skip_the_whole_first_chunk_is_decoded():
+    decoder, calls = _decoder(skip=False)
+    out = decoder._decode_stream_first_chunk(
+        torch.zeros(1, 16, 3, dtype=torch.long),
+        {"skip_first_audio": False},
+    )
+    assert calls == ["full"] and out.shape[-1] == 3 * UPSAMPLE
+
+
+def test_request_without_first_audio_uses_regular_decoder_even_when_graph_skip_is_enabled():
+    decoder, calls = _decoder(skip=True)
+    out = decoder._decode_stream_first_chunk(torch.zeros(1, 16, 1, dtype=torch.long), {})
+    assert calls == ["full"] and out.shape[-1] == UPSAMPLE

tests/model_executor/models/qwen3_tts/test_time_major_conv.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+import torch.nn.functional as F
+
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import time_major_conv as tmc
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _snake_nct(x, exp_alpha, inv_beta):
+    return x + inv_beta[None, :, None] * torch.sin(x * exp_alpha[None, :, None]).square()
+
+
+@pytest.mark.parametrize(("kernel_size", "dilation"), [(1, 1), (7, 1), (7, 3), (3, 9)])
+@pytest.mark.parametrize("with_bias_and_snake", [False, True])
+def test_im2col_gemm_matches_causal_conv1d(kernel_size, dilation, with_bias_and_snake):
+    torch.manual_seed(0)
+    bsz, t, c_in, c_out = 2, 11, 6, 5
+    conv = torch.nn.Conv1d(c_in, c_out, kernel_size, dilation=dilation)
+    x = torch.randn(bsz, c_in, t)
+    in_bias = torch.randn(c_in) if with_bias_and_snake else None
+    snake = (torch.rand(c_in) + 0.5, torch.rand(c_in) + 0.5) if with_bias_and_snake else None
+
+    ref_in = x if in_bias is None else x + in_bias[None, :, None]
+    if snake is not None:
+        ref_in = _snake_nct(ref_in, *snake)
+    expected = conv(F.pad(ref_in, ((kernel_size - 1) * dilation, 0)))
+
+    cols = tmc.snake_im2col(x.transpose(1, 2).contiguous(), kernel_size, dilation, in_bias, snake)
+    actual = torch.addmm(conv.bias, cols, tmc.conv_gemm_weight(conv.weight)).view(bsz, t, c_out)
+    torch.testing.assert_close(actual.transpose(1, 2), expected, rtol=1e-5, atol=1e-5)
+
+
+@pytest.mark.parametrize(("stride", "kernel_size"), [(2, 2), (3, 6), (8, 16)])
+def test_transposed_overlap_add_matches_trimmed_conv_transpose(stride, kernel_size):
+    torch.manual_seed(0)
+    bsz, t, c_in, c_out = 2, 5, 4, 3
+    conv = torch.nn.ConvTranspose1d(c_in, c_out, kernel_size, stride=stride)
+    x = torch.randn(bsz, c_in, t)
+    expected = conv(x)[..., : t * stride]  # causal: drop the kernel - stride right tail
+
+    z = (x.transpose(1, 2).reshape(bsz * t, c_in) @ tmc.transposed_gemm_weight(conv.weight)).view(
+        bsz, t, kernel_size, c_out
+    )
+    actual = tmc.transposed_overlap_add(z, stride, conv.bias)
+    torch.testing.assert_close(actual.transpose(1, 2), expected, rtol=1e-5, atol=1e-5)
+
+
+def test_transposed_overlap_add_rejects_other_kernels():
+    with pytest.raises(ValueError, match="kernel stride or 2"):
+        tmc.transposed_overlap_add(torch.zeros(1, 2, 5, 3), 2, torch.zeros(3))
+
+
+@pytest.mark.parametrize(("kernel_size", "dilation"), [(1, 1), (7, 3)])
+@pytest.mark.parametrize("with_residual", [False, True])
+def test_causal_conv_matches_conv1d(kernel_size, dilation, with_residual):
+    torch.manual_seed(0)
+    bsz, t, c_in, c_out = 2, 13, 6, 5
+    conv = torch.nn.Conv1d(c_in, c_out, kernel_size, dilation=dilation)
+    x = torch.randn(bsz, c_in, t)
+    residual = torch.randn(bsz, t, c_out) if with_residual else None
+    expected = conv(F.pad(x, ((kernel_size - 1) * dilation, 0))).transpose(1, 2)
+    if residual is not None:
+        expected = expected + residual
+    actual = tmc.causal_conv(
+        x.transpose(1, 2).contiguous(), tmc.conv_gemm_weight(conv.weight), conv.bias, kernel_size, dilation, residual
+    )
+    torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
+
+
+def _gpu_tolerances(dtype):
+    # bf16 results may differ by one ulp (sin rounding, fp32 vs bf16 accumulation order).
+    return {"rtol": 1.6e-2, "atol": 1e-2} if dtype == torch.bfloat16 else {"rtol": 1e-4, "atol": 1e-4}
+
+
+@pytest.mark.skipif(not torch.cuda.is_available() or not tmc.HAS_TRITON, reason="needs CUDA and Triton")
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+def test_triton_kernels_match_reference(dtype):
+    torch.manual_seed(0)
+    tol = _gpu_tolerances(dtype)
+    x = torch.randn(3, 37, 96, device="cuda", dtype=dtype)
+    bias = torch.randn(96, device="cuda", dtype=dtype)
+    snake = (torch.rand(96, device="cuda", dtype=dtype) + 0.5, torch.rand(96, device="cuda", dtype=dtype) + 0.5)
+    for kernel_size, dilation in ((1, 1), (7, 9)):
+        torch.testing.assert_close(
+            tmc.snake_im2col(x, kernel_size, dilation, bias, snake),
+            tmc.snake_im2col_reference(x, kernel_size, dilation, bias, snake),
+            **tol,
+        )
+    z = torch.randn(3, 37, 10, 48, device="cuda", dtype=dtype)
+    torch.testing.assert_close(
+        tmc.transposed_overlap_add(z, 5, bias[:48]), tmc.transposed_overlap_add_reference(z, 5, bias[:48]), **tol
+    )
+
+
+@pytest.mark.skipif(not torch.cuda.is_available() or not tmc.HAS_TRITON, reason="needs CUDA and Triton")
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@pytest.mark.parametrize(("c_in", "c_out", "kernel_size", "dilation"), [(96, 96, 7, 9), (192, 96, 1, 1), (96, 1, 7, 1)])
+def test_triton_causal_conv_matches_reference(dtype, c_in, c_out, kernel_size, dilation):
+    torch.manual_seed(0)
+    x = torch.randn(3, 101, c_in, device="cuda", dtype=dtype)
+    weight = torch.randn(kernel_size * c_in, c_out, device="cuda", dtype=dtype) / (kernel_size * c_in) ** 0.5
+    bias = torch.randn(c_out, device="cuda", dtype=dtype)
+    residual = torch.randn(3, 101, c_out, device="cuda", dtype=dtype)
+    expected = tmc.causal_conv_reference(
+        x.float(), weight.float(), bias.float(), kernel_size, dilation, residual.float()
+    )
+    actual = tmc.causal_conv(x, weight, bias, kernel_size, dilation, residual)
+    torch.testing.assert_close(actual.float(), expected, **_gpu_tolerances(dtype))

tests/model_executor/models/qwen3_tts/test_time_major_decoder.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2DecoderConfig,
+)
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2Decoder,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_decoder() -> Qwen3TTSTokenizerV2Decoder:
+    config = Qwen3TTSTokenizerV2DecoderConfig(
+        codebook_size=32,
+        hidden_size=16,
+        latent_dim=16,
+        codebook_dim=16,
+        num_attention_heads=2,
+        num_key_value_heads=2,
+        intermediate_size=32,
+        num_hidden_layers=1,
+        num_quantizers=2,
+        decoder_dim=32,
+        upsample_rates=(8, 5, 4, 3),
+        upsampling_ratios=(2, 2),
+        sliding_window=72,
+    )
+    torch.manual_seed(0)
+    decoder = Qwen3TTSTokenizerV2Decoder(config).eval()
+    with torch.no_grad():
+        for name, param in decoder.named_parameters():
+            # Non-trivial SnakeBeta and ConvNeXt scales so every term matters.
+            if name.endswith(("alpha", "beta")):
+                param.uniform_(-0.5, 0.5)
+            elif name.endswith("gamma"):
+                param.uniform_(0.5, 1.0)
+    decoder.precompute_snake_caches()
+    return decoder
+
+
+@torch.inference_mode()
+def test_time_major_conv_decode_matches_nct_stack() -> None:
+    decoder = _make_decoder()
+    hidden = torch.randn(2, 5, decoder.config.latent_dim)
+    expected = decoder._conv_decode(hidden)
+
+    decoder.enable_time_major_conv()
+    actual = decoder._conv_decode(hidden)
+
+    assert actual.shape == expected.shape == (2, 1, 5 * int(decoder.total_upsample))
+    scale = expected.abs().max()
+    torch.testing.assert_close(actual / scale, expected / scale, rtol=0, atol=1e-5)
+    # The stack holds weight copies outside the module tree.
+    assert not any("time_major" in key for key in decoder.state_dict())
+
+
+@torch.inference_mode()
+def test_time_major_conv_decode_takes_sliced_hidden() -> None:
+    decoder = _make_decoder()
+    hidden = torch.randn(1, 9, decoder.config.latent_dim)
+    expected = decoder._conv_decode(hidden[:, -4:, :])
+    decoder.enable_time_major_conv()
+    actual = decoder._conv_decode(hidden[:, -4:, :])  # non-contiguous, like the suffix decode's slice
+    scale = expected.abs().max()
+    torch.testing.assert_close(actual / scale, expected / scale, rtol=0, atol=1e-5)
+
+
+@torch.inference_mode()
+def test_first_frame_single_key_attention_matches_codec_without_graphs():
+    from vllm_omni.model_executor.models.qwen3_tts.first_frame_decoder import Qwen3TTSFirstFrameDecoder
+
+    decoder = _make_decoder()
+    codes = torch.randint(0, 32, (3, 2, 1))
+    expected = decoder._decode_xvec_first_chunk(codes, {})[:, 0, :]
+    first = Qwen3TTSFirstFrameDecoder.__new__(Qwen3TTSFirstFrameDecoder)
+    torch.nn.Module.__init__(first)
+    first.decoder = decoder
+    first._graphs = {}
+    actual = first.decode(codes[:, :, 0])
+    torch.testing.assert_close(actual, expected, rtol=1e-4, atol=1e-5)

tests/worker_v2/test_first_audio_sender.py

@@ -0,0 +1,237 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import queue
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.engine import OmniEngineCoreOutputs
+from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _scheduler(**client_by_request):
+    return SimpleNamespace(
+        requests={request_id: SimpleNamespace(client_index=client) for request_id, client in client_by_request.items()}
+    )
+
+
+def test_sink_groups_outputs_by_client_and_carries_audio():
+    output_queue: queue.Queue = queue.Queue()
+    sink = engine_output_queue_sink(output_queue, _scheduler(a=0, b=1, c=0))
+    sample_rate = torch.tensor(24000, dtype=torch.int32)
+    rows = [torch.full((4,), float(i)) for i in range(3)]
+
+    sink.prepare(["a", "b", "c"])(["a", "b", "c"], rows, sample_rate)
+
+    delivered = {}
+    while not output_queue.empty():
+        client_index, outputs = output_queue.get_nowait()
+        assert isinstance(outputs, OmniEngineCoreOutputs)
+        delivered[client_index] = outputs.outputs
+    assert sorted(delivered) == [0, 1]
+    assert [o.request_id for o in delivered[0]] == ["a", "c"]
+    assert [o.request_id for o in delivered[1]] == ["b"]
+    for output in delivered[0] + delivered[1]:
+        # First audio carries no tokens and does not finish the request.
+        assert output.new_token_ids == []
+        assert output.finish_reason is None
+        assert output.multimodal_output["sr"] is sample_rate
+    assert torch.equal(delivered[0][1].multimodal_output["model_outputs"], rows[2])
+
+
+def test_sink_skips_requests_that_already_left_the_scheduler():
+    output_queue: queue.Queue = queue.Queue()
+    sink = engine_output_queue_sink(output_queue, _scheduler(kept=0))
+
+    sink.prepare(["gone", "kept"])(["gone", "kept"], [torch.zeros(2), torch.ones(2)], torch.tensor(24000))
+    sink.prepare(["gone"])(["gone"], [torch.zeros(2)], torch.tensor(24000))
+
+    client_index, outputs = output_queue.get_nowait()
+    assert client_index == 0
+    assert [o.request_id for o in outputs.outputs] == ["kept"]
+    assert output_queue.empty()
+
+
+class _DoneEvent:
+    def record(self):
+        pass
+
+    def synchronize(self):
+        pass
+
+
+def _run_sender(items):
+    from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender
+
+    output_queue: queue.Queue = queue.Queue()
+    sink = engine_output_queue_sink(output_queue, _scheduler(**{rid: 0 for item in items for rid in item[2]}))
+    sender = FirstAudioSender(sink)
+    for item in items:
+        sender._queue.put((*item, sink.prepare(item[2])))
+    sender.close()
+    delivered = []
+    while not output_queue.empty():
+        _, outputs = output_queue.get_nowait()
+        delivered.append(
+            ([o.request_id for o in outputs.outputs], [o.multimodal_output["model_outputs"] for o in outputs.outputs])
+        )
+    return delivered
+
+
+def test_sender_drops_rows_whose_frame_is_not_audio():
+    pcm = torch.arange(6, dtype=torch.float32).reshape(3, 2)
+    sample_rate = torch.tensor(24000, dtype=torch.int32)
+    valid = torch.tensor([True, False, True])
+
+    delivered = _run_sender(
+        [
+            (_DoneEvent(), pcm, ["a", "b", "c"], sample_rate, valid),
+            # Every row invalid: nothing reaches the sink.
+            (_DoneEvent(), pcm[:1], ["d"], sample_rate, torch.tensor([False])),
+            # No validity: every row is delivered (previous behaviour).
+            (_DoneEvent(), pcm[:1], ["e"], sample_rate, None),
+        ]
+    )
+
+    assert [ids for ids, _rows in delivered] == [["a", "c"], ["e"]]
+    assert torch.equal(delivered[0][1][1], pcm[2])
+
+
+def test_prepared_delivery_survives_scheduler_completion():
+    output_queue: queue.Queue = queue.Queue()
+    scheduler = _scheduler(r=3)
+    sink = engine_output_queue_sink(output_queue, scheduler)
+    deliver = sink.prepare(["r"])
+    scheduler.requests.clear()
+    deliver(["r"], [torch.ones(2)], torch.tensor(24000))
+    client, outputs = output_queue.get_nowait()
+    assert client == 3
+    assert outputs.outputs[0].multimodal_output["_omni_first_audio"]
+
+
+def test_delivery_failure_emits_request_error():
+    from vllm.v1.engine import FinishReason
+
+    output_queue: queue.Queue = queue.Queue()
+    delivery = engine_output_queue_sink(output_queue, _scheduler(r=0)).prepare(["r"])
+    delivery.fail(["r"])
+    _, outputs = output_queue.get_nowait()
+    assert outputs.outputs[0].finish_reason == FinishReason.ERROR
+
+
+@pytest.mark.parametrize("client_index", [0, 3])
+def test_submit_reports_only_requests_with_a_prepared_route(monkeypatch, client_index):
+    from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender
+
+    # Exercise the routing contract without allocating CUDA or pinned memory.
+    original_empty = torch.empty
+    monkeypatch.setattr(torch, "empty", lambda *args, pin_memory=False, **kwargs: original_empty(*args, **kwargs))
+    monkeypatch.setattr(torch.cuda, "Event", _DoneEvent)
+    output_queue: queue.Queue = queue.Queue()
+    sender = FirstAudioSender(engine_output_queue_sink(output_queue, _scheduler(kept=client_index)))
+    pcm = torch.arange(6, dtype=torch.float32).reshape(3, 2)
+    try:
+        accepted = sender.submit(["gone", "kept", "also-gone"], pcm, torch.tensor(24000))
+    finally:
+        sender.close()
+
+    assert accepted == ["kept"]
+    client, outputs = output_queue.get_nowait()
+    assert client == client_index
+    assert [output.request_id for output in outputs.outputs] == ["kept"]
+    assert torch.equal(outputs.outputs[0].multimodal_output["model_outputs"], pcm[1])
+    assert output_queue.empty()
+
+
+def test_submit_without_any_route_leaves_audio_to_the_codec(monkeypatch):
+    from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender
+
+    def unexpected_copy(*args, **kwargs):
+        pytest.fail("No delivery is possible; do not start a device copy")
+
+    monkeypatch.setattr(torch, "empty", unexpected_copy)
+    monkeypatch.setattr(torch.cuda, "Event", unexpected_copy)
+    output_queue: queue.Queue = queue.Queue()
+    sender = FirstAudioSender(engine_output_queue_sink(output_queue, _scheduler()))
+    try:
+        accepted = sender.submit(["gone"], torch.ones(1, 2), torch.tensor(24000))
+    finally:
+        sender.close()
+
+    assert accepted == []
+    assert output_queue.empty()
+
+
+@pytest.mark.parametrize("sink", [lambda ids, rows, sr: None, SimpleNamespace(prepare=None)])
+def test_sender_rejects_sink_without_route_preparation(sink):
+    from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender
+
+    with pytest.raises(TypeError, match="prepare"):
+        FirstAudioSender(sink)
+
+
+def test_sender_failure_only_aborts_valid_undelivered_rows():
+    from vllm.v1.engine import FinishReason
+
+    from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender
+
+    class FailSecondClientOnce(queue.Queue):
+        def put_nowait(self, item):
+            client, outputs = item
+            if client == 1 and outputs.outputs[0].finish_reason is None:
+                raise RuntimeError("client enqueue failed")
+            super().put_nowait(item)
+
+    output_queue = FailSecondClientOnce()
+    sink = engine_output_queue_sink(output_queue, _scheduler(eos=1, sent=0, failed=1, pending=2, next=3))
+    sender = FirstAudioSender(sink)
+    ids = ["eos", "sent", "failed", "pending", "gone"]
+    sender._queue.put(
+        (
+            _DoneEvent(),
+            torch.ones(5, 2),
+            ids,
+            torch.tensor(24000),
+            torch.tensor([False, True, True, True, True]),
+            sink.prepare(ids),
+        )
+    )
+    sender._queue.put((_DoneEvent(), torch.ones(1, 2), ["next"], torch.tensor(24000), None, sink.prepare(["next"])))
+    sender.close()
+
+    actual: list[tuple[int, str, FinishReason | None]] = []
+    while not output_queue.empty():
+        client, outputs = output_queue.get_nowait()
+        actual.extend((client, o.request_id, o.finish_reason) for o in outputs.outputs)
+    assert actual == [
+        (0, "sent", None),
+        (1, "failed", FinishReason.ERROR),
+        (2, "pending", FinishReason.ERROR),
+        (3, "next", None),
+    ]
+
+
+def test_sender_copy_failure_reports_all_prepared_routes():
+    from vllm.v1.engine import FinishReason
+
+    from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender
+
+    class FailedEvent:
+        def synchronize(self):
+            raise RuntimeError("copy failed before validity became readable")
+
+    output_queue: queue.Queue = queue.Queue()
+    sink = engine_output_queue_sink(output_queue, _scheduler(r=3))
+    sender = FirstAudioSender(sink)
+    sender._queue.put(
+        (FailedEvent(), torch.ones(2, 2), ["r", "gone"], torch.tensor(24000), None, sink.prepare(["r", "gone"]))
+    )
+    sender.close()
+    client, outputs = output_queue.get_nowait()
+    assert client == 3
+    assert [(o.request_id, o.finish_reason) for o in outputs.outputs] == [("r", FinishReason.ERROR)]
+    assert output_queue.empty()

tests/worker_v2/test_omni_model_state.py

@@ -391,3 +391,204 @@ def test_mtp_requires_model_declared_output_key(output_key):
     model = SimpleNamespace(mtp=lambda: None, mtp_output_key=output_key)
     with pytest.raises(TypeError, match="must declare mtp_output_key"):
         state._init_mtp_runner(model)
+
+
+_EAGER_DIM = 3
+_CODEBOOK = 2048
+_EOS = 2150
+
+
+class _EagerBatch(_DummyInputBatch):
+    def __init__(self, spans, indices=None):
+        indices = list(range(len(spans))) if indices is None else indices
+        super().__init__(indices)
+        starts = [0]
+        for n_tok in spans:
+            starts.append(starts[-1] + n_tok)
+        self.num_scheduled_tokens = list(spans)
+        self.query_start_loc_np = starts[:-1]
+        self.query_start_loc = torch.tensor(starts, dtype=torch.int32)
+
+
+def _make_eager_state(max_num_reqs=4):
+    state = _make_state(max_num_reqs=max_num_reqs, has_preprocess=True)
+    _init_static(state, max_num_reqs, dim=_EAGER_DIM)
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+
+    state._eager_state = EagerMTPState(state)
+    state._first_audio_requests = set()
+    state._eager_mtp = True
+    state._eager_fastpath = False
+    state._eager_rows = None
+    state._eager_ready = {}
+    state._eager_settled = {}
+    state._eager_embeds = torch.zeros((max_num_reqs, _EAGER_DIM))
+    state._decode_preprocess = None
+    state._mtp_sample_uniforms = None
+    state.vllm_config.cache_config = SimpleNamespace(enable_prefix_caching=False)
+    model = state.model
+    model.first_frame_decoder = None
+    model.mtp_frame_valid = lambda layer0: (layer0 >= 0) & (layer0 < _CODEBOOK)
+    model.embed_input_ids = lambda ids: ids.float().reshape(-1, 1, 1).expand(-1, 1, _EAGER_DIM)
+    model.mtp_calls = []
+
+    def mtp(input_ids, input_embeds, last_hidden, text_step, **kwargs):
+        model.mtp_calls.append((input_ids.clone(), last_hidden.clone(), text_step.clone()))
+        codes = torch.stack([input_ids, input_ids + 1, input_ids + 2], dim=1)
+        # Frame embedding sum plus the text step, like the Qwen3 talkers.
+        return input_embeds + 100 + text_step, codes
+
+    model.mtp = mtp
+    return state
+
+
+def _eager_outputs(num_tokens):
+    return {
+        "codes": {"audio": torch.zeros((num_tokens, 3), dtype=torch.long)},
+        "meta": {"codec_frame_valid": torch.zeros((num_tokens,), dtype=torch.int8)},
+    }
+
+
+def test_eager_mtp_publishes_frames_in_the_step_that_sampled_cb0():
+    state = _make_eager_state()
+    _fill_buffers(state, "prefill", "decode", "ended")
+    # prefill: final 3-token chunk; decode: input CB0 42; ended: input CB0 was
+    # EOS (async scheduling ran the finished request once more).
+    batch = _EagerBatch([3, 1, 1])
+    input_ids = torch.tensor([0, 0, 0, 42, _EOS])
+    state._eager_rows = (batch, [(0, 0, "prefill", True), (1, 1, "decode", False), (2, 2, "ended", False)], input_ids)
+    text_hidden = torch.arange(5 * _EAGER_DIM, dtype=torch.float32).reshape(5, _EAGER_DIM)
+    sampled = torch.tensor([[7], [_EOS], [9]])
+    outputs = _eager_outputs(5)
+
+    state.run_eager_mtp(batch, text_hidden, sampled, outputs)
+
+    ids, hidden, text_step = state.model.mtp_calls[0]
+    assert ids.tolist() == [7, _EOS, 9]
+    # The frame's hidden is this step's last-token hidden of each span.
+    assert torch.equal(hidden, text_hidden[[2, 3, 4]])
+    # The next step adds its own text step.
+    assert not text_step.any()
+    codes = outputs["codes"]["audio"]
+    assert codes[2].tolist() == [7, 8, 9]
+    assert codes[3].tolist() == [_EOS, _EOS + 1, _EOS + 2]
+    assert codes[4].tolist() == [9, 10, 11]
+    assert not codes[:2].any()
+    # Valid only when both the sampled and (for decode) the input CB0 are codec ids.
+    assert outputs["meta"]["codec_frame_valid"].tolist() == [0, 0, 1, 0, 0]
+    assert state._eager_ready == {0: "prefill", 1: "decode", 2: "ended"}
+    assert torch.equal(state._eager_embeds[0], torch.full((_EAGER_DIM,), 107.0))
+    assert state._eager_rows is None
+
+
+def test_eager_mtp_ignores_rows_recorded_for_another_batch():
+    state = _make_eager_state()
+    _fill_buffers(state, "r0")
+    state._eager_rows = (_EagerBatch([1]), [(0, 0, "r0", False)], torch.tensor([5]))
+    outputs = _eager_outputs(1)
+    state.run_eager_mtp(_EagerBatch([1]), torch.zeros(1, _EAGER_DIM), torch.tensor([[3]]), outputs)
+    assert state.model.mtp_calls == []
+    assert not outputs["codes"]["audio"].any()
+
+
+def test_eager_mtp_requires_retained_frame_outputs():
+    state = _make_eager_state()
+    _fill_buffers(state, "r0")
+    batch = _EagerBatch([1])
+    state._eager_rows = (batch, [(0, 0, "r0", False)], torch.tensor([5]))
+    with pytest.raises(RuntimeError, match="requires retained codes.audio"):
+        state.run_eager_mtp(batch, torch.zeros(1, _EAGER_DIM), torch.tensor([[3]]), {"codes": {"audio": None}})
+
+
+def test_eager_decode_input_is_previous_frame_plus_text_step():
+    state = _make_eager_state()
+    _fill_buffers(state, "r0", "r1")
+    state._eager_ready = {0: "r0", 1: "r1"}
+    state._eager_embeds[0] = 5.0
+    state._eager_embeds[1] = 7.0
+    embeds = torch.zeros((2, _EAGER_DIM))
+    batches = [
+        (0, 0, (torch.zeros(_EAGER_DIM), torch.full((_EAGER_DIM,), 2.0))),
+        (1, 1, (None, torch.ones(_EAGER_DIM))),
+    ]
+    state._eager_state._apply_eager_frames(batches, embeds, _EagerBatch([1, 1]), None)
+    assert embeds.tolist() == [[7.0] * _EAGER_DIM, [8.0] * _EAGER_DIM]
+    # The deferred MTP must not run for an eager decode row.
+    assert state.model.mtp_calls == []
+
+
+def test_eager_decode_without_a_frame_fails_loudly():
+    state = _make_eager_state()
+    _fill_buffers(state, "r0")
+    batches = [(0, 0, (torch.zeros(_EAGER_DIM), torch.zeros(_EAGER_DIM)))]
+    with pytest.raises(RuntimeError, match="frame missing"):
+        state._eager_state._apply_eager_frames(batches, torch.zeros((1, _EAGER_DIM)), _EagerBatch([1]), None)
+
+
+def test_run_preprocess_records_rows_that_keep_a_sample():
+    state = _make_eager_state()
+    _fill_buffers(state, "chunk", "final", "decode")
+    state._eager_ready = {2: "decode"}
+    state._eager_embeds[2] = 4.0
+    text_step = torch.ones(_EAGER_DIM)
+
+    def preprocess(input_ids, input_embeds, **info):
+        updates = {"mtp_inputs": (torch.zeros(_EAGER_DIM), text_step)} if input_ids.shape[0] == 1 else {}
+        return input_ids, input_embeds, updates
+
+    state.model.preprocess = preprocess
+    batch = _EagerBatch([3, 3, 1])
+    model_inputs = {"input_ids": torch.zeros(7, dtype=torch.long), "inputs_embeds": torch.zeros((7, _EAGER_DIM))}
+    req_states = SimpleNamespace(
+        prompt_len=np.array([10, 10, 10], dtype=np.int32),
+        num_computed_tokens=np.array([0, 7, 10], dtype=np.int32),
+    )
+    state.run_preprocess(batch, model_inputs, req_states)
+
+    recorded_batch, entries, _ids = state._eager_rows
+    assert recorded_batch is batch
+    # A non-final prefill chunk samples nothing that is kept.
+    assert entries == [(1, 1, "final", True), (2, 2, "decode", False)]
+    assert model_inputs["inputs_embeds"][6].tolist() == [5.0] * _EAGER_DIM
+    assert state.model.mtp_calls == []
+
+
+@pytest.mark.parametrize("first_was_valid", [False, True])
+def test_first_audio_requirement_survives_eos_only_if_audio_was_queued(first_was_valid):
+    state = _make_eager_state()
+    _fill_buffers(state, "r0")
+    state._first_audio_requests.add("r0")
+    state._eager_state._first_audio_valid = torch.tensor([first_was_valid, False, False, False])
+    batch = _EagerBatch([1])
+    state._eager_rows = (batch, [(0, 0, "r0", False)], torch.tensor([5]))
+    outputs = _eager_outputs(1)
+    outputs["meta"]["first_audio"] = torch.zeros(1, dtype=torch.bool)
+    state.run_eager_mtp(batch, torch.zeros(1, _EAGER_DIM), torch.tensor([[_EOS]]), outputs)
+    assert outputs["meta"]["first_audio"].tolist() == [first_was_valid]
+
+
+@pytest.mark.parametrize("accepted", [[], ["r1"], ["r0", "r1"]])
+def test_first_audio_marker_requires_accepted_delivery(monkeypatch, accepted):
+    from contextlib import nullcontext
+
+    state = _make_eager_state()
+    _fill_buffers(state, "r0", "r1")
+    state.model.first_frame_decoder = SimpleNamespace(
+        sample_rate=24000, decode=lambda codes: torch.ones(codes.shape[0], 2)
+    )
+    state._first_audio_stream = SimpleNamespace(wait_stream=lambda stream: None)
+    state._first_audio_sender = SimpleNamespace(submit=lambda ids, pcm, sr, valid: accepted)
+    monkeypatch.setattr(torch.cuda, "current_stream", lambda device: None)
+    monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext())
+    monkeypatch.setattr(torch.Tensor, "record_stream", lambda tensor, stream: None)
+    batch = _EagerBatch([1, 1])
+    state._eager_rows = (batch, [(0, 0, "r0", True), (1, 1, "r1", True)], torch.zeros(2, dtype=torch.long))
+    outputs = _eager_outputs(2)
+    outputs["meta"]["first_audio"] = torch.zeros(2, dtype=torch.bool)
+
+    state.run_eager_mtp(batch, torch.zeros(2, _EAGER_DIM), torch.tensor([[7], [8]]), outputs)
+
+    assert state._first_audio_requests == set(accepted)
+    # A missing route must leave the normal codec path responsible for frame 0;
+    # otherwise it skips the frame and the orchestrator waits forever for it.
+    assert outputs["meta"]["first_audio"].tolist() == ["r0" in accepted, "r1" in accepted]

vllm_omni/config/environment_variable_inventory.py

@@ -112,6 +112,8 @@ _INHERITED_VLLM = (
 _PLATFORM_EXTERNAL = (
     "ASCEND_LAUNCH_BLOCKING",
     "ASCEND_RT_VISIBLE_DEVICES",
+    "CUDA_MPS_LOG_DIRECTORY",
+    "CUDA_MPS_PIPE_DIRECTORY",
     "FLASHINFER_DISABLE_VERSION_CHECK",
     "HF_HOME",
     "HF_HUB_DISABLE_XET",

vllm_omni/config/omni_config.py

@@ -629,6 +629,8 @@ class OmniStageConnectorConfig:
 class OmniStageRuntimeConfig:
     """Per-stage process placement and backend runtime behavior."""
 
+    cuda_mps: bool = False
+
     # LLM backend extensions; diffusion owns these in its config projection.
     additional_config: dict[str, Any] | None = None
     distributed_executor_backend: Any = None
@@ -2127,6 +2129,7 @@ def _build_runtime_config(
         kwargs["num_replicas"] = stage_deploy.num_replicas
     if "env" not in kwargs and stage_deploy is not None and stage_deploy.env is not None:
         kwargs["env"] = _copy_value(stage_deploy.env)
+    kwargs["cuda_mps"] = deploy.cuda_mps
     kwargs["num_gpus"] = parallel_config.world_size
     return OmniStageRuntimeConfig(**kwargs)
 

vllm_omni/config/stage_config.py

@@ -582,6 +582,8 @@ class DeployConfig:
     model_runner: Literal["v1", "v2"] = "v1"
     # Stage-1 active stream slots; 0 preserves legacy all-stream cycling.
     active_stream_window: int = 0
+    # Experimental local NVIDIA MPS; disabled unless a deploy explicitly opts in.
+    cuda_mps: bool = False
     duplex_session: DuplexSessionRuntimeConfig = field(default_factory=DuplexSessionRuntimeConfig)
     connectors: dict[str, Any] | None = None
     edges: list[dict[str, Any]] | None = None
@@ -804,7 +806,11 @@ def load_deploy_config(path: str | Path) -> DeployConfig:
     if model_runner not in ("v1", "v2"):
         raise ValueError(f"model_runner must be one of ('v1', 'v2'), got {model_runner!r}")
 
+    if not isinstance(raw_dict.get("cuda_mps", False), bool):
+        raise ValueError("cuda_mps must be a boolean")
+
     kwargs: dict[str, Any] = {
+        "cuda_mps": raw_dict.get("cuda_mps", False),
         "async_chunk": raw_dict.get("async_chunk", True),
         "session_mode": raw_dict.get("session_mode", "turn"),
         "model_runner": model_runner,
@@ -936,7 +942,11 @@ def _resolve_execution_mode(
 
 
 def resolve_stage_async_chunk(deploy: DeployConfig, stage: StageDeployConfig | None) -> bool:
-    return bool(deploy.async_chunk and (stage is None or stage.async_chunk is not False))
+    if not isinstance(deploy.async_chunk, bool):
+        raise ValueError("async_chunk must be a boolean")
+    if stage is not None and stage.async_chunk is not None and not isinstance(stage.async_chunk, bool):
+        raise ValueError(f"Stage {stage.stage_id} async_chunk must be a boolean or null")
+    return deploy.async_chunk and (stage is None or stage.async_chunk is not False)
 
 
 def validate_stage_async_chunk_edges(pipeline: PipelineConfig, deploy: DeployConfig) -> None:
@@ -1178,6 +1188,8 @@ def merge_pipeline_deploy(
             runtime["num_replicas"] = ds.num_replicas
             if ds.env is not None:
                 runtime["env"] = ds.env
+        if deploy.cuda_mps:
+            runtime["cuda_mps"] = True
         runtime["requires_multimodal_data"] = ps.requires_multimodal_data
 
         result.append(

vllm_omni/data_entry_keys.py

@@ -19,6 +19,10 @@ import msgspec
 import numpy as np
 import torch
 
+# Internal output routing markers shared by first-frame producers and orchestration.
+FIRST_AUDIO_KEY = "_omni_first_audio"
+FIRST_AUDIO_REQUIRED_KEY = "_omni_first_audio_required"
+
 REQUEST_ARTIFACT_DIRS_KEY = "_omni_request_artifact_dirs"
 TRANSFORM_OWNED_META_KEYS = frozenset({"minimax_h3_prepared_reference_videos"})
 
@@ -101,6 +105,7 @@ class OmniPayloadMeta(TypedDict, total=False):
     width: int
     decode_flag: bool
     codec_streaming: bool
+    first_audio: bool | torch.Tensor
     codec_frame_valid: bool | torch.Tensor
     ref_code_len: int
     ref_context_size: int
@@ -214,6 +219,7 @@ class MetaStruct(_StructBase):
     width: int | None = None
     decode_flag: bool | None = None
     codec_streaming: bool | None = None
+    first_audio: torch.Tensor | None = None
     codec_frame_valid: torch.Tensor | None = None
     ref_code_len: int | None = None
     # Expected FINAL length of a growing async-chunk sequence, when the

vllm_omni/deploy/qwen3_tts.yaml

@@ -24,6 +24,7 @@ connectors:
   connector_of_shared_memory:
     name: SharedMemoryConnector
     extra:
+      talker_first_audio: true
       codec_streaming: true
       connector_get_sleep_s: 0.01
       connector_get_max_wait_first_chunk: 3000

vllm_omni/deploy/qwen3_tts_high_concurrency.yaml

@@ -24,6 +24,7 @@ connectors:
   connector_of_shared_memory:
     name: SharedMemoryConnector
     extra:
+      talker_first_audio: true
       codec_streaming: true
       connector_get_sleep_s: 0.01
       connector_get_max_wait_first_chunk: 3000

vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml

@@ -15,6 +15,7 @@ platforms:
 connectors:
   connector_of_shared_memory:
     extra:
+      talker_first_audio: true
       # Keep unqualified prefix-graph optimizations disabled in this profile.
       code_predictor_prefix_graphs: false
       ref_audio_artifact_cache_max_entries: 1024

vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml

@@ -1,15 +1,36 @@
-# Opt-in single-GPU throughput profile, measured on H200 at concurrency 128
-# with Qwen3-TTS-12Hz-1.7B-Base. Warm the full pipeline before benchmarking.
-# B8 uses more graph memory than B4; this is not a low-latency profile.
-# MPS is optional and is not enabled by this file.
+# Experimental single-GPU CUDA/BF16 throughput profile.
+# Optimizations: frame-local predictor KV, fused sampling, time-major codec
+# convolutions, eager residual prediction and first-frame delivery; B1-B8 codec
+# graphs, FULL Talker graphs and 128 Talker slots. Floating-point results differ:
+# validate audio quality for your workload. Ordinary profiles do not enable MPS.
+# cuda_mps starts one private NVIDIA MPS daemon shared by this pipeline's stages
+# and stops it after they exit. An existing CUDA_MPS_PIPE_DIRECTORY is reused
+# without taking ownership. Requires nvidia-cuda-mps-control on PATH.
+# Set cuda_mps: false in an overlay to disable it.
+# MPS can increase throughput under load but also increase first-packet and
+# first-audible-audio latency. Compare cuda_mps: false for latency-sensitive
+# workloads, especially at low request rates.
+# Warm the full pipeline before measuring latency or throughput.
 #
-# vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \
-#   --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml
+# vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+#   --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml \
+#   --api-server-count 4 --stage-init-timeout 1800 --init-timeout 2400
 base_config: qwen3_tts_high_concurrency_mrv2.yaml
+cuda_mps: true
 
 connectors:
   connector_of_shared_memory:
     extra:
+      talker_first_audio: true
+      codec_chunk_ramp: null
+      codec_chunk_adaptive: false
+      initial_codec_chunk_frames: 1
+      codec_chunk_frames: 25
+      codec_left_context_frames: 72
+      code_predictor_kv_cache: true
+      code_predictor_fused_sampling: true
+      decode_time_major_conv: true
+      decode_cudnn_benchmark: true
       decode_batch_max_size: 8
       decode_cudagraph_batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8]
 
@@ -17,7 +38,11 @@ stages:
   - stage_id: 0
     devices: "0"
     max_num_seqs: 128
+    compilation_config:
+      cudagraph_mode: FULL
+      cudagraph_capture_sizes: [1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128, 144, 160, 176, 192, 208, 224, 240, 256, 272, 288, 304, 320, 336, 352, 368, 384, 400, 416, 432, 448, 464, 480, 496, 512]
+      max_cudagraph_capture_size: 512
   - stage_id: 1
     devices: "0"
-    # Admission candidates are grouped by codec phase before B8 execution.
-    max_num_seqs: 32
+    max_num_seqs: 64
+    async_scheduling: false

vllm_omni/deploy/qwen3_tts_mrv2.yaml

@@ -18,3 +18,8 @@ stages:
   - stage_id: 0
     # Bound per-step prefill work, matching the high-concurrency TTS profile.
     max_num_batched_tokens: 512
+
+connectors:
+  connector_of_shared_memory:
+    extra:
+      talker_first_audio: true

vllm_omni/engine/cuda_mps.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Private NVIDIA MPS lifetime for explicitly opted-in local stage processes."""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+
+def physical_gpu_uuid(device: str) -> str:
+    """Resolve a physical ordinal/UUID through NVML without initializing CUDA."""
+    from vllm.utils.import_utils import import_pynvml
+
+    nvml = import_pynvml()
+    nvml.nvmlInit()
+    try:
+        if device.isdigit():
+            handle = nvml.nvmlDeviceGetHandleByIndex(int(device))
+        elif device.startswith("GPU-"):
+            handle = nvml.nvmlDeviceGetHandleByUUID(device)
+        else:
+            raise ValueError("cuda_mps requires one physical NVIDIA GPU ordinal or GPU UUID")
+        uuid = nvml.nvmlDeviceGetUUID(handle)
+        return uuid.decode() if isinstance(uuid, bytes) else str(uuid)
+    finally:
+        nvml.nvmlShutdown()
+
+
+class CudaMPSServer:
+    """Own a private control socket; never stop an operator's existing daemon."""
+
+    def __init__(self, gpu_uuid: str, *, pipe_directory: str | None = None) -> None:
+        control = shutil.which("nvidia-cuda-mps-control")
+        if control is None:
+            raise RuntimeError("cuda_mps requires nvidia-cuda-mps-control on PATH")
+        self._control = control
+        self._directory: Path | None = None
+        self._closed = False
+        # None inherits the parent; an explicit empty stage value selects a
+        # private daemon even when the parent names an operator-managed one.
+        self.operator_pipe_directory = (
+            os.environ.get("CUDA_MPS_PIPE_DIRECTORY") if pipe_directory is None else pipe_directory
+        ) or None
+        if self.operator_pipe_directory:
+            self.env = {"CUDA_MPS_PIPE_DIRECTORY": self.operator_pipe_directory, "CUDA_VISIBLE_DEVICES": gpu_uuid}
+            self._run(input="get_server_list\n")
+            logger.info("Using operator-managed MPS at %s for %s", self.operator_pipe_directory, gpu_uuid)
+            return
+
+        self._directory = root = Path(tempfile.mkdtemp(prefix="vllm-omni-mps-"))
+        (root / "pipe").mkdir()
+        (root / "log").mkdir()
+        self.env = {
+            "CUDA_VISIBLE_DEVICES": gpu_uuid,
+            "CUDA_MPS_PIPE_DIRECTORY": str(root / "pipe"),
+            "CUDA_MPS_LOG_DIRECTORY": str(root / "log"),
+        }
+        try:
+            self._run("-d")
+            self._run(input="get_server_list\n")
+        except BaseException:
+            # A failed/timed-out launcher may already have forked the daemon.
+            # Always address our private socket; keep it if cleanup also fails.
+            try:
+                self.close()
+            except Exception:
+                logger.exception("MPS startup cleanup failed; control files remain at %s", root)
+            raise
+        logger.info("Started private MPS for %s at %s", gpu_uuid, self.env["CUDA_MPS_PIPE_DIRECTORY"])
+
+    def _run(self, *args: str, input: str | None = None) -> subprocess.CompletedProcess[str]:
+        return subprocess.run(
+            [self._control, *args],
+            input=input,
+            env=os.environ | self.env,
+            check=True,
+            text=True,
+            capture_output=True,
+            timeout=15,
+        )
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        if self._directory is not None:
+            # Keep the control socket if quit fails, so cleanup can be retried.
+            self._run(input="quit\n")
+            shutil.rmtree(self._directory)
+            self._directory = None
+        self._closed = True

vllm_omni/engine/orchestrator.py

@@ -21,7 +21,7 @@ import time as _time
 from collections.abc import Awaitable, Callable, Coroutine
 from dataclasses import dataclass, field
 from http import HTTPStatus
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import janus
 import torch
@@ -34,6 +34,7 @@ from vllm.v1.engine import EngineCoreOutputs, FinishReason
 from vllm.v1.engine.exceptions import EngineDeadError
 from vllm.v1.metrics.stats import IterationStats
 
+from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY, FIRST_AUDIO_REQUIRED_KEY
 from vllm_omni.diffusion.data import is_diffusion_request_started_output
 from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks
 from vllm_omni.engine import OmniEngineCoreRequest
@@ -66,6 +67,9 @@ from vllm_omni.metrics.stat_logger import OmniPrometheusStatLogger
 from vllm_omni.metrics.utils import DIFFUSION_METRICS_ONLY_REQUEST_ID
 from vllm_omni.outputs import OmniRequestOutput
 
+if TYPE_CHECKING:
+    from vllm_omni.engine import OmniEngineCoreOutput
+
 logger = init_logger(__name__)
 
 
@@ -226,6 +230,10 @@ class OrchestratorRequestState:
     running_counter_registered: bool = False
     request_artifact_dirs: set[str] = field(default_factory=set)
     native_kv_transfer_id: str | None = None
+    # First-frame delivery and final-stage outputs use different channels.
+    upstream_first_audio: bool = False
+    pending_upstream_first_audio: tuple[int, OmniEngineCoreOutput] | None = None
+    pending_first_audio_outputs: list[tuple[int, int, Any]] = field(default_factory=list)
 
 
 @dataclass
@@ -272,7 +280,7 @@ class StreamingInputState:
 
         Fresh rather than shared: ``output_metadata`` is handed out by reference.
         """
-        segment = self.segments.get(stage_id)
+        segment = self.segments.get(stage_id) if stage_id is not None else None
         return segment if segment is not None else StreamingSegmentState()
 
 
@@ -440,9 +448,13 @@ class OrchestratorBase:
         # Start membership watcher if distributed mode is active.
         membership_watcher: asyncio.Task[None] | None = None
         if self._membership is not None:
+
+            async def cleanup_member_requests(ids: list[str]) -> None:
+                await self._cleanup_request_ids(ids, abort=True)
+
             self._membership.install_unregister_handlers(
                 output_queue=self.output_async_queue,
-                cleanup_callback=lambda ids: self._cleanup_request_ids(ids, abort=True),
+                cleanup_callback=cleanup_member_requests,
                 replica_removed_callback=self._remove_stage_replica_waiting,
             )
             membership_watcher = self._membership.start()
@@ -528,7 +540,7 @@ class OrchestratorBase:
                 # process exit as EngineDeadError during teardown.
                 for pool in self.stage_pools:
                     for client in pool.clients:
-                        if hasattr(client, "_shutting_down"):
+                        if client is not None and hasattr(client, "_shutting_down"):
                             client._shutting_down = True
                 # Stage teardown runs once in run()'s finally after the
                 # orchestration loop observes _shutdown_event and exits.
@@ -786,6 +798,7 @@ class OrchestratorBase:
         ``_finish_raw_terminal_requests`` once routing is done.
         """
         pool = self.stage_pools[stage_id]
+        await self._route_upstream_first_audio(stage_id, replica_id, raw_outputs)
         await self._handle_kv_ready_raw_outputs(stage_id, raw_outputs)
         for eco in raw_outputs.outputs:
             # Emit kv_wait_s before _handle_kv_ready_raw_outputs'
@@ -800,7 +813,8 @@ class OrchestratorBase:
                     kv_params.get("connector_type") or "unknown",
                     float(kv_wait_s),
                 )
-            req_state = self.request_states.get(getattr(eco, "request_id", None))
+            request_id = getattr(eco, "request_id", None)
+            req_state = self.request_states.get(request_id) if request_id is not None else None
             if req_state is None:
                 continue
             if (
@@ -861,6 +875,97 @@ class OrchestratorBase:
             )
         return processed
 
+    async def _route_upstream_first_audio(self, stage_id: int, replica_id: int, raw_outputs: Any) -> None:
+        """Order explicitly marked first audio ahead of a codec that skipped it.
+
+        The sender and codec run independently. Keep overtaking codec outputs,
+        including terminal outputs, on the request until its first PCM arrives.
+        Cancellation and stage errors use ordinary request cleanup.
+        """
+        kept = []
+        final_output = self.stage_pools[stage_id].final_output
+        for eco in raw_outputs.outputs:
+            mm = getattr(eco, "multimodal_output", None)
+            is_first = isinstance(mm, dict) and bool(mm.pop(FIRST_AUDIO_KEY, False))
+            requires_first = isinstance(mm, dict) and bool(mm.pop(FIRST_AUDIO_REQUIRED_KEY, False))
+            request_id = getattr(eco, "request_id", None)
+            req_state = self.request_states.get(request_id) if request_id is not None else None
+            if is_first:
+                assert isinstance(mm, dict)
+                # Only the dedicated first-frame channel is intercepted.
+                if (
+                    req_state is None
+                    or req_state.upstream_first_audio
+                    or req_state.pending_upstream_first_audio is not None
+                    or stage_id + 1 > req_state.final_stage_id
+                ):
+                    continue
+                audio = mm.get("model_outputs")
+                if not isinstance(audio, torch.Tensor) or audio.numel() == 0:
+                    await self._handle_stage_error(stage_id, eco, error="Invalid or empty first-audio payload")
+                    continue
+                req_state.pending_upstream_first_audio = (stage_id, eco)
+                await self._flush_upstream_first_audio(req_state)
+                continue
+            if (
+                req_state is not None
+                and final_output
+                and not req_state.upstream_first_audio
+                and (requires_first or req_state.pending_first_audio_outputs)
+                and getattr(eco, "finish_reason", None) != FinishReason.ERROR
+            ):
+                req_state.pending_first_audio_outputs.append((stage_id, replica_id, eco))
+                continue
+            kept.append(eco)
+        raw_outputs.outputs = kept
+
+    async def _flush_upstream_first_audio(self, req_state: OrchestratorRequestState) -> None:
+        """Feed upstream PCM through the codec's normal output-kind contract.
+
+        The source can finish its first frame while downstream prewarm is
+        still selecting/registering a replica. Retain that frame until the
+        codec owns an output-processor state, then accumulate it before any
+        codec output that omitted it. This preserves FINAL_ONLY/CUMULATIVE
+        waveforms and downstream consumers such as forced alignment.
+        """
+        from vllm_omni.engine import OmniEngineCoreOutputs
+
+        pending = req_state.pending_upstream_first_audio
+        if pending is None or self.request_states.get(req_state.request_id) is not req_state:
+            return
+        source_stage, first_output = pending
+        codec_stage = source_stage + 1
+        pool = self.stage_pools[codec_stage]
+        replica_id = pool.get_bound_replica_id(req_state.request_id)
+        if replica_id is None or req_state.request_id not in pool.output_processor.request_states:
+            return
+        req_state.pending_upstream_first_audio = None
+        req_state.upstream_first_audio = True
+        try:
+            processed = await pool.process_llm_raw_outputs(replica_id, OmniEngineCoreOutputs(outputs=[first_output]))
+            await self._handle_processed_outputs(codec_stage, replica_id, processed)
+        except EngineDeadError as error:
+            await self._handle_dead_replica(codec_stage, replica_id, error)
+            return
+        except Exception as error:
+            # A promised first frame must either enter the codec accumulator
+            # or fail its request; leaving it pending would stall completion.
+            logger.exception("First-audio output processing failed for %s", req_state.request_id)
+            await self._handle_stage_error(codec_stage, first_output, error=str(error))
+            return
+        if self.request_states.get(req_state.request_id) is not req_state:
+            return
+        ready, req_state.pending_first_audio_outputs = req_state.pending_first_audio_outputs, []
+        for pending_stage, pending_replica, output in ready:
+            if self.request_states.get(req_state.request_id) is not req_state:
+                break
+            terminal: set[str] = set()
+            processed = await self._process_llm_stage_outputs(
+                pending_stage, pending_replica, OmniEngineCoreOutputs(outputs=[output]), terminal
+            )
+            await self._handle_processed_outputs(pending_stage, pending_replica, processed)
+            await self._finish_raw_terminal_requests(pending_stage, pending_replica, terminal)
+
     def _absorb_diffusion_metrics(self, stage_id: int, replica_id: int, diffusion_output: Any) -> bool:
         """Drain a diffusion output's piggybacked metrics.
 
@@ -1543,7 +1648,7 @@ class OrchestratorBase:
             self._pd_kv_params.pop(request_id, None)
             req_state = self.request_states.pop(request_id, None)
             if req_state is not None:
-                cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", ()))
+                cleanup_request_artifact_dirs(getattr(req_state, "request_artifact_dirs", set()))
             if req_state is not None and req_state.running_counter_registered and self._running_counter is not None:
                 self._running_counter.decrement()
                 req_state.running_counter_registered = False
@@ -2201,7 +2306,7 @@ class OrchestratorBase:
                 )
                 return
             diffusion_source_outputs = [output, *companion_outputs]
-            if next_client.custom_process_input_func is not None:
+            if next_client is not None and next_client.custom_process_input_func is not None:
                 _t_ar2d = _time.perf_counter()
                 _fn = next_client.custom_process_input_func
                 _extra_kwargs: dict[str, Any] = {}
@@ -2386,7 +2491,9 @@ class OrchestratorBase:
             req_state.streaming.source_token_decoder = decode
 
         try:
-            next_inputs = next_client.process_engine_inputs(
+            assert next_client is not None
+            process_inputs = getattr(next_client, "process_engine_inputs")
+            next_inputs = process_inputs(
                 source_outputs,
                 req_state.prompt,
                 streaming_context=req_state.streaming,
@@ -2637,6 +2744,11 @@ class OrchestratorBase:
                 req_state,
             )
 
+            if req_state.pending_upstream_first_audio is not None or req_state.upstream_first_audio:
+                await self._flush_upstream_first_audio(req_state)
+                if self.request_states.get(request_id) is not req_state:
+                    return False
+
             # async_chunk pre-submit fires per stage edge (N-1 -> N). Source
             # replica is stage 0's bound replica (single-replica thinker in
             # all current configs); fall back to 0 if unknown.

vllm_omni/engine/stage_engine_core_proc.py

@@ -79,6 +79,23 @@ def _signal_exit_code(signum: int) -> int:
     return _SIGNAL_EXIT_BASE + signum
 
 
+def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: Any) -> bool:
+    """Let a TP1 in-process runner that decodes first audio deliver it as its own output."""
+    if not isinstance(model_executor, UniProcExecutor):
+        return False
+    worker = getattr(getattr(model_executor, "driver_worker", None), "worker", None)
+    model_runner = getattr(worker, "model_runner", None)
+    model_state = getattr(model_runner, "model_state", None)
+    model = getattr(model_runner, "model", None)
+    if getattr(model, "first_frame_decoder", None) is None or not hasattr(model_state, "set_first_audio_sink"):
+        return False
+    from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink
+
+    assert model_state is not None
+    model_state.set_first_audio_sink(engine_output_queue_sink(output_queue, scheduler))
+    return True
+
+
 def _bind_native_data_plane_ready_sink(model_executor: Any, scheduler: Any) -> bool:
     """Bind the TP1 in-process runner control plane directly to its scheduler."""
     if not isinstance(model_executor, UniProcExecutor):
@@ -107,6 +124,7 @@ class StageEngineCoreProc(EngineCoreProc):
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
         super().__init__(*args, **kwargs)
+        _bind_first_audio_sink(self.model_executor, self.output_queue, self.scheduler)
         if _bind_native_data_plane_ready_sink(self.model_executor, self.scheduler):
             logger.info("Bound native MRv2 connector readiness directly to the scheduler inbox.")
 

vllm_omni/engine/stage_init_utils.py

@@ -1041,14 +1041,17 @@ def stage_runtime_env(stage_id: int, runtime_cfg: Any) -> Generator[None, None,
             return
 
     previous_env: dict[str, str | None] = {}
-    for key, value in runtime_env.items():
-        env_key = str(key)
-        previous_env[env_key] = os.environ.get(env_key)
-        os.environ[env_key] = str(value)
-
-    if previous_env:
-        logger.info("[stage_init] Stage-%s applied runtime env keys: %s", stage_id, sorted(previous_env))
     try:
+        for key, value in runtime_env.items():
+            env_key = str(key)
+            old_value = os.environ.get(env_key)
+            os.environ[env_key] = str(value)
+            # Track only successful writes, preserving the original value
+            # if distinct keys normalize to the same environment name.
+            previous_env.setdefault(env_key, old_value)
+
+        if previous_env:
+            logger.info("[stage_init] Stage-%s applied runtime env keys: %s", stage_id, sorted(previous_env))
         yield
     finally:
         for key, old_value in previous_env.items():
@@ -1107,7 +1110,7 @@ def _project_omni_stage_engine_args(
         "default_sampling_params",
         "has_sampling_extra_args",
     }
-    runtime_excluded_fields = {"devices", "num_replicas", "env", "num_gpus"}
+    runtime_excluded_fields = {"devices", "num_replicas", "env", "num_gpus", "cuda_mps"}
     if not is_diffusion:
         # These values configure OmniDiffusionConfig or its worker process;
         # OmniEngineArgs has no matching fields for LLM stages.

vllm_omni/engine/stage_runtime.py

@@ -30,6 +30,7 @@ from vllm_omni.distributed.omni_coordinator import (
     RandomBalancer,
     RoundRobinBalancer,
 )
+from vllm_omni.engine.cuda_mps import CudaMPSServer, physical_gpu_uuid
 from vllm_omni.engine.messages import (
     EngineQueueMessage,
     RegisterRemoteReplicaMessage,
@@ -205,6 +206,7 @@ class StageRuntime:
         # device groups can initialize concurrently.
         self._replica_launch_lock = threading.Lock()
         self._init_visible_devices_baseline: str | None = None
+        self._mps_servers: dict[str, CudaMPSServer] = {}
 
     @staticmethod
     def _client_addresses_from_zmq(addresses: Any) -> dict[str, str]:
@@ -423,6 +425,8 @@ class StageRuntime:
                             )
                         )
 
+                    with self._replica_launch_lock:
+                        mps_env = self._mps_environment(physical_devices, replica.metadata.runtime_cfg)
                     launch_context = launch_stage_replica(
                         vllm_config=replica.stage_vllm_config,
                         executor_class=replica.executor_class,
@@ -432,7 +436,7 @@ class StageRuntime:
                         stage_config=replica.stage_cfg,
                         omni_master_server=self._get_omni_master_server(),
                         omni_coordinator_address=self._get_coordinator_address(),
-                        stage_visible_devices=physical_devices,
+                        stage_visible_devices=mps_env.get("CUDA_VISIBLE_DEVICES", physical_devices),
                         spawn_device_lock=self._spawn_device_lock,
                         omni_parallel_stage_init=self._parallel_stage_init,
                         num_api_servers=num_api_servers,
@@ -441,7 +445,10 @@ class StageRuntime:
                     # Environment overlays are process-global. Serialize spawn;
                     # parallel initialization waits for READY outside this lock.
                     with self._replica_launch_lock:
-                        with stage_runtime_env(replica.metadata.stage_id, replica.metadata.runtime_cfg):
+                        with (
+                            stage_runtime_env(replica.metadata.stage_id, replica.metadata.runtime_cfg),
+                            stage_runtime_env(replica.metadata.stage_id, {"env": mps_env}),
+                        ):
                             stage_resources = launch_context.__enter__()
                         entered_contexts.append(launch_context)
                         if stage_resources is None:
@@ -514,6 +521,55 @@ class StageRuntime:
             if lock_fds:
                 release_device_locks(lock_fds)
 
+    @staticmethod
+    def _mps_enabled(runtime_cfg: Any) -> bool:
+        return bool(
+            runtime_cfg.get("cuda_mps", False)
+            if isinstance(runtime_cfg, Mapping)
+            else getattr(runtime_cfg, "cuda_mps", False)
+        )
+
+    def _validate_mps_topology(self, stage_plans: Sequence[LogicalStageInitPlan]) -> None:
+        for plan in stage_plans:
+            for replica in plan.replicas:
+                runtime_cfg = getattr(replica.stage_cfg, "runtime_config", getattr(replica.stage_cfg, "runtime", None))
+                if not self._mps_enabled(runtime_cfg):
+                    continue
+                if replica.launch_mode != "local" or replica.metadata.stage_type == "diffusion":
+                    raise ValueError("cuda_mps currently supports local EngineCore stages only")
+                if self._parallel_stage_init:
+                    raise ValueError(
+                        "cuda_mps requires parallel_stage_init=false for physical GPU initialization locks"
+                    )
+
+    def _mps_environment(self, devices: str | None, runtime_cfg: Any) -> dict[str, str]:
+        if not self._mps_enabled(runtime_cfg):
+            return {}
+        if not current_omni_platform.is_cuda() or devices is None or len(devices.split(",")) != 1:
+            raise ValueError("cuda_mps currently requires a local CUDA stage on exactly one explicit GPU")
+        runtime_env = runtime_cfg.get("env") if isinstance(runtime_cfg, Mapping) else getattr(runtime_cfg, "env", None)
+        pipe_directory = (
+            str(runtime_env["CUDA_MPS_PIPE_DIRECTORY"])
+            if isinstance(runtime_env, Mapping) and "CUDA_MPS_PIPE_DIRECTORY" in runtime_env
+            else None
+        )
+        uuid = physical_gpu_uuid(devices.strip())
+        if uuid not in self._mps_servers:
+            self._mps_servers[uuid] = CudaMPSServer(uuid, pipe_directory=pipe_directory)
+        server = self._mps_servers[uuid]
+        if pipe_directory is not None and (pipe_directory or None) != server.operator_pipe_directory:
+            raise ValueError(f"Conflicting CUDA_MPS_PIPE_DIRECTORY settings for stages sharing GPU {uuid}")
+        return server.env
+
+    def _close_mps_servers(self) -> None:
+        for uuid, server in list(self._mps_servers.items()):
+            try:
+                server.close()
+            except Exception:
+                logger.exception("Failed to close private MPS server for %s; retaining its control directory", uuid)
+            else:
+                del self._mps_servers[uuid]
+
     def shutdown(self) -> None:
         for pool in self.stage_pools:
             for client in pool.clients:
@@ -525,6 +581,7 @@ class StageRuntime:
         if self._stage_init_executor is not None:
             self._stage_init_executor.shutdown(wait=True, cancel_futures=True)
             self._stage_init_executor = None
+        self._close_mps_servers()
 
     def create_membership_controller(self) -> Any | None:
         """Return a distributed membership controller, if this runtime needs one."""
@@ -540,6 +597,7 @@ class StageRuntime:
             replicas_per_stage,
             replica_devices_map,
         )
+        self._validate_mps_topology(stage_plans)
         self._validate_native_kv_topology(stage_plans)
         return stage_plans
 
@@ -557,17 +615,20 @@ class StageRuntime:
 
     def _cleanup_after_initialize_failure(self) -> None:
         """Hook for runtimes that own extra infrastructure during init."""
-        return None
+        self._close_mps_servers()
 
     def _resolve_replica_physical_devices(self, stage_id: int, runtime_cfg: Any) -> str | None:
         if runtime_cfg is None:
             runtime_cfg = {}
         devices = runtime_cfg.get("devices") if hasattr(runtime_cfg, "get") else getattr(runtime_cfg, "devices", None)
-        return resolve_stage_physical_devices(
+        physical = resolve_stage_physical_devices(
             stage_id,
             devices,
             visible_baseline=self._init_visible_devices_baseline,
         )
+        if self._mps_enabled(runtime_cfg) and (physical is None or not physical.isdigit()):
+            raise ValueError("cuda_mps requires one numeric physical GPU for initialization locking")
+        return physical
 
     # ---- Internal methods ----
 
@@ -664,7 +725,7 @@ class StageRuntime:
             port = int(extra.get("bootstrap_port", 8998)) + replica_metadata.replica_id
             extra["bootstrap_addr"] = f"http://{kv_config.kv_ip}:{port}"
             if isinstance(replica_cfg, BaseVllmOmniStageConfig):
-                runtime_cfg = replica_cfg.runtime_config
+                runtime_cfg: Any = replica_cfg.runtime_config
                 runtime_cfg.env = {
                     **(runtime_cfg.env or {}),
                     "VLLM_MOONCAKE_BOOTSTRAP_PORT": str(port),
@@ -1374,6 +1435,7 @@ class DistStageRuntime(StageRuntime):
         self._stage_remote_factory_contexts = self._capture_stage_factory_contexts(stage_plans)
 
     def _cleanup_after_initialize_failure(self) -> None:
+        super()._cleanup_after_initialize_failure()
         self._cleanup_distributed_infra()
 
     def shutdown(self) -> None:

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -1077,6 +1077,64 @@ class CodePredictorWrapper(nn.Module):
             ]
         )
 
+    def _predict_step_logits(
+        self,
+        proj_buf: torch.Tensor,
+        bsz: int,
+        padded_bsz: int,
+        step: int,
+        is_npu_capturing: bool,
+    ) -> torch.Tensor:
+        """Run one residual-codebook step; model wrappers may specialize it."""
+        device = proj_buf.device
+        max_seq = self._num_groups + 1
+        model_fwd = self._compiled_model_fwd
+        lm_heads = self._lm_heads_list
+        graph_key: int | tuple[int, int] = padded_bsz
+        seq_len = max_seq
+        if self._prefix_reprefill_enabled:
+            actual_seq_len = step + 1
+            for prefix_seq_len in self._prefix_reprefill_seq_lens:
+                if prefix_seq_len < actual_seq_len:
+                    continue
+                prefix_key = (padded_bsz, prefix_seq_len)
+                if prefix_key in self._bucket_pos_ids:
+                    graph_key = prefix_key
+                    seq_len = prefix_seq_len
+                    break
+        pos_ids = self._bucket_pos_ids.get(graph_key)
+        if pos_ids is None:
+            pos_ids = (
+                torch.arange(seq_len, device=device, dtype=torch.long).unsqueeze(0).expand(padded_bsz, -1).contiguous()
+            )
+
+        # Use captured device graph if available, otherwise call compiled fn.
+        device_graph_entry = self._device_graphs.get(graph_key)
+
+        # Let the outer graph record the regular forward during capture;
+        # normal inference still uses the inner graph replay fast path.
+        if device_graph_entry is not None and not is_npu_capturing:
+            device_graph_entry[0].replay()
+            hidden_out = device_graph_entry[1]
+        else:
+            hidden_out = model_fwd(proj_buf[:padded_bsz, :seq_len, :], pos_ids)
+
+        return lm_heads[step - 1](hidden_out[:bsz, step, :])
+
+    def _sample_per_call(
+        self,
+        logits: torch.Tensor,
+        inv_temperature: float,
+        top_k: int,
+        generator: _GeneratorLike,
+        uniforms: torch.Tensor | None,
+    ) -> torch.Tensor:
+        scaled = logits * inv_temperature
+        if top_k > 0:
+            topk_vals, _ = scaled.topk(top_k, dim=-1)
+            scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
+        return self._sample_codes_gumbel(scaled, generator=generator, uniforms=uniforms)
+
     @torch.inference_mode()
     def forward(
         self,
@@ -1114,10 +1172,7 @@ class CodePredictorWrapper(nn.Module):
         self._ensure_buffers(device, dtype, padded_bsz)
 
         proj_buf = self._proj_buf
-        max_seq = num_groups + 1
         projection = self.small_to_mtp_projection
-        model_fwd = self._compiled_model_fwd
-        lm_heads = self._lm_heads_list
         codec_embeds = self._codec_embeds_list
         # torch-npu cannot replay an inner NPUGraph while the outer talker_mtp
         # graph is being captured. Capture state is constant for this forward.
@@ -1153,39 +1208,7 @@ class CodePredictorWrapper(nn.Module):
 
         # Autoregressive loop: predict layers 1..G-1
         for step in range(1, num_groups):
-            graph_key: int | tuple[int, int] = padded_bsz
-            seq_len = max_seq
-            if self._prefix_reprefill_enabled:
-                actual_seq_len = step + 1
-                for prefix_seq_len in self._prefix_reprefill_seq_lens:
-                    if prefix_seq_len < actual_seq_len:
-                        continue
-                    prefix_key = (padded_bsz, prefix_seq_len)
-                    if prefix_key in self._bucket_pos_ids:
-                        graph_key = prefix_key
-                        seq_len = prefix_seq_len
-                        break
-            pos_ids = self._bucket_pos_ids.get(graph_key)
-            if pos_ids is None:
-                pos_ids = (
-                    torch.arange(seq_len, device=device, dtype=torch.long)
-                    .unsqueeze(0)
-                    .expand(padded_bsz, -1)
-                    .contiguous()
-                )
-
-            # Use captured device graph if available, otherwise call compiled fn.
-            device_graph_entry = self._device_graphs.get(graph_key)
-
-            # Let the outer graph record the regular forward during capture;
-            # normal inference still uses the inner graph replay fast path.
-            if device_graph_entry is not None and not is_npu_capturing:
-                device_graph_entry[0].replay()
-                hidden_out = device_graph_entry[1]
-            else:
-                hidden_out = model_fwd(proj_buf[:padded_bsz, :seq_len, :], pos_ids)
-
-            logits = lm_heads[step - 1](hidden_out[:bsz, step, :])
+            logits = self._predict_step_logits(proj_buf, bsz, padded_bsz, step, is_npu_capturing)
 
             # Sample next code via Gumbel-max.
             #
@@ -1216,12 +1239,8 @@ class CodePredictorWrapper(nn.Module):
             else:
                 # "per_call" mode: temperature-scaled + top-k -> Gumbel-max
                 if use_sampling:
-                    scaled = logits * inv_temperature
-                    if top_k > 0:
-                        topk_vals, _ = scaled.topk(top_k, dim=-1)
-                        scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
                     step_uniforms = sample_uniforms[:, step - 1, :] if sample_uniforms is not None else None
-                    code = self._sample_codes_gumbel(scaled, generator=sample_generator, uniforms=step_uniforms)
+                    code = self._sample_per_call(logits, inv_temperature, top_k, sample_generator, step_uniforms)
                 else:
                     code = logits.argmax(dim=-1, keepdim=True)
 

vllm_omni/model_executor/models/qwen3_tts/cached_code_predictor.py

@@ -0,0 +1,133 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Frame-local KV reuse for the Qwen3-TTS residual codebooks.
+
+The scratch buffer belongs to one predictor worker, just like its projection
+buffer. Calls are serialized on the worker's CUDA stream. Every frame starts
+by overwriting positions 0 and 1; later queries can only read that frame's
+written positions. No request state survives between frames.
+"""
+
+from collections.abc import Sequence
+
+import torch
+
+from vllm_omni.model_executor.models.common.qwen3_code_predictor import (
+    CodePredictorAttention,
+    CodePredictorBaseModel,
+    CodePredictorDecoderLayer,
+    _rotate_half,
+)
+
+from .short_kv_attention import short_kv_attention
+
+
+def _attention_cached(
+    attn: CodePredictorAttention,
+    hidden_states: torch.Tensor,
+    position_embeddings: tuple[torch.Tensor, torch.Tensor],
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    write_positions: torch.Tensor,
+) -> torch.Tensor:
+    """Attend only the new positions, reusing earlier positions' keys and values.
+
+    ``k_cache``/``v_cache`` are ``[B, num_kv_heads, max_seq, head_dim]``;
+    the new positions' keys and values are written at ``write_positions``
+    (``[n]``). The attention kernel masks keys after each query position,
+    so unwritten slots are never read.
+    """
+    bsz, seq_len, _ = hidden_states.shape
+    qkv = attn.qkv_proj(hidden_states)
+    q_raw, k_raw, v_raw = attn._split_qkv(qkv)
+    q = attn.q_norm(q_raw.view(bsz, seq_len, attn.num_heads, attn.head_dim)).transpose(1, 2)
+    k = attn.k_norm(k_raw.view(bsz, seq_len, attn.num_kv_heads, attn.head_dim)).transpose(1, 2)
+    v = v_raw.view(bsz, seq_len, attn.num_kv_heads, attn.head_dim).transpose(1, 2)
+
+    cos, sin = position_embeddings
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+    q = (q * cos) + (_rotate_half(q) * sin)
+    k = (k * cos) + (_rotate_half(k) * sin)
+    k_cache.index_copy_(2, write_positions, k.to(k_cache.dtype))
+    v_cache.index_copy_(2, write_positions, v.to(v_cache.dtype))
+    attn_out = short_kv_attention(q, k_cache, v_cache, write_positions, attn.scaling)
+    attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1)
+    return attn.o_proj(attn_out)
+
+
+def _layer_cached(
+    layer: CodePredictorDecoderLayer,
+    hidden_states: torch.Tensor,
+    position_embeddings: tuple[torch.Tensor, torch.Tensor],
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    write_positions: torch.Tensor,
+) -> torch.Tensor:
+    residual = hidden_states
+    hidden_states = layer.input_layernorm(hidden_states)
+    hidden_states = _attention_cached(
+        layer.self_attn, hidden_states, position_embeddings, k_cache, v_cache, write_positions
+    )
+    hidden_states = residual + hidden_states
+    residual = hidden_states
+    hidden_states = layer.post_attention_layernorm(hidden_states)
+    hidden_states = layer.mlp(hidden_states)
+    return residual + hidden_states
+
+
+def forward_cached(
+    model: CodePredictorBaseModel,
+    inputs_embeds: torch.Tensor,
+    position_ids: torch.Tensor,
+    kv_caches: Sequence[tuple[torch.Tensor, torch.Tensor]],
+    write_positions: torch.Tensor,
+) -> torch.Tensor:
+    """``forward`` for the new positions only (``position_ids`` ``[B, n]``), reusing earlier K/V."""
+    hidden_states = inputs_embeds
+    with torch.amp.autocast(inputs_embeds.device.type, enabled=False):
+        position_embeddings = model.rotary_emb(hidden_states, position_ids)
+        for layer, (k_cache, v_cache) in zip(model.layers, kv_caches):
+            hidden_states = _layer_cached(layer, hidden_states, position_embeddings, k_cache, v_cache, write_positions)
+        hidden_states = model.norm(hidden_states)
+    return hidden_states
+
+
+class FrameLocalKVCache:
+    """Non-module execution state; checkpoint parameters remain on the predictor."""
+
+    def __init__(self, model: CodePredictorBaseModel, max_batch: int) -> None:
+        self.model = model
+        weight = next(model.parameters())
+        attn = model.layers[0].self_attn
+        self.buffer = torch.zeros(
+            len(model.layers),
+            2,
+            max_batch,
+            attn.num_kv_heads,
+            attn.max_seq,
+            attn.head_dim,
+            device=weight.device,
+            dtype=weight.dtype,
+        )
+        self.positions = {
+            step: torch.arange(0 if step == 1 else step, step + 1, device=weight.device)
+            for step in range(1, attn.max_seq - 1)
+        }
+        # Batch buckets share compiled kernels; outer MTP graphs still capture
+        # fixed addresses. This helper never creates nested CUDA graphs.
+        self.forward = torch.compile(forward_cached, dynamic=True, options={"epilogue_fusion": False})
+
+    def __call__(self, proj_buf: torch.Tensor, batch: int, step: int) -> torch.Tensor:
+        if batch > self.buffer.shape[2]:
+            raise ValueError("Qwen3-TTS cached predictor batch exceeds its configured capacity")
+        first = 0 if step == 1 else step
+        caches = [(layer[0, :batch], layer[1, :batch]) for layer in self.buffer]
+        positions = self.positions[step]
+        return self.forward(
+            self.model,
+            proj_buf[:batch, first : step + 1],
+            positions.unsqueeze(0).expand(batch, -1),
+            caches,
+            positions,
+        )

vllm_omni/model_executor/models/qwen3_tts/code_predictor_sampling.py

@@ -0,0 +1,75 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Top-k Gumbel sampling with the outer MTP graph's supplied uniforms."""
+
+import torch
+from vllm.triton_utils import tl, tldevice, triton
+
+
+@triton.jit
+def _sample(
+    logits,
+    uniforms,
+    output,
+    ls: tl.constexpr,
+    us: tl.constexpr,
+    lc: tl.constexpr,
+    uc: tl.constexpr,
+    vocab: tl.constexpr,
+    top_k: tl.constexpr,
+    topk_block: tl.constexpr,
+    inv_temperature: tl.constexpr,
+    block: tl.constexpr,
+):
+    row = tl.program_id(0)
+    idx = tl.arange(0, block)
+    value = tl.load(logits + row * ls + idx * lc, idx < vocab, other=-float("inf"))
+    scaled = (value.to(tl.float32) * inv_temperature).to(value.dtype).to(tl.float32)
+    if top_k > 0:
+        if top_k == 1:
+            threshold = tl.max(scaled, 0)
+        else:
+            ordered = tl.topk(scaled, topk_block)
+            threshold = tl.sum(tl.where(tl.arange(0, topk_block) == top_k - 1, ordered, 0))
+        scaled = tl.where(scaled < threshold, -float("inf"), scaled)
+    u = tl.load(uniforms + row * us + idx * uc, idx < vocab, other=0.5)
+    score = scaled - tldevice.log(-tldevice.log(u))
+    score = tl.where(idx < vocab, score, -float("inf"))
+    max_score = tl.max(score, 0)
+    chosen = tl.min(tl.where(score == max_score, idx, block), 0)
+    first_nan = tl.min(tl.where(score != score, idx, block), 0)
+    tl.store(output + row, tl.where(first_nan < block, first_nan, chosen))
+
+
+def sample_code_topk_gumbel(
+    logits: torch.Tensor, uniforms: torch.Tensor, top_k: int, inv_temperature: float
+) -> torch.Tensor:
+    """Match BF16 scale rounding, top-k ties, and first-index Gumbel argmax.
+
+    Random values are supplied by the outer graph, so RNG state and seeded
+    per-request streams are unchanged. CUDA libdevice logs match ATen's logs.
+    """
+    if logits.shape != uniforms.shape:
+        raise ValueError("sampling uniforms must match logits shape")
+    if logits.device != uniforms.device or not logits.is_cuda:
+        raise ValueError("logits and uniforms must be on the same CUDA device")
+    if top_k > logits.shape[-1]:
+        raise ValueError("top_k must not exceed the vocabulary size")
+    b, v = logits.shape
+    out = torch.empty((b, 1), device=logits.device, dtype=torch.int64)
+    _sample[(b,)](
+        logits,
+        uniforms,
+        out,
+        logits.stride(0),
+        uniforms.stride(0),
+        logits.stride(1),
+        uniforms.stride(1),
+        v,
+        top_k,
+        triton.next_power_of_2(max(1, top_k)),
+        inv_temperature,
+        triton.next_power_of_2(v),
+        num_warps=4,
+    )
+    return out

vllm_omni/model_executor/models/qwen3_tts/first_audio.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Qwen3-TTS first-frame decode eligibility."""
+
+from typing import Any
+
+import torch
+
+from vllm_omni.platforms import current_omni_platform
+
+
+def talker_first_audio_enabled(vllm_config: Any) -> bool:
+    """First-frame delivery requires CUDA MRv2 streaming with an in-process Talker.
+
+    Unsupported runners retain the regular codec path. Prefix-cache replay
+    and distributed Talkers likewise retain the regular path.
+    """
+    from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor
+
+    extra = Predictor._stage_connector_extra_config(vllm_config)
+    model = vllm_config.model_config
+    parallel = vllm_config.parallel_config
+    return (
+        Predictor._parse_bool_config(extra.get("talker_first_audio"))
+        and current_omni_platform.is_cuda()
+        and torch.device(vllm_config.device_config.device).type == "cuda"
+        and bool(getattr(model, "use_v2_model_runner", False))
+        and bool(getattr(model, "async_chunk", False))
+        and parallel.tensor_parallel_size == 1
+        and parallel.pipeline_parallel_size == 1
+        and parallel.distributed_executor_backend in (None, "uni")
+        and not vllm_config.cache_config.enable_prefix_caching
+    )

vllm_omni/model_executor/models/qwen3_tts/first_frame_decoder.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""First codec frame -> PCM inside the Talker process.
+
+A stream's first frame has no decoder context, so its waveform does not depend
+on the Code2Wav stage's streaming state. Decoding it where the frame is
+produced removes the stage hop from time to first audio. Code2Wav advances
+its causal state without emitting that delivered frame again. Only streams
+without reference-code context use this path.
+
+The decode runs beside the Talker's next step, so its latency is set by how
+many dependent kernels it issues. A single frame with no context attends to
+itself only: each softmax is over one key and equals 1 exactly, so attention
+returns the value projection. The single-frame path therefore
+skips q/k, RoPE, masks and softmax (about two thirds of the kernels).
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+import torch.nn as nn
+from torch.cuda import CUDAGraph
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.model_loader import DefaultModelLoader
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
+
+from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config
+from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Decoder
+
+logger = init_logger(__name__)
+
+_SUBFOLDER = "speech_tokenizer"
+
+
+class Qwen3TTSFirstFrameDecoder(nn.Module):
+    def __init__(self, model_path: str) -> None:
+        super().__init__()
+        self.model_path = model_path
+        config = Qwen3TTSTokenizerV2Config.from_pretrained(model_path, subfolder=_SUBFOLDER)
+        self.decoder = Qwen3TTSTokenizerV2Decoder._from_config(config.decoder_config)
+        self.decoder.eval()
+        self.num_quantizers = int(config.decoder_config.num_quantizers)
+        self.sample_rate = int(config.output_sample_rate)
+        self._graphs: dict[int, tuple[CUDAGraph, torch.Tensor, torch.Tensor]] = {}
+
+    def load(self, vllm_config: VllmConfig) -> set[str]:
+        """Load decoder weights exactly as Code2Wav does; returns the loaded parameter names."""
+        loader = DefaultModelLoader(vllm_config.load_config)
+        source = DefaultModelLoader.Source(
+            model_or_path=self.model_path,
+            revision=vllm_config.model_config.revision,
+            subfolder=_SUBFOLDER,
+        )
+        loaded = AutoWeightsLoader(self).load_weights(
+            loader._get_weights_iterator(source),
+            mapper=WeightsMapper(orig_to_new_prefix={"encoder.": None}),
+        )
+        self.decoder.to(device=vllm_config.device_config.device, dtype=vllm_config.model_config.dtype)
+        if hasattr(self.decoder, "precompute_snake_caches"):
+            self.decoder.precompute_snake_caches()
+        from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor
+
+        extra = Predictor._stage_connector_extra_config(vllm_config)
+        if Predictor._parse_bool_config(extra.get("decode_time_major_conv")):
+            self.decoder.enable_time_major_conv()
+        return loaded
+
+    def _decode_single_frame(self, codes: torch.Tensor) -> torch.Tensor:
+        """``codes`` [n, num_quantizers, 1] -> waveform; equals ``_decode_xvec_first_chunk``."""
+        d = self.decoder
+        hidden = d.pre_conv(d.quantizer.decode(codes)).transpose(1, 2)
+        t = d.pre_transformer
+        h = t.input_proj(hidden)
+        for layer in t.layers[: t.config.num_hidden_layers]:
+            attention = layer.self_attn
+            values = attention.v_proj(layer.input_layernorm(h))
+            if attention.num_key_value_groups != 1:
+                values = values.reshape(*values.shape[:-1], -1, attention.head_dim)
+                values = values.repeat_interleave(attention.num_key_value_groups, dim=-2).flatten(-2)
+            h = h + layer.self_attn_layer_scale(attention.o_proj(values))
+            h = h + layer.mlp_layer_scale(layer.mlp(layer.post_attention_layernorm(h)))
+        return d._conv_decode(t.output_proj(t.norm(h))).clamp(min=-1, max=1)
+
+    @torch.inference_mode()
+    def capture(self, batch_sizes: Sequence[int] = (1, 2, 4, 8)) -> None:
+        device = next(self.decoder.parameters()).device
+        for batch_size in sorted(set(batch_sizes)):
+            static_input = torch.zeros(batch_size, self.num_quantizers, 1, dtype=torch.long, device=device)
+            self._decode_single_frame(static_input)
+            torch.accelerator.synchronize(device)
+            graph = CUDAGraph()
+            # A separate pool prevents overlap with Talker/MTP graph allocations.
+            with torch.cuda.graph(graph, pool=torch.cuda.graph_pool_handle()):
+                static_output = self._decode_single_frame(static_input)
+            self._graphs[batch_size] = (graph, static_input, static_output)
+        logger.info("Captured Talker first-frame decoder graphs for batch sizes %s", sorted(self._graphs))
+
+    @torch.inference_mode()
+    def decode(self, codes: torch.Tensor) -> torch.Tensor:
+        """``codes`` [n, num_quantizers] long on device -> float32 PCM [n, samples] (new tensor)."""
+        if not self._graphs:
+            return self._decode_single_frame(codes.unsqueeze(-1))[:, 0, :].float()
+        outputs = []
+        largest = max(self._graphs)
+        for start in range(0, int(codes.shape[0]), largest):
+            chunk = codes[start : start + largest]
+            rows = int(chunk.shape[0])
+            batch_size = min(size for size in self._graphs if size >= rows)
+            graph, static_input, static_output = self._graphs[batch_size]
+            static_input.zero_()
+            static_input[:rows, :, 0].copy_(chunk)
+            graph.replay()
+            outputs.append(static_output[:rows, 0, :].to(dtype=torch.float32, copy=True))
+        return torch.cat(outputs, dim=0)

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py

@@ -6,6 +6,7 @@ from __future__ import annotations
 import os
 from collections import Counter
 from collections.abc import Iterable
+from contextlib import nullcontext
 from typing import Any
 
 import torch
@@ -16,6 +17,7 @@ from vllm.logger import init_logger
 from vllm.model_executor.model_loader import DefaultModelLoader
 from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
 
+from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import parse_chunk_ramp
 
@@ -112,6 +114,7 @@ class Qwen3TTSCode2Wav(nn.Module):
         dec_config = tok_config.decoder_config
         self.decoder = Qwen3TTSTokenizerV2Decoder._from_config(dec_config)
         self.decoder.eval()
+
         self._num_quantizers = int(dec_config.num_quantizers)
         self._output_sample_rate = int(tok_config.output_sample_rate)
         self._total_upsample = int(self.decoder.total_upsample)
@@ -302,6 +305,7 @@ class Qwen3TTSCode2Wav(nn.Module):
         valid_indices: list[int] = []
         left_context_size = [0] * len(request_ids_list)
         ref_context_size = [0] * len(request_ids_list)
+        first_audio_flags = [False] * len(request_ids_list)
         segment_finished_flags = [False] * len(request_ids_list)
         request_state_ids: list[str | None] = [None] * len(request_ids_list)
         ref_context_request_ids: list[str | None] = [None] * len(request_ids_list)
@@ -336,6 +340,7 @@ class Qwen3TTSCode2Wav(nn.Module):
                 if not isinstance(info, dict):
                     continue
                 meta = info.get("meta", {})
+                first_audio_flags[i] = _meta_bool(meta.get("first_audio", False))
                 if "is_segment_finished" in meta:
                     segment_finished_flags[i] = _meta_bool(meta["is_segment_finished"])
                 if "left_context_size" in meta:
@@ -404,6 +409,7 @@ class Qwen3TTSCode2Wav(nn.Module):
                         )
                     state = {}
                     self._decoder_state_cache[state_req_id] = state
+                state.setdefault("skip_first_audio", first_audio_flags[i])
                 state.setdefault("prefix_frames", 0)
                 if state_req_id.startswith(_DUMMY_REQUEST_ID):
                     state["_is_dummy_run"] = True
@@ -515,6 +521,10 @@ class Qwen3TTSCode2Wav(nn.Module):
         if self._batch_stats_log_every > 0 and self._batch_stats_forwards % self._batch_stats_log_every == 0:
             self.log_decode_batch_stats()
 
+        first_audio_required = [torch.tensor(False)] * num_req
+        if request_states is not None:
+            for row, idx in enumerate(valid_indices):
+                first_audio_required[idx] = torch.tensor(bool(request_states[row].get("skip_first_audio", False)))
         audios: list[torch.Tensor] = [empty] * num_req
         srs = [sr_tensor] * num_req
 
@@ -538,7 +548,7 @@ class Qwen3TTSCode2Wav(nn.Module):
 
         return OmniOutput(
             text_hidden_states=None,
-            multimodal_outputs={"model_outputs": audios, "sr": srs},
+            multimodal_outputs={"model_outputs": audios, "sr": srs, FIRST_AUDIO_REQUIRED_KEY: first_audio_required},
         )
 
     def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput | tuple, **kwargs: Any) -> OmniOutput:
@@ -714,6 +724,11 @@ class Qwen3TTSCode2Wav(nn.Module):
                 raise ValueError(f"Invalid Qwen3-TTS Code2Wav config decode_batch_max_size={decode_batch_max_size}")
             self._decode_batch_max_size = decode_batch_max_size
             decode_enable_tf32 = _get_bool_config("decode_enable_tf32", False)
+            decode_time_major_conv = _get_bool_config("decode_time_major_conv", False)
+            self.decoder.capture_first_audio_state_only = self._async_chunk and _get_bool_config(
+                "talker_first_audio", False
+            )
+            decode_cudnn_benchmark = _get_bool_config("decode_cudnn_benchmark", False)
         else:
             codec_chunk_frames = 0
             codec_left_context_frames = 0
@@ -722,6 +737,8 @@ class Qwen3TTSCode2Wav(nn.Module):
             decode_cudagraph_batch_sizes = None
             decode_cudagraph_capture_sizes = None
             decode_enable_tf32 = False
+            decode_cudnn_benchmark = False
+            decode_time_major_conv = False
 
         if decode_enable_tf32 and device.type == "cuda":
             # PyTorch exposes TF32 controls as process-wide CUDA backend
@@ -738,21 +755,39 @@ class Qwen3TTSCode2Wav(nn.Module):
                 torch.get_float32_matmul_precision(),
             )
 
+        if decode_time_major_conv and device.type == "cuda":
+            self.decoder.enable_time_major_conv()
+
         self.decoder._initial_codec_chunk_frames = initial_codec_chunk_frames
         self.decoder._incremental_chunk_frames = codec_chunk_frames or 25
         self.decoder._incremental_chunk_ramp = list(codec_chunk_ramp or ())
 
         if hasattr(self.decoder, "enable_cudagraph") and device.type == "cuda":
             try:
-                self._maybe_enable_decoder_cudagraph(
-                    device=device,
-                    codec_chunk_frames=codec_chunk_frames,
-                    codec_left_context_frames=codec_left_context_frames,
-                    initial_codec_chunk_frames=initial_codec_chunk_frames,
-                    codec_chunk_ramp=codec_chunk_ramp,
-                    decode_cudagraph_batch_sizes=decode_cudagraph_batch_sizes,
-                    decode_cudagraph_capture_sizes=decode_cudagraph_capture_sizes,
+                # Autotune only during warmup/capture, then restore the process
+                # flags. The captured convolution algorithms remain in the
+                # graphs without changing later models' cuDNN policy.
+                autotune = (
+                    torch.backends.cudnn.flags(
+                        enabled=torch.backends.cudnn.enabled,
+                        benchmark=True,
+                        benchmark_limit=10,
+                        deterministic=torch.backends.cudnn.deterministic,
+                        allow_tf32=torch.backends.cudnn.allow_tf32,
+                    )
+                    if decode_cudnn_benchmark
+                    else nullcontext()
                 )
+                with autotune:
+                    self._maybe_enable_decoder_cudagraph(
+                        device=device,
+                        codec_chunk_frames=codec_chunk_frames,
+                        codec_left_context_frames=codec_left_context_frames,
+                        initial_codec_chunk_frames=initial_codec_chunk_frames,
+                        codec_chunk_ramp=codec_chunk_ramp,
+                        decode_cudagraph_batch_sizes=decode_cudagraph_batch_sizes,
+                        decode_cudagraph_capture_sizes=decode_cudagraph_capture_sizes,
+                    )
             except Exception:
                 logger.warning(
                     "Failed to enable CUDA Graph for Code2Wav decoder",

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py

@@ -1,12 +1,17 @@
-"""Qwen3-TTS Code Predictor -- thin wrapper over CodePredictorWrapper."""
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Qwen3-TTS predictor: per-call sampling and opt-in CUDA frame-local KV reuse."""
 
 from __future__ import annotations
 
-from collections.abc import Iterable
+from collections.abc import Iterable, Sequence
+from typing import TYPE_CHECKING
 
 import torch
 from vllm.config import VllmConfig
 from vllm.config.vllm import set_current_vllm_config
+from vllm.logger import init_logger
 
 from vllm_omni.model_executor.models.common.qwen3_code_predictor import (
     CodePredictorBaseModel,
@@ -17,6 +22,11 @@ from vllm_omni.platforms import current_omni_platform
 
 from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig
 
+if TYPE_CHECKING:
+    from .cached_code_predictor import FrameLocalKVCache
+
+logger = init_logger(__name__)
+
 # Backward-compat alias used by tests
 Qwen3TTSTalkerCodePredictorModelVLLM = CodePredictorBaseModel
 
@@ -48,6 +58,82 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp
         # Store talker_config for backward compat (accessed by some callers)
         self.talker_config = talker_config
         self._vllm_config = vllm_config
+        extra = self._stage_connector_extra_config(vllm_config)
+        self._kv_requested = self._parse_bool_config(extra.get("code_predictor_kv_cache"))
+        self._fused_sampling = self._parse_bool_config(extra.get("code_predictor_fused_sampling"))
+        self._frame_cache: FrameLocalKVCache | None = None
+
+    def _setup_compile(self) -> None:
+        if self._compiled_model_fwd is not None:
+            return
+        weight = next(self.model.parameters())
+        supported = (
+            current_omni_platform.is_cuda()
+            and weight.device.type == "cuda"
+            and weight.dtype == torch.bfloat16
+            and current_omni_platform.supports_torch_inductor()
+        )
+        if not self._kv_requested or not supported:
+            if self._kv_requested:
+                logger.warning("Qwen3-TTS frame-local KV requires CUDA BF16; using re-prefill")
+            return super()._setup_compile()
+
+        # Import CUDA-only kernels only when this model explicitly opts in.
+        from .cached_code_predictor import FrameLocalKVCache
+
+        self._model_dtype = weight.dtype
+        self._lm_heads_list = list(self.lm_head)
+        self._codec_embeds_list = list(self.model.codec_embedding)
+        self._bucket_sizes = self._batch_bucket_sizes()
+        max_batch = max(self._bucket_sizes)
+        self._ensure_buffers(weight.device, weight.dtype, max_batch)
+        cache = FrameLocalKVCache(self.model, max_batch)
+        with torch._dynamo.config.patch(
+            cache_size_limit=max(torch._dynamo.config.cache_size_limit, 2 * len(self._bucket_sizes))
+        ):
+            for batch in self._bucket_sizes:
+                for step in range(1, min(3, self._num_groups)):
+                    for _ in range(2):
+                        cache(self._proj_buf, batch, step)
+        self._synchronize_warmup(weight.device)
+        self._frame_cache = cache
+        # Preserve the shared wrapper's initialization sentinel. Its full
+        # re-prefill callable is unused while _frame_cache is active.
+        self._compiled_model_fwd = self.model.forward
+        logger.info("Qwen3-TTS frame-local KV warmed for batches %s", self._bucket_sizes)
+
+    def _predict_step_logits(
+        self,
+        proj_buf: torch.Tensor,
+        bsz: int,
+        padded_bsz: int,
+        step: int,
+        is_npu_capturing: bool,
+    ) -> torch.Tensor:
+        if self._frame_cache is None:
+            return super()._predict_step_logits(proj_buf, bsz, padded_bsz, step, is_npu_capturing)
+        hidden = self._frame_cache(proj_buf, padded_bsz, step)
+        return self._lm_heads_list[step - 1](hidden[:bsz, -1])
+
+    def _sample_per_call(
+        self,
+        logits: torch.Tensor,
+        inv_temperature: float,
+        top_k: int,
+        generator: torch.Generator | Sequence[torch.Generator | None] | None,
+        uniforms: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if (
+            self._fused_sampling
+            and current_omni_platform.is_cuda()
+            and logits.is_cuda
+            and uniforms is not None
+            and uniforms.dtype == torch.float32
+        ):
+            from .code_predictor_sampling import sample_code_topk_gumbel
+
+            return sample_code_topk_gumbel(logits, uniforms, top_k, inv_temperature)
+        return super()._sample_per_call(logits, inv_temperature, top_k, generator, uniforms)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
         """Load weights with vllm config context (required for VocabParallelEmbedding)."""

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -36,6 +36,7 @@ from vllm_omni.utils.speaker_cache import (
 from vllm_omni.worker.sampling_utils import get_tts_local_seed
 
 from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig
+from .first_audio import talker_first_audio_enabled
 from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder, resolve_x_vector_only
 from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
 from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config
@@ -361,6 +362,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         }
     )
 
+    # Set by the MRV2 model state once eager frames are enabled.
+    eager_frames_active = False
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
         self.vllm_config = vllm_config
@@ -404,6 +408,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # per-request additional_information.
         self.talker_mtp_output_key = ("codes", "audio")
         self.talker_mtp_graph_safe = True
+        # A frame's residual codebooks depend only on the Talker hidden and
+        # CB0 of the step that sampled it, so MRV2 may complete the frame at
+        # the end of that step instead of in the next step's preprocess.
+        self.mtp_eager_frames = talker_first_audio_enabled(vllm_config)
         # The runners bypass only the outer whole-MTP graph when explicit
         # generators are present, so seeded requests can still share one raw
         # batched MTP call with independent per-row streams.
@@ -722,6 +730,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
         model_config = getattr(getattr(self, "vllm_config", None), "model_config", None)
         async_chunk = bool(getattr(model_config, "async_chunk", False))
+        spans = kwargs.get("request_token_spans")
+        if async_chunk and self.eager_frames_active and spans is not None:
+            return self._make_eager_omni_output(hidden, info_dicts, spans)
         audio_codes_list: list[torch.Tensor] = []
         ref_code_len_segments: list[tuple[int | torch.Tensor, int, torch.device]] = []
         ref_code_list: list[torch.Tensor] = []
@@ -863,6 +874,44 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             mm.setdefault("meta", {})["codec_frame_valid"] = codec_frame_valid[:span_len]
         return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm)
 
+    def _make_eager_omni_output(
+        self,
+        hidden: torch.Tensor,
+        info_dicts: list[Any],
+        spans: list[tuple[int, int]],
+    ) -> OmniOutput:
+        """Async-chunk output when the runner fills frames after sampling.
+
+        Every row that sampled this step gets its codes and validity from the
+        runner's eager MTP; the rest (non-final prefill chunks) carry no frame,
+        so the per-request buffer rows need not be gathered here.
+        """
+        num_tokens = spans[-1][1] if spans else 0
+        q = int(self.talker_config.num_code_groups)
+        mm: OmniPayload = {
+            "codes": {"audio": torch.zeros((num_tokens, q), dtype=torch.long, device=hidden.device)},
+            "meta": {
+                "codec_frame_valid": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device),
+                "first_audio": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device),
+            },
+        }
+        ref_rows = [
+            index
+            for index, info in enumerate(info_dicts)
+            if isinstance(info, dict)
+            and isinstance(ref := info.get("codes", {}).get("ref"), torch.Tensor)
+            and ref.numel() > 0
+            and _should_publish_async_ref_codes(info)
+        ]
+        if ref_rows:
+            empty = torch.empty(0, dtype=torch.long)
+            ref_code_list = [empty] * len(info_dicts)
+            for index in ref_rows:
+                ref_code_list[index] = info_dicts[index]["codes"]["ref"]
+                _mark_async_ref_codes_published(info_dicts[index])
+            mm["codes"]["ref"] = ref_code_list
+        return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm)
+
     # -------------------- preprocess / postprocess --------------------
 
     def preprocess(
@@ -1087,7 +1136,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 f"preprocess_decode_batch expected {len(req_infos)} input ids, got {int(input_ids_flat.numel())}"
             )
 
-        codec_frame_valid_batch = (input_ids_flat >= 0) & (input_ids_flat < self._codebook_vocab_size)
+        eager = self.eager_frames_active
+        codec_frame_valid_batch = (
+            None if eager else (input_ids_flat >= 0) & (input_ids_flat < self._codebook_vocab_size)
+        )
         device = input_ids_flat.device
         dtype = self._embedding_dtype
         # Request-independent constant (see :meth:`_init_runtime_buffers`) —
@@ -1150,28 +1202,29 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 text_step = tts_pad_embed
                 next_text_offset = text_offset
 
-            last_hidden = hs.get("last")
-            if isinstance(last_hidden, torch.Tensor):
-                past_hidden = last_hidden.to(device=device, dtype=dtype).reshape(1, -1)
-            else:
-                # Match scalar preprocess(): EOS / async-scheduling races can
-                # arrive before postprocess has written last; zeros are filtered
-                # downstream the same way as the single-request path.
-                logger.warning_once(
-                    "Qwen3-TTS batched decode missing last hidden; zero-filling past_hidden "
-                    "to match scalar preprocess()"
-                )
-                past_hidden = torch.zeros_like(text_step)
-            past_hidden_list.append(past_hidden)
             text_step_list.append(text_step)
-
             info_update: dict[str, Any] = {
                 "meta": {
                     "talker_text_offset": int(next_text_offset),
                     "codec_streaming": codec_streaming,
-                    "codec_frame_valid": codec_frame_valid_batch[len(updates)].reshape(()),
                 },
             }
+            if not eager:
+                # Eager frames finished MTP last step: no hidden input or row validity needed.
+                last_hidden = hs.get("last")
+                if isinstance(last_hidden, torch.Tensor):
+                    past_hidden = last_hidden.to(device=device, dtype=dtype).reshape(1, -1)
+                else:
+                    # Match scalar preprocess(): EOS / async-scheduling races can
+                    # arrive before postprocess has written last; zeros are filtered
+                    # downstream the same way as the single-request path.
+                    logger.warning_once(
+                        "Qwen3-TTS batched decode missing last hidden; zero-filling past_hidden "
+                        "to match scalar preprocess()"
+                    )
+                    past_hidden = torch.zeros_like(text_step)
+                past_hidden_list.append(past_hidden)
+                info_update["meta"]["codec_frame_valid"] = codec_frame_valid_batch[len(updates)].reshape(())
             if trailing_text_update is not None:
                 info_update["hidden_states"] = {"trailing_text": trailing_text_update.detach()}
             updates.append(info_update)
@@ -1179,11 +1232,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         if input_embeds is None:
             input_embeds = self.embed_input_ids(input_ids_flat.reshape(-1, 1).to(torch.long))
         inputs_embeds_out = input_embeds.to(device=device, dtype=dtype).reshape(len(req_infos), -1)
+        text_steps = torch.cat(text_step_list, dim=0)
         return (
             input_ids_flat,
             inputs_embeds_out,
-            torch.cat(past_hidden_list, dim=0),
-            torch.cat(text_step_list, dim=0),
+            text_steps if eager else torch.cat(past_hidden_list, dim=0),
+            text_steps,
             updates,
         )
 
@@ -1382,6 +1436,12 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
         logger.info("Loaded %d weights for Qwen3TTSTalkerForConditionalGeneration", len(loaded))
         self._build_stacked_codec_embed()
+        if talker_first_audio_enabled(self.vllm_config):
+            from .first_frame_decoder import Qwen3TTSFirstFrameDecoder
+
+            self.first_frame_decoder = Qwen3TTSFirstFrameDecoder(self.model_path)
+            decoder_loaded = self.first_frame_decoder.load(self.vllm_config)
+            loaded = set(loaded) | {f"first_frame_decoder.{name}" for name in decoder_loaded}
         return loaded
 
     def _build_stacked_codec_embed(self) -> None:
@@ -1569,6 +1629,35 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
     mtp = talker_mtp
     get_mtp_seed = staticmethod(get_tts_local_seed)
 
+    # Postprocess only stores the last hidden for the deferred decode MTP.
+    eager_frames_need_postprocess = False
+
+    def eager_decode_settled(self, info: dict[str, Any]) -> bool:
+        """Decode preprocess is a no-op for this row: its trailing text is spent.
+
+        From then on every decode step feeds the frame embedding plus the pad
+        text step and changes no per-request state (the non-streaming
+        CustomVoice/VoiceDesign case after the first decode).
+        """
+        hs = info.get("hidden_states")
+        tail = hs.get("trailing_text") if isinstance(hs, dict) else None
+        meta = info.get("meta")
+        offset = meta.get("talker_text_offset", 0) if isinstance(meta, dict) else 0
+        return isinstance(tail, torch.Tensor) and tail.numel() == 0 and not offset
+
+    def eager_settled_text_step(self) -> torch.Tensor:
+        return self._tts_pad_embed.reshape(1, -1)
+
+    def capture_first_frame_graphs(self) -> None:
+        # Set by load_weights when talker_first_audio_enabled accepts the connector option.
+        decoder = getattr(self, "first_frame_decoder", None)
+        if decoder is not None:
+            decoder.capture()
+
+    def mtp_frame_valid(self, layer0: torch.Tensor) -> torch.Tensor:
+        """Rows whose CB0 is a codec id rather than EOS/special; matches decode preprocess."""
+        return (layer0 >= 0) & (layer0 < self._codebook_vocab_size)
+
     @property
     def mtp_output_key(self) -> tuple[str, str]:
         return self.talker_mtp_output_key

vllm_omni/model_executor/models/qwen3_tts/segmented_graph_wrapper.py

@@ -71,6 +71,7 @@ class CUDAGraphDecoderWrapper:
         self.combined_states: dict[tuple[str, int, int], dict] = {}
         self.icl_prefix_states: dict[int, dict] = {}
         self.xvec_prefix_states: dict[int, dict] = {}
+        self.xvec_prefix_state_only_states: dict[int, dict] = {}
         self.stateless_states: dict[tuple[int, int], dict] = {}
 
         self._device = None
@@ -230,6 +231,7 @@ class CUDAGraphDecoderWrapper:
             *self._retained_kv_caches(self.combined_states),
             *self._retained_kv_caches(self.icl_prefix_states),
             *self._retained_kv_caches(self.xvec_prefix_states),
+            *self._retained_kv_caches(self.xvec_prefix_state_only_states),
         ]
         batched_captures = sum(rows for rows, _ in retained)
 
@@ -248,7 +250,7 @@ class CUDAGraphDecoderWrapper:
                 rows_reason=(
                     f"rows across {len(retained)} retained caches in "
                     f"{len(self.combined_states)} suffix / {len(self.icl_prefix_states)} icl-prefix / "
-                    f"{len(self.xvec_prefix_states)} xvec-prefix captures"
+                    f"{len(self.xvec_prefix_states) + len(self.xvec_prefix_state_only_states)} xvec-prefix captures"
                 ),
                 allocation_note=(
                     "counted from caches transformers has materialized; captures whose tensors are still "
@@ -424,6 +426,14 @@ class CUDAGraphDecoderWrapper:
                     logger.info("Captured xvec prefix CUDA Graph for batch=%d", batch_size)
                 except Exception:
                     logger.warning("Failed to capture xvec prefix CUDA Graph for batch=%d", batch_size, exc_info=True)
+                if getattr(self.decoder, "capture_first_audio_state_only", False) and self.initial_chunk_frames == 1:
+                    try:
+                        self._capture_xvec_prefix(batch_size, device, dtype, state_only=True)
+                        logger.info("Captured xvec state-only CUDA Graph for batch=%d", batch_size)
+                    except Exception:
+                        logger.warning(
+                            "Failed to capture xvec state-only graph for batch=%d", batch_size, exc_info=True
+                        )
 
         logger.info(
             "Starting CUDA Graph warmup for %d shapes: batch_sizes=%s seq_lens=%s",
@@ -631,7 +641,9 @@ class CUDAGraphDecoderWrapper:
             "cache": caches,
         }
 
-    def _capture_xvec_prefix(self, batch_size: int, device: torch.device, dtype: torch.dtype) -> None:
+    def _capture_xvec_prefix(
+        self, batch_size: int, device: torch.device, dtype: torch.dtype, *, state_only: bool = False
+    ) -> None:
         static_input = torch.zeros(
             batch_size,
             self.num_quantizers,
@@ -639,8 +651,13 @@ class CUDAGraphDecoderWrapper:
             dtype=dtype,
             device=device,
         )
+        # The Talker delivers exactly one frame upstream: a one-frame first
+        # chunk then only needs its state; a longer one keeps the other frames.
+        first_chunk = (
+            self.decoder._decode_xvec_first_chunk_state_only if state_only else self.decoder._decode_xvec_first_chunk
+        )
         with torch.no_grad():
-            _ = self.decoder._decode_xvec_first_chunk(static_input, {})
+            _ = first_chunk(static_input, {})
         torch.accelerator.synchronize(device)
 
         config = self.decoder.config
@@ -662,13 +679,15 @@ class CUDAGraphDecoderWrapper:
         graph = CUDAGraph()
         with torch.no_grad():
             with torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()):
-                static_output = self.decoder._decode_xvec_first_chunk(static_input, caches)
+                static_output = first_chunk(static_input, caches)
 
-        self.xvec_prefix_states[batch_size] = {
+        states = self.xvec_prefix_state_only_states if state_only else self.xvec_prefix_states
+        states[batch_size] = {
             "graph": graph,
             "input": {"codes": static_input},
             "output": static_output,
             "cache": caches,
+            "state_only": state_only,
         }
 
     def _ensure_suffix_buffers(self, caches: dict) -> None:
@@ -808,7 +827,14 @@ class CUDAGraphDecoderWrapper:
         codes_list: list[torch.Tensor],
         request_caches: list[dict],
     ) -> list[torch.Tensor] | None:
-        available = set(getattr(self, "xvec_prefix_states", {}))
+        skip_flags = [bool(cache.get("skip_first_audio", False)) for cache in request_caches]
+        # Full PCM graphs remain available even when the Talker cannot use
+        # direct delivery (or only some requests qualify). The option is only
+        # a capture hint; request metadata decides which audio is still owed.
+        states = getattr(self, "xvec_prefix_states", {})
+        if skip_flags and all(skip_flags):
+            states = getattr(self, "xvec_prefix_state_only_states", {}) or states
+        available = set(states)
         if not available:
             self._record_graph_fallback("xvec_prefix:no_graph", len(codes_list))
             return None
@@ -835,7 +861,7 @@ class CUDAGraphDecoderWrapper:
             self._record_graph_fallback("xvec_prefix:no_graph", len(codes_list))
             return None
 
-        state = self.xvec_prefix_states[batch_size]
+        state = states[batch_size]
         static_input = state["input"]["codes"]
         static_input.zero_()
         for row, codes in enumerate(codes_list):
@@ -855,7 +881,12 @@ class CUDAGraphDecoderWrapper:
             cache["suffix_conv"] = static_caches["suffix_conv"][row : row + 1].clone()
             cache["suffix_frames"] = self.initial_chunk_frames
             self._ensure_suffix_buffers(cache)
-            outputs.append(state["output"][row : row + 1].clone())
+            if state.get("state_only"):
+                outputs.append(static_input.new_zeros((1, 1, 0), dtype=torch.float32))
+            elif skip_flags[row]:
+                outputs.append(state["output"][row : row + 1, :, self.decoder.total_upsample :].clone())
+            else:
+                outputs.append(state["output"][row : row + 1].clone())
         return outputs
 
     def _decode_suffix_batch(
@@ -983,6 +1014,8 @@ class CUDAGraphDecoderWrapper:
     def _decode_request_fallback(self, codes: torch.Tensor, cache: dict) -> torch.Tensor:
         if "suffix_quantized" not in cache:
             if int(cache["prefix_frames"]) == 0:
+                if cache.get("skip_first_audio", False):
+                    return self.decoder._decode_stream_first_chunk(codes, cache)
                 return self.decoder._decode_xvec_first_chunk(codes, cache)
             prefix_frames = int(cache["prefix_frames"])
             output = self.decoder._decode_icl_first_chunk(codes, cache, prefix_frames)

vllm_omni/model_executor/models/qwen3_tts/short_kv_attention.py

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Short causal attention for the scratch KV cache inside a code predictor."""
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _short_kv_attention(
+    q_ptr,
+    k_ptr,
+    v_ptr,
+    positions_ptr,
+    out_ptr,
+    q_stride_b: tl.constexpr,
+    q_stride_h: tl.constexpr,
+    q_stride_n: tl.constexpr,
+    k_stride_b: tl.constexpr,
+    k_stride_h: tl.constexpr,
+    k_stride_n: tl.constexpr,
+    v_stride_b: tl.constexpr,
+    v_stride_h: tl.constexpr,
+    v_stride_n: tl.constexpr,
+    heads: tl.constexpr,
+    num_queries: tl.constexpr,
+    num_keys: tl.constexpr,
+    head_dim: tl.constexpr,
+    queries_per_kv: tl.constexpr,
+    scale: tl.constexpr,
+    block_n: tl.constexpr,
+    block_d: tl.constexpr,
+):
+    b = tl.program_id(0)
+    h = tl.program_id(1)
+    qi = tl.program_id(2)
+    kvh = h // queries_per_kv
+    d = tl.arange(0, block_d)
+    n = tl.arange(0, block_n)
+    last = tl.load(positions_ptr + qi)
+    q = tl.load(q_ptr + b * q_stride_b + h * q_stride_h + qi * q_stride_n + d, d < head_dim, other=0).to(tl.float32)
+    k = tl.load(
+        k_ptr + b * k_stride_b + kvh * k_stride_h + n[:, None] * k_stride_n + d[None, :],
+        (n[:, None] < num_keys) & (n[:, None] <= last) & (d[None, :] < head_dim),
+        other=0,
+    ).to(tl.float32)
+    score = tl.sum(k * q[None, :], 1) * scale
+    score = tl.where((n < num_keys) & (n <= last), score, -float("inf"))
+    p = tl.exp(score - tl.max(score, 0))
+    p = p / tl.sum(p, 0)
+    v = tl.load(
+        v_ptr + b * v_stride_b + kvh * v_stride_h + n[:, None] * v_stride_n + d[None, :],
+        (n[:, None] < num_keys) & (n[:, None] <= last) & (d[None, :] < head_dim),
+        other=0,
+    ).to(tl.float32)
+    out = tl.sum(p[:, None] * v, 0)
+    tl.store(out_ptr + ((b * num_queries + qi) * heads + h) * head_dim + d, out, d < head_dim)
+
+
+def short_kv_attention(
+    q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, positions: torch.Tensor, scale: float
+) -> torch.Tensor:
+    """One/two-query GQA over a frame-local scratch cache with contiguous head dimensions.
+
+    The returned [B, H, Q, D] tensor has a contiguous [B, Q, H, D] view,
+    avoiding a separate transpose-copy before the output projection.
+    """
+    b, h, nq, d = q.shape
+    nk = k.shape[2]
+    out = torch.empty((b, nq, h, d), device=q.device, dtype=q.dtype)
+    _short_kv_attention[(b, h, nq)](
+        q,
+        k,
+        v,
+        positions,
+        out,
+        *q.stride()[:3],
+        *k.stride()[:3],
+        *v.stride()[:3],
+        h,
+        nq,
+        nk,
+        d,
+        h // k.shape[1],
+        scale,
+        triton.next_power_of_2(nk),
+        triton.next_power_of_2(d),
+        num_warps=4,
+    )
+    return out.transpose(1, 2)

vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/modeling_qwen3_tts_tokenizer_v2.py

@@ -997,6 +997,33 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
         if count > 0:
             logger.info("Precomputed exp caches for %d SnakeBeta activations", count)
 
+    # Set by enable_time_major_conv(); a plain object so it stays out of state_dict().
+    _time_major_stack = None
+
+    def enable_time_major_conv(self) -> None:
+        """Run the upsample + decoder convs time-major: one tensor-core GEMM per conv.
+
+        Call after the decoder reaches its device and dtype (the stack keeps
+        GEMM-layout weight copies) and before CUDA graph capture.
+        """
+        from .time_major_stack import TimeMajorConvStack
+
+        self._time_major_stack = TimeMajorConvStack(self)
+        logger.info("Qwen3-TTS decoder: time-major (GEMM) conv stack enabled")
+
+    def _conv_decode(self, hidden: torch.Tensor) -> torch.Tensor:
+        """Pre-transformer output ``[B, T, latent]`` -> unclamped waveform ``[B, 1, T * total_upsample]``."""
+        if self._time_major_stack is not None:
+            return self._time_major_stack(hidden)
+        hidden = hidden.permute(0, 2, 1)
+        for blocks in self.upsample:
+            for block in blocks:
+                hidden = block(hidden)
+        wav = hidden
+        for block in self.decoder:
+            wav = block(wav)
+        return wav
+
     def enable_cudagraph(
         self,
         capture_modes: tuple[str, ...] = ("icl", "xvec"),
@@ -1056,15 +1083,7 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
         hidden = self.pre_conv(hidden).transpose(1, 2)
 
         hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state
-        hidden = hidden.permute(0, 2, 1)
-        for blocks in self.upsample:
-            for block in blocks:
-                hidden = block(hidden)
-
-        wav = hidden
-        for block in self.decoder:
-            wav = block(wav)
-        return wav.clamp(min=-1, max=1)
+        return self._conv_decode(hidden).clamp(min=-1, max=1)
 
     def _decode_icl_first_chunk(
         self,
@@ -1126,20 +1145,32 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
         caches.update({"past_key_values": prefix_cache})
         caches.update({"prefix_hidden": prefix_hidden})
 
-        hidden = hidden.permute(0, 2, 1)
-
-        for blocks in self.upsample:
-            for block in blocks:
-                hidden = block(hidden)
-        wav = hidden
-
-        for block in self.decoder:
-            wav = block(wav)
-        wav = wav.clamp(min=-1, max=1)
-        return wav
+        return self._conv_decode(hidden).clamp(min=-1, max=1)
 
     def _decode_xvec_first_chunk(self, codes: torch.Tensor, caches: dict) -> torch.Tensor:
         """Initialize prefixless incremental state from the first codec chunk."""
+        suffix_hidden = self._init_xvec_first_chunk_state(codes, caches)
+        return self._conv_decode(suffix_hidden).clamp(min=-1, max=1)
+
+    def _decode_xvec_first_chunk_state_only(self, codes: torch.Tensor, caches: dict) -> torch.Tensor:
+        """State for later chunks when the first chunk's audio was delivered upstream."""
+        return self._init_xvec_first_chunk_state(codes, caches)
+
+    # Capture hint only: the raw connector option does not guarantee that
+    # a request received upstream audio. Only its skip_first_audio marker does.
+    capture_first_audio_state_only = False
+
+    def _decode_stream_first_chunk(self, codes: torch.Tensor, caches: dict) -> torch.Tensor:
+        if not caches.get("skip_first_audio", False):
+            return self._decode_xvec_first_chunk(codes, caches)
+        if int(codes.shape[-1]) <= 1:
+            self._decode_xvec_first_chunk_state_only(codes, caches)
+            return codes.new_zeros((codes.shape[0], 1, 0), dtype=torch.float32)
+        # The Talker delivered only the first frame; a longer first chunk
+        # (adaptive or ramped schedules) still owes the audio of the rest.
+        return self._decode_xvec_first_chunk(codes, caches)[..., self.total_upsample :]
+
+    def _init_xvec_first_chunk_state(self, codes: torch.Tensor, caches: dict) -> torch.Tensor:
         hidden = self.quantizer.decode(codes)
         caches["decoder_prefix_frames"] = 0
         caches["ref_hidden"] = hidden[:, :, :0]
@@ -1172,16 +1203,8 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
         ).last_hidden_state
         caches["past_key_values"] = empty_prefix_cache
         caches["prefix_hidden"] = suffix_hidden[:, :0, :]
-
-        hidden = suffix_hidden.permute(0, 2, 1)
-        for blocks in self.upsample:
-            for block in blocks:
-                hidden = block(hidden)
-        wav = hidden
-        for block in self.decoder:
-            wav = block(wav)
         caches["suffix_frames"] = int(codes.shape[-1])
-        return wav.clamp(min=-1, max=1)
+        return suffix_hidden
 
     def _decode_suffix(
         self,
@@ -1229,14 +1252,7 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
         hidden = torch.cat([caches["prefix_hidden"], suffix_hidden], dim=1)
 
         required_frames = min(hidden.shape[1], new_frames + _DOWNSTREAM_CONTEXT_FRAME)
-        hidden = hidden[:, -required_frames:, :].permute(0, 2, 1)
-        for blocks in self.upsample:
-            for block in blocks:
-                hidden = block(hidden)
-
-        wav = hidden
-        for block in self.decoder:
-            wav = block(wav)
+        wav = self._conv_decode(hidden[:, -required_frames:, :])
         wav = wav[..., -new_frames * self.total_upsample :].clamp(min=-1, max=1)
         return wav, next_quantized, next_conv
 
@@ -1300,7 +1316,7 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
                     f"prefix_frames={prefix_frames}, total_frames={codes.shape[-1]}"
                 )
             if prefix_frames == 0:
-                return self._decode_xvec_first_chunk(codes, caches)
+                return self._decode_stream_first_chunk(codes, caches)
             return self._decode_icl_first_chunk(codes, caches, prefix_frames)
         else:
             decoder_prefix_frames = int(caches.get("decoder_prefix_frames", prefix_frames))
@@ -1546,7 +1562,11 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
             return None
         batched_codes = torch.cat(codes_list, dim=0)
         batched_cache: dict[str, Any] = {}
-        output = self._decode_xvec_first_chunk(batched_codes, batched_cache)
+        skip_flags = {bool(cache.get("skip_first_audio", False)) for cache in request_caches}
+        if len(skip_flags) != 1:
+            return None
+        batched_cache["skip_first_audio"] = skip_flags.pop()
+        output = self._decode_stream_first_chunk(batched_codes, batched_cache)
         outputs: list[torch.Tensor] = []
         for row, cache in enumerate(request_caches):
             for key in ("ref_hidden", "ref_conv", "prefix_hidden", "suffix_quantized", "suffix_conv"):

vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/time_major_conv.py

@@ -0,0 +1,347 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Time-major causal 1-D convolution for streaming vocoders.
+
+Activations are ``[B, T, C]`` so that every dense convolution is one GEMM over
+``B * T`` rows (bf16 tensor cores through cuBLAS) instead of an NCT cuDNN
+convolution, which for these long, thin decoder layers picks non-tensor-core
+kernels and adds layout transforms around each call.
+
+- :func:`snake_im2col` gathers each output step's causal, dilated window into a
+  ``[B * T, K * C]`` row (tap-major), optionally applying ``bias`` and SnakeBeta
+  on the way in; ``K = 1`` is a standalone (bias +) SnakeBeta.
+- :func:`causal_conv` is the dense causal (dilated) convolution as an implicit
+  GEMM: each K-block of the reduction reads the input rows shifted by one tap,
+  so the ``K * C`` im2col matrix is never materialized; bias and an optional
+  residual are added in the epilogue.
+- :func:`transposed_overlap_add` finishes a transposed convolution of kernel
+  ``stride`` or ``2 * stride`` whose GEMM produced ``[B, T, K, C_out]``: the
+  overlap-add of consecutive inputs, the bias and the causal right trim.
+
+Weights come from ``nn.Conv1d``/``nn.ConvTranspose1d`` via :func:`conv_gemm_weight`
+and :func:`transposed_gemm_weight`.
+"""
+
+from __future__ import annotations
+
+import torch
+from vllm.triton_utils import HAS_TRITON, tl, triton
+
+
+def conv_gemm_weight(weight: torch.Tensor) -> torch.Tensor:
+    """``nn.Conv1d`` weight ``[C_out, C_in, K]`` -> tap-major GEMM weight ``[K * C_in, C_out]``."""
+    c_out, c_in, k = weight.shape
+    return weight.detach().permute(2, 1, 0).reshape(k * c_in, c_out).contiguous()
+
+
+def transposed_gemm_weight(weight: torch.Tensor) -> torch.Tensor:
+    """``nn.ConvTranspose1d`` weight ``[C_in, C_out, K]`` -> GEMM weight ``[C_in, K * C_out]``."""
+    c_in, c_out, k = weight.shape
+    return weight.detach().permute(0, 2, 1).reshape(c_in, k * c_out).contiguous()
+
+
+def snake_im2col_reference(
+    x: torch.Tensor,
+    kernel_size: int,
+    dilation: int = 1,
+    bias: torch.Tensor | None = None,
+    snake: tuple[torch.Tensor, torch.Tensor] | None = None,
+) -> torch.Tensor:
+    """Eager equivalent of :func:`snake_im2col` (also the CPU path)."""
+    if bias is not None:
+        x = x + bias
+    if snake is not None:
+        exp_alpha, inv_beta = snake
+        x = x + (inv_beta.float() * torch.sin(x.float() * exp_alpha.float()).square()).to(x.dtype)
+    if kernel_size == 1:
+        return x.reshape(-1, x.shape[-1])
+    t = x.shape[1]
+    xp = torch.nn.functional.pad(x, (0, 0, (kernel_size - 1) * dilation, 0))
+    cols = torch.cat([xp[:, j * dilation : j * dilation + t] for j in range(kernel_size)], dim=-1)
+    return cols.reshape(-1, cols.shape[-1])
+
+
+def causal_conv_reference(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    kernel_size: int,
+    dilation: int = 1,
+    residual: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Eager equivalent of :func:`causal_conv`: ``[B, T, C]`` -> ``[B, T, C_out]``."""
+    bsz, t, _ = x.shape
+    cols = snake_im2col_reference(x, kernel_size, dilation)
+    out = cols @ weight if bias is None else torch.addmm(bias, cols, weight)
+    out = out.view(bsz, t, -1)
+    return out if residual is None else out + residual
+
+
+def transposed_overlap_add_reference(z: torch.Tensor, stride: int, bias: torch.Tensor) -> torch.Tensor:
+    """Eager equivalent of :func:`transposed_overlap_add`."""
+    bsz, t, k, c_out = z.shape
+    out = z[:, :, :stride]
+    if k > stride:
+        out = out + torch.nn.functional.pad(z[:, :-1, stride:], (0, 0, 0, 0, 1, 0))
+    return out.reshape(bsz, t * stride, c_out) + bias
+
+
+if HAS_TRITON:
+
+    @triton.jit
+    def _snake_im2col_kernel(
+        x_ptr,
+        bias_ptr,
+        exp_alpha_ptr,
+        inv_beta_ptr,
+        out_ptr,
+        n_rows,
+        t_len,
+        c_len,
+        dilation,
+        K: tl.constexpr,  # noqa: N803
+        HAS_BIAS: tl.constexpr,  # noqa: N803
+        HAS_SNAKE: tl.constexpr,  # noqa: N803
+        BLOCK_R: tl.constexpr,  # noqa: N803
+        BLOCK_C: tl.constexpr,  # noqa: N803
+    ):
+        rows = tl.program_id(0) * BLOCK_R + tl.arange(0, BLOCK_R)
+        chans = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C)
+        row_mask = rows < n_rows
+        chan_mask = chans < c_len
+        t = rows % t_len
+        if HAS_BIAS:
+            bias = tl.load(bias_ptr + chans, mask=chan_mask, other=0.0)
+        if HAS_SNAKE:
+            exp_alpha = tl.load(exp_alpha_ptr + chans, mask=chan_mask, other=0.0).to(tl.float32)
+            inv_beta = tl.load(inv_beta_ptr + chans, mask=chan_mask, other=0.0).to(tl.float32)
+        out_mask = row_mask[:, None] & chan_mask[None, :]
+        for j in tl.static_range(K):
+            shift = (K - 1 - j) * dilation
+            valid = row_mask & (t >= shift)
+            src = rows - shift
+            x = tl.load(
+                x_ptr + src[:, None].to(tl.int64) * c_len + chans[None, :],
+                mask=valid[:, None] & chan_mask[None, :],
+                other=0.0,
+            )
+            if HAS_BIAS:
+                # Causal padding is applied to the conv input, after bias.
+                x = tl.where(valid[:, None], x + bias[None, :], 0.0).to(x_ptr.dtype.element_ty)
+            if HAS_SNAKE:
+                s = tl.sin(x.to(tl.float32) * exp_alpha[None, :])
+                x = x + (inv_beta[None, :] * s * s).to(x.dtype)
+            tl.store(out_ptr + rows[:, None].to(tl.int64) * (K * c_len) + j * c_len + chans[None, :], x, mask=out_mask)
+
+    @triton.autotune(
+        configs=[
+            triton.Config({"BLOCK_M": bm, "BLOCK_N": bn, "BLOCK_K": bk}, num_warps=w, num_stages=st)
+            for bm, bn, bk, w, st in (
+                (128, 128, 64, 8, 3),
+                (128, 64, 64, 4, 4),
+                (64, 128, 64, 4, 4),
+                (128, 32, 32, 4, 4),
+                (64, 64, 32, 4, 4),
+                (256, 32, 32, 8, 3),
+                (256, 64, 32, 8, 3),
+                (128, 16, 32, 4, 4),
+            )
+        ],
+        key=["row_bucket", "n_cols", "c_len", "K", "dilation"],
+    )
+    @triton.jit
+    def _causal_conv_gemm_kernel(
+        x_ptr,
+        w_ptr,
+        bias_ptr,
+        res_ptr,
+        out_ptr,
+        n_rows,
+        n_cols,
+        c_len,
+        t_len,
+        dilation,
+        row_bucket,  # autotune key only: power-of-two bucket of n_rows
+        K: tl.constexpr,  # noqa: N803
+        HAS_BIAS: tl.constexpr,  # noqa: N803
+        HAS_RES: tl.constexpr,  # noqa: N803
+        PRECISION: tl.constexpr,  # noqa: N803
+        BLOCK_M: tl.constexpr,  # noqa: N803
+        BLOCK_N: tl.constexpr,  # noqa: N803
+        BLOCK_K: tl.constexpr,  # noqa: N803
+    ):
+        # Row-major tiles, N fastest: neighbouring programs share the input rows.
+        pid = tl.program_id(0)
+        num_n = tl.cdiv(n_cols, BLOCK_N)
+        pid_m = pid // num_n
+        pid_n = pid % num_n
+        rows = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+        cols = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+        row_mask = rows < n_rows
+        col_mask = cols < n_cols
+        t = rows % t_len
+        k_blocks = tl.cdiv(c_len, BLOCK_K)
+        acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+        for kk in range(0, K * k_blocks):
+            tap = kk // k_blocks
+            chans = (kk % k_blocks) * BLOCK_K + tl.arange(0, BLOCK_K)
+            chan_mask = chans < c_len
+            shift = (K - 1 - tap) * dilation
+            a = tl.load(
+                x_ptr + (rows - shift)[:, None] * c_len + chans[None, :],
+                mask=(row_mask & (t >= shift))[:, None] & chan_mask[None, :],
+                other=0.0,
+            )
+            b = tl.load(
+                w_ptr + (tap * c_len + chans)[:, None] * n_cols + cols[None, :],
+                mask=chan_mask[:, None] & col_mask[None, :],
+                other=0.0,
+            )
+            acc = tl.dot(a, b, acc, input_precision=PRECISION)
+        if HAS_BIAS:
+            acc += tl.load(bias_ptr + cols, mask=col_mask, other=0.0).to(tl.float32)[None, :]
+        out_offsets = rows[:, None] * n_cols + cols[None, :]
+        out_mask = row_mask[:, None] & col_mask[None, :]
+        if HAS_RES:
+            acc += tl.load(res_ptr + out_offsets, mask=out_mask, other=0.0).to(tl.float32)
+        tl.store(out_ptr + out_offsets, acc.to(out_ptr.dtype.element_ty), mask=out_mask)
+
+    @triton.jit
+    def _transposed_overlap_add_kernel(
+        z_ptr,
+        bias_ptr,
+        out_ptr,
+        n_out_rows,
+        t_len,
+        c_len,
+        STRIDE: tl.constexpr,  # noqa: N803
+        K: tl.constexpr,  # noqa: N803
+        BLOCK_R: tl.constexpr,  # noqa: N803
+        BLOCK_C: tl.constexpr,  # noqa: N803
+    ):
+        rows = tl.program_id(0) * BLOCK_R + tl.arange(0, BLOCK_R)  # over B * T * STRIDE output steps
+        chans = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C)
+        row_mask = rows < n_out_rows
+        chan_mask = chans < c_len
+        mask = row_mask[:, None] & chan_mask[None, :]
+        tap = rows % STRIDE
+        src = rows // STRIDE  # input row b * T + t
+        acc = tl.load(
+            z_ptr + (src[:, None].to(tl.int64) * K + tap[:, None]) * c_len + chans[None, :], mask=mask, other=0.0
+        ).to(tl.float32)
+        if K > STRIDE:
+            prev_valid = row_mask & (src % t_len > 0)
+            acc += tl.load(
+                z_ptr + ((src[:, None].to(tl.int64) - 1) * K + tap[:, None] + STRIDE) * c_len + chans[None, :],
+                mask=prev_valid[:, None] & chan_mask[None, :],
+                other=0.0,
+            ).to(tl.float32)
+        acc += tl.load(bias_ptr + chans, mask=chan_mask, other=0.0).to(tl.float32)[None, :]
+        tl.store(
+            out_ptr + rows[:, None].to(tl.int64) * c_len + chans[None, :], acc.to(out_ptr.dtype.element_ty), mask=mask
+        )
+
+
+def _blocks(c_len: int) -> tuple[int, int]:
+    block_c = min(128, triton.next_power_of_2(c_len))
+    return max(1, 4096 // block_c), block_c
+
+
+def snake_im2col(
+    x: torch.Tensor,
+    kernel_size: int,
+    dilation: int = 1,
+    bias: torch.Tensor | None = None,
+    snake: tuple[torch.Tensor, torch.Tensor] | None = None,
+) -> torch.Tensor:
+    """``[B, T, C]`` -> ``[B * T, K * C]`` causal im2col with optional bias + SnakeBeta."""
+    if not HAS_TRITON or not x.is_cuda:
+        return snake_im2col_reference(x, kernel_size, dilation, bias, snake)
+    x = x.contiguous()
+    bsz, t, c = x.shape
+    n_rows = bsz * t
+    out = torch.empty(n_rows, kernel_size * c, dtype=x.dtype, device=x.device)
+    block_r, block_c = _blocks(c)
+    exp_alpha, inv_beta = snake if snake is not None else (x, x)
+    _snake_im2col_kernel[(triton.cdiv(n_rows, block_r), triton.cdiv(c, block_c))](
+        x,
+        bias if bias is not None else x,
+        exp_alpha,
+        inv_beta,
+        out,
+        n_rows,
+        t,
+        c,
+        dilation,
+        K=kernel_size,
+        HAS_BIAS=bias is not None,
+        HAS_SNAKE=snake is not None,
+        BLOCK_R=block_r,
+        BLOCK_C=block_c,
+    )
+    return out
+
+
+def causal_conv(
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor | None,
+    kernel_size: int,
+    dilation: int = 1,
+    residual: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Causal dilated conv ``[B, T, C]`` -> ``[B, T, C_out]`` with a :func:`conv_gemm_weight` weight.
+
+    ``residual`` (``[B, T, C_out]``) is added in the epilogue.
+    """
+    if not HAS_TRITON or not x.is_cuda:
+        return causal_conv_reference(x, weight, bias, kernel_size, dilation, residual)
+    x = x.contiguous()
+    bsz, t, c = x.shape
+    n_rows, n_cols = bsz * t, weight.shape[1]
+    if weight.shape[0] != kernel_size * c:
+        raise ValueError(f"conv weight has {weight.shape[0]} rows, expected {kernel_size} taps x {c} channels")
+    if max(n_rows * max(c, n_cols), weight.numel()) >= 2**31:
+        raise ValueError("causal_conv offsets are 32-bit; split the batch")
+    out = torch.empty(bsz, t, n_cols, dtype=x.dtype, device=x.device)
+    if residual is not None:
+        residual = residual.contiguous()
+
+    def grid(meta):
+        return (triton.cdiv(n_rows, meta["BLOCK_M"]) * triton.cdiv(n_cols, meta["BLOCK_N"]),)
+
+    _causal_conv_gemm_kernel[grid](
+        x,
+        weight,
+        bias if bias is not None else x,
+        residual if residual is not None else x,
+        out,
+        n_rows,
+        n_cols,
+        c,
+        t,
+        dilation,
+        triton.next_power_of_2(triton.cdiv(n_rows, 1024)),
+        K=kernel_size,
+        HAS_BIAS=bias is not None,
+        HAS_RES=residual is not None,
+        PRECISION="ieee" if x.dtype == torch.float32 else "tf32",
+    )
+    return out
+
+
+def transposed_overlap_add(z: torch.Tensor, stride: int, bias: torch.Tensor) -> torch.Tensor:
+    """``[B, T, K, C_out]`` GEMM output of a kernel-``K`` transposed conv -> ``[B, T * stride, C_out]``."""
+    bsz, t, k, c_out = z.shape
+    if k not in (stride, 2 * stride):
+        raise ValueError(f"transposed_overlap_add supports kernel stride or 2*stride, got kernel={k} stride={stride}")
+    if not HAS_TRITON or not z.is_cuda:
+        return transposed_overlap_add_reference(z, stride, bias)
+    z = z.contiguous()
+    n_out_rows = bsz * t * stride
+    out = torch.empty(bsz, t * stride, c_out, dtype=z.dtype, device=z.device)
+    block_r, block_c = _blocks(c_out)
+    _transposed_overlap_add_kernel[(triton.cdiv(n_out_rows, block_r), triton.cdiv(c_out, block_c))](
+        z, bias, out, n_out_rows, t, c_out, STRIDE=stride, K=k, BLOCK_R=block_r, BLOCK_C=block_c
+    )
+    return out

vllm_omni/model_executor/models/qwen3_tts/tokenizer_12hz/time_major_stack.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Time-major upsample + decoder conv stack for the Qwen3-TTS 12 Hz codec.
+
+Same math as ``Qwen3TTSTokenizerV2Decoder.upsample``/``.decoder`` on ``[B, T, C]``
+activations: dense causal convs are implicit GEMMs (tensor cores, no im2col
+buffer, residual adds in the epilogue), SnakeBeta one pass per activation,
+transposed convs a GEMM plus an overlap-add. See ``time_major_conv.py``.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.time_major_conv import (
+    causal_conv,
+    conv_gemm_weight,
+    snake_im2col,
+    transposed_gemm_weight,
+    transposed_overlap_add,
+)
+
+
+def _snake(module) -> tuple[torch.Tensor, torch.Tensor]:
+    if module._exp_alpha is None:
+        module.precompute_exp_cache()
+    return module._exp_alpha, module._inv_beta
+
+
+class _Conv:
+    def __init__(self, causal_conv) -> None:
+        conv = causal_conv.conv
+        if conv.stride[0] != 1 or conv.groups != 1:
+            raise ValueError("time-major decoder convs must be dense with stride 1")
+        self.weight = conv_gemm_weight(conv.weight)
+        self.bias = conv.bias.detach()
+        self.kernel_size = conv.kernel_size[0]
+        self.dilation = conv.dilation[0]
+
+    def __call__(self, x: torch.Tensor, residual: torch.Tensor | None = None) -> torch.Tensor:
+        return causal_conv(x, self.weight, self.bias, self.kernel_size, self.dilation, residual)
+
+
+class _Transposed:
+    def __init__(self, causal_trans_conv) -> None:
+        conv = causal_trans_conv.conv
+        self.weight = transposed_gemm_weight(conv.weight)
+        self.bias = conv.bias.detach()
+        self.kernel_size = conv.kernel_size[0]
+        self.stride = conv.stride[0]
+        self.out_channels = conv.out_channels
+
+    def __call__(self, x2d: torch.Tensor, bsz: int, t: int) -> torch.Tensor:
+        z = torch.mm(x2d, self.weight).view(bsz, t, self.kernel_size, self.out_channels)
+        return transposed_overlap_add(z, self.stride, self.bias)
+
+
+class TimeMajorConvStack:
+    """Callable replacing the NCT conv loops; holds GEMM-layout copies of the weights."""
+
+    def __init__(self, decoder: nn.Module) -> None:
+        self.upsample = [(_Transposed(trans), block) for trans, block in decoder.upsample]
+        self.conv_in = _Conv(decoder.decoder[0])
+        self.blocks = []
+        for block in decoder.decoder[1:-2]:
+            snake, trans, *units = block.block
+            residuals = [(_snake(u.act1), _Conv(u.conv1), _snake(u.act2), _Conv(u.conv2)) for u in units]
+            self.blocks.append((_snake(snake), _Transposed(trans), residuals))
+        self.snake_out = _snake(decoder.decoder[-2])
+        self.conv_out = _Conv(decoder.decoder[-1])
+
+    @staticmethod
+    def _convnext(x: torch.Tensor, block) -> torch.Tensor:
+        dw = block.dwconv.conv
+        k, t = dw.kernel_size[0], x.shape[1]
+        xp = torch.nn.functional.pad(x, (0, 0, k - 1, 0))
+        weight = dw.weight[:, 0, :].float()
+        h = dw.bias.float() + sum(xp[:, j : j + t].float() * weight[:, j] for j in range(k))
+        h = block.pwconv2(block.act(block.pwconv1(block.norm(h.to(x.dtype)))))
+        return x + block.gamma * h
+
+    def __call__(self, hidden: torch.Tensor) -> torch.Tensor:
+        """``hidden`` ``[B, T, latent]`` -> unclamped waveform ``[B, 1, T * upsample]``."""
+        bsz = hidden.shape[0]
+
+        def snake(x, params):
+            return snake_im2col(x, 1, snake=params).view(x.shape)
+
+        for trans, block in self.upsample:
+            t = hidden.shape[1]
+            hidden = self._convnext(trans(hidden.reshape(bsz * t, -1), bsz, t), block)
+        hidden = self.conv_in(hidden)
+        for block_snake, trans, residuals in self.blocks:
+            t = hidden.shape[1]
+            hidden = trans(snake_im2col(hidden, 1, snake=block_snake), bsz, t)
+            for act1, conv1, act2, conv2 in residuals:
+                hidden = conv2(snake(conv1(snake(hidden, act1)), act2), residual=hidden)
+        wav = self.conv_out(snake(hidden, self.snake_out))
+        return wav.view(bsz, 1, -1)

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -407,6 +407,12 @@ def talker2code2wav_async_chunk(
         left_context_size=left_context_size,
         finished=torch.tensor(finished, dtype=torch.bool),
     )
+    source_meta = multimodal_output.get("meta", {}) if isinstance(multimodal_output, Mapping) else {}
+    first_audio = source_meta.get("first_audio", False)
+    if isinstance(first_audio, torch.Tensor):
+        first_audio = bool(first_audio.numel() and first_audio.reshape(-1)[-1].item())
+    if first_audio:
+        meta.first_audio = torch.tensor(True, dtype=torch.bool)
     if ref_context_size > 0 and ref_context_request_id is not None:
         meta.ref_context_size = ref_context_size
         meta.ref_context_request_id = ref_context_request_id
@@ -494,11 +500,13 @@ def talker2code2wav_async_chunk_batch(
             request_payload[request_id] = value.contiguous()
 
     payloads: list[OmniPayloadStruct | None] = []
-    for request, finished in zip(requests, is_finished):
+    for request, finished, pooling_output in zip(requests, is_finished, pooling_outputs):
         payloads.append(
             talker2code2wav_async_chunk(
                 transfer_manager=transfer_manager,
-                multimodal_output={},
+                multimodal_output={"meta": pooling_output.get("meta", {})}
+                if isinstance(pooling_output, Mapping)
+                else {},
                 request=request,
                 is_finished=finished,
             )

vllm_omni/utils/device_copy.py

@@ -0,0 +1,27 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Host-to-device copies that do not block the host on queued GPU work.
+
+A copy from pageable host memory (``tensor.to("cuda")``, ``torch.tensor(...,
+device="cuda")``, ``torch.as_tensor(..., device="cuda")``) makes the CPU wait
+until every kernel already queued on the stream has finished. In a per-step
+path under async scheduling that removes the CPU/GPU overlap the batch queue
+exists for. Staging the data in
+pinned memory and copying with ``non_blocking`` keeps the host running ahead;
+the caching host allocator keeps the pinned source alive until the copy's
+stream event completes.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+
+def index_to_device(values: Sequence[int], device: torch.device | str, dtype: torch.dtype = torch.long) -> torch.Tensor:
+    """A host index list as a device tensor, without a host sync."""
+    if torch.device(device).type != "cuda":
+        return torch.tensor(values, dtype=dtype, device=device)
+    return torch.tensor(values, dtype=dtype, pin_memory=True).to(device, non_blocking=True)

vllm_omni/worker_v2/first_audio_sender.py

@@ -0,0 +1,167 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Deliver in-stage first audio as soon as its device copy completes.
+
+The regular output of a step is published only after the whole step's
+outputs are copied, materialized for every request, popped by the engine loop
+and serialized together. A stream's first audio chunk does not need any of
+that, so a dedicated thread waits for the chunk's own copy event and hands a
+one-request output straight to the engine core's output queue.
+"""
+
+from __future__ import annotations
+
+import queue
+import threading
+from dataclasses import dataclass
+from typing import Any, Protocol
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY
+
+logger = init_logger(__name__)
+
+
+class FirstAudioSink(Protocol):
+    def prepare(self, request_ids: list[str]) -> _PreparedDelivery:
+        """Freeze request routes on the submitting thread before accepting delivery."""
+        ...
+
+
+class FirstAudioSender:
+    def __init__(self, sink: FirstAudioSink) -> None:
+        if not callable(getattr(sink, "prepare", None)):
+            raise TypeError("FirstAudioSender requires a sink with prepare(request_ids)")
+        self._sink = sink
+        self._queue: queue.SimpleQueue[
+            tuple[torch.cuda.Event, torch.Tensor, list[str], torch.Tensor, torch.Tensor | None, _PreparedDelivery]
+            | None
+        ]
+        self._queue = queue.SimpleQueue()
+        self._thread = threading.Thread(target=self._run, daemon=True, name="omni-first-audio-sender")
+        self._thread.start()
+
+    def submit(
+        self,
+        request_ids: list[str],
+        pcm: torch.Tensor,
+        sample_rate: torch.Tensor,
+        valid: torch.Tensor | None = None,
+    ) -> list[str]:
+        """Queue a D2H copy of ``pcm`` [n, samples] on the current stream and deliver it once done.
+
+        ``valid`` [n] (on device) drops rows whose frame turned out not to be
+        audio (e.g. a codec EOS sample); it is read only after the copy.
+
+        Return the request IDs whose delivery was accepted. Rows without an
+        engine output route must retain regular codec delivery instead of
+        promising the orchestrator a first frame that can never arrive.
+        """
+        delivery = self._sink.prepare(request_ids)
+        accepted = list(delivery.routes)
+        if not accepted:
+            return []
+        host = torch.empty(pcm.shape, dtype=pcm.dtype, pin_memory=True)
+        host.copy_(pcm, non_blocking=True)
+        host_valid = None
+        if valid is not None:
+            host_valid = torch.empty(valid.shape, dtype=torch.bool, pin_memory=True)
+            host_valid.copy_(valid, non_blocking=True)
+        copied = torch.cuda.Event()
+        copied.record()
+        self._queue.put((copied, host, list(request_ids), sample_rate, host_valid, delivery))
+        return accepted
+
+    def close(self) -> None:
+        self._queue.put(None)
+        self._thread.join(timeout=5)
+
+    def _run(self) -> None:
+        while (item := self._queue.get()) is not None:
+            copied, host, request_ids, sample_rate, host_valid, delivery = item
+            pending = request_ids
+            try:
+                copied.synchronize()
+                rows = [row for row in range(len(request_ids)) if host_valid is None or bool(host_valid[row])]
+                pending = [request_ids[row] for row in rows]
+                if rows:
+                    delivery(pending, [host[row].clone() for row in rows], sample_rate)
+            except Exception:
+                # If synchronization failed, validity is not readable yet.
+                # Otherwise only valid, not-yet-delivered rows owe first audio.
+                logger.exception("First-audio delivery failed for %s", pending)
+                delivery.fail(pending)
+
+
+@dataclass
+class _PreparedDelivery:
+    sink: _EngineOutputSink
+    # Frozen client indices for outstanding deliveries. Successful enqueues
+    # retire their routes so a later failure cannot abort completed deliveries.
+    routes: dict[str, int]
+
+    def __call__(self, request_ids: list[str], pcm_rows: list[torch.Tensor], sample_rate: torch.Tensor) -> None:
+        self.sink._send(self.routes, request_ids, pcm_rows, sample_rate)
+
+    def fail(self, request_ids: list[str]) -> None:
+        self.sink._fail(self.routes, request_ids)
+
+
+class _EngineOutputSink:
+    def __init__(self, output_queue: Any, scheduler: Any) -> None:
+        self.output_queue = output_queue
+        self.scheduler = scheduler
+
+    def prepare(self, request_ids: list[str]) -> _PreparedDelivery:
+        routes = {}
+        for request_id in request_ids:
+            request = self.scheduler.requests.get(request_id)
+            if request is not None:
+                routes[request_id] = int(getattr(request, "client_index", 0) or 0)
+        return _PreparedDelivery(self, routes)
+
+    def _send(self, routes: dict[str, int], request_ids, pcm_rows, sample_rate) -> None:
+        from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs
+
+        by_client: dict[int, list[OmniEngineCoreOutput]] = {}
+        for request_id, pcm in zip(request_ids, pcm_rows, strict=True):
+            if request_id not in routes:
+                continue
+            by_client.setdefault(routes[request_id], []).append(
+                OmniEngineCoreOutput(
+                    request_id=request_id,
+                    new_token_ids=[],
+                    multimodal_output={
+                        "model_outputs": pcm,
+                        "sr": sample_rate,
+                        FIRST_AUDIO_KEY: torch.tensor(True),
+                    },
+                )
+            )
+        for client_index, outputs in by_client.items():
+            self.output_queue.put_nowait((client_index, OmniEngineCoreOutputs(outputs=outputs)))
+            for output in outputs:
+                routes.pop(output.request_id)
+
+    def _fail(self, routes: dict[str, int], request_ids: list[str]) -> None:
+        from vllm.v1.engine import FinishReason
+
+        from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs
+
+        for request_id in request_ids:
+            if request_id in routes:
+                output = OmniEngineCoreOutput(
+                    request_id=request_id,
+                    new_token_ids=[],
+                    finish_reason=FinishReason.ERROR,
+                    stop_reason="First-audio delivery failed",
+                )
+                self.output_queue.put_nowait((routes[request_id], OmniEngineCoreOutputs(outputs=[output])))
+                routes.pop(request_id)
+
+
+def engine_output_queue_sink(output_queue: Any, scheduler: Any) -> _EngineOutputSink:
+    return _EngineOutputSink(output_queue, scheduler)

vllm_omni/worker_v2/model_states/eager_mtp.py

@@ -0,0 +1,226 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Opt-in completion of a Talker frame immediately after sampling codebook zero."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+import torch
+from vllm.logger import init_logger
+from vllm.v1.worker.gpu.input_batch import InputBatch
+
+from vllm_omni.utils.device_copy import index_to_device
+
+if TYPE_CHECKING:
+    from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+
+logger = init_logger(__name__)
+
+
+def _has_ref_codes(buffer: dict[str, Any]) -> bool:
+    ref = buffer.get("codes", {}).get("ref") if isinstance(buffer.get("codes"), dict) else None
+    return isinstance(ref, torch.Tensor) and ref.numel() > 0
+
+
+class EagerMTPState:
+    """Used only by models explicitly declaring ``mtp_eager_frames``."""
+
+    def __init__(self, owner: OmniModelState) -> None:
+        self.owner = owner
+        self._first_audio_valid: torch.Tensor | None = None
+
+    def set_first_audio_sink(self, sink: Any) -> None:
+        from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender
+
+        self.owner._first_audio_sender = FirstAudioSender(sink)
+
+    def _apply_eager_frames(
+        self,
+        mtp_batches: list[tuple[int, int, tuple[torch.Tensor, torch.Tensor]]],
+        embeds: torch.Tensor,
+        input_batch: InputBatch,
+        prepacked_mtp_inputs: tuple[torch.Tensor, torch.Tensor] | None,
+    ) -> None:
+        """Decode input = codec embeddings of the previous (eager) frame + this step's text."""
+        req_indices = [int(input_batch.idx_mapping_np[i]) for i, _start, _mtp in mtp_batches]
+        for req_idx in req_indices:
+            req_id = self.owner.intermediate_buffer.buffers[req_idx].get("req_id")
+            if self.owner._eager_ready.get(req_idx) != req_id:
+                # Running the deferred MTP here would re-emit or drop a frame.
+                raise RuntimeError(f"Eager Talker-MTP frame missing for request {req_id!r}")
+        assert self.owner._eager_embeds is not None
+        device = embeds.device
+        if prepacked_mtp_inputs is None:
+            text_step = torch.cat([step.reshape(1, -1) for _i, _start, (_hidden, step) in mtp_batches], dim=0)
+        else:
+            text_step = prepacked_mtp_inputs[1].reshape(len(mtp_batches), -1)
+        rows = index_to_device(req_indices, device)
+        offsets = self.owner._mtp_batch_offsets(mtp_batches, input_batch, device)
+        frame_embeds = self.owner._eager_embeds.index_select(0, rows)
+        embeds.index_copy_(0, offsets, (frame_embeds + text_step.to(frame_embeds.dtype)).to(embeds.dtype))
+
+    def _apply_settled_frames(self, settled_rows: list[tuple[int, int, int, str]], embeds: torch.Tensor) -> None:
+        """Settled decode rows: previous eager frame embedding plus the constant text step."""
+        for _i, req_idx, _start, req_id in settled_rows:
+            if self.owner._eager_ready.get(req_idx) != req_id:
+                raise RuntimeError(f"Eager Talker-MTP frame missing for request {req_id!r}")
+        assert self.owner._eager_embeds is not None
+        device = embeds.device
+        rows = index_to_device([req_idx for _i, req_idx, _start, _req_id in settled_rows], device)
+        offsets = index_to_device([start for _i, _req_idx, start, _req_id in settled_rows], device)
+        text_step = self.owner.model.eager_settled_text_step().to(device=device, dtype=self.owner._eager_embeds.dtype)
+        frame_embeds = self.owner._eager_embeds.index_select(0, rows)
+        embeds.index_copy_(0, offsets, (frame_embeds + text_step).to(embeds.dtype))
+
+    def run_eager_mtp(
+        self,
+        input_batch: InputBatch,
+        text_hidden: torch.Tensor,
+        sampled_token_ids: torch.Tensor,
+        multimodal_outputs: dict[str, Any],
+        mtp_batch_descriptor_dispatcher: Callable[[int], Any] | None = None,
+    ) -> None:
+        """Complete each sampled row's frame in this step and publish it with this step's output.
+
+        Consumes the rows recorded by ``run_preprocess``. Writes the frame codes
+        and validity into the last token row of each request span of the
+        retained multimodal output, and keeps the frame's codec embedding sum
+        for the next step's input. Issued on the main stream before the async
+        output copy, so no extra synchronization is needed.
+        """
+        if not self.owner._eager_mtp:
+            return
+        recorded, self.owner._eager_rows = self.owner._eager_rows, None
+        if recorded is None or recorded[0] is not input_batch or not recorded[1]:
+            return
+        entries, input_ids = recorded[1], recorded[2]
+        codes_out = multimodal_outputs.get("codes", {}).get("audio")
+        valid_out = multimodal_outputs.get("meta", {}).get("codec_frame_valid")
+        if not isinstance(codes_out, torch.Tensor) or not isinstance(valid_out, torch.Tensor):
+            raise RuntimeError("Eager Talker-MTP requires retained codes.audio and codec_frame_valid outputs")
+        assert self.owner._mtp_input_ids is not None and self.owner._mtp_input_embeds is not None
+        assert self.owner._mtp_hidden is not None and self.owner._mtp_text_step is not None
+        assert self.owner._eager_embeds is not None
+
+        bsz = len(entries)
+        device = text_hidden.device
+        batch_rows = index_to_device([i for i, _req_idx, _req_id, _prefill in entries], device)
+        req_indices = [req_idx for _i, req_idx, _req_id, _prefill in entries]
+        last_tokens = input_batch.query_start_loc.index_select(0, batch_rows + 1).long() - 1
+        layer0 = sampled_token_ids.reshape(input_batch.num_reqs, -1)[:, 0].index_select(0, batch_rows)
+
+        batch_ids = self.owner._mtp_input_ids[:bsz]
+        batch_emb = self.owner._mtp_input_embeds[:bsz]
+        batch_hidden = self.owner._mtp_hidden[:bsz]
+        batch_step = self.owner._mtp_text_step[:bsz]
+        batch_ids.copy_(layer0.to(batch_ids.dtype))
+        batch_emb.copy_(self.owner.model.embed_input_ids(batch_ids.reshape(-1, 1)).reshape(bsz, -1))
+        torch.index_select(text_hidden, 0, last_tokens, out=batch_hidden)
+        # The text step is added by the next step's preprocess.
+        batch_step.zero_()
+        frame_embeds, codes = self.owner._mtp_forward(
+            req_indices,
+            batch_ids,
+            batch_emb,
+            batch_hidden,
+            batch_step,
+            mtp_batch_descriptor_dispatcher,
+        )
+        assert codes is not None
+        rows = index_to_device(req_indices, device)
+        self.owner._eager_embeds.index_copy_(
+            0, rows, frame_embeds[:bsz].reshape(bsz, -1).to(self.owner._eager_embeds.dtype)
+        )
+        if self.owner.vllm_config.cache_config.enable_prefix_caching:
+            # Prefix-hit reconstruction reads the latest frame from the buffer.
+            self.owner.intermediate_buffer.update_gpu_tensor_rows(
+                req_indices, self.owner.model.mtp_output_key, codes[:bsz]
+            )
+        if codes_out.ndim != 2:
+            raise RuntimeError(f"Eager Talker-MTP expects token-major codes.audio, got {tuple(codes_out.shape)}")
+        codes_out.index_copy_(0, last_tokens, codes[:bsz].to(codes_out.dtype))
+        # A decode row whose input CB0 was EOS belongs to a request that already
+        # ended (async scheduling runs it once more); its new sample must not
+        # become a frame. Prefill rows have no codec input.
+        prefill_rows = index_to_device([int(prefill) for _i, _req_idx, _req_id, prefill in entries], device).bool()
+        input_valid = self.owner.model.mtp_frame_valid(input_ids.index_select(0, last_tokens)) | prefill_rows
+        valid = self.owner.model.mtp_frame_valid(layer0) & input_valid
+        valid_out.index_copy_(0, last_tokens, valid.to(valid_out.dtype))
+        for _i, req_idx, req_id, _prefill in entries:
+            self.owner._eager_ready[req_idx] = req_id
+        self._publish_first_audio(input_batch, entries, codes[:bsz], valid)
+        first_audio = multimodal_outputs.get("meta", {}).get("first_audio")
+        if isinstance(first_audio, torch.Tensor):
+            scheduled = index_to_device(
+                [int(req_id in self.owner._first_audio_requests) for _, _, req_id, _ in entries], device
+            )
+            if self._first_audio_valid is None:
+                scheduled.zero_()
+            else:
+                scheduled *= self._first_audio_valid.index_select(0, rows)
+            first_audio.index_copy_(0, last_tokens, scheduled.to(first_audio.dtype))
+
+    def _publish_first_audio(
+        self,
+        input_batch: InputBatch,
+        entries: list[tuple[int, int, str, bool]],
+        codes: torch.Tensor,
+        valid: torch.Tensor | None = None,
+    ) -> None:
+        """Decode first frames here and hand the PCM to the client output.
+
+        Only streams without reference codes qualify: their first chunk is
+        decoded with no context, so the waveform equals what Code2Wav will
+        produce for it. ``valid`` (per entry, on device) marks rows whose frame
+        is real; direct delivery drops the others once their copy is done, so a
+        stream whose first sample is codec EOS gets no first audio from here.
+        """
+        decoder = getattr(self.owner.model, "first_frame_decoder", None)
+        if decoder is None:
+            logger.info_once("Talker first-frame audio: no decoder on %s", type(self.owner.model).__name__)
+            return
+        first = [
+            (row, i)
+            for row, (i, req_idx, _req_id, prefill) in enumerate(entries)
+            if prefill
+            and _req_id not in self.owner._first_audio_requests
+            and not _has_ref_codes(self.owner.intermediate_buffer.buffers[req_idx])
+        ]
+        if not first:
+            return
+        # Decode on a side stream so the Talker's next step does not queue
+        # behind it; the async output copy waits on its event instead.
+        stream = self.owner._first_audio_stream
+        if stream is None:
+            _least, greatest = torch.cuda.Stream.priority_range()
+            stream = self.owner._first_audio_stream = torch.cuda.Stream(device=codes.device, priority=greatest)
+        frame_rows = index_to_device([row for row, _i in first], codes.device)
+        frame_codes = codes.index_select(0, frame_rows)
+        frame_valid = valid.index_select(0, frame_rows) if valid is not None else None
+        if frame_valid is not None:
+            frame_codes = torch.where(frame_valid[:, None], frame_codes, 0)
+        if self._first_audio_valid is None:
+            assert self.owner._eager_embeds is not None
+            self._first_audio_valid = torch.zeros(
+                self.owner._eager_embeds.shape[0], dtype=torch.bool, device=codes.device
+            )
+        request_rows = index_to_device([entries[row][1] for row, _ in first], codes.device)
+        if frame_valid is not None:
+            self._first_audio_valid.index_copy_(0, request_rows, frame_valid)
+        else:
+            self._first_audio_valid.index_fill_(0, request_rows, True)
+        stream.wait_stream(torch.cuda.current_stream(codes.device))
+        sr = torch.tensor(decoder.sample_rate, dtype=torch.int32)
+        with torch.cuda.stream(stream):
+            frame_codes.record_stream(stream)
+            pcm = decoder.decode(frame_codes)
+            sender = self.owner._first_audio_sender
+            if sender is None:
+                raise RuntimeError("First-frame decoder requires an engine output sink")
+            if frame_valid is not None:
+                frame_valid.record_stream(stream)
+            request_ids = [entries[row][2] for row, _i in first]
+            accepted = sender.submit(request_ids, pcm, sr, valid=frame_valid)
+            self.owner._first_audio_requests.update(accepted)

vllm_omni/worker_v2/model_states/omni_model_state.py

@@ -32,6 +32,7 @@ from vllm.v1.worker.gpu.states import RequestState
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
 from vllm_omni.worker_v2.model_states.intermediate_buffer import (
     OmniIntermediateBuffer,
 )
@@ -85,6 +86,14 @@ class OmniModelState(DefaultModelState):
     Model-owned preprocess and postprocess hooks supply stage-specific behavior.
     """
 
+    # Eager-frame defaults for states built without __init__ (tests, adapters).
+    _eager_mtp = False
+    _eager_fastpath = False
+    _eager_rows: tuple[InputBatch, list[tuple[int, int, str, bool]], torch.Tensor] | None = None
+    _first_audio_stream: torch.cuda.Stream | None = None
+    # Set by the stage engine process when it can take one-request outputs directly.
+    _first_audio_sender: Any = None
+
     def __init__(
         self,
         vllm_config: VllmConfig,
@@ -167,6 +176,40 @@ class OmniModelState(DefaultModelState):
                     device=device,
                 )
             self._mtp_runner = self._init_mtp_runner(model)
+        # Eager frames: finish MTP in the step that sampled CB0 (see
+        # run_eager_mtp). The next step only adds its text-step embedding.
+        self._eager_mtp = self._mtp_input_ids is not None and bool(getattr(model, "mtp_eager_frames", False))
+        self._eager_rows = None
+        self._eager_state = EagerMTPState(self)
+        self._eager_embeds: torch.Tensor | None = None
+        self._eager_ready: dict[int, str] = {}
+        self._first_audio_requests: set[str] = set()
+        # Slots whose decode preprocess no longer changes state (see
+        # model.eager_decode_settled); they skip the per-row path.
+        self._eager_settled: dict[int, str] = {}
+        self._eager_fastpath = self._eager_mtp
+        if self._eager_mtp:
+            self._eager_embeds = torch.zeros((max_num_reqs, self._embed_dim), dtype=self.dtype, device=device)
+            # The model then leaves frame codes/validity and the decode MTP
+            # inputs to run_eager_mtp instead of building them per row.
+            model.eager_frames_active = True
+            logger.info("Eager Talker-MTP frames enabled")
+
+    def set_first_audio_sink(self, sink: Any) -> None:
+        self._eager_state.set_first_audio_sink(sink)
+
+    def run_eager_mtp(
+        self,
+        input_batch: InputBatch,
+        text_hidden: torch.Tensor,
+        sampled_token_ids: torch.Tensor,
+        multimodal_outputs: dict[str, Any],
+        mtp_batch_descriptor_dispatcher: Callable[[int], Any] | None = None,
+    ) -> None:
+        if self._eager_mtp:
+            self._eager_state.run_eager_mtp(
+                input_batch, text_hidden, sampled_token_ids, multimodal_outputs, mtp_batch_descriptor_dispatcher
+            )
 
     @staticmethod
     def _resolve_decode_preprocess(model: nn.Module) -> Callable | None:
@@ -350,6 +393,9 @@ class OmniModelState(DefaultModelState):
         req_id = self.intermediate_buffer.buffers[req_index].get("req_id")
         if req_id is not None:
             getattr(self, "_mtp_generators", {}).pop(req_id, None)
+            getattr(self, "_first_audio_requests", set()).discard(req_id)
+        getattr(self, "_eager_ready", {}).pop(req_index, None)
+        getattr(self, "_eager_settled", {}).pop(req_index, None)
         self.intermediate_buffer.remove_request(req_index)
 
     # ------------------------------------------------------------------
@@ -555,15 +601,19 @@ class OmniModelState(DefaultModelState):
         self._stage_batched_preprocess_inputs(req_indices, embeds.device)
 
         preprocess_entries: list[tuple[int, int, int, int, dict[str, Any], bool]] = []
+        settled = self._eager_settled if self._eager_fastpath else None
+        settled_rows: list[tuple[int, int, int, str]] = []
         for i, req_idx in enumerate(req_indices):
             buf = self.intermediate_buffer.buffers[req_idx]
             if not buf or "req_id" not in buf:
                 continue
-            if str(buf["req_id"]).startswith("_warmup_"):
-                continue
-
             start = int(input_batch.query_start_loc_np[i])
             n_tok = int(input_batch.num_scheduled_tokens[i])
+            if settled and n_tok == 1 and settled.get(req_idx) == buf["req_id"]:
+                settled_rows.append((i, req_idx, start, buf["req_id"]))
+                continue
+            if str(buf["req_id"]).startswith("_warmup_"):
+                continue
 
             info = {key: value for key, value in buf.items() if isinstance(key, str)}
             prompt_len = None
@@ -590,6 +640,20 @@ class OmniModelState(DefaultModelState):
             )
             preprocess_entries.append((i, req_idx, start, n_tok, info, is_prefill))
 
+        if self._eager_mtp:
+            # Rows whose sample is kept this step: decode, or the final prefill chunk.
+            self._eager_rows = (
+                input_batch,
+                [
+                    (i, req_idx, str(info["req_id"]), is_prefill)
+                    for i, req_idx, _start, n_tok, info, is_prefill in preprocess_entries
+                    if not is_prefill
+                    or int(info.get("_omni_num_computed_tokens", 0)) + n_tok >= int(info.get("_omni_prompt_len", 0))
+                ]
+                + [(i, req_idx, str(req_id), False) for i, req_idx, _start, req_id in settled_rows],
+                input_ids,
+            )
+
         preprocess_batch_mrv2 = getattr(self.model, "preprocess_batch_mrv2", None)
         if callable(preprocess_batch_mrv2):
             prefill_infos = [
@@ -620,6 +684,7 @@ class OmniModelState(DefaultModelState):
                 batch_embeds = embeds[:batch_size]
                 batch_offsets = None
             else:
+                # A pageable host tensor would block on the stream here.
                 batch_offsets = torch.as_tensor(starts, device=input_ids.device, dtype=torch.long)
                 batch_ids = input_ids.index_select(0, batch_offsets)
                 batch_embeds = embeds.index_select(0, batch_offsets)
@@ -659,13 +724,13 @@ class OmniModelState(DefaultModelState):
             )
             for row, (entry, updates) in enumerate(zip(decode_entries, updates_by_req, strict=True)):
                 i, req_idx, start, _n_tok, _info, _is_prefill = entry
-                mtp_batches.append(
-                    (
-                        i,
-                        start,
-                        (batch_hidden[row : row + 1], batch_text_step[row : row + 1]),
-                    )
+                # Eager frames read the packed text steps; skip per-row views.
+                row_inputs = (
+                    (batch_hidden, batch_text_step)
+                    if self._eager_mtp
+                    else (batch_hidden[row : row + 1], batch_text_step[row : row + 1])
                 )
+                mtp_batches.append((i, start, row_inputs))
                 batched_decode_indices.add(i)
             state_updates = [(entry[1], updates) for entry, updates in zip(decode_entries, updates_by_req, strict=True)]
             for req_idx, updates in state_updates:
@@ -694,7 +759,18 @@ class OmniModelState(DefaultModelState):
 
             self.intermediate_buffer.update(req_idx, updates, gpu_keys)
 
-        if mtp_batches and hasattr(self.model, "mtp"):
+        if self._eager_mtp:
+            settled_check = getattr(self.model, "eager_decode_settled", None)
+            if self._eager_fastpath and callable(settled_check):
+                for _i, req_idx, _start, n_tok, info, is_prefill in preprocess_entries:
+                    buf = self.intermediate_buffer.buffers[req_idx]
+                    if not is_prefill and n_tok == 1 and settled_check(buf):
+                        self._eager_settled[req_idx] = buf["req_id"]
+            if settled_rows:
+                self._eager_state._apply_settled_frames(settled_rows, embeds)
+        if mtp_batches and self._eager_mtp:
+            self._eager_state._apply_eager_frames(mtp_batches, embeds, input_batch, prepacked_mtp_inputs)
+        elif mtp_batches and hasattr(self.model, "mtp"):
             if prepacked_mtp_inputs is None:
                 self._run_batched_mtp(
                     mtp_batches,
@@ -727,9 +803,7 @@ class OmniModelState(DefaultModelState):
         bsz = len(mtp_batches)
         if offsets is None:
             offsets = torch.as_tensor(
-                [start for _i, start, _mtp in mtp_batches],
-                device=input_ids.device,
-                dtype=torch.long,
+                [start for _i, start, _mtp in mtp_batches], device=input_ids.device, dtype=torch.long
             )
         if prepacked_mtp_inputs is None:
             hidden_rows = [past_hidden.reshape(1, -1) for _i, _start, (past_hidden, _step) in mtp_batches]
@@ -802,11 +876,7 @@ class OmniModelState(DefaultModelState):
             offsets.copy_(query_start_loc[:bsz], non_blocking=True)
             return offsets
 
-        return torch.as_tensor(
-            [start for _i, start, _mtp in mtp_batches],
-            device=device,
-            dtype=torch.long,
-        )
+        return torch.as_tensor([start for _i, start, _mtp in mtp_batches], device=device, dtype=torch.long)
 
     def _run_batched_mtp(
         self,
@@ -823,8 +893,6 @@ class OmniModelState(DefaultModelState):
         Uses pre-allocated static buffers to avoid per-step torch.cat
         memory allocations.
         """
-        from vllm.forward_context import set_forward_context
-
         bsz = len(mtp_batches)
         batch_offsets = self._mtp_batch_offsets(
             mtp_batches,
@@ -840,6 +908,72 @@ class OmniModelState(DefaultModelState):
         )
 
         req_indices = [int(input_batch.idx_mapping_np[i]) for i, _start, _mtp in mtp_batches]
+        new_emb, codes = self._mtp_forward(
+            req_indices,
+            batch_ids,
+            batch_emb,
+            batch_hidden,
+            batch_step,
+            mtp_batch_descriptor_dispatcher,
+        )
+
+        embeds.index_copy_(0, batch_offsets, new_emb[:bsz].reshape(bsz, -1))
+        audio_key = getattr(self.model, "mtp_output_key", None)
+        validity_key = getattr(self.model, "mtp_validity_key", None)
+        valid_rows = None
+        if codes is not None and validity_key is not None:
+            valid_rows = torch.ones((bsz,), dtype=torch.bool, device=codes.device)
+        if codes is not None and audio_key in gpu_keys:
+            self.intermediate_buffer.update_gpu_tensor_rows(
+                req_indices,
+                audio_key,
+                codes[:bsz],
+            )
+            if valid_rows is not None:
+                self.intermediate_buffer.update_gpu_tensor_rows(
+                    req_indices,
+                    validity_key,
+                    valid_rows,
+                    keepdim=False,
+                )
+            return
+        for j, (i, _start, _) in enumerate(mtp_batches):
+            if codes is None:
+                continue
+            req_idx = int(input_batch.idx_mapping_np[i])
+            if isinstance(audio_key, tuple) and len(audio_key) == 2:
+                updates = {audio_key[0]: {audio_key[1]: codes[j : j + 1]}}
+            elif isinstance(audio_key, str):
+                updates = {audio_key: codes[j : j + 1]}
+            else:
+                raise TypeError(
+                    f"mtp_output_key must be a string or 2-tuple, got {type(audio_key).__name__}: {audio_key!r}"
+                )
+            if valid_rows is not None:
+                if isinstance(validity_key, tuple) and len(validity_key) == 2:
+                    updates.setdefault(validity_key[0], {})[validity_key[1]] = valid_rows[j]
+                elif isinstance(validity_key, str):
+                    updates[validity_key] = valid_rows[j]
+                else:
+                    raise TypeError(
+                        "mtp_validity_key must be a string or 2-tuple, "
+                        f"got {type(validity_key).__name__}: {validity_key!r}"
+                    )
+            self.intermediate_buffer.update(req_idx, updates, gpu_keys)
+
+    def _mtp_forward(
+        self,
+        req_indices: list[int],
+        batch_ids: torch.Tensor,
+        batch_emb: torch.Tensor,
+        batch_hidden: torch.Tensor,
+        batch_step: torch.Tensor,
+        mtp_batch_descriptor_dispatcher: Callable[[int], Any] | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        """Run the residual-codebook predictor on packed rows; returns graph-owned outputs."""
+        from vllm.forward_context import set_forward_context
+
+        bsz = len(req_indices)
         buffers = [self.intermediate_buffer.buffers[req_idx] for req_idx in req_indices]
         req_ids = [str(buffer.get("req_id")) for buffer in buffers]
         generators = [
@@ -894,50 +1028,7 @@ class OmniModelState(DefaultModelState):
                 generators=generators,
                 sample_uniforms=sample_uniforms,
             )
-
-        embeds.index_copy_(0, batch_offsets, new_emb[:bsz].reshape(bsz, -1))
-        audio_key = getattr(self.model, "mtp_output_key", None)
-        validity_key = getattr(self.model, "mtp_validity_key", None)
-        valid_rows = None
-        if codes is not None and validity_key is not None:
-            valid_rows = torch.ones((bsz,), dtype=torch.bool, device=codes.device)
-        if codes is not None and audio_key in gpu_keys:
-            self.intermediate_buffer.update_gpu_tensor_rows(
-                req_indices,
-                audio_key,
-                codes[:bsz],
-            )
-            if valid_rows is not None:
-                self.intermediate_buffer.update_gpu_tensor_rows(
-                    req_indices,
-                    validity_key,
-                    valid_rows,
-                    keepdim=False,
-                )
-            return
-        for j, (i, _start, _) in enumerate(mtp_batches):
-            if codes is None:
-                continue
-            req_idx = int(input_batch.idx_mapping_np[i])
-            if isinstance(audio_key, tuple) and len(audio_key) == 2:
-                updates = {audio_key[0]: {audio_key[1]: codes[j : j + 1]}}
-            elif isinstance(audio_key, str):
-                updates = {audio_key: codes[j : j + 1]}
-            else:
-                raise TypeError(
-                    f"mtp_output_key must be a string or 2-tuple, got {type(audio_key).__name__}: {audio_key!r}"
-                )
-            if valid_rows is not None:
-                if isinstance(validity_key, tuple) and len(validity_key) == 2:
-                    updates.setdefault(validity_key[0], {})[validity_key[1]] = valid_rows[j]
-                elif isinstance(validity_key, str):
-                    updates[validity_key] = valid_rows[j]
-                else:
-                    raise TypeError(
-                        "mtp_validity_key must be a string or 2-tuple, "
-                        f"got {type(validity_key).__name__}: {validity_key!r}"
-                    )
-            self.intermediate_buffer.update(req_idx, updates, gpu_keys)
+        return new_emb, codes
 
     def _get_mtp_base_sampling_kwargs(self) -> dict[str, Any]:
         sampling_params = getattr(self.model, "mtp_sampling_params", None)
@@ -1107,6 +1198,9 @@ class OmniModelState(DefaultModelState):
         """
         if not self.has_postprocess:
             return
+        if self._eager_fastpath and not getattr(self.model, "eager_frames_need_postprocess", True):
+            # The last hidden it stores only fed the deferred decode MTP.
+            return
         gpu_keys: set[str] = getattr(self.model, "gpu_resident_buffer_keys", set())
         batch_postprocess = getattr(self.model, "postprocess_batch_mrv2", None)
         if callable(batch_postprocess) and input_batch.num_reqs:

vllm_omni/worker_v2/omni_ar_model_runner.py

@@ -192,6 +192,15 @@ class OmniARModelRunner(OmniGPUModelRunner):
                 input_batch,
                 grammar_output,
             )
+        run_eager_mtp = getattr(self.model_state, "run_eager_mtp", None)
+        if multimodal_outputs and run_eager_mtp is not None:
+            run_eager_mtp(
+                input_batch,
+                text_hidden,
+                sampler_output.sampled_token_ids,
+                multimodal_outputs,
+                self._dispatch_mtp_batch_descriptor,
+            )
         if self.pp_handler is not None:
             self.pp_handler.broadcast(
                 sampler_output.sampled_token_ids,

vllm_omni/worker_v2/omni_model_runner.py

@@ -186,6 +186,9 @@ class OmniGPUModelRunner(GPUModelRunner):
         super().add_requests(scheduler_output)
 
     def shutdown(self) -> None:
+        sender = getattr(getattr(self, "model_state", None), "_first_audio_sender", None)
+        if sender is not None:
+            sender.close()
         # Every owner must release its resources even if an earlier drain fails.
         try:
             worker = getattr(self, "_native_output_materializer", None)
@@ -352,6 +355,9 @@ class OmniGPUModelRunner(GPUModelRunner):
         capture_mtp = getattr(getattr(self, "model_state", None), "capture_mtp_graphs", None)
         if callable(capture_mtp):
             capture_mtp(self._dispatch_mtp_batch_descriptor)
+        capture_first_frame = getattr(self.model, "capture_first_frame_graphs", None)
+        if callable(capture_first_frame):
+            capture_first_frame()
         return result
 
     def _dispatch_mtp_batch_descriptor(self, num_mtp_reqs: int) -> Any: