| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import pytest |
| 7 | |
| 8 | from vllm import SamplingParams |
| 9 | from vllm.config.diffusion import DiffusionConfig |
| 10 | from vllm.exceptions import VLLMValidationError |
| 11 | from vllm.utils.diffusion import validate_diffusion_sampling_params |
| 12 | from vllm.v1.engine.input_processor import InputProcessor |
| 13 | |
| 14 | |
| 15 | @dataclass |
| 16 | class MockModelConfig: |
| 17 | is_diffusion: bool = False |
| 18 | max_logprobs: int = 20 |
| 19 | logits_processors: list | None = None |
| 20 | return_sampling_mask: bool = False |
| 21 | |
| 22 | def get_vocab_size(self) -> int: |
| 23 | return 1024 |
| 24 | |
| 25 | |
| 26 | @pytest.mark.parametrize( |
| 27 | "kwargs", |
| 28 | [ |
| 29 | {"temperature": 0.7}, |
| 30 | {"temperature": 0.0}, |
| 31 | {"min_p": 0.1}, |
| 32 | {"seed": 42}, |
| 33 | {"min_tokens": 5}, |
| 34 | {"logit_bias": {0: 1.0}}, |
| 35 | {"bad_words": ["foo"]}, |
| 36 | {"allowed_token_ids": [0, 1]}, |
| 37 | ], |
| 38 | ) |
| 39 | def test_diffusion_rejects_unsupported_params(kwargs: dict): |
| 40 | params = SamplingParams(**kwargs) |
| 41 | with pytest.raises(VLLMValidationError, match="not yet supported with diffusion"): |
| 42 | params.verify(MockModelConfig(is_diffusion=True), None, None, None) |
| 43 | |
| 44 | |
| 45 | def test_diffusion_accepts_default_params(): |
| 46 | SamplingParams().verify(MockModelConfig(is_diffusion=True), None, None, None) |
| 47 | |
| 48 | |
| 49 | def test_diffusion_accepts_top_k_top_p(): |
| 50 | params = SamplingParams(top_p=0.9, top_k=10) |
| 51 | params.verify(MockModelConfig(is_diffusion=True), None, None, None) |
| 52 | |
| 53 | |
| 54 | def test_non_diffusion_models_unaffected(): |
| 55 | params = SamplingParams(temperature=0.7, top_k=10, seed=42) |
| 56 | params.verify(MockModelConfig(), None, None, None) |
| 57 | |
| 58 | |
| 59 | def test_verify_leaves_logits_processors_to_admission(): |
| 60 | """verify() is runner-agnostic; LP validation lives in the admission |
| 61 | layer, so an unimportable FQCN must not fail verify().""" |
| 62 | SamplingParams().verify( |
| 63 | MockModelConfig(logits_processors=["no.such:Cls"]), None, None, None |
| 64 | ) |
| 65 | |
| 66 | |
| 67 | def _verify_diffusion(params: SamplingParams, canvas_length: int | None = None): |
| 68 | model_config = MockModelConfig(is_diffusion=True) |
| 69 | params.verify(model_config, None, None, None) |
| 70 | validate_diffusion_sampling_params( |
| 71 | params, |
| 72 | canvas_length=canvas_length, |
| 73 | vocab_size=model_config.get_vocab_size(), |
| 74 | ) |
| 75 | |
| 76 | |
| 77 | def test_diffusion_extra_args_are_validated_without_a_served_canvas(): |
| 78 | # No --diffusion-config: the canvas is unknown, the ids are still checked |
| 79 | # and a read-only request is still normalised. |
| 80 | params = SamplingParams( |
| 81 | max_tokens=64, |
| 82 | extra_args={"diffusion_seed_canvas": [0, 1], "diffusion_read_only": True}, |
| 83 | ) |
| 84 | _verify_diffusion(params, canvas_length=None) |
| 85 | assert params.ignore_eos is True |
| 86 | |
| 87 | bad = SamplingParams(extra_args={"diffusion_seed_canvas": [0, 10**9]}) |
| 88 | with pytest.raises(VLLMValidationError, match="ids must be in"): |
| 89 | _verify_diffusion(bad, canvas_length=None) |
| 90 | |
| 91 | |
| 92 | @pytest.mark.parametrize("async_scheduling", [False, True]) |
| 93 | @pytest.mark.parametrize( |
| 94 | "extra_args", |
| 95 | [ |
| 96 | {}, |
| 97 | {"diffusion_canvas_length": None}, |
| 98 | {"diffusion_canvas_length": 4}, |
| 99 | {"diffusion_canvas_length": 8}, |
| 100 | ], |
| 101 | ) |
| 102 | def test_diffusion_canvas_width_is_accepted_with_either_scheduler( |
| 103 | async_scheduling, extra_args |
| 104 | ): |
| 105 | processor = SimpleNamespace( |
| 106 | model_config=MockModelConfig(is_diffusion=True), |
| 107 | vllm_config=SimpleNamespace( |
| 108 | scheduler_config=SimpleNamespace(async_scheduling=async_scheduling) |
| 109 | ), |
| 110 | speculative_config=None, |
| 111 | structured_outputs_config=None, |
| 112 | diffusion_config=DiffusionConfig(canvas_length=8), |
| 113 | tokenizer=None, |
| 114 | validate_logits_processors_params=lambda params: None, |
| 115 | ) |
| 116 | params = SamplingParams(extra_args=extra_args) |
| 117 | InputProcessor._validate_params(processor, params, ("generate",)) |
| 118 | |
| 119 | |
| 120 | @pytest.mark.parametrize( |
| 121 | "extra_args, match", |
| 122 | [ |
| 123 | ({"diffusion_seed_canvas": "abc"}, "list of token ids"), |
| 124 | ({"diffusion_seed_canvas": [1, 2.5]}, "list of token ids"), |
| 125 | ({"diffusion_seed_canvas": [1, True]}, "list of token ids"), |
| 126 | ({"diffusion_seed_canvas": [0, 1024]}, r"in \[0, 1024\)"), |
| 127 | ({"diffusion_seed_canvas": [0, -1]}, r"in \[0, 1024\)"), |
| 128 | ({"diffusion_max_steps": 0}, "positive integer"), |
| 129 | ({"diffusion_max_steps": "1"}, "positive integer"), |
| 130 | ({"diffusion_max_steps": True}, "positive integer"), |
| 131 | ({"diffusion_read_only": "yes"}, "boolean"), |
| 132 | ({"diffusion_read_only": 2}, "boolean"), |
| 133 | ({"diffusion_pinned": "0,1"}, "list of canvas positions"), |
| 134 | ({"diffusion_pinned": [0, True]}, "list of canvas positions"), |
| 135 | ({"diffusion_pinned": [0]}, "needs a diffusion_seed_canvas"), |
| 136 | ], |
| 137 | ) |
| 138 | def test_diffusion_rejects_bad_extra_args(extra_args: dict, match: str): |
| 139 | with pytest.raises(VLLMValidationError, match=match): |
| 140 | _verify_diffusion(SamplingParams(extra_args=extra_args)) |
| 141 | |
| 142 | |
| 143 | @pytest.mark.parametrize("flag", [True, 1]) |
| 144 | def test_diffusion_constrained_needs_logprob_token_ids(flag): |
| 145 | bad = SamplingParams(extra_args={"diffusion_constrained": flag}) |
| 146 | with pytest.raises(VLLMValidationError, match="needs logprob_token_ids"): |
| 147 | _verify_diffusion(bad) |
| 148 | |
| 149 | ok = SamplingParams( |
| 150 | logprob_token_ids=[3, 5], extra_args={"diffusion_constrained": flag} |
| 151 | ) |
| 152 | _verify_diffusion(ok) |
| 153 | |
| 154 | # An unset or false flag needs no ids. |
| 155 | _verify_diffusion(SamplingParams(extra_args={"diffusion_constrained": False})) |
| 156 | _verify_diffusion(SamplingParams(extra_args={"diffusion_constrained": 0})) |
| 157 | |
| 158 | |
| 159 | @pytest.mark.parametrize("value", ["yes", 2, [1]]) |
| 160 | def test_diffusion_constrained_must_be_a_bool(value): |
| 161 | params = SamplingParams(extra_args={"diffusion_constrained": value}) |
| 162 | with pytest.raises(VLLMValidationError, match="must be a boolean"): |
| 163 | _verify_diffusion(params) |
| 164 | |
| 165 | |
| 166 | def test_diffusion_seed_canvas_must_fill_the_canvas(): |
| 167 | params = SamplingParams(extra_args={"diffusion_seed_canvas": [0] * 7}) |
| 168 | with pytest.raises(VLLMValidationError, match="exactly 8 ids, got 7"): |
| 169 | _verify_diffusion(params, canvas_length=8) |
| 170 | # Without the served diffusion config the canvas length is unknown. |
| 171 | _verify_diffusion(params) |
| 172 | |
| 173 | |
| 174 | @pytest.mark.parametrize("max_tokens, expected", [(100, 8), (5, 5), (None, 8)]) |
| 175 | @pytest.mark.parametrize("flag", [True, 1]) |
| 176 | def test_diffusion_read_only_ends_after_one_canvas(max_tokens, expected, flag): |
| 177 | params = SamplingParams( |
| 178 | max_tokens=max_tokens, extra_args={"diffusion_read_only": flag} |
| 179 | ) |
| 180 | _verify_diffusion(params, canvas_length=8) |
| 181 | assert params.max_tokens == expected |
| 182 | assert params.ignore_eos |
| 183 | |
| 184 | |
| 185 | @pytest.mark.parametrize( |
| 186 | "width, match", |
| 187 | [ |
| 188 | (0, "positive integer"), |
| 189 | ("4", "positive integer"), |
| 190 | (True, "positive integer"), |
| 191 | (9, "no larger"), |
| 192 | ], |
| 193 | ) |
| 194 | def test_diffusion_rejects_bad_canvas_length(width, match): |
| 195 | with pytest.raises(VLLMValidationError, match=match): |
| 196 | _verify_diffusion( |
| 197 | SamplingParams(extra_args={"diffusion_canvas_length": width}), |
| 198 | canvas_length=8, |
| 199 | ) |
| 200 | |
| 201 | |
| 202 | def test_diffusion_canvas_length_sizes_the_seed_and_the_read(): |
| 203 | params = SamplingParams( |
| 204 | max_tokens=100, |
| 205 | extra_args={ |
| 206 | "diffusion_canvas_length": 4, |
| 207 | "diffusion_seed_canvas": [1, 2, 3, 4], |
| 208 | "diffusion_read_only": True, |
| 209 | }, |
| 210 | ) |
| 211 | _verify_diffusion(params, canvas_length=8) |
| 212 | assert params.max_tokens == 4 |
| 213 | |
| 214 | params = SamplingParams( |
| 215 | extra_args={"diffusion_canvas_length": 4, "diffusion_seed_canvas": [0] * 8} |
| 216 | ) |
| 217 | with pytest.raises(VLLMValidationError, match="exactly 4 ids, got 8"): |
| 218 | _verify_diffusion(params, canvas_length=8) |
| 219 | |
| 220 | |
| 221 | def test_diffusion_pinned_positions_stay_inside_the_canvas(): |
| 222 | seed = [0] * 8 |
| 223 | with pytest.raises(VLLMValidationError, match="inside the canvas"): |
| 224 | _verify_diffusion( |
| 225 | SamplingParams( |
| 226 | extra_args={"diffusion_seed_canvas": seed, "diffusion_pinned": [7, 8]} |
| 227 | ), |
| 228 | canvas_length=8, |
| 229 | ) |
| 230 | with pytest.raises(VLLMValidationError, match="inside the canvas"): |
| 231 | _verify_diffusion( |
| 232 | SamplingParams( |
| 233 | extra_args={"diffusion_seed_canvas": seed, "diffusion_pinned": [-1]} |
| 234 | ) |
| 235 | ) |
| 236 | # A narrower request canvas bounds the positions. |
| 237 | with pytest.raises(VLLMValidationError, match="inside the canvas"): |
| 238 | _verify_diffusion( |
| 239 | SamplingParams( |
| 240 | extra_args={ |
| 241 | "diffusion_canvas_length": 4, |
| 242 | "diffusion_seed_canvas": seed[:4], |
| 243 | "diffusion_pinned": [4], |
| 244 | } |
| 245 | ), |
| 246 | canvas_length=8, |
| 247 | ) |
| 248 | |
| 249 | |
| 250 | def test_diffusion_accepts_extra_args(): |
| 251 | params = SamplingParams( |
| 252 | extra_args={ |
| 253 | "diffusion_seed_canvas": list(range(8)), |
| 254 | "diffusion_pinned": [0, 1, 7], |
| 255 | "diffusion_max_steps": 4, |
| 256 | "diffusion_read_only": True, |
| 257 | } |
| 258 | ) |
| 259 | _verify_diffusion(params, canvas_length=8) |
| 260 | |
| 261 | |
| 262 | @pytest.mark.parametrize("value", [-(2**63) - 1, 2**64]) |
| 263 | def test_extra_args_rejects_nested_integer_overflow(value): |
| 264 | """Reject extension values before they reach the engine transport.""" |
| 265 | with pytest.raises(VLLMValidationError, match="extra_args integers"): |
| 266 | SamplingParams(extra_args={"ec_transfer_params": {"nested": [{"x": value}]}}) |
| 267 | |
| 268 | |
| 269 | @pytest.mark.parametrize("value", [-(2**63), 2**63 - 1, 2**63, 2**64 - 1, True]) |
| 270 | def test_extra_args_accepts_messagepack_integer_boundaries(value): |
| 271 | extra_args = {"kv_transfer_params": {"nested": [{"x": value}]}} |
| 272 | assert SamplingParams(extra_args=extra_args).extra_args == extra_args |
| 273 | |
| 274 | |
| 275 | def test_extra_args_preserves_custom_objects_and_shared_containers(): |
| 276 | custom = object() |
| 277 | shared = [custom, (None, "value", 1.5)] |
| 278 | extra_args = {"first": shared, "second": shared} |
| 279 | params = SamplingParams(extra_args=extra_args) |
| 280 | assert params.extra_args["first"][0] is custom |
| 281 | assert params.extra_args["first"] is params.extra_args["second"] |