234 lines
11 KiB
Python
234 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
from collections import defaultdict
|
|
|
|
from vllm.logger import init_logger
|
|
from vllm.v1.core.sched.output import SchedulerOutput
|
|
from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler
|
|
from vllm.v1.core.sched.utils import check_stop, remove_all
|
|
from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs
|
|
from vllm.v1.outputs import ModelRunnerOutput
|
|
from vllm.v1.request import Request, RequestStatus
|
|
from vllm.v1.spec_decode.metrics import SpecDecodingStats
|
|
|
|
|
|
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.
|
|
"""
|
|
|
|
# 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_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),
|
|
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
|
|
|
|
def update_from_output(
|
|
self,
|
|
scheduler_output: SchedulerOutput,
|
|
model_runner_output: ModelRunnerOutput,
|
|
) -> dict[int, EngineCoreOutputs]:
|
|
"""Override update_from_output to fix check_stop bug when output_token_ids is empty.
|
|
|
|
The original vLLM implementation calls check_stop when pooler_outputs exist,
|
|
but doesn't check if output_token_ids is empty. This causes IndexError during
|
|
prefill phase when pooler_outputs exist but output_token_ids is still empty.
|
|
|
|
This override adds a safety check before calling check_stop.
|
|
"""
|
|
sampled_token_ids = model_runner_output.sampled_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
|
|
kv_connector_output = model_runner_output.kv_connector_output
|
|
|
|
outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
|
|
spec_decoding_stats: SpecDecodingStats | None = None
|
|
kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else 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_draft_tokens = len(scheduled_spec_token_ids)
|
|
num_accepted = len(generated_token_ids) - 1
|
|
num_rejected = num_draft_tokens - num_accepted
|
|
# 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.
|
|
request.num_computed_tokens -= num_rejected
|
|
spec_decoding_stats = self.make_spec_decoding_stats(
|
|
spec_decoding_stats, num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted
|
|
)
|
|
|
|
stopped = False
|
|
new_logprobs = None
|
|
new_token_ids = generated_token_ids
|
|
kv_transfer_params = None
|
|
status_before_stop = request.status
|
|
|
|
# Check for stop and update request status.
|
|
if new_token_ids:
|
|
new_token_ids, stopped = self._update_request_with_output(request, new_token_ids)
|
|
|
|
"""
|
|
# Stop checking for pooler models.
|
|
pooler_output = None
|
|
if pooler_outputs:
|
|
pooler_output = pooler_outputs[req_index]
|
|
stopped = check_stop(request, self.max_model_len,
|
|
pooler_output)
|
|
"""
|
|
|
|
# Stop checking for pooler models.
|
|
# FIX: Check if output_token_ids is empty before calling check_stop
|
|
# to prevent IndexError during prefill phase when pooler_outputs exist
|
|
# but output_token_ids is still empty.
|
|
pooler_output = None
|
|
if pooler_outputs:
|
|
pooler_output = pooler_outputs[req_index]
|
|
# Only call check_stop if output_token_ids is not empty
|
|
# This prevents IndexError when accessing output_token_ids[-1]
|
|
if request.output_token_ids:
|
|
stopped = check_stop(request, self.max_model_len, pooler_output)
|
|
# If output_token_ids is empty, we don't call check_stop
|
|
# and keep stopped=False (don't stop during prefill phase)
|
|
|
|
if stopped:
|
|
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
|
|
# checked above, so safe to ignore type warning
|
|
request.structured_output_request.grammar.accept_tokens( # type: ignore[union-attr]
|
|
req_id, new_token_ids
|
|
)
|
|
|
|
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]
|
|
|
|
# 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,
|
|
trace_headers=request.trace_headers,
|
|
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 = 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)
|
|
|
|
# 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 (stats := self.make_stats(spec_decoding_stats, kv_connector_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
|
|
|
|
return engine_core_outputs
|