| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Request validation for discrete diffusion models.""" |
| 4 | |
| 5 | from typing import TYPE_CHECKING |
| 6 | |
| 7 | from vllm.exceptions import VLLMValidationError |
| 8 | |
| 9 | if TYPE_CHECKING: |
| 10 | from vllm.sampling_params import SamplingParams |
| 11 | |
| 12 | |
| 13 | def validate_diffusion_sampling_params( |
| 14 | params: "SamplingParams", |
| 15 | *, |
| 16 | canvas_length: int | None, |
| 17 | vocab_size: int, |
| 18 | ) -> None: |
| 19 | """Validate request options and normalize one-canvas reads at admission.""" |
| 20 | extra = params.extra_args |
| 21 | if not extra: |
| 22 | return |
| 23 | |
| 24 | def check(valid: bool, message: str) -> None: |
| 25 | if not valid: |
| 26 | raise VLLMValidationError(message, parameter="extra_args") |
| 27 | |
| 28 | width = extra.get("diffusion_canvas_length") |
| 29 | if width is not None: |
| 30 | check( |
| 31 | isinstance(width, int) |
| 32 | and not isinstance(width, bool) |
| 33 | and width > 0 |
| 34 | and (canvas_length is None or width <= canvas_length), |
| 35 | "diffusion_canvas_length must be a positive integer no larger " |
| 36 | "than the served canvas.", |
| 37 | ) |
| 38 | # The request's canvas: its own width, else the served one. |
| 39 | expected_len = width or canvas_length |
| 40 | |
| 41 | seed = extra.get("diffusion_seed_canvas") |
| 42 | if seed is not None: |
| 43 | check( |
| 44 | isinstance(seed, (list, tuple)) |
| 45 | and all(isinstance(t, int) and not isinstance(t, bool) for t in seed), |
| 46 | "diffusion_seed_canvas must be a list of token ids.", |
| 47 | ) |
| 48 | check( |
| 49 | all(0 <= t < vocab_size for t in seed), |
| 50 | f"diffusion_seed_canvas ids must be in [0, {vocab_size}).", |
| 51 | ) |
| 52 | check( |
| 53 | expected_len is None or len(seed) == expected_len, |
| 54 | f"diffusion_seed_canvas must hold exactly {expected_len} ids, " |
| 55 | f"got {len(seed)}.", |
| 56 | ) |
| 57 | |
| 58 | pins = extra.get("diffusion_pinned") |
| 59 | if pins is not None: |
| 60 | check( |
| 61 | isinstance(pins, (list, tuple)) |
| 62 | and all(isinstance(p, int) and not isinstance(p, bool) for p in pins), |
| 63 | "diffusion_pinned must be a list of canvas positions.", |
| 64 | ) |
| 65 | check( |
| 66 | seed is not None, |
| 67 | "diffusion_pinned needs a diffusion_seed_canvas to hold.", |
| 68 | ) |
| 69 | check( |
| 70 | all(p >= 0 and (expected_len is None or p < expected_len) for p in pins), |
| 71 | "diffusion_pinned positions must lie inside the canvas.", |
| 72 | ) |
| 73 | |
| 74 | cap = extra.get("diffusion_max_steps") |
| 75 | if cap is not None: |
| 76 | check( |
| 77 | isinstance(cap, int) and not isinstance(cap, bool) and cap > 0, |
| 78 | "diffusion_max_steps must be a positive integer.", |
| 79 | ) |
| 80 | |
| 81 | # The OpenAI server's vllm_xargs narrows JSON booleans to 0/1. |
| 82 | read_only = extra.get("diffusion_read_only") |
| 83 | if read_only is not None: |
| 84 | check( |
| 85 | isinstance(read_only, (bool, int)) and read_only in (0, 1), |
| 86 | "diffusion_read_only must be a boolean (or 0/1).", |
| 87 | ) |
| 88 | if read_only: |
| 89 | # One canvas is the whole output. Cap max_tokens so the scheduler |
| 90 | # ends the request there, and ignore EOS so an end-of-turn token |
| 91 | # drawn into a noise slot does not end it early. |
| 92 | if expected_len is not None: |
| 93 | params.max_tokens = min(params.max_tokens or expected_len, expected_len) |
| 94 | params.ignore_eos = True |
| 95 | |
| 96 | constrained = extra.get("diffusion_constrained") |
| 97 | if constrained is not None: |
| 98 | check( |
| 99 | isinstance(constrained, (bool, int)) and constrained in (0, 1), |
| 100 | "diffusion_constrained must be a boolean (or 0/1).", |
| 101 | ) |
| 102 | if constrained: |
| 103 | check( |
| 104 | bool(params.logprob_token_ids), |
| 105 | "diffusion_constrained needs logprob_token_ids: they are the allowed set.", |
| 106 | ) |