[Core][Model] Enable Qwen3-TTS MRV2 and optimize the TTS pipeline (#7781)

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

.buildkite/common/ci_source_file_dependencies.yml

@@ -23,6 +23,16 @@ x-qwen3-tts: &qwen3_tts
   - vllm_omni/entrypoints/openai/serving_speech.py
   - vllm_omni/entrypoints/openai/tts_adapters/
   - vllm_omni/engine/stage_pool.py
+  - vllm_omni/worker_v2/
+  - vllm_omni/core/sched/
+  - vllm_omni/distributed/omni_connectors/
+  - vllm_omni/config/stage_config.py
+  - vllm_omni/config/omni_config.py
+  - vllm_omni/engine/orchestrator.py
+  - vllm_omni/engine/stage_engine_core_proc.py
+  - vllm_omni/deploy/qwen3_tts_mrv2.yaml
+  - vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml
+  - vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_b4.yaml
 
 x-minicpmo-4-5: &minicpmo_4_5
   - vllm_omni/engine/async_omni_engine.py

.buildkite/cuda/test-ready.yml

@@ -102,7 +102,7 @@ steps:
     commands:
       - |
         timeout 15m bash -c "
-                  pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py -m 'core_model and cuda' --run-level 'core_model'
+                  pytest -sv tests/engine/ tests/model_executor/ tests/worker/test_batched_omni_output.py tests/worker_v2/ -m 'core_model and cuda' --run-level 'core_model'
         "
     mirror_hardwares: l4_1
 

docs/configuration/stage_configs.md

@@ -290,3 +290,90 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --deploy-config /path/to/depl
 
 !!! important
     We are actively iterating on the definition of deployment configurations, and we welcome feedback from users and developers.
+
+## Qwen3-TTS with Model Runner V2
+
+Qwen3-TTS can opt into the native CUDA Model Runner V2 pipeline on vLLM
+0.29.0. Select one of these deployment profiles:
+
+| Profile | Runner | Code2Wav graph batches | Intended use |
+| --- | --- | --- | --- |
+| `qwen3_tts.yaml` | V1 | Existing defaults | Existing deployment / regression control |
+| `qwen3_tts_mrv2.yaml` | V2 | B1 | Native runner validation |
+| `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 |
+
+```bash
+vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \
+  --deploy-config vllm_omni/deploy/qwen3_tts_mrv2.yaml
+```
+
+The V2 profiles bound Talker prefill to 512 tokens per step and select the
+Talker AR runner and the Code2Wav generation
+runner together. Native inter-stage delivery carries codec payloads directly;
+request-owned snapshots preserve buffers through asynchronous completion and
+CUDA graph reuse. Terminal completion waits for upstream stage metrics before
+releasing request state. Platform sections retain V1 on NPU, XPU, ROCm and MUSA;
+this change does not qualify MRV2 on those backends or enable other model families.
+Selecting V2 for a stage without `supports_native_mrv2_data_plane` emits a warning
+when pipeline and deployment settings are merged. That stage retains the legacy
+transport path; the warning does not establish support for that combination.
+
+MRV2 model hooks use capability declarations rather than architecture or stage
+names. Omni lifecycle flags select the model state, and `_returns_tuple` declares
+the capture output contract. The optional batched predictor hook is `mtp`, with
+an explicit `mtp_output_key` (a string or two-part payload key), optional
+`mtp_validity_key`, and `mtp_graph_safe`/`mtp_disable_graph` capture controls.
+Its inputs are token IDs, embeddings, previous hidden states and per-row
+conditioning; it returns updated embeddings and prediction codes. Models may
+supply `mtp_sampling_params` and `get_mtp_seed(sampling_params)` for model-local explicit seeds, and declare
+`mtp_accepts_per_row_generators`, `mtp_accepts_req_infos`, or `mtp_sample_uniforms`
+(with `mtp_sample_steps` and `mtp_sample_vocab_size`) as needed. Qwen3-TTS retains
+its existing `talker_mtp` entry point for V1.
+
+### Included performance work
+
+- Request snapshots have a fast path for immutable scalar leaves, including
+  waveform lists. Mutable containers, aliases and cycles keep deep-copy
+  semantics. This is not the historical shallow `dict(prompt)` experiment.
+- The API reference cache holds owned float32 arrays with a default capacity
+  of 1024 entries and a 512 MiB waveform-payload budget. Cache hits return
+  independent lists. Qwen3-TTS model artifact caches also default to 1024
+  entries. These caches have different owners and eviction policies; equal
+  entry limits do not make them a single coherent cache.
+- Talker state stays on GPU where its lifecycle allows it. BOS/EOS projections
+  are cached in their projection dtype and invalidated on weight loading.
+- Code2Wav packs CPU codec inputs before device transfer. Optional B2/B4 graph
+  buckets batch compatible requests without sharing their per-request state.
+- Native output materialization runs in a bounded worker and drains before
+  closing the data plane. The existing control-only shortcut avoids launching
+  model work for a step that has only lifecycle events.
+
+MTP prefix re-prefill remains disabled in the high-concurrency V2 profile.
+Adding more graph shapes increases compilation cost and has not established a
+stable end-to-end gain for this extracted version.
+
+### Decoder batches and playback
+
+Changing `decode_cudagraph_batch_sizes` selects the captured batch buckets.
+Keep `decode_batch_max_size` consistent with the intended maximum too; the
+stateful graph path currently groups according to the captured buckets, while
+the stateless path also uses the explicit maximum.
+
+B4 remains experimental. First-packet latency alone does not establish
+uninterrupted playback. Validate inter-chunk arrival times, buffering, WER and
+speaker similarity before adopting either batching preset for a production
+workload. Floating-point decoder outputs can differ across batch sizes; this PR
+does not claim bitwise or quality equivalence.
+
+### Optional MPS deployment
+
+NVIDIA MPS is an optional operator setting for colocated CUDA processes, not a
+YAML option or a library default. This PR does not establish a throughput or
+first-packet latency benefit from MPS. Measure the exact deployment with and
+without MPS before enabling it.
+
+Use only assigned GPUs and an independent MPS pipe directory. A private MPS
+server does not provide exclusive GPU ownership or MIG isolation. For a
+single-GPU deployment, explicitly place both stages on that GPU; the supplied
+high-concurrency profile places its two stages on different GPUs by default.

tests/config/test_config_factory.py

@@ -1429,6 +1429,9 @@ stages:
         assert deploy.async_chunk is True
         assert deploy.connectors is not None
         assert deploy.platforms is not None
+        connector = deploy.connectors["connector_of_shared_memory"]
+        assert "async_chunk_batch_min_size" not in connector["extra"]
+        assert "async_chunk_batch_max_wait_ms" not in connector["extra"]
 
     @pytest.mark.parametrize(
         ("hf_config", "model"),
@@ -2704,6 +2707,56 @@ class TestPlatformOverrides:
 
         assert deploy.stages[0].engine_extras["block_size"] == 128
 
+    @pytest.mark.parametrize("platform", ["cuda", "npu"])
+    def test_recommended_native_runner_platform_defaults(self, platform):
+        filename, pipeline_key = "qwen3_tts_mrv2.yaml", "qwen3_tts"
+        deploy = load_deploy_config(Path(get_deploy_config_path(filename)))
+        deploy = _apply_platform_overrides(deploy, platform=platform)
+        expect_v2 = platform == "cuda"
+        assert deploy.model_runner == ("v2" if expect_v2 else "v1")
+
+        # The selection must reach the live worker-dispatch consumer, not just
+        # the transport-level DeployConfig field.
+        pipeline = resolve_pipeline_config(pipeline_key)
+        stages = merge_pipeline_deploy(pipeline, deploy)
+        assert [stage.yaml_engine_args["use_v2_model_runner"] for stage in stages] == [expect_v2] * len(stages)
+
+        if expect_v2 and pipeline.stages and deploy.async_chunk:
+            # v2 only engages the native plane on stages declaring support.
+            assert all(ps.supports_native_mrv2_data_plane for ps in pipeline.stages)
+
+    @pytest.mark.parametrize("runner,native", [("v1", False), ("v2", False), ("v2", True)])
+    def test_mrv2_undeclared_transport_warns(self, monkeypatch, runner, native):
+        from unittest.mock import Mock
+
+        warning = Mock()
+        monkeypatch.setattr("vllm_omni.config.stage_config.logger.warning", warning)
+        pipeline = PipelineConfig(
+            model_type="custom_pipeline",
+            stages=[
+                StagePipelineConfig(
+                    stage_id=0, model_stage="custom_ar", final_output=True, supports_native_mrv2_data_plane=native
+                )
+            ],
+        )
+        merge_pipeline_deploy(pipeline, DeployConfig(model_runner=runner))
+        assert warning.call_count == int(runner == "v2" and not native)
+        if warning.called:
+            assert "legacy transport path" in warning.call_args.args[0]
+
+    def test_runner_selection_rejects_engine_extras_override(self):
+        pipeline = PipelineConfig(
+            model_type="runner_selection_reserved",
+            stages=SINGLE_STAGE_PIPE_CFG,
+        )
+        for reserved in ("use_v2_model_runner", "supports_native_mrv2_data_plane"):
+            deploy = DeployConfig(
+                model_runner="v2",
+                stages=[StageDeployConfig(stage_id=0, engine_extras={reserved: False})],
+            )
+            with pytest.raises(ValueError, match=f"{reserved!r} must not be set"):
+                merge_pipeline_deploy(pipeline, deploy)
+
     def test_npu_overrides(self):
         deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
         if not deploy_path.exists():

tests/config/test_omni_config.py

@@ -50,8 +50,10 @@ from vllm_omni.config.stage_config import (
     StageDeployConfig,
     StageExecutionType,
     StagePipelineConfig,
+    _apply_platform_overrides,
     load_deploy_config,
     merge_pipeline_deploy,
+    resolve_deploy_yaml,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
 from vllm_omni.engine.stage_engine_startup import _serialize_stage_config
@@ -783,6 +785,67 @@ def test_joyai_code2wav_waits_for_full_payload():
     assert code2wav.model_config.requires_full_payload_input is True
 
 
+@pytest.mark.parametrize("model_runner", ["v1", "v2"])
+def test_deploy_model_runner_selection_propagates_to_every_stage(tmp_path: Path, model_runner: str):
+    deploy_path = tmp_path / "qwen3_tts_runner.yaml"
+    deploy_path.write_text(
+        f"""\
+pipeline: qwen3_tts
+model_runner: {model_runner}
+async_chunk: true
+stages:
+  - stage_id: 0
+  - stage_id: 1
+"""
+    )
+
+    deploy = load_deploy_config(deploy_path)
+    pipeline = _resolve_pipeline_or_skip("qwen3_tts")
+    legacy_stages = merge_pipeline_deploy(pipeline, deploy)
+    structured = _from_pipeline_key("qwen3_tts", deploy_config_path=str(deploy_path))
+    expect_v2 = model_runner == "v2"
+
+    assert deploy.model_runner == model_runner
+    assert all(stage.yaml_engine_args["use_v2_model_runner"] is expect_v2 for stage in legacy_stages)
+    assert all(stage.model_config.use_v2_model_runner is expect_v2 for stage in structured.stage_configs)
+
+
+@pytest.mark.parametrize("platform", ["npu", "xpu"])
+def test_mrv2_fails_fast_on_platforms_without_native_workers(platform: str):
+    with pytest.raises(NotImplementedError, match="Model Runner V2"):
+        _apply_platform_overrides(DeployConfig(model_runner="v2"), platform=platform)
+
+
+@pytest.mark.parametrize(
+    ("default_name", "mrv2_name"),
+    [
+        ("qwen3_tts.yaml", "qwen3_tts_mrv2.yaml"),
+        (
+            "qwen3_tts_high_concurrency.yaml",
+            "qwen3_tts_high_concurrency_mrv2.yaml",
+        ),
+    ],
+)
+def test_qwen3_mrv2_profiles_are_explicit_opt_in(default_name: str, mrv2_name: str):
+    assert load_deploy_config(_DEPLOY_DIR / default_name).model_runner == "v1"
+    assert load_deploy_config(_DEPLOY_DIR / mrv2_name).model_runner == "v2"
+
+
+def test_qwen3_tts_mrv2_retunes_do_not_change_default_mrv1_profile():
+    default = resolve_deploy_yaml(_DEPLOY_DIR / "qwen3_tts_high_concurrency.yaml")
+    mrv2 = resolve_deploy_yaml(_DEPLOY_DIR / "qwen3_tts_high_concurrency_mrv2.yaml")
+    default_extra = default["connectors"]["connector_of_shared_memory"]["extra"]
+    mrv2_extra = mrv2["connectors"]["connector_of_shared_memory"]["extra"]
+
+    # B2 batch stays consistent with the graph buckets; V1 profile untouched.
+    assert default_extra["decode_batch_max_size"] == 1
+    assert mrv2_extra["decode_batch_max_size"] == 2
+    assert mrv2_extra["decode_cudagraph_batch_sizes"] == [1, 2]
+    assert mrv2_extra["code_predictor_prefix_graphs"] is False
+    assert mrv2["stages"][1]["enforce_eager"] is False
+    assert default["stages"][1]["enforce_eager"] is True
+
+
 def test_vllm_omni_stage_config_public_fields_use_typed_stage_realizations():
     assert not hasattr(BaseVllmOmniStageConfig, "from_stage_config")
     assert not hasattr(BaseVllmOmniStageConfig, "to_legacy_stage_config")
@@ -879,9 +942,11 @@ def test_sub_config_fields_match_structured_scopes():
         "limit_mm_per_prompt",
         "interleave_mm_strings",
         "media_io_kwargs",
+        "final_output",
         "active_stream_window",
         "session_mode",
         "duplex_max_sessions",
+        "use_v2_model_runner",
         "enable_sleep_mode",
         "default_sampling_params",
         "subtalker_sampling_params",
@@ -890,6 +955,7 @@ def test_sub_config_fields_match_structured_scopes():
         "custom_voice_dir",
         "task_type",
         "codec_frame_rate_hz",
+        "supports_native_mrv2_data_plane",
         "enforce_eager",
         "max_cudagraph_capture_size",
         "enable_flashinfer_autotune",

tests/core/sched/test_generation_scheduler_restore.py

@@ -1,14 +1,22 @@
-"""Test that OmniGenerationScheduler restores chunk-waiting requests
-even when the OmniNewRequestData rewrapping fails.
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
-Regression test: if process_pending_chunks() moves requests into
-internal deques but restore_queues() is not called due to an exception,
-those requests are permanently orphaned.
-"""
+"""Chunk-lifecycle coverage for OmniGenerationScheduler: restore on error,
+native terminal state, in-flight no-resubmission, slot release on finish,
+and single scheduling of a terminal empty-prompt chunk."""
 
 from collections import deque
+from types import SimpleNamespace
 
 import pytest
+import torch
+from vllm import SamplingParams
+from vllm.v1.core.sched.interface import PauseState
+from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
+from vllm.v1.request import Request, RequestStatus
+
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler, _has_async_chunk_payload_to_run
+from vllm_omni.engine.serialization import serialize_additional_information
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -20,16 +28,20 @@ class FakeAdapter:
         self.waiting_for_chunk_waiting_requests = deque()
         self.waiting_for_chunk_running_requests = deque()
         self.restore_called = False
+        self.done_request_ids = set()
 
-    def process_pending_chunks(self, waiting, running):
-        """Simulate moving requests out of the scheduler queues."""
-        # Move one request from running into internal deque
+    def process_pending_chunks(self, waiting, running, scheduler_requests=None):
         if running:
             req = running.pop()
             self.waiting_for_chunk_running_requests.append(req)
 
+    def is_done_receiving_chunks(self, request_id):
+        return request_id in self.done_request_ids
+
+    def collect_failed_send_request_ids(self):
+        return {}
+
     def restore_queues(self, waiting, running, scheduler_requests=None):
-        """Put requests back."""
         self.restore_called = True
         running.extend(self.waiting_for_chunk_running_requests)
         self.waiting_for_chunk_running_requests = deque()
@@ -38,6 +50,131 @@ class FakeAdapter:
         pass
 
 
+def _make_generation_scheduler(waiting_request, *, use_v2_model_runner=False):
+    scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
+    scheduler.max_num_scheduled_tokens = 8
+    scheduler.max_num_running_reqs = 1
+    scheduler._pause_state = PauseState.UNPAUSED
+    scheduler.running = []
+    scheduler.waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.waiting.add_request(waiting_request)
+    scheduler.skipped_waiting = create_request_queue(SchedulingPolicy.FCFS)
+    scheduler.requests = {waiting_request.request_id: waiting_request}
+    scheduler.policy = SchedulingPolicy.FCFS
+    scheduler.chunk_transfer_adapter = FakeAdapter()
+    scheduler.input_coordinator = None
+    scheduler.log_stats = False
+    scheduler.scheduler_config = SimpleNamespace(enable_chunked_prefill=True)
+    scheduler.num_lookahead_tokens = 0
+    scheduler.kv_cache_manager = SimpleNamespace(
+        new_step_starts=lambda: None,
+        allocate_slots=lambda *args, **kwargs: SimpleNamespace(get_block_ids=lambda: ([1],)),
+        get_num_common_prefix_blocks=lambda request_id: [0],
+        take_new_block_ids=lambda: None,
+    )
+    scheduler.kv_cache_config = SimpleNamespace(kv_cache_groups=[object()])
+    scheduler.use_v2_model_runner = use_v2_model_runner
+    scheduler._retains_state_across_chunks = False
+    scheduler.needs_kv_cache_zeroing = False
+    scheduler.finished_req_ids = set()
+    scheduler.encoder_cache_manager = SimpleNamespace(get_freed_mm_hashes=lambda: [])
+    scheduler.connector = None
+    scheduler.ec_connector = None
+    scheduler.prev_step_scheduled_req_ids = set()
+    scheduler._pending_finish_reqs = []
+    scheduler._consume_pending_connector_output = lambda model_mode: None
+    scheduler._process_pending_input_timeouts = lambda: None
+    scheduler._make_cached_request_data = lambda **kwargs: SimpleNamespace(
+        req_ids=[],
+        resumed_req_ids=[],
+        new_token_ids=[],
+        all_token_ids=[],
+        new_block_ids=[],
+        num_computed_tokens=[],
+        num_output_tokens=[],
+    )
+    scheduler._update_after_schedule = lambda output: None
+    scheduler._wrap_omni_scheduler_output = lambda output: output
+    return scheduler
+
+
+def _chunk_request(request_id, **kwargs):
+    defaults = dict(
+        request_id=request_id,
+        prompt_token_ids=[1],
+        num_computed_tokens=0,
+        num_in_flight_tokens=0,
+        status=None,
+        sampling_params=None,
+        pooling_params=None,
+        mm_features=None,
+        lora_request=None,
+        prompt_is_token_ids=True,
+        additional_information=None,
+        external_req_id=request_id,
+        prefill_stats=None,
+        record_event=lambda *args, **kwargs: None,
+    )
+    defaults.update(kwargs)
+    return SimpleNamespace(**defaults)
+
+
+def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) -> None:
+    # Native plane: a completed payload without a new chunk is never executed
+    # again; requeued requests keep WAITING_FOR_CHUNK and their token counts.
+    scheduler = _make_generation_scheduler(_chunk_request("unused"))
+    monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.create_request_queue", create_request_queue)
+    scheduler._native_data_plane = True
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = SimpleNamespace(finished_requests=set(), restore_queues=lambda *args, **kwargs: None)
+    scheduler.running = []
+    scheduler.waiting = create_request_queue(scheduler.policy)
+    scheduler.skipped_waiting = create_request_queue(scheduler.policy)
+
+    completed = Request("completed", [1, 2], SamplingParams(max_tokens=4), pooling_params=None)
+    completed.status = RequestStatus.RUNNING
+    completed.num_computed_tokens = 2
+    scheduler.running = [completed]
+    scheduler.requests = {"completed": completed}
+    scheduler.input_coordinator.finished_requests.add("completed")  # terminal payload, no new chunk
+
+    assert scheduler._is_done_receiving_chunks("completed")
+    assert not scheduler._is_done_receiving_chunks("unknown")
+    output = scheduler.schedule()
+    assert not output.num_scheduled_tokens
+    assert scheduler._pending_finish_reqs == [completed]
+
+    completed.status = RequestStatus.WAITING_FOR_CHUNK
+    scheduler._pending_finish_reqs = []
+    scheduler._requeue_completed_native_chunks()
+    assert completed.status == RequestStatus.WAITING_FOR_CHUNK
+    assert completed.num_computed_tokens == 2
+    assert not scheduler.running
+
+
+def test_generation_scheduler_schedules_terminal_empty_prompt_chunk_once(monkeypatch: pytest.MonkeyPatch) -> None:
+    waiting = _chunk_request(
+        "terminal",
+        prompt_token_ids=[],
+        num_prompt_tokens=0,
+        additional_information=serialize_additional_information(
+            {"codes": {"audio": torch.tensor([[1, 2, 3]], dtype=torch.long)}}
+        ),
+    )
+    scheduler = _make_generation_scheduler(waiting)
+    scheduler.chunk_transfer_adapter.done_request_ids.add("terminal")
+    scheduler.requests = {"terminal": waiting}
+    assert _has_async_chunk_payload_to_run(waiting)
+    # A terminal codec request with an empty prompt is still scheduled exactly once.
+    monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.create_request_queue", create_request_queue)
+
+    output = OmniGenerationScheduler.schedule(scheduler)
+
+    assert output.num_scheduled_tokens == {"terminal": 1}
+    assert output.scheduled_new_reqs[0].req_id == "terminal"
+    assert scheduler._pending_finish_reqs == []
+
+
 class TestRestoreQueuesOnError:
     """Verify that restore_queues is called even when rewrapping raises."""
 

tests/core/sched/test_omni_ar_scheduler_aborted_queue_sweep.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Tests for ``OmniARScheduler._drop_aborted_queued_requests``.
 
 ``schedule()`` sweeps ``FINISHED_ABORTED`` requests out of the queues before
@@ -68,6 +71,7 @@ def test_schedule_sweeps_skipped_waiting_before_upstream_selection(
     aborted = _StubRequest("req-aborted", RequestStatus.FINISHED_ABORTED)
     scheduler = _make_scheduler(policy, skipped_waiting=[aborted])
     scheduler.chunk_transfer_adapter = None
+    scheduler.max_num_running_reqs = 8
     scheduler.input_coordinator = None
     scheduler.num_waiting_for_streaming_input = 0
     scheduler._process_pending_omni_inputs = lambda model_mode: None

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Tests for the AR sampled-token logprob contract."""
 
 from __future__ import annotations
@@ -93,6 +96,7 @@ class _Request:
         self.num_computed_tokens = 0
         self.num_in_flight_tokens = 0
         self.num_output_placeholders = 0
+        self.output_token_ids: list[int] = []
         # vLLM 0.27 (a0c092ee72): Request gained num_stale_output_tokens to
         # track in-flight outputs discarded at preemption/streaming-stop.
         self.num_stale_output_tokens = 0
@@ -145,6 +149,14 @@ def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
         kv_cache_manager=SimpleNamespace(take_events=lambda: None),
         kv_event_publisher=SimpleNamespace(publish=lambda _events: None),
         recompute_kv_load_failures=False,
+        _native_data_plane=False,
+        vllm_config=SimpleNamespace(
+            model_config=SimpleNamespace(
+                use_v2_model_runner=False,
+                async_chunk=False,
+                final_output=True,
+            )
+        ),
     )
     for name in _MIXIN_UPDATE_HELPERS:
         setattr(scheduler, name, MethodType(getattr(OmniSchedulerMixin, name), scheduler))
@@ -167,10 +179,15 @@ def _bind_request_lifecycle(
         scheduler.requests.pop(request.request_id, None)
         scheduler.finished_req_ids.add(request.request_id)
         scheduler.finished_req_ids_dict[request.client_index].add(request.request_id)
-        # vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
+        # Current vLLM contract: (kv_xfer_params, ec_xfer_params)
         return None, None
 
-    scheduler._update_request_with_output = update_request
+    def update_with_history(request, token_ids):
+        accepted, stopped = update_request(request, token_ids)
+        request.output_token_ids.extend(accepted)
+        return accepted, stopped
+
+    scheduler._update_request_with_output = update_with_history
     scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
     scheduler._handle_stopped_request = handle_stopped or (lambda _request: True)
     scheduler._free_request = free_request
@@ -219,6 +236,7 @@ def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
     (output,) = outputs[0].outputs
 
     assert output.new_token_ids == [7]
+    assert output.num_generation_tokens == 1
     assert output.finish_reason is FinishReason.STOP
     token_rows = np.asarray(output.new_logprobs.logprob_token_ids)
     value_rows = np.asarray(output.new_logprobs.logprobs)

tests/core/sched/test_omni_ar_scheduler_output_policy.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.core.sched.omni_ar_scheduler import _should_emit_engine_output
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize(
+    ("async_chunk", "final_output", "use_v2", "stopped", "has_control", "expected"),
+    [
+        (True, False, True, False, False, False),  # async+v2 suppresses data
+        (True, False, False, False, False, True),  # v1 continues to emit
+        (True, False, True, True, False, True),  # stopped always emits
+        (True, False, True, False, True, True),  # control-only still emits
+    ],
+)
+def test_async_chunk_intermediate_stage_emits_only_control_outputs(
+    async_chunk: bool,
+    final_output: bool,
+    use_v2: bool,
+    stopped: bool,
+    has_control: bool,
+    expected: bool,
+) -> None:
+    model_config = SimpleNamespace(
+        async_chunk=async_chunk,
+        final_output=final_output,
+        use_v2_model_runner=use_v2,
+    )
+
+    assert (
+        _should_emit_engine_output(
+            model_config,
+            stopped=stopped,
+            has_control=has_control,
+        )
+        is expected
+    )

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -17,6 +17,7 @@ import pytest
 import vllm_omni  # noqa: F401 - import for side effects (patch vLLM)
 from vllm.sampling_params import SamplingParams
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.outputs import ModelRunnerOutput
 from vllm.v1.engine import FinishReason
 from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
@@ -1261,3 +1262,47 @@ def test_context_overflow_emits_an_error_output_with_the_reason() -> None:
 
     sched._emit_streaming_context_overflow_outputs(outputs)
     assert len(outputs[2]) == 1
+
+
+@pytest.mark.parametrize("native", [True, False])
+def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, native) -> None:
+    sched = _make_scheduler(stage_id=1)
+    parked = SimpleNamespace(request_id="parked")
+    sched.requests = {"parked": parked}
+    sched.waiting = []
+    sched.running = []
+    sched._native_data_plane = native
+    sched.use_v2_model_runner = native
+    sched.max_num_running_reqs = 8
+    sched.input_coordinator = (
+        SimpleNamespace(_waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None) if native else None
+    )
+    sched.chunk_transfer_adapter = (
+        None
+        if native
+        else SimpleNamespace(
+            waiting_for_chunk_running_requests=[parked],
+            _held_non_active=[],
+            process_pending_chunks=lambda *_a, **_kw: None,
+            collect_failed_send_request_ids=lambda: {},
+            restore_queues=lambda *_a, **_kw: None,
+            postprocess_scheduler_output=lambda *_a, **_kw: None,
+        )
+    )
+    sched._consume_pending_connector_output = lambda model_mode: None
+    sched._process_pending_input_timeouts = lambda: None
+    sched._should_defer_waiting_admission = lambda: False
+    sched.get_finished_requests_needing_kv_transfer = lambda: {}
+    sched._wrap_omni_scheduler_output = lambda output, **_kwargs: output
+    observed_limits: list[int] = []
+
+    def fake_schedule(self, _throttle_prefills=False):
+        observed_limits.append(self.max_num_running_reqs)
+        return SimpleNamespace(scheduled_new_reqs=[])
+
+    monkeypatch.setattr(VLLMScheduler, "schedule", fake_schedule)
+
+    sched.schedule()
+
+    assert observed_limits == [7 if native else 8]
+    assert sched.max_num_running_reqs == 8

tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from __future__ import annotations
 
 from types import SimpleNamespace
@@ -141,6 +144,7 @@ def test_voxcpm2_unified_decode_graph_does_not_defer_without_cuda_graph(monkeypa
 
 def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None:
     scheduler = _make_scheduler()
+    scheduler.max_num_running_reqs = 8
     scheduler.running = [_MockRequest("decode")]
     original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING)
     deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING)

tests/core/sched/test_omni_sched_deferred_free_fence.py

@@ -123,6 +123,9 @@ def _make_fast_path_sched(*, defer_block_free: bool) -> tuple[MagicMock, Request
 
     sched = MagicMock()
     sched.max_num_scheduled_tokens = 8
+    sched.max_num_running_reqs = 8
+    sched._native_data_plane = False
+    sched._retains_state_across_chunks = False
     sched._pause_state = PauseState.UNPAUSED
     sched.requests = {request.request_id: request}
     sched.running = [request]

tests/core/sched/test_omni_scheduler_mixin_shared.py

@@ -1,5 +1,9 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from collections import defaultdict
 from types import SimpleNamespace
+from typing import Any
 
 import pytest
 from vllm.v1.engine import FinishReason
@@ -87,7 +91,7 @@ def test_full_payload_coordinator_matches_legacy_gate(monkeypatch, stage_id, asy
 
 
 def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
-    calls = []
+    calls: list[tuple[Any, ...]] = []
     scheduler = _Scheduler()
     scheduler.waiting = ["waiting"]
     scheduler.running = ["running"]
@@ -115,7 +119,7 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
         collect_failed_send_request_ids=_collect_failed_sends,
     )
     scheduler.input_coordinator = SimpleNamespace(
-        restore_queues=lambda waiting: calls.append(("restore-full", waiting))
+        restore_queues=lambda waiting, running: calls.append(("restore-full", waiting, running))
     )
 
     scheduler._process_pending_omni_inputs("ar")
@@ -130,7 +134,7 @@ def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
         ("chunk-timeouts", omni_scheduler_mixin.DEFAULT_INPUT_WAIT_TIMEOUT_S),
         ("failed-sends",),
         ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
-        ("restore-full", scheduler.waiting),
+        ("restore-full", scheduler.waiting, scheduler.running),
     ]
 
 

tests/core/sched/test_omni_scheduler_ready_inbox.py

@@ -0,0 +1,38 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Drain ready events once, merging live requests and filtering cancellations."""
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+from vllm_omni.outputs import OmniConnectorOutput
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_ready_inbox_coalesces_live_events_and_drops_cancelled(mocker):
+    scheduler = OmniSchedulerMixin()
+    scheduler.requests, scheduler.waiting, scheduler.running = {"r": object()}, [], []
+    coordinator = SimpleNamespace(
+        _async_chunk=True, update_request_metadata=mocker.Mock(), process_pending_chunks=mocker.Mock()
+    )
+    scheduler.input_coordinator = coordinator
+    scheduler._init_omni_connector_output_inbox()
+    scheduler.enqueue_omni_connector_output(OmniConnectorOutput(chunk_ready_req_ids={"r"}))
+    scheduler.enqueue_omni_connector_output(
+        OmniConnectorOutput(
+            chunk_ready_req_ids={"r", "aborted"},
+            chunk_finished_req_ids={"r", "aborted"},
+            request_metadata={"r": {"decode_token_end": 2}, "aborted": {"decode_token_end": 99}},
+        )
+    )
+    scheduler._consume_pending_connector_output(model_mode="ar")
+    coordinator.update_request_metadata.assert_called_once_with(
+        scheduler.requests, {"r": {"decode_token_end": 2}}, model_mode="ar"
+    )
+    coordinator.process_pending_chunks.assert_called_once_with([], [], {"r"}, {"r"})
+    scheduler._consume_pending_connector_output(model_mode="ar")
+    assert coordinator.update_request_metadata.call_count == 1
+    coordinator.process_pending_chunks.assert_called_with([], [], set(), set())

tests/core/sched/test_omni_scheduler_streaming_input_counter.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Tests for the ``num_waiting_for_streaming_input`` resync placed in
 ``OmniARScheduler.finish_requests`` / ``OmniGenerationScheduler.finish_requests``
 and in ``OmniARScheduler.schedule``.
@@ -128,6 +131,7 @@ def test_ar_schedule_resyncs_before_delegating_upstream(monkeypatch: pytest.Monk
     scheduler.skipped_waiting = []
     scheduler.num_waiting_for_streaming_input = 1  # leaked by the status rewrite
     scheduler.policy = "fcfs"
+    scheduler.max_num_running_reqs = 8
 
     seen: dict[str, int] = {}
 

tests/core/sched/test_omni_scheduling_coordinator.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for OmniSchedulingCoordinator.
 
 These tests use mock request objects and mock queues.  They do not require
@@ -20,7 +20,9 @@ import torch
 import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod
 from vllm_omni.core.sched.omni_scheduling_coordinator import (
     OmniSchedulingCoordinator,
+    uses_native_mrv2_data_plane,
 )
+from vllm_omni.core.sched.output import OmniChunkRecvHandle
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -56,6 +58,7 @@ def _make_request(req_id: str, status: str = "waiting") -> SimpleNamespace:
         prompt_token_ids=[],
         num_prompt_tokens=0,
         num_computed_tokens=0,
+        num_output_placeholders=0,
         _all_token_ids=[],
         _output_token_ids=[],
         payload_sender_info=None,
@@ -83,7 +86,7 @@ class MockQueue:
     def prepend_requests(self, requests):
         self._items = list(requests) + self._items
 
-    def remove(self, request):
+    def remove_request(self, request):
         self._items.remove(request)
 
     def remove_requests(self, requests):
@@ -96,7 +99,47 @@ class MockQueue:
 # ------------------------------------------------------------------ #
 
 
-class TestCoordinatorUpdateRequestMetadata(unittest.TestCase):
+class TestNativeMRV2DataPlaneSelection(unittest.TestCase):
+    def test_native_plane_requires_v2_and_async_chunk_capability(self):
+        native = SimpleNamespace(async_chunk=True, supports_native_mrv2_data_plane=True)
+
+        self.assertTrue(uses_native_mrv2_data_plane(native, use_v2_model_runner=True))
+        self.assertFalse(uses_native_mrv2_data_plane(native, use_v2_model_runner=False))
+        self.assertFalse(uses_native_mrv2_data_plane(SimpleNamespace(async_chunk=False), use_v2_model_runner=True))
+
+
+def test_chunk_registration_ready_and_terminal_lifecycle():
+    coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
+    req = _make_request("internal", status=RequestStatus.WAITING)
+    req.external_req_id = "external"
+    waiting = MockQueue([req])
+    coord.process_pending_chunks(waiting, [], set(), set())
+    assert req.status == RequestStatus.WAITING_FOR_CHUNK
+    [handle] = coord.pending_chunk_registrations
+    assert isinstance(handle, OmniChunkRecvHandle)
+    assert (handle.request_id, handle.external_req_id) == ("internal", "external")
+    coord.restore_queues(waiting, [])
+    coord.process_pending_chunks(waiting, [], {"internal"}, set())
+    assert req.status == RequestStatus.WAITING
+    assert "internal" in coord.requests_with_ready_chunks
+    req.status = RequestStatus.WAITING_FOR_CHUNK
+    coord.process_pending_chunks(waiting, [], set(), {"internal"})
+    assert "internal" in coord.finished_requests
+
+
+def test_chunk_waiting_removes_request_from_running_list():
+    coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
+    req = _make_request("running", status=RequestStatus.RUNNING)
+    running = [req]
+
+    coord.process_pending_chunks(MockQueue(), running, set(), set())
+
+    assert running == []
+    assert req.status == RequestStatus.WAITING_FOR_CHUNK
+    assert list(coord._waiting_for_chunk_running) == [req]
+
+
+class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
     """Test update_request_metadata applies scheduling metadata to requests."""
 
     def test_ar_mode_no_longer_sets_additional_information(self):
@@ -382,17 +425,18 @@ class TestTimeoutDetection(unittest.TestCase):
         self.assertEqual(result, {"r1"})
         self.assertEqual(len(coord._waiting_for_input), 0)
 
-    def test_free_finished_request_clears_waiting_since(self):
-        """free_finished_request clears coordinator lifecycle markers."""
-        coord = OmniSchedulingCoordinator(stage_id=1)
-        coord._waiting_since["r1"] = 0.0
-        coord._full_payload_input_received.add("r1")
+    def test_free_finished_request_clears_all_lifecycle_state(self):
+        """free_finished_request makes stale connector events harmless."""
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+        r1, r2 = _make_request("r1"), _make_request("r2")
         coord.finished_requests.add("r1")
+        coord.requests_with_ready_chunks.add("r1")
+        coord._waiting_for_chunk_running.extend([r1, r2])
+        coord.pending_chunk_registrations = [OmniChunkRecvHandle(request_id="r1"), OmniChunkRecvHandle(request_id="r2")]
+
         coord.free_finished_request("r1")
-        self.assertNotIn("r1", coord._waiting_since)
-        self.assertNotIn("r1", coord._full_payload_input_received)
+
         self.assertNotIn("r1", coord.finished_requests)
-
-
-if __name__ == "__main__":
-    unittest.main()
+        self.assertNotIn("r1", coord.requests_with_ready_chunks)
+        self.assertEqual([r.request_id for r in coord._waiting_for_chunk_running], ["r2"])
+        self.assertEqual([h.request_id for h in coord.pending_chunk_registrations], ["r2"])

tests/e2e/offline_inference/test_qwen3_tts_base.py

@@ -13,6 +13,7 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 import pytest
+from vllm.platforms import current_platform
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import get_asset_path
@@ -57,7 +58,12 @@ tts_server_params = [
     pytest.param(
         (MODEL, get_cuda_graph_config()),
         id="no_cuda_graph",
-    )
+    ),
+    pytest.param(
+        (MODEL, get_deploy_config_path("qwen3_tts_mrv2.yaml")),
+        id="async_chunk_mrv2",
+        marks=pytest.mark.skipif(not current_platform.is_cuda(), reason="MRV2 validation is CUDA-only"),
+    ),
 ]
 
 

tests/e2e/online_serving/test_qwen3_tts_base.py

@@ -12,6 +12,7 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 import pytest
+from vllm.platforms import current_platform
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.media import get_asset_path
@@ -50,7 +51,16 @@ tts_server_params = [
             server_args=["--trust-remote-code"],
         ),
         id="async_chunk",
-    )
+    ),
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("qwen3_tts_mrv2.yaml"),
+            server_args=["--trust-remote-code"],
+        ),
+        id="async_chunk_mrv2",
+        marks=pytest.mark.skipif(not current_platform.is_cuda(), reason="MRV2 validation is CUDA-only"),
+    ),
 ]
 
 

tests/engine/test_orchestrator.py

@@ -311,7 +311,7 @@ def _engine_core_outputs(tag: str, timestamp: float) -> SimpleNamespace:
     return SimpleNamespace(outputs=[tag], timestamp=timestamp, scheduler_stats=None, finished_requests=None)
 
 
-def _terminal_engine_core_outputs(request_id: str) -> EngineCoreOutputs:
+def _terminal_engine_core_outputs(request_id: str, timestamp: float = 1.0) -> EngineCoreOutputs:
     return EngineCoreOutputs(
         outputs=[
             EngineCoreOutput(
@@ -320,7 +320,7 @@ def _terminal_engine_core_outputs(request_id: str) -> EngineCoreOutputs:
                 finish_reason=FinishReason.STOP,
             )
         ],
-        timestamp=1.0,
+        timestamp=timestamp,
         finished_requests={request_id},
     )
 
@@ -520,6 +520,7 @@ async def _enqueue_add_request(
     original_prompt,
     sampling_params_list,
     final_stage_id: int,
+    final_output_stage_ids: list[int] | None = None,
 ) -> None:
     orchestrator_fixture.request_sync_q.put_nowait(
         StageSubmissionMessage(
@@ -530,6 +531,7 @@ async def _enqueue_add_request(
             output_prompt_text=None,
             sampling_params_list=sampling_params_list,
             final_stage_id=final_stage_id,
+            final_output_stage_ids=final_output_stage_ids,
             preprocess_ms=0.0,
             request_timestamp=time.time(),
             enqueue_ts=time.perf_counter(),
@@ -789,6 +791,10 @@ async def test_run_async_chunk(orchestrator_factory) -> None:
 
         stage1.push_engine_core_outputs(_engine_core_outputs("stage1-final", 3.0))
 
+        await _wait_for(
+            lambda: orchestrator_fixture.orchestrator.request_states["req-async"].pending_final_output is not None
+        )
+        stage0.push_engine_core_outputs(_engine_core_outputs("stage0-final", 3.1))
         output_msg = await _get_output_message(orchestrator_fixture)
 
         assert output_msg.request_id == "req-async"

tests/engine/test_output_processor.py

@@ -259,6 +259,17 @@ def test_native_text_tpot_only_accepts_positive_finished_metric(
     assert processor.pop_native_text_metrics("r")["vllm_tpot_ms"] == expected_tpot_ms
 
 
+def test_suppressed_native_tokens_use_authoritative_count_even_without_stats(monkeypatch):
+    monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None)
+    processor = object.__new__(MultimodalOutputProcessor)
+    processor._native_text_metrics_by_request = {}
+    state = _make_state(RequestOutputKind.FINAL_ONLY)
+    processor._update_stats_from_output(
+        state, SimpleNamespace(new_token_ids=[2150], num_generation_tokens=2048), None, None
+    )
+    assert processor.pop_native_text_metrics("r")["num_generation_tokens"] == 2048
+
+
 def test_delta_drains_output_modality_per_step():
     """DELTA drains the mm_type key (output modality) but preserves hidden-state keys."""
     s = _make_state(RequestOutputKind.DELTA)
@@ -616,6 +627,7 @@ def test_no_detokenizer_process_outputs_returns_nonterminal_audio_chunk(monkeypa
         finish_reason=None,
         stop_reason=None,
         kv_transfer_params=None,
+        ec_transfer_params=None,
         routed_experts=None,
         num_cached_tokens=0,
     )
@@ -631,6 +643,7 @@ def _make_mm_only_output_processor(monkeypatch):
     processor = object.__new__(MultimodalOutputProcessor)
     processor.output_modality = OutputModality.AUDIO
     processor.request_states = {"r": _make_no_detok_state(RequestOutputKind.DELTA)}
+    processor.tracing_enabled = False
     monkeypatch.setattr(
         VLLMOutputProcessor,
         "process_outputs",
@@ -660,6 +673,7 @@ def _audio_engine_output(*, is_segment_finished: bool, is_last_chunk: bool):
         finish_reason=FinishReason.STOP,
         stop_reason=None,
         kv_transfer_params=None,
+        ec_transfer_params=None,
         routed_experts=None,
         num_cached_tokens=0,
         is_segment_finished=is_segment_finished,

tests/engine/test_output_processor_mrv2_text.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Regression tests for MRv2 Omni AR text detach and generation payload outputs."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+from vllm.sampling_params import RequestOutputKind
+from vllm.v1.engine import FinishReason
+
+from vllm_omni.engine import OmniEngineCoreOutput
+from vllm_omni.engine.output_modality import OutputModality
+from vllm_omni.engine.output_processor import (
+    MultimodalOutputProcessor,
+    OmniRequestState,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Detokenizer:
+    def __init__(self):
+        self.output_token_ids = []
+        self.output_text = ""
+        self.update_calls = []
+
+    def update(self, token_ids, stop):
+        self.update_calls.append((list(token_ids), stop))
+        self.output_token_ids.extend(token_ids)
+        self.output_text += "X" * len(token_ids)
+        return None
+
+    def get_next_output_text(self, finished, delta):
+        return self.output_text
+
+    def num_output_tokens(self):
+        return len(self.output_token_ids)
+
+
+def _make_processor(output_modality, detokenizer=None):
+    state = OmniRequestState(
+        request_id="r",
+        external_req_id="r",
+        parent_req=None,
+        request_index=0,
+        lora_request=None,
+        output_kind=RequestOutputKind.FINAL_ONLY,
+        prompt="prompt" if detokenizer else None,
+        prompt_token_ids=[1] if detokenizer else [],
+        prompt_embeds=None,
+        logprobs_processor=MagicMock(logprobs=None, cumulative_logprob=None, prompt_logprobs=None)
+        if detokenizer
+        else None,
+        detokenizer=detokenizer,
+        max_tokens_param=None,
+        arrival_time=0.0,
+        queue=None,
+        log_stats=False,
+        stream_interval=1,
+    )
+    processor = MultimodalOutputProcessor(tokenizer=None, log_stats=False, output_modality=output_modality)
+    processor.request_states["r"] = state
+    processor.external_req_ids["r"].append("r")
+    return processor, state
+
+
+def _rehydrated_output(*, pooling_output=None, finish_reason=FinishReason.STOP):
+    """Build the post-StagePool shape without violating vLLM's wire schema."""
+    return SimpleNamespace(
+        request_id="r",
+        new_token_ids=[42],
+        pooling_output=pooling_output,
+        finish_reason=finish_reason,
+        stop_reason=None,
+        kv_transfer_params=None,
+        ec_transfer_params=None,
+        routed_experts=None,
+        trace_headers=None,
+        prefill_stats=None,
+        spec_decode_metrics=None,
+        new_sampling_mask=None,
+        num_nans_in_logits=0,
+        multimodal_output=None,
+        is_segment_finished=False,
+        is_non_final_audio_chunk=False,
+        output_type=None,
+        num_generation_tokens=None,
+        finished=True,
+    )
+
+
+def test_text_tokens_are_detokenized_when_mrv2_ar_output_has_pooling_payload():
+    detokenizer = _Detokenizer()
+    processor, state = _make_processor(OutputModality.LATENT, detokenizer)
+
+    # StagePool has already rehydrated the MRv2 carrier.  Use a structural
+    # output here because upstream EngineCoreOutput strictly validates
+    # pooling_output as torch.Tensor | None on construction.
+    output = _rehydrated_output(pooling_output={"hidden": torch.ones(1, 4)})
+
+    processed = processor.process_outputs([output])
+
+    assert detokenizer.update_calls == [([42], True)]
+    completion = processed.request_outputs[0].outputs[0]
+    assert list(completion.token_ids) == [42]
+    assert completion.text == "X"
+    assert not state.mm_accumulated.is_empty
+
+
+@pytest.mark.parametrize("streaming", [False, True])
+def test_generation_stage_mm_delivered_only_with_terminal(streaming):
+    # Non-terminal chunks accumulate; the terminal output carries the audio.
+    # When the first output is already terminal it is delivered directly.
+    processor, _ = _make_processor(OutputModality.AUDIO)
+    audio = torch.ones(1, 320)
+    sr = torch.tensor(24000, dtype=torch.int32)
+    mm_output = OmniEngineCoreOutput(
+        request_id="r",
+        new_token_ids=[],
+        multimodal_output={"model_outputs": audio, "sr": sr},
+        finish_reason=None if streaming else FinishReason.STOP,
+    )
+
+    processed = processor.process_outputs([mm_output])
+    if streaming:
+        assert processed.request_outputs == []
+        processed = processor.process_outputs(
+            [OmniEngineCoreOutput(request_id="r", new_token_ids=[], finish_reason=FinishReason.STOP)]
+        )
+
+    assert len(processed.request_outputs) == 1
+    completion = processed.request_outputs[0].outputs[0]
+    assert completion.text == ""
+    assert torch.equal(completion.multimodal_output["audio"], audio)
+    assert completion.multimodal_output["sr"].item() == 24000

tests/engine/test_request_snapshot.py

@@ -0,0 +1,21 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import copy
+
+import pytest
+
+from vllm_omni.engine.request_snapshot import copy_request_snapshot
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_full_snapshot_preserves_fields_aliases_and_isolation():
+    samples = [0.25, 0.5]
+    source = {"prompt_token_ids": [1, 2], "additional_information": {"samples": samples}, "alias": samples}
+    result = copy_request_snapshot(source)
+    assert result == copy.deepcopy(source)
+    assert result["alias"] is result["additional_information"]["samples"]
+    result["alias"].append(0.75)
+    result["prompt_token_ids"][0] = 9
+    assert samples == [0.25, 0.5]
+    assert source["prompt_token_ids"] == [1, 2]

tests/engine/test_single_stage_mode.py

@@ -1393,7 +1393,7 @@ class TestConnectRemoteEngineCoresCoordinator:
             yield mocker.Mock()
 
         mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx())
-        mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup")
+        mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True)
         with connect_remote_engine_cores(
             vllm_config=vllm_config,
             omni_master_server=omni_master_server,
@@ -1428,7 +1428,7 @@ class TestConnectRemoteEngineCoresCoordinator:
             yield mocker.Mock()
 
         mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx())
-        mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup")
+        mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True)
         with connect_remote_engine_cores(
             vllm_config=vllm_config,
             omni_master_server=omni_master_server,
@@ -1454,7 +1454,7 @@ class TestConnectRemoteEngineCoresCoordinator:
             yield mocker.Mock()
 
         mocker.patch("vllm_omni.engine.stage_engine_startup.zmq_socket_ctx", return_value=fake_socket_ctx())
-        mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup")
+        mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True)
 
         with connect_remote_diffusion_proc(
             omni_master_server=omni_master_server,
@@ -1466,7 +1466,8 @@ class TestConnectRemoteEngineCoresCoordinator:
         omni_master_server.get_zmq_addresses.assert_called_once_with(7, replica_id=2)
         omni_master_server.get_allocation.assert_called_once_with(7, replica_id=2)
         mock_wait.assert_called_once()
-        _, core_engines, parallel_config, *_ = mock_wait.call_args.args
+        _, core_engines, parallel_config, _, _, launch = mock_wait.call_args.args
+        assert launch.addresses is omni_master_server.get_zmq_addresses.return_value
         assert core_engines[0].local is False
         assert parallel_config.data_parallel_size_local == 0
 
@@ -1507,7 +1508,7 @@ class TestLaunchOmniCoreEngines:
             "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager",
             return_value=local_engine_manager,
         )
-        mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup")
+        mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True)
         with _launch_omni_core_engines(
             vllm_config=vllm_config,
             executor_class=mocker.Mock(),
@@ -1584,7 +1585,7 @@ class TestLaunchOmniCoreEngines:
             "vllm_omni.engine.stage_engine_startup.CoreEngineProcManager",
             return_value=mocker.Mock(),
         )
-        mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup")
+        mock_wait = mocker.patch("vllm_omni.engine.stage_engine_startup.wait_for_engine_startup", autospec=True)
         with _launch_omni_core_engines(
             vllm_config=vllm_config,
             executor_class=mocker.Mock(),

tests/engine/test_stage_engine_args.py

@@ -8,6 +8,7 @@ import sys
 import types
 from dataclasses import fields, replace
 from pathlib import Path
+from unittest.mock import patch
 
 import pytest
 from pydantic.fields import FieldInfo
@@ -61,9 +62,14 @@ _DEPLOY_DIR = Path(__file__).parents[2] / "vllm_omni" / "deploy"
 def _effective_backend_values(config_cls: type, engine_args: dict) -> dict[str, object]:
     backend_fields = fields(config_cls)
     backend_field_names = {backend_field.name for backend_field in backend_fields}
-    backend_config = config_cls(
-        **{name: copy.deepcopy(value) for name, value in engine_args.items() if name in backend_field_names}
-    )
+    # This is a schema/normalization comparison, not a Hub download test.
+    with (
+        patch("vllm.engine.arg_utils.get_model_path", side_effect=lambda model, *_: model),
+        patch("vllm_omni.diffusion.data.get_model_path", side_effect=lambda model, *_: model),
+    ):
+        backend_config = config_cls(
+            **{name: copy.deepcopy(value) for name, value in engine_args.items() if name in backend_field_names}
+        )
     effective_values: dict[str, object] = {}
     for backend_field in backend_fields:
         value = getattr(backend_config, backend_field.name)
@@ -78,6 +84,9 @@ _DIFFUSION_BACKEND_FIELDS = frozenset(field.name for field in fields(OmniDiffusi
 _TOPOLOGY_ONLY_ENGINE_ARGS = frozenset({"inline_diffusion"})
 _OMNI_ONLY_LLM_STAGE_ENGINE_FIELDS = frozenset(
     {
+        "final_output",
+        "use_v2_model_runner",
+        "supports_native_mrv2_data_plane",
         "active_stream_window",
         "codec_frame_rate_hz",
         "custom_voice_dir",

tests/engine/test_terminal_stage_metrics_order.py

@@ -0,0 +1,82 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Terminal audio must not retire a request before upstream terminal metrics."""
+
+import asyncio
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def build():
+    o = Orchestrator.__new__(Orchestrator)
+    o.stage_pools = [
+        SimpleNamespace(
+            final_output=False, build_stage_metrics=lambda *a, **k: SimpleNamespace(stage_id=0, num_tokens_out=42)
+        ),
+        SimpleNamespace(
+            final_output=True, build_stage_metrics=lambda *a, **k: SimpleNamespace(stage_id=1, num_tokens_out=0)
+        ),
+    ]
+    state = OrchestratorRequestState(
+        request_id="r",
+        final_stage_id=1,
+        final_output_stage_ids={1},
+        stage_submit_ts={0: 0, 1: 0},
+        sampling_params_list=[None, None],
+    )
+    o.request_states = {"r": state}
+    o.output_async_queue = asyncio.Queue()
+    o.async_chunk = True
+    o._pd_pair = None
+    o._cfg_tracker = MagicMock()
+    o._cfg_tracker.is_companion.return_value = False
+    o._cfg_tracker.cleanup_parent.return_value = []
+    o._is_duplex_session_request = lambda s: False
+    o._duplex_output_decision = lambda *a: None
+    o._stage_receives_async_chunks = lambda s: True
+    cleaned = []
+
+    async def cleanup(ids):
+        cleaned.extend(ids)
+        for r in ids:
+            o.request_states.pop(r, None)
+
+    o._cleanup_request_ids = cleanup
+    return o, state, cleaned
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("order", [(0, 1), (1, 0)])
+async def test_terminal_metrics_precede_final_audio_for_either_arrival_order(order):
+    o, state, cleaned = build()
+    for stage in order:
+        await o._handle_processed_outputs(stage, 0, [SimpleNamespace(request_id="r", finished=True, error=None)])
+        if stage == 1 and not state.finished_stage_ids.issuperset({0, 1}):
+            assert o.output_async_queue.empty()
+            assert "r" in o.request_states
+            assert cleaned == []
+    metrics = o.output_async_queue.get_nowait()
+    final = o.output_async_queue.get_nowait()
+    assert type(metrics).__name__ == "StageMetricsMessage"
+    assert metrics.stage_id == 0 and metrics.metrics.num_tokens_out == 42
+    assert type(final).__name__ == "OutputMessage" and final.finished
+    assert o.output_async_queue.empty()
+    assert cleaned == ["r"]
+    assert state.pending_final_output is None
+
+
+@pytest.mark.asyncio
+async def test_cancelled_pending_output_is_not_published_by_late_upstream():
+    o, state, cleaned = build()
+    await o._handle_processed_outputs(1, 0, [SimpleNamespace(request_id="r", finished=True, error=None)])
+    assert state.pending_final_output is not None
+    await o._cleanup_request_ids(["r"])
+    await o._handle_processed_outputs(0, 0, [SimpleNamespace(request_id="r", finished=True, error=None)])
+    assert o.output_async_queue.empty()
+    assert cleaned == ["r"]

tests/entrypoints/openai_api/test_ref_audio_cache.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+import asyncio
+import hashlib
+from collections import OrderedDict
+
+import numpy as np
+import pytest
+
+from vllm_omni.entrypoints.openai.serving_speech import (
+    _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES,
+    _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES,
+    OmniOpenAIServingSpeech,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.fixture
+def cache():
+    server = object.__new__(OmniOpenAIServingSpeech)
+    server._diffusion_mode = False
+    server.model_config = None
+    server._ref_audio_resolve_cache = OrderedDict()
+    server._ref_audio_resolve_cache_bytes = 0
+    server._ref_audio_resolve_cache_max_entries = _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES
+    server._ref_audio_resolve_cache_max_bytes = _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES
+    server._ref_audio_model_artifact_ready = set()
+    server._request_ref_audio_artifact_keys = {}
+    return server
+
+
+def _put(cache, key, samples, artifact=None):
+    cache._put_resolved_ref_audio(key, samples, 24000, artifact or key)
+
+
+def test_compact_storage_and_default_capacity(cache):
+    assert cache._ref_audio_resolve_cache_max_entries == 1024
+    assert cache._ref_audio_resolve_cache_max_bytes == 512 * 1024 * 1024
+    _put(cache, "a", [0.0, 0.25, -0.5, 1.0])
+    value = cache._ref_audio_resolve_cache["a"]
+    assert value[0].dtype == np.float32
+    assert value[2] == cache._ref_audio_resolve_cache_bytes == 16
+
+
+def test_cached_float32_roundtrip_and_mutation_isolation(cache):
+    source = "data:audio/wav;base64,already-cached"
+    key = hashlib.sha1(source.encode()).hexdigest()
+    values = np.random.default_rng(42).standard_normal(48000).astype(np.float32).tolist()
+    expected = values.copy()
+    _put(cache, key, values)
+    values[0] = 99.0
+    first, sample_rate, _ = asyncio.run(cache._resolve_ref_audio(source))
+    assert first == expected
+    first[0] = 100.0
+    second, _, _ = asyncio.run(cache._resolve_ref_audio(source))
+    assert sample_rate == 24000
+    assert second == expected
+
+
+def test_entry_eviction_invalidates_artifact_readiness(cache):
+    cache._ref_audio_resolve_cache_max_entries = 2
+    _put(cache, "a", [0.0])
+    cache._ref_audio_model_artifact_ready.add(("a", False))
+    _put(cache, "b", [0.0])
+    _put(cache, "c", [0.0])
+    assert list(cache._ref_audio_resolve_cache) == ["b", "c"]
+    assert ("a", False) not in cache._ref_audio_model_artifact_ready
+
+
+def test_byte_budget_evicts_and_rejects_oversized_entries(cache):
+    cache._ref_audio_resolve_cache_max_bytes = 16
+    _put(cache, "a", [0.0] * 3)
+    _put(cache, "b", [0.0] * 3)
+    assert list(cache._ref_audio_resolve_cache) == ["b"]
+    assert cache._ref_audio_resolve_cache_bytes == 12
+    _put(cache, "too-large", [0.0] * 5)
+    assert "too-large" not in cache._ref_audio_resolve_cache
+
+
+def test_replacement_accounts_bytes_and_preserves_referenced_artifact(cache):
+    _put(cache, "a", [0.0] * 3, "same")
+    _put(cache, "alias", [0.0] * 2, "same")
+    cache._ref_audio_model_artifact_ready.add(("same", False))
+    _put(cache, "a", [1.0] * 4, "new")
+    assert cache._ref_audio_resolve_cache_bytes == 24
+    assert ("same", False) in cache._ref_audio_model_artifact_ready
+
+
+def test_pool_accepts_666_references(cache):
+    for index in range(666):
+        _put(cache, str(index), [0.0] * 8)
+    assert len(cache._ref_audio_resolve_cache) == 666
+    assert cache._ref_audio_resolve_cache_bytes == 666 * 8 * 4

tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py

@@ -682,6 +682,7 @@ class TestCodePredictorGraphReplay:
             return_proj_buf=False,
         )
         predictor._prefix_graphs_enabled = False
+        predictor._prefix_reprefill_enabled = False
         predictor._bucket_pos_ids = {1: torch.arange(4).unsqueeze(0)}
         graph = mocker.Mock()
         graph_output = torch.zeros(1, 4, 4)
@@ -1362,3 +1363,38 @@ class TestCodePredictorFusedProjections:
 
         with pytest.raises(RuntimeError, match="missing fused parameters"):
             model.load_weights(weights)
+
+
+class TestMTPExecutionBuckets:
+    """Outer MRv2 MTP graph buckets must align with the predictor bucket set."""
+
+    def test_configure_mtp_execution_buckets(self, mocker: MockerFixture, loaded_target_classes) -> None:
+        _ = loaded_target_classes
+        common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
+        mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
+
+        cp_config, _ = _make_tiny_config(loaded_target_classes)
+        vllm_config = _make_vllm_config(mocker, max_num_seqs=4)
+        wrapper = common_mod.CodePredictorWrapper(
+            vllm_config=vllm_config,
+            cp_config=cp_config,
+            wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=False),
+            talker_hidden_size=cp_config.hidden_size,
+        )
+
+        # Default keeps the legacy power-of-two derivation.
+        assert wrapper._execution_batch_buckets is None
+        assert wrapper._batch_bucket_sizes() == [1, 2, 4]
+
+        wrapper.configure_mtp_execution_buckets([3, 1, 3, 0, 9])
+        assert wrapper._batch_bucket_sizes() == [1, 3, 4]  # deduplicate/filter declaration and retain max
+
+        with pytest.raises(RuntimeError, match="before the first warmup"):
+            wrapper2 = common_mod.CodePredictorWrapper(
+                vllm_config=vllm_config,
+                cp_config=cp_config,
+                wrapper_config=common_mod.CodePredictorWrapperConfig(use_cuda_graphs=False),
+                talker_hidden_size=cp_config.hidden_size,
+            )
+            wrapper2._bucket_sizes = [1, 2, 4]  # stand-in for warmed-up state
+            wrapper2.configure_mtp_execution_buckets([3])

tests/model_executor/models/qwen3_tts/test_cuda_graph_decoder.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """
 Tests for CUDA Graph decoder wrapper numerical equivalence.
 
@@ -142,6 +142,28 @@ def test_padded_output_shape_and_length(decoder, wrapper, seq_len):
     assert graph_out.shape[-1] == expected_len
 
 
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_padded_batch_size_uses_larger_capture_bucket(decoder):
+    """Non-exact batch sizes should replay the next captured batch bucket."""
+    w = CUDAGraphDecoderWrapper(
+        decoder=decoder,
+        capture_sizes=[25],
+        capture_batch_sizes=[1, 4],
+        num_quantizers=NUM_QUANTIZERS,
+        enabled=True,
+    )
+    w.warmup(DEVICE)
+    codes = _random_codes(25, batch_size=3)
+
+    with torch.no_grad():
+        eager_out = decoder(codes)
+        graph_out = w.decode(codes)
+
+    assert (4, 25) in w.graphs
+    assert graph_out.shape == eager_out.shape
+    torch.testing.assert_close(graph_out, eager_out, atol=0, rtol=0)
+
+
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 @pytest.mark.parametrize("seq_len", [10, 30, 47, 73, 99])
 def test_padded_interior_positions_close(decoder, wrapper, seq_len):

tests/model_executor/models/qwen3_tts/test_qwen3_tts_output_async_cuda.py

@@ -0,0 +1,55 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Output metadata must not wait for the preceding model forward on CUDA."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
+    Qwen3TTSTalkerForConditionalGeneration,
+)
+from vllm_omni.worker.gpu_ar_model_runner import _snapshot_tensor_payload_to_cpu_async
+
+pytestmark = [pytest.mark.core_model]
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("validity_device", ["cpu", "cuda"])
+@torch.inference_mode()
+def test_output_validity_preserves_values_without_waiting_for_forward(validity_device):
+    talker = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
+    talker.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=True))
+    hidden = torch.zeros((3, 8), device="cuda")
+    infos = [
+        {
+            "codes": {"audio": torch.zeros((frames, 16), device="cuda", dtype=torch.long)},
+            "meta": {"codec_frame_valid": torch.tensor(valid, device=validity_device)},
+        }
+        for frames, valid in [(2, True), (1, False)]
+    ]
+    # Warm allocation paths before checking the stream-ordering contract.
+    talker.make_omni_output(hidden, model_intermediate_buffer=infos)
+    copy_stream = torch.cuda.Stream()
+    torch.accelerator.synchronize()
+
+    preceding_forward = torch.cuda.Event()
+    torch.cuda._sleep(500_000_000)
+    preceding_forward.record()
+    output = talker.make_omni_output(hidden, model_intermediate_buffer=infos)
+    waited_for_forward = preceding_forward.query()
+    validity = output.multimodal_outputs["meta"]["codec_frame_valid"]
+    snapshot = _snapshot_tensor_payload_to_cpu_async(
+        output.multimodal_outputs, copy_stream=copy_stream, pin_memory=True
+    )
+    snapshot.wait()
+
+    assert not waited_for_forward, "output construction synchronized the preceding forward"
+    assert validity.device.type == "cuda"
+    assert validity.dtype == torch.int8
+    cpu_validity = snapshot.payload["meta"]["codec_frame_valid"]
+    assert cpu_validity.device.type == "cpu"
+    assert cpu_validity.tolist() == [1, 1, 0]
+    assert all(info["meta"]["codec_frame_valid"].device.type == validity_device for info in infos)

tests/model_executor/models/qwen3_tts/test_qwen3_tts_stateless_capture.py

@@ -0,0 +1,100 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Exercise actual stateless decoder capture, not a mocked capture indicator."""
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_marks
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2DecoderConfig,
+)
+from vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import (
+    Qwen3TTSTokenizerV2DecoderTransformerModel,
+)
+
+pytestmark = [pytest.mark.core_model, *hardware_marks(res={"cuda": "L4"}, num_cards=1)]
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+@pytest.mark.parametrize("input_mode", ["implicit", "positions"])
+@torch.inference_mode()
+def test_stateless_decoder_capture_replays_new_inputs(batch_size, input_mode):
+    sequence_length = 8
+    torch.manual_seed(42)
+    config = Qwen3TTSTokenizerV2DecoderConfig(
+        hidden_size=32,
+        latent_dim=16,
+        max_position_embeddings=512,
+        num_attention_heads=4,
+        num_key_value_heads=4,
+        intermediate_size=64,
+        num_hidden_layers=1,
+        num_quantizers=2,
+        sliding_window=16,
+        decoder_dim=32,
+        upsample_rates=(2,),
+        upsampling_ratios=(2,),
+    )
+    config._attn_implementation = "sdpa"
+    model = Qwen3TTSTokenizerV2DecoderTransformerModel._from_config(config).cuda().eval()
+    inputs = torch.randn(batch_size, sequence_length, config.latent_dim, device="cuda")
+    positions = {}
+    if input_mode == "positions":
+        cache_position = torch.arange(sequence_length, device="cuda")
+        positions = {"cache_position": cache_position, "position_ids": cache_position.unsqueeze(0)}
+    stream = torch.cuda.Stream()
+    stream.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(stream):
+        for _ in range(3):
+            model(inputs_embeds=inputs, use_cache=False, **positions)
+    torch.cuda.current_stream().wait_stream(stream)
+    torch.accelerator.synchronize()
+
+    # No wrapper or eager fallback: an illegal capture must fail this test.
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured_output = model(inputs_embeds=inputs, use_cache=False, **positions).last_hidden_state
+
+    outputs = []
+    for _ in range(3):
+        new_inputs = torch.randn_like(inputs)
+        expected = model(inputs_embeds=new_inputs, use_cache=False, **positions).last_hidden_state
+        if input_mode == "implicit":
+            # Compare the cached mask against the independent live-mask path,
+            # not only against eager execution using the same cached mask.
+            reference = model(
+                inputs_embeds=new_inputs,
+                use_cache=False,
+                position_ids=torch.arange(sequence_length, device="cuda").unsqueeze(0),
+            ).last_hidden_state
+            torch.testing.assert_close(expected, reference, atol=1e-5, rtol=1e-4)
+        inputs.copy_(new_inputs)
+        graph.replay()
+        torch.accelerator.synchronize()
+        torch.testing.assert_close(captured_output, expected, atol=1e-5, rtol=1e-4)
+        outputs.append(captured_output.clone())
+    assert not torch.equal(outputs[0], outputs[1]), "replay must consume the new input"
+
+
+@pytest.mark.parametrize("batch_size", [1, 2])
+@torch.inference_mode()
+def test_stateless_codec_wrapper_captures_and_replays_waveform(batch_size):
+    from tests.model_executor.models.qwen3_tts.test_qwen3_tts_incremental_decode import _make_small_decoder
+    from vllm_omni.model_executor.models.qwen3_tts.segmented_graph_wrapper import CUDAGraphDecoderWrapper
+
+    torch.manual_seed(42)
+    decoder = _make_small_decoder().cuda().eval()
+    wrapper = CUDAGraphDecoderWrapper(decoder, num_quantizers=2, async_chunk=False)
+    # Exercise the production capture method, including quantization and
+    # waveform convolutions, without accepting warmup's eager fallback.
+    wrapper._capture_stateless(batch_size, 8, torch.device("cuda"), torch.long)
+    state = wrapper.stateless_states[(batch_size, 8)]
+    for _ in range(3):
+        codes = torch.randint(0, 32, (batch_size, 2, 8), device="cuda")
+        expected = decoder._forward_exact(codes)
+        state["input"]["codes"].copy_(codes)
+        state["graph"].replay()
+        torch.accelerator.synchronize()
+        torch.testing.assert_close(state["output"], expected, atol=1e-5, rtol=1e-4)
+        assert torch.isfinite(state["output"]).all()

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_preprocess.py

@@ -1,5 +1,9 @@
-from collections import OrderedDict
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from collections import OrderedDict, UserDict
 from types import SimpleNamespace
+from typing import Any
 
 import numpy as np
 import pytest
@@ -36,6 +40,7 @@ def _make_minimal_talker(
     """
     model = Qwen3TTSTalkerForConditionalGeneration.__new__(Qwen3TTSTalkerForConditionalGeneration)
     model.talker_config = SimpleNamespace(codec_pad_id=7, num_code_groups=16)
+    model._codebook_vocab_size = 2048
     model._embedding_dtype = torch.bfloat16
     if tts_pad_embed is None:
         tts_pad_embed = torch.zeros((1, 4), dtype=torch.bfloat16)
@@ -51,6 +56,22 @@ def _make_minimal_talker(
     return model
 
 
+def test_postprocess_batch_gathers_each_request_tail():
+    from vllm_omni.model_executor.models.output_templates import OwnedBatchTensor
+
+    model = _make_minimal_talker()
+    hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+
+    key, values = model.postprocess_batch_mrv2(
+        hidden_states=hidden,
+        last_token_indices=torch.tensor([1, 3]),
+    )
+
+    assert key == ("hidden_states", "last")
+    assert isinstance(values, OwnedBatchTensor)  # freshly allocated rows transfer ownership
+    assert torch.equal(values.tensor, torch.tensor([[3.0, 4.0], [7.0, 8.0]]))
+
+
 def _make_minimal_builder(
     *,
     config: SimpleNamespace | None = None,
@@ -95,6 +116,7 @@ def _make_minimal_builder(
     )
     builder._speaker_cache = None
     builder._text_tokenizer = None
+    builder._projected_token_cache = {}
     builder._embedding_dtype = torch.bfloat16
     builder._ref_audio_artifact_cache_max_entries = 256
     builder._ref_audio_artifact_cache = OrderedDict()
@@ -259,6 +281,30 @@ def test_decode_compacts_long_trailing_text_after_large_offset():
     assert torch.equal(update["hidden_states"]["trailing_text"], trailing_text[65:])
 
 
+def test_decode_marks_codec_frame_validity_from_processed_input_token():
+    model = _make_minimal_talker()
+    model.embed_input_ids = lambda input_ids: input_ids.to(torch.float32).reshape(1, 1, 1).expand(1, 1, 4)
+    common = {
+        "input_embeds": None,
+        "text": ["hello"],
+        "task_type": ["CustomVoice"],
+        "hidden_states": {
+            "trailing_text": torch.ones((2, 4)),
+            "last": torch.ones(4),
+        },
+        "meta": {"talker_text_offset": 0},
+        "_omni_is_prefill": False,
+        "_omni_num_computed_tokens": 2,
+        "_omni_prompt_len": 2,
+    }
+
+    _, _, valid_update = model.preprocess(input_ids=torch.tensor([123]), **common)
+    _, _, eos_update = model.preprocess(input_ids=torch.tensor([4198]), **common)
+
+    assert valid_update["meta"]["codec_frame_valid"].item() is True
+    assert eos_update["meta"]["codec_frame_valid"].item() is False
+
+
 def test_decode_replay_span_embeds_all_tokens_without_mutating_decode_state():
     model = _make_minimal_talker()
 
@@ -286,14 +332,17 @@ def test_decode_replay_span_embeds_all_tokens_without_mutating_decode_state():
         out_embeds.cpu(),
         torch.tensor([[101.0] * 4, [202.0] * 4, [303.0] * 4], dtype=torch.bfloat16),
     )
-    assert update == {"meta": {"codec_streaming": True}}
+    assert update["meta"]["codec_streaming"] is True
+    assert update["meta"]["codec_frame_valid"].item() is False
 
 
-def test_decode_batch_preprocess_matches_decode_state_updates():
+@pytest.mark.parametrize("mrv2", [False, True])
+def test_decode_batch_preprocess_matches_decode_state_updates(mrv2):
     tts_pad = torch.full((1, 4), -1.0, dtype=torch.bfloat16)
     model = _make_minimal_talker(tts_pad_embed=tts_pad)
 
     def fake_embed_input_ids(input_ids):
+        assert not mrv2, "MRV2 must reuse the prepared token embeddings"
         return input_ids.to(torch.float32).reshape(-1, 1, 1).expand(-1, 1, 4)
 
     model.embed_input_ids = fake_embed_input_ids
@@ -302,7 +351,10 @@ def test_decode_batch_preprocess_matches_decode_state_updates():
     last_a = torch.full((4,), 2.0, dtype=torch.float32)
     last_b = torch.full((4,), 3.0, dtype=torch.float32)
 
-    out_ids, out_embeds, past_hidden, text_step, updates = model.preprocess_decode_batch(
+    preprocess = model.preprocess_decode_batch_mrv2 if mrv2 else model.preprocess_decode_batch
+    prepared = torch.tensor([[101.0] * 4, [202.0] * 4], dtype=torch.bfloat16)
+    out_ids, out_embeds, past_hidden, text_step, updates = preprocess(
+        **({"input_embeds": prepared} if mrv2 else {}),
         input_ids=torch.tensor([101, 202], dtype=torch.long),
         req_infos=[
             {
@@ -321,15 +373,19 @@ def test_decode_batch_preprocess_matches_decode_state_updates():
     )
 
     assert out_ids.tolist() == [101, 202]
+    if mrv2:
+        assert out_embeds.data_ptr() == prepared.data_ptr()
     assert torch.equal(out_embeds.cpu(), torch.tensor([[101.0] * 4, [202.0] * 4], dtype=torch.bfloat16))
     assert torch.equal(past_hidden.cpu(), torch.stack([last_a, last_b]).to(torch.bfloat16))
     assert torch.equal(text_step[0].cpu(), trailing_a[1].to(torch.bfloat16))
     assert torch.equal(text_step[1].cpu(), tts_pad.reshape(-1).to(torch.bfloat16))
     assert updates[0]["meta"]["talker_text_offset"] == 2
     assert updates[0]["meta"]["codec_streaming"] is True
+    assert updates[0]["meta"]["codec_frame_valid"].item() is True
     assert "hidden_states" not in updates[0]
     assert updates[1]["meta"]["talker_text_offset"] == 0
     assert updates[1]["meta"]["codec_streaming"] is False
+    assert updates[1]["meta"]["codec_frame_valid"].item() is True
     assert updates[1]["hidden_states"]["trailing_text"].numel() == 0
 
 
@@ -449,11 +505,11 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
 
         def __call__(self, texts, *, padding=False):
             self.calls.append((texts, padding))
-            return {"input_ids": [[idx + 1, idx + 2, idx + 3] for idx, _ in enumerate(texts)]}
+            return UserDict({"input_ids": [[1, 2, 3], [2, 3, 4, 5]]})
 
     text_tok = FakeTextTokenizer()
     builder._text_tokenizer = text_tok
-    buf = {
+    buf: dict[str, dict[str, Any]] = {
         "r1": {
             "task_type": ["Base"],
             "text": ["one"],
@@ -488,9 +544,17 @@ def test_base_voice_clone_batch_preprocess_encodes_ref_code_by_sample_rate():
     assert buf["r2"][NORMALIZED_REF_AUDIO_KEY][0] is wav2
     assert len(text_tok.calls) == 2
     assert torch.equal(buf["r1"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([1, 2, 3]))
-    assert torch.equal(buf["r2"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4]))
+    assert torch.equal(buf["r2"][PRECOMPUTED_TEXT_IDS_KEY], torch.tensor([2, 3, 4, 5]))
     assert torch.equal(buf["r1"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([1, 2, 3]))
-    assert torch.equal(buf["r2"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4]))
+    assert torch.equal(buf["r2"][PRECOMPUTED_REF_IDS_KEY], torch.tensor([2, 3, 4, 5]))
+    assert (
+        buf["r1"][PRECOMPUTED_TEXT_IDS_KEY].untyped_storage().data_ptr()
+        == buf["r2"][PRECOMPUTED_TEXT_IDS_KEY].untyped_storage().data_ptr()
+    )
+    assert (
+        buf["r1"][PRECOMPUTED_REF_IDS_KEY].untyped_storage().data_ptr()
+        == buf["r2"][PRECOMPUTED_REF_IDS_KEY].untyped_storage().data_ptr()
+    )
 
 
 def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_without_speech_tokenizer():
@@ -506,7 +570,7 @@ def test_base_voice_clone_batch_preprocess_reuses_singleton_normalized_audio_wit
             return {"input_ids": [[7, 8, 9] for _ in texts]}
 
     builder._text_tokenizer = FakeTextTokenizer()
-    buf = {
+    buf: dict[str, dict[str, Any]] = {
         "r1": {
             "task_type": ["Base"],
             "text": ["one"],
@@ -539,7 +603,7 @@ def test_base_voice_clone_batch_preprocess_skips_after_initial_prefill_state_exi
     builder._encode_ref_audio_batch_fn = lambda *a, **kw: (_ for _ in ()).throw(
         AssertionError("speech tokenizer not expected")
     )
-    buf = {
+    buf: dict[str, dict[str, Any]] = {
         "r1": {
             "task_type": ["Base"],
             "text": ["one"],
@@ -574,7 +638,7 @@ def test_base_voice_clone_batch_preprocess_uses_serving_artifact_cache_key_witho
             return {"input_ids": [[7, 8, 9] for _ in texts]}
 
     builder._text_tokenizer = FakeTextTokenizer()
-    buf = {
+    buf: dict[str, dict[str, Any]] = {
         "r1": {
             "task_type": ["Base"],
             "text": ["one"],

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_ref_codes.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Regression tests for cross-request ``codes.ref`` leakage (#4370).
 
 ``make_omni_output`` used to collapse every request's reference codec
@@ -11,11 +11,14 @@ vocoder chunks then decoded with another request's reference voice as
 context (audible onset timbre deformation at any concurrency > 1).
 """
 
+from types import SimpleNamespace
+
 import pytest
 import torch
 
 from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import (
     Qwen3TTSTalkerForConditionalGeneration,
+    _qwen3_tts_gpu_resident_buffer_keys,
 )
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -23,6 +26,18 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 _NUM_CODE_GROUPS = 16
 
 
+def test_large_tts_prefill_artifacts_stay_gpu_resident_only_on_mrv2() -> None:
+    v1_keys = _qwen3_tts_gpu_resident_buffer_keys(False)
+    v2_keys = _qwen3_tts_gpu_resident_buffer_keys(True)
+
+    assert ("embed", "prefill") not in v1_keys
+    assert ("codes", "ref") not in v1_keys
+    assert ("meta", "codec_frame_valid") not in v1_keys
+    assert ("embed", "prefill") in v2_keys
+    assert ("codes", "ref") in v2_keys
+    assert ("meta", "codec_frame_valid") in v2_keys
+
+
 def _make_talker() -> Qwen3TTSTalkerForConditionalGeneration:
     # make_omni_output reads no instance state, so skip __init__ (no real
     # config / checkpoint needed), same as the preprocess tests.
@@ -89,3 +104,42 @@ def test_make_omni_output_keeps_hidden_states_for_replay_spans_without_audio_cod
 
     assert torch.equal(out.text_hidden_states, hidden)
     assert out.multimodal_outputs["codes"]["audio"].shape == (1, _NUM_CODE_GROUPS)
+
+
+def test_make_omni_output_preserves_per_request_codec_frame_validity() -> None:
+    talker = _make_talker()
+    talker.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=True))
+    valid = _info(1, None)
+    invalid = _info(1, None)
+    valid["meta"]["codec_frame_valid"] = True
+    invalid["meta"]["codec_frame_valid"] = False
+
+    out = talker.make_omni_output(
+        torch.zeros((2, 8)),
+        model_intermediate_buffer=[valid, invalid],
+    )
+
+    assert out.multimodal_outputs["meta"]["codec_frame_valid"].tolist() == [1, 0]
+
+
+def test_make_omni_output_publishes_ref_once_for_kv_resumed_first_decode() -> None:
+    talker = _make_talker()
+    talker.vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(async_chunk=True),
+    )
+    ref_code = torch.ones((5, _NUM_CODE_GROUPS), dtype=torch.long)
+    info = _info(1, ref_code)
+    info["_omni_num_computed_tokens"] = 9
+    info["_omni_prompt_len"] = 5
+
+    first = talker.make_omni_output(
+        torch.zeros((1, 8)),
+        model_intermediate_buffer=[info],
+    )
+    second = talker.make_omni_output(
+        torch.zeros((1, 8)),
+        model_intermediate_buffer=[info],
+    )
+
+    assert first.multimodal_outputs["codes"]["ref"][0] is ref_code
+    assert "ref" not in second.multimodal_outputs["codes"]

tests/model_executor/models/qwen3_tts/test_qwen3_tts_talker_silence_ban.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Unit tests for the Qwen3-TTS leading-silence ban (#4966, #5048).
 
 ``compute_logits`` suppresses the checkpoint-derived silence codec vocabulary
@@ -59,6 +59,7 @@ def _make_talker(*, ban_frames: int, logits: torch.Tensor, modes=None):
     object.__setattr__(talker, "_silence_mask", silence_mask)
     object.__setattr__(talker, "_codec_disallowed_mask", torch.zeros((_VOCAB,), dtype=torch.bool))
     object.__setattr__(talker, "_batch_req_ids", req_ids)
+    object.__setattr__(talker, "_mrv2_silence_ban_mask", None)
     object.__setattr__(
         talker,
         "_req_x_vector_only",
@@ -78,7 +79,26 @@ def _banned(row: torch.Tensor) -> set[int]:
     return {int(i) for i in torch.nonzero(torch.isinf(row) & (row < 0)).flatten()}
 
 
-# -------------------- the step gate --------------------
+def test_mrv2_silence_ban_tracks_batch_order_scope_and_decode_window() -> None:
+    talker = _make_talker(ban_frames=3, logits=torch.zeros((3, _VOCAB)), modes=[True, False, True])
+    with talker.mrv2_sampling_context(
+        req_ids=["req2", "req1", "req0"],
+        num_output_tokens=torch.tensor([3, 0, 1]),
+    ):
+        logits = talker.compute_logits(torch.zeros((3, 4)))
+
+    assert _banned(logits[0]) == set()
+    assert _banned(logits[1]) == set()
+    assert _banned(logits[2]) == set(_SILENCE_IDS)
+
+    # Prompt logprobs are computed after sampling and must remain unmasked.
+    assert not torch.isneginf(talker.compute_logits(torch.zeros((3, 4)))).any()
+
+    with talker.mrv2_sampling_context(
+        req_ids=["req2", "req1", "req0"],
+        num_output_tokens=torch.tensor([4, 1, 3]),
+    ):
+        assert not torch.isneginf(talker.compute_logits(torch.zeros((3, 4)))).any()
 
 
 def test_ban_disabled_leaves_logits_unchanged() -> None:

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -7,6 +7,9 @@ from types import SimpleNamespace
 import pytest
 import torch
 
+from vllm_omni.model_executor.stage_input_processors import (
+    qwen3_tts as qwen3_tts_processors,
+)
 from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
     AdaptiveChunkController,
     compute_adaptive_emit,
@@ -32,7 +35,52 @@ _FRAME = [1, 2, 3, 4]
 _Q = len(_FRAME)
 
 
-def _req(rid, *, finished, initial_codec_chunk_frames=None, non_streaming_mode=None, full_utterance_decode=None):
+@pytest.mark.parametrize("frames", [1, 26])
+def test_terminal_without_new_frames_does_not_replay_last_chunk(frames):
+    tm = _tm(initial_chunk_frames=1)
+    emitted = 0
+    for count in range(1, frames + 1):
+        payload = _call(tm, "r", n_frames=count)
+        if payload is not None:
+            emitted += payload.codes.audio.numel() // _Q
+            tm.put_req_chunk["r"] += 1
+    terminal = _call(tm, "r", n_frames=frames, finished=True)
+    assert terminal is not None and terminal.meta.finished.item()
+    assert terminal.codes.audio.numel() == 0
+    assert emitted == frames
+
+
+def test_repeated_nonterminal_callback_does_not_republish_boundary():
+    tm = _tm(initial_chunk_frames=1)
+    first = _call(tm, "r", n_frames=1)
+    assert first is not None
+    tm.put_req_chunk["r"] += 1
+    assert _call(tm, "r", n_frames=1) is None
+    next_chunk = _call(tm, "r", n_frames=26)
+    assert next_chunk.codes.audio.numel() // _Q == 25
+
+
+def test_terminal_flush_contains_only_frames_not_previously_emitted():
+    tm = _tm(initial_chunk_frames=1)
+    _call(tm, "r", n_frames=1)
+    tm.put_req_chunk["r"] += 1
+    _call(tm, "r", n_frames=26)
+    tm.put_req_chunk["r"] += 1
+    final = _call(tm, "r", n_frames=30, finished=True)
+    assert final.codes.audio.numel() // _Q == 4
+    tm.put_req_chunk["r"] += 1
+    assert _call(tm, "r", n_frames=30, finished=True).codes.audio.numel() == 0
+
+
+def _req(
+    rid,
+    *,
+    finished,
+    initial_codec_chunk_frames=None,
+    non_streaming_mode=None,
+    full_utterance_decode=None,
+    output_token_ids=None,
+):
     entries = {}
     if initial_codec_chunk_frames is not None:
         entries["initial_codec_chunk_frames"] = SimpleNamespace(list_data=[initial_codec_chunk_frames])
@@ -40,11 +88,11 @@ def _req(rid, *, finished, initial_codec_chunk_frames=None, non_streaming_mode=N
         entries["non_streaming_mode"] = SimpleNamespace(list_data=[non_streaming_mode])
     if full_utterance_decode is not None:
         entries["full_utterance_decode"] = SimpleNamespace(list_data=[full_utterance_decode])
-    ai = SimpleNamespace(entries=entries) if entries else None
     return SimpleNamespace(
         external_req_id=rid,
         is_finished=lambda: finished,
-        additional_information=ai,
+        additional_information=SimpleNamespace(entries=entries) if entries else None,
+        output_token_ids=output_token_ids or [],
     )
 
 
@@ -206,6 +254,112 @@ def test_voicedesign_prompt_mode_with_streaming_keeps_windowed_emit():
     assert len(p.codes.audio) == _Q * 25
 
 
+def test_async_chunk_skips_stop_token_zero_frame_without_gpu_value_check():
+    tm = _tm()
+    rid = "r-stop"
+    tm.code_prompt_token_ids[rid] = [_FRAME[:]]
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        multimodal_output={"codes": {"audio": torch.zeros((1, _Q), dtype=torch.long)}},
+        request=_req(rid, finished=True, initial_codec_chunk_frames=2, output_token_ids=[2150]),
+        is_finished=True,
+    )
+
+    assert payload is not None
+    assert payload.codes.audio.tolist() == _FRAME
+    assert len(tm.code_prompt_token_ids[rid]) == 1
+
+
+def test_async_chunk_uses_authoritative_frame_validity_without_token_history():
+    tm = _tm()
+    rid = "r-native-validity"
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        multimodal_output={
+            "codes": {"audio": torch.tensor([[1, 2, 3, 4]])},
+            "meta": {"codec_frame_valid": False},
+        },
+        request=_req(rid, finished=False, initial_codec_chunk_frames=1, output_token_ids=[]),
+        is_finished=False,
+    )
+
+    assert payload is None
+    assert tm.code_prompt_token_ids[rid] == []
+
+
+def test_async_chunk_keeps_final_real_frame_when_eos_was_just_sampled():
+    tm = _tm()
+    rid = "r-final-real-frame"
+
+    payload = talker2code2wav_async_chunk(
+        transfer_manager=tm,
+        multimodal_output={
+            "codes": {"audio": torch.tensor([[1, 2, 3, 4]])},
+            "meta": {"codec_frame_valid": True},
+        },
+        request=_req(rid, finished=True, initial_codec_chunk_frames=1, output_token_ids=[4198]),
+        is_finished=True,
+    )
+
+    assert payload is not None
+    assert payload.codes.audio.tolist() == [1, 2, 3, 4]
+
+
+def test_async_chunk_batch_matches_scalar_payloads_and_state():
+    batch_builder = getattr(qwen3_tts_processors, "talker2code2wav_async_chunk_batch", None)
+    assert callable(batch_builder), "Qwen3-TTS must provide the native MRv2 batch payload builder"
+
+    scalar_tm = _tm(chunk_frames=4, left_context=3, initial_chunk_frames=1)
+    batch_tm = _tm(chunk_frames=4, left_context=3, initial_chunk_frames=1)
+    requests = [
+        _req("r-batch-0", finished=False, output_token_ids=[7]),
+        _req("r-batch-1", finished=False, output_token_ids=[8]),
+    ]
+    ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
+    # Request 2 has a multi-token span: the last validity flag wins the chunk.
+    outputs = [
+        {"codes": {"audio": torch.tensor([[1, 2, 3, 4]]), "ref": ref_code}},
+        {
+            "codes": {"audio": torch.tensor([[9, 9, 9, 9], [5, 6, 7, 8]])},
+            "meta": {"codec_frame_valid": torch.tensor([False, True])},
+        },
+    ]
+
+    scalar_payloads = [
+        talker2code2wav_async_chunk(
+            transfer_manager=scalar_tm,
+            multimodal_output=output,
+            request=request,
+            is_finished=False,
+        )
+        for request, output in zip(requests, outputs)
+    ]
+    batch_payloads = batch_builder(
+        transfer_manager=batch_tm,
+        pooling_outputs=outputs,
+        requests=requests,
+        is_finished=[False, False],
+    )
+
+    assert len(batch_payloads) == len(scalar_payloads)
+    for scalar, batched in zip(scalar_payloads, batch_payloads):
+        assert scalar is not None and batched is not None
+        torch.testing.assert_close(batched.codes.audio, scalar.codes.audio)
+        assert batched.meta.left_context_size == scalar.meta.left_context_size
+        assert batched.meta.finished.item() == scalar.meta.finished.item()
+        assert batched.speaker == scalar.speaker
+        assert batched.language == scalar.language
+    for request in requests:
+        rid = request.external_req_id
+        torch.testing.assert_close(
+            torch.stack(batch_tm.code_prompt_token_ids[rid]),
+            torch.stack(scalar_tm.code_prompt_token_ids[rid]),
+        )
+    torch.testing.assert_close(batch_tm.request_payload["r-batch-0"], scalar_tm.request_payload["r-batch-0"])
+
+
 _CASES = [
     # ── IC boundary rule ──────────────────────────────────────────────
     # initial_codec_chunk_frames only controls the first emitted chunk.
@@ -703,35 +857,27 @@ class TestRampHelpers:
     def test_parse_ramp_mixed_list_returns_none(self):
         assert parse_chunk_ramp({"codec_chunk_ramp": [4, "x"]}) is None
 
-    def test_parse_ramp_warns_on_tail_mismatch(self, caplog):
-        import logging
+    def test_parse_ramp_warns_on_tail_mismatch(self, monkeypatch):
+        warnings = []
+        monkeypatch.setattr(
+            "vllm_omni.model_executor.stage_input_processors.chunk_size_utils.logger.warning",
+            lambda message, *args: warnings.append(message % args),
+        )
 
-        target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.chunk_size_utils")
-        target_logger.addHandler(caplog.handler)
-        prev_level = target_logger.level
-        target_logger.setLevel(logging.WARNING)
-        try:
-            result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8]}, steady=25)
-        finally:
-            target_logger.removeHandler(caplog.handler)
-            target_logger.setLevel(prev_level)
+        result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8]}, steady=25)
         assert result == [4, 4, 8]
-        assert "reintroduces" in caplog.text
+        assert any("reintroduces" in message for message in warnings)
 
-    def test_parse_ramp_no_warn_on_tail_match(self, caplog):
-        import logging
+    def test_parse_ramp_no_warn_on_tail_match(self, monkeypatch):
+        warnings = []
+        monkeypatch.setattr(
+            "vllm_omni.model_executor.stage_input_processors.chunk_size_utils.logger.warning",
+            lambda message, *args: warnings.append(message % args),
+        )
 
-        target_logger = logging.getLogger("vllm_omni.model_executor.stage_input_processors.chunk_size_utils")
-        target_logger.addHandler(caplog.handler)
-        prev_level = target_logger.level
-        target_logger.setLevel(logging.WARNING)
-        try:
-            result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8, 16, 25]}, steady=25)
-        finally:
-            target_logger.removeHandler(caplog.handler)
-            target_logger.setLevel(prev_level)
+        result = parse_chunk_ramp({"codec_chunk_ramp": [4, 4, 8, 16, 25]}, steady=25)
         assert result == [4, 4, 8, 16, 25]
-        assert "reintroduces" not in caplog.text
+        assert warnings == []
 
     @pytest.mark.parametrize(
         "index,ramp,steady,expected",
@@ -1028,6 +1174,7 @@ class TestChunkRampEmission:
         tm.put_req_chunk[rid] = 2
 
         tm.code_prompt_token_ids[rid] = []
+        tm._qwen3_tts_emitted_frames.pop(rid, None)
         tm.ramp_chunk_count.pop(rid, None)
 
         p_seg2_0 = self._emit(tm, rid, 4)
@@ -1051,6 +1198,7 @@ class TestChunkRampEmission:
         tm.put_req_chunk[rid] = 1
 
         tm.code_prompt_token_ids[rid] = []
+        tm._qwen3_tts_emitted_frames.pop(rid, None)
         tm.ramp_chunk_count.pop(rid, None)
 
         p_seg2_0 = self._emit(tm, rid, 4)
@@ -1429,6 +1577,7 @@ class TestAdaptiveEmission:
         tm.put_req_chunk[rid] = 1
 
         tm.code_prompt_token_ids[rid] = []
+        tm._qwen3_tts_emitted_frames.pop(rid, None)
         tm.ramp_chunk_count.pop(rid, None)
         tm._adaptive_states.pop(rid, None)
 

tests/worker/test_min_tokens_stop_ids_guard.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Guard against out-of-vocabulary stop ids in min-tokens masking (#4962).
 
 vLLM's input processor folds the stage tokenizer's EOS id into
@@ -19,7 +19,7 @@ from vllm.v1.sample.logits_processor import (
     MinTokensLogitsProcessor,
 )
 
-from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids
+from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids, sanitize_sampling_params_min_tokens_stop_ids
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -28,6 +28,12 @@ CODEC_EOS = 2150
 TEXT_EOS = 151645
 
 
+@pytest.fixture(autouse=True)
+def _disable_vllm_pin_memory_for_cpu(monkeypatch: pytest.MonkeyPatch):
+    """Keep the module-level H2D helper in CPU-only mode."""
+    monkeypatch.setattr("vllm.utils.torch_utils.PIN_MEMORY", False)
+
+
 def _make_min_tokens_proc(stop_token_ids: list[int], eos_token_id: int | None) -> MinTokensLogitsProcessor:
     """Build a MinTokensLogitsProcessor with one active request, mirroring
     the production path (engine folds the tokenizer EOS via
@@ -113,3 +119,14 @@ def test_guard_handles_multiple_requests():
     assert out[0, CODEC_EOS] == -float("inf")
     assert out[1, 7] == -float("inf")
     assert int((out == -float("inf")).sum()) == 2
+
+
+def test_sampling_params_guard_filters_only_unreachable_min_tokens_ids():
+    params = SamplingParams(min_tokens=2, stop_token_ids=[CODEC_EOS])
+    params.update_from_generation_config({}, TEXT_EOS)
+
+    sanitize_sampling_params_min_tokens_stop_ids(params, TALKER_VOCAB)
+
+    assert params.all_stop_token_ids == {CODEC_EOS}
+    assert params.stop_token_ids == [CODEC_EOS]
+    assert params.eos_token_id == TEXT_EOS

tests/worker/test_omni_connector_mixin.py

@@ -1724,3 +1724,16 @@ class TestSendRetry(unittest.TestCase):
 
 if __name__ == "__main__":
     unittest.main()
+
+
+@pytest.mark.parametrize(
+    "fallback,external,expected", [("r", None, "mapped"), (None, "ext", "ext"), (None, None, None)]
+)
+def test_resolve_request_id_requires_internal_or_external_id(fallback, external, expected):
+    host = SimpleNamespace(_request_ids_mapping={"r": "mapped"})
+    request = SimpleNamespace(external_req_id=external)
+    if expected is None:
+        with pytest.raises(ValueError, match="request ID"):
+            OmniConnectorModelRunnerMixin._resolve_external_req_id(host, request, fallback)
+    else:
+        assert OmniConnectorModelRunnerMixin._resolve_external_req_id(host, request, fallback) == expected

tests/worker/test_omni_gpu_model_runner.py

@@ -770,8 +770,10 @@ def test_accumulate_full_payload_output_keeps_all_zero_qwen3_omni_prefill_placeh
 def test_full_payload_output_accumulation_hook_matrix():
     assert _make_full_payload_accumulation_runner(model_stage="thinker")._should_accumulate_full_payload_output()
     assert _make_full_payload_accumulation_runner(model_stage="talker")._should_accumulate_full_payload_output()
-    assert not _make_full_payload_accumulation_runner(
-        model_stage="code2wav", final_output=True
+    # A stage may publish a user-visible output and still feed another stage
+    # (Qwen3-Omni thinker publishes text while sending audio state to talker).
+    assert _make_full_payload_accumulation_runner(
+        model_stage="thinker", final_output=True
     )._should_accumulate_full_payload_output()
     assert not _make_full_payload_accumulation_runner(
         model_stage="token2audio",

tests/worker/test_omni_v2_profile_dispatch.py

@@ -0,0 +1,60 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+from vllm.compilation.cuda_graph import CUDAGraphMode
+
+from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+@pytest.mark.parametrize("dp_size", [1, 2])
+def test_dispatch_eager_and_manager_paths(monkeypatch, dp_size):
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.cudagraph_manager = None
+    runner.dp_size = dp_size
+    runner.dp_rank = 1
+
+    import vllm.v1.worker.gpu.dp_utils as dp_utils
+
+    expected_tokens = object()
+    sync_calls = []
+
+    def sync_padding(*args, **kwargs):
+        sync_calls.append((args, kwargs))
+        return args[1], expected_tokens
+
+    monkeypatch.setattr(
+        dp_utils,
+        "sync_cudagraph_and_dp_padding",
+        sync_padding,
+    )
+
+    batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor(
+        num_reqs=1, num_toks=8, uniform_tok_count=8, num_active_loras=0, use_eager=True, max_query_len=8
+    )
+    assert batch_desc.cg_mode == CUDAGraphMode.NONE
+    if dp_size == 1:
+        assert num_tokens_across_dp is None
+        assert sync_calls == []
+    else:
+        assert num_tokens_across_dp is expected_tokens
+
+    # Non-eager dispatch goes through the cudagraph manager.
+    expected = SimpleNamespace(cg_mode=CUDAGraphMode.PIECEWISE, num_tokens=8, num_reqs=1)
+    dispatch_calls = []
+
+    def dispatch(*args, **kwargs):
+        dispatch_calls.append((args, kwargs))
+        return expected
+
+    runner.cudagraph_manager = SimpleNamespace(dispatch=dispatch)
+    runner.dp_size = 1
+    batch_desc, _ = runner._dispatch_batch_descriptor(
+        num_reqs=1, num_toks=8, uniform_tok_count=8, num_active_loras=0, use_eager=False, max_query_len=8
+    )
+    assert batch_desc is expected
+    assert dispatch_calls == [((1, 8, 8), {"num_active_loras": 0, "max_query_len": 8})]

tests/worker/test_tts_sampling_seed.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.worker import sampling_utils
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_fixed_seed_contract_separates_stage_sampler_from_talker_mtp_rng() -> None:
+    assert hasattr(sampling_utils, "apply_fixed_seed_to_sampling_params")
+    assert hasattr(sampling_utils, "get_tts_local_seed")
+
+    stage_only = SimpleNamespace(seed=None, extra_args=None)
+    sampling_utils.apply_fixed_seed_to_sampling_params(
+        stage_only,
+        17,
+        seed_talker_mtp=False,
+    )
+    assert stage_only.seed == 17
+    assert sampling_utils.get_tts_local_seed(stage_only) is None
+
+    talker = SimpleNamespace(seed=None, extra_args={"preserved": True})
+    sampling_utils.apply_fixed_seed_to_sampling_params(
+        talker,
+        23,
+        seed_talker_mtp=True,
+    )
+    assert talker.seed == 23
+    assert talker.extra_args == {"preserved": True, "tts_local_seed": 23}
+    assert sampling_utils.get_tts_local_seed(talker) == 23

tests/worker/test_worker_base.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Characterization tests for ``vllm_omni.worker.base.OmniGPUWorkerBase``.
 
 Pins the CURRENT behaviour of ``determine_available_memory`` (device-level

tests/worker_v2/test_intermediate_buffer.py

@@ -0,0 +1,96 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Unit tests for OmniIntermediateBuffer lifecycle.
+
+Retained contracts: wire decode with malformed-payload rejection, slot
+lifecycle with index migration on reuse, gather in batch order, update merge
+semantics (cpu/gpu-resident keys), single owned snapshot for
+update_gpu_tensor_rows, and decode-span merge until ack.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.engine import (
+    AdditionalInformationEntry,
+    AdditionalInformationPayload,
+    PromptEmbedsPayload,
+)
+from vllm_omni.worker_v2.model_states.intermediate_buffer import (
+    OmniIntermediateBuffer,
+    _resolve_additional_information,
+    _resolve_prompt_embeds,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_new_req_data(req_id="r1", **kwargs):
+    kwargs.setdefault("mm_features", [])
+    return SimpleNamespace(req_id=req_id, **kwargs)
+
+
+def test_wire_decode_passthrough_and_malformed_rejection():
+    tensor = torch.randn(3, 4)
+    assert _resolve_prompt_embeds(tensor) is not None
+    assert _resolve_prompt_embeds(None) is None
+    with pytest.raises(ValueError, match="Failed to decode prompt_embeds payload"):
+        _resolve_prompt_embeds(PromptEmbedsPayload(data=b"\x00", shape=[2, 2], dtype="float32"))
+
+    assert _resolve_additional_information({"key": "val"}) == {"key": "val"}
+    assert _resolve_additional_information(None) == {}
+    payload = AdditionalInformationPayload(
+        entries={
+            "embed.decode": AdditionalInformationEntry(tensor_data=b"\x00", tensor_shape=[2], tensor_dtype="float32")
+        }
+    )
+    with pytest.raises(ValueError, match="Failed to decode additional_information payload"):
+        _resolve_additional_information(payload)
+
+
+def test_slot_lifecycle_fields_reuse_and_remove_idempotent():
+    buf = OmniIntermediateBuffer(max_num_reqs=4)
+    embeds = torch.randn(2, 8)
+    features = [torch.randn(3, 4)]
+    params = SimpleNamespace(extra_args={"tts_local_seed": 1234})
+    buf.add_request(0, _make_new_req_data("r0", prompt_embeds=embeds, mm_features=features, sampling_params=params))
+    assert torch.equal(buf.buffers[0]["prompt_embeds_cpu"], embeds)
+    assert buf.buffers[0]["mm_features"] is features
+    assert buf.buffers[0]["sampling_params"] is params
+
+    buf.add_request(2, _make_new_req_data("r2"))
+    buf.add_request(3, _make_new_req_data("r3"))
+    batch = SimpleNamespace(idx_mapping_np=[2, 0, 3])
+    assert [g["req_id"] for g in buf.gather(batch)] == ["r2", "r0", "r3"]
+
+    # Slot reuse: req_id mapping migrates; duplicate removal is a no-op.
+    buf.add_request(0, _make_new_req_data("new"))
+    assert buf.req_id_to_index["new"] == 0
+    buf.remove_request(0)
+    buf.remove_request(0)
+    assert "new" not in buf.req_id_to_index
+    assert buf.buffers[0] == {}
+
+
+def test_update_merge_semantics(monkeypatch):
+    buf = OmniIntermediateBuffer(max_num_reqs=1)
+    buf.add_request(0, _make_new_req_data("r0"))
+
+    buf.update(0, {"hidden": torch.randn(4)})
+    assert buf.buffers[0]["hidden"].device == torch.device("cpu")  # default migrates to cpu
+    buf.update(0, {})
+    buf.update(0, {"a": torch.tensor([1.0])})
+    buf.update(0, {"b": torch.tensor([2.0])})
+    assert "a" in buf.buffers[0] and "b" in buf.buffers[0]
+
+    # GPU-resident keys keep values (including list items) off the cpu path.
+    tensor = torch.randn(2)
+    monkeypatch.setattr(torch.Tensor, "cpu", lambda *_a, **_k: pytest.fail("GPU-resident value must not move to CPU"))
+    buf.update(0, {"items": [tensor], "kv": torch.randn(4)}, gpu_resident_keys={"items", "kv"})
+    stored = buf.buffers[0]["items"][0]
+    assert stored is not tensor
+    assert torch.equal(stored, tensor)
+    assert "kv" in buf.buffers[0]

tests/worker_v2/test_intermediate_buffer_cuda_residency.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""CUDA-lane checks for GPU-resident intermediate buffer values."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.parametrize("key,dtype", [(("embed", "prefill"), torch.float32), (("codes", "ref"), torch.long)])
+def test_gpu_resident_values_stay_on_device_after_update(key, dtype):
+    # Declared gpu-resident keys must keep tensors on device (CLA snapshot real).
+    buffer = OmniIntermediateBuffer(1)
+    value = torch.ones(2, 3, device="cuda", dtype=dtype)
+    buffer.add_request(0, SimpleNamespace(req_id="r1", mm_features=[]))
+
+    buffer.update(0, {key[0]: {key[1]: value}}, {key})
+
+    stored = buffer.buffers[0][key[0]][key[1]]
+    assert stored.device.type == "cuda"
+    assert stored.dtype == dtype
+    assert torch.equal(stored.cpu(), value.cpu())

tests/worker_v2/test_omni_ar_model_runner.py

@@ -0,0 +1,253 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Unit tests for OmniARModelRunner v2: async output staging, snapshot ownership, payload slicing."""
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+from typing import Any
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+import torch
+from vllm.v1.outputs import RoutedExpertsTensors
+from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors
+
+import vllm_omni.worker_v2.omni_ar_model_runner as omni_ar_model_runner
+from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner, OmniAsyncOutput
+from vllm_omni.worker_v2.output_snapshot import pack_output_snapshot
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeStream:
+    def wait_stream(self, _stream) -> None:
+        pass
+
+
+class _FakeEvent:
+    def record(self, _stream) -> None:
+        pass
+
+    def synchronize(self) -> None:
+        pass
+
+
+def _async_output(req_ids=("req-0",), **overrides) -> OmniAsyncOutput:
+    rid2idx = {rid: i for i, rid in enumerate(req_ids)}
+    mro = omni_ar_model_runner.OmniModelRunnerOutput(list(req_ids), rid2idx, None, prompt_logprobs_dict={})
+    sampler_output = SamplerOutput(torch.tensor([[123]]), None, None, None)
+    kwargs = dict(model_runner_output=mro, sampler_output=sampler_output)
+    kwargs.update(num_sampled_tokens=torch.tensor([1] * len(req_ids)), copy_event=_FakeEvent())
+    kwargs.update(main_stream=_FakeStream(), copy_stream=_FakeStream())
+    return OmniAsyncOutput(**(kwargs | overrides))
+
+
+def test_async_output_blocking_event_and_routing_masks(monkeypatch) -> None:
+    event_kwargs = []
+    monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
+
+    def make_event(**kwargs):
+        event_kwargs.append(kwargs)
+        return _FakeEvent()
+
+    monkeypatch.setattr(torch.cuda, "Event", make_event)
+    masks = SamplingMaskTensors(torch.tensor([[5], [0]], dtype=torch.uint8), torch.tensor([2, 0]), 4)
+    sampler_output = SamplerOutput(torch.tensor([[2], [0]]), None, None, torch.tensor([1, 0]), None, masks)
+    routed = RoutedExpertsTensors(torch.tensor([[[2, 3]], [[4, 5]]]), torch.tensor([7, 9]))
+
+    output = _async_output(
+        req_ids=["decode", "prefill"],
+        sampler_output=sampler_output,
+        num_sampled_tokens=torch.tensor([1, 0]),
+        copy_event=None,  # None → constructor builds the default blocking event
+        routed_experts=routed,
+    ).get_output()
+    assert event_kwargs == [{"blocking": True}]  # blocking event by default
+    assert output.sampled_token_ids == [[2], []] and output.sampling_masks.cu_num_generated_tokens == [0, 1, 1]
+    np.testing.assert_array_equal(output.routed_experts.routing_data, [[[2, 3]], [[4, 5]]])
+    np.testing.assert_array_equal(output.routed_experts.slot_mapping, [7, 9])
+    np.testing.assert_array_equal(output.sampling_masks.token_ids, [0, 2])
+
+
+@pytest.mark.parametrize("needs_history", [False, True])
+def test_last_pp_rank_orchestration_and_kv_resolver(monkeypatch, needs_history) -> None:
+    runner = OmniARModelRunner.__new__(OmniARModelRunner)
+    input_batch = SimpleNamespace(req_ids=["req"], num_reqs=1, seq_lens=torch.tensor([3]))
+    input_batch.idx_mapping, input_batch.query_start_loc = torch.tensor([0]), torch.tensor([0, 1])
+    state = SimpleNamespace(input_batch=input_batch, hidden_states=torch.zeros(1, 2))
+    state.finished_req_ids, state.ec_connector_output, state.routed_experts = {"finished"}, None, None
+    runner.execute_model_state = state
+    runner._kv_extracted_req_ids = runner._last_aux_output = runner._last_multimodal_outputs = None
+    runner.is_last_pp_rank, runner.pp_handler, runner.check_ep_fault = True, None, False
+    runner.model_config = SimpleNamespace(async_chunk=False)
+    runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(engine_output_type="text"))
+    runner.model_state = SimpleNamespace(postprocess_model_output=MagicMock(return_value=(torch.zeros(1, 2), None)))
+    runner.model_state.intermediate_buffer = SimpleNamespace(buffers={0: {"global_request_id": "global-req"}})
+    runner.req_states = SimpleNamespace(req_id_to_index={"req": 0})
+    runner.req_states.all_token_ids = SimpleNamespace(gpu=torch.tensor([[1]]))
+    runner.req_states.num_computed_tokens = SimpleNamespace(gpu=torch.tensor([0]))
+    runner.req_states.prompt_len = SimpleNamespace(np=np.array([1]), gpu=torch.tensor([1]))
+    runner.main_stream = runner.output_copy_stream = MagicMock()
+    runner.eplb = runner._finalize_native_data_plane_output = runner._reserve_native_data_plane_outputs = MagicMock()
+    sampler_out = (SimpleNamespace(sampled_token_ids=torch.tensor([[2]])), MagicMock(), MagicMock())
+    sampling_active = False
+
+    @contextmanager
+    def sampling_context(*, req_ids, num_output_tokens):
+        nonlocal sampling_active
+        assert needs_history and req_ids == ["req"]
+        sampling_active = True
+        yield
+        sampling_active = False
+
+    runner.model = SimpleNamespace(compute_logits=None, logitsprocs_need_output_token_ids=needs_history)
+    runner.model.mrv2_sampling_context = sampling_context
+    runner.sample = MagicMock(return_value=sampler_out)
+    runner.sample.side_effect = lambda *_: sampler_out if sampling_active is needs_history else pytest.fail()
+    logprobs_mock = MagicMock(side_effect=lambda *_: pytest.fail("inside ctx") if sampling_active else {})
+    runner.prompt_logprobs_worker = SimpleNamespace(compute_prompt_logprobs=logprobs_mock)
+    runner.postprocess_sampled, connector_output = MagicMock(), object()
+
+    def post_forward(finished_req_ids):
+        runner.postprocess_sampled.assert_called_once()  # postprocess precedes kv_connector.post_forward
+        assert finished_req_ids == {"finished"}
+        return connector_output
+
+    runner.kv_connector = SimpleNamespace(post_forward=MagicMock(side_effect=post_forward))
+    mock_out = SimpleNamespace(copy_event=None)
+    monkeypatch.setattr(omni_ar_model_runner, "OmniAsyncOutput", MagicMock(return_value=mock_out))
+
+    assert runner.sample_tokens(None) is mock_out
+    built = omni_ar_model_runner.OmniAsyncOutput.call_args.kwargs["model_runner_output"]
+    assert built.kv_connector_output is connector_output
+    assert runner._resolve_global_request_id("req") == "global-req"  # from the intermediate buffer
+    assert runner._resolve_global_request_id("unknown") == "unknown"  # fallback to the local id
+
+
+def test_async_mm_snapshot_owns_output_until_copy_finishes() -> None:
+    runner = OmniARModelRunner.__new__(OmniARModelRunner)
+    runner.model_config = SimpleNamespace(async_chunk=True)
+    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] = []
+    runner.main_stream = SimpleNamespace(wait_event=waited.append)
+    source = torch.tensor([[7, 8]], dtype=torch.long)
+
+    snapshot = runner._retain_multimodal_outputs({"codes": {"audio": source}})
+    source.fill_(99)
+
+    # Snapshot owns the data (graph replay cannot overwrite it)...
+    snap = snapshot["codes"]["audio"]
+    assert snap.tolist() == [[7, 8]] and snap.data_ptr() != source.data_ptr()
+    assert runner._last_multimodal_snapshot_slot == 0
+    # ...and slot reuse waits for the previous D2H copy event.
+    runner._release_multimodal_snapshot(0, copy_event := object())
+    runner._retain_multimodal_outputs({"codes": {"audio": torch.zeros(1, 2)}})
+    assert waited == [copy_event]
+
+
+def test_snapshot_slots_bounded_by_shape_and_packed_grouping_isolation(monkeypatch) -> None:
+    slot: dict[tuple[Any, ...], torch.Tensor] = {}
+    omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(1, 2), slot)
+    omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(4, 2), slot)
+    assert len(slot) == 2  # separate shapes do not share a buffer
+    monkeypatch.setattr(omni_ar_model_runner, "_ASYNC_MM_SNAPSHOT_MAX_BUCKETS_PER_SLOT", 1)
+    bounded: dict[tuple[Any, ...], torch.Tensor] = {}
+    kept = omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(2, 2), bounded)
+    overflow = omni_ar_model_runner._copy_mm_to_snapshot_slot(torch.ones(5, 2), bounded)
+    assert len(bounded) == 1 and overflow.data_ptr() != kept.data_ptr()  # overflow clones without evicting
+
+    source = torch.arange(12, dtype=torch.int64).view(3, 4).t()
+    payload = {"noncontiguous": source, "nested": [torch.tensor(7), (torch.tensor([1.5]),)], "meta": "ok"}
+    pack_slot: dict[tuple[Any, ...], torch.Tensor] = {}
+    snapshot = pack_output_snapshot(payload, pack_slot, max_buckets=4)
+    source.fill_(99)  # ownership: snapshot keeps pre-mutation values
+    copies = []
+
+    def copy(tensor):
+        copies.append(tensor.numel())
+        return tensor.clone()
+
+    host = snapshot.copy_to_cpu(copy)
+    assert len(copies) == 2  # grouped by dtype (int64, float32), not tensor count
+    assert host["noncontiguous"].tolist() == torch.arange(12).view(3, 4).t().tolist()
+    assert host["nested"][0].item() == 7 and isinstance(host["nested"][1], tuple) and host["meta"] == "ok"
+    pack_output_snapshot(payload, pack_slot, max_buckets=4)
+    assert host["noncontiguous"][0, 0].item() == 0  # repacking must not corrupt published host data
+
+
+@pytest.mark.parametrize("need_pooler,async_chunk", [(False, False), (False, True), (True, False), (True, True)])
+def test_guard_graph_replay_for_pooler_copy(need_pooler, async_chunk) -> None:
+    main_stream = MagicMock()
+    omni_ar_model_runner._guard_graph_replay_for_pooler_copy(
+        main_stream, object(), need_pooler=need_pooler, async_chunk=async_chunk
+    )
+    # Non-async pooler copies must gate the next graph replay on the copy event.
+    assert main_stream.wait_event.call_count == (1 if need_pooler and not async_chunk else 0)
+
+
+def test_build_async_chunk_outputs_slices_padded_axis_and_splits_channels() -> None:
+    # Graph-padded batch: padded_total_tokens > total_tokens; slice by the real token axis.
+    padded_codes = torch.arange(16, dtype=torch.long).reshape(8, 2)
+    build = OmniARModelRunner._build_async_chunk_outputs_from_mm
+    inter_stage, client = build({"codes": {"audio": padded_codes}}, np.array([0, 1, 2]), np.array([1, 1]), 2, 2, 8)
+    assert client is None
+    assert torch.equal(inter_stage[0]["codes.audio"], padded_codes[0:1])
+    assert torch.equal(inter_stage[1]["codes.audio"], padded_codes[1:2])
+
+    codes, audio = torch.arange(16, dtype=torch.long).reshape(4, 4), torch.randn(4, 8)
+    req_codes = [torch.arange(16 * i, 16 * (i + 1), dtype=torch.long).reshape(1, 16) for i in range(2)]
+    inter_stage, client = build(
+        {"codes": {"audio": codes}, "audio": audio}, np.array([0, 2, 4]), np.array([2, 2]), 2, 4
+    )
+    assert "hidden" not in inter_stage[0] and torch.equal(client[1]["audio"], audio[2:])  # channel split
+    assert torch.equal(inter_stage[1]["codes.audio"], codes[2:])  # inter-stage channel
+    # Per-request code lists pass through unsliced.
+    inter_stage, client = build({"codes": {"audio": req_codes}}, np.array([0, 1, 2]), np.array([1, 1]), 2, 2)
+    assert client is None and torch.equal(inter_stage[0]["codes.audio"], req_codes[0])
+
+
+def test_async_chunk_output_stages_mm_on_copy_stream_before_get_output(monkeypatch) -> None:
+    monkeypatch.setattr(torch.cuda, "set_stream", lambda _stream: None)
+    calls = []
+
+    def copy_mm(mm_outputs, total_tokens, **ctx):
+        calls.append((total_tokens, ctx))
+        return {"codes": {"audio": mm_outputs["codes"]["audio"].clone()}}
+
+    monkeypatch.setattr(omni_ar_model_runner, "_async_copy_mm", copy_mm)
+    source_codes = torch.tensor([[7, 8]], dtype=torch.long)
+    input_batch = SimpleNamespace(query_start_loc_np=np.array([0, 1]), num_scheduled_tokens=[1], num_reqs=1)
+    output = _async_output(
+        multimodal_outputs={"codes": {"audio": source_codes}},
+        input_batch=input_batch,
+        copy_stream=(copy_stream := _FakeStream()),
+        async_chunk=True,
+    )
+    # Staged once on the copy stream during construction, with one resolved pin-memory context for all D2H helpers.
+    [(total_tokens, ctx)] = calls
+    assert total_tokens == 1 and ctx["copy_stream"] is copy_stream and ctx["pin_memory"] is not None
+    assert output._mm_snapshot["codes"]["audio"].device.type == "cpu"
+    source_codes.fill_(99)  # a later graph replay cannot leak into the snapshot
+    assert torch.equal(output.get_output().inter_stage_outputs[0]["codes.audio"], torch.tensor([[7, 8]]))
+
+
+@pytest.mark.parametrize("prefill_first", [False, True])
+@pytest.mark.parametrize("padded", [False, True])
+def test_request_reference_codes_preserve_local_axis(prefill_first, padded):
+    lengths = np.array([255, 1] if prefill_first else [1, 255])
+    offsets = np.array([0, lengths[0], 256])
+    size = 512 if padded else 256
+    ref = torch.arange(size * 16).reshape(size, 16)
+    refs = [ref, torch.empty(0)] if prefill_first else [torch.empty(0), ref]
+    codes = torch.arange(size * 16).reshape(size, 16)
+    outputs, _ = OmniARModelRunner._build_async_chunk_outputs_from_mm(
+        {"codes": {"audio": codes, "ref": refs}}, offsets, lengths, 2, 256, size
+    )
+    index = 0 if prefill_first else 1
+    assert torch.equal(outputs[index]["codes.ref"], ref)
+    assert outputs[index]["codes.ref"].data_ptr() != ref.data_ptr()
+    for i in range(2):
+        assert torch.equal(outputs[i]["codes.audio"], codes[offsets[i] : offsets[i + 1]])

tests/worker_v2/test_omni_data_plane.py

@@ -0,0 +1,481 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Lifecycle contracts for OmniRunnerDataPlane/NativeOutputWorker."""
+
+import threading
+import time
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+
+from vllm_omni.worker_v2.delivery import DeliveryCancelledError, DeliveryState, OmniDeliveryManager
+from vllm_omni.worker_v2.native_output_worker import NativeOutputWorker
+from vllm_omni.worker_v2.omni_data_plane import OmniRunnerDataPlane
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _complete(plane, payloads, token=21, req="internal"):
+    return plane.complete_outputs(
+        req_ids=[req], inter_stage_outputs=payloads, sampled_token_ids=[[token]] * len(payloads)
+    )
+
+
+def _new_request(req_id="internal", external_req_id="external"):
+    return SimpleNamespace(
+        req_id=req_id,
+        external_req_id=external_req_id,
+        prompt_token_ids=[10, 11],
+        num_computed_tokens=2,
+        resumable=True,
+        additional_information=SimpleNamespace(entries={}),
+        sampling_params=SimpleNamespace(stop_token_ids=[2150]),
+    )
+
+
+def _bare_plane(*, delivery_timeout_s=1.0, shutdown_timeout_s=1.0):
+    plane = object.__new__(OmniRunnerDataPlane)
+    plane.__dict__.update(
+        _native_requests={},
+        _native_output_lock=threading.RLock(),
+        _native_send_lock=threading.Lock(),
+        _native_output_error=None,
+        _native_output_error_lock=threading.Lock(),
+        _native_output_closed=False,
+        _native_outputs_in_flight=defaultdict(int),
+        _native_terminal_pending=set(),
+        _put_req_chunk=defaultdict(int),
+        _ramp_chunk_count=defaultdict(int),
+        _delivery_manager=OmniDeliveryManager(
+            delivery_timeout_s=delivery_timeout_s, shutdown_timeout_s=shutdown_timeout_s
+        ),
+    )
+    return plane
+
+
+def _make_plane(replace_send):
+    p = _bare_plane()
+    p.record = SimpleNamespace(batches=[], cleaned=[])
+    if replace_send:
+        p.send_chunks = lambda entries, **_kw: p.record.batches.append(entries) or len(entries)
+    p.get_omni_connector_output, p.cleanup_finished_request = (lambda: None), p.record.cleaned.append
+    return p
+
+
+@pytest.fixture
+def raw_plane():  # keeps the real connector send path
+    p = _make_plane(replace_send=False)
+    yield p
+    if getattr(p, "_native_output_worker", None) is not None:
+        p._stop_output_worker()
+
+
+@pytest.fixture
+def plane():
+    p = _make_plane(replace_send=True)
+    yield p
+    if getattr(p, "_native_output_worker", None) is not None:
+        p._stop_output_worker()
+
+
+def _gated_connector():
+    # put() blocks its first call until released (or forever without release).
+    conn = SimpleNamespace(
+        put_keys=[], put_started=threading.Event(), release=threading.Event(), close_called=threading.Event()
+    )
+
+    def put(**kwargs):
+        conn.put_keys.append(kwargs["put_key"])
+        if len(conn.put_keys) == 1:
+            conn.put_started.set()
+            assert conn.release.wait(timeout=2)
+        return True, 1, None
+
+    conn.put, conn.close = put, conn.close_called.set
+    return conn
+
+
+def _start_save_thread(plane, connector):
+    plane.__dict__.update(
+        _omni_connector=connector,
+        _stage_id=0,
+        _next_stage_id=1,
+        _lock=threading.Lock(),
+        _pending_save_reqs={},
+        _pending_save_counts=defaultdict(int),
+        _deferred_send_cleanup=set(),
+        _request_ids_mapping={},
+        _cached_ic={},
+        _send_side_request_payload={},
+        _code_prompt_token_ids=defaultdict(list),
+        _work_available=threading.Event(),
+        _stop_event=threading.Event(),
+        _MAX_SEND_RETRIES=0,
+        _can_send=True,
+        _recv_thread=None,
+        _custom_process_batch_func=lambda **kwargs: kwargs["pooling_outputs"],
+        is_data_transfer_rank=lambda: True,
+        _connector_send_error_sink=plane._record_output_error,
+    )
+    plane._save_thread = threading.Thread(target=plane._save_loop, daemon=True)
+    plane._save_thread.start()
+
+
+def _stop_save_thread(plane):
+    plane._stop_event.set()
+    plane._work_available.set()
+    plane._save_thread.join(timeout=2)
+    assert not plane._save_thread.is_alive()
+
+
+class _Output:
+    copy_event = None
+
+    def __init__(self, rid="internal", gate=None, started=None, error=None, payload=None, token=1):
+        self.rid, self.gate, self.started = rid, gate, started
+        self.error, self.payload, self.token = error, payload, token
+
+    def get_output(self):
+        if self.started is not None:
+            self.started.set()
+        if self.gate is not None:
+            assert self.gate.wait(timeout=3)
+        if self.error is not None:
+            raise self.error
+        return SimpleNamespace(
+            req_ids=[self.rid],
+            inter_stage_outputs=[self.payload] if self.payload is not None else None,
+            sampled_token_ids=[[self.token]],
+        )
+
+
+def test_emit_cohort_one_batch_and_trims_non_resumable_history(plane):
+    plane.send_chunk = lambda *_a, **_kw: pytest.fail("must not fall back to per-request send_chunk")
+    plane.register_request(_new_request("r0", "ext-0"))
+    plane.register_request(_new_request("r1", "ext-1"))
+    emitted = plane.emit_chunks(
+        req_ids=["r0", "r1"],
+        inter_stage_outputs=[{"codes.audio": "c0"}, {"codes.audio": "c1"}],
+        sampled_token_ids=[[21]],
+        terminal_req_ids={"r1"},
+    )
+    assert emitted == 2 and len(plane.record.batches) == 1  # the whole step cohort goes out in one batch
+    (req0, _), (req1, _) = plane.record.batches[0]
+    assert (
+        req0.all_token_ids == [10, 11, 21]
+        and not req0.is_finished()
+        and req1.is_finished()
+        and plane.record.cleaned == ["r1"]
+    )
+    # A non-resumable request omits token history after its first chunk.
+    plane.register_request(_new_request("r2", "ext-2"))
+    state = plane._native_requests["r2"]
+    state.resumable, state.output_token_ids, plane._put_req_chunk["ext-2"] = False, list(range(1024)), 1
+    plane.emit_chunks(
+        req_ids=["r2"], inter_stage_outputs=[{"h": "d"}], sampled_token_ids=[[1024]], terminal_req_ids={"r2"}
+    )
+    request, _ = plane.record.batches[1][0]
+    assert request.output_token_count == 1025 and request.all_token_ids == []
+
+
+def test_terminal_reserved_abort_and_cleanup_lifecycle(plane):
+    batches, cleaned = plane.record.batches, plane.record.cleaned
+    plane.register_request(_new_request())
+    plane.reserve_outputs(["internal"])
+    plane.reserve_outputs(["internal"])
+    # The terminal defers while reservations are in flight, and is idempotent.
+    assert plane.request_terminal({"internal"}) == 0
+    assert plane.request_terminal({"internal"}) == 0 and batches == []
+    assert _complete(plane, [{"codes.audio": "c0"}]) == 1 and not batches[0][0][0].is_finished()
+    assert _complete(plane, [{"codes.audio": "c1"}], token=22) == 2  # data chunk, then the released terminal
+    chunk, terminal = batches[1][0], batches[2][0]
+    assert chunk[1] == {"codes": {"audio": "c1"}} and chunk[0].output_token_ids == [21, 22]
+    assert terminal[1] is None and terminal[0].is_finished() and cleaned == ["internal"]
+    # Stale terminal/complete after cleanup are no-ops; cleanup ran exactly once.
+    assert plane.request_terminal({"internal"}) == 0 and _complete(plane, [{"codes.audio": "stale"}]) == 0
+    # Abort emits exactly one terminal and silently drops the stale deferred output.
+    plane.register_request(_new_request("ra"))
+    plane.reserve_outputs(["ra"])
+    plane.request_terminal({"ra"})
+    assert plane.abort_requests({"ra"}) == 1
+    request, payload = batches[-1][0]
+    assert request.is_finished() and payload is None and cleaned == ["internal", "ra"]
+    assert (
+        plane.abort_requests({"ra"}) == 0 and _complete(plane, [{"s": 1}], req="ra") == 0
+    )  # duplicate abort, stale output
+    assert sum(b[0][0].request_id == "ra" for b in batches) == 1  # terminal exactly once
+
+
+def test_abort_cannot_overtake_committed_output(plane):
+    in_lock, release, sent = threading.Event(), threading.Event(), threading.Event()
+
+    def gated(entries, _orig=plane._send_chunk_entries, **kw):
+        terminal = entries[0][0].is_finished()
+        if not terminal:
+            in_lock.set()
+            assert release.wait(timeout=2)  # hold the committed data send inside _native_send_lock
+        result = _orig(entries, **kw)
+        if terminal:
+            sent.set()
+        return result
+
+    plane._send_chunk_entries = gated
+    plane.register_request(_new_request())
+    plane.reserve_outputs(["internal"])
+    threads = [
+        threading.Thread(target=lambda: _complete(plane, [{"c": 0}])),
+        threading.Thread(target=plane.abort_requests, args=({"internal"},)),
+    ]
+    threads[0].start()
+    try:
+        assert in_lock.wait(timeout=2)
+        threads[1].start()
+        overtook = sent.wait(timeout=0.1)  # the terminal must not pass the committed data send
+    finally:
+        release.set()
+        for t in threads:
+            if t.ident is not None:
+                t.join(timeout=2)
+                assert not t.is_alive()
+    assert not overtook
+
+
+def test_output_worker_fifo_and_scheduler_not_blocked(plane):
+    batches, send_started, release_send = [], threading.Event(), threading.Event()
+
+    def send_chunks(entries, **_kw):
+        if not entries[0][0].is_finished():
+            send_started.set()
+            assert release_send.wait(timeout=2)
+        batches.append(entries)
+        return len(entries)
+
+    plane.send_chunks = send_chunks
+    plane._start_output_worker(max_pending_batches=2)
+    plane.register_request(_new_request())
+    plane.reserve_outputs(["internal"])
+    plane.enqueue_outputs(req_ids=["internal"], inter_stage_outputs=[{"c": 0}], sampled_token_ids=[[21]])
+    assert send_started.wait(timeout=2)
+    # A stuck deferred send must not block the scheduler path.
+    assert plane.request_terminal({"internal"}) == 0 and batches == []
+    release_send.set()
+    plane.drain_outputs()
+    assert [e[0].is_finished() for batch in batches for e in batch] == [False, True]
+
+
+def test_worker_failure_surfaces_and_close_drains_in_order(plane):
+    events = []
+    plane.send_chunks = lambda _e, **_kw: (_ for _ in ()).throw(RuntimeError("enqueue failed"))
+    plane.shutdown_omni_connectors = lambda: events.append("shutdown")
+    plane._start_output_worker(max_pending_batches=2)
+    plane.register_request(_new_request())
+    plane.enqueue_outputs(req_ids=["internal"], inter_stage_outputs=[{"c": 0}], sampled_token_ids=[[21]])
+    with pytest.raises(RuntimeError, match="enqueue failed"):
+        plane.drain_outputs()
+    with pytest.raises(RuntimeError, match="enqueue failed"):  # an errored worker rejects new submissions
+        plane.enqueue_outputs(req_ids=["internal"], inter_stage_outputs=[{"c": 1}], sampled_token_ids=[[22]])
+    # close() drains pending batches before shutting connectors down.
+    plane._native_output_error = None
+
+    def send_chunks(entries, **_kw):
+        events.append("send")
+        return len(entries)
+
+    plane.send_chunks = send_chunks
+    plane.enqueue_outputs(req_ids=["internal"], inter_stage_outputs=[{"c": 1}], sampled_token_ids=[[22]])
+    plane.close()
+    assert events == ["send", "shutdown"]
+
+
+def test_permanent_put_failure_quarantines_and_holds_state(raw_plane):
+    puts = []
+
+    def put(**_kw):
+        puts.append(1)
+        return False, 0, None
+
+    connector = SimpleNamespace(put=put, close=lambda: None)
+    _start_save_thread(raw_plane, connector)
+    raw_plane.register_request(_new_request())
+    raw_plane.reserve_outputs(["internal"])
+    raw_plane.request_terminal({"internal"})
+    with pytest.raises(RuntimeError, match="connector send failed"):
+        raw_plane.complete_outputs(req_ids=["internal"], inter_stage_outputs=[{"c": 0}], sampled_token_ids=[[21]])
+    assert len(puts) == 1 and raw_plane._native_output_error is not None  # fail fast, no retry
+    assert raw_plane._native_outputs_in_flight["internal"] == 1 and "internal" in raw_plane._native_terminal_pending
+    with pytest.raises(RuntimeError, match="quarantined"):  # the quarantine rejects new tickets
+        raw_plane._delivery_manager.create_ticket(request_id="other", put_key="other_0_0")
+    _stop_save_thread(raw_plane)
+
+
+def test_delivery_timeout_drops_queued_ticket(raw_plane):
+    raw_plane._delivery_manager = OmniDeliveryManager(delivery_timeout_s=0.05, shutdown_timeout_s=0.5)
+    connector = _gated_connector()
+    _start_save_thread(raw_plane, connector)
+    req = SimpleNamespace(request_id="external")
+    _, first = raw_plane._enqueue_chunk_payload(req, {"c": 1}, wait_for_delivery=True)
+    _, second = raw_plane._enqueue_chunk_payload(req, {"c": 2}, wait_for_delivery=True)
+    assert connector.put_started.wait(timeout=1)
+    with pytest.raises(TimeoutError, match="delivery timed out"):
+        first.wait()
+    assert raw_plane._delivery_manager.is_quarantined
+    connector.release.set()
+    deadline = time.monotonic() + 1
+    while time.monotonic() < deadline and "external" in raw_plane._pending_save_counts:  # the save loop drops it
+        time.sleep(0.001)
+    # After quarantine the queued ticket is never put and is marked FAILED.
+    assert connector.put_keys == [first.put_key] and second.state is DeliveryState.FAILED
+    assert "external" not in raw_plane._pending_save_counts and "external" not in raw_plane._pending_save_reqs
+    _stop_save_thread(raw_plane)
+
+
+def test_close_cancels_waiting_delivery_once_and_shutdown_stays_bounded():
+    # A queued delivery waiter is cancelled exactly once on shutdown.
+    manager = OmniDeliveryManager(delivery_timeout_s=30.0, shutdown_timeout_s=0.5)
+    ticket = manager.create_ticket(request_id="r", put_key="k")
+    manager.shutdown(RuntimeError("close"))
+    manager.shutdown(RuntimeError("again"))  # idempotent: cancelled at most once
+    with pytest.raises(DeliveryCancelledError, match="cancelled"):
+        ticket.wait()
+    # A connector put() that cannot be cancelled must not hang shutdown.
+    plane = _bare_plane(delivery_timeout_s=0.05, shutdown_timeout_s=0.2)
+    connector = _gated_connector()
+    _start_save_thread(plane, connector)
+    assert plane._enqueue_chunk_payload(SimpleNamespace(request_id="external"), {"c": 0}, wait_for_delivery=True)[0]
+    assert connector.put_started.wait(timeout=1)
+    start = time.monotonic()
+    with pytest.raises(RuntimeError, match="shutdown exceeded"):
+        plane.shutdown_omni_connectors()
+    assert time.monotonic() - start < 1.5 and connector.close_called.is_set()
+    connector.release.set()
+    _stop_save_thread(plane)
+
+
+@pytest.mark.parametrize("abort", [False, True])
+def test_materialization_fences_terminal_and_abort(plane, abort):
+    batches, cleaned, gate, started = plane.record.batches, plane.record.cleaned, threading.Event(), threading.Event()
+    plane._start_output_worker(max_pending_batches=2)
+    materializer = NativeOutputWorker(2)
+    try:
+        plane.register_request(_new_request())
+        plane.reserve_outputs(["internal"])
+        output = materializer.submit(
+            _Output(gate=gate, started=started, payload={"codes.audio": "c0"}, token=21), plane
+        )
+        assert started.wait(timeout=3)
+        assert plane.request_terminal({"internal"}) == 0  # fenced behind the in-flight materialization
+        assert not abort or plane.abort_requests({"internal"}) == 1
+        gate.set()
+        output.get_output()
+        materializer.close()
+        plane.drain_outputs()
+        assert cleaned == ["internal"] and len(batches) == (1 if abort else 2)
+        assert abort or (batches[0][0][1] == {"codes": {"audio": "c0"}} and not batches[0][0][0].is_finished())
+        assert batches[-1][0][0].is_finished()
+    finally:
+        gate.set()
+        materializer.close()
+
+
+def test_native_worker_fifo_thread_affinity_and_signals():
+    owner, gate = threading.get_ident(), threading.Event()
+    calls: list[tuple[Any, ...]] = []
+
+    def get_signals():
+        calls.append(("signals", threading.get_ident()))
+        return "ready"
+
+    plane = SimpleNamespace(
+        enqueue_outputs=lambda **kw: calls.append(("publish", kw["req_ids"], threading.get_ident())),
+        get_omni_connector_output=get_signals,
+    )
+    worker = NativeOutputWorker(2)
+    try:
+        a = worker.submit(_Output("a", gate=gate, payload={"v": 1}), plane)
+        b = worker.submit(_Output("b", payload={"v": 1}), plane)
+        gate.set()
+        b.get_output()  # publication of b implies a already published (FIFO, single worker thread)
+        publishes = [c for c in calls if c[0] == "publish"]
+        assert [c[1] for c in publishes] == [["a"], ["b"]]
+        assert all(c[2] != owner for c in publishes)  # published on the worker thread
+        result = a.get_output()
+        assert a.get_output() is result and result.inter_stage_outputs is None  # resolution is cached
+        assert result.omni_connector_output == "ready"
+        assert calls[-2:] == [("signals", owner)] * 2  # signals resolve lazily on the caller thread
+    finally:
+        gate.set()
+        worker.close()
+
+
+def test_native_worker_capacity_error_cache_and_tp_gates():
+    from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
+
+    plane = SimpleNamespace(enqueue_outputs=lambda **kw: None, get_omni_connector_output=lambda: None)
+    worker, gate = NativeOutputWorker(1), threading.Event()
+    try:
+        first = worker.submit(_Output("a", gate=gate), plane)
+        with ThreadPoolExecutor(max_workers=1) as pool:
+            future = pool.submit(worker.submit, _Output("b"), plane)
+            time.sleep(0.05)
+            assert not future.done()  # capacity 1: pending materialization is bounded
+            gate.set()
+            second = future.result(3)
+        assert first.get_output().req_ids == ["a"] and second.get_output().req_ids == ["b"]
+        bad = worker.submit(_Output("x", error=ValueError("crash")), plane)
+    finally:
+        gate.set()
+        worker.close()
+    # A failed materialization is cached: repeated reads raise the same error.
+    for _ in range(2):
+        with pytest.raises(ValueError, match="crash"):
+            bad.get_output()
+    runner = SimpleNamespace(
+        _omni_data_plane=object(), vllm_config=SimpleNamespace(parallel_config=SimpleNamespace(tensor_parallel_size=2))
+    )
+    assert OmniGPUModelRunner._uses_native_output_materializer(runner) is False  # TP>1 keeps TP consumers
+    runner.vllm_config.parallel_config.tensor_parallel_size = 1
+    assert OmniGPUModelRunner._uses_native_output_materializer(runner) is True
+
+
+def test_abort_before_first_chunk_cleans_receiver_state():
+    plane = _bare_plane()
+    plane._stage_id, plane._lock, plane._work_available = 1, threading.Lock(), threading.Event()
+    for name in (
+        "_pending_full_payload_send",
+        "_request_ids_mapping",
+        "_pending_save_counts",
+        "_send_side_request_payload",
+        "_code_prompt_token_ids",
+        "_cached_ic",
+        "_adaptive_states",
+        "_kv_pending_transfers",
+        "_get_req_chunk",
+        "_pending_load_reqs",
+        "_local_stage_payload_cache",
+        "_local_request_metadata",
+    ):
+        setattr(plane, name, {})
+    for name in (
+        "_deferred_send_cleanup",
+        "_kv_active_transfers",
+        "_kv_completed_transfers",
+        "_kv_triggered_requests",
+        "_finished_load_reqs",
+        "_chunk_ready_req_ids",
+        "_chunk_finished_req_ids",
+        "_chunk_stream_completed",
+        "_stage_recv_req_ids",
+        "_full_payload_pending_broadcast_req_ids",
+        "_async_chunk_updated_req_ids",
+    ):
+        setattr(plane, name, set())
+    plane.register_receivers([SimpleNamespace(request_id="r", external_req_id="external")])
+    plane._local_stage_payload_cache["r"] = {"codes": torch.ones(1)}
+    assert "r" in plane._pending_load_reqs and plane._request_ids_mapping["r"] == "external"
+    assert plane.abort_requests({"r"}) == 0
+    assert plane.abort_requests({"r"}) == 0
+    assert not plane._pending_load_reqs and not plane._request_ids_mapping and not plane._local_stage_payload_cache

tests/worker_v2/test_omni_generation_model_runner.py

@@ -0,0 +1,184 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""OmniGenerationModelRunner contracts: empty-step lifecycle, output partition,
+CPU-sync vs CUDA-async dispatch, and async-chunk slot recycling."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import numpy as np
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _FakeStagedField:
+    def __init__(self, data):
+        self.np = data
+        self._staged = []
+
+    def stage_write_elem(self, idx, value):
+        self._staged.append((idx, value))
+
+    def apply_write(self):
+        for idx, value in self._staged:
+            self.np[idx] = value
+        self._staged.clear()
+
+
+def _make_runner(model_output, num_reqs=1, prompt_len=10):
+    runner = object.__new__(OmniGenerationModelRunner)
+    runner.device = torch.device("cpu")
+    runner.model_config = MagicMock(hf_text_config=None)
+    del runner.model_config.eos_token_id
+    runner.postprocess = lambda *a, **kw: None
+    runner._gen_model_output = model_output
+    runner._gen_input_batch = SimpleNamespace(
+        num_reqs=num_reqs, req_ids=[f"req-{i}" for i in range(num_reqs)], idx_mapping_np=np.arange(num_reqs)
+    )
+    runner.execute_model_state = SimpleNamespace(finished_req_ids={"finished"}, ec_connector_output=None)
+    runner.kv_connector = SimpleNamespace(post_forward=MagicMock(return_value=None))
+    runner.check_ep_fault = False
+    runner.req_states = MagicMock(
+        prompt_len=SimpleNamespace(np=np.full(num_reqs, prompt_len, dtype=np.int32)),
+        num_computed_tokens=_FakeStagedField(np.zeros(num_reqs, dtype=np.int32)),
+    )
+    return runner
+
+
+def test_control_only_step_keeps_lifecycle_and_skips_input_construction():
+    runner = object.__new__(OmniGenerationModelRunner)
+    order = []
+    for name in (
+        "_prepare_native_data_plane",
+        "finish_requests",
+        "free_states",
+        "_apply_block_table_staged_writes_if_available",
+    ):
+        setattr(runner, name, lambda *args, name=name: order.append(name))
+    for name in ("_handle_async_chunk_updates", "add_requests", "update_requests", "_sync_native_data_plane_payloads"):
+        setattr(runner, name, MagicMock(side_effect=AssertionError("control step built model inputs")))
+    output = object()
+    runner.kv_connector = SimpleNamespace(no_forward=lambda _s: output)
+    runner._merge_ec_connector_no_forward = lambda _s, value: value
+    runner._attach_native_data_plane_signals = lambda value: value
+    scheduler_output = SimpleNamespace(
+        total_num_scheduled_tokens=0, scheduled_new_reqs=[], scheduled_cached_reqs=SimpleNamespace(req_ids=[])
+    )
+    assert runner.execute_model(scheduler_output) is output
+    assert order == [
+        "_prepare_native_data_plane",
+        "finish_requests",
+        "free_states",
+        "_apply_block_table_staged_writes_if_available",
+    ]
+
+
+def test_released_chunk_reuses_scheduler_output_and_slot_recycle_clears_state():
+    from vllm.v1.core.sched.output import SchedulerOutput
+
+    from vllm_omni.core.sched.output import OmniCachedRequestData
+
+    runner = object.__new__(OmniGenerationModelRunner)
+    runner.model = SimpleNamespace(logits_processor=None)
+    runner._remove_request = MagicMock()
+    req_id_to_index = MagicMock()
+    req_id_to_index.get.return_value = None
+    runner.req_states = SimpleNamespace(
+        add_request=MagicMock(), req_id_to_index=req_id_to_index, apply_staged_writes=MagicMock()
+    )
+    runner.pooling_runner = None
+    runner.encoder_cache = None
+    runner.model_state = SimpleNamespace(
+        add_request=MagicMock(),
+        apply_staged_writes=MagicMock(),
+        remove_request=MagicMock(),
+        intermediate_buffer=SimpleNamespace(remove_request=MagicMock()),
+    )
+    runner.block_tables = SimpleNamespace(append_block_ids=MagicMock())
+    runner.lora_state = SimpleNamespace(add_request=MagicMock())
+    runner.is_last_pp_rank = False
+    runner.adaptive_verification = None
+    runner.sampler = None
+    cached = OmniCachedRequestData(
+        req_ids=["req"],
+        resumed_req_ids=set(),
+        new_token_ids=[[]],
+        all_token_ids={"req": [1]},
+        new_block_ids=[()],
+        num_computed_tokens=[0],
+        num_output_tokens=[0],
+        prompt_token_ids={"req": [1]},
+        additional_information={"req": None},
+    )
+    scheduler_output = SchedulerOutput(
+        scheduled_new_reqs=[],
+        scheduled_cached_reqs=cached,
+        num_scheduled_tokens={"req": 1},
+        total_num_scheduled_tokens=1,
+        scheduled_spec_decode_tokens={},
+        scheduled_encoder_inputs={},
+        num_common_prefix_blocks=[0],
+        finished_req_ids=set(),
+        free_encoder_mm_hashes=[],
+    )
+    runner._handle_async_chunk_updates(scheduler_output)
+    added = runner.model_state.add_request.call_args.args[1]
+    assert (added.req_id, added.prompt_token_ids) == ("req", [1])
+
+
+def test_sample_tokens_cpu_sync_owns_outputs_and_runs_connector_last():
+    waveform = torch.arange(4, dtype=torch.float32)
+    output = OmniOutput(text_hidden_states=torch.empty(0), multimodal_outputs={"codes": {"audio": [waveform]}})
+    runner = _make_runner(output, num_reqs=1)
+    runner._reserve_native_data_plane_outputs = MagicMock()
+    runner._finalize_native_data_plane_output = MagicMock(side_effect=lambda value: value)
+
+    def post_forward(_finished):
+        assert runner.req_states.num_computed_tokens.np.tolist() == [10]
+
+    runner.kv_connector.post_forward.side_effect = post_forward
+    result = OmniGenerationModelRunner.sample_tokens(runner)
+    assert isinstance(result, OmniModelRunnerOutput)
+    # Reserve precedes finalize; connector runs after token state is applied.
+    runner._reserve_native_data_plane_outputs.assert_called_once_with(["req-0"])
+    runner.kv_connector.post_forward.assert_called_once_with({"finished"})
+    waveform.fill_(-1)
+    assert torch.equal(result.multimodal_outputs[0]["codes.audio"], torch.arange(4, dtype=torch.float32))
+
+
+def test_sample_tokens_uses_async_output_for_cuda_and_snapshots_req_ids(monkeypatch):
+    from vllm_omni.worker_v2 import omni_generation_model_runner as generation_runner
+
+    output = OmniOutput(text_hidden_states=torch.zeros(1), multimodal_outputs={"model_outputs": [torch.randn(4)]})
+    runner = _make_runner(output, num_reqs=1)
+    runner.device = SimpleNamespace(type="cuda")
+    runner.main_stream = object()
+    runner.output_copy_stream = object()
+    runner.model_config.async_chunk = True
+    runner._release_generation_slots = MagicMock()
+    runner._finalize_native_data_plane_output = MagicMock()
+    monkeypatch.setattr(generation_runner, "_contains_cuda_tensor", lambda _: True)
+    captured = {}
+
+    class _FakeAsyncOutput:
+        def __init__(self, **kwargs):
+            captured.update(kwargs)
+
+    monkeypatch.setattr(generation_runner, "OmniGenerationAsyncOutput", _FakeAsyncOutput)
+    input_batch = runner._gen_input_batch
+    result = generation_runner.OmniGenerationModelRunner.sample_tokens(runner)
+
+    assert isinstance(result, _FakeAsyncOutput)
+    assert captured["finalize_output"] is runner._finalize_native_data_plane_output
+    assert captured["model_runner_output"].sampled_token_ids == [[]]
+    runner._release_generation_slots.assert_called_once()
+    output_req_ids = captured["model_runner_output"].req_ids
+    input_batch.req_ids[0] = "reused"  # snapshot must insulate the published output
+    assert output_req_ids == ["req-0"]

tests/worker_v2/test_omni_gpu_model_runner.py

@@ -0,0 +1,166 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""MRV2 admission, capture, dispatch and request lifecycle contracts."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+from vllm import SamplingParams
+from vllm.config.compilation import CUDAGraphMode
+from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+from vllm.v1.worker.gpu.sample.sampler import Sampler
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.worker_v2.model_states import init_omni_model_state
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _make_runner():
+    """Create an OmniGPUModelRunner without calling __init__."""
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.model = MagicMock()
+    runner.req_states = SimpleNamespace(req_id_to_index={"r1": 0, "r2": 1})
+    runner.execute_model_state = None
+    return runner
+
+
+def test_add_requests_empty_admission_and_stop_id_sanitization():
+    runner = _make_runner()
+    runner.sampler = Sampler.__new__(Sampler)
+    with patch.object(GPUModelRunner, "add_requests", return_value=None) as parent:
+        runner.add_requests(SchedulerOutput.make_empty())
+        parent.assert_not_called()
+
+        # New requests always go upstream; narrow logits heads sanitize stop ids.
+        sampling_params = SamplingParams(min_tokens=2, stop_token_ids=[2150])
+        sampling_params.update_from_generation_config({}, 151645)
+        runner.model = SimpleNamespace(logits_processor=SimpleNamespace(vocab_size=3072))
+        output = SchedulerOutput.make_empty()
+        output.scheduled_new_reqs = [SimpleNamespace(sampling_params=sampling_params)]
+        runner.add_requests(output)
+        parent.assert_called_once_with(output)
+    assert sampling_params.all_stop_token_ids == {2150}
+    assert sampling_params.eos_token_id == 151645
+
+
+def test_prepare_native_data_plane_terminal_abort_split_and_warmup_skip():
+    runner = _make_runner()
+    plane = SimpleNamespace(
+        register_request=MagicMock(),
+        register_receivers=MagicMock(),
+        request_terminal=MagicMock(),
+        abort_requests=MagicMock(),
+    )
+    runner._omni_data_plane = plane
+    new_req = SimpleNamespace(req_id="r1")
+    warmup_req = SimpleNamespace(req_id="_warmup_0_")
+    handle = SimpleNamespace(request_id="r2", external_req_id="ext-r2")
+    scheduler_output = SimpleNamespace(
+        scheduled_new_reqs=[new_req, warmup_req],
+        pending_input_registrations=[handle],
+        data_plane_terminal_req_ids={"r0"},
+        finished_req_ids={"r0", "aborted"},
+    )
+
+    runner._prepare_native_data_plane(scheduler_output)
+
+    plane.register_request.assert_called_once_with(new_req)  # warmup skipped
+    plane.register_receivers.assert_called_once_with([handle])
+    plane.request_terminal.assert_called_once_with({"r0"})
+    plane.abort_requests.assert_called_once_with({"aborted"})
+
+
+@pytest.mark.parametrize("output_form", ["tuple", "omni"])
+def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form):
+    runner = object.__new__(OmniGPUModelRunner)
+    hidden = torch.ones(1, 2)
+
+    def original_forward():
+        if output_form == "tuple":
+            return hidden, {"layers": {}}
+        return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
+
+    runner.model = SimpleNamespace(forward=original_forward)
+    runner._model_returns_tuple = True
+    runner._exclude_full_graph = True
+    runner.use_aux_hidden_state_outputs = False
+    piecewise = SimpleNamespace(cg_mode=CUDAGraphMode.PIECEWISE)
+    full = SimpleNamespace(cg_mode=CUDAGraphMode.FULL)
+    runner.cudagraph_manager = SimpleNamespace(
+        _capture_descs={CUDAGraphMode.PIECEWISE: [piecewise], CUDAGraphMode.FULL: [full]},
+        _candidates={(1, 0): [piecewise, full]},
+    )
+    runner.model_state = SimpleNamespace(capture_mtp_graphs=MagicMock())
+    runner._dispatch_mtp_batch_descriptor = MagicMock(return_value="desc")
+
+    def assert_unwrapped(_self):
+        assert torch.equal(runner.model.forward(), hidden)  # forward unwrapped during capture
+        return 3
+
+    with patch.object(GPUModelRunner, "capture_model", assert_unwrapped):
+        assert runner.capture_model() == 3
+
+    assert runner.model.forward is original_forward  # restored after capture
+    assert runner.cudagraph_manager._capture_descs == {CUDAGraphMode.PIECEWISE: [piecewise]}
+    runner.model_state.capture_mtp_graphs.assert_called_once_with(runner._dispatch_mtp_batch_descriptor)
+
+
+@pytest.mark.parametrize(
+    "parallel_config,match",
+    [
+        (dict(pipeline_parallel_size=2, prefill_context_parallel_size=1), "pipeline parallel"),
+        (dict(pipeline_parallel_size=1, prefill_context_parallel_size=2), "prefill context parallelism"),
+    ],
+)
+def test_mrv2_rejects_parallel_modes_at_startup(parallel_config, match):
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.vllm_config = SimpleNamespace(parallel_config=SimpleNamespace(**parallel_config))
+    with pytest.raises(NotImplementedError, match=match):
+        runner._validate_parallel_support()
+
+
+@pytest.mark.parametrize(
+    "flag,expect_omni",
+    [(None, False), ("has_preprocess", True), ("has_postprocess", True), ("have_multimodal_outputs", True)],
+)
+def test_init_model_state_factory_dispatches_omni_only(monkeypatch, flag, expect_omni):
+    upstream = MagicMock(return_value=object())
+    monkeypatch.setattr("vllm_omni.worker_v2.model_states._upstream_init_model_state", upstream)
+    monkeypatch.setattr(OmniModelState, "__init__", lambda *args: None)
+    cfg = SimpleNamespace(model_config=SimpleNamespace(architectures=["CustomModel"]))
+    model = SimpleNamespace(**({flag: True} if flag else {}))
+    state = init_omni_model_state(cfg, model, None, torch.device("cpu"))
+
+    if expect_omni:
+        assert isinstance(state, OmniModelState)
+        upstream.assert_not_called()
+    else:
+        upstream.assert_called_once()
+        assert state is upstream.return_value
+
+
+def test_finish_requests_notifies_model_and_cleans_only_known_slots(monkeypatch):
+    runner = _make_runner()
+    calls = []
+    runner.model = SimpleNamespace(on_requests_finished=lambda ids: calls.append(set(ids)))
+    runner.model_state = MagicMock()
+    monkeypatch.setattr(GPUModelRunner, "finish_requests", lambda *args: None)
+    runner.finish_requests(SimpleNamespace(finished_req_ids={"released"}, preempted_req_ids={"r1"}))
+    assert calls == [{"released"}]
+    assert sorted(c.args[0] for c in runner.model_state.remove_request.call_args_list) == [0]
+
+
+@pytest.mark.parametrize("stage,declared", [("thinker", False), ("custom_ar", True)])
+def test_capture_contract_uses_model_declaration(stage, declared):
+    runner = _make_runner()
+    runner.model = SimpleNamespace(model_stage=stage, _returns_tuple=declared)
+    runner._configure_cudagraph_output_contract()
+    assert runner._model_returns_tuple is declared
+    assert runner._exclude_full_graph is declared

tests/worker_v2/test_omni_model_state.py

@@ -0,0 +1,361 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""OmniModelState: mixed-batch reorder, state isolation/slot reuse, MTP
+graph/eager paths with per-request seed independence, async snapshot ownership."""
+
+from contextlib import nullcontext
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import numpy as np
+import pytest
+import torch
+from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState, _make_safe_get_rope
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyInputBatch:
+    input_ids: SimpleNamespace
+    query_start_loc: torch.Tensor
+
+    def __init__(self, indices, *, num_computed_tokens_cpu=None):
+        self.idx_mapping_np = indices
+        self.num_reqs = len(indices)
+        self.num_scheduled_tokens = [1] * len(indices)
+        self.query_start_loc_np = list(range(len(indices)))
+        if num_computed_tokens_cpu is not None:
+            self.num_computed_tokens_cpu = np.array(num_computed_tokens_cpu, dtype=np.int32)
+
+
+def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, have_multimodal_outputs=False):
+    state = object.__new__(OmniModelState)
+    model = MagicMock()
+    model.has_preprocess = has_preprocess
+    model.has_postprocess = has_postprocess
+    model.have_multimodal_outputs = have_multimodal_outputs
+    model.gpu_resident_buffer_keys = set()
+    model.batched_gpu_staging_keys = set()
+    model.mtp_accepts_per_row_generators = False
+    model.mtp_accepts_req_infos = False
+    model.mtp_output_key = ("codes", "audio")
+    model.mtp_validity_key = None
+    model.mtp_sampling_params = {}
+    model.get_mtp_seed = lambda params: (getattr(params, "extra_args", None) or {}).get("test_seed")
+    model.preprocess_batch_mrv2 = None
+    model.preprocess_decode_batch_mrv2 = None
+    model.preprocess_decode_batch = None
+    model.postprocess_batch_mrv2 = None
+    state.model = model
+    state.scheduler_config = SimpleNamespace(max_num_seqs=max_num_reqs)
+    state.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
+    state.has_preprocess = has_preprocess
+    state.has_postprocess = has_postprocess
+    state.have_multimodal_outputs = have_multimodal_outputs
+
+    from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer
+
+    state.intermediate_buffer = OmniIntermediateBuffer(max_num_reqs)
+    state._static_inputs_embeds = None
+    state._mtp_generators = {}
+    state._mtp_runner = None
+    for name in ("_mtp_input_ids", "_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step", "_mtp_offsets"):
+        setattr(state, name, None)
+    return state
+
+
+def _add(state, req_id="r1", idx=0, **req_kwargs):
+    with patch.object(DefaultModelState, "add_request", return_value=None):
+        state.add_request(idx, SimpleNamespace(req_id=req_id, mm_features=[], **req_kwargs))
+
+
+def _fill_buffers(state, *req_ids):
+    for idx, req_id in enumerate(req_ids):
+        state.intermediate_buffer.buffers[idx] = {"req_id": req_id}
+
+
+def _seeded(seed):
+    return SimpleNamespace(extra_args={"test_seed": seed}, seed=None)
+
+
+def _init_static(state, bsz, dim=3):
+    state._mtp_input_ids = torch.zeros(bsz, dtype=torch.long)
+    state._mtp_input_embeds = torch.zeros((bsz, dim))
+    state._mtp_hidden = torch.zeros((bsz, dim))
+    state._mtp_text_step = torch.zeros((bsz, dim))
+    state._mtp_offsets = torch.zeros(bsz, dtype=torch.long)
+
+
+def _mtp_batches(dim=3):
+    return [(0, 0, (torch.ones(dim), torch.ones(dim) * 2)), (1, 1, (torch.ones(dim) * 3, torch.ones(dim) * 4))]
+
+
+@pytest.fixture(autouse=True)
+def _fwd_ctx():
+    with patch("vllm.forward_context.set_forward_context", return_value=nullcontext()):
+        yield
+
+
+@pytest.mark.parametrize("req_id,expect_validity", [("_warmup_0_", True), ("r1", False)])
+def test_add_request_declared_validity_only_for_warmup(req_id, expect_validity):
+    # Real requests must not be masked with a fabricated validity key.
+    state = _make_state()
+    state.model.mtp_validity_key = ("meta", "codec_frame_valid")
+    _add(state, req_id, 0)
+    buf = state.intermediate_buffer.buffers[0]
+    assert ("meta" in buf) is expect_validity
+    if expect_validity:
+        assert buf["meta"]["codec_frame_valid"].item() is False
+
+
+@pytest.mark.parametrize("mode", ["index", "req_id", "unknown"])
+def test_remove_request_state_isolation(mode):
+    state = _make_state()
+    _add(state, "r1", 0)
+    state._mtp_generators["r1"] = torch.Generator(device="cpu")
+    if mode == "unknown":
+        state.remove_request("missing")
+        assert state.intermediate_buffer.buffers[0]["req_id"] == "r1"
+        assert "r1" in state._mtp_generators
+        return
+    state.remove_request(0 if mode == "index" else "r1")
+    assert state.intermediate_buffer.buffers[0] == {}
+    assert "r1" not in state._mtp_generators  # recycled slot inherits no seed stream
+
+
+def test_output_spans_follow_reordered_mixed_batch():
+    # batch=[2, 0]: prefill (3 tokens) reordered ahead of decode (1 token).
+    state = _make_state(have_multimodal_outputs=True)
+    batch = _DummyInputBatch([2, 0])
+    batch.num_scheduled_tokens = [3, 1]
+    batch.query_start_loc_np = [0, 3]
+    state.intermediate_buffer.buffers[2] = {"req_id": "prefill"}
+    state.intermediate_buffer.buffers[0] = {"req_id": "decode"}
+    seen = {}
+    state.model.make_omni_output = lambda hidden, **kwargs: (
+        seen.update(kwargs) or OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
+    )
+    state.postprocess_model_output(torch.zeros(4, 2), batch, SimpleNamespace())
+    assert seen["request_token_spans"] == [(0, 3), (3, 4)]
+    assert [info["req_id"] for info in seen["model_intermediate_buffer"]] == ["prefill", "decode"]
+
+
+@pytest.mark.parametrize("hook_name", ["preprocess_decode_batch_mrv2", "preprocess_decode_batch", "legacy"])
+def test_decode_cohort_uses_batch_hook(hook_name):
+    state = _make_state(max_num_reqs=2, has_preprocess=True)
+    state.intermediate_buffer.buffers[0] = {"req_id": "r1", "meta": {"step": 1}}
+    state.intermediate_buffer.buffers[1] = {"req_id": "r2", "meta": {"step": 2}}
+    batch_calls = []
+
+    def preprocess_decode_batch(*, input_ids, input_embeds, req_infos):
+        batch_calls.append(list(req_infos))
+        return (
+            input_ids + 10,
+            input_embeds + 20,
+            torch.tensor([[31.0, 32.0], [41.0, 42.0]]),
+            torch.tensor([[51.0, 52.0], [61.0, 62.0]]),
+            [{"meta": {"step": 11}}, {"meta": {"step": 12}}],
+        )
+
+    if hook_name == "legacy":
+        # Legacy V1 hook without input_embeds is wrapped by the adapter.
+        state.model.preprocess_decode_batch = lambda *, input_ids, req_infos: preprocess_decode_batch(
+            input_ids=input_ids, input_embeds=model_inputs["inputs_embeds"], req_infos=req_infos
+        )
+    else:
+        setattr(state.model, hook_name, preprocess_decode_batch)
+        if hook_name == "preprocess_decode_batch_mrv2":
+            state.model.preprocess_decode_batch = MagicMock(
+                side_effect=AssertionError("MRV2 hook must take precedence")
+            )
+    state.model.preprocess = MagicMock(side_effect=AssertionError("decode requests must use the batch hook"))
+    seen_prepacked = []
+
+    def capture(batches, *_a, prepacked_mtp_inputs=None):
+        seen_prepacked.append(prepacked_mtp_inputs)
+
+    state._run_batched_mtp = capture
+    model_inputs = {
+        "input_ids": torch.tensor([101, 202], dtype=torch.long),
+        "inputs_embeds": torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
+    }
+    state.run_preprocess(
+        _DummyInputBatch([0, 1], num_computed_tokens_cpu=[1, 1]),
+        model_inputs,
+        SimpleNamespace(prompt_len=np.array([1, 1], dtype=np.int32)),
+    )
+    state.model.preprocess.assert_not_called()
+    assert [info["req_id"] for info in batch_calls[0]] == ["r1", "r2"]
+    assert torch.equal(model_inputs["input_ids"], torch.tensor([111, 212]))
+    assert state.intermediate_buffer.buffers[0]["meta"]["step"] == 11
+    assert torch.equal(seen_prepacked[0][0], torch.tensor([[31.0, 32.0], [41.0, 42.0]]))
+
+
+def test_static_decode_embeddings_refresh_from_input_ids():
+    # FULL-graph replay reads the static buffer: it must hold fresh embeddings.
+    state = _make_state(has_preprocess=True)
+    static_embeds = torch.zeros(1, 4)
+    state._static_inputs_embeds = static_embeds
+    expected = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
+    state.model.embed_input_ids = lambda input_ids: expected.to(input_ids.device)
+    state.model.preprocess = lambda input_ids, input_embeds, **_info: (input_ids, input_embeds, {})
+    state.intermediate_buffer.buffers[0] = {"req_id": "r1"}
+    state.run_preprocess(
+        _DummyInputBatch([0]),
+        {"input_ids": torch.tensor([1049], dtype=torch.long), "inputs_embeds": static_embeds[:1]},
+    )
+    assert torch.equal(static_embeds, expected)
+    original = torch.tensor([[1.0, 2.0]])
+    assert OmniModelState._preprocess_result_needs_writeback(original, original) is False
+    assert OmniModelState._preprocess_result_needs_writeback(original, original.view_as(original)) is True
+
+
+@pytest.mark.parametrize("owned", [False, True])
+def test_batched_postprocess_gpu_snapshot_writeback(owned):
+    # batch=[1, 0]: last-token indices follow the reordered batch, the scalar
+    # hook is not called, and both rows share one snapshot storage so a later
+    # graph replay cannot overwrite published rows. Owned producer batches
+    # skip the extra snapshot while borrowed graph output still gets one.
+    state = _make_state(max_num_reqs=2, has_postprocess=True)
+    state.model.gpu_resident_buffer_keys = {("hidden_states", "last")}
+    _fill_buffers(state, "r0", "r1")
+    state.model.postprocess = MagicMock(side_effect=AssertionError("batch hook must replace scalar postprocess"))
+    calls = []
+    produced = []
+
+    def postprocess_batch(*, hidden_states, last_token_indices):
+        calls.append(last_token_indices.clone())
+        gathered = hidden_states.index_select(0, last_token_indices)
+        produced.append(gathered)
+        if owned:
+            return (("hidden_states", "last"), OwnedBatchTensor(gathered))
+        return (("hidden_states", "last"), gathered)
+
+    state.model.postprocess_batch_mrv2 = postprocess_batch
+    batch = _DummyInputBatch([1, 0])
+    batch.num_scheduled_tokens = [2, 2]
+    batch.query_start_loc_np = [0, 2]
+    batch.query_start_loc = torch.tensor([0, 2, 4], dtype=torch.int32)
+    hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
+    state.run_postprocess(hidden, batch)
+
+    state.model.postprocess.assert_not_called()
+    assert torch.equal(calls[0], torch.tensor([1, 3], dtype=torch.int32))
+    first = state.intermediate_buffer.buffers[0]["hidden_states"]["last"]
+    second = state.intermediate_buffer.buffers[1]["hidden_states"]["last"]
+    assert torch.equal(first, hidden[3])
+    assert torch.equal(second, hidden[1])
+    assert first.untyped_storage().data_ptr() == second.untyped_storage().data_ptr()
+    if owned:
+        # Owned output: rows view into the producer batch, no extra snapshot.
+        assert first.untyped_storage().data_ptr() == produced[0].untyped_storage().data_ptr()
+    else:
+        # Borrowed/graph output: the snapshot owns its storage, not the producer.
+        assert first.untyped_storage().data_ptr() != produced[0].untyped_storage().data_ptr()
+
+
+def test_seed_independence_resolve_once_and_sampling_kwargs():
+    state = _make_state(max_num_reqs=2)
+    # 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
+    # Same model-local seed reproduces identical uniforms regardless of batch makeup.
+    state._mtp_sample_uniforms = torch.empty((2, 2, 4))
+    assert torch.equal(
+        state._prepare_mtp_sample_uniforms([torch.Generator().manual_seed(11)], 1).clone(),
+        state._prepare_mtp_sample_uniforms([torch.Generator().manual_seed(11)], 1),
+    )
+    # Per-request generators are resolved once per step and ride sampling kwargs.
+    state.model.mtp_sampling_params = {"do_sample": True, "temperature": 0.7}
+    state.model.mtp_accepts_per_row_generators = True
+    _init_static(state, 2)
+    state.intermediate_buffer.buffers[0] = {"req_id": "r0", "sampling_params": _seeded(11)}
+    state.intermediate_buffer.buffers[1] = {"req_id": "r1", "sampling_params": _seeded(22)}
+    resolver = MagicMock(wraps=state._get_mtp_generator)
+    state._get_mtp_generator = resolver
+
+    def mtp(input_ids, input_embeds, last_hidden, text_step, **kwargs):
+        assert (kwargs["do_sample"], kwargs["temperature"]) == (True, 0.7)
+        assert all(isinstance(g, torch.Generator) for g in kwargs["generators"])
+        return input_embeds + 10, torch.tensor([[1, 2, 3], [4, 5, 6]])
+
+    state.model.mtp = mtp
+    state._run_batched_mtp(
+        _mtp_batches(), torch.tensor([101, 202]), torch.zeros(2, 3), _DummyInputBatch([0, 1]), {("codes", "audio")}
+    )
+    assert resolver.call_count == 2
+    assert torch.equal(state.intermediate_buffer.buffers[0]["codes"]["audio"], torch.tensor([[1, 2, 3]]))
+    assert torch.equal(state.intermediate_buffer.buffers[1]["codes"]["audio"], torch.tensor([[4, 5, 6]]))
+
+
+def test_seeded_mtp_bypasses_outer_graph_runner():
+    state = _make_state(max_num_reqs=2)
+    raw_calls = []
+
+    def mtp(*args, **kwargs):
+        raw_calls.append(kwargs)
+        return args[1], args[0].reshape(-1, 1)
+
+    state.model.mtp = mtp
+    state._mtp_runner = MagicMock(side_effect=AssertionError("seeded sampling must not replay the graph"))
+    generators = [torch.Generator().manual_seed(11), torch.Generator().manual_seed(22)]
+    state._call_mtp_runner(
+        torch.tensor([101, 202]), torch.zeros(2, 3), torch.zeros(2, 3), torch.zeros(2, 3), generators=generators
+    )
+    assert raw_calls == [{"generators": generators}]
+    state._mtp_runner.assert_not_called()
+
+
+def test_run_batched_mtp_uses_dispatched_graph_descriptor():
+    state = _make_state(max_num_reqs=4)
+
+    class _FakeGraphRunner:
+        def __call__(self, input_ids, input_embeds, last_hidden, text_step, **kwargs):
+            assert input_ids.shape[0] == 4  # graph runs the padded descriptor size
+            return input_embeds + 2, torch.arange(12, dtype=torch.long).reshape(4, 3)
+
+    state._mtp_runner = _FakeGraphRunner()
+    _init_static(state, 4)
+    _fill_buffers(state, "r0", "r1")
+    graph_desc = SimpleNamespace(cg_mode="FULL", num_tokens=4)
+    set_ctx = MagicMock(return_value=nullcontext())
+    with (
+        patch(
+            "vllm_omni.worker_v2.model_states.omni_model_state.current_omni_platform.get_graph_wrapper_cls",
+            return_value=_FakeGraphRunner,
+        ),
+        patch("vllm.forward_context.set_forward_context", set_ctx),
+    ):
+        state._run_batched_mtp(
+            _mtp_batches(),
+            torch.tensor([101, 202]),
+            torch.zeros(2, 3),
+            _DummyInputBatch([0, 1]),
+            {("codes", "audio")},
+            lambda bsz: graph_desc,
+        )
+    _, kwargs = set_ctx.call_args
+    assert kwargs["batch_descriptor"] is graph_desc
+    assert kwargs["num_tokens"] == 4
+    assert torch.equal(state.intermediate_buffer.buffers[0]["codes"]["audio"], torch.tensor([[0, 1, 2]]))
+    assert torch.equal(state.intermediate_buffer.buffers[1]["codes"]["audio"], torch.tensor([[3, 4, 5]]))
+
+
+def test_rope_shim_propagates_type_error():
+    # vLLM rope API drift (TypeError) must not be swallowed by the shim.
+    def broken_get_rope(*_args, **_kwargs):
+        raise TypeError("vLLM rope API drift")
+
+    with pytest.raises(TypeError, match="rope API drift"):
+        _make_safe_get_rope(broken_get_rope)(SimpleNamespace(uses_mrope=False), object())
+
+
+@pytest.mark.parametrize("output_key", [None, ("only_one",)])
+def test_mtp_requires_model_declared_output_key(output_key):
+    state = _make_state()
+    model = SimpleNamespace(mtp=lambda: None, mtp_output_key=output_key)
+    with pytest.raises(TypeError, match="must declare mtp_output_key"):
+        state._init_mtp_runner(model)

vllm_omni/config/model.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from dataclasses import MISSING, field
 from typing import Any
 
@@ -124,6 +127,8 @@ class OmniModelConfig(ModelConfig):
     async_chunk: bool = False
     session_mode: str = "turn"
     retains_state_across_chunks: bool = False
+    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.
     active_stream_window: int = 0
     duplex_max_sessions: int = 1
@@ -134,6 +139,7 @@ class OmniModelConfig(ModelConfig):
     # Optional dotted path of a per-stage pooling-output decoder applied
     # worker-side before IPC. Read by the AR scheduler.
     pooling_output_decoder: str | None = None
+    final_output: bool = False
     hf_config_name: str | None = None
     custom_process_next_stage_input_func: str | None = None
     stage_connector_config: dict[str, Any] = field(

vllm_omni/config/omni_config.py

@@ -170,7 +170,10 @@ class _ModelEngineOverrides(TypedDict, total=False):
     limit_mm_per_prompt: dict[str, Any]
     interleave_mm_strings: bool
     media_io_kwargs: dict[str, Any]
+    final_output: bool
     active_stream_window: int
+    use_v2_model_runner: bool
+    supports_native_mrv2_data_plane: bool
     enable_sleep_mode: bool
     subtalker_sampling_params: dict[str, Any]
     silence_ban_frames: int
@@ -493,9 +496,12 @@ class OmniStageModelConfig(_TrackExplicitConfigFields):
     # MiniCPM interleaved AV packing and media decode knobs (Daily-Omni).
     interleave_mm_strings: bool | None = None
     media_io_kwargs: dict[str, Any] | None = None
+    final_output: bool = False
     active_stream_window: int = Field(default=0, ge=0)
     session_mode: str = "turn"
     duplex_max_sessions: int = Field(default=1, ge=1)
+    use_v2_model_runner: bool = False
+    supports_native_mrv2_data_plane: bool = False
     enable_sleep_mode: bool = False
     default_sampling_params: dict[str, Any] | None = None
     subtalker_sampling_params: dict[str, Any] | None = None
@@ -1965,8 +1971,14 @@ def _build_model_config(
         kwargs["dtype"] = _copy_value(deploy.dtype)
     if "active_stream_window" not in kwargs:
         kwargs["active_stream_window"] = _copy_value(deploy.active_stream_window)
+    kwargs["final_output"] = topology.final_output
     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)
+    kwargs.setdefault("use_v2_model_runner", deploy.model_runner == "v2")
+    kwargs.setdefault(
+        "supports_native_mrv2_data_plane",
+        topology.supports_native_mrv2_data_plane,
+    )
     if "has_sampling_extra_args" not in kwargs:
         kwargs["has_sampling_extra_args"] = bool((default_sampling_params or {}).get("extra_args"))
     if "model_subdir" not in kwargs and topology.model_subdir is not None:

vllm_omni/config/stage_config.py

@@ -11,7 +11,7 @@ from collections.abc import Callable, Mapping
 from dataclasses import asdict, dataclass, field, fields
 from enum import Enum
 from pathlib import Path
-from typing import Any, NamedTuple
+from typing import Any, Literal, NamedTuple
 
 from transformers import PretrainedConfig
 from vllm.logger import init_logger
@@ -250,6 +250,7 @@ class StagePipelineConfig:
     # Alternates picked by ``merge_pipeline_deploy`` based on ``deploy.async_chunk``.
     async_chunk_process_next_stage_input_func: str | None = None
     sync_process_input_func: str | None = None
+    supports_native_mrv2_data_plane: bool = False
     # Rewrites the Stage-0 view of a raw prompt before vLLM input processing.
     # The callable receives ``(prompt, sampling_params_list)``; downstream
     # stages continue to receive the original prompt.
@@ -566,6 +567,7 @@ class DeployConfig:
 
     async_chunk: bool = True
     session_mode: str = "turn"
+    model_runner: Literal["v1", "v2"] = "v1"
     # Stage-1 active stream slots; 0 preserves legacy all-stream cycling.
     active_stream_window: int = 0
     duplex_session: DuplexSessionRuntimeConfig = field(default_factory=DuplexSessionRuntimeConfig)
@@ -736,6 +738,16 @@ def _merge_platforms(
     return merged
 
 
+def _merge_connectors(
+    base: dict[str, Any] | None,
+    overlay: dict[str, Any] | None,
+) -> dict[str, Any] | None:
+    """Deep-merge named connector definitions from a deploy overlay."""
+    if not base and not overlay:
+        return None
+    return _get_recursively_merged_dict(base or {}, overlay or {})
+
+
 def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]:
     """Load a deploy YAML with optional ``base_config`` inheritance."""
     raw_dict = to_dict(load_yaml_config(path))
@@ -748,12 +760,15 @@ def resolve_deploy_yaml(path: str | Path) -> dict[str, Any]:
     base_path = Path(path).parent / base_path
     base_dict = resolve_deploy_yaml(base_path)
 
-    # Merge top-level scalars: overlay wins. ``stages:`` and ``platforms:``
-    # are deep-merged below so an overlay can layer on top of the base.
+    # Merge top-level scalars: overlay wins. Structured sections are merged
+    # below so a thin overlay does not discard inherited runtime contracts.
     merged = {
         **base_dict,
-        **{k: v for k, v in raw_dict.items() if k not in ("stages", "platforms")},
+        **{k: v for k, v in raw_dict.items() if k not in ("connectors", "stages", "platforms")},
     }
+    merged_connectors = _merge_connectors(base_dict.get("connectors"), raw_dict.get("connectors"))
+    if merged_connectors is not None:
+        merged["connectors"] = merged_connectors
     merged["stages"] = _merge_stage_lists(base_dict.get("stages"), raw_dict.get("stages"))
     merged_platforms = _merge_platforms(base_dict.get("platforms"), raw_dict.get("platforms"))
     if merged_platforms is not None:
@@ -773,9 +788,14 @@ def load_deploy_config(path: str | Path) -> DeployConfig:
 
     stages = [_parse_stage_deploy(s) for s in raw_dict.get("stages", [])]
 
+    model_runner = raw_dict.get("model_runner", "v1")
+    if model_runner not in ("v1", "v2"):
+        raise ValueError(f"model_runner must be one of ('v1', 'v2'), got {model_runner!r}")
+
     kwargs: dict[str, Any] = {
         "async_chunk": raw_dict.get("async_chunk", True),
         "session_mode": raw_dict.get("session_mode", "turn"),
+        "model_runner": model_runner,
         "active_stream_window": int(raw_dict.get("active_stream_window", 0) or 0),
         "duplex_session": DuplexSessionRuntimeConfig(**(raw_dict.get("duplex_session") or {})),
         "connectors": raw_dict.get("connectors", None),
@@ -828,15 +848,25 @@ def _apply_platform_overrides(
     deploy: DeployConfig,
     platform: str | None = None,
 ) -> DeployConfig:
-    """Merge platform-specific stage overrides into deploy config."""
+    """Merge platform-specific runner and stage overrides into deploy config."""
     if platform is None:
         from vllm_omni.platforms import current_omni_platform
 
         device_name = current_omni_platform.device_name
         platform = device_name.lower() if device_name is not None else None
+    platform_section = (deploy.platforms or {}).get(platform) if platform is not None else None
+    if platform_section is not None and "model_runner" in platform_section:
+        model_runner = platform_section["model_runner"]
+        if model_runner not in ("v1", "v2"):
+            raise ValueError(f"platform model_runner must be one of ('v1', 'v2'), got {model_runner!r}")
+        deploy.model_runner = model_runner
+    if deploy.model_runner == "v2" and platform in {"npu", "xpu"}:
+        raise NotImplementedError(
+            f"Model Runner V2 is not supported on {platform.upper()}: "
+            "the platform worker still uses the legacy chunk-transfer data plane."
+        )
     if platform is None or deploy.platforms is None:
         return deploy
-    platform_section = deploy.platforms.get(platform)
     if platform_section is None:
         return deploy
 
@@ -998,6 +1028,28 @@ def _build_engine_args(
         # model-owned streaming state. Propagate it to every stage instead of
         # making individual models duplicate the value in connector extras.
         engine_args["duplex_max_sessions"] = deploy.duplex_session.max_sessions
+    # The runner selection is a deploy-topology decision owned by the
+    # ``model_runner`` field; do not let an opaque ``engine_extras`` entry
+    # silently veto or force it per stage.
+    if ds is not None:
+        for reserved in ("use_v2_model_runner", "supports_native_mrv2_data_plane"):
+            if reserved in ds.engine_extras:
+                raise ValueError(
+                    f"stage {ds.stage_id}: {reserved!r} must not be set via engine_extras; "
+                    "it is derived from the deploy-level `model_runner` field and the "
+                    "pipeline's `supports_native_mrv2_data_plane` declaration."
+                )
+    engine_args["use_v2_model_runner"] = deploy.model_runner == "v2"
+    engine_args["supports_native_mrv2_data_plane"] = bool(ps.supports_native_mrv2_data_plane)
+    if deploy.model_runner == "v2" and not ps.supports_native_mrv2_data_plane:
+        logger.warning(
+            "Stage %s (%s) selects model_runner=v2 without declaring "
+            "supports_native_mrv2_data_plane. It will use the legacy transport path; "
+            "MRV2 support for this pipeline has not been validated. Use model_runner=v1 "
+            "unless you are validating a new MRV2 integration.",
+            ps.stage_id,
+            ps.model_arch or pipeline.model_arch or pipeline.model_type,
+        )
     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
@@ -1196,6 +1248,9 @@ class StageConfig:
                 else:
                     engine_args[key] = value
 
+        # Terminal-stage ownership comes from topology, not engine overrides.
+        engine_args["final_output"] = self.final_output
+
         # Build runtime config from YAML defaults + CLI overrides
         runtime: dict[str, Any] = dict(self.yaml_runtime)
         runtime.setdefault("process", True)

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -33,6 +33,21 @@ from vllm_omni.engine.serialization import deserialize_additional_information
 logger = init_logger(__name__)
 
 
+def _should_emit_engine_output(
+    model_config: Any,
+    *,
+    stopped: bool,
+    has_control: bool,
+) -> bool:
+    if stopped or has_control:
+        return True
+    return not (
+        bool(getattr(model_config, "use_v2_model_runner", False))
+        and bool(getattr(model_config, "async_chunk", False))
+        and not bool(getattr(model_config, "final_output", False))
+    )
+
+
 class SampledLogprobContractError(RuntimeError):
     """The model runner returned unusable sampled-token logprobs."""
 
@@ -84,6 +99,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
     core scheduling logic.
     """
 
+    max_num_running_reqs: int
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         # Track requests that need KV cache transfer when finished
@@ -200,9 +217,10 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         finished_status: RequestStatus,
     ) -> list[Request]:
         """Finish requests and discard any incomplete KV-wait timing."""
+        cleanup_ids: Iterable[str]
         if isinstance(request_ids, str):
             cleanup_ids = (request_ids,)
-            finish_request_ids: str | tuple[str, ...] | None = request_ids
+            finish_request_ids: str | Iterable[str] | None = request_ids
         elif request_ids is None:
             cleanup_ids = ()
             finish_request_ids = None
@@ -329,15 +347,22 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._process_pending_omni_inputs(model_mode="ar")
         self._drop_aborted_queued_requests()
         self._resync_streaming_input_counter()
-
         original_waiting = None
         if self._should_defer_waiting_admission():
             original_waiting = waiting
             self.waiting = create_request_queue(self.policy)
 
+        original_max_num_running_reqs = self.max_num_running_reqs
+        async_chunk_transport = self._async_chunk_transport_enabled()
+        reserved_running_slots = (
+            self._get_async_chunk_reserved_running_slots() if async_chunk_transport and self.use_v2_model_runner else 0
+        )
+        if reserved_running_slots:
+            self.max_num_running_reqs = max(0, original_max_num_running_reqs - reserved_running_slots)
         try:
             scheduler_output = super().schedule(throttle_prefills)
         finally:
+            self.max_num_running_reqs = original_max_num_running_reqs
             if original_waiting is not None:
                 deferred_waiting = list(self.waiting)
                 if deferred_waiting:
@@ -441,7 +466,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             assert num_tokens_scheduled > 0
             request = self.requests.get(req_id)
             if request is not None:
-                # vLLM 0.26: settle the in-flight tokens counted in schedule().
+                # Settle the in-flight tokens counted in schedule().
                 # Must happen before the skips below — failed-KV-load and
                 # already-finished requests were incremented too, and the two
                 # readers (allocate_slots, _connector_finished) clamp with
@@ -609,6 +634,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             confirmed_num_computed_tokens = None
             boundary_generation = None
+            # Capture before resumable stop handling can clear token history.
+            output_token_ids: Any = getattr(request, "output_token_ids", None)
+            if output_token_ids is None:
+                output_token_ids = getattr(request, "_output_token_ids", ())
+            num_generation_tokens = len(output_token_ids)
             if stopped:
                 if self.chunk_transfer_adapter is not None:
                     confirmed_num_computed_tokens = self.chunk_transfer_adapter._confirmed_num_computed_tokens(request)
@@ -627,6 +657,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 is_segment_finished = not finished
                 if finished:
                     request.resumable = False
+                    if self._native_data_plane:
+                        self._pending_data_plane_terminal_req_ids.add(req_id)
                 if not finished:
                     # for streaming input request only
                     if self.chunk_transfer_adapter:
@@ -683,7 +715,18 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
-            if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped:
+            has_stage_output = (
+                bool(new_token_ids)
+                or mm_output is not None
+                or pooler_output is not None
+                or kv_transfer_params
+                or stopped
+            )
+            if has_stage_output and _should_emit_engine_output(
+                self.vllm_config.model_config,
+                stopped=stopped,
+                has_control=kv_transfer_params is not None,
+            ):
                 OmniSchedulerMixin._append_request_output(
                     self,
                     outputs,
@@ -702,6 +745,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     num_nans_in_logits=request.num_nans_in_logits,
                     is_segment_finished=is_segment_finished,
                     new_prompt_len_snapshot=self._new_prompt_len_snapshot.get(req_id),
+                    num_generation_tokens=num_generation_tokens,
                 )
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
@@ -877,7 +921,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 # This streaming update has already been dequeued. Report the
                 # permanent contract failure so the next scheduling pass
                 # finishes only this request instead of crashing EngineCore.
-                self.chunk_transfer_adapter.record_receive_failure(req_id, str(exc))
+                if self.chunk_transfer_adapter is not None:
+                    self.chunk_transfer_adapter.record_receive_failure(req_id, str(exc))
                 return
             if replaced is not None:
                 if replaced:

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -5,7 +5,7 @@ from __future__ import annotations
 
 import time
 from collections import defaultdict
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.logger import init_logger
@@ -15,11 +15,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.request_queue import create_request_queue
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
 from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import (
-    EngineCoreEventType,
-    EngineCoreOutput,
-    EngineCoreOutputs,
-)
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
 from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
@@ -27,12 +23,33 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
 from vllm_omni.core.sched.utils import omni_routed_experts_for_request
+from vllm_omni.engine.serialization import deserialize_additional_information
 from vllm_omni.outputs import OmniModelRunnerOutput
 
+if TYPE_CHECKING:
+    from vllm.v1.core.sched.request_queue import RequestQueue
+
+
 logger = init_logger(__name__)
 
 
+def _has_async_chunk_payload_to_run(request: Request) -> bool:
+    additional_information = getattr(request, "additional_information", None)
+    if not isinstance(additional_information, dict):
+        additional_information = deserialize_additional_information(additional_information)
+
+    codes = additional_information.get("codes")
+    audio = codes.get("audio") if isinstance(codes, dict) else additional_information.get("codes.audio")
+    if audio is None:
+        return False
+    if hasattr(audio, "numel"):
+        return bool(audio.numel())
+    return bool(audio)
+
+
 class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
+    waiting: RequestQueue
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         model_config = self.vllm_config.model_config
@@ -40,6 +57,25 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False))
         self._pending_finish_reqs: list[Request] = []
 
+    def _is_done_receiving_chunks(self, request_id: str) -> bool:
+        if self.chunk_transfer_adapter is not None:
+            return self.chunk_transfer_adapter.is_done_receiving_chunks(request_id)
+        coordinator = self.input_coordinator
+        return bool(self._native_data_plane and coordinator is not None and request_id in coordinator.finished_requests)
+
+    def _input_execution_is_terminal(
+        self,
+        request_id: str,
+        scheduler_output: SchedulerOutput,
+    ) -> bool:
+        if self._native_data_plane:
+            return request_id in getattr(
+                scheduler_output,
+                "input_terminal_req_ids",
+                set(),
+            )
+        return self._is_done_receiving_chunks(request_id)
+
     @staticmethod
     def _record_prefill_stats(request: Request) -> None:
         """Mirror upstream first-prefill prompt accounting for generation stages.
@@ -73,6 +109,29 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             return False
         return super()._handle_stopped_request(request)
 
+    def _requeue_completed_native_chunks(self) -> None:
+        """Return completed native generation chunks to the admission queue.
+
+        A generation batch slot belongs to one chunk, not the entire stream.
+        Keep requests with outstanding output in ``running`` so their state
+        cannot be reused before completion. Other live requests join the tail
+        of ``waiting`` without freeing their decoder state or KV blocks.
+        Preserve WAITING_FOR_CHUNK until the coordinator observes fresh input.
+        Stateful codecs retain lifetime admission until their decoder state is
+        released, even when no chunk is currently executing.
+        """
+        if not getattr(self, "_native_data_plane", False) or self._retains_state_across_chunks:
+            return
+        in_flight: list[Request] = []
+        for request in self.running:
+            if request.num_in_flight_tokens > 0 or request.is_finished():
+                in_flight.append(request)
+                continue
+            if request.status == RequestStatus.RUNNING:
+                request.status = RequestStatus.WAITING
+            self.waiting.add_request(request)
+        self.running = in_flight
+
     def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
         """One-shot generation fast path:
         - Feed all input tokens of the request at once
@@ -82,6 +141,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         """
 
         token_budget = self.max_num_scheduled_tokens
+        execution_batch_size = self.max_num_running_reqs
         if self._pause_state == PauseState.PAUSED_ALL:
             token_budget = 0
         scheduled_timestamp = time.monotonic()
@@ -102,14 +162,22 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         skipped_waiting_requests = create_request_queue(self.policy)
         req_index = 0
         self._drop_aborted_queued_requests()
+        self._requeue_completed_native_chunks()
         self._process_pending_omni_inputs(model_mode="generation")
         self._drop_aborted_queued_requests()
         self._resync_streaming_input_counter()
+        async_chunk_transport = self._async_chunk_transport_enabled()
+        native_chunks = bool(getattr(self, "_native_data_plane", False))
+        # Parking releases an execution slot, but stateful codecs retain their
+        # request-owned state until completion or abort, in either runner.
+        reserved_running_slots = (
+            self._get_async_chunk_reserved_running_slots() if self._retains_state_across_chunks else 0
+        )
 
         # OMNI: Track requests that are already finished (e.g., marked by connector)
         # These should be removed from running and not scheduled
         already_finished_reqs: set[Request] = set()
-        while req_index < len(self.running) and token_budget > 0:
+        while req_index < len(self.running) and token_budget > 0 and len(num_scheduled_tokens) < execution_batch_size:
             request = self.running[req_index]
             # OMNI: Skip requests that are not in self.requests
             if request.request_id not in self.requests or (
@@ -119,6 +187,12 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 req_index += 1
                 continue
 
+            if native_chunks and request.num_in_flight_tokens > 0:
+                # Overlap different streams, never two stateful chunks from
+                # the same stream. Readiness remains pending for the next step.
+                req_index += 1
+                continue
+
             num_computed_tokens = request.num_computed_tokens
             required_tokens = len(request.prompt_token_ids) - num_computed_tokens
             if not self.scheduler_config.enable_chunked_prefill and required_tokens > token_budget:
@@ -127,12 +201,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 break
             # async_chunk: don't schedule placeholder tokens when no new chunk is available.
             if required_tokens <= 0:
-                if self.chunk_transfer_adapter is not None and self.chunk_transfer_adapter.is_done_receiving_chunks(
-                    request.request_id
-                ):
-                    self._pending_finish_reqs.append(request)
-                req_index += 1
-                continue
+                if async_chunk_transport and self._is_done_receiving_chunks(request.request_id):
+                    if len(request.prompt_token_ids) == 0 and _has_async_chunk_payload_to_run(request):
+                        required_tokens = 1
+                    else:
+                        self._pending_finish_reqs.append(request)
+                        req_index += 1
+                        continue
+                else:
+                    req_index += 1
+                    continue
             num_new_tokens = min(required_tokens, token_budget)
             new_blocks = self.kv_cache_manager.allocate_slots(
                 request,
@@ -161,20 +239,27 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         if already_finished_reqs:
             self.running = remove_all(self.running, already_finished_reqs)
 
-        # Fast path selection and scheduling for one-shot generation requests,
-        # independent of pooling_params.
-        while self.waiting and token_budget > 0 and self._pause_state == PauseState.UNPAUSED:
-            # Requests waiting for their next chunk are temporarily absent
-            # from `running`, but stateful models still retain their model
-            # runner slot. Mirror vLLM's treatment of
-            # `num_waiting_for_streaming_input` when enforcing max_num_seqs.
-            num_running = len(self.running)
-            if self._retains_state_across_chunks and self.chunk_transfer_adapter is not None:
-                num_running += self.chunk_transfer_adapter.num_running_waiting_for_chunk
-            if num_running >= self.max_num_running_reqs:
-                break
-
+        # Fast path selection and scheduling (treat all as diffusion requests,
+        # independent of pooling_params)
+        while (
+            self.waiting
+            and token_budget > 0
+            and len(num_scheduled_tokens) < execution_batch_size
+            and (
+                len(num_scheduled_tokens)
+                if native_chunks and not self._retains_state_across_chunks
+                else len(self.running) + reserved_running_slots
+            )
+            < self.max_num_running_reqs
+            and self._pause_state == PauseState.UNPAUSED
+        ):
             request = self.waiting.peek_request()
+            if native_chunks and request.num_in_flight_tokens > 0:
+                # A restored waiting entry can still own an outstanding batch.
+                # Do not let it block other ready streams or execute it twice.
+                self.waiting.pop_request()
+                skipped_waiting_requests.add_request(request)
+                continue
             # OMNI: Skip requests that are not in self.requests
             if request.request_id not in self.requests or (
                 self.chunk_transfer_adapter is None and request.status == RequestStatus.FINISHED_STOPPED
@@ -184,11 +269,12 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 continue
 
             # async_chunk: wait for the first upstream chunk (don't start with placeholders).
-            if self.chunk_transfer_adapter is not None and len(request.prompt_token_ids) == 0:
-                if self.chunk_transfer_adapter.is_done_receiving_chunks(request.request_id):
-                    self.waiting.pop_request()
-                    self._pending_finish_reqs.append(request)
-                    continue
+            if async_chunk_transport and len(request.prompt_token_ids) == 0:
+                if self._is_done_receiving_chunks(request.request_id):
+                    if not _has_async_chunk_payload_to_run(request):
+                        self.waiting.pop_request()
+                        self._pending_finish_reqs.append(request)
+                        continue
                 else:
                     self.waiting.pop_request()
                     skipped_waiting_requests.prepend_request(request)
@@ -197,6 +283,22 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             # Allocate all input tokens for the request in one shot
             # (allocate 1 placeholder if zero)
             required_tokens = max(len(request.prompt_token_ids), 1)
+            if native_chunks:
+                required_tokens = len(request.prompt_token_ids) - request.num_computed_tokens
+                if required_tokens <= 0:
+                    if (
+                        self._is_done_receiving_chunks(request.request_id)
+                        and not request.prompt_token_ids
+                        and _has_async_chunk_payload_to_run(request)
+                    ):
+                        required_tokens = 1
+                    else:
+                        self.waiting.pop_request()
+                        if self._is_done_receiving_chunks(request.request_id):
+                            self._pending_finish_reqs.append(request)
+                        else:
+                            skipped_waiting_requests.add_request(request)
+                        continue
             num_new_tokens = min(required_tokens, token_budget)
             new_blocks = self.kv_cache_manager.allocate_slots(
                 request,
@@ -229,7 +331,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         # If fast path scheduled none, fall back to the original scheduling
         if not num_scheduled_tokens:
-            if self.chunk_transfer_adapter:
+            if async_chunk_transport:
                 # Don't fall back: base scheduler doesn't handle async_chunk
                 # requests with empty prompt_token_ids.
                 self._restore_omni_wait_queues()
@@ -245,7 +347,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             any_request = self.running[0]
             num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(any_request.request_id)
 
-        # Assemble SchedulerOutput (align with v0.14.0)
+        # Assemble the current SchedulerOutput contract.
         if self.use_v2_model_runner:
             # No resumed reqs in fast path; pass prefill_token_ids for new reqs.
             new_reqs_data = [
@@ -284,7 +386,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         total_num_scheduled_tokens = sum(num_scheduled_tokens.values())
 
-        # Record the request ids scheduled in this step (v0.14.0 behavior).
+        # Record request ids scheduled in this step for the next update.
         self.prev_step_scheduled_req_ids.clear()
         self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys())
 
@@ -404,7 +506,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             assert num_tokens_scheduled > 0
             request = self.requests.get(req_id)
             if request is not None:
-                # vLLM 0.26: settle the in-flight tokens counted in schedule().
+                # Settle the in-flight tokens counted in schedule().
                 # Must happen before the skips below — failed-KV-load and
                 # already-finished requests were incremented too, and the two
                 # readers (allocate_slots, _connector_finished) clamp with
@@ -496,10 +598,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             if (
                 request.status == RequestStatus.FINISHED_ERROR
                 or request.status == RequestStatus.FINISHED_STOPPED
-                or (self.chunk_transfer_adapter is None and request.num_computed_tokens >= request.num_prompt_tokens)
                 or (
-                    self.chunk_transfer_adapter is not None
-                    and self.chunk_transfer_adapter.is_done_receiving_chunks(request.request_id)
+                    not self._async_chunk_transport_enabled()
+                    and request.num_computed_tokens >= request.num_prompt_tokens
+                )
+                or (
+                    self._async_chunk_transport_enabled()
+                    and self._input_execution_is_terminal(
+                        request.request_id,
+                        scheduler_output,
+                    )
                     and request.num_computed_tokens >= len(request.prompt_token_ids)
                 )
             ):
@@ -526,6 +634,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 finish_reason = request.get_finished_reason()
                 finished = self._handle_stopped_request(request)
                 is_segment_finished = not finished
+                if finished and self._native_data_plane:
+                    self._pending_data_plane_terminal_req_ids.add(req_id)
                 if not finished:
                     # for streaming input request only
                     if self.chunk_transfer_adapter:

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -5,17 +5,20 @@ from __future__ import annotations
 
 import math
 import os
+import queue
 import time
 from collections.abc import Iterable, Iterator
 from typing import Any
 
 import torch
 from vllm.compilation.cuda_graph import CUDAGraphStat
+from vllm.config import VllmConfig
 from vllm.distributed.kv_events import KVEventBatch
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
 from vllm.utils.import_utils import resolve_obj_by_qualname
 from vllm.v1.core.sched.output import SchedulerOutput
+from vllm.v1.core.sched.request_queue import RequestQueue
 from vllm.v1.core.sched.utils import remove_all
 from vllm.v1.engine import (
     EngineCoreEventType,
@@ -31,6 +34,7 @@ from vllm.v1.spec_decode.metrics import SpecDecodingStats
 from vllm_omni.core.sched.omni_scheduling_coordinator import (
     OmniSchedulingCoordinator,
     uses_full_payload_input_coordinator,
+    uses_native_mrv2_data_plane,
 )
 from vllm_omni.core.sched.output import (
     OmniChunkRecvHandle,
@@ -41,6 +45,8 @@ from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapt
     OmniChunkTransferAdapter,
 )
 from vllm_omni.engine import OmniEngineCoreOutput
+from vllm_omni.engine.serialization import serialize_additional_information
+from vllm_omni.outputs import OmniConnectorOutput
 
 logger = init_logger(__name__)
 
@@ -101,6 +107,19 @@ elif DEFAULT_INPUT_WAIT_TIMEOUT_S == 0:
 class OmniSchedulerMixin:
     """Shared scheduler helpers for omni-specific request handling."""
 
+    # Provided by the concrete vLLM scheduler.
+    vllm_config: VllmConfig
+    requests: dict[str, Request]
+    waiting: RequestQueue
+    running: list[Request]
+
+    def _init_omni_connector_output_inbox(self) -> None:
+        self._omni_connector_output_inbox: queue.SimpleQueue[OmniConnectorOutput] = queue.SimpleQueue()
+
+    def enqueue_omni_connector_output(self, output: OmniConnectorOutput) -> None:
+        """Accept a lightweight readiness event from the local runner thread."""
+        self._omni_connector_output_inbox.put(output)
+
     # ------------------------------------------------------------------ #
     #  Shared scheduler/output helpers (lift the AR / generation duplicates)
     # ------------------------------------------------------------------ #
@@ -119,15 +138,30 @@ class OmniSchedulerMixin:
             self.use_v2_model_runner = False
 
         model_config = self.vllm_config.model_config
-        self.chunk_transfer_adapter = (
-            OmniChunkTransferAdapter(self.vllm_config) if getattr(model_config, "async_chunk", False) else None
+        self.use_v2_model_runner = bool(getattr(model_config, "use_v2_model_runner", False))
+        self._native_data_plane = uses_native_mrv2_data_plane(
+            model_config,
+            use_v2_model_runner=self.use_v2_model_runner,
         )
-        self.input_coordinator = (
-            OmniSchedulingCoordinator(stage_id=getattr(model_config, "stage_id", 0))
-            if uses_full_payload_input_coordinator(model_config)
+        self.chunk_transfer_adapter = (
+            OmniChunkTransferAdapter(self.vllm_config)
+            if getattr(model_config, "async_chunk", False) and not self._native_data_plane
             else None
         )
-        self._latest_omni_connector_output = None
+        self.input_coordinator: OmniSchedulingCoordinator | None = None
+        if self._native_data_plane:
+            self.input_coordinator = OmniSchedulingCoordinator(
+                scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs,
+                stage_id=getattr(model_config, "stage_id", 0),
+                async_chunk=True,
+            )
+        elif uses_full_payload_input_coordinator(model_config):
+            self.input_coordinator = OmniSchedulingCoordinator(
+                stage_id=getattr(model_config, "stage_id", 0),
+            )
+        self._latest_omni_connector_output: OmniConnectorOutput | None = None
+        self._init_omni_connector_output_inbox()
+        self._pending_data_plane_terminal_req_ids: set[str] = set()
         # Optional per-stage pooling-output decoder hook (dotted path in
         # model_config); applied worker-side before IPC.
         self._pooling_output_decoder = None
@@ -214,6 +248,35 @@ class OmniSchedulerMixin:
         session.num_prompt_tokens = len(new_prompt)
         self._finish_streaming_session_update(session, update)
 
+    def _async_chunk_transport_enabled(self) -> bool:
+        return getattr(self, "chunk_transfer_adapter", None) is not None or bool(
+            getattr(self, "_native_data_plane", False)
+        )
+
+    def _get_async_chunk_reserved_running_slots(self) -> int:
+        adapter = getattr(self, "chunk_transfer_adapter", None)
+        live_requests = self.requests
+        reserved_ids: set[str] = set()
+        if adapter is not None:
+            parked = getattr(adapter, "waiting_for_chunk_running_requests", ())
+            held = getattr(adapter, "_held_non_active", ())
+        else:
+            coordinator = getattr(self, "input_coordinator", None)
+            parked = getattr(coordinator, "_waiting_for_chunk_running", ())
+            held = ()
+        for request in parked:
+            req_id = getattr(request, "request_id", None)
+            if req_id is not None and req_id in live_requests:
+                reserved_ids.add(req_id)
+        for request in held:
+            req_id = getattr(request, "request_id", None)
+            if req_id is not None and req_id in live_requests:
+                reserved_ids.add(req_id)
+        return len(reserved_ids)
+
+    # ------------------------------------------------------------------ #
+    #  Shared scheduler/output helpers (lift the AR / generation duplicates)
+    # ------------------------------------------------------------------ #
     def _release_replaced_streaming_prompt_cache(self, session: Request) -> None:
         """Discard cache state that belongs to a replaced prompt."""
         # A prompt replacement is not a normal streaming extension: none of
@@ -259,19 +322,55 @@ class OmniSchedulerMixin:
         AR and generation schedulers except for the ``model_mode`` argument
         forwarded to ``update_request_metadata``.
         """
+        connector_outputs: list[OmniConnectorOutput] = []
+        inbox = getattr(self, "_omni_connector_output_inbox", None)
+        if inbox is not None:
+            while True:
+                try:
+                    connector_outputs.append(inbox.get_nowait())
+                except queue.Empty:
+                    break
         connector_output = getattr(self, "_latest_omni_connector_output", None)
         self._latest_omni_connector_output = None
+        if connector_output is not None:
+            connector_outputs.append(connector_output)
         input_coordinator = getattr(self, "input_coordinator", None)
         if input_coordinator is None:
             return
-        if connector_output and connector_output.request_metadata:
+        request_metadata: dict[str, dict[str, Any]] = {}
+        chunk_ready_req_ids: set[str] = set()
+        chunk_finished_req_ids: set[str] = set()
+        stage_recv_req_ids: set[str] = set()
+        for output in connector_outputs:
+            request_metadata.update(output.request_metadata)
+            chunk_ready_req_ids.update(output.chunk_ready_req_ids)
+            chunk_finished_req_ids.update(output.chunk_finished_req_ids)
+            stage_recv_req_ids.update(output.stage_recv_req_ids)
+        live_request_ids = self.requests.keys()
+        request_metadata = {
+            req_id: metadata for req_id, metadata in request_metadata.items() if req_id in live_request_ids
+        }
+        chunk_ready_req_ids.intersection_update(live_request_ids)
+        chunk_finished_req_ids.intersection_update(live_request_ids)
+        stage_recv_req_ids.intersection_update(live_request_ids)
+        if request_metadata:
             input_coordinator.update_request_metadata(
-                self.requests, connector_output.request_metadata, model_mode=model_mode
+                self.requests,
+                request_metadata,
+                model_mode=model_mode,
+            )
+        if input_coordinator._async_chunk:
+            input_coordinator.process_pending_chunks(
+                self.waiting,
+                self.running,
+                chunk_ready_req_ids,
+                chunk_finished_req_ids,
+            )
+        else:
+            input_coordinator.process_pending_full_payload_inputs(
+                self.waiting,
+                stage_recv_req_ids,
             )
-        input_coordinator.process_pending_full_payload_inputs(
-            self.waiting,
-            connector_output.stage_recv_req_ids if connector_output else set(),
-        )
 
     def _process_pending_omni_inputs(self, model_mode: str) -> None:
         """Apply pending connector inputs, timeouts, and async chunks."""
@@ -316,7 +415,7 @@ class OmniSchedulerMixin:
                 scheduler_requests=self.requests,
             )
         if self.input_coordinator:
-            self.input_coordinator.restore_queues(self.waiting)
+            self.input_coordinator.restore_queues(self.waiting, self.running)
 
     def _process_pending_input_timeouts(self) -> None:
         """Force-fail requests waiting on the full-payload coordinator too long.
@@ -458,6 +557,8 @@ class OmniSchedulerMixin:
         *,
         finished_requests_needing_kv_transfer: dict | None = None,
         pending_input_registrations: list[OmniChunkRecvHandle] | None = None,
+        data_plane_terminal_req_ids: set[str] | None = None,
+        input_terminal_req_ids: set[str] | None = None,
     ) -> OmniSchedulerOutput:
         """Wrap a base ``SchedulerOutput`` in ``OmniSchedulerOutput``.
 
@@ -468,11 +569,31 @@ class OmniSchedulerMixin:
         base_data = {name: getattr(base, name) for name in SchedulerOutput.__dataclass_fields__}
         input_coordinator = getattr(self, "input_coordinator", None)
         if pending_input_registrations is None:
-            pending_input_registrations = input_coordinator.pending_input_registrations if input_coordinator else []
+            if input_coordinator is None:
+                pending_input_registrations = []
+            elif input_coordinator._async_chunk:
+                pending_input_registrations = input_coordinator.pending_chunk_registrations
+            else:
+                pending_input_registrations = input_coordinator.pending_input_registrations
+        if data_plane_terminal_req_ids is None:
+            pending_terminal: set[str] = getattr(self, "_pending_data_plane_terminal_req_ids", set())
+            data_plane_terminal_req_ids = set(pending_terminal)
+            pending_terminal.clear()
+        if input_coordinator is not None:
+            scheduled_terminal_req_ids = input_coordinator.get_scheduled_input_terminal_req_ids(base)
+            if input_terminal_req_ids is None:
+                input_terminal_req_ids = scheduled_terminal_req_ids
+            else:
+                input_terminal_req_ids = set(input_terminal_req_ids).union(
+                    scheduled_terminal_req_ids,
+                )
+            input_coordinator.postprocess_scheduler_output(base)
         return OmniSchedulerOutput(
             **base_data,
             finished_requests_needing_kv_transfer=finished_requests_needing_kv_transfer or {},
             pending_input_registrations=pending_input_registrations,
+            data_plane_terminal_req_ids=data_plane_terminal_req_ids,
+            input_terminal_req_ids=input_terminal_req_ids or set(),
         )
 
     def _rewrap_scheduled_new_reqs(self, scheduler_output: SchedulerOutput) -> None:
@@ -520,8 +641,14 @@ class OmniSchedulerMixin:
         num_nans_in_logits: int = 0,
         is_segment_finished: bool | None = False,
         new_prompt_len_snapshot: int | None = None,
+        num_generation_tokens: int | None = None,
     ) -> OmniEngineCoreOutput:
         """Build the common request-output envelope used by LLM schedulers."""
+        pooling_output_payload = None
+        if isinstance(pooling_output, dict):
+            pooling_output_payload = serialize_additional_information(pooling_output)
+            pooling_output = None
+
         return OmniEngineCoreOutput(
             request_id=request.request_id,
             new_token_ids=new_token_ids,
@@ -529,6 +656,7 @@ class OmniSchedulerMixin:
             new_logprobs=new_logprobs,
             new_prompt_logprobs_tensors=new_prompt_logprobs_tensors,
             pooling_output=pooling_output,
+            pooling_output_payload=pooling_output_payload,
             multimodal_output=multimodal_output,
             stop_reason=stop_reason,
             events=request.take_events(),
@@ -540,6 +668,7 @@ class OmniSchedulerMixin:
             num_nans_in_logits=num_nans_in_logits,
             is_segment_finished=is_segment_finished,
             new_prompt_len_snapshot=new_prompt_len_snapshot,
+            num_generation_tokens=num_generation_tokens,
         )
 
     def _append_request_output(
@@ -953,10 +1082,10 @@ class OmniSchedulerMixin:
         self.running[:] = [req for req in self.running if keep_running(req)]
 
     def _drop_aborted_queued_requests(self) -> None:
-        for queue in (self.waiting, self.skipped_waiting):
-            aborted = [req for req in queue if req.status == RequestStatus.FINISHED_ABORTED]
+        for request_queue in (self.waiting, self.skipped_waiting):
+            aborted = [req for req in request_queue if req.status == RequestStatus.FINISHED_ABORTED]
             if aborted:
-                queue.remove_requests(aborted)
+                request_queue.remove_requests(aborted)
         self.running[:] = [req for req in self.running if req.status != RequestStatus.FINISHED_ABORTED]
 
     def _resync_streaming_input_counter(self) -> None:

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -1,11 +1,9 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Scheduling-side coordination for full_payload input waiting.
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Scheduling-side coordination for chunk and full-payload input waiting.
 
-Manages WAITING_FOR_INPUT state transitions based on readiness signals
-from OmniConnectorOutput, without ever calling connector.put()/get().
-
-Chunk waiting (WAITING_FOR_CHUNK) lives on OmniChunkTransferAdapter.
+Manages WAITING_FOR_CHUNK and WAITING_FOR_INPUT state transitions based on
+readiness signals from OmniConnectorOutput, without calling connector I/O.
 """
 
 from __future__ import annotations
@@ -22,6 +20,18 @@ from vllm_omni.core.sched.output import OmniChunkRecvHandle
 logger = init_logger(__name__)
 
 
+def uses_native_mrv2_data_plane(
+    model_config: Any,
+    *,
+    use_v2_model_runner: bool,
+) -> bool:
+    return bool(
+        use_v2_model_runner
+        and getattr(model_config, "async_chunk", False)
+        and getattr(model_config, "supports_native_mrv2_data_plane", False)
+    )
+
+
 def uses_full_payload_input_coordinator(model_config: Any) -> bool:
     """Returns True if this stage parks pending requests in
     WAITING_FOR_INPUT awaiting a full_payload delivery on the worker connector.
@@ -38,7 +48,7 @@ def uses_full_payload_input_coordinator(model_config: Any) -> bool:
 
 
 class OmniSchedulingCoordinator:
-    """Pure-scheduling coordinator for full_payload input waiting.
+    """Pure-scheduling coordinator for chunk and full-payload input waiting.
 
     The Scheduler owns an instance of this class.  It consumes readiness
     signals produced by the Model Runner's ``OmniConnectorModelRunnerMixin``
@@ -46,12 +56,28 @@ class OmniSchedulingCoordinator:
     transitions accordingly.
     """
 
-    def __init__(self, stage_id: int = 0):
+    def __init__(
+        self,
+        scheduler_max_num_seqs: int = 0,
+        stage_id: int = 0,
+        async_chunk: bool = False,
+    ) -> None:
         self._stage_id = stage_id
+        self._scheduler_max_num_seqs = scheduler_max_num_seqs
+        self._async_chunk = async_chunk
 
         self.finished_requests: set[str] = set()
+        self.requests_with_ready_chunks: set[str] = set()
+        self.input_terminal_req_ids: set[str] = set()
         self._full_payload_input_received: set[str] = set()
 
+        self._waiting_for_chunk_waiting: deque[Any] = deque()
+        self._waiting_for_chunk_running: deque[Any] = deque()
+
+        # Request IDs that were newly registered for chunk recv this cycle.
+        # The engine/Model Runner should call register_chunk_recv() for these
+        # so the bg thread starts polling.
+        self.pending_chunk_registrations: list[OmniChunkRecvHandle] = []
         # Requests waiting for full_payload stage input (WAITING_FOR_INPUT).
         self._waiting_for_input: deque[Any] = deque()
         # Per-cycle list of minimal handles to ship to the model runner so it
@@ -69,6 +95,47 @@ class OmniSchedulingCoordinator:
     #  Core scheduling methods
     # ------------------------------------------------------------------ #
 
+    def process_pending_chunks(
+        self,
+        waiting_queue: Any,
+        running_queue: list[Request],
+        chunk_ready_req_ids: set[str],
+        chunk_finished_req_ids: set[str],
+    ) -> None:
+        """Transition requests whose chunks have arrived.
+
+        Args:
+            waiting_queue: Scheduler's waiting request queue.
+            running_queue: Scheduler's running request list.
+            chunk_ready_req_ids: IDs with a newly arrived chunk this cycle.
+            chunk_finished_req_ids: IDs whose final chunk has arrived.
+        """
+        if self._stage_id == 0 or not self._async_chunk:
+            return
+
+        self.finished_requests.update(chunk_finished_req_ids)
+        self.pending_chunk_registrations = []
+
+        self._process_chunk_queue(
+            waiting_queue,
+            self._waiting_for_chunk_waiting,
+            RequestStatus.WAITING,
+            chunk_ready_req_ids,
+        )
+        self._process_chunk_queue(
+            running_queue,
+            self._waiting_for_chunk_running,
+            RequestStatus.RUNNING,
+            chunk_ready_req_ids,
+        )
+        while len(running_queue) > self._scheduler_max_num_seqs:
+            request = running_queue.pop()
+            # Must reset status to WAITING so the scheduler treats it as
+            # schedulable work.  KV blocks are NOT freed here (unlike a
+            # real preemption), so PREEMPTED would be incorrect.
+            request.status = RequestStatus.WAITING
+            waiting_queue.prepend_requests([request])
+
     def process_pending_full_payload_inputs(
         self,
         waiting_queue: Any,
@@ -144,10 +211,29 @@ class OmniSchedulingCoordinator:
             waiting_queue.remove_requests(to_remove)
 
     def free_finished_request(self, request_id: str) -> None:
-        """Prune internal tracking sets for a freed request to prevent unbounded growth."""
+        """Prune all coordinator state owned by a freed request."""
         self._full_payload_input_received.discard(request_id)
         self.finished_requests.discard(request_id)
+        self.requests_with_ready_chunks.discard(request_id)
+        self.input_terminal_req_ids.discard(request_id)
         self._waiting_since.pop(request_id, None)
+        for queue_attr in (
+            "_waiting_for_chunk_waiting",
+            "_waiting_for_chunk_running",
+            "_waiting_for_input",
+        ):
+            queue = getattr(self, queue_attr)
+            setattr(
+                self,
+                queue_attr,
+                deque(request for request in queue if request.request_id != request_id),
+            )
+        self.pending_chunk_registrations = [
+            handle for handle in self.pending_chunk_registrations if handle.request_id != request_id
+        ]
+        self.pending_input_registrations = [
+            handle for handle in self.pending_input_registrations if handle.request_id != request_id
+        ]
 
     def collect_timed_out_request_ids(
         self,
@@ -178,11 +264,14 @@ class OmniSchedulingCoordinator:
 
         # Defensively remove from coordinator internal queues (may already
         # be empty if restore_queues() has run).
-        remaining: deque[Any] = deque()
-        for request in self._waiting_for_input:
-            if request.request_id not in timed_out_ids:
-                remaining.append(request)
-        self._waiting_for_input = remaining
+        for queue_attr in (
+            "_waiting_for_chunk_waiting",
+            "_waiting_for_chunk_running",
+            "_waiting_for_input",
+        ):
+            queue = getattr(self, queue_attr)
+            remaining = deque(request for request in queue if request.request_id not in timed_out_ids)
+            setattr(self, queue_attr, remaining)
 
         for req_id in timed_out_ids:
             self._waiting_since.pop(req_id, None)
@@ -198,8 +287,17 @@ class OmniSchedulingCoordinator:
     def restore_queues(
         self,
         waiting_queue: Any,
+        running_queue: list[Request] | None = None,
     ) -> None:
-        """Return waiting-for-input requests to the waiting queue."""
+        """Return waiting-for-chunk/input requests to scheduling queues."""
+        for request in self._waiting_for_chunk_waiting:
+            waiting_queue.add_request(request)
+        self._waiting_for_chunk_waiting = deque()
+
+        if running_queue is not None and self._waiting_for_chunk_running:
+            running_queue.extend(self._waiting_for_chunk_running)
+        self._waiting_for_chunk_running = deque()
+
         for request in self._waiting_for_input:
             waiting_queue.add_request(request)
         self._waiting_for_input = deque()
@@ -247,6 +345,9 @@ class OmniSchedulingCoordinator:
             if request is None:
                 continue
 
+            if metadata.get("input_terminal") is True:
+                self.input_terminal_req_ids.add(req_id)
+
             # Handle next_stage_prompt_len if present (for models like Qwen3-Omni).
             # Only apply when the request has not started decoding yet
             # (no output tokens). Resetting a mid-decode request would
@@ -255,8 +356,7 @@ class OmniSchedulingCoordinator:
                 next_len = metadata["next_stage_prompt_len"]
                 if isinstance(next_len, int) and next_len > 0:
                     output_token_ids = getattr(request, "_output_token_ids", None)
-                    has_decode_output = output_token_ids is not None and len(output_token_ids) > 0
-                    if has_decode_output:
+                    if output_token_ids is not None and len(output_token_ids) > 0:
                         logger.debug(
                             "[Coordinator stage-%s] Skipping prompt resize for req %s: "
                             "request already has %s output tokens",
@@ -284,6 +384,12 @@ class OmniSchedulingCoordinator:
 
             if model_mode != "ar":
                 new_ids = self._flatten_prompt_token_ids(metadata.get("code_predictor_codes"))
+                runtime_seed = None
+                if "left_context_size" in metadata:
+                    runtime_seed = {
+                        "meta": {"left_context_size": metadata["left_context_size"]},
+                    }
+                request._omni_initial_model_buffer = runtime_seed
                 if new_ids:
                     request.prompt_token_ids = new_ids
                     request.num_prompt_tokens = len(new_ids)
@@ -291,3 +397,103 @@ class OmniSchedulingCoordinator:
                     request._all_token_ids.extend(new_ids)
                     request._output_token_ids.clear()
                     request.num_computed_tokens = 0
+
+    def postprocess_scheduler_output(
+        self,
+        scheduler_output: Any,
+        requests: dict[str, Request] | None = None,
+    ) -> None:
+        """Clear per-cycle ready state after scheduler output is materialized."""
+        self._clear_chunk_ready(scheduler_output)
+        self.input_terminal_req_ids.difference_update(
+            self._scheduled_request_ids(scheduler_output),
+        )
+
+    def get_scheduled_input_terminal_req_ids(
+        self,
+        scheduler_output: Any,
+    ) -> set[str]:
+        """Return terminal inputs executed by this immutable scheduler step."""
+        return self.input_terminal_req_ids.intersection(
+            self._scheduled_request_ids(scheduler_output),
+        )
+
+    @staticmethod
+    def _scheduled_request_ids(scheduler_output: Any) -> set[str]:
+        num_scheduled_tokens = getattr(
+            scheduler_output,
+            "num_scheduled_tokens",
+            None,
+        )
+        if num_scheduled_tokens is not None:
+            return set(num_scheduled_tokens)
+
+        request_ids = {req.req_id for req in getattr(scheduler_output, "scheduled_new_reqs", ())}
+        cached_reqs = getattr(scheduler_output, "scheduled_cached_reqs", None)
+        request_ids.update(getattr(cached_reqs, "req_ids", ()))
+        return request_ids
+
+    # ------------------------------------------------------------------ #
+    #  Internal helpers
+    # ------------------------------------------------------------------ #
+
+    def _process_chunk_queue(
+        self,
+        queue: Any,
+        waiting_for_chunk_list: deque[Any],
+        target_status: RequestStatus,
+        chunk_ready_req_ids: set[str],
+    ) -> None:
+        queue_snapshot = list(queue)
+        for request in queue_snapshot:
+            request_can_run = self._request_can_run(request, chunk_ready_req_ids)
+            if request.status != RequestStatus.WAITING_FOR_CHUNK:
+                if request_can_run:
+                    if request.request_id in chunk_ready_req_ids:
+                        self.requests_with_ready_chunks.add(request.request_id)
+                    continue
+                if request.status == RequestStatus.WAITING_FOR_INPUT:
+                    continue
+                self.pending_chunk_registrations.append(
+                    OmniChunkRecvHandle(
+                        request_id=request.request_id,
+                        external_req_id=getattr(request, "external_req_id", None),
+                    )
+                )
+                request.status = RequestStatus.WAITING_FOR_CHUNK
+                self._waiting_since.setdefault(request.request_id, time.monotonic())
+            else:
+                if request_can_run:
+                    request.status = target_status
+                    if request.request_id in chunk_ready_req_ids:
+                        self.requests_with_ready_chunks.add(request.request_id)
+                    self._waiting_since.pop(request.request_id, None)
+                    continue
+            if isinstance(queue, list):
+                queue.remove(request)
+            else:
+                queue.remove_request(request)
+            waiting_for_chunk_list.append(request)
+
+    def _request_can_run(
+        self,
+        request: Request,
+        chunk_ready_req_ids: set[str],
+    ) -> bool:
+        req_id = request.request_id
+        if req_id in self.finished_requests:
+            # Admit the final payload or terminal marker for cleanup.
+            return True
+
+        return req_id in chunk_ready_req_ids or req_id in self.requests_with_ready_chunks
+
+    def _clear_chunk_ready(self, scheduler_output: Any) -> None:
+        if scheduler_output.scheduled_new_reqs:
+            for req_data in scheduler_output.scheduled_new_reqs:
+                req_id = getattr(req_data, "req_id", None)
+                if req_id is not None:
+                    self.requests_with_ready_chunks.discard(req_id)
+
+        if scheduler_output.scheduled_cached_reqs:
+            for req_id in scheduler_output.scheduled_cached_reqs.req_ids:
+                self.requests_with_ready_chunks.discard(req_id)

vllm_omni/core/sched/output.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from dataclasses import dataclass, field, fields
 
 from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
@@ -18,14 +21,14 @@ class OmniNewRequestData(NewRequestData):
 
     Args:
         external_req_id: Optional external request ID for tracking
-        additional_information: Optional serialized additional information
+        additional_information: Optional serialized or materialized additional information
             dictionary containing tensors or lists
         model_intermediate_buffer: Optional runner-owned payload for
             GPUModelRunner.model_intermediate_buffer
     """
 
     external_req_id: str | None = None
-    additional_information: AdditionalInformationPayload | None = None
+    additional_information: AdditionalInformationPayload | dict[str, object] | None = None
     model_intermediate_buffer: dict[str, object] | None = None
 
     @classmethod
@@ -114,3 +117,5 @@ class OmniSchedulerOutput(SchedulerOutput):
 
     finished_requests_needing_kv_transfer: dict[str, dict] = field(default_factory=dict)
     pending_input_registrations: list[OmniChunkRecvHandle] = field(default_factory=list)
+    data_plane_terminal_req_ids: set[str] = field(default_factory=set)
+    input_terminal_req_ids: set[str] = field(default_factory=set)

vllm_omni/data_entry_keys.py

@@ -34,9 +34,14 @@ class HiddenStates(TypedDict, total=False):
 
 
 class Embeddings(TypedDict, total=False):
+    prepared_prefill: torch.Tensor
     prefill: torch.Tensor
     decode: torch.Tensor
+    decode_token_start: int
+    decode_token_end: int
     cached_decode: torch.Tensor
+    cached_decode_token_start: int
+    cached_decode_token_end: int
     tts_bos: torch.Tensor
     tts_eos: torch.Tensor
     tts_pad: torch.Tensor
@@ -54,6 +59,7 @@ class Codes(TypedDict, total=False):
 
 
 class Ids(TypedDict, total=False):
+    prepared_prefill: torch.Tensor
     all: list[int]
     prompt: list[int]
     output: list[int]
@@ -81,6 +87,7 @@ class OmniPayloadMeta(TypedDict, total=False):
     next_stage_prompt_len: int
     next_stage_generation_tokens: int
     replace_streaming_prompt: bool
+    next_stage_prompt_ids: list[int]
     streaming_prompt_recompute: bool
     streaming_condition_seq: int
     replace_runtime_additional_information: bool
@@ -94,6 +101,7 @@ class OmniPayloadMeta(TypedDict, total=False):
     width: int
     decode_flag: bool
     codec_streaming: bool
+    codec_frame_valid: bool | torch.Tensor
     ref_code_len: int
     ref_context_size: int
     ref_context_request_id: str
@@ -148,6 +156,8 @@ class EmbeddingsStruct(_StructBase):
     decode_token_start: int | None = None
     decode_token_end: int | None = None
     cached_decode: torch.Tensor | None = None
+    cached_decode_token_start: int | None = None
+    cached_decode_token_end: int | None = None
     tts_bos: torch.Tensor | None = None
     tts_eos: torch.Tensor | None = None
     tts_pad: torch.Tensor | None = None
@@ -190,6 +200,7 @@ class MetaStruct(_StructBase):
     next_stage_prompt_len: int | None = None
     next_stage_generation_tokens: int | None = None
     replace_streaming_prompt: bool | None = None
+    next_stage_prompt_ids: list[int] | None = None
     streaming_prompt_recompute: bool | None = None
     streaming_condition_seq: int | None = None
     replace_runtime_additional_information: bool | None = None
@@ -203,6 +214,7 @@ class MetaStruct(_StructBase):
     width: int | None = None
     decode_flag: bool | None = None
     codec_streaming: bool | None = None
+    codec_frame_valid: torch.Tensor | None = None
     ref_code_len: int | None = None
     # Expected FINAL length of a growing async-chunk sequence, when the
     # producer knows it up front (e.g. a frame-locked AR stage whose

vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2.yaml

@@ -0,0 +1,28 @@
+# Opt-in CUDA Model Runner V2 profile. The base profile retains V1.
+base_config: qwen3_tts_high_concurrency.yaml
+model_runner: v2
+
+platforms:
+  npu:
+    model_runner: v1
+  xpu:
+    model_runner: v1
+  rocm:
+    model_runner: v1
+  musa:
+    model_runner: v1
+
+connectors:
+  connector_of_shared_memory:
+    extra:
+      # Keep unqualified prefix-graph optimizations disabled in this profile.
+      code_predictor_prefix_graphs: false
+      ref_audio_artifact_cache_max_entries: 1024
+      # Opt-in B2 codec graphs. Validate quality for the deployment workload.
+      decode_batch_max_size: 2
+      decode_cudagraph_batch_sizes: [1, 2]
+
+stages:
+  - stage_id: 1
+    # Enable the decoder's inner graphs; the generation runner stays eager.
+    enforce_eager: false

vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_b4.yaml

@@ -0,0 +1,8 @@
+# Experimental B4 profile. Early first audio can require startup buffering.
+# MPS is a separate operator setting and is not enabled by this file.
+base_config: qwen3_tts_high_concurrency_mrv2.yaml
+connectors:
+  connector_of_shared_memory:
+    extra:
+      decode_batch_max_size: 4
+      decode_cudagraph_batch_sizes: [1, 2, 3, 4]

vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml

@@ -0,0 +1,23 @@
+# Opt-in single-GPU throughput profile, measured on H200 at concurrency 128
+# with Qwen3-TTS-12Hz-1.7B-Base. Warm the full pipeline before benchmarking.
+# B8 uses more graph memory than B4; this is not a low-latency profile.
+# MPS is optional and is not enabled by this file.
+#
+# vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-Base --omni \
+#   --deploy-config vllm_omni/deploy/qwen3_tts_high_concurrency_mrv2_single_gpu.yaml
+base_config: qwen3_tts_high_concurrency_mrv2.yaml
+
+connectors:
+  connector_of_shared_memory:
+    extra:
+      decode_batch_max_size: 8
+      decode_cudagraph_batch_sizes: [1, 2, 3, 4, 5, 6, 7, 8]
+
+stages:
+  - stage_id: 0
+    devices: "0"
+    max_num_seqs: 128
+  - stage_id: 1
+    devices: "0"
+    # Admission candidates are grouped by codec phase before B8 execution.
+    max_num_seqs: 32

vllm_omni/deploy/qwen3_tts_mrv2.yaml

@@ -0,0 +1,18 @@
+# Opt-in CUDA Model Runner V2 profile. The base profile retains V1.
+base_config: qwen3_tts.yaml
+model_runner: v2
+
+platforms:
+  npu:
+    model_runner: v1
+  xpu:
+    model_runner: v1
+  rocm:
+    model_runner: v1
+  musa:
+    model_runner: v1
+
+stages:
+  - stage_id: 0
+    # Bound per-step prefill work, matching the high-concurrency TTS profile.
+    max_num_batched_tokens: 512

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import fcntl
 import os

vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 import importlib
 import inspect
 from collections import deque
+from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, cast
 
 import torch
@@ -65,7 +66,6 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         extracted: dict[str, Any] = {}
         meta = payload.get("meta") if isinstance(payload, dict) else None
         meta = meta if isinstance(meta, dict) else {}
-
         if "next_stage_prompt_len" in meta:
             extracted["next_stage_prompt_len"] = meta["next_stage_prompt_len"]
         else:
@@ -85,6 +85,9 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         elif "left_context_size" in payload:
             logger.warning_once("legacy flat 'left_context_size' key in payload; expected 'meta.left_context_size'")
 
+        if cls._payload_finished(payload):
+            extracted["input_terminal"] = True
+
         return extracted
 
     _NON_CONSUMABLE_PAYLOAD_KEYS: set[tuple[str, str]] = {
@@ -331,67 +334,11 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         """
         if not hasattr(self, "_lock"):
             return OmniConnectorOutput()
-
-        tp_group = self._get_local_tp_group()
-        if self._async_chunk and tp_group is not None and getattr(tp_group, "world_size", 1) > 1:
-            if self.is_data_transfer_rank():
-                with self._lock:
-                    fanout_packet = self._collect_async_chunk_fanout_packet_locked()
-            else:
-                fanout_packet = None
-            fanout_packet = self._broadcast_tp_payload_packet(fanout_packet)
-            if fanout_packet is None:
-                newly_finished = set()
-                chunk_finished = set()
-                request_metadata = {}
-            else:
-                if not self.is_data_transfer_rank():
-                    self._apply_async_chunk_fanout_packet(fanout_packet)
-                newly_finished = set(fanout_packet["newly_finished"])
-                chunk_finished = set(fanout_packet["chunk_finished"])
-                request_metadata = dict(fanout_packet["request_metadata"])
-        else:
-            with self._lock:
-                newly_finished = set(self._finished_load_reqs)
-                self._finished_load_reqs.clear()
-                chunk_finished = set(self._chunk_finished_req_ids)
-                self._chunk_finished_req_ids.clear()
-                request_metadata = dict(self._local_request_metadata)
-                self._local_request_metadata.clear()
-                # _send_side_request_payload is the async accumulation buffer for
-                # future recv chunks. Clearing it on every consumable wake-up drops
-                # intermediate
-                # thinker decode spans before the model side can consume them.
-                # Only terminal chunk_finished requests may release that buffer.
-                for req_id in chunk_finished:
-                    if req_id not in self._local_stage_payload_cache:
-                        continue
-                    ext_req_id = self._request_ids_mapping.get(req_id, req_id)
-                    self._send_side_request_payload.pop(ext_req_id, None)
-                    if ext_req_id != req_id:
-                        self._send_side_request_payload.pop(req_id, None)
-        self._chunk_ready_req_ids.update(newly_finished)
-
-        output = OmniConnectorOutput(
-            chunk_ready_req_ids=set(self._chunk_ready_req_ids),
-            chunk_finished_req_ids=chunk_finished,
-            request_metadata=request_metadata,
-            kv_sent_req_ids=list(self._kv_sent_req_ids),
-            stage_recv_req_ids=set(self._stage_recv_req_ids),
-            has_pending_kv_work=self.has_pending_kv_work(),
-        )
-        if output.stage_recv_req_ids or chunk_finished or newly_finished:
-            logger.debug(
-                "[Stage-%s] get_omni_connector_output: stage_recv=%s, chunk_finished=%s, chunk_ready=%s",
-                self._stage_id,
-                output.stage_recv_req_ids,
-                chunk_finished,
-                output.chunk_ready_req_ids,
-            )
-        self._chunk_ready_req_ids.clear()
-        self._kv_sent_req_ids.clear()
-        self._stage_recv_req_ids.clear()
-        return output
+        drain_lock = getattr(self, "_omni_connector_output_drain_lock", None)
+        if drain_lock is None:
+            return self._drain_omni_connector_output()
+        with drain_lock:
+            return self._drain_omni_connector_output()
 
     @staticmethod
     def _connector_output_has_signals(output: OmniConnectorOutput) -> bool:
@@ -854,6 +801,9 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         self,
         request: Any,
         pooling_output: Any | None = None,
+        *,
+        propagate_errors: bool = False,
+        wait_for_delivery: bool = False,
     ) -> bool:
         """Derive and enqueue one chunk for async sending.
 
@@ -882,10 +832,11 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
             request_id=request_id,
             request=request,
             pooling_output=pooling_output,
+            propagate_errors=propagate_errors,
         )
         if payload_data is None:
             if chunk_id == 0:
-                logger.warning(
+                logger.debug(
                     "[Stage-%s] send_chunk: payload is None for req=%s chunk=%s (process_func=%s)",
                     self._stage_id,
                     request_id,
@@ -894,31 +845,14 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 )
             return False
 
-        self._put_req_chunk[request_id] += 1
-        self._ramp_chunk_count[request_id] += 1
-        next_stage_id = self._next_stage_id
-        connector_put_key = f"{request_id}_{self._stage_id}_{chunk_id}"
-
-        if chunk_id == 0:
-            logger.debug(
-                "[Stage-%s] send_chunk: first chunk enqueued, req=%s key=%s",
-                self._stage_id,
-                request_id,
-                connector_put_key,
-            )
-
-        task = {
-            "stage_id": self._stage_id,
-            "next_stage_id": next_stage_id,
-            "put_key": connector_put_key,
-            "data": payload_data,
-            "request_id": request_id,
-        }
-        with self._lock:
-            self._pending_save_reqs.setdefault(request_id, deque()).append(task)
-            self._pending_save_counts[request_id] += 1
-        self._work_available.set()
-        return True
+        enqueued, completion = self._enqueue_chunk_payload(
+            request,
+            payload_data,
+            wait_for_delivery=wait_for_delivery,
+        )
+        if enqueued and completion is not None:
+            completion.wait()
+        return enqueued
 
     # ------------------------------------------------------------------ #
     #  Background I/O threads
@@ -946,14 +880,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                     _recv_poll_count,
                 )
 
-            made_progress = False
-            for req_id in pending_ids:
-                if self._stop_event.is_set():
-                    break
-                try:
-                    made_progress = self._poll_single_request(req_id) or made_progress
-                except Exception:
-                    logger.warning("Error receiving data for %s", req_id, exc_info=True)
+            made_progress = self._poll_pending_requests_once(pending_ids)
 
             if not made_progress and not self._stop_event.is_set():
                 self._work_available.wait(timeout=0.005)
@@ -977,22 +904,29 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
 
             if task is not None:
                 success = False
+                send_error: BaseException | None = None
                 try:
                     success = self._send_single_request(task)
-                except Exception:
+                except Exception as error:
+                    send_error = error
                     logger.error(
                         "Error saving data for %s",
                         task.get("request_id"),
                         exc_info=True,
                     )
                 if not success:
-                    self._requeue_or_drop_failed_send(task)
+                    self._requeue_or_drop_failed_send(task, error=send_error)
                 continue
 
             self._work_available.wait(timeout=0.01)
             self._work_available.clear()
 
-    def _requeue_or_drop_failed_send(self, task: dict) -> None:
+    def _requeue_or_drop_failed_send(
+        self,
+        task: dict,
+        *,
+        error: BaseException | None = None,
+    ) -> None:
         """Re-enqueue a failed send task or drop it after max retries."""
         retry_count = task.get("_retry_count", 0) + 1
         req_id = task.get("request_id")
@@ -1002,6 +936,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 getattr(self, "_stage_id", "?"),
             )
             return
+        if error is not None:
+            task["_last_send_error"] = error
         if retry_count <= self._MAX_SEND_RETRIES:
             task["_retry_count"] = retry_count
             logger.warning(
@@ -1015,19 +951,36 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 dq = self._pending_save_reqs.setdefault(req_id, deque())
                 dq.appendleft(task)
         else:
+            cause = error or task.get("_last_send_error")
+            failure = RuntimeError(
+                "connector send failed permanently for "
+                f"request {req_id!r}, key {task.get('put_key')!r} "
+                f"after {self._MAX_SEND_RETRIES + 1} attempts"
+            )
+            if cause is not None:
+                failure.__cause__ = cause
             logger.error(
                 "[Stage-%s] Giving up on send for %s after %d retries",
                 getattr(self, "_stage_id", "?"),
                 req_id,
                 self._MAX_SEND_RETRIES,
             )
+            sink = getattr(self, "_connector_send_error_sink", None)
+            if sink is not None:
+                try:
+                    sink(failure)
+                except Exception:
+                    logger.exception("Connector send error sink failed for %s", req_id)
+            completion = task.get("completion")
+            if completion is not None:
+                completion.set_error(failure)
             self._decrement_pending_save_count(req_id)
 
     # ------------------------------------------------------------------ #
     #  Chunk-level poll / send  (ported from OmniChunkTransferAdapter)
     # ------------------------------------------------------------------ #
 
-    def _poll_single_request(self, req_id: str) -> bool:
+    def _poll_single_request(self, req_id: str, *, publish_ready: bool = True) -> bool:
         """Poll connector for one chunk of a request (non-blocking)."""
         connector = self._omni_connector
         if connector is None:
@@ -1099,16 +1052,20 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
             is_finished = self._payload_finished(payload_data)
             incoming_payload_consumable = self._payload_is_consumable(payload_data)
 
-            if self._model_mode == "ar":
-                payload_data = self._accumulate_payload(external_req_id, payload_data)
-                payload_consumable = incoming_payload_consumable
-            else:
-                new_ids = self._payload_audio_codes(payload_data) or []
-                if not new_ids and not is_finished:
+            if self._model_mode != "ar":
+                new_ids = self._payload_audio_codes(payload_data)
+                if not self._payload_value_has_content(new_ids) and not is_finished:
                     return False
                 payload_consumable = self._payload_is_consumable(payload_data)
 
             with self._lock:
+                if self._model_mode == "ar":
+                    # Accumulation, staging, and model-side consume/ack share
+                    # this lock. Keeping the transition atomic prevents the
+                    # model thread from clearing a just-merged decode span
+                    # before the recv thread publishes it to the local cache.
+                    payload_data = self._accumulate_payload(external_req_id, payload_data)
+                    payload_consumable = incoming_payload_consumable
                 if is_finished:
                     self._chunk_finished_req_ids.add(req_id)
                     self._chunk_stream_completed.add(req_id)
@@ -1167,6 +1124,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 type(engine_inputs).__name__,
             )
 
+        if publish_ready and self._async_chunk and (payload_consumable or is_finished):
+            self.publish_omni_connector_output_to_sink()
         logger.debug("[Stage-%s] Received data for key %s", self._stage_id, connector_get_key)
         return True
 
@@ -1175,14 +1134,27 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         request_id: str | None,
         request: Any | None,
         pooling_output: Any | None,
+        *,
+        propagate_errors: bool = False,
     ) -> Any | None:
         """Run the custom process hook with a best-effort finished kwarg."""
         if self._custom_process_func is None:
             return None
 
+        payload_kwarg = getattr(self, "_custom_process_payload_kwarg", None)
+        if payload_kwarg is None:
+            payload_kwarg = self._connector_payload_kwarg(self._custom_process_func)
+            self._custom_process_payload_kwarg = payload_kwarg
+        if payload_kwarg is None:
+            message = f"custom process hook {self._custom_process_func} has no connector payload argument"
+            logger.error(message)
+            if propagate_errors:
+                raise RuntimeError(message)
+            return None
+
         kwargs = {
             "transfer_manager": self,
-            "pooling_output": pooling_output,
+            payload_kwarg: pooling_output,
             "request": request,
         }
         supports_is_finished = getattr(
@@ -1197,21 +1169,29 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                     kwargs["is_finished"] = bool(is_finished_fn())
             except Exception:
                 logger.debug("request.is_finished() failed for %s", request_id, exc_info=True)
+                if propagate_errors:
+                    raise
 
         try:
             return self._custom_process_func(**kwargs)
         except TypeError as exc:
             if "is_finished" not in kwargs or not self._is_unexpected_is_finished_kwarg_error(exc):
                 logger.exception("custom_process_stage_input_func failed for chunk %s", request_id)
+                if propagate_errors:
+                    raise
                 return None
             kwargs.pop("is_finished", None)
             try:
                 return self._custom_process_func(**kwargs)
             except Exception:
                 logger.exception("custom_process_stage_input_func failed for chunk %s", request_id)
+                if propagate_errors:
+                    raise
                 return None
         except Exception:
             logger.exception("custom_process_stage_input_func failed for chunk %s", request_id)
+            if propagate_errors:
+                raise
             return None
 
     def _custom_process_supports_is_finished_kwarg(self) -> bool | None:
@@ -1247,15 +1227,17 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
     def _send_single_request(self, task: dict) -> bool:
         """Send one queued task via connector.put().
 
-        Returns True on success.  On failure (put() raises or returns
-        ``success=False``), returns False **without** decrementing
-        ``_pending_save_counts`` so the caller can retry or clean up.
+        Returns True when the task is consumed, either by a successful send or
+        because its delivery completion was already terminal.  On transport
+        failure (put() raises or returns ``success=False``), returns False
+        **without** decrementing ``_pending_save_counts`` so the caller can
+        retry or clean up.
         """
         connector = self._omni_connector
         if connector is None:
             return True
 
-        request_id = task.get("request_id")
+        request_id = task["request_id"]
         payload_data = task.get("data")
         if payload_data is None and task.get("request") is not None:
             payload_data = self._build_custom_process_payload(
@@ -1264,6 +1246,15 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 pooling_output=task.get("pooling_output"),
             )
         put_key = task.get("put_key")
+        completion = task.get("completion")
+        if completion is not None and not completion.mark_in_flight():
+            logger.debug(
+                "[Stage-%s] Skipping terminal connector delivery: put_key=%s",
+                task["stage_id"],
+                put_key,
+            )
+            self._decrement_pending_save_count(request_id)
+            return True
 
         success, _size, _metadata = connector.put(
             from_stage=str(task["stage_id"]),
@@ -1282,8 +1273,9 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         if not success:
             return False
 
-        if isinstance(request_id, str):
-            self._decrement_pending_save_count(request_id)
+        self._decrement_pending_save_count(request_id)
+        if completion is not None:
+            completion.set_result()
         return True
 
     def _decrement_pending_save_count(self, request_id: str) -> None:
@@ -1302,6 +1294,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 self._put_req_chunk.pop(cleanup_req_id, None)
                 self._send_side_request_payload.pop(cleanup_req_id, None)
                 self._code_prompt_token_ids.pop(cleanup_req_id, None)
+                getattr(self, "_qwen3_tts_emitted_frames", {}).pop(cleanup_req_id, None)
                 self._cached_ic.pop(cleanup_req_id, None)
                 self._ramp_chunk_count.pop(cleanup_req_id, None)
                 self._adaptive_states.pop(cleanup_req_id, None)
@@ -1354,3 +1347,238 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
 
         self._send_side_request_payload[req_id] = merged
         return cast(OmniPayload, dict(merged))
+
+    def _drain_omni_connector_output(self) -> OmniConnectorOutput:
+        tp_group = self._get_local_tp_group()
+        if self._async_chunk and tp_group is not None and getattr(tp_group, "world_size", 1) > 1:
+            if self.is_data_transfer_rank():
+                with self._lock:
+                    fanout_packet = self._collect_async_chunk_fanout_packet_locked()
+            else:
+                fanout_packet = None
+            fanout_packet = self._broadcast_tp_payload_packet(fanout_packet)
+            if fanout_packet is None:
+                newly_finished = set()
+                chunk_finished = set()
+                request_metadata = {}
+            else:
+                if not self.is_data_transfer_rank():
+                    self._apply_async_chunk_fanout_packet(fanout_packet)
+                newly_finished = set(fanout_packet["newly_finished"])
+                chunk_finished = set(fanout_packet["chunk_finished"])
+                request_metadata = dict(fanout_packet["request_metadata"])
+        else:
+            with self._lock:
+                newly_finished = set(self._finished_load_reqs)
+                self._finished_load_reqs.clear()
+                chunk_finished = set(self._chunk_finished_req_ids)
+                self._chunk_finished_req_ids.clear()
+                request_metadata = dict(self._local_request_metadata)
+                self._local_request_metadata.clear()
+                # _send_side_request_payload is the async accumulation buffer for
+                # future recv chunks. Clearing it on every consumable wake-up drops
+                # intermediate
+                # thinker decode spans before the model side can consume them.
+                # Only terminal chunk_finished requests may release that buffer.
+                for req_id in chunk_finished:
+                    if req_id not in self._local_stage_payload_cache:
+                        continue
+                    ext_req_id = self._request_ids_mapping.get(req_id, req_id)
+                    self._send_side_request_payload.pop(ext_req_id, None)
+                    if ext_req_id != req_id:
+                        self._send_side_request_payload.pop(req_id, None)
+        self._chunk_ready_req_ids.update(newly_finished)
+
+        output = OmniConnectorOutput(
+            chunk_ready_req_ids=set(self._chunk_ready_req_ids),
+            chunk_finished_req_ids=chunk_finished,
+            request_metadata=request_metadata,
+            kv_sent_req_ids=list(self._kv_sent_req_ids),
+            stage_recv_req_ids=set(self._stage_recv_req_ids),
+            has_pending_kv_work=self.has_pending_kv_work(),
+        )
+        if output.stage_recv_req_ids or chunk_finished or newly_finished:
+            logger.debug(
+                "[Stage-%s] get_omni_connector_output: stage_recv=%s, chunk_finished=%s, chunk_ready=%s",
+                self._stage_id,
+                output.stage_recv_req_ids,
+                chunk_finished,
+                output.chunk_ready_req_ids,
+            )
+        self._chunk_ready_req_ids.clear()
+        self._kv_sent_req_ids.clear()
+        self._stage_recv_req_ids.clear()
+        return output
+
+    def _enqueue_chunk_payload(
+        self,
+        request: Any,
+        payload_data: Any,
+        *,
+        wait_for_delivery: bool = False,
+    ) -> tuple[bool, Any | None]:
+        """Enqueue an already-built payload under its per-request chunk key."""
+        raw_req_id = getattr(request, "request_id", None) or getattr(request, "req_id", None)
+        request_id = self._resolve_external_req_id(request, raw_req_id)
+        if raw_req_id and raw_req_id != request_id:
+            self._request_ids_mapping.setdefault(raw_req_id, request_id)
+        chunk_id = self._put_req_chunk[request_id]
+        self._put_req_chunk[request_id] += 1
+        self._ramp_chunk_count[request_id] += 1
+        next_stage_id = self._next_stage_id
+        connector_put_key = f"{request_id}_{self._stage_id}_{chunk_id}"
+        completion = (
+            self._create_send_completion(
+                request_id=request_id,
+                put_key=connector_put_key,
+            )
+            if wait_for_delivery
+            else None
+        )
+
+        if chunk_id == 0:
+            logger.debug(
+                "[Stage-%s] send_chunk: first chunk enqueued, req=%s key=%s",
+                self._stage_id,
+                request_id,
+                connector_put_key,
+            )
+
+        task = {
+            "stage_id": self._stage_id,
+            "next_stage_id": next_stage_id,
+            "put_key": connector_put_key,
+            "data": payload_data,
+            "request_id": request_id,
+            "completion": completion,
+        }
+        with self._lock:
+            self._pending_save_reqs.setdefault(request_id, deque()).append(task)
+            self._pending_save_counts[request_id] += 1
+        self._work_available.set()
+        return True, completion
+
+    def _poll_pending_requests_once(self, pending_ids: list[str]) -> bool:
+        """Poll one receiver pass and publish its ready requests as a cohort."""
+        made_progress = False
+        for req_id in pending_ids:
+            if self._stop_event.is_set():
+                break
+            try:
+                made_progress = self._poll_single_request(req_id, publish_ready=False) or made_progress
+            except Exception:
+                logger.warning("Error receiving data for %s", req_id, exc_info=True)
+        if made_progress and self._async_chunk:
+            self.publish_omni_connector_output_to_sink()
+        return made_progress
+
+    def _publish_chunk_cohort(self, entries: list[tuple[Any, Any]], *, wait_for_delivery: bool) -> int:
+        entries = [(request, payload) for request, payload in entries if payload is not None]
+        if not entries:
+            return 0
+        emitted = 0
+        completions: list[Any] = []
+        for request, payload in entries:
+            enqueued, completion = self._enqueue_chunk_payload(request, payload, wait_for_delivery=wait_for_delivery)
+            if enqueued:
+                emitted += 1
+                if completion is not None:
+                    completions.append(completion)
+        for completion in completions:
+            completion.wait()
+        return emitted
+
+    def publish_omni_connector_output_to_sink(self) -> bool:
+        """Drain pending readiness into the same-process scheduler inbox."""
+        sink = getattr(self, "_omni_connector_output_sink", None)
+        if sink is None:
+            return False
+        output = self.get_omni_connector_output()
+        if not self._connector_output_has_signals(output):
+            return False
+        sink(output)
+        return True
+
+    def send_chunks(
+        self,
+        entries: list[tuple[Any, Any | None]],
+        *,
+        propagate_errors: bool = False,
+        wait_for_delivery: bool = False,
+    ) -> int:
+        """Build and enqueue all request payloads from one model step.
+
+        A model-specific batch builder may coalesce GPU post-processing and
+        D2H copies across requests. Models without one retain the scalar
+        builder contract.
+        """
+        if not entries or self._omni_connector is None:
+            return 0
+        if not self.is_data_transfer_rank():
+            return len(entries)
+
+        batch_func = getattr(self, "_custom_process_batch_func", None)
+        if batch_func is None:
+            payloads = []
+            for request, pooling_output in entries:
+                raw_req_id = getattr(request, "request_id", None) or getattr(request, "req_id", None)
+                request_id = self._resolve_external_req_id(request, raw_req_id)
+                payloads.append(
+                    self._build_custom_process_payload(
+                        request_id=request_id,
+                        request=request,
+                        pooling_output=pooling_output,
+                        propagate_errors=propagate_errors,
+                    )
+                )
+            return self._publish_chunk_cohort(
+                [(request, payload) for (request, _), payload in zip(entries, payloads, strict=True)],
+                wait_for_delivery=wait_for_delivery,
+            )
+
+        requests = [request for request, _ in entries]
+        pooling_outputs = [pooling_output for _, pooling_output in entries]
+        is_finished: list[bool] = []
+        for request in requests:
+            finished_fn = getattr(request, "is_finished", None)
+            is_finished.append(bool(finished_fn()) if callable(finished_fn) else False)
+
+        try:
+            payloads = batch_func(
+                transfer_manager=self,
+                pooling_outputs=pooling_outputs,
+                requests=requests,
+                is_finished=is_finished,
+            )
+        except Exception:
+            logger.exception(
+                "batch custom_process_stage_input_func failed for %s requests",
+                len(entries),
+            )
+            if propagate_errors:
+                raise
+            return 0
+        if not isinstance(payloads, list) or len(payloads) != len(entries):
+            message = (
+                "batch custom process hook "
+                f"{batch_func} returned "
+                f"{len(payloads) if isinstance(payloads, list) else type(payloads).__name__} "
+                f"payloads for {len(entries)} requests"
+            )
+            logger.error(
+                "batch custom process hook %s returned %s payloads for %s requests",
+                batch_func,
+                len(payloads) if isinstance(payloads, list) else type(payloads).__name__,
+                len(entries),
+            )
+            if propagate_errors:
+                raise RuntimeError(message)
+            return 0
+
+        return self._publish_chunk_cohort(list(zip(requests, payloads)), wait_for_delivery=wait_for_delivery)
+
+    def set_omni_connector_output_sink(
+        self,
+        sink: Callable[[OmniConnectorOutput], None] | None,
+    ) -> None:
+        self._omni_connector_output_sink = sink

vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py

@@ -9,6 +9,7 @@ import inspect
 import os
 import threading
 from collections import defaultdict, deque
+from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
 import torch
@@ -21,6 +22,7 @@ from vllm_omni.distributed.omni_connectors.utils.config import (
 )
 from vllm_omni.distributed.omni_connectors.utils.initialization import resolve_connector_spec
 from vllm_omni.distributed.omni_connectors.utils.kv_utils import get_local_tp_rank, get_omni_replica_id
+from vllm_omni.outputs import OmniConnectorOutput
 
 logger = init_logger("vllm_omni.worker.omni_connector_model_runner_mixin")
 
@@ -69,14 +71,35 @@ def should_accumulate_full_payload_output(model_config, custom_process_func) ->
         return False
     if getattr(model_config, "async_chunk", False):
         return False
-    if getattr(model_config, "final_output", False):
-        return False
     next_stage_func = getattr(model_config, "custom_process_next_stage_input_func", None)
     if not isinstance(next_stage_func, str) or not next_stage_func:
         return False
     return getattr(model_config, "model_stage", None) is not None
 
 
+class _SendCompletion:
+    """One connector task's delivery acknowledgement."""
+
+    def __init__(self) -> None:
+        self._event = threading.Event()
+        self._error: BaseException | None = None
+
+    def mark_in_flight(self) -> bool:
+        return True
+
+    def set_result(self) -> None:
+        self._event.set()
+
+    def set_error(self, error: BaseException) -> None:
+        self._error = error
+        self._event.set()
+
+    def wait(self) -> None:
+        self._event.wait()
+        if self._error is not None:
+            raise self._error
+
+
 class _OmniConnectorRuntimeMixin:
     """Own connector lifecycle, shared state, and KV transfer delegation."""
 
@@ -166,6 +189,8 @@ class _OmniConnectorRuntimeMixin:
         self._stage_id: int = stage_id if isinstance(stage_id, int) else 0
 
         self._custom_process_func_path, self._custom_process_func = self._load_custom_func(model_config)
+        self._custom_process_batch_func = self._load_custom_batch_func(self._custom_process_func)
+        self._custom_process_payload_kwarg = self._connector_payload_kwarg(self._custom_process_func)
         self._custom_process_supports_is_finished = self._custom_process_supports_is_finished_kwarg()
         logger.debug(
             "[Stage-%s] init_omni_connectors: async_chunk=%s, custom_process_func=%s, connector=%s, func_path=%s",
@@ -229,6 +254,7 @@ class _OmniConnectorRuntimeMixin:
         self._pending_save_reqs: dict[str, deque] = {}
         self._pending_save_counts: dict[str, int] = defaultdict(int)
         self._deferred_send_cleanup: set[str] = set()
+        self._connector_send_error_sink: Callable[[BaseException], None] | None = None
         # -- per-cycle output accumulator --
         self._chunk_ready_req_ids: set[str] = set()
         self._chunk_finished_req_ids: set[str] = set()
@@ -247,6 +273,9 @@ class _OmniConnectorRuntimeMixin:
         self._local_stage_payload_cache: dict[str, dict[str, Any]] = {}
         # Lightweight scheduling metadata pending delivery to the Scheduler.
         self._local_request_metadata: dict[str, dict[str, Any]] = {}
+        # Optional same-process control-plane fast path. Payload tensors remain
+        # runner-owned; only OmniConnectorOutput readiness is published.
+        self._omni_connector_output_sink: Callable[[OmniConnectorOutput], None] | None = None
 
         # -- persistent set of request IDs whose chunk stream is complete --
         # Prevents re-registration after the finish sentinel has been received.
@@ -270,6 +299,7 @@ class _OmniConnectorRuntimeMixin:
         self._kv_triggered_requests: set[str] = set()
 
         self._lock = threading.Lock()
+        self._omni_connector_output_drain_lock = threading.Lock()
         self._stop_event = threading.Event()
         self._work_available = threading.Event()
 
@@ -364,6 +394,7 @@ class _OmniConnectorRuntimeMixin:
                 self._put_req_chunk.pop(k, None)
                 self._send_side_request_payload.pop(k, None)
                 self._code_prompt_token_ids.pop(k, None)
+                getattr(self, "_qwen3_tts_emitted_frames", {}).pop(k, None)
                 self._cached_ic.pop(k, None)
                 self._ramp_chunk_count.pop(k, None)
                 self._adaptive_states.pop(k, None)
@@ -384,6 +415,10 @@ class _OmniConnectorRuntimeMixin:
         self._cleanup_recv_delivery_state(req_id)
 
     def _drop_send_side_payload_state(self, req_id: str, ext_id: str | None) -> None:
+        watermark = getattr(self, "_qwen3_tts_emitted_frames", {})
+        watermark.pop(req_id, None)
+        if ext_id is not None:
+            watermark.pop(ext_id, None)
         if ext_id is not None:
             self._send_side_request_payload.pop(ext_id, None)
             self._cached_ic.pop(ext_id, None)
@@ -594,31 +629,12 @@ class _OmniConnectorRuntimeMixin:
 
         return None, None
 
-    @staticmethod
-    def _is_connector_payload_builder(func: Any) -> bool:
+    @classmethod
+    def _is_connector_payload_builder(cls, func: Any) -> bool:
         """Whether *func* matches the mixin payload-builder contract."""
-        try:
-            signature = inspect.signature(func)
-        except (TypeError, ValueError):
-            return False
+        return cls._connector_payload_kwarg(func) is not None
 
-        params = signature.parameters
-        if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()):
-            return True
-
-        required = {"transfer_manager", "pooling_output", "request"}
-        supported = {
-            name
-            for name, param in params.items()
-            if param.kind
-            in (
-                inspect.Parameter.POSITIONAL_OR_KEYWORD,
-                inspect.Parameter.KEYWORD_ONLY,
-            )
-        }
-        return required.issubset(supported)
-
-    def _resolve_external_req_id(self, request: Any, fallback_req_id: str) -> str:
+    def _resolve_external_req_id(self, request: Any, fallback_req_id: str | None) -> str:
         """Resolve the external request ID consistently.
 
         Checks ``_request_ids_mapping`` first (populated by
@@ -626,7 +642,7 @@ class _OmniConnectorRuntimeMixin:
         ``external_req_id`` attribute, and finally to the given
         ``fallback_req_id``.
         """
-        mapped = self._request_ids_mapping.get(fallback_req_id)
+        mapped = self._request_ids_mapping.get(fallback_req_id) if fallback_req_id is not None else None
         if mapped is not None:
             return mapped
         if request is not None:
@@ -634,6 +650,8 @@ class _OmniConnectorRuntimeMixin:
             ext = getattr(request, "external_req_id", None)
             if ext is not None:
                 return ext
+        if fallback_req_id is None:
+            raise ValueError("Connector request has neither an external nor an internal request ID")
         return fallback_req_id
 
     def _resolve_next_stage_id(self, model_config: Any) -> int:
@@ -1073,3 +1091,67 @@ class _OmniConnectorRuntimeMixin:
     ) -> str:
         """Build connector key that includes rank info for KV transfers."""
         return f"{req_id}_{from_stage}_{chunk_id}_{from_rank}_{to_rank}"
+
+    @staticmethod
+    def _connector_payload_kwarg(func: Any) -> str | None:
+        """Resolve the model output argument used by a connector builder.
+
+        Full-payload builders historically name it ``pooling_output`` while
+        async-chunk builders name it ``multimodal_output``. Both represent the
+        runner-owned per-step inter-stage payload and share the same transport
+        contract.
+        """
+        if func is None:
+            return None
+        try:
+            signature = inspect.signature(func)
+        except (TypeError, ValueError):
+            return None
+
+        params = signature.parameters
+        if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in params.values()):
+            return "pooling_output"
+
+        supported = {
+            name
+            for name, param in params.items()
+            if param.kind
+            in (
+                inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                inspect.Parameter.KEYWORD_ONLY,
+            )
+        }
+        if not {"transfer_manager", "request"}.issubset(supported):
+            return None
+        if "pooling_output" in supported:
+            return "pooling_output"
+        if "multimodal_output" in supported:
+            return "multimodal_output"
+        return None
+
+    @staticmethod
+    def _load_custom_batch_func(custom_process_func: Any | None) -> Any | None:
+        """Resolve an optional ``<scalar_builder>_batch`` peer."""
+        if custom_process_func is None:
+            return None
+        module_name = getattr(custom_process_func, "__module__", None)
+        func_name = getattr(custom_process_func, "__name__", None)
+        if not module_name or not func_name:
+            return None
+        try:
+            module = importlib.import_module(module_name)
+            batch_func = getattr(module, f"{func_name}_batch", None)
+        except ImportError:
+            logger.debug(
+                "Could not import batch connector payload builder for %s.%s",
+                module_name,
+                func_name,
+                exc_info=True,
+            )
+            return None
+        return batch_func if callable(batch_func) else None
+
+    def _create_send_completion(self, *, request_id: str, put_key: str) -> Any:
+        """Create a delivery completion; MRv2 overrides this with a ticket."""
+        del request_id, put_key
+        return _SendCompletion()

vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py

@@ -273,6 +273,8 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
                 frozen_ids = token_ids.copy()
                 setattr(snapshot, private_name, frozen_ids)
                 setattr(snapshot, public_name, ConstantList(frozen_ids))
+                if public_name == "output_token_ids":
+                    snapshot.output_token_count = len(frozen_ids)
         return snapshot
 
     @staticmethod
@@ -878,6 +880,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
 
         if is_segment_finished:
             self.code_prompt_token_ids.pop(external_req_id, None)
+            getattr(self, "_qwen3_tts_emitted_frames", {}).pop(external_req_id, None)
             self.ramp_chunk_count.pop(external_req_id, None)
             self._adaptive_states.pop(external_req_id, None)
             cached_ic = getattr(self, "_cached_ic", None)
@@ -992,6 +995,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase):
         self.put_req_chunk.pop(external_req_id, None)
         self.request_payload.pop(external_req_id, None)
         self.code_prompt_token_ids.pop(external_req_id, None)
+        getattr(self, "_qwen3_tts_emitted_frames", {}).pop(external_req_id, None)
         self.requests_num_chunks_sent.pop(external_req_id, None)
         self._segment_generation.pop(external_req_id, None)
         self.ramp_chunk_count.pop(external_req_id, None)

vllm_omni/engine/__init__.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 Engine components for vLLM-Omni.
 """
@@ -134,10 +137,19 @@ class OmniEngineCoreOutput(EngineCoreOutput):
     # pooling_output is inherited from EngineCoreOutput as torch.Tensor | None
     # and retains its original vLLM semantics for pooling/embedding tasks.
     multimodal_output: dict[str, torch.Tensor] | None = None
+    # Bytes-serialized carrier for a dict-shaped pooling_output (MR V2 cross-stage
+    # hidden/codec handoff). vLLM strictly decodes ``pooling_output`` as a
+    # ``torch.Tensor``, so a per-request dict payload cannot ride that field. The
+    # producer (OmniARScheduler) sets ``pooling_output=None`` and serializes the
+    # dict here; the consumer (StagePool._poll_stage_raw) rehydrates it back into
+    # ``pooling_output`` after decode. Tensor pooling outputs are unaffected.
+    pooling_output_payload: AdditionalInformationPayload | None = None
     # Finished flag for streaming input segment
     is_segment_finished: bool | None = False
     # Streaming update prompt length
     new_prompt_len_snapshot: int | None = None
+    # Authoritative segment count when the native plane suppresses token IPC.
+    num_generation_tokens: int | None = None
 
 
 class OmniEngineCoreOutputs(EngineCoreOutputs):

vllm_omni/engine/arg_utils.py

@@ -190,6 +190,7 @@ class OmniEngineArgs(EngineArgs):
     model_stage: str = "thinker"
     model_arch: str | None = None
     engine_output_type: str | None = None
+    final_output: bool = False
     hf_config_name: str | None = None
     custom_process_next_stage_input_func: str | None = None
     requires_full_payload_input: bool = False
@@ -199,6 +200,8 @@ class OmniEngineArgs(EngineArgs):
     async_chunk: bool = False
     session_mode: str = "turn"
     retains_state_across_chunks: bool = False
+    use_v2_model_runner: bool = False
+    supports_native_mrv2_data_plane: bool = False
     # WS-A: Stage-1 active stream slots. 0 = legacy preempt-everything.
     # Must be declared here so engine_args dict propagation does not silently
     # drop the value when constructing OmniEngineArgs from kwargs.
@@ -423,6 +426,8 @@ class OmniEngineArgs(EngineArgs):
             async_chunk=self.async_chunk,
             session_mode=self.session_mode,
             retains_state_across_chunks=self.retains_state_across_chunks,
+            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,
             duplex_max_sessions=self.duplex_max_sessions,
             model_stage=self.model_stage,
@@ -430,6 +435,7 @@ class OmniEngineArgs(EngineArgs):
             worker_type=self.worker_type,
             pooling_output_decoder=self.pooling_output_decoder,
             engine_output_type=self.engine_output_type,
+            final_output=self.final_output,
             hf_config_name=self.hf_config_name,
             custom_process_next_stage_input_func=self.custom_process_next_stage_input_func,
             requires_full_payload_input=self.requires_full_payload_input,

vllm_omni/engine/orchestrator.py

@@ -192,6 +192,8 @@ class OrchestratorRequestState:
     final_stage_id: int = -1
     final_output_stage_ids: set[int] = field(default_factory=set)
     finished_final_output_stage_ids: set[int] = field(default_factory=set)
+    finished_stage_ids: set[int] = field(default_factory=set)
+    pending_final_output: OutputMessage | None = None
 
     # Wall-clock timestamp when the client-facing engine request was accepted.
     request_timestamp: float = 0.0
@@ -763,20 +765,21 @@ class OrchestratorBase:
                     float(kv_wait_s),
                 )
             req_state = self.request_states.get(getattr(eco, "request_id", None))
-            if req_state is None or not req_state.streaming.enabled:
+            if req_state is None:
                 continue
-            segment_finished = bool(getattr(eco, "is_segment_finished", False))
-            raw_mm = self._completion_multimodal_output(eco, None)
-            req_state.streaming.segments[stage_id] = StreamingSegmentState(
-                finished=segment_finished,
-                token_ids=(self._coerce_int_list(getattr(eco, "new_token_ids", None)) if segment_finished else []),
-                output_metadata=(dict(raw_mm) if segment_finished and isinstance(raw_mm, dict) else {}),
-            )
-            req_state.streaming.new_prompt_len_snapshot = getattr(
-                eco,
-                "new_prompt_len_snapshot",
-                None,
-            )
+            if req_state.streaming.enabled:
+                segment_finished = bool(getattr(eco, "is_segment_finished", False))
+                raw_mm = self._completion_multimodal_output(eco, None)
+                req_state.streaming.segments[stage_id] = StreamingSegmentState(
+                    finished=segment_finished,
+                    token_ids=(self._coerce_int_list(getattr(eco, "new_token_ids", None)) if segment_finished else []),
+                    output_metadata=(dict(raw_mm) if segment_finished and isinstance(raw_mm, dict) else {}),
+                )
+                req_state.streaming.new_prompt_len_snapshot = getattr(
+                    eco,
+                    "new_prompt_len_snapshot",
+                    None,
+                )
             if await self._apply_raw_terminal_stage_finish(stage_id, eco, req_state):
                 raw_terminal_request_ids.add(req_state.request_id)
             await self._report_duplex_session_request_error(stage_id, replica_id, eco, req_state)
@@ -1506,18 +1509,19 @@ class OrchestratorBase:
         ``is_segment_finished=False``, but vLLM's output processor may remove the
         request state before that EngineCoreOutput is processed.
 
-        Only update ``finished_final_output_stage_ids`` here. Request cleanup stays
-        in ``_route_output`` so downstream async-chunk stages can still deliver
-        outputs after stage-0 session end.
+        Record raw stage completion here. Client completion is
+        resolved after this raw batch passes through the output processor, so a
+        real processed output wins over the swallowed-terminal fallback.
         """
         if getattr(eco, "finish_reason", None) is None:
             return False
         if getattr(eco, "is_segment_finished", False):
             return False
 
+        req_state.finished_stage_ids.add(stage_id)
         final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id}
         if stage_id not in final_output_stage_ids:
-            return False
+            return True
         req_state.finished_final_output_stage_ids.add(stage_id)
         return True
 
@@ -1570,16 +1574,23 @@ class OrchestratorBase:
         replica_id: int,
         request_ids: set[str],
     ) -> None:
-        """Finish streaming requests whose raw terminal had no processed output."""
+        """Finish requests whose raw terminal had no processed output."""
         pool = self.stage_pools[stage_id]
-        if not pool.final_output:
-            return
-
         for request_id in request_ids:
             req_state = self.request_states.get(request_id)
             if req_state is None or req_state.session_owned:
                 continue
-
+            pending = req_state.pending_final_output
+            if pending is not None:
+                if set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids):
+                    req_state.pending_final_output = None
+                    await self.output_async_queue.put(pending)
+                    await self._cleanup_request_ids([request_id, *self._cfg_tracker.cleanup_parent(request_id)])
+                # A real terminal output is pending; do not replace it with
+                # the swallowed-output fallback before upstream completion.
+                continue
+            if not pool.final_output:
+                continue
             final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id}
             if not final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids):
                 continue
@@ -1653,27 +1664,41 @@ class OrchestratorBase:
             await self._cleanup_request_ids([req_id])
             return
 
+        if finished and not segment_finished:
+            req_state.finished_stage_ids.add(stage_id)
+
         request_finished = False
         if finished and self.stage_pools[stage_id].final_output and not segment_finished:
             req_state.finished_final_output_stage_ids.add(stage_id)
             final_output_stage_ids = req_state.final_output_stage_ids or {req_state.final_stage_id}
             request_finished = final_output_stage_ids.issubset(req_state.finished_final_output_stage_ids)
         if req_state.session_owned:
-            # Session-owned outputs are delivered to their session through
-            # ``_intercept_stage_output`` below, never to a request queue.
+            # Session-owned outputs are delivered through the session interceptor.
             pass
         elif self.stage_pools[stage_id].final_output:
-            await self.output_async_queue.put(
-                OutputMessage(
-                    request_id=req_id,
-                    stage_id=stage_id,
-                    replica_id=replica_id,
-                    engine_outputs=output,
-                    metrics=stage_metrics,
-                    finished=request_finished,
-                    stage_submit_ts=submit_ts,
-                )
+            message = OutputMessage(
+                request_id=req_id,
+                stage_id=stage_id,
+                replica_id=replica_id,
+                engine_outputs=output,
+                metrics=stage_metrics,
+                finished=request_finished,
+                stage_submit_ts=submit_ts,
             )
+            if (
+                request_finished
+                and self.async_chunk
+                and not req_state.streaming.enabled
+                and not set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids)
+            ):
+                # Direct chunk delivery can finish the decoder before the
+                # upstream engine's terminal output/metrics reaches this loop.
+                # Keep request state until those messages are routed, otherwise
+                # the frontend publishes incomplete usage and loses stop reasons.
+                req_state.pending_final_output = message
+                request_finished = False
+            else:
+                await self.output_async_queue.put(message)
         elif stage_metrics is not None:
             await self.output_async_queue.put(
                 StageMetricsMessage(
@@ -1741,6 +1766,12 @@ class OrchestratorBase:
                         is_final_update=True,
                     )
 
+        pending = req_state.pending_final_output
+        if pending is not None and set(req_state.stage_submit_ts).issubset(req_state.finished_stage_ids):
+            req_state.pending_final_output = None
+            await self.output_async_queue.put(pending)
+            request_finished = True
+
         if request_finished and not req_state.session_owned:
             await self._cleanup_request_ids([req_id, *self._cfg_tracker.cleanup_parent(req_id)])
 
@@ -2483,8 +2514,6 @@ class OrchestratorBase:
                     operation="async-chunk prewarm",
                 )
             else:
-                import copy
-
                 from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ids_length
 
                 try:
@@ -2507,7 +2536,9 @@ class OrchestratorBase:
 
                 original_prompt = req_state.prompt
                 if isinstance(original_prompt, dict):
-                    base_input = copy.deepcopy(original_prompt)
+                    from vllm_omni.engine.request_snapshot import copy_request_snapshot
+
+                    base_input = copy_request_snapshot(original_prompt)
                 else:
                     base_input = {}
 

vllm_omni/engine/request_snapshot.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Full request snapshots with a fast path for immutable scalar leaves."""
+
+import copy
+from typing import Any
+
+_ATOMIC_TYPES = frozenset((type(None), bool, int, float, complex, str, bytes))
+
+
+def copy_request_snapshot(value: Any, memo: dict[int, Any] | None = None) -> Any:
+    """Deep-copy all fields, preserving aliases and isolating mutable containers.
+
+    Plain numeric/text trees do not need generic object dispatch or an identity
+    lookup for each immutable scalar. Subclasses and other objects retain their
+    regular deepcopy contracts. No model fields are filtered or shared here.
+    """
+    kind = type(value)
+    if kind in _ATOMIC_TYPES:
+        return value
+    if memo is None:
+        memo = {}
+    key = id(value)
+    if key in memo:
+        return memo[key]
+    if kind is list:
+        atomic = all(type(item) in _ATOMIC_TYPES for item in value)
+        result = value.copy() if atomic else []
+        memo[key] = result
+        if not atomic:
+            result.extend(copy_request_snapshot(item, memo) for item in value)
+    elif kind is dict:
+        result = {}
+        memo[key] = result
+        for k, item in value.items():
+            result[copy_request_snapshot(k, memo)] = copy_request_snapshot(item, memo)
+    else:
+        return copy.deepcopy(value, memo)
+    # Keep sources alive for the duration of a copy, including custom deepcopy
+    # implementations that mutate a source container while being visited.
+    memo.setdefault(id(memo), []).append(value)
+    return result

vllm_omni/engine/stage_engine_core_client.py

@@ -19,7 +19,7 @@ import vllm.v1.engine as _vllm_engine_module
 import vllm.v1.engine.core_client as _vllm_core_client_module
 from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreRequest
-from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient
+from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient, MPClient
 from vllm.v1.engine.exceptions import EngineDeadError
 
 from vllm_omni.distributed.omni_connectors.utils.config import (
@@ -232,6 +232,9 @@ class StageEngineCoreClientBase(StageClientBase):
         if self.resources.engine_dead:
             raise EngineDeadError(f"Stage-{self.stage_id} engine core is dead")
 
+    def _apply_ready_response(self, payload: bytes) -> None:
+        MPClient._apply_ready_response(self, payload)
+
     # ==================== Overrides ====================
 
     async def add_request_async(self, request: EngineCoreRequest) -> None:

vllm_omni/engine/stage_engine_core_proc.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """
 Stage Core Process for vLLM-Omni V1 architecture.
 
@@ -27,6 +30,7 @@ from vllm.v1.engine.utils import (
     EngineZmqAddresses,
     SignalCallback,
 )
+from vllm.v1.executor.uniproc_executor import UniProcExecutor
 
 from vllm_omni.distributed.omni_coordinator import create_stage_coord_client
 from vllm_omni.engine import OmniEngineCoreRequest
@@ -75,6 +79,24 @@ def _signal_exit_code(signum: int) -> int:
     return _SIGNAL_EXIT_BASE + signum
 
 
+def _bind_native_data_plane_ready_sink(model_executor: Any, scheduler: Any) -> bool:
+    """Bind the TP1 in-process runner control plane directly to its scheduler."""
+    if not isinstance(model_executor, UniProcExecutor):
+        return False
+    parallel_config = model_executor.vllm_config.parallel_config
+    if parallel_config.tensor_parallel_size != 1 or parallel_config.pipeline_parallel_size != 1:
+        return False
+    driver_worker = getattr(model_executor, "driver_worker", None)
+    worker = getattr(driver_worker, "worker", None)
+    model_runner = getattr(worker, "model_runner", None)
+    data_plane = getattr(model_runner, "_omni_data_plane", None)
+    sink = getattr(scheduler, "enqueue_omni_connector_output", None)
+    if data_plane is None or not callable(sink):
+        return False
+    data_plane.set_omni_connector_output_sink(sink)
+    return True
+
+
 class StageEngineCoreProc(EngineCoreProc):
     """Stage-specific engine core process for vLLM-Omni.
 
@@ -83,6 +105,11 @@ class StageEngineCoreProc(EngineCoreProc):
     ``EngineCoreProc.run_engine_core()``.
     """
 
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        if _bind_native_data_plane_ready_sink(self.model_executor, self.scheduler):
+            logger.info("Bound native MRv2 connector readiness directly to the scheduler inbox.")
+
     def preprocess_add_request(self, request: OmniEngineCoreRequest) -> tuple[Any, int]:
         """Preserve omni payloads when vLLM builds its scheduler request."""
         scheduler_request, current_wave = super().preprocess_add_request(request)

vllm_omni/engine/stage_pool.py

@@ -15,6 +15,7 @@ from vllm.logger import init_logger
 from vllm.v1.engine import EngineCoreOutputs
 from vllm.v1.metrics.stats import IterationStats
 
+from vllm_omni.data_entry_keys import flatten_payload
 from vllm_omni.distributed.omni_coordinator import (
     LoadBalancer,
     OmniCoordClientForHub,
@@ -22,6 +23,7 @@ from vllm_omni.distributed.omni_coordinator import (
     ReplicaStatus,
 )
 from vllm_omni.distributed.omni_coordinator.load_balancer import Task
+from vllm_omni.engine.serialization import deserialize_additional_information
 from vllm_omni.engine.stage_client import (
     StagePoolClient,
     StagePoolDiffusionClient,
@@ -1136,8 +1138,30 @@ class StagePool:
         # gauges for that interval.
         if not outputs.outputs and outputs.scheduler_stats is None and not outputs.finished_requests:
             return None
+        self._rehydrate_pooling_output_payloads(outputs)
         return outputs
 
+    @staticmethod
+    def _rehydrate_pooling_output_payloads(outputs: EngineCoreOutputs) -> None:
+        """Restore dict-shaped pooling_output from its bytes carrier (MR V2).
+
+        vLLM decodes EngineCoreOutput.pooling_output as a torch.Tensor, so MR
+        V2 runners ship the per-request dict handoff via ``pooling_output_payload``
+        with ``pooling_output=None``. Decode it back here so downstream stage-input
+        processors see the same ``pooling_output`` shape as the legacy runner.
+        """
+        for eco in outputs.outputs:
+            payload = getattr(eco, "pooling_output_payload", None)
+            if payload is None:
+                continue
+            if getattr(eco, "pooling_output", None) is None:
+                # deserialize_* returns the nested OmniPayload form; the producer
+                # serialized a flat (dotted-key) pooler dict, and downstream
+                # consumers (e.g. talker2code2wav_full_payload) read flat keys like
+                # "codes.audio". Re-flatten to restore the exact on-wire shape.
+                eco.pooling_output = flatten_payload(deserialize_additional_information(payload))
+            eco.pooling_output_payload = None
+
     async def process_llm_raw_outputs(
         self,
         replica_id: int,

vllm_omni/entrypoints/openai/serving_speech.py

@@ -110,8 +110,8 @@ _REF_AUDIO_MIN_DURATION = 1.0  # seconds
 _REF_AUDIO_MAX_DURATION = 30.0  # seconds
 _REF_AUDIO_METADATA_FETCH_ATTEMPTS = 3
 _REMOTE_REF_AUDIO_SCHEMES = frozenset({"http", "https", "data"})
-_REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 256
-_REF_AUDIO_RESOLVE_CACHE_MAX_BYTES = 256 * 1024 * 1024
+_REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 1024
+_REF_AUDIO_RESOLVE_CACHE_MAX_BYTES = 512 * 1024 * 1024
 _TTS_MAX_INSTRUCTIONS_LENGTH = 500
 _DEFAULT_VOICE_NAME = "default"
 
@@ -1359,10 +1359,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin):
         self._ref_audio_resolve_cache.move_to_end(cache_key)
         return cached[3]
 
-    def _put_resolved_ref_audio(self, cache_key: str, waveform: np.ndarray, sr: int, artifact_key: str) -> None:
+    def _put_resolved_ref_audio(
+        self, cache_key: str, waveform: np.ndarray | list[float], sr: int, artifact_key: str
+    ) -> None:
         if self._ref_audio_resolve_cache_max_entries <= 0 or self._ref_audio_resolve_cache_max_bytes <= 0:
             return
-        size = waveform.nbytes
+        # Keep the finalized ndarray's ownership and identity for the array
+        # resolver. Legacy list callers are materialized into a compact array;
+        # their Python list remains independent from cached storage.
+        waveform = np.asarray(waveform, dtype=np.float32)
+        if not waveform.flags.c_contiguous:
+            waveform = np.ascontiguousarray(waveform)
+        size = int(waveform.nbytes)
         if size > self._ref_audio_resolve_cache_max_bytes:
             return
         previous = self._ref_audio_resolve_cache.pop(cache_key, None)

vllm_omni/model_executor/models/common/qwen3_code_predictor.py

@@ -682,14 +682,20 @@ class CodePredictorWrapper(nn.Module):
         self._lm_heads_list: list[nn.Module] | None = None
         self._codec_embeds_list: list[nn.Module] | None = None
         self._device_graphs: dict[int | tuple[int, int], tuple] = {}  # (graph, static_output) per bucket
+        # Outer-runner execution buckets when a model host declares them (MRv2);
+        # ``None`` keeps the legacy power-of-two bucket derivation.
+        self._execution_batch_buckets: list[int] | None = None
         prefix_graph_cfg = self._stage_connector_extra_config(vllm_config)
         prefix_graphs_requested = self._parse_bool_config(prefix_graph_cfg.get("code_predictor_prefix_graphs"))
         is_npu = current_omni_platform.is_npu()
-        self._prefix_graphs_enabled = prefix_graphs_requested and wrapper_config.use_cuda_graphs and not is_npu
-        if prefix_graphs_requested and not self._prefix_graphs_enabled:
+        # MRv2 captures the whole Talker MTP call. It can therefore capture
+        # shorter compiled re-prefill forwards without nesting predictor-owned
+        # CUDA graphs inside the outer graph.
+        self._prefix_reprefill_enabled = prefix_graphs_requested and not is_npu
+        self._prefix_graphs_enabled = self._prefix_reprefill_enabled and wrapper_config.use_cuda_graphs
+        if prefix_graphs_requested and not self._prefix_reprefill_enabled:
             logger.info_once(
-                "code_predictor: prefix CUDA graphs requested but disabled because use_cuda_graphs=%s is_npu=%s",
-                wrapper_config.use_cuda_graphs,
+                "code_predictor: prefix re-prefill requested but disabled because is_npu=%s",
                 is_npu,
             )
         self._prefix_graph_buckets = self._parse_positive_int_set(
@@ -698,6 +704,7 @@ class CodePredictorWrapper(nn.Module):
         self._prefix_graph_seq_lens = self._parse_positive_int_set(
             prefix_graph_cfg.get("code_predictor_prefix_graph_seq_lens")
         )
+        self._prefix_reprefill_seq_lens = tuple(self._prefix_seq_lens(self._num_groups + 1))
 
     def get_input_embeddings(self) -> nn.ModuleList:
         return self.model.get_input_embeddings()
@@ -761,7 +768,8 @@ class CodePredictorWrapper(nn.Module):
             dynamic=False,
             options={"epilogue_fusion": False},
         )
-        self._warmup_buckets()
+        with torch._dynamo.config.patch(cache_size_limit=self._compile_cache_size_limit()):
+            self._warmup_buckets()
 
         if self._wrapper_config.use_cuda_graphs:
             self._capture_cuda_graphs()
@@ -770,7 +778,7 @@ class CodePredictorWrapper(nn.Module):
             logger.info("code_predictor: torch.compile (dynamic=False, no epilogue fusion)")
 
     def _padded_bsz(self, bsz: int) -> int:
-        """Round batch size up to nearest power-of-2 bucket."""
+        """Round batch size up to the nearest warmed bucket."""
         for bucket in self._bucket_sizes:
             if bsz <= bucket:
                 return bucket
@@ -832,15 +840,28 @@ class CodePredictorWrapper(nn.Module):
         return row_generators
 
     @classmethod
-    def _sample_codes_gumbel(cls, logits: torch.Tensor, generator: _GeneratorLike = None) -> torch.Tensor:
-        """Sample ``logits`` via Gumbel-max with optional per-row generators."""
-        row_generators = cls._normalize_generators(generator, int(logits.shape[0]))
-        u = torch.empty_like(logits, dtype=torch.float32)
-        if isinstance(row_generators, list):
-            for row, row_generator in enumerate(row_generators):
-                u[row : row + 1].uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generator)
+    def _sample_codes_gumbel(
+        cls,
+        logits: torch.Tensor,
+        generator: _GeneratorLike = None,
+        uniforms: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Sample ``logits`` via Gumbel-max with optional precomputed noise."""
+        if uniforms is None:
+            row_generators = cls._normalize_generators(generator, int(logits.shape[0]))
+            u = torch.empty_like(logits, dtype=torch.float32)
+            if isinstance(row_generators, list):
+                for row, row_generator in enumerate(row_generators):
+                    u[row : row + 1].uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generator)
+            else:
+                u.uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generators)
         else:
-            u.uniform_(_UNIFORM_EPS, 1.0 - _UNIFORM_EPS, generator=row_generators)
+            if uniforms.shape != logits.shape:
+                raise ValueError(
+                    f"precomputed sampling uniforms must match logits: uniforms={tuple(uniforms.shape)} "
+                    f"logits={tuple(logits.shape)}"
+                )
+            u = uniforms
         return (logits.float() - torch.log(-torch.log(u))).argmax(dim=-1, keepdim=True)
 
     def _prefix_seq_lens(self, max_seq: int) -> list[int]:
@@ -850,13 +871,64 @@ class CodePredictorWrapper(nn.Module):
         allowed = set(all_seq_lens)
         return sorted(seq_len for seq_len in self._prefix_graph_seq_lens if seq_len in allowed)
 
-    def _warmup_buckets(self) -> None:
-        """Warmup power-of-2 batch-size buckets to front-load Inductor compilation."""
+    @staticmethod
+    def _synchronize_warmup(device: torch.device) -> None:
+        if device.type != "cpu":
+            current_omni_platform.synchronize()
+
+    def configure_mtp_execution_buckets(self, sizes: Iterable[int]) -> None:
+        """Declare the batch buckets the outer runner can actually reach.
+
+        Must be called before the first warmup; once buckets are captured the
+        existing set is only acceptable when it already covers the declaration
+        (never clear active graphs to rebuild them).
+        """
+        max_bsz = int(self._vllm_config.scheduler_config.max_num_seqs)
+        cleaned = sorted({int(size) for size in sizes if 0 < int(size) <= max_bsz} | {max_bsz})
+        if self._bucket_sizes:
+            missing = [size for size in cleaned if size not in self._bucket_sizes]
+            if missing:
+                raise RuntimeError(
+                    "mtp execution buckets must be configured before the first warmup; "
+                    f"missing={missing} existing={self._bucket_sizes}"
+                )
+            return
+        self._execution_batch_buckets = cleaned
+
+    def _batch_bucket_sizes(self) -> list[int]:
         max_bsz = self._vllm_config.scheduler_config.max_num_seqs
+        if self._execution_batch_buckets is not None:
+            # Outer-reachable buckets plus prefix-graph capture buckets and the
+            # full-size eager fallback row already merged at configure time.
+            bucket_sizes = list(self._execution_batch_buckets)
+            bucket_sizes.extend(bucket for bucket in self._prefix_graph_buckets if bucket <= max_bsz)
+            return sorted(set(bucket_sizes))
         bucket_sizes = [1 << i for i in range(max_bsz.bit_length()) if (1 << i) <= max_bsz]
-        if max_bsz not in bucket_sizes:
-            bucket_sizes.append(max_bsz)
-        self._bucket_sizes = sorted(bucket_sizes)
+        bucket_sizes.extend(bucket for bucket in self._prefix_graph_buckets if bucket <= max_bsz)
+        bucket_sizes.append(max_bsz)
+        return sorted(set(bucket_sizes))
+
+    def _compile_cache_size_limit(self) -> int:
+        bucket_sizes = self._batch_bucket_sizes()
+        if not self._prefix_reprefill_enabled:
+            required_entries = len(bucket_sizes)
+        else:
+            prefix_buckets = (
+                self._prefix_graph_buckets.intersection(bucket_sizes)
+                if self._prefix_graph_buckets
+                else set(bucket_sizes)
+            )
+            prefix_seq_lens = self._prefix_reprefill_seq_lens
+            needs_full_graph = set(prefix_seq_lens) != set(range(2, self._num_groups + 1))
+            full_graph_entries = len(bucket_sizes) - len(prefix_buckets)
+            if needs_full_graph:
+                full_graph_entries += len(prefix_buckets)
+            required_entries = full_graph_entries + len(prefix_buckets) * len(prefix_seq_lens)
+        return max(torch._dynamo.config.cache_size_limit, required_entries)
+
+    def _warmup_buckets(self) -> None:
+        """Warm up batch-size buckets to front-load Inductor compilation."""
+        self._bucket_sizes = self._batch_bucket_sizes()
 
         max_seq = self._num_groups + 1
         device = next(self.model.parameters()).device
@@ -866,8 +938,8 @@ class CodePredictorWrapper(nn.Module):
         self._ensure_buffers(device, self._model_dtype, max(self._bucket_sizes))
         proj_buf = self._proj_buf
 
-        if self._prefix_graphs_enabled:
-            prefix_seq_lens = self._prefix_seq_lens(max_seq)
+        if self._prefix_reprefill_enabled:
+            prefix_seq_lens = self._prefix_reprefill_seq_lens
             needs_full_graph = set(prefix_seq_lens) != set(range(2, max_seq))
             for bsz in self._bucket_sizes:
                 capture_prefixes = not self._prefix_graph_buckets or bsz in self._prefix_graph_buckets
@@ -905,6 +977,10 @@ class CodePredictorWrapper(nn.Module):
                     self._compiled_model_fwd(proj_buf[:bsz, :max_seq, :], pos_ids)
             logger.info("code_predictor: warmup done for buckets %s", self._bucket_sizes)
 
+        # Compiled attention warmup is asynchronous. Complete it before graph
+        # capture and serving reuse the same static buffers.
+        self._synchronize_warmup(device)
+
     def _capture_cuda_graphs(self) -> None:
         """Capture a CUDA graph per bucket using vLLM's global graph pool."""
         from vllm.platforms import current_platform
@@ -914,7 +990,7 @@ class CodePredictorWrapper(nn.Module):
         proj_buf = self._proj_buf
 
         if self._prefix_graphs_enabled:
-            prefix_seq_lens = self._prefix_seq_lens(max_seq)
+            prefix_seq_lens = self._prefix_reprefill_seq_lens
             needs_full_graph = set(prefix_seq_lens) != set(range(2, max_seq))
             for bsz in self._bucket_sizes:
                 capture_prefixes = not self._prefix_graph_buckets or bsz in self._prefix_graph_buckets
@@ -1013,13 +1089,20 @@ class CodePredictorWrapper(nn.Module):
         top_p: float = 1.0,
         generator: torch.Generator | None = None,
         generators: Sequence[torch.Generator | None] | None = None,
+        sample_uniforms: torch.Tensor | None = None,
     ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
         """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)}"
+                )
         sample_generator: _GeneratorLike = generators if generators is not None else generator
-        num_groups = self._num_groups
         device = layer0_code.device
 
         # _setup_compile caches _model_dtype on first call; use it for buffers
@@ -1072,11 +1155,16 @@ class CodePredictorWrapper(nn.Module):
         for step in range(1, num_groups):
             graph_key: int | tuple[int, int] = padded_bsz
             seq_len = max_seq
-            if self._prefix_graphs_enabled:
-                prefix_key = (padded_bsz, step + 1)
-                if prefix_key in self._device_graphs:
-                    graph_key = prefix_key
-                    seq_len = step + 1
+            if self._prefix_reprefill_enabled:
+                actual_seq_len = step + 1
+                for prefix_seq_len in self._prefix_reprefill_seq_lens:
+                    if prefix_seq_len < actual_seq_len:
+                        continue
+                    prefix_key = (padded_bsz, prefix_seq_len)
+                    if prefix_key in self._bucket_pos_ids:
+                        graph_key = prefix_key
+                        seq_len = prefix_seq_len
+                        break
             pos_ids = self._bucket_pos_ids.get(graph_key)
             if pos_ids is None:
                 pos_ids = (
@@ -1123,7 +1211,8 @@ class CodePredictorWrapper(nn.Module):
                     remove_mask = (cumulative_probs - sorted_probs) >= s_top_p
                     sorted_logits[remove_mask] = float("-inf")
                     logits = sorted_logits.scatter(1, sorted_idx, sorted_logits)
-                code = self._sample_codes_gumbel(logits, generator=sample_generator)
+                step_uniforms = sample_uniforms[:, step - 1, :] if sample_uniforms is not None else None
+                code = self._sample_codes_gumbel(logits, generator=sample_generator, uniforms=step_uniforms)
             else:
                 # "per_call" mode: temperature-scaled + top-k -> Gumbel-max
                 if use_sampling:
@@ -1131,7 +1220,8 @@ class CodePredictorWrapper(nn.Module):
                     if top_k > 0:
                         topk_vals, _ = scaled.topk(top_k, dim=-1)
                         scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf"))
-                    code = self._sample_codes_gumbel(scaled, generator=sample_generator)
+                    step_uniforms = sample_uniforms[:, step - 1, :] if sample_uniforms is not None else None
+                    code = self._sample_codes_gumbel(scaled, generator=sample_generator, uniforms=step_uniforms)
                 else:
                     code = logits.argmax(dim=-1, keepdim=True)
 

vllm_omni/model_executor/models/output_templates.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 from typing import NamedTuple
 
 import torch
@@ -6,6 +9,18 @@ from vllm.sequence import IntermediateTensors
 from vllm_omni.data_entry_keys import OmniPayload
 
 
+class OwnedBatchTensor(NamedTuple):
+    """A batch tensor whose producer transfers ownership to the buffer.
+
+    The producer must have freshly allocated the tensor (e.g. an index_select
+    result) and must not write it afterwards. The buffer keeps row views
+    instead of taking a second snapshot. Tensors without this marker are
+    treated as borrowed (graph output / shared scratch) and still get a clone.
+    """
+
+    tensor: torch.Tensor
+
+
 class OmniOutput(NamedTuple):
     """Output from the merged Omni model containing both text and audio."""
 

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -179,6 +179,7 @@ class Qwen3OmniMoeForConditionalGeneration(
         # vLLM run with no talker stage downstream, so no one consumes captured
         # thinker layers and the forward must return what stock vLLM expects.
         self.is_staged_run = getattr(vllm_config.model_config, "model_stage", None) is not None
+        self._returns_tuple = self.model_stage == "thinker" and self.is_staged_run
 
         if self.model_stage == "thinker":
             self.use_async_omni_output = True

vllm_omni/model_executor/models/qwen3_tts/configuration_qwen3_tts.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -533,6 +536,10 @@ class Qwen3TTSConfig(PretrainedConfig):
         if any("Code2Wav" in str(a) for a in archs):
             if hasattr(config, "rope_parameters"):
                 delattr(config, "rope_parameters")
+            # Code2Wav doesn't use position embeddings, but vLLM validates
+            # max_model_len <= max_position_embeddings. Set a large value
+            # to allow the pipeline.yaml max_model_len (65536) to pass validation.
+            config.max_position_embeddings = 131072
         return config
 
 

vllm_omni/model_executor/models/qwen3_tts/cuda_graph_decoder_wrapper.py

@@ -1,5 +1,6 @@
 # Copyright 2026 The Alibaba Qwen team.
 # SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """
 CUDA Graph wrapper for Qwen3TTSTokenizerV2Decoder.
 
@@ -8,10 +9,8 @@ reducing kernel launch overhead during inference.
 """
 
 import bisect
-import os
 import time
-from collections import Counter
-from collections.abc import Callable, Sequence
+from collections.abc import Callable, Collection, Sequence
 
 import torch
 from torch.cuda import CUDAGraph
@@ -187,22 +186,6 @@ class CUDAGraphDecoderWrapper:
 
         self._warmed_up = False
         self._device = None
-        self._stats_enabled = os.environ.get("VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS", "").lower() in (
-            "1",
-            "true",
-            "yes",
-            "on",
-        )
-        self._stats_log_every = int(os.environ.get("VLLM_OMNI_QWEN3_CODE2WAV_CUDAGRAPH_STATS_LOG_EVERY", "0") or 0)
-        self._stats_total = 0
-        self._stats_hits = 0
-        self._stats_compiled_hits = 0
-        self._stats_fallbacks = 0
-        self._stats_stream_capture_fallbacks = 0
-        self._stats_requests: Counter[tuple[int, int]] = Counter()
-        self._stats_hit_shapes: Counter[tuple[int, int, int]] = Counter()
-        self._stats_compiled_shapes: Counter[tuple[int, int]] = Counter()
-        self._stats_fallback_shapes: Counter[tuple[int, int, int]] = Counter()
 
     @staticmethod
     def compute_capture_sizes(
@@ -244,57 +227,19 @@ class CUDAGraphDecoderWrapper:
         shapes.update(self.extra_capture_shapes)
         return sorted(shapes)
 
-    def _record_decode_stats(
-        self,
-        *,
-        hit: bool,
-        batch_size: int,
+    @staticmethod
+    def _select_replay_shape(
+        actual_batch_size: int,
         actual_size: int,
-        padded_size: int | None,
-        stream_capture: bool = False,
-        compiled: bool = False,
-    ) -> None:
-        if not self._stats_enabled:
-            return
-
-        padded_key = int(padded_size) if padded_size is not None else -1
-        self._stats_total += 1
-        self._stats_requests[(batch_size, actual_size)] += 1
-        if hit:
-            self._stats_hits += 1
-            if compiled:
-                self._stats_compiled_hits += 1
-                self._stats_compiled_shapes[(batch_size, actual_size)] += 1
-            else:
-                self._stats_hit_shapes[(batch_size, actual_size, padded_key)] += 1
-        else:
-            self._stats_fallbacks += 1
-            self._stats_fallback_shapes[(batch_size, actual_size, padded_key)] += 1
-            if stream_capture:
-                self._stats_stream_capture_fallbacks += 1
-
-        if self._stats_log_every > 0 and self._stats_total % self._stats_log_every == 0:
-            self.log_decode_stats()
-
-    def log_decode_stats(self) -> None:
-        if not self._stats_enabled or self._stats_total == 0:
-            return
-        hit_rate = 100.0 * self._stats_hits / self._stats_total
-        logger.info(
-            "Code2Wav CUDA Graph stats: total=%d hits=%d fallbacks=%d "
-            "compiled_hits=%d stream_capture_fallbacks=%d hit_rate=%.2f%% "
-            "top_requests=%s top_compiled=%s top_hits=%s top_fallbacks=%s",
-            self._stats_total,
-            self._stats_hits,
-            self._stats_fallbacks,
-            self._stats_compiled_hits,
-            self._stats_stream_capture_fallbacks,
-            hit_rate,
-            self._stats_requests.most_common(12),
-            self._stats_compiled_shapes.most_common(12),
-            self._stats_hit_shapes.most_common(12),
-            self._stats_fallback_shapes.most_common(12),
+        available_shapes: Collection[tuple[int, int]],
+    ) -> tuple[int, int] | None:
+        """Select the lowest-cost captured shape that contains the request."""
+        candidates = (
+            (batch_size, size)
+            for batch_size, size in available_shapes
+            if batch_size >= actual_batch_size and size >= actual_size
         )
+        return min(candidates, key=lambda shape: (shape[0] * shape[1], shape[0], shape[1]), default=None)
 
     def warmup(
         self,
@@ -486,68 +431,42 @@ class CUDAGraphDecoderWrapper:
         # outside the startup capture window, so normal inference still hits
         # the graph fast path.
         if torch.cuda.is_current_stream_capturing():
-            self._record_decode_stats(
-                hit=False,
-                batch_size=int(codes.shape[0]),
-                actual_size=int(codes.shape[-1]),
-                padded_size=None,
-                stream_capture=True,
-            )
             return self.decoder(codes)
 
         batch_size = int(codes.shape[0])
         actual_size = int(codes.shape[-1])
-        padded_size = self._get_padded_size(actual_size)
-        compile_key = (batch_size, actual_size)
-        if compile_key not in self._compiled_shapes and padded_size is not None:
-            compile_key = (batch_size, padded_size)
-        if compile_key in self._compiled_shapes:
+        compile_key = self._select_replay_shape(batch_size, actual_size, self._compiled_shapes)
+        if compile_key is not None:
             compiled_size = compile_key[1]
-            self._record_decode_stats(
-                hit=True,
-                batch_size=batch_size,
-                actual_size=actual_size,
-                padded_size=compiled_size,
-                compiled=True,
-            )
             static_input = self._compiled_static_inputs[compile_key]
-            if actual_size == compiled_size:
+            if actual_size == compiled_size and batch_size == compile_key[0]:
                 static_input.copy_(codes)
             else:
                 static_input.zero_()
-                static_input[:, :, :actual_size] = codes
+                static_input[:batch_size, :, :actual_size] = codes
             self._compiled_graphs[compile_key].replay()
             output = self._trim_replay_output(self._compiled_static_outputs[compile_key], actual_size, compiled_size)
+            output = output[:batch_size]
             if clone_graph_output:
                 return output.clone()
             return output
 
-        graph_key = (batch_size, padded_size) if padded_size is not None else None
+        graph_key = self._select_replay_shape(batch_size, actual_size, self.graphs)
 
-        if graph_key is None or graph_key not in self.graphs:
-            self._record_decode_stats(
-                hit=False,
-                batch_size=batch_size,
-                actual_size=actual_size,
-                padded_size=padded_size,
-            )
+        if graph_key is None:
             return self.decoder(codes)
 
-        self._record_decode_stats(
-            hit=True,
-            batch_size=batch_size,
-            actual_size=actual_size,
-            padded_size=padded_size,
-        )
         static_input = self.static_inputs[graph_key]
-        if actual_size == padded_size:
+        graph_size = graph_key[1]
+        if actual_size == graph_size and batch_size == graph_key[0]:
             static_input.copy_(codes)
         else:
             static_input.zero_()
-            static_input[:, :, :actual_size] = codes
+            static_input[:batch_size, :, :actual_size] = codes
         self.graphs[graph_key].replay()
 
-        output = self._trim_replay_output(self.static_outputs[graph_key], actual_size, padded_size)
+        output = self._trim_replay_output(self.static_outputs[graph_key], actual_size, graph_size)
+        output = output[:batch_size]
         if clone_graph_output:
             return output.clone()
         return output

vllm_omni/model_executor/models/qwen3_tts/pipeline.py

@@ -1,5 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 """Qwen3-TTS pipeline: Talker (text → RVQ codec) → Code2Wav (codec → audio).
 
 Chunked vs end-to-end mode is dispatched from ``deploy.async_chunk``.
@@ -41,6 +41,7 @@ QWEN3_TTS_PIPELINE = PipelineConfig(
             owns_tokenizer=True,
             engine_output_type="latent",
             async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
+            supports_native_mrv2_data_plane=True,
             custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload",
             sampling_constraints={
                 "detokenize": False,
@@ -56,6 +57,7 @@ QWEN3_TTS_PIPELINE = PipelineConfig(
             final_output_type="audio",
             engine_output_type="audio",
             model_arch="Qwen3TTSCode2Wav",
+            supports_native_mrv2_data_plane=True,
             # ``sync_process_input_func`` is the only input-proc override for
             # this stage in sync (non-async-chunk) mode: a length-only
             # ``_token_only`` placeholder.  The bulk codec payload itself

vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py

@@ -198,10 +198,6 @@ def mel_spectrogram(
     hann_window: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Mel spectrogram via torch STFT and a (cached) mel filterbank."""
-    if torch.min(y) < -1.0:
-        logger.warning("Min value of input waveform signal is %s", torch.min(y))
-    if torch.max(y) > 1.0:
-        logger.warning("Max value of input waveform signal is %s", torch.max(y))
     device = y.device
     if mel_basis is None:
         mel_basis = _cached_mel_filter_bank(sampling_rate, n_fft, num_mels, fmin, fmax).to(device)
@@ -350,7 +346,7 @@ class Qwen3TTSPromptEmbedsBuilder:
         tts_pad_embed: torch.Tensor,
         encode_ref_audio_batch: Callable[..., list[torch.Tensor]],
         speaker_cache: Any | None = None,
-        ref_audio_artifact_cache_max_entries: int = 256,
+        ref_audio_artifact_cache_max_entries: int = 1024,
     ):
         self._config = config
         self._talker_config = talker_config
@@ -366,9 +362,14 @@ class Qwen3TTSPromptEmbedsBuilder:
         self._embedding_dtype = torch.bfloat16
 
         self._text_tokenizer: Any | None = None
+        # Projected special-token embeddings are immutable model constants.
+        # 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] = {}
 
         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()
+        self._long_tensor_cache: dict[tuple[str, tuple[int, ...]], torch.Tensor] = {}
 
         # Bounded LRU; caller-supplied orig_sr can otherwise grow this without limit.
         self._resampler_cache: OrderedDict[tuple[int, int], AudioResampler] = OrderedDict()
@@ -384,6 +385,29 @@ class Qwen3TTSPromptEmbedsBuilder:
         """Return the pad-token embedding on the requested device/dtype, shape ``[1, 1, H]``."""
         return self._tts_pad_embed_buffer.to(device=device, dtype=dtype).reshape(1, 1, -1)
 
+    def _long_tensor(self, values: Sequence[int], device: torch.device) -> torch.Tensor:
+        device = torch.device(device)
+        if device.type == "cuda" and device.index is None:
+            device = torch.device("cuda", torch.accelerator.current_device_index())
+        key = (str(device), tuple(int(v) for v in values))
+        cache = getattr(self, "_long_tensor_cache", None)
+        if cache is None:
+            cache = self._long_tensor_cache = {}
+        cached = cache.get(key)
+        if cached is None or cached.device != device:
+            cached = torch.tensor([list(key[1])], device=device, dtype=torch.long)
+            cache[key] = cached
+        return cached
+
+    def _projected_tokens(self, values: Sequence[int], device: torch.device) -> torch.Tensor:
+        device = torch.device(device)
+        key = (str(device), tuple(int(v) for v in values))
+        cached = self._projected_token_cache.get(key)
+        if cached is None:
+            cached = self._text_projection(self._text_embedding(self._long_tensor(values, device)))
+            self._projected_token_cache[key] = cached
+        return cached
+
     def _get_resampler(self, orig_sr: int, target_sr: int) -> AudioResampler:
         key = (int(orig_sr), int(target_sr))
         cache = self._resampler_cache
@@ -631,9 +655,9 @@ class Qwen3TTSPromptEmbedsBuilder:
         if entry is None:
             entry = {}
         if isinstance(ref_code, torch.Tensor):
-            entry["ref_code"] = ref_code.detach().to("cpu", dtype=torch.long).contiguous()
+            entry["ref_code"] = ref_code.detach().to(dtype=torch.long).contiguous()
         if isinstance(ref_spk_embedding, torch.Tensor):
-            entry["ref_spk_embedding"] = ref_spk_embedding.detach().to("cpu", dtype=self._embedding_dtype).reshape(-1)
+            entry["ref_spk_embedding"] = ref_spk_embedding.detach().to(dtype=self._embedding_dtype).reshape(-1)
         if not entry:
             return
         self._ref_audio_artifact_cache[cache_key] = entry
@@ -659,7 +683,6 @@ class Qwen3TTSPromptEmbedsBuilder:
                 self._speaker_encoder.to(device=dev, dtype=dtype)
         except StopIteration:
             pass
-
         # Resample to 24kHz for speaker encoder.
         target_sr = int(getattr(self._config.speaker_encoder_config, "sample_rate", 24000))
         if sr != target_sr:
@@ -761,6 +784,23 @@ class Qwen3TTSPromptEmbedsBuilder:
         req_ids: list[str],
         model_intermediate_buffer: dict[str, dict[str, Any]],
         device: torch.device,
+    ) -> None:
+        """MRv1 adapter for the request-info batch implementation."""
+        self.preprocess_infos_batch(
+            req_infos=[
+                info_dict
+                for req_id in req_ids
+                if isinstance((info_dict := model_intermediate_buffer.get(req_id)), dict)
+            ],
+            device=device,
+        )
+
+    @torch.inference_mode()
+    def preprocess_infos_batch(
+        self,
+        *,
+        req_infos: list[dict[str, Any]],
+        device: torch.device,
     ) -> None:
         """Batch Base voice-clone ref-audio codec extraction for current prefill requests.
 
@@ -775,11 +815,7 @@ class Qwen3TTSPromptEmbedsBuilder:
         pending_text: list[tuple[dict[str, Any], str]] = []
         pending_ref_text: list[tuple[dict[str, Any], str]] = []
         groups: dict[int, list[tuple[dict[str, Any], np.ndarray, int]]] = {}
-        for req_id in req_ids:
-            info_dict = model_intermediate_buffer.get(req_id)
-            if not isinstance(info_dict, dict):
-                continue
-
+        for info_dict in req_infos:
             if (
                 self._needs_initial_prompt_preprocess(info_dict)
                 and first_value(info_dict.get("task_type"), "CustomVoice") == "Base"
@@ -837,17 +873,31 @@ class Qwen3TTSPromptEmbedsBuilder:
                 except Exception as exc:
                     logger.debug("Qwen3-TTS batched text tokenization failed; falling back to serial path: %s", exc)
                     continue
-                input_ids = tokenized.get("input_ids") if isinstance(tokenized, dict) else None
+                input_ids = tokenized.get("input_ids") if isinstance(tokenized, Mapping) else None
                 if not isinstance(input_ids, list) or len(input_ids) != len(items):
                     continue
+                valid_items: list[tuple[dict[str, Any], int]] = []
+                flat_ids: list[int] = []
                 for (info_dict, _), ids in zip(items, input_ids, strict=True):
-                    if isinstance(ids, list) and ids:
-                        info_dict[key] = torch.tensor(ids, dtype=torch.long)
+                    if not (
+                        isinstance(ids, list)
+                        and ids
+                        and all(isinstance(token_id, (int, np.integer)) for token_id in ids)
+                    ):
+                        continue
+                    valid_items.append((info_dict, len(ids)))
+                    flat_ids.extend(int(token_id) for token_id in ids)
+                if not valid_items:
+                    continue
+                use_async_h2d = device.type == "cuda"
+                packed_ids_cpu = torch.tensor(flat_ids, dtype=torch.long, pin_memory=use_async_h2d)
+                packed_ids = packed_ids_cpu.to(device=device, non_blocking=use_async_h2d)
+                offset = 0
+                for info_dict, num_ids in valid_items:
+                    info_dict[key] = packed_ids.narrow(0, offset, num_ids)
+                    offset += num_ids
 
         groups = {sr: items for sr, items in groups.items() if len(items) >= 2}
-        if not groups:
-            return
-
         for sr, items in groups.items():
             wavs = [wav for _, wav, _ in items]
             try:
@@ -892,9 +942,7 @@ class Qwen3TTSPromptEmbedsBuilder:
         codec_embed_sum = torch.cat(codec_embed, dim=1).sum(1).unsqueeze(0)  # [1,T,H]
         codec_embed_sum = torch.cat(
             [
-                self._codec_embed(
-                    torch.tensor([[self._talker_config.codec_bos_id]], device=codec_embed_sum.device, dtype=torch.long)
-                ),
+                self._codec_embed(self._long_tensor([self._talker_config.codec_bos_id], codec_embed_sum.device)),
                 codec_embed_sum,
             ],
             dim=1,
@@ -990,12 +1038,8 @@ class Qwen3TTSPromptEmbedsBuilder:
         # tts special token embeds (projected into talker hidden).
         # ``tts_pad_embed`` is precomputed (request-independent), so we only
         # need bos/eos here.
-        tts_tokens = torch.tensor(
-            [[config.tts_bos_token_id, config.tts_eos_token_id]],
-            device=input_ids.device,
-            dtype=input_ids.dtype,
-        )
-        tts_bos_embed, tts_eos_embed = text_projection(text_embedding(tts_tokens)).chunk(2, dim=1)
+        projected_tts = self._projected_tokens([config.tts_bos_token_id, config.tts_eos_token_id], input_ids.device)
+        tts_bos_embed, tts_eos_embed = projected_tts.chunk(2, dim=1)
         tts_pad_embed = self._pad_embed(input_ids.device, tts_bos_embed.dtype)
 
         # Codec prefill tags.
@@ -1032,9 +1076,9 @@ class Qwen3TTSPromptEmbedsBuilder:
                 ]
             ]
 
-        codec_input_0 = codec_embed(torch.tensor(codec_prefill_list, device=input_ids.device, dtype=torch.long))
+        codec_input_0 = codec_embed(self._long_tensor(codec_prefill_list[0], input_ids.device))
         codec_input_1 = codec_embed(
-            torch.tensor([[talker_config.codec_pad_id, talker_config.codec_bos_id]], device=input_ids.device)
+            self._long_tensor([talker_config.codec_pad_id, talker_config.codec_bos_id], input_ids.device)
         )
 
         # Speaker embedding/token (task-dependent)
@@ -1339,7 +1383,7 @@ class Qwen3TTSPromptEmbedsBuilder:
                             talker_prompt,
                             text_all + codec_embed(pad_ids),
                             tts_pad_embed
-                            + codec_embed(torch.tensor([[talker_config.codec_bos_id]], device=input_ids.device)),
+                            + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device)),
                         ],
                         dim=1,
                     )
@@ -1392,8 +1436,7 @@ class Qwen3TTSPromptEmbedsBuilder:
                     [
                         talker_prompt,
                         text_all + codec_embed(pad_ids),
-                        tts_pad_embed
-                        + codec_embed(torch.tensor([[talker_config.codec_bos_id]], device=input_ids.device)),
+                        tts_pad_embed + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device)),
                     ],
                     dim=1,
                 )
@@ -1431,8 +1474,7 @@ class Qwen3TTSPromptEmbedsBuilder:
                     [
                         talker_prompt,
                         text_all + codec_embed(pad_ids),
-                        tts_pad_embed
-                        + codec_embed(torch.tensor([[talker_config.codec_bos_id]], device=input_ids.device)),
+                        tts_pad_embed + codec_embed(self._long_tensor([talker_config.codec_bos_id], input_ids.device)),
                     ],
                     dim=1,
                 )

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code2wav.py

@@ -46,9 +46,9 @@ def _codec_ids_from_payload_or_input(
         if isinstance(codes, dict):
             audio = codes.get("audio")
             if isinstance(audio, torch.Tensor) and audio.numel() > 0:
-                return audio.reshape(-1).to(device=input_ids.device, dtype=torch.long)
+                return audio.reshape(-1).to(dtype=torch.long)
             if isinstance(audio, (list, tuple)) and audio:
-                return torch.as_tensor(audio, device=input_ids.device, dtype=torch.long).reshape(-1)
+                return torch.as_tensor(audio, dtype=torch.long).reshape(-1)
     return input_ids.reshape(-1).to(dtype=torch.long)
 
 
@@ -63,6 +63,10 @@ class Qwen3TTSCode2Wav(nn.Module):
     # decoder caches must use the same IDs delivered by on_requests_finished;
     # payload metadata carries an external ID which may differ.
     requires_request_ids = True
+    # A nonempty native codes.audio payload is authoritative; token IDs may
+    # serve only as per-request control slots for the generation scheduler.
+    supports_native_payload_input = True
+    batched_gpu_staging_keys = {("codes", "audio")}
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -427,12 +431,9 @@ class Qwen3TTSCode2Wav(nn.Module):
             try:
                 _, c = valid_codes_qf[0]
                 logger.info(
-                    "Code2Wav codec: frames=%d q=%d uniq=%d range=[%d,%d] batch=%d",
+                    "Code2Wav codec: frames=%d q=%d batch=%d",
                     c.shape[1],
                     q,
-                    int(torch.unique(c).numel()),
-                    int(c.min().item()),
-                    int(c.max().item()),
                     len(valid_codes_qf),
                 )
             except Exception:
@@ -454,9 +455,22 @@ class Qwen3TTSCode2Wav(nn.Module):
 
         request_lengths = [int(codes_qf.shape[-1]) for _, codes_qf in valid_codes_qf]
         max_request_length = max(request_lengths)
-        request_codes = valid_codes_qf[0][1].new_zeros((len(valid_codes_qf), q, max_request_length))
-        for row, (_, codes_qf) in enumerate(valid_codes_qf):
-            request_codes[row, :, : codes_qf.shape[-1]].copy_(codes_qf)
+        request_codes_shape = (len(valid_codes_qf), q, max_request_length)
+        target_device = ids.device
+        if target_device.type == "cuda" and all(codes_qf.device.type == "cpu" for _, codes_qf in valid_codes_qf):
+            staged_codes = torch.zeros(
+                request_codes_shape,
+                dtype=torch.long,
+                device="cpu",
+                pin_memory=True,
+            )
+            for row, (_, codes_qf) in enumerate(valid_codes_qf):
+                staged_codes[row, :, : codes_qf.shape[-1]].copy_(codes_qf)
+            request_codes = staged_codes.to(device=target_device, non_blocking=True)
+        else:
+            request_codes = torch.zeros(request_codes_shape, dtype=torch.long, device=target_device)
+            for row, (_, codes_qf) in enumerate(valid_codes_qf):
+                request_codes[row, :, : codes_qf.shape[-1]].copy_(codes_qf)
 
         self._record_decode_batch_stats(
             group_size=len(valid_codes_qf),

vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_talker.py

@@ -4,7 +4,8 @@
 from __future__ import annotations
 
 import copy
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from contextlib import contextmanager
 from typing import Any
 
 import numpy as np
@@ -25,13 +26,14 @@ from vllm.multimodal.audio import AudioResampler
 from vllm.sequence import IntermediateTensors
 
 from vllm_omni.data_entry_keys import OmniPayload
-from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor
 from vllm_omni.utils.speaker_cache import (
     get_speaker_cache,
     iter_custom_voice_profiles,
     load_validated_profile_tensors,
     validate_qwen3_tts_profile,
 )
+from vllm_omni.worker.sampling_utils import get_tts_local_seed
 
 from .configuration_qwen3_tts import Qwen3TTSConfig, Qwen3TTSSpeakerEncoderConfig, Qwen3TTSTalkerConfig
 from .prompt_embeds_builder import PRECOMPUTED_TEXT_IDS_KEY, Qwen3TTSPromptEmbedsBuilder, resolve_x_vector_only
@@ -42,6 +44,80 @@ from .tokenizer_12hz.modeling_qwen3_tts_tokenizer_v2 import Qwen3TTSTokenizerV2E
 logger = init_logger(__name__)
 
 _TRAILING_TEXT_COMPACT_MIN_FRAMES = 64
+_ASYNC_REF_CODES_PUBLISHED_KEY = "_qwen3_tts_ref_codes_published"
+
+
+def _qwen3_tts_gpu_resident_buffer_keys(
+    use_v2_model_runner: bool,
+) -> set[tuple[str, str]]:
+    keys = {
+        ("codes", "audio"),
+        ("hidden_states", "last"),
+        ("hidden_states", "trailing_text"),
+    }
+    if use_v2_model_runner:
+        keys.update(
+            {
+                ("codes", "ref"),
+                ("embed", "prefill"),
+                ("meta", "codec_frame_valid"),
+            }
+        )
+    return keys
+
+
+def _materialize_span_scalars(
+    values: Sequence[int],
+    span_lengths: Sequence[int],
+    *,
+    dtype: torch.dtype,
+    device: torch.device,
+) -> torch.Tensor | None:
+    """Expand a uniform scalar cohort with one GPU allocation."""
+    if len(values) != len(span_lengths):
+        raise ValueError(
+            f"values and span_lengths must have the same length, got {len(values)} and {len(span_lengths)}"
+        )
+    total = sum(int(length) for length in span_lengths)
+    if not values:
+        return torch.empty((0,), dtype=dtype, device=device)
+    first = int(values[0])
+    if not all(int(value) == first for value in values[1:]):
+        return None
+    return torch.full((total,), first, dtype=dtype, device=device)
+
+
+def _should_publish_async_ref_codes(info: dict[str, Any]) -> bool:
+    """Publish ref codes exactly once, including after a KV-resumed prefill."""
+    meta = info.get("meta")
+    return not (isinstance(meta, dict) and bool(meta.get(_ASYNC_REF_CODES_PUBLISHED_KEY)))
+
+
+def _mark_async_ref_codes_published(info: dict[str, Any]) -> None:
+    meta = info.setdefault("meta", {})
+    if not isinstance(meta, dict):
+        raise TypeError("Qwen3-TTS payload meta must be a dict")
+    meta[_ASYNC_REF_CODES_PUBLISHED_KEY] = True
+
+
+def _ref_audio_artifact_cache_capacity(vllm_config: VllmConfig) -> int:
+    model_cfg = getattr(vllm_config, "model_config", None)
+    connector_cfg = getattr(model_cfg, "stage_connector_config", None)
+    connector_extra = (
+        connector_cfg.get("extra", connector_cfg)
+        if isinstance(connector_cfg, dict)
+        else getattr(connector_cfg, "extra", None)
+    )
+    if not isinstance(connector_extra, dict):
+        connector_extra = {}
+    raw_capacity = connector_extra.get("ref_audio_artifact_cache_max_entries", 1024)
+    try:
+        capacity = int(raw_capacity)
+    except (TypeError, ValueError) as exc:
+        raise ValueError(f"Invalid ref_audio_artifact_cache_max_entries={raw_capacity!r}") from exc
+    if capacity < 0:
+        raise ValueError(f"ref_audio_artifact_cache_max_entries must be non-negative, got {capacity}")
+    return capacity
 
 
 def _has_tts_text_conditioning(info_dict: dict[str, Any], hidden_states: Any | None = None) -> bool:
@@ -267,8 +343,7 @@ class Qwen3TTSSpeakerEncoder(torch.nn.Module):
 
 
 class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
-    """vLLM-AR talker: step-wise layer-0 codec decoding.
-    Predicts residual codebooks (1..Q-1) into `audio_codes` and streams text via `tailing_text_hidden`."""
+    """vLLM-AR talker with step-wise layer-0 codec decoding."""
 
     hf_to_vllm_mapper = WeightsMapper(
         orig_to_new_prefix={
@@ -329,29 +404,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # per-request additional_information.
         self.talker_mtp_output_key = ("codes", "audio")
         self.talker_mtp_graph_safe = True
-        # talker_mtp samples with per-row generators, so explicitly-seeded
-        # requests stay batched instead of one scalar forward per row (#4883).
-        # This is only valid while talker_mtp receives the *unpadded* active
-        # batch, i.e. when it is NOT graph-wrapped. The runner wraps talker_mtp
-        # (padded batches + a single captured RNG stream) whenever full
-        # cudagraphs are enabled, so disable per-row generators in that case.
-        # Consequence: under full cudagraphs, per-request ``tts_local_seed`` is
-        # not reproducible (matching Qwen3-Omni, which has no per-row seeding);
-        # eager mode keeps the per-row generator fast-path.
-        # TODO(#4923): the model should not read ``cudagraph_mode`` — it is an
-        # upstream (runner) concern, and the runner already re-derives the same
-        # wrap decision in ``_init_talker_mtp``. Once all TTS models declare
-        # ``talker_mtp_graph_safe`` we should move this gating into the model
-        # runner (models only expose the static capability flag, and users can
-        # then configure ``cudagraph_mode`` freely). This also means mtp seeding
-        # is silently dropped when the decode batch > 1 under full cudagraphs;
-        # the follow-up should refactor the per-row generator in the runner to
-        # be cudagraph-safe so seeded batches stay reproducible.
-        cudagraph_mode = getattr(vllm_config.compilation_config, "cudagraph_mode", None)
-        talker_mtp_graph_wrapped = (
-            self.talker_mtp_graph_safe and cudagraph_mode is not None and cudagraph_mode.has_full_cudagraphs()
-        )
-        self.talker_mtp_accepts_per_row_generators = not talker_mtp_graph_wrapped
+        # 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.
+        self.talker_mtp_accepts_per_row_generators = True
+        self.mtp_sample_uniforms = True
+        self.mtp_sample_steps = max(0, int(self.talker_config.num_code_groups) - 1)
+        self.mtp_sample_vocab_size = self._codebook_vocab_size
         self.use_async_omni_output = True
         self.eager_omni_postprocess_before_async_output = True
         self.omni_pooler_payload_include_hidden = False
@@ -445,14 +504,24 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # the same req_id ordering that indexes sampling_metadata.
         self._req_x_vector_only: dict[str, bool | None] = {}
         self._batch_req_ids: list[str] = []
+        self._mrv2_silence_ban_mask: torch.Tensor | None = None
 
         # Keys that should stay on GPU in model_intermediate_buffer to avoid
         # CPU-to-GPU round-trips on every decode step.
-        self.gpu_resident_buffer_keys: set[tuple[str, str]] = {
-            ("codes", "audio"),
-            ("hidden_states", "last"),
-            ("hidden_states", "trailing_text"),
-        }
+        self._use_v2_model_runner = bool(getattr(vllm_config.model_config, "use_v2_model_runner", False))
+        self.gpu_resident_buffer_keys = _qwen3_tts_gpu_resident_buffer_keys(self._use_v2_model_runner)
+
+        if self._use_v2_model_runner:
+            # Align predictor warmup/capture buckets with the outer MRv2 MTP
+            # graph buckets so the compile cache covers every reachable size.
+            outer_max_bsz = int(vllm_config.scheduler_config.max_num_seqs)
+            outer_buckets = {
+                int(size)
+                for size in vllm_config.compilation_config.cudagraph_capture_sizes
+                if 0 < int(size) <= outer_max_bsz
+            }
+            outer_buckets.update({1, outer_max_bsz})  # eager fallback rows
+            self.code_predictor.configure_mtp_execution_buckets(sorted(outer_buckets))
 
         # ``text_proj(text_emb(tts_pad_token_id))`` is request-independent —
         # it depends only on frozen ``text_embedding`` / ``text_projection``
@@ -492,6 +561,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         self._subtalker_sampling_params: dict[str, Any] = (
             dict(raw_subtalker_sampling) if isinstance(raw_subtalker_sampling, Mapping) else {}
         )
+        self.mtp_sampling_params = self._subtalker_sampling_params
 
         self._stacked_codec_embed: torch.Tensor | None = None
 
@@ -514,6 +584,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             tts_pad_embed=self._tts_pad_embed,
             encode_ref_audio_batch=self._encode_ref_audio_batch,
             speaker_cache=self._speaker_cache,
+            ref_audio_artifact_cache_max_entries=_ref_audio_artifact_cache_capacity(vllm_config),
         )
         self._load_custom_voice_profiles()
 
@@ -592,6 +663,10 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         # Suppress silence-region tokens for the first N decode frames (#4966). The
         # per-request decode step is the length of its generated-token history.
         ban_n = self._silence_ban_frames
+        if ban_n > 0 and sampling_metadata is None and self._mrv2_silence_ban_mask is not None:
+            if self._mrv2_silence_ban_mask.shape[0] != logits.shape[0]:
+                raise ValueError("MRV2 silence-ban batch must match the sampled logits rows")
+            return logits.masked_fill(self._mrv2_silence_ban_mask & self._silence_mask, float("-inf"))
         if ban_n > 0 and sampling_metadata is not None:
             output_token_ids = getattr(sampling_metadata, "output_token_ids", None)
             if output_token_ids is None or len(output_token_ids) != logits.shape[0]:
@@ -645,11 +720,17 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             info_dicts = kwargs.get("runtime_additional_information") or []
         if "runtime_additional_information" in kwargs and "model_intermediate_buffer" not in kwargs:
             logger.warning_once("runtime_additional_information is deprecated, use model_intermediate_buffer")
+        model_config = getattr(getattr(self, "vllm_config", None), "model_config", None)
+        async_chunk = bool(getattr(model_config, "async_chunk", False))
         audio_codes_list: list[torch.Tensor] = []
-        ref_code_len_list: list[torch.Tensor] = []
+        ref_code_len_segments: list[tuple[int | torch.Tensor, int, torch.device]] = []
         ref_code_list: list[torch.Tensor] = []
         has_ref_code = False
-        codec_streaming_list: list[torch.Tensor] = []
+        codec_streaming_values: list[int] = []
+        codec_streaming_span_lengths: list[int] = []
+        codec_streaming_device: torch.device | None = None
+        codec_frame_valid_device: torch.device | None = None
+        codec_frame_valid_parts: list[torch.Tensor] = []
         for info in info_dicts:
             if not isinstance(info, dict):
                 ref_code_list.append(torch.empty(0, dtype=torch.long))
@@ -659,17 +740,46 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             ac = codes.get("audio")
             if isinstance(ac, torch.Tensor):
                 audio_codes_list.append(ac)
-                cs = meta.get("codec_streaming")
-                if isinstance(cs, bool):
-                    codec_streaming_list.append(
-                        torch.full((int(ac.shape[0]),), int(cs), dtype=torch.int8, device=ac.device)
+                frame_valid = meta.get("codec_frame_valid")
+                if isinstance(frame_valid, bool):
+                    codec_frame_valid_device = ac.device
+                    codec_frame_valid_parts.append(
+                        torch.full(
+                            (int(ac.shape[0]),),
+                            int(frame_valid),
+                            dtype=torch.int8,
+                            device=ac.device,
+                        )
                     )
+                elif isinstance(frame_valid, torch.Tensor):
+                    if frame_valid.numel() != 1:
+                        raise ValueError("codec_frame_valid must be scalar per request")
+                    codec_frame_valid_device = ac.device
+                    codec_frame_valid_parts.append(
+                        frame_valid.to(device=ac.device, dtype=torch.int8, non_blocking=True)
+                        .reshape(1)
+                        .expand(int(ac.shape[0]))
+                    )
+                cs = meta.get("codec_streaming")
+                if not async_chunk and isinstance(cs, bool):
+                    codec_streaming_values.append(int(cs))
+                    codec_streaming_span_lengths.append(int(ac.shape[0]))
+                    codec_streaming_device = ac.device
             ref_code = codes.get("ref")
-            if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
+            publish_ref_code = not async_chunk or _should_publish_async_ref_codes(info)
+            if publish_ref_code and isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
                 ref_code_list.append(ref_code)
                 has_ref_code = True
+                if async_chunk:
+                    _mark_async_ref_codes_published(info)
             else:
                 ref_code_list.append(torch.empty(0, dtype=torch.long))
+            if async_chunk:
+                # The streaming consumer uses codes.ref.shape[0] directly and
+                # does not read ref_code_len or codec_streaming. Avoid creating
+                # GPU control tensors that would only be snapshotted and copied
+                # back to the CPU on every decode step.
+                continue
             ref_len = meta.get("ref_code_len")
             if ref_len is None:
                 continue
@@ -680,7 +790,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 if ref_len.numel() == 0:
                     raise ValueError("ref_code_len is an empty tensor")
                 ref_len_tail = ref_len.reshape(-1)[-1:].to(dtype=torch.int32, device=ac.device)
-                ref_code_len_list.append(ref_len_tail.expand(span_len).contiguous())
+                ref_code_len_segments.append((ref_len_tail, span_len, ac.device))
                 continue
             if isinstance(ref_len, list):
                 if len(ref_len) != 1:
@@ -688,7 +798,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 ref_len_val = int(ref_len[0])
             else:
                 ref_len_val = int(ref_len)
-            ref_code_len_list.append(torch.full((span_len,), ref_len_val, dtype=torch.int32, device=ac.device))
+            ref_code_len_segments.append((ref_len_val, span_len, ac.device))
 
         if not audio_codes_list:
             return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
@@ -696,15 +806,61 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         audio_codes = torch.cat(audio_codes_list, dim=0)
         span_len = int(audio_codes.shape[0])
         mm: OmniPayload = {"codes": {"audio": audio_codes}}
-        if ref_code_len_list:
-            mm.setdefault("meta", {})["ref_code_len"] = torch.cat(ref_code_len_list, dim=0)[:span_len]
+        if ref_code_len_segments:
+            if all(isinstance(value, int) for value, _length, _device in ref_code_len_segments):
+                ref_code_len = _materialize_span_scalars(
+                    [int(value) for value, _length, _device in ref_code_len_segments],
+                    [length for _value, length, _device in ref_code_len_segments],
+                    dtype=torch.int32,
+                    device=ref_code_len_segments[0][2],
+                )
+                if ref_code_len is None:
+                    ref_code_len = torch.cat(
+                        [
+                            torch.full((length,), int(value), dtype=torch.int32, device=device)
+                            for value, length, device in ref_code_len_segments
+                        ],
+                        dim=0,
+                    )
+            else:
+                ref_code_len_parts = [
+                    value.expand(length)
+                    if isinstance(value, torch.Tensor)
+                    else torch.full((length,), value, dtype=torch.int32, device=device)
+                    for value, length, device in ref_code_len_segments
+                ]
+                ref_code_len = torch.cat(ref_code_len_parts, dim=0)
+            mm.setdefault("meta", {})["ref_code_len"] = ref_code_len[:span_len]
         if has_ref_code:
             # Batch-aligned, one entry per request: ``to_payload_element``
             # indexes ``element[idx]`` per request, and a shorter list would
             # silently broadcast one request's reference codes to the others.
             mm.setdefault("codes", {})["ref"] = ref_code_list
-        if codec_streaming_list:
-            mm.setdefault("meta", {})["codec_streaming"] = torch.cat(codec_streaming_list, dim=0)[:span_len]
+        if codec_streaming_values:
+            assert codec_streaming_device is not None
+            codec_streaming = _materialize_span_scalars(
+                codec_streaming_values,
+                codec_streaming_span_lengths,
+                dtype=torch.int8,
+                device=codec_streaming_device,
+            )
+            if codec_streaming is None:
+                codec_streaming = torch.cat(
+                    [
+                        torch.full((length,), value, dtype=torch.int8, device=codec_streaming_device)
+                        for value, length in zip(
+                            codec_streaming_values,
+                            codec_streaming_span_lengths,
+                            strict=True,
+                        )
+                    ],
+                    dim=0,
+                )
+            mm.setdefault("meta", {})["codec_streaming"] = codec_streaming[:span_len]
+        if codec_frame_valid_parts:
+            assert codec_frame_valid_device is not None
+            codec_frame_valid = torch.cat(codec_frame_valid_parts, dim=0)
+            mm.setdefault("meta", {})["codec_frame_valid"] = codec_frame_valid[:span_len]
         return OmniOutput(text_hidden_states=hidden, multimodal_outputs=mm)
 
     # -------------------- preprocess / postprocess --------------------
@@ -763,27 +919,32 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
 
         if is_prefill:
             # Prefill (prompt embeddings)
-            prompt_embeds_cpu = embed.get("prefill")
-            # First prefill round: prompt_embeds_cpu is not yet populated.
-            # Subsequent prefill rounds (multi-chunk): prompt_embeds_cpu is a Tensor stored by the first round.
-            is_first_prefill = not isinstance(prompt_embeds_cpu, torch.Tensor) or prompt_embeds_cpu.ndim != 2
+            prompt_embeds_buffer = embed.get("prefill")
+            # First prefill round: the retained prompt buffer is not populated.
+            # Subsequent chunks slice the buffer stored by the first round.
+            is_first_prefill = not isinstance(prompt_embeds_buffer, torch.Tensor) or prompt_embeds_buffer.ndim != 2
             if is_first_prefill:
                 full_prompt_embeds, tailing_text_hidden, ref_code_len, ref_code = (
                     self._prompt_builder.build_prompt_embeds(task_type=task_type, info_dict=info_dict)
                 )
-                # Store full prompt embeddings on CPU (large, prefill-only).
-                # tailing_text_hidden stays on GPU (gpu_resident_buffer_keys).
-                prompt_embeds_cpu = full_prompt_embeds.detach().to("cpu").contiguous()
+                prompt_embeds_buffer = full_prompt_embeds.detach()
+                if not bool(getattr(self, "_use_v2_model_runner", False)):
+                    prompt_embeds_buffer = prompt_embeds_buffer.to("cpu")
+                prompt_embeds_buffer = prompt_embeds_buffer.contiguous()
                 info_update: OmniPayload = {
-                    "embed": {"prefill": prompt_embeds_cpu},
+                    "embed": {"prefill": prompt_embeds_buffer},
                     "hidden_states": {"trailing_text": tailing_text_hidden.detach()},
                     "meta": {
                         "talker_text_offset": 0,
                         "codec_streaming": codec_streaming,
+                        "codec_frame_valid": torch.zeros((), dtype=torch.bool, device=input_ids.device),
                     },
                 }
                 if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
-                    info_update.setdefault("codes", {})["ref"] = ref_code.detach().to("cpu").contiguous()
+                    ref_code_buffer = ref_code.detach()
+                    if not bool(getattr(self, "_use_v2_model_runner", False)):
+                        ref_code_buffer = ref_code_buffer.to("cpu")
+                    info_update.setdefault("codes", {})["ref"] = ref_code_buffer.contiguous()
                 if ref_code_len is not None:
                     info_update["meta"]["ref_code_len"] = int(ref_code_len)
                 # First prefill: source the slice offset from `_omni_num_computed_tokens`
@@ -794,16 +955,21 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             else:
                 # Subsequent prefill chunk: slice from stored embeddings at running offset.
                 offset = max(0, int(meta.get("talker_prefill_offset", 0) or 0))
-                info_update = {"meta": {"codec_streaming": codec_streaming}}
+                info_update = {
+                    "meta": {
+                        "codec_streaming": codec_streaming,
+                        "codec_frame_valid": torch.zeros((), dtype=torch.bool, device=input_ids.device),
+                    }
+                }
 
             # Always return a span_len slice; if the scheduled placeholder is longer than what
             # the prompt actually fills, pad with tts_pad_embed (preserves placeholder/embedding alignment).
-            s = max(0, min(offset, int(prompt_embeds_cpu.shape[0])))
-            e = max(0, min(offset + span_len, int(prompt_embeds_cpu.shape[0])))
-            take = prompt_embeds_cpu[s:e]
+            s = max(0, min(offset, int(prompt_embeds_buffer.shape[0])))
+            e = max(0, min(offset + span_len, int(prompt_embeds_buffer.shape[0])))
+            take = prompt_embeds_buffer[s:e]
             if int(take.shape[0]) < span_len:
                 pad_n = int(span_len - int(take.shape[0]))
-                pad_rows = tts_pad_embed.reshape(1, -1).to("cpu").expand(pad_n, -1)
+                pad_rows = tts_pad_embed.to(device=take.device).reshape(1, -1).expand(pad_n, -1)
                 take = torch.cat([take, pad_rows], dim=0)
             prompt_embeds = take.to(device=input_ids.device, dtype=dtype)
             info_update["meta"]["talker_prefill_offset"] = int(offset + span_len)
@@ -826,7 +992,16 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 .to(device=input_ids.device, dtype=dtype)
                 .reshape(span_len, -1)
             )
-            return input_ids, inputs_embeds_out, {"meta": {"codec_streaming": codec_streaming}}
+            return (
+                input_ids,
+                inputs_embeds_out,
+                {
+                    "meta": {
+                        "codec_streaming": codec_streaming,
+                        "codec_frame_valid": torch.zeros((), dtype=torch.bool, device=input_ids.device),
+                    }
+                },
+            )
 
         # Decode: span_len == 1
         # Pop one text-step vector from tailing_text_hidden queue.
@@ -884,6 +1059,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             "meta": {
                 "talker_text_offset": int(next_text_offset),
                 "codec_streaming": codec_streaming,
+                "codec_frame_valid": (
+                    (input_ids.reshape(-1)[-1] >= 0) & (input_ids.reshape(-1)[-1] < self._codebook_vocab_size)
+                ).reshape(()),
             },
         }
         if trailing_text_update is not None:
@@ -895,11 +1073,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         *,
         input_ids: torch.Tensor,
         req_infos: list[dict[str, Any]],
+        input_embeds: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[dict[str, Any]]]:
         """Batch the decode-only preprocess path for Qwen3-TTS.
 
         This mirrors the scalar decode branch in ``preprocess()``, but performs
-        the token embedding lookup once for the whole decode batch.
+        the token embedding lookup once for the whole decode batch, unless
+        the runner has already prepared embeddings for these input ids.
         """
         input_ids_flat = input_ids.reshape(-1)
         if int(input_ids_flat.numel()) != len(req_infos):
@@ -907,6 +1087,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 f"preprocess_decode_batch expected {len(req_infos)} input ids, got {int(input_ids_flat.numel())}"
             )
 
+        codec_frame_valid_batch = (input_ids_flat >= 0) & (input_ids_flat < self._codebook_vocab_size)
         device = input_ids_flat.device
         dtype = self._embedding_dtype
         # Request-independent constant (see :meth:`_init_runtime_buffers`) —
@@ -979,17 +1160,16 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
                 "meta": {
                     "talker_text_offset": int(next_text_offset),
                     "codec_streaming": codec_streaming,
+                    "codec_frame_valid": codec_frame_valid_batch[len(updates)].reshape(()),
                 },
             }
             if trailing_text_update is not None:
                 info_update["hidden_states"] = {"trailing_text": trailing_text_update.detach()}
             updates.append(info_update)
 
-        inputs_embeds_out = self.embed_input_ids(input_ids_flat.reshape(-1, 1).to(torch.long)).to(
-            device=device,
-            dtype=dtype,
-        )
-        inputs_embeds_out = inputs_embeds_out.reshape(len(req_infos), -1)
+        if input_embeds is None:
+            input_embeds = self.embed_input_ids(input_ids_flat.reshape(-1, 1).to(torch.long))
+        inputs_embeds_out = input_embeds.to(device=device, dtype=dtype).reshape(len(req_infos), -1)
         return (
             input_ids_flat,
             inputs_embeds_out,
@@ -998,6 +1178,16 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             updates,
         )
 
+    def preprocess_decode_batch_mrv2(
+        self,
+        *,
+        input_ids: torch.Tensor,
+        input_embeds: torch.Tensor,
+        req_infos: list[dict[str, Any]],
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, list[dict[str, Any]]]:
+        """Expose the existing TTS decode batcher through the MRv2 contract."""
+        return self.preprocess_decode_batch(input_ids=input_ids, req_infos=req_infos, input_embeds=input_embeds)
+
     def postprocess(self, hidden_states: torch.Tensor, **_: Any) -> dict[str, Any]:
         # Keep the last token hidden for the next decode step's code predictor.
         # Stays on GPU - gpu_resident_buffer_keys avoids the CPU round-trip.
@@ -1006,6 +1196,23 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         last = hidden_states[-1, :].detach()
         return {"hidden_states": {"last": last}}
 
+    def postprocess_batch_mrv2(
+        self,
+        *,
+        hidden_states: torch.Tensor,
+        last_token_indices: torch.Tensor,
+    ) -> tuple[tuple[str, str], OwnedBatchTensor]:
+        """Gather every request's Talker tail in one GPU operation.
+
+        ``index_select`` freshly allocates the result, so the batch is marked
+        as owned and the intermediate buffer stores row views without taking
+        a second snapshot.
+        """
+        return (
+            ("hidden_states", "last"),
+            OwnedBatchTensor(hidden_states.index_select(0, last_token_indices)),
+        )
+
     @torch.inference_mode()
     def preprocess_batch(
         self,
@@ -1026,6 +1233,39 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             device=device,
         )
 
+    @torch.inference_mode()
+    def preprocess_batch_mrv2(
+        self,
+        *,
+        req_infos: list[dict[str, Any]],
+        device: torch.device,
+    ) -> None:
+        """Batch new MRv2 prefill slots without rebuilding a req-id map."""
+        if self._silence_ban_frames > 0:
+            for info in req_infos:
+                self._req_x_vector_only[info["req_id"]] = resolve_x_vector_only(info)
+        self._prompt_builder.preprocess_infos_batch(
+            req_infos=req_infos,
+            device=device,
+        )
+
+    @contextmanager
+    def mrv2_sampling_context(self, *, req_ids: list[str], num_output_tokens: torch.Tensor) -> Iterator[None]:
+        """Scope the silence gate to sampling, never prompt-logprob computation."""
+        self.set_batch_req_ids(req_ids)
+        self._mrv2_silence_ban_mask = None
+        if self._silence_ban_frames > 0:
+            in_scope = torch.tensor(
+                [self._req_x_vector_only.get(req_id) is True for req_id in req_ids],
+                device=num_output_tokens.device,
+                dtype=torch.bool,
+            )
+            self._mrv2_silence_ban_mask = (in_scope & (num_output_tokens < self._silence_ban_frames)).unsqueeze(1)
+        try:
+            yield
+        finally:
+            self._mrv2_silence_ban_mask = None
+
     def set_batch_req_ids(self, req_ids: Sequence[str]) -> None:
         """Record the current batch's req_ids, in the order that indexes logits.
 
@@ -1051,10 +1291,9 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             sampling_rate=target_sr,
             return_tensors="pt",
         )
-        inputs = inputs.to(device).to(torch.bfloat16)
-
-        input_values = inputs["input_values"].squeeze(1)
-        padding_mask = inputs["padding_mask"].squeeze(1)
+        input_values = inputs["input_values"].squeeze(1).to(device=device, dtype=torch.bfloat16)
+        downsample = self._encoder_downsample_rate
+        valid_code_lengths = [(int(wav.size) + downsample - 1) // downsample for wav in wavs]
 
         with torch.inference_mode():
             encoded = self.encoder.encode(
@@ -1063,13 +1302,13 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             )
 
         audio_codes = encoded.audio_codes[:, : self._encoder_valid_num_quantizers]
-        downsample = self._encoder_downsample_rate
         return [
-            code[..., : -(-mask.sum() // downsample)].transpose(0, 1).to(device=device, dtype=torch.long)
-            for code, mask in zip(audio_codes, padding_mask)
+            code[:, : int(valid_len)].transpose(0, 1).to(dtype=torch.long).contiguous()
+            for code, valid_len in zip(audio_codes, valid_code_lengths, strict=True)
         ]
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        self._prompt_builder._projected_token_cache.clear()
         # Consume talker weights, and conditionally consume speaker encoder
         # weights only if they are present in the checkpoint.
         speaker_weights: list[tuple[str, torch.Tensor]] = []
@@ -1256,6 +1495,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         top_p: float | None = None,
         generator: torch.Generator | None = None,
         generators: Sequence[torch.Generator | None] | None = None,
+        sample_uniforms: torch.Tensor | None = None,
         **kwargs: Any,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         """GPU fast-path used by OmniGPUModelRunner to predict residual codebooks (1..Q-1).
@@ -1296,6 +1536,7 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
             top_p=top_p,
             generator=generator,
             generators=generators,
+            sample_uniforms=sample_uniforms,
         )  # [B, Q]
 
         # Map invalid layer-0 ids (e.g. EOS) to PAD=0 so SpeechTokenizer sees only real codes.
@@ -1313,3 +1554,20 @@ class Qwen3TTSTalkerForConditionalGeneration(nn.Module):
         summed = (last_id_hidden.squeeze(1) + gathered.sum(dim=1)).unsqueeze(1)
         inputs_embeds_out = (summed + text_step).reshape(bsz, -1)
         return inputs_embeds_out, audio_codes.to(dtype=torch.long)
+
+    # MRV2 capability names read the V1 canonical values so platform patches
+    # such as the NPU 310P graph-safety override apply to both runners.
+    mtp = talker_mtp
+    get_mtp_seed = staticmethod(get_tts_local_seed)
+
+    @property
+    def mtp_output_key(self) -> tuple[str, str]:
+        return self.talker_mtp_output_key
+
+    @property
+    def mtp_graph_safe(self) -> bool:
+        return self.talker_mtp_graph_safe
+
+    @property
+    def mtp_accepts_per_row_generators(self) -> bool:
+        return self.talker_mtp_accepts_per_row_generators

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

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 # Copyright 2026 The Qwen team, Alibaba Group and the HuggingFace Inc. team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -28,7 +31,7 @@ from transformers import MimiConfig, MimiModel
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
 from transformers.integrations import use_kernel_forward_from_hub
-from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
+from transformers.masking_utils import create_sliding_window_causal_mask
 from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
 from transformers.modeling_layers import GradientCheckpointingLayer
 from transformers.modeling_outputs import BaseModelOutputWithPast
@@ -521,15 +524,63 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
         self.norm = Qwen3TTSTokenizerV2DecoderRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.rotary_emb = Qwen3TTSTokenizerV2DecoderRotatoryEmbedding(config=config)
         self.gradient_checkpointing = False
-        self.has_sliding_layers = "sliding_attention" in self.config.layer_types
         self.window_size = config.sliding_window
 
         self.input_proj = nn.Linear(config.latent_dim, config.hidden_size)
         self.output_proj = nn.Linear(config.hidden_size, config.latent_dim)
+        # Stateless calls with implicit positions share one broadcastable mask
+        # per warmup shape. Incremental or explicit-position calls build a live mask.
+        self._sliding_attention_mask_cache: dict[tuple[int, torch.dtype, torch.device, str], torch.Tensor | None] = {}
 
         # Initialize weights and apply final processing
         self.post_init()
 
+    def _apply(self, fn, recurse=True):
+        # A placement or dtype change invalidates every cached mask.
+        self._sliding_attention_mask_cache.clear()
+        return super()._apply(fn, recurse=recurse)
+
+    def _get_sliding_attention_mask(
+        self,
+        inputs_embeds: torch.Tensor,
+    ) -> torch.Tensor | None:
+        sequence_length = int(inputs_embeds.shape[1])
+        max_position_embeddings = int(self.config.max_position_embeddings)
+        if sequence_length > max_position_embeddings:
+            raise ValueError(
+                f"Input length {sequence_length} exceeds max_position_embeddings {max_position_embeddings}"
+            )
+
+        cache_key = (
+            sequence_length,
+            inputs_embeds.dtype,
+            inputs_embeds.device,
+            str(self.config._attn_implementation),
+        )
+        if cache_key not in self._sliding_attention_mask_cache:
+            mask_inputs = inputs_embeds[:1]
+            mask_kwargs = {
+                "config": self.config,
+                "attention_mask": None,
+                "past_key_values": None,
+                "position_ids": None,
+            }
+            sig = inspect.signature(create_sliding_window_causal_mask)
+            if "input_embeds" in sig.parameters:
+                mask_kwargs["input_embeds"] = mask_inputs
+                mask_kwargs["cache_position"] = torch.arange(
+                    sequence_length,
+                    device=inputs_embeds.device,
+                )
+            else:
+                mask_kwargs["inputs_embeds"] = mask_inputs
+            sliding_attention_mask = create_sliding_window_causal_mask(**mask_kwargs)
+            if sliding_attention_mask is not None:
+                sliding_attention_mask = sliding_attention_mask.contiguous()
+            self._sliding_attention_mask_cache[cache_key] = sliding_attention_mask
+
+        return self._sliding_attention_mask_cache[cache_key]
+
     # Note: @check_model_inputs decorator removed for vLLM compatibility
     # The decorator causes "unexpected keyword argument 'inputs_embeds'" error
     @auto_docstring
@@ -558,6 +609,17 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
 
         inputs_embeds = self.input_proj(inputs_embeds)
 
+        sequence_length = inputs_embeds.shape[1]
+        # Only internally generated positions are known to be contiguous and
+        # zero-based without reading CUDA tensor values back to the host.
+        # Explicit positions use the live mask path, including during capture.
+        uses_prepared_or_incremental_mask = (
+            attention_mask is not None
+            or bool(use_cache)
+            or past_key_values is not None
+            or cache_position is not None
+            or position_ids is not None
+        )
         if use_cache and past_key_values is None:
             past_key_values = DynamicCache(config=self.config)
 
@@ -565,7 +627,7 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
             past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
             cache_position = torch.arange(
                 past_seen_tokens,
-                past_seen_tokens + inputs_embeds.shape[1],
+                past_seen_tokens + sequence_length,
                 device=inputs_embeds.device,
             )
 
@@ -574,22 +636,24 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
 
         hidden_states = inputs_embeds
 
-        if not isinstance(causal_mask_mapping := attention_mask, dict):
+        if isinstance(attention_mask, dict):
+            sliding_attention_mask = attention_mask["sliding_attention"]
+        elif uses_prepared_or_incremental_mask:
             mask_kwargs = {
                 "config": self.config,
                 "attention_mask": attention_mask,
                 "past_key_values": past_key_values,
                 "position_ids": position_ids,
             }
-            sig = inspect.signature(create_causal_mask)
+            sig = inspect.signature(create_sliding_window_causal_mask)
             if "input_embeds" in sig.parameters:
                 mask_kwargs["input_embeds"] = inputs_embeds
                 mask_kwargs["cache_position"] = cache_position
             else:
                 mask_kwargs["inputs_embeds"] = inputs_embeds
-            causal_mask_mapping = {"full_attention": create_causal_mask(**mask_kwargs)}
-            if self.has_sliding_layers:
-                causal_mask_mapping["sliding_attention"] = create_sliding_window_causal_mask(**mask_kwargs)
+            sliding_attention_mask = create_sliding_window_causal_mask(**mask_kwargs)
+        else:
+            sliding_attention_mask = self._get_sliding_attention_mask(inputs_embeds)
 
         # create position embeddings to be shared across the decoder layers
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
@@ -597,7 +661,7 @@ class Qwen3TTSTokenizerV2DecoderTransformerModel(Qwen3TTSTokenizerV2DecoderPreTr
         for decoder_layer in self.layers[: self.config.num_hidden_layers]:
             hidden_states = decoder_layer(
                 hidden_states,
-                attention_mask=causal_mask_mapping[decoder_layer.attention_type],
+                attention_mask=sliding_attention_mask,
                 position_ids=position_ids,
                 past_key_values=past_key_values,
                 use_cache=use_cache,
@@ -1400,13 +1464,15 @@ class Qwen3TTSTokenizerV2Decoder(Qwen3TTSTokenizerV2DecoderPreTrainedModel):
 
     @staticmethod
     def _slice_dynamic_cache(cache: DynamicCache, row: int) -> DynamicCache:
-        request_cache = copy.deepcopy(cache)
-        for layer in request_cache.layers:
-            if layer.keys is not None:
-                layer.keys = layer.keys[row : row + 1].clone()
-            if layer.values is not None:
-                layer.values = layer.values[row : row + 1].clone()
-        return request_cache
+        # Preserve cache metadata and independent ownership without cloning the
+        # full batch for every request. Seed deepcopy with only the selected KV
+        # rows; it will copy the remaining cache structure normally.
+        memo: dict[int, Any] = {}
+        for layer in cache.layers:
+            for tensor in (layer.keys, layer.values):
+                if tensor is not None and id(tensor) not in memo:
+                    memo[id(tensor)] = tensor[row : row + 1].clone()
+        return copy.deepcopy(cache, memo)
 
     @staticmethod
     def _cache_tensors_are_batchable(request_caches: list[dict[str, Any]], keys: tuple[str, ...]) -> bool:

vllm_omni/model_executor/stage_input_processors/qwen3_tts.py

@@ -70,7 +70,7 @@ def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torc
         return None
     if audio_codes.ndim == 2:
         frame = audio_codes[-1]
-        if frame.numel() == 0 or not bool(frame.any().item()):
+        if frame.numel() == 0:
             return None
         return frame.to(torch.long).reshape(-1)
     if audio_codes.ndim == 1:
@@ -78,6 +78,31 @@ def _extract_last_frame(multimodal_output: OmniPayload | dict[str, Any]) -> torc
     raise ValueError(f"Invalid audio_codes shape for Qwen3-TTS async_chunk: {tuple(audio_codes.shape)}")
 
 
+def _should_append_async_frame(
+    multimodal_output: OmniPayload | dict[str, Any],
+    frame: torch.Tensor,
+) -> bool:
+    meta = multimodal_output.get("meta", {})
+    frame_valid = meta.get("codec_frame_valid") if isinstance(meta, Mapping) else None
+    if isinstance(frame_valid, torch.Tensor):
+        if frame_valid.numel() == 0:
+            raise ValueError("codec_frame_valid must not be empty")
+        # A single request can carry a multi-token prefill span. The model
+        # emits one validity entry per codec row, while _extract_last_frame()
+        # selects the final row from that span.
+        return bool(frame_valid.reshape(-1)[-1].item())
+    if frame_valid is not None:
+        return bool(frame_valid)
+    # Compatibility for payload producers predating the explicit validity
+    # contract. Real Qwen3-TTS outputs carry codec_frame_valid and avoid this
+    # value-dependent synchronization.
+    return bool(frame.any().item())
+
+
+def _append_tensor_frame(transfer_manager: Any, request_id: str, frame: torch.Tensor) -> None:
+    transfer_manager.code_prompt_token_ids[request_id].append(frame.detach().to(torch.long).reshape(-1).contiguous())
+
+
 def talker2code2wav_async_chunk(
     transfer_manager: Any,
     multimodal_output: OmniPayload | dict[str, Any] | None,
@@ -93,12 +118,11 @@ def talker2code2wav_async_chunk(
 
     if isinstance(multimodal_output, Mapping):
         frame = _extract_last_frame(multimodal_output)
-        if frame is not None:
-            codec_codes = frame.cpu().tolist()
-            transfer_manager.code_prompt_token_ids[request_id].append(codec_codes)
+        if frame is not None and _should_append_async_frame(multimodal_output, frame):
+            _append_tensor_frame(transfer_manager, request_id, frame)
         ref_code = multimodal_output.get("codes", {}).get("ref")
         if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0 and request_payload.get(request_id) is None:
-            request_payload[request_id] = ref_code.to(torch.long).cpu().contiguous()
+            request_payload[request_id] = ref_code.detach().to(torch.long).contiguous()
     elif not finished:
         return None
 
@@ -183,6 +207,17 @@ def talker2code2wav_async_chunk(
         )
         initial_chunk_size = chunk_size
     length = len(transfer_manager.code_prompt_token_ids[request_id])
+    emitted_frames = getattr(transfer_manager, "_qwen3_tts_emitted_frames", None)
+    if emitted_frames is None:
+        emitted_frames = {}
+        transfer_manager._qwen3_tts_emitted_frames = emitted_frames
+    if request_id in emitted_frames and length <= emitted_frames[request_id]:
+        if finished:
+            return OmniPayloadStruct(
+                codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+                meta=MetaStruct(request_id=request_id, left_context_size=0, finished=torch.tensor(True)),
+            )
+        return None
 
     # Full-utterance Code2Wav is gated by ``full_utterance_decode`` (explicit
     # opt-in via additional_information). Do not reuse ``non_streaming_mode``:
@@ -305,6 +340,8 @@ def talker2code2wav_async_chunk(
 
     if finished and first_chunk:
         context_length = length
+    if request_id in emitted_frames:
+        context_length = min(context_length, length - emitted_frames[request_id])
 
     if adaptive_enabled:
         ctrl = transfer_manager._adaptive_states.get(request_id)
@@ -343,7 +380,7 @@ def talker2code2wav_async_chunk(
     if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0:
         ref_context = ref_code
         if ref_code_context_frames > 0 and int(ref_context.shape[0]) > ref_code_context_frames:
-            logger.info_once(
+            logger.debug(
                 "Qwen3-TTS async chunk uses the last %d/%d ref_code frames as bounded Code2Wav context.",
                 ref_code_context_frames,
                 int(ref_context.shape[0]),
@@ -375,6 +412,7 @@ def talker2code2wav_async_chunk(
         meta.ref_context_request_id = ref_context_request_id
         meta.ref_context_included = ref_context_included
 
+    emitted_frames[request_id] = length
     return OmniPayloadStruct(
         codes=CodesStruct(audio=code_predictor_codes),
         meta=meta,
@@ -383,6 +421,91 @@ def talker2code2wav_async_chunk(
     )
 
 
+def _copy_async_chunk_tensors_to_cpu(
+    tensors: list[torch.Tensor],
+) -> list[torch.Tensor]:
+    """Materialize one producer cohort with at most one D2H per device."""
+    if not tensors:
+        return []
+
+    outputs: list[torch.Tensor | None] = [None] * len(tensors)
+    by_device: dict[torch.device, list[tuple[int, torch.Tensor]]] = {}
+    for index, tensor in enumerate(tensors):
+        value = tensor.detach().to(dtype=torch.long).reshape(-1).contiguous()
+        if value.device.type == "cpu":
+            outputs[index] = value
+        else:
+            by_device.setdefault(value.device, []).append((index, value))
+
+    for entries in by_device.values():
+        lengths = [int(value.numel()) for _, value in entries]
+        combined = torch.cat([value for _, value in entries], dim=0).to(device="cpu")
+        offset = 0
+        for (index, source), length in zip(entries, lengths):
+            outputs[index] = combined[offset : offset + length].reshape(source.shape).contiguous()
+            offset += length
+
+    return [value for value in outputs if value is not None]
+
+
+def talker2code2wav_async_chunk_batch(
+    transfer_manager: Any,
+    pooling_outputs: list[OmniPayload | dict[str, Any] | None],
+    requests: list[Any],
+    is_finished: list[bool],
+) -> list[OmniPayloadStruct | None]:
+    """Build one Talker model-step worth of Code2Wav payloads.
+
+    MRv2 invokes this hook from its ordered native output worker. Codec frames
+    and first-chunk reference codes from the same model step are copied as a
+    cohort, while chunk boundaries and request metadata remain delegated to the
+    scalar builder below.
+    """
+    if not (len(pooling_outputs) == len(requests) == len(is_finished)):
+        raise ValueError("batch codec inputs must have identical lengths")
+
+    request_payload = getattr(transfer_manager, "request_payload", None)
+    if request_payload is None:
+        request_payload = {}
+        transfer_manager.request_payload = request_payload
+
+    selected: list[torch.Tensor] = []
+    destinations: list[tuple[str, str, tuple[int, ...]]] = []
+    for pooling_output, request in zip(pooling_outputs, requests):
+        if not isinstance(pooling_output, Mapping):
+            continue
+        request_id = request.external_req_id
+        frame = _extract_last_frame(pooling_output)
+        if frame is not None and _should_append_async_frame(pooling_output, frame):
+            selected.append(frame)
+            destinations.append(("frame", request_id, tuple(frame.shape)))
+
+        ref_code = pooling_output.get("codes", {}).get("ref")
+        if isinstance(ref_code, torch.Tensor) and ref_code.numel() > 0 and request_payload.get(request_id) is None:
+            selected.append(ref_code)
+            destinations.append(("ref", request_id, tuple(ref_code.shape)))
+
+    copied = _copy_async_chunk_tensors_to_cpu(selected)
+    for value, (kind, request_id, shape) in zip(copied, destinations):
+        value = value.reshape(shape)
+        if kind == "frame":
+            _append_tensor_frame(transfer_manager, request_id, value)
+        else:
+            request_payload[request_id] = value.contiguous()
+
+    payloads: list[OmniPayloadStruct | None] = []
+    for request, finished in zip(requests, is_finished):
+        payloads.append(
+            talker2code2wav_async_chunk(
+                transfer_manager=transfer_manager,
+                multimodal_output={},
+                request=request,
+                is_finished=finished,
+            )
+        )
+    return payloads
+
+
 # ============================================================================
 # Worker-connector data plane (non-async-chunk path).
 # AR runner's `flatten_payload` converts the model emit

vllm_omni/outputs/__init__.py

@@ -60,6 +60,10 @@ class OmniModelRunnerOutput(ModelRunnerOutput):
     # The Scheduler can safely free the block tables for these requests.
     kv_extracted_req_ids: list[str] | None = None
     omni_connector_output: OmniConnectorOutput | None = None
+    # True when sampled_token_ids has already been materialized on the host.
+    # MRv2 consumers must treat the value as immutable and must not perform
+    # another device-to-host conversion or rebuild it from sampler tensors.
+    sampled_token_ids_materialized: bool = False
 
     @classmethod
     def with_kv_conn_output_only(cls, kv_connector_output: Any) -> "OmniModelRunnerOutput":

vllm_omni/outputs/output_processor.py

@@ -684,11 +684,19 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
 
             # Accumulate multimodal tensors regardless of path.
             if isinstance(req_state, OmniRequestState):
+                mm_type = getattr(eco, "output_type", None) or default_mm_type
+
                 mm_output = getattr(eco, "multimodal_output", None)
                 if mm_output is not None:
-                    mm_type = getattr(eco, "output_type", None) or default_mm_type
                     req_state.add_multimodal_tensor(mm_output, mm_type)
 
+                # Omni AR stages carry model payloads alongside text tokens in
+                # pooling_output. Preserve the payload, then clear the field so
+                # upstream still runs the normal detokenization path.
+                if eco.pooling_output is not None and req_state.detokenizer is not None:
+                    req_state.add_multimodal_tensor(eco.pooling_output, mm_type)
+                    eco.pooling_output = None
+
             # Route: if no detokenizer and no pooling output, handle locally
             # to avoid upstream's assert on detokenizer.
             if req_state.detokenizer is None and eco.pooling_output is None:
@@ -697,20 +705,22 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
                 upstream_outputs.append(eco)
 
         # Handle multimodal-only outputs (generation stages) locally.
-        mm_request_outputs = self._process_mm_only_outputs(
+        mm_request_outputs, mm_reqs_to_abort = self._process_mm_only_outputs(
             mm_only_outputs,
             engine_core_timestamp=engine_core_timestamp,
             iteration_stats=iteration_stats,
         )
 
         # Delegate text/pooling outputs to upstream.
-        processed = super().process_outputs(
+        upstream_processed = super().process_outputs(
             upstream_outputs,
             engine_core_timestamp=engine_core_timestamp,
             iteration_stats=iteration_stats,
         )
-        processed.request_outputs.extend(mm_request_outputs)
-        return processed
+        return OutputProcessorOutput(
+            request_outputs=mm_request_outputs + upstream_processed.request_outputs,
+            reqs_to_abort=mm_reqs_to_abort + upstream_processed.reqs_to_abort,
+        )
 
     def _process_mm_only_outputs(
         self,
@@ -718,30 +728,32 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
         *,
         engine_core_timestamp: float | None = None,
         iteration_stats: IterationStats | None = None,
-    ) -> list[OmniRequestOutput | PoolingRequestOutput]:
+    ) -> tuple[list[OmniRequestOutput | PoolingRequestOutput], list[str]]:
         """Handle outputs from generation stages that have no detokenizer.
 
         These cannot go through upstream process_outputs because it asserts
         detokenizer is not None when pooling_output is None.
         """
         request_outputs: list[OmniRequestOutput | PoolingRequestOutput] = []
+        reqs_to_abort: list[str] = []
         for eco in engine_core_outputs:
             req_state = self.request_states.get(eco.request_id)
             if req_state is None or not isinstance(req_state, OmniRequestState):
                 continue
 
-            new_token_ids = eco.new_token_ids
-            finish_reason = eco.finish_reason
-            stop_reason = eco.stop_reason
-            kv_transfer_params = eco.kv_transfer_params
-            ec_transfer_params = getattr(eco, "ec_transfer_params", None)
-            routed_experts = eco.routed_experts
             self._update_stats_from_output(
                 req_state,
                 eco,
                 engine_core_timestamp,
                 iteration_stats,
             )
+
+            new_token_ids = eco.new_token_ids
+            finish_reason = eco.finish_reason
+            stop_reason = eco.stop_reason
+            kv_transfer_params = eco.kv_transfer_params
+            ec_transfer_params = eco.ec_transfer_params
+            routed_experts = eco.routed_experts
             prefill_stats = getattr(eco, "prefill_stats", None)
             if prefill_stats is not None:
                 req_state.num_cached_tokens = prefill_stats.num_cached_tokens
@@ -770,12 +782,17 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             is_segment_finished = bool(getattr(eco, "is_segment_finished", False))
             if finish_reason is not None and not is_segment_finished and not is_non_final_audio_chunk:
                 self._finish_request(req_state)
+                if not getattr(eco, "finished", True):
+                    reqs_to_abort.append(eco.request_id)
                 self._update_stats_from_finished(
                     req_state,
                     finish_reason,
                     iteration_stats,
                 )
-        return request_outputs
+                if self.tracing_enabled:
+                    self.do_tracing(eco, req_state, iteration_stats)
+
+        return request_outputs, reqs_to_abort
 
     def _update_stats_from_output(
         self,
@@ -804,6 +821,10 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             None,
         )
 
+        reported_count = getattr(engine_core_output, "num_generation_tokens", None)
+        if isinstance(reported_count, int) and not isinstance(reported_count, bool) and reported_count >= 0:
+            self._native_text_metric_record(req_state.external_req_id)["num_generation_tokens"] = reported_count
+
         if iteration_stats is None or engine_core_timestamp is None or native_stats is None:
             return
 
@@ -816,7 +837,11 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
             req_state.lora_name,
         )
         record = self._native_text_metric_record(req_state.external_req_id)
-        record["num_generation_tokens"] = int(native_stats.num_generation_tokens)
+        record["num_generation_tokens"] = (
+            reported_count
+            if isinstance(reported_count, int) and not isinstance(reported_count, bool) and reported_count >= 0
+            else int(native_stats.num_generation_tokens)
+        )
         if was_prefilling:
             record["vllm_ttft_ms"] = max(float(native_stats.first_token_latency) * 1000.0, 0.0)
             return

vllm_omni/worker/base.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 """Base worker class for vLLM-Omni with device-level GPU memory profiling."""
 
 from __future__ import annotations

vllm_omni/worker/gpu_ar_worker.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import gc
 import os
 
@@ -75,7 +78,6 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
                     f"DP adjusted local rank {self.local_rank} is out of "
                     f"bounds for {torch.accelerator.device_count()} devices."
                 )
-
             visible_device_index = current_platform.logical_device_id_to_visible_device_id(self.local_rank)
             self.device = current_omni_platform.get_torch_device(visible_device_index)
             torch.accelerator.set_device_index(self.device)
@@ -113,13 +115,24 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase):
         num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
         init_workspace_manager(self.device, num_ubatches)
 
-        if self.use_v2_model_runner:
-            # OMNI: v2 model runner does not yet include omni hooks.
-            logger.warning("OMNI GPUARWorker forces v1 model runner for omni hooks.")
-            self.use_v2_model_runner = False
+        model_stage = getattr(self.model_config, "model_stage", None)
+        stage_id = getattr(self.model_config, "stage_id", None)
 
-        # Construct the model runner
-        self.model_runner = self.model_runner_cls(self.vllm_config, self.device)
+        self.use_v2_model_runner = bool(getattr(self.model_config, "use_v2_model_runner", False))
+        if self.use_v2_model_runner:
+            from vllm_omni.worker_v2.omni_ar_model_runner import (
+                OmniARModelRunner,
+            )
+
+            logger.info(
+                "Using MR v2 OmniARModelRunner for omni AR stage (stage_id=%s model_stage=%s).",
+                stage_id,
+                model_stage,
+            )
+            self.model_runner = OmniARModelRunner(self.vllm_config, self.device)
+        else:
+            self.use_v2_model_runner = False
+            self.model_runner = self.model_runner_cls(self.vllm_config, self.device)
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.

vllm_omni/worker/gpu_generation_worker.py

@@ -1,3 +1,6 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
 import gc
 import os
 
@@ -8,7 +11,10 @@ from vllm.tracing import instrument
 from vllm.utils.mem_utils import MemorySnapshot, format_gib
 from vllm.utils.torch_utils import set_random_seed
 from vllm.v1.utils import report_usage_stats
-from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
+from vllm.v1.worker.gpu_worker import (
+    CompilationTimes,
+    init_worker_distributed_environment,
+)
 from vllm.v1.worker.workspace import init_workspace_manager
 
 from vllm_omni.platforms import current_omni_platform
@@ -96,13 +102,28 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase):
         num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
         init_workspace_manager(self.device, num_ubatches)
 
+        self.use_v2_model_runner = bool(getattr(self.model_config, "use_v2_model_runner", False))
         if self.use_v2_model_runner:
-            # OMNI: v2 model runner does not yet include omni hooks.
-            logger.warning("OMNI GPUGenerationWorker forces v1 model runner for omni hooks.")
-            self.use_v2_model_runner = False
+            from vllm_omni.worker_v2.omni_generation_model_runner import (
+                OmniGenerationModelRunner,
+            )
 
-        self.model_runner = self.model_runner_cls(self.vllm_config, self.device)
+            logger.info("Using MR v2 OmniGenerationModelRunner for generation stage.")
+            self.model_runner = OmniGenerationModelRunner(self.vllm_config, self.device)
+        else:
+            self.model_runner = self.model_runner_cls(self.vllm_config, self.device)
 
         if self.rank == 0:
             # If usage stat is enabled, collect relevant info.
             report_usage_stats(self.vllm_config)
+
+    @instrument(span_name="Compile/warmup")
+    def compile_or_warm_up_model(self) -> CompilationTimes:
+        """Generation stages have no KV cache or sampler — skip warmup_kernels."""
+        if not self.use_v2_model_runner:
+            return super().compile_or_warm_up_model()
+        import time
+
+        start = time.perf_counter()
+        self.model_runner.profile_run()
+        return CompilationTimes(language_model=time.perf_counter() - start, encoder=0.0)

vllm_omni/worker/gpu_model_runner.py

@@ -2042,10 +2042,20 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner):
                     if req_state is not None:
                         existing = self.model_intermediate_buffer.setdefault(req_id, {})
                         existing.setdefault(out_key[0], {})[out_key[1]] = row
+                        validity_key = getattr(self.model, "talker_mtp_validity_key", None)
+                        if validity_key is not None:
+                            existing.setdefault(validity_key[0], {})[validity_key[1]] = torch.ones(
+                                (), dtype=torch.bool, device=code_predictor_codes.device
+                            )
                         req_state.additional_information_cpu = existing
             else:
                 for idx, req_id in enumerate(decode_req_ids):
                     update_dict = {out_key[0]: {out_key[1]: code_predictor_codes[idx : idx + 1]}}
+                    validity_key = getattr(self.model, "talker_mtp_validity_key", None)
+                    if validity_key is not None:
+                        update_dict.setdefault(validity_key[0], {})[validity_key[1]] = torch.ones(
+                            (), dtype=torch.bool, device=code_predictor_codes.device
+                        )
                     self._update_intermediate_buffer(req_id, update_dict)
 
     def _model_forward(

vllm_omni/worker/payload_span.py

@@ -0,0 +1,48 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Helpers for explicit thinker decode span metadata."""
+
+from collections.abc import Mapping
+from typing import Any
+
+import torch
+
+TensorSpan = tuple[torch.Tensor, int, int]
+
+
+def get_tensor_span(payload: Mapping[str, Any], *, tensor_key: str, start_key: str, end_key: str) -> TensorSpan | None:
+    tensor = payload.get(tensor_key)
+    start = payload.get(start_key)
+    end = payload.get(end_key)
+    if not isinstance(tensor, torch.Tensor):
+        return None
+    if not isinstance(start, int) or not isinstance(end, int):
+        return None
+    if start < 0 or end < start or (end - start) != int(tensor.shape[0]):
+        return None
+    return tensor, start, end
+
+
+def merge_tensor_spans(existing_span: TensorSpan | None, incoming_span: TensorSpan | None) -> TensorSpan | None:
+    if existing_span is None or incoming_span is None:
+        return None
+
+    existing_tensor, existing_start, existing_end = existing_span
+    incoming_tensor, incoming_start, incoming_end = incoming_span
+    if incoming_tensor.device != existing_tensor.device or incoming_tensor.dtype != existing_tensor.dtype:
+        incoming_tensor = incoming_tensor.to(device=existing_tensor.device, dtype=existing_tensor.dtype)
+    if incoming_start == existing_end:
+        if existing_start == existing_end:
+            return incoming_tensor, incoming_start, incoming_end
+        return torch.cat([existing_tensor, incoming_tensor], dim=0), existing_start, incoming_end
+    if incoming_start < existing_end:
+        overlap = existing_end - incoming_start
+        if overlap >= int(incoming_tensor.shape[0]):
+            return existing_tensor, existing_start, existing_end
+        trimmed_tensor = incoming_tensor[overlap:]
+        return (
+            torch.cat([existing_tensor, trimmed_tensor], dim=0),
+            existing_start,
+            existing_end + int(trimmed_tensor.shape[0]),
+        )
+    return None

vllm_omni/worker/sampling_utils.py

@@ -6,18 +6,72 @@ from typing import Any
 
 import torch
 from vllm.logger import init_logger
+from vllm.sampling_params import SamplingParams
 from vllm.v1.sample.logits_processor import LogitsProcessors, MinTokensLogitsProcessor
 
 logger = init_logger(__name__)
 
 __all__ = [
-    "build_model_sampler_extra_args",
-    "call_model_sampler",
+    "apply_fixed_seed_to_sampling_params",
+    "get_tts_local_seed",
     "clamp_prompt_ids_to_penalty_padding",
     "sanitize_min_tokens_stop_ids",
+    "sanitize_sampling_params_min_tokens_stop_ids",
+    "build_model_sampler_extra_args",
+    "call_model_sampler",
 ]
 
 
+def get_tts_local_seed(sampling_params: Any) -> int | None:
+    """Return the explicit per-request Talker MTP seed, if configured."""
+    extra_args = getattr(sampling_params, "extra_args", None) if sampling_params is not None else None
+    seed = extra_args.get("tts_local_seed") if isinstance(extra_args, dict) else None
+    return int(seed) if seed is not None else None
+
+
+def apply_fixed_seed_to_sampling_params(
+    sampling_params: Any,
+    seed: int,
+    *,
+    seed_talker_mtp: bool,
+) -> None:
+    """Apply one request seed to a stage sampler and, for Talkers, its MTP RNG."""
+    seed = int(seed)
+    sampling_params.seed = seed
+    if not seed_talker_mtp:
+        return
+    extra_args = dict(getattr(sampling_params, "extra_args", None) or {})
+    extra_args["tts_local_seed"] = seed
+    sampling_params.extra_args = extra_args
+
+
+def sanitize_sampling_params_min_tokens_stop_ids(
+    sampling_params: SamplingParams,
+    logits_vocab: int,
+) -> None:
+    """Remove unreachable ids before MRv2 builds persistent min-token state.
+
+    MRv2 copies ``SamplingParams.all_stop_token_ids`` into GPU-resident
+    ``LogitBiasState`` when a request is registered. Narrow codec heads cannot
+    safely carry the text tokenizer EOS in that state. The engine-facing EOS
+    and explicit ``stop_token_ids`` fields remain unchanged; only the set used
+    by the sampler's min-token mask is normalized.
+    """
+    if sampling_params.min_tokens <= 0:
+        return
+    stop_token_ids = sampling_params.all_stop_token_ids
+    unreachable = {token_id for token_id in stop_token_ids if token_id < 0 or token_id >= logits_vocab}
+    if not unreachable:
+        return
+    stop_token_ids.difference_update(unreachable)
+    logger.warning_once(
+        "min_tokens: dropped stop token ids %s that exceed the logits vocabulary (%d); "
+        "the model head cannot emit them.",
+        str(sorted(unreachable)),
+        logits_vocab,
+    )
+
+
 def build_model_sampler_extra_args(input_batch: Any, requests: Any) -> list[dict | None]:
     """Return per-request ``SamplingParams.extra_args`` in batch-row order."""
     request_states = requests or {}

vllm_omni/worker_v2/delivery.py

@@ -0,0 +1,218 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import threading
+import time
+from enum import Enum, auto
+
+
+class DeliveryState(Enum):
+    QUEUED = auto()
+    IN_FLIGHT = auto()
+    DELIVERED = auto()
+    FAILED = auto()
+    CANCELLED = auto()
+
+
+class DeliveryTimeoutError(TimeoutError):
+    pass
+
+
+class DeliveryCancelledError(RuntimeError):
+    pass
+
+
+_TERMINAL_STATES = {
+    DeliveryState.DELIVERED,
+    DeliveryState.FAILED,
+    DeliveryState.CANCELLED,
+}
+
+
+class DeliveryTicket:
+    """Exactly-once completion for one MRv2 connector delivery."""
+
+    def __init__(
+        self,
+        manager: OmniDeliveryManager,
+        *,
+        request_id: str,
+        put_key: str,
+        deadline: float,
+    ) -> None:
+        self._manager = manager
+        self.request_id = request_id
+        self.put_key = put_key
+        self.deadline = deadline
+        self._state = DeliveryState.QUEUED
+        self._error: BaseException | None = None
+        self._event = threading.Event()
+
+    @property
+    def state(self) -> DeliveryState:
+        with self._manager._lock:
+            return self._state
+
+    def mark_in_flight(self) -> bool:
+        return self._manager._mark_in_flight(self)
+
+    def set_delivered(self) -> bool:
+        return self._manager._set_delivered(self)
+
+    def set_failed(self, error: BaseException) -> bool:
+        return self._manager._set_failed(self, error)
+
+    # Compatibility with the connector mixin's completion interface.
+    def set_result(self) -> bool:
+        return self.set_delivered()
+
+    def set_error(self, error: BaseException) -> bool:
+        return self.set_failed(error)
+
+    def wait(self) -> None:
+        while True:
+            with self._manager._lock:
+                state = self._state
+                error = self._error
+                remaining = self.deadline - time.monotonic()
+            if state is DeliveryState.DELIVERED:
+                return
+            if state in (DeliveryState.FAILED, DeliveryState.CANCELLED):
+                assert error is not None
+                raise error
+            if remaining <= 0:
+                self._manager._timeout(self)
+                continue
+            self._event.wait(timeout=remaining)
+
+
+class OmniDeliveryManager:
+    """MRv2-only connector delivery state machine.
+
+    A timeout or permanent delivery failure quarantines the whole manager.
+    This is a stage-fatal condition: no new payload may be accepted after the
+    transport can no longer prove ordered, exactly-once completion.
+    """
+
+    def __init__(
+        self,
+        *,
+        delivery_timeout_s: float,
+        shutdown_timeout_s: float = 5.0,
+    ) -> None:
+        if delivery_timeout_s <= 0:
+            raise ValueError("delivery_timeout_s must be positive")
+        if shutdown_timeout_s <= 0:
+            raise ValueError("shutdown_timeout_s must be positive")
+        self.delivery_timeout_s = float(delivery_timeout_s)
+        self.shutdown_timeout_s = float(shutdown_timeout_s)
+        self._lock = threading.Lock()
+        self._tickets: set[DeliveryTicket] = set()
+        self._quarantine_error: BaseException | None = None
+        self._closed = False
+
+    @property
+    def is_quarantined(self) -> bool:
+        with self._lock:
+            return self._quarantine_error is not None
+
+    def create_ticket(self, *, request_id: str, put_key: str) -> DeliveryTicket:
+        with self._lock:
+            if self._closed:
+                raise RuntimeError("MRv2 delivery manager is closed")
+            if self._quarantine_error is not None:
+                raise RuntimeError("MRv2 delivery manager is quarantined") from self._quarantine_error
+            ticket = DeliveryTicket(
+                self,
+                request_id=request_id,
+                put_key=put_key,
+                deadline=time.monotonic() + self.delivery_timeout_s,
+            )
+            self._tickets.add(ticket)
+            return ticket
+
+    def quarantine(self, error: BaseException) -> None:
+        with self._lock:
+            self._quarantine_locked(error, primary_ticket=None)
+
+    def shutdown(self, cause: BaseException | None = None) -> None:
+        with self._lock:
+            if self._closed:
+                return
+            self._closed = True
+            if self._quarantine_error is None:
+                self._quarantine_error = cause or RuntimeError("MRv2 stage shutdown")
+            for ticket in tuple(self._tickets):
+                if ticket._state in _TERMINAL_STATES:
+                    continue
+                message = f"connector delivery cancelled for request {ticket.request_id!r}, key {ticket.put_key!r}"
+                if cause is not None:
+                    message += f": {cause}"
+                self._finish_locked(
+                    ticket,
+                    DeliveryState.CANCELLED,
+                    DeliveryCancelledError(message),
+                )
+
+    def _mark_in_flight(self, ticket: DeliveryTicket) -> bool:
+        with self._lock:
+            if ticket._state is not DeliveryState.QUEUED:
+                return False
+            ticket._state = DeliveryState.IN_FLIGHT
+            return True
+
+    def _set_delivered(self, ticket: DeliveryTicket) -> bool:
+        with self._lock:
+            if ticket._state in _TERMINAL_STATES:
+                return False
+            return self._finish_locked(ticket, DeliveryState.DELIVERED, None)
+
+    def _set_failed(self, ticket: DeliveryTicket, error: BaseException) -> bool:
+        with self._lock:
+            if ticket._state in _TERMINAL_STATES:
+                return False
+            self._quarantine_locked(error, primary_ticket=ticket)
+            return True
+
+    def _timeout(self, ticket: DeliveryTicket) -> bool:
+        error = DeliveryTimeoutError(
+            "connector delivery timed out after "
+            f"{self.delivery_timeout_s:.3f}s for request "
+            f"{ticket.request_id!r}, key {ticket.put_key!r}"
+        )
+        with self._lock:
+            if ticket._state in _TERMINAL_STATES:
+                return False
+            self._quarantine_locked(error, primary_ticket=ticket)
+            return True
+
+    def _quarantine_locked(
+        self,
+        error: BaseException,
+        *,
+        primary_ticket: DeliveryTicket | None,
+    ) -> None:
+        if self._quarantine_error is None:
+            self._quarantine_error = error
+        stage_error = self._quarantine_error
+        for ticket in tuple(self._tickets):
+            if ticket._state in _TERMINAL_STATES:
+                continue
+            ticket_error = error if ticket is primary_ticket else stage_error
+            self._finish_locked(ticket, DeliveryState.FAILED, ticket_error)
+
+    def _finish_locked(
+        self,
+        ticket: DeliveryTicket,
+        state: DeliveryState,
+        error: BaseException | None,
+    ) -> bool:
+        if ticket._state in _TERMINAL_STATES:
+            return False
+        ticket._state = state
+        ticket._error = error
+        self._tickets.discard(ticket)
+        ticket._event.set()
+        return True

vllm_omni/worker_v2/model_states/__init__.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""Omni-aware ``init_model_state`` factory.
+
+Extends the upstream v2 factory with Omni architecture dispatch.
+Non-Omni architectures fall through to the upstream ``init_model_state``.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from vllm.config import VllmConfig
+from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
+from vllm.v1.worker.gpu.model_states import (
+    init_model_state as _upstream_init_model_state,
+)
+from vllm.v1.worker.gpu.model_states.interface import ModelState
+
+
+def init_omni_model_state(
+    vllm_config: VllmConfig,
+    model: nn.Module,
+    encoder_cache: EncoderCache | None,
+    device: torch.device,
+) -> ModelState:
+    """Create the appropriate ``ModelState`` for *model*.
+
+    Returns an ``OmniModelState`` when the model declares Omni lifecycle
+    capabilities; otherwise
+    delegates to the upstream v2 factory.
+    """
+    uses_omni_lifecycle = any(
+        getattr(model, flag, False) is True for flag in ("has_preprocess", "has_postprocess", "have_multimodal_outputs")
+    )
+    if uses_omni_lifecycle:
+        from vllm_omni.worker_v2.model_states.omni_model_state import (
+            OmniModelState,
+        )
+
+        return OmniModelState(vllm_config, model, encoder_cache, device)
+    return _upstream_init_model_state(vllm_config, model, encoder_cache, device)

vllm_omni/worker_v2/model_states/intermediate_buffer.py

@@ -0,0 +1,265 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""OmniIntermediateBuffer — per-request cross-stage state for Omni pipelines.
+
+Uses ``req_index`` (not ``req_id``) for O(1) access, aligned with v2's
+``RequestState`` slot management.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+import numpy as np
+import torch
+from vllm.v1.core.sched.output import NewRequestData
+from vllm.v1.worker.gpu.input_batch import InputBatch
+
+from vllm_omni.worker.payload_span import get_tensor_span, merge_tensor_spans
+
+
+def _resolve_prompt_embeds(pe: Any) -> torch.Tensor | None:
+    """Convert a prompt_embeds payload to a contiguous CPU tensor."""
+    if pe is None:
+        return None
+    if isinstance(pe, torch.Tensor):
+        return pe.detach().cpu().contiguous()
+    try:
+        data = getattr(pe, "data", None)
+        shape = getattr(pe, "shape", None)
+        if data is not None and shape is not None:
+            dt = np.dtype(getattr(pe, "dtype", "float32"))
+            arr = np.frombuffer(data, dtype=dt).reshape(shape)
+            return torch.from_numpy(arr.copy())
+    except Exception as exc:
+        raise ValueError("Failed to decode prompt_embeds payload") from exc
+    raise TypeError(f"Unsupported prompt_embeds payload type: {type(pe).__name__}")
+
+
+def _resolve_additional_information(payload: Any) -> dict[str, Any]:
+    """Convert an additional_information payload to a plain dict."""
+    if payload is None:
+        return {}
+    try:
+        from vllm_omni.data_entry_keys import unflatten_payload
+        from vllm_omni.engine.serialization import deserialize_additional_information
+
+        info = deserialize_additional_information(payload)
+        if isinstance(info, dict) and any(isinstance(key, str) and "." in key for key in info):
+            info = unflatten_payload(info)
+    except Exception as exc:
+        raise ValueError("Failed to decode additional_information payload") from exc
+    if not isinstance(info, dict):
+        raise TypeError(f"Decoded additional_information must be a dict, got {type(info).__name__}")
+    return info
+
+
+class OmniIntermediateBuffer:
+    """Per-request intermediate state for multi-stage Omni pipelines.
+
+    Stores ``prompt_embeds``, ``additional_information``, ``mm_features``,
+    ``req_id`` and any runtime updates written back by model postprocess.
+    """
+
+    def __init__(self, max_num_reqs: int):
+        self.buffers: list[dict[str, Any]] = [{} for _ in range(max_num_reqs)]
+        self.req_id_to_index: dict[str, int] = {}
+
+    def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
+        info: dict[str, Any] = {}
+
+        pe = getattr(new_req_data, "prompt_embeds", None)
+        if pe is not None:
+            pe_cpu = _resolve_prompt_embeds(pe)
+            if pe_cpu is not None:
+                info["prompt_embeds_cpu"] = pe_cpu
+
+        ai = getattr(new_req_data, "additional_information", None)
+        if ai is not None:
+            info.update(_resolve_additional_information(ai))
+
+        if new_req_data.mm_features:
+            info["mm_features"] = new_req_data.mm_features
+
+        sampling_params = getattr(new_req_data, "sampling_params", None)
+        if sampling_params is not None:
+            info["sampling_params"] = sampling_params
+
+        info["req_id"] = new_req_data.req_id
+        previous_id = self.buffers[req_index].get("req_id")
+        if previous_id is not None and self.req_id_to_index.get(previous_id) == req_index:
+            self.req_id_to_index.pop(previous_id, None)
+        self.buffers[req_index] = info
+        self.req_id_to_index[new_req_data.req_id] = req_index
+
+    def remove_request(self, req_index: int) -> None:
+        req_id = self.buffers[req_index].get("req_id")
+        if req_id is not None and self.req_id_to_index.get(req_id) == req_index:
+            self.req_id_to_index.pop(req_id)
+        self.buffers[req_index] = {}
+
+    def gather(self, input_batch: InputBatch) -> list[dict[str, Any]]:
+        """Return buffer dicts in current batch order (via ``idx_mapping_np``)."""
+        return [self.buffers[idx] for idx in input_batch.idx_mapping_np]
+
+    @staticmethod
+    def _split_gpu_keys(gpu_resident_keys: set[Any] | None) -> tuple[set[Any], dict[Any, set[Any]]]:
+        top_level: set[Any] = set()
+        nested: dict[Any, set[Any]] = {}
+        for key in gpu_resident_keys or set():
+            if isinstance(key, tuple) and len(key) == 2:
+                nested.setdefault(key[0], set()).add(key[1])
+            else:
+                top_level.add(key)
+        return top_level, nested
+
+    @staticmethod
+    def _store_value(dest: dict[Any, Any], key: Any, value: Any, gpu_keys: set[Any]) -> None:
+        if isinstance(value, torch.Tensor):
+            if key in gpu_keys:
+                dest[key] = value.detach().clone()
+            else:
+                dest[key] = value.detach().cpu().contiguous()
+        elif isinstance(value, list):
+            dest[key] = [
+                (
+                    item.detach().clone()
+                    if key in gpu_keys and isinstance(item, torch.Tensor)
+                    else item.detach().cpu().contiguous()
+                    if isinstance(item, torch.Tensor)
+                    else item
+                )
+                for item in value
+            ]
+        else:
+            dest[key] = value
+
+    def update(
+        self,
+        req_index: int,
+        updates: dict[Any, Any],
+        gpu_resident_keys: set[Any] | None = None,
+    ) -> None:
+        """Merge *updates* into the buffer at *req_index*.
+
+        Tensors are detached; those whose key is **not** in
+        *gpu_resident_keys* are moved to CPU.
+        """
+        if not updates:
+            return
+        top_gpu_keys, nested_gpu_keys = self._split_gpu_keys(gpu_resident_keys)
+        existing = self.buffers[req_index]
+        for k, v in updates.items():
+            if isinstance(k, tuple) and len(k) == 2:
+                type_key, qualifier = k
+                existing_sub = existing.setdefault(type_key, {})
+                if not isinstance(existing_sub, dict):
+                    existing_sub = {}
+                    existing[type_key] = existing_sub
+                self._store_value(existing_sub, qualifier, v, nested_gpu_keys.get(type_key, set()))
+            elif isinstance(v, dict):
+                existing_sub = existing.setdefault(k, {})
+                if not isinstance(existing_sub, dict):
+                    existing_sub = {}
+                    existing[k] = existing_sub
+                merged_qualifiers: set[str] = set()
+                if k == "embed" and isinstance(v.get("decode"), torch.Tensor):
+                    span_keys = {"decode_token_start", "decode_token_end"}
+                    if span_keys.intersection(v):
+                        incoming_span = get_tensor_span(
+                            v,
+                            tensor_key="decode",
+                            start_key="decode_token_start",
+                            end_key="decode_token_end",
+                        )
+                        if incoming_span is None:
+                            raise RuntimeError("Incoming Thinker decode delta has an invalid absolute span")
+                        existing_span = get_tensor_span(
+                            existing_sub,
+                            tensor_key="decode",
+                            start_key="decode_token_start",
+                            end_key="decode_token_end",
+                        )
+                        merged_span: tuple[Any, int, int] | None = incoming_span
+                        if existing_span is not None:
+                            merged_span = merge_tensor_spans(existing_span, incoming_span)
+                            if merged_span is None:
+                                raise RuntimeError(
+                                    "Non-contiguous pending Thinker decode spans in model buffer: "
+                                    f"existing=[{existing_span[1]}, {existing_span[2]}) "
+                                    f"incoming=[{incoming_span[1]}, {incoming_span[2]})"
+                                )
+                        assert merged_span is not None
+                        merged_tensor, merged_start, merged_end = merged_span
+                        self._store_value(
+                            existing_sub,
+                            "decode",
+                            merged_tensor,
+                            nested_gpu_keys.get(k, set()),
+                        )
+                        existing_sub["decode_token_start"] = merged_start
+                        existing_sub["decode_token_end"] = merged_end
+                        merged_qualifiers.update({"decode", "decode_token_start", "decode_token_end"})
+                for qualifier, value in v.items():
+                    if qualifier in merged_qualifiers:
+                        continue
+                    self._store_value(existing_sub, qualifier, value, nested_gpu_keys.get(k, set()))
+            else:
+                self._store_value(existing, k, v, top_gpu_keys)
+
+    def update_gpu_tensor_rows(
+        self,
+        req_indices: list[int],
+        key: Any,
+        values: torch.Tensor,
+        *,
+        keepdim: bool = True,
+    ) -> None:
+        """Snapshot a batch-first tensor once and retain owned row views."""
+        self._store_gpu_tensor_rows(req_indices, key, values.detach().clone(), keepdim=keepdim)
+
+    def update_owned_gpu_tensor_rows(
+        self,
+        req_indices: list[int],
+        key: Any,
+        owned_values: torch.Tensor,
+        *,
+        keepdim: bool = True,
+    ) -> None:
+        """Store row views of an already-owned batch tensor without cloning.
+
+        The caller must guarantee ``owned_values`` is not written after this
+        call (the model output ownership contract). Row views keep the owned storage
+        alive until each request row is replaced or the slot is freed.
+        """
+        self._store_gpu_tensor_rows(req_indices, key, owned_values, keepdim=keepdim)
+
+    def _store_gpu_tensor_rows(
+        self,
+        req_indices: list[int],
+        key: Any,
+        owned: torch.Tensor,
+        *,
+        keepdim: bool,
+    ) -> None:
+        num_rows = len(req_indices)
+        if owned.ndim == 0 or owned.shape[0] != num_rows:
+            raise ValueError(
+                "Batched GPU state update changed the request axis: "
+                f"key={key!r} expected={num_rows} "
+                f"actual={owned.shape[0] if owned.ndim else 0}"
+            )
+
+        for row, req_index in enumerate(req_indices):
+            existing = self.buffers[req_index]
+            row_value = owned[row : row + 1] if keepdim else owned[row]
+            if isinstance(key, tuple) and len(key) == 2:
+                type_key, qualifier = key
+                existing_sub = existing.setdefault(type_key, {})
+                if not isinstance(existing_sub, dict):
+                    existing_sub = {}
+                    existing[type_key] = existing_sub
+                existing_sub[qualifier] = row_value
+            else:
+                existing[key] = row_value

vllm_omni/worker_v2/native_output_worker.py

@@ -0,0 +1,99 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Materialize native producer output independently of scheduler consumption."""
+
+from concurrent.futures import Future, ThreadPoolExecutor
+from threading import BoundedSemaphore, Lock
+from typing import Any
+
+from vllm.v1.outputs import AsyncModelRunnerOutput
+
+
+class NativeAsyncOutput(AsyncModelRunnerOutput):
+    """Consume an ordered publication on the engine's owner thread."""
+
+    def __init__(self, future: Future, plane: Any, copy_event: Any) -> None:
+        self._future: Future | None = future
+        self._plane = plane
+        self.copy_event = copy_event
+        self._resolved: Future = Future()
+
+    def get_output(self) -> Any:
+        # Like the upstream lazy AsyncOutputFuture, consumption belongs to one
+        # engine thread. Cache finalization and errors for repeated consumption.
+        if not self._resolved.done():
+            try:
+                assert self._future is not None
+                output = self._future.result()
+                output.omni_connector_output = self._plane.get_omni_connector_output()
+                self._resolved.set_result(output)
+            except Exception as error:
+                self._resolved.set_exception(error)
+            finally:
+                if self._resolved.done():
+                    # Do not retain raw snapshots or a callback bound to self.
+                    self._future = None
+                    self._plane = None
+        return self._resolved.result()
+
+
+class NativeOutputWorker:
+    """Bounded FIFO CPU materialization and native-queue publication.
+
+    Input outputs must have no owner-thread finalizer installed. Reservations
+    and metadata writes must be complete before submit. This worker never
+    updates scheduler state or drains connector notifications.
+    """
+
+    def __init__(self, capacity: int, device: Any | None = None) -> None:
+        if capacity < 1:
+            raise ValueError("Native output capacity must be positive")
+        self._slots = BoundedSemaphore(capacity)
+        self._lock = Lock()
+        self._closed = False
+        self._executor = ThreadPoolExecutor(
+            max_workers=1,
+            thread_name_prefix="omni-native-materialize",
+            initializer=self._initialize_device,
+            initargs=(device,),
+        )
+
+    @staticmethod
+    def _initialize_device(device: Any | None) -> None:
+        if device is not None:
+            from vllm.platforms import current_platform
+
+            current_platform.set_device(device)
+
+    @staticmethod
+    def _materialize(output: AsyncModelRunnerOutput, plane: Any) -> Any:
+        result = output.get_output()
+        plane.enqueue_outputs(
+            req_ids=list(result.req_ids),
+            inter_stage_outputs=getattr(result, "inter_stage_outputs", None),
+            sampled_token_ids=getattr(result, "sampled_token_ids", None),
+        )
+        result.inter_stage_outputs = None
+        return result
+
+    def submit(self, output: AsyncModelRunnerOutput, plane: Any) -> NativeAsyncOutput:
+        with self._lock:
+            if self._closed:
+                raise RuntimeError("Native output worker is closed")
+        self._slots.acquire()
+        try:
+            with self._lock:
+                if self._closed:
+                    raise RuntimeError("Native output worker is closed")
+                future = self._executor.submit(self._materialize, output, plane)
+        except BaseException:
+            self._slots.release()
+            raise
+        future.add_done_callback(lambda _: self._slots.release())
+        return NativeAsyncOutput(future, plane, output.copy_event)
+
+    def close(self) -> None:
+        with self._lock:
+            self._closed = True
+        # Drain submitted publications before the native data plane is closed.
+        self._executor.shutdown(wait=True)

vllm_omni/worker_v2/omni_ar_model_runner.py

@@ -0,0 +1,818 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+"""OmniARModelRunner — autoregressive stage runner on MR V2.
+
+Extends ``OmniGPUModelRunner`` with:
+
+* ``OmniOutput`` post-processing in ``sample_tokens``
+* Per-request ``pooler_output`` construction (hidden + multimodal slices)
+* Async D2H copy via ``OmniAsyncOutput`` for non-blocking output transfer
+* Cross-stage KV extraction before state cleanup
+"""
+
+from __future__ import annotations
+
+from contextlib import nullcontext
+from typing import Any, cast
+
+import numpy as np
+import torch
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.all2all_utils import (
+    get_ep_all2all_manager,
+)
+from vllm.utils.torch_utils import PIN_MEMORY
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
+from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput, RoutedExpertsTensors
+from vllm.v1.worker.gpu.eplb_utils import step_eplb_after
+
+from vllm_omni.data_entry_keys import OmniPayload, flatten_payload
+from vllm_omni.distributed.omni_connectors.kv_transfer_manager import (
+    OmniKVTransferManager,
+)
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.outputs import OmniModelRunnerOutput
+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.output_snapshot import PackedOutputSnapshot, pack_output_snapshot
+
+logger = init_logger(__name__)
+_ASYNC_MM_SNAPSHOT_MAX_BUCKETS_PER_SLOT = 64
+
+
+def _copy_mm_to_snapshot_slot(value: Any, slot: dict[tuple[Any, ...], torch.Tensor], path: tuple[Any, ...] = ()) -> Any:
+    if not path and isinstance(value, dict):
+        packed = pack_output_snapshot(value, slot, max_buckets=_ASYNC_MM_SNAPSHOT_MAX_BUCKETS_PER_SLOT)
+        if packed is not None:
+            return packed
+    if isinstance(value, torch.Tensor):
+        bucket_key = path + (tuple(value.shape), value.dtype, value.device)
+        cached = slot.get(bucket_key)
+        if cached is None:
+            if len(slot) >= _ASYNC_MM_SNAPSHOT_MAX_BUCKETS_PER_SLOT:
+                return value.detach().clone()
+            cached = torch.empty_like(value)
+            slot[bucket_key] = cached
+        cached.copy_(value)
+        return cached
+    if isinstance(value, dict):
+        return {key: _copy_mm_to_snapshot_slot(item, slot, path + (key,)) for key, item in value.items()}
+    if isinstance(value, list):
+        return [_copy_mm_to_snapshot_slot(item, slot, path + (index,)) for index, item in enumerate(value)]
+    if isinstance(value, tuple):
+        return tuple(_copy_mm_to_snapshot_slot(item, slot, path + (index,)) for index, item in enumerate(value))
+    return value
+
+
+def _has_cuda_tensor(value: Any) -> bool:
+    if isinstance(value, torch.Tensor):
+        return value.device.type == "cuda"
+    if isinstance(value, dict):
+        return any(_has_cuda_tensor(item) for item in value.values())
+    if isinstance(value, (list, tuple)):
+        return any(_has_cuda_tensor(item) for item in value)
+    return False
+
+
+def _guard_graph_replay_for_pooler_copy(
+    main_stream: Any,
+    copy_event: Any,
+    *,
+    need_pooler: bool,
+    async_chunk: bool,
+) -> None:
+    """Keep static graph outputs alive until non-snapshotted D2H completes."""
+    if need_pooler and not async_chunk:
+        main_stream.wait_event(copy_event)
+
+
+class OmniARModelRunner(OmniGPUModelRunner):
+    """AR stage runner. Produces per-request hidden states + multimodal outputs."""
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.kv_transfer_manager: OmniKVTransferManager | None = None
+        self._kv_extracted_req_ids: list[str] | None = None
+        self._async_mm_snapshot_slots: list[dict[tuple[Any, ...], torch.Tensor]] = [{} for _ in range(4)]
+        self._async_mm_snapshot_events: list[torch.cuda.Event | None] = [None] * 4
+        self._async_mm_snapshot_pending = [False] * 4
+        self._async_mm_snapshot_cursor = 0
+        self._last_multimodal_snapshot_slot: int | None = None
+
+    def _ensure_kv_transfer_manager(self) -> OmniKVTransferManager:
+        if self.kv_transfer_manager is None:
+            self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
+        return self.kv_transfer_manager
+
+    # ------------------------------------------------------------------
+    # execute_model: KV transfer pre-hook + delegate to super
+    # ------------------------------------------------------------------
+
+    def execute_model(
+        self,
+        scheduler_output: SchedulerOutput,
+        intermediate_tensors: Any | None = None,
+        dummy_run: bool = False,
+        skip_attn_for_dummy_run: bool = False,
+        is_profile: bool = False,
+        context_len: int = 0,
+    ) -> Any:
+        if not dummy_run:
+            self._handle_kv_transfer_pre(scheduler_output)
+        return super().execute_model(
+            scheduler_output,
+            intermediate_tensors,
+            dummy_run=dummy_run,
+            skip_attn_for_dummy_run=skip_attn_for_dummy_run,
+            is_profile=is_profile,
+            context_len=context_len,
+        )
+
+    # ------------------------------------------------------------------
+    # sample_tokens: OmniOutput handling + pooler_output + async D2H
+    # ------------------------------------------------------------------
+
+    @torch.inference_mode()
+    @step_eplb_after()
+    def sample_tokens(
+        self, grammar_output: GrammarOutput | None
+    ) -> OmniAsyncOutput | OmniModelRunnerOutput | ModelRunnerOutput | None:
+        kv_extracted = self._kv_extracted_req_ids
+        self._kv_extracted_req_ids = None
+
+        if self.execute_model_state is None:
+            return None
+
+        input_batch = self.execute_model_state.input_batch
+        hidden_states = self.execute_model_state.hidden_states
+        finished_req_ids = self.execute_model_state.finished_req_ids
+        ec_connector_output = self.execute_model_state.ec_connector_output
+        routed_experts = self.execute_model_state.routed_experts
+        self.execute_model_state = None
+
+        # --- Omni: reconstruct raw model output and post-process ---
+        aux = self._last_aux_output
+        self._last_aux_output = None
+        multimodal_outputs = self._last_multimodal_outputs
+        self._last_multimodal_outputs = None
+        snapshot_slot = None
+        self._last_multimodal_snapshot_slot = None
+        raw_output = self._reconstruct_raw_model_output(
+            hidden_states=hidden_states,
+            multimodal_outputs=multimodal_outputs,
+            aux=aux,
+        )
+        text_hidden, multimodal_outputs = self.model_state.postprocess_model_output(
+            raw_output, input_batch, self.req_states
+        )
+        if bool(getattr(self.model_config, "async_chunk", False)) and multimodal_outputs:
+            multimodal_outputs = self._retain_multimodal_outputs(multimodal_outputs)
+            snapshot_slot = self._last_multimodal_snapshot_slot
+            self._last_multimodal_snapshot_slot = None
+
+        # --- Standard v2 sampling ---
+        sampling_context = nullcontext()
+        model_sampling_context = getattr(self.model, "mrv2_sampling_context", None)
+        if callable(model_sampling_context) and getattr(self.model, "logitsprocs_need_output_token_ids", False):
+            # seq_lens includes this forward's input, but not its sampled token:
+            # final prefill has zero prior outputs; the first decode has one.
+            sampling_context = model_sampling_context(
+                req_ids=input_batch.req_ids,
+                num_output_tokens=(
+                    input_batch.seq_lens[: input_batch.num_reqs]
+                    - self.req_states.prompt_len.gpu[input_batch.idx_mapping]
+                ),
+            )
+        with sampling_context:
+            sampler_output, num_sampled, num_rejected = self.sample(
+                text_hidden,
+                input_batch,
+                grammar_output,
+            )
+        if self.pp_handler is not None:
+            self.pp_handler.broadcast(
+                sampler_output.sampled_token_ids,
+                num_sampled,
+                num_rejected,
+                input_batch,
+            )
+
+        # --- Omni: prompt logprobs ---
+        assert self.prompt_logprobs_worker is not None
+        # Mirror the current parent GPUModelRunner call exactly.
+        prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
+            self.model.compute_logits,
+            text_hidden,
+            input_batch,
+            self.req_states.all_token_ids.gpu,
+            self.req_states.num_computed_tokens.gpu,
+            self.req_states.prompt_len.np,
+        )
+
+        # --- Omni: pooler_output ---
+        engine_output_type = getattr(self.vllm_config.model_config, "engine_output_type", "text")
+        need_pooler = engine_output_type != "text"
+
+        # --- Build base output ---
+        model_runner_output = OmniModelRunnerOutput(
+            req_ids=input_batch.req_ids,
+            req_id_to_index={rid: i for i, rid in enumerate(input_batch.req_ids)},
+            sampled_token_ids=None,  # type: ignore[arg-type]
+            prompt_logprobs_dict=prompt_logprobs_dict,
+            kv_connector_output=None,
+        )
+        model_runner_output.kv_extracted_req_ids = kv_extracted
+        model_runner_output._async_chunk = bool(getattr(self.model_config, "async_chunk", False))
+
+        # --- Async D2H via OmniAsyncOutput ---
+        materialize_native = self._uses_native_output_materializer()
+        async_output = OmniAsyncOutput(
+            model_runner_output=model_runner_output,
+            sampler_output=sampler_output,
+            num_sampled_tokens=num_sampled,
+            main_stream=self.main_stream,
+            copy_stream=self.output_copy_stream,
+            text_hidden=text_hidden if need_pooler else None,
+            multimodal_outputs=multimodal_outputs if need_pooler else None,
+            input_batch=input_batch if need_pooler else None,
+            async_chunk=bool(getattr(self.model_config, "async_chunk", False)),
+            finalize_output=(None if materialize_native else self._finalize_native_data_plane_output),
+            check_ep_fault=self.check_ep_fault,
+            routed_experts=routed_experts,
+        )
+        self._release_multimodal_snapshot(snapshot_slot, async_output.copy_event)
+        _guard_graph_replay_for_pooler_copy(
+            self.main_stream,
+            async_output.copy_event,
+            need_pooler=need_pooler,
+            async_chunk=bool(getattr(self.model_config, "async_chunk", False)),
+        )
+
+        # Postprocess AFTER creating async output (so copy_event is
+        # recorded before postprocess, matching upstream pattern).
+        self.postprocess_sampled(
+            input_batch.idx_mapping,
+            sampler_output.sampled_token_ids,
+            num_sampled,
+            num_rejected,
+            input_batch.query_start_loc,
+        )
+        model_runner_output.kv_connector_output = self.kv_connector.post_forward(finished_req_ids)
+        model_runner_output.ec_connector_output = ec_connector_output
+
+        self._reserve_native_data_plane_outputs(list(model_runner_output.req_ids))
+        if materialize_native:
+            return self._materialize_native_output(async_output)
+        return async_output
+
+    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
+        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().")
+        event = self._async_mm_snapshot_events[slot_index]
+        if event is not None:
+            self.main_stream.wait_event(event)
+        if _has_cuda_tensor(outputs):
+            with torch.cuda.stream(self.main_stream):
+                snapshot = _copy_mm_to_snapshot_slot(outputs, self._async_mm_snapshot_slots[slot_index])
+        else:
+            snapshot = _copy_mm_to_snapshot_slot(outputs, self._async_mm_snapshot_slots[slot_index])
+        self._async_mm_snapshot_pending[slot_index] = True
+        self._last_multimodal_snapshot_slot = slot_index
+        self._async_mm_snapshot_cursor = (slot_index + 1) % len(self._async_mm_snapshot_slots)
+        return snapshot
+
+    def _release_multimodal_snapshot(self, slot_index: int | None, copy_event: torch.cuda.Event) -> None:
+        if slot_index is None:
+            return
+        self._async_mm_snapshot_events[slot_index] = copy_event
+        self._async_mm_snapshot_pending[slot_index] = False
+
+    # ------------------------------------------------------------------
+    # pooler_output construction
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _reconstruct_raw_model_output(
+        *,
+        hidden_states: torch.Tensor,
+        multimodal_outputs: dict[str, Any] | None,
+        aux: Any | None,
+    ) -> Any:
+        if multimodal_outputs:
+            return OmniOutput(
+                text_hidden_states=hidden_states,
+                multimodal_outputs=cast(OmniPayload, multimodal_outputs),
+            )
+        if aux is not None:
+            return (hidden_states, aux)
+        return hidden_states
+
+    @staticmethod
+    def _build_pooler_output_from_cpu(
+        hidden_cpu: torch.Tensor,
+        mm_cpu: dict[str, Any],
+        query_start_loc_np: np.ndarray,
+        num_scheduled_tokens: np.ndarray,
+        num_reqs: int,
+    ) -> list[dict[str, Any]]:
+        """Build pooler_output from already-CPU tensors."""
+        total = hidden_cpu.shape[0]
+        pooler: list[dict[str, Any]] = []
+        for i in range(num_reqs):
+            start = int(query_start_loc_np[i])
+            end = start + int(num_scheduled_tokens[i])
+            payload: dict[str, Any] = {"hidden": hidden_cpu[start:end].clone()}
+            for k, v in mm_cpu.items():
+                payload[k] = _slice_pooler_value(
+                    v,
+                    req_index=i,
+                    start=start,
+                    end=end,
+                    total_tokens=total,
+                )
+            pooler.append(flatten_payload(payload))
+        return pooler
+
+    @staticmethod
+    def _build_async_chunk_outputs_from_mm(
+        mm_outputs: dict[str, Any],
+        query_start_loc_np: np.ndarray,
+        num_scheduled_tokens: np.ndarray,
+        num_reqs: int,
+        total_tokens: int,
+        padded_total_tokens: int | None = None,
+    ) -> tuple[list[dict[str, Any] | None] | None, list[dict[str, Any] | None] | None]:
+        """Build async-chunk payloads without materializing hidden on CPU."""
+        if not mm_outputs:
+            return None, None
+
+        inter_stage_list: list[dict[str, Any] | None] = []
+        client_mm_list: list[dict[str, Any] | None] = []
+        for i in range(num_reqs):
+            start = int(query_start_loc_np[i])
+            end = start + int(num_scheduled_tokens[i])
+            payload = {
+                key: _slice_pooler_value(
+                    value,
+                    req_index=i,
+                    start=start,
+                    end=end,
+                    total_tokens=total_tokens,
+                    padded_total_tokens=padded_total_tokens,
+                )
+                for key, value in mm_outputs.items()
+            }
+            inter_stage, client_mm = partition_flat_payload(flatten_payload(payload))
+            inter_stage_list.append(inter_stage or None)
+            client_mm_list.append(client_mm or None)
+
+        return (
+            None if all(item is None for item in inter_stage_list) else inter_stage_list,
+            None if all(item is None for item in client_mm_list) else client_mm_list,
+        )
+
+    # ------------------------------------------------------------------
+    # KV transfer
+    # ------------------------------------------------------------------
+
+    def _handle_kv_transfer_pre(self, scheduler_output: SchedulerOutput) -> None:
+        finished: dict = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {})
+        if not finished:
+            return
+
+        kv_caches = getattr(self, "kv_caches", None)
+        if kv_caches is None:
+            return
+
+        if hasattr(self.model, "get_kv_transfer_metadata"):
+            for req_id, data in finished.items():
+                try:
+                    meta = self.model.get_kv_transfer_metadata(req_id)
+                    if meta:
+                        existing = data.get("custom_metadata") or {}
+                        existing.update(meta)
+                        data["custom_metadata"] = existing
+                except Exception:
+                    logger.warning(
+                        "Failed to get KV transfer metadata for %s",
+                        req_id,
+                        exc_info=True,
+                    )
+
+        mgr = self._ensure_kv_transfer_manager()
+        self._kv_extracted_req_ids = mgr.handle_finished_requests_kv_transfer(
+            finished_reqs=finished,
+            kv_caches=kv_caches,
+            block_size=self.cache_config.block_size,
+            cache_dtype=str(self.cache_config.cache_dtype),
+            request_id_resolver=self._resolve_global_request_id,
+        )
+
+    def _resolve_global_request_id(self, req_id: str) -> str:
+        req_idx = self.req_states.req_id_to_index.get(req_id)
+        if req_idx is None:
+            return req_id
+        info = self.model_state.intermediate_buffer.buffers[req_idx]
+        global_id = info.get("global_request_id")
+        if isinstance(global_id, list):
+            global_id = global_id[0] if global_id else None
+        if isinstance(global_id, bytes):
+            return global_id.decode("utf-8")
+        return str(global_id) if global_id else req_id
+
+
+# ======================================================================
+# OmniAsyncOutput — async D2H for Omni AR outputs
+# ======================================================================
+
+
+def _async_copy_to_np(
+    x: torch.Tensor,
+    *,
+    copy_stream: torch.cuda.Stream | None = None,
+    pin_memory: bool | None = None,
+) -> np.ndarray:
+    return _async_copy_tensor(
+        x,
+        copy_stream=copy_stream,
+        pin_memory=pin_memory,
+    ).numpy()
+
+
+def _async_copy_tensor(
+    x: torch.Tensor,
+    *,
+    copy_stream: torch.cuda.Stream | None = None,
+    pin_memory: bool | None = None,
+) -> torch.Tensor:
+    x = x.detach()
+    if x.device.type == "cpu":
+        return x.clone()
+    if pin_memory is None:
+        pin_memory = PIN_MEMORY
+    cpu = torch.empty_like(x, device="cpu", pin_memory=pin_memory)
+    cpu.copy_(x, non_blocking=pin_memory)
+    if copy_stream is None:
+        copy_stream = torch.cuda.current_stream()
+    x.record_stream(copy_stream)
+    return cpu
+
+
+def _async_copy_mm_value(
+    value: Any,
+    *,
+    copy_stream: torch.cuda.Stream | None = None,
+    pin_memory: bool | None = None,
+) -> Any:
+    if isinstance(value, torch.Tensor):
+        return _async_copy_tensor(
+            value,
+            copy_stream=copy_stream,
+            pin_memory=pin_memory,
+        )
+    if isinstance(value, dict):
+        return {
+            key: _async_copy_mm_value(
+                val,
+                copy_stream=copy_stream,
+                pin_memory=pin_memory,
+            )
+            for key, val in value.items()
+        }
+    if isinstance(value, list):
+        return [
+            _async_copy_mm_value(
+                val,
+                copy_stream=copy_stream,
+                pin_memory=pin_memory,
+            )
+            for val in value
+        ]
+    if isinstance(value, tuple):
+        return tuple(
+            _async_copy_mm_value(
+                val,
+                copy_stream=copy_stream,
+                pin_memory=pin_memory,
+            )
+            for val in value
+        )
+    return value
+
+
+def _async_copy_mm(
+    mm_outputs: dict | None,
+    total_tokens: int,
+    *,
+    copy_stream: torch.cuda.Stream | None = None,
+    pin_memory: bool | None = None,
+) -> dict[str, Any]:
+    """Non-blocking D2H copy of multimodal output tensors."""
+    if not mm_outputs:
+        return {}
+    if isinstance(mm_outputs, PackedOutputSnapshot):
+        return mm_outputs.copy_to_cpu(
+            lambda tensor: _async_copy_tensor(tensor, copy_stream=copy_stream, pin_memory=pin_memory)
+        )
+    return {
+        key: _async_copy_mm_value(
+            value,
+            copy_stream=copy_stream,
+            pin_memory=pin_memory,
+        )
+        for key, value in mm_outputs.items()
+    }
+
+
+def _slice_pooler_value(
+    value: Any,
+    *,
+    req_index: int,
+    start: int,
+    end: int,
+    total_tokens: int,
+    padded_total_tokens: int | None = None,
+    request_scoped: bool = False,
+) -> Any:
+    if isinstance(value, torch.Tensor):
+        token_axis_sizes = {total_tokens}
+        if padded_total_tokens is not None:
+            token_axis_sizes.add(padded_total_tokens)
+        if not request_scoped and value.dim() > 0 and value.shape[0] in token_axis_sizes:
+            return value[start:end].contiguous()
+        return value.clone()
+    if isinstance(value, dict):
+        return {
+            key: _slice_pooler_value(
+                val,
+                req_index=req_index,
+                start=start,
+                end=end,
+                total_tokens=total_tokens,
+                padded_total_tokens=padded_total_tokens,
+                request_scoped=request_scoped,
+            )
+            for key, val in value.items()
+        }
+    if isinstance(value, list):
+        if not value:
+            return []
+        # Lists at the batch level contain request-owned payloads. Their
+        # tensors use local axes even when a size coincides with batch tokens.
+        values = value if request_scoped else [value[req_index] if req_index < len(value) else value[0]]
+        owned = [
+            _slice_pooler_value(
+                item,
+                req_index=req_index,
+                start=start,
+                end=end,
+                total_tokens=total_tokens,
+                request_scoped=True,
+            )
+            for item in values
+        ]
+        return owned if request_scoped else owned[0]
+    return value
+
+
+def _ensure_tensor_values(payload: dict[str, Any]) -> dict[str, torch.Tensor]:
+    """Convert a flattened payload to strictly ``dict[str, torch.Tensor]``.
+
+    Non-tensor scalars (int/float/bool) are wrapped with ``torch.tensor()``;
+    values that cannot be safely converted are dropped. Enforces the tensor-only
+    invariant required by ``OmniEngineCoreOutput.multimodal_output`` (the channel
+    the async_chunk stage-input processor reads). Mirrors the V1 runner helper.
+    """
+    result: dict[str, torch.Tensor] = {}
+    for key, val in payload.items():
+        if isinstance(val, torch.Tensor):
+            result[key] = val
+        elif isinstance(val, (int, float, bool)):
+            result[key] = torch.tensor(val)
+        elif isinstance(val, (list, tuple)):
+            try:
+                result[key] = torch.tensor(val)
+            except (ValueError, TypeError, RuntimeError):
+                logger.warning(
+                    "Dropping non-tensorizable multimodal output key '%s' (type=%s).",
+                    key,
+                    type(val).__name__,
+                )
+        else:
+            logger.warning(
+                "Dropping non-tensor multimodal output key '%s' (type=%s).",
+                key,
+                type(val).__name__,
+            )
+    return result
+
+
+class OmniAsyncOutput(AsyncModelRunnerOutput):
+    """Async D2H copy for Omni AR model outputs.
+
+    Mirrors upstream ``AsyncOutput`` but additionally handles
+    ``pooler_output`` (hidden states + multimodal outputs) via
+    non-blocking copies on the copy stream.
+    """
+
+    def __init__(
+        self,
+        model_runner_output: OmniModelRunnerOutput,
+        sampler_output: Any,
+        num_sampled_tokens: torch.Tensor,
+        main_stream: torch.cuda.Stream,
+        copy_stream: torch.cuda.Stream,
+        copy_event: torch.cuda.Event | None = None,
+        text_hidden: torch.Tensor | None = None,
+        multimodal_outputs: dict | None = None,
+        input_batch: Any | None = None,
+        async_chunk: bool = False,
+        finalize_output: Any | None = None,
+        check_ep_fault: bool = False,
+        routed_experts: RoutedExpertsTensors | None = None,
+    ):
+        self.model_runner_output = model_runner_output
+        self.sampler_output = sampler_output
+        self.num_sampled_tokens = num_sampled_tokens
+        self.routed_experts = routed_experts
+        self.copy_event = copy_event if copy_event is not None else torch.cuda.Event(blocking=True)
+        self._async_chunk = bool(async_chunk)
+        self._finalize_output = finalize_output
+        self._has_fault: torch.Tensor | None = None
+
+        # Snapshot input_batch metadata needed for pooler_output slicing
+        self._need_pooler = text_hidden is not None or (self._async_chunk and bool(multimodal_outputs))
+        self._query_start_loc_np: np.ndarray | None = None
+        self._num_scheduled_tokens: np.ndarray | None = None
+        self._num_reqs: int = 0
+        self._total_tokens: int = 0
+        self._padded_total_tokens: int = 0
+        if self._need_pooler and input_batch is not None:
+            self._query_start_loc_np = input_batch.query_start_loc_np.copy()
+            self._num_scheduled_tokens = np.array(input_batch.num_scheduled_tokens, dtype=np.int32)
+            self._num_reqs = input_batch.num_reqs
+            if self._query_start_loc_np.shape[0] > self._num_reqs:
+                self._total_tokens = int(self._query_start_loc_np[self._num_reqs])
+            else:
+                self._total_tokens = int(self._num_scheduled_tokens[: self._num_reqs].sum())
+            self._padded_total_tokens = int(getattr(input_batch, "num_tokens_after_padding", self._total_tokens))
+
+        # Perform all D2H copies on the copy stream (non-blocking).
+        import contextlib
+
+        @contextlib.contextmanager
+        def _stream(to_stream, from_stream):
+            try:
+                torch.cuda.set_stream(to_stream)
+                yield
+            finally:
+                torch.cuda.set_stream(from_stream)
+
+        with _stream(copy_stream, main_stream):
+            copy_stream.wait_stream(main_stream)
+            pin_memory = PIN_MEMORY
+
+            # Sampled token ids
+            self.sampled_token_ids_np = _async_copy_to_np(
+                sampler_output.sampled_token_ids,
+                copy_stream=copy_stream,
+                pin_memory=pin_memory,
+            )
+            self.num_sampled_tokens_np = _async_copy_to_np(
+                num_sampled_tokens,
+                copy_stream=copy_stream,
+                pin_memory=pin_memory,
+            )
+            self.sampling_mask_tensors = None
+            if sampler_output.sampling_mask_tensors is not None:
+                self.sampling_mask_tensors = sampler_output.sampling_mask_tensors.to_cpu_nonblocking()
+            self.routed_experts_cpu = None
+            if routed_experts is not None:
+                self.routed_experts_cpu = routed_experts.to_cpu_nonblocking()
+
+            # Logprobs
+            self.logprobs_tensors = None
+            if sampler_output.logprobs_tensors is not None:
+                self.logprobs_tensors = sampler_output.logprobs_tensors.to_cpu_nonblocking()
+            self.num_nans: np.ndarray | None = None
+            if sampler_output.num_nans is not None:
+                self.num_nans = _async_copy_to_np(
+                    sampler_output.num_nans,
+                    copy_stream=copy_stream,
+                    pin_memory=pin_memory,
+                )
+
+            # Prompt logprobs
+            self.prompt_logprobs_dict = {
+                k: v.to_cpu_nonblocking() if v is not None else None
+                for k, v in self.model_runner_output.prompt_logprobs_dict.items()
+            }
+            if check_ep_fault:
+                has_fault = get_ep_all2all_manager().query_fault()
+                self._has_fault = has_fault.to("cpu", non_blocking=True)
+
+            # Pooler output (hidden + multimodal) — async D2H
+            self._hidden_cpu: torch.Tensor | None = None
+            self._mm_cpu: dict[str, Any] = {}
+            self._mm_snapshot: dict[str, Any] = {}
+            if self._need_pooler and self._async_chunk:
+                # 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,
+                )
+            elif self._need_pooler and text_hidden is not None:
+                self._hidden_cpu = _async_copy_tensor(
+                    text_hidden,
+                    copy_stream=copy_stream,
+                    pin_memory=pin_memory,
+                )
+                total_tokens = text_hidden.shape[0]
+                self._mm_cpu = _async_copy_mm(
+                    multimodal_outputs,
+                    total_tokens,
+                    copy_stream=copy_stream,
+                    pin_memory=pin_memory,
+                )
+
+            self.copy_event.record(copy_stream)
+
+    def get_output(self) -> OmniModelRunnerOutput:
+        self.copy_event.synchronize()
+
+        # Sampled token ids
+        sampled_token_ids: list[list[int]] = self.sampled_token_ids_np.tolist()
+        num_sampled_tokens: list[int] = self.num_sampled_tokens_np.tolist()
+        for token_ids, num_tokens in zip(sampled_token_ids, num_sampled_tokens):
+            del token_ids[num_tokens:]
+        self.model_runner_output.sampled_token_ids = sampled_token_ids
+        if self.sampling_mask_tensors is not None:
+            self.model_runner_output.sampling_masks = self.sampling_mask_tensors.tolists(self.num_sampled_tokens_np)
+        if self.routed_experts_cpu is not None:
+            self.model_runner_output.routed_experts = self.routed_experts_cpu.tolists()
+        self.model_runner_output.sampled_token_ids_materialized = True
+
+        # Logprobs
+        if self.num_nans is not None:
+            self.model_runner_output.num_nans_in_logits = dict(
+                zip(self.model_runner_output.req_ids, self.num_nans.tolist())
+            )
+        if self.logprobs_tensors is not None:
+            self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
+        self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict
+
+        if self._has_fault is not None and self._has_fault.item():
+            mask = get_ep_all2all_manager().query_active_mask()
+            raise RuntimeError(
+                "Fault detected in EP all2all communication: one or more ranks "
+                f"timed out during dispatch/combine. Mask: {mask.cpu().tolist()}"
+            )
+
+        # Pooler output. Populate two channels from the same per-request payloads,
+        # mirroring the V1 runner:
+        #   * pooler_output  -> sync/full-payload path (inline pooling_output bridge)
+        #   * multimodal_outputs -> wire multimodal_output, which the async_chunk
+        #     stage-input processor (talker2code2wav_async_chunk) reads for codes.
+        if self._need_pooler and self._async_chunk:
+            pooler_inter, pooler_client = OmniARModelRunner._build_async_chunk_outputs_from_mm(
+                self._mm_snapshot,
+                self._query_start_loc_np,
+                self._num_scheduled_tokens,
+                self._num_reqs,
+                self._total_tokens,
+                self._padded_total_tokens,
+            )
+            self.model_runner_output.pooler_output = None
+            self.model_runner_output.inter_stage_outputs = pooler_inter
+            self.model_runner_output.multimodal_outputs = (
+                [_ensure_tensor_values(_async_copy_mm_value(p)) if p else {} for p in pooler_client]
+                if pooler_client
+                else None
+            )
+        elif self._need_pooler and self._hidden_cpu is not None:
+            pooler_output = OmniARModelRunner._build_pooler_output_from_cpu(
+                self._hidden_cpu,
+                self._mm_cpu,
+                self._query_start_loc_np,
+                self._num_scheduled_tokens,
+                self._num_reqs,
+            )
+            pooler_payload = cast(list[dict[str, Any] | None], pooler_output) if pooler_output else None
+            self.model_runner_output.pooler_output = pooler_payload
+            self.model_runner_output.inter_stage_outputs = pooler_payload
+            self.model_runner_output.multimodal_outputs = (
+                [_ensure_tensor_values(p) if p else {} for p in pooler_payload] if pooler_payload else None
+            )
+
+        if self._finalize_output is not None:
+            return self._finalize_output(self.model_runner_output)
+        return self.model_runner_output

vllm_omni/worker_v2/omni_data_plane.py

@@ -0,0 +1,604 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+
+from __future__ import annotations
+
+import threading
+import time
+from collections import defaultdict
+from contextlib import AbstractContextManager, nullcontext
+from dataclasses import dataclass, field
+from queue import Full, Queue
+from types import SimpleNamespace
+from typing import Any, cast
+
+import torch
+from vllm.logger import init_logger
+
+from vllm_omni.data_entry_keys import OmniPayload, unflatten_payload
+from vllm_omni.worker.omni_connector_model_runner_mixin import (
+    OmniConnectorModelRunnerMixin,
+)
+from vllm_omni.worker.payload_span import get_tensor_span, merge_tensor_spans
+from vllm_omni.worker_v2.delivery import (
+    DeliveryTimeoutError,
+    OmniDeliveryManager,
+)
+
+
+@dataclass
+class _NativeRequestState:
+    request_id: str
+    external_req_id: str
+    prompt_token_ids: list[int]
+    additional_information: Any = None
+    sampling_params: Any = None
+    num_computed_tokens: int = 0
+    resumable: bool = False
+    output_token_ids: list[int] = field(default_factory=list)
+    finished: bool = False
+
+    def snapshot(self, *, include_token_history: bool) -> SimpleNamespace:
+        prompt = list(self.prompt_token_ids) if include_token_history else []
+        output = list(self.output_token_ids) if include_token_history else []
+        finished = self.finished
+        request = SimpleNamespace(
+            request_id=self.request_id,
+            req_id=self.request_id,
+            external_req_id=self.external_req_id,
+            prompt_token_ids=prompt,
+            output_token_ids=output,
+            all_token_ids=prompt + output,
+            output_token_count=len(self.output_token_ids),
+            additional_information=self.additional_information,
+            sampling_params=self.sampling_params,
+            num_computed_tokens=self.num_computed_tokens,
+            resumable=self.resumable,
+        )
+        request.is_finished = lambda: finished
+        return request
+
+
+@dataclass(frozen=True)
+class _NativeOutputBatch:
+    req_ids: list[str]
+    inter_stage_outputs: list[Any | None] | None
+    sampled_token_ids: list[list[int]] | None
+
+
+_NATIVE_OUTPUT_STOP = object()
+_NATIVE_OUTPUT_QUEUE_DEPTH = 8
+_DEFAULT_DELIVERY_TIMEOUT_S = 30.0
+_DEFAULT_SHUTDOWN_TIMEOUT_S = 5.0
+logger = init_logger(__name__)
+
+
+class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin):
+    """MRv2-owned stage transport and request-side payload state."""
+
+    def __init__(self, vllm_config: Any, model_config: Any) -> None:
+        self._native_requests: dict[str, _NativeRequestState] = {}
+        self._native_output_lock = threading.RLock()
+        self._native_send_lock = threading.Lock()
+        self._native_outputs_in_flight: dict[str, int] = defaultdict(int)
+        self._native_terminal_pending: set[str] = set()
+        self.init_omni_connectors(model_config=model_config)
+        self._delivery_manager = OmniDeliveryManager(
+            delivery_timeout_s=self._connector_delivery_timeout(),
+            shutdown_timeout_s=_DEFAULT_SHUTDOWN_TIMEOUT_S,
+        )
+        self._can_send = self._custom_process_func is not None
+        self._start_output_worker(max_pending_batches=_NATIVE_OUTPUT_QUEUE_DEPTH)
+
+    def _connector_delivery_timeout(self) -> float:
+        config = getattr(getattr(self, "_omni_connector", None), "config", None)
+        extra = config.get("extra", {}) if isinstance(config, dict) else {}
+        value = extra.get("delivery_timeout_s", _DEFAULT_DELIVERY_TIMEOUT_S)
+        try:
+            timeout = float(value)
+        except (TypeError, ValueError) as error:
+            raise ValueError(f"invalid connector delivery_timeout_s: {value!r}") from error
+        if timeout <= 0:
+            raise ValueError("connector delivery_timeout_s must be positive")
+        return timeout
+
+    def _create_send_completion(self, *, request_id: str, put_key: str) -> Any:
+        return self._delivery_manager.create_ticket(
+            request_id=request_id,
+            put_key=put_key,
+        )
+
+    def _start_output_worker(self, *, max_pending_batches: int) -> None:
+        if max_pending_batches <= 0:
+            raise ValueError("max_pending_batches must be positive")
+        self._native_output_queue: Queue[_NativeOutputBatch | object] = Queue(maxsize=max_pending_batches)
+        self._native_output_error_lock = threading.Lock()
+        self._native_output_error: BaseException | None = None
+        self._connector_send_error_sink = self._record_output_error
+        self._native_output_closed = False
+        stage_id = getattr(self, "_stage_id", "unknown")
+        self._native_output_worker: threading.Thread | None = threading.Thread(
+            target=self._output_worker_loop,
+            name=f"omni-native-output-{stage_id}",
+            daemon=True,
+        )
+        self._native_output_worker.start()
+
+    def _record_output_error(self, error: BaseException) -> None:
+        first_error = False
+        with self._native_output_error_lock:
+            if self._native_output_error is None:
+                self._native_output_error = error
+                first_error = True
+        if first_error:
+            manager = getattr(self, "_delivery_manager", None)
+            if manager is not None:
+                manager.quarantine(error)
+
+    def _raise_output_error(self) -> None:
+        with self._native_output_error_lock:
+            error = self._native_output_error
+        if error is not None:
+            raise RuntimeError(f"native output worker failed: {error}") from error
+
+    def _output_worker_loop(self) -> None:
+        while True:
+            batch = self._native_output_queue.get()
+            try:
+                if batch is _NATIVE_OUTPUT_STOP:
+                    return
+                if self._native_output_error is not None:
+                    continue
+                assert isinstance(batch, _NativeOutputBatch)
+                self.complete_outputs(
+                    req_ids=batch.req_ids,
+                    inter_stage_outputs=batch.inter_stage_outputs,
+                    sampled_token_ids=batch.sampled_token_ids,
+                )
+            except BaseException as error:
+                self._record_output_error(error)
+            finally:
+                self._native_output_queue.task_done()
+
+    def enqueue_outputs(
+        self,
+        *,
+        req_ids: list[str],
+        inter_stage_outputs: list[Any | None] | None,
+        sampled_token_ids: list[list[int]] | None,
+    ) -> None:
+        """Transfer one completed model batch to the ordered output worker."""
+        self._raise_output_error()
+        if self._native_output_closed:
+            raise RuntimeError("native output worker is closed")
+        batch = _NativeOutputBatch(
+            req_ids=req_ids,
+            inter_stage_outputs=inter_stage_outputs,
+            sampled_token_ids=sampled_token_ids,
+        )
+        while True:
+            try:
+                self._native_output_queue.put(batch, timeout=0.1)
+                break
+            except Full:
+                self._raise_output_error()
+        self._raise_output_error()
+
+    def drain_outputs(self) -> None:
+        queue = self._native_output_queue
+        pending_batches = max(1, queue.qsize() + 1)
+        timeout_s = self._delivery_manager.delivery_timeout_s * pending_batches
+        deadline = time.monotonic() + timeout_s
+        while queue.unfinished_tasks:
+            self._raise_output_error()
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                error = DeliveryTimeoutError(
+                    "native output queue drain timed out after "
+                    f"{timeout_s:.3f}s with {queue.unfinished_tasks} "
+                    "unfinished batch(es)"
+                )
+                self._record_output_error(error)
+                self._raise_output_error()
+            with queue.all_tasks_done:
+                if queue.unfinished_tasks:
+                    queue.all_tasks_done.wait(timeout=min(0.05, remaining))
+        self._raise_output_error()
+
+    def _stop_output_worker(self) -> None:
+        worker = getattr(self, "_native_output_worker", None)
+        if worker is None:
+            return
+        self._native_output_closed = True
+        self._delivery_manager.shutdown(RuntimeError("MRv2 output worker shutdown"))
+        deadline = time.monotonic() + self._delivery_manager.shutdown_timeout_s
+        while worker.is_alive():
+            remaining = deadline - time.monotonic()
+            if remaining <= 0:
+                break
+            try:
+                self._native_output_queue.put(
+                    _NATIVE_OUTPUT_STOP,
+                    timeout=min(0.05, remaining),
+                )
+                break
+            except Full:
+                continue
+        worker.join(timeout=max(0.0, deadline - time.monotonic()))
+        if worker.is_alive():
+            raise RuntimeError("MRv2 output worker did not stop within its shutdown deadline")
+        self._native_output_worker = None
+
+    def register_request(self, request_data: Any) -> None:
+        self._raise_output_error()
+        req_id = str(request_data.req_id)
+        external_req_id = str(getattr(request_data, "external_req_id", None) or req_id)
+        with self._native_output_lock:
+            self._native_outputs_in_flight.pop(req_id, None)
+            self._native_terminal_pending.discard(req_id)
+            self._native_requests[req_id] = _NativeRequestState(
+                request_id=req_id,
+                external_req_id=external_req_id,
+                prompt_token_ids=list(getattr(request_data, "prompt_token_ids", None) or []),
+                additional_information=getattr(request_data, "additional_information", None),
+                sampling_params=getattr(request_data, "sampling_params", None),
+                num_computed_tokens=int(getattr(request_data, "num_computed_tokens", 0) or 0),
+                resumable=bool(getattr(request_data, "resumable", False)),
+            )
+
+    def reserve_outputs(self, req_ids: list[str]) -> None:
+        """Keep request state alive until deferred runner outputs are consumed."""
+        self._raise_output_error()
+        with self._native_output_lock:
+            for req_id in dict.fromkeys(req_ids):
+                if req_id in self._native_requests:
+                    self._native_outputs_in_flight[req_id] += 1
+
+    def _ready_terminal_requests(self) -> set[str]:
+        return {
+            req_id for req_id in self._native_terminal_pending if self._native_outputs_in_flight.get(req_id, 0) == 0
+        }
+
+    def _emit_ready_terminals(self) -> int:
+        ready = self._ready_terminal_requests()
+        if not ready:
+            return 0
+        emitted = self.emit_chunks(
+            req_ids=[],
+            inter_stage_outputs=None,
+            sampled_token_ids=None,
+            terminal_req_ids=ready,
+        )
+        self._native_terminal_pending.difference_update(ready)
+        for req_id in ready:
+            self._native_outputs_in_flight.pop(req_id, None)
+        return emitted
+
+    def request_terminal(self, req_ids: set[str]) -> int:
+        """Emit terminal markers only after all deferred outputs are enqueued."""
+        self._raise_output_error()
+        with self._native_output_lock:
+            self._native_terminal_pending.update(req_id for req_id in req_ids if req_id in self._native_requests)
+            return self._emit_ready_terminals()
+
+    def abort_requests(self, req_ids: set[str]) -> int:
+        """Cancel deferred outputs and terminate each live request once."""
+        with self._native_output_lock:
+            active_req_ids = {req_id for req_id in req_ids if req_id in self._native_requests}
+            # Receivers are registered before requests acquire model slots.
+            # Cancelling during that interval must still stop recv polling.
+            for req_id in req_ids - active_req_ids:
+                self.cleanup_finished_request(req_id)
+            if not active_req_ids:
+                return 0
+            for req_id in active_req_ids:
+                self._native_outputs_in_flight.pop(req_id, None)
+                self._native_terminal_pending.discard(req_id)
+            return self.emit_chunks(
+                req_ids=[],
+                inter_stage_outputs=None,
+                sampled_token_ids=None,
+                terminal_req_ids=active_req_ids,
+            )
+
+    def complete_outputs(
+        self,
+        *,
+        req_ids: list[str],
+        inter_stage_outputs: list[Any | None] | None,
+        sampled_token_ids: list[list[int]] | None,
+    ) -> int:
+        """Commit one deferred output batch, then release its lifecycle holds."""
+        send_lock_held = False
+        with self._native_output_lock:
+            entries, finished_req_ids = self._build_chunk_entries(
+                req_ids=req_ids,
+                inter_stage_outputs=inter_stage_outputs,
+                sampled_token_ids=sampled_token_ids,
+                terminal_req_ids=set(),
+            )
+            # Claim connector ordering before publishing the updated request
+            # state. request_terminal() can still register a pending terminal
+            # while the enqueue runs, but abort cannot overtake this output.
+            if entries:
+                self._native_send_lock.acquire()
+                send_lock_held = True
+        try:
+            emitted = self._send_chunk_entries(
+                entries,
+                send_lock_held=send_lock_held,
+            )
+        finally:
+            if send_lock_held:
+                self._native_send_lock.release()
+        with self._native_output_lock:
+            self._cleanup_finished_entries(finished_req_ids)
+            for req_id in dict.fromkeys(req_ids):
+                in_flight = self._native_outputs_in_flight.get(req_id, 0)
+                if in_flight > 1:
+                    self._native_outputs_in_flight[req_id] = in_flight - 1
+                elif in_flight == 1:
+                    self._native_outputs_in_flight.pop(req_id, None)
+            return emitted + self._emit_ready_terminals()
+
+    def register_receivers(self, handles: list[Any]) -> None:
+        for handle in handles:
+            self.register_chunk_recv(handle)
+
+    @classmethod
+    def _copy_payload_structure(cls, value: Any) -> Any:
+        if isinstance(value, dict):
+            return {key: cls._copy_payload_structure(item) for key, item in value.items()}
+        if isinstance(value, list):
+            return list(value)
+        return value
+
+    def _accumulate_payload(self, req_id: str, payload_data: OmniPayload) -> OmniPayload:
+        """Accumulate chunk payloads (concat tensors, extend lists)."""
+        if req_id not in self._send_side_request_payload:
+            self._send_side_request_payload[req_id] = dict(payload_data)
+            return cast(OmniPayload, dict(self._send_side_request_payload[req_id]))
+
+        origin = self._send_side_request_payload[req_id]
+        merged = dict(origin)
+        raw_ok = payload_data.get("meta", {}).get("override_keys", []) if isinstance(payload_data, dict) else []
+        override_keys = {tuple(k) if isinstance(k, list) else k for k in raw_ok}
+
+        for key, value in payload_data.items():
+            if isinstance(value, dict):
+                origin_sub = origin.get(key)
+                merged_sub = dict(origin_sub) if isinstance(origin_sub, dict) else {}
+                span_handled: set[str] = set()
+                if key == "embed" and isinstance(origin_sub, dict):
+                    for tk, sk, ek in (("decode", "decode_token_start", "decode_token_end"),):
+                        if tk not in value or (key, tk) in override_keys:
+                            continue
+                        span = merge_tensor_spans(
+                            get_tensor_span(origin_sub, tensor_key=tk, start_key=sk, end_key=ek),
+                            get_tensor_span(value, tensor_key=tk, start_key=sk, end_key=ek),
+                        )
+                        if span is None:
+                            continue
+                        t, s, e = span
+                        merged_sub[tk] = t
+                        merged_sub[sk] = s
+                        merged_sub[ek] = e
+                        span_handled |= {tk, sk, ek}
+                for qual, qval in value.items():
+                    if qual in span_handled:
+                        continue
+                    if key == "meta" and qual == "finished":
+                        merged_sub[qual] = qval
+                        continue
+                    if (key, qual) in override_keys:
+                        merged_sub[qual] = qval
+                        continue
+                    osv = merged_sub.get(qual)
+                    if isinstance(qval, torch.Tensor) and isinstance(osv, torch.Tensor):
+                        merged_sub[qual] = torch.cat([osv, qval], dim=0)
+                    elif isinstance(qval, list) and isinstance(osv, list):
+                        merged_sub[qual] = osv + qval
+                    else:
+                        merged_sub[qual] = qval
+                merged[key] = merged_sub
+            else:
+                if key in override_keys:
+                    merged[key] = value
+                    continue
+                ov = origin.get(key)
+                if isinstance(value, torch.Tensor) and isinstance(ov, torch.Tensor):
+                    merged[key] = torch.cat([ov, value], dim=0)
+                elif isinstance(value, list) and isinstance(ov, list):
+                    merged[key] = ov + value
+                else:
+                    merged[key] = value
+
+        self._send_side_request_payload[req_id] = merged
+        return cast(OmniPayload, dict(merged))
+
+    def pop_local_stage_payload(self, req_id: str) -> Any:
+        """Hand one accumulated delta to the model and acknowledge its rows.
+
+        Connector accumulation and ``model_intermediate_buffer`` are separate
+        ownership domains. Once decode rows are handed to the model, remove
+        them from connector accumulation so the next chunk cannot replay the
+        same rows into Talker's ``cached_decode``.
+        """
+        with self._lock:
+            payload = self._local_stage_payload_cache.pop(req_id, None)
+            if payload is None:
+                return None
+            delivered = self._copy_payload_structure(payload)
+            external_req_id = self._request_ids_mapping.get(req_id, req_id)
+            accumulated = self._send_side_request_payload.get(external_req_id)
+            if isinstance(accumulated, dict):
+                embed = accumulated.get("embed")
+                if isinstance(embed, dict):
+                    embed.pop("decode", None)
+                    embed.pop("decode_token_start", None)
+                    embed.pop("decode_token_end", None)
+                hidden_states = accumulated.get("hidden_states")
+                if isinstance(hidden_states, dict):
+                    hidden_states.pop("output", None)
+                ids = accumulated.get("ids")
+                if isinstance(ids, dict):
+                    ids.pop("output", None)
+            return delivered
+
+    def emit_chunks(
+        self,
+        *,
+        req_ids: list[str],
+        inter_stage_outputs: list[Any | None] | None,
+        sampled_token_ids: list[list[int]] | None,
+        terminal_req_ids: set[str],
+    ) -> int:
+        with self._native_output_lock:
+            entries, finished_req_ids = self._build_chunk_entries(
+                req_ids=req_ids,
+                inter_stage_outputs=inter_stage_outputs,
+                sampled_token_ids=sampled_token_ids,
+                terminal_req_ids=terminal_req_ids,
+            )
+            emitted = self._send_chunk_entries(entries)
+            self._cleanup_finished_entries(finished_req_ids)
+            return emitted
+
+    def _build_chunk_entries(
+        self,
+        *,
+        req_ids: list[str],
+        inter_stage_outputs: list[Any | None] | None,
+        sampled_token_ids: list[list[int]] | None,
+        terminal_req_ids: set[str],
+    ) -> tuple[list[tuple[Any, Any | None]], list[str]]:
+        if not getattr(self, "_can_send", True):
+            receive_finished_ids = [req_id for req_id in sorted(terminal_req_ids) if req_id in self._native_requests]
+            return [], receive_finished_ids
+        payload_by_req = {
+            req_id: inter_stage_outputs[index]
+            for index, req_id in enumerate(req_ids)
+            if inter_stage_outputs is not None and index < len(inter_stage_outputs)
+        }
+        sampled_by_req = {
+            req_id: sampled_token_ids[index]
+            for index, req_id in enumerate(req_ids)
+            if sampled_token_ids is not None and index < len(sampled_token_ids)
+        }
+        entries: list[tuple[Any, Any | None]] = []
+        finished_req_ids: list[str] = []
+        for req_id in dict.fromkeys([*req_ids, *sorted(terminal_req_ids)]):
+            state = self._native_requests.get(req_id)
+            if state is None:
+                continue
+            state.output_token_ids.extend(int(token_id) for token_id in sampled_by_req.get(req_id, []))
+            state.finished = req_id in terminal_req_ids
+            payload = payload_by_req.get(req_id)
+            if isinstance(payload, dict):
+                payload = unflatten_payload(payload)
+            if payload is None and not state.finished:
+                continue
+            include_token_history = bool(state.resumable or self._put_req_chunk.get(state.external_req_id, 0) == 0)
+            entries.append(
+                (
+                    state.snapshot(include_token_history=include_token_history),
+                    payload,
+                )
+            )
+            if state.finished:
+                finished_req_ids.append(req_id)
+        return entries, finished_req_ids
+
+    def _send_chunk_entries(
+        self,
+        entries: list[tuple[Any, Any | None]],
+        *,
+        send_lock_held: bool = False,
+    ) -> int:
+        if not entries:
+            return 0
+        context: AbstractContextManager
+        if send_lock_held:
+            context = nullcontext()
+        else:
+            context = self._native_send_lock
+        with context:
+            return self.send_chunks(entries, propagate_errors=True, wait_for_delivery=True)
+
+    def _cleanup_finished_entries(self, finished_req_ids: list[str]) -> None:
+        for req_id in finished_req_ids:
+            self._native_requests.pop(req_id, None)
+            self.cleanup_finished_request(req_id)
+
+    def shutdown_omni_connectors(self) -> None:
+        """Bound MRv2 shutdown even if a connector call cannot be cancelled.
+
+        Connector ``close()`` runs concurrently with the I/O threads so a
+        backend that supports cancellation can release a blocked ``put()``.
+        Backends that do not support it are quarantined and left only on
+        daemon threads after the bounded deadline; V1 keeps its existing
+        shutdown implementation in the shared mixin.
+        """
+        timeout_s = self._delivery_manager.shutdown_timeout_s
+        deadline = time.monotonic() + timeout_s
+        self._delivery_manager.shutdown(RuntimeError("MRv2 connector shutdown"))
+        self._stop_event.set()
+        self._work_available.set()
+
+        close_errors: list[BaseException] = []
+        connector = getattr(self, "_omni_connector", None)
+
+        def close_connector() -> None:
+            if connector is None:
+                return
+            try:
+                connector.close()
+            except BaseException as error:
+                close_errors.append(error)
+
+        close_thread = threading.Thread(
+            target=close_connector,
+            name=f"omni-native-connector-close-{getattr(self, '_stage_id', 'unknown')}",
+            daemon=True,
+        )
+        close_thread.start()
+
+        live_threads: list[str] = []
+        for name, thread in (
+            ("recv", getattr(self, "_recv_thread", None)),
+            ("save", getattr(self, "_save_thread", None)),
+            ("connector-close", close_thread),
+        ):
+            if thread is None:
+                continue
+            thread.join(timeout=max(0.0, deadline - time.monotonic()))
+            if thread.is_alive():
+                live_threads.append(name)
+
+        if close_errors:
+            raise RuntimeError("MRv2 connector close failed") from close_errors[0]
+        if live_threads:
+            raise RuntimeError(
+                f"MRv2 connector shutdown exceeded {timeout_s:.3f}s; blocked thread(s): {', '.join(live_threads)}"
+            )
+
+    def close(self) -> None:
+        error: BaseException | None = None
+        try:
+            self.drain_outputs()
+        except BaseException as exc:
+            error = exc
+        finally:
+            self._delivery_manager.shutdown(error or RuntimeError("MRv2 data plane close"))
+            try:
+                self._stop_output_worker()
+            except BaseException as exc:
+                if error is None:
+                    error = exc
+            try:
+                self.shutdown_omni_connectors()
+            except BaseException as exc:
+                if error is None:
+                    error = exc
+        if error is not None:
+            raise error