[Model] Add single-stage Qwen3-TTS with in-Talker streaming codec decode (#8259)

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

docs/configuration/stage_configs.md

@@ -319,7 +319,9 @@ ROCm and MUSA. Select one of these deployment profiles:
 | `qwen3_tts_mrv2.yaml` | V2 | B1 | Explicit MRV2 profile (same runner selection as the default) |
 | `qwen3_tts_high_concurrency_mrv2.yaml` | V2 | B1, B2 | Opt-in throughput tuning |
 | `qwen3_tts_high_concurrency_mrv2_b4.yaml` | V2 | B1, B2, B3, B4 | Experimental throughput / buffered playback |
+| `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` | V2 | B1–B8 | Experimental two-stage deployment on one GPU, with MPS |
 | `qwen3_tts_high_concurrency.yaml` | V1 | Existing defaults | V1 high-concurrency control |
+| `qwen3_tts_fused_single_gpu.yaml` | V2 | N/A (in-Talker decoder) | Opt-in single-stage CUDA pipeline on one GPU |
 
 ```bash
 # MRV2 is the default; pass a copy with `model_runner: v1` to force V1.
@@ -350,12 +352,73 @@ 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.
 
+### Selecting one or two Qwen3-TTS stages
+
+Choose the deployment profile explicitly with `--deploy-config`. The default
+`qwen3_tts` pipeline retains two stages: the Talker publishes codec chunks to
+Code2Wav, which produces audio. Its stage devices can be placed on the same
+GPU or on separate GPUs. `qwen3_tts_fused` is a separate, opt-in pipeline for
+one CUDA GPU: the Talker owns the stateful codec decoder and emits PCM directly,
+without a Code2Wav engine or stage connector.
+
+```bash
+# Two stages on one GPU; the two-stage pipeline remains independently selectable.
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml
+
+# One stage on one GPU; no inter-stage transport or Code2Wav scheduling.
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml
+```
+
+The single-stage profile requires CUDA, V2, asynchronous chunks, TP=PP=1,
+an in-process worker and disabled prefix caching. Select the full profile;
+adding `talker_stream_decode: true` to a two-stage Talker is rejected because
+its latent output belongs to the Code2Wav input contract. Single-stage options
+live in `stages[].additional_config`; two-stage transport and first-frame
+options live in `connectors.*.extra`.
+
+The single-stage profile defaults `talker_stream_first_audio` to `false` for
+throughput. To opt into a separate first-frame delivery queue, save this
+overlay beside the bundled deployment files and pass its path to
+`--deploy-config`:
+
+```yaml
+base_config: qwen3_tts_fused_single_gpu.yaml
+stages:
+  - stage_id: 0
+    additional_config:
+      talker_stream_decode: true
+      talker_stream_first_audio: true
+      ref_code_context_frames: 72
+      code_predictor_kv_cache: true
+      code_predictor_fused_sampling: true
+      code_predictor_fused: true
+```
+
+This queue uses PCM from the single-stage decoder. The two-stage
+`talker_first_audio` option below uses its own first-frame decoder and does
+not require single-stage execution. The two options cannot be enabled
+together. Single-stage prefix-cache reset with
+`reset_running_requests=True` returns `False` while requests are active;
+wait for completion or abort them before retrying.
+
+Reference-cloning context is selected independently by each profile. The
+single-stage profile explicitly uses 72 reference frames, matching the
+`qwen3_tts.yaml` context boundary. A shorter
+`additional_config.ref_code_context_frames` (for example 25) can change
+timbre continuity even when WER is similar;
+the throughput measurements for CustomVoice do not qualify Base voice
+cloning quality. Reference layout, grouping and priming are model hooks;
+the shared worker does not select Qwen's context length or codebook layout.
+
 ### 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:
+two-stage `qwen3_tts_high_concurrency_mrv2_single_gpu.yaml` 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.

tests/config/test_config_merge.py

@@ -35,3 +35,161 @@ def test_recursive_deploy_inheritance(tmp_path):
     merged = resolve_deploy_yaml(leaf)
     assert merged["speech_cache"] == {"nested": {"a": 1, "b": 0}}
     assert merged["async_chunk"] is False
+
+
+def test_single_stage_streaming_capability_uses_central_async_resolver():
+    from vllm_omni.config.stage_config import DeployConfig, _resolve_pipeline_async_chunk_enabled
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE
+
+    assert _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig())
+    with pytest.raises(ValueError, match="requires async_chunk=True"):
+        _resolve_pipeline_async_chunk_enabled(QWEN3_TTS_FUSED_PIPELINE, DeployConfig(async_chunk=False))
+
+
+@pytest.mark.parametrize(
+    "pipeline_key,profile,stage_count",
+    [
+        ("qwen3_tts", "qwen3_tts_high_concurrency_mrv2_single_gpu.yaml", 2),
+        ("qwen3_tts_fused", "qwen3_tts_fused_single_gpu.yaml", 1),
+    ],
+)
+def test_qwen3_tts_deployment_selection_preserves_output_and_transport_contract(pipeline_key, profile, stage_count):
+    from tests.helpers.stage_config import get_deploy_config_path
+    from vllm_omni.config.pipeline_registry import resolve_pipeline_config
+    from vllm_omni.config.stage_config import load_deploy_config, merge_pipeline_deploy
+
+    pipeline = resolve_pipeline_config(pipeline_key)
+    assert pipeline is not None
+    deploy = load_deploy_config(get_deploy_config_path(profile))
+    stages = merge_pipeline_deploy(pipeline, deploy)
+    assert len(stages) == stage_count
+    talker = stages[0]
+    if stage_count == 1:
+        assert talker.final_output and talker.final_output_type == "audio"
+        assert talker.yaml_engine_args["engine_output_type"] == "audio"
+        assert "custom_process_next_stage_input_func" not in talker.yaml_engine_args
+        assert not talker.yaml_extras.get("output_connectors")
+        assert talker.yaml_engine_args["supports_running_prefix_cache_reset"] is False
+        assert talker.yaml_engine_args["additional_config"]["ref_code_context_frames"] == 72
+    else:
+        assert not talker.final_output
+        assert talker.yaml_engine_args["engine_output_type"] == "latent"
+        assert talker.yaml_engine_args["custom_process_next_stage_input_func"].endswith("talker2code2wav_async_chunk")
+        assert talker.yaml_extras["output_connectors"]
+        assert stages[1].input_sources == [0] and stages[1].final_output_type == "audio"
+        assert talker.yaml_engine_args.get("supports_running_prefix_cache_reset", True)
+
+
+def test_single_stage_running_reset_capability_cannot_be_overridden():
+    from vllm_omni.config.omni_config import VllmOmniConfig
+    from vllm_omni.config.stage_config import DeployConfig, StageDeployConfig, merge_pipeline_deploy
+    from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE
+
+    deploy = DeployConfig(
+        stages=[StageDeployConfig(stage_id=0, engine_extras={"supports_running_prefix_cache_reset": True})]
+    )
+    legacy = merge_pipeline_deploy(QWEN3_TTS_FUSED_PIPELINE, deploy)[0]
+    legacy.runtime_overrides["supports_running_prefix_cache_reset"] = True
+    assert legacy.to_omegaconf().engine_args.supports_running_prefix_cache_reset is False
+    with pytest.raises(ValueError, match="no structured config owner: supports_running_prefix_cache_reset"):
+        VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE, user_deploy_config=deploy)
+    structured = VllmOmniConfig.from_pipeline_config(QWEN3_TTS_FUSED_PIPELINE)
+    assert structured.stage_by_id(0).model_config.supports_running_prefix_cache_reset is False
+
+
+@pytest.fixture
+def single_stage_tts_config(mocker):
+    from vllm.config import VllmConfig
+
+    config = mocker.Mock(spec=VllmConfig)
+    config.additional_config = {"talker_stream_decode": True}
+    config.model_config = mocker.Mock(
+        stage_connector_config={},
+        engine_output_type="audio",
+        final_output=True,
+        supports_running_prefix_cache_reset=False,
+        use_v2_model_runner=True,
+        async_chunk=True,
+    )
+    config.parallel_config = mocker.Mock(
+        tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend="uni"
+    )
+    config.cache_config = mocker.Mock(enable_prefix_caching=False)
+    config.device_config = mocker.Mock(device="cuda")
+    mocker.patch(
+        "vllm_omni.model_executor.models.qwen3_tts.stream_decode.current_omni_platform.is_cuda", return_value=True
+    )
+    return config
+
+
+def test_single_stage_tts_rejects_unsupported_decoder_configuration(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.cache_config.enable_prefix_caching = True
+    with pytest.raises(ValueError, match="enable_prefix_caching=False"):
+        talker_stream_decode_enabled(config)
+    config.cache_config.enable_prefix_caching = False
+    assert talker_stream_decode_enabled(config)
+
+
+@pytest.mark.parametrize("option_location", ["additional_config", "connector"])
+@pytest.mark.parametrize("output_type", ["latent", "audio"])
+def test_two_stage_talker_rejects_single_stage_pcm_option(single_stage_tts_config, option_location, output_type):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.model_config.engine_output_type = output_type
+    config.model_config.final_output = False
+    if option_location == "connector":
+        config.model_config.stage_connector_config = {"extra": config.additional_config}
+        config.additional_config = {}
+    with pytest.raises(ValueError, match="final audio-output Talker"):
+        talker_stream_decode_enabled(config)
+
+
+def test_two_stage_first_audio_remains_independent(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.first_audio import talker_first_audio_enabled
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.model_config.engine_output_type = "latent"
+    config.additional_config = {}
+    config.model_config.stage_connector_config = {"extra": {"talker_first_audio": True}}
+    assert not talker_stream_decode_enabled(config)
+    assert talker_first_audio_enabled(config)
+    config.model_config.stage_connector_config["extra"]["talker_stream_first_audio"] = True
+    with pytest.raises(ValueError, match="requires the single-stage"):
+        talker_stream_decode_enabled(config)
+
+
+def test_single_stage_pcm_requires_cuda_device_and_own_first_audio_option(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    config = single_stage_tts_config
+    config.device_config.device = "cpu"
+    with pytest.raises(ValueError, match="requires CUDA"):
+        talker_stream_decode_enabled(config)
+    config.device_config.device = "cuda"
+    config.additional_config["talker_first_audio"] = True
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        talker_stream_decode_enabled(config)
+    config.additional_config.pop("talker_first_audio")
+    config.additional_config["talker_stream_first_audio"] = True
+    assert talker_stream_decode_enabled(config)
+
+
+def test_single_stage_audio_talker_requires_pcm_option(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    single_stage_tts_config.additional_config = {}
+    with pytest.raises(ValueError, match="requires talker_stream_decode=True"):
+        talker_stream_decode_enabled(single_stage_tts_config)
+
+
+def test_single_stage_pcm_requires_reset_capability_declaration(single_stage_tts_config):
+    from vllm_omni.model_executor.models.qwen3_tts.stream_decode import talker_stream_decode_enabled
+
+    single_stage_tts_config.model_config.supports_running_prefix_cache_reset = True
+    with pytest.raises(ValueError, match="must disable running prefix-cache reset"):
+        talker_stream_decode_enabled(single_stage_tts_config)

tests/config/test_omni_config.py

@@ -1046,6 +1046,7 @@ def test_sub_config_fields_match_structured_scopes():
         "interleave_mm_strings",
         "media_io_kwargs",
         "final_output",
+        "supports_running_prefix_cache_reset",
         "active_stream_window",
         "session_mode",
         "duplex_max_sessions",

tests/core/sched/test_omni_ar_scheduler_stale_drain.py

@@ -106,7 +106,7 @@ def _make_drain_sched(session: Request) -> MagicMock:
     return sched
 
 
-def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int) -> bool:
+def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token: int, spec=None) -> bool:
     """Feed one model-runner frame through update_from_output.
 
     Returns True when the frame's tokens were delivered (the append ran),
@@ -114,7 +114,7 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
     """
     scheduler_output = MagicMock(spec=SchedulerOutput)
     scheduler_output.num_scheduled_tokens = {session.request_id: num_scheduled}
-    scheduler_output.scheduled_spec_decode_tokens = {}
+    scheduler_output.scheduled_spec_decode_tokens = {session.request_id: spec} if spec else {}
     scheduler_output.num_invalid_spec_tokens = 0
 
     model_runner_output = MagicMock(spec=ModelRunnerOutput)
@@ -134,6 +134,27 @@ def _run_step(sched: MagicMock, session: Request, *, num_scheduled: int, token:
     return sched._update_request_with_output.called
 
 
+def test_connector_prompt_replacement_drops_old_frame_and_delivers_new_frame() -> None:
+    from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+
+    session = _make_session()
+    session.num_computed_tokens = 6
+    session.num_in_flight_tokens = 1
+    session.num_output_placeholders = 1
+    sched = _make_drain_sched(session)
+    sched.chunk_transfer_adapter.replaced_streaming_prompt_ids = {session.request_id}
+    sched.chunk_transfer_adapter.requests_with_ready_chunks = {session.request_id}
+    sched.chunk_transfer_adapter.requests_num_chunks_sent = {session.external_req_id: 2}
+
+    OmniSchedulerMixin._reset_ready_async_chunk_replacements(sched)
+    assert session.num_stale_output_tokens == 1 and session.drop_stale_output
+    assert session.num_output_placeholders == 0
+    sched._release_replaced_streaming_prompt_cache.assert_called_once_with(session)
+    assert _run_step(sched, session, num_scheduled=1, token=42) is False
+    assert session.num_stale_output_tokens == 0
+    assert _run_step(sched, session, num_scheduled=1, token=43) is True
+
+
 def test_exact_drain_delivers_new_segment_frame() -> None:
     """One in-flight decode frame at replacement: its late output drains the
     counter exactly and the new segment's first frame is delivered."""
@@ -225,3 +246,80 @@ def test_in_flight_prefill_chunk_drains_exactly_without_underflow() -> None:
     assert _run_step(sched, session, num_scheduled=5, token=42) is False  # late prefill chunk dropped
     assert session.num_stale_output_tokens == 0
     assert _run_step(sched, session, num_scheduled=1, token=43) is True  # new segment survives
+
+
+@pytest.mark.parametrize("drop", [False, True])
+def test_preemption_stale_output_follows_upstream_delivery_policy(drop):
+    request = _make_session()
+    request.status = RequestStatus.PREEMPTED
+    request.num_stale_output_tokens = 1
+    request.num_in_flight_tokens = 1
+    request.drop_stale_output = drop
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=1, token=42) is (not drop)
+    assert request.num_stale_output_tokens == 0
+    if not drop:
+        sched._update_request_with_output.assert_called_once_with(request, [42], is_stale=True)
+
+
+def test_preemption_stale_spec_rejection_does_not_roll_back_resumed_counters():
+    request = _make_session()
+    request.num_stale_output_tokens = 2
+    request.num_in_flight_tokens = 2
+    request.num_computed_tokens = 3
+    request.num_output_placeholders = 2
+    request.drop_stale_output = False
+    sched = _make_drain_sched(request)
+    assert _run_step(sched, request, num_scheduled=2, token=42, spec=[41])
+    assert request.num_computed_tokens == 3
+    assert request.num_output_placeholders == 2
+
+
+@pytest.mark.parametrize(
+    "arch,supports_reset,running,reset_running,blocked",
+    [
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, True, True),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, False, True, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", False, True, False, False),
+        ("Qwen3TTSTalkerForConditionalGeneration", True, True, True, False),
+        ("AnotherStatefulAudioModel", False, True, True, True),
+        ("AnotherStatefulAudioModel", True, True, True, False),
+    ],
+)
+def test_reset_running_streaming_codec_is_rejected_before_preemption(
+    mocker, arch, supports_reset, running, reset_running, blocked
+):
+    from vllm.config import VllmConfig
+    from vllm.v1.core.sched.scheduler import Scheduler
+
+    from vllm_omni.config.model import OmniModelConfig
+
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.vllm_config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=mocker.Mock(
+            spec=OmniModelConfig,
+            model_arch=arch,
+            supports_running_prefix_cache_reset=supports_reset,
+        ),
+    )
+    request = _make_session()
+    request.num_in_flight_tokens = 1
+    request.num_output_placeholders = 1
+    scheduler.running = [request] if running else []
+
+    def reset_cache(_self, reset_running_requests=False, reset_connector=False):
+        assert reset_running_requests is reset_running and reset_connector is True
+        return True
+
+    reset = mocker.patch.object(Scheduler, "reset_prefix_cache", autospec=True, side_effect=reset_cache)
+
+    assert scheduler.reset_prefix_cache(reset_running, reset_connector=True) is (not blocked)
+    if blocked:
+        reset.assert_not_called()
+        assert scheduler.running == [request]
+        assert request.status == RequestStatus.RUNNING
+        assert request.num_in_flight_tokens == request.num_output_placeholders == 1
+        assert not request.drop_stale_output
+    else:
+        reset.assert_called_once()

tests/engine/test_orchestrator_first_audio.py

@@ -114,10 +114,10 @@ def _raw(*, first=False, required=False, terminal=False, samples=(1, 1, 1, 1)):
     )
 
 
-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)
+async def _codec_output(obj, raw, stage=1):
+    replica = obj.stage_pools[stage].replica_id
+    processed = await obj._process_llm_stage_outputs(stage, replica, raw, set())
+    await obj._handle_processed_outputs(stage, replica, processed)
 
 
 def _messages(obj):
@@ -132,6 +132,57 @@ def _audio(output):
     return torch.cat(audio) if isinstance(audio, list) else audio
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("kind", list(RequestOutputKind))
+@pytest.mark.parametrize("suffix_overtook", [False, True])
+async def test_single_stage_first_audio_obeys_output_kind_and_terminal_order(kind, suffix_overtook):
+    obj = _orchestrator(kind, final_stage_id=0)
+    obj.stage_pools = [obj.stage_pools[1]]
+    suffix = _raw(required=True, samples=(20, 21))
+    terminal = _raw(required=True, terminal=True, samples=(30, 31))
+    if suffix_overtook:
+        await _codec_output(obj, suffix, stage=0)
+        await _codec_output(obj, terminal, stage=0)
+        assert obj.output_async_queue.empty()
+    await obj._route_upstream_first_audio(0, 2, _raw(first=True, samples=(10, 11)))
+    if not suffix_overtook:
+        await _codec_output(obj, suffix, stage=0)
+        await _codec_output(obj, terminal, stage=0)
+    messages = _messages(obj)
+    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 [_audio(message.engine_outputs).tolist() for message in messages] == expected[kind]
+    assert all((message.stage_id, message.replica_id) == (0, 2) for message in messages)
+    assert messages[-1].finished
+    assert not obj.request_states["r"].pending_first_audio_outputs
+
+
+@pytest.mark.asyncio
+async def test_single_stage_cancellation_discards_overtaking_terminal_and_late_first_audio():
+    obj = _orchestrator(final_stage_id=0)
+    obj.stage_pools = [obj.stage_pools[1]]
+    await _codec_output(obj, _raw(required=True, terminal=True), stage=0)
+    assert obj.request_states["r"].pending_first_audio_outputs
+    await obj._cleanup_request_ids(["r"], abort=True)
+    await obj._route_upstream_first_audio(0, 2, _raw(first=True))
+    assert not obj.request_states and obj.output_async_queue.empty()
+
+
+@pytest.mark.asyncio
+async def test_single_stage_first_audio_processing_failure_releases_waiting_terminal():
+    obj = _orchestrator(final_stage_id=0)
+    obj.stage_pools = [obj.stage_pools[1]]
+    await _codec_output(obj, _raw(required=True, terminal=True), stage=0)
+    obj.stage_pools[0].process_llm_raw_outputs = AsyncMock(side_effect=RuntimeError("processing failed"))
+    await obj._route_upstream_first_audio(0, 2, _raw(first=True))
+    error = obj.output_async_queue.get_nowait()
+    assert isinstance(error, ErrorMessage) and error.stage_id == 0
+    assert not obj.request_states
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize("kind", list(RequestOutputKind))
 @pytest.mark.parametrize("codec_overtook", [False, True])

tests/engine/test_stage_engine_core_proc.py

@@ -1,9 +1,37 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from types import SimpleNamespace
 from unittest.mock import patch
 
+import pytest
 from vllm.v1.engine.core import EngineCoreProc
+from vllm.v1.executor.uniproc_executor import UniProcExecutor
 
-from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc
+from vllm_omni.engine.stage_engine_core_proc import StageEngineCoreProc, _bind_first_audio_sink
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    "first_decoder,stream_decoder,stream_first_audio,expected",
+    [(True, False, False, True), (False, True, False, False), (False, True, True, True), (False, False, True, False)],
+)
+def test_first_audio_sink_preserves_first_decoder_and_opts_in_stream_decoder(
+    mocker,
+    first_decoder,
+    stream_decoder,
+    stream_first_audio,
+    expected,
+):
+    executor = UniProcExecutor.__new__(UniProcExecutor)
+    executor.driver_worker = mocker.Mock()
+    runner = executor.driver_worker.worker.model_runner
+    runner.model.first_frame_decoder = object() if first_decoder else None
+    runner.model.stream_decoder = object() if stream_decoder else None
+    runner.model.stream_first_audio = stream_first_audio
+    assert _bind_first_audio_sink(executor, mocker.Mock(), mocker.Mock()) is expected
+    assert runner.model_state.set_first_audio_sink.call_count == int(expected)
 
 
 def test_preprocess_add_request_preserves_omni_fields():

tests/entrypoints/openai_api/test_audio_format.py

@@ -346,3 +346,19 @@ def test_chat_audio_metadata_survives_response_serialization(stream, fmt):
             assert audio_file.samplerate == 22050
     else:
         assert len(raw) == frame_count * 2
+
+
+def test_float32_pcm_fast_path_matches_soundfile():
+    from vllm_omni.entrypoints.openai.audio_utils_mixin import _float32_to_pcm16_bytes
+
+    rng = np.random.default_rng(0)
+    audio = np.concatenate(
+        [
+            rng.uniform(-1.2, 1.2, 4096),
+            np.array([-1.0, 1.0, 0.0, -0.0, 32767 / 32768, -1 / 32768, 0.5 / 32768, -0.5 / 32768]),
+        ]
+    ).astype(np.float32)
+    with BytesIO() as buffer:
+        soundfile.write(buffer, audio, 24000, format="RAW", subtype="PCM_16")
+        expected = buffer.getvalue()
+    assert _float32_to_pcm16_bytes(audio) == expected

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -569,6 +569,27 @@ class TestCodePredictorPerRowGenerators:
             )
 
 
+@pytest.mark.parametrize("fused_requested", [False, True])
+@pytest.mark.parametrize("invalid_input", ["uniforms", "generators"])
+def test_sampling_contract_rejected_before_dispatch(mocker, loaded_target_classes, fused_requested, invalid_input):
+    predictor, talker_config = TestCodePredictorPerRowGenerators()._make_predictor(mocker, loaded_target_classes)
+    predictor._fused_requested = fused_requested
+    predictor._fused = mocker.Mock()
+    setup = mocker.patch.object(predictor, "_setup_compile")
+    uniforms = torch.full((2, 3, 1 if invalid_input == "uniforms" else 64), 0.5)
+    generators = [torch.Generator()] if invalid_input == "generators" else None
+    with pytest.raises(ValueError, match="sample_uniforms must have shape|one entry per row"):
+        predictor(
+            torch.zeros(2, dtype=torch.long),
+            torch.randn(2, talker_config.hidden_size),
+            torch.randn(2, talker_config.hidden_size),
+            sample_uniforms=uniforms,
+            generators=generators,
+        )
+    setup.assert_not_called()
+    predictor._fused.assert_not_called()
+
+
 class TestCodePredictorModelDtype:
     """Test the inner model forward with different dtypes."""
 
@@ -663,6 +684,8 @@ class TestCodePredictorGraphReplay:
 
         predictor = object.__new__(code_predictor_wrapper)
         torch.nn.Module.__init__(predictor)
+        predictor._fused_requested = False
+        predictor._fused = None
         predictor._num_groups = 3
         predictor._model_dtype = torch.float32
         predictor._setup_compile = mocker.Mock()

tests/model_executor/models/qwen3_tts/test_code_predictor_sampling.py

@@ -2,6 +2,8 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Supplied-uniform samples must match the existing ATen path exactly."""
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
@@ -12,6 +14,79 @@ pytestmark = [
 ]
 
 
+@pytest.mark.parametrize("top_k", [0, 1, 50, 2048])
+@torch.inference_mode()
+def test_fused_sample_kernel_preserves_signed_zero_cutoff_ties(top_k):
+    from vllm_omni.model_executor.models.qwen3_tts.fused_code_predictor import _cp_sample_kernel
+
+    vocab, hidden, groups = 2048, 64, 4
+    logits = torch.zeros(2, vocab, device="cuda", dtype=torch.bfloat16)
+    logits[0, 1] = -0.0
+    logits[1, ::2] = -0.0
+    uniforms = torch.full((2, groups - 1, vocab), 0.01, device="cuda")[:, 1, :]
+    uniforms[0, 1] = uniforms[1, 2] = 0.99
+    codes = torch.full((2, groups), -1, device="cuda", dtype=torch.int64)
+    table = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16)
+    next_input = torch.empty(2, hidden, device="cuda", dtype=table.dtype)
+
+    def sample():
+        _cp_sample_kernel[(2,)](
+            logits, uniforms, uniforms.stride(0), codes, 1, groups, table, next_input, 1.0,
+            V=vocab, TOPK=top_k, HID=hidden, HAS_NEXT=True, num_warps=4,
+        )  # fmt: skip
+
+    scaled = logits.float()
+    if top_k > 0:
+        scaled = scaled.masked_fill(scaled < scaled.topk(top_k).values[:, -1:], -torch.inf)
+    expected = (scaled - torch.log(-torch.log(uniforms))).argmax(-1)
+    assert expected.tolist() == [1, 2]
+    sample()
+    torch.testing.assert_close(codes[:, 1], expected)
+    torch.testing.assert_close(next_input, table[expected])
+    assert (codes[:, [0, 2, 3]] == -1).all()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        sample()
+    codes.fill_(-1)
+    graph.replay()
+    torch.testing.assert_close(codes[:, 1], expected)
+    torch.testing.assert_close(next_input, table[expected])
+
+
+@torch.inference_mode()
+def test_fused_predictor_warmup_with_capacity_one():
+    from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+        Qwen3TTSTalkerCodePredictorConfig,
+        Qwen3TTSTalkerConfig,
+    )
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM,
+    )
+
+    cp = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1,
+        num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4,
+    )  # fmt: skip
+    talker = Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4)
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(stage_connector_config={}),
+        additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True},
+        scheduler_config=SimpleNamespace(max_num_seqs=1),
+    )
+    predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(
+        vllm_config=config, config=cp, talker_config=talker
+    ).to(device="cuda", dtype=torch.bfloat16)
+    predictor._setup_compile()
+    assert predictor._fused is not None
+    hidden = torch.randn(1, 1, 64, device="cuda", dtype=torch.bfloat16)
+    result = predictor(
+        torch.zeros(1, 1, device="cuda", dtype=torch.long), hidden, hidden,
+        sample_uniforms=torch.full((1, 3, 64), 0.5, device="cuda"),
+    )  # fmt: skip
+    assert result.shape == (1, 4)
+    assert ((result >= 0) & (result < 64)).all()
+
+
 @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])
@@ -66,3 +141,47 @@ def test_topk_gumbel_accepts_strided_supplied_uniforms(top_k):
     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)
+
+
+@pytest.mark.parametrize("batch", [1, 4])
+@pytest.mark.parametrize("top_k", [1, 50])
+@torch.inference_mode()
+def test_fused_predictor_matches_full_reference_with_fixed_uniforms(mocker, batch, top_k):
+    from vllm.config import VllmConfig
+
+    from vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts import (
+        Qwen3TTSTalkerCodePredictorConfig,
+        Qwen3TTSTalkerConfig,
+    )
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
+        Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM,
+    )
+
+    torch.manual_seed(42)
+    cp = Qwen3TTSTalkerCodePredictorConfig(
+        vocab_size=64, hidden_size=64, intermediate_size=128, num_hidden_layers=1,
+        num_attention_heads=4, num_key_value_heads=2, head_dim=16, num_code_groups=4,
+    )  # fmt: skip
+    config = mocker.Mock(
+        spec=VllmConfig,
+        model_config=mocker.Mock(stage_connector_config={}),
+        additional_config={"code_predictor_kv_cache": True, "code_predictor_fused": True},
+        scheduler_config=mocker.Mock(max_num_seqs=4),
+    )
+    predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(
+        vllm_config=config, config=cp, talker_config=Qwen3TTSTalkerConfig(hidden_size=64, num_code_groups=4)
+    ).to(device="cuda", dtype=torch.bfloat16)
+    predictor._setup_compile()
+    assert predictor._fused is not None
+    ids = torch.randint(0, 64, (batch, 1), device="cuda")
+    embeds = torch.randn(batch, 1, 64, device="cuda", dtype=torch.bfloat16)
+    hidden = torch.randn_like(embeds)
+    uniforms = torch.rand(batch, 3, 64, device="cuda").clamp_(1e-6, 1 - 1e-6)
+    actual = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms).clone()
+    # Use identical parameters through the full re-prefill and ATen sampling
+    # path, bypassing both the fused predictor and its frame-local KV cache.
+    predictor._fused_requested = False
+    predictor._fused = None
+    predictor._frame_cache = None
+    expected = predictor(ids, embeds, hidden, top_k=top_k, sample_uniforms=uniforms)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)

tests/model_executor/models/qwen3_tts/test_first_audio_graphs.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Prefix graphs must follow per-request first-audio delivery, not a capture hint."""
 
+import math
 from types import SimpleNamespace
 
 import pytest
@@ -18,6 +19,38 @@ pytestmark = [pytest.mark.core_model]
 DEVICE = torch.device("cuda:0")
 
 
+@pytest.mark.cpu
+@pytest.mark.parametrize("capacity,model_type", [(1, "base"), (128, "base"), (128, "custom_voice")])
+def test_stream_priming_graph_workspace_respects_token_budget(monkeypatch, capacity, model_type):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz import streaming_decoder
+
+    captures = []
+
+    def capture(stream, sizes, frames=1):
+        captures.append((sizes, frames))
+        return SimpleNamespace(sizes=sizes, frames=frames)
+
+    monkeypatch.setattr(streaming_decoder, "StreamingDecodeGraphs", capture)
+    model = SimpleNamespace(
+        stream_decoder=SimpleNamespace(max_frames=33),
+        stream_graphs=None,
+        stream_prime_graphs=None,
+        stream_ref_context_frames=25,
+        stream_chunk_frames=25,
+        config=SimpleNamespace(tts_model_type=model_type),
+        vllm_config=SimpleNamespace(scheduler_config=SimpleNamespace(max_num_batched_tokens=512)),
+    )
+    sizes = [size for size in [1, 2, 4, 8, 16, 32, 64, 128] if size <= capacity]
+    Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes)
+    expected = [(sizes, 1)]
+    if model_type == "base":
+        expected.append(([size for size in sizes if size * 25 <= 512], 25))
+    assert captures == expected
+    Qwen3TTSTalkerForConditionalGeneration.capture_stream_decode_graphs(model, sizes)
+    assert captures == expected
+
+
 @pytest.mark.cpu
 @pytest.mark.parametrize("skip_flags", [(False, False), (True, True), (False, True), (True, False)])
 @pytest.mark.parametrize("initial_frames", [1, 3])
@@ -124,3 +157,134 @@ def test_xvec_graphs_follow_request_first_audio_flags(initial_frames, time_major
             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)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.inference_mode()
+def test_stream_decoder_preemption_restores_all_state_into_a_different_slot():
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import StreamingCodecDecoder
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    stream = StreamingCodecDecoder(decoder, num_slots=2, dtype=torch.bfloat16)
+    codes = torch.randint(0, 32, (1, 8, 2), device=DEVICE)
+    slot = torch.tensor([0], device=DEVICE, dtype=torch.int32)
+    for pos in range(7):
+        stream(codes[:, pos : pos + 1], slot, torch.tensor([pos], device=DEVICE, dtype=torch.int32))
+    saved = stream.save_slot(0)
+    expected = stream(codes[:, 7:], slot, torch.tensor([7], device=DEVICE, dtype=torch.int32)).clone()
+    # Both the freed slot and the destination can hold an unrelated request.
+    other_codes = (codes[:, :3] + 1) % 32
+    for pos in range(3):
+        stream(
+            other_codes[:, pos : pos + 1].expand(2, -1, -1).contiguous(),
+            torch.tensor([0, 1], device=DEVICE, dtype=torch.int32),
+            torch.full((2,), pos, device=DEVICE, dtype=torch.int32),
+        )
+    without_restore = stream(
+        codes[:, 7:],
+        torch.tensor([1], device=DEVICE, dtype=torch.int32),
+        torch.tensor([7], device=DEVICE, dtype=torch.int32),
+    ).clone()
+    assert not torch.equal(without_restore, expected)
+    stream.restore_slot(1, saved)
+    actual = stream(
+        codes[:, 7:],
+        torch.tensor([1], device=DEVICE, dtype=torch.int32),
+        torch.tensor([7], device=DEVICE, dtype=torch.int32),
+    )
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("use_graph", [False, True])
+@torch.inference_mode()
+def test_batched_reference_priming_matches_serial_decoder_continuation(use_graph):
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import (
+        StreamingCodecDecoder,
+        StreamingDecodeGraphs,
+    )
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=torch.bfloat16)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    stream = StreamingCodecDecoder(decoder, num_slots=5, dtype=torch.bfloat16)
+    graphs = StreamingDecodeGraphs(stream, [2], frames=25) if use_graph else None
+    model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+    torch.nn.Module.__init__(model)
+    model.stream_decoder = stream
+    model.stream_prime_graphs = graphs
+    model.stream_chunk_frames = 25
+    refs = [torch.randint(0, 32, (length, 2), device=DEVICE) for length in [25, 3, 51, 25]]
+    slots = [2, 0, 3, 4]
+    frame = torch.randint(0, 32, (4, 1, 2), device=DEVICE)
+    slot_tensor = torch.tensor(slots, device=DEVICE, dtype=torch.int32)
+    positions = torch.tensor([len(ref) for ref in refs], device=DEVICE, dtype=torch.int32)
+    for slot, ref in zip(slots, refs, strict=True):
+        for t0 in range(0, len(ref), 25):
+            stream(
+                ref[None, t0 : t0 + 25].to(torch.int32).contiguous(),
+                torch.tensor([slot], device=DEVICE, dtype=torch.int32),
+                torch.tensor([t0], device=DEVICE, dtype=torch.int32),
+            )
+    expected = stream(frame, slot_tensor, positions).clone()
+    untouched = stream.save_slot(1)
+    # Reuse dirty slots; pos=0 must reset every causal layer for each group.
+    model.prime_stream_decoder(list(zip(slots, refs, strict=True)))
+    actual = stream(frame, slot_tensor, positions)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=3e-5)
+    for got, want in zip(stream.save_slot(1), untouched, strict=True):
+        torch.testing.assert_close(got, want, rtol=0, atol=0)
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("use_graph", [False, True])
+@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
+@torch.inference_mode()
+def test_stream_frames_match_fp32_exact_decoder(use_graph, dtype):
+    from tests.model_executor.models.qwen3_tts.test_time_major_decoder import _make_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import (
+        StreamingCodecDecoder,
+        StreamingDecodeGraphs,
+    )
+
+    decoder = _make_decoder().to(device=DEVICE, dtype=dtype)
+    # The tiny HF default initialization shrinks the deep conv stack to
+    # near-zero PCM. Keep the oracle signal large enough to catch lost state.
+    for module in decoder.modules():
+        if isinstance(module, (torch.nn.Conv1d, torch.nn.ConvTranspose1d, torch.nn.Linear)):
+            torch.nn.init.kaiming_uniform_(module.weight, a=math.sqrt(5))
+            if module.bias is not None:
+                module.bias.uniform_(-0.01, 0.01)
+    for name, param in decoder.named_parameters():
+        if name.endswith("embedding_sum"):
+            param.normal_()
+    decoder.config.head_dim = decoder.config.hidden_size // decoder.config.num_attention_heads
+    reference = _make_decoder().to(device=DEVICE, dtype=torch.float32)
+    reference.load_state_dict(decoder.state_dict())
+    reference.precompute_snake_caches()
+    # Cross both the attention window and KV ring wrap, checking every frame.
+    codes = torch.randint(0, 32, (2, 136, 2), device=DEVICE)
+    expected = reference._forward_exact(codes.transpose(1, 2))[:, 0]
+    assert expected.square().sum() > 1e-3
+    stream = StreamingCodecDecoder(decoder, num_slots=3, dtype=dtype)
+    decode = StreamingDecodeGraphs(stream, [2]) if use_graph else stream
+    slots = torch.tensor([2, 0], device=DEVICE, dtype=torch.int32)
+    spf = int(stream.spf)
+    for frame in range(codes.shape[1]):
+        positions = torch.full((2,), frame, device=DEVICE, dtype=torch.int32)
+        actual = decode(codes[:, frame : frame + 1].contiguous().to(torch.int32), slots, positions).reshape(2, -1)
+        exact_frame = expected[:, frame * spf : (frame + 1) * spf]
+        relative_rms = (actual - exact_frame).square().mean(1).sqrt() / exact_frame.square().mean(1).sqrt()
+        # BF16 kernels are approximate; compare every frame to the independent
+        # FP32 full decoder, not only one streaming implementation to itself.
+        tolerance = 0.05 if dtype == torch.bfloat16 else 5e-4
+        assert torch.all(relative_rms < tolerance), (frame, relative_rms)

tests/model_executor/models/qwen3_tts/test_streaming_col2im_cuda.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Streaming col2im preserves PCM inputs and ping-pong slot state."""
+
+import pytest
+import torch
+from vllm.triton_utils import triton
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.streaming_decoder import _col2im_kernel
+
+pytestmark = [pytest.mark.core_model, pytest.mark.tts]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("batch,frames,channels,rate", [(1, 1, 5, 2), (3, 3, 129, 3), (9, 1, 64, 4)])
+def test_col2im_partial_tiles_and_reordered_slot_state(batch: int, frames: int, channels: int, rate: int):
+    torch.manual_seed(23)
+    capacity = batch + 2
+    z = torch.randn(batch, frames, 2 * rate, channels).to(torch.bfloat16)
+    bias = torch.randn(channels).to(torch.bfloat16)
+    state = torch.randn(2, capacity, rate, channels).to(torch.bfloat16)
+    expected_state = state.clone()
+    slots = torch.arange(batch - 1, -1, -1, dtype=torch.int32)
+    positions = torch.arange(batch, dtype=torch.int32)
+    parity = torch.arange(capacity, dtype=torch.int32) % 2
+    expected = torch.empty(batch, frames, rate, channels, dtype=torch.bfloat16)
+    for b in range(batch):
+        slot = int(slots[b])
+        side = int(parity[slot])
+        for t in range(frames):
+            previous = z[b, t - 1, rate:] if t else state[side, slot] if positions[b] else torch.zeros(rate, channels)
+            expected[b, t] = (z[b, t, :rate].float() + previous.float() + bias.float()).to(torch.bfloat16)
+        expected_state[1 - side, slot] = z[b, -1, rate:]
+
+    z_gpu, bias_gpu, state_gpu = z.cuda(), bias.cuda(), state.cuda()
+    slots_gpu, positions_gpu, parity_gpu = slots.cuda(), positions.cuda(), parity.cuda()
+    actual = torch.empty_like(expected, device="cuda")
+    width = min(128, triton.next_power_of_2(channels))
+    rows = batch * frames
+    _col2im_kernel[(rows, rate, triton.cdiv(channels, width))](
+        z_gpu,
+        bias_gpu,
+        state_gpu,
+        slots_gpu,
+        positions_gpu,
+        parity_gpu,
+        actual,
+        frames,
+        capacity,
+        R=rate,
+        C=channels,
+        BC=width,
+        num_warps=4,
+    )
+    torch.testing.assert_close(actual.cpu(), expected, rtol=0, atol=0)
+    torch.testing.assert_close(state_gpu.cpu(), expected_state, rtol=0, atol=0)

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -594,6 +594,40 @@ def test_streaming_ref_code_context_is_bounded_for_batchable_shapes():
     torch.testing.assert_close(frames[:3], ref_code[-3:])
 
 
+@pytest.mark.parametrize("tensor_frames", [False, True])
+def test_async_chunks_preserve_codebook_order_with_reference_and_short_tail(tensor_frames):
+    tm = _tm(left_context=3, initial_chunk_frames=1)
+    rid = "r-codebook-order"
+    quantizers = 16
+    reference = [[1500 + f * quantizers + q for q in range(quantizers)] for f in range(5)]
+    tm.request_payload[rid] = torch.tensor(reference, dtype=torch.long)
+    generated = [[1 + f * quantizers + q for q in range(quantizers)] for f in range(30)]
+
+    for start, end, finished in [(0, 1, False), (1, 26, False), (26, 30, True)]:
+        rows = generated[:end]
+        if tensor_frames:
+            # Non-contiguous int32 rows exercise conversion independently of
+            # the native producer's contiguous CPU-long representation.
+            storage = torch.tensor(rows, dtype=torch.int32).repeat_interleave(2, dim=1)
+            tm.code_prompt_token_ids[rid] = list(storage[:, ::2].unbind())
+        else:
+            tm.code_prompt_token_ids[rid] = rows
+        payload = talker2code2wav_async_chunk(
+            transfer_manager=tm,
+            multimodal_output=None if finished else {},
+            request=_req(rid, finished=finished),
+            is_finished=finished,
+        )
+        assert payload is not None
+        expected = (reference[-3:] if start == 0 else []) + generated[start:end]
+        codes = payload.codes.audio
+        assert codes.dtype == torch.long and codes.device.type == "cpu" and codes.is_contiguous()
+        torch.testing.assert_close(codes.reshape(quantizers, -1).T, torch.tensor(expected))
+        assert payload.meta.left_context_size == (3 if start == 0 else 0)
+        assert payload.meta.finished.item() == finished
+        tm.put_req_chunk[rid] += 1
+
+
 def test_ref_code_context_can_be_buffered_before_first_emit():
     tm = _tm()
     rid = "r-ref-buffered"

tests/worker_v2/test_index_staging.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+
+from vllm_omni.utils.device_copy import index_to_device
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+def test_index_staging_survives_later_copies_and_graph_replay():
+    stream = torch.cuda.Stream()
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.stream(stream):
+        with torch.cuda.graph(graph, stream=stream):
+            captured = index_to_device([17] * 6, "cuda")
+    for _ in range(32):
+        index_to_device([99] * 6, "cuda")
+    graph.replay()
+    torch.accelerator.synchronize()
+    assert captured.tolist() == [17] * 6

tests/worker_v2/test_omni_ar_model_runner.py

@@ -25,6 +25,9 @@ class _FakeStream:
     def wait_stream(self, _stream) -> None:
         pass
 
+    def wait_event(self, _event) -> None:
+        pass
+
 
 class _FakeEvent:
     def record(self, _stream) -> None:
@@ -138,6 +141,7 @@ def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history)
 def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None:
     runner = OmniARModelRunner.__new__(OmniARModelRunner)
     runner.model_config = SimpleNamespace(async_chunk=True)
+    runner.model = SimpleNamespace()
     runner._async_mm_snapshot_slots, runner._async_mm_snapshot_events = [{}], [None]
     runner._async_mm_snapshot_pending, runner._async_mm_snapshot_cursor = [False], 0
     waited: list[object] = []
@@ -307,6 +311,231 @@ def test_request_reference_codes_preserve_local_axis(prefill_first, padded):
         assert torch.equal(outputs[i]["codes.audio"], codes[offsets[i] : offsets[i + 1]])
 
 
+def test_stream_audio_cancel_rejects_inflight_output_and_id_reuse() -> None:
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    cancelled = buffer.add("request")
+    buffer.finish({"request"})
+    replacement = buffer.add("request")
+    assert cancelled.push(torch.ones(4), False) is None
+    assert not cancelled.pending
+    assert replacement.push(torch.full((4,), 2.0), False).tolist() == [2.0] * 4
+    buffer.finish({"request"})
+    assert not buffer.requests
+
+
+def test_stream_audio_length_end_flushes_first_and_partial_chunks() -> None:
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    state = buffer.add("request")
+    output = StreamingAudioOutput(
+        torch.ones(1, 4),
+        torch.tensor([True]),
+        torch.tensor(24000),
+        None,
+        np.array([0, 1]),
+        [state],
+        np.array([True]),
+    )
+    assert output.get_output()[0]["model_outputs"].numel() == 4
+    assert not state.active
+    state = buffer.add("another")
+    assert state.push(torch.ones(4), False).numel() == 4
+    assert state.push(torch.ones(4), False) is None
+    assert state.push(torch.ones(4), True).numel() == 8
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.float32, torch.float64, torch.bfloat16])
+def test_stream_audio_batch_views_retain_pcm_after_output_release(dtype):
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    states = [buffer.add("a"), buffer.add("b")]
+    expected: list[list[torch.Tensor]] = [[], []]
+    emitted: list[list[torch.Tensor]] = [[], []]
+    for step in range(5):
+        # Mixed prefill/decode spans, invalid rows, and a final partial chunk.
+        wav = (torch.arange(16).reshape(4, 4) + 100 * step).to(dtype)
+        valid = torch.tensor([False, True, True, step != 4])
+        output = StreamingAudioOutput(
+            wav,
+            valid,
+            torch.tensor(24000),
+            None,
+            np.array([0, 3, 4]),
+            states,
+            np.array([step == 4, step == 4]),
+        )
+        for i, (start, end) in enumerate([(0, 3), (3, 4)]):
+            expected[i].append(wav[start:end][valid[start:end]].reshape(-1).clone())
+        for i, payload in enumerate(output.get_output()):
+            if payload is not None:
+                emitted[i].append(payload["model_outputs"])
+        del output, wav  # Pending views must keep the old host allocation alive.
+    for i in range(2):
+        assert torch.equal(torch.cat(emitted[i]), torch.cat(expected[i]))
+        assert torch.cat(emitted[i]).dtype == dtype
+        assert not states[i].active and not states[i].pending
+        assert states[i].pending_samples == 0
+
+
+def test_stream_audio_numpy_direct_first_cancel_and_mixed_fallback():
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    state = buffer.add("r")
+    state.accept_first_audio()
+    assert state.push(np.arange(8, dtype=np.float32), False) is None
+    assert state.pending_samples == 4
+    # A dtype fallback can arrive while a NumPy frame is pending.
+    chunk = state.push(torch.arange(8, 12, dtype=torch.bfloat16), True)
+    assert chunk.tolist() == list(range(4, 12))
+    assert state.pending_samples == 0
+    buffer.finish({"r"})
+    replacement = buffer.add("r")
+    assert state.push(np.ones(4, dtype=np.float32), False) is None
+    assert replacement.push(np.full(4, 99, dtype=np.float32), False).tolist() == [99] * 4
+    replacement.push(np.ones(4, dtype=np.float32), False)
+    assert replacement.pending_samples == 4
+    buffer.finish({"r"})
+    assert replacement.pending_samples == 0 and not replacement.pending
+
+
+def test_stream_audio_prefill_length_end_reaches_output(mocker) -> None:
+    from vllm.v1.worker.gpu.input_batch import InputBatch
+    from vllm.v1.worker.gpu.states import RequestState
+
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer
+
+    owner = mocker.Mock()
+    owner.vllm_config.model_config.max_model_len = 4096
+    owner._stream_decode_event = None
+    eager = EagerMTPState(owner)
+    eager._audio_buffer = StreamingAudioBuffer(4, 25)
+    eager._audio_buffer.add("request")
+    eager._frame_requests = {"request"}
+    batch = InputBatch.__new__(InputBatch)
+    batch.num_reqs, batch.req_ids = 1, ["request"]
+    batch.num_computed_tokens_np = np.array([0])
+    batch.num_scheduled_tokens = np.array([3])
+    batch.idx_mapping_np = np.array([0])
+    batch.query_start_loc_np = np.array([0, 3])
+    batch.is_prefilling_np = np.array([True])
+    states = RequestState.__new__(RequestState)
+    states.max_seq_len = np.array([4])
+    mm = {
+        "model_outputs": torch.ones(3, 4),
+        "meta": {"codec_frame_valid": torch.tensor([0, 0, 1])},
+        "sr": [torch.tensor(24000)],
+    }
+    output = eager.prepare_audio_output(batch, states, mm)
+    assert output.get_output()[0]["model_outputs"].numel() == 4
+    assert "model_outputs" not in mm
+
+
+def test_stream_audio_preemption_preserves_position_rng_and_pending_pcm(mocker) -> None:
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState, TalkerInputs
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer
+
+    owner = mocker.Mock()
+    owner._stream_pos = {"request": 37}
+    owner._eager_embeds = torch.arange(18).reshape(6, 3).float()
+    owner._eager_ready = {}
+    owner.intermediate_buffer.buffers = [{} for _ in range(6)]
+    generator = torch.Generator().manual_seed(42)
+    owner._mtp_generators = {"request": generator}
+    saved = [torch.arange(6)]
+    owner.model.stream_decoder.save_slot.return_value = saved
+    eager = EagerMTPState(owner)
+    eager._talker_inputs["request"] = TalkerInputs(torch.ones(20, 3), 10)
+    eager._side_stream = mocker.Mock()
+    eager._audio_buffer = StreamingAudioBuffer(4, 25)
+    request = eager._audio_buffer.add("request")
+    request.push(torch.ones(4), False)
+    request.push(torch.ones(4), False)
+    eager.suspend_audio("request", 2)
+    owner._stream_pos.clear()
+    owner._mtp_generators.clear()
+    eager.resume_audio("request", 5)
+    assert owner._stream_pos["request"] == 37
+    assert owner._mtp_generators["request"] is generator
+    assert eager._restore_audio["request"] is saved
+    assert sum(part.numel() for part in request.pending) == 4
+    # A second preemption can happen before the restored slot is used.
+    eager.suspend_audio("request", 5)
+    owner.model.stream_decoder.save_slot.assert_called_once_with(2)
+    eager.finish_audio({"request"})
+    assert not eager._suspended_audio and not eager._restore_audio
+    assert not eager._audio_buffer.requests and not request.active
+
+
+def test_stream_audio_replay_uses_exact_inputs_and_consumes_last_frame_once(mocker):
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState, TalkerInputs
+
+    owner = mocker.Mock()
+    owner._eager_embeds = torch.tensor([[20.0, 30.0]])
+    owner.model.preprocess.return_value = (None, None, {"mtp_inputs": (None, torch.tensor([[2.0, 3.0]]))})
+    owner.intermediate_buffer.buffers = [{"req_id": "request"}]
+    eager = EagerMTPState(owner)
+    eager._restore_audio["request"] = []
+    history = torch.arange(8).reshape(4, 2).float()
+    eager._talker_inputs["request"] = TalkerInputs(history, 4)
+    first = torch.empty(2, 2)
+    assert eager.replay_inputs("request", 0, 0, torch.ones(2), first)
+    torch.testing.assert_close(first, history[:2])
+    owner.model.preprocess.assert_not_called()
+    last = torch.empty(3, 2)
+    assert eager.replay_inputs("request", 0, 2, torch.ones(3), last)
+    torch.testing.assert_close(last, torch.cat((history[2:], torch.tensor([[22.0, 33.0]]))))
+    owner.model.preprocess.assert_called_once()
+    with pytest.raises(RuntimeError, match="exceeds saved inputs"):
+        eager.replay_inputs("request", 0, 4, torch.ones(2), torch.empty(2, 2))
+
+
+def test_stream_reference_priming_groups_lengths_without_padding_or_slot_aliasing(mocker):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+
+    primes = [(4, torch.full((26, 2), 4)), (1, torch.full((3, 2), 1)), (7, torch.full((26, 2), 7))]
+    calls = []
+
+    def stream(codes, slots, pos):
+        calls.append((codes.clone(), slots.clone(), pos.clone()))
+
+    decoder = mocker.Mock(side_effect=stream, device=torch.device("cpu"))
+    model = mocker.Mock(stream_decoder=decoder, stream_prime_graphs=None, stream_chunk_frames=25)
+    Qwen3TTSTalkerForConditionalGeneration.prime_stream_decoder(model, primes)
+    assert len(calls) == 3
+    for (codes, slots, pos), frames, indices, position in zip(
+        calls, [25, 1, 3], [[4, 7], [4, 7], [1]], [0, 25, 0], strict=True
+    ):
+        assert codes.shape == (len(indices), frames, 2)
+        assert codes.dtype == slots.dtype == pos.dtype == torch.int32
+        assert slots.tolist() == indices and pos.tolist() == [position] * len(indices)
+        for row, idx in enumerate(indices):
+            assert torch.all(codes[row] == idx)
+    Qwen3TTSTalkerForConditionalGeneration.prime_stream_decoder(model, [])
+    assert len(calls) == 3
+
+
+@pytest.mark.parametrize("context_frames", [25, 72])
+def test_model_owned_reference_context_respects_configured_boundary(mocker, context_frames):
+    from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import Qwen3TTSTalkerForConditionalGeneration
+
+    model = mocker.Mock(stream_ref_context_frames=context_frames, talker_config=mocker.Mock(num_code_groups=16))
+    reference = torch.arange(100 * 16).reshape(100, 16)
+    info = {"codes": {"ref": reference}}
+    got = Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, info)
+    torch.testing.assert_close(got, reference[-context_frames:])
+    assert info["codes"]["ref"] is reference
+    assert Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, {}) is None
+    assert (
+        Qwen3TTSTalkerForConditionalGeneration.get_stream_ref_context(model, {"codes": {"ref": torch.empty(0)}}) is None
+    )
+
+
 @pytest.mark.parametrize("streaming", [False, True])
 def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch, streaming):
     monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
@@ -340,6 +569,136 @@ def test_model_owned_audio_finalizer_runs_after_copy_without_hidden(monkeypatch,
     assert (result.pooler_output is None) == streaming
 
 
+def test_stream_audio_history_replays_exact_inputs_after_batch_reordering(mocker):
+    from vllm.v1.worker.gpu.input_batch import InputBatch
+
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+
+    owner = mocker.Mock()
+    owner.vllm_config.model_config.max_model_len = 12
+    eager = EagerMTPState(owner)
+    first = torch.arange(9, dtype=torch.float32).reshape(3, 3)
+    second = torch.arange(9, 18, dtype=torch.float32).reshape(3, 3)
+    expected = {"a": torch.cat((first[:2], second[2:])), "b": torch.cat((first[2:], second[:2]))}
+    batch = InputBatch.__new__(InputBatch)
+    batch.req_ids = ["a", "b"]
+    batch.query_start_loc_np = np.array([0, 2, 3])
+    batch.num_computed_tokens_np = np.array([0, 0])
+    eager.record_inputs(batch, first)
+    batch.req_ids = ["b", "a"]
+    batch.num_computed_tokens_np = np.array([1, 2])
+    eager.record_inputs(batch, second)
+    # Later input-buffer reuse must not change the conditioned replay history.
+    first.fill_(-1)
+    second.fill_(-2)
+    for request_id, wanted in expected.items():
+        eager._restore_audio[request_id] = []
+        replayed = torch.empty_like(wanted)
+        assert eager.replay_inputs(request_id, 0, 0, torch.zeros(3, dtype=torch.long), replayed)
+        torch.testing.assert_close(replayed, wanted)
+
+
+@pytest.mark.parametrize("valid", [False, True])
+def test_stream_audio_direct_first_requires_only_real_pcm_on_terminal(valid):
+    from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
+
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    state = buffer.add("r")
+    state.accept_first_audio()
+    first = StreamingAudioOutput(
+        torch.arange(4).reshape(1, 4),
+        torch.tensor([valid]),
+        torch.tensor(24000),
+        None,
+        np.array([0, 1]),
+        [state],
+        np.array([True]),
+    )
+    payload = first.get_output()[0]
+    if valid:
+        assert payload == {FIRST_AUDIO_REQUIRED_KEY: torch.tensor(True)}
+    else:
+        assert payload is None  # EOS has no promised first PCM to wait for.
+    assert not state.active and not state.pending
+    buffer.finish({"r"})
+    replacement = buffer.add("r")
+    assert replacement.push(torch.full((4,), 99), False).tolist() == [99] * 4
+
+
+def test_stream_audio_direct_first_drops_one_frame_and_retains_partial_suffix():
+    from vllm_omni.worker_v2.streaming_audio import AudioRequest
+
+    state = AudioRequest(samples_per_frame=4, chunk_frames=3)
+    state.accept_first_audio()
+    assert state.push(torch.arange(4), False) is None
+    assert state.first_audio_required and state.emitted
+    assert state.push(torch.arange(4, 8), False) is None
+    chunk = state.push(torch.arange(8, 12), True)
+    assert chunk.tolist() == list(range(4, 12))
+    assert not state.active and not state.pending
+
+
+@pytest.mark.parametrize("consumer", ["pcm_only", "finalizer", "extra"])
+def test_stream_audio_snapshot_skips_discarded_code_partition(monkeypatch, mocker, consumer):
+    from vllm.v1.worker.gpu.input_batch import InputBatch
+
+    from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot
+    from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
+
+    monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
+    monkeypatch.setattr(
+        OmniARModelRunner,
+        "_build_async_chunk_outputs_from_mm",
+        lambda *args: pytest.fail("in-stage PCM unnecessarily partitioned code outputs"),
+    )
+    batch = InputBatch.__new__(InputBatch)
+    batch.num_reqs = 3
+    batch.query_start_loc_np = np.array([0, 1, 2, 4])
+    batch.num_scheduled_tokens = np.array([1, 1, 2])
+    batch.num_tokens_after_padding = 4
+    buffer = StreamingAudioBuffer(samples_per_frame=4, chunk_frames=3)
+    states = [buffer.add(request_id) for request_id in ["frame", "eos", "prefill"]]
+    wav = torch.arange(16, dtype=torch.float32).reshape(4, 4)
+    streaming_audio = StreamingAudioOutput(
+        wav,
+        torch.tensor([True, False, True, False]),
+        torch.tensor(24000),
+        None,
+        batch.query_start_loc_np.copy(),
+        states,
+        np.array([False, False, True]),
+    )
+    copy_mm = mocker.spy(omni_ar_model_runner, "_async_copy_mm")
+
+    def finalize(payload, counts):
+        assert payload["codes"]["audio"].shape == (4, 2)
+        return RequestOutputSnapshot([None] * batch.num_reqs)
+
+    finalizer = mocker.Mock(side_effect=finalize) if consumer == "finalizer" else None
+    extra = ({"extra": torch.tensor([7])}, _FakeEvent()) if consumer == "extra" else None
+    result = _async_output(
+        req_ids=["frame", "eos", "prefill"],
+        sampler_output=SamplerOutput(torch.tensor([[1], [2], [3]]), None, None, None),
+        num_sampled_tokens=torch.tensor([1, 1, 1]),
+        multimodal_outputs={"codes": {"audio": torch.ones(4, 2)}},
+        input_batch=batch,
+        async_chunk=True,
+        streaming_audio=streaming_audio,
+        finalize_multimodal=finalizer,
+        extra_multimodal_outputs=extra,
+    ).get_output()
+    assert copy_mm.call_count == {"pcm_only": 0, "finalizer": 1, "extra": 2}[consumer]
+    if finalizer is not None:
+        finalizer.assert_called_once()
+    wav.fill_(-1)
+    assert result.inter_stage_outputs is None
+    assert result.multimodal_outputs[0]["model_outputs"].tolist() == [0, 1, 2, 3]
+    assert result.multimodal_outputs[1] == {}
+    assert result.multimodal_outputs[2]["model_outputs"].tolist() == [8, 9, 10, 11]
+    assert not states[1].active and not states[2].active
+
+
 @pytest.mark.parametrize("streaming", [False, True])
 def test_request_owned_snapshot_skips_generic_partition(monkeypatch, streaming):
     from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot

tests/worker_v2/test_omni_model_state.py

@@ -36,6 +36,7 @@ class _DummyInputBatch:
 def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, have_multimodal_outputs=False):
     state = object.__new__(OmniModelState)
     model = MagicMock()
+    model.stream_decoder = None
     model.has_preprocess = has_preprocess
     model.has_postprocess = has_postprocess
     model.have_multimodal_outputs = have_multimodal_outputs
@@ -62,6 +63,10 @@ def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, hav
 
     state.intermediate_buffer = OmniIntermediateBuffer(max_num_reqs)
     state._static_inputs_embeds = None
+    from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
+
+    state._eager_state = EagerMTPState(state)
+    state._stream_pos = {}
     state._mtp_generators = {}
     state._mtp_runner = None
     for name in ("_mtp_input_ids", "_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step", "_mtp_offsets"):
@@ -264,6 +269,7 @@ def test_seed_independence_resolve_once_and_sampling_kwargs():
     # vLLM sampling seed must not produce a talker generator.
     cpu = torch.device("cpu")
     assert state._get_mtp_generator("r1", SimpleNamespace(extra_args={}, seed=42), cpu) is None
+    state._mtp_generators.clear()  # The following rows are new requests.
     # Same model-local seed reproduces identical uniforms regardless of batch makeup.
     state._mtp_sample_uniforms = torch.empty((2, 2, 4))
     assert torch.equal(

vllm_omni/config/model.py

@@ -5,6 +5,7 @@ from dataclasses import MISSING, field
 from typing import Any
 
 from pydantic import ConfigDict, TypeAdapter
+from transformers import PretrainedConfig
 from vllm.config import ModelConfig
 from vllm.config.utils import config
 from vllm.logger import init_logger
@@ -124,9 +125,12 @@ class OmniModelConfig(ModelConfig):
     """
 
     stage_id: int = 0
+    # Match the upstream derived field while its value is replaced per stage.
+    hf_text_config: PretrainedConfig = field(init=False)
     async_chunk: bool = False
     session_mode: str = "turn"
     retains_state_across_chunks: bool = False
+    supports_running_prefix_cache_reset: bool = True
     use_v2_model_runner: bool = False
     supports_native_mrv2_data_plane: bool = False
     # Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin.
@@ -218,7 +222,7 @@ class OmniModelConfig(ModelConfig):
             return convertor.convert()
         return super().get_model_arch_config()
 
-    def draw_hf_text_config(self):
+    def draw_hf_text_config(self) -> PretrainedConfig:
         # transformers' get_text_config method is used to get the text config from thinker_config.
         # to handle the case that each model stage has their own text config,
         # we need to draw the text config from the corresponding model stage.

vllm_omni/config/omni_config.py

@@ -482,6 +482,7 @@ class OmniStageModelConfig(_TrackExplicitConfigFields):
     interleave_mm_strings: bool | None = None
     media_io_kwargs: dict[str, Any] | None = None
     final_output: bool = False
+    supports_running_prefix_cache_reset: bool = True
     active_stream_window: int = Field(default=0, ge=0)
     session_mode: str = "turn"
     duplex_max_sessions: int = Field(default=1, ge=1)
@@ -1325,6 +1326,7 @@ _DIFFUSION_STAGE_METADATA_FIELDS = frozenset(
         "model_arch",
         "model_stage",
         "retains_state_across_chunks",
+        "supports_running_prefix_cache_reset",
         "scheduler_cls",
         "stage_connector_spec",
         "worker_type",
@@ -1972,6 +1974,8 @@ def _build_model_config(
     if "active_stream_window" not in kwargs:
         kwargs["active_stream_window"] = _copy_value(deploy.active_stream_window)
     kwargs["final_output"] = topology.final_output
+    if not topology.supports_running_prefix_cache_reset:
+        kwargs["supports_running_prefix_cache_reset"] = False
     if "custom_voice_dir" not in kwargs and deploy.custom_voice_dir is not None:
         kwargs["custom_voice_dir"] = _copy_value(deploy.custom_voice_dir)
     stage_runner = resolve_stage_model_runner(deploy, stage_deploy)

vllm_omni/config/pipeline_registry.py

@@ -117,7 +117,7 @@ from vllm_omni.model_executor.models.qwen3_omni.pipeline import (
     QWEN3_OMNI_THINKER_ONLY_PIPELINE,
     resolve_qwen3_omni_pipeline,
 )
-from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_PIPELINE
+from vllm_omni.model_executor.models.qwen3_tts.pipeline import QWEN3_TTS_FUSED_PIPELINE, QWEN3_TTS_PIPELINE
 from vllm_omni.model_executor.models.step_audio2.pipeline import (
     STEP_AUDIO2_ASR_PIPELINE,
     STEP_AUDIO2_PIPELINE,
@@ -204,6 +204,7 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "qwen3_omni_moe": resolve_qwen3_omni_pipeline,
     "qwen3_omni_moe_thinker_only": QWEN3_OMNI_THINKER_ONLY_PIPELINE,
     "qwen3_tts": QWEN3_TTS_PIPELINE,
+    "qwen3_tts_fused": QWEN3_TTS_FUSED_PIPELINE,
     "step_audio_2": STEP_AUDIO2_PIPELINE,
     "step_audio_2_asr": STEP_AUDIO2_ASR_PIPELINE,
     "voxcpm2": VOXCPM2_PIPELINE,

vllm_omni/config/stage_config.py

@@ -245,6 +245,8 @@ class StagePipelineConfig:
     # The model keeps per-request execution state while awaiting the next
     # async chunk, so the parked request continues to consume model capacity.
     retains_state_across_chunks: bool = False
+    # Some stateful audio stages cannot roll back already consumed frames.
+    supports_running_prefix_cache_reset: bool = True
     sampling_constraints: dict[str, Any] = field(default_factory=dict)
     custom_process_input_func: str | None = None
     custom_process_next_stage_input_func: str | None = None
@@ -295,6 +297,9 @@ class PipelineConfig:
     model_type: str
     model_arch: str = ""
     stages: tuple[StagePipelineConfig, ...] = ()
+    # A single stage that streams its own final output in async-chunk mode
+    # (e.g. a Talker decoding audio in-stage) keeps deploy.async_chunk.
+    single_stage_async_chunk: bool = False
     # HF architecture aliases: used by StageConfigFactory when the model's
     # HF config reports a generic model_type that collides with a different
     # model (e.g. MiMo Audio reports model_type="qwen2"). The factory
@@ -1107,6 +1112,8 @@ def _build_engine_args(
     if ps.omni_kv_config:
         engine_args["omni_kv_config"] = dict(ps.omni_kv_config)
     engine_args["requires_full_payload_input"] = ps.requires_full_payload_input
+    if not ps.supports_running_prefix_cache_reset:
+        engine_args["supports_running_prefix_cache_reset"] = False
     return engine_args
 
 
@@ -1172,8 +1179,11 @@ def _resolve_pipeline_async_chunk_enabled(
     and False otherwise. If the user tried to enable async chunk through the deploy
     config, but it's inapplicable or unsupported, it will be disabled with a warning.
     """
-    # Single stage should never use async chunk
     if len(pipeline.stages) <= 1:
+        if pipeline.single_stage_async_chunk:
+            if deploy.async_chunk is False:
+                raise ValueError(f"Pipeline {pipeline.model_type!r} requires async_chunk=True")
+            return True
         if deploy.async_chunk:
             logger.warning(
                 "Deploy config set async_chunk=True, but async chunk is inapplicable "
@@ -1350,6 +1360,8 @@ class StageConfig:
 
         # Terminal-stage ownership comes from topology, not engine overrides.
         engine_args["final_output"] = self.final_output
+        if self.yaml_engine_args.get("supports_running_prefix_cache_reset") is False:
+            engine_args["supports_running_prefix_cache_reset"] = False
 
         # Build runtime config from YAML defaults + CLI overrides
         runtime: dict[str, Any] = dict(self.yaml_runtime)

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -101,6 +101,23 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
     max_num_running_reqs: int
 
+    def reset_prefix_cache(self, reset_running_requests: bool = False, reset_connector: bool = False) -> bool:
+        model_config = self.vllm_config.model_config
+        if (
+            reset_running_requests
+            and self.running
+            and not getattr(model_config, "supports_running_prefix_cache_reset", True)
+        ):
+            # Reset preempts and resumes in the same step while discarding
+            # in-flight tokens. The stateful codec and queued PCM have already
+            # consumed those frames and cannot roll back to that boundary.
+            logger.warning(
+                "This stage cannot reset running requests; wait for "
+                "completion or abort them before resetting the prefix cache."
+            )
+            return False
+        return super().reset_prefix_cache(reset_running_requests, reset_connector)
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         # Track requests that need KV cache transfer when finished
@@ -517,7 +534,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # domains must consume the old frame before new output passes.
                 request.async_tokens_to_discard = max(0, stale_async_tokens - len(generated_token_ids))
 
-            if output_is_stale or async_output_is_stale:
+            if async_output_is_stale or (output_is_stale and request.drop_stale_output):
                 # Output of a step scheduled before the request's in-flight
                 # tokens were discarded (segment stop / session replacement).
                 # num_computed_tokens was rolled back at the discard site, so
@@ -556,12 +573,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # tokens and rejections. If some tokens are rejected,
                 # num_computed_tokens is decreased by the number of rejected
                 # tokens.
-                if request.num_computed_tokens > 0:
-                    request.num_computed_tokens -= num_rejected
-                # If async scheduling, num_output_placeholders also includes
-                # the scheduled spec tokens count and so is similarly adjusted.
-                if request.num_output_placeholders > 0:
-                    request.num_output_placeholders -= num_rejected
+                if not output_is_stale:
+                    if request.num_computed_tokens > 0:
+                        request.num_computed_tokens -= num_rejected
+                    # If async scheduling, num_output_placeholders also includes
+                    # the scheduled spec tokens count and so is similarly adjusted.
+                    if request.num_output_placeholders > 0:
+                        request.num_output_placeholders -= num_rejected
                 spec_decoding_stats = self.make_spec_decoding_stats(
                     spec_decoding_stats,
                     num_draft_tokens=num_draft_tokens,
@@ -600,7 +618,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             # Check for stop and update request status.
             if new_token_ids:
                 num_sampled_tokens = len(new_token_ids)
-                new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
+                if output_is_stale:
+                    new_token_ids, stopped = self._update_request_with_output(request, new_token_ids, is_stale=True)
+                else:
+                    new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
                 if new_logprobs is not None and len(new_token_ids) < num_sampled_tokens:
                     # A mid-step stop (e.g. spec-decode tokens sampled past
                     # EOS) trims new_token_ids after the validation slice
@@ -704,6 +725,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     # a queued streaming update, and adding twice swallows the
                     # next duplex unit's listen/speak under async scheduling.
                     if request.num_in_flight_tokens > 0:
+                        request.drop_stale_output = True
                         request.num_stale_output_tokens = request.num_in_flight_tokens
                     if outstanding_async_tokens > 0:
                         # Discard only outputs that are already in flight and
@@ -884,6 +906,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # a resumable stop applies a queued update through this helper.
         in_flight_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0)
         if in_flight_tokens > 0:
+            session.drop_stale_output = True
             session.num_stale_output_tokens = in_flight_tokens
         if outstanding_async_tokens > 0:
             # Async scheduling may already have sampled the previous
@@ -1405,6 +1428,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if status is not None:
                 request.status = status
         if native_transfer and connector_delay_free_blocks:
+            assert transfer_params is not None
             kv_xfer_params = {
                 **(kv_xfer_params or {}),
                 "transfer_id": transfer_params["transfer_id"],

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -222,6 +222,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
         # num_in_flight_tokens already includes any undrained stale share.
         # Assign instead of accumulating so callers that fenced the same
         # rollover before entering this helper do not count it twice.
+        session.drop_stale_output = True
         session.num_stale_output_tokens = int(getattr(session, "num_in_flight_tokens", 0) or 0)
         session.num_output_placeholders = 0
         session.spec_token_ids = []
@@ -316,6 +317,7 @@ class OmniSchedulerMixin(_SchedulerMixinBase):
             # The streaming update may already have fenced this same in-flight
             # frame. Seed idempotently so the replacement does not count it twice.
             request.num_stale_output_tokens = int(getattr(request, "num_in_flight_tokens", 0) or 0)
+            request.drop_stale_output = True
             request.num_output_placeholders = 0
             request.spec_token_ids = []
             self._release_replaced_streaming_prompt_cache(request)

vllm_omni/deploy/qwen3_tts_fused_single_gpu.yaml

@@ -0,0 +1,54 @@
+# Single-stage Qwen3-TTS: the Talker decodes every frame with the stateful
+# streaming codec decoder and returns PCM itself (no Code2Wav stage).
+# reset_prefix_cache(reset_running_requests=True) returns False while requests
+# are running; wait for completion or abort them before resetting the cache.
+pipeline: qwen3_tts_fused
+model_runner: v2
+async_chunk: true
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 128
+    # Exact conditioned Talker inputs for preemption replay are allocated
+    # after KV profiling: 16 MiB per active request at 1.7B/max_model_len=4096
+    # (~2 GiB at 128 requests), in addition to codec state and output buffers.
+    # Keep headroom when increasing gpu_memory_utilization or capacity.
+    gpu_memory_utilization: 0.5
+    trust_remote_code: true
+    enable_prefix_caching: false
+    # Model options that multi-stage deployments put on the connector: a
+    # single stage has no connector edge.
+    additional_config:
+      talker_stream_decode: true
+      # Opt in to first-frame delivery on its own queue for lower first-audio
+      # latency; leave disabled when prioritizing peak throughput.
+      talker_stream_first_audio: false
+      # Match the standard two-stage profile's voice-clone context boundary.
+      # A shorter context (e.g. 25) can change timbre continuity.
+      ref_code_context_frames: 72
+      code_predictor_kv_cache: true
+      code_predictor_fused_sampling: true
+      code_predictor_fused: true
+    # Talker decode contexts are short (~100-700 tokens): FA3 split-KV plus
+    # its combine kernel costs more than one pass per head.
+    attention_config:
+      flash_attn_max_num_splits_for_cuda_graph: 1
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    max_model_len: 4096
+    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
+    default_sampling_params:
+      temperature: 0.9
+      top_k: 50
+      max_tokens: 4096
+      min_tokens: 2
+      repetition_penalty: 1.05
+    subtalker_sampling_params:
+      do_sample: true
+      temperature: 0.9
+      top_k: 50
+      top_p: 1.0

vllm_omni/engine/arg_utils.py

@@ -198,6 +198,7 @@ class OmniEngineArgs(EngineArgs):
     async_chunk: bool = False
     session_mode: str = "turn"
     retains_state_across_chunks: bool = False
+    supports_running_prefix_cache_reset: bool = True
     use_v2_model_runner: bool = False
     supports_native_mrv2_data_plane: bool = False
     # WS-A: Stage-1 active stream slots. 0 = legacy preempt-everything.
@@ -430,6 +431,7 @@ class OmniEngineArgs(EngineArgs):
             async_chunk=self.async_chunk,
             session_mode=self.session_mode,
             retains_state_across_chunks=self.retains_state_across_chunks,
+            supports_running_prefix_cache_reset=self.supports_running_prefix_cache_reset,
             use_v2_model_runner=self.use_v2_model_runner,
             supports_native_mrv2_data_plane=self.supports_native_mrv2_data_plane,
             active_stream_window=self.active_stream_window,

vllm_omni/engine/orchestrator.py

@@ -898,7 +898,7 @@ class OrchestratorBase:
                     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
+                    or (stage_id if final_output else stage_id + 1) > req_state.final_stage_id
                 ):
                     continue
                 audio = mm.get("model_outputs")
@@ -935,7 +935,7 @@ class OrchestratorBase:
         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
+        codec_stage = source_stage if self.stage_pools[source_stage].final_output else 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:

vllm_omni/engine/stage_engine_core_proc.py

@@ -87,7 +87,10 @@ def _bind_first_audio_sink(model_executor: Any, output_queue: Any, scheduler: An
     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"):
+    decodes_audio = getattr(model, "first_frame_decoder", None) is not None or (
+        getattr(model, "stream_decoder", None) is not None and bool(getattr(model, "stream_first_audio", False))
+    )
+    if not decodes_audio or not hasattr(model_state, "set_first_audio_sink"):
         return False
     from vllm_omni.worker_v2.first_audio_sender import engine_output_queue_sink
 

vllm_omni/entrypoints/openai/api_server.py

@@ -249,6 +249,33 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
     await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
 
 
+def _own_reuseport_socket(sock: socket.socket, client_count: int) -> socket.socket:
+    """Give this API server its own SO_REUSEPORT listener on the shared address.
+
+    All API servers accepting from the one inherited socket lets whichever
+    wakes first take a whole burst of connections (the event loop accepts
+    in a loop), and keep-alive then pins that imbalance for the connections'
+    lifetime. Separate reuseport listeners let the kernel spread connections
+    by hash. The inherited socket is bound but never listened on, so it
+    receives nothing.
+    """
+    if client_count <= 1:
+        return sock
+    if sock.family not in (socket.AF_INET, socket.AF_INET6):
+        return sock
+    try:
+        if not sock.getsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT):
+            return sock
+        own = socket.socket(family=sock.family, type=socket.SOCK_STREAM)
+        own.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        own.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEPORT, 1)
+        own.bind(sock.getsockname()[:2])
+    except OSError as e:
+        logger.warning("Per-server reuseport listener unavailable (%s); sharing the inherited socket", e)
+        return sock
+    return own
+
+
 def run_omni_api_server_worker_proc(
     listen_address: str,
     sock: socket.socket,
@@ -278,6 +305,7 @@ def run_omni_api_server_worker_proc(
 
     set_process_title("APIServer", str(client_index))
     decorate_logs("APIServer", skip_if_decorated=True)
+    sock = _own_reuseport_socket(sock, int(omni_client_config["client_count"]))
     uvloop.run(
         omni_run_server_worker(
             listen_address,
@@ -512,7 +540,7 @@ async def build_async_omni(
     # Ensures everything is shutdown and cleaned up on error/exit
     async with build_async_omni_from_stage_config(
         args,
-        disable_frontend_multiprocessing=disable_frontend_multiprocessing,
+        disable_frontend_multiprocessing=bool(disable_frontend_multiprocessing),
         client_config=client_config,
     ) as async_omni:
         yield async_omni

vllm_omni/entrypoints/openai/audio_utils_mixin.py

@@ -46,6 +46,15 @@ except ImportError:
 logger = init_logger(__name__)
 
 
+def _float32_to_pcm16_bytes(audio: np.ndarray) -> bytes:
+    """PCM_16 bytes of float32 samples, bit-identical to libsndfile's RAW writer.
+
+    libsndfile (1.2) converts with ``floor(x * 0x8000)`` clipped to int16.
+    """
+    scaled = np.floor(audio.astype(np.float64) * 32768.0)
+    return np.clip(scaled, -32768, 32767).astype("<i2").tobytes()
+
+
 class AudioMixin:
     """Mixin class to add audio-related utilities."""
 
@@ -93,9 +102,15 @@ class AudioMixin:
 
         soundfile_format, media_type, kwargs = supported_formats[response_format]
 
-        with BytesIO() as buffer:
-            soundfile.write(buffer, audio_tensor, sample_rate, format=soundfile_format, **kwargs)
-            audio_data = buffer.getvalue()
+        audio_data: bytes | str
+        if response_format == "pcm" and isinstance(audio_tensor, np.ndarray) and audio_tensor.dtype == np.float32:
+            # Through BytesIO, soundfile makes a Python callback per block;
+            # streamed chunks made that the API workers' largest cost.
+            audio_data = _float32_to_pcm16_bytes(audio_tensor)
+        else:
+            with BytesIO() as buffer:
+                soundfile.write(buffer, audio_tensor, sample_rate, format=soundfile_format, **kwargs)
+                audio_data = buffer.getvalue()
 
         if base64_encode:
             import base64

vllm_omni/entrypoints/openai/tts_adapters/qwen3_tts.py

@@ -397,34 +397,60 @@ class Qwen3TTSAdapter(ARTTSAdapter):
 
         return params
 
+    def _load_text_tokenizer(self) -> None:
+        server = self.ctx.server
+        if server._tts_tokenizer is not None:
+            return
+        import transformers
+        from transformers import AutoTokenizer
+
+        # Same construction as the engine's text tokenizer
+        # (Qwen3TTSPromptEmbedsBuilder.get_text_tokenizer), so the ids computed
+        # here can be handed to the engine.
+        kwargs: dict[str, Any] = dict(trust_remote_code=True, use_fast=True, padding_side="left")
+        if transformers.__version__ < "5":
+            kwargs["fix_mistral_regex"] = True
+        assert self.ctx.engine_client is not None
+        model_name = self.ctx.engine_client.model_config.model
+        server._tts_tokenizer = AutoTokenizer.from_pretrained(model_name, **kwargs)
+
     def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int:
         """Estimate prompt length so the placeholder matches model-side embeddings."""
         try:
             from vllm_omni.model_executor.models.qwen3_tts.prompt_embeds_builder import (
+                PRECOMPUTED_TEXT_IDS_KEY,
                 Qwen3TTSPromptEmbedsBuilder,
+                build_assistant_text,
             )
 
             server = self.ctx.server
-            if server._tts_tokenizer is None:
-                from transformers import AutoTokenizer
-
-                model_name = self.ctx.engine_client.model_config.model
-                server._tts_tokenizer = AutoTokenizer.from_pretrained(
-                    model_name,
-                    trust_remote_code=True,
-                    padding_side="left",
-                )
+            self._load_text_tokenizer()
+            assert self.ctx.engine_client is not None
             hf_config = self.ctx.engine_client.model_config.hf_config
             talker_config = hf_config.talker_config
             task_type = (tts_params.get("task_type") or ["CustomVoice"])[0]
-            return Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
+            text = (tts_params.get("text") or [""])[0]
+            assistant_text = build_assistant_text(text) if isinstance(text, str) else None
+            assistant_ids: list[list[int]] = []
+
+            def tokenize_prompt(t: str) -> list[int]:
+                ids = server._tts_tokenizer(t, padding=False)["input_ids"]
+                if t == assistant_text:
+                    assistant_ids.append(ids)
+                return ids
+
+            prompt_len = Qwen3TTSPromptEmbedsBuilder.estimate_prompt_len_from_additional_information(
                 additional_information=tts_params,
                 task_type=task_type,
-                tokenize_prompt=lambda t: server._tts_tokenizer(t, padding=False)["input_ids"],
+                tokenize_prompt=tokenize_prompt,
                 codec_language_id=getattr(talker_config, "codec_language_id", None),
                 spk_is_dialect=getattr(talker_config, "spk_is_dialect", None),
                 estimate_ref_code_len=self._estimate_ref_code_len,
             )
+            if assistant_ids:
+                # Spares the engine tokenizing the same text on its step thread.
+                tts_params[PRECOMPUTED_TEXT_IDS_KEY] = [assistant_ids[0]]
+            return prompt_len
         except Exception as e:
             logger.warning("Failed to estimate TTS prompt length, using fallback 2048: %s", e)
             return 2048

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -954,7 +954,13 @@ class CodePredictorWrapper(nn.Module):
             extra_cfg = connector_cfg.get("extra", connector_cfg)
         else:
             extra_cfg = getattr(connector_cfg, "extra", None)
-        return extra_cfg if isinstance(extra_cfg, dict) else {}
+        extra_cfg = extra_cfg if isinstance(extra_cfg, dict) else {}
+        # A single-stage deployment has no connector edge to carry model
+        # options; it sets them in the stage's ``additional_config``.
+        additional = getattr(vllm_config, "additional_config", None)
+        if isinstance(additional, dict) and additional:
+            return {**additional, **extra_cfg}
+        return extra_cfg
 
     @staticmethod
     def _parse_bool_config(value: object) -> bool:
@@ -1332,6 +1338,21 @@ class CodePredictorWrapper(nn.Module):
             scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
         return self._sample_codes_gumbel(scaled, generator=generator, uniforms=uniforms)
 
+    def _validate_sampling_inputs(
+        self,
+        bsz: int,
+        generators: Sequence[torch.Generator | None] | None,
+        sample_uniforms: torch.Tensor | None,
+    ) -> None:
+        if generators is not None and len(generators) != bsz:
+            raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}")
+        if sample_uniforms is not None:
+            expected_shape = (bsz, self._num_groups - 1, int(self.config.vocab_size))
+            if tuple(sample_uniforms.shape) != expected_shape:
+                raise ValueError(
+                    f"sample_uniforms must have shape {expected_shape}, got {tuple(sample_uniforms.shape)}"
+                )
+
     @torch.inference_mode()
     def forward(
         self,
@@ -1349,14 +1370,7 @@ class CodePredictorWrapper(nn.Module):
         """Predict residual codebooks 1..G-1 autoregressively via re-prefill."""
         bsz = int(layer0_code.shape[0])
         num_groups = self._num_groups
-        if generators is not None and len(generators) != bsz:
-            raise ValueError(f"generators must have one entry per row: got {len(generators)} for batch {bsz}")
-        if sample_uniforms is not None:
-            expected_shape = (bsz, num_groups - 1, int(self.config.vocab_size))
-            if tuple(sample_uniforms.shape) != expected_shape:
-                raise ValueError(
-                    f"sample_uniforms must have shape {expected_shape}, got {tuple(sample_uniforms.shape)}"
-                )
+        self._validate_sampling_inputs(bsz, generators, sample_uniforms)
         sample_generator: _GeneratorLike = generators if generators is not None else generator
         device = layer0_code.device
 

vllm_omni/model_executor/models/qwen3_tts/fused_code_predictor.py

@@ -0,0 +1,312 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# ruff: noqa: N803 - Triton constexpr parameters use kernel-style capitals.
+"""Fused residual-codebook predictor for Qwen3-TTS (CUDA, BF16).
+
+Same computation as the frame-local-KV predictor with per-call top-k Gumbel
+sampling, restructured to issue few kernels per residual step, since at serving
+batch sizes each of the 15 dependent steps is launch- and latency-bound:
+
+- Per layer: RMSNorm, one QKV GEMM, one attention kernel (q/k RMSNorm, RoPE,
+  this call's K/V write and attention over the frame's keys), O projection
+  accumulated into the residual stream (``addmm_``), RMSNorm, gate/up GEMM,
+  SiLU-mul, down projection accumulated into the stream.
+- The embedding of each sampled code followed by ``small_to_mtp_projection``
+  is linear, so it folds at load into one ``[vocab, hidden]`` table per
+  codebook; the sampling kernel writes the next step's input row directly.
+- Sampling finds the top-k threshold by bisection over the BF16 order keys
+  (exact k-th largest, ties kept) instead of a sort.
+"""
+
+from __future__ import annotations
+
+import torch
+from vllm.triton_utils import tl, tldevice, triton
+
+from .tokenizer_12hz.streaming_decoder import _silu_mul_kernel
+
+
+@triton.jit
+def _rms_rows_kernel(x_ptr, x_stride, w_ptr, out_ptr, eps, N: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.arange(0, N)
+    x = tl.load(x_ptr + row * x_stride + offs).to(tl.float32)
+    var = tl.sum(x * x, axis=0) / N
+    dt = out_ptr.dtype.element_ty
+    y = (x * tl.rsqrt(var + eps)).to(dt).to(tl.float32)
+    w = tl.load(w_ptr + offs).to(tl.float32)
+    tl.store(out_ptr + row * N + offs, (w * y).to(dt))
+
+
+@triton.jit
+def _bf(x):
+    return x.to(tl.bfloat16).to(tl.float32)
+
+
+@triton.jit
+def _head_norm_rope(x1, x2, w1, w2, c1, c2, s1, s2, eps, HD: tl.constexpr):
+    """HF RMSNorm (``weight * normed.to(bf16)``) then RoPE in BF16 elementwise steps, on the two halves."""
+    var = (tl.sum(x1 * x1, axis=0) + tl.sum(x2 * x2, axis=0)) / HD
+    r = tl.rsqrt(var + eps)
+    n1 = _bf(w1 * _bf(x1 * r))
+    n2 = _bf(w2 * _bf(x2 * r))
+    # rotate_half(n) = [-n2, n1]
+    o1 = _bf(_bf(n1 * c1) + _bf(-n2 * s1))
+    o2 = _bf(_bf(n2 * c2) + _bf(n1 * s2))
+    return o1, o2
+
+
+@triton.jit
+def _cp_attention_kernel(
+    qkv_ptr,  # [B * NQ, (H + 2 * KVH) * HD]
+    kc_ptr,  # [Bmax, KVH, MAXP, HD]
+    vc_ptr,
+    qw_ptr,
+    kw_ptr,
+    cos_ptr,  # [MAXP, HD] bf16
+    sin_ptr,
+    out_ptr,  # [B * NQ, H * HD]
+    first,
+    eps,
+    scale,
+    NQ: tl.constexpr,
+    H: tl.constexpr,
+    KVH: tl.constexpr,
+    HD: tl.constexpr,
+    MAXP: tl.constexpr,
+    BK: tl.constexpr,
+):
+    b = tl.program_id(0)
+    h = tl.program_id(1)
+    qi = tl.program_id(2)
+    kvh = h // (H // KVH)
+    W: tl.constexpr = (H + 2 * KVH) * HD
+    HALF: tl.constexpr = HD // 2
+    d = tl.arange(0, HALF)
+    p = first + qi
+    row = (b * NQ + qi).to(tl.int64)
+    qw1 = tl.load(qw_ptr + d).to(tl.float32)
+    qw2 = tl.load(qw_ptr + HALF + d).to(tl.float32)
+    kw1 = tl.load(kw_ptr + d).to(tl.float32)
+    kw2 = tl.load(kw_ptr + HALF + d).to(tl.float32)
+    qb = qkv_ptr + row * W + h * HD
+    q1, q2 = _head_norm_rope(
+        tl.load(qb + d).to(tl.float32), tl.load(qb + HALF + d).to(tl.float32), qw1, qw2,
+        tl.load(cos_ptr + p * HD + d).to(tl.float32), tl.load(cos_ptr + p * HD + HALF + d).to(tl.float32),
+        tl.load(sin_ptr + p * HD + d).to(tl.float32), tl.load(sin_ptr + p * HD + HALF + d).to(tl.float32),
+        eps, HD,
+    )  # fmt: skip
+    j = tl.arange(0, BK)
+    cache = ((b * KVH + kvh) * MAXP).to(tl.int64) * HD
+    old = (j < first)[:, None]
+    kp = kc_ptr + cache + j[:, None] * HD + d[None, :]
+    vp = vc_ptr + cache + j[:, None] * HD + d[None, :]
+    k1 = tl.load(kp, mask=old, other=0.0).to(tl.float32)
+    k2 = tl.load(kp + HALF, mask=old, other=0.0).to(tl.float32)
+    v1 = tl.load(vp, mask=old, other=0.0).to(tl.float32)
+    v2 = tl.load(vp + HALF, mask=old, other=0.0).to(tl.float32)
+    for n in tl.static_range(NQ):
+        # this call's position first + n, computed from its own QKV row
+        on = first + n
+        kb = qkv_ptr + (b * NQ + n).to(tl.int64) * W + (H + kvh) * HD
+        vb = qkv_ptr + (b * NQ + n).to(tl.int64) * W + (H + KVH + kvh) * HD
+        kn1, kn2 = _head_norm_rope(
+            tl.load(kb + d).to(tl.float32), tl.load(kb + HALF + d).to(tl.float32), kw1, kw2,
+            tl.load(cos_ptr + on * HD + d).to(tl.float32), tl.load(cos_ptr + on * HD + HALF + d).to(tl.float32),
+            tl.load(sin_ptr + on * HD + d).to(tl.float32), tl.load(sin_ptr + on * HD + HALF + d).to(tl.float32),
+            eps, HD,
+        )  # fmt: skip
+        vn1 = tl.load(vb + d).to(tl.float32)
+        vn2 = tl.load(vb + HALF + d).to(tl.float32)
+        sel = (j == on)[:, None]
+        k1 = tl.where(sel, kn1[None, :], k1)
+        k2 = tl.where(sel, kn2[None, :], k2)
+        v1 = tl.where(sel, vn1[None, :], v1)
+        v2 = tl.where(sel, vn2[None, :], v2)
+        wm = (d < HALF) & (qi == n) & (h % (H // KVH) == 0)
+        dtc = kc_ptr.dtype.element_ty
+        tl.store(kc_ptr + cache + on * HD + d, kn1.to(dtc), mask=wm)
+        tl.store(kc_ptr + cache + on * HD + HALF + d, kn2.to(dtc), mask=wm)
+        tl.store(vc_ptr + cache + on * HD + d, vn1.to(dtc), mask=wm)
+        tl.store(vc_ptr + cache + on * HD + HALF + d, vn2.to(dtc), mask=wm)
+    km = j <= p
+    s = (tl.sum(k1 * q1[None, :], axis=1) + tl.sum(k2 * q2[None, :], axis=1)) * scale
+    s = tl.where(km, s, -float("inf"))
+    e = tl.exp(s - tl.max(s, axis=0))
+    e = tl.where(km, e, 0.0)
+    w = e / tl.sum(e, axis=0)
+    ob = out_ptr + row * (H * HD) + h * HD
+    dto = out_ptr.dtype.element_ty
+    tl.store(ob + d, tl.sum(v1 * w[:, None], axis=0).to(dto))
+    tl.store(ob + HALF + d, tl.sum(v2 * w[:, None], axis=0).to(dto))
+
+
+@triton.jit
+def _cp_sample_kernel(
+    logits_ptr,  # [B, V] bf16
+    u_ptr,  # uniforms, row stride u_stride, contiguous over V
+    u_stride,
+    codes_ptr,  # [B, G] int64
+    step,
+    G: tl.constexpr,
+    table_ptr,  # [V, HID] next-step input table, or dummy
+    next_ptr,  # [B, HID]
+    inv_temperature,
+    V: tl.constexpr,
+    TOPK: tl.constexpr,
+    HID: tl.constexpr,
+    HAS_NEXT: tl.constexpr,
+):
+    row = tl.program_id(0).to(tl.int64)
+    idx = tl.arange(0, V)
+    value = tl.load(logits_ptr + row * V + idx)
+    scaled = (value.to(tl.float32) * inv_temperature).to(value.dtype)
+    if TOPK > 0:
+        # order-preserving 16-bit keys of the BF16 values
+        bits = scaled.to(tl.int16, bitcast=True).to(tl.int32) & 0xFFFF
+        # Top-k keeps all numeric ties at the cutoff, including both zeros.
+        bits = tl.where((bits & 0x7FFF) == 0, 0, bits)
+        key = tl.where(bits >= 0x8000, 0xFFFF - bits, bits | 0x8000)
+        the = 0
+        for bit in tl.static_range(15, -1, -1):
+            cand = the | (1 << bit)
+            cnt = tl.sum((key >= cand).to(tl.int32), axis=0)
+            the = tl.where(cnt >= TOPK, cand, the)
+        sf = tl.where(key >= the, scaled.to(tl.float32), -float("inf"))
+    else:
+        sf = scaled.to(tl.float32)
+    u = tl.load(u_ptr + row * u_stride + idx)
+    score = sf - tldevice.log(-tldevice.log(u))
+    max_score = tl.max(score, 0)
+    chosen = tl.min(tl.where(score == max_score, idx, V), 0)
+    first_nan = tl.min(tl.where(score != score, idx, V), 0)
+    code = tl.where(first_nan < V, first_nan, chosen)
+    tl.store(codes_ptr + row * G + step, code.to(tl.int64))
+    if HAS_NEXT:
+        c = tl.arange(0, HID)
+        tl.store(next_ptr + row * HID + c, tl.load(table_ptr + code.to(tl.int64) * HID + c))
+
+
+class FusedCodePredictor:
+    """Execution object built from a loaded ``Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM``."""
+
+    def __init__(self, predictor, max_batch: int) -> None:
+        model = predictor.model
+        weight = next(model.parameters())
+        dev, dt = weight.device, weight.dtype
+        if dt != torch.bfloat16:
+            raise ValueError("fused code predictor requires BF16 weights")
+        cfg = predictor.config
+        self.groups = int(cfg.num_code_groups)
+        self.vocab = int(cfg.vocab_size)
+        self.hidden = int(cfg.hidden_size)
+        attn0 = model.layers[0].self_attn
+        self.heads, self.kv_heads, self.head_dim = attn0.num_heads, attn0.num_kv_heads, attn0.head_dim
+        self.scale = attn0.scaling
+        self.eps = float(model.norm.variance_epsilon)
+        self.max_pos = attn0.max_seq
+        self.max_batch = max_batch
+        f32 = lambda t: t.detach().to(device=dev, dtype=torch.float32)  # noqa: E731
+        c = lambda t: t.to(dt).contiguous()  # noqa: E731
+        proj = predictor.small_to_mtp_projection
+        if isinstance(proj, torch.nn.Identity):
+            pw, pb = None, None
+        else:
+            pw, pb = f32(proj.weight), f32(proj.bias)
+        self.proj_wt = c(pw.t()) if pw is not None else None
+        self.proj_b = c(pb) if pb is not None else None
+        tables = []
+        for emb in model.codec_embedding:
+            e = f32(emb.weight)
+            tables.append(c(e @ pw.t() + pb) if pw is not None else c(e))
+        self.tables = tables
+        self.layers = []
+        for layer in model.layers:
+            a = layer.self_attn
+            self.layers.append(
+                dict(
+                    ln1=c(f32(layer.input_layernorm.weight)),
+                    ln2=c(f32(layer.post_attention_layernorm.weight)),
+                    eps1=float(layer.input_layernorm.variance_epsilon),
+                    eps2=float(layer.post_attention_layernorm.variance_epsilon),
+                    qkv=c(f32(a.qkv_proj.weight).t()),
+                    qn=c(f32(a.q_norm.weight)),
+                    kn=c(f32(a.k_norm.weight)),
+                    head_eps=float(a.q_norm.variance_epsilon),
+                    o=c(f32(a.o_proj.weight).t()),
+                    gu=c(f32(layer.mlp.gate_up_proj.weight).t()),
+                    down=c(f32(layer.mlp.down_proj.weight).t()),
+                )
+            )
+            if a.qkv_proj.bias is not None:
+                raise ValueError("fused code predictor assumes bias-free attention")
+        self.norm_w = c(f32(model.norm.weight))
+        self.lm_heads = [c(f32(h.weight).t()) for h in predictor.lm_head]
+        rot = model.rotary_emb
+        self.cos = rot.cos_cached.to(device=dev, dtype=dt).contiguous()
+        self.sin = rot.sin_cached.to(device=dev, dtype=dt).contiguous()
+        kv = (len(self.layers), 2, max_batch, self.kv_heads, self.max_pos, self.head_dim)
+        self.cache = torch.zeros(kv, device=dev, dtype=dt)
+        self._dummy = torch.zeros(1, device=dev, dtype=dt)
+
+    def _layer(self, x: torch.Tensor, li: int, B: int, nq: int, first: int) -> None:
+        L = self.layers[li]
+        rows = x.shape[0]
+        h = torch.empty_like(x)
+        _rms_rows_kernel[(rows,)](x, x.stride(0), L["ln1"], h, L["eps1"], N=self.hidden)
+        qkv = torch.mm(h, L["qkv"])
+        att = torch.empty(rows, self.heads * self.head_dim, device=x.device, dtype=x.dtype)
+        _cp_attention_kernel[(B, self.heads, nq)](
+            qkv, self.cache[li, 0], self.cache[li, 1], L["qn"], L["kn"], self.cos, self.sin, att,
+            first, L["head_eps"], self.scale,
+            NQ=nq, H=self.heads, KVH=self.kv_heads, HD=self.head_dim, MAXP=self.max_pos,
+            BK=triton.next_power_of_2(self.max_pos), num_warps=1,
+        )  # fmt: skip
+        x.addmm_(att, L["o"])
+        _rms_rows_kernel[(rows,)](x, x.stride(0), L["ln2"], h, L["eps2"], N=self.hidden)
+        gu = torch.mm(h, L["gu"])
+        inter = gu.shape[1] // 2
+        a = torch.empty(rows, inter, device=x.device, dtype=x.dtype)
+        _silu_mul_kernel[(rows, triton.cdiv(inter, 1024))](gu, a, N=inter, BN=1024)
+        x.addmm_(a, L["down"])
+
+    @torch.inference_mode()
+    def __call__(
+        self,
+        layer0_code: torch.Tensor,
+        layer0_embed: torch.Tensor,
+        last_talker_hidden: torch.Tensor,
+        inv_temperature: float,
+        top_k: int,
+        sample_uniforms: torch.Tensor,
+    ) -> torch.Tensor:
+        B = int(layer0_code.shape[0])
+        if B > self.max_batch:
+            raise ValueError("fused code predictor batch exceeds its capacity")
+        dev, dt = layer0_embed.device, self.norm_w.dtype
+        inp = torch.stack([last_talker_hidden.reshape(B, -1), layer0_embed.reshape(B, -1)], 1).to(dt)
+        inp = inp.reshape(2 * B, -1)
+        x = torch.addmm(self.proj_b, inp, self.proj_wt) if self.proj_wt is not None else inp.clone()
+        codes = torch.empty(B, self.groups, device=dev, dtype=torch.int64)
+        codes[:, 0] = layer0_code.reshape(B)
+        u = sample_uniforms
+        nq, first = 2, 0
+        for step in range(1, self.groups):
+            for li in range(len(self.layers)):
+                self._layer(x, li, B, nq, first)
+            last = x[nq - 1 :: nq] if nq > 1 else x
+            hf = torch.empty(B, self.hidden, device=dev, dtype=dt)
+            _rms_rows_kernel[(B,)](last, last.stride(0), self.norm_w, hf, self.eps, N=self.hidden)
+            logits = torch.mm(hf, self.lm_heads[step - 1])
+            has_next = step < self.groups - 1
+            nxt = torch.empty(B, self.hidden, device=dev, dtype=dt) if has_next else self._dummy
+            us = u[:, step - 1]
+            _cp_sample_kernel[(B,)](
+                logits, us, us.stride(0), codes, step, self.groups,
+                self.tables[step - 1] if has_next else self._dummy, nxt, inv_temperature,
+                V=self.vocab, TOPK=top_k, HID=self.hidden, HAS_NEXT=has_next, num_warps=4,
+            )  # fmt: skip
+            x = nxt
+            first = step + 1
+            nq = 1
+        return codes

vllm_omni/model_executor/models/qwen3_tts/pipeline.py

@@ -72,3 +72,33 @@ QWEN3_TTS_PIPELINE = PipelineConfig(
         ),
     ),
 )
+
+
+# Single-stage variant: the Talker decodes each frame with the stateful
+# streaming codec decoder and emits PCM as the final output (enable with the
+# ``talker_stream_decode`` model option). No Code2Wav stage, no connector.
+QWEN3_TTS_FUSED_PIPELINE = PipelineConfig(
+    model_type="qwen3_tts_fused",
+    default_deploy_config_name="qwen3_tts_fused_single_gpu.yaml",
+    single_stage_async_chunk=True,
+    model_arch="Qwen3TTSTalkerForConditionalGeneration",
+    stages=(
+        StagePipelineConfig(
+            stage_id=0,
+            model_stage="qwen3_tts",
+            execution_type=StageExecutionType.LLM_AR,
+            input_sources=(),
+            owns_tokenizer=True,
+            final_output=True,
+            final_output_type="audio",
+            # The API output processor tags ``model_outputs`` by this modality.
+            engine_output_type="audio",
+            supports_running_prefix_cache_reset=False,
+            supports_native_mrv2_data_plane=True,
+            sampling_constraints={
+                "detokenize": False,
+                "stop_token_ids": [2150],
+            },
+        ),
+    ),
+)

vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py

@@ -42,6 +42,7 @@ from vllm.logger import init_logger
 from vllm.multimodal.audio import AudioResampler
 
 from vllm_omni.utils.audio import mel_filter_bank
+from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking
 
 if TYPE_CHECKING:
     from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSTalkerConfig
@@ -160,7 +161,7 @@ def coerce_token_ids(value: object, *, device: torch.device) -> torch.Tensor | N
         ids = ids.unsqueeze(0)
     if ids.ndim != 2 or ids.numel() == 0:
         return None
-    return ids.to(device=device, dtype=torch.long).contiguous()
+    return to_device_nonblocking(ids.to(dtype=torch.long).contiguous(), device)
 
 
 # ---------------------------------------------------------------------------
@@ -366,6 +367,9 @@ class Qwen3TTSPromptEmbedsBuilder:
         # Reusing them avoids launching the embedding/projection pair once per
         # request during prompt construction.
         self._projected_token_cache: dict[tuple[str, tuple[int, ...]], torch.Tensor] = {}
+        # CustomVoice prompt pieces that do not depend on the text, per
+        # (device, speaker, language): bounded by the model's voices.
+        self._custom_voice_prompt_cache: dict[tuple, tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = {}
 
         self._ref_audio_artifact_cache_max_entries = int(ref_audio_artifact_cache_max_entries)
         self._ref_audio_artifact_cache: OrderedDict[str, dict[str, torch.Tensor | bool]] = OrderedDict()
@@ -1024,7 +1028,9 @@ class Qwen3TTSPromptEmbedsBuilder:
         dev = self._device()
         input_ids = coerce_token_ids(info_dict.pop(PRECOMPUTED_TEXT_IDS_KEY, None), device=dev)
         if input_ids is None:
-            input_ids = tok(build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"].to(device=dev)
+            input_ids = to_device_nonblocking(
+                tok(build_assistant_text(text), return_tensors="pt", padding=False)["input_ids"], dev
+            )
 
         # Optional instruct prefix.
         instruct = (info_dict.get("instruct") or [""])[0]
@@ -1411,35 +1417,44 @@ class Qwen3TTSPromptEmbedsBuilder:
                 raise ValueError(f"Unsupported speaker: {speaker}")
             spk_id = spk_id_map[speaker]
             # Keep it at least 1D; embedding on a 0-d tensor can return 1D.
-            spk_tensor = torch.tensor([spk_id], device=input_ids.device, dtype=torch.long)
+            spk_tensor = index_to_device([spk_id], input_ids.device)
             spk_embed = codec_embed(spk_tensor)
             if spk_embed.ndim in (1, 2):
                 spk_embed = spk_embed.view(1, 1, -1)
             speaker_embed = spk_embed
             codec_input = torch.cat([codec_input_0, speaker_embed, codec_input_1], dim=1)
 
-            role_embed = text_projection(text_embedding(input_ids[:, :3]))
-            codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1)
-            codec_prefix = codec_prefix + codec_input[:, :-1]
-            talker_prompt = torch.cat((role_embed, codec_prefix), dim=1)
+            # The role tokens, speaker/language codec prefix, codec pad row and
+            # closing row do not depend on the text: cache them per speaker and
+            # language so a request only embeds and projects its own text.
+            # A text starting with whitespace could merge with the template's
+            # newline token, so it keeps the per-request role embedding.
+            cache_key = None
+            if non_streaming_mode and instruct_embed is None and not text[:1].isspace():
+                cache_key = (
+                    str(input_ids.device),
+                    int(spk_id),
+                    codec_prefill_list[0][2] if len(codec_prefill_list[0]) == 4 else None,
+                )
+            cv_cache = self._custom_voice_prompt_cache
+            cached = cv_cache.get(cache_key) if cache_key is not None else None
+            if cached is None:
+                role_embed = text_projection(text_embedding(input_ids[:, :3]))
+                codec_prefix = torch.cat((tts_pad_embed.expand(-1, codec_input.shape[1] - 2, -1), tts_bos_embed), dim=1)
+                codec_prefix = codec_prefix + codec_input[:, :-1]
+                prefix = torch.cat((role_embed, codec_prefix), dim=1)
+                pad_row = codec_embed(self._long_tensor([talker_config.codec_pad_id], input_ids.device)).view(1, 1, -1)
+                tail = tts_pad_embed + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device))
+                if cache_key is not None:
+                    cv_cache[cache_key] = (prefix, pad_row, tail)
+            else:
+                prefix, pad_row, tail = cached
+            talker_prompt = prefix
 
             if non_streaming_mode:
                 text_all = text_projection(text_embedding(input_ids[:, 3:-5]))
                 text_all = torch.cat([text_all, tts_eos_embed], dim=1)
-                pad_ids = torch.full(
-                    (1, int(text_all.shape[1])),
-                    int(talker_config.codec_pad_id),
-                    device=input_ids.device,
-                    dtype=torch.long,
-                )
-                talker_prompt = torch.cat(
-                    [
-                        talker_prompt,
-                        text_all + codec_embed(pad_ids),
-                        tts_pad_embed + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device)),
-                    ],
-                    dim=1,
-                )
+                talker_prompt = torch.cat([talker_prompt, text_all + pad_row, tail], dim=1)
                 trailing_text_hidden = tts_pad_embed
             else:
                 first_text = text_projection(text_embedding(input_ids[:, 3:4])) + codec_input[:, -1:]

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py

@@ -65,6 +65,8 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp
         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
+        self._fused_requested = self._parse_bool_config(extra.get("code_predictor_fused"))
+        self._fused = None
 
     def _setup_compile(self) -> None:
         if self._compiled_model_fwd is not None:
@@ -100,11 +102,73 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapp
                         cache(self._proj_buf, batch, step)
         self._synchronize_warmup(weight.device)
         self._frame_cache = cache
+        if self._fused_requested:
+            from .fused_code_predictor import FusedCodePredictor
+
+            self._fused = FusedCodePredictor(self, max_batch)
+            # Compile every kernel variant outside any graph capture.
+            hidden = int(self.talker_config.hidden_size)
+            for batch in range(1, min(2, max_batch) + 1):
+                zeros = torch.zeros(batch, 1, hidden, device=weight.device, dtype=weight.dtype)
+                uniforms = torch.full(
+                    (batch, self._num_groups - 1, int(self.config.vocab_size)), 0.5, device=weight.device
+                )
+                codes = torch.zeros(batch, 1, device=weight.device, dtype=torch.long)
+                self._fused(codes, zeros, zeros, 1.0, 50, uniforms)
+            self._synchronize_warmup(weight.device)
+            logger.info("Qwen3-TTS fused code predictor enabled (capacity %d)", max_batch)
         # 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)
 
+    @torch.inference_mode()
+    def forward(
+        self,
+        layer0_code: torch.Tensor,
+        layer0_embed: torch.Tensor,
+        last_talker_hidden: torch.Tensor,
+        do_sample: bool = True,
+        temperature: float = 0.9,
+        top_k: int = 50,
+        top_p: float = 1.0,
+        generator: torch.Generator | None = None,
+        generators: Sequence[torch.Generator | None] | None = None,
+        sample_uniforms: torch.Tensor | None = None,
+    ):
+        self._validate_sampling_inputs(int(layer0_code.shape[0]), generators, sample_uniforms)
+        if self._fused_requested:
+            self._setup_compile()
+        if (
+            self._fused is not None
+            and do_sample
+            and temperature > 0
+            and top_p == 1.0
+            and sample_uniforms is not None
+            and sample_uniforms.dtype == torch.float32
+            and sample_uniforms.is_contiguous()
+        ):
+            return self._fused(
+                layer0_code,
+                layer0_embed,
+                last_talker_hidden,
+                1.0 / temperature,
+                int(top_k),
+                sample_uniforms,
+            )
+        return super().forward(
+            layer0_code,
+            layer0_embed,
+            last_talker_hidden,
+            do_sample=do_sample,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            generator=generator,
+            generators=generators,
+            sample_uniforms=sample_uniforms,
+        )
+
     def _predict_step_logits(
         self,
         proj_buf: torch.Tensor,

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -27,6 +27,7 @@ from vllm.sequence import IntermediateTensors
 
 from vllm_omni.data_entry_keys import OmniPayload
 from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor
+from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking
 from vllm_omni.utils.speaker_cache import (
     get_speaker_cache,
     iter_custom_voice_profiles,
@@ -39,6 +40,7 @@ from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfi
 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 .stream_decode import stream_ref_context_frames, talker_stream_decode_enabled
 from .tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Config
 from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2Encoder
 
@@ -411,7 +413,21 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # 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)
+        self.stream_decode = talker_stream_decode_enabled(vllm_config)
+        predictor = Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM
+        extra = predictor._stage_connector_extra_config(vllm_config)
+        self.stream_first_audio = self.stream_decode and predictor._parse_bool_config(
+            extra.get("talker_stream_first_audio")
+        )
+        # Model-local voice-clone priming boundary; each deployment selects
+        # its reference context independently of the two-stage codec profile.
+        self.stream_ref_context_frames = stream_ref_context_frames(vllm_config) if self.stream_decode else 0
+        self.mtp_eager_frames = talker_first_audio_enabled(vllm_config) or self.stream_decode
+        self.stream_decoder = None
+        self.stream_graphs = None
+        self.stream_prime_graphs = None
+        self.stream_sample_rate = 0
+        self.stream_chunk_frames = 25
         # 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.
@@ -718,6 +734,11 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
     # -------------------- Omni multimodal output plumbing --------------------
 
+    @property
+    def mm_outputs_fresh_per_step(self) -> bool:
+        """Stream-decode outputs come from :meth:`_make_eager_omni_output`, allocated per step."""
+        return self.stream_decoder is not None and self.eager_frames_active
+
     def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput:
         if isinstance(model_outputs, OmniOutput):
             return model_outputs
@@ -895,6 +916,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 "first_audio": torch.zeros((num_tokens,), dtype=torch.int8, device=hidden.device),
             },
         }
+        if self.stream_decoder is not None:
+            # Filled per sampled row by the runner's eager MTP; token-major so
+            # each request's slice is its frames of this step.
+            spf = int(self.stream_decoder.spf)
+            mm["model_outputs"] = torch.zeros((num_tokens, spf), dtype=torch.float32, device=hidden.device)
+            sr = torch.tensor(self.stream_sample_rate, dtype=torch.int32)
+            mm["sr"] = [sr] * len(info_dicts)
         ref_rows = [
             index
             for index, info in enumerate(info_dicts)
@@ -1442,8 +1470,72 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             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}
+        elif self.stream_decode:
+            # The in-stage decoder owns every frame. Optional first-audio
+            # delivery uses its PCM; later chunks use regular step outputs.
+            from .first_frame_decoder import Qwen3TTSFirstFrameDecoder
+            from .tokenizer_12hz.streaming_decoder import StreamingCodecDecoder
+
+            holder = Qwen3TTSFirstFrameDecoder(self.model_path)
+            decoder_loaded = holder.load(self.vllm_config)
+            loaded = set(loaded) | {f"stream_decoder.{name}" for name in decoder_loaded}
+            slots = int(self.vllm_config.scheduler_config.max_num_seqs)
+            self.stream_decoder = StreamingCodecDecoder(holder.decoder, num_slots=slots, dtype=torch.bfloat16)
+            self.stream_sample_rate = int(holder.sample_rate)
+            logger.info("Qwen3-TTS Talker stream decode enabled (%d slots)", slots)
         return loaded
 
+    def capture_stream_decode_graphs(self, batch_sizes: list[int]) -> None:
+        if self.stream_decoder is None or self.stream_graphs is not None:
+            return
+        from .tokenizer_12hz.streaming_decoder import StreamingDecodeGraphs
+
+        with torch.inference_mode():
+            self.stream_graphs = StreamingDecodeGraphs(self.stream_decoder, batch_sizes)
+            frames = min(self.stream_ref_context_frames, self.stream_chunk_frames, self.stream_decoder.max_frames)
+            if self.config.tts_model_type == "base" and frames > 1:
+                # Bound the multi-frame graph workspace by the configured
+                # token budget; larger priming groups keep the eager path.
+                budget = max(1, self.vllm_config.scheduler_config.max_num_batched_tokens // frames)
+                prime_sizes = sorted({1, *(size for size in batch_sizes if size <= budget)})
+                self.stream_prime_graphs = StreamingDecodeGraphs(self.stream_decoder, prime_sizes, frames=frames)
+        logger.info("Captured Talker stream decode graphs for batch sizes %s", self.stream_graphs.sizes)
+        if self.stream_prime_graphs is not None:
+            logger.info(
+                "Captured Talker reference priming graphs for %d frames and batch sizes %s",
+                self.stream_prime_graphs.frames,
+                self.stream_prime_graphs.sizes,
+            )
+
+    def get_stream_ref_context(self, info: OmniPayload) -> torch.Tensor | None:
+        """Model-owned reference-code layout and context boundary for a PCM slot."""
+        codes = info.get("codes")
+        ref = codes.get("ref") if isinstance(codes, dict) else None
+        if not isinstance(ref, torch.Tensor) or not ref.numel():
+            return None
+        return ref.reshape(-1, int(self.talker_config.num_code_groups))[-self.stream_ref_context_frames :]
+
+    def prime_stream_decoder(self, primes: list[tuple[int, torch.Tensor]]) -> None:
+        """Prime equal-length reference groups without padding or aliasing live slots."""
+        if not primes:
+            return
+        stream = self.stream_decoder
+        if stream is None:
+            raise RuntimeError("Reference priming requires the single-stage PCM decoder")
+        groups: dict[int, list[tuple[int, torch.Tensor]]] = {}
+        for req_idx, ref in primes:
+            groups.setdefault(int(ref.shape[0]), []).append((req_idx, ref))
+        for items in groups.values():
+            codes = to_device_nonblocking(torch.stack([ref for _idx, ref in items]), stream.device).to(torch.int32)
+            n = len(items)
+            slots = index_to_device([idx for idx, _ref in items], stream.device, dtype=torch.int32)
+            for t0 in range(0, int(codes.shape[1]), self.stream_chunk_frames):
+                pos = index_to_device([t0] * n, stream.device, dtype=torch.int32)
+                chunk = codes[:, t0 : t0 + self.stream_chunk_frames].contiguous()
+                graphs = self.stream_prime_graphs
+                decode = graphs if graphs is not None and chunk.shape[1] == graphs.frames else stream
+                decode(chunk, slots, pos)
+
     def _build_stacked_codec_embed(self) -> None:
         embeds = self.code_predictor.get_input_embeddings()
         if not embeds:

vllm_omni/model_executor/models/qwen3_tts/stream_decode.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Eligibility and model options for the single-stage CUDA PCM pipeline."""
+
+import torch
+from vllm.config import VllmConfig
+
+from vllm_omni.platforms import current_omni_platform
+
+
+def talker_stream_decode_enabled(vllm_config: VllmConfig) -> bool:
+    """The Talker decodes every frame itself and emits PCM as its final output."""
+    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
+    requested = Predictor._parse_bool_config(extra.get("talker_stream_decode"))
+    if not requested:
+        if getattr(model, "engine_output_type", None) == "audio":
+            raise ValueError("The single-stage Qwen3-TTS pipeline requires talker_stream_decode=True")
+        if Predictor._parse_bool_config(extra.get("talker_stream_first_audio")):
+            raise ValueError("talker_stream_first_audio requires the single-stage talker_stream_decode path")
+        return False
+    if getattr(model, "engine_output_type", None) != "audio" or not getattr(model, "final_output", False):
+        raise ValueError(
+            "talker_stream_decode requires a final audio-output Talker; "
+            "select qwen3_tts_fused_single_gpu.yaml instead of enabling it on the two-stage pipeline"
+        )
+    if getattr(model, "supports_running_prefix_cache_reset", True):
+        raise ValueError(
+            "The single-stage PCM pipeline must disable running prefix-cache reset; "
+            "select qwen3_tts_fused_single_gpu.yaml"
+        )
+    if not (
+        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
+    ):
+        raise ValueError(
+            "talker_stream_decode requires CUDA, MRv2, async_chunk=True, TP=PP=1, "
+            "an in-process worker and enable_prefix_caching=False; "
+            "use qwen3_tts_fused_single_gpu.yaml"
+        )
+    if Predictor._parse_bool_config(extra.get("talker_first_audio")):
+        raise ValueError("talker_stream_decode and talker_first_audio are mutually exclusive")
+    return True
+
+
+def stream_ref_context_frames(vllm_config: VllmConfig) -> int:
+    """Reference-code context frames for voice-clone priming in the Talker."""
+    from .qwen3_tts_code_predictor_vllm import Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor
+
+    extra = Predictor._stage_connector_extra_config(vllm_config)
+    return int(extra.get("ref_code_context_frames") or extra.get("codec_left_context_frames", 25))

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

@@ -0,0 +1,799 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# ruff: noqa: N803 - Triton constexpr parameters use kernel-style capitals.
+"""Stateful, fused streaming decoder for the Qwen3-TTS 12 Hz codec.
+
+Every layer of the codec decoder is causal, so what a new frame needs from the
+past is small and fixed: each causal conv's last ``(k - 1) * dilation`` input
+rows, each overlapping transposed conv's last GEMM row, and each attention
+layer's last ``sliding_window`` keys and values. A request owns a slot of
+per-layer state holding exactly that, and a call decodes ``T`` new frames for
+each of ``B`` requests, reading and advancing the state. The output equals the
+whole-utterance decode (``Qwen3TTSTokenizerV2Decoder._forward_exact``) frame by
+frame, up to floating point rounding: nothing is recomputed and nothing is
+truncated.
+
+Layout and fusion:
+
+- Activations are time-major ``[rows, C]``; every dense conv is a tap-major
+  im2col (one Triton kernel that also applies the pending bias and SnakeBeta,
+  and reads/writes the history) followed by one cuBLAS GEMM.
+- Linear chains fold at load: the RVQ output projections, ``pre_conv`` and the
+  transformer's ``input_proj`` become one ``k = 3`` conv over the concatenated
+  codebook embeddings; LayerScale folds into ``o_proj``/``down_proj``; ConvNeXt
+  ``gamma`` folds into ``pwconv2``.
+- A bias that feeds a residual stream is carried to that stream's consumer
+  (the next im2col or transposed-conv GEMM) instead of being added in place.
+- Residual branches accumulate into their stream with in-place ``addmm_``.
+
+History buffers are double-buffered per slot: a call reads half ``parity[slot]``
+and writes the other half, then flips the parity, so no kernel overwrites rows
+another program of the same launch still reads. A row whose first frame index
+``pos`` is 0 reads zeros, which are the convs' causal padding, and attention
+only looks back ``min(pos + 1, window)`` positions, so a new request needs no
+state reset.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from vllm.triton_utils import tl, triton
+
+from vllm_omni.platforms import current_omni_platform
+
+_KERNEL = 7
+
+
+# --------------------------------------------------------------------------- kernels
+
+
+@triton.jit
+def _rvq_gather_kernel(codes_ptr, books_ptr, out_ptr, NQ: tl.constexpr, CB: tl.constexpr, HALF: tl.constexpr):
+    row = tl.program_id(0)
+    offs = tl.arange(0, HALF)
+    c0 = tl.load(codes_ptr + row * NQ).to(tl.int64)
+    c0 = tl.where((c0 >= 0) & (c0 < CB), c0, 0)  # EOS/special rows decode as code 0; callers drop them
+    first = tl.load(books_ptr + c0 * HALF + offs).to(tl.float32)
+    acc = tl.zeros([HALF], dtype=tl.float32)
+    for q in range(1, NQ):
+        c = tl.load(codes_ptr + row * NQ + q).to(tl.int64)
+        c = tl.where((c >= 0) & (c < CB), c, 0)
+        acc += tl.load(books_ptr + (q * CB + c) * HALF + offs).to(tl.float32)
+    dt = out_ptr.dtype.element_ty
+    tl.store(out_ptr + row * 2 * HALF + offs, first.to(dt))
+    tl.store(out_ptr + row * 2 * HALF + HALF + offs, acc.to(dt))
+
+
+@triton.jit
+def _fast_sin(y):
+    # Two-step reduction to [-pi, pi], then the hardware approximation
+    # (absolute error ~4e-7 there): accurate sin's instruction count, not
+    # memory, bounds every kernel that applies SnakeBeta.
+    k = tl.extra.cuda.libdevice.rint(y * 0.15915494309189535)
+    r = tl.fma(-k, -1.7484555314695172e-07, tl.fma(-k, 6.2831854820251465, y))
+    return tl.inline_asm_elementwise("sin.approx.f32 $0, $1;", "=r,r", [r], dtype=tl.float32, is_pure=True, pack=1)
+
+
+@triton.jit
+def _snake(v, a, ib):
+    s = _fast_sin(v * a)
+    return v + ib * s * s
+
+
+@triton.jit
+def _im2col_kernel(
+    x_ptr,  # [B * T, C] layer input rows
+    bias_ptr,  # [C] pending bias (added to x rows only) or dummy
+    alpha_ptr,  # [C] exp(alpha) or dummy
+    ibeta_ptr,  # [C] 1 / (exp(beta) + eps) or dummy
+    hist_ptr,  # [2, S, H, C] transformed history rows
+    slots_ptr,
+    pos_ptr,
+    par_ptr,
+    col_ptr,  # [B * T, K * C]
+    T,
+    S,
+    C: tl.constexpr,
+    K: tl.constexpr,
+    D: tl.constexpr,
+    H: tl.constexpr,
+    HAS_BIAS: tl.constexpr,
+    HAS_SNAKE: tl.constexpr,
+    BU: tl.constexpr,
+    BC: tl.constexpr,
+    EXT: tl.constexpr = False,
+):
+    """Source row ``u`` of ``[history (H) | x (T)]``: scatter to every tap that reads it, save the last H."""
+    b = tl.program_id(0)
+    u = tl.program_id(1) * BU + tl.arange(0, BU)
+    c = tl.program_id(2) * BC + tl.arange(0, BC)
+    cm = c < C
+    n_src = H + T
+    um = u < n_src
+    slot = tl.load(slots_ptr + b).to(tl.int64)
+    live = tl.load(pos_ptr + b) != 0
+    par = tl.load(par_ptr + slot).to(tl.int64)
+    # history rows
+    from_hist = u < H
+    hrow = tl.where(from_hist, u, 0)
+    hv = tl.load(
+        hist_ptr + ((par * S + slot) * H + hrow)[:, None] * C + c[None, :],
+        mask=(from_hist & um & live)[:, None] & cm[None, :],
+        other=0.0,
+    ).to(tl.float32)
+    xrow = tl.where(from_hist, 0, u - H)
+    xv = tl.load(
+        x_ptr + (b * T + xrow).to(tl.int64)[:, None] * C + c[None, :],
+        mask=((~from_hist) & um)[:, None] & cm[None, :],
+        other=0.0,
+    ).to(tl.float32)
+    if HAS_BIAS:
+        xv += tl.load(bias_ptr + c, mask=cm, other=0.0).to(tl.float32)[None, :]
+    if HAS_SNAKE:
+        a = tl.load(alpha_ptr + c, mask=cm, other=0.0).to(tl.float32)[None, :]
+        ib = tl.load(ibeta_ptr + c, mask=cm, other=0.0).to(tl.float32)[None, :]
+        xv = _snake(xv, a, ib)
+    dt = col_ptr.dtype.element_ty
+    v = tl.where(from_hist[:, None], hv, xv.to(dt).to(tl.float32)).to(dt)
+    if EXT:
+        tl.store(col_ptr + (b * n_src + u).to(tl.int64)[:, None] * C + c[None, :], v, mask=um[:, None] & cm[None, :])
+    else:
+        for j in tl.static_range(K):
+            t = u - j * D
+            tm = (t >= 0) & (t < T) & um
+            tl.store(
+                col_ptr + (b * T + tl.where(tm, t, 0)).to(tl.int64)[:, None] * (K * C) + j * C + c[None, :],
+                v,
+                mask=tm[:, None] & cm[None, :],
+            )
+    if H > 0:
+        hm = (u >= T) & um
+        tl.store(
+            hist_ptr + (((1 - par) * S + slot) * H + tl.where(hm, u - T, 0))[:, None] * C + c[None, :],
+            v,
+            mask=hm[:, None] & cm[None, :],
+        )
+
+
+@triton.jit
+def _rms_norm_kernel(x_ptr, w_ptr, out_ptr, eps, N: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.arange(0, N)
+    x = tl.load(x_ptr + row * N + offs).to(tl.float32)
+    var = tl.sum(x * x, axis=0) / N
+    dt = out_ptr.dtype.element_ty
+    y = (x * tl.rsqrt(var + eps)).to(dt).to(tl.float32)
+    w = tl.load(w_ptr + offs).to(tl.float32)
+    tl.store(out_ptr + row * N + offs, (w * y).to(dt))
+
+
+@triton.jit
+def _rope_kv_write_kernel(
+    qkv_ptr,  # [B * T, 3 * NH * HD]
+    q_ptr,  # [B * T, NH * HD]
+    k_ring_ptr,  # [S, RING, NH * HD]
+    v_ring_ptr,
+    slots_ptr,
+    pos_ptr,
+    T,
+    theta,
+    NH: tl.constexpr,
+    HD: tl.constexpr,
+    RING: tl.constexpr,
+):
+    r = tl.program_id(0)
+    h = tl.program_id(1)
+    b = r // T
+    p = tl.load(pos_ptr + b) + r % T
+    slot = tl.load(slots_ptr + b).to(tl.int64)
+    half: tl.constexpr = HD // 2
+    i = tl.arange(0, half)
+    inv_freq = tl.exp(-(2.0 * i.to(tl.float32) / HD) * tl.log(theta))
+    ang = p.to(tl.float32) * inv_freq
+    dt = q_ptr.dtype.element_ty
+    cos = tl.cos(ang).to(dt).to(tl.float32)
+    sin = tl.sin(ang).to(dt).to(tl.float32)
+    W: tl.constexpr = NH * HD
+    base = qkv_ptr + r.to(tl.int64) * 3 * W + h * HD
+    q1 = tl.load(base + i).to(tl.float32)
+    q2 = tl.load(base + half + i).to(tl.float32)
+    k1 = tl.load(base + W + i).to(tl.float32)
+    k2 = tl.load(base + W + half + i).to(tl.float32)
+    vv = tl.arange(0, HD)
+    v = tl.load(base + 2 * W + vv)
+    qo = q_ptr + r.to(tl.int64) * W + h * HD
+    tl.store(qo + i, (q1 * cos - q2 * sin).to(dt))
+    tl.store(qo + half + i, (q2 * cos + q1 * sin).to(dt))
+    ring = (slot * RING + p % RING) * W + h * HD
+    tl.store(k_ring_ptr + ring + i, (k1 * cos - k2 * sin).to(dt))
+    tl.store(k_ring_ptr + ring + half + i, (k2 * cos + k1 * sin).to(dt))
+    tl.store(v_ring_ptr + ring + vv, v)
+
+
+@triton.jit
+def _attend_kernel(
+    q_ptr,  # [B * T, NH * HD]
+    k_ring_ptr,
+    v_ring_ptr,
+    out_ptr,  # [B * T, NH * HD]
+    slots_ptr,
+    pos_ptr,
+    T,
+    scale,
+    NH: tl.constexpr,
+    HD: tl.constexpr,
+    RING: tl.constexpr,
+    WINDOW: tl.constexpr,
+    BK: tl.constexpr,
+):
+    r = tl.program_id(0)
+    h = tl.program_id(1)
+    b = r // T
+    p = tl.load(pos_ptr + b) + r % T
+    slot = tl.load(slots_ptr + b).to(tl.int64)
+    W: tl.constexpr = NH * HD
+    d = tl.arange(0, HD)
+    q = tl.load(q_ptr + r.to(tl.int64) * W + h * HD + d).to(tl.float32)
+    j = tl.arange(0, BK)
+    n = tl.minimum(p + 1, WINDOW)
+    km = j < n
+    kp = p - j  # key positions p, p-1, ...
+    idx = (slot * RING + tl.where(km, kp, 0) % RING) * W + h * HD
+    k = tl.load(k_ring_ptr + idx[:, None] + d[None, :], mask=km[:, None], other=0.0).to(tl.float32)
+    s = tl.sum(k * q[None, :], axis=1) * scale
+    s = tl.where(km, s, -float("inf"))
+    m = tl.max(s, axis=0)
+    e = tl.exp(s - m)
+    e = tl.where(km, e, 0.0)
+    w = e / tl.sum(e, axis=0)
+    dt = out_ptr.dtype.element_ty
+    w = w.to(dt).to(tl.float32)
+    v = tl.load(v_ring_ptr + idx[:, None] + d[None, :], mask=km[:, None], other=0.0).to(tl.float32)
+    o = tl.sum(v * w[:, None], axis=0)
+    tl.store(out_ptr + r.to(tl.int64) * W + h * HD + d, o.to(dt))
+
+
+@triton.jit
+def _silu_mul_kernel(gu_ptr, out_ptr, N: tl.constexpr, BN: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.program_id(1) * BN + tl.arange(0, BN)
+    m = offs < N
+    g = tl.load(gu_ptr + row * 2 * N + offs, mask=m).to(tl.float32)
+    u = tl.load(gu_ptr + row * 2 * N + N + offs, mask=m).to(tl.float32)
+    dt = out_ptr.dtype.element_ty
+    sg = (g / (1.0 + tl.exp(-g))).to(dt).to(tl.float32)
+    tl.store(out_ptr + row * N + offs, (sg * u).to(dt), mask=m)
+
+
+@triton.jit
+def _dwconv_ln_kernel(
+    z_ptr,  # [B * T, C] ConvNeXt input rows (true values)
+    w_ptr,  # [K, C]
+    b_ptr,  # [C]
+    lnw_ptr,
+    lnb_ptr,
+    hist_ptr,  # [2, S, K - 1, C]
+    slots_ptr,
+    pos_ptr,
+    par_ptr,
+    out_ptr,  # [B * T, C] LayerNorm output
+    T,
+    S,
+    eps,
+    C: tl.constexpr,
+    K: tl.constexpr,
+):
+    b = tl.program_id(0)
+    t = tl.program_id(1)  # 0 .. max(T, K - 1)
+    H: tl.constexpr = K - 1
+    c = tl.arange(0, C)
+    slot = tl.load(slots_ptr + b).to(tl.int64)
+    live = tl.load(pos_ptr + b) != 0
+    par = tl.load(par_ptr + slot).to(tl.int64)
+    dt = out_ptr.dtype.element_ty
+    hbase = hist_ptr + (par * S + slot) * H * C
+    acc = tl.load(b_ptr + c).to(tl.float32)
+    for j in tl.static_range(K):
+        u = t + j  # index into [hist (H) | z (T)]
+        fh = u < H
+        hv = tl.load(hbase + tl.where(fh, u, 0) * C + c, mask=(c < C) & fh & live, other=0.0).to(tl.float32)
+        zrow = (b * T + tl.where(fh, 0, u - H)).to(tl.int64)
+        zv = tl.load(z_ptr + zrow * C + c, mask=(c < C) & (~fh) & (t < T), other=0.0).to(tl.float32)
+        acc += tl.where(fh, hv, zv) * tl.load(w_ptr + j * C + c).to(tl.float32)
+    h = acc.to(dt).to(tl.float32)
+    mean = tl.sum(h, axis=0) / C
+    hc = h - mean
+    var = tl.sum(hc * hc, axis=0) / C
+    y = hc * tl.rsqrt(var + eps) * tl.load(lnw_ptr + c).to(tl.float32) + tl.load(lnb_ptr + c).to(tl.float32)
+    tl.store(out_ptr + (b * T + t).to(tl.int64) * C + c, y.to(dt), mask=(c < C) & (t < T))
+    # new history row t: source row T + t of [hist | z]
+    u = T + t
+    fh = u < H
+    hv = tl.load(hbase + tl.where(fh, u, 0) * C + c, mask=(c < C) & fh & live & (t < H), other=0.0)
+    zrow = (b * T + tl.where(fh, 0, u - H)).to(tl.int64)
+    zv = tl.load(z_ptr + zrow * C + c, mask=(c < C) & (~fh) & (t < H), other=0.0)
+    tl.store(
+        hist_ptr + (((1 - par) * S + slot) * H + tl.where(t < H, t, 0)) * C + c,
+        tl.where(fh, hv, zv).to(dt),
+        mask=(c < C) & (t < H),
+    )
+
+
+@triton.jit
+def _bias_gelu_kernel(x_ptr, b_ptr, N: tl.constexpr, BN: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.program_id(1) * BN + tl.arange(0, BN)
+    m = offs < N
+    x = tl.load(x_ptr + row * N + offs, mask=m).to(tl.float32) + tl.load(b_ptr + offs, mask=m).to(tl.float32)
+    dt = x_ptr.dtype.element_ty
+    x = x.to(dt).to(tl.float32)
+    y = 0.5 * x * (1.0 + tl.erf(x * 0.7071067811865476))
+    tl.store(x_ptr + row * N + offs, y.to(dt), mask=m)
+
+
+@triton.jit
+def _col2im_kernel(
+    z_ptr,  # [B * T, 2 * R, C] transposed-conv GEMM rows
+    bias_ptr,  # [C]
+    prev_ptr,  # [2, S, R, C] previous call's last GEMM row, second half
+    slots_ptr,
+    pos_ptr,
+    par_ptr,
+    out_ptr,  # [B * T * R, C]
+    T,
+    S,
+    R: tl.constexpr,
+    C: tl.constexpr,
+    BC: tl.constexpr,
+):
+    row = tl.program_id(0)  # b * T + t
+    j = tl.program_id(1)  # 0 .. R - 1
+    c = tl.program_id(2) * BC + tl.arange(0, BC)
+    cm = c < C
+    b = row // T
+    t = row % T
+    slot = tl.load(slots_ptr + b).to(tl.int64)
+    live = tl.load(pos_ptr + b) != 0
+    par = tl.load(par_ptr + slot).to(tl.int64)
+    zr = z_ptr + row.to(tl.int64) * 2 * R * C
+    cur = tl.load(zr + j * C + c, mask=cm).to(tl.float32)
+    prev_in = tl.load(zr - 2 * R * C + (R + j) * C + c, mask=cm & (t > 0), other=0.0).to(tl.float32)
+    prev_st = tl.load(prev_ptr + ((par * S + slot) * R + j) * C + c, mask=cm & (t == 0) & live, other=0.0).to(
+        tl.float32
+    )
+    y = cur + tl.where(t > 0, prev_in, prev_st) + tl.load(bias_ptr + c, mask=cm).to(tl.float32)
+    dt = out_ptr.dtype.element_ty
+    tl.store(out_ptr + (row.to(tl.int64) * R + j) * C + c, y.to(dt), mask=cm)
+    last = tl.load(zr + (R + j) * C + c, mask=cm & (t == T - 1))
+    tl.store(prev_ptr + (((1 - par) * S + slot) * R + j) * C + c, last, mask=cm & (t == T - 1))
+
+
+@triton.jit
+def _conv_out_kernel(
+    ext_ptr,  # [B, H + T, C] transformed rows: history then this call's
+    w_ptr,  # [K, C]
+    wb,  # conv bias (float)
+    out_ptr,  # [B, T] float32 waveform
+    T,
+    C: tl.constexpr,
+    K: tl.constexpr,
+    BT: tl.constexpr,
+    CP: tl.constexpr,
+):
+    b = tl.program_id(0)
+    t = tl.program_id(1) * BT + tl.arange(0, BT)
+    c = tl.arange(0, CP)
+    cm = c < C
+    tm = t < T
+    base = ext_ptr + b.to(tl.int64) * (T + K - 1) * C
+    acc = tl.zeros([BT], dtype=tl.float32)
+    for j in tl.static_range(K):
+        v = tl.load(base + (t + j)[:, None] * C + c[None, :], mask=tm[:, None] & cm[None, :], other=0.0).to(tl.float32)
+        acc += tl.sum(v * tl.load(w_ptr + j * C + c, mask=cm, other=0.0).to(tl.float32)[None, :], axis=1)
+    dt = ext_ptr.dtype.element_ty
+    y = acc.to(dt).to(tl.float32) + wb
+    y = tl.minimum(tl.maximum(y.to(dt).to(tl.float32), -1.0), 1.0)
+    tl.store(out_ptr + b.to(tl.int64) * T + t, y, mask=tm)
+
+
+@triton.jit
+def _flip_kernel(par_ptr, slots_ptr):
+    slot = tl.load(slots_ptr + tl.program_id(0)).to(tl.int64)
+    tl.store(par_ptr + slot, 1 - tl.load(par_ptr + slot))
+
+
+# --------------------------------------------------------------------------- module
+
+
+@dataclass
+class _ConvSpec:
+    weight: torch.Tensor  # [K * C_in, C_out] tap-major
+    c_in: int
+    k: int
+    d: int
+    hist: torch.Tensor | None = None  # [2, S, H, C_in]
+
+
+def _conv_taps(w: torch.Tensor) -> torch.Tensor:
+    """``[C_out, C_in, K]`` -> ``[K * C_in, C_out]``."""
+    c_out, c_in, k = w.shape
+    return w.permute(2, 1, 0).reshape(k * c_in, c_out)
+
+
+def _snake_params(module) -> tuple[torch.Tensor, torch.Tensor]:
+    # SnakeBeta: x + 1 / (exp(beta) + 1e-9) * sin(exp(alpha) * x) ** 2
+    alpha = module.alpha.detach().float()
+    beta = module.beta.detach().float()
+    return torch.exp(alpha).contiguous(), (1.0 / (torch.exp(beta) + 1e-9)).contiguous()
+
+
+class StreamingCodecDecoder(nn.Module):
+    """Frame-exact stateful decoder built from a loaded ``Qwen3TTSTokenizerV2Decoder``."""
+
+    RING = 128  # KV ring positions; a call may carry up to RING - window + 1 frames
+
+    def __init__(self, decoder: nn.Module, num_slots: int, dtype: torch.dtype | None = None) -> None:
+        super().__init__()
+        cfg = decoder.config
+        dev = next(decoder.parameters()).device
+        dtype = dtype or next(decoder.parameters()).dtype
+        self.dtype, self.device = dtype, dev
+        self.num_slots = num_slots + 1  # last slot: scratch for padding rows
+        self.scratch_slot = num_slots
+        S = self.num_slots
+        self.nq = int(cfg.num_quantizers)
+        self.cb = int(cfg.codebook_size)
+        self.half = int(cfg.codebook_dim) // 2
+        self.hidden = int(cfg.hidden_size)
+        self.heads = int(cfg.num_attention_heads)
+        self.head_dim = int(cfg.head_dim)
+        self.window = int(cfg.sliding_window)
+        self.theta = float(cfg.rope_theta)
+        self.eps = float(cfg.rms_norm_eps)
+        self.max_frames = self.RING - self.window + 1
+        self.spf = int(decoder.total_upsample)
+        if self.heads != int(cfg.num_key_value_heads):
+            raise ValueError("streaming codec decoder assumes no grouped KV")
+
+        f32 = lambda t: t.detach().to(device=dev, dtype=torch.float32)  # noqa: E731
+        cast = lambda t: t.to(dtype).contiguous()  # noqa: E731
+
+        def hist(h: int, c: int) -> torch.Tensor:
+            return torch.zeros(2, S, h, c, device=dev, dtype=dtype)
+
+        # RVQ books, pre-divided by usage: [NQ * CB, HALF]
+        q = decoder.quantizer
+        books = []
+        for vq in [q.rvq_first.vq.layers[0]] + list(q.rvq_rest.vq.layers):
+            cbk = vq._codebook
+            books.append(f32(cbk.embedding_sum) / f32(cbk.cluster_usage).clamp(min=cbk.epsilon)[:, None])
+        self.books = cast(torch.cat(books))
+        w_rvq = torch.cat(
+            [f32(q.rvq_first.output_proj.weight)[:, :, 0], f32(q.rvq_rest.output_proj.weight)[:, :, 0]], 1
+        )
+        # pre_conv (k=3) and input_proj folded onto the embedding concat
+        pre_w, pre_b = f32(decoder.pre_conv.conv.weight), f32(decoder.pre_conv.conv.bias)
+        pt = decoder.pre_transformer
+        in_w, in_b = f32(pt.input_proj.weight), f32(pt.input_proj.bias)
+        taps = torch.stack([in_w @ pre_w[:, :, j] @ w_rvq for j in range(pre_w.shape[2])], dim=2)  # [hid, 2H, 3]
+        self.pre = _ConvSpec(cast(_conv_taps(taps)), 2 * self.half, 3, 1, hist(2, 2 * self.half))
+        self.pre_b = cast(in_w @ pre_b + in_b)
+
+        # transformer
+        self.layers = []
+        qd = self.heads * self.head_dim
+        for layer in pt.layers:
+            a = layer.self_attn
+            qkv = torch.cat([f32(a.q_proj.weight), f32(a.k_proj.weight), f32(a.v_proj.weight)], 0)
+            o = f32(a.o_proj.weight) * f32(layer.self_attn_layer_scale.scale)[:, None]
+            gu = torch.cat([f32(layer.mlp.gate_proj.weight), f32(layer.mlp.up_proj.weight)], 0)
+            down = f32(layer.mlp.down_proj.weight) * f32(layer.mlp_layer_scale.scale)[:, None]
+            self.layers.append(
+                dict(
+                    ln1=cast(f32(layer.input_layernorm.weight)),
+                    ln2=cast(f32(layer.post_attention_layernorm.weight)),
+                    qkv=cast(qkv.t()),
+                    o=cast(o.t()),
+                    gu=cast(gu.t()),
+                    down=cast(down.t()),
+                    k_ring=torch.zeros(S, self.RING, qd, device=dev, dtype=dtype),
+                    v_ring=torch.zeros(S, self.RING, qd, device=dev, dtype=dtype),
+                )
+            )
+        self.norm_w = cast(f32(pt.norm.weight))
+        self.out_w = cast(f32(pt.output_proj.weight).t())
+        self.out_b = cast(f32(pt.output_proj.bias))
+
+        # upsample: transposed conv (kernel = stride, no overlap) + ConvNeXt
+        self.ups = []
+        latent = int(cfg.latent_dim)
+        pending = torch.zeros(latent, device=dev, dtype=torch.float32)
+        for trans, block in decoder.upsample:
+            tw = f32(trans.conv.weight)  # [C_in, C_out, r]
+            r = tw.shape[2]
+            w = tw.permute(0, 2, 1).reshape(latent, r * latent)  # [C_in, r * C_out]
+            bias = (pending @ w).view(r, latent) + f32(trans.conv.bias)[None, :]
+            dw = block.dwconv.conv
+            gamma = f32(block.gamma)
+            self.ups.append(
+                dict(
+                    r=r,
+                    w=cast(w),
+                    b=cast(bias.reshape(-1)),
+                    dw_w=cast(f32(dw.weight)[:, 0, :].t()),
+                    dw_b=cast(f32(dw.bias)),
+                    ln_w=cast(f32(block.norm.weight)),
+                    ln_b=cast(f32(block.norm.bias)),
+                    ln_eps=float(block.norm.eps),
+                    pw1=cast(f32(block.pwconv1.weight).t()),
+                    pw1_b=cast(f32(block.pwconv1.bias)),
+                    pw2=cast((f32(block.pwconv2.weight) * gamma[:, None]).t()),
+                    hist=hist(_KERNEL - 1, latent),
+                )
+            )
+            pending = f32(block.pwconv2.bias) * gamma
+
+        # conv_in (pending bias of the last ConvNeXt added on load)
+        dec = decoder.decoder
+        cin = dec[0].conv
+        self.conv_in = _ConvSpec(cast(_conv_taps(f32(cin.weight))), latent, _KERNEL, 1, hist(_KERNEL - 1, latent))
+        self.conv_in_pb = cast(pending)
+        pending = f32(cin.bias)
+
+        # decoder blocks
+        self.blocks = []
+        for i, blk in enumerate(dec[1:-2]):
+            snake, trans, *units = blk.block
+            c_in = trans.conv.in_channels
+            c_out = trans.conv.out_channels
+            rate = trans.conv.stride[0]
+            tw = f32(trans.conv.weight)  # [C_in, C_out, 2 * rate]
+            entry = dict(
+                pb=cast(pending),
+                snake=_snake_params(snake),
+                rate=rate,
+                c_in=c_in,
+                c_out=c_out,
+                w=cast(tw.permute(0, 2, 1).reshape(c_in, 2 * rate * c_out)),
+                b=cast(f32(trans.conv.bias)),
+                prev=torch.zeros(2, S, rate, c_out, device=dev, dtype=dtype),
+                units=[],
+            )
+            pending = torch.zeros(c_out, device=dev, dtype=torch.float32)
+            for u in units:
+                c1, c2 = u.conv1.conv, u.conv2.conv
+                entry["units"].append(
+                    dict(
+                        pb=cast(pending),
+                        act1=_snake_params(u.act1),
+                        conv1=_ConvSpec(
+                            cast(_conv_taps(f32(c1.weight))),
+                            c_out,
+                            _KERNEL,
+                            c1.dilation[0],
+                            hist((_KERNEL - 1) * c1.dilation[0], c_out),
+                        ),
+                        b1=cast(f32(c1.bias)),
+                        conv1_w4=(
+                            cast(f32(c1.weight)).unsqueeze(2).contiguous(memory_format=torch.channels_last)
+                            if c1.dilation[0] == 1
+                            else None
+                        ),
+                        act2=_snake_params(u.act2),
+                        w2=cast(f32(c2.weight)[:, :, 0].t()),
+                    )
+                )
+                pending = pending + f32(c2.bias)
+            self.blocks.append(entry)
+        self.out_pb = cast(pending)
+        self.out_snake = _snake_params(dec[-2])
+        co = dec[-1].conv
+        self.out_c = co.in_channels
+        self.conv_out_w = cast(f32(co.weight)[0].t())  # [K, C]
+        self.conv_out_b = float(f32(co.bias)[0])
+        self.conv_out = _ConvSpec(self.conv_out_w, self.out_c, _KERNEL, 1, hist(_KERNEL - 1, self.out_c))
+        self.parity = torch.zeros(S, device=dev, dtype=torch.int32)
+        self._dummy = torch.zeros(1, device=dev, dtype=dtype)
+
+    def _slot_state(self, slot: int) -> list[torch.Tensor]:
+        state = [self.parity[slot : slot + 1]]
+        for layer in self.layers:
+            state.extend((layer["k_ring"][slot], layer["v_ring"][slot]))
+        for spec in (self.pre, self.conv_in, self.conv_out):
+            assert spec.hist is not None
+            state.append(spec.hist[:, slot])
+        for up in self.ups:
+            state.append(up["hist"][:, slot])
+        for block in self.blocks:
+            state.append(block["prev"][:, slot])
+            for unit in block["units"]:
+                state.append(unit["conv1"].hist[:, slot])
+        return state
+
+    def save_slot(self, slot: int) -> list[torch.Tensor]:
+        """Offload a preempted request before its runner slot is recycled."""
+        return [tensor.to(device="cpu", copy=True) for tensor in self._slot_state(slot)]
+
+    def restore_slot(self, slot: int, state: list[torch.Tensor]) -> None:
+        for target, saved in zip(self._slot_state(slot), state, strict=True):
+            target.copy_(saved)
+
+    # ------------------------------------------------------------------ helpers
+    def _im2col(self, x, spec: _ConvSpec, T, slots, pos, bias=None, snake=None, ext=False):
+        k = spec.k
+        d = spec.d
+        c = x.shape[1]
+        h = (k - 1) * d
+        B = slots.shape[0]
+        if ext:
+            col = torch.empty(B * (h + T), c, device=x.device, dtype=self.dtype)
+        else:
+            col = torch.empty(B * T, k * c, device=x.device, dtype=self.dtype)
+        bu = 32
+        bc = min(128, triton.next_power_of_2(c))
+        grid = (B, triton.cdiv(h + T, bu), triton.cdiv(c, bc))
+        a, ib = snake if snake is not None else (self._dummy, self._dummy)
+        _im2col_kernel[grid](
+            x, bias if bias is not None else self._dummy, a, ib,
+            spec.hist if h > 0 else self._dummy, slots, pos, self.parity, col,
+            T, self.num_slots, C=c, K=k, D=d, H=h,
+            HAS_BIAS=bias is not None, HAS_SNAKE=snake is not None, BU=bu, BC=bc, EXT=ext,
+        )  # fmt: skip
+        return col
+
+    def _conv_cudnn(self, x, u, T, slots, pos):
+        """Undilated unit conv: one source-row pass, then cuDNN on its channels-last view.
+
+        The im2col path writes a ``K``-times patch matrix; cuDNN reads the
+        ``[history | new]`` rows directly (same result up to summation order).
+        """
+        spec = u["conv1"]
+        ext = self._im2col(x, spec, T, slots, pos, bias=u["pb"], snake=u["act1"], ext=True)
+        b = slots.shape[0]
+        c = x.shape[1]
+        x4 = ext.view(b, -1, c).permute(0, 2, 1).unsqueeze(2)  # [B, C, 1, H + T], channels-last memory
+        y = torch.nn.functional.conv2d(x4, u["conv1_w4"])  # [B, N, 1, T], channels-last memory
+        return y.permute(0, 2, 3, 1).reshape(b * T, -1)
+
+    def _snake(self, x, snake, bias):
+        n, c = x.shape
+        out = torch.empty_like(x)
+        bu, bc = 32, min(128, triton.next_power_of_2(c))
+        one = torch.zeros(1, device=x.device, dtype=torch.int32)
+        _im2col_kernel[(1, triton.cdiv(n, bu), triton.cdiv(c, bc))](
+            x, bias, snake[0], snake[1], self._dummy, one, one, self.parity, out,
+            n, self.num_slots, C=c, K=1, D=1, H=0, HAS_BIAS=True, HAS_SNAKE=True, BU=bu, BC=bc,
+        )  # fmt: skip
+        return out
+
+    # ------------------------------------------------------------------ forward
+    @torch.no_grad()
+    def forward(self, codes: torch.Tensor, slots: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        """``codes`` [B, T, NQ] int, ``slots``/``pos`` [B] int32 -> waveform [B, T * spf] float32.
+
+        ``pos[b]`` is the index of the row's first frame in its stream; 0 starts
+        a new stream in ``slots[b]``. Frames beyond ``max_frames`` per call must
+        be split by the caller.
+        """
+        B, T, _ = codes.shape
+        if T > self.max_frames:
+            raise ValueError(f"at most {self.max_frames} frames per call, got {T}")
+        dev, dt = codes.device, self.dtype
+        codes = codes.reshape(B * T, self.nq).to(torch.int32).contiguous()
+        emb = torch.empty(B * T, 2 * self.half, device=dev, dtype=dt)
+        _rvq_gather_kernel[(B * T,)](codes, self.books, emb, NQ=self.nq, CB=self.cb, HALF=self.half)
+        x = torch.addmm(self.pre_b, self._im2col(emb, self.pre, T, slots, pos), self.pre.weight)
+
+        hid, qd = self.hidden, self.heads * self.head_dim
+        rows = B * T
+        h = torch.empty(rows, hid, device=dev, dtype=dt)
+        q = torch.empty(rows, qd, device=dev, dtype=dt)
+        att = torch.empty(rows, qd, device=dev, dtype=dt)
+        a = torch.empty(rows, self.layers[0]["down"].shape[0], device=dev, dtype=dt)
+        inter = a.shape[1]
+        scale = self.head_dim**-0.5
+        for L in self.layers:
+            _rms_norm_kernel[(rows,)](x, L["ln1"], h, self.eps, N=hid)
+            qkv = torch.mm(h, L["qkv"])
+            _rope_kv_write_kernel[(rows, self.heads)](
+                qkv, q, L["k_ring"], L["v_ring"], slots, pos, T, self.theta,
+                NH=self.heads, HD=self.head_dim, RING=self.RING,
+            )  # fmt: skip
+            _attend_kernel[(rows, self.heads)](
+                q, L["k_ring"], L["v_ring"], att, slots, pos, T, scale,
+                NH=self.heads, HD=self.head_dim, RING=self.RING, WINDOW=self.window,
+                BK=triton.next_power_of_2(self.window),
+            )  # fmt: skip
+            x.addmm_(att, L["o"])
+            _rms_norm_kernel[(rows,)](x, L["ln2"], h, self.eps, N=hid)
+            gu = torch.mm(h, L["gu"])
+            _silu_mul_kernel[(rows, triton.cdiv(inter, 1024))](gu, a, N=inter, BN=1024)
+            x.addmm_(a, L["down"])
+        _rms_norm_kernel[(rows,)](x, self.norm_w, h, self.eps, N=hid)
+        y = torch.addmm(self.out_b, h, self.out_w)  # [rows, latent]
+
+        t_rows = T
+        for U in self.ups:
+            r = U["r"]
+            z = torch.addmm(U["b"], y, U["w"]).view(rows * r, -1)
+            rows, t_rows = rows * r, t_rows * r
+            c = z.shape[1]
+            n = torch.empty_like(z)
+            _dwconv_ln_kernel[(B, max(t_rows, _KERNEL - 1))](
+                z, U["dw_w"], U["dw_b"], U["ln_w"], U["ln_b"], U["hist"], slots, pos, self.parity, n,
+                t_rows, self.num_slots, U["ln_eps"], C=c, K=_KERNEL,
+            )  # fmt: skip
+            p1 = torch.mm(n, U["pw1"])
+            _bias_gelu_kernel[(rows, triton.cdiv(p1.shape[1], 1024))](p1, U["pw1_b"], N=p1.shape[1], BN=1024)
+            z.addmm_(p1, U["pw2"])
+            y = z
+
+        x = torch.mm(self._im2col(y, self.conv_in, t_rows, slots, pos, bias=self.conv_in_pb), self.conv_in.weight)
+        for blk in self.blocks:
+            s = self._snake(x, blk["snake"], blk["pb"])
+            zt = torch.mm(s, blk["w"])  # [rows, 2 * rate * c_out]
+            r, c_out = blk["rate"], blk["c_out"]
+            x = torch.empty(rows * r, c_out, device=dev, dtype=dt)
+            bc = min(128, triton.next_power_of_2(c_out))
+            _col2im_kernel[(rows, r, triton.cdiv(c_out, bc))](
+                zt, blk["b"], blk["prev"], slots, pos, self.parity, x, t_rows, self.num_slots,
+                R=r, C=c_out, BC=bc,
+            )  # fmt: skip
+            rows, t_rows = rows * r, t_rows * r
+            for u in blk["units"]:
+                if u.get("conv1_w4") is not None:
+                    c1 = self._conv_cudnn(x, u, t_rows, slots, pos)
+                else:
+                    col = self._im2col(x, u["conv1"], t_rows, slots, pos, bias=u["pb"], snake=u["act1"])
+                    c1 = torch.mm(col, u["conv1"].weight)
+                s2 = self._snake(c1, u["act2"], u["b1"])
+                x.addmm_(s2, u["w2"])
+
+        ext = self._im2col(x, self.conv_out, t_rows, slots, pos, bias=self.out_pb, snake=self.out_snake, ext=True)
+        wav = torch.empty(B, t_rows, device=dev, dtype=torch.float32)
+        _conv_out_kernel[(B, triton.cdiv(t_rows, 256))](
+            ext, self.conv_out_w, self.conv_out_b, wav, t_rows,
+            C=self.out_c, K=_KERNEL, BT=256, CP=triton.next_power_of_2(self.out_c),
+        )  # fmt: skip
+        _flip_kernel[(B,)](self.parity, slots)
+        return wav
+
+
+class StreamingDecodeGraphs:
+    """CUDA graphs of ``StreamingCodecDecoder`` for a fixed frame count per batch bucket."""
+
+    def __init__(self, sd: StreamingCodecDecoder, batch_sizes: list[int], frames: int = 1) -> None:
+        self.sd, self.frames = sd, frames
+        self.graphs: dict[int, tuple] = {}
+        dev = sd.device
+        pool = torch.cuda.graph_pool_handle()
+        for bsz in sorted(set(batch_sizes)):
+            codes = torch.zeros(bsz, frames, sd.nq, dtype=torch.int32, device=dev)
+            slots = torch.full((bsz,), sd.scratch_slot, dtype=torch.int32, device=dev)
+            pos = torch.zeros(bsz, dtype=torch.int32, device=dev)
+            sd(codes, slots, pos)
+            current_omni_platform.synchronize()
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, pool=pool):
+                out = sd(codes, slots, pos)
+            self.graphs[bsz] = (graph, codes, slots, pos, out)
+        self.sizes = sorted(self.graphs)
+
+    def __call__(self, codes: torch.Tensor, slots: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        """``codes`` [n, frames, NQ], ``slots``/``pos`` [n] on device -> waveform view [n, frames * spf]."""
+        n = int(codes.shape[0])
+        size = next((b for b in self.sizes if b >= n), None)
+        if size is None:
+            return self.sd(codes, slots.to(torch.int32), pos.to(torch.int32))
+        graph, s_codes, s_slots, s_pos, out = self.graphs[size]
+        s_codes[:n].copy_(codes)
+        s_slots[:n].copy_(slots)
+        s_pos[:n].copy_(pos)
+        if n < size:
+            s_slots[n:].fill_(self.sd.scratch_slot)
+            s_pos[n:].zero_()
+        graph.replay()
+        return out[:n]

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -390,16 +390,17 @@ def talker2code2wav_async_chunk(
         if ref_context_size > 0:
             if emitted_chunks <= 0:
                 ref_context_request_id = request_id
-                ref_frames = ref_context.tolist()
+                ref_frames = list(ref_context.to(device="cpu").unbind(0))
                 window_frames = ref_frames + window_frames
                 ref_context_included = True
                 left_context_size = ref_context_size
 
-    num_quantizers = len(window_frames[0])
-    num_frames = len(window_frames)
-    code_predictor_codes = torch.tensor(
-        [window_frames[f][q] for q in range(num_quantizers) for f in range(num_frames)],
-        dtype=torch.long,
+    # Keep the codebook-major CPU wire format without extracting every
+    # tensor element as a Python scalar. Legacy buffered frames can be lists.
+    code_predictor_codes = (
+        torch.stack([torch.as_tensor(frame, dtype=torch.long, device="cpu") for frame in window_frames])
+        .transpose(0, 1)
+        .reshape(-1)
     )
 
     meta = MetaStruct(

vllm_omni/worker_v2/model_states/eager_mtp.py

@@ -5,13 +5,17 @@
 from __future__ import annotations
 
 from collections.abc import Callable
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
+import numpy as np
 import torch
 from vllm.logger import init_logger
 from vllm.v1.worker.gpu.input_batch import InputBatch
+from vllm.v1.worker.gpu.states import RequestState
 
 from vllm_omni.utils.device_copy import index_to_device
+from vllm_omni.worker_v2.streaming_audio import StreamingAudioBuffer, StreamingAudioOutput
 
 if TYPE_CHECKING:
     from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
@@ -24,12 +28,175 @@ def _has_ref_codes(buffer: dict[str, Any]) -> bool:
     return isinstance(ref, torch.Tensor) and ref.numel() > 0
 
 
+def _snapshot_runtime(value: Any) -> Any:
+    if isinstance(value, torch.Tensor):
+        return value.to("cpu", copy=True)
+    if isinstance(value, dict):
+        return {key: _snapshot_runtime(item) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_snapshot_runtime(item) for item in value]
+    if isinstance(value, tuple):
+        return tuple(_snapshot_runtime(item) for item in value)
+    return value
+
+
+@dataclass
+class TalkerInputs:
+    embeds: torch.Tensor
+    length: int = 0
+
+
+@dataclass
+class SuspendedAudio:
+    position: int
+    decoder_state: list[torch.Tensor]
+    generator: torch.Generator | None
+    runtime: dict[str, Any]
+    frame_embed: torch.Tensor
+
+
 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
+        self._side_stream: torch.cuda.Stream | None = None
+        self._fast_ok: bool | None = None
+        self._audio_buffer: StreamingAudioBuffer | None = None
+        self._frame_requests: set[str] = set()
+        self._suspended_audio: dict[str, SuspendedAudio] = {}
+        self._restore_audio: dict[str, list[torch.Tensor]] = {}
+        self._talker_inputs: dict[str, TalkerInputs] = {}
+
+    def suspend_audio(self, req_id: str, req_idx: int) -> None:
+        decoder = getattr(self.owner.model, "stream_decoder", None)
+        position = self.owner._stream_pos.get(req_id)
+        if decoder is None or position is None:
+            return
+        # Slot reuse is rare and must wait for the last side-stream write.
+        saved = self._restore_audio.pop(req_id, None)
+        if saved is None:
+            self._decode_stream.synchronize()
+            saved = decoder.save_slot(req_idx)
+        history = self._talker_inputs[req_id]
+        history.embeds = history.embeds[: history.length].to("cpu", copy=True)
+        runtime = self.owner.intermediate_buffer.buffers[req_idx]
+        assert self.owner._eager_embeds is not None
+        self._suspended_audio[req_id] = SuspendedAudio(
+            position,
+            saved,
+            self.owner._mtp_generators.get(req_id),
+            _snapshot_runtime(runtime),
+            self.owner._eager_embeds[req_idx].to("cpu", copy=True),
+        )
+
+    def resume_audio(self, req_id: str, req_idx: int) -> None:
+        saved = self._suspended_audio.pop(req_id, None)
+        if saved is not None:
+            self.owner._stream_pos[req_id] = saved.position
+            self.owner._mtp_generators[req_id] = saved.generator
+            self._restore_audio[req_id] = saved.decoder_state
+            self.owner.intermediate_buffer.buffers[req_idx] = saved.runtime
+            assert self.owner._eager_embeds is not None
+            self.owner._eager_embeds[req_idx].copy_(saved.frame_embed)
+            self.owner._eager_ready[req_idx] = req_id
+
+    def finish_audio(self, req_ids: set[str]) -> None:
+        for req_id in req_ids:
+            self._suspended_audio.pop(req_id, None)
+            self._restore_audio.pop(req_id, None)
+            self.owner._stream_pos.pop(req_id, None)
+            self._talker_inputs.pop(req_id, None)
+        if self._audio_buffer is not None:
+            self._audio_buffer.finish(req_ids)
+
+    def replay_inputs(self, req_id: str, req_idx: int, offset: int, ids: torch.Tensor, embeds: torch.Tensor) -> bool:
+        """Rebuild freed Talker KV from the exact conditioned inputs, without rerunning MTP."""
+        if req_id not in self._restore_audio:
+            return False
+        history = self._talker_inputs[req_id]
+        count = min(len(ids), max(0, history.length - offset))
+        if count:
+            embeds[:count].copy_(history.embeds[offset : offset + count])
+        if count < len(ids):
+            # Recompute ends with the last accepted CB0, whose full eager frame
+            # was saved but has not yet been consumed as a Talker input.
+            if len(ids) - count != 1 or offset + count != history.length:
+                raise RuntimeError(f"Talker replay exceeds saved inputs for {req_id}")
+            info = self.owner.intermediate_buffer.buffers[req_idx]
+            _, _, updates = self.owner.model.preprocess(ids[count:], embeds[count:], **info, _omni_is_prefill=False)
+            mtp_inputs = updates.pop("mtp_inputs", None)
+            if mtp_inputs is None:
+                raise RuntimeError(f"Talker replay has no text conditioning for {req_id}")
+            assert self.owner._eager_embeds is not None
+            embeds[count:].copy_(self.owner._eager_embeds[req_idx] + mtp_inputs[1].reshape(1, -1))
+            self.owner.intermediate_buffer.update(req_idx, updates, self.owner.model.gpu_resident_buffer_keys)
+        return True
+
+    def record_inputs(self, input_batch: InputBatch, embeds: torch.Tensor) -> None:
+        if getattr(self.owner.model, "stream_decoder", None) is None:
+            return
+        history_slices: list[torch.Tensor] = []
+        input_slices: list[torch.Tensor] = []
+        for i, req_id in enumerate(input_batch.req_ids):
+            start, end = input_batch.query_start_loc_np[i : i + 2]
+            offset = int(input_batch.num_computed_tokens_np[i])
+            length = offset + int(end - start)
+            history = self._talker_inputs.get(req_id)
+            if history is None or history.embeds.device != embeds.device:
+                storage = torch.empty(
+                    self.owner.vllm_config.model_config.max_model_len,
+                    embeds.shape[-1],
+                    device=embeds.device,
+                    dtype=embeds.dtype,
+                )
+                if history is not None:
+                    storage[: history.length].copy_(history.embeds)
+                history = TalkerInputs(storage, history.length if history else 0)
+                self._talker_inputs[req_id] = history
+            history_slices.append(history.embeds[offset:length])
+            input_slices.append(embeds[start:end])
+            history.length = max(history.length, length)
+        if history_slices:
+            # Keep exact inputs for replay without launching one copy per row.
+            torch._foreach_copy_(history_slices, input_slices)
+
+    def prepare_audio_output(
+        self, input_batch: InputBatch, req_states: RequestState, outputs: dict[str, Any]
+    ) -> StreamingAudioOutput | None:
+        if self._audio_buffer is None:
+            return None
+        wav = outputs.pop("model_outputs", None)
+        if not isinstance(wav, torch.Tensor):
+            return None
+        n = input_batch.num_reqs
+        total_after = input_batch.num_computed_tokens_np[:n] + np.asarray(input_batch.num_scheduled_tokens[:n]) + 1
+        limits = np.minimum(
+            req_states.max_seq_len[input_batch.idx_mapping_np[:n]], self.owner.vllm_config.model_config.max_model_len
+        )
+        requests = [
+            self._audio_buffer.requests.get(rid) if rid in self._frame_requests else None for rid in input_batch.req_ids
+        ]
+        self._frame_requests.clear()
+        event = self.owner._stream_decode_event
+        self.owner._stream_decode_event = None
+        return StreamingAudioOutput(
+            wav,
+            outputs["meta"]["codec_frame_valid"],
+            outputs["sr"][0],
+            event,
+            input_batch.query_start_loc_np[: n + 1].copy(),
+            requests,
+            total_after >= limits,
+        )
+
+    @property
+    def _decode_stream(self) -> torch.cuda.Stream:
+        """Side stream for Talker stream decode, so the codec overlaps the next step."""
+        if self._side_stream is None:
+            self._side_stream = torch.cuda.Stream()
+        return self._side_stream
 
     def set_first_audio_sink(self, sink: Any) -> None:
         from vllm_omni.worker_v2.first_audio_sender import FirstAudioSender
@@ -106,6 +273,14 @@ class EagerMTPState:
 
         bsz = len(entries)
         device = text_hidden.device
+        if self._fast_path_ok(codes_out, text_hidden):
+            self._run_eager_mtp_fused(
+                input_batch, text_hidden, sampled_token_ids, multimodal_outputs, entries, input_ids,
+                codes_out, valid_out, mtp_batch_descriptor_dispatcher,
+            )  # fmt: skip
+            return
+        if getattr(self.owner.model, "stream_decoder", None) is not None:
+            raise RuntimeError("Talker stream decode requires the fused eager-MTP path")
         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
@@ -162,6 +337,168 @@ class EagerMTPState:
                 scheduled *= self._first_audio_valid.index_select(0, rows)
             first_audio.index_copy_(0, last_tokens, scheduled.to(first_audio.dtype))
 
+    def _embedding_weight(self) -> torch.Tensor | None:
+        inner = getattr(self.owner.model, "model", None)
+        embed = getattr(inner, "embed_tokens", None)
+        weight = getattr(embed, "weight", None)
+        method = getattr(embed, "quant_method", None)
+        unquantized = method is None or type(method).__name__ == "UnquantizedEmbeddingMethod"
+        tp = getattr(embed, "tp_size", 1)
+        if isinstance(weight, torch.Tensor) and weight.ndim == 2 and unquantized and tp == 1:
+            return weight
+        return None
+
+    def _fast_path_ok(self, codes_out: torch.Tensor, text_hidden: torch.Tensor) -> bool:
+        if self._fast_ok is None:
+            self._fast_ok = (
+                text_hidden.is_cuda
+                and not self.owner.vllm_config.cache_config.enable_prefix_caching
+                and self._embedding_weight() is not None
+                and hasattr(self.owner.model, "_codebook_vocab_size")
+            )
+        return self._fast_ok and codes_out.ndim == 2 and text_hidden.is_contiguous()
+
+    def _run_eager_mtp_fused(
+        self,
+        input_batch: InputBatch,
+        text_hidden: torch.Tensor,
+        sampled_token_ids: torch.Tensor,
+        multimodal_outputs: dict[str, Any],
+        entries: list[tuple[int, int, str, bool]],
+        input_ids: torch.Tensor,
+        codes_out: torch.Tensor,
+        valid_out: torch.Tensor,
+        mtp_batch_descriptor_dispatcher: Callable[[int], Any] | None,
+    ) -> None:
+        from vllm_omni.worker_v2.model_states.eager_mtp_kernels import eager_post, eager_pre
+
+        owner = self.owner
+        bsz = len(entries)
+        device = text_hidden.device
+        first_audio = multimodal_outputs.get("meta", {}).get("first_audio")
+        fa_requests = owner._first_audio_requests
+        model = owner.model
+        stream = getattr(model, "stream_decoder", None)
+        stream_out = multimodal_outputs.get("model_outputs") if stream is not None else None
+        pos_list: list[int] = []
+        first_rows: list[int] = []
+        primes: list[tuple[int, torch.Tensor]] = []
+        if isinstance(stream_out, torch.Tensor):
+            assert stream is not None
+            # Frame index of each row's frame in its request's stream; 0 starts
+            # a fresh decoder state in the request's slot. A voice-clone
+            # request's stream starts after its reference codes, which prime
+            # the decoder state first (as Code2Wav's first-chunk context does).
+            if self._audio_buffer is None:
+                self._audio_buffer = StreamingAudioBuffer(int(stream.spf), int(model.stream_chunk_frames))
+            self._frame_requests = {req_id for _i, _idx, req_id, _p in entries}
+            positions = owner._stream_pos
+            for row, (_i, req_idx, req_id, _p) in enumerate(entries):
+                self._audio_buffer.add(req_id)
+                p = positions.get(req_id)
+                if p is None:
+                    first_rows.append(row)
+                    ref = model.get_stream_ref_context(owner.intermediate_buffer.buffers[req_idx])
+                    p = 0
+                    if ref is not None:
+                        primes.append((req_idx, ref))
+                        p = int(ref.shape[0])
+                pos_list.append(p)
+                positions[req_id] = p + 1
+        meta_list = (
+            [i for i, _r, _q, _p in entries]
+            + [r for _i, r, _q, _p in entries]
+            + [int(p) for _i, _r, _q, p in entries]
+            + [int(q in fa_requests) for _i, _r, q, _p in entries]
+            + pos_list
+        )
+        meta = index_to_device(meta_list, device)
+        rows = meta[bsz : 2 * bsz]
+        sampled = sampled_token_ids.reshape(input_batch.num_reqs, -1)
+        assert owner._mtp_input_ids is not None and owner._mtp_input_embeds is not None
+        assert owner._mtp_hidden is not None and owner._mtp_text_step is not None
+        batch_ids = owner._mtp_input_ids[:bsz]
+        batch_emb = owner._mtp_input_embeds[:bsz]
+        batch_hidden = owner._mtp_hidden[:bsz]
+        batch_step = owner._mtp_text_step[:bsz]
+        last_tokens, layer0 = eager_pre(
+            meta, bsz, input_batch.query_start_loc, sampled, text_hidden, self._embedding_weight(),
+            batch_ids, batch_emb, batch_hidden, batch_step,
+        )  # fmt: skip
+        req_indices = meta_list[bsz : 2 * bsz]
+        frame_embeds, codes = owner._mtp_forward(
+            req_indices, batch_ids, batch_emb, batch_hidden, batch_step, mtp_batch_descriptor_dispatcher
+        )
+        assert codes is not None
+        has_prefill = any(p for _i, _r, _q, p in entries)
+        fa_in_kernel = isinstance(first_audio, torch.Tensor) and not has_prefill
+        valid = eager_post(
+            meta, bsz, last_tokens, layer0, frame_embeds.reshape(frame_embeds.shape[0], -1).contiguous(),
+            owner._eager_embeds, codes.contiguous(), codes_out, input_ids, valid_out,
+            first_audio if fa_in_kernel else None, self._first_audio_valid if fa_in_kernel else None,
+            owner.model._codebook_vocab_size,
+        )  # fmt: skip
+        for _i, req_idx, req_id, _prefill in entries:
+            owner._eager_ready[req_idx] = req_id
+        if pos_list:
+            assert stream is not None
+            assert isinstance(stream_out, torch.Tensor)
+            decode = model.stream_graphs if model.stream_graphs is not None else stream
+            frame_codes = codes[:bsz].reshape(bsz, 1, -1).to(torch.int32)
+            # The next Talker step does not read the PCM, so the codec runs
+            # beside it; only the output copies wait for it.
+            side = self._decode_stream
+            side.wait_stream(torch.cuda.current_stream())
+            for t in (frame_codes, meta, last_tokens, valid, stream_out):
+                t.record_stream(side)
+            with torch.cuda.stream(side):
+                for _i, req_idx, req_id, _p in entries:
+                    saved = self._restore_audio.pop(req_id, None)
+                    if saved is not None:
+                        stream.restore_slot(req_idx, saved)
+                model.prime_stream_decoder(primes)
+                pcm = decode(frame_codes, rows, meta[4 * bsz : 5 * bsz])
+                stream_out.index_copy_(0, last_tokens, pcm.reshape(bsz, -1).to(stream_out.dtype))
+                self._send_stream_first_frames(entries, pcm, valid, first_rows)
+                done = torch.cuda.Event()
+                done.record(side)
+            owner._stream_decode_event = done
+        if has_prefill:
+            self._publish_first_audio(input_batch, entries, codes[:bsz], valid)
+            if isinstance(first_audio, torch.Tensor):
+                # Requests accepted by this step's publish count as scheduled.
+                scheduled = index_to_device(
+                    [int(q in owner._first_audio_requests) for _i, _r, q, _p 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 _send_stream_first_frames(
+        self,
+        entries: list[tuple[int, int, str, bool]],
+        pcm: torch.Tensor,
+        valid: torch.Tensor,
+        first_rows: list[int],
+    ) -> None:
+        sender = self.owner._first_audio_sender
+        if sender is None or not first_rows:
+            return
+        rows = index_to_device(first_rows, pcm.device)
+        request_ids = [entries[row][2] for row in first_rows]
+        sample_rate = torch.tensor(int(self.owner.model.stream_sample_rate), dtype=torch.int32)
+        accepted = sender.submit(
+            request_ids,
+            pcm.reshape(pcm.shape[0], -1).index_select(0, rows).float(),
+            sample_rate,
+            valid=valid.index_select(0, rows),
+        )
+        assert self._audio_buffer is not None
+        for request_id in accepted:
+            self._audio_buffer.requests[request_id].accept_first_audio()
+
     def _publish_first_audio(
         self,
         input_batch: InputBatch,

vllm_omni/worker_v2/model_states/eager_mtp_kernels.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+# ruff: noqa: N803 - Triton constexpr parameters use kernel-style capitals.
+"""Gather/scatter kernels around the eager Talker-MTP call.
+
+Each engine step, the rows sampled this step need their CB0, codec embedding
+and last hidden state gathered into the MTP inputs, and the MTP results
+scattered back into the eager-frame embeddings and the retained outputs. As
+separate torch ops that is a few dozen small launches per step, and the engine
+loop is host-bound, so each costs the whole batch. These two kernels do the
+same indexing in one launch each; the row metadata arrives in one H2D copy.
+"""
+
+from __future__ import annotations
+
+import torch
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _pre_kernel(
+    meta_ptr,  # int64 [4, n]: batch row, request index, prefill flag, first-audio scheduled
+    n,
+    qsl_ptr,  # query_start_loc
+    sampled_ptr,  # [num_reqs, s_stride]
+    s_stride,
+    hidden_ptr,  # [T, H]
+    emb_w_ptr,  # [V, H]
+    ids_ptr,  # [n] MTP input ids
+    emb_ptr,  # [n, H]
+    mtp_hidden_ptr,  # [n, H]
+    step_ptr,  # [n, H]
+    last_ptr,  # int64 [n] out
+    layer0_ptr,  # int64 [n] out
+    H: tl.constexpr,
+    BH: tl.constexpr,
+):
+    b = tl.program_id(0)
+    row = tl.load(meta_ptr + b)
+    last = tl.load(qsl_ptr + row + 1).to(tl.int64) - 1
+    l0 = tl.load(sampled_ptr + row * s_stride).to(tl.int64)
+    tl.store(ids_ptr + b, l0.to(ids_ptr.dtype.element_ty))
+    tl.store(last_ptr + b, last)
+    tl.store(layer0_ptr + b, l0)
+    for h0 in tl.static_range(0, H, BH):
+        c = h0 + tl.arange(0, BH)
+        m = c < H
+        tl.store(emb_ptr + b * H + c, tl.load(emb_w_ptr + l0 * H + c, mask=m).to(emb_ptr.dtype.element_ty), mask=m)
+        tl.store(mtp_hidden_ptr + b * H + c, tl.load(hidden_ptr + last * H + c, mask=m), mask=m)
+        tl.store(step_ptr + b * H + c, tl.zeros([BH], dtype=step_ptr.dtype.element_ty), mask=m)
+
+
+@triton.jit
+def _post_kernel(
+    meta_ptr,
+    n,
+    last_ptr,
+    layer0_ptr,
+    frame_ptr,  # [>=n, H] MTP frame embeddings
+    eager_ptr,  # [R, H]
+    codes_ptr,  # [>=n, Q]
+    codes_out_ptr,  # [T, Q]
+    input_ids_ptr,  # [T] this step's input ids
+    valid_out_ptr,  # [T]
+    valid_ptr,  # [n] bool out
+    fa_ptr,  # [T] first_audio output or dummy
+    fa_valid_ptr,  # [R] bool or dummy
+    vocab,
+    H: tl.constexpr,
+    Q: tl.constexpr,
+    QP: tl.constexpr,
+    BH: tl.constexpr,
+    WRITE_FA: tl.constexpr,
+    HAS_FA_VALID: tl.constexpr,
+):
+    b = tl.program_id(0)
+    req = tl.load(meta_ptr + n + b)
+    prefill = tl.load(meta_ptr + 2 * n + b) != 0
+    last = tl.load(last_ptr + b)
+    l0 = tl.load(layer0_ptr + b)
+    for h0 in tl.static_range(0, H, BH):
+        c = h0 + tl.arange(0, BH)
+        m = c < H
+        v = tl.load(frame_ptr + b * H + c, mask=m)
+        tl.store(eager_ptr + req * H + c, v.to(eager_ptr.dtype.element_ty), mask=m)
+    q = tl.arange(0, QP)
+    qm = q < Q
+    cv = tl.load(codes_ptr + b * Q + q, mask=qm)
+    tl.store(codes_out_ptr + last * Q + q, cv.to(codes_out_ptr.dtype.element_ty), mask=qm)
+    inp = tl.load(input_ids_ptr + last).to(tl.int64)
+    in_valid = ((inp >= 0) & (inp < vocab)) | prefill
+    valid = (l0 >= 0) & (l0 < vocab) & in_valid
+    tl.store(valid_out_ptr + last, valid.to(valid_out_ptr.dtype.element_ty))
+    tl.store(valid_ptr + b, valid)
+    if WRITE_FA:
+        scheduled = tl.load(meta_ptr + 3 * n + b) != 0
+        if HAS_FA_VALID:
+            scheduled = scheduled & (tl.load(fa_valid_ptr + req) != 0)
+        else:
+            scheduled = scheduled & False
+        tl.store(fa_ptr + last, scheduled.to(fa_ptr.dtype.element_ty))
+
+
+def eager_pre(meta, n, qsl, sampled, hidden, emb_w, ids, emb, mtp_hidden, step):
+    last = torch.empty(n, dtype=torch.long, device=hidden.device)
+    layer0 = torch.empty(n, dtype=torch.long, device=hidden.device)
+    H = hidden.shape[-1]
+    _pre_kernel[(n,)](
+        meta, n, qsl, sampled, sampled.stride(0), hidden, emb_w, ids, emb, mtp_hidden, step, last, layer0,
+        H=H, BH=min(1024, triton.next_power_of_2(H)),
+    )  # fmt: skip
+    return last, layer0
+
+
+def eager_post(meta, n, last, layer0, frames, eager, codes, codes_out, input_ids, valid_out, fa, fa_valid, vocab):
+    valid = torch.empty(n, dtype=torch.bool, device=last.device)
+    H = eager.shape[-1]
+    Q = codes_out.shape[-1]
+    dummy = valid
+    _post_kernel[(n,)](
+        meta, n, last, layer0, frames, eager, codes, codes_out, input_ids, valid_out, valid,
+        fa if fa is not None else dummy, fa_valid if fa_valid is not None else dummy, int(vocab),
+        H=H, Q=Q, QP=triton.next_power_of_2(Q), BH=min(1024, triton.next_power_of_2(H)),
+        WRITE_FA=fa is not None, HAS_FA_VALID=fa_valid is not None,
+    )  # fmt: skip
+    return valid

vllm_omni/worker_v2/model_states/omni_model_state.py

@@ -41,6 +41,7 @@ from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
 from vllm_omni.worker_v2.model_states.intermediate_buffer import (
     OmniIntermediateBuffer,
 )
+from vllm_omni.worker_v2.streaming_audio import StreamingAudioOutput
 
 logger = init_logger(__name__)
 
@@ -85,6 +86,9 @@ def _make_safe_get_rope(orig_get_rope):
     return _safe_get_rope
 
 
+_NO_SEED = object()
+
+
 class OmniModelState(DefaultModelState):
     """Generic Omni ``ModelState`` — works for **all** Omni model stages.
 
@@ -146,7 +150,7 @@ class OmniModelState(DefaultModelState):
         self.have_multimodal_outputs: bool = getattr(model, "have_multimodal_outputs", False)
         self._decode_preprocess = self._resolve_decode_preprocess(model)
         self._decode_preprocess_is_identity = bool(getattr(model, "mrv2_decode_preprocess_is_identity", False))
-        self._mtp_generators: dict[str, torch.Generator] = {}
+        self._mtp_generators: dict[str, torch.Generator | None] = {}
         # Talker's codec_embedding dim may differ from hf_text_config.hidden_size; probe real dim.
         self._embed_dim = self._get_embed_dim(model, device) if self.has_preprocess else 0
 
@@ -167,6 +171,10 @@ class OmniModelState(DefaultModelState):
         self._mtp_offsets: torch.Tensor | None = None
         self._mtp_sample_uniforms: torch.Tensor | None = None
         self._mtp_runner: Any | None = None
+        # Talker stream decode: each request's next frame position, and the
+        # side-stream event the step's PCM output copy must wait on.
+        self._stream_pos: dict[str, int] = {}
+        self._stream_decode_event: torch.cuda.Event | None = None
         if self._embed_dim > 0 and hasattr(model, "mtp"):
             max_bs = max_num_reqs
             self._mtp_input_ids = torch.zeros(max_bs, dtype=torch.long, device=device)
@@ -218,6 +226,17 @@ class OmniModelState(DefaultModelState):
                 input_batch, text_hidden, sampled_token_ids, multimodal_outputs, mtp_batch_descriptor_dispatcher
             )
 
+    def prepare_streaming_audio_output(
+        self, input_batch: InputBatch, req_states: RequestState, outputs: dict[str, Any]
+    ) -> StreamingAudioOutput | None:
+        return self._eager_state.prepare_audio_output(input_batch, req_states, outputs)
+
+    def on_requests_finished(self, req_ids: set[str]) -> None:
+        self._eager_state.finish_audio(req_ids)
+
+    def on_request_preempted(self, req_id: str, req_index: int) -> None:
+        self._eager_state.suspend_audio(req_id, req_index)
+
     def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None:
         """Let the model replace or wrap the MRv2 sampler (``mrv2_custom_sampler``)."""
         hook = getattr(self.model, "mrv2_custom_sampler", None)
@@ -372,6 +391,7 @@ class OmniModelState(DefaultModelState):
     def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
         super().add_request(req_index, new_req_data)
         self.intermediate_buffer.add_request(req_index, new_req_data)
+        self._eager_state.resume_audio(new_req_data.req_id, req_index)
         self._initialize_upstream_warmup_buffer(req_index, new_req_data.req_id)
 
     def _initialize_upstream_warmup_buffer(self, req_index: int, req_id: str) -> None:
@@ -407,6 +427,7 @@ 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)
+            self._stream_pos.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)
@@ -673,6 +694,7 @@ class OmniModelState(DefaultModelState):
         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]] = []
+        replay_rows: list[tuple[int, int, str, bool]] = []
         # Declared identity decode has no per-row updates or hook work.
         skip_decode_rows = self._decode_preprocess_is_identity and not self._eager_mtp
         is_prefilling_np = getattr(input_batch, "is_prefilling_np", None) if skip_decode_rows else None
@@ -684,6 +706,20 @@ class OmniModelState(DefaultModelState):
                 continue
             start = int(input_batch.query_start_loc_np[i])
             n_tok = int(input_batch.num_scheduled_tokens[i])
+            if (
+                self._eager_mtp
+                and self.model.stream_decoder is not None
+                and self._eager_state.replay_inputs(
+                    buf["req_id"],
+                    req_idx,
+                    int(input_batch.num_computed_tokens_np[i]),
+                    input_ids[start : start + n_tok],
+                    embeds[start : start + n_tok],
+                )
+            ):
+                if input_batch.num_computed_prefill_tokens_np[i] + n_tok >= input_batch.prefill_len_np[i]:
+                    replay_rows.append((i, req_idx, buf["req_id"], True))
+                continue
             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
@@ -725,7 +761,8 @@ class OmniModelState(DefaultModelState):
                     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],
+                + [(i, req_idx, str(req_id), False) for i, req_idx, _start, req_id in settled_rows]
+                + replay_rows,
                 input_ids,
             )
 
@@ -760,7 +797,7 @@ class OmniModelState(DefaultModelState):
                 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_offsets = index_to_device(starts, input_ids.device)
                 batch_ids = input_ids.index_select(0, batch_offsets)
                 batch_embeds = embeds.index_select(0, batch_offsets)
 
@@ -866,6 +903,9 @@ class OmniModelState(DefaultModelState):
                     prepacked_mtp_inputs=prepacked_mtp_inputs,
                 )
 
+        if self._eager_mtp:
+            self._eager_state.record_inputs(input_batch, embeds)
+
     def publish_sampled_embeddings(
         self, input_batch: InputBatch, sampled_token_ids: torch.Tensor
     ) -> tuple[dict[str, Any], torch.cuda.Event] | None:
@@ -989,7 +1029,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 index_to_device([start for _i, start, _mtp in mtp_batches], device)
 
     def _run_batched_mtp(
         self,
@@ -1180,15 +1220,21 @@ class OmniModelState(DefaultModelState):
         sampling_params: Any,
         device: torch.device,
     ) -> torch.Generator | None:
-        resolve_seed = getattr(self.model, "get_mtp_seed", None)
-        seed = resolve_seed(sampling_params) if resolve_seed is not None else None
-        if seed is None:
-            return None
-
         cache = getattr(self, "_mtp_generators", None)
         if cache is None:
             cache = {}
             self._mtp_generators = cache
+        generator = cache.get(req_id, _NO_SEED)
+        if generator is None:
+            # Resolved earlier: this request has no seed.
+            return None
+        if generator is not _NO_SEED and generator.device == device:
+            return generator
+        resolve_seed = getattr(self.model, "get_mtp_seed", None)
+        seed = resolve_seed(sampling_params) if resolve_seed is not None else None
+        if seed is None:
+            cache[req_id] = None
+            return None
         generator = cache.get(req_id)
         if generator is None or generator.device != device:
             generator = torch.Generator(device=device)

vllm_omni/worker_v2/omni_ar_model_runner.py

@@ -14,6 +14,7 @@ Extends ``OmniGPUModelRunner`` with:
 from __future__ import annotations
 
 from contextlib import nullcontext
+from functools import partial
 from typing import Any, cast
 
 import numpy as np
@@ -37,6 +38,7 @@ from vllm_omni.utils.mm_outputs import partition_flat_payload
 from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
 from vllm_omni.worker_v2.omni_sampler import sample_with_output
 from vllm_omni.worker_v2.output_snapshot import PackedOutputSnapshot, RequestOutputSnapshot, pack_output_snapshot
+from vllm_omni.worker_v2.streaming_audio import StreamingAudioOutput
 
 logger = init_logger(__name__)
 _ASYNC_MM_SNAPSHOT_MAX_BUCKETS_PER_SLOT = 64
@@ -250,6 +252,13 @@ class OmniARModelRunner(OmniGPUModelRunner):
         model_runner_output.kv_extracted_req_ids = kv_extracted
         model_runner_output._async_chunk = bool(getattr(self.model_config, "async_chunk", False))
 
+        prepare_streaming = getattr(self.model_state, "prepare_streaming_audio_output", None)
+        streaming_audio = (
+            prepare_streaming(input_batch, self.req_states, multimodal_outputs)
+            if need_pooler and multimodal_outputs and callable(prepare_streaming)
+            else None
+        )
+
         # --- Async D2H via OmniAsyncOutput ---
         materialize_native = self._uses_native_output_materializer()
         async_output = OmniAsyncOutput(
@@ -266,6 +275,7 @@ class OmniARModelRunner(OmniGPUModelRunner):
             finalize_multimodal=sampling_output.finalize_multimodal,
             check_ep_fault=self.check_ep_fault,
             routed_experts=routed_experts,
+            streaming_audio=streaming_audio,
             extra_multimodal_outputs=extra_outputs,
         )
         self._release_multimodal_snapshot(snapshot_slot, async_output.copy_event)
@@ -297,6 +307,10 @@ class OmniARModelRunner(OmniGPUModelRunner):
     def _retain_multimodal_outputs(self, outputs: dict[str, Any]) -> dict[str, Any]:
         if not bool(getattr(self.model_config, "async_chunk", False)) or not outputs:
             return outputs
+        if getattr(self.model, "mm_outputs_fresh_per_step", False):
+            # Freshly allocated each step and filled after sampling; nothing
+            # overwrites them before the host copy.
+            return outputs
         slot_index = self._async_mm_snapshot_cursor
         if self._async_mm_snapshot_pending[slot_index]:
             raise RuntimeError("Async multimodal snapshot ring exhausted before sample_tokens().")
@@ -516,6 +530,12 @@ def _async_copy_mm_value(
             for key, val in value.items()
         }
     if isinstance(value, list):
+        if len(value) > 1 and isinstance(value[0], torch.Tensor) and value[0].device.type == "cpu":
+            first = value[0]
+            if all(val is first for val in value):
+                # One shared per-step host tensor (e.g. every request's sample
+                # rate): one copy serves all entries.
+                return [_async_copy_tensor(first)] * len(value)
         return [
             _async_copy_mm_value(
                 val,
@@ -679,6 +699,7 @@ class OmniAsyncOutput(AsyncModelRunnerOutput):
         finalize_multimodal: Any | None = None,
         check_ep_fault: bool = False,
         routed_experts: RoutedExpertsTensors | None = None,
+        streaming_audio: StreamingAudioOutput | None = None,
         extra_multimodal_outputs: tuple[dict[str, Any], torch.cuda.Event] | None = None,
     ):
         self.model_runner_output = model_runner_output
@@ -768,17 +789,31 @@ class OmniAsyncOutput(AsyncModelRunnerOutput):
             self._hidden_cpu: torch.Tensor | None = None
             self._mm_cpu: dict[str, Any] = {}
             self._mm_snapshot: dict[str, Any] = {}
+            self._streaming_audio = (
+                streaming_audio.to_cpu(
+                    copy_stream, partial(_async_copy_tensor, copy_stream=copy_stream, pin_memory=pin_memory)
+                )
+                if streaming_audio is not None
+                else None
+            )
             if self._need_pooler and (self._async_chunk or self._finalize_multimodal is not None):
                 # CUDA graph replay reuses the model's output buffers. Take
                 # ownership directly in pinned host memory on the output copy
                 # stream so deferred finalization never performs a blocking
                 # D2H copy on the runner thread.
-                self._mm_snapshot = _async_copy_mm(
-                    multimodal_outputs,
-                    self._total_tokens,
-                    copy_stream=copy_stream,
-                    pin_memory=pin_memory,
-                )
+                if (
+                    streaming_audio is None
+                    or self._finalize_multimodal is not None
+                    or extra_multimodal_outputs is not None
+                ):
+                    # PCM owns its request partition. Generic codes/meta are
+                    # discarded unless a finalizer or extra payload needs them.
+                    self._mm_snapshot = _async_copy_mm(
+                        multimodal_outputs,
+                        self._total_tokens,
+                        copy_stream=copy_stream,
+                        pin_memory=pin_memory,
+                    )
                 if extra_multimodal_outputs:
                     # Produced after sampling on the producer stream; copy only
                     # once its completion event has been observed.
@@ -856,7 +891,9 @@ class OmniAsyncOutput(AsyncModelRunnerOutput):
                     pooler_client is not None and len(pooler_client) != self._num_reqs
                 ):
                     raise ValueError("Model-owned output snapshot does not match the request batch")
-            else:
+            elif self._streaming_audio is None:
+                # In-stage PCM already owns its request partition. Building
+                # generic code payloads here would immediately discard them.
                 pooler_inter, pooler_client = OmniARModelRunner._build_async_chunk_outputs_from_mm(
                     self._mm_snapshot,
                     self._query_start_loc_np,
@@ -865,6 +902,8 @@ class OmniAsyncOutput(AsyncModelRunnerOutput):
                     self._total_tokens,
                     self._padded_total_tokens,
                 )
+            if self._streaming_audio is not None:
+                pooler_inter, pooler_client = None, self._streaming_audio.get_output()
             self.model_runner_output.pooler_output = None if self._async_chunk else pooler_inter
             self.model_runner_output.inter_stage_outputs = pooler_inter
             self.model_runner_output.multimodal_outputs = (

vllm_omni/worker_v2/omni_model_runner.py

@@ -437,6 +437,10 @@ class OmniGPUModelRunner(GPUModelRunner):
         capture_first_frame = getattr(self.model, "capture_first_frame_graphs", None)
         if callable(capture_first_frame):
             capture_first_frame()
+        capture_stream = getattr(self.model, "capture_stream_decode_graphs", None)
+        if callable(capture_stream):
+            sizes = [int(s) for s in self.model_state._get_mtp_capture_sizes()]
+            capture_stream(sorted(set(sizes + [1])))
         return result
 
     def _dispatch_mtp_batch_descriptor(self, num_mtp_reqs: int) -> Any:
@@ -746,7 +750,16 @@ class OmniGPUModelRunner(GPUModelRunner):
             # Request-owned state outlives generation slots and chunk boundaries.
             # Notify even if the last chunk already released its runner slot.
             on_finished(finished)
+        finish_state = getattr(self.model_state, "on_requests_finished", None)
+        if finished and callable(finish_state):
+            finish_state(finished)
         preempted = scheduler_output.preempted_req_ids
+        suspend_state = getattr(self.model_state, "on_request_preempted", None)
+        if preempted and callable(suspend_state):
+            for req_id in preempted - finished:
+                idx = self.req_states.req_id_to_index.get(req_id)
+                if idx is not None:
+                    suspend_state(req_id, idx)
         all_done = finished | preempted if preempted else finished
         for req_id in all_done:
             idx = self.req_states.req_id_to_index.get(req_id)

vllm_omni/worker_v2/streaming_audio.py

@@ -0,0 +1,154 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Request-owned PCM accumulation for a stage that decodes audio in place."""
+
+from __future__ import annotations
+
+import threading
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+
+from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
+
+
+@dataclass
+class AudioRequest:
+    samples_per_frame: int
+    chunk_frames: int
+    active: bool = True
+    emitted: bool = False
+    direct_first_pending: bool = False
+    first_audio_required: bool = False
+    pending: list[torch.Tensor | np.ndarray] = field(default_factory=list)
+    pending_samples: int = 0
+    lock: threading.Lock = field(default_factory=threading.Lock)
+
+    def accept_first_audio(self) -> None:
+        with self.lock:
+            if self.active:
+                self.direct_first_pending = True
+
+    def push(self, pcm: torch.Tensor | np.ndarray, ended: bool) -> torch.Tensor | None:
+        with self.lock:
+            if not self.active:
+                return None
+            size = pcm.numel() if isinstance(pcm, torch.Tensor) else pcm.size
+            if self.direct_first_pending:
+                self.direct_first_pending = False
+                if size:
+                    pcm = pcm[self.samples_per_frame :]
+                    size = max(0, size - self.samples_per_frame)
+                    self.emitted = True
+                    self.first_audio_required = True
+            if size:
+                self.pending.append(pcm)
+                self.pending_samples += size
+            frames = self.pending_samples // self.samples_per_frame
+            out = None
+            if frames and (ended or not self.emitted or frames >= self.chunk_frames):
+                if all(isinstance(part, np.ndarray) for part in self.pending):
+                    array = np.concatenate(self.pending) if len(self.pending) > 1 else self.pending[0]
+                    out = torch.from_numpy(array)
+                else:
+                    parts = [torch.from_numpy(part) if isinstance(part, np.ndarray) else part for part in self.pending]
+                    out = torch.cat(parts) if len(parts) > 1 else parts[0]
+                self.pending.clear()
+                self.pending_samples = 0
+                self.emitted = True
+            if ended:
+                self.active = False
+                self.pending.clear()
+                self.pending_samples = 0
+            return out
+
+    def finish(self) -> None:
+        with self.lock:
+            self.active = False
+            self.pending.clear()
+            self.pending_samples = 0
+
+
+class StreamingAudioBuffer:
+    def __init__(self, samples_per_frame: int, chunk_frames: int) -> None:
+        self.samples_per_frame = samples_per_frame
+        self.chunk_frames = chunk_frames
+        self.requests: dict[str, AudioRequest] = {}
+
+    def add(self, request_id: str) -> AudioRequest:
+        if request_id not in self.requests:
+            self.requests[request_id] = AudioRequest(self.samples_per_frame, self.chunk_frames)
+        return self.requests[request_id]
+
+    def finish(self, request_ids: Iterable[str]) -> None:
+        for request_id in request_ids:
+            state = self.requests.pop(request_id, None)
+            if state is not None:
+                # Already queued outputs retain this object, never a fresh lookup
+                # by ID. A late copy cannot resurrect a cancelled request.
+                state.finish()
+
+
+@dataclass
+class StreamingAudioOutput:
+    wav: torch.Tensor
+    valid: torch.Tensor
+    sample_rate: torch.Tensor
+    event: torch.cuda.Event | None
+    query_start_loc: np.ndarray
+    requests: list[AudioRequest | None]
+    length_end: np.ndarray
+
+    def to_cpu(
+        self,
+        copy_stream: torch.cuda.Stream,
+        copy_tensor: Callable[[torch.Tensor], torch.Tensor],
+    ) -> StreamingAudioOutput:
+        if self.event is not None:
+            copy_stream.wait_event(self.event)
+        return StreamingAudioOutput(
+            copy_tensor(self.wav),
+            copy_tensor(self.valid),
+            self.sample_rate,
+            None,
+            self.query_start_loc,
+            self.requests,
+            self.length_end,
+        )
+
+    def get_output(self) -> list[dict[str, torch.Tensor] | None]:
+        valid = self.valid.numpy().astype(bool)
+        # Each D2H batch owns its host allocation. NumPy slices retain that
+        # backing tensor while a request accumulates PCM across steps; no
+        # per-request Torch views are needed until an audio chunk is emitted.
+        # Keep the tensor path for dtypes NumPy cannot represent (e.g. BF16).
+        wav = self.wav.numpy() if self.wav.dtype in (torch.float16, torch.float32, torch.float64) else None
+        output: list[dict[str, torch.Tensor] | None] = []
+        required_marker = None
+        for i, state in enumerate(self.requests):
+            payload = None
+            start, end = self.query_start_loc[i : i + 2]
+            if state is not None and end > start:
+                if wav is not None:
+                    if end - start == 1:
+                        pcm = wav[start:end].reshape(-1) if valid[start] else wav[0:0].reshape(-1)
+                    else:
+                        pcm = wav[start:end][valid[start:end]].reshape(-1)
+                elif end - start == 1:
+                    pcm = self.wav[start:end].reshape(-1) if valid[start] else self.wav.new_empty(0)
+                else:
+                    pcm = self.wav[start:end][torch.from_numpy(valid[start:end])].reshape(-1)
+                ended = bool(self.length_end[i]) or not valid[end - 1]
+                chunk = state.push(pcm, ended)
+                if chunk is not None:
+                    payload = {"model_outputs": chunk, "sr": self.sample_rate}
+                if state.first_audio_required and (chunk is not None or ended):
+                    if required_marker is None:
+                        required_marker = torch.tensor(True)
+                    if payload is None:
+                        payload = {}
+                    payload[FIRST_AUDIO_REQUIRED_KEY] = required_marker
+            output.append(payload)
+        return output