| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import enum |
| 4 | from abc import ABC, abstractmethod |
| 5 | from collections.abc import Iterable |
| 6 | from typing import TYPE_CHECKING |
| 7 | |
| 8 | from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry |
| 9 | |
| 10 | if TYPE_CHECKING: |
| 11 | from vllm.config import VllmConfig |
| 12 | from vllm.config.kv_events import KVEventsConfig |
| 13 | from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorBase |
| 14 | from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1 |
| 15 | from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput |
| 16 | from vllm.v1.engine import EngineCoreOutputs |
| 17 | from vllm.v1.kv_cache_interface import KVCacheConfig |
| 18 | from vllm.v1.metrics.stats import SchedulerStats |
| 19 | from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput |
| 20 | from vllm.v1.request import Request, RequestStatus |
| 21 | from vllm.v1.structured_output import StructuredOutputManager |
| 22 | |
| 23 | |
| 24 | class PauseState(enum.IntEnum): |
| 25 | """Scheduler pause state. |
| 26 | |
| 27 | - UNPAUSED: Normal operation |
| 28 | - PAUSE_NEW: No new requests are scheduled, requests already in |
| 29 | running state are scheduled. |
| 30 | - PAUSE_ALL: No requests are scheduled |
| 31 | """ |
| 32 | |
| 33 | UNPAUSED = 0 |
| 34 | PAUSED_NEW = 1 |
| 35 | PAUSED_ALL = 2 |
| 36 | |
| 37 | |
| 38 | class SchedulerInterface(ABC): |
| 39 | @abstractmethod |
| 40 | def __init__( |
| 41 | self, |
| 42 | vllm_config: "VllmConfig", |
| 43 | kv_cache_config: "KVCacheConfig", |
| 44 | structured_output_manager: "StructuredOutputManager", |
| 45 | block_size: int, |
| 46 | hash_block_size: int, |
| 47 | mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, |
| 48 | include_finished_set: bool = False, |
| 49 | log_stats: bool = False, |
| 50 | ) -> None: |
| 51 | raise NotImplementedError |
| 52 | |
| 53 | @abstractmethod |
| 54 | def schedule(self, throttle_prefills: bool = False) -> "SchedulerOutput": |
| 55 | """Schedule the requests to process in this scheduling step. |
| 56 | |
| 57 | The scheduling decision is made at the iteration level. Each scheduling |
| 58 | step corresponds to a single forward pass of the model. Therefore, this |
| 59 | method is called repeatedly by a busy loop in the engine. |
| 60 | |
| 61 | Essentially, the scheduler produces a dictionary of {req_id: num_tokens} |
| 62 | that specifies how many tokens to process for each request in this |
| 63 | scheduling step. For example, num_tokens can be as large as the number |
| 64 | of prompt tokens for new requests, or it can be 1 for the requests that |
| 65 | are auto-regressively generating new tokens one by one. Otherwise, it |
| 66 | can be somewhere in between in case of chunked prefills, prefix caching, |
| 67 | speculative decoding, etc. |
| 68 | |
| 69 | Additionally, the scheduler also returns useful data about each request |
| 70 | or the batch as a whole. The model runner will use this information in |
| 71 | preparing inputs to the model. |
| 72 | |
| 73 | Args: |
| 74 | throttle_prefills: DP prefill balancing. When True (set by the DP |
| 75 | engine core on non-cadence-aligned steps), new prefill compute is |
| 76 | deferred to a later step so prefills stay aligned across DP ranks; |
| 77 | automatically overridden when the rank is saturated. |
| 78 | |
| 79 | Returns: |
| 80 | A SchedulerOutput object containing information about the scheduled |
| 81 | requests. |
| 82 | |
| 83 | """ |
| 84 | raise NotImplementedError |
| 85 | |
| 86 | @abstractmethod |
| 87 | def get_grammar_bitmask( |
| 88 | self, scheduler_output: "SchedulerOutput" |
| 89 | ) -> "GrammarOutput | None": |
| 90 | raise NotImplementedError |
| 91 | |
| 92 | @abstractmethod |
| 93 | def update_from_output( |
| 94 | self, |
| 95 | scheduler_output: "SchedulerOutput", |
| 96 | model_runner_output: "ModelRunnerOutput", |
| 97 | ) -> dict[int, "EngineCoreOutputs"]: |
| 98 | """Update the scheduler state based on the model runner output. |
| 99 | |
| 100 | This method is called after the model runner has processed the scheduled |
| 101 | requests. The model runner output includes generated token ids, draft |
| 102 | token ids for next step, etc. The scheduler uses this information to |
| 103 | update its states, checks the finished requests, and returns the output |
| 104 | for each request. |
| 105 | |
| 106 | Returns: |
| 107 | A dict of client index to EngineCoreOutputs object containing the |
| 108 | outputs for each request originating from that client. |
| 109 | |
| 110 | """ |
| 111 | raise NotImplementedError |
| 112 | |
| 113 | @abstractmethod |
| 114 | def update_draft_token_ids(self, draft_token_ids: "DraftTokenIds") -> None: |
| 115 | """Update requests with newly generated draft token ids, applying |
| 116 | structured output grammar validation if needed. |
| 117 | |
| 118 | Args: |
| 119 | draft_token_ids: The input draft token ids for each request. |
| 120 | |
| 121 | """ |
| 122 | raise NotImplementedError |
| 123 | |
| 124 | @abstractmethod |
| 125 | def update_draft_token_ids_in_output( |
| 126 | self, draft_token_ids: "DraftTokenIds", scheduler_output: "SchedulerOutput" |
| 127 | ) -> None: |
| 128 | """Update scheduler output with newly generated draft token ids, applying |
| 129 | structured output grammar validation if needed. |
| 130 | |
| 131 | Args: |
| 132 | draft_token_ids: The input draft token ids for each request. |
| 133 | scheduler_output: Update the given scheduler_output |
| 134 | with the corresponding draft token ids. |
| 135 | |
| 136 | """ |
| 137 | raise NotImplementedError |
| 138 | |
| 139 | @abstractmethod |
| 140 | def add_request(self, request: "Request") -> None: |
| 141 | """Add a new request to the scheduler's internal queue. |
| 142 | |
| 143 | Args: |
| 144 | request: The new request being added. |
| 145 | |
| 146 | """ |
| 147 | raise NotImplementedError |
| 148 | |
| 149 | @abstractmethod |
| 150 | def finish_requests( |
| 151 | self, |
| 152 | request_ids: str | Iterable[str] | None, |
| 153 | finished_status: "RequestStatus", |
| 154 | ) -> "list[Request]": |
| 155 | """Finish the requests in the scheduler's internal queue. If the request |
| 156 | is not in the queue, this method will do nothing for that request. |
| 157 | |
| 158 | This method is called in two cases: |
| 159 | 1. When the request is aborted by the client. |
| 160 | 2. When the frontend process detects a stop string of the request after |
| 161 | de-tokenizing its generated tokens. |
| 162 | |
| 163 | Args: |
| 164 | request_ids: A single or a list of request IDs, or None to finish all. |
| 165 | finished_status: The finished status of the given requests. |
| 166 | |
| 167 | Returns: |
| 168 | List of requests that were aborted. Will not include any that were |
| 169 | already finished. |
| 170 | |
| 171 | """ |
| 172 | raise NotImplementedError |
| 173 | |
| 174 | @abstractmethod |
| 175 | def get_num_unfinished_requests(self) -> int: |
| 176 | """Number of unfinished requests in the scheduler's internal queue.""" |
| 177 | raise NotImplementedError |
| 178 | |
| 179 | def has_unfinished_requests(self) -> bool: |
| 180 | """Returns True if there are unfinished requests in the scheduler's |
| 181 | internal queue.""" |
| 182 | return self.get_num_unfinished_requests() > 0 |
| 183 | |
| 184 | @abstractmethod |
| 185 | def has_finished_requests(self) -> bool: |
| 186 | """Returns True if there are finished requests that need to be cleared. |
| 187 | NOTE: This is different from `not self.has_unfinished_requests()`. |
| 188 | |
| 189 | The scheduler maintains an internal list of the requests finished in the |
| 190 | previous step. This list is returned from the next call to schedule(), |
| 191 | to be sent to the model runner in the next step to clear cached states |
| 192 | for these finished requests. |
| 193 | |
| 194 | This method checks if this internal list of finished requests is |
| 195 | non-empty. This information is useful for DP attention. |
| 196 | """ |
| 197 | raise NotImplementedError |
| 198 | |
| 199 | def has_requests(self) -> bool: |
| 200 | """Returns True if there are unfinished requests, or finished requests |
| 201 | not yet returned in SchedulerOutputs.""" |
| 202 | return self.has_unfinished_requests() or self.has_finished_requests() |
| 203 | |
| 204 | @property |
| 205 | @abstractmethod |
| 206 | def pause_state(self) -> PauseState: |
| 207 | """Current pause state of the scheduler.""" |
| 208 | raise NotImplementedError |
| 209 | |
| 210 | @abstractmethod |
| 211 | def set_pause_state(self, pause_state: PauseState) -> None: |
| 212 | raise NotImplementedError |
| 213 | |
| 214 | @abstractmethod |
| 215 | def reset_prefix_cache( |
| 216 | self, reset_running_requests: bool = False, reset_connector: bool = False |
| 217 | ) -> bool: |
| 218 | """Reset the prefix cache for KV cache. |
| 219 | |
| 220 | This is particularly required when the model weights are live-updated. |
| 221 | |
| 222 | Args: |
| 223 | reset_running_requests: If True, all the running requests will be |
| 224 | preempted and moved to the waiting queue. Otherwise, this method |
| 225 | will only reset the KV prefix cache when there is no running request |
| 226 | taking KV cache. |
| 227 | reset_connector: If True, also reset any KV connector state. |
| 228 | |
| 229 | """ |
| 230 | raise NotImplementedError |
| 231 | |
| 232 | @abstractmethod |
| 233 | def reset_encoder_cache(self) -> None: |
| 234 | """Reset the encoder cache to invalidate all cached encoder outputs. |
| 235 | |
| 236 | This should be called when model weights are updated to ensure |
| 237 | stale vision embeddings are not reused. |
| 238 | """ |
| 239 | raise NotImplementedError |
| 240 | |
| 241 | @abstractmethod |
| 242 | def get_request_counts(self) -> tuple[int, int]: |
| 243 | """Returns (num_running_reqs, num_waiting_reqs).""" |
| 244 | raise NotImplementedError |
| 245 | |
| 246 | def get_kv_cache_usage(self) -> float: |
| 247 | """Returns the fraction of the KV cache currently in use (0.0-1.0).""" |
| 248 | return 0.0 |
| 249 | |
| 250 | @abstractmethod |
| 251 | def make_stats(self) -> "SchedulerStats | None": |
| 252 | """Make a SchedulerStats object for logging. |
| 253 | |
| 254 | The SchedulerStats object is created for every scheduling step. |
| 255 | """ |
| 256 | raise NotImplementedError |
| 257 | |
| 258 | @abstractmethod |
| 259 | def shutdown(self) -> None: |
| 260 | """Shutdown the scheduler.""" |
| 261 | raise NotImplementedError |
| 262 | |
| 263 | def get_kv_connector(self) -> "KVConnectorBase_V1 | None": |
| 264 | return None |
| 265 | |
| 266 | def get_ec_connector(self) -> "ECConnectorBase | None": |
| 267 | return None |
| 268 | |
| 269 | def get_kv_event_publisher_config(self) -> "KVEventsConfig | None": |
| 270 | return None |