From 19220aae82f2f28f97d40c34793480b179061173 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Wed, 22 Oct 2025 15:05:23 +0800 Subject: [PATCH 1/7] Add scheduling components for vLLM-omni - Introduced `DiffusionScheduler` to handle scheduling for diffusion models, allowing immediate request completion. - Added `OmniNewRequestData` to encapsulate additional request data, including prompt embeddings and extra information. - Updated `OmniScheduler` to ensure compatibility with new request data structures. - Created `output.py` to define new data classes for enriched request handling. - Enhanced `__init__.py` in the engine module to support new payload structures for prompt embeddings and additional information. --- vllm_omni/core/sched/__init__.py | 14 + vllm_omni/core/sched/diffusion_scheduler.py | 315 ++++++++++++++++++++ vllm_omni/core/sched/output.py | 36 +++ vllm_omni/core/sched/scheduler.py | 78 +++-- vllm_omni/engine/__init__.py | 54 +++- 5 files changed, 463 insertions(+), 34 deletions(-) create mode 100644 vllm_omni/core/sched/diffusion_scheduler.py create mode 100644 vllm_omni/core/sched/output.py 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..104beec02 --- /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, defaultdict, 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: + """扩散快速通道： + - 一次性喂入该请求的全部输入 token（若为 0，则分配 1 个占位 token）。 + - 若无法一次性满足 token 预算，则退回上游 vLLM 的默认调度。 + """ + + # 选出零 prompt 且使用 pooling（扩散结果经 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] = {} + + # 临时队列：保持等待队列顺序，不破坏非扩散请求 + skipped_waiting_requests = create_request_queue(self.policy) + + # 快速通道挑选并调度（所有请求都视为扩散请求，不依赖 pooling_params） + while self.waiting and token_budget > 0 and capacity > 0: + request = self.waiting.peek_request() + # 统一按扩散处理。若未来需要条件开关，可接入配置或请求标记。 + is_diffusion = True + if not is_diffusion: + # 暂存到跳过队列，稍后归还到等待队列头部 + self.waiting.pop_request() + skipped_waiting_requests.prepend_request(request) + continue + + # 一次性为该请求分配全部输入 token（若为 0，则分配 1 个占位 token） + required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1) + if required_tokens > token_budget: + # 无足够预算一次性完成该请求的输入处理，停止快速通道尝试 + 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: + # 无法分配（显存紧张等），停止快速通道尝试，回退到默认调度 + # 将当前 request 放回等待队列头 + # 注意：此处不改变原队列顺序 + break + + # 正式调度该请求 + 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) + + # 归还被跳过的等待请求 + if skipped_waiting_requests: + self.waiting.prepend_requests(skipped_waiting_requests) + + # 若快速通道未调度任何请求，则回退到原始调度逻辑 + if not num_scheduled_tokens: + return super().schedule() + + # 计算公共前缀块（与 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, + ) + + # 组装 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：封装元信息 + if self.connector is not None: + meta = self.connector.build_connector_meta(scheduler_output) + scheduler_output.kv_connector_metadata = meta + + # 发布 KV 事件（与 v1 对齐） + events = self.kv_cache_manager.take_events() + if events: + batch = KVEventBatch(ts=time.time(), events=events) + self.kv_event_publisher.publish(batch) + + # 更新内部状态（推进 num_computed_tokens，释放 encoder 输入等） + 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] + + # 扩散请求：单步完成，直接标记完成并释放资源 + request.status = RequestStatus.FINISHED_STOPPED + # 可选：标注停止原因，便于前端区分（不影响协议） + 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..bfb395fcc --- /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], ...], + ) -> NewRequestData: + 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..2e9ba59b4 100644 --- a/vllm_omni/core/sched/scheduler.py +++ b/vllm_omni/core/sched/scheduler.py @@ -1,40 +1,58 @@ +from __future__ import annotations -from vllm_omni.request import OmniRequest -from typing import List -from threading import Lock, Condition -from vllm_omni.config import OmniConfig +from collections import defaultdict +from typing import Optional -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 +from vllm.v1.core.sched.output import ( + SchedulerOutput, ) +from vllm.v1.core.sched.utils import check_stop +from vllm.v1.engine import EngineCoreOutputs +from vllm.v1.request import Request, RequestStatus +from vllm.v1.spec_decode.metrics import SpecDecodingStats +from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler -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 + pass + + return scheduler_output \ No newline at end of file diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index 406c26b39..f0a906709 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -1,9 +1,55 @@ """ Engine components for vLLM-omni. """ +import msgspec +from typing import Any -from .output_processor import MultimodalOutputProcessor +from vllm.v1.engine import EngineCoreRequest +from typing import Optional -__all__ = [ - "MultimodalOutputProcessor", -] +class PromptEmbedsPayload(msgspec.Struct): + """Serialized prompt embeddings payload for direct transfer. + + data: raw bytes of the tensor in row-major order + shape: [seq_len, hidden_size] + dtype: torch dtype name (e.g., "float16", "float32") + """ + + data: bytes + shape: list[int] + dtype: str + + +class AdditionalInformationEntry(msgspec.Struct): + """One entry of additional_information. + + Two supported forms are encoded: + - tensor: data/shape/dtype + - list: a Python list (msgspec-serializable) + Exactly one of (tensor_data, list_data) should be non-None. + """ + + # Tensor form + tensor_data: Optional[bytes] = None + tensor_shape: Optional[list[int]] = None + tensor_dtype: Optional[str] = None + + # List form + list_data: Optional[list[Any]] = None + + +class AdditionalInformationPayload(msgspec.Struct): + """Serialized dictionary payload for additional_information. + + Keys are strings; values are encoded as AdditionalInformationEntry. + """ + + entries: dict[str, AdditionalInformationEntry] + + +class OmniEngineCoreRequest(EngineCoreRequest): + # Optional prompt embeddings (direct-transfer version) + prompt_embeds: Optional[PromptEmbedsPayload] = None + # Optional additional information dictionary (serialized) + additional_information: Optional[AdditionalInformationPayload] = None + \ No newline at end of file From 8d6a360f96b35daeb1646cccff6f73e713bf24e3 Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Wed, 22 Oct 2025 17:13:45 +0800 Subject: [PATCH 2/7] Update vllm_omni/core/sched/diffusion_scheduler.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Zhou Taichang --- vllm_omni/core/sched/diffusion_scheduler.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm_omni/core/sched/diffusion_scheduler.py b/vllm_omni/core/sched/diffusion_scheduler.py index 104beec02..30ba0c4e0 100644 --- a/vllm_omni/core/sched/diffusion_scheduler.py +++ b/vllm_omni/core/sched/diffusion_scheduler.py @@ -11,9 +11,9 @@ from vllm_omni.core.sched.output import OmniNewRequestData class DiffusionScheduler(OmniScheduler): def schedule(self) -> SchedulerOutput: - """扩散快速通道： - - 一次性喂入该请求的全部输入 token（若为 0，则分配 1 个占位 token）。 - - 若无法一次性满足 token 预算，则退回上游 vLLM 的默认调度。 + """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. """ # 选出零 prompt 且使用 pooling（扩散结果经 pooler_output 回传）的请求 From b207bf89f2c63da50aff0eebb318765edc6aa485 Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Wed, 22 Oct 2025 17:44:20 +0800 Subject: [PATCH 3/7] Update vllm_omni/core/sched/diffusion_scheduler.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Zhou Taichang --- vllm_omni/core/sched/diffusion_scheduler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm_omni/core/sched/diffusion_scheduler.py b/vllm_omni/core/sched/diffusion_scheduler.py index 30ba0c4e0..2c9e10c99 100644 --- a/vllm_omni/core/sched/diffusion_scheduler.py +++ b/vllm_omni/core/sched/diffusion_scheduler.py @@ -1,4 +1,6 @@ -from vllm.v1.core.sched.scheduler import SchedulerOutput, EngineCoreOutputs, Request, RequestStatus, SpecDecodingStats, defaultdict, Optional +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 From c37cd6f0e3725a2c1076d7ffcd6d613ac41fcfd5 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Fri, 24 Oct 2025 09:41:20 +0800 Subject: [PATCH 4/7] delete unused imports --- vllm_omni/core/sched/scheduler.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/vllm_omni/core/sched/scheduler.py b/vllm_omni/core/sched/scheduler.py index 2e9ba59b4..5a021949d 100644 --- a/vllm_omni/core/sched/scheduler.py +++ b/vllm_omni/core/sched/scheduler.py @@ -1,14 +1,5 @@ from __future__ import annotations - -from collections import defaultdict -from typing import Optional - -from vllm.v1.core.sched.output import ( - SchedulerOutput, ) -from vllm.v1.core.sched.utils import check_stop -from vllm.v1.engine import EngineCoreOutputs -from vllm.v1.request import Request, RequestStatus -from vllm.v1.spec_decode.metrics import SpecDecodingStats +from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler From 9ebc1dd652b2005f7e441349ff228e8b5f5c6842 Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Fri, 24 Oct 2025 10:41:54 +0800 Subject: [PATCH 5/7] Apply suggestion from @Copilot Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> Signed-off-by: Zhou Taichang --- vllm_omni/core/sched/output.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index bfb395fcc..45e29239e 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -19,7 +19,7 @@ class OmniNewRequestData(NewRequestData): cls, request: Request, block_ids: tuple[list[int], ...], - ) -> NewRequestData: + ) -> "OmniNewRequestData": return cls( req_id=request.request_id, prompt_token_ids=request.prompt_token_ids, From c23d24ca198378582090f435535b8d559ed7ceeb Mon Sep 17 00:00:00 2001 From: tzhouam Date: Fri, 24 Oct 2025 11:57:39 +0800 Subject: [PATCH 6/7] Refactor comments in diffusion_scheduler.py for clarity and consistency - Translated comments from Chinese to English for better accessibility. - Improved clarity of comments related to request handling and scheduling logic. - Ensured comments align with the code functionality and maintain a uniform style. --- vllm_omni/core/sched/diffusion_scheduler.py | 48 ++++++++++----------- 1 file changed, 23 insertions(+), 25 deletions(-) diff --git a/vllm_omni/core/sched/diffusion_scheduler.py b/vllm_omni/core/sched/diffusion_scheduler.py index 2c9e10c99..c076fde7b 100644 --- a/vllm_omni/core/sched/diffusion_scheduler.py +++ b/vllm_omni/core/sched/diffusion_scheduler.py @@ -18,7 +18,8 @@ class DiffusionScheduler(OmniScheduler): - If the token budget cannot be satisfied at once, fall back to the default vLLM scheduling. """ - # 选出零 prompt 且使用 pooling（扩散结果经 pooler_output 回传）的请求 + # 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() @@ -33,24 +34,20 @@ class DiffusionScheduler(OmniScheduler): 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) - # 快速通道挑选并调度（所有请求都视为扩散请求，不依赖 pooling_params） + # 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() - # 统一按扩散处理。若未来需要条件开关，可接入配置或请求标记。 - is_diffusion = True - if not is_diffusion: - # 暂存到跳过队列，稍后归还到等待队列头部 - self.waiting.pop_request() - skipped_waiting_requests.prepend_request(request) - continue + # Uniformly treat as diffusion. A feature flag can be added later via config or request tag. - # 一次性为该请求分配全部输入 token（若为 0，则分配 1 个占位 token） + # 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( @@ -59,12 +56,13 @@ class DiffusionScheduler(OmniScheduler): num_lookahead_tokens=self.num_lookahead_tokens, ) if new_blocks is None: - # 无法分配（显存紧张等），停止快速通道尝试，回退到默认调度 - # 将当前 request 放回等待队列头 - # 注意：此处不改变原队列顺序 + # 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 @@ -78,15 +76,15 @@ class DiffusionScheduler(OmniScheduler): 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() - # 计算公共前缀块（与 v1 对齐） + # 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] @@ -100,7 +98,7 @@ class DiffusionScheduler(OmniScheduler): scheduled_spec_decode_tokens, ) - # 组装 SchedulerOutput + # Assemble SchedulerOutput new_reqs_data = [ OmniNewRequestData.from_request(req, req_to_new_block_ids[req.request_id]) for req in scheduled_new_reqs @@ -128,18 +126,18 @@ class DiffusionScheduler(OmniScheduler): grammar_bitmask=grammar_bitmask, ) - # KVTransfer：封装元信息 + # KVTransfer: package metadata if self.connector is not None: meta = self.connector.build_connector_meta(scheduler_output) scheduler_output.kv_connector_metadata = meta - # 发布 KV 事件（与 v1 对齐） + # 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) - # 更新内部状态（推进 num_computed_tokens，释放 encoder 输入等） + # Update internal state (advance num_computed_tokens, free encoder inputs, etc.) self._update_after_schedule(scheduler_output) return scheduler_output """ @@ -212,9 +210,9 @@ class DiffusionScheduler(OmniScheduler): 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: From eb006f66e5ad96715b2e5690a90a6747087f806b Mon Sep 17 00:00:00 2001 From: tzhouam Date: Fri, 24 Oct 2025 12:00:36 +0800 Subject: [PATCH 7/7] Enhance OmniScheduler error handling with logging - Added logging for exceptions when wrapping scheduled requests in OmniScheduler. - Removed unnecessary blank lines for cleaner code structure. --- vllm_omni/core/sched/scheduler.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/vllm_omni/core/sched/scheduler.py b/vllm_omni/core/sched/scheduler.py index 5a021949d..c2ab49b6b 100644 --- a/vllm_omni/core/sched/scheduler.py +++ b/vllm_omni/core/sched/scheduler.py @@ -1,8 +1,7 @@ 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 class OmniScheduler(VLLMScheduler): """ @@ -44,6 +43,6 @@ class OmniScheduler(VLLMScheduler): scheduler_output.scheduled_new_reqs = new_list # type: ignore[assignment] except Exception: # If anything goes wrong, leave the original output unchanged - pass + init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData") return scheduler_output \ No newline at end of file