[2/5] [core]refactor communication layer: PR 2 of 5 Qwen3 Omni non async (#2677)

Signed-off-by: natureofnature <wzliu@connect.hku.hk>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.request import RequestStatus
+
+import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
+import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
+from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class FakeInputCoordinator:
+    def __init__(self) -> None:
+        self.freed: list[str] = []
+
+    def free_finished_request(self, request_id: str) -> None:
+        self.freed.append(request_id)
+
+
+class DummyRequest:
+    request_id = "req-free"
+    client_index = 0
+    additional_information = None
+
+    def is_finished(self) -> bool:
+        return True
+
+
+@pytest.mark.parametrize(
+    ("scheduler_cls", "scheduler_mod"),
+    [
+        (OmniARScheduler, ar_sched_mod),
+        (OmniGenerationScheduler, gen_sched_mod),
+    ],
+)
+def test_finish_requests_cleans_input_coordinator_for_finished_ids(
+    monkeypatch: pytest.MonkeyPatch,
+    scheduler_cls,
+    scheduler_mod,
+) -> None:
+    coordinator = FakeInputCoordinator()
+    scheduler = scheduler_cls.__new__(scheduler_cls)
+    scheduler.chunk_transfer_adapter = None
+    scheduler.input_coordinator = coordinator
+
+    def fake_finish_requests(self, request_ids, finished_status):
+        assert request_ids == ["req-a", "req-b"]
+        assert finished_status == RequestStatus.FINISHED_STOPPED
+        return [("req-a", 0), ("req-b", 1)]
+
+    monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
+
+    result = scheduler_cls.finish_requests(
+        scheduler,
+        ["req-a", "req-b"],
+        RequestStatus.FINISHED_STOPPED,
+    )
+
+    assert result == [("req-a", 0), ("req-b", 1)]
+    assert coordinator.freed == ["req-a", "req-b"]
+
+
+def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
+    coordinator = FakeInputCoordinator()
+    scheduler = OmniARScheduler.__new__(OmniARScheduler)
+    scheduler.input_coordinator = coordinator
+    scheduler._omits_kv_transfer_cache = {"req-free": True}
+    scheduler.encoder_cache_manager = SimpleNamespace(free=lambda request: None)
+    scheduler.finished_req_ids = set()
+    scheduler.finished_req_ids_dict = None
+    scheduler._new_prompt_len_snapshot = {"req-free": 3}
+    scheduler._connector_finished = lambda request: (False, None)
+    scheduler._should_transfer_kv_for_request = lambda request_id: False
+    scheduler._free_blocks = lambda request: None
+
+    result = OmniARScheduler._free_request(scheduler, DummyRequest())
+
+    assert result is None
+    assert coordinator.freed == ["req-free"]
+    assert scheduler._omits_kv_transfer_cache == {}
+    assert scheduler._new_prompt_len_snapshot == {}
+
+
+def test_generation_free_request_cleans_input_coordinator(monkeypatch: pytest.MonkeyPatch) -> None:
+    coordinator = FakeInputCoordinator()
+    scheduler = OmniGenerationScheduler.__new__(OmniGenerationScheduler)
+    scheduler.input_coordinator = coordinator
+
+    def fake_free_request(self, request, delay_free_blocks=False):
+        assert delay_free_blocks is True

tests/core/sched/test_omni_scheduling_coordinator.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for OmniSchedulingCoordinator (formerly ChunkSchedulingCoordinator).
+"""Unit tests for OmniSchedulingCoordinator.
 
 These tests use mock request objects and mock queues.  They do not require
 GPU, vLLM runtime, or any connector.
@@ -11,10 +11,12 @@ from __future__ import annotations
 import unittest
 from types import SimpleNamespace
 
+import torch
+
 import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod
 from vllm_omni.core.sched.omni_scheduling_coordinator import (
-    ChunkSchedulingCoordinator,
     OmniSchedulingCoordinator,
+    uses_qwen3_omni_full_payload_input_coordinator,
 )
 
 # ------------------------------------------------------------------ #
@@ -89,11 +91,55 @@ class MockQueue:
 # ------------------------------------------------------------------ #
 
 
+class TestFullPayloadCoordinatorSelection(unittest.TestCase):
+    def test_qwen3_omni_talker_and_code2wav_use_full_payload_input_coordinator(self):
+        for model_stage in ("talker", "code2wav"):
+            model_config = SimpleNamespace(
+                stage_id=1,
+                async_chunk=False,
+                model_arch="Qwen3OmniMoeForConditionalGeneration",
+                model_stage=model_stage,
+            )
+
+            self.assertTrue(uses_qwen3_omni_full_payload_input_coordinator(model_config))
+
+    def test_async_chunk_and_non_qwen3_omni_do_not_use_full_payload_input_coordinator(self):
+        cases = [
+            SimpleNamespace(
+                stage_id=1,
+                async_chunk=True,
+                model_arch="Qwen3OmniMoeForConditionalGeneration",
+                model_stage="talker",
+            ),
+            SimpleNamespace(
+                stage_id=1,
+                async_chunk=False,
+                model_arch="Qwen3TTSForConditionalGeneration",
+                model_stage="code2wav",
+            ),
+            SimpleNamespace(
+                stage_id=1,
+                async_chunk=False,
+                model_arch="Qwen2_5OmniForConditionalGeneration",
+                model_stage="talker",
+            ),
+            SimpleNamespace(
+                stage_id=0,
+                async_chunk=False,
+                model_arch="Qwen3OmniMoeForConditionalGeneration",
+                model_stage="thinker",
+            ),
+        ]
+
+        for model_config in cases:
+            self.assertFalse(uses_qwen3_omni_full_payload_input_coordinator(model_config))
+
+
 class TestChunkCoordinatorStateTransition(unittest.TestCase):
     """Test 5: process_pending_chunks transitions WAITING_FOR_CHUNK → target."""
 
     def test_ready_request_transitions_to_waiting(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
 
         req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
         waiting = MockQueue([req])
@@ -110,7 +156,7 @@ class TestChunkCoordinatorStateTransition(unittest.TestCase):
         self.assertIn("r1", coord.requests_with_ready_chunks)
 
     def test_non_ready_stays_waiting_for_chunk(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
 
         req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
         waiting = MockQueue([req])
@@ -126,7 +172,7 @@ class TestChunkCoordinatorStateTransition(unittest.TestCase):
         self.assertEqual(req.status, RequestStatus.WAITING_FOR_CHUNK)
 
     def test_stage_0_is_noop(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=0)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=0)
         req = _make_request("r1")
         waiting = MockQueue([req])
         running: list = []
@@ -144,7 +190,7 @@ class TestChunkCoordinatorRestoreQueues(unittest.TestCase):
     """Test 6: restore_queues returns waiting-for-chunk requests."""
 
     def test_restore(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
 
         r1 = _make_request("r1")
         r2 = _make_request("r2")
@@ -166,7 +212,7 @@ class TestChunkCoordinatorFinishedSignal(unittest.TestCase):
     """Test 8: chunk_finished_req_ids → finished_requests."""
 
     def test_finished_signal(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1, async_chunk=True)
 
         req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
         waiting = MockQueue([req])
@@ -187,7 +233,7 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
 
     def test_ar_mode_no_longer_sets_additional_information(self):
         """AR mode only processes scheduling metadata, not full payloads."""
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
 
         req = _make_request("r1")
         requests = {"r1": req}
@@ -204,10 +250,14 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
         self.assertIsNone(getattr(req, "additional_information", None))
 
     def test_generation_mode(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
 
         req = _make_request("r1")
         req.prompt_token_ids = [0, 0, 0]
+        req.num_prompt_tokens = 3
+        req.num_computed_tokens = 3
+        req._all_token_ids = [0, 0, 0, 99]
+        req._output_token_ids = [99]
         requests = {"r1": req}
 
         request_metadata = {
@@ -220,16 +270,61 @@ class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
         coord.update_request_metadata(requests, request_metadata, model_mode="generation")
 
         self.assertEqual(req.prompt_token_ids, [10, 20, 30])
+        self.assertEqual(req.num_prompt_tokens, 3)
         self.assertEqual(req.num_computed_tokens, 0)
+        self.assertEqual(req._all_token_ids, [10, 20, 30])
+        self.assertEqual(req._output_token_ids, [])
         self.assertIsNone(req.additional_information)
         self.assertEqual(req._omni_initial_model_buffer, {"meta": {"left_context_size": 25}})
 
+    def test_generation_mode_flattens_tensor_code_predictor_codes(self):
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+
+        req = _make_request("r1")
+        req.prompt_token_ids = [9]
+        req.num_prompt_tokens = 1
+        req._all_token_ids = [9, 8]
+        req._output_token_ids = [8]
+        requests = {"r1": req}
+
+        coord.update_request_metadata(
+            requests,
+            {"r1": {"code_predictor_codes": torch.tensor([[1, 2, 3]], dtype=torch.long)}},
+            model_mode="generation",
+        )
+
+        self.assertEqual(req.prompt_token_ids, [1, 2, 3])
+        self.assertEqual(req.num_prompt_tokens, 3)
+        self.assertEqual(req._all_token_ids, [1, 2, 3])
+        self.assertEqual(req._output_token_ids, [])
+
+    def test_generation_mode_flattens_nested_code_predictor_codes(self):
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+
+        req = _make_request("r1")
+        req.prompt_token_ids = [9]
+        req.num_prompt_tokens = 1
+        req._all_token_ids = [9, 8]
+        req._output_token_ids = [8]
+        requests = {"r1": req}
+
+        coord.update_request_metadata(
+            requests,
+            {"r1": {"code_predictor_codes": [[1, 2], [3, 4]]}},
+            model_mode="generation",
+        )
+
+        self.assertEqual(req.prompt_token_ids, [1, 2, 3, 4])
+        self.assertEqual(req.num_prompt_tokens, 4)
+        self.assertEqual(req._all_token_ids, [1, 2, 3, 4])
+        self.assertEqual(req._output_token_ids, [])
+
 
 class TestChunkCoordinatorPostprocess(unittest.TestCase):
     """Test postprocess_scheduler_output clears ready chunks."""
 
     def test_clear_ready(self):
-        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
         coord.requests_with_ready_chunks = {"r1", "r2"}
 
         new_req = SimpleNamespace(req_id="r1")
@@ -366,6 +461,33 @@ class TestWaitingForInputTransition(unittest.TestCase):
         self.assertEqual(len(coord.pending_input_registrations), 1)
         self.assertEqual(coord.pending_input_registrations[0].request_id, "r1")
 
+    def test_idle_cycles_retain_received_marker_before_request_appears(self):
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+            async_chunk=False,
+        )
+        coord._full_payload_input_received.add("late")
+        coord.finished_requests.add("late")
+
+        waiting = MockQueue()
+        running: list = []
+
+        coord.process_pending_full_payload_inputs(waiting, running, stage_recv_req_ids=set())
+
+        self.assertIn("late", coord._full_payload_input_received)
+        self.assertIn("late", coord.finished_requests)
+
+        late_req = _make_request("late", status=RequestStatus.WAITING)
+        waiting.add_request(late_req)
+
+        coord.process_pending_full_payload_inputs(waiting, running, stage_recv_req_ids=set())
+
+        self.assertEqual(late_req.status, RequestStatus.WAITING)
+        self.assertEqual(coord.pending_input_registrations, [])
+        self.assertIn("late", coord._full_payload_input_received)
+        self.assertIn("late", coord.finished_requests)
+
 
 class TestTimeoutDetection(unittest.TestCase):
     """Regression tests for orphaned pending-recv timeout detection.
@@ -508,14 +630,18 @@ class TestTimeoutDetection(unittest.TestCase):
         self.assertEqual(len(coord._waiting_for_input), 0)
 
     def test_free_finished_request_clears_waiting_since(self):
-        """free_finished_request clears _waiting_since."""
+        """free_finished_request clears coordinator lifecycle markers."""
         coord = OmniSchedulingCoordinator(
             scheduler_max_num_seqs=10,
             stage_id=1,
         )
         coord._waiting_since["r1"] = 0.0
+        coord._full_payload_input_received.add("r1")
+        coord.finished_requests.add("r1")
         coord.free_finished_request("r1")
         self.assertNotIn("r1", coord._waiting_since)
+        self.assertNotIn("r1", coord._full_payload_input_received)
+        self.assertNotIn("r1", coord.finished_requests)
 
     def test_timeout_from_running_queue_full_lifecycle(self):
         """End-to-end: request from running → WAITING_FOR_CHUNK → restore →

tests/model_executor/stage_input_processors/test_qwen3_omni_streaming_helpers.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 from types import SimpleNamespace
 
 import pytest
+import torch
 
 import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3
 
@@ -79,3 +80,110 @@ def test_get_streaming_codec_delta_len_increments_and_finishes(_streaming_contex
     assert d3 == 1
     state = q3._get_qwen3_streaming_state("c1", _streaming_context)
     assert state.talker2code2wav_last_seq_len == 0
+
+
+def test_talker2code2wav_full_payload_filters_by_output_token_ids() -> None:
+    request = SimpleNamespace(
+        request_id="codec",
+        output_token_ids=[4197, 1, 2, 4198, -1, 2048],
+    )
+    rows = torch.tensor(
+        [
+            [100, 101, 102],
+            [10, 11, 12],
+            [20, 21, 22],
+            [30, 31, 32],
+            [40, 41, 42],
+            [50, 51, 52],
+        ],
+        dtype=torch.long,
+    )
+
+    payload = q3.talker2code2wav_full_payload(None, {"codes.audio": rows}, request)
+
+    assert payload is not None
+    assert payload["codes"]["audio"] == [10, 20, 11, 21, 12, 22]
+    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
+
+
+def test_talker2code2wav_full_payload_drops_count_matched_terminal_row() -> None:
+    request = SimpleNamespace(
+        request_id="codec_terminal_row",
+        output_token_ids=[0, 4198],
+    )
+    rows = torch.tensor(
+        [
+            [10, 11, 12],
+        ],
+        dtype=torch.long,
+    )
+
+    payload = q3.talker2code2wav_full_payload(None, {"codes.audio": rows}, request)
+
+    assert payload is None
+
+
+def test_talker2code2wav_full_payload_drops_rows_aligned_to_non_codec_ids() -> None:
+    request = SimpleNamespace(
+        request_id="codec_invalid_ids",
+        output_token_ids=[4197, 0, 4198, 4196, -1, 2048],
+    )
+    rows = torch.tensor(
+        [
+            [91, 92, 93],
+            [0, 0, 0],
+            [81, 82, 83],
+            [71, 72, 73],
+            [61, 62, 63],
+            [51, 52, 53],
+        ],
+        dtype=torch.long,
+    )
+
+    payload = q3.talker2code2wav_full_payload(None, {"codes.audio": rows}, request)
+
+    assert payload is not None
+    assert payload["codes"]["audio"] == [0, 0, 0]
+    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
+
+
+def test_talker2code2wav_full_payload_keeps_all_zero_codec_rows() -> None:
+    request = SimpleNamespace(
+        request_id="codec_zero",
+        output_token_ids=[0, 1],
+    )
+    rows = torch.tensor(
+        [
+            [0, 0, 0],
+            [7, 8, 9],
+        ],
+        dtype=torch.long,
+    )
+
+    payload = q3.talker2code2wav_full_payload(None, {"codes.audio": rows}, request)
+
+    assert payload is not None
+    assert payload["codes"]["audio"] == [0, 7, 0, 8, 0, 9]
+    assert payload["code_predictor_codes"] == payload["codes"]["audio"]
+
+
+def test_thinker2talker_full_payload_packs_complete_tensors() -> None:
+    request = SimpleNamespace(
+        request_id="thinker",
+        prompt_token_ids=[151644, 872],
+        output_token_ids=[3],
+        all_token_ids=[151644, 872, 3],
+    )
+    pooling_output = {
+        "hidden_states.layer_0": torch.ones(3, 2),
+        "hidden_states.layer_24": torch.full((3, 2), 2.0),
+        "embed.tts_bos": torch.zeros(1, 2),
+    }
+
+    payload = q3.thinker2talker_full_payload(None, pooling_output, request)
+
+    assert payload is not None
+    assert payload["ids"]["all"] == [151644, 872, 3]
+    assert payload["embed"]["prefill"].device.type == "cpu"
+    assert payload["hidden_states"]["output"].device.type == "cpu"
+    assert payload["next_stage_prompt_len"] > 0

tests/test_config_factory.py

@@ -1504,7 +1504,7 @@ class TestSentinelDefaultPrecedence:
     def test_none_value_skipped_yaml_wins(self):
         stages = self._stages({"max_num_seqs": None})
         assert stages[2].runtime_overrides.get("max_num_seqs") is None
-        assert "max_num_seqs" not in stages[2].yaml_engine_args
+        assert stages[2].yaml_engine_args["max_num_seqs"] == 64
 
     def test_empty_kwargs_yaml_only(self):
         stages = self._stages({})
@@ -1598,6 +1598,41 @@ class TestSentinelDefaultPrecedence:
         assert sync_stages[1].custom_process_input_func is not None
         assert sync_stages[1].custom_process_input_func.endswith("talker2code2wav")
 
+    def test_async_chunk_dispatches_qwen3_omni_processors(self):
+        import runpy
+        from pathlib import Path
+
+        from vllm_omni.config.stage_config import DeployConfig, merge_pipeline_deploy
+
+        pipeline_path = (
+            Path(__file__).parent.parent / "vllm_omni" / "model_executor" / "models" / "qwen3_omni" / "pipeline.py"
+        )
+        pipeline = runpy.run_path(str(pipeline_path))["QWEN3_OMNI_PIPELINE"]
+
+        async_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=True))
+        assert (
+            async_stages[0]
+            .yaml_engine_args["custom_process_next_stage_input_func"]
+            .endswith("thinker2talker_async_chunk")
+        )
+        assert (
+            async_stages[1]
+            .yaml_engine_args["custom_process_next_stage_input_func"]
+            .endswith("talker2code2wav_async_chunk")
+        )
+
+        sync_stages = merge_pipeline_deploy(pipeline, DeployConfig(async_chunk=False))
+        assert (
+            sync_stages[0]
+            .yaml_engine_args["custom_process_next_stage_input_func"]
+            .endswith("thinker2talker_full_payload")
+        )
+        assert (
+            sync_stages[1]
+            .yaml_engine_args["custom_process_next_stage_input_func"]
+            .endswith("talker2code2wav_full_payload")
+        )
+
 
 class TestSamplingConstraintsPrecedence:
     """Test that pipeline sampling_constraints override deploy defaults."""

tests/worker/test_omni_gpu_model_runner.py

@@ -5,6 +5,7 @@ import pytest
 import torch
 
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -330,6 +331,117 @@ def test_update_intermediate_buffer_skips_unknown_req_id():
     assert "unknown_req" not in runner.model_intermediate_buffer
 
 
+def _make_full_payload_accumulation_runner(
+    model_arch="Qwen3OmniMoeForConditionalGeneration",
+    model_stage="talker",
+    async_chunk=False,
+):
+    runner = object.__new__(OmniConnectorModelRunnerMixin)
+    runner.model_config = SimpleNamespace(
+        model_arch=model_arch,
+        model_stage=model_stage,
+        async_chunk=async_chunk,
+    )
+    runner._custom_process_func = object()
+    runner._pending_full_payload_send = {}
+    runner._stage_id = 1
+    return runner
+
+
+def test_accumulate_full_payload_output_preserves_aligned_all_zero_qwen3_omni_codec_rows():
+    runner = _make_full_payload_accumulation_runner()
+    request = SimpleNamespace(output_token_ids=[0, 1])
+    codes = torch.zeros((2, 3), dtype=torch.long)
+
+    OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
+
+    stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
+    assert torch.equal(stored["codes.audio"], codes)
+
+
+def test_accumulate_full_payload_output_keeps_misaligned_all_zero_qwen3_omni_codec_rows():
+    # After removing the sender-side zero filter, the accumulator keeps every
+    # codec row including misaligned all-zero rows. The downstream consumer
+    # (_extract_qwen3_full_payload_codec_rows) is the authoritative crop and
+    # filters by output_token_ids.
+    runner = _make_full_payload_accumulation_runner()
+    request = SimpleNamespace(output_token_ids=[0, 1])
+    codes = torch.zeros((1, 3), dtype=torch.long)
+
+    OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
+
+    stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
+    assert "codes.audio" in stored
+    assert torch.equal(stored["codes.audio"], codes)
+
+
+def test_accumulate_full_payload_output_preserves_incremental_aligned_all_zero_qwen3_omni_codec_rows():
+    runner = _make_full_payload_accumulation_runner()
+    request = SimpleNamespace(output_token_ids=[0, 1])
+    runner._pending_full_payload_send["r1"] = (
+        {"codes.audio": torch.ones((1, 3), dtype=torch.long)},
+        request,
+    )
+    codes = torch.zeros((1, 3), dtype=torch.long)
+
+    OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
+
+    stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
+    assert stored["codes.audio"].shape == (2, 3)
+    assert torch.equal(stored["codes.audio"][1], torch.zeros(3, dtype=torch.long))
+
+
+def test_accumulate_full_payload_output_keeps_all_zero_qwen3_omni_prefill_placeholder():
+    # Prefill placeholder rows (output_token_ids empty) are no longer dropped
+    # at the sender. The consumer-side crop trims them off using
+    # output_token_ids, so the end-to-end semantics are unchanged.
+    runner = _make_full_payload_accumulation_runner()
+    request = SimpleNamespace(output_token_ids=[])
+    codes = torch.zeros((2, 3), dtype=torch.long)
+
+    OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
+
+    stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
+    assert "codes.audio" in stored
+    assert torch.equal(stored["codes.audio"], codes)
+
+
+def test_full_payload_output_accumulation_hook_matrix():
+    assert _make_full_payload_accumulation_runner(model_stage="thinker")._should_accumulate_full_payload_output()
+    assert _make_full_payload_accumulation_runner(model_stage="talker")._should_accumulate_full_payload_output()
+    assert not _make_full_payload_accumulation_runner(model_stage="code2wav")._should_accumulate_full_payload_output()
+    assert not _make_full_payload_accumulation_runner(
+        model_stage="talker", async_chunk=True
+    )._should_accumulate_full_payload_output()
+    assert not _make_full_payload_accumulation_runner(
+        model_arch="Qwen3TTSForConditionalGeneration"
+    )._should_accumulate_full_payload_output()
+    assert not _make_full_payload_accumulation_runner(
+        model_arch="Qwen2_5OmniForConditionalGeneration"
+    )._should_accumulate_full_payload_output()
+
+
+def test_sync_local_stage_payloads_retains_payload_until_request_is_active():
+    runner = object.__new__(OmniGPUModelRunner)
+    payload = {"codes": {"audio": [1, 2, 3]}}
+    runner._local_stage_payload_cache = {"late": payload}
+    runner._full_payload_pending_broadcast_req_ids = set()
+    runner.requests = {}
+    runner.model_intermediate_buffer = {}
+
+    OmniGPUModelRunner._sync_local_stage_payloads(runner)
+
+    assert runner._local_stage_payload_cache == {"late": payload}
+    assert runner.model_intermediate_buffer == {}
+
+    runner.requests = {"late": DummyReqState()}
+    OmniGPUModelRunner._sync_local_stage_payloads(runner)
+
+    assert runner._local_stage_payload_cache == {}
+    assert runner.model_intermediate_buffer["late"] == payload
+    assert runner.requests["late"].additional_information_cpu == payload
+
+
 def test_maybe_attach_mimo_audio_req_infos_enriches_dict():
     runner = _make_runner_for_mimo()
     req_id = "r_mimo"

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -20,11 +20,16 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+from vllm_omni.core.sched.omni_scheduling_coordinator import (
+    OmniSchedulingCoordinator,
+    uses_qwen3_omni_full_payload_input_coordinator,
+)
 from vllm_omni.core.sched.output import OmniSchedulerOutput
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
 from vllm_omni.engine.serialization import deserialize_additional_information
+from vllm_omni.outputs import OmniConnectorOutput
 
 logger = init_logger(__name__)
 
@@ -75,6 +80,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         self.chunk_transfer_adapter = None
         if getattr(model_config, "async_chunk", False):
             self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config)
+        self.input_coordinator: OmniSchedulingCoordinator | None = None
+        if uses_qwen3_omni_full_payload_input_coordinator(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),
+                async_chunk=False,
+            )
+        self._latest_omni_connector_output: OmniConnectorOutput | None = None
         # Snapshot prompt length for each streaming input update
         self._new_prompt_len_snapshot: dict[str, int] = {}
 
@@ -190,6 +203,19 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         return False
 
     def schedule(self) -> SchedulerOutput:  # type: ignore[override]
+        connector_output = self._latest_omni_connector_output
+        self._latest_omni_connector_output = None
+        if self.input_coordinator:
+            if connector_output and connector_output.request_metadata:
+                self.input_coordinator.update_request_metadata(
+                    self.requests, connector_output.request_metadata, model_mode="ar"
+                )
+            self.input_coordinator.process_pending_full_payload_inputs(
+                self.waiting,
+                self.running,
+                connector_output.stage_recv_req_ids if connector_output else set(),
+            )
+
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running)
 
@@ -199,6 +225,8 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if self.chunk_transfer_adapter:
                 # Add request waiting for chunk to the waiting and running queue
                 self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
+            if self.input_coordinator:
+                self.input_coordinator.restore_queues(self.waiting, self.running)
         try:
             # Late import to avoid circulars in some launch modes
             from .output import OmniNewRequestData
@@ -239,9 +267,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Wrap in omni scheduler output to carry transfer metadata.
         base_fields = SchedulerOutput.__dataclass_fields__.keys()
         base_data = {name: getattr(scheduler_output, name) for name in base_fields}
+        input_regs = self.input_coordinator.pending_input_registrations if self.input_coordinator else []
         return OmniSchedulerOutput(
             **base_data,
             finished_requests_needing_kv_transfer=finished_reqs,
+            pending_input_registrations=input_regs,
         )
 
     def update_from_output(
@@ -534,6 +564,16 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 engine_core_outputs[0] = eco = EngineCoreOutputs()
             eco.scheduler_stats = stats
 
+        omni_output = getattr(model_runner_output, "omni_connector_output", None)
+        if omni_output is not None:
+            self._latest_omni_connector_output = omni_output
+            if self.input_coordinator and omni_output.request_metadata:
+                self.input_coordinator.update_request_metadata(
+                    self.requests,
+                    omni_output.request_metadata,
+                    model_mode="ar",
+                )
+
         # Free blocks that were held for transfer (kv_ready and
         # active_kv_transfers updates already done before the per-request loop).
         if kv_extracted_ids:
@@ -572,7 +612,12 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
 
-        return super().finish_requests(request_ids, finished_status)
+        finished = super().finish_requests(request_ids, finished_status)
+        input_coordinator = getattr(self, "input_coordinator", None)
+        if input_coordinator is not None:
+            for request_id, _ in finished:
+                self._free_input_coordinator_request(request_id)
+        return finished
 
     def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None:
         """
@@ -616,10 +661,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     # It triggered but hasn't finished yet. We MUST wait.
                     logger.debug(f"[Omni] Request {request_id} finished but transfer is still ACTIVE. Waiting.")
                     self.waiting_for_transfer_free.add(request_id)
+                    if self.input_coordinator is not None:
+                        self._free_input_coordinator_request(request_id)
                     kv_xfer_params = None
                     return kv_xfer_params
                 elif request_id in self.waiting_for_transfer_free:
                     # Blocks held until KV extraction completes in a future step.
+                    if self.input_coordinator is not None:
+                        self._free_input_coordinator_request(request_id)
                     return None
                 else:
                     logger.debug(
@@ -654,10 +703,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     if isinstance(add_info, dict):
                         add_info.update(kv_xfer_params)
 
+                if self.input_coordinator is not None:
+                    self._free_input_coordinator_request(request_id)
                 return kv_xfer_params
 
         # 3. Standard Freeing
         delay_free_blocks |= connector_delay_free_blocks
+        if self.input_coordinator is not None:
+            self._free_input_coordinator_request(request_id)
         if not delay_free_blocks:
             self._free_blocks(request)
 

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import time
 from collections import defaultdict
+from typing import Any
 
 from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.distributed.kv_events import KVEventBatch
@@ -23,11 +24,15 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
 from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
-from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
+from vllm_omni.core.sched.omni_scheduling_coordinator import (
+    OmniSchedulingCoordinator,
+    uses_qwen3_omni_full_payload_input_coordinator,
+)
+from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData, OmniSchedulerOutput
 from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
     OmniChunkTransferAdapter,
 )
-from vllm_omni.outputs import OmniModelRunnerOutput
+from vllm_omni.outputs import OmniConnectorOutput, OmniModelRunnerOutput
 
 logger = init_logger(__name__)
 
@@ -39,6 +44,14 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self.chunk_transfer_adapter = None
         if getattr(model_config, "async_chunk", False):
             self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config)
+        self.input_coordinator: OmniSchedulingCoordinator | None = None
+        if uses_qwen3_omni_full_payload_input_coordinator(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),
+                async_chunk=False,
+            )
+        self._latest_omni_connector_output: OmniConnectorOutput | None = None
 
     def schedule(self) -> SchedulerOutput:
         """Diffusion fast path:
@@ -68,6 +81,18 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Temporary queue: preserve waiting order, do not disturb non-diffusion requests
         skipped_waiting_requests = create_request_queue(self.policy)
         req_index = 0
+        connector_output = self._latest_omni_connector_output
+        self._latest_omni_connector_output = None
+        if self.input_coordinator:
+            if connector_output and connector_output.request_metadata:
+                self.input_coordinator.update_request_metadata(
+                    self.requests, connector_output.request_metadata, model_mode="generation"
+                )
+            self.input_coordinator.process_pending_full_payload_inputs(
+                self.waiting,
+                self.running,
+                connector_output.stage_recv_req_ids if connector_output else set(),
+            )
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running)
 
@@ -201,7 +226,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
             else:
                 res = super().schedule()
-                return res
+                if self.input_coordinator:
+                    self.input_coordinator.restore_queues(self.waiting, self.running)
+                base_fields = SchedulerOutput.__dataclass_fields__.keys()
+                base_data = {name: getattr(res, name) for name in base_fields}
+                return OmniSchedulerOutput(
+                    **base_data,
+                    pending_input_registrations=(
+                        self.input_coordinator.pending_input_registrations if self.input_coordinator else []
+                    ),
+                )
 
         # Compute common prefix blocks (aligned with v1)
         num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
@@ -321,8 +355,17 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             # internal deques and never scheduled again.
             if self.chunk_transfer_adapter:
                 self.chunk_transfer_adapter.restore_queues(self.waiting, self.running)
+            if self.input_coordinator:
+                self.input_coordinator.restore_queues(self.waiting, self.running)
 
-        return scheduler_output
+        base_fields = SchedulerOutput.__dataclass_fields__.keys()
+        base_data = {name: getattr(scheduler_output, name) for name in base_fields}
+        return OmniSchedulerOutput(
+            **base_data,
+            pending_input_registrations=(
+                self.input_coordinator.pending_input_registrations if self.input_coordinator else []
+            ),
+        )
 
     def finish_requests(self, request_ids, finished_status: RequestStatus) -> list[tuple[str, int]]:
         """Handles the finish signal from outside the scheduler.
@@ -340,7 +383,20 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         if self.chunk_transfer_adapter:
             self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
 
-        return super().finish_requests(request_ids, finished_status)
+        finished = super().finish_requests(request_ids, finished_status)
+        if self.input_coordinator is not None:
+            for request_id, _ in finished:
+                self._free_input_coordinator_request(request_id)
+        return finished
+
+    def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None:
+        if self.input_coordinator is None:
+            return super()._free_request(request, delay_free_blocks)
+
+        try:
+            return super()._free_request(request, delay_free_blocks)
+        finally:
+            self._free_input_coordinator_request(request.request_id)
 
     """
     Scheduler for the diffusion model.
@@ -589,6 +645,16 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                 engine_core_outputs[0] = eco = EngineCoreOutputs()
             eco.scheduler_stats = stats
 
+        omni_output = getattr(model_runner_output, "omni_connector_output", None)
+        if omni_output is not None:
+            self._latest_omni_connector_output = omni_output
+            if self.input_coordinator and omni_output.request_metadata:
+                self.input_coordinator.update_request_metadata(
+                    self.requests,
+                    omni_output.request_metadata,
+                    model_mode="generation",
+                )
+
         return engine_core_outputs
 
     def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None:

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -7,6 +7,12 @@ from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 class OmniSchedulerMixin:
     """Shared scheduler helpers for omni-specific request handling."""
 
+    def _free_input_coordinator_request(self, request_id: str) -> None:
+        """Prune full-payload coordinator state for a completed request."""
+        input_coordinator = getattr(self, "input_coordinator", None)
+        if input_coordinator is not None:
+            input_coordinator.free_finished_request(request_id)
+
     def _replace_session_with_streaming_update(
         self,
         session: Request,

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -22,6 +22,15 @@ from vllm.v1.request import Request, RequestStatus
 logger = init_logger(__name__)
 
 
+def uses_qwen3_omni_full_payload_input_coordinator(model_config: Any) -> bool:
+    return (
+        getattr(model_config, "stage_id", 0) > 0
+        and not getattr(model_config, "async_chunk", False)
+        and getattr(model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration"
+        and getattr(model_config, "model_stage", None) in {"talker", "code2wav"}
+    )
+
+
 class OmniSchedulingCoordinator:
     """Pure-scheduling coordinator for chunk and full_payload input waiting.
 
@@ -166,8 +175,10 @@ class OmniSchedulingCoordinator:
                         to_remove.append(request)
                         self._waiting_for_input.append(request)
                         self.pending_input_registrations.append(request)
-            for request in to_remove:
-                waiting_queue.remove(request)
+            if to_remove:
+                # Use the bulk-remove helper: one O(N) sweep instead of N
+                # repeated O(N) removes from a list-backed queue.
+                waiting_queue.remove_requests(to_remove)
 
     def process_pending_full_payload_inputs_legacy(
         self,
@@ -255,6 +266,30 @@ class OmniSchedulingCoordinator:
             waiting_queue.add_request(request)
         self._waiting_for_input = deque()
 
+    @staticmethod
+    def _flatten_prompt_token_ids(value: Any) -> list[int]:
+        """Normalize connector metadata into flat prompt token ids."""
+        if value is None:
+            return []
+        if hasattr(value, "detach") and hasattr(value, "cpu") and hasattr(value, "tolist"):
+            value = value.detach().cpu().tolist()
+        elif hasattr(value, "tolist") and not isinstance(value, (list, tuple)):
+            value = value.tolist()
+
+        if isinstance(value, (list, tuple)):
+            flattened: list[int] = []
+            for item in value:
+                if hasattr(item, "detach") and hasattr(item, "cpu") and hasattr(item, "tolist"):
+                    item = item.detach().cpu().tolist()
+                elif hasattr(item, "tolist") and not isinstance(item, (list, tuple)):
+                    item = item.tolist()
+                if isinstance(item, (list, tuple)):
+                    flattened.extend(int(token_id) for token_id in item)
+                else:
+                    flattened.append(int(item))
+            return flattened
+        return [int(value)]
+
     def update_request_metadata(
         self,
         requests: dict[str, Request],
@@ -310,7 +345,7 @@ class OmniSchedulingCoordinator:
                             )
 
             if model_mode != "ar":
-                new_ids = metadata.get("code_predictor_codes", [])
+                new_ids = self._flatten_prompt_token_ids(metadata.get("code_predictor_codes"))
                 runtime_seed = None
                 if "left_context_size" in metadata:
                     runtime_seed = {
@@ -319,6 +354,10 @@ class OmniSchedulingCoordinator:
                 request._omni_initial_model_buffer = runtime_seed
                 if new_ids:
                     request.prompt_token_ids = new_ids
+                    request.num_prompt_tokens = len(new_ids)
+                    request._all_token_ids.clear()
+                    request._all_token_ids.extend(new_ids)
+                    request._output_token_ids.clear()
                     request.num_computed_tokens = 0
 
     def postprocess_scheduler_output(
@@ -374,7 +413,3 @@ class OmniSchedulingCoordinator:
         if scheduler_output.scheduled_cached_reqs:
             for req_id in scheduler_output.scheduled_cached_reqs.req_ids:
                 self.requests_with_ready_chunks.discard(req_id)
-
-
-# Backward-compatible alias
-ChunkSchedulingCoordinator = OmniSchedulingCoordinator

vllm_omni/core/sched/output.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass, field
+from typing import Any
 
 from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
 from vllm.v1.request import Request
@@ -75,3 +76,4 @@ class OmniSchedulerOutput(SchedulerOutput):
     """Scheduler output with omni-specific transfer metadata."""
 
     finished_requests_needing_kv_transfer: dict[str, dict] = field(default_factory=dict)
+    pending_input_registrations: list[Any] = field(default_factory=list)

vllm_omni/distributed/omni_connectors/utils/initialization.py

@@ -5,6 +5,7 @@
 
 import json
 import sys
+import warnings
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
@@ -403,7 +404,16 @@ def build_stage_connectors(
     connectors_config: dict[str, Any],
     purpose: str = "request_forwarding",
 ) -> dict[tuple[str, str], Any] | None:
-    """Instantiate OmniConnectors for a stage based on config."""
+    """Instantiate OmniConnectors for a stage based on config.
+
+    Deprecated: prefer ``get_stage_connector_config`` plus the unified
+    connector factory. Kept as a thin shim so legacy callers keep working.
+    """
+    warnings.warn(
+        "build_stage_connectors is deprecated; use get_stage_connector_config and OmniConnectorFactory instead.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
     if not connectors_config:
         return {}
 

vllm_omni/entrypoints/openai/realtime_connection.py

@@ -120,6 +120,7 @@ class RealtimeConnection(VllmRealtimeConnection):
         full_text = ""
         prompt_token_ids_len = 0
         completion_tokens_len = 0
+        last_prompt_token_ids_len = 0  # detect Stage-0 segment rollover
         self._realtime_audio_ref = None
 
         # Coerce cumulative outputs to delta outputs; this ensures
@@ -144,6 +145,18 @@ class RealtimeConnection(VllmRealtimeConnection):
                     new_token_ids = list(first_output.token_ids)
                     new_tokens_len = len(new_token_ids)
 
+                    cur_prompt_token_ids_len = len(output.prompt_token_ids or [])
+                    # Stage-0 segment rollover: buffer_realtime_audio may yield
+                    # multiple TokensPrompt segments and the second segment's
+                    # prompt_token_ids include the first segment's decoded output.
+                    # Clear accumulated full_text at the boundary so the
+                    # carried-over prefix is not duplicated in the final
+                    # TranscriptionDone.text.
+                    stage_id = getattr(output, "stage_id", None)
+                    if stage_id == 0 and cur_prompt_token_ids_len > last_prompt_token_ids_len > 0:
+                        full_text = ""
+                    last_prompt_token_ids_len = cur_prompt_token_ids_len
+
                     if not prompt_token_ids_len and output.prompt_token_ids:
                         prompt_token_ids_len = len(output.prompt_token_ids)
 

vllm_omni/model_executor/models/qwen3_omni/pipeline.py

@@ -30,7 +30,8 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             requires_multimodal_data=True,
             hf_config_name="thinker_config",
             engine_output_type="latent",
-            custom_process_next_stage_input_func=(f"{_PROC}.thinker2talker_async_chunk"),
+            custom_process_next_stage_input_func=(f"{_PROC}.thinker2talker_full_payload"),
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.thinker2talker_async_chunk"),
             sampling_constraints={"detokenize": True},
         ),
         StagePipelineConfig(
@@ -41,7 +42,9 @@ QWEN3_OMNI_PIPELINE = PipelineConfig(
             hf_config_name="talker_config",
             engine_output_type="latent",
             custom_process_input_func=f"{_PROC}.thinker2talker",
-            custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
+            sync_process_input_func=f"{_PROC}.thinker2talker_token_only",
+            custom_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_full_payload"),
+            async_chunk_process_next_stage_input_func=(f"{_PROC}.talker2code2wav_async_chunk"),
             sampling_constraints={
                 "detokenize": False,
                 "stop_token_ids": [2150],

vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py

@@ -237,14 +237,22 @@ class Qwen3OmniMoeForConditionalGeneration(
 
         prompt_token_ids = tokenizer.encode(prompt_template)
 
+        # In non-async-chunk (full-payload) mode the engine treats each
+        # streaming TokensPrompt as a fresh decode, so mid-stream segment
+        # yields cause the thinker to emit duplicate responses and the
+        # talker to emit duplicate audio. Defer all audio to the final
+        # flush so the thinker sees one complete prompt.
+        async_chunk = getattr(model_config, "async_chunk", False)
+
         async for audio_chunk in audio_stream:
             buffer.write_audio(audio_chunk)
 
-            while (segment := buffer.read_audio()) is not None:
-                yield TokensPrompt(
-                    prompt_token_ids=prompt_token_ids,
-                    multi_modal_data={"audio": segment},
-                )
+            if async_chunk:
+                while (segment := buffer.read_audio()) is not None:
+                    yield TokensPrompt(
+                        prompt_token_ids=prompt_token_ids,
+                        multi_modal_data={"audio": segment},
+                    )
 
         remaining = buffer.flush()
         if remaining is not None and len(remaining) > 0:

vllm_omni/model_executor/stage_input_processors/qwen3_omni.py

@@ -35,6 +35,10 @@ logger = logging.getLogger(__name__)
 # Pooling output layer keys: "0" = word embedding, "24" = accept_hidden_layer
 _EMBED_LAYER_KEY = "0"
 _HIDDEN_LAYER_KEY = "24"
+_QWEN3_CODEC_CODEBOOK_SIZE = 2048
+_QWEN3_CODEC_PAD_TOKEN_ID = 4196
+_QWEN3_CODEC_BOS_TOKEN_ID = 4197
+_QWEN3_CODEC_EOS_TOKEN_ID = 4198
 
 
 def _layer_tensor(layers: dict[Any, Any], key: str) -> torch.Tensor | None:
@@ -99,6 +103,84 @@ def _ensure_list(x):
     return list(x)
 
 
+def _as_tensor_or_none(value: Any) -> torch.Tensor | None:
+    if isinstance(value, torch.Tensor):
+        return value.detach().cpu()
+    if isinstance(value, list) and value and isinstance(value[0], torch.Tensor):
+        return value[0].detach().cpu()
+    return None
+
+
+def _is_valid_qwen3_codec_token_id(token_id: Any) -> bool:
+    try:
+        token_id = int(token_id)
+    except (TypeError, ValueError):
+        return False
+    return 0 <= token_id < _QWEN3_CODEC_CODEBOOK_SIZE
+
+
+def should_accumulate_qwen3_omni_full_payload_output(
+    model_config: Any,
+    custom_process_func: Any,
+) -> bool:
+    """Return whether Qwen3-Omni should accumulate full-payload outputs."""
+    return (
+        custom_process_func is not None
+        and not getattr(model_config, "async_chunk", False)
+        and getattr(model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration"
+        and getattr(model_config, "model_stage", None) in {"thinker", "talker"}
+    )
+
+
+def _extract_qwen3_full_payload_codec_rows(
+    code_predictor_codes: torch.Tensor,
+    output_token_ids: list[int],
+) -> tuple[torch.Tensor, dict[str, int]]:
+    """Filter full-payload codec rows by the authoritative output ids."""
+    if code_predictor_codes.ndim != 2 or code_predictor_codes.numel() == 0:
+        return code_predictor_codes, {
+            "raw_rows": int(code_predictor_codes.shape[0]) if code_predictor_codes.ndim > 0 else 0,
+            "aligned_rows": 0,
+            "valid_rows": 0,
+            "trailing_placeholder_count": 0,
+        }
+
+    trailing_placeholder_count = 0
+    while (
+        trailing_placeholder_count < len(output_token_ids) and output_token_ids[-1 - trailing_placeholder_count] == -1
+    ):
+        trailing_placeholder_count += 1
+
+    aligned_len = min(int(code_predictor_codes.shape[0]), len(output_token_ids))
+    if aligned_len <= 0:
+        return code_predictor_codes[:0], {
+            "raw_rows": int(code_predictor_codes.shape[0]),
+            "aligned_rows": 0,
+            "valid_rows": 0,
+            "trailing_placeholder_count": trailing_placeholder_count,
+        }
+
+    aligned_rows = code_predictor_codes[-aligned_len:]
+    aligned_token_ids = output_token_ids[-aligned_len:]
+    aligned_token_mask = torch.tensor(
+        [_is_valid_qwen3_codec_token_id(token_id) for token_id in aligned_token_ids],
+        dtype=torch.bool,
+        device=aligned_rows.device,
+    )
+    row_valid_mask = (aligned_rows.max(dim=1).values < _QWEN3_CODEC_CODEBOOK_SIZE) & (
+        aligned_rows.min(dim=1).values >= 0
+    )
+    filtered_rows = aligned_rows[aligned_token_mask & row_valid_mask]
+    if filtered_rows.numel() == 0:
+        filtered_rows = aligned_rows[:0]
+    return filtered_rows, {
+        "raw_rows": int(code_predictor_codes.shape[0]),
+        "aligned_rows": aligned_len,
+        "valid_rows": int(filtered_rows.shape[0]) if filtered_rows.ndim > 0 else 0,
+        "trailing_placeholder_count": trailing_placeholder_count,
+    }
+
+
 # =========================
 # PD disaggregation helpers
 # =========================
@@ -373,6 +455,79 @@ def thinker2talker_async_chunk(
     return payload
 
 
+def thinker2talker_full_payload(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any],
+    request: OmniEngineCoreRequest,
+) -> dict[str, Any] | None:
+    """Pack complete thinker output for the non-async connector path."""
+    if not isinstance(pooling_output, dict):
+        return None
+
+    layers = {
+        0: pooling_output.get("hidden_states.layer_0"),
+        24: pooling_output.get("hidden_states.layer_24"),
+    }
+    thinker_emb = _layer_tensor(layers, _EMBED_LAYER_KEY)
+    thinker_hid = _layer_tensor(layers, _HIDDEN_LAYER_KEY)
+    if thinker_emb is None:
+        hidden = pooling_output.get("hidden")
+        thinker_emb = hidden if isinstance(hidden, torch.Tensor) else None
+    if thinker_emb is None or thinker_hid is None:
+        logger.debug(
+            "thinker2talker_full_payload: missing thinker tensors for req=%s (embed=%s hidden=%s)",
+            getattr(request, "request_id", None),
+            thinker_emb is not None,
+            thinker_hid is not None,
+        )
+        return None
+
+    prompt_token_ids = _ensure_list(getattr(request, "prompt_token_ids", []) or [])
+    all_token_ids = _ensure_list(getattr(request, "all_token_ids", None) or [])
+    if not all_token_ids:
+        output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or [])
+        all_token_ids = list(prompt_token_ids) + list(output_token_ids)
+
+    # Trim the trailing stop-token row from the accumulated thinker output.
+    # The accumulator captures one hidden-state row per executed thinker
+    # forward (prefill + every decode step including the one that emitted
+    # the stop_token), so for a finished request thinker_emb has exactly one
+    # row more than the rows the talker should consume.  async_chunk's
+    # chunk-0 path naturally captures only the prefill / non-stop portion,
+    # which is why the [async_chunk] parametrization passes while [default]
+    # over-generates one codec frame on short outputs (e.g.
+    # test_one_word_prompt_001[default]: audio extends "London" with
+    # spurious phonemes).
+    if isinstance(thinker_emb, torch.Tensor) and thinker_emb.shape[0] > 0:
+        thinker_emb_prefill = thinker_emb[:-1]
+    else:
+        thinker_emb_prefill = thinker_emb
+    if isinstance(thinker_hid, torch.Tensor) and thinker_hid.shape[0] > 0:
+        thinker_hid_prefill = thinker_hid[:-1]
+    else:
+        thinker_hid_prefill = thinker_hid
+
+    payload: OmniPayload = {
+        "embed": {
+            "prefill": thinker_emb_prefill.detach().cpu(),
+            "tts_bos": _as_tensor_or_none(pooling_output.get("embed.tts_bos")),
+            "tts_eos": _as_tensor_or_none(pooling_output.get("embed.tts_eos")),
+            "tts_pad": _as_tensor_or_none(pooling_output.get("embed.tts_pad")),
+        },
+        "hidden_states": {"output": thinker_hid_prefill.detach().cpu()},
+        "ids": {"all": list(all_token_ids), "prompt": list(prompt_token_ids)},
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
+    payload["next_stage_prompt_len"] = _compute_talker_prompt_ids_length(payload, device="cpu")
+    speaker = extract_speaker_from_request(request)
+    if speaker is not None:
+        payload["speaker"] = speaker
+    language = extract_language_from_request(request)
+    if language is not None:
+        payload["language"] = language
+    return payload
+
+
 def thinker2talker(
     source_outputs: list[Any],
     prompt: OmniTokensPrompt | TextPrompt | None = None,
@@ -482,6 +637,85 @@ def thinker2talker(
     return talker_inputs
 
 
+def thinker2talker_token_only(
+    source_outputs: list[Any],
+    prompt: OmniTokensPrompt | TextPrompt | None = None,
+    requires_multimodal_data: bool = False,
+    streaming_context: Any | None = None,
+) -> list[OmniTokensPrompt]:
+    """Non-async-chunk Stage-1 input builder for the connector data plane.
+
+    The worker connector (Stage-0 ``thinker2talker_full_payload`` →
+    ``_sync_local_stage_payloads``) supplies the bulk talker conditioning
+    tensors (embed / hidden_states / ids) via ``model_intermediate_buffer``.
+    The orchestrator only needs to ship a placeholder prefill prompt of the
+    correct length so the scheduler can allocate KV-cache slots.
+
+    Small per-request voice metadata (``speaker`` / ``language``) is forwarded
+    here from the user prompt so the worker's line-408 buffer seed picks it
+    up. The connector-side ``extract_speaker_from_request`` reads the
+    strongly-typed ``request.additional_information.entries["speaker"]`` which
+    currently does not always round-trip the user-supplied voice; until that
+    plumbing is normalized, providing the small fields directly preserves
+    voice selection (regression discovered on Buildkite 9668:
+    ``test_speaker_002[default]`` lost the preset voice).
+    """
+    talker_inputs: list[OmniTokensPrompt] = []
+    for i, thinker_output in enumerate(source_outputs):
+        output = thinker_output.outputs[0]
+        req_id = str(getattr(thinker_output, "request_id", f"idx-{i}"))
+        # Skip-on-missing parity with thinker2talker_full_payload: if the
+        # connector builder would drop this request (no MM dict or missing
+        # hidden-state layers), do the same here so the worker buffer
+        # presence agrees with the orchestrator's scheduling decision.
+        thinker_mm_raw = getattr(output, "multimodal_output", None)
+        if not isinstance(thinker_mm_raw, dict):
+            logger.debug("thinker2talker_token_only: skip req=%s due to empty multimodal_output", req_id)
+            continue
+        mm_hs = thinker_mm_raw.get("hidden_states", {})
+        mm_layers = mm_hs.get("layers", {}) if isinstance(mm_hs, dict) else {}
+        if _layer_tensor(mm_layers, _EMBED_LAYER_KEY) is None or _layer_tensor(mm_layers, _HIDDEN_LAYER_KEY) is None:
+            logger.debug("thinker2talker_token_only: skip req=%s due to missing hidden-state layers", req_id)
+            continue
+        prompt_token_ids = _ensure_list(thinker_output.prompt_token_ids)
+        output_ids = _ensure_list(output.cumulative_token_ids)
+        is_streaming_session = bool(getattr(streaming_context, "enabled", False))
+        if is_streaming_session:
+            prompt_token_ids, output_ids, thinker_sequences, thinker_input_ids = _get_streaming_talker_tokens(
+                req_id,
+                prompt_token_ids,
+                output_ids,
+                getattr(streaming_context, "new_prompt_len_snapshot", None),
+                streaming_context,
+                clear_state=bool(getattr(thinker_output, "finished", False)),
+            )
+        else:
+            thinker_sequences = prompt_token_ids + output_ids
+            thinker_input_ids = prompt_token_ids
+        info_for_len = {"ids": {"all": thinker_sequences, "prompt": thinker_input_ids}}
+        prompt_len = _compute_talker_prompt_ids_length(info_for_len, device="cpu")
+
+        # Forward only small voice metadata; bulk tensors come from the
+        # connector path via _sync_local_stage_payloads.
+        small_info: dict[str, Any] = {}
+        speaker = extract_speaker_from_prompt(prompt, index=i)
+        if speaker is not None:
+            small_info["speaker"] = speaker
+        language = extract_language_from_prompt(prompt, index=i)
+        if language is not None:
+            small_info["language"] = language
+
+        talker_inputs.append(
+            OmniTokensPrompt(
+                prompt_token_ids=[0] * prompt_len,
+                additional_information=(small_info if small_info else None),
+                multi_modal_data=None,
+                mm_processor_kwargs=None,
+            )
+        )
+    return talker_inputs
+
+
 # =========================
 # Talker -> Code2Wav
 # =========================
@@ -556,6 +790,56 @@ def talker2code2wav_async_chunk(
     )
 
 
+def talker2code2wav_full_payload(
+    transfer_manager: Any,
+    pooling_output: dict[str, Any],
+    request: OmniEngineCoreRequest,
+) -> dict[str, Any] | None:
+    """Pack complete talker codec output for the non-async connector path."""
+    if not isinstance(pooling_output, dict):
+        return None
+    code_predictor_codes = pooling_output.get("codes.audio")
+    if code_predictor_codes is None:
+        codes = pooling_output.get("codes")
+        if isinstance(codes, dict):
+            code_predictor_codes = codes.get("audio")
+    if code_predictor_codes is None:
+        return None
+    if not isinstance(code_predictor_codes, torch.Tensor):
+        code_predictor_codes = torch.as_tensor(code_predictor_codes)
+    if code_predictor_codes.numel() == 0:
+        return None
+
+    output_token_ids = _ensure_list(getattr(request, "output_token_ids", []) or [])
+    raw_shape = tuple(code_predictor_codes.shape)
+    code_predictor_codes, codec_stats = _extract_qwen3_full_payload_codec_rows(
+        code_predictor_codes.to(torch.long),
+        list(output_token_ids),
+    )
+    if code_predictor_codes.numel() == 0:
+        return None
+
+    codec_codes = code_predictor_codes.transpose(0, 1).cpu().reshape(-1).tolist()
+    logger.debug(
+        "talker2code2wav_full_payload: raw_shape=%s output_ids_len=%s aligned_rows=%s "
+        "valid_rows=%s placeholders=%s flattened_len=%s pad4196=%s bos4197=%s eos4198=%s",
+        raw_shape,
+        len(output_token_ids),
+        codec_stats["aligned_rows"],
+        codec_stats["valid_rows"],
+        codec_stats["trailing_placeholder_count"],
+        len(codec_codes),
+        sum(1 for tid in output_token_ids if tid == _QWEN3_CODEC_PAD_TOKEN_ID),
+        sum(1 for tid in output_token_ids if tid == _QWEN3_CODEC_BOS_TOKEN_ID),
+        sum(1 for tid in output_token_ids if tid == _QWEN3_CODEC_EOS_TOKEN_ID),
+    )
+    return {
+        "codes": {"audio": codec_codes},
+        "code_predictor_codes": codec_codes,
+        "meta": {"finished": torch.tensor(True, dtype=torch.bool)},
+    }
+
+
 def talker2code2wav(
     source_outputs: list[Any],
     _prompt: OmniTokensPrompt | TextPrompt | None = None,

vllm_omni/patch.py

@@ -102,6 +102,11 @@ if not hasattr(RequestStatus, "WAITING_FOR_CHUNK"):
     # as a non-finished state and remains compatible with existing comparisons.
     extend_enum(RequestStatus, "WAITING_FOR_CHUNK", -1)
 
+if not hasattr(RequestStatus, "WAITING_FOR_INPUT"):
+    # Full-payload stage handoff uses a distinct waiting state so the
+    # scheduler can restore the request once non-stage-0 inputs arrive.
+    extend_enum(RequestStatus, "WAITING_FOR_INPUT", -2)
+
 # Snapshot sys.modules: `hasattr` below can trigger lazy submodule imports
 # (e.g. transformers' `_LazyModule.__getattr__`), which mutate sys.modules
 # during iteration and raise `dictionary changed size during iteration`.

vllm_omni/worker/gpu_ar_model_runner.py

@@ -81,6 +81,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
         self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False)
         # Initialize KV cache manager (preserve vllm_config fallback behavior)
         self.kv_transfer_manager = OmniKVTransferManager.from_vllm_config(self.vllm_config, self.model_config)
+        # Only Qwen3-Omni currently consumes the connector-based full-payload
+        # handoff added in this PR. Other model architectures (e.g. Bagel
+        # diffusion) retain their pre-existing runner behavior so this PR
+        # does not perturb them.
+        if getattr(self.model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration":
+            self.init_omni_connectors(
+                vllm_config=self.vllm_config,
+                model_config=self.model_config,
+                kv_transfer_manager=self.kv_transfer_manager,
+            )
         self._downstream_payload_cache: dict[str, bool] = {}
 
     def _make_buffer(self, *size, dtype, numpy=True):
@@ -337,6 +347,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             request_id_resolver=self._resolve_global_request_id,
         )
 
+        if hasattr(self, "_omni_connector"):
+            for request in getattr(scheduler_output, "pending_input_registrations", []):
+                self.register_chunk_recv(request)
+            self.recv_full_payload_inputs(scheduler_output)
+            if self._pending_full_payload_send:
+                flush_ids = set(getattr(scheduler_output, "finished_req_ids", set()))
+                flush_ids.update({rid for rid in self._pending_full_payload_send if rid not in self.requests})
+                if flush_ids:
+                    self.flush_full_payload_outputs(flush_ids)
+
         if self.routed_experts_initialized:
             capturer = RoutedExpertsCapturer.get_instance()
             if capturer is not None:
@@ -375,7 +395,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     if kv_ids:
                         output = copy(output)
                         output.kv_extracted_req_ids = kv_ids
-                    return output
+                    return self.attach_omni_connector_output(output)
 
             if not num_scheduled_tokens:
                 if (
@@ -399,7 +419,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                     output = copy(output)
                     output.kv_extracted_req_ids = kv_ids
 
-                return output
+                return self.attach_omni_connector_output(output)
 
             if self.cache_config.kv_sharing_fast_prefill:
                 assert not self.num_prompt_logprobs, (
@@ -735,11 +755,11 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
             # In case of PP with kv transfer, we need to pass through the
             # kv_connector_output
             if kv_connector_output.is_empty():
-                return EMPTY_MODEL_RUNNER_OUTPUT
+                return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
 
             output = copy(EMPTY_MODEL_RUNNER_OUTPUT)
             output.kv_connector_output = kv_connector_output
-            return output
+            return self.attach_omni_connector_output(output)
 
         # Unpack ephemeral state.
         (
@@ -985,6 +1005,13 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 # Flatten nested dicts to dotted keys so pooling_output
                 # stays dict[str, torch.Tensor] for msgspec serialization.
                 pooler_output.append(flatten_payload(payload))
+
+        if pooler_output and self._should_accumulate_full_payload_output():
+            for i, rid in enumerate(req_ids_output_copy):
+                req_state = self.requests.get(rid)
+                if req_state is not None and pooler_output[i]:
+                    self.accumulate_full_payload_output(rid, pooler_output[i], req_state)
+
         with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
             if self.routed_experts_initialized:
                 capturer = RoutedExpertsCapturer.get_instance()
@@ -1005,6 +1032,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
                 cudagraph_stats=cudagraph_stats,
             )
             output.kv_extracted_req_ids = kv_extracted_req_ids
+            output.omni_connector_output = self.get_omni_connector_output()
 
         if not self.use_async_scheduling:
             return output

vllm_omni/worker/gpu_generation_model_runner.py

@@ -53,6 +53,17 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
     - Executes generation process and returns tensors via `pooler_output`.
     """
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # Scope full-payload connector init to Qwen3-Omni: other generation
+        # models (e.g. Bagel DiT) retain their pre-existing runner setup
+        # so this refactor does not perturb them.
+        if getattr(self.model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration":
+            self.init_omni_connectors(
+                vllm_config=self.vllm_config,
+                model_config=self.model_config,
+            )
+
     def _update_request_states(self, scheduler_output: SchedulerOutput):
         # remove requests
         for req_id in scheduler_output.finished_req_ids:
@@ -89,6 +100,16 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         if self.execute_model_state is not None:
             raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.")
 
+        if hasattr(self, "_omni_connector"):
+            for request in getattr(scheduler_output, "pending_input_registrations", []):
+                self.register_chunk_recv(request)
+            self.recv_full_payload_inputs(scheduler_output)
+            if self._pending_full_payload_send:
+                flush_ids = set(getattr(scheduler_output, "finished_req_ids", set()))
+                flush_ids.update({rid for rid in self._pending_full_payload_send if rid not in self.requests})
+                if flush_ids:
+                    self.flush_full_payload_outputs(flush_ids)
+
         if self.routed_experts_initialized:
             capturer = RoutedExpertsCapturer.get_instance()
             if capturer is not None:
@@ -110,7 +131,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                 self._update_request_states(scheduler_output)
             deferred_state_corrections_fn = self._update_states(scheduler_output)
             if not scheduler_output.total_num_scheduled_tokens:
-                return EMPTY_MODEL_RUNNER_OUTPUT
+                return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
 
             if has_ec_transfer() and not get_ec_transfer().is_consumer:
                 with self.maybe_get_ec_connector_output(
@@ -118,7 +139,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                     encoder_cache=self.encoder_cache,
                 ) as ec_connector_output:
                     self._execute_mm_encoder(scheduler_output)
-                    return make_empty_encoder_model_runner_output(scheduler_output)
+                    return self.attach_omni_connector_output(make_empty_encoder_model_runner_output(scheduler_output))
 
             if not num_scheduled_tokens:
                 if (
@@ -134,9 +155,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
                     self._dummy_run(1)
                 if not has_kv_transfer_group():
                     # Return empty ModelRunnerOutput if no work to do.
-                    return EMPTY_MODEL_RUNNER_OUTPUT
+                    return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
 
-                return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
+                result = self.kv_connector_no_forward(scheduler_output, self.vllm_config)
+                return self.attach_omni_connector_output(result)
 
             if self.cache_config.kv_sharing_fast_prefill:
                 assert not self.num_prompt_logprobs, (
@@ -338,11 +360,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             # In case of PP with kv transfer, we need to pass through the
             # kv_connector_output
             if kv_connector_output.is_empty():
-                return EMPTY_MODEL_RUNNER_OUTPUT
+                return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT)
 
             output = copy(EMPTY_MODEL_RUNNER_OUTPUT)
             output.kv_connector_output = kv_connector_output
-            return output
+            return self.attach_omni_connector_output(output)
 
         # Unpack ephemeral state.
         (
@@ -403,6 +425,12 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
         # [Omni] Copy req_id mappings to avoid async scheduling mutation.
         req_ids_output_copy = self.input_batch.req_ids.copy()
         req_id_to_index_output_copy = self.input_batch.req_id_to_index.copy()
+        if self._should_accumulate_full_payload_output():
+            for i, rid in enumerate(req_ids_output_copy):
+                req_state = self.requests.get(rid)
+                if req_state is not None and pooler_output[i]:
+                    self.accumulate_full_payload_output(rid, pooler_output[i], req_state)
+
         output = OmniModelRunnerOutput(
             req_ids=req_ids_output_copy,
             req_id_to_index=req_id_to_index_output_copy,
@@ -415,6 +443,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin
             cudagraph_stats=cudagraph_stats,
             ec_connector_output=ec_connector_output if self.supports_mm_inputs else None,
         )
+        output.omni_connector_output = self.get_omni_connector_output()
 
         if not self.use_async_scheduling:
             return output

vllm_omni/worker/gpu_model_runner.py

@@ -1,3 +1,4 @@
+import contextlib
 from typing import TYPE_CHECKING, Any, cast
 
 import numpy as np
@@ -264,6 +265,14 @@ class OmniGPUModelRunner(GPUModelRunner):
             self.omni_prefix_cache.reset_prefix_cached_new_req_ids()
 
         # Remove finished requests from the cached states.
+        # cleanup_finished_request lives on OmniConnectorModelRunnerMixin and
+        # is only safe to call once init_omni_connectors() has populated the
+        # mixin state. Archs that inherit the method via MRO without running
+        # that init must be skipped, so probe a mixin-owned attribute as the
+        # "state initialized" gate.
+        cleanup_finished_request = (
+            getattr(self, "cleanup_finished_request", None) if hasattr(self, "_request_ids_mapping") else None
+        )
         for req_id in scheduler_output.finished_req_ids:
             self.requests.pop(req_id, None)
             self.model_intermediate_buffer.pop(req_id, None)
@@ -272,6 +281,8 @@ class OmniGPUModelRunner(GPUModelRunner):
                 self._downstream_payload_cache.pop(req_id, None)
             if hasattr(self, "_talker_mtp_generators"):
                 self._talker_mtp_generators.pop(req_id, None)
+            if cleanup_finished_request is not None:
+                cleanup_finished_request(req_id)
 
         if hasattr(self, "late_interaction_runner"):
             self.late_interaction_runner.on_requests_finished(scheduler_output.finished_req_ids)
@@ -1033,8 +1044,31 @@ class OmniGPUModelRunner(GPUModelRunner):
             req_token_spans.append((start_offset, start_offset + sched_tokens))
         return req_token_spans
 
+    def _sync_local_stage_payloads(self) -> None:
+        """Move received full-payload stage inputs into model_intermediate_buffer."""
+        cache = getattr(self, "_local_stage_payload_cache", None)
+        if not cache:
+            return
+        lock = getattr(self, "_lock", None)
+        ctx = lock if lock is not None else contextlib.nullcontext()
+        with ctx:
+            if not cache:
+                return
+            active_req_ids = set(getattr(self, "requests", {}))
+            pending = set(getattr(self, "_full_payload_pending_broadcast_req_ids", set()))
+            staged = {
+                req_id: payload
+                for req_id, payload in cache.items()
+                if req_id not in pending and req_id in active_req_ids and isinstance(payload, dict)
+            }
+            for req_id in staged:
+                cache.pop(req_id, None)
+        for req_id, payload in staged.items():
+            self._update_intermediate_buffer(req_id, payload)
+
     def _build_model_kwargs_extra(self) -> dict:
         """Build extra keyword arguments passed to the model for this step."""
+        self._sync_local_stage_payloads()
         model_kwargs_extra: dict[str, object] = {}
         try:
             buffer_map = self._gather_runtime_additional_information()
@@ -1299,6 +1333,13 @@ class OmniGPUModelRunner(GPUModelRunner):
         if hasattr(self.model, "has_preprocess") or hasattr(self.model, "enable_update_additional_information"):
             if self.vllm_config.model_config.async_chunk:
                 self._update_additional_information(scheduler_output)
+            else:
+                # In full-payload (non-async-chunk) mode, connector-delivered
+                # stage payloads must override any earlier engine-level
+                # additional_information written by the legacy
+                # custom_process_input_func codec, so talker_preprocess reads
+                # the full thinker payload.
+                self._sync_local_stage_payloads()
 
         if hasattr(self.model, "has_preprocess") and self.model.has_preprocess:
             # Overlay custom prompt_embeds per request for the prompt portion;

vllm_omni/worker/omni_connector_model_runner_mixin.py

@@ -153,7 +153,7 @@ class OmniConnectorModelRunnerMixin:
 
         # -- full_payload_mode: accumulate latest pooler_output per request,
         #    send only when the request finishes (next-cycle flush) --
-        self._pending_full_payload_send: dict[str, tuple[Any, Any]] = {}
+        self._pending_full_payload_send: dict[str, tuple[Any, ...]] = {}
 
         # -- KV sent accumulator --
         self._kv_sent_req_ids: list[str] = []
@@ -206,21 +206,33 @@ class OmniConnectorModelRunnerMixin:
         """Clean up per-request state after a request is fully finished.
 
         Call this when a request is freed from the model runner to prevent
-        memory leaks in the mixin's tracking dicts/sets.  The external
-        request ID is resolved before cleaning up ``_put_req_chunk`` which
-        is keyed by external ID.
+        memory leaks in the mixin's tracking dicts/sets.
+
+        Two senders use different keys: ``send_chunk`` keys per-request
+        state under the EXTERNAL id (after mapping resolution), while
+        ``send_full_payload_outputs`` keys under the INTERNAL id. To cover
+        both modes (and forward compat with id-rename scenarios) we attempt
+        cleanup against both keys; the entry that doesn't exist for the
+        active mode is a no-op pop.  Only the key that actually has pending
+        saves is added to ``_deferred_send_cleanup`` so the bg save's
+        decrement path drains it without leaving orphans.
         """
         ext_id = self._request_ids_mapping.pop(req_id, None)
-        send_req_id = ext_id if ext_id is not None else req_id
+        keys_to_clean: list[str] = [req_id]
+        if ext_id is not None and ext_id != req_id:
+            keys_to_clean.append(ext_id)
 
         with self._lock:
-            if self._pending_save_counts.get(send_req_id, 0):
-                self._deferred_send_cleanup.add(send_req_id)
-            else:
-                self._put_req_chunk.pop(send_req_id, None)
-                self._send_side_request_payload.pop(send_req_id, None)
-                self._code_prompt_token_ids.pop(send_req_id, None)
-                self._cached_ic.pop(send_req_id, None)
+            keys_pending = [k for k in keys_to_clean if self._pending_save_counts.get(k, 0)]
+            for k in keys_pending:
+                self._deferred_send_cleanup.add(k)
+            for k in keys_to_clean:
+                if k in keys_pending:
+                    continue
+                self._put_req_chunk.pop(k, None)
+                self._send_side_request_payload.pop(k, None)
+                self._code_prompt_token_ids.pop(k, None)
+                self._cached_ic.pop(k, None)
             self._kv_pending_transfers.pop(req_id, None)
             self._kv_active_transfers.discard(req_id)
             self._kv_completed_transfers.discard(req_id)
@@ -628,6 +640,16 @@ class OmniConnectorModelRunnerMixin:
         that has arrived, or ``None`` if nothing is ready.  Stores full
         payloads in the local cache and extracts scheduling metadata.
         """
+        # Fast path: when TP is trivial (no peer ranks waiting on a broadcast)
+        # and the bg recv thread has not staged anything, skip the lock + TP
+        # broadcast cycle entirely. _broadcast_tp_payload_packet already
+        # returns its input unchanged under the same world_size<=1 condition,
+        # so the original code path was a no-op here on every empty step.
+        tp_group = self._get_local_tp_group()
+        if (
+            tp_group is None or getattr(tp_group, "world_size", 1) <= 1
+        ) and not self._full_payload_pending_broadcast_req_ids:
+            return None
         with self._lock:
             results = self._collect_full_payload_results_locked() if self.is_data_transfer_rank() else None
         results = self._broadcast_tp_payload_packet(results)
@@ -649,10 +671,63 @@ class OmniConnectorModelRunnerMixin:
         )
         return results
 
+    def _get_model_config(self) -> Any:
+        model_config = getattr(self, "model_config", None)
+        if model_config is not None:
+            return model_config
+        return getattr(getattr(self, "vllm_config", None), "model_config", None)
+
+    def _should_accumulate_full_payload_output(self) -> bool:
+        """Gate send-side full-payload output accumulation only.
+
+        Cached per instance: the result depends only on model_config /
+        _custom_process_func, both of which are set at init time. Avoid
+        the per-step dynamic import inside the model decode loop.
+        """
+        cached = getattr(self, "_should_accumulate_full_payload_output_cached", None)
+        if cached is not None:
+            return cached
+        model_config = self._get_model_config()
+        if model_config is None:
+            self._should_accumulate_full_payload_output_cached = False
+            return False
+        if getattr(model_config, "model_arch", None) == "Qwen3OmniMoeForConditionalGeneration":
+            from vllm_omni.model_executor.stage_input_processors.qwen3_omni import (
+                should_accumulate_qwen3_omni_full_payload_output,
+            )
+
+            result = should_accumulate_qwen3_omni_full_payload_output(
+                model_config,
+                getattr(self, "_custom_process_func", None),
+            )
+            self._should_accumulate_full_payload_output_cached = result
+            return result
+        self._should_accumulate_full_payload_output_cached = False
+        return False
+
     @staticmethod
-    def _is_all_zero_tensor(t: Any) -> bool:
-        """Return True if *t* is a torch.Tensor whose elements are all zero."""
-        return isinstance(t, torch.Tensor) and t.numel() > 0 and not t.any()
+    def _new_full_payload_accumulator(output: dict[str, Any]):
+        chunks: dict[str, list[torch.Tensor]] = {}
+        latest: dict[str, Any] = {}
+        rows: dict[str, int] = {}
+        for k, v in output.items():
+            if isinstance(v, torch.Tensor) and v.dim() >= 2:
+                chunks[k] = [v]
+                rows[k] = int(v.shape[0])
+            else:
+                latest[k] = v
+        return chunks, latest, rows
+
+    @staticmethod
+    def _materialize_full_payload_entry(entry):
+        if len(entry) == 2:
+            return entry
+        chunks, latest, _rows, request = entry
+        output = dict(latest)
+        for k, tensors in chunks.items():
+            if tensors:
+                output[k] = tensors[0] if len(tensors) == 1 else torch.cat(tensors, dim=0)
+        return output, request
 
     def accumulate_full_payload_output(
         self,
@@ -666,55 +741,45 @@ class OmniConnectorModelRunnerMixin:
         along dim-0.  Scalar / global tensors (1-D or 0-D) are replaced
         with the latest value.
 
-        All-zero tensors (e.g. ``code_predictor_codes`` emitted during
-        prefill) are dropped so that they do not pollute downstream stages
-        with garbage / noise frames.
+        Note: codec rows are NOT filtered for zero placeholders here. The
+        downstream consumer ``_extract_qwen3_full_payload_codec_rows`` crops
+        codec rows using ``output_token_ids`` as the authoritative source,
+        which makes any sender-side zero filtering redundant. Skipping the
+        sender-side ``t.any()`` scan also avoids a per-tensor GPU->CPU device
+        sync that stalled the decode pipeline.
 
         The data is actually sent when ``flush_full_payload_outputs`` is called
         with the finished request IDs from the next scheduler cycle.
         """
-        # ---- Filter out all-zero tensors from the incoming pooler_output ----
-        filtered: dict[str, Any] = {}
-        dropped_zero_keys: list[tuple[str, tuple[int, ...]]] = []
-        for k, v in pooler_output.items():
-            if self._is_all_zero_tensor(v):
-                dropped_zero_keys.append((k, tuple(v.shape)))
-                continue  # skip prefill zero-filled placeholders
-            filtered[k] = v
-        if dropped_zero_keys:
-            logger.info(
-                "[Stage-%s] accumulate_full_payload_output: req=%s dropped_zero_keys=%s",
-                self._stage_id,
-                req_id,
-                dropped_zero_keys,
-            )
-        pooler_output = filtered
-
         existing = self._pending_full_payload_send.get(req_id)
+
         if existing is None:
-            self._pending_full_payload_send[req_id] = (pooler_output, request)
+            chunks, latest, rows = self._new_full_payload_accumulator(pooler_output)
+            self._pending_full_payload_send[req_id] = (chunks, latest, rows, request)
             return
 
-        prev_output, _ = existing
-        merged: dict[str, Any] = {}
-        for k in set(prev_output) | set(pooler_output):
-            v_new = pooler_output.get(k)
-            v_old = prev_output.get(k)
-            if v_new is None:
-                merged[k] = v_old
-            elif v_old is None:
-                merged[k] = v_new
-            elif (
-                isinstance(v_new, torch.Tensor)
-                and isinstance(v_old, torch.Tensor)
-                and v_new.dim() >= 2
-                and v_old.dim() >= 2
-                and v_new.shape[1:] == v_old.shape[1:]
-            ):
-                merged[k] = torch.cat([v_old, v_new], dim=0)
+        if len(existing) == 2:
+            chunks, latest, rows = self._new_full_payload_accumulator(existing[0])
+        else:
+            chunks, latest, rows, _ = existing
+
+        for k, v in pooler_output.items():
+            if v is None:
+                continue
+            if isinstance(v, torch.Tensor) and v.dim() >= 2:
+                if k in chunks and chunks[k] and v.shape[1:] == chunks[k][0].shape[1:]:
+                    chunks[k].append(v)
+                    rows[k] += int(v.shape[0])
+                else:
+                    latest.pop(k, None)
+                    chunks[k] = [v]
+                    rows[k] = int(v.shape[0])
             else:
-                merged[k] = v_new
-        self._pending_full_payload_send[req_id] = (merged, request)
+                chunks.pop(k, None)
+                rows.pop(k, None)
+                latest[k] = v
+
+        self._pending_full_payload_send[req_id] = (chunks, latest, rows, request)
 
     def flush_full_payload_outputs(self, finished_req_ids: set[str]) -> None:
         """Send accumulated full_payload outputs for requests that just finished."""
@@ -728,7 +793,7 @@ class OmniConnectorModelRunnerMixin:
         for req_id in finished_req_ids:
             entry = self._pending_full_payload_send.pop(req_id, None)
             if entry is not None:
-                to_send[req_id] = entry
+                to_send[req_id] = self._materialize_full_payload_entry(entry)
         logger.info("[Stage-%s] flush_full_payload_outputs: to_send=%s", self._stage_id, list(to_send.keys()))
         if to_send:
             self.send_full_payload_outputs(scheduler_output=None, outputs=to_send)
@@ -1476,7 +1541,7 @@ class OmniConnectorModelRunnerMixin:
                     logger.warning("Error receiving data for %s", req_id, exc_info=True)
 
             if not made_progress and not self._stop_event.is_set():
-                self._work_available.wait(timeout=0.001)
+                self._work_available.wait(timeout=0.005)
                 self._work_available.clear()
 
     _MAX_SEND_RETRIES = 3