[Refactor][1/N Scheduler]Remove duplicated AR/generation scheduler plumbing and establish explicit shared lifecycle contracts. (#5461)

Signed-off-by: R2-Y <ruiruyang2@gmail.com>

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -3,7 +3,8 @@
 from __future__ import annotations
 
 from collections import defaultdict
-from types import SimpleNamespace
+from collections.abc import Callable
+from types import MethodType, SimpleNamespace
 
 import numpy as np
 import pytest
@@ -12,9 +13,22 @@ from vllm.v1.outputs import LogprobsLists
 from vllm.v1.request import RequestStatus
 
 from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler, _slice_sampled_logprobs
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
+# Shared update_from_output helpers moved onto OmniSchedulerMixin; bind them on
+# the lightweight stub so direct method calls still exercise the real paths.
+_MIXIN_UPDATE_HELPERS = (
+    "_remove_stopped_requests_from_queues",
+    "_handle_failed_kv_load_outputs",
+    "_aggregate_kv_connector_stats",
+    "_publish_kv_cache_events",
+    "_attach_finished_request_sets",
+    "_attach_scheduler_stats",
+    "_capture_omni_connector_output",
+)
+
 
 def _logprobs(
     token_ids: list[list[int]],
@@ -105,18 +119,14 @@ class _RequestQueue(list):
                 self.remove(request)
 
 
-def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
-    """Spec decode can sample tokens past a stop; the trailing tokens are
-    trimmed by _update_request_with_output, and the emitted logprob rows must
-    be trimmed with them (upstream slices logprobs after the trim)."""
-    request = _Request("req")
-    request.num_computed_tokens = 8
+def _make_scheduler_stub(requests: list[_Request]) -> SimpleNamespace:
+    """Build a minimal stub for OmniARScheduler.update_from_output()."""
     scheduler = SimpleNamespace(
         perf_metrics=None,
         connector=None,
         chunk_transfer_adapter=None,
-        requests={request.request_id: request},
-        running=[request],
+        requests={request.request_id: request for request in requests},
+        running=list(requests),
         waiting=_RequestQueue(),
         skipped_waiting=_RequestQueue(),
         structured_output_manager=SimpleNamespace(should_advance=lambda _request: False),
@@ -129,7 +139,42 @@ def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
         finished_req_ids_dict=defaultdict(set),
         kv_cache_manager=SimpleNamespace(take_events=lambda: None),
         kv_event_publisher=SimpleNamespace(publish=lambda _events: None),
+        recompute_kv_load_failures=False,
     )
+    for name in _MIXIN_UPDATE_HELPERS:
+        setattr(scheduler, name, MethodType(getattr(OmniSchedulerMixin, name), scheduler))
+    scheduler._cleanup_kv_tracking = MethodType(OmniARScheduler._cleanup_kv_tracking, scheduler)
+    scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
+    scheduler.make_stats = lambda *args, **kwargs: None
+    return scheduler
+
+
+def _bind_request_lifecycle(
+    scheduler: SimpleNamespace,
+    *,
+    update_request: Callable,
+    handle_stopped: Callable | None = None,
+) -> None:
+    def free_request(request):
+        scheduler.requests.pop(request.request_id, None)
+        scheduler.finished_req_ids.add(request.request_id)
+        scheduler.finished_req_ids_dict[request.client_index].add(request.request_id)
+        # vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
+        return None, None
+
+    scheduler._update_request_with_output = update_request
+    scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
+    scheduler._handle_stopped_request = handle_stopped or (lambda _request: True)
+    scheduler._free_request = free_request
+
+
+def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
+    """Spec decode can sample tokens past a stop; the trailing tokens are
+    trimmed by _update_request_with_output, and the emitted logprob rows must
+    be trimmed with them (upstream slices logprobs after the trim)."""
+    request = _Request("req")
+    request.num_computed_tokens = 8
+    scheduler = _make_scheduler_stub([request])
 
     def update_request_trimming(req, token_ids):
         # Mirror vllm Scheduler._update_request_with_output: the first sampled
@@ -139,19 +184,7 @@ def test_mid_step_stop_trims_logprob_rows_with_token_ids() -> None:
         req.status = RequestStatus.FINISHED_STOPPED
         return token_ids, True
 
-    def free_request(req):
-        scheduler.requests.pop(req.request_id, None)
-        scheduler.finished_req_ids.add(req.request_id)
-        scheduler.finished_req_ids_dict[req.client_index].add(req.request_id)
-        return None, None
-
-    scheduler._update_request_with_output = update_request_trimming
-    scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
-    scheduler._handle_stopped_request = lambda _request: True
-    scheduler._free_request = free_request
-    scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
-    scheduler.make_stats = lambda *args, **kwargs: None
-    scheduler._capture_omni_connector_output = lambda _output: None
+    _bind_request_lifecycle(scheduler, update_request=update_request_trimming)
 
     scheduler_output = SimpleNamespace(
         num_scheduled_tokens={"req": 3},
@@ -190,45 +223,14 @@ def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
     """A bad runner response must not bring down unrelated batch requests."""
     bad = _Request("bad")
     good = _Request("good")
-    scheduler = SimpleNamespace(
-        perf_metrics=None,
-        connector=None,
-        chunk_transfer_adapter=None,
-        requests={bad.request_id: bad, good.request_id: good},
-        running=[bad, good],
-        waiting=_RequestQueue(),
-        skipped_waiting=_RequestQueue(),
-        structured_output_manager=SimpleNamespace(should_advance=lambda _request: False),
-        transfer_triggered_requests=set(),
-        active_kv_transfers=set(),
-        pending_stop_after_extraction=set(),
-        waiting_for_transfer_free=set(),
-        _new_prompt_len_snapshot={},
-        finished_req_ids=set(),
-        finished_req_ids_dict=defaultdict(set),
-        kv_cache_manager=SimpleNamespace(take_events=lambda: None),
-        kv_event_publisher=SimpleNamespace(publish=lambda _events: None),
-    )
+    scheduler = _make_scheduler_stub([bad, good])
     update_calls: list[str] = []
 
     def update_request(request, token_ids):
         update_calls.append(request.request_id)
         return token_ids, False
 
-    def free_request(request):
-        scheduler.requests.pop(request.request_id, None)
-        scheduler.finished_req_ids.add(request.request_id)
-        scheduler.finished_req_ids_dict[request.client_index].add(request.request_id)
-        # vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
-        return None, None
-
-    scheduler._update_request_with_output = update_request
-    scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
-    scheduler._handle_stopped_request = lambda _request: True
-    scheduler._free_request = free_request
-    scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
-    scheduler.make_stats = lambda *args, **kwargs: None
-    scheduler._capture_omni_connector_output = lambda _output: None
+    _bind_request_lifecycle(scheduler, update_request=update_request)
 
     scheduler_output = SimpleNamespace(
         num_scheduled_tokens={"bad": 1, "good": 1},

tests/core/sched/test_omni_scheduler_mixin_shared.py

@@ -0,0 +1,90 @@
+from collections import defaultdict
+from types import SimpleNamespace
+
+import pytest
+from vllm.v1.engine import FinishReason
+
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+from vllm_omni.core.sched.output import OmniChunkRecvHandle
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Scheduler(OmniSchedulerMixin):
+    pass
+
+
+def test_schedule_lifecycle_helpers_process_and_restore_both_input_paths():
+    calls = []
+    scheduler = _Scheduler()
+    scheduler.waiting = ["waiting"]
+    scheduler.running = ["running"]
+    scheduler.requests = {"request": object()}
+    scheduler._consume_pending_connector_output = lambda mode: calls.append(("consume", mode))
+    scheduler._process_pending_input_timeouts = lambda: calls.append(("timeouts",))
+    scheduler.chunk_transfer_adapter = SimpleNamespace(
+        process_pending_chunks=lambda waiting, running, scheduler_requests: calls.append(
+            ("process", waiting, running, scheduler_requests)
+        ),
+        restore_queues=lambda waiting, running, scheduler_requests: calls.append(
+            ("restore-chunks", waiting, running, scheduler_requests)
+        ),
+    )
+    scheduler.input_coordinator = SimpleNamespace(
+        restore_queues=lambda waiting: calls.append(("restore-full", waiting))
+    )
+
+    scheduler._process_pending_omni_inputs("ar")
+    scheduler._restore_omni_wait_queues()
+
+    assert calls == [
+        ("consume", "ar"),
+        ("timeouts",),
+        ("process", scheduler.waiting, scheduler.running, scheduler.requests),
+        ("restore-chunks", scheduler.waiting, scheduler.running, scheduler.requests),
+        ("restore-full", scheduler.waiting),
+    ]
+
+
+@pytest.mark.parametrize(
+    ("synthesize_abort_outputs", "expected_finish_reason"),
+    [(False, None), (True, FinishReason.ABORT)],
+)
+def test_finished_request_attachment_keeps_ar_abort_policy_explicit(
+    synthesize_abort_outputs,
+    expected_finish_reason,
+):
+    scheduler = _Scheduler()
+    scheduler.finished_req_ids_dict = defaultdict(set, {2: {"req-finished"}})
+    outputs = {}
+
+    scheduler._attach_finished_request_sets(
+        outputs,
+        synthesize_abort_outputs=synthesize_abort_outputs,
+    )
+
+    assert outputs[2].finished_requests == {"req-finished"}
+    if expected_finish_reason is None:
+        assert outputs[2].outputs == []
+    else:
+        assert outputs[2].outputs[0].finish_reason == expected_finish_reason
+    assert scheduler.finished_req_ids_dict == {}
+
+
+def test_chunk_receive_handle_carries_minimal_registration_fields():
+    handle = OmniChunkRecvHandle(request_id="req", external_req_id="external")
+    assert handle.request_id == "req"
+    assert handle.external_req_id == "external"
+
+
+def test_output_helper_preserves_required_nan_counter_default():
+    scheduler = _Scheduler()
+    request = SimpleNamespace(
+        request_id="req-output",
+        trace_headers=None,
+        take_events=lambda: [],
+    )
+
+    output = scheduler._make_omni_engine_output(request, new_token_ids=[])
+
+    assert output.num_nans_in_logits == 0

tests/entrypoints/test_omni_new_request_data.py

@@ -2,7 +2,9 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from vllm.v1.core.sched.output import NewRequestData
 
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
 from vllm_omni.core.sched.output import OmniNewRequestData
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -52,3 +54,80 @@ def test_omni_new_request_data_allows_missing_payloads():
 
     assert data.prompt_embeds is None
     assert data.additional_information is None
+
+
+def test_from_base_preserves_every_upstream_field_and_adds_omni_payloads():
+    prompt_embeds = torch.randn(2, 3)
+    base = NewRequestData(
+        req_id="req-base",
+        prompt_token_ids=[1, 2],
+        mm_features=[],
+        sampling_params=None,
+        pooling_params=None,
+        block_ids=([3, 4],),
+        num_computed_tokens=1,
+        lora_request=None,
+        prompt_embeds=prompt_embeds,
+        prompt_is_token_ids=False,
+        prefill_token_ids=[1],
+    )
+    request = SimpleNamespace(
+        external_req_id="external-base",
+        additional_information={"payload": "value"},
+        model_intermediate_buffer={"hidden": "state"},
+    )
+
+    data = OmniNewRequestData.from_base(base, request)
+
+    for field_name in NewRequestData.__dataclass_fields__:
+        assert getattr(data, field_name) is getattr(base, field_name)
+    assert data.external_req_id == "external-base"
+    assert data.additional_information == {"payload": "value"}
+    assert data.model_intermediate_buffer == {"hidden": "state"}
+
+
+def test_rewrap_converts_fallback_entries_and_keeps_fast_path_entries():
+    base = NewRequestData(
+        req_id="req-fallback",
+        prompt_token_ids=[1],
+        mm_features=[],
+        sampling_params=None,
+        pooling_params=None,
+        block_ids=([2],),
+        num_computed_tokens=0,
+        lora_request=None,
+        prompt_embeds=None,
+        prompt_is_token_ids=True,
+        prefill_token_ids=[1],
+    )
+    fast_path = OmniNewRequestData(
+        req_id="req-fast",
+        prompt_token_ids=[3],
+        mm_features=[],
+        sampling_params=None,
+        pooling_params=None,
+        block_ids=([4],),
+        num_computed_tokens=0,
+        lora_request=None,
+        prompt_embeds=None,
+        prompt_is_token_ids=True,
+        prefill_token_ids=[3],
+        additional_information={"fast": True},
+    )
+    scheduler = OmniSchedulerMixin()
+    scheduler.requests = {
+        "req-fallback": SimpleNamespace(
+            external_req_id="external-fallback",
+            additional_information={"fallback": True},
+            model_intermediate_buffer=None,
+        )
+    }
+    output = SimpleNamespace(scheduled_new_reqs=[base, fast_path])
+
+    scheduler._rewrap_scheduled_new_reqs(output)
+
+    fallback_data, fast_path_data = output.scheduled_new_reqs
+    assert isinstance(fallback_data, OmniNewRequestData)
+    assert fallback_data.prefill_token_ids == [1]
+    assert fallback_data.additional_information == {"fallback": True}
+    assert fast_path_data is fast_path

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -2,37 +2,25 @@ from __future__ import annotations
 
 from collections import defaultdict
 from collections.abc import Iterable
-from time import time
 from typing import Any
 
 import numpy as np
 from vllm.compilation.cuda_graph import CUDAGraphStat
-from vllm.distributed.kv_events import KVEventBatch
-from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
 from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMScheduler
 from vllm.v1.core.sched.output import SchedulerOutput
 from vllm.v1.core.sched.request_queue import create_request_queue
 from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
-from vllm.v1.core.sched.utils import remove_all
-from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs, FinishReason
+from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
 from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.outputs import ModelRunnerOutput
 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_full_payload_input_coordinator,
-)
 from vllm_omni.core.sched.utils import omni_routed_experts_for_request
-from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
-    OmniChunkTransferAdapter,
-)
 from vllm_omni.engine import OmniEngineCoreOutput
 from vllm_omni.engine.serialization import deserialize_additional_information
-from vllm_omni.outputs import OmniConnectorOutput
 
 logger = init_logger(__name__)
 
@@ -108,6 +96,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         self.finished_req_ids_dict = defaultdict(set)
 
         # [Omni] Pre-parse KV transfer criteria
+        self._omni_kv_config = getattr(self.vllm_config.model_config, "omni_kv_config", None)
         self.kv_transfer_criteria = self._get_kv_transfer_criteria()
 
         # Track requests that have already triggered prefill transfer to avoid duplicates
@@ -115,16 +104,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         # Cache per-request flag to avoid repeated deserialization of additional_information
         self._omits_kv_transfer_cache: dict[str, bool] = {}
-        model_config = self.vllm_config.model_config
-        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_full_payload_input_coordinator(model_config):
-            self.input_coordinator = OmniSchedulingCoordinator(
-                stage_id=getattr(model_config, "stage_id", 0),
-            )
-        self._latest_omni_connector_output: OmniConnectorOutput | None = None
+        self._init_omni_io_scheduling_state()
         # Snapshot prompt length for each streaming input update
         self._new_prompt_len_snapshot: dict[str, int] = {}
 
@@ -134,17 +114,15 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
         return request.num_computed_tokens - request.num_output_placeholders
 
     def _get_kv_transfer_criteria(self) -> dict | None:
-        # Note: vllm_config is available in Scheduler after super().__init__
-        if not hasattr(self, "vllm_config"):
-            return None
+        return self._get_omni_kv_config_value("kv_transfer_criteria")
 
-        omni_kv_config = getattr(self.vllm_config.model_config, "omni_kv_config", None)
-        if omni_kv_config:
-            if isinstance(omni_kv_config, dict):
-                return omni_kv_config.get("kv_transfer_criteria", None)
-            else:
-                return getattr(omni_kv_config, "kv_transfer_criteria", None)
-        return None
+    def _get_omni_kv_config_value(self, key: str, default: Any = None) -> Any:
+        config = getattr(self, "_omni_kv_config", None)
+        if config is None and hasattr(self, "vllm_config"):
+            config = getattr(self.vllm_config.model_config, "omni_kv_config", None)
+        if isinstance(config, dict):
+            return config.get(key, default)
+        return getattr(config, key, default) if config is not None else default
 
     def _request_omits_kv_transfer_to_next_stage(self, request: Request) -> bool:
         """True when this stage-zero-final request does not need downstream KV.
@@ -207,42 +185,39 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         if criteria_type == "prefill_finished":
             if confirmed_computed >= request.num_prompt_tokens:
-                self.transfer_triggered_requests.add(request.request_id)
-
-                self._mark_request_for_kv_transfer(request.request_id, confirmed_computed)
-                actually_queued = request.request_id in self.requests_needing_kv_transfer
-
-                if stop_decode_on_trigger and actually_queued:
-                    # Defer the stop until KV extraction completes so that
-                    # the kv_ready signal can be emitted while the request
-                    # is still alive.  The request will be stopped on the
-                    # next scheduler step after extraction ack arrives.
-                    self.pending_stop_after_extraction.add(request.request_id)
-
+                self._commit_kv_transfer_trigger(
+                    request.request_id,
+                    confirmed_computed,
+                    stop_decode_on_trigger,
+                )
                 return False
 
         elif criteria_type == "special_token":
             target_token_id = self.kv_transfer_criteria.get("token_id")
             if target_token_id is not None and target_token_id in new_token_ids:
-                self.transfer_triggered_requests.add(request.request_id)
-
-                try:
-                    idx = new_token_ids.index(target_token_id)
-                    tokens_to_exclude = len(new_token_ids) - (idx + 1)
-                    snapshot_len = confirmed_computed - tokens_to_exclude
-                except ValueError:
-                    snapshot_len = confirmed_computed
-
-                self._mark_request_for_kv_transfer(request.request_id, snapshot_len)
-                actually_queued = request.request_id in self.requests_needing_kv_transfer
-
-                if stop_decode_on_trigger and actually_queued:
-                    self.pending_stop_after_extraction.add(request.request_id)
-
+                idx = new_token_ids.index(target_token_id)
+                tokens_to_exclude = len(new_token_ids) - (idx + 1)
+                snapshot_len = confirmed_computed - tokens_to_exclude
+                self._commit_kv_transfer_trigger(
+                    request.request_id,
+                    snapshot_len,
+                    stop_decode_on_trigger,
+                )
                 return False
 
         return False
 
+    def _commit_kv_transfer_trigger(
+        self,
+        req_id: str,
+        seq_len: int,
+        stop_after_transfer: bool,
+    ) -> None:
+        self.transfer_triggered_requests.add(req_id)
+        self._mark_request_for_kv_transfer(req_id, seq_len)
+        if stop_after_transfer and req_id in self.requests_needing_kv_transfer:
+            self.pending_stop_after_extraction.add(req_id)
+
     def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
         # Remove FINISHED_ABORTED requests before the upstream scheduler sees
         # them. Upstream vllm raises RuntimeError on this status; omni allows
@@ -252,12 +227,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             for req in list(queue):
                 if getattr(req, "status", None) == RequestStatus.FINISHED_ABORTED:
                     queue.remove(req)
-        self._consume_pending_connector_output(model_mode="ar")
-        self._process_pending_input_timeouts()
-        if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.process_pending_chunks(
-                self.waiting, self.running, scheduler_requests=self.requests
-            )
+        self._process_pending_omni_inputs(model_mode="ar")
 
         original_waiting = None
         if self._should_defer_waiting_admission():
@@ -272,55 +242,13 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 if deferred_waiting:
                     original_waiting.prepend_requests(deferred_waiting)
                 self.waiting = original_waiting
-            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,
-                    scheduler_requests=self.requests,
-                )
-            if self.input_coordinator:
-                self.input_coordinator.restore_queues(self.waiting)
-        try:
-            # Late import to avoid circulars in some launch modes
-            from .output import OmniNewRequestData
+            self._restore_omni_wait_queues()
 
-            # Rewrap base NewRequestData entries with OmniNewRequestData,
-            # enriching with request-level payloads
-            new_list = []
-            for nr in scheduler_output.scheduled_new_reqs:
-                req_id = getattr(nr, "req_id", None)
-                request = self.requests.get(req_id) if req_id else None
-                # Build omni entry preserving all base fields
-                omni_nr = OmniNewRequestData(
-                    req_id=nr.req_id,
-                    external_req_id=(getattr(request, "external_req_id", None) if request else None),
-                    prompt_token_ids=nr.prompt_token_ids,
-                    mm_features=nr.mm_features,
-                    sampling_params=nr.sampling_params,
-                    pooling_params=nr.pooling_params,
-                    block_ids=nr.block_ids,
-                    num_computed_tokens=nr.num_computed_tokens,
-                    lora_request=nr.lora_request,
-                    # Enrich with omni payloads from the live request object
-                    prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
-                    prompt_is_token_ids=nr.prompt_is_token_ids,
-                    additional_information=(getattr(request, "additional_information", None) if request else None),
-                    model_intermediate_buffer=(
-                        getattr(request, "model_intermediate_buffer", None) if request else None
-                    ),
-                )
-                new_list.append(omni_nr)
-
-            scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
-            if self.chunk_transfer_adapter:
-                self.chunk_transfer_adapter.postprocess_scheduler_output(scheduler_output, self.requests)
-            # Add information about requests needing KV cache transfer
-            finished_reqs = self.get_finished_requests_needing_kv_transfer()
-        except Exception:
-            # If anything goes wrong, leave the original output unchanged
-            init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
-            finished_reqs = {}
+        self._postprocess_omni_schedule_output(
+            scheduler_output,
+            include_cached_payloads=True,
+        )
+        finished_reqs = self.get_finished_requests_needing_kv_transfer()
 
         # Wrap in omni scheduler output to carry transfer metadata.
         return self._wrap_omni_scheduler_output(
@@ -550,26 +478,23 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
             if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped:
-                # Add EngineCoreOutput for this Request.
-                outputs[request.client_index].append(
-                    OmniEngineCoreOutput(
-                        request_id=req_id,
-                        new_token_ids=new_token_ids,
-                        finish_reason=finish_reason,
-                        new_logprobs=new_logprobs,
-                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
-                        pooling_output=pooler_output,
-                        multimodal_output=mm_output,
-                        stop_reason=request.stop_reason,
-                        events=request.take_events(),
-                        prefill_stats=request.take_prefill_stats(),
-                        kv_transfer_params=kv_transfer_params,
-                        trace_headers=request.trace_headers,
-                        routed_experts=routed_experts,
-                        num_nans_in_logits=request.num_nans_in_logits,
-                        is_segment_finished=is_segment_finished,
-                        new_prompt_len_snapshot=self._new_prompt_len_snapshot.get(req_id, None),
-                    )
+                OmniSchedulerMixin._append_request_output(
+                    self,
+                    outputs,
+                    request,
+                    new_token_ids=new_token_ids,
+                    finish_reason=finish_reason,
+                    new_logprobs=new_logprobs,
+                    new_prompt_logprobs_tensors=prompt_logprobs_tensors,
+                    pooling_output=pooler_output,
+                    multimodal_output=mm_output,
+                    stop_reason=request.stop_reason,
+                    prefill_stats=request.take_prefill_stats(),
+                    kv_transfer_params=kv_transfer_params,
+                    routed_experts=routed_experts,
+                    num_nans_in_logits=request.num_nans_in_logits,
+                    is_segment_finished=is_segment_finished,
+                    new_prompt_len_snapshot=self._new_prompt_len_snapshot.get(req_id),
                 )
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
@@ -584,117 +509,44 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                     is_segment_finished,
                 )
 
-        # Remove the stopped requests from the running and waiting queues.
-        if stopped_running_reqs:
-            self.running = remove_all(self.running, stopped_running_reqs)
-        if stopped_preempted_reqs:
-            # This is a rare case and unlikely to impact performance.
-            self.waiting.remove_requests(stopped_preempted_reqs)
-            self.skipped_waiting.remove_requests(stopped_preempted_reqs)
+        self._remove_stopped_requests_from_queues(
+            stopped_running_reqs,
+            stopped_preempted_reqs,
+        )
 
-        # [Main] Handle failed KV load requests
-        if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
-            requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids]
-            self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
-            for request in requests:
-                outputs[request.client_index].append(
-                    OmniEngineCoreOutput(
-                        request_id=request.request_id,
-                        new_token_ids=[],
-                        finish_reason=request.get_finished_reason(),
-                        events=request.take_events(),
-                        trace_headers=request.trace_headers,
-                    )
-                )
-                if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.cleanup_receiver(
-                        request.request_id,
-                    )
+        failed_requests = self._handle_failed_kv_load_outputs(
+            failed_kv_load_req_ids,
+            outputs,
+        )
+        if self.chunk_transfer_adapter is not None:
+            for request in failed_requests:
+                self.chunk_transfer_adapter.cleanup_receiver(request.request_id)
 
-        # [Omni] Cleanup state for finished requests
-        for req in stopped_running_reqs:
-            if req.request_id not in self.waiting_for_transfer_free:
-                if req.request_id in self.transfer_triggered_requests:
-                    self.transfer_triggered_requests.remove(req.request_id)
-                if req.request_id in self.active_kv_transfers:
-                    self.active_kv_transfers.remove(req.request_id)
-                self.pending_stop_after_extraction.discard(req.request_id)
-
-        # Same for preempted
-        for req in stopped_preempted_reqs:
-            if req.request_id not in self.waiting_for_transfer_free:
-                if req.request_id in self.transfer_triggered_requests:
-                    self.transfer_triggered_requests.remove(req.request_id)
-                if req.request_id in self.active_kv_transfers:
-                    self.active_kv_transfers.remove(req.request_id)
-                self.pending_stop_after_extraction.discard(req.request_id)
+        self._cleanup_kv_tracking(req.request_id for req in stopped_running_reqs | stopped_preempted_reqs)
 
         # KV Connector: update state for finished KV Transfers.
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
-        # Worker-side KV connector stats from the model runner output.
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
-        if self.connector:
-            # Scheduler-side KV connector stats collected after connector update.
-            scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
-            if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty():
-                kv_connector_stats = (
-                    kv_connector_stats.aggregate(scheduler_kv_connector_stats)
-                    if kv_connector_stats is not None
-                    else scheduler_kv_connector_stats
-                )
-
-        # collect KV cache events from KV cache manager
-        events = self.kv_cache_manager.take_events()
-
-        # collect KV cache events from connector
-        if self.connector is not None:
-            connector_events = self.connector.take_events()
-            if connector_events:
-                if events is None:
-                    events = list(connector_events)
-                else:
-                    events.extend(connector_events)
-
-        # publish collected KV cache events
-        if events:
-            batch = KVEventBatch(ts=time(), events=events)
-            self.kv_event_publisher.publish(batch)
+        kv_connector_stats = self._aggregate_kv_connector_stats(kv_connector_output)
+        self._publish_kv_cache_events()
 
         # Create EngineCoreOutputs for all clients that have requests with
         # outputs in this step.
         engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()}
 
-        # FIXME: finished_req_ids_dict is unconditionally initialized as
-        # defaultdict(set) in __init__ (not gated by include_finished_set).
-        # This branch is therefore always eligible once any client_index is
-        # populated; revisit when wiring streaming-only / upstream semantics.
-        finished_req_ids = self.finished_req_ids_dict
-        if finished_req_ids:
-            # Include ids of requests that finished since last outputs
-            # were sent.
-            for client_index, finished_set in finished_req_ids.items():
-                eco = engine_core_outputs.get(client_index)
-                if eco is None:
-                    eco = EngineCoreOutputs()
-                    engine_core_outputs[client_index] = eco
-                emitted = {o.request_id for o in eco.outputs}
-                for req_id in finished_set:
-                    if req_id not in emitted:
-                        eco.outputs.append(EngineCoreOutput(req_id, [], finish_reason=FinishReason.ABORT))
-                eco.finished_requests = finished_set
-            finished_req_ids.clear()
+        self._attach_finished_request_sets(
+            engine_core_outputs,
+            synthesize_abort_outputs=True,
+        )
 
-        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats)) is not None:
-            # Return stats to only one of the front-ends.
-            if (eco := next(iter(engine_core_outputs.values()), None)) is None:
-                # We must return the stats even if there are no request
-                # outputs this step.
-                engine_core_outputs[0] = eco = EngineCoreOutputs()
-            eco.scheduler_stats = stats
+        self._attach_scheduler_stats(
+            engine_core_outputs,
+            spec_decoding_stats,
+            kv_connector_stats,
+            cudagraph_stats,
+            perf_stats,
+        )
 
         self._capture_omni_connector_output(model_runner_output)
 
@@ -720,56 +572,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         return engine_core_outputs
 
-    def finish_requests(self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus) -> list[Request]:
-        """Handles the finish signal from outside the scheduler.
-
-        For example, the API server can abort a request when the client
-        disconnects.
-
-        If request_ids is None, all requests will be finished.
-
-        Returns:
-            The Request objects that were aborted. Will not include any that
-            were already finished.
-        """
-        # TODO(yrr): chunk transfer adapter & input_coordinator unified to one
-        if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
-
-        # Realign stale ``request.status`` (chunk-transfer-adapter's
-        # ``requests_origin_status`` table doesn't follow the
-        # ``waiting → running`` admit transition; without this, an abort
-        # arriving between admit and the next deque round-trip leaves
-        # the request in ``self.running`` with ``status=WAITING`` and
-        # upstream ``Scheduler.finish_requests`` silently fails to
-        # release the worker's ``input_batch`` slot -- after
-        # ``max_num_seqs`` such aborts new requests hang at
-        # ``chunks=0``). Only the ``async_chunk`` path triggers the
-        # staleness; with ``async_chunk`` disabled this is a cheap O(n)
-        # no-op over an already-aligned set, kept unconditional so the
-        # abort path stays uniform across configurations. See
-        # ``OmniSchedulerMixin._realign_request_status_to_queues`` and
-        # #3774 discussion.
-        self._realign_request_status_to_queues(request_ids)
-
-        finished = super().finish_requests(request_ids, finished_status)
-
-        # Defensive post-finish purge: belt-and-suspenders to the
-        # realignment above. Even after realign + ``super()``, corner
-        # cases (mid-transition status, connector cleanups that pop
-        # from ``self.requests`` without unwinding ``self.running``)
-        # can leave already-finished or untracked entries in
-        # ``self.running``. Sweep them now so the worker's
-        # ``input_batch`` slot never pins a freed request and starves
-        # new admissions. See ``OmniSchedulerMixin._purge_finished_from_running``.
-        self._purge_finished_from_running()
-
-        input_coordinator = getattr(self, "input_coordinator", None)
-        if input_coordinator is not None:
-            for request in finished:
-                self._free_input_coordinator_request(request.request_id)
-        return finished
-
     def _update_request_as_session(self, session: Request, update: StreamingUpdate) -> None:
         """
         Override: Only extend prompt at stage 0, and replace
@@ -967,44 +769,35 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
     def _should_transfer_kv_for_request(self, req_id: str) -> bool:
         """Determine if a request should trigger KV cache transfer."""
-        need_send = False
-        # Try to read from vLLM Config (where YAML config is typically loaded)
-        # Check for omni_kv_config attribute
-        omni_kv_config = getattr(self.vllm_config.model_config, "omni_kv_config", None)
-        if omni_kv_config:
-            # omni_kv_config could be an object or a dict
-            if isinstance(omni_kv_config, dict):
-                need_send = omni_kv_config.get("need_send_cache", False)
-            else:
-                need_send = getattr(omni_kv_config, "need_send_cache", False)
-        if not need_send:
+        if not self._get_omni_kv_config_value("need_send_cache", False):
             return False
         request = self.requests.get(req_id)
         if request is not None and self._request_omits_kv_transfer_to_next_stage(request):
             return False
         return True
 
+    def _cleanup_kv_tracking(self, request_ids: Iterable[str]) -> None:
+        for req_id in request_ids:
+            if req_id in self.waiting_for_transfer_free:
+                continue
+            self.transfer_triggered_requests.discard(req_id)
+            self.active_kv_transfers.discard(req_id)
+            self.pending_stop_after_extraction.discard(req_id)
+
+    def _has_pending_kv_work(self) -> bool:
+        return bool(self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free)
+
     def has_requests(self) -> bool:
         """Check if there are any requests to process, including KV transfers."""
-        # [Omni] Also check for pending KV transfers
-        if self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free:
-            return True
-        return super().has_requests()
+        return self._has_pending_kv_work() or super().has_requests()
 
     def has_finished_requests(self) -> bool:
         """Check if there are any finished requests (including those needing KV transfer)."""
-        if self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free:
-            return True
-        return super().has_finished_requests()
+        return self._has_pending_kv_work() or super().has_finished_requests()
 
     def has_unfinished_requests(self) -> bool:
         """Check if there are any unfinished requests (including those needing KV transfer)."""
-        # [Omni] Also check for pending KV transfers to ensure the engine loop continues
-        # MUST verify waiting_for_transfer_free and active_kv_transfers
-        # Otherwise engine loop might exit before transfer Ack is received.
-        if self.requests_needing_kv_transfer or self.active_kv_transfers or self.waiting_for_transfer_free:
-            return True
-        return super().has_unfinished_requests()
+        return self._has_pending_kv_work() or super().has_unfinished_requests()
 
     def get_finished_requests_needing_kv_transfer(self) -> dict[str, dict]:
         """Get and clear the list of requests needing KV cache transfer.

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -2,13 +2,9 @@ from __future__ import annotations
 
 import time
 from collections import defaultdict
-from collections.abc import Iterable
 from typing import Any
 
 from vllm.compilation.cuda_graph import CUDAGraphStat
-from vllm.distributed.kv_events import KVEventBatch
-from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
-from vllm.logger import init_logger
 from vllm.v1.core.kv_cache_manager import KVCacheBlocks
 from vllm.v1.core.sched.interface import PauseState
 from vllm.v1.core.sched.output import SchedulerOutput
@@ -25,36 +21,18 @@ 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_full_payload_input_coordinator,
-)
 from vllm_omni.core.sched.output import OmniCachedRequestData, OmniNewRequestData
 from vllm_omni.core.sched.utils import omni_routed_experts_for_request
-from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
-    OmniChunkTransferAdapter,
-)
-from vllm_omni.engine import OmniEngineCoreOutput
-from vllm_omni.outputs import OmniConnectorOutput, OmniModelRunnerOutput
-
-logger = init_logger(__name__)
+from vllm_omni.outputs import OmniModelRunnerOutput
 
 
 class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         model_config = self.vllm_config.model_config
-        self.chunk_transfer_adapter = None
-        if getattr(model_config, "async_chunk", False):
-            self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config)
+        self._init_omni_io_scheduling_state()
         self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False))
         self._pending_finish_reqs: list[Request] = []
-        self.input_coordinator: OmniSchedulingCoordinator | None = None
-        if uses_full_payload_input_coordinator(model_config):
-            self.input_coordinator = OmniSchedulingCoordinator(
-                stage_id=getattr(model_config, "stage_id", 0),
-            )
-        self._latest_omni_connector_output: OmniConnectorOutput | None = None
 
     @staticmethod
     def _record_prefill_stats(request: Request) -> None:
@@ -117,12 +95,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         # Temporary queue: preserve waiting order while requests await input.
         skipped_waiting_requests = create_request_queue(self.policy)
         req_index = 0
-        self._consume_pending_connector_output(model_mode="generation")
-        self._process_pending_input_timeouts()
-        if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.process_pending_chunks(
-                self.waiting, self.running, scheduler_requests=self.requests
-            )
+        self._process_pending_omni_inputs(model_mode="generation")
 
         # OMNI: Track requests that are already finished (e.g., marked by connector)
         # These should be removed from running and not scheduled
@@ -250,15 +223,11 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             if self.chunk_transfer_adapter:
                 # Don't fall back: base scheduler doesn't handle async_chunk
                 # requests with empty prompt_token_ids.
-                self.chunk_transfer_adapter.restore_queues(
-                    self.waiting,
-                    self.running,
-                    scheduler_requests=self.requests,
-                )
+                self._restore_omni_wait_queues()
             else:
                 res = super().schedule(throttle_prefills)
-                if self.input_coordinator:
-                    self.input_coordinator.restore_queues(self.waiting)
+                self._restore_omni_wait_queues()
+                self._postprocess_omni_schedule_output(res)
                 return self._wrap_omni_scheduler_output(res)
 
         # Compute common prefix blocks (aligned with v1)
@@ -342,94 +311,12 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self._update_after_schedule(scheduler_output)
 
         try:
-            # Rewrap base NewRequestData entries with OmniNewRequestData,
-            # enriching with request-level payloads
-            new_list = []
-            for nr in scheduler_output.scheduled_new_reqs:
-                req_id = getattr(nr, "req_id", None)
-                request = self.requests.get(req_id) if req_id else None
-                # Build omni entry preserving all base fields
-                omni_nr = OmniNewRequestData(
-                    req_id=nr.req_id,
-                    external_req_id=(getattr(request, "external_req_id", None) if request else None),
-                    prompt_token_ids=nr.prompt_token_ids,
-                    mm_features=nr.mm_features,
-                    sampling_params=nr.sampling_params,
-                    pooling_params=nr.pooling_params,
-                    block_ids=nr.block_ids,
-                    num_computed_tokens=nr.num_computed_tokens,
-                    lora_request=nr.lora_request,
-                    # Enrich with omni payloads from the live request object
-                    prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
-                    prompt_is_token_ids=nr.prompt_is_token_ids,
-                    additional_information=(getattr(request, "additional_information", None) if request else None),
-                    model_intermediate_buffer=(
-                        getattr(request, "model_intermediate_buffer", None) if request else None
-                    ),
-                )
-                new_list.append(omni_nr)
-
-            scheduler_output.scheduled_new_reqs = new_list  # type: ignore[assignment]
-
-            if self.chunk_transfer_adapter:
-                self.chunk_transfer_adapter.postprocess_scheduler_output(scheduler_output)
-
-        except Exception:
-            # If anything goes wrong, leave the original output unchanged
-            logger.exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
+            self._postprocess_omni_schedule_output(scheduler_output)
         finally:
-            # Ensure chunk-waiting requests are restored even on error,
-            # otherwise they are permanently orphaned in the adapter's
-            # internal deques and never scheduled again.
-            if self.chunk_transfer_adapter:
-                self.chunk_transfer_adapter.restore_queues(
-                    self.waiting,
-                    self.running,
-                    scheduler_requests=self.requests,
-                )
-            if self.input_coordinator:
-                self.input_coordinator.restore_queues(self.waiting)
+            self._restore_omni_wait_queues()
 
         return self._wrap_omni_scheduler_output(scheduler_output)
 
-    def finish_requests(self, request_ids: str | Iterable[str] | None, finished_status: RequestStatus) -> list[Request]:
-        """Handles the finish signal from outside the scheduler.
-
-        For example, the API server can abort a request when the client
-        disconnects.
-
-        If request_ids is None, all requests will be finished.
-
-        Returns:
-            The Request objects that were aborted. Will not include any that
-            were already finished.
-        """
-
-        if self.chunk_transfer_adapter:
-            self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
-
-        # See ``OmniSchedulerMixin._realign_request_status_to_queues`` --
-        # closes the residual hang on the
-        # ``waiting → running → abort-before-next-deque-round-trip`` race
-        # surfaced by #3774's reproduction matrix. Only the
-        # ``async_chunk`` path actually triggers the staleness; with
-        # ``async_chunk`` disabled this is a cheap O(n) no-op, kept
-        # unconditional so the abort path stays uniform.
-        self._realign_request_status_to_queues(request_ids)
-
-        finished = super().finish_requests(request_ids, finished_status)
-
-        # See ``OmniSchedulerMixin._purge_finished_from_running`` --
-        # defensive belt-and-suspenders sweep paired with the realign
-        # above. Closes residual ``self.running`` slot leaks if any
-        # corner case slips past upstream's status-driven removal.
-        self._purge_finished_from_running()
-
-        if self.input_coordinator is not None:
-            for request in finished:
-                self._free_input_coordinator_request(request.request_id)
-        return finished
-
     def _free_request(
         self, request: Request, delay_free_blocks: bool = False
     ) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
@@ -595,25 +482,22 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
             # Get prompt logprobs for this request.
             prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
             if new_token_ids or mm_output is not None or pooler_output is not None or kv_transfer_params or stopped:
-                # Add EngineCoreOutput for this Request.
-                outputs[request.client_index].append(
-                    OmniEngineCoreOutput(
-                        request_id=req_id,
-                        new_token_ids=new_token_ids,
-                        finish_reason=finish_reason,
-                        new_logprobs=new_logprobs,
-                        new_prompt_logprobs_tensors=prompt_logprobs_tensors,
-                        pooling_output=pooler_output,
-                        multimodal_output=mm_output,
-                        stop_reason=request.stop_reason,
-                        events=request.take_events(),
-                        prefill_stats=request.take_prefill_stats(),
-                        kv_transfer_params=kv_transfer_params,
-                        trace_headers=request.trace_headers,
-                        routed_experts=routed_experts,
-                        num_nans_in_logits=request.num_nans_in_logits,
-                        is_segment_finished=is_segment_finished,
-                    )
+                OmniSchedulerMixin._append_request_output(
+                    self,
+                    outputs,
+                    request,
+                    new_token_ids=new_token_ids,
+                    finish_reason=finish_reason,
+                    new_logprobs=new_logprobs,
+                    new_prompt_logprobs_tensors=prompt_logprobs_tensors,
+                    pooling_output=pooler_output,
+                    multimodal_output=mm_output,
+                    stop_reason=request.stop_reason,
+                    prefill_stats=request.take_prefill_stats(),
+                    kv_transfer_params=kv_transfer_params,
+                    routed_experts=routed_experts,
+                    num_nans_in_logits=request.num_nans_in_logits,
+                    is_segment_finished=is_segment_finished,
                 )
             else:
                 # Invariant: EngineCore returns no partial prefill outputs.
@@ -636,103 +520,58 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
                         request.request_id,
                         getattr(request, "external_req_id", None),
                     )
-            outputs[request.client_index].append(
-                OmniEngineCoreOutput(
-                    request_id=request.request_id,
-                    new_token_ids=[],
-                    finish_reason=finish_reason,
-                    stop_reason=request.stop_reason,
-                    events=request.take_events(),
-                    kv_transfer_params=kv_transfer_params,
-                    trace_headers=request.trace_headers,
-                    is_segment_finished=is_segment_finished,
-                )
+            OmniSchedulerMixin._append_request_output(
+                self,
+                outputs,
+                request,
+                new_token_ids=[],
+                finish_reason=finish_reason,
+                stop_reason=request.stop_reason,
+                kv_transfer_params=kv_transfer_params,
+                is_segment_finished=is_segment_finished,
             )
             stopped_running_reqs.add(request)
         self._pending_finish_reqs.clear()
 
-        # Remove the stopped requests from the running and waiting queues.
-        if stopped_running_reqs:
-            self.running = remove_all(self.running, stopped_running_reqs)
-        if stopped_preempted_reqs:
-            # This is a rare case and unlikely to impact performance.
-            self.waiting.remove_requests(stopped_preempted_reqs)
-            self.skipped_waiting.remove_requests(stopped_preempted_reqs)
+        self._remove_stopped_requests_from_queues(
+            stopped_running_reqs,
+            stopped_preempted_reqs,
+        )
 
-        # Handle failed KV load requests
-        if failed_kv_load_req_ids and not self.recompute_kv_load_failures:
-            requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids]
-            self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR)
-            for request in requests:
-                outputs[request.client_index].append(
-                    OmniEngineCoreOutput(
-                        request_id=request.request_id,
-                        new_token_ids=[],
-                        finish_reason=request.get_finished_reason(),
-                        events=request.take_events(),
-                        trace_headers=request.trace_headers,
-                    )
+        failed_requests = self._handle_failed_kv_load_outputs(
+            failed_kv_load_req_ids,
+            outputs,
+        )
+        if self.chunk_transfer_adapter is not None:
+            for request in failed_requests:
+                self.chunk_transfer_adapter.cleanup(
+                    request.request_id,
+                    getattr(request, "external_req_id", None),
                 )
-                if self.chunk_transfer_adapter is not None:
-                    self.chunk_transfer_adapter.cleanup(
-                        request.request_id,
-                        getattr(request, "external_req_id", None),
-                    )
 
         # KV Connector: update state for finished KV Transfers.
         if kv_connector_output:
             self._update_from_kv_xfer_finished(kv_connector_output)
 
-        # Worker-side KV connector stats from the model runner output.
-        kv_connector_stats: KVConnectorStats | None = (
-            kv_connector_output.kv_connector_stats if kv_connector_output else None
-        )
-        if self.connector:
-            # Scheduler-side KV connector stats collected after connector update.
-            scheduler_kv_connector_stats = self.connector.get_kv_connector_stats()
-            if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty():
-                kv_connector_stats = (
-                    kv_connector_stats.aggregate(scheduler_kv_connector_stats)
-                    if kv_connector_stats is not None
-                    else scheduler_kv_connector_stats
-                )
-
-        # Collect and publish KV cache events (align with v0.14.0)
-        events = self.kv_cache_manager.take_events()
-        if self.connector is not None:
-            connector_events = self.connector.take_events()
-            if connector_events:
-                if events is None:
-                    events = list(connector_events)
-                else:
-                    events.extend(connector_events)
-        if events:
-            batch = KVEventBatch(ts=time.time(), events=events)
-            self.kv_event_publisher.publish(batch)
+        kv_connector_stats = self._aggregate_kv_connector_stats(kv_connector_output)
+        self._publish_kv_cache_events()
 
         # Create EngineCoreOutputs for all clients that have requests with
         # outputs in this step.
         engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()}
 
-        finished_req_ids = self.finished_req_ids_dict
-        if finished_req_ids:
-            # Include ids of requests that finished since last outputs
-            # were sent.
-            for client_index, finished_set in finished_req_ids.items():
-                # Set finished request set in EngineCoreOutputs for this client.
-                if (eco := engine_core_outputs.get(client_index)) is not None:
-                    eco.finished_requests = finished_set
-                else:
-                    engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
-            finished_req_ids.clear()
+        self._attach_finished_request_sets(
+            engine_core_outputs,
+            synthesize_abort_outputs=False,
+        )
 
-        if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats)) is not None:
-            # Return stats to only one of the front-ends.
-            if (eco := next(iter(engine_core_outputs.values()), None)) is None:
-                # We must return the stats even if there are no request
-                # outputs this step.
-                engine_core_outputs[0] = eco = EngineCoreOutputs()
-            eco.scheduler_stats = stats
+        self._attach_scheduler_stats(
+            engine_core_outputs,
+            spec_decoding_stats,
+            kv_connector_stats,
+            cudagraph_stats,
+            perf_stats,
+        )
 
         self._capture_omni_connector_output(model_runner_output)
 

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -5,13 +5,36 @@ import time
 from collections.abc import Iterable
 from typing import Any
 
+from vllm.compilation.cuda_graph import CUDAGraphStat
+from vllm.distributed.kv_events import KVEventBatch
+from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
 from vllm.logger import init_logger
 from vllm.v1.core.sched.output import SchedulerOutput
-from vllm.v1.engine import EngineCoreEventType
+from vllm.v1.core.sched.utils import remove_all
+from vllm.v1.engine import (
+    EngineCoreEventType,
+    EngineCoreOutput,
+    EngineCoreOutputs,
+    FinishReason,
+)
+from vllm.v1.metrics.perf import PerfStats
 from vllm.v1.metrics.stats import SchedulerStats
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
+from vllm.v1.spec_decode.metrics import SpecDecodingStats
 
-from vllm_omni.core.sched.output import OmniChunkRecvHandle, OmniSchedulerOutput
+from vllm_omni.core.sched.omni_scheduling_coordinator import (
+    OmniSchedulingCoordinator,
+    uses_full_payload_input_coordinator,
+)
+from vllm_omni.core.sched.output import (
+    OmniChunkRecvHandle,
+    OmniNewRequestData,
+    OmniSchedulerOutput,
+)
+from vllm_omni.distributed.omni_connectors.transfer_adapter.chunk_transfer_adapter import (
+    OmniChunkTransferAdapter,
+)
+from vllm_omni.engine import OmniEngineCoreOutput
 
 logger = init_logger(__name__)
 
@@ -45,6 +68,19 @@ class OmniSchedulerMixin:
     #  Shared scheduler/output helpers (lift the AR / generation duplicates)
     # ------------------------------------------------------------------ #
 
+    def _init_omni_io_scheduling_state(self) -> None:
+        """Initialize scheduler state shared by AR and generation stages."""
+        model_config = self.vllm_config.model_config
+        self.chunk_transfer_adapter = (
+            OmniChunkTransferAdapter(self.vllm_config) if getattr(model_config, "async_chunk", False) else None
+        )
+        self.input_coordinator = (
+            OmniSchedulingCoordinator(stage_id=getattr(model_config, "stage_id", 0))
+            if uses_full_payload_input_coordinator(model_config)
+            else None
+        )
+        self._latest_omni_connector_output = None
+
     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)
@@ -102,6 +138,28 @@ class OmniSchedulerMixin:
             connector_output.stage_recv_req_ids if connector_output else set(),
         )
 
+    def _process_pending_omni_inputs(self, model_mode: str) -> None:
+        """Apply pending connector inputs, timeouts, and async chunks."""
+        self._consume_pending_connector_output(model_mode)
+        self._process_pending_input_timeouts()
+        if self.chunk_transfer_adapter:
+            self.chunk_transfer_adapter.process_pending_chunks(
+                self.waiting,
+                self.running,
+                scheduler_requests=self.requests,
+            )
+
+    def _restore_omni_wait_queues(self) -> None:
+        """Restore requests temporarily parked by Omni input gates."""
+        if self.chunk_transfer_adapter:
+            self.chunk_transfer_adapter.restore_queues(
+                self.waiting,
+                self.running,
+                scheduler_requests=self.requests,
+            )
+        if self.input_coordinator:
+            self.input_coordinator.restore_queues(self.waiting)
+
     def _process_pending_input_timeouts(self) -> None:
         """Force-fail requests waiting on the full-payload coordinator too long.
 
@@ -181,6 +239,201 @@ class OmniSchedulerMixin:
             pending_input_registrations=pending_input_registrations,
         )
 
+    def _rewrap_scheduled_new_reqs(self, scheduler_output: SchedulerOutput) -> None:
+        """Attach Omni payloads without reconstructing existing Omni entries."""
+        scheduler_output.scheduled_new_reqs = [  # type: ignore[assignment]
+            data
+            if isinstance(data, OmniNewRequestData)
+            else OmniNewRequestData.from_base(data, self.requests.get(data.req_id))
+            for data in scheduler_output.scheduled_new_reqs
+        ]
+
+    def _postprocess_omni_schedule_output(
+        self,
+        scheduler_output: SchedulerOutput,
+        *,
+        include_cached_payloads: bool = False,
+    ) -> None:
+        """Enrich new requests and apply async-chunk output bookkeeping."""
+        self._rewrap_scheduled_new_reqs(scheduler_output)
+        if not self.chunk_transfer_adapter:
+            return
+        if include_cached_payloads:
+            self.chunk_transfer_adapter.postprocess_scheduler_output(
+                scheduler_output,
+                self.requests,
+            )
+        else:
+            self.chunk_transfer_adapter.postprocess_scheduler_output(scheduler_output)
+
+    def _make_omni_engine_output(
+        self,
+        request: Request,
+        *,
+        new_token_ids: list[int],
+        finish_reason: Any = None,
+        new_logprobs: Any = None,
+        new_prompt_logprobs_tensors: Any = None,
+        pooling_output: Any = None,
+        multimodal_output: Any = None,
+        stop_reason: Any = None,
+        prefill_stats: Any = None,
+        kv_transfer_params: Any = None,
+        routed_experts: Any = None,
+        num_nans_in_logits: int = 0,
+        is_segment_finished: bool | None = False,
+        new_prompt_len_snapshot: int | None = None,
+    ) -> OmniEngineCoreOutput:
+        """Build the common request-output envelope used by LLM schedulers."""
+        return OmniEngineCoreOutput(
+            request_id=request.request_id,
+            new_token_ids=new_token_ids,
+            finish_reason=finish_reason,
+            new_logprobs=new_logprobs,
+            new_prompt_logprobs_tensors=new_prompt_logprobs_tensors,
+            pooling_output=pooling_output,
+            multimodal_output=multimodal_output,
+            stop_reason=stop_reason,
+            events=request.take_events(),
+            prefill_stats=prefill_stats,
+            kv_transfer_params=kv_transfer_params,
+            trace_headers=request.trace_headers,
+            routed_experts=routed_experts,
+            num_nans_in_logits=num_nans_in_logits,
+            is_segment_finished=is_segment_finished,
+            new_prompt_len_snapshot=new_prompt_len_snapshot,
+        )
+
+    def _append_request_output(
+        self,
+        outputs: dict[int, list[EngineCoreOutput]],
+        request: Request,
+        **output_fields: Any,
+    ) -> None:
+        outputs[request.client_index].append(
+            OmniSchedulerMixin._make_omni_engine_output(
+                self,
+                request,
+                **output_fields,
+            )
+        )
+
+    def _handle_failed_kv_load_outputs(
+        self,
+        failed_request_ids: set[str] | None,
+        outputs: dict[int, list[EngineCoreOutput]],
+    ) -> list[Request]:
+        """Finish unrecoverable KV loads and emit their terminal outputs."""
+        if not failed_request_ids or self.recompute_kv_load_failures:
+            return []
+        requests = [self.requests[req_id] for req_id in failed_request_ids]
+        self.finish_requests(failed_request_ids, RequestStatus.FINISHED_ERROR)
+        for request in requests:
+            OmniSchedulerMixin._append_request_output(
+                self,
+                outputs,
+                request,
+                new_token_ids=[],
+                finish_reason=request.get_finished_reason(),
+            )
+        return requests
+
+    def _attach_finished_request_sets(
+        self,
+        engine_core_outputs: dict[int, EngineCoreOutputs],
+        *,
+        synthesize_abort_outputs: bool,
+    ) -> None:
+        """Attach finished IDs while keeping AR's synthetic-abort policy explicit."""
+        finished_req_ids = self.finished_req_ids_dict
+        if not finished_req_ids:
+            return
+        for client_index, finished_set in finished_req_ids.items():
+            output = engine_core_outputs.get(client_index)
+            if output is None:
+                output = EngineCoreOutputs()
+                engine_core_outputs[client_index] = output
+            if synthesize_abort_outputs:
+                emitted = {item.request_id for item in output.outputs}
+                output.outputs.extend(
+                    EngineCoreOutput(req_id, [], finish_reason=FinishReason.ABORT)
+                    for req_id in finished_set
+                    if req_id not in emitted
+                )
+            output.finished_requests = finished_set
+        finished_req_ids.clear()
+
+    def _remove_stopped_requests_from_queues(
+        self,
+        stopped_running_reqs: set[Request],
+        stopped_preempted_reqs: set[Request],
+    ) -> None:
+        if stopped_running_reqs:
+            self.running = remove_all(self.running, stopped_running_reqs)
+        if stopped_preempted_reqs:
+            self.waiting.remove_requests(stopped_preempted_reqs)
+            self.skipped_waiting.remove_requests(stopped_preempted_reqs)
+
+    def _aggregate_kv_connector_stats(
+        self,
+        kv_connector_output: Any,
+    ) -> KVConnectorStats | None:
+        stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
+        if self.connector:
+            scheduler_stats = self.connector.get_kv_connector_stats()
+            if scheduler_stats is not None and not scheduler_stats.is_empty():
+                stats = stats.aggregate(scheduler_stats) if stats is not None else scheduler_stats
+        return stats
+
+    def _publish_kv_cache_events(self) -> None:
+        events = self.kv_cache_manager.take_events()
+        if self.connector is not None:
+            connector_events = self.connector.take_events()
+            if connector_events:
+                if events is None:
+                    events = list(connector_events)
+                else:
+                    events.extend(connector_events)
+        if events:
+            self.kv_event_publisher.publish(KVEventBatch(ts=time.time(), events=events))
+
+    def _attach_scheduler_stats(
+        self,
+        engine_core_outputs: dict[int, EngineCoreOutputs],
+        spec_decoding_stats: SpecDecodingStats | None,
+        kv_connector_stats: KVConnectorStats | None,
+        cudagraph_stats: CUDAGraphStat | None,
+        perf_stats: PerfStats | None,
+    ) -> None:
+        stats = self.make_stats(
+            spec_decoding_stats,
+            kv_connector_stats,
+            cudagraph_stats,
+            perf_stats,
+        )
+        if stats is None:
+            return
+        if (output := next(iter(engine_core_outputs.values()), None)) is None:
+            engine_core_outputs[0] = output = EngineCoreOutputs()
+        output.scheduler_stats = stats
+
+    def finish_requests(
+        self,
+        request_ids: str | Iterable[str] | None,
+        finished_status: RequestStatus,
+    ) -> list[Request]:
+        """Finish requests and clean all Omni-owned queue/coordinator state."""
+        if self.chunk_transfer_adapter:
+            self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
+
+        self._realign_request_status_to_queues(request_ids)
+        finished = super().finish_requests(request_ids, finished_status)
+        self._purge_finished_from_running()
+
+        for request in finished:
+            self._free_input_coordinator_request(request.request_id)
+        return finished
+
     def make_stats(self, *args, **kwargs) -> SchedulerStats | None:
         now = time.monotonic()
         if now - getattr(self, "_last_stats_time", 0.0) < _STATS_INTERVAL_S:

vllm_omni/core/sched/output.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, fields
 
 from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
 from vllm.v1.request import Request
@@ -28,6 +28,21 @@ class OmniNewRequestData(NewRequestData):
     additional_information: AdditionalInformationPayload | None = None
     model_intermediate_buffer: dict[str, object] | None = None
 
+    @classmethod
+    def from_base(
+        cls,
+        data: NewRequestData,
+        request: Request | None,
+    ) -> "OmniNewRequestData":
+        """Preserve upstream request data while attaching Omni payloads."""
+        base_data = {field.name: getattr(data, field.name) for field in fields(NewRequestData)}
+        return cls(
+            **base_data,
+            external_req_id=getattr(request, "external_req_id", None),
+            additional_information=getattr(request, "additional_information", None),
+            model_intermediate_buffer=getattr(request, "model_intermediate_buffer", None),
+        )
+
     @classmethod
     def from_request(
         cls,