diff --git a/vllm_omni/core/sched/__init__.py b/vllm_omni/core/sched/__init__.py index e69de29bb..3ef88c25d 100644 --- a/vllm_omni/core/sched/__init__.py +++ b/vllm_omni/core/sched/__init__.py @@ -0,0 +1,14 @@ +""" +Scheduling components for vLLM-omni. +""" + +from .scheduler import OmniScheduler +from .diffusion_scheduler import DiffusionScheduler +from .output import OmniNewRequestData + +__all__ = [ + "OmniScheduler", + "DiffusionScheduler", + "OmniNewRequestData", +] + diff --git a/vllm_omni/core/sched/diffusion_scheduler.py b/vllm_omni/core/sched/diffusion_scheduler.py new file mode 100644 index 000000000..c076fde7b --- /dev/null +++ b/vllm_omni/core/sched/diffusion_scheduler.py @@ -0,0 +1,315 @@ +from vllm.v1.core.sched.scheduler import SchedulerOutput, EngineCoreOutputs, Request, RequestStatus, SpecDecodingStats +from collections import defaultdict +from typing import Optional +from vllm.v1.core.sched.request_queue import create_request_queue +from vllm.v1.engine import EngineCoreEventType +from vllm.distributed.kv_events import KVEventBatch +import time +from vllm_omni.core.sched.scheduler import OmniScheduler +from vllm_omni.outputs import OmniModelRunnerOutput +from vllm.v1.engine import EngineCoreOutput +from vllm_omni.core.sched.output import OmniNewRequestData + + +class DiffusionScheduler(OmniScheduler): + def schedule(self) -> SchedulerOutput: + """Diffusion fast path: + - Feed all input tokens of the request at once (if 0, allocate 1 placeholder token). + - If the token budget cannot be satisfied at once, fall back to the default vLLM scheduling. + """ + + # Select requests with zero prompt and using pooling + # (diffusion results returned via pooler_output) + token_budget = self.max_num_scheduled_tokens + capacity = self.max_num_running_reqs - len(self.running) + scheduled_timestamp = time.monotonic() + + scheduled_new_reqs: list[Request] = [] + scheduled_resumed_reqs: list[Request] = [] + scheduled_running_reqs: list[Request] = [] + + req_to_new_block_ids: dict[str, tuple[list[int], ...]] = {} + num_scheduled_tokens: dict[str, int] = {} + scheduled_spec_decode_tokens: dict[str, list[int]] = {} + scheduled_encoder_inputs: dict[str, list[int]] = {} + structured_output_request_ids: dict[str, int] = {} + + # Temporary queue: preserve waiting order, do not disturb non-diffusion requests + skipped_waiting_requests = create_request_queue(self.policy) + + # Fast path selection and scheduling (treat all as diffusion requests, + # independent of pooling_params) + while self.waiting and token_budget > 0 and capacity > 0: + request = self.waiting.peek_request() + # Uniformly treat as diffusion. A feature flag can be added later via config or request tag. + + # Allocate all input tokens for the request in one shot + # (allocate 1 placeholder if zero) + required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1) + if required_tokens > token_budget: + # Insufficient budget to process all inputs at once; stop fast path attempt + break + num_new_tokens = required_tokens + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_new_tokens, + num_lookahead_tokens=self.num_lookahead_tokens, + ) + if new_blocks is None: + # Allocation failed (e.g., VRAM pressure); stop fast path and + # fall back to default scheduling + # Put the current request back to the head of the waiting queue + # Note: the original queue order is preserved + break + + # Officially schedule this request + request = self.waiting.pop_request() + self.running.append(request) + request.status = RequestStatus.RUNNING + if self.log_stats: + request.record_event(EngineCoreEventType.SCHEDULED, + scheduled_timestamp) + + req_to_new_block_ids[request.request_id] = new_blocks.get_block_ids() + num_scheduled_tokens[request.request_id] = num_new_tokens + token_budget -= num_new_tokens + capacity -= 1 + scheduled_new_reqs.append(request) + + # Return skipped waiting requests + if skipped_waiting_requests: + self.waiting.prepend_requests(skipped_waiting_requests) + + # If fast path scheduled none, fall back to the original scheduling + if not num_scheduled_tokens: + return super().schedule() + + # Compute common prefix blocks (aligned with v1) + num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) + if self.running: + any_request = self.running[0] + num_common_prefix_blocks = ( + self.kv_cache_manager.get_num_common_prefix_blocks( + any_request, len(self.running))) + + grammar_bitmask = self.structured_output_manager.grammar_bitmask( + self.requests, + structured_output_request_ids, + scheduled_spec_decode_tokens, + ) + + # Assemble SchedulerOutput + new_reqs_data = [ + OmniNewRequestData.from_request(req, req_to_new_block_ids[req.request_id]) + for req in scheduled_new_reqs + ] + cached_reqs_data = self._make_cached_request_data( + scheduled_running_reqs, + scheduled_resumed_reqs, + num_scheduled_tokens, + scheduled_spec_decode_tokens, + req_to_new_block_ids, + ) + + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) + scheduler_output = SchedulerOutput( + scheduled_new_reqs=new_reqs_data, + scheduled_cached_reqs=cached_reqs_data, + num_scheduled_tokens=num_scheduled_tokens, + total_num_scheduled_tokens=total_num_scheduled_tokens, + scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, + scheduled_encoder_inputs=scheduled_encoder_inputs, + num_common_prefix_blocks=num_common_prefix_blocks, + finished_req_ids=self.finished_req_ids, + free_encoder_input_ids=self.encoder_cache_manager.get_freed_ids(), + structured_output_request_ids=structured_output_request_ids, + grammar_bitmask=grammar_bitmask, + ) + + # KVTransfer: package metadata + if self.connector is not None: + meta = self.connector.build_connector_meta(scheduler_output) + scheduler_output.kv_connector_metadata = meta + + # Publish KV events (aligned with v1) + events = self.kv_cache_manager.take_events() + if events: + batch = KVEventBatch(ts=time.time(), events=events) + self.kv_event_publisher.publish(batch) + + # Update internal state (advance num_computed_tokens, free encoder inputs, etc.) + self._update_after_schedule(scheduler_output) + return scheduler_output + """ + Scheduler for the diffusion model. + This scheduler is modified to stop the request immediately for the diffusion model. + This is because the diffusion model can generate the final image/audio in one step. + Note: This is just a minimal modification to the original scheduler, and there should be some further efforts to optimize the scheduler. + The original scheduler is still used for the AR model. + """ + def update_from_output( + self, + scheduler_output: SchedulerOutput, + model_runner_output: OmniModelRunnerOutput, + ) -> dict[int, EngineCoreOutputs]: + """Update the scheduler state based on the model runner output. + + This method is modified to stop the request immediately for the diffusion model. + """ + sampled_token_ids = model_runner_output.sampled_token_ids + spec_token_ids = model_runner_output.spec_token_ids + logprobs = model_runner_output.logprobs + prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict + num_scheduled_tokens = scheduler_output.num_scheduled_tokens + pooler_outputs = model_runner_output.pooler_output + num_nans_in_logits = model_runner_output.num_nans_in_logits + + outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) + spec_decoding_stats: Optional[SpecDecodingStats] = None + + # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more, + # the below loop can be a performance bottleneck. We should do our best + # to avoid expensive operations inside the loop. + stopped_running_reqs: set[Request] = set() + stopped_preempted_reqs: set[Request] = set() + for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): + assert num_tokens_scheduled > 0 + request = self.requests.get(req_id) + if request is None: + # The request is already finished. This can happen if the + # request is aborted while the model is executing it (e.g., + # in pipeline parallelism). + continue + + req_index = model_runner_output.req_id_to_index[req_id] + generated_token_ids = sampled_token_ids[ + req_index] if sampled_token_ids else [] + + scheduled_spec_token_ids = ( + scheduler_output.scheduled_spec_decode_tokens.get(req_id)) + if scheduled_spec_token_ids: + # num_computed_tokens represents the number of tokens + # processed in the current step, considering scheduled + # tokens and rejections. If some tokens are rejected, + # num_computed_tokens is decreased by the number of rejected + # tokens, where is given by: + # len(scheduled_spec_token_ids) + 1 - len(generated_token_ids). + num_tokens_rejected = (len(scheduled_spec_token_ids) + 1 - + len(generated_token_ids)) + request.num_computed_tokens -= num_tokens_rejected + spec_decoding_stats = self.make_spec_decoding_stats( + spec_decoding_stats, + num_draft_tokens=len(scheduled_spec_token_ids), + num_accepted_tokens=len(generated_token_ids) - 1) + + new_logprobs = None + new_token_ids = generated_token_ids + kv_transfer_params = None + status_before_stop = request.status + pooler_output = None + if pooler_outputs: + pooler_output = pooler_outputs[req_index] + + # Diffusion request: completes in one step; mark finished and free resources + request.status = RequestStatus.FINISHED_STOPPED + # Optional: set a stop_reason for front-end clarity (does not affect protocol) + request.stop_reason = request.stop_reason or "diffusion_done" + kv_transfer_params = self._free_request(request) + if status_before_stop == RequestStatus.RUNNING: + stopped_running_reqs.add(request) + else: + stopped_preempted_reqs.add(request) + + # Extract sample logprobs if needed. + if request.sampling_params is not None \ + and request.sampling_params.logprobs is not None and logprobs: + # NOTE: once we support N tokens per step (spec decode), + # the outer lists can be of length > 1. + new_logprobs = logprobs.slice(req_index, req_index + 1) + + if new_token_ids and self.structured_output_manager.should_advance( + request): + # NOTE: structured_output_request + # should not be None if use_structured_output, we have + # check above, so safe to ignore type warning + request.structured_output_request.grammar.accept_tokens( # type: ignore[union-attr] + req_id, new_token_ids) + + # spec_token_ids comes from the model runner output + if num_nans_in_logits is not None and req_id in num_nans_in_logits: + request.num_nans_in_logits = num_nans_in_logits[req_id] + + # Add newly generated spec token ids to the request. + if spec_token_ids is not None: + if self.structured_output_manager.should_advance(request): + metadata = request.structured_output_request + # Needs to happen after new_token_ids are accepted. + request.spec_token_ids = metadata.grammar.validate_tokens( # type: ignore[union-attr] + spec_token_ids[req_index]) + else: + request.spec_token_ids = spec_token_ids[req_index] + + # Get prompt logprobs for this request. + prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) + if new_token_ids or pooler_output is not None \ + or kv_transfer_params: + + # Add EngineCoreOutput for this Request. + outputs[request.client_index].append( + EngineCoreOutput( + request_id=req_id, + new_token_ids=new_token_ids, + finish_reason=request.get_finished_reason(), + new_logprobs=new_logprobs, + new_prompt_logprobs_tensors=prompt_logprobs_tensors, + pooling_output=pooler_output, + stop_reason=request.stop_reason, + events=request.take_events(), + kv_transfer_params=kv_transfer_params, + num_cached_tokens=request.num_cached_tokens, + )) + + else: + # Invariant: EngineCore returns no partial prefill outputs. + assert not prompt_logprobs_tensors + + # Remove the stopped requests from the running and waiting queues. + if stopped_running_reqs: + self.running = [ + req for req in self.running if req not in stopped_running_reqs + ] + if stopped_preempted_reqs: + # This is a rare case and unlikely to impact performance. + self.waiting.remove_requests(stopped_preempted_reqs) + + # KV Connector: update state for finished KV Transfers. + if model_runner_output.kv_connector_output: + self._update_from_kv_xfer_finished( + model_runner_output.kv_connector_output) + + # 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() + + if engine_core_outputs: + # Return stats to only one of the front-ends. + next(iter(engine_core_outputs.values())).scheduler_stats = ( + self.make_stats(spec_decoding_stats)) + + return engine_core_outputs \ No newline at end of file diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py new file mode 100644 index 000000000..45e29239e --- /dev/null +++ b/vllm_omni/core/sched/output.py @@ -0,0 +1,36 @@ +from dataclasses import dataclass +from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.request import Request +from typing import Optional + +from vllm_omni.engine import PromptEmbedsPayload, AdditionalInformationPayload + + + +@dataclass +class OmniNewRequestData(NewRequestData): + # Optional serialized prompt embeddings + prompt_embeds: Optional[PromptEmbedsPayload] = None + # Optional serialized additional information + additional_information: Optional[AdditionalInformationPayload] = None + + @classmethod + def from_request( + cls, + request: Request, + block_ids: tuple[list[int], ...], + ) -> "OmniNewRequestData": + return cls( + req_id=request.request_id, + prompt_token_ids=request.prompt_token_ids, + mm_inputs=request.mm_inputs, + mm_hashes=request.mm_hashes, + mm_positions=request.mm_positions, + sampling_params=request.sampling_params, + pooling_params=request.pooling_params, + block_ids=block_ids, + num_computed_tokens=request.num_computed_tokens, + lora_request=request.lora_request, + prompt_embeds=request.prompt_embeds, + additional_information=request.additional_information, + ) \ No newline at end of file diff --git a/vllm_omni/core/sched/scheduler.py b/vllm_omni/core/sched/scheduler.py index 5bbbfeb05..c2ab49b6b 100644 --- a/vllm_omni/core/sched/scheduler.py +++ b/vllm_omni/core/sched/scheduler.py @@ -1,40 +1,48 @@ +from __future__ import annotations +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler +from vllm.logger import init_logger -from vllm_omni.request import OmniRequest -from typing import List -from threading import Lock, Condition -from vllm_omni.config import OmniConfig - -from vllm.v1.kv_cache_interface import KVCacheConfig -from vllm.v1.core.sched import SchedulerInterface -from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry -from vllm.v1.structured_output import StructuredOutputManager - - -class OmniScheduler(SchedulerInterface): +class OmniScheduler(VLLMScheduler): """ OmniScheduler: Scheduler for vLLM-omni multimodal processing. This scheduler extends vLLM's scheduler to support multimodal and non-autoregressive processing with additional fields and methods specific to vLLM-omni. """ - def __init__(self, - omni_config: OmniConfig, - kv_cache_config: KVCacheConfig, - structured_output_manager: StructuredOutputManager, - mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, - include_finished_set: bool = False, - log_stats: bool = False, - ): - super().__init__( - vllm_config=omni_config.vllm_config, - kv_cache_config=kv_cache_config, - multimodal_registry=mm_registry, - structured_output_manager=structured_output_manager, - include_finished_set=include_finished_set, - log_stats=log_stats, - ) - self.omni_config = omni_config - - def schedule(self, requests: List[OmniRequest]) -> List[OmniRequest]: - # TODO: Implement scheduling logic - pass \ No newline at end of file + + # Ensure scheduled_new_reqs carry omni-specific payloads (e.g., additional_information) + def schedule(self) -> SchedulerOutput: # type: ignore[override] + scheduler_output = super().schedule() + try: + # Late import to avoid circulars in some launch modes + from .output import OmniNewRequestData + # 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, + prompt_token_ids=nr.prompt_token_ids, + mm_inputs=nr.mm_inputs, + mm_hashes=nr.mm_hashes, + mm_positions=nr.mm_positions, + 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, + additional_information=getattr(request, "additional_information", None) if request else None, + ) + new_list.append(omni_nr) + + scheduler_output.scheduled_new_reqs = new_list # type: ignore[assignment] + except Exception: + # If anything goes wrong, leave the original output unchanged + init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData") + + return scheduler_output \ No newline at end of file