| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import time |
| 5 | from dataclasses import dataclass |
| 6 | |
| 7 | import pytest |
| 8 | from transformers import AutoTokenizer |
| 9 | |
| 10 | from vllm.config import DeviceConfig, StructuredOutputsConfig, VllmConfig |
| 11 | from vllm.config.model import ModelConfig |
| 12 | from vllm.config.speculative import SpeculativeConfig |
| 13 | from vllm.config.structured_outputs import StructuredOutputsBackend |
| 14 | from vllm.sampling_params import SamplingParams, StructuredOutputsParams |
| 15 | from vllm.v1.request import Request |
| 16 | from vllm.v1.structured_output import StructuredOutputManager |
| 17 | from vllm.v1.structured_output.backend_outlines import OutlinesGrammar |
| 18 | |
| 19 | TOKENIZER = "gpt2" |
| 20 | THINK_END = "\n" # reasoning-end marker (single GPT-2 token) |
| 21 | EOS = "<|eos|>" # resolved to tokenizer.eos_token_id |
| 22 | JSON_SCHEMA = '{"type": "object"}' |
| 23 | BACKENDS = ("xgrammar", "guidance") |
| 24 | REGEX_BACKENDS = (*BACKENDS, "outlines") |
| 25 | DIGITS_REGEX = "[0-9]+" |
| 26 | MAX_WAIT_SECONDS = 5 |
| 27 | NUM_SPEC_TOKENS = 8 |
| 28 | |
| 29 | |
| 30 | @pytest.fixture(scope="module") |
| 31 | def tokenizer(): |
| 32 | return AutoTokenizer.from_pretrained(TOKENIZER) |
| 33 | |
| 34 | |
| 35 | @dataclass(frozen=True) |
| 36 | class FlowCase: |
| 37 | """One validate -> bitmask -> accept -> post-accept bitmask scenario.""" |
| 38 | |
| 39 | raw_drafts: tuple[str, ...] |
| 40 | expected_row_pattern: str |
| 41 | expected_reasoning: bool | None |
| 42 | expect_terminated: bool |
| 43 | prefix: str = "" |
| 44 | expected_validated: tuple[str, ...] | None = None |
| 45 | reasoning_ended: bool | None = False |
| 46 | xfail_guidance: str | None = None |
| 47 | |
| 48 | |
| 49 | class MockReasoner: |
| 50 | def __init__(self, tokenizer, marker: int | None = None): |
| 51 | self.marker = marker |
| 52 | |
| 53 | def is_reasoning_end(self, input_ids): |
| 54 | if self.marker is None: |
| 55 | return True |
| 56 | return self.marker in list(input_ids) |
| 57 | |
| 58 | def is_reasoning_end_streaming(self, input_ids, delta_ids): |
| 59 | return self.is_reasoning_end(delta_ids) |
| 60 | |
| 61 | |
| 62 | def _single_token(tokenizer, text: str) -> int: |
| 63 | token_ids = tokenizer.encode(text) |
| 64 | assert len(token_ids) == 1, (text, token_ids) |
| 65 | return token_ids[0] |
| 66 | |
| 67 | |
| 68 | def _to_token_ids(tokenizer, texts: tuple[str, ...]) -> list[int]: |
| 69 | """Convert literal text (or EOS sentinel) to token IDs.""" |
| 70 | return [ |
| 71 | tokenizer.eos_token_id if t == EOS else _single_token(tokenizer, t) |
| 72 | for t in texts |
| 73 | ] |
| 74 | |
| 75 | |
| 76 | def _wait_for_grammar(request: Request) -> None: |
| 77 | structured_req = request.structured_output_request |
| 78 | assert structured_req is not None |
| 79 | |
| 80 | deadline = time.time() + MAX_WAIT_SECONDS |
| 81 | while not structured_req._check_grammar_completion(): |
| 82 | if time.time() > deadline: |
| 83 | pytest.fail("Grammar compilation timed out") |
| 84 | time.sleep(0.01) |
| 85 | |
| 86 | |
| 87 | def _build_harness( |
| 88 | tokenizer, |
| 89 | backend: StructuredOutputsBackend, |
| 90 | prefix: str = "", |
| 91 | use_reasoner: bool = True, |
| 92 | reasoning_ended: bool | None = None, |
| 93 | enable_in_reasoning: bool = False, |
| 94 | reasoning_parser_kwargs: dict | None = None, |
| 95 | structured_outputs: StructuredOutputsParams | None = None, |
| 96 | ) -> tuple[StructuredOutputManager, Request]: |
| 97 | vllm_config = VllmConfig( |
| 98 | model_config=ModelConfig(tokenizer=TOKENIZER), |
| 99 | device_config=DeviceConfig(device="cpu"), |
| 100 | structured_outputs_config=StructuredOutputsConfig( |
| 101 | backend=backend, |
| 102 | enable_in_reasoning=enable_in_reasoning, |
| 103 | ), |
| 104 | speculative_config=SpeculativeConfig( |
| 105 | model="[ngram]", |
| 106 | num_speculative_tokens=NUM_SPEC_TOKENS, |
| 107 | ), |
| 108 | ) |
| 109 | manager = StructuredOutputManager(vllm_config) |
| 110 | if use_reasoner: |
| 111 | manager.reasoner_cls = MockReasoner |
| 112 | |
| 113 | sampling_params = SamplingParams( |
| 114 | structured_outputs=structured_outputs |
| 115 | or StructuredOutputsParams(json=JSON_SCHEMA) |
| 116 | ) |
| 117 | sampling_params.structured_outputs._backend = backend # type: ignore[union-attr] |
| 118 | sampling_params.update_from_generation_config({}, tokenizer.eos_token_id) |
| 119 | |
| 120 | prompt_ids = tokenizer.encode(prefix) if prefix else [] |
| 121 | request = Request( |
| 122 | request_id=f"{backend}-flow", |
| 123 | prompt_token_ids=prompt_ids, |
| 124 | sampling_params=sampling_params, |
| 125 | pooling_params=None, |
| 126 | reasoning_ended=reasoning_ended, |
| 127 | reasoning_parser_kwargs=reasoning_parser_kwargs, |
| 128 | ) |
| 129 | manager.grammar_init(request) |
| 130 | _wait_for_grammar(request) |
| 131 | |
| 132 | structured_req = request.structured_output_request |
| 133 | assert request.prompt_token_ids is not None |
| 134 | assert structured_req is not None |
| 135 | grammar = structured_req.grammar |
| 136 | assert grammar is not None and not isinstance(grammar, Exception) |
| 137 | if prompt_ids: |
| 138 | assert grammar.accept_tokens( |
| 139 | request.request_id, |
| 140 | prompt_ids, |
| 141 | ) |
| 142 | return manager, request |
| 143 | |
| 144 | |
| 145 | def _row_pattern(bitmask) -> str: |
| 146 | return "".join("U" if (row == -1).all() else "C" for row in bitmask) |
| 147 | |
| 148 | |
| 149 | def _run_real_flow( |
| 150 | manager: StructuredOutputManager, |
| 151 | request: Request, |
| 152 | raw_drafts: list[int], |
| 153 | expected_validated: list[int], |
| 154 | expected_row_pattern: str, |
| 155 | expected_reasoning: bool | None, |
| 156 | expect_terminated: bool, |
| 157 | ) -> None: |
| 158 | assert len(expected_validated) <= len(raw_drafts) |
| 159 | |
| 160 | validated = manager.validate_tokens(request, list(raw_drafts)) |
| 161 | assert validated == expected_validated |
| 162 | |
| 163 | padded = validated + [-1] * (len(raw_drafts) - len(validated)) |
| 164 | bitmask = manager.grammar_bitmask( |
| 165 | requests={request.request_id: request}, |
| 166 | structured_output_request_ids=[request.request_id], |
| 167 | scheduled_spec_decode_tokens={request.request_id: padded}, |
| 168 | ) |
| 169 | assert bitmask is not None |
| 170 | assert bitmask.shape[0] == len(raw_drafts) + 1 |
| 171 | assert _row_pattern(bitmask) == expected_row_pattern |
| 172 | |
| 173 | # grammar_bitmask() must rollback any speculative state it advanced. |
| 174 | assert manager.validate_tokens(request, list(raw_drafts)) == expected_validated |
| 175 | |
| 176 | structured_req = request.structured_output_request |
| 177 | assert structured_req is not None |
| 178 | grammar = structured_req.grammar |
| 179 | assert grammar is not None and not isinstance(grammar, Exception) |
| 180 | assert not grammar.is_terminated() |
| 181 | # `bitmask[i]` is the grammar state after the first `i` scheduled tokens. |
| 182 | # These tests commit the validated prefix, so the matching post-accept |
| 183 | # state is already present in the scheduled bitmask. |
| 184 | expected_post_accept_row = bitmask[len(expected_validated)].copy() |
| 185 | |
| 186 | # Mirror the scheduler flow: sampled tokens are appended before |
| 187 | # StructuredOutputManager.accept_tokens() trims them to the |
| 188 | # grammar-constrained suffix. |
| 189 | request.append_output_token_ids(list(expected_validated)) |
| 190 | assert manager.accept_tokens(request, list(expected_validated)) |
| 191 | |
| 192 | post_accept_bitmask = manager.grammar_bitmask( |
| 193 | requests={request.request_id: request}, |
| 194 | structured_output_request_ids=[request.request_id], |
| 195 | scheduled_spec_decode_tokens={}, |
| 196 | ) |
| 197 | assert post_accept_bitmask is not None |
| 198 | assert post_accept_bitmask.shape[0] == 1 |
| 199 | assert (post_accept_bitmask[0] == expected_post_accept_row).all() |
| 200 | assert structured_req.reasoning_ended == expected_reasoning |
| 201 | assert grammar.is_terminated() is expect_terminated |
| 202 | |
| 203 | |
| 204 | FLOW_CASES = [ |
| 205 | pytest.param( |
| 206 | FlowCase( |
| 207 | raw_drafts=(" ", "z", " "), |
| 208 | expected_row_pattern="UUUU", |
| 209 | expected_reasoning=False, |
| 210 | expect_terminated=False, |
| 211 | ), |
| 212 | id="inactive", |
| 213 | ), |
| 214 | pytest.param( |
| 215 | FlowCase( |
| 216 | prefix='{"a"', |
| 217 | raw_drafts=(":", ' "', "b", '"}'), |
| 218 | expected_row_pattern="CCCCC", |
| 219 | expected_reasoning=True, |
| 220 | expect_terminated=False, |
| 221 | reasoning_ended=True, |
| 222 | ), |
| 223 | id="active_all_valid", |
| 224 | ), |
| 225 | pytest.param( |
| 226 | FlowCase( |
| 227 | prefix='{"a"', |
| 228 | raw_drafts=("z", ":", ' "', "b"), |
| 229 | expected_validated=(), |
| 230 | expected_row_pattern="CUUUU", |
| 231 | expected_reasoning=True, |
| 232 | expect_terminated=False, |
| 233 | reasoning_ended=True, |
| 234 | ), |
| 235 | id="active_first_token_invalid", |
| 236 | ), |
| 237 | pytest.param( |
| 238 | FlowCase( |
| 239 | prefix='{"a"', |
| 240 | raw_drafts=(":", "z", ' "', "b"), |
| 241 | expected_validated=(":",), |
| 242 | expected_row_pattern="CCUUU", |
| 243 | expected_reasoning=True, |
| 244 | expect_terminated=False, |
| 245 | reasoning_ended=True, |
| 246 | ), |
| 247 | id="active_later_token_invalid", |
| 248 | ), |
| 249 | pytest.param( |
| 250 | FlowCase( |
| 251 | prefix='{"a": "b"}', |
| 252 | raw_drafts=(EOS, " ", " "), |
| 253 | expected_validated=(EOS,), |
| 254 | expected_row_pattern="CUUU", |
| 255 | expected_reasoning=True, |
| 256 | expect_terminated=True, |
| 257 | reasoning_ended=True, |
| 258 | xfail_guidance=( |
| 259 | "guidance validate_tokens rejects EOS on an already-complete " |
| 260 | "object even though accept_tokens would accept it" |
| 261 | ), |
| 262 | ), |
| 263 | id="active_first_token_terminates", |
| 264 | ), |
| 265 | pytest.param( |
| 266 | FlowCase( |
| 267 | prefix='{"a"', |
| 268 | raw_drafts=(":", ' "', "b", '"}', EOS, " "), |
| 269 | expected_validated=(":", ' "', "b", '"}', EOS), |
| 270 | expected_row_pattern="CCCCCUU", |
| 271 | expected_reasoning=True, |
| 272 | expect_terminated=True, |
| 273 | reasoning_ended=True, |
| 274 | ), |
| 275 | id="active_later_token_terminates", |
| 276 | ), |
| 277 | pytest.param( |
| 278 | FlowCase( |
| 279 | raw_drafts=(THINK_END, "{", ' "', "b"), |
| 280 | expected_row_pattern="UCCCC", |
| 281 | expected_reasoning=True, |
| 282 | expect_terminated=False, |
| 283 | ), |
| 284 | id="becomes_active_on_first_token", |
| 285 | ), |
| 286 | pytest.param( |
| 287 | FlowCase( |
| 288 | raw_drafts=(" ", THINK_END, "{", ' "'), |
| 289 | expected_row_pattern="UUCCC", |
| 290 | expected_reasoning=True, |
| 291 | expect_terminated=False, |
| 292 | ), |
| 293 | id="becomes_active_on_middle_token", |
| 294 | ), |
| 295 | pytest.param( |
| 296 | FlowCase( |
| 297 | raw_drafts=(" ", " ", THINK_END), |
| 298 | expected_row_pattern="UUUC", |
| 299 | expected_reasoning=True, |
| 300 | expect_terminated=False, |
| 301 | ), |
| 302 | id="becomes_active_on_last_token", |
| 303 | ), |
| 304 | pytest.param( |
| 305 | FlowCase( |
| 306 | raw_drafts=(" ", THINK_END, "z", "{"), |
| 307 | expected_validated=(" ", THINK_END), |
| 308 | expected_row_pattern="UUCUU", |
| 309 | expected_reasoning=True, |
| 310 | expect_terminated=False, |
| 311 | ), |
| 312 | id="becomes_active_first_token_invalid", |
| 313 | ), |
| 314 | pytest.param( |
| 315 | FlowCase( |
| 316 | raw_drafts=(" ", THINK_END, "{", "z"), |
| 317 | expected_validated=(" ", THINK_END, "{"), |
| 318 | expected_row_pattern="UUCCU", |
| 319 | expected_reasoning=True, |
| 320 | expect_terminated=False, |
| 321 | ), |
| 322 | id="becomes_active_later_token_invalid", |
| 323 | ), |
| 324 | pytest.param( |
| 325 | FlowCase( |
| 326 | raw_drafts=(" ", THINK_END, "{", "}", EOS, " "), |
| 327 | expected_validated=(" ", THINK_END, "{", "}", EOS), |
| 328 | expected_row_pattern="UUCCCUU", |
| 329 | expected_reasoning=True, |
| 330 | expect_terminated=True, |
| 331 | ), |
| 332 | id="becomes_active_terminates", |
| 333 | ), |
| 334 | ] |
| 335 | |
| 336 | |
| 337 | @pytest.mark.parametrize("backend", BACKENDS) |
| 338 | @pytest.mark.parametrize("case", FLOW_CASES) |
| 339 | def test_real_flow( |
| 340 | tokenizer, |
| 341 | backend: StructuredOutputsBackend, |
| 342 | case: FlowCase, |
| 343 | ): |
| 344 | if backend == "guidance" and case.xfail_guidance: |
| 345 | pytest.xfail(case.xfail_guidance) |
| 346 | |
| 347 | reasoner_kwargs = {"marker": _single_token(tokenizer, THINK_END)} |
| 348 | manager, request = _build_harness( |
| 349 | tokenizer, |
| 350 | backend, |
| 351 | prefix=case.prefix, |
| 352 | reasoning_ended=case.reasoning_ended, |
| 353 | reasoning_parser_kwargs=reasoner_kwargs, |
| 354 | ) |
| 355 | |
| 356 | raw_drafts = _to_token_ids(tokenizer, case.raw_drafts) |
| 357 | expected_texts = ( |
| 358 | case.raw_drafts if case.expected_validated is None else case.expected_validated |
| 359 | ) |
| 360 | expected_validated = _to_token_ids(tokenizer, expected_texts) |
| 361 | |
| 362 | _run_real_flow( |
| 363 | manager, |
| 364 | request, |
| 365 | raw_drafts=raw_drafts, |
| 366 | expected_validated=expected_validated, |
| 367 | expected_row_pattern=case.expected_row_pattern, |
| 368 | expected_reasoning=case.expected_reasoning, |
| 369 | expect_terminated=case.expect_terminated, |
| 370 | ) |
| 371 | |
| 372 | |
| 373 | @pytest.mark.parametrize("backend", BACKENDS) |
| 374 | @pytest.mark.parametrize( |
| 375 | ("use_reasoner", "reasoning_ended", "enable_in_reasoning"), |
| 376 | [ |
| 377 | pytest.param(True, False, True, id="enable_in_reasoning"), |
| 378 | pytest.param(True, True, False, id="reasoning_ended"), |
| 379 | pytest.param(False, False, False, id="no_reasoner"), |
| 380 | pytest.param(True, None, False, id="inferred"), |
| 381 | ], |
| 382 | ) |
| 383 | def test_initial_constraint_activation( |
| 384 | tokenizer, |
| 385 | backend: StructuredOutputsBackend, |
| 386 | use_reasoner: bool, |
| 387 | reasoning_ended: bool | None, |
| 388 | enable_in_reasoning: bool, |
| 389 | ): |
| 390 | manager, request = _build_harness( |
| 391 | tokenizer, |
| 392 | backend, |
| 393 | use_reasoner=use_reasoner, |
| 394 | reasoning_ended=reasoning_ended, |
| 395 | enable_in_reasoning=enable_in_reasoning, |
| 396 | reasoning_parser_kwargs={ |
| 397 | "marker": _single_token(tokenizer, THINK_END) |
| 398 | if reasoning_ended is not None |
| 399 | else None |
| 400 | }, |
| 401 | ) |
| 402 | |
| 403 | # "{" is valid JSON start; "z" is not. Truncation proves grammar is active. |
| 404 | open_brace = _single_token(tokenizer, "{") |
| 405 | z = _single_token(tokenizer, "z") |
| 406 | assert manager.validate_tokens(request, [open_brace, z]) == [open_brace] |
| 407 | |
| 408 | bitmask = manager.grammar_bitmask( |
| 409 | requests={request.request_id: request}, |
| 410 | structured_output_request_ids=[request.request_id], |
| 411 | scheduled_spec_decode_tokens={}, |
| 412 | ) |
| 413 | assert bitmask is not None |
| 414 | request.append_output_token_ids([open_brace]) |
| 415 | assert manager.accept_tokens(request, [open_brace]) |
| 416 | |
| 417 | # Short-circuit paths must not construct a request-local reasoner (#52573). |
| 418 | if enable_in_reasoning or reasoning_ended is True: |
| 419 | structured_req = request.structured_output_request |
| 420 | assert structured_req is not None |
| 421 | assert structured_req.reasoner is None |
| 422 | |
| 423 | |
| 424 | @pytest.mark.parametrize("backend", REGEX_BACKENDS) |
| 425 | @pytest.mark.parametrize( |
| 426 | ("raw_drafts", "expected_validated", "expected_row_pattern", "terminated"), |
| 427 | [ |
| 428 | pytest.param(("1", "2"), None, "CCC", False, id="all_valid"), |
| 429 | # EOS is only valid once the regex has matched; rows after EOS are |
| 430 | # unconstrained since the request stops there. |
| 431 | pytest.param(("1", EOS, "2"), ("1", EOS), "CCUU", True, id="terminates"), |
| 432 | pytest.param((EOS, "1"), (), "CUU", False, id="early_eos"), |
| 433 | ], |
| 434 | ) |
| 435 | def test_regex_flow( |
| 436 | tokenizer, |
| 437 | backend: StructuredOutputsBackend, |
| 438 | raw_drafts: tuple[str, ...], |
| 439 | expected_validated: tuple[str, ...] | None, |
| 440 | expected_row_pattern: str, |
| 441 | terminated: bool, |
| 442 | ): |
| 443 | manager, request = _build_harness( |
| 444 | tokenizer, |
| 445 | backend, |
| 446 | use_reasoner=False, |
| 447 | structured_outputs=StructuredOutputsParams(regex=DIGITS_REGEX), |
| 448 | ) |
| 449 | validated = _to_token_ids( |
| 450 | tokenizer, raw_drafts if expected_validated is None else expected_validated |
| 451 | ) |
| 452 | _run_real_flow( |
| 453 | manager, |
| 454 | request, |
| 455 | raw_drafts=_to_token_ids(tokenizer, raw_drafts), |
| 456 | expected_validated=validated, |
| 457 | expected_row_pattern=expected_row_pattern, |
| 458 | # accept_tokens() latches reasoning_ended once it accepts any token. |
| 459 | expected_reasoning=True if validated else None, |
| 460 | expect_terminated=terminated, |
| 461 | ) |
| 462 | |
| 463 | |
| 464 | @pytest.mark.parametrize("backend", REGEX_BACKENDS) |
| 465 | def test_rejected_draft_keeps_later_rows_constrained( |
| 466 | tokenizer, backend: StructuredOutputsBackend |
| 467 | ): |
| 468 | """A draft the grammar rejects during bitmask fill must not unconstrain |
| 469 | the remaining rows or the bonus row.""" |
| 470 | manager, request = _build_harness( |
| 471 | tokenizer, |
| 472 | backend, |
| 473 | use_reasoner=False, |
| 474 | structured_outputs=StructuredOutputsParams(regex=DIGITS_REGEX), |
| 475 | ) |
| 476 | drafts = _to_token_ids(tokenizer, ("1", "z", "2")) |
| 477 | bitmask = manager.grammar_bitmask( |
| 478 | requests={request.request_id: request}, |
| 479 | structured_output_request_ids=[request.request_id], |
| 480 | scheduled_spec_decode_tokens={request.request_id: drafts}, |
| 481 | ) |
| 482 | assert bitmask is not None |
| 483 | assert _row_pattern(bitmask) == "CCCC" |
| 484 | assert (bitmask[2] == bitmask[1]).all() |
| 485 | assert (bitmask[3] == bitmask[1]).all() |
| 486 | |
| 487 | |
| 488 | def test_outlines_termination(tokenizer): |
| 489 | """outlines_core never advances on EOS; termination must still be |
| 490 | tracked exactly once per accepted EOS and undone by rollback.""" |
| 491 | manager, request = _build_harness( |
| 492 | tokenizer, |
| 493 | "outlines", |
| 494 | use_reasoner=False, |
| 495 | structured_outputs=StructuredOutputsParams(regex=DIGITS_REGEX), |
| 496 | ) |
| 497 | structured_req = request.structured_output_request |
| 498 | assert structured_req is not None |
| 499 | grammar = structured_req.grammar |
| 500 | assert isinstance(grammar, OutlinesGrammar) |
| 501 | one, eos = _to_token_ids(tokenizer, ("1", EOS)) |
| 502 | |
| 503 | assert not grammar.accept_tokens(request.request_id, [eos]) |
| 504 | assert grammar.accept_tokens(request.request_id, [one]) |
| 505 | # Polling must not change the result. |
| 506 | assert not any(grammar.is_terminated() for _ in range(3)) |
| 507 | assert grammar.accept_tokens(request.request_id, [eos]) |
| 508 | assert all(grammar.is_terminated() for _ in range(3)) |
| 509 | assert grammar.num_processed_tokens == 2 |
| 510 | |
| 511 | grammar.rollback(2) |
| 512 | assert not grammar.is_terminated() |
| 513 | assert grammar.num_processed_tokens == 0 |
| 514 | assert grammar.validate_tokens([eos]) == [] |
| 515 | assert grammar.accept_tokens(request.request_id, [one, eos, one]) |
| 516 | assert grammar.is_terminated() |