[core]refactor communication layer: PR1(Added Refactor Infra Only) (#1555)

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

tests/core/sched/test_chunk_scheduling_coordinator.py

@@ -0,0 +1,690 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for OmniSchedulingCoordinator (formerly ChunkSchedulingCoordinator).
+
+These tests use mock request objects and mock queues.  They do not require
+GPU, vLLM runtime, or any connector.
+"""
+
+from __future__ import annotations
+
+import unittest
+from types import SimpleNamespace
+
+import vllm_omni.core.sched.omni_scheduling_coordinator as coord_mod
+from vllm_omni.core.sched.omni_scheduling_coordinator import (
+    ChunkSchedulingCoordinator,
+    OmniSchedulingCoordinator,
+)
+
+# ------------------------------------------------------------------ #
+#  Mock helpers
+# ------------------------------------------------------------------ #
+
+
+class _RequestStatus:
+    WAITING = "waiting"
+    RUNNING = "running"
+    WAITING_FOR_CHUNK = "waiting_for_chunk"
+    WAITING_FOR_INPUT = "waiting_for_input"
+    FINISHED_STOPPED = "finished_stopped"
+
+
+# Patch RequestStatus for tests that don't import vllm
+try:
+    from vllm.v1.request import RequestStatus
+except ImportError:
+    RequestStatus = _RequestStatus  # type: ignore[misc,assignment]
+
+if not hasattr(RequestStatus, "WAITING_FOR_INPUT"):
+    coord_mod.RequestStatus = _RequestStatus  # type: ignore[assignment]
+    RequestStatus = _RequestStatus  # type: ignore[misc,assignment]
+
+
+def _make_request(req_id: str, status: str = "waiting") -> SimpleNamespace:
+    return SimpleNamespace(
+        request_id=req_id,
+        external_req_id=req_id,
+        status=status,
+        additional_information=None,
+        prompt_token_ids=[],
+        num_prompt_tokens=0,
+        num_computed_tokens=0,
+        _all_token_ids=[],
+        _output_token_ids=[],
+    )
+
+
+class MockQueue:
+    """Simplified queue that mimics the Scheduler waiting queue interface."""
+
+    def __init__(self, items: list | None = None):
+        self._items: list = list(items or [])
+
+    def __iter__(self):
+        return iter(self._items)
+
+    def __len__(self):
+        return len(self._items)
+
+    def __contains__(self, item):
+        return item in self._items
+
+    def add_request(self, request):
+        self._items.append(request)
+
+    def prepend_requests(self, requests):
+        self._items = list(requests) + self._items
+
+    def remove(self, request):
+        self._items.remove(request)
+
+    def remove_requests(self, requests):
+        remove_set = set(id(r) for r in requests)
+        self._items = [r for r in self._items if id(r) not in remove_set]
+
+
+# ------------------------------------------------------------------ #
+#  Tests
+# ------------------------------------------------------------------ #
+
+
+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)
+
+        req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
+        waiting = MockQueue([req])
+        running: list = []
+
+        coord.process_pending_chunks(
+            waiting,
+            running,
+            chunk_ready_req_ids={"r1"},
+            chunk_finished_req_ids=set(),
+        )
+
+        self.assertEqual(req.status, RequestStatus.WAITING)
+        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)
+
+        req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
+        waiting = MockQueue([req])
+        running: list = []
+
+        coord.process_pending_chunks(
+            waiting,
+            running,
+            chunk_ready_req_ids=set(),
+            chunk_finished_req_ids=set(),
+        )
+
+        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)
+        req = _make_request("r1")
+        waiting = MockQueue([req])
+        running: list = []
+
+        coord.process_pending_chunks(
+            waiting,
+            running,
+            chunk_ready_req_ids={"r1"},
+            chunk_finished_req_ids=set(),
+        )
+        self.assertNotEqual(req.status, RequestStatus.WAITING_FOR_CHUNK)
+
+
+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)
+
+        r1 = _make_request("r1")
+        r2 = _make_request("r2")
+        coord._waiting_for_chunk_waiting.append(r1)
+        coord._waiting_for_chunk_running.append(r2)
+
+        waiting = MockQueue()
+        running: list = []
+
+        coord.restore_queues(waiting, running)
+
+        self.assertIn(r1, waiting)
+        self.assertIn(r2, running)
+        self.assertEqual(len(coord._waiting_for_chunk_waiting), 0)
+        self.assertEqual(len(coord._waiting_for_chunk_running), 0)
+
+
+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)
+
+        req = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
+        waiting = MockQueue([req])
+        running: list = []
+
+        coord.process_pending_chunks(
+            waiting,
+            running,
+            chunk_ready_req_ids={"r1"},
+            chunk_finished_req_ids={"r1"},
+        )
+
+        self.assertIn("r1", coord.finished_requests)
+
+
+class TestChunkCoordinatorUpdateRequestMetadata(unittest.TestCase):
+    """Test update_request_metadata applies scheduling metadata to requests."""
+
+    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)
+
+        req = _make_request("r1")
+        requests = {"r1": req}
+
+        # Only scheduling metadata is passed now (full payload stays in model runner)
+        request_metadata = {"r1": {"next_stage_prompt_len": 50}}
+
+        coord.update_request_metadata(requests, request_metadata, model_mode="ar")
+
+        # next_stage_prompt_len should update prompt_token_ids
+        self.assertEqual(len(req.prompt_token_ids), 50)
+        self.assertEqual(req.num_prompt_tokens, 50)
+        # additional_information should NOT be set
+        self.assertIsNone(getattr(req, "additional_information", None))
+
+    def test_generation_mode(self):
+        coord = ChunkSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+
+        req = _make_request("r1")
+        req.prompt_token_ids = [0, 0, 0]
+        requests = {"r1": req}
+
+        request_metadata = {
+            "r1": {
+                "code_predictor_codes": [10, 20, 30],
+                "left_context_size": 25,
+            }
+        }
+
+        coord.update_request_metadata(requests, request_metadata, model_mode="generation")
+
+        self.assertEqual(req.prompt_token_ids, [10, 20, 30])
+        self.assertEqual(req.num_computed_tokens, 0)
+        self.assertIsNone(req.additional_information)
+        self.assertEqual(req._omni_initial_model_buffer, {"left_context_size": 25})
+
+
+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.requests_with_ready_chunks = {"r1", "r2"}
+
+        new_req = SimpleNamespace(req_id="r1")
+        cached_reqs = SimpleNamespace(req_ids=["r2"])
+        scheduler_output = SimpleNamespace(
+            scheduled_new_reqs=[new_req],
+            scheduled_cached_reqs=cached_reqs,
+        )
+
+        coord.postprocess_scheduler_output(scheduler_output)
+
+        self.assertEqual(coord.requests_with_ready_chunks, set())
+
+
+class TestWaitingForInputTransition(unittest.TestCase):
+    """Test B8: process_pending_full_payload_inputs transitions WAITING_FOR_INPUT."""
+
+    def test_transition_on_recv(self):
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+
+        req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT)
+        waiting = MockQueue([req])
+        running: list = []
+
+        coord.process_pending_full_payload_inputs(
+            waiting,
+            running,
+            stage_recv_req_ids={"r1"},
+        )
+
+        self.assertEqual(req.status, RequestStatus.WAITING)
+
+    def test_stays_waiting_for_input_if_not_received(self):
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+
+        req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT)
+        waiting = MockQueue([req])
+        running: list = []
+
+        coord.process_pending_full_payload_inputs(
+            waiting,
+            running,
+            stage_recv_req_ids=set(),
+        )
+
+        self.assertEqual(req.status, RequestStatus.WAITING_FOR_INPUT)
+        self.assertEqual(len(coord._waiting_for_input), 1)
+
+    def test_stage_0_is_noop(self):
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=0)
+
+        req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT)
+        waiting = MockQueue([req])
+        running: list = []
+
+        coord.process_pending_full_payload_inputs(
+            waiting,
+            running,
+            stage_recv_req_ids={"r1"},
+        )
+        self.assertEqual(req.status, RequestStatus.WAITING_FOR_INPUT)
+
+    def test_restore_queues_includes_waiting_for_input(self):
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+
+        r1 = _make_request("r1")
+        coord._waiting_for_input.append(r1)
+
+        waiting = MockQueue()
+        running: list = []
+
+        coord.restore_queues(waiting, running)
+
+        self.assertIn(r1, waiting)
+        self.assertEqual(len(coord._waiting_for_input), 0)
+
+    def test_full_payload_mode_auto_transitions_waiting_to_waiting_for_input(self):
+        """In full_payload_mode (async_chunk=False), fresh WAITING requests on
+        non-Stage-0 should be transitioned to WAITING_FOR_INPUT."""
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+            async_chunk=False,
+        )
+
+        req = _make_request("r1", status=RequestStatus.WAITING)
+        waiting = MockQueue([req])
+        running: list = []
+
+        coord.process_pending_full_payload_inputs(
+            waiting,
+            running,
+            stage_recv_req_ids=set(),
+        )
+
+        self.assertEqual(req.status, RequestStatus.WAITING_FOR_INPUT)
+        self.assertEqual(len(coord._waiting_for_input), 1)
+        self.assertEqual(len(coord.pending_input_registrations), 1)
+
+    def test_async_chunk_mode_does_not_auto_transition(self):
+        """In async_chunk mode, fresh WAITING requests should NOT be
+        transitioned to WAITING_FOR_INPUT."""
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+            async_chunk=True,
+        )
+
+        req = _make_request("r1", status=RequestStatus.WAITING)
+        waiting = MockQueue([req])
+        running: list = []
+
+        coord.process_pending_full_payload_inputs(
+            waiting,
+            running,
+            stage_recv_req_ids=set(),
+        )
+
+        self.assertEqual(req.status, RequestStatus.WAITING)
+
+    def test_pending_input_registrations(self):
+        coord = OmniSchedulingCoordinator(scheduler_max_num_seqs=10, stage_id=1)
+
+        req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT)
+        waiting = MockQueue([req])
+        running: list = []
+
+        coord.process_pending_full_payload_inputs(
+            waiting,
+            running,
+            stage_recv_req_ids=set(),
+        )
+
+        self.assertEqual(len(coord.pending_input_registrations), 1)
+        self.assertEqual(coord.pending_input_registrations[0].request_id, "r1")
+
+
+class TestTimeoutDetection(unittest.TestCase):
+    """Regression tests for orphaned pending-recv timeout detection.
+
+    Covers the full lifecycle:
+      1. Request enters WAITING_FOR_CHUNK from either waiting or running queue
+      2. restore_queues() moves it back to the scheduler queue
+      3. Timeout fires via collect_timed_out_request_ids()
+      4. Scheduler removes from both queues and calls _free_request()
+    """
+
+    def test_waiting_since_recorded_on_chunk_wait(self):
+        """_waiting_since is set when a request enters WAITING_FOR_CHUNK."""
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+            async_chunk=True,
+        )
+        req = _make_request("r1", status=RequestStatus.WAITING)
+        waiting = MockQueue([req])
+
+        coord.process_pending_chunks(
+            waiting,
+            [],
+            chunk_ready_req_ids=set(),
+            chunk_finished_req_ids=set(),
+        )
+
+        self.assertIn("r1", coord._waiting_since)
+        self.assertEqual(req.status, RequestStatus.WAITING_FOR_CHUNK)
+
+    def test_waiting_since_cleared_on_chunk_arrival(self):
+        """_waiting_since is cleared when a chunk arrives."""
+        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])
+
+        coord.process_pending_chunks(
+            waiting,
+            [],
+            chunk_ready_req_ids={"r1"},
+            chunk_finished_req_ids=set(),
+        )
+
+        self.assertNotIn("r1", coord._waiting_since)
+
+    def test_waiting_since_recorded_on_input_wait(self):
+        """_waiting_since is set when a request enters WAITING_FOR_INPUT."""
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+            async_chunk=False,
+        )
+        req = _make_request("r1", status=RequestStatus.WAITING)
+        waiting = MockQueue([req])
+
+        coord.process_pending_full_payload_inputs(
+            waiting,
+            [],
+            stage_recv_req_ids=set(),
+        )
+
+        self.assertIn("r1", coord._waiting_since)
+
+    def test_waiting_since_cleared_on_input_arrival(self):
+        """_waiting_since is cleared when input data arrives."""
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+            async_chunk=False,
+        )
+        req = _make_request("r1", status=RequestStatus.WAITING_FOR_INPUT)
+        coord._waiting_for_input.append(req)
+        coord._waiting_since["r1"] = 0.0
+
+        waiting = MockQueue()
+        coord.process_pending_full_payload_inputs(
+            waiting,
+            [],
+            stage_recv_req_ids={"r1"},
+        )
+
+        self.assertNotIn("r1", coord._waiting_since)
+        self.assertEqual(req.status, RequestStatus.WAITING)
+
+    def test_collect_timed_out_request_ids_no_timeout(self):
+        """No IDs returned when nothing has timed out."""
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+        )
+        import time
+
+        coord._waiting_since["r1"] = time.monotonic()
+
+        result = coord.collect_timed_out_request_ids(timeout_s=300.0)
+        self.assertEqual(result, set())
+
+    def test_collect_timed_out_request_ids_expired(self):
+        """Timed-out IDs are returned and _waiting_since is cleared."""
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+        )
+        coord._waiting_since["r1"] = 0.0  # epoch → definitely expired
+        coord._waiting_since["r2"] = 0.0
+
+        import time
+
+        coord._waiting_since["r3"] = time.monotonic() + 9999  # far future
+
+        result = coord.collect_timed_out_request_ids(timeout_s=1.0)
+
+        self.assertEqual(result, {"r1", "r2"})
+        self.assertNotIn("r1", coord._waiting_since)
+        self.assertNotIn("r2", coord._waiting_since)
+        self.assertIn("r3", coord._waiting_since)
+
+    def test_collect_removes_from_coordinator_queues(self):
+        """Timed-out requests are defensively removed from internal queues."""
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+        )
+        r1 = _make_request("r1")
+        r2 = _make_request("r2")
+        coord._waiting_for_chunk_waiting.append(r1)
+        coord._waiting_for_input.append(r2)
+        coord._waiting_since["r1"] = 0.0
+        coord._waiting_since["r2"] = 0.0
+
+        result = coord.collect_timed_out_request_ids(timeout_s=1.0)
+
+        self.assertEqual(result, {"r1", "r2"})
+        self.assertEqual(len(coord._waiting_for_chunk_waiting), 0)
+        self.assertEqual(len(coord._waiting_for_input), 0)
+
+    def test_free_finished_request_clears_waiting_since(self):
+        """free_finished_request clears _waiting_since."""
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+        )
+        coord._waiting_since["r1"] = 0.0
+        coord.free_finished_request("r1")
+        self.assertNotIn("r1", coord._waiting_since)
+
+    def test_timeout_from_running_queue_full_lifecycle(self):
+        """End-to-end: request from running → WAITING_FOR_CHUNK → restore →
+        timeout → removed from running list.
+
+        This is the critical regression case: WAITING_FOR_CHUNK requests
+        that originated from self.running are placed back into self.running
+        by restore_queues(), but their status remains WAITING_FOR_CHUNK.
+        The scheduler must remove from BOTH queues unconditionally.
+        """
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+            async_chunk=True,
+        )
+
+        # 1) Request starts in running queue with WAITING status
+        req = _make_request("r1", status=RequestStatus.WAITING)
+        running = [req]
+        waiting = MockQueue()
+
+        # 2) process_pending_chunks: moves to WAITING_FOR_CHUNK
+        coord.process_pending_chunks(
+            waiting,
+            running,
+            chunk_ready_req_ids=set(),
+            chunk_finished_req_ids=set(),
+        )
+        self.assertEqual(req.status, RequestStatus.WAITING_FOR_CHUNK)
+        self.assertIn("r1", coord._waiting_since)
+        self.assertEqual(len(coord._waiting_for_chunk_running), 1)
+
+        # 3) restore_queues: back to running (status stays WAITING_FOR_CHUNK)
+        coord.restore_queues(waiting, running)
+        self.assertIn(req, running)
+        self.assertEqual(len(coord._waiting_for_chunk_running), 0)
+        self.assertEqual(req.status, RequestStatus.WAITING_FOR_CHUNK)
+
+        # 4) Force timeout by setting _waiting_since to epoch
+        coord._waiting_since["r1"] = 0.0
+
+        timed_out_ids = coord.collect_timed_out_request_ids(timeout_s=1.0)
+        self.assertEqual(timed_out_ids, {"r1"})
+
+        # 5) Scheduler removes from both queues (simulating the scheduler path)
+        timed_out_id_set = {id(req)}
+        running = [r for r in running if id(r) not in timed_out_id_set]
+        waiting.remove_requests([req])
+
+        self.assertNotIn(req, running)
+        self.assertEqual(len(waiting), 0)
+
+    def test_timeout_from_waiting_queue_full_lifecycle(self):
+        """End-to-end: request from waiting → WAITING_FOR_CHUNK → restore →
+        timeout → removed from waiting queue."""
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=10,
+            stage_id=1,
+            async_chunk=True,
+        )
+
+        req = _make_request("r1", status=RequestStatus.WAITING)
+        waiting = MockQueue([req])
+        running: list = []
+
+        coord.process_pending_chunks(
+            waiting,
+            running,
+            chunk_ready_req_ids=set(),
+            chunk_finished_req_ids=set(),
+        )
+        self.assertEqual(len(coord._waiting_for_chunk_waiting), 1)
+
+        coord.restore_queues(waiting, running)
+        self.assertIn(req, waiting)
+
+        coord._waiting_since["r1"] = 0.0
+        timed_out_ids = coord.collect_timed_out_request_ids(timeout_s=1.0)
+        self.assertEqual(timed_out_ids, {"r1"})
+
+        waiting.remove_requests([req])
+        self.assertEqual(len(waiting), 0)
+
+
+class TestOverflowPreemption(unittest.TestCase):
+    """Tests for P1-1: overflow requests must get WAITING status.
+
+    Overflow happens when multiple WAITING_FOR_CHUNK requests in
+    ``_waiting_for_chunk_running`` receive their chunk in the same cycle.
+    ``_process_chunk_queue`` restores them to RUNNING (``continue``
+    path) while RUNNING requests without chunks are moved out.  If the
+    net result exceeds ``scheduler_max_num_seqs``, the tail is pushed
+    to ``waiting_queue`` and must have status == WAITING.
+    """
+
+    def test_overflow_sets_waiting_status(self):
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=1,
+            stage_id=1,
+            async_chunk=True,
+        )
+
+        # r1 is currently RUNNING in the queue.
+        # r2, r3 were previously moved to _waiting_for_chunk_running.
+        r1 = _make_request("r1", status=RequestStatus.RUNNING)
+        r2 = _make_request("r2", status=RequestStatus.WAITING_FOR_CHUNK)
+        r3 = _make_request("r3", status=RequestStatus.WAITING_FOR_CHUNK)
+
+        running = [r1]
+        waiting = MockQueue([])
+        coord._waiting_for_chunk_running.extend([r2, r3])
+
+        # restore_queues puts r2, r3 back into running
+        coord.restore_queues(waiting, running)
+        self.assertEqual(len(running), 3)
+
+        # Now process_pending_chunks with r2, r3 chunks ready:
+        # _process_chunk_queue will:
+        #   r1 (RUNNING) → no chunk → move to _waiting_for_chunk_running
+        #   r2 (WAITING_FOR_CHUNK, chunk ready) → set RUNNING, stay in running
+        #   r3 (WAITING_FOR_CHUNK, chunk ready) → set RUNNING, stay in running
+        # running = [r2, r3], len=2 > max=1 → overflow
+        coord.process_pending_chunks(
+            waiting,
+            running,
+            chunk_ready_req_ids={"r2", "r3"},
+            chunk_finished_req_ids=set(),
+        )
+
+        self.assertEqual(len(running), 1)
+        self.assertEqual(len(waiting), 1)
+        overflow_req = list(waiting)[0]
+        self.assertEqual(
+            overflow_req.status,
+            RequestStatus.WAITING,
+            f"Overflowed request should have WAITING status, got {overflow_req.status}",
+        )
+
+    def test_overflow_does_not_strand_request(self):
+        """Without the fix, the overflowed request would keep its
+        RUNNING status in the waiting queue and never be re-scheduled."""
+        coord = OmniSchedulingCoordinator(
+            scheduler_max_num_seqs=1,
+            stage_id=1,
+            async_chunk=True,
+        )
+
+        r1 = _make_request("r1", status=RequestStatus.WAITING_FOR_CHUNK)
+        r2 = _make_request("r2", status=RequestStatus.WAITING_FOR_CHUNK)
+        coord._waiting_for_chunk_running.extend([r1, r2])
+
+        running: list = []
+        waiting = MockQueue([])
+
+        coord.restore_queues(waiting, running)
+        self.assertEqual(len(running), 2)
+
+        coord.process_pending_chunks(
+            waiting,
+            running,
+            chunk_ready_req_ids={"r1", "r2"},
+            chunk_finished_req_ids=set(),
+        )
+
+        self.assertEqual(len(running), 1)
+        self.assertEqual(len(waiting), 1)
+        for req in waiting:
+            self.assertNotEqual(req.status, RequestStatus.RUNNING, "Overflowed request must not keep RUNNING status")
+
+
+if __name__ == "__main__":
+    unittest.main()

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -0,0 +1,380 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Scheduling-side coordination for chunk and full_payload input waiting.
+
+Manages WAITING_FOR_CHUNK and WAITING_FOR_INPUT state transitions
+based on readiness signals from OmniConnectorOutput, without ever
+calling connector.put()/get().
+
+This replaces the scheduling half of OmniChunkTransferAdapter; the
+transport half lives in OmniConnectorModelRunnerMixin.
+"""
+
+from __future__ import annotations
+
+import time
+from collections import deque
+from typing import Any
+
+from vllm.logger import init_logger
+from vllm.v1.request import Request, RequestStatus
+
+logger = init_logger(__name__)
+
+
+class OmniSchedulingCoordinator:
+    """Pure-scheduling coordinator for chunk and full_payload input waiting.
+
+    The Scheduler owns an instance of this class.  It consumes readiness
+    signals produced by the Model Runner's ``OmniConnectorModelRunnerMixin``
+    (via ``OmniConnectorOutput``) and manages ``WAITING_FOR_CHUNK`` and
+    ``WAITING_FOR_INPUT`` state transitions accordingly.
+    """
+
+    def __init__(self, scheduler_max_num_seqs: int, stage_id: int = 0, async_chunk: bool = False):
+        self._stage_id = stage_id
+        self._scheduler_max_num_seqs = scheduler_max_num_seqs
+        self._async_chunk = async_chunk
+
+        self.finished_requests: set[str] = set()
+        self.requests_with_ready_chunks: set[str] = set()
+        self._full_payload_input_received: set[str] = set()
+
+        self._waiting_for_chunk_waiting: deque[Any] = deque()
+        self._waiting_for_chunk_running: deque[Any] = deque()
+
+        # Request IDs that were newly registered for chunk recv this cycle.
+        # The engine/Model Runner should call register_chunk_recv() for these
+        # so the bg thread starts polling.
+        self.pending_chunk_registrations: list[Any] = []
+
+        # Requests waiting for full_payload stage input (WAITING_FOR_INPUT).
+        self._waiting_for_input: deque[Any] = deque()
+        self.pending_input_registrations: list[Any] = []
+
+        # Monotonic timestamp recording when each request first entered
+        # WAITING_FOR_CHUNK or WAITING_FOR_INPUT.  Used by
+        # collect_timed_out_request_ids() to detect orphaned waits.
+        self._waiting_since: dict[str, float] = {}
+
+    # ------------------------------------------------------------------ #
+    #  Core scheduling methods
+    # ------------------------------------------------------------------ #
+
+    def process_pending_chunks(
+        self,
+        waiting_queue: Any,
+        running_queue: list[Request],
+        chunk_ready_req_ids: set[str],
+        chunk_finished_req_ids: set[str],
+    ) -> None:
+        """Transition requests whose chunks have arrived.
+
+        Args:
+            waiting_queue: Scheduler's waiting request queue.
+            running_queue: Scheduler's running request list.
+            chunk_ready_req_ids: IDs with a newly arrived chunk this cycle.
+            chunk_finished_req_ids: IDs whose final chunk has arrived.
+        """
+        if self._stage_id == 0 or not self._async_chunk:
+            return
+
+        terminal_ready_req_ids = chunk_ready_req_ids.intersection(chunk_finished_req_ids)
+        self.finished_requests.update(chunk_finished_req_ids - terminal_ready_req_ids)
+        self.pending_chunk_registrations = []
+
+        self._process_chunk_queue(
+            waiting_queue,
+            self._waiting_for_chunk_waiting,
+            RequestStatus.WAITING,
+            chunk_ready_req_ids,
+        )
+        self._process_chunk_queue(
+            running_queue,
+            self._waiting_for_chunk_running,
+            RequestStatus.RUNNING,
+            chunk_ready_req_ids,
+        )
+        self.finished_requests.update(terminal_ready_req_ids)
+
+        while len(running_queue) > self._scheduler_max_num_seqs:
+            request = running_queue.pop()
+            # Must reset status to WAITING so the scheduler treats it as
+            # schedulable work.  KV blocks are NOT freed here (unlike a
+            # real preemption), so PREEMPTED would be incorrect.
+            request.status = RequestStatus.WAITING
+            waiting_queue.prepend_requests([request])
+
+    def process_pending_full_payload_inputs(
+        self,
+        waiting_queue: Any,
+        running_queue: list[Request],
+        stage_recv_req_ids: set[str],
+    ) -> None:
+        """Manage WAITING_FOR_INPUT lifecycle for full_payload_mode.
+
+        For non-Stage-0 stages in full_payload_mode (``async_chunk=False``):
+        1. Fresh WAITING requests are transitioned to WAITING_FOR_INPUT
+           and registered for bg-thread polling.
+        2. WAITING_FOR_INPUT requests whose data has arrived (in
+           ``stage_recv_req_ids``) are transitioned back to WAITING.
+        """
+        if self._stage_id == 0:
+            return
+
+        self._full_payload_input_received.update(stage_recv_req_ids)
+        if not self._async_chunk and stage_recv_req_ids:
+            self.finished_requests.update(stage_recv_req_ids)
+            logger.debug(
+                "[Coordinator stage-%s] full_payload recv -> finished_requests: %s",
+                self._stage_id,
+                stage_recv_req_ids,
+            )
+        self.pending_input_registrations = []
+
+        remaining: deque[Any] = deque()
+        for request in self._waiting_for_input:
+            if request.request_id in stage_recv_req_ids:
+                request.status = RequestStatus.WAITING
+                self._waiting_since.pop(request.request_id, None)
+                waiting_queue.add_request(request)
+            else:
+                remaining.append(request)
+        self._waiting_for_input = remaining
+
+        if not self._async_chunk:
+            to_remove: list[Any] = []
+            queue_snapshot = list(waiting_queue)
+            for request in queue_snapshot:
+                if request.status == RequestStatus.WAITING:
+                    if request.request_id in self._full_payload_input_received:
+                        continue
+                    if request.request_id in self.requests_with_ready_chunks:
+                        continue
+                    if request.request_id in self.finished_requests:
+                        continue
+                    request.status = RequestStatus.WAITING_FOR_INPUT
+                    self._waiting_since.setdefault(request.request_id, time.monotonic())
+                    to_remove.append(request)
+                    self._waiting_for_input.append(request)
+                    self.pending_input_registrations.append(request)
+                elif request.status == RequestStatus.WAITING_FOR_INPUT:
+                    if request.request_id in stage_recv_req_ids:
+                        request.status = RequestStatus.WAITING
+                        self._waiting_since.pop(request.request_id, None)
+                    else:
+                        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)
+
+    def process_pending_full_payload_inputs_legacy(
+        self,
+        waiting_queue: Any,
+        running_queue: list[Request],
+        stage_recv_req_ids: set[str],
+    ) -> None:
+        """Compatibility wrapper for ``process_pending_full_payload_inputs``."""
+        self.process_pending_full_payload_inputs(waiting_queue, running_queue, stage_recv_req_ids)
+
+    def free_finished_request(self, request_id: str) -> None:
+        """Prune internal tracking sets for a freed request to prevent unbounded growth."""
+        self._full_payload_input_received.discard(request_id)
+        self.finished_requests.discard(request_id)
+        self.requests_with_ready_chunks.discard(request_id)
+        self._waiting_since.pop(request_id, None)
+
+    def collect_timed_out_request_ids(
+        self,
+        timeout_s: float,
+    ) -> set[str]:
+        """Return IDs of requests that have been waiting longer than *timeout_s*.
+
+        Uses ``_waiting_since`` timestamps (always up-to-date) to detect
+        timed-out requests.  This method is safe to call at any point in
+        the scheduling cycle — it does **not** rely on coordinator internal
+        queues (which are empty after ``restore_queues()``).
+
+        Clears ``_waiting_since`` for timed-out IDs and defensively removes
+        them from coordinator internal queues if present.  The caller
+        (scheduler) should then remove the requests from its queues,
+        set ``FINISHED_ERROR``, and call ``_free_request()`` so that
+        ``cleanup_finished_request()`` fires in the model runner mixin.
+        """
+        if timeout_s <= 0:
+            return set()
+        now = time.monotonic()
+        timed_out_ids: set[str] = set()
+        for req_id, start_time in self._waiting_since.items():
+            if now - start_time > timeout_s:
+                timed_out_ids.add(req_id)
+        if not timed_out_ids:
+            return set()
+
+        # Defensively remove from coordinator internal queues (may already
+        # be empty if restore_queues() has run).
+        for queue_attr in (
+            "_waiting_for_chunk_waiting",
+            "_waiting_for_chunk_running",
+            "_waiting_for_input",
+        ):
+            queue = getattr(self, queue_attr)
+            remaining: deque[Any] = deque()
+            for request in queue:
+                if request.request_id not in timed_out_ids:
+                    remaining.append(request)
+            setattr(self, queue_attr, remaining)
+
+        for req_id in timed_out_ids:
+            self._waiting_since.pop(req_id, None)
+            logger.warning(
+                "[Coordinator stage-%s] Request %s timed out waiting for chunk/input (waited > %.0fs)",
+                self._stage_id,
+                req_id,
+                timeout_s,
+            )
+
+        return timed_out_ids
+
+    def restore_queues(
+        self,
+        waiting_queue: Any,
+        running_queue: list[Request],
+    ) -> None:
+        """Return waiting-for-chunk/input requests to scheduling queues."""
+        for request in self._waiting_for_chunk_waiting:
+            waiting_queue.add_request(request)
+        self._waiting_for_chunk_waiting = deque()
+
+        if self._waiting_for_chunk_running:
+            running_queue.extend(self._waiting_for_chunk_running)
+        self._waiting_for_chunk_running = deque()
+
+        for request in self._waiting_for_input:
+            waiting_queue.add_request(request)
+        self._waiting_for_input = deque()
+
+    def update_request_metadata(
+        self,
+        requests: dict[str, Request],
+        request_metadata: dict[str, dict[str, Any]],
+        model_mode: str = "ar",
+    ) -> None:
+        """Apply received scheduling metadata to request objects.
+
+        For AR mode: only scheduler-visible metadata is applied locally.
+        For Generation mode: updates ``request.prompt_token_ids``.
+
+        Additionally, if the payload contains ``next_stage_prompt_len``,
+        updates the request's ``prompt_token_ids`` to the correct length.
+        """
+        for req_id, metadata in request_metadata.items():
+            request = requests.get(req_id)
+            if request is None:
+                continue
+
+            # Handle next_stage_prompt_len if present (for models like Qwen3-Omni).
+            # Only apply when the request has not started decoding yet
+            # (no output tokens). Resetting a mid-decode request would
+            # destroy generated tokens and desync KV cache state.
+            if "next_stage_prompt_len" in metadata:
+                next_len = metadata["next_stage_prompt_len"]
+                if isinstance(next_len, int) and next_len > 0:
+                    output_token_ids = getattr(request, "_output_token_ids", None)
+                    has_decode_output = output_token_ids is not None and len(output_token_ids) > 0
+                    if has_decode_output:
+                        logger.debug(
+                            "[Coordinator stage-%s] Skipping prompt resize for req %s: "
+                            "request already has %s output tokens",
+                            self._stage_id,
+                            req_id,
+                            len(output_token_ids),
+                        )
+                    else:
+                        current_prompt_ids = getattr(request, "prompt_token_ids", []) or []
+                        current_prompt_len = len(current_prompt_ids)
+                        if current_prompt_len != next_len or getattr(request, "num_prompt_tokens", None) != next_len:
+                            new_prompt = [0] * next_len
+                            request.prompt_token_ids = new_prompt
+                            request.num_prompt_tokens = next_len
+                            request._all_token_ids.clear()
+                            request._all_token_ids.extend(new_prompt)
+                            request._output_token_ids.clear()
+                            request.num_computed_tokens = 0
+                            logger.debug(
+                                "[Coordinator stage-%s] Updated prompt_token_ids length to %s for req %s",
+                                self._stage_id,
+                                next_len,
+                                req_id,
+                            )
+
+            if model_mode != "ar":
+                new_ids = metadata.get("code_predictor_codes", [])
+                runtime_seed = None
+                if "left_context_size" in metadata:
+                    runtime_seed = {
+                        "left_context_size": metadata["left_context_size"],
+                    }
+                request._omni_initial_model_buffer = runtime_seed
+                if new_ids:
+                    request.prompt_token_ids = new_ids
+                    request.num_computed_tokens = 0
+
+    def postprocess_scheduler_output(
+        self,
+        scheduler_output: Any,
+        requests: dict[str, Request] | None = None,
+    ) -> None:
+        """Clear per-cycle ready state after scheduler output is materialized."""
+        self._clear_chunk_ready(scheduler_output)
+
+    # ------------------------------------------------------------------ #
+    #  Internal helpers
+    # ------------------------------------------------------------------ #
+
+    def _process_chunk_queue(
+        self,
+        queue: Any,
+        waiting_for_chunk_list: deque[Any],
+        target_status: RequestStatus,
+        chunk_ready_req_ids: set[str],
+    ) -> None:
+        queue_snapshot = list(queue)
+        for request in queue_snapshot:
+            if request.status != RequestStatus.WAITING_FOR_CHUNK:
+                if request.request_id in self.requests_with_ready_chunks:
+                    continue
+                if request.request_id in self.finished_requests:
+                    continue
+                if request.status == RequestStatus.WAITING_FOR_INPUT:
+                    continue
+                if request.request_id in chunk_ready_req_ids:
+                    self.requests_with_ready_chunks.add(request.request_id)
+                    continue
+                self.pending_chunk_registrations.append(request)
+                request.status = RequestStatus.WAITING_FOR_CHUNK
+                self._waiting_since.setdefault(request.request_id, time.monotonic())
+            else:
+                if request.request_id in chunk_ready_req_ids:
+                    request.status = target_status
+                    self.requests_with_ready_chunks.add(request.request_id)
+                    self._waiting_since.pop(request.request_id, None)
+                    continue
+            queue.remove(request)
+            waiting_for_chunk_list.append(request)
+
+    def _clear_chunk_ready(self, scheduler_output: Any) -> None:
+        if scheduler_output.scheduled_new_reqs:
+            for req_data in scheduler_output.scheduled_new_reqs:
+                self.requests_with_ready_chunks.discard(
+                    getattr(req_data, "req_id", None),
+                )
+
+        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/diffusion/worker/diffusion_model_runner.py

@@ -35,11 +35,12 @@ from vllm_omni.diffusion.sched.interface import DiffusionSchedulerOutput
 from vllm_omni.diffusion.worker.utils import DiffusionRequestState, RunnerOutput
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.platforms import current_omni_platform
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 logger = init_logger(__name__)
 
 
-class DiffusionModelRunner:
+class DiffusionModelRunner(OmniConnectorModelRunnerMixin):
     """
     Model runner that handles model loading and execution for diffusion models.
 

vllm_omni/outputs.py

@@ -9,6 +9,33 @@ from vllm.v1.outputs import ModelRunnerOutput
 from vllm_omni.inputs.data import OmniPromptType
 
 
+@dataclass
+class OmniConnectorOutput:
+    """Communication results from Model Runner to Scheduler.
+
+    Carries transfer readiness signals so the Scheduler can make scheduling
+    decisions without ever calling connector.put()/get() directly.
+
+    Attributes:
+        chunk_ready_req_ids: Request IDs with newly arrived chunks this cycle.
+        chunk_finished_req_ids: Request IDs whose final chunk has arrived.
+        request_metadata: Lightweight scheduling metadata keyed by request ID
+            (e.g. next_stage_prompt_len, code_predictor_codes, left_context_size).
+            Full payloads are owned by the Model Runner's local cache.
+        kv_sent_req_ids: Request IDs whose KV cache was successfully sent.
+        stage_recv_req_ids: Request IDs that received batch stage inputs.
+        has_pending_kv_work: True if the mixin has pending, active, or
+            completed KV transfers that the scheduler should account for.
+    """
+
+    chunk_ready_req_ids: set[str] = field(default_factory=set)
+    chunk_finished_req_ids: set[str] = field(default_factory=set)
+    request_metadata: dict[str, dict[str, Any]] = field(default_factory=dict)
+    kv_sent_req_ids: list[str] = field(default_factory=list)
+    stage_recv_req_ids: set[str] = field(default_factory=set)
+    has_pending_kv_work: bool = False
+
+
 class OmniModelRunnerOutput(ModelRunnerOutput):
     """Model runner output for omni models.
 
@@ -24,6 +51,7 @@ class OmniModelRunnerOutput(ModelRunnerOutput):
     # IDs of requests whose KV cache has been extracted from GPU/NPU to CPU.
     # The Scheduler can safely free the block tables for these requests.
     kv_extracted_req_ids: list[str] | None = None
+    omni_connector_output: OmniConnectorOutput | None = None
 
 
 @dataclass

vllm_omni/worker/gpu_ar_model_runner.py

@@ -40,6 +40,7 @@ from vllm.v1.worker.utils import is_residual_scattered_for_sp
 from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 logger = init_logger(__name__)
 
@@ -60,7 +61,7 @@ class ExecuteModelState(NamedTuple):
     slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None
 
 
-class GPUARModelRunner(OmniGPUModelRunner):
+class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
     """Autoregressive GPU model runner that returns hidden states per request.
 
     Follows the v0.12 two-phase execute/sample flow from GPUModelRunner, and

vllm_omni/worker/gpu_generation_model_runner.py

@@ -39,11 +39,12 @@ from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
 from vllm_omni.outputs import OmniModelRunnerOutput
 from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState
 from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
+from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
 
 logger = logging.getLogger(__name__)
 
 
-class GPUGenerationModelRunner(OmniGPUModelRunner):
+class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin):
     """Generation model runner for vLLM-Omni (non-autoregressive).
 
     - Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue.

vllm_omni/worker/payload_span.py

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