| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Canvas-width handling and async read deferral for diffusion requests.""" |
| 4 | |
| 5 | from vllm.v1.core.sched.async_scheduler import AsyncScheduler |
| 6 | from vllm.v1.core.sched.output import SchedulerOutput |
| 7 | from vllm.v1.core.sched.scheduler import Scheduler |
| 8 | from vllm.v1.outputs import DraftTokenIds |
| 9 | from vllm.v1.request import Request |
| 10 | |
| 11 | |
| 12 | def diffusion_canvas_width(request: Request, canvas_length: int) -> int: |
| 13 | """The canvas width a diffusion request asked for, else the served one.""" |
| 14 | params = request.sampling_params |
| 15 | extra = params.extra_args if params is not None else None |
| 16 | width = extra.get("diffusion_canvas_length") if extra else None |
| 17 | return int(width) if width else canvas_length |
| 18 | |
| 19 | |
| 20 | def _read_in_flight(request: Request, width: int) -> bool: |
| 21 | """True when a read-only request has all its denoise steps in flight. |
| 22 | |
| 23 | The request emits its canvas on the last step and ends, so a further |
| 24 | step is discarded. |
| 25 | """ |
| 26 | params = request.sampling_params |
| 27 | extra = params.extra_args if params is not None else None |
| 28 | if not extra or not extra.get("diffusion_read_only"): |
| 29 | return False |
| 30 | steps = extra.get("diffusion_max_steps") |
| 31 | if not steps: |
| 32 | return False |
| 33 | # Each in-flight denoise step holds one canvas of placeholders. |
| 34 | return request.num_output_placeholders >= steps * width |
| 35 | |
| 36 | |
| 37 | class DiffusionScheduler(Scheduler): |
| 38 | def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None: |
| 39 | for i, (req_id, token_ids) in enumerate( |
| 40 | zip(draft_token_ids.req_ids, draft_token_ids.draft_token_ids) |
| 41 | ): |
| 42 | request = self.requests.get(req_id) |
| 43 | if request is None: |
| 44 | continue |
| 45 | width = diffusion_canvas_width(request, self.num_spec_tokens) |
| 46 | # Remove padded positions before the base scheduler validates grammar. |
| 47 | if len(token_ids) > width: |
| 48 | draft_token_ids.draft_token_ids[i] = token_ids[:width] |
| 49 | super().update_draft_token_ids(draft_token_ids) |
| 50 | |
| 51 | |
| 52 | class DiffusionAsyncScheduler(DiffusionScheduler, AsyncScheduler): |
| 53 | def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: |
| 54 | for request in self.running: |
| 55 | width = diffusion_canvas_width(request, self.num_spec_tokens) |
| 56 | # The placeholders are as wide as the served canvas. |
| 57 | if len(request.spec_token_ids) > width: |
| 58 | request.spec_token_ids = request.spec_token_ids[:width] |
| 59 | if _read_in_flight(request, width): |
| 60 | # Scheduler.schedule()'s max_tokens guard cannot be reached |
| 61 | # from a subclass. schedule() advances current_step before it |
| 62 | # checks decode eligibility, so current_step + 2 skips this |
| 63 | # call alone. max keeps a longer pipeline-parallel wait. |
| 64 | request.next_decode_eligible_step = max( |
| 65 | request.next_decode_eligible_step, self.current_step + 2 |
| 66 | ) |
| 67 | return super().schedule(throttle_prefills) |