[Bugfix] Restore MRV2 full-payload transport for offline Qwen3-TTS (#8010)

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

tests/core/sched/test_generation_scheduler_restore.py

@@ -133,7 +133,9 @@ def test_chunk_lifecycle_no_resubmit_and_state_survives_requeue(monkeypatch) ->
     monkeypatch.setattr("vllm_omni.core.sched.omni_generation_scheduler.create_request_queue", create_request_queue)
     scheduler._native_data_plane = True
     scheduler.chunk_transfer_adapter = None
-    scheduler.input_coordinator = SimpleNamespace(finished_requests=set(), restore_queues=lambda *args, **kwargs: None)
+    scheduler.input_coordinator = SimpleNamespace(
+        _async_chunk=True, finished_requests=set(), restore_queues=lambda *args, **kwargs: None
+    )
     scheduler.running = []
     scheduler.waiting = create_request_queue(scheduler.policy)
     scheduler.skipped_waiting = create_request_queue(scheduler.policy)

tests/core/sched/test_omni_ar_scheduler_streaming.py

@@ -1527,7 +1527,9 @@ def test_async_chunk_reserves_parked_slots_during_ar_admission(monkeypatch, nati
     sched.use_v2_model_runner = native
     sched.max_num_running_reqs = 8
     sched.input_coordinator = (
-        SimpleNamespace(_waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None) if native else None
+        SimpleNamespace(_async_chunk=True, _waiting_for_chunk_running=[parked], restore_queues=lambda _w, _r: None)
+        if native
+        else None
     )
     sched.chunk_transfer_adapter = (
         None

tests/core/sched/test_omni_scheduler_mixin_shared.py

@@ -6,8 +6,10 @@ from types import SimpleNamespace
 from typing import Any
 
 import pytest
+from vllm.config import SchedulerConfig, VllmConfig
 from vllm.v1.engine import FinishReason
 
+from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.core.sched import omni_scheduler_mixin
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.core.sched.output import OmniChunkRecvHandle
@@ -90,6 +92,31 @@ def test_full_payload_coordinator_matches_legacy_gate(monkeypatch, stage_id, asy
     assert (scheduler.input_coordinator is not None) is enabled
 
 
+@pytest.mark.parametrize("async_chunk", [False, True])
+@pytest.mark.parametrize(("stage_id", "required"), [(0, False), (1, False), (1, True)])
+def test_native_data_plane_uses_the_selected_input_protocol(async_chunk, stage_id, required):
+    scheduler = _Scheduler()
+    model_config = object.__new__(OmniModelConfig)
+    model_config.stage_id = stage_id
+    model_config.async_chunk = async_chunk
+    model_config.requires_full_payload_input = required
+    model_config.supports_native_mrv2_data_plane = True
+    model_config.use_v2_model_runner = True
+    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
+    if async_chunk or (stage_id > 0 and required):
+        assert scheduler.input_coordinator._async_chunk is async_chunk
+    else:
+        assert scheduler.input_coordinator is None
+    assert scheduler._async_chunk_transport_enabled() is async_chunk
+    assert scheduler.chunk_transfer_adapter is None
+
+
 def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
     calls: list[tuple[Any, ...]] = []
     scheduler = _Scheduler()

tests/core/sched/test_omni_scheduler_mixin_timeouts.py

@@ -6,12 +6,8 @@
 ``_process_pending_chunk_timeouts`` covers async-chunk. Both share
 ``VLLM_OMNI_INPUT_WAIT_TIMEOUT_S``.
 
-Verifies that the mixin correctly *delegates* timed-out requests to the
-base scheduler's ``finish_requests`` API with ``RequestStatus.FINISHED_ERROR``.
-The end-to-end effect (queue removal + status set + per-request cleanup +
-client-facing FINISHED_ERROR emission) is the responsibility of upstream
-vLLM's ``finish_requests`` implementation and is covered by upstream tests;
-this file only asserts the wiring from the mixin to that API.
+Verifies scheduler cleanup and explicit client-facing ERROR outputs. Upstream
+``finish_requests`` frees requests but does not emit their terminal output.
 """
 
 from __future__ import annotations
@@ -23,6 +19,9 @@ import os
 from types import SimpleNamespace
 
 import pytest
+from vllm import SamplingParams
+from vllm.v1.engine import EngineCoreOutputs
+from vllm.v1.request import Request
 
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 
@@ -52,7 +51,7 @@ class _FakeScheduler(OmniSchedulerMixin):
 def test_process_pending_input_timeouts_delegates_to_finish_requests():
     """Timed-out request present in self.requests is forwarded to finish_requests."""
     req_id = "stuck-req"
-    requests = {req_id: SimpleNamespace(request_id=req_id)}
+    requests = {req_id: SimpleNamespace(request_id=req_id, client_index=0)}
     coord = _FakeCoordinator(timed_out_ids={req_id})
     scheduler = _FakeScheduler(requests, coord)
 
@@ -142,7 +141,7 @@ def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
     """
     req_id = "stalled-stream"
     adapter = _FakeChunkAdapter(timed_out_ids={req_id})
-    scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id)}, adapter)
+    scheduler = _FakeChunkScheduler({req_id: SimpleNamespace(request_id=req_id, client_index=0)}, adapter)
 
     scheduler._process_pending_chunk_timeouts()
 
@@ -154,6 +153,41 @@ def test_process_pending_chunk_timeouts_delegates_to_finish_requests():
     assert getattr(status, "name", str(status)).endswith("FINISHED_ERROR")
 
 
+@pytest.mark.parametrize("chunked", [False, True])
+@pytest.mark.parametrize("synthesize_abort_outputs", [False, True])
+def test_timeout_emits_error_once_even_without_model_outputs(chunked, synthesize_abort_outputs):
+    from vllm.v1.engine import FinishReason
+
+    request = Request(
+        request_id="stuck",
+        prompt_token_ids=[1],
+        sampling_params=SamplingParams(max_tokens=1),
+        pooling_params=None,
+        client_index=3,
+    )
+    if chunked:
+        scheduler = _FakeChunkScheduler({"stuck": request}, _FakeChunkAdapter({"stuck"}))
+        scheduler._process_pending_chunk_timeouts()
+    else:
+        scheduler = _FakeScheduler({"stuck": request}, _FakeCoordinator({"stuck"}))
+        scheduler._process_pending_input_timeouts()
+    # Model upstream's finish_requests cleanup. Generation must still emit
+    # ERROR, and AR's synthetic ABORT must not replace or duplicate it.
+    scheduler.requests.clear()
+    scheduler.finished_req_ids_dict = {3: {"stuck"}}
+    outputs: dict[int, EngineCoreOutputs] = {}
+    scheduler._attach_finished_request_sets(outputs, synthesize_abort_outputs=synthesize_abort_outputs)
+
+    [error] = outputs[3].outputs
+    assert error.request_id == "stuck"
+    assert error.finish_reason == FinishReason.ERROR
+    assert "connector input" in error.stop_reason
+    assert outputs[3].finished_requests == {"stuck"}
+    outputs = {}
+    scheduler._attach_finished_request_sets(outputs, synthesize_abort_outputs=synthesize_abort_outputs)
+    assert outputs == {}
+
+
 def test_process_pending_chunk_timeouts_skips_already_freed_request():
     adapter = _FakeChunkAdapter(timed_out_ids={"already-freed"})
     scheduler = _FakeChunkScheduler(requests={}, adapter=adapter)

tests/engine/test_orchestrator.py

@@ -2000,6 +2000,47 @@ async def test_abort_retry_does_not_repeat_successful_stage_abort():
     assert second.bound == set()
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("stage_id", [0, 1])
+async def test_raw_stage_error_reaches_caller_before_normal_terminal_routing(mocker, stage_id):
+    pools = _build_stage_pools([[FakeStageClient()], [FakeStageClient()]])
+    for pool in pools:
+        mocker.patch.object(pool, "process_llm_raw_outputs", new_callable=mocker.AsyncMock, return_value=[])
+    output_queue: asyncio.Queue[ErrorMessage] = asyncio.Queue()
+    orchestrator = Orchestrator(
+        request_async_queue=asyncio.Queue(),
+        output_async_queue=output_queue,
+        rpc_async_queue=asyncio.Queue(),
+        stage_pools=[],
+    )
+    orchestrator.stage_pools = pools
+    orchestrator.request_states["stuck"] = OrchestratorRequestState(request_id="stuck", final_stage_id=1)
+    mocker.patch.object(orchestrator, "_handle_kv_ready_raw_outputs", new_callable=mocker.AsyncMock)
+    cleanup = mocker.patch.object(orchestrator, "_cleanup_request_ids", new_callable=mocker.AsyncMock)
+    normal_terminal = mocker.patch.object(
+        orchestrator, "_apply_raw_terminal_stage_finish", new_callable=mocker.AsyncMock
+    )
+    reason = "Timed out waiting for connector input after 5s"
+    raw = EngineCoreOutputs(
+        outputs=[
+            OmniEngineCoreOutput(
+                request_id="stuck", new_token_ids=[], finish_reason=FinishReason.ERROR, stop_reason=reason
+            )
+        ]
+    )
+    terminal_ids: set[str] = set()
+
+    await orchestrator._process_llm_stage_outputs(stage_id, 0, raw, terminal_ids)
+
+    error = output_queue.get_nowait()
+    assert isinstance(error, ErrorMessage)
+    assert (error.request_id, error.stage_id, error.error) == ("stuck", stage_id, reason)
+    cleanup.assert_awaited_once_with(["stuck"], abort=True, release_owners=True)
+    normal_terminal.assert_not_awaited()
+    assert not terminal_ids
+    assert output_queue.empty()
+
+
 @pytest.mark.asyncio
 async def test_duplex_session_request_error_finish_is_delivered_as_request_error() -> None:
     """A session-owned request the scheduler finished with FinishReason.ERROR

tests/worker_v2/test_omni_data_plane.py

@@ -48,6 +48,8 @@ def _bare_plane(*, delivery_timeout_s=1.0, shutdown_timeout_s=1.0):
         _native_output_closed=False,
         _native_outputs_in_flight=defaultdict(int),
         _native_terminal_pending=set(),
+        _async_chunk=True,
+        _pending_full_payload_send={},
         _put_req_chunk=defaultdict(int),
         _ramp_chunk_count=defaultdict(int),
         _delivery_manager=OmniDeliveryManager(
@@ -82,6 +84,76 @@ def plane():
         p._stop_output_worker()
 
 
+def test_full_payload_waits_for_terminal_and_last_deferred_frame(plane):
+    plane._async_chunk = False
+    plane._full_payload_replace_keys_cached = frozenset({"codes.ref"})
+    request = _new_request()
+    request.resumable = False
+    plane.register_request(request)
+    first = torch.tensor([[1, 2]])
+    last = torch.tensor([[3, 4]])
+    ref = torch.tensor([[5, 6]])
+    plane.reserve_outputs(["internal"])
+    assert _complete(plane, [{"codes.audio": first, "codes.ref": ref}], token=21) == 0
+    assert not plane.record.batches
+
+    plane.reserve_outputs(["internal"])
+    assert plane.request_terminal({"internal"}) == 0
+    assert not plane.record.cleaned
+    assert _complete(plane, [{"codes.audio": last, "codes.ref": ref}], token=2150) == 1
+
+    [(snapshot, payload)] = plane.record.batches[0]
+    assert snapshot.is_finished()
+    assert snapshot.output_token_ids == [21, 2150]
+    torch.testing.assert_close(payload["codes.audio"], torch.cat([first, last]))
+    torch.testing.assert_close(payload["codes.ref"], ref)
+    assert plane.record.cleaned == ["internal"]
+    assert not plane._pending_full_payload_send
+    assert plane.request_terminal({"internal"}) == 0
+
+
+def test_full_payload_abort_discards_partial_and_late_outputs(plane):
+    plane._async_chunk = False
+    plane._full_payload_replace_keys_cached = frozenset()
+    plane.register_request(_new_request())
+    assert _complete(plane, [{"codes.audio": torch.tensor([[1, 2]])}]) == 0
+    plane.reserve_outputs(["internal"])
+    assert plane.abort_requests({"internal"}) == 1
+    [(snapshot, payload)] = plane.record.batches[0]
+    assert snapshot.is_finished() and payload is None
+    assert _complete(plane, [{"codes.audio": torch.tensor([[3, 4]])}]) == 0
+    assert len(plane.record.batches) == 1
+    assert not plane._pending_full_payload_send
+
+
+def test_full_payload_qwen_builder_receives_complete_codec(raw_plane, monkeypatch):
+    from vllm_omni.model_executor.stage_input_processors.qwen3_tts import talker2code2wav_full_payload
+
+    raw_plane._async_chunk = False
+    raw_plane._omni_connector = object()
+    raw_plane._request_ids_mapping = {}
+    raw_plane._custom_process_func = talker2code2wav_full_payload
+    raw_plane._custom_process_batch_func = None
+    monkeypatch.setattr(raw_plane, "is_data_transfer_rank", lambda: True)
+    sent = []
+
+    def publish(entries, **kwargs):
+        sent.extend(entries)
+        return len(entries)
+
+    monkeypatch.setattr(raw_plane, "_publish_chunk_cohort", publish)
+    raw_plane.register_request(_new_request())
+    for token in (21, 22, 2150):
+        _complete(raw_plane, [{"codes.audio": torch.full((1, 16), token)}], token=token)
+    assert not sent
+    assert raw_plane.request_terminal({"internal"}) == 1
+    [(request, payload)] = sent
+    assert request.output_token_ids == [21, 22, 2150]
+    assert payload["meta"]["finished"].item()
+    # EOS is filtered, and the two remaining frames are codebook-major.
+    assert payload["codes"]["audio"].tolist() == [21, 22] * 16
+
+
 def _gated_connector():
     # put() blocks its first call until released (or forever without release).
     conn = SimpleNamespace(

tests/worker_v2/test_omni_gpu_model_runner.py

@@ -15,6 +15,7 @@ from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
 from vllm.v1.worker.gpu.sample.sampler import Sampler
 
+from vllm_omni.config.model import OmniModelConfig
 from vllm_omni.model_executor.models.output_templates import OmniOutput
 from vllm_omni.worker_v2.model_states import init_omni_model_state
 from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState
@@ -53,6 +54,8 @@ def test_add_requests_empty_admission_and_stop_id_sanitization():
 
 def test_prepare_native_data_plane_terminal_abort_split_and_warmup_skip():
     runner = _make_runner()
+    runner.model_config = object.__new__(OmniModelConfig)
+    runner.model_config.async_chunk = True
     plane = SimpleNamespace(
         register_request=MagicMock(),
         register_receivers=MagicMock(),
@@ -78,6 +81,19 @@ def test_prepare_native_data_plane_terminal_abort_split_and_warmup_skip():
     plane.abort_requests.assert_called_once_with({"aborted"})
 
 
+def test_full_payload_receive_is_polled_without_scheduled_tokens(mocker):
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.model_config = object.__new__(OmniModelConfig)
+    runner.model_config.async_chunk = False
+    plane = mocker.Mock()
+    runner._omni_data_plane = plane
+    scheduler_output = SchedulerOutput.make_empty()
+
+    runner._prepare_native_data_plane(scheduler_output)
+
+    plane.recv_full_payload_inputs.assert_called_once_with(scheduler_output)
+
+
 @pytest.mark.parametrize("output_form", ["tuple", "omni"])
 def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form):
     runner = object.__new__(OmniGPUModelRunner)

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -189,7 +189,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         Stateful codecs retain lifetime admission until their decoder state is
         released, even when no chunk is currently executing.
         """
-        if not getattr(self, "_native_data_plane", False) or self._retains_state_across_chunks:
+        if (
+            not getattr(self, "_native_data_plane", False)
+            or not self._async_chunk_transport_enabled()
+            or self._retains_state_across_chunks
+        ):
             return
         in_flight: list[Request] = []
         for request in self.running:
@@ -236,7 +240,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._drop_aborted_queued_requests()
         self._resync_streaming_input_counter()
         async_chunk_transport = self._async_chunk_transport_enabled()
-        native_chunks = bool(getattr(self, "_native_data_plane", False))
+        native_chunks = bool(getattr(self, "_native_data_plane", False) and async_chunk_transport)
         # Parking releases an execution slot, but stateful codecs retain their
         # request-owned state until completion or abort, in either runner.
         reserved_running_slots = (

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -112,6 +112,7 @@ class OmniSchedulerMixin:
     requests: dict[str, Request]
     waiting: RequestQueue
     running: list[Request]
+    _pending_input_timeout_outputs: dict[str, tuple[int, OmniEngineCoreOutput]]
 
     def _init_omni_connector_output_inbox(self) -> None:
         self._omni_connector_output_inbox: queue.SimpleQueue[OmniConnectorOutput] = queue.SimpleQueue()
@@ -149,7 +150,7 @@ class OmniSchedulerMixin:
             else None
         )
         self.input_coordinator: OmniSchedulingCoordinator | None = None
-        if self._native_data_plane:
+        if self._native_data_plane and getattr(model_config, "async_chunk", False):
             self.input_coordinator = OmniSchedulingCoordinator(
                 scheduler_max_num_seqs=self.vllm_config.scheduler_config.max_num_seqs,
                 stage_id=getattr(model_config, "stage_id", 0),
@@ -251,6 +252,7 @@ class OmniSchedulerMixin:
     def _async_chunk_transport_enabled(self) -> bool:
         return getattr(self, "chunk_transfer_adapter", None) is not None or bool(
             getattr(self, "_native_data_plane", False)
+            and getattr(getattr(self, "input_coordinator", None), "_async_chunk", False)
         )
 
     def _get_async_chunk_reserved_running_slots(self) -> int:
@@ -453,7 +455,29 @@ class OmniSchedulerMixin:
             DEFAULT_INPUT_WAIT_TIMEOUT_S,
             sorted(present_ids),
         )
-        self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR)
+        self._finish_input_timeout_requests(present_ids)
+
+    def _finish_input_timeout_requests(self, request_ids: set[str]) -> None:
+        # Upstream finish_requests frees scheduler state, but does not emit
+        # EngineCoreOutput. Keep an explicit ERROR for both AR and generation
+        # callers, including a tick with no scheduled model work.
+        pending = getattr(self, "_pending_input_timeout_outputs", None)
+        if pending is None:
+            pending = self._pending_input_timeout_outputs = {}
+        for request_id in request_ids:
+            request = self.requests[request_id]
+            reason = f"Timed out waiting for connector input after {DEFAULT_INPUT_WAIT_TIMEOUT_S:g}s"
+            request.stop_reason = reason
+            pending[request_id] = (
+                request.client_index,
+                OmniEngineCoreOutput(
+                    request_id=request_id,
+                    new_token_ids=[],
+                    finish_reason=FinishReason.ERROR,
+                    stop_reason=reason,
+                ),
+            )
+        self.finish_requests(request_ids, RequestStatus.FINISHED_ERROR)
 
     def _log_failed_chunk_sends(self) -> None:
         """Surface chunks the sender gave up on (R1.2 of #4855).
@@ -531,7 +555,7 @@ class OmniSchedulerMixin:
             DEFAULT_INPUT_WAIT_TIMEOUT_S,
             sorted(present_ids),
         )
-        self.finish_requests(present_ids, RequestStatus.FINISHED_ERROR)
+        self._finish_input_timeout_requests(present_ids)
 
     def _capture_omni_connector_output(self, model_runner_output: Any) -> None:
         """Stash the model runner's omni_connector_output for next schedule().
@@ -712,6 +736,13 @@ class OmniSchedulerMixin:
         synthesize_abort_outputs: bool,
     ) -> None:
         """Attach finished IDs while keeping AR's synthetic-abort policy explicit."""
+        pending = getattr(self, "_pending_input_timeout_outputs", None)
+        if pending:
+            for client_index, error_output in pending.values():
+                output = engine_core_outputs.setdefault(client_index, EngineCoreOutputs())
+                if not any(item.request_id == error_output.request_id for item in output.outputs):
+                    output.outputs.append(error_output)
+            pending.clear()
         finished_req_ids = self.finished_req_ids_dict
         if not finished_req_ids:
             return

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -25,11 +25,7 @@ def uses_native_mrv2_data_plane(
     *,
     use_v2_model_runner: bool,
 ) -> bool:
-    return bool(
-        use_v2_model_runner
-        and getattr(model_config, "async_chunk", False)
-        and getattr(model_config, "supports_native_mrv2_data_plane", False)
-    )
+    return bool(use_v2_model_runner and getattr(model_config, "supports_native_mrv2_data_plane", False))
 
 
 def uses_full_payload_input_coordinator(model_config: Any) -> bool:

vllm_omni/engine/orchestrator.py

@@ -789,6 +789,18 @@ class OrchestratorBase:
             req_state = self.request_states.get(getattr(eco, "request_id", None))
             if req_state is None:
                 continue
+            if (
+                getattr(eco, "finish_reason", None) == FinishReason.ERROR
+                and not getattr(eco, "is_segment_finished", False)
+                and not req_state.session_owned
+            ):
+                reason = getattr(eco, "stop_reason", None)
+                await self._handle_stage_error(
+                    stage_id,
+                    eco,
+                    error=reason if isinstance(reason, str) and reason else "Stage request failed",
+                )
+                continue
             if req_state.streaming.enabled:
                 segment_finished = bool(getattr(eco, "is_segment_finished", False))
                 raw_mm = self._completion_multimodal_output(eco, None)
@@ -1200,7 +1212,7 @@ class OrchestratorBase:
 
             await self._route_output(stage_id, replica_id, output, req_state, stage_metrics)
 
-    async def _handle_stage_error(self, stage_id: int, output: Any) -> None:
+    async def _handle_stage_error(self, stage_id: int, output: Any, *, error: str | None = None) -> None:
         """Emit a frontend-visible error and clean up request state."""
         if self._cfg_tracker.is_companion(output.request_id):
             parent_id = self._cfg_tracker.get_parent_id(output.request_id) or output.request_id
@@ -1210,7 +1222,7 @@ class OrchestratorBase:
             ErrorMessage(
                 request_id=parent_id,
                 stage_id=stage_id,
-                error=output.error,
+                error=error if error is not None else output.error,
                 status_code=getattr(output, "error_status_code", None),
                 error_type=getattr(output, "error_type", None),
             )

vllm_omni/worker_v2/omni_data_plane.py

@@ -15,9 +15,10 @@ from typing import Any, cast
 import torch
 from vllm.logger import init_logger
 
-from vllm_omni.data_entry_keys import OmniPayload, unflatten_payload
+from vllm_omni.data_entry_keys import OmniPayload, flatten_payload, unflatten_payload
 from vllm_omni.worker.omni_connector_model_runner_mixin import (
     OmniConnectorModelRunnerMixin,
+    should_accumulate_full_payload_output,
 )
 from vllm_omni.worker.payload_span import get_tensor_span, merge_tensor_spans
 from vllm_omni.worker_v2.delivery import (
@@ -87,7 +88,11 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin):
             delivery_timeout_s=self._connector_delivery_timeout(),
             shutdown_timeout_s=_DEFAULT_SHUTDOWN_TIMEOUT_S,
         )
-        self._can_send = self._custom_process_func is not None
+        self._can_send = (
+            self._custom_process_func is not None
+            if self._async_chunk
+            else should_accumulate_full_payload_output(model_config, self._custom_process_func)
+        )
         self._start_output_worker(max_pending_batches=_NATIVE_OUTPUT_QUEUE_DEPTH)
 
     def _connector_delivery_timeout(self) -> float:
@@ -294,6 +299,8 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin):
             for req_id in active_req_ids:
                 self._native_outputs_in_flight.pop(req_id, None)
                 self._native_terminal_pending.discard(req_id)
+                # Cancellation must not flush a partially accumulated utterance.
+                self._pending_full_payload_send.pop(req_id, None)
             return self.emit_chunks(
                 req_ids=[],
                 inter_stage_outputs=None,
@@ -494,11 +501,23 @@ class OmniRunnerDataPlane(OmniConnectorModelRunnerMixin):
             state.output_token_ids.extend(int(token_id) for token_id in sampled_by_req.get(req_id, []))
             state.finished = req_id in terminal_req_ids
             payload = payload_by_req.get(req_id)
-            if isinstance(payload, dict):
+            if not self._async_chunk:
+                # Full-payload consumers execute once, after the producer has
+                # finished. Reuse V1's concat/replace contract for flat outputs
+                # and retain the final token history for codec-row cropping.
+                if isinstance(payload, dict):
+                    self.accumulate_full_payload_output(req_id, flatten_payload(payload), state)
+                if not state.finished:
+                    continue
+                accumulated = self._pending_full_payload_send.pop(req_id, None)
+                payload = self._materialize_full_payload_entry(accumulated)[0] if accumulated is not None else None
+            if self._async_chunk and isinstance(payload, dict):
                 payload = unflatten_payload(payload)
             if payload is None and not state.finished:
                 continue
-            include_token_history = bool(state.resumable or self._put_req_chunk.get(state.external_req_id, 0) == 0)
+            include_token_history = bool(
+                not self._async_chunk or state.resumable or self._put_req_chunk.get(state.external_req_id, 0) == 0
+            )
             entries.append(
                 (
                     state.snapshot(include_token_history=include_token_history),

vllm_omni/worker_v2/omni_model_runner.py

@@ -208,6 +208,8 @@ class OmniGPUModelRunner(GPUModelRunner):
                 continue
             plane.register_request(request_data)
         plane.register_receivers(list(getattr(scheduler_output, "pending_input_registrations", [])))
+        if not getattr(self.model_config, "async_chunk", False):
+            plane.recv_full_payload_inputs(scheduler_output)
         natural_terminal_req_ids = set(getattr(scheduler_output, "data_plane_terminal_req_ids", set()))
         aborted_req_ids = set(getattr(scheduler_output, "finished_req_ids", set())).difference(natural_terminal_req_ids)
         if natural_terminal_req_ids: