diff --git a/tests/core/sched/test_omni_ar_scheduler_logprobs.py b/tests/core/sched/test_omni_ar_scheduler_logprobs.py index 92bb30a27..b8dc2ee9e 100644 --- a/tests/core/sched/test_omni_ar_scheduler_logprobs.py +++ b/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}, diff --git a/tests/core/sched/test_omni_scheduler_mixin_shared.py b/tests/core/sched/test_omni_scheduler_mixin_shared.py new file mode 100644 index 000000000..9aaa4c656 --- /dev/null +++ b/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 diff --git a/tests/entrypoints/test_omni_new_request_data.py b/tests/entrypoints/test_omni_new_request_data.py index b1ad56cdd..abbb04d8f 100644 --- a/tests/entrypoints/test_omni_new_request_data.py +++ b/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 diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 514cc9b94..5a36d969a 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/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. diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 47e6bb4d2..8a165e2a1 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/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) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index 684b1f260..7b0b37e7e 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/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: diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index be9c867b1..b73fbd6aa 100644 --- a/vllm_omni/core/sched/output.py +++ b/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,