[Core] Extend the shared multi-stage execution framework for mixed V1/MRv2 pipelines (#8184)

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

docs/configuration/stage_configs.md

@@ -80,6 +80,19 @@ Note: for the diffusion path, an omitted `distributed_executor_backend` selects
 segments) and `mp` when `num_gpus > 1`. Set `mp` explicitly to keep a worker
 subprocess on one GPU. `ray` / `external_launcher` are not fully supported yet.
 
+### Stage-level runner selection
+
+`model_runner: v1` or `v2` at the deploy level sets the default runner.
+A `model_runner` on an individual stage overrides that default, allowing
+one stage to migrate or roll back independently. This does not change
+`PipelineConfig` topology, input adapters, or the connector payload contract.
+An omitted stage override preserves the deploy-level selection.
+
+MRv2 native downstream receivers currently support turn-based requests only;
+streaming sessions and resumable input require the V1 prompt-replacement path.
+Selecting a runner does not add the session capabilities it lacks. Platform
+fallbacks and stage overrides are validated after configuration resolution.
+
 ### Stage fields
 
 Each entry under `stages:` accepts any `StageDeployConfig` field directly (no nested `engine_args:`). Only fields whose value legitimately varies across stages live here; pipeline-wide settings (trust_remote_code, distributed_executor_backend, dtype, quantization, prefix/chunked prefill, DP/PP sizes) are declared at the top level and applied to every stage. Unknown keys fall through to `engine_extras:` and are forwarded to the engine. Frequently used fields are listed below; the source-of-truth schema is `StageDeployConfig` in `vllm_omni/config/stage_config.py`.

tests/config/test_environment_variables.py

@@ -169,7 +169,7 @@ def test_inventory_matches_reviewed_snapshot_counts():
         EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29,
         EnvironmentVariableCategory.MODEL_SPECIFIC: 66,
         EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21,
-        EnvironmentVariableCategory.INTERNAL: 2,
+        EnvironmentVariableCategory.INTERNAL: 6,
     }
 
     disposition_counts = Counter(

tests/config/test_omni_config.py

@@ -822,6 +822,62 @@ def test_mrv2_fails_fast_on_platforms_without_native_workers(platform: str):
         _apply_platform_overrides(DeployConfig(model_runner="v2"), platform=platform)
 
 
+def _qwen3_omni_pipeline() -> PipelineConfig:
+    from dataclasses import replace
+
+    pipeline = _resolve_pipeline_or_skip("qwen3_omni_moe", Qwen3OmniMoeConfig())
+    return replace(pipeline, stages=tuple(replace(s, supports_native_mrv2_data_plane=True) for s in pipeline.stages))
+
+
+def test_stage_model_runner_overrides_deploy_runner(tmp_path: Path):
+    deploy_path = tmp_path / "qwen3_omni_mixed_runner.yaml"
+    deploy_path.write_text(
+        """\
+model_runner: v1
+async_chunk: true
+stages:
+  - stage_id: 0
+  - stage_id: 1
+    model_runner: v2
+  - stage_id: 2
+    model_runner: v2
+"""
+    )
+    deploy = load_deploy_config(deploy_path)
+    assert [stage.model_runner for stage in deploy.stages] == [None, "v2", "v2"]
+
+    pipeline = _qwen3_omni_pipeline()
+    legacy = merge_pipeline_deploy(pipeline, deploy)
+    assert [stage.yaml_engine_args["use_v2_model_runner"] for stage in legacy] == [False, True, True]
+    # The test pipeline declares native support; only v2 stages use it.
+    assert [stage.yaml_engine_args["supports_native_mrv2_data_plane"] for stage in legacy] == [True, True, True]
+
+    structured = VllmOmniConfig.from_pipeline_config(pipeline, deploy_config_path=str(deploy_path))
+    assert [stage.model_config.use_v2_model_runner for stage in structured.stage_configs] == [False, True, True]
+
+
+def test_platform_v1_fallback_overrides_stage_v2():
+    deploy = DeployConfig(
+        stages=[StageDeployConfig(stage_id=0), StageDeployConfig(stage_id=1, model_runner="v2")],
+        platforms={"rocm": {"model_runner": "v1"}},
+    )
+    deploy = _apply_platform_overrides(deploy, platform="rocm")
+    assert deploy.model_runner == "v1"
+    assert deploy.stages[1].model_runner is None
+
+
+def test_downstream_mrv2_stage_requires_turn_sessions():
+    stages = [StageDeployConfig(stage_id=0), StageDeployConfig(stage_id=1, model_runner="v2")]
+    with pytest.raises(ValueError, match="session_mode 'turn' only"):
+        merge_pipeline_deploy(_qwen3_omni_pipeline(), DeployConfig(stages=stages, session_mode="duplex"))
+
+
+def test_v1_downstream_stage_keeps_duplex_sessions():
+    deploy = DeployConfig(stages=[StageDeployConfig(stage_id=0)], session_mode="duplex")
+    stages = merge_pipeline_deploy(_qwen3_omni_pipeline(), deploy)
+    assert not any(stage.yaml_engine_args["use_v2_model_runner"] for stage in stages)
+
+
 def test_qwen3_tts_high_concurrency_mrv2_profile_is_explicit_opt_in():
     assert load_deploy_config(_DEPLOY_DIR / "qwen3_tts_high_concurrency.yaml").model_runner == "v1"
     assert load_deploy_config(_DEPLOY_DIR / "qwen3_tts_high_concurrency_mrv2.yaml").model_runner == "v2"
@@ -2220,3 +2276,12 @@ def test_mps_stays_in_runtime_instead_of_engine_arguments():
         assert stage.runtime_config.cuda_mps
         args = build_engine_args_dict_from_omni_stage_config(stage, model="test-model")
         assert "cuda_mps" not in args
+
+
+def test_platform_stage_overlay_rejects_invalid_model_runner():
+    deploy = DeployConfig(
+        stages=[StageDeployConfig(stage_id=0)],
+        platforms={"cuda": {"stages": [{"stage_id": 0, "model_runner": "v3"}]}},
+    )
+    with pytest.raises(ValueError, match="model_runner must be 'v1' or 'v2'"):
+        _apply_platform_overrides(deploy, platform="cuda")

tests/core/sched/test_generation_scheduler_restore.py

@@ -25,8 +25,8 @@ class FakeAdapter:
     """Minimal mock of OmniChunkTransferAdapter tracking restore calls."""
 
     def __init__(self):
-        self.waiting_for_chunk_waiting_requests = deque()
-        self.waiting_for_chunk_running_requests = deque()
+        self.waiting_for_chunk_waiting_requests: deque = deque()
+        self.waiting_for_chunk_running_requests: deque = deque()
         self.restore_called = False
         self.done_request_ids = set()
 
@@ -249,3 +249,103 @@ class TestRestoreQueuesOnError:
 
         assert adapter.restore_called is True
         assert "req-B" in running
+
+
+def test_first_chunk_express_slack_guard_tracks_emitted_audio(monkeypatch):
+    """Express steps wait while a ready later chunk's stream is close to underrun."""
+    import vllm_omni.core.sched.omni_generation_scheduler as module
+
+    scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
+    scheduler._first_chunk_express = True
+    scheduler._express_min_slack_s = 0.5
+    scheduler._express_skipped_for_slack = 0
+    scheduler._stream_audio = {}
+    scheduler._chunk_started = {"started", "idle"}
+    ready = SimpleNamespace(
+        request_id="started", num_in_flight_tokens=0, prompt_token_ids=[0] * 28, num_computed_tokens=0
+    )
+    # Started, but no chunk to decode now: never blocks an express step.
+    idle = SimpleNamespace(request_id="idle", num_in_flight_tokens=0, prompt_token_ids=[0] * 28, num_computed_tokens=28)
+    scheduler.running = [ready, idle]
+    scheduler.waiting = []
+    clock = [100.0]
+    monkeypatch.setattr(module.time, "monotonic", lambda: clock[0])
+
+    # No measured audio credit: do not delay the ready stream.
+    assert not scheduler._continuations_have_slack()
+    # 1 s of audio at t=100; at t=100.2 it holds 0.8 s, at t=100.7 only 0.3 s.
+    scheduler._record_stream_audio("started", {"model_outputs": torch.zeros(24000), "sr": torch.tensor(24000)})
+    scheduler._record_stream_audio("started", {"model_outputs": torch.zeros(0), "sr": torch.tensor(24000)})
+    clock[0] = 100.2
+    assert scheduler._continuations_have_slack()
+    clock[0] = 100.7
+    assert not scheduler._continuations_have_slack()
+    # Another second emitted: slack is back to 1.3 s.
+    scheduler._record_stream_audio("started", {"model_outputs": torch.zeros(24000), "sr": torch.tensor(24000)})
+    assert scheduler._continuations_have_slack()
+    assert scheduler._stream_audio["started"] == [100.0, 2.0]
+
+
+def _express_scheduler():
+    continuation = Request("continuation", [1, 2], SamplingParams(max_tokens=4), pooling_params=None)
+    first = Request("first", [1], SamplingParams(max_tokens=4), pooling_params=None)
+    scheduler = _make_generation_scheduler(continuation, use_v2_model_runner=True)
+    scheduler.waiting.add_request(first)
+    scheduler.requests[first.request_id] = first
+    scheduler.max_num_running_reqs = 4
+    scheduler._native_data_plane = True
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = SimpleNamespace(
+        _async_chunk=True, finished_requests=set(), restore_queues=lambda *args, **kwargs: None
+    )
+    scheduler._first_chunk_express = True
+    scheduler._last_step_express = False
+    scheduler._chunk_started = {"continuation"}
+    scheduler._express_min_slack_s = 0
+    scheduler._stream_audio = {}
+    return scheduler, continuation, first
+
+
+def test_express_schedules_only_first_chunks_then_allows_continuations():
+    scheduler, continuation, first = _express_scheduler()
+    output = scheduler.schedule()
+    assert output.num_scheduled_tokens == {"first": 1}
+    assert scheduler._last_step_express
+    assert continuation in list(scheduler.waiting)
+    # Another first chunk arrives before the continuation runs.
+    first.num_computed_tokens = len(first.prompt_token_ids)
+    second = Request("second", [1], SamplingParams(max_tokens=4), pooling_params=None)
+    scheduler.waiting.add_request(second)
+    scheduler.requests[second.request_id] = second
+    output = scheduler.schedule()
+    assert not scheduler._last_step_express
+    assert "continuation" in output.num_scheduled_tokens
+
+
+def test_express_does_not_delay_continuation_for_unready_first_chunk():
+    scheduler, continuation, first = _express_scheduler()
+    first.num_in_flight_tokens = 1
+    output = scheduler.schedule()
+    assert not scheduler._last_step_express
+    assert output.num_scheduled_tokens == {"continuation": 2}
+
+
+def test_express_slack_guard_is_used_by_schedule():
+    scheduler, continuation, first = _express_scheduler()
+    scheduler._express_min_slack_s = 0.5
+    scheduler._express_skipped_for_slack = 0
+    output = scheduler.schedule()
+    assert not scheduler._last_step_express
+    assert "continuation" in output.num_scheduled_tokens
+
+
+def test_express_cancel_cleans_playback_state(monkeypatch):
+    from vllm.v1.core.sched.scheduler import Scheduler
+
+    scheduler, continuation, _ = _express_scheduler()
+    scheduler.input_coordinator = None
+    scheduler._stream_audio["continuation"] = [1.0, 2.0]
+    monkeypatch.setattr(Scheduler, "_free_request", lambda *args: (None, None))
+    scheduler._free_request(continuation)
+    assert "continuation" not in scheduler._chunk_started
+    assert "continuation" not in scheduler._stream_audio

tests/core/sched/test_omni_generation_scheduler_update_session.py

@@ -118,6 +118,8 @@ def test_resumable_generation_stop_marks_segment_boundary() -> None:
     session.num_in_flight_tokens = 1
 
     sched = MagicMock()
+    sched._first_chunk_express = False
+    sched._express_min_slack_s = 0.0
     sched.requests = {session.request_id: session}
     sched.perf_metrics = None
     sched.chunk_transfer_adapter = SimpleNamespace(
@@ -185,6 +187,8 @@ def test_rejected_grammar_finishes_and_frees_generation_request(prompt_complete:
     session.num_in_flight_tokens = 1
 
     sched = MagicMock()
+    sched._first_chunk_express = False
+    sched._express_min_slack_s = 0.0
     sched.requests = {session.request_id: session}
     sched.perf_metrics = None
     sched.chunk_transfer_adapter = SimpleNamespace(

tests/core/sched/test_omni_sched_ec_request_finish.py

@@ -120,6 +120,8 @@ def _make_finish_sched(request: Request) -> MagicMock:
     """MagicMock scheduler driving the real ``update_from_output`` through a
     finishing request; ``_free_request`` reports EC transfer params."""
     sched = MagicMock()
+    sched._first_chunk_express = False
+    sched._express_min_slack_s = 0.0
     sched.requests = {request.request_id: request}
     sched.perf_metrics = None
     sched.defer_block_free = False

tests/core/sched/test_omni_sched_prefill_stats_finalize.py

@@ -53,6 +53,8 @@ def _make_request() -> Request:
 
 def _make_sched(request: Request) -> MagicMock:
     sched = MagicMock()
+    sched._first_chunk_express = False
+    sched._express_min_slack_s = 0.0
     sched.requests = {request.request_id: request}
     sched.perf_metrics = None
     sched.defer_block_free = False

tests/core/sched/test_omni_scheduler_mixin_shared.py

@@ -175,7 +175,7 @@ def test_finished_request_attachment_keeps_ar_abort_policy_explicit(
 ):
     scheduler = _Scheduler()
     scheduler.finished_req_ids_dict = defaultdict(set, {2: {"req-finished"}})
-    outputs = {}
+    outputs: dict = {}
 
     scheduler._attach_finished_request_sets(
         outputs,
@@ -207,3 +207,26 @@ def test_output_helper_preserves_required_nan_counter_default():
     output = scheduler._make_omni_engine_output(request, new_token_ids=[])
 
     assert output.num_nans_in_logits == 0
+
+
+@pytest.mark.parametrize(("role", "coordinated"), [("sender", False), ("receiver", True), (None, True)])
+def test_native_downstream_sender_stage_does_not_wait_for_chunks(role, coordinated):
+    """An orchestrator-fed downstream stage (MiniCPM-o's Talker) owns only an
+    outgoing connector; the native plane must not park it for input chunks."""
+    scheduler = _Scheduler()
+    model_config = object.__new__(OmniModelConfig)
+    model_config.stage_id = 1
+    model_config.async_chunk = True
+    model_config.requires_full_payload_input = False
+    model_config.supports_native_mrv2_data_plane = True
+    model_config.use_v2_model_runner = True
+    model_config.stage_connector_config = {"name": "SharedMemoryConnector", "extra": {"role": role} if role else {}}
+    scheduler.vllm_config = object.__new__(VllmConfig)
+    scheduler.vllm_config.model_config = model_config
+    scheduler.vllm_config.scheduler_config = object.__new__(SchedulerConfig)
+    scheduler.vllm_config.scheduler_config.max_num_seqs = 1
+    scheduler._init_omni_io_scheduling_state()
+
+    assert scheduler._native_data_plane
+    assert (scheduler.input_coordinator is not None) is coordinated
+    assert scheduler._async_chunk_transport_enabled() is coordinated

tests/diffusion/models/qwen_image/test_qwen_image_pipeline_device.py

@@ -64,7 +64,7 @@ class _MoveableModule(nn.Module):
         return [True, True, True]
 
 
-def _run_pipeline_init(monkeypatch, *, enable_cpu_offload, loader_device):
+def _run_pipeline_init(monkeypatch, *, enable_cpu_offload, loader_device, diffusion_offload_config=None):
     import vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image as pipe_mod
     from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImagePipeline
 
@@ -110,14 +110,21 @@ def _run_pipeline_init(monkeypatch, *, enable_cpu_offload, loader_device):
         tf_model_config={},
         quantization_config=None,
         enable_cpu_offload=enable_cpu_offload,
+        diffusion_offload_config=diffusion_offload_config,
         enable_diffusion_pipeline_profiler=False,
     )
     return QwenImagePipeline(od_config=od_config)
 
 
-def test_pipeline_cpu_offload_parks_encoder_vae_on_cpu(monkeypatch):
+@pytest.mark.parametrize("compact", [False, True])
+def test_pipeline_cpu_offload_parks_encoder_vae_on_cpu(monkeypatch, compact):
     """#7555: cpu_offload keeps BF16 encoder/VAE off the GPU during DiT init."""
-    pipe = _run_pipeline_init(monkeypatch, enable_cpu_offload=True, loader_device="cuda")
+    pipe = _run_pipeline_init(
+        monkeypatch,
+        enable_cpu_offload=not compact,
+        loader_device="cuda",
+        diffusion_offload_config={"mode": "module", "components": ["dit"]} if compact else None,
+    )
     assert pipe.text_encoder.placed_device.type == "cpu"
     assert pipe.vae.placed_device.type == "cpu"
     assert pipe.transformer.probe_device_type == "cuda"

tests/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -224,7 +224,7 @@ def test_streaming_payload_can_replace_placeholder_prompt(mocker: MockerFixture)
 
 def test_turn_start_replacement_ignores_accumulated_prompt_capacity(mocker: MockerFixture) -> None:
     request = _streaming_request(mocker, num_computed_tokens=4064)
-    payload = {
+    payload: dict = {
         "ids": {"prompt": [1]},
         "meta": {
             "replace_streaming_prompt": True,
@@ -336,7 +336,7 @@ def test_streaming_window_builds_one_chunk_recompute_recipe(mocker: MockerFixtur
         num_output_placeholders=1,
         update_block_hashes=mocker.Mock(),
     )
-    payload = {
+    payload: dict = {
         "ids": {"prompt": [1]},
         "meta": {
             "next_stage_prompt_len": 10,
@@ -406,7 +406,7 @@ def test_streaming_window_appends_until_capacity_then_recomputes(mocker: MockerF
         num_output_placeholders=0,
         update_block_hashes=mocker.Mock(),
     )
-    second = {
+    second: dict = {
         "ids": {"prompt": [1]},
         "meta": {
             "next_stage_prompt_len": 12,
@@ -432,7 +432,7 @@ def test_streaming_window_appends_until_capacity_then_recomputes(mocker: MockerF
     request._all_token_ids.extend([201, 202, 203])
     request._output_token_ids.extend([201, 202, 203])
     request.num_computed_tokens = 4017
-    third = {
+    third: dict = {
         "ids": {"prompt": [2]},
         "meta": {
             "next_stage_prompt_len": 60,
@@ -458,7 +458,7 @@ def test_streaming_window_appends_until_capacity_then_recomputes(mocker: MockerF
     request._all_token_ids.extend([301])
     request._output_token_ids.extend([301])
     request.num_computed_tokens = 76
-    fourth = {
+    fourth: dict = {
         "ids": {"prompt": [3]},
         "meta": {
             "next_stage_prompt_len": 8,
@@ -1806,7 +1806,7 @@ def test_sender_only_adapter_does_not_park_or_clear_requests(build_adapter):
     request = _req("req-1", RequestStatus.WAITING)
     request.additional_information = {"tts_token_ids": torch.tensor([1])}
     waiting_queue = DummyWaitingQueue([request])
-    running_queue = []
+    running_queue: list = []
 
     adapter.load_async(request)
     adapter.process_pending_chunks(
@@ -1846,7 +1846,7 @@ def test_fifo_promotion(build_adapter):
     adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=2)
     reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)]
     waiting_queue = DummyWaitingQueue(reqs)
-    running_queue = []
+    running_queue: list = []
 
     adapter.process_pending_chunks(waiting_queue, running_queue)
 
@@ -1876,7 +1876,7 @@ def test_non_active_waiting_request_is_held_off_scheduler(build_adapter):
     active = _req("req-active", RequestStatus.WAITING)
     non_active = _req("req-non-active", RequestStatus.WAITING)
     waiting_queue = DummyWaitingQueue([active, non_active])
-    running_queue = []
+    running_queue: list = []
 
     adapter.process_pending_chunks(waiting_queue, running_queue)
 
@@ -1911,7 +1911,7 @@ def test_finished_releases_slot(build_adapter):
     req_1 = _req("req-1", RequestStatus.WAITING)
     req_2 = _req("req-2", RequestStatus.WAITING)
     waiting_queue = DummyWaitingQueue([req_1, req_2])
-    running_queue = []
+    running_queue: list = []
 
     adapter.process_pending_chunks(waiting_queue, running_queue)
     assert list(adapter._active_streams) == ["req-1"]
@@ -2004,7 +2004,7 @@ def test_cleanup_receiver_releases_multiple_slots_in_sequence(build_adapter):
     adapter, _ = build_adapter(stage_id=1, model_mode="generation", max_num_seqs=2, active_stream_window=2)
     reqs = [_req(f"req-{idx}", RequestStatus.WAITING) for idx in range(1, 5)]
     waiting_queue = DummyWaitingQueue(reqs)
-    running_queue = []
+    running_queue: list = []
 
     adapter.process_pending_chunks(waiting_queue, running_queue)
     assert list(adapter._active_streams) == ["req-1", "req-2"]
@@ -2364,7 +2364,7 @@ def test_finish_requests_releases_active_stream_slot(build_adapter):
     aborted = _req("req-aborted", RequestStatus.RUNNING)
     waiting = _req("req-waiting", RequestStatus.WAITING)
     waiting_queue = DummyWaitingQueue([waiting])
-    running_queue = []
+    running_queue: list = []
     adapter._active_streams[aborted.request_id] = aborted
     adapter._held_non_active.append(aborted)
 
@@ -2488,7 +2488,7 @@ def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_ada
     zombie = _req("req-zombie", RequestStatus.WAITING_FOR_CHUNK)
     live = _req("req-live", RequestStatus.WAITING_FOR_CHUNK)
     waiting_queue = DummyWaitingQueue()
-    running_queue = []
+    running_queue: list = []
     adapter.waiting_for_chunk_waiting_requests = deque([zombie, live])
     adapter.waiting_for_chunk_running_requests = deque([zombie, live])
 
@@ -2522,7 +2522,7 @@ class _HashableRequest(SimpleNamespace):
     # double needs it too.
     num_stale_output_tokens = 0
 
-    def __hash__(self):
+    def __hash__(self):  # type: ignore[override]  # Test requests are hashable by request ID.
         return hash(self.request_id)
 
     def __eq__(self, other):
@@ -2540,6 +2540,8 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
     from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
 
     scheduler = mocker.MagicMock()
+    scheduler._first_chunk_express = False
+    scheduler._express_min_slack_s = 0.0
     scheduler.chunk_transfer_adapter = adapter_mock
     scheduler.connector = None
     scheduler.ec_connector = None

tests/distributed/omni_connectors/test_shm_connector.py

@@ -294,3 +294,280 @@ def test_consumed_bad_payload_removes_lock_file(connector, monkeypatch):
     assert not os.path.exists(f"/dev/shm/shm_{key}_lockfile.lock")
     connector.reap_consumed()
     assert key not in connector._pending_keys
+
+
+# ── Arrival wakeups ───────────────────────────────────────────────────
+
+
+def _stage_connector(stage_id):
+    return SharedMemoryConnector({"stage_id": stage_id})
+
+
+def test_wakeup_discovery_only_scans_deployment_directory(monkeypatch, tmp_path):
+    from vllm_omni.distributed.omni_connectors.connectors import shm_connector as module
+
+    directory = tmp_path / "deployment"
+    monkeypatch.setattr(module, "_wakeup_directory", lambda: str(directory))
+    original_scandir = os.scandir
+    scanned = []
+
+    def scoped_scandir(path):
+        # A parent-directory scan would recreate the performance regression.
+        assert os.fspath(path) == str(directory)
+        scanned.append(path)
+        return original_scandir(path)
+
+    monkeypatch.setattr(module.glob.os, "scandir", scoped_scandir)
+    sender, receiver = _stage_connector(86), _stage_connector(87)
+    try:
+        generation = receiver.get_wakeup_generation()
+        assert generation is not None
+        assert os.path.dirname(receiver._wake_path) == str(directory)
+        sender._wake_receiver(87)
+        assert receiver.wait_for_change(generation, timeout=0.5)
+        assert scanned
+    finally:
+        sender.close()
+        receiver.close()
+
+
+def test_missing_wakeup_directory_keeps_payload_transport_working(monkeypatch, tmp_path):
+    from vllm_omni.distributed.omni_connectors.connectors import shm_connector as module
+
+    directory = tmp_path / "missing"
+    monkeypatch.setattr(module, "_wakeup_directory", lambda: str(directory))
+    sender, receiver = _stage_connector(88), _stage_connector(89)
+    key = f"wake_{uuid.uuid4().hex}"
+    try:
+        assert sender.put("88", "89", key, {"value": 1})[0]
+        assert not directory.exists()
+        assert receiver.get("88", "89", key)[0] == {"value": 1}
+    finally:
+        sender.close()
+        receiver.close()
+
+
+def test_put_wakes_waiting_receiver():
+    sender, receiver = _stage_connector(70), _stage_connector(71)
+    try:
+        generation = receiver.get_wakeup_generation()
+        assert generation is not None
+        start = time.monotonic()
+        assert receiver.wait_for_change(generation, timeout=0.05) is False
+        assert time.monotonic() - start >= 0.04  # no arrival: blocks for the timeout
+
+        key = f"wake_{uuid.uuid4().hex}"
+        sender.put("70", "71", key, {"value": 1})
+        start = time.monotonic()
+        assert receiver.wait_for_change(generation, timeout=5) is True
+        assert time.monotonic() - start < 0.5
+        assert receiver.get("70", "71", key)[0] == {"value": 1}
+        # Drained: the next wait blocks again.
+        assert receiver.wait_for_change(receiver.get_wakeup_generation(), timeout=0.02) is False
+    finally:
+        sender.close()
+        receiver.close()
+
+
+def test_receiver_does_not_spin_after_sender_closes():
+    sender, receiver = _stage_connector(72), _stage_connector(73)
+    try:
+        receiver.get_wakeup_generation()
+        sender.put("72", "73", f"wake_{uuid.uuid4().hex}", {"value": 1})
+        sender.close()  # the only external writer goes away
+        generation = receiver.get_wakeup_generation()
+        receiver.wait_for_change(generation, timeout=1)
+        generation = receiver.get_wakeup_generation()
+        start = time.monotonic()
+        assert receiver.wait_for_change(generation, timeout=0.05) is False
+        assert time.monotonic() - start >= 0.04
+    finally:
+        receiver.close()
+
+
+def test_wakeups_can_be_disabled_and_close_unlinks_fifo(monkeypatch):
+    receiver = _stage_connector(74)
+    assert receiver.get_wakeup_generation() is not None
+    path = receiver._wake_path
+    assert os.path.exists(path)
+    receiver.close()
+    assert not os.path.exists(path)
+
+    monkeypatch.setenv("VLLM_OMNI_SHM_WAKEUP", "0")
+    sender, receiver = _stage_connector(75), _stage_connector(76)
+    try:
+        assert receiver.get_wakeup_generation() is None  # caller keeps its timed poll
+        key = f"wake_{uuid.uuid4().hex}"
+        assert sender.put("75", "76", key, {"value": 2})[0]
+        assert receiver.get("75", "76", key)[0] == {"value": 2}
+    finally:
+        sender.close()
+        receiver.close()
+
+
+def _sibling_receiver(stage_id, ready, result):
+    receiver = _stage_connector(stage_id)
+    generation = receiver.get_wakeup_generation()
+    ready.set()
+    start = time.monotonic()
+    woke = receiver.wait_for_change(generation, timeout=5)
+    result.put((woke, time.monotonic() - start))
+    receiver.close()
+
+
+def _sibling_sender(stage_id, key):
+    sender = _stage_connector(stage_id - 1)
+    sender.put(str(stage_id - 1), str(stage_id), key, {"value": 3})
+
+
+def test_sibling_stage_processes_share_wakeups():
+    # Stage engine processes are siblings under the launching process.
+    import multiprocessing
+
+    ctx = multiprocessing.get_context("fork")
+    ready, result = ctx.Event(), ctx.Queue()
+    key = f"wake_{uuid.uuid4().hex}"
+    receiver = ctx.Process(target=_sibling_receiver, args=(78, ready, result))
+    receiver.start()
+    assert ready.wait(5)
+    time.sleep(0.05)
+    sender = ctx.Process(target=_sibling_sender, args=(78, key))
+    sender.start()
+    sender.join(5)
+    woke, waited = result.get(timeout=5)
+    receiver.join(5)
+    assert woke and waited < 1.0
+    cleanup = SharedMemoryConnector({})
+    cleanup.cleanup(key)
+
+
+def test_same_stage_receivers_are_independent_across_close_and_restart():
+    sender = _stage_connector(80)
+    first, second = _stage_connector(81), _stage_connector(81)
+    replacement = None
+    try:
+        first_generation = first.get_wakeup_generation()
+        second_generation = second.get_wakeup_generation()
+        assert first._wake_path != second._wake_path
+        sender._wake_receiver(81)
+        assert first.wait_for_change(first_generation, timeout=0.5)
+        assert second.wait_for_change(second_generation, timeout=0.5)
+        second_generation = second.get_wakeup_generation()
+        first.close()
+        assert first.get_wakeup_generation() is None
+        assert os.path.exists(second._wake_path)
+        replacement = _stage_connector(81)
+        replacement_generation = replacement.get_wakeup_generation()
+        sender._wake_receiver(81)
+        assert second.wait_for_change(second_generation, timeout=0.5)
+        assert replacement.wait_for_change(replacement_generation, timeout=0.5)
+    finally:
+        sender.close()
+        first.close()
+        second.close()
+        if replacement is not None:
+            replacement.close()
+
+
+def test_partial_receiver_open_closes_fd_and_removes_only_own_fifo(monkeypatch, tmp_path):
+    import errno
+
+    from vllm_omni.distributed.omni_connectors.connectors import shm_connector as module
+
+    monkeypatch.setattr(module, "_wakeup_directory", lambda: str(tmp_path))
+    receiver = _stage_connector(82)
+    peer_path = tmp_path / "stage_82_peer"
+    os.mkfifo(peer_path, 0o600)
+    original_open = os.open
+    opened = []
+
+    def fail_hold_open(path, flags, *args, **kwargs):
+        if flags & os.O_WRONLY:
+            raise OSError(errno.EMFILE, "injected fd exhaustion")
+        fd = original_open(path, flags, *args, **kwargs)
+        opened.append(fd)
+        return fd
+
+    monkeypatch.setattr(module.os, "open", fail_hold_open)
+    try:
+        assert receiver.get_wakeup_generation() is None
+        assert list(tmp_path.iterdir()) == [peer_path]
+        assert len(opened) == 1
+        with pytest.raises(OSError, match="Bad file descriptor"):
+            os.fstat(opened[0])
+    finally:
+        receiver.close()
+
+
+def test_wakeup_ignores_stale_fifo_and_non_fifo_paths(monkeypatch, tmp_path):
+    from vllm_omni.distributed.omni_connectors.connectors import shm_connector as module
+
+    monkeypatch.setattr(module, "_wakeup_directory", lambda: str(tmp_path))
+    sender, receiver = _stage_connector(83), _stage_connector(84)
+    plain = tmp_path / "84_plain"
+    plain.write_text("unchanged")
+    (tmp_path / "84_link").symlink_to(plain)
+    os.mkfifo(tmp_path / "84_stale", 0o600)
+    try:
+        generation = receiver.get_wakeup_generation()
+        sender._wake_receiver(84)
+        assert receiver.wait_for_change(generation, timeout=0.5)
+        assert plain.read_text() == "unchanged"
+        assert (tmp_path / "84_stale").exists()
+    finally:
+        sender.close()
+        receiver.close()
+
+
+def test_waiter_does_not_read_fd_reused_after_close(monkeypatch, tmp_path):
+    from vllm_omni.distributed.omni_connectors.connectors import shm_connector as module
+
+    receiver = _stage_connector(85)
+    generation = receiver.get_wakeup_generation()
+    read_fd = receiver._wake_read_fd
+    path = tmp_path / "unrelated_ipc"
+    path.write_bytes(b"untouched")
+    reused = []
+
+    def close_during_select(*args):
+        receiver.close()
+        fd = os.open(path, os.O_RDONLY)
+        if fd != read_fd:
+            os.dup2(fd, read_fd)
+            os.close(fd)
+        reused.append(read_fd)
+        return [read_fd], [], []
+
+    monkeypatch.setattr(module.select, "select", close_during_select)
+    try:
+        assert not receiver.wait_for_change(generation, timeout=0.1)
+        assert os.lseek(read_fd, 0, os.SEEK_CUR) == 0
+    finally:
+        for fd in reused:
+            os.close(fd)
+        receiver.close()
+
+
+def test_deployment_scope_survives_different_worker_parents(monkeypatch):
+    config = {"stage_id": 1, "extra": {"wakeup_scope": uuid.uuid4().hex}}
+    monkeypatch.setattr(os, "getppid", lambda: 111)
+    receiver = SharedMemoryConnector(config)
+    monkeypatch.setattr(os, "getppid", lambda: 222)
+    sender = SharedMemoryConnector(config)
+    try:
+        generation = receiver.get_wakeup_generation()
+        sender._wake_receiver(1)
+        assert receiver.wait_for_change(generation, timeout=0.1)
+        directory = receiver._wake_directory
+    finally:
+        receiver.close()
+        sender.close()
+    assert not os.path.exists(directory)
+
+
+@pytest.mark.parametrize("value,expected", [("bad", 0.005), ("nan", 0.005), ("-1", 0.005), ("20", 0.02)])
+def test_receive_poll_interval_validation(monkeypatch, value, expected):
+    from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_payload_transport import _recv_poll_seconds
+
+    monkeypatch.setenv("VLLM_OMNI_CONNECTOR_RECV_POLL_MS", value)
+    assert _recv_poll_seconds() == expected

tests/engine/test_orchestrator_error_handling.py

@@ -519,6 +519,40 @@ async def test_async_chunk_prewarm_without_prompt_token_ids_fails_only_that_requ
             orchestrator_fixture.thread.join(timeout=5)
 
 
+@pytest.mark.asyncio
+async def test_streaming_input_to_mrv2_downstream_stage_fails_with_client_error(orchestrator_factory) -> None:
+    """A downstream MRv2 native-data-plane stage supports turn-based requests
+    only: a resumable (realtime) request must fail up front, not hang."""
+    stage0 = FakeStageClient(stage_type="llm", final_output=False)
+    stage1 = FakeStageClient(stage_type="llm", final_output=True)
+    v1 = SimpleNamespace(model_config=SimpleNamespace(max_model_len=64))
+    mrv2 = SimpleNamespace(
+        model_config=SimpleNamespace(max_model_len=64, use_v2_model_runner=True, supports_native_mrv2_data_plane=True)
+    )
+    fixture = orchestrator_factory([stage0, stage1], stage_vllm_configs=[v1, mrv2], async_chunk=True)
+
+    try:
+        await _enqueue_add_request(
+            fixture,
+            request_id="req-stream",
+            prompt=SimpleNamespace(request_id="req-stream", prompt_token_ids=[1, 2], resumable=True),
+            original_prompt={"prompt": "stream"},
+            sampling_params_list=[_sampling_params(), _sampling_params()],
+            final_stage_id=1,
+        )
+
+        error_msg = await _wait_for_error_message(fixture, request_id="req-stream")
+        assert error_msg.fatal is False
+        assert error_msg.status_code == 400
+        assert "stage 1 runs on model_runner v2" in error_msg.error
+        assert stage0.add_request_calls == [] and stage1.add_request_calls == []
+        assert fixture.thread.is_alive()
+    finally:
+        if fixture.thread.is_alive():
+            fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
+            fixture.thread.join(timeout=5)
+
+
 # ───────── Direct unit tests for the fault-isolation helpers ─────────
 
 

tests/worker/test_device_copy.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from vllm_omni.utils.device_copy import to_device_nonblocking
+
+pytestmark = pytest.mark.core_model
+
+
+@pytest.mark.cpu
+def test_tensor_copy_cpu_preserves_view_and_dtype():
+    source = torch.arange(12, dtype=torch.float64).reshape(3, 4).T
+    result = to_device_nonblocking(source, "cpu")
+    assert result is source
+    assert result.dtype == torch.float64
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_tensor_copy_cuda_handles_pageable_and_pinned_sources():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    source = torch.arange(12, dtype=torch.float32).reshape(3, 4).T
+    for host in (source, source.pin_memory()):
+        result = to_device_nonblocking(host, "cuda")
+        torch.testing.assert_close(result.cpu(), host)
+        assert to_device_nonblocking(result, result.device) is result

tests/worker/test_omni_gpu_model_runner.py

@@ -607,6 +607,39 @@ def test_update_additional_information_deserializes_new_request_payload():
     )
 
 
+@pytest.mark.parametrize("replace_payload", [False, True])
+def test_resumed_streaming_payload_converts_to_typed_model_input(replace_payload):
+    from vllm_omni.core.sched.output import OmniNewRequestData
+    from vllm_omni.data_entry_keys import to_struct
+
+    runner = _make_runner(req_ids=("r1",), hidden_size=4)
+    runner.model.replace_runtime_additional_information = replace_payload
+    payload = {
+        "codes": {"audio": torch.tensor([1, 2])},
+        "meta": {"req_id": ["r1"], "stream_finished": torch.tensor(False)},
+    }
+    new_req = OmniNewRequestData(
+        req_id="r1",
+        prompt_token_ids=[1, 2],
+        mm_features=[],
+        sampling_params=None,
+        pooling_params=None,
+        block_ids=([],),
+        num_computed_tokens=0,
+        lora_request=None,
+        model_intermediate_buffer=payload,
+    )
+
+    OmniGPUModelRunner._update_streaming_input_additional_info(runner, new_req, "r1")
+    # CosyVoice and other typed-payload consumers perform this conversion
+    # after the runner adds its resume metadata.
+    result = to_struct(runner.model_intermediate_buffer["r1"])
+    assert result.meta.resumable is True
+    assert result.meta.num_processed_tokens == 0
+    assert result.meta.req_id == ["r1"]
+    torch.testing.assert_close(result.codes.audio, payload["codes"]["audio"])
+
+
 def test_streaming_new_request_marker_replaces_terminal_chunk_snapshot():
     from vllm_omni.engine.serialization import serialize_additional_information
 
@@ -1222,3 +1255,30 @@ def test_decode_batch_without_mtp_rejects_incomplete_outputs(monkeypatch, invali
     with pytest.raises(ValueError, match="decode preprocessing"):
         runner._preprocess(scheduled, 2)
     assert torch.equal(runner.inputs_embeds.gpu, before)
+
+
+@pytest.mark.parametrize("supports_boundaries", [False, True])
+def test_multimodal_embeddings_receive_opt_in_host_request_boundaries(supports_boundaries):
+    runner = object.__new__(OmniGPUModelRunner)
+    received = {}
+
+    def embed(ids, **kwargs):
+        received.update(kwargs)
+        return ids
+
+    runner.model = SimpleNamespace(supports_embed_input_ids_query_start_loc=supports_boundaries, embed_input_ids=embed)
+    runner.input_batch = SimpleNamespace(num_reqs=3)
+    runner.query_start_loc = SimpleNamespace(cpu=torch.tensor([0, 6, 7, 11, 999]))
+    runner.input_ids = SimpleNamespace(gpu=torch.arange(16))
+    mm = [torch.ones(2, 4), torch.ones(1, 4)]
+    mask = torch.tensor([True] * 4 + [False] * 3 + [True] * 3 + [False])
+
+    result = runner._embed_multimodal_input_ids(11, mm, mask)
+
+    torch.testing.assert_close(result, torch.arange(11))
+    assert received["multimodal_embeddings"] is mm
+    assert received["is_multimodal"] is mask
+    if supports_boundaries:
+        assert received["query_start_loc"] == [0, 6, 7, 11]
+    else:
+        assert "query_start_loc" not in received

tests/worker_v2/test_intermediate_buffer.py

@@ -94,3 +94,25 @@ def test_update_merge_semantics(monkeypatch):
     assert stored is not tensor
     assert torch.equal(stored, tensor)
     assert "kv" in buf.buffers[0]
+
+
+def test_add_request_merges_orchestrator_model_intermediate_buffer():
+    """Stage bridges such as MiniCPM-o's llm2tts hand over ``model_intermediate_buffer``."""
+    buf = OmniIntermediateBuffer(max_num_reqs=2)
+    hidden = torch.randn(3, 4)
+    data = _make_new_req_data(
+        "talker-1",
+        additional_information={"meta": {"a": 1}},
+        model_intermediate_buffer={
+            "ids": {"tts": [1, 2, 3]},
+            "hidden_states": {"tts": hidden},
+            "meta": {"next_stage_prompt_len": 5},
+            "request_id": "thinker-1",
+        },
+    )
+    buf.add_request(0, data)
+    info = buf.buffers[0]
+    assert info["ids"]["tts"] == [1, 2, 3]
+    assert torch.equal(info["hidden_states"]["tts"], hidden)
+    assert info["meta"] == {"a": 1, "next_stage_prompt_len": 5}
+    assert info["request_id"] == "thinker-1" and info["req_id"] == "talker-1"

tests/worker_v2/test_omni_data_plane.py

@@ -185,6 +185,7 @@ def _start_save_thread(plane, connector):
         _send_side_request_payload={},
         _code_prompt_token_ids=defaultdict(list),
         _work_available=threading.Event(),
+        _save_work_available=threading.Event(),
         _stop_event=threading.Event(),
         _MAX_SEND_RETRIES=0,
         _can_send=True,
@@ -199,7 +200,7 @@ def _start_save_thread(plane, connector):
 
 def _stop_save_thread(plane):
     plane._stop_event.set()
-    plane._work_available.set()
+    plane._save_work_available.set()
     plane._save_thread.join(timeout=2)
     assert not plane._save_thread.is_alive()
 
@@ -551,3 +552,71 @@ def test_abort_before_first_chunk_cleans_receiver_state():
     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
+
+
+def test_ar_receiver_accumulates_rows_and_replaces_chunk_metadata():
+    # A V1 upstream sender stamps scalar flags on every chunk; they describe
+    # the latest chunk and must not be concatenated like decode rows.
+    plane = _bare_plane()
+    plane._send_side_request_payload = {}
+
+    def chunk(rows, finished):
+        return {
+            "embed": {"decode": torch.ones((rows, 2))},
+            "meta": {"finished": torch.tensor(finished), "is_segment_finished": torch.tensor(False)},
+        }
+
+    plane._accumulate_payload("ext", chunk(1, False))
+    merged = plane._accumulate_payload("ext", chunk(2, True))
+
+    assert merged["embed"]["decode"].shape == (3, 2)
+    assert merged["meta"]["finished"].item() is True
+    assert merged["meta"]["is_segment_finished"].ndim == 0
+
+
+def test_terminal_without_processor_payload_still_sends_finish_marker(raw_plane, monkeypatch):
+    # Qwen3-Omni's Talker processor returns None on the terminal call when the
+    # last frames ended exactly on a chunk boundary; Code2Wav must still learn
+    # that the stream finished, as with the V1 chunk adapter.
+    raw_plane._async_chunk = True
+    raw_plane._omni_connector = object()
+    raw_plane._request_ids_mapping = {}
+    raw_plane._custom_process_func = lambda transfer_manager, multimodal_output, request, is_finished=False: None
+    raw_plane._custom_process_batch_func = None
+    monkeypatch.setattr(raw_plane, "is_data_transfer_rank", lambda: True)
+    enqueued = []
+
+    def enqueue(request, payload, **_kw):
+        enqueued.append((request.request_id, payload))
+        return True, None
+
+    monkeypatch.setattr(raw_plane, "_enqueue_chunk_payload", enqueue)
+    raw_plane.register_request(_new_request())
+
+    _complete(raw_plane, [{"codes": {"audio": torch.ones((1, 16))}}])
+    assert enqueued == []  # a non-terminal None payload sends nothing
+    assert raw_plane.request_terminal({"internal"}) == 1
+    [(req_id, payload)] = enqueued
+    assert req_id == "internal"
+    assert payload["meta"]["finished"].item() is True
+
+
+def test_payload_builders_see_stage_model_config(monkeypatch):
+    # Producer-side builders (Qwen3-Omni's Thinker reads
+    # hf_config.talker_config.accept_hidden_layer) resolve the stage config
+    # through _get_model_config(), as they do on the V1 runner mixin.
+    from vllm_omni.model_executor.stage_input_processors.qwen3_omni import _get_accept_hidden_layer_index
+
+    def init_connectors(self, *, model_config):
+        self._async_chunk = True
+        self._custom_process_func = None
+
+    monkeypatch.setattr(OmniRunnerDataPlane, "init_omni_connectors", init_connectors)
+    monkeypatch.setattr(OmniRunnerDataPlane, "_start_output_worker", lambda self, **_kw: None)
+    model_config = SimpleNamespace(hf_config=SimpleNamespace(talker_config=SimpleNamespace(accept_hidden_layer=24)))
+    vllm_config = SimpleNamespace(model_config=model_config)
+
+    plane = OmniRunnerDataPlane(vllm_config, model_config)
+
+    assert plane._get_model_config() is model_config
+    assert _get_accept_hidden_layer_index(plane) == 24

tests/worker_v2/test_omni_generation_model_runner.py

@@ -13,7 +13,7 @@ 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
+from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner, check_exact_input_shape
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -182,3 +182,16 @@ def test_sample_tokens_uses_async_output_for_cuda_and_snapshots_req_ids(monkeypa
     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"]
+
+
+@pytest.mark.parametrize("exact", [True, False])
+def test_exact_input_shape_models_reject_padded_batches(exact):
+    # Qwen3-Omni Code2Wav splits input_ids by seq_token_counts (#6712).
+    model = SimpleNamespace(requires_exact_input_shape=exact)
+    check_exact_input_shape(model, torch.zeros(32, dtype=torch.long), 32)
+    check_exact_input_shape(model, None, 16)
+    if exact:
+        with pytest.raises(RuntimeError, match="exact input shape"):
+            check_exact_input_shape(model, torch.zeros(48, dtype=torch.long), 32)
+    else:
+        check_exact_input_shape(model, torch.zeros(48, dtype=torch.long), 32)

tests/worker_v2/test_omni_gpu_model_runner.py

@@ -181,6 +181,84 @@ def test_capture_contract_uses_model_declaration(stage, declared):
     runner._configure_cudagraph_output_contract()
     assert runner._model_returns_tuple is declared
     assert runner._exclude_full_graph is declared
+    assert runner._full_graph_aux_outputs is False
+
+
+def _aux_output(num_tokens):
+    hidden = torch.arange(num_tokens * 2, dtype=torch.float32).reshape(num_tokens, 2)
+    return hidden, {"hidden_states": {"layers": {0: hidden + 1, 24: hidden + 2}}}
+
+
+def test_full_graph_aux_contract_keeps_full_and_round_trips_leaves():
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.model = SimpleNamespace(_returns_tuple=True, supports_mrv2_full_graph_aux_outputs=True)
+    runner._configure_cudagraph_output_contract()
+    assert runner._full_graph_aux_outputs and not runner._exclude_full_graph
+
+    calls = []
+
+    def forward(num_tokens=3):
+        calls.append(num_tokens)
+        return _aux_output(num_tokens)
+
+    runner.model.forward = forward
+    original_forward = runner.model.forward
+    runner.use_aux_hidden_state_outputs = False
+    full = SimpleNamespace(cg_mode=CUDAGraphMode.FULL)
+    runner.cudagraph_manager = SimpleNamespace(_capture_descs={CUDAGraphMode.FULL: [full]}, _candidates={})
+    runner.model_state = SimpleNamespace()
+    captured = {}
+
+    def capture(_self):
+        # The graph manager stores (hidden, leaves) in its aux buffers.
+        assert runner.use_aux_hidden_state_outputs is True
+        hidden, leaves = runner.model.forward(4)
+        captured["hidden"], captured["leaves"] = hidden, leaves
+        assert [tuple(leaf.shape) for leaf in leaves] == [(4, 2), (4, 2)]
+        runner.model.forward(2)  # smaller capture, same structure
+        return 1
+
+    with patch.object(GPUModelRunner, "capture_model", capture):
+        assert runner.capture_model() == 1
+    assert runner.model.forward is original_forward and runner.use_aux_hidden_state_outputs is False
+    assert runner.cudagraph_manager._capture_descs == {CUDAGraphMode.FULL: [full]}  # FULL kept
+
+    hidden, aux = runner._split_fullgraph_output((captured["hidden"], captured["leaves"]))
+    reference_hidden, reference_aux = _aux_output(4)
+    assert torch.equal(hidden, reference_hidden)
+    assert set(aux["hidden_states"]["layers"]) == {0, 24}
+    assert torch.equal(aux["hidden_states"]["layers"][24], reference_aux["hidden_states"]["layers"][24])
+
+
+@pytest.mark.parametrize(
+    "aux",
+    [
+        {},  # no leaves
+        {"layers": {0: torch.zeros(5, 2)}},  # leaf not on the token axis
+        {"layers": {0: 1.0}},  # non-tensor leaf
+    ],
+)
+def test_full_graph_aux_contract_rejects_invalid_outputs(aux):
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.model = SimpleNamespace(_returns_tuple=True, supports_mrv2_full_graph_aux_outputs=True)
+    runner._configure_cudagraph_output_contract()
+    with pytest.raises(RuntimeError, match="supports_mrv2_full_graph_aux_outputs"):
+        runner._flatten_capture_aux_output((torch.zeros(3, 2), aux))
+
+
+def test_full_graph_aux_contract_rejects_structure_change():
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.model = SimpleNamespace(_returns_tuple=True, supports_mrv2_full_graph_aux_outputs=True)
+    runner._configure_cudagraph_output_contract()
+    runner._flatten_capture_aux_output((torch.zeros(3, 2), {"layers": {0: torch.zeros(3, 2)}}))
+    with pytest.raises(RuntimeError, match="structure changed"):
+        runner._flatten_capture_aux_output((torch.zeros(3, 2), {"layers": {1: torch.zeros(3, 2)}}))
+
+
+def test_full_graph_output_without_contract_is_hidden_only():
+    runner = object.__new__(OmniGPUModelRunner)
+    hidden = torch.ones(2, 2)
+    assert runner._split_fullgraph_output(hidden) == (hidden, None)
 
 
 @pytest.mark.parametrize("runner_kind", ["gpu", "ar", "generation"])
@@ -202,6 +280,7 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch):
         else hidden
     )
     runner._dummy_hidden = hidden
+    runner.model.requires_request_ids = False
     runner.model_config = SimpleNamespace()
     runner.vllm_config = SimpleNamespace()
     runner.req_states = SimpleNamespace()

tests/worker_v2/test_omni_model_state.py

@@ -19,6 +19,8 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 class _DummyInputBatch:
+    is_prefilling_np: np.ndarray
+
     input_ids: SimpleNamespace
     query_start_loc: torch.Tensor
 
@@ -592,3 +594,58 @@ def test_first_audio_marker_requires_accepted_delivery(monkeypatch, accepted):
     # A missing route must leave the normal codec path responsible for frame 0;
     # otherwise it skips the frame and the orchestrator waits forever for it.
     assert outputs["meta"]["first_audio"].tolist() == ["r0" in accepted, "r1" in accepted]
+
+
+class _FakeEvent:
+    def record(self, stream=None) -> None:
+        self.stream = stream
+
+
+def test_publish_sampled_embeddings_for_rows_whose_sample_is_kept(monkeypatch) -> None:
+    monkeypatch.setattr(torch.cuda, "Event", _FakeEvent)
+    state = _make_state()
+    state.model.publishes_sampled_embeddings = True
+    state.model.embed_input_ids = lambda ids: ids.to(torch.float32).reshape(-1, 1).repeat(1, 3)
+    # Rows: final prefill chunk, decode, non-final prefill chunk, one-chunk prefill.
+    batch = SimpleNamespace(
+        num_reqs=4,
+        num_scheduled_tokens=[4, 1, 2, 6],
+        num_computed_prefill_tokens_np=np.array([2, 9, 0, 0], dtype=np.int32),
+        prefill_len_np=np.array([6, 9, 5, 6], dtype=np.int32),
+        is_prefilling_np=np.array([True, False, True, True]),
+    )
+    extra, _done = state.publish_sampled_embeddings(batch, torch.tensor([[11], [12], [13], [14]]))
+    sampled = extra["embed"]["sampled"]
+    # A non-final prefill chunk's sample is discarded, so it publishes nothing.
+    assert [tuple(row.shape) for row in sampled] == [(1, 3), (1, 3), (0,), (1, 3)]
+    assert [row.tolist() for row in sampled if row.numel()] == [[[v] * 3] for v in (11.0, 12.0, 14.0)]
+
+    # Speculative steps sample several tokens per row: not published.
+    assert state.publish_sampled_embeddings(batch, torch.tensor([[11, 1], [12, 1], [13, 1], [14, 1]])) is None
+
+
+def test_publish_sampled_embeddings_is_opt_in() -> None:
+    state = _make_state()
+    state.model.publishes_sampled_embeddings = False
+    assert state.publish_sampled_embeddings(SimpleNamespace(num_reqs=1), torch.tensor([[1]])) is None
+
+
+@pytest.mark.parametrize("prefilling", [False, True])
+def test_identity_preprocess_skips_only_decode_rows(prefilling):
+    state = _make_state(has_preprocess=True)
+    state._decode_preprocess_is_identity = True
+    state._decode_preprocess = None
+    state.intermediate_buffer.buffers[0] = {"req_id": "r1"}
+    seen = []
+
+    def preprocess(input_ids, input_embeds, **info):
+        seen.append(info["req_id"])
+        return input_ids, input_embeds, {}
+
+    state.model.preprocess = preprocess
+    batch = _DummyInputBatch([0])
+    batch.is_prefilling_np = np.array([prefilling])
+    embeds = torch.ones(1, 4)
+    state.run_preprocess(batch, {"input_ids": torch.tensor([1]), "inputs_embeds": embeds})
+    assert seen == (["r1"] if prefilling else [])
+    assert torch.equal(embeds, torch.ones(1, 4))

tests/worker_v2/test_shared_stage_contracts.py

@@ -0,0 +1,79 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Generic model opt-in contracts; no model checkpoints required."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
+from vllm_omni.worker_v2.omni_ar_model_runner import _merge_payload_trees
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_device_output_hook_receives_live_batch_and_request_state():
+    state = object.__new__(OmniModelState)
+    state.have_multimodal_outputs = True
+    batch, requests, buffers = object(), object(), [{"req_id": "r"}]
+    hidden = torch.ones(1, 2)
+    seen = []
+
+    def make_output(value, **kwargs):
+        seen.append((value, kwargs))
+        return OmniOutput(text_hidden_states=value, multimodal_outputs={"codes": {"audio": value}})
+
+    state.model = SimpleNamespace(make_omni_output_mrv2=make_output)
+    state.intermediate_buffer = SimpleNamespace(gather=lambda value: buffers)
+    result_hidden, payload = state.postprocess_model_output(hidden, batch, requests)
+    assert result_hidden is hidden
+    assert payload["codes"]["audio"] is hidden
+    assert seen == [(hidden, {"input_batch": batch, "req_states": requests, "model_intermediate_buffer": buffers})]
+
+
+def test_sampled_embeddings_merge_preserves_sibling_fields_and_input():
+    prefill, sampled = torch.ones(2, 3), torch.zeros(1, 3)
+    base = {"embed": {"prefill": prefill}, "meta": {"finished": False}}
+    merged = _merge_payload_trees(base, {"embed": {"sampled": [sampled]}})
+    assert merged["embed"]["prefill"] is prefill
+    assert merged["embed"]["sampled"] == [sampled]
+    assert "sampled" not in base["embed"]
+    assert merged["meta"] == {"finished": False}
+
+
+def test_custom_sampler_falls_back_to_parent(monkeypatch):
+    from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+
+    state = object.__new__(OmniModelState)
+    state.model = SimpleNamespace()
+    sampler, marker = object(), object()
+    monkeypatch.setattr(DefaultModelState, "custom_sampler", lambda self, original: (original, marker))
+    assert state.custom_sampler(sampler) == (sampler, marker)
+
+
+def test_logits_vocab_is_adopted_before_parent_sampler_initialization(monkeypatch):
+    import vllm.v1.worker.gpu.model_runner as upstream
+
+    from vllm_omni.worker_v2 import omni_model_runner as module
+
+    runner = object.__new__(module.OmniGPUModelRunner)
+    runner.vocab_size = 100
+    runner.req_states = SimpleNamespace(vocab_size=100)
+    runner.model_config = SimpleNamespace(model_arch_config=SimpleNamespace(vocab_size=100))
+    model = SimpleNamespace(logits_vocab_size=32)
+    original_initializer = upstream.init_model_state
+    sentinel = object()
+    monkeypatch.setattr(module, "init_omni_model_state", lambda *args: sentinel)
+
+    def load(parent):
+        assert upstream.init_model_state(None, model, None, None) is sentinel
+        assert parent.vocab_size == parent.req_states.vocab_size == 32
+        assert parent.model_config.model_arch_config.vocab_size == 32
+        raise RuntimeError("stop before sampler allocation")
+
+    monkeypatch.setattr(upstream.GPUModelRunner, "load_model", load)
+    with pytest.raises(RuntimeError, match="stop before sampler allocation"):
+        runner.load_model()
+    assert upstream.init_model_state is original_initializer

tests/worker_v2/test_shared_stage_cuda.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Small CUDA correctness checks for shared stage output contracts."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.worker.gpu_generation_model_runner import _HostCopyBatch
+from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_host_copy_batch_preserves_noncontiguous_outputs_on_current_stream():
+    stream = torch.cuda.Stream()
+    with torch.cuda.stream(stream):
+        source = torch.arange(24, device="cuda", dtype=torch.float32).reshape(4, 6)
+        copies = _HostCopyBatch(pin_memory=True)
+        first = copies.copy(source.t())
+        second = copies.copy(source + 10)
+        copies.wait()
+        assert first.is_contiguous() and first.is_pinned()
+        assert torch.equal(first, torch.arange(24, dtype=torch.float32).reshape(4, 6).t())
+        assert torch.equal(second, torch.arange(24, dtype=torch.float32).reshape(4, 6) + 10)
+        copies.wait()  # no outstanding copies
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_aux_pytree_tracks_replayed_graph_values():
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.model = SimpleNamespace(_returns_tuple=True, supports_mrv2_full_graph_aux_outputs=True)
+    runner._configure_cudagraph_output_contract()
+    source = torch.ones((4, 3), device="cuda")
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = runner._flatten_capture_aux_output((source * 2, {"layers": [source + 1, source + 2]}))
+    source.fill_(7)
+    graph.replay()
+    hidden, aux = runner._split_fullgraph_output(captured)
+    assert torch.equal(hidden, torch.full_like(source, 14))
+    assert torch.equal(aux["layers"][0], torch.full_like(source, 8))
+    assert torch.equal(aux["layers"][1], torch.full_like(source, 9))

vllm_omni/config/environment_variable_inventory.py

@@ -251,6 +251,10 @@ _BENCHMARK_TRANSITIONAL = (
 )
 
 _INTERNAL = (
+    "VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS",
+    "VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS_SLACK_S",
+    "VLLM_OMNI_CONNECTOR_RECV_POLL_MS",
+    "VLLM_OMNI_SHM_WAKEUP",
     "VLLM_OMNI_DLO_DP_WAVE_TIMEOUT",
     "VLLM_OMNI_REPLICA_ID",
 )

vllm_omni/config/omni_config.py

@@ -57,6 +57,8 @@ from vllm_omni.config.stage_config import (
     normalize_pipeline_cli_overrides,
     reconcile_diffusion_attention_overrides,
     resolve_stage_async_chunk,
+    resolve_stage_model_runner,
+    validate_native_mrv2_session,
     validate_stage_async_chunk_edges,
 )
 from vllm_omni.diffusion.diffusion_kv.config import DiffusionKVCacheMode
@@ -1988,7 +1990,9 @@ def _build_model_config(
     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")
+    stage_runner = resolve_stage_model_runner(deploy, stage_deploy)
+    validate_native_mrv2_session(deploy, topology, stage_runner)
+    kwargs.setdefault("use_v2_model_runner", stage_runner == "v2")
     kwargs.setdefault(
         "supports_native_mrv2_data_plane",
         topology.supports_native_mrv2_data_plane,

vllm_omni/config/stage_config.py

@@ -403,6 +403,9 @@ class StageDeployConfig:
 
     # False opts this stage out of pipeline-wide async chunking.
     async_chunk: bool | None = None
+    # Overrides the deploy-level ``model_runner`` for this stage, so a pipeline
+    # can run e.g. its LLM stage on v1 and its codec stages on MRv2.
+    model_runner: Literal["v1", "v2"] | None = None
 
     # Inter-stage connector wiring and request defaults.
     output_connectors: dict[str, str] | None = None
@@ -609,6 +612,7 @@ class DeployConfig:
 _STAGE_RESERVED_KEYS = frozenset(
     {
         "async_chunk",
+        "model_runner",
         "stage_id",
         "devices",
         "num_replicas",
@@ -657,6 +661,9 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
             else:
                 flat_args[k] = v
 
+    stage_runner = stage_data.get("model_runner")
+    if stage_runner is not None and stage_runner not in ("v1", "v2"):
+        raise ValueError(f"stage {stage_data['stage_id']}: model_runner must be 'v1' or 'v2', got {stage_runner!r}")
     kwargs: dict[str, Any] = {
         "stage_id": stage_data["stage_id"],
         "devices": devices,
@@ -668,6 +675,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig:
             kwargs[name] = flat_args.pop(name)
 
     kwargs["async_chunk"] = stage_data.get("async_chunk")
+    kwargs["model_runner"] = stage_runner
     kwargs["output_connectors"] = stage_data.get("output_connectors")
     kwargs["input_connectors"] = stage_data.get("input_connectors")
     kwargs["default_sampling_params"] = stage_data.get("default_sampling_params")
@@ -889,17 +897,12 @@ def _apply_platform_overrides(
         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
-    if platform_section is None:
-        return deploy
-
-    platform_stages = platform_section.get("stages", [])
+        if model_runner == "v1":
+            # A platform's V1 fallback also covers stages that opt into V2.
+            for stage in deploy.stages:
+                if stage.model_runner == "v2":
+                    stage.model_runner = None
+    platform_stages = platform_section.get("stages", []) if platform_section else []
     base_by_id = {s.stage_id: s for s in deploy.stages}
 
     for ps in platform_stages:
@@ -935,6 +938,18 @@ def _apply_platform_overrides(
             else:
                 base.engine_extras[key] = val
 
+    # Validate the final values, including stage entries from the platform
+    # overlay. A global V2 default may still apply to omitted pipeline stages.
+    for stage in deploy.stages:
+        if stage.model_runner is not None and stage.model_runner not in ("v1", "v2"):
+            raise ValueError(f"stage {stage.stage_id}: model_runner must be 'v1' or 'v2', got {stage.model_runner!r}")
+    uses_v2 = deploy.model_runner == "v2" or any(stage.model_runner == "v2" for stage in deploy.stages)
+    if uses_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."
+        )
+
     return deploy
 
 
@@ -952,6 +967,28 @@ def _resolve_execution_mode(
     return _EXECUTION_TYPE_TO_STAGE_WORKER.get(execution_type, (StageType.LLM, None))
 
 
+def resolve_stage_model_runner(deploy: DeployConfig, stage: StageDeployConfig | None) -> str:
+    """The stage's own ``model_runner`` if set, else the deploy-level one."""
+    runner = getattr(stage, "model_runner", None) if stage is not None else None
+    return runner or deploy.model_runner
+
+
+def validate_native_mrv2_session(deploy: DeployConfig, ps: StagePipelineConfig, stage_runner: str) -> None:
+    """Reject session modes a downstream MRv2 native-data-plane stage lacks.
+
+    Streaming-session prompt replacement exists only in the V1 chunk adapter,
+    so a stage that receives from an upstream stage on MRv2 supports
+    turn-based sessions only.
+    """
+    if stage_runner != "v2" or not ps.supports_native_mrv2_data_plane or not ps.input_sources:
+        return
+    if deploy.session_mode != "turn":
+        raise ValueError(
+            f"stage {ps.stage_id}: model_runner v2 supports session_mode 'turn' only, got "
+            f"{deploy.session_mode!r}; run this stage with model_runner: v1 for streaming sessions."
+        )
+
+
 def resolve_stage_async_chunk(deploy: DeployConfig, stage: StageDeployConfig | None) -> bool:
     if not isinstance(deploy.async_chunk, bool):
         raise ValueError("async_chunk must be a boolean")
@@ -1076,9 +1113,11 @@ def _build_engine_args(
                     "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"
+    stage_runner = resolve_stage_model_runner(deploy, ds)
+    validate_native_mrv2_session(deploy, ps, stage_runner)
+    engine_args["use_v2_model_runner"] = stage_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:
+    if stage_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; "

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -3,10 +3,12 @@
 
 from __future__ import annotations
 
+import os
 import time
 from collections import defaultdict
 from typing import TYPE_CHECKING, Any
 
+import torch
 from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.logger import init_logger
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
@@ -56,6 +58,23 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._init_omni_io_scheduling_state()
         self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False))
         self._pending_finish_reqs: list[Request] = []
+        self._first_chunk_express = os.environ.get("VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS") == "1"
+        self._last_step_express = False
+        # Streams that already had a chunk scheduled; the rest await their first.
+        self._chunk_started: set[str] = set()
+        # Optional guard: take an express step only while every ready later
+        # chunk's stream still holds at least this much unplayed audio, so a
+        # loaded stage does not trade first packets for playback underruns.
+        # Playback is assumed to start with a stream's first emitted audio.
+        self._express_min_slack_s = float(os.environ.get("VLLM_OMNI_CODEC_FIRST_CHUNK_EXPRESS_SLACK_S", "0") or 0)
+        # request id -> [monotonic time of first emitted audio, emitted seconds]
+        self._stream_audio: dict[str, list[float]] = {}
+        self._express_skipped_for_slack = 0
+        if self._first_chunk_express:
+            logger.info(
+                "Generation scheduler decodes first chunks in express steps (min continuation slack %.2f s)",
+                self._express_min_slack_s,
+            )
 
     def _build_generation_scheduler_output(
         self,
@@ -250,7 +269,31 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         # 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 and len(num_scheduled_tokens) < execution_batch_size:
+        # An express step isolates first chunks from ready continuations.
+        # Two express steps never run back to back; the optional slack guard
+        # requires playback headroom before delaying a continuation.
+        first_chunk_express = getattr(self, "_first_chunk_express", False)
+        express = (
+            first_chunk_express
+            and native_chunks
+            and not self._last_step_express
+            and any(
+                r.request_id not in self._chunk_started
+                and r.num_in_flight_tokens == 0
+                and len(r.prompt_token_ids) > r.num_computed_tokens
+                for r in self.waiting
+            )
+        )
+        if express and getattr(self, "_express_min_slack_s", 0) > 0 and not self._continuations_have_slack():
+            express = False
+        if first_chunk_express:
+            self._last_step_express = express
+        while (
+            not express
+            and 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 (
@@ -327,6 +370,10 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             and self._pause_state == PauseState.UNPAUSED
         ):
             request = self.waiting.peek_request()
+            if express and request.request_id in self._chunk_started:
+                self.waiting.pop_request()
+                skipped_waiting_requests.add_request(request)
+                continue
             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.
@@ -398,6 +445,9 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             token_budget -= num_new_tokens
             scheduled_new_reqs.append(request)
 
+        if first_chunk_express:
+            self._chunk_started.update(num_scheduled_tokens)
+
         # Return skipped waiting requests
         if skipped_waiting_requests:
             self.waiting.prepend_requests(skipped_waiting_requests)
@@ -493,9 +543,52 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         return self._wrap_omni_scheduler_output(scheduler_output)
 
+    def _continuations_have_slack(self) -> bool:
+        """Whether every started stream with a chunk ready to decode can wait one express step."""
+        now = time.monotonic()
+        for request in (*self.running, *self.waiting):
+            if request.request_id not in self._chunk_started or request.num_in_flight_tokens > 0:
+                continue
+            if len(request.prompt_token_ids) <= request.num_computed_tokens:
+                continue
+            audio = self._stream_audio.get(request.request_id)
+            if audio is None:
+                # No known playback credit: do not delay this continuation.
+                return False
+            if audio[0] + audio[1] - now < self._express_min_slack_s:
+                self._express_skipped_for_slack += 1
+                if self._express_skipped_for_slack in (1, 100, 10000):
+                    logger.info(
+                        "First-chunk express deferred: a stream holds under %.2f s of audio (%d times)",
+                        self._express_min_slack_s,
+                        self._express_skipped_for_slack,
+                    )
+                return False
+        return True
+
+    def _record_stream_audio(self, request_id: str, mm_output: Any) -> None:
+        """Accumulate the seconds of audio each stream has emitted (``model_outputs`` at ``sr``)."""
+        if not isinstance(mm_output, dict):
+            return
+        audio = mm_output.get("model_outputs")
+        rate = mm_output.get("sr")
+        # Only mono sample vectors have an unambiguous duration here.
+        # Unknown/multichannel layouts earn no slack credit.
+        if not isinstance(audio, torch.Tensor) or audio.ndim != 1 or audio.numel() == 0:
+            return
+        if isinstance(rate, torch.Tensor):
+            rate = int(rate.reshape(-1)[0]) if rate.numel() else 0
+        if not isinstance(rate, int) or rate <= 0:
+            return
+        entry = self._stream_audio.setdefault(request_id, [time.monotonic(), 0.0])
+        entry[1] += audio.numel() / rate
+
     def _free_request(
         self, request: Request, delay_free_blocks: bool = False
     ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
+        if getattr(self, "_first_chunk_express", False):
+            self._chunk_started.discard(request.request_id)
+            getattr(self, "_stream_audio", {}).pop(request.request_id, None)
         if self.input_coordinator is None:
             return super()._free_request(request, delay_free_blocks)
 
@@ -641,6 +734,8 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             ec_transfer_params = None
             pooler_output = pooler_outputs[req_index] if pooler_outputs else None
             mm_output = mm_outputs[req_index] if mm_outputs else None
+            if getattr(self, "_express_min_slack_s", 0) > 0 and self._first_chunk_express:
+                self._record_stream_audio(req_id, mm_output)
             status_before_stop = request.status
             finish_reason = None
             is_segment_finished = False

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -8,7 +8,7 @@ import os
 import queue
 import time
 from collections.abc import Iterable, Iterator
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import torch
 from vllm.compilation.cuda_graph import CUDAGraphStat
@@ -44,6 +44,7 @@ from vllm_omni.core.sched.output import (
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
+from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks
 from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.serialization import serialize_additional_information
 from vllm_omni.outputs import OmniConnectorOutput
@@ -104,7 +105,13 @@ elif DEFAULT_INPUT_WAIT_TIMEOUT_S == 0:
     )
 
 
-class OmniSchedulerMixin:
+if TYPE_CHECKING:
+    from vllm.v1.core.sched.scheduler import Scheduler as _SchedulerMixinBase
+else:
+    _SchedulerMixinBase = object
+
+
+class OmniSchedulerMixin(_SchedulerMixinBase):
     """Shared scheduler helpers for omni-specific request handling."""
 
     # Provided by the concrete vLLM scheduler.
@@ -127,17 +134,12 @@ class OmniSchedulerMixin:
 
     def _init_omni_io_scheduling_state(self) -> None:
         """Initialize scheduler state shared by AR and generation stages."""
-        # Every omni worker forces the v1 model runner (the v2 runner carries no
-        # omni hooks), so the scheduler has to agree or the two disagree about
-        # what the SchedulerOutput carries. vLLM 0.29 defaults this to v2 where
-        # 0.28 did not, so the scheduler took its v2 fast path -- the one that
-        # deliberately omits resumed-request token ids -- while the v1 runner
-        # still read them, and resuming a request raised ``KeyError: <req_id>``
-        # in ``_update_states``.
-        if getattr(self, "use_v2_model_runner", False):
-            logger.warning("OMNI scheduler forces v1 model runner for omni hooks.")
-            self.use_v2_model_runner = False
-
+        # The scheduler must agree with the stage's worker about which runner
+        # consumes its SchedulerOutput: vLLM's own default (v2 since 0.29) would
+        # take the v2 fast path, which omits resumed-request token ids, while a
+        # v1 omni runner still reads them (``KeyError: <req_id>`` in
+        # ``_update_states``). The stage's model config carries the choice
+        # (deploy ``model_runner``, optionally per stage).
         model_config = self.vllm_config.model_config
         self.use_v2_model_runner = bool(getattr(model_config, "use_v2_model_runner", False))
         self._native_data_plane = uses_native_mrv2_data_plane(
@@ -150,7 +152,14 @@ class OmniSchedulerMixin:
             else None
         )
         self.input_coordinator: OmniSchedulingCoordinator | None = None
-        if self._native_data_plane and getattr(model_config, "async_chunk", False):
+        if (
+            self._native_data_plane
+            and getattr(model_config, "async_chunk", False)
+            # A downstream sender-only stage (e.g. MiniCPM-o's Talker, fed by
+            # the orchestrator) has no connector input to wait for; parking
+            # its requests for chunks would never release them.
+            and (getattr(model_config, "stage_id", 0) == 0 or stage_receives_chunks(model_config))
+        ):
             self.input_coordinator = OmniSchedulingCoordinator(
                 scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs,
                 stage_id=getattr(model_config, "stage_id", 0),
@@ -825,6 +834,8 @@ class OmniSchedulerMixin:
     def _resume_downstream_chunk_receiver(self, request: Request) -> None:
         """Resume duplex connector polling without an external update."""
         adapter = self.chunk_transfer_adapter
+        if adapter is None:
+            return
         adapter.segment_finished_requests.discard(request.request_id)
         if (
             not adapter.receives_chunks

vllm_omni/data_entry_keys.py

@@ -41,6 +41,8 @@ class Embeddings(TypedDict, total=False):
     prepared_prefill: torch.Tensor
     prefill: torch.Tensor
     decode: torch.Tensor
+    # [1, H] embedding of the token sampled by a row's final prefill step (MRv2 producers).
+    sampled: torch.Tensor
     decode_token_start: int
     decode_token_end: int
     cached_decode: torch.Tensor
@@ -88,6 +90,7 @@ class OmniPayloadMeta(TypedDict, total=False):
     right_holdback_size: int
     override_keys: list[tuple[str, str]]
     num_processed_tokens: int
+    resumable: bool
     next_stage_prompt_len: int
     next_stage_generation_tokens: int
     replace_streaming_prompt: bool
@@ -202,6 +205,8 @@ class MetaStruct(_StructBase):
     right_holdback_size: int | None = None
     override_keys: list[tuple[str, str]] | None = None
     num_processed_tokens: int | None = None
+    # The model runner sets this when a streaming request is resumed.
+    resumable: bool | None = None
     next_stage_prompt_len: int | None = None
     next_stage_generation_tokens: int | None = None
     replace_streaming_prompt: bool | None = None

vllm_omni/distributed/omni_connectors/connectors/shm_connector.py

@@ -2,9 +2,14 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 import fcntl
+import glob
+import hashlib
 import os
+import select
+import stat
 import threading
 import time
+import uuid
 from collections import OrderedDict
 from multiprocessing import shared_memory as shm_pkg
 from typing import Any
@@ -17,6 +22,15 @@ from .base import OmniConnectorBase
 logger = get_connector_logger(__name__)
 
 
+def _wakeup_enabled() -> bool:
+    return os.environ.get("VLLM_OMNI_SHM_WAKEUP", "1") == "1"
+
+
+def _wakeup_directory() -> str:
+    # Stage engine processes of one deployment share their launching parent.
+    return f"/dev/shm/omni_shm_wake_{os.getuid()}_{os.getppid()}"
+
+
 class SharedMemoryConnector(OmniConnectorBase):
     """Key-addressed local shared-memory connector.
 
@@ -25,6 +39,12 @@ class SharedMemoryConnector(OmniConnectorBase):
     remote-transport metadata such as ``source_host`` / ``source_port``
     (that is the RDMA connector's job).  When such metadata is passed in,
     the connector silently falls back to key-based lookup.
+
+    Arrival wakeups: each receiving connector owns a unique named FIFO. A
+    ``put`` broadcasts one byte to every FIFO for the destination stage, so the receive loop blocks in
+    ``wait_for_change`` until data arrives instead of re-polling on a fixed
+    interval. Wakeups are hints: a stage that cannot reach the FIFO (another
+    deployment layout, or ``VLLM_OMNI_SHM_WAKEUP=0``) keeps the timed poll.
     """
 
     def get_with_deadline(self, from_stage, to_stage, get_key, metadata=None, *, deadline):
@@ -35,6 +55,12 @@ class SharedMemoryConnector(OmniConnectorBase):
     def __init__(self, config: dict[str, Any]):
         self.config = config
         self.stage_id = config.get("stage_id", -1)
+        scope = config.get("extra", {}).get("wakeup_scope")
+        self._wake_directory = (
+            f"/dev/shm/omni_shm_wake_{os.getuid()}_{hashlib.sha256(str(scope).encode()).hexdigest()[:24]}"
+            if scope is not None
+            else _wakeup_directory()
+        )
         self._pending_keys: OrderedDict[str, None] = OrderedDict()
         self._pending_keys_lock = threading.Lock()
         self._metrics = {
@@ -42,6 +68,129 @@ class SharedMemoryConnector(OmniConnectorBase):
             "gets": 0,
             "bytes_transferred": 0,
         }
+        # Receiver side: FIFO read end (plus a write end of our own, so the
+        # FIFO never reports EOF when no sender has it open) and a counter of
+        # drained wakeups. Each receiver owns its path, including across restarts.
+        self._wake_lock = threading.Lock()
+        self._wake_read_fd: int | None = None
+        self._wake_hold_fd: int | None = None
+        self._wake_path: str | None = None
+        self._wake_generation = 0
+        self._wake_closed = False
+
+    def _open_wakeup_receiver(self) -> bool:
+        if self._wake_closed:
+            return False
+        if self._wake_read_fd is not None:
+            return True
+        if not _wakeup_enabled():
+            return False
+        try:
+            if int(self.stage_id) < 0:
+                return False
+            path = f"{self._wake_directory}/{int(self.stage_id)}_{uuid.uuid4().hex}"
+        except (TypeError, ValueError):
+            return False
+        read_fd = hold_fd = None
+        created = False
+        try:
+            os.makedirs(os.path.dirname(path), mode=0o700, exist_ok=True)
+            os.mkfifo(path, 0o600)
+            created = True
+            read_fd = os.open(path, os.O_RDONLY | os.O_NONBLOCK)
+            hold_fd = os.open(path, os.O_WRONLY | os.O_NONBLOCK)
+        except OSError as e:
+            for fd in (read_fd, hold_fd):
+                if fd is not None:
+                    os.close(fd)
+            if created:
+                try:
+                    os.unlink(path)
+                except FileNotFoundError:
+                    pass
+            logger.debug("SHM wakeup FIFO unavailable at %s: %s", path, e)
+            return False
+        self._wake_read_fd, self._wake_hold_fd, self._wake_path = read_fd, hold_fd, path
+        return True
+
+    def get_wakeup_generation(self) -> int | None:
+        """Snapshot before polling; a later arrival changes it (see ``wait_for_change``).
+
+        ``None`` when this stage has no wakeup FIFO: the caller keeps its timed poll.
+        """
+        with self._wake_lock:
+            return self._wake_generation if self._open_wakeup_receiver() else None
+
+    def wait_for_change(self, generation: int, *, timeout: float) -> bool:
+        """Block until a ``put`` to this stage since ``generation``, or ``timeout``."""
+        with self._wake_lock:
+            read_fd = self._wake_read_fd
+            if self._wake_closed or read_fd is None:
+                return False
+        if self._wake_generation == generation:
+            try:
+                readable, _, _ = select.select([read_fd], [], [], timeout)
+            except (OSError, ValueError):
+                return False
+            if not readable:
+                return False
+            with self._wake_lock:
+                if self._wake_closed or self._wake_read_fd != read_fd:
+                    return False
+                try:
+                    while os.read(read_fd, 4096):
+                        pass
+                except BlockingIOError:
+                    pass
+                except OSError:
+                    return False
+                self._wake_generation += 1
+        return self._wake_generation != generation
+
+    def _wake_receiver(self, to_stage: Any) -> None:
+        if not _wakeup_enabled():
+            return
+        try:
+            pattern = f"{self._wake_directory}/{int(to_stage)}_*"
+        except (TypeError, ValueError):
+            return  # non-numeric stage names use the ordinary polling path
+        with self._wake_lock:
+            if self._wake_closed:
+                return
+            # Discover every replica, including newly restarted receivers. Do
+            # not cache writers across receiver lifetimes or unlink peers' paths.
+            for path in glob.iglob(pattern):
+                fd = None
+                try:
+                    fd = os.open(path, os.O_WRONLY | os.O_NONBLOCK | os.O_NOFOLLOW)
+                    if stat.S_ISFIFO(os.fstat(fd).st_mode):
+                        os.write(fd, b"\0")
+                except OSError:
+                    # Missing reader/file or full pipe: timed polling remains
+                    # the correctness fallback. A full FIFO already has hints.
+                    pass
+                finally:
+                    if fd is not None:
+                        os.close(fd)
+
+    def _close_wakeups(self) -> None:
+        with self._wake_lock:
+            self._wake_closed = True
+            for fd in (self._wake_read_fd, self._wake_hold_fd):
+                if fd is not None:
+                    os.close(fd)
+            self._wake_read_fd = self._wake_hold_fd = None
+            if self._wake_path is not None:
+                try:
+                    os.unlink(self._wake_path)
+                except FileNotFoundError:
+                    pass
+                self._wake_path = None
+                try:
+                    os.rmdir(self._wake_directory)
+                except OSError:
+                    # Other receivers may still own FIFOs in this deployment.
+                    pass
 
     def put(
         self,
@@ -65,6 +214,7 @@ class SharedMemoryConnector(OmniConnectorBase):
             metadata = {"shm": meta, "size": size}
             with self._pending_keys_lock:
                 self._pending_keys[put_key] = None
+            self._wake_receiver(to_stage)
 
             self._metrics["puts"] += 1
             self._metrics["bytes_transferred"] += size
@@ -184,6 +334,7 @@ class SharedMemoryConnector(OmniConnectorBase):
             keys = list(self._pending_keys)
         for key in keys:
             self.cleanup(key)
+        self._close_wakeups()
 
     def reap_consumed(self) -> None:
         """Bounded round-robin sweep; receivers unlink SHM in another process."""

vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py

@@ -6,6 +6,8 @@ from __future__ import annotations
 
 import importlib
 import inspect
+import math
+import os
 import time
 from collections import deque
 from collections.abc import Callable
@@ -22,6 +24,23 @@ from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_runtime i
 )
 from vllm_omni.outputs import OmniConnectorOutput
 
+
+# No-progress recheck for connectors without a change notification (SHM polls).
+def _recv_poll_seconds() -> float:
+    """Resolve a finite positive poll interval, falling back for invalid input."""
+    raw = os.environ.get("VLLM_OMNI_CONNECTOR_RECV_POLL_MS", "5")
+    try:
+        value = float(raw) / 1000
+        if math.isfinite(value) and value > 0:
+            return value
+    except ValueError:
+        pass
+    logger.warning("Invalid VLLM_OMNI_CONNECTOR_RECV_POLL_MS=%r; using 5 ms", raw)
+    return 0.005
+
+
+_RECV_POLL_S = _recv_poll_seconds()
+
 if TYPE_CHECKING:
     from vllm_omni.distributed.omni_connectors.connectors.base import (
         OmniConnectorBase,
@@ -189,10 +208,12 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         connector: OmniConnectorBase,
         from_stage: str,
         to_stage: str,
-        connector_get_key: str,
+        connector_get_key: str | None,
         metadata: dict[str, Any] | None = None,
     ) -> Any:
         """Receive one ordinary non-KV stage payload on the local leader rank only."""
+        if connector_get_key is None:
+            return None
         if not self._stage_payload_broadcast_groups():
             if metadata is None:
                 return connector.get(from_stage, to_stage, connector_get_key)
@@ -208,7 +229,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         connector: OmniConnectorBase,
         from_stage: str,
         to_stage: str,
-        connector_get_key: str,
+        connector_get_key: str | None,
         metadata: dict[str, Any] | None = None,
     ) -> Any:
         """Receive one full-payload transfer on the local leader rank only."""
@@ -225,7 +246,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         connector: OmniConnectorBase,
         from_stage: str,
         to_stage: str,
-        connector_get_key: str,
+        connector_get_key: str | None,
         metadata: dict[str, Any] | None = None,
     ) -> Any:
         """Receive one ordinary async chunk on the local leader rank only."""
@@ -803,7 +824,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                 self._pending_save_counts[req_id] += 1
             sent_ids.append(req_id)
         if sent_ids:
-            self._work_available.set()
+            self._save_work_available.set()
         return sent_ids
 
     # ------------------------------------------------------------------ #
@@ -930,6 +951,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         """Background thread: poll connector for incoming data."""
         _recv_poll_count = 0
         while not self._stop_event.is_set():
+            snapshot = getattr(self._omni_connector, "get_wakeup_generation", None)
+            generation = snapshot() if callable(snapshot) else None
             with self._lock:
                 pending_ids = list(self._pending_load_reqs.keys())
 
@@ -951,8 +974,16 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
             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)
-                self._work_available.clear()
+                # Wait for a *new* arrival: existing future chunks may be
+                # blocked by scheduler ownership of the previous payload.
+                # A key-presence predicate would spin and compete for the
+                # GIL. Honor the configured readiness recheck interval.
+                wait_for_change = getattr(self._omni_connector, "wait_for_change", None)
+                if generation is not None and callable(wait_for_change):
+                    wait_for_change(generation, timeout=_RECV_POLL_S)
+                else:
+                    self._work_available.wait(timeout=_RECV_POLL_S)
+                    self._work_available.clear()
 
     _MAX_SEND_RETRIES = 3
 
@@ -986,8 +1017,8 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
                     self._requeue_or_drop_failed_send(task, error=send_error)
                 continue
 
-            self._work_available.wait(timeout=0.01)
-            self._work_available.clear()
+            self._save_work_available.wait(timeout=0.01)
+            self._save_work_available.clear()
 
     def _requeue_or_drop_failed_send(
         self,
@@ -1547,7 +1578,7 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
         with self._lock:
             self._pending_save_reqs.setdefault(request_id, deque()).append(task)
             self._pending_save_counts[request_id] += 1
-        self._work_available.set()
+        self._save_work_available.set()
         return True, completion
 
     def _poll_pending_requests_once(self, pending_ids: list[str]) -> bool:
@@ -1564,8 +1595,25 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin):
             self.publish_omni_connector_output_to_sink()
         return made_progress
 
+    @staticmethod
+    def _finish_marker_payload() -> OmniPayload:
+        return {"meta": {"finished": torch.tensor(True, dtype=torch.bool)}}
+
+    @staticmethod
+    def _request_is_finished(request: Any) -> bool:
+        is_finished = getattr(request, "is_finished", None)
+        return bool(is_finished()) if callable(is_finished) else False
+
     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]
+        # A request's last chunk must carry its finish marker even when the
+        # processor has nothing left to send (e.g. its frames ended exactly on
+        # a chunk boundary), as the V1 chunk adapter guarantees; otherwise the
+        # receiver waits for input until its deadline.
+        entries = [
+            (request, payload if payload is not None else self._finish_marker_payload())
+            for request, payload in entries
+            if payload is not None or self._request_is_finished(request)
+        ]
         if not entries:
             return 0
         emitted = 0

vllm_omni/distributed/omni_connectors/model_runner/omni_connector_runtime.py

@@ -103,6 +103,10 @@ class _SendCompletion:
 class _OmniConnectorRuntimeMixin:
     """Own connector lifecycle, shared state, and KV transfer delegation."""
 
+    if TYPE_CHECKING:
+
+        def _stage_payload_broadcast_groups(self) -> tuple[Any, ...]: ...
+
     _omni_connector: Any
     _kv_transfer_manager: Any
     _async_chunk: bool
@@ -145,6 +149,7 @@ class _OmniConnectorRuntimeMixin:
     _lock: Any
     _stop_event: Any
     _work_available: Any
+    _save_work_available: Any
     _recv_thread: Any
     _save_thread: Any
     _omni_connector_initialized: bool
@@ -313,6 +318,10 @@ class _OmniConnectorRuntimeMixin:
         self._omni_connector_output_drain_lock = threading.Lock()
         self._stop_event = threading.Event()
         self._work_available = threading.Event()
+        # The save loop has its own event: both loops clear theirs after a
+        # wait, so a shared event let the receiver swallow a send wakeup and
+        # left the payload waiting for the saver's 10 ms timeout.
+        self._save_work_available = threading.Event()
 
         # Start background threads only when there's a connector
         self._recv_thread: threading.Thread | None = None
@@ -343,6 +352,8 @@ class _OmniConnectorRuntimeMixin:
         if getattr(self, "_synchronous_payload_transport", False):
             return
         self._stop_event.set()
+        self._work_available.set()
+        self._save_work_available.set()
         if self._recv_thread is not None:
             self._recv_thread.join(timeout=5)
         if self._save_thread is not None:

vllm_omni/distributed/omni_connectors/utils/initialization.py

@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
 
 """Utilities for OmniConnector configuration and validation."""
 
 import json
 import os
+import uuid
 import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Literal
@@ -195,7 +196,9 @@ def create_connectors_from_config(
                 replica_id = max(int(os.environ.get("VLLM_OMNI_REPLICA_ID", "0")), 0)
                 resolved_spec = resolve_connector_spec(
                     ConnectorSpec(name=connector_spec.name, extra=extra),
-                    stage_id=int(caller_stage_id) if str(caller_stage_id).isdigit() else from_stage_id,
+                    stage_id=int(caller_stage_id)
+                    if caller_stage_id is not None and str(caller_stage_id).isdigit()
+                    else from_stage_id,
                     role=role,
                     purpose=resolved_purpose,
                     local_rank=local_rank,
@@ -308,7 +311,7 @@ def load_omni_transfer_config(
         config_dict = normalized
 
     # Parse connectors
-    connectors = {}
+    connectors: dict[tuple[str, str], ConnectorSpec] = {}
     runtime_config = config_dict.get("runtime", {})
 
     # Parse global connectors (from runtime.connectors)
@@ -429,6 +432,12 @@ def load_omni_transfer_config(
             f"{missing_str}. Define connectors or allow auto SHM creation for these edges."
         )
 
+    # Resolve once in the launcher; serialized connector extras carry the same
+    # namespace through EngineCore and multiprocess/TP workers.
+    wakeup_scope = uuid.uuid4().hex
+    for spec in connectors.values():
+        if spec.name == "SharedMemoryConnector":
+            spec.extra.setdefault("wakeup_scope", wakeup_scope)
     config = OmniTransferConfig(connectors=connectors)
 
     logger.info(f"Loaded OmniTransferConfig with {len(connectors)} connector configurations")

vllm_omni/engine/orchestrator.py

@@ -34,6 +34,7 @@ from vllm.v1.engine import EngineCoreOutputs, FinishReason
 from vllm.v1.engine.exceptions import EngineDeadError
 from vllm.v1.metrics.stats import IterationStats
 
+from vllm_omni.core.sched.omni_scheduling_coordinator import uses_native_mrv2_data_plane
 from vllm_omni.data_entry_keys import FIRST_AUDIO_KEY, FIRST_AUDIO_REQUIRED_KEY
 from vllm_omni.diffusion.data import is_diffusion_request_started_output
 from vllm_omni.distributed.omni_connectors.utils.config import stage_receives_chunks
@@ -2907,6 +2908,18 @@ class Orchestrator(OrchestratorBase):
             return False
         return True
 
+    def _native_mrv2_receiver_stage(self, final_stage_id: int) -> int | None:
+        """First downstream stage up to ``final_stage_id`` that receives on MRv2's native data plane."""
+        for stage_id in range(1, min(final_stage_id, len(self.stage_pools) - 1) + 1):
+            vllm_config = getattr(self.stage_pools[stage_id], "stage_vllm_config", None)
+            model_config = getattr(vllm_config, "model_config", None)
+            if uses_native_mrv2_data_plane(
+                model_config,
+                use_v2_model_runner=bool(getattr(model_config, "use_v2_model_runner", False)),
+            ):
+                return stage_id
+        return None
+
     async def _handle_add_request(self, msg: StageSubmissionMessage) -> None:
         """Handle an add_request message from the main thread."""
         stage_id = 0
@@ -2926,6 +2939,20 @@ class Orchestrator(OrchestratorBase):
             await self._fail_request_dead_stage(request_id, stage_id)
             return
 
+        if getattr(prompt, "resumable", False):
+            mrv2_stage_id = self._native_mrv2_receiver_stage(final_stage_id)
+            if mrv2_stage_id is not None:
+                # Streaming-input sessions replace downstream prompts through
+                # the V1 chunk adapter, which MRv2's native data plane lacks.
+                await self._fail_request_client_error(
+                    request_id,
+                    stage_id,
+                    f"streaming (resumable) input is not supported: stage {mrv2_stage_id} runs on "
+                    "model_runner v2, which supports turn-based requests only; use model_runner: v1 "
+                    "for that stage",
+                )
+                return
+
         logger.debug(
             "[Orchestrator] _handle_add_request: stage=%s req=%s "
             "prompt_type=%s original_prompt_type=%s final_stage=%s "

vllm_omni/utils/device_copy.py

@@ -20,6 +20,14 @@ from collections.abc import Sequence
 import torch
 
 
+def to_device_nonblocking(tensor: torch.Tensor, device: torch.device | str) -> torch.Tensor:
+    """``tensor.to(device)`` without a host sync when ``tensor`` is on the CPU and ``device`` is CUDA."""
+    device = torch.device(device)
+    if tensor.device.type != "cpu" or device.type != "cuda":
+        return tensor.to(device)
+    return tensor.pin_memory().to(device, non_blocking=True)
+
+
 def index_to_device(values: Sequence[int], device: torch.device | str, dtype: torch.dtype = torch.long) -> torch.Tensor:
     """A host index list as a device tensor, without a host sync."""
     if torch.device(device).type != "cuda":

vllm_omni/worker/gpu_ar_model_runner.py

@@ -309,6 +309,10 @@ def _ensure_tensor_values(payload: dict[str, object]) -> dict[str, torch.Tensor]
     """
     result: dict[str, torch.Tensor] = {}
     for key, val in payload.items():
+        # Sparse per-request conditioning uses None for requests with no
+        # update. Absence is expected, not an unsupported wire value.
+        if val is None:
+            continue
         if isinstance(val, torch.Tensor):
             result[key] = val
         elif isinstance(val, (int, float, bool)):
@@ -1817,6 +1821,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex
         if not needs_pooler_payload and prefix_cache_step_id is not None:
             # No consumer for this step's merge: consume the step context by
             # id (exactly-once contract). The cache write still lands.
+            assert self.omni_prefix_cache is not None
             self.omni_prefix_cache.discard_step(prefix_cache_step_id)
             prefix_cache_step_id = None
         if self.omni_prefix_cache is None and needs_scheduled_hidden_payload and not audio_sparse_output:

vllm_omni/worker/gpu_generation_model_runner.py

@@ -13,6 +13,7 @@ import gc
 import logging
 from collections.abc import Mapping
 from dataclasses import replace
+from typing import cast
 
 import numpy as np
 import torch
@@ -23,6 +24,7 @@ from vllm.distributed.parallel_state import get_pp_group
 from vllm.forward_context import set_forward_context
 from vllm.sequence import IntermediateTensors
 from vllm.utils.math_utils import cdiv
+from vllm.utils.platform_utils import is_pin_memory_available
 from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
 from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, AsyncModelRunnerOutput, make_empty_encoder_model_runner_output
 from vllm.v1.spec_decode.dflash import DFlashProposer
@@ -51,6 +53,36 @@ from vllm_omni.worker.omni_connector_model_runner_mixin import (
 logger = logging.getLogger(__name__)
 
 
+class _HostCopyBatch:
+    """Device-to-host copies of one step's outputs behind a single host sync.
+
+    ``tensor.to("cpu")`` blocks the host once per tensor, so a step returning
+    one waveform per request paid one sync per request. Here each CUDA tensor
+    is copied into pinned host memory without blocking and ``wait`` blocks
+    once, after which every returned tensor holds its data, exactly like the
+    per-tensor ``.detach().to("cpu").contiguous()`` it replaces. Without
+    pinned memory, or off CUDA, the per-tensor blocking copy is kept.
+    """
+
+    def __init__(self, pin_memory: bool) -> None:
+        self._pin_memory = bool(pin_memory)
+        self._pending = False
+
+    def copy(self, tensor: torch.Tensor) -> torch.Tensor:
+        tensor = tensor.detach()
+        if tensor.device.type != "cuda" or not self._pin_memory:
+            return tensor.to("cpu").contiguous()
+        host = torch.empty(tensor.shape, dtype=tensor.dtype, device="cpu", pin_memory=True)
+        host.copy_(tensor, non_blocking=True)
+        self._pending = True
+        return host
+
+    def wait(self) -> None:
+        if self._pending:
+            torch.cuda.current_stream().synchronize()
+            self._pending = False
+
+
 class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
     """Generation model runner for vLLM-Omni (non-autoregressive).
 
@@ -59,6 +91,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
     - Executes generation process and returns tensors via `pooler_output`.
     """
 
+    execute_model_state: ExecuteModelState | None
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self._async_chunk = getattr(self.model_config, "async_chunk", False)
@@ -419,6 +453,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         # Build per-request multimodal_outputs list (dedicated channel).
         # pooler_output is no longer used for multimodal data.
         per_req_payloads: list[dict[str, object]] = []
+        # One host sync for the whole step instead of one per request tensor.
+        to_host = _HostCopyBatch(is_pin_memory_available())
         if isinstance(multimodal_outputs_raw, torch.Tensor):
             # One row per request. The old asserts (`shape[0] == 1` AND
             # `shape[0] == num_reqs`) jointly forced num_reqs == 1, silently
@@ -431,7 +467,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                     "to return multiple tensors per request, use a dict)."
                 )
             for i in range(num_reqs):
-                per_req_payloads.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()})
+                per_req_payloads.append({"model_outputs": to_host.copy(multimodal_outputs_raw[i])})
         elif isinstance(multimodal_outputs_raw, list):
             # One entry per request. The old `len == 1` assert did not check
             # num_reqs, so a batched step built a length-1 payload list that
@@ -444,9 +480,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                     "to return multiple lists per request, use a dict)."
                 )
             for out in multimodal_outputs_raw:
-                per_req_payloads.append(
-                    {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None}
-                )
+                per_req_payloads.append({"model_outputs": to_host.copy(out) if out is not None else None})
         elif isinstance(multimodal_outputs_raw, Mapping):
             num_reqs = self.input_batch.num_reqs
             for i in range(num_reqs):
@@ -458,21 +492,26 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                                 f"Multimodal output list for key '{key}' has length {len(out)} "
                                 f"but expected {num_reqs} (one entry per request)."
                             )
-                        mm_payload[key] = out[i].detach().to("cpu").contiguous()
+                        mm_payload[key] = to_host.copy(out[i])
                     elif isinstance(out, torch.Tensor):
-                        mm_payload[key] = out.detach().to("cpu").contiguous()
+                        mm_payload[key] = to_host.copy(out)
                     else:
                         logger.warning(f"Unsupported multimodal output type for key '{key}': {type(out)}")
                 per_req_payloads.append(_ensure_tensor_values(mm_payload))
         else:
             raise RuntimeError("Unsupported diffusion output type")
+        # Sync mode hands these payloads out as finished host tensors.
+        to_host.wait()
 
+        inter_stage_outputs: list[dict[str, object] | None] | None
+        multimodal_outputs: list[dict[str, object] | None] | None
         if self._async_chunk:
             inter_stage_outputs, multimodal_outputs = partition_payload_list(per_req_payloads)
         else:
             # See gpu_ar_model_runner: non-async-chunk ships the full payload to the next
             # stage; #4527's (None, per_req_payloads) starved the downstream stage. (PR #4792)
-            inter_stage_outputs, multimodal_outputs = per_req_payloads, per_req_payloads
+            inter_stage_outputs = cast(list[dict[str, object] | None], per_req_payloads)
+            multimodal_outputs = inter_stage_outputs
 
         # [Omni] Copy req_id mappings to avoid async scheduling mutation.
         req_ids_output_copy = self.input_batch.req_ids.copy()
@@ -493,7 +532,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             logprobs=None,
             prompt_logprobs_dict={},
             pooler_output=None,
-            multimodal_outputs=multimodal_outputs,
+            # Preserve absent per-row client payloads; the scheduler distinguishes None from {}.
+            multimodal_outputs=multimodal_outputs,  # type: ignore[arg-type]
             inter_stage_outputs=inter_stage_outputs,
             kv_connector_output=kv_connector_output,
             num_nans_in_logits={},

vllm_omni/worker/gpu_model_runner.py

@@ -41,6 +41,7 @@ from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEm
 from vllm_omni.model_executor.models.model_local_kv import collect_model_local_kv_specs
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.device_copy import index_to_device
 
 if TYPE_CHECKING:
     from vllm.v1.core.sched.output import SchedulerOutput
@@ -210,6 +211,32 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner):
         self._init_talker_mtp()
         self._prewarm_attention_capture_workspaces()
         self._report_model_local_kv()
+        self._warn_unexposed_stage_hooks(model)
+
+    # Read on the model this runner holds. A multi-stage wrapper that builds its
+    # stage module as a child must re-export them, or the runner silently takes
+    # the per-row, host-synchronizing default paths.
+    _STAGE_HOOKS = ("gpu_resident_buffer_keys", "preprocess_decode_batch", "use_async_omni_output")
+
+    @classmethod
+    def _warn_unexposed_stage_hooks(cls, model: Any) -> None:
+        if model is None or not hasattr(model, "named_children"):
+            return
+        for child_name, child in model.named_children():
+            missing = [
+                hook
+                for hook in cls._STAGE_HOOKS
+                if getattr(child, hook, None) not in (None, False) and getattr(model, hook, None) in (None, False)
+            ]
+            if missing:
+                logger.warning(
+                    "%s.%s defines %s but %s does not expose them; the runner reads these on %s.",
+                    type(model).__name__,
+                    child_name,
+                    ", ".join(missing),
+                    type(model).__name__,
+                    type(model).__name__,
+                )
 
     def _report_model_local_kv(self) -> None:
         """Log attention KV this model holds outside the paged manager.
@@ -1608,6 +1635,21 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner):
             device=device,
         )
 
+    def _embed_multimodal_input_ids(self, num_scheduled_tokens, mm_embeds, is_mm_embed):
+        embedding_kwargs = {}
+        if mm_embeds and getattr(self.model, "supports_embed_input_ids_query_start_loc", False):
+            # The batch has already been reordered. Read host-owned boundaries
+            # so model prompt rearrangement does not assume a request order or
+            # introduce a GPU-to-host transfer on every decode step.
+            num_reqs = self.input_batch.num_reqs
+            embedding_kwargs["query_start_loc"] = self.query_start_loc.cpu[: num_reqs + 1].tolist()
+        return self.model.embed_input_ids(
+            self.input_ids.gpu[:num_scheduled_tokens],
+            multimodal_embeddings=mm_embeds,
+            is_multimodal=is_mm_embed,
+            **embedding_kwargs,
+        )
+
     def _preprocess(
         self,
         scheduler_output: "SchedulerOutput",
@@ -1644,11 +1686,7 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner):
             # NOTE(woosuk): To unify token ids and soft tokens (vision
             # embeddings), we always use embeddings (rather than token ids)
             # as input to the multimodal model, even when the input is text.
-            inputs_embeds_scheduled = self.model.embed_input_ids(
-                self.input_ids.gpu[:num_scheduled_tokens],
-                multimodal_embeddings=mm_embeds,
-                is_multimodal=is_mm_embed,
-            )
+            inputs_embeds_scheduled = self._embed_multimodal_input_ids(num_scheduled_tokens, mm_embeds, is_mm_embed)
 
             # TODO(woosuk): Avoid the copy. Optimize.
             self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled)
@@ -1821,7 +1859,8 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner):
                         dtype=req_embeds.dtype,
                     )
 
-                offsets_t = torch.tensor(start_offsets_b, device=req_embeds.device, dtype=torch.long)
+                # A pageable H2D here would sync the host every decode step.
+                offsets_t = index_to_device(start_offsets_b, req_embeds.device)
                 inputs_embeds.index_copy_(0, offsets_t, req_embeds)
                 preprocess_input_ids.index_copy_(
                     0,

vllm_omni/worker_v2/model_states/intermediate_buffer.py

@@ -79,6 +79,22 @@ class OmniIntermediateBuffer:
         if ai is not None:
             info.update(_resolve_additional_information(ai))
 
+        # Orchestrator stage bridges (e.g. MiniCPM-o's llm2tts) hand the
+        # request's runner payload over directly, like the V1 runner's
+        # ``_update_intermediate_buffer``: nested sections merge key by key.
+        mib = getattr(new_req_data, "model_intermediate_buffer", None)
+        if isinstance(mib, dict) and mib:
+            for key, value in mib.items():
+                if isinstance(value, dict):
+                    section = info.get(key)
+                    if not isinstance(section, dict):
+                        section = {}
+                        info[key] = section
+                    for qualifier, item in value.items():
+                        self._store_value(section, qualifier, item, set())
+                else:
+                    self._store_value(info, key, value, set())
+
         if new_req_data.mm_features:
             info["mm_features"] = new_req_data.mm_features
 

vllm_omni/worker_v2/model_states/omni_model_state.py

@@ -19,6 +19,7 @@ import types
 from collections.abc import Callable
 from typing import Any, cast
 
+import numpy as np
 import torch
 import torch.nn as nn
 from vllm.config import VllmConfig
@@ -32,6 +33,7 @@ from vllm.v1.worker.gpu.states import RequestState
 
 from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.utils.device_copy import index_to_device
 from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
 from vllm_omni.worker_v2.model_states.intermediate_buffer import (
     OmniIntermediateBuffer,
@@ -87,6 +89,7 @@ class OmniModelState(DefaultModelState):
     """
 
     # Eager-frame defaults for states built without __init__ (tests, adapters).
+    _decode_preprocess_is_identity = False
     _eager_mtp = False
     _eager_fastpath = False
     _eager_rows: tuple[InputBatch, list[tuple[int, int, str, bool]], torch.Tensor] | None = None
@@ -139,6 +142,7 @@ class OmniModelState(DefaultModelState):
         self.has_postprocess: bool = getattr(model, "has_postprocess", False)
         self.have_multimodal_outputs: bool = getattr(model, "have_multimodal_outputs", False)
         self._decode_preprocess = self._resolve_decode_preprocess(model)
+        self._decode_preprocess_is_identity = bool(getattr(model, "mrv2_decode_preprocess_is_identity", False))
         self._mtp_generators: dict[str, torch.Generator] = {}
         # Talker's codec_embedding dim may differ from hf_text_config.hidden_size; probe real dim.
         self._embed_dim = self._get_embed_dim(model, device) if self.has_preprocess else 0
@@ -211,6 +215,13 @@ class OmniModelState(DefaultModelState):
                 input_batch, text_hidden, sampled_token_ids, multimodal_outputs, mtp_batch_descriptor_dispatcher
             )
 
+    def custom_sampler(self, sampler: Any) -> tuple[Any, Any] | None:
+        """Let the model replace or wrap the MRv2 sampler (``mrv2_custom_sampler``)."""
+        hook = getattr(self.model, "mrv2_custom_sampler", None)
+        if callable(hook):
+            return hook(sampler)
+        return super().custom_sampler(sampler)
+
     @staticmethod
     def _resolve_decode_preprocess(model: nn.Module) -> Callable | None:
         """Adapt existing V1 hooks once, outside the per-step hot path.
@@ -603,7 +614,12 @@ class OmniModelState(DefaultModelState):
         preprocess_entries: list[tuple[int, int, int, int, dict[str, Any], bool]] = []
         settled = self._eager_settled if self._eager_fastpath else None
         settled_rows: list[tuple[int, int, int, str]] = []
+        # Declared identity decode has no per-row updates or hook work.
+        skip_decode_rows = self._decode_preprocess_is_identity and not self._eager_mtp
+        is_prefilling_np = getattr(input_batch, "is_prefilling_np", None) if skip_decode_rows else None
         for i, req_idx in enumerate(req_indices):
+            if is_prefilling_np is not None and not is_prefilling_np[i]:
+                continue
             buf = self.intermediate_buffer.buffers[req_idx]
             if not buf or "req_id" not in buf:
                 continue
@@ -791,6 +807,44 @@ class OmniModelState(DefaultModelState):
                     prepacked_mtp_inputs=prepacked_mtp_inputs,
                 )
 
+    def publish_sampled_embeddings(
+        self, input_batch: InputBatch, sampled_token_ids: torch.Tensor
+    ) -> tuple[dict[str, Any], torch.cuda.Event] | None:
+        """Embedding of the token each row sampled this step.
+
+        A producer whose consumer is fed the embeddings of its generated tokens
+        (``publishes_sampled_embeddings`` on the model) can then hand each token
+        over with the step that sampled it instead of one step later, when that
+        token's forward would capture the same embedding. Published beside the
+        packed snapshot as ``embed.sampled``: one ``[1, H]`` row per row whose
+        sample is kept, an empty tensor for rows still inside their prefill.
+        """
+        if not getattr(self.model, "publishes_sampled_embeddings", False):
+            return None
+        num_reqs = input_batch.num_reqs
+        if not num_reqs:
+            return None
+        if sampled_token_ids.reshape(num_reqs, -1).shape[1] != 1:
+            return None  # speculative steps sample several tokens per row
+        computed = input_batch.num_computed_prefill_tokens_np[:num_reqs]
+        scheduled = np.asarray(input_batch.num_scheduled_tokens[:num_reqs])
+        finishing = computed + scheduled >= input_batch.prefill_len_np[:num_reqs]
+        kept = ~input_batch.is_prefilling_np[:num_reqs] | finishing
+        rows = np.flatnonzero(kept).tolist()
+        if not rows:
+            return None
+        device = sampled_token_ids.device
+        first = sampled_token_ids.reshape(num_reqs, -1)[:, 0].index_select(0, index_to_device(rows, device))
+        embeds = self.model.embed_input_ids(first.long())
+        empty = torch.empty(0, dtype=embeds.dtype)
+        sampled: list[torch.Tensor] = [empty] * num_reqs
+        for k, row in enumerate(rows):
+            sampled[row] = embeds[k : k + 1]
+        done = torch.cuda.Event()
+        done.record()
+        extra: dict[str, Any] = {"embed": {"sampled": sampled}}
+        return extra, done
+
     def _pack_mtp_batch(
         self,
         mtp_batches: list[tuple[int, int, tuple[torch.Tensor, torch.Tensor]]],
@@ -1288,6 +1342,16 @@ class OmniModelState(DefaultModelState):
         Handles ``OmniOutput`` unwrapping and ``make_omni_output``
         conversion.
         """
+        make_output_mrv2 = getattr(self.model, "make_omni_output_mrv2", None)
+        if not isinstance(model_output, OmniOutput) and callable(make_output_mrv2) and self.have_multimodal_outputs:
+            # Device-side output: the hook reads this step's batch and request
+            # state directly instead of per-request host views.
+            model_output = make_output_mrv2(
+                model_output,
+                input_batch=input_batch,
+                req_states=req_states,
+                model_intermediate_buffer=self.intermediate_buffer.gather(input_batch),
+            )
         if not isinstance(model_output, OmniOutput) and hasattr(self.model, "make_omni_output"):
             if isinstance(model_output, (list, tuple)) or self.have_multimodal_outputs:
                 buffer_list = self.intermediate_buffer.gather(input_batch)

vllm_omni/worker_v2/omni_ar_model_runner.py

@@ -201,6 +201,17 @@ class OmniARModelRunner(OmniGPUModelRunner):
                 multimodal_outputs,
                 self._dispatch_mtp_batch_descriptor,
             )
+        publish_sampled = getattr(self.model_state, "publish_sampled_embeddings", None)
+        extra_outputs = (
+            publish_sampled(input_batch, sampler_output.sampled_token_ids)
+            if (
+                multimodal_outputs
+                and callable(publish_sampled)
+                and bool(getattr(self.model_config, "async_chunk", False))
+                and getattr(self.model_config, "engine_output_type", "text") != "text"
+            )
+            else None
+        )
         if self.pp_handler is not None:
             self.pp_handler.broadcast(
                 sampler_output.sampled_token_ids,
@@ -251,6 +262,7 @@ class OmniARModelRunner(OmniGPUModelRunner):
             finalize_output=(None if materialize_native else self._finalize_native_data_plane_output),
             check_ep_fault=self.check_ep_fault,
             routed_experts=routed_experts,
+            extra_multimodal_outputs=extra_outputs,
         )
         self._release_multimodal_snapshot(snapshot_slot, async_output.copy_event)
         _guard_graph_replay_for_pooler_copy(
@@ -543,6 +555,18 @@ def _async_copy_mm(
     }
 
 
+def _merge_payload_trees(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]:
+    """``base`` with ``extra`` merged in; nested mappings merge key by key."""
+    merged = dict(base)
+    for key, value in extra.items():
+        current = merged.get(key)
+        if isinstance(current, dict) and isinstance(value, dict):
+            merged[key] = _merge_payload_trees(current, value)
+        else:
+            merged[key] = value
+    return merged
+
+
 def _slice_pooler_value(
     value: Any,
     *,
@@ -649,6 +673,7 @@ class OmniAsyncOutput(AsyncModelRunnerOutput):
         finalize_output: Any | None = None,
         check_ep_fault: bool = False,
         routed_experts: RoutedExpertsTensors | None = None,
+        extra_multimodal_outputs: tuple[dict[str, Any], torch.cuda.Event] | None = None,
     ):
         self.model_runner_output = model_runner_output
         self.sampler_output = sampler_output
@@ -745,6 +770,20 @@ class OmniAsyncOutput(AsyncModelRunnerOutput):
                     copy_stream=copy_stream,
                     pin_memory=pin_memory,
                 )
+                if extra_multimodal_outputs:
+                    # Produced after sampling on the producer stream; copy only
+                    # once its completion event has been observed.
+                    extra_outputs, extra_ready = extra_multimodal_outputs
+                    copy_stream.wait_event(extra_ready)
+                    self._mm_snapshot = _merge_payload_trees(
+                        self._mm_snapshot,
+                        _async_copy_mm(
+                            extra_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,

vllm_omni/worker_v2/omni_data_plane.py

@@ -51,6 +51,7 @@ class _NativeRequestState:
             output_token_ids=output,
             all_token_ids=prompt + output,
             output_token_count=len(self.output_token_ids),
+            last_output_token_id=self.output_token_ids[-1] if self.output_token_ids else None,
             additional_information=self.additional_information,
             sampling_params=self.sampling_params,
             num_computed_tokens=self.num_computed_tokens,
@@ -83,6 +84,11 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin):
         self._native_send_lock = threading.Lock()
         self._native_outputs_in_flight: dict[str, int] = defaultdict(int)
         self._native_terminal_pending: set[str] = set()
+        # Stage payload builders read the stage config through
+        # ``_get_model_config()`` (e.g. Qwen3-Omni's Thinker reads
+        # ``hf_config.talker_config.accept_hidden_layer``).
+        self.vllm_config = vllm_config
+        self.model_config = model_config
         self.init_omni_connectors(model_config=model_config)
         self._delivery_manager = OmniDeliveryManager(
             delivery_timeout_s=self._connector_delivery_timeout(),
@@ -394,7 +400,9 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin):
                 for qual, qval in value.items():
                     if qual in span_handled:
                         continue
-                    if key == "meta" and qual == "finished":
+                    if key == "meta":
+                        # Per-chunk state (``finished``, the V1 sender's
+                        # ``is_segment_finished`` flag, ...), not rows.
                         merged_sub[qual] = qval
                         continue
                     if (key, qual) in override_keys:
@@ -563,6 +571,7 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin):
         self._delivery_manager.shutdown(RuntimeError("MRv2 connector shutdown"))
         self._stop_event.set()
         self._work_available.set()
+        self._save_work_available.set()
 
         close_errors: list[BaseException] = []
         connector = getattr(self, "_omni_connector", None)

vllm_omni/worker_v2/omni_generation_model_runner.py

@@ -56,6 +56,23 @@ def _contains_cuda_tensor(value: Any) -> bool:
     return False
 
 
+def check_exact_input_shape(model: Any, input_ids: torch.Tensor | None, num_tokens: int) -> None:
+    """Enforce ``requires_exact_input_shape`` (#6712).
+
+    Such models split the flat ``input_ids`` by per-request
+    ``seq_token_counts``, so padding rows would misalign every request. V1
+    trims cudagraph padding for them (``maybe_unpad_input_ids``); this runner
+    dispatches eagerly and must already hand over the exact batch.
+    """
+    if input_ids is None or not getattr(model, "requires_exact_input_shape", False):
+        return
+    if int(input_ids.shape[0]) != int(num_tokens):
+        raise RuntimeError(
+            f"{type(model).__name__} requires an exact input shape: got {int(input_ids.shape[0])} "
+            f"input ids for {int(num_tokens)} scheduled tokens"
+        )
+
+
 def _materialize_generation_value(
     value: Any,
     req_index: int,
@@ -368,6 +385,8 @@ class OmniGenerationModelRunner(OmniGPUModelRunner):
             input_batch,
             dummy_run=dummy_run,
         )
+        if not dummy_run:
+            check_exact_input_shape(self.model, input_ids, input_batch.num_tokens)
 
         model_inputs = {
             "input_ids": input_ids,
@@ -376,6 +395,10 @@ class OmniGenerationModelRunner(OmniGPUModelRunner):
             "intermediate_tensors": intermediate_tensors,
             **self.model_state.prepare_inputs(input_batch, self.req_states),
         }
+        if getattr(self.model, "requires_request_ids", False):
+            # Request-owned model state (e.g. MiniCPM-o's vocoder caches) is
+            # keyed by the scheduler request id, in batch order.
+            model_inputs["request_ids"] = list(input_batch.req_ids)
         self._add_legacy_forward_inputs(model_inputs, input_batch)
 
         eplb = getattr(self, "eplb", None)

vllm_omni/worker_v2/omni_model_runner.py

@@ -9,6 +9,7 @@ import threading
 from typing import TYPE_CHECKING, Any, cast
 
 import torch
+from torch.utils._pytree import tree_flatten, tree_unflatten
 from vllm.config.compilation import CUDAGraphMode
 from vllm.forward_context import (
     set_forward_context,
@@ -55,6 +56,16 @@ def _needs_capture_tensor_unwrap(model: Any) -> bool:
     return bool(getattr(model, "_returns_tuple", False))
 
 
+def _supports_full_graph_aux_outputs(model: Any) -> bool:
+    """Tuple-returning model whose auxiliary output FULL graphs may replay.
+
+    The model declares ``supports_mrv2_full_graph_aux_outputs`` when its
+    auxiliary output is a tensor-only pytree with the same structure for every
+    batch shape and each leaf's leading dimension is the token axis.
+    """
+    return _needs_capture_tensor_unwrap(model) and bool(getattr(model, "supports_mrv2_full_graph_aux_outputs", False))
+
+
 class OmniGPUModelRunner(GPUModelRunner):
     """Thin layer over v2 ``GPUModelRunner`` for Omni lifecycle hooks."""
 
@@ -62,14 +73,48 @@ class OmniGPUModelRunner(GPUModelRunner):
     encoder_cache: EncoderCache | None
     _native_output_materializer: NativeOutputWorker | None
     model_state: OmniModelState
+    vocab_size: int
+    use_aux_hidden_state_outputs: bool
     _last_aux_output: Any
     _last_multimodal_outputs: dict[str, Any] | None
     _model_returns_tuple: bool
+    # FULL-graph aux-output contract (see _configure_cudagraph_output_contract).
+    _full_graph_aux_outputs = False
+    _aux_output_spec: Any = None
 
     def _configure_cudagraph_output_contract(self) -> None:
         """Select the CUDA graph output contract declared by the model."""
         self._model_returns_tuple = _needs_capture_tensor_unwrap(self.model)
-        self._exclude_full_graph = self._model_returns_tuple or hasattr(self.model, "_last_captured_layers")
+        # FULL replay returns what capture stored: the hidden states and, for a
+        # declared tuple model, its auxiliary pytree flattened to leaves.
+        self._full_graph_aux_outputs = _supports_full_graph_aux_outputs(self.model)
+        self._aux_output_spec = None
+        self._exclude_full_graph = (self._model_returns_tuple and not self._full_graph_aux_outputs) or hasattr(
+            self.model, "_last_captured_layers"
+        )
+
+    def _split_fullgraph_output(self, replayed: Any) -> tuple[Any, Any]:
+        """FULL replay output -> ``(hidden_states, aux pytree or None)``."""
+        if self._full_graph_aux_outputs and isinstance(replayed, tuple):
+            hidden, leaves = replayed
+            return hidden, tree_unflatten(list(leaves), self._aux_output_spec)
+        return replayed, None
+
+    def _flatten_capture_aux_output(self, model_output: Any) -> tuple[torch.Tensor, list[torch.Tensor]]:
+        """``(hidden, aux pytree)`` -> ``(hidden, leaves)`` for the graph manager's aux buffers."""
+        hidden, aux = model_output
+        leaves, spec = tree_flatten(aux)
+        num_tokens = hidden.shape[0]
+        if not leaves or any(not isinstance(leaf, torch.Tensor) or leaf.shape[:1] != (num_tokens,) for leaf in leaves):
+            raise RuntimeError(
+                f"{type(self.model).__name__} declares supports_mrv2_full_graph_aux_outputs, but its auxiliary "
+                "output is not a non-empty tensor pytree on the token axis"
+            )
+        if self._aux_output_spec is None:
+            self._aux_output_spec = spec
+        elif spec != self._aux_output_spec:
+            raise RuntimeError(f"{type(self.model).__name__} auxiliary output structure changed between captures")
+        return hidden, leaves
 
     @staticmethod
     def _prepare_cudagraph_capture_output(model_output: Any) -> Any:
@@ -171,10 +216,16 @@ class OmniGPUModelRunner(GPUModelRunner):
         # With no new requests, its UVA snapshots remain valid and immutable.
         # Keep the upstream path for custom/speculative samplers, whose staged
         # writes may have a different lifecycle.
-        if not scheduler_output.scheduled_new_reqs and type(getattr(self, "sampler", None)) is Sampler:
+        sampler = getattr(self, "sampler", None)
+        if not scheduler_output.scheduled_new_reqs and (
+            type(sampler) is Sampler or getattr(sampler, "omni_static_staged_writes", False)
+        ):
             return
         logits_processor = getattr(self.model, "logits_processor", None)
         logits_vocab = getattr(logits_processor, "vocab_size", None)
+        if logits_vocab is None:
+            # Narrow model heads without a vLLM LogitsProcessor (codec heads).
+            logits_vocab = getattr(self.model, "logits_vocab_size", None)
         if isinstance(logits_vocab, int) and logits_vocab > 0:
             for request_data in scheduler_output.scheduled_new_reqs:
                 sampling_params = request_data.sampling_params
@@ -289,9 +340,17 @@ class OmniGPUModelRunner(GPUModelRunner):
     def load_model(self, *args: Any, **kwargs: Any) -> None:
         import vllm.v1.worker.gpu.model_runner as _mr_module
 
+        def _init_model_state(vllm_config: Any, model: Any, encoder_cache: Any, device: Any) -> Any:
+            # Runs after the model is built and before the sampler is: a model
+            # whose logits head is narrower than its HF config reports (the
+            # MiniCPM-o codec Talker's config has no vocab_size at all) sizes
+            # the MRv2 sampler and request state by that head.
+            self._adopt_model_logits_vocab(model)
+            return init_omni_model_state(vllm_config, model, encoder_cache, device)
+
         with _model_state_patch_lock:
             _orig = _mr_module.init_model_state
-            _mr_module.init_model_state = init_omni_model_state
+            _mr_module.init_model_state = _init_model_state
             try:
                 super().load_model(*args, **kwargs)
             finally:
@@ -305,6 +364,18 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.supports_mm_inputs = False
             self.encoder_cache = None
 
+    def _adopt_model_logits_vocab(self, model: Any) -> None:
+        vocab = getattr(model, "logits_vocab_size", None)
+        if not isinstance(vocab, int) or vocab <= 0 or vocab == self.vocab_size:
+            return
+        logger.info("Sizing the MRv2 sampler by the model's logits head: vocab %s -> %d", self.vocab_size, vocab)
+        self.vocab_size = vocab
+        self.req_states.vocab_size = vocab
+        # Kernel warmup sizes its structured-output bitmask from the config.
+        arch_config = getattr(self.model_config, "model_arch_config", None)
+        if arch_config is not None and getattr(arch_config, "vocab_size", vocab) != vocab:
+            arch_config.vocab_size = vocab
+
     # ------------------------------------------------------------------
     # CUDA Graph: conditionally exclude FULL mode
     # ------------------------------------------------------------------
@@ -336,19 +407,27 @@ class OmniGPUModelRunner(GPUModelRunner):
             self._exclude_unsupported_full_graphs()
 
         # Wrap model forward during capture so tuple returns don't crash
-        # torch.empty_like() in the PIECEWISE warmup pass.
+        # torch.empty_like() in the PIECEWISE warmup pass. A model with the
+        # FULL aux-output contract instead hands the graph manager its leaves,
+        # which it keeps in persistent buffers and returns on replay.
         if self._model_returns_tuple:
             original_forward = self.model.forward
+            aux_outputs = self._full_graph_aux_outputs
 
             def _capture_forward(*args: Any, **kwargs: Any) -> Any:
                 output = original_forward(*args, **kwargs)
+                if aux_outputs:
+                    return self._flatten_capture_aux_output(output)
                 return self._prepare_cudagraph_capture_output(output)
 
             self.model.forward = _capture_forward  # type: ignore[assignment]
+            use_aux = self.use_aux_hidden_state_outputs
+            self.use_aux_hidden_state_outputs = use_aux or aux_outputs
             try:
                 result = super().capture_model()
             finally:
                 self.model.forward = original_forward  # type: ignore[assignment]
+                self.use_aux_hidden_state_outputs = use_aux
         else:
             result = super().capture_model()
 
@@ -554,9 +633,9 @@ class OmniGPUModelRunner(GPUModelRunner):
             # inputs_embeds buffer above.
             assert self.cudagraph_manager is not None
             self.kv_connector.pre_forward(scheduler_output)
-            hidden_states = self.cudagraph_manager.run_fullgraph(batch_desc)
-            self._last_aux_output = None
+            replayed = self.cudagraph_manager.run_fullgraph(batch_desc)
             self._last_multimodal_outputs = None
+            hidden_states, self._last_aux_output = self._split_fullgraph_output(replayed)
             if hasattr(self.model, "_last_captured_layers"):
                 self.model._last_captured_layers = self._last_aux_output
         else: