| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import vllm.envs as envs |
| 5 | from vllm.logger import init_logger |
| 6 | from vllm.platforms import current_platform |
| 7 | from vllm.v1.core.sched.output import SchedulerOutput |
| 8 | from vllm.v1.core.sched.scheduler import Scheduler |
| 9 | from vllm.v1.request import Request, RequestStatus |
| 10 | |
| 11 | logger = init_logger(__name__) |
| 12 | |
| 13 | |
| 14 | class AsyncScheduler(Scheduler): |
| 15 | def __init__(self, *args, **kwargs) -> None: |
| 16 | super().__init__(*args, **kwargs) |
| 17 | # reusable read-only placeholder list for speculative decoding. |
| 18 | self._spec_token_placeholders: list[int] = [-1] * self.num_spec_tokens |
| 19 | self.pp_size = self.parallel_config.pipeline_parallel_size |
| 20 | # Only XPU can disable PP decode microbatching via VLLM_XPU_PP_MICROBATCH. |
| 21 | self.decode_stagger = self.pp_size |
| 22 | if current_platform.is_xpu() and not envs.VLLM_XPU_PP_MICROBATCH: |
| 23 | self.decode_stagger = 1 |
| 24 | |
| 25 | def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None: |
| 26 | super()._update_after_schedule(scheduler_output) |
| 27 | spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens |
| 28 | # Use the latest num of scheduled draft tokens in next step as placeholder. |
| 29 | self._spec_token_placeholders = [ |
| 30 | -1 |
| 31 | ] * scheduler_output.num_spec_tokens_to_schedule |
| 32 | for req_id in scheduler_output.num_scheduled_tokens: |
| 33 | request = self.requests[req_id] |
| 34 | if request.is_prefill_chunk: |
| 35 | continue |
| 36 | |
| 37 | scheduler_output.pending_structured_output_tokens |= ( |
| 38 | request.use_structured_output and request.num_output_placeholders > 0 |
| 39 | ) |
| 40 | # The request will generate num_sampled_tokens_per_step new tokens |
| 41 | # plus num_spec_tokens in this scheduling step. Diffusion has no AR |
| 42 | # bonus token (num_sampled_tokens_per_step == 0) — only the canvas |
| 43 | # (spec) tokens. |
| 44 | cur_num_spec_tokens = len(spec_decode_tokens.get(req_id, ())) |
| 45 | request.num_output_placeholders += ( |
| 46 | self.num_sampled_tokens_per_step + cur_num_spec_tokens |
| 47 | ) |
| 48 | # Add placeholders for the new draft/spec tokens. |
| 49 | # We will update the actual spec token ids in the worker process. |
| 50 | request.spec_token_ids = self._spec_token_placeholders |
| 51 | |
| 52 | if self.use_v2_model_runner: |
| 53 | # Set the next step index in which this request is eligible to be |
| 54 | # scheduled for decode (for PP microbatching). |
| 55 | request.next_decode_eligible_step = ( |
| 56 | self.current_step + self.decode_stagger |
| 57 | ) |
| 58 | |
| 59 | def _update_request_with_output( |
| 60 | self, request: Request, new_token_ids: list[int], is_stale: bool = False |
| 61 | ) -> tuple[list[int], bool]: |
| 62 | status_before_update = request.status |
| 63 | new_token_ids, stopped = super()._update_request_with_output( |
| 64 | request, new_token_ids |
| 65 | ) |
| 66 | |
| 67 | # Placeholders were zeroed at preemption; a stale delivery must not |
| 68 | # decrement them (it would underflow). |
| 69 | if not is_stale: |
| 70 | request.num_output_placeholders -= len(new_token_ids) |
| 71 | assert request.num_output_placeholders >= 0 |
| 72 | |
| 73 | # Cache the new tokens. Preempted requests should be skipped. |
| 74 | if status_before_update == RequestStatus.RUNNING: |
| 75 | self.kv_cache_manager.cache_blocks( |
| 76 | request, request.num_computed_tokens - request.num_output_placeholders |
| 77 | ) |
| 78 | return new_token_ids, stopped |