| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Regression tests for issue #42403. |
| 4 | |
| 5 | A request's stop-token set (the generation_config eos list plus any |
| 6 | user-supplied ``stop_token_ids``) is invisible to xgrammar, which only knows |
| 7 | the tokenizer's single eos. Such tokens can therefore escape the grammar |
| 8 | bitmask while the FSM is still mid-object and truncate structured output. |
| 9 | ``compile_grammar`` now forwards ``all_stop_token_ids`` to the matcher as |
| 10 | ``override_stop_tokens`` so xgrammar masks them until the grammar completes. |
| 11 | """ |
| 12 | |
| 13 | import pytest |
| 14 | from transformers import AutoTokenizer |
| 15 | |
| 16 | from vllm.config import StructuredOutputsConfig, VllmConfig |
| 17 | from vllm.v1.structured_output.backend_types import StructuredOutputOptions |
| 18 | from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend |
| 19 | |
| 20 | TOKENIZER = "openai-community/gpt2" |
| 21 | VOCAB_SIZE = 50257 |
| 22 | |
| 23 | # gpt2 token ids used to drive a `{"type": "string"}` grammar deterministically. |
| 24 | EOS = 50256 # <|endoftext|> -- the tokenizer's only default stop token |
| 25 | QUOTE = 1 # standalone `"`; opens then closes the JSON string |
| 26 | LETTER = 55 # `X`: valid string content, not a special/stop token by default |
| 27 | |
| 28 | |
| 29 | def _token_allowed(row, token_id: int) -> bool: |
| 30 | word = int(row[token_id // 32].item()) & 0xFFFFFFFF |
| 31 | return bool(word & (1 << (token_id % 32))) |
| 32 | |
| 33 | |
| 34 | @pytest.fixture(scope="module") |
| 35 | def backend() -> XgrammarBackend: |
| 36 | vllm_config = VllmConfig( |
| 37 | structured_outputs_config=StructuredOutputsConfig(backend="xgrammar") |
| 38 | ) |
| 39 | return XgrammarBackend( |
| 40 | vllm_config, |
| 41 | tokenizer=AutoTokenizer.from_pretrained(TOKENIZER), |
| 42 | vocab_size=VOCAB_SIZE, |
| 43 | ) |
| 44 | |
| 45 | |
| 46 | def test_request_stop_tokens_gated_to_grammar_terminal(backend: XgrammarBackend): |
| 47 | schema = '{"type": "string"}' |
| 48 | default = backend.compile_grammar(StructuredOutputOptions.JSON, schema) |
| 49 | override = backend.compile_grammar( |
| 50 | StructuredOutputOptions.JSON, schema, stop_token_ids={EOS, LETTER} |
| 51 | ) |
| 52 | |
| 53 | # Open the string: both grammars are now in a non-terminal state. |
| 54 | for grammar in (default, override): |
| 55 | assert grammar.accept_tokens("req", [QUOTE]) |
| 56 | |
| 57 | bm_default = backend.allocate_token_bitmask(1) |
| 58 | bm_override = backend.allocate_token_bitmask(1) |
| 59 | default.fill_bitmask(bm_default, 0) |
| 60 | override.fill_bitmask(bm_override, 0) |
| 61 | |
| 62 | # Mid-string, the plain token is valid content, so the default grammar |
| 63 | # leaves it samplable -- this is the leak. Registering it as a stop token |
| 64 | # masks it until the grammar can terminate. |
| 65 | assert _token_allowed(bm_default[0], LETTER) |
| 66 | assert not _token_allowed(bm_override[0], LETTER) |
| 67 | |
| 68 | # Close the string -> accepting state (grammar complete, not yet terminated). |
| 69 | for grammar in (default, override): |
| 70 | assert grammar.accept_tokens("req", [QUOTE]) |
| 71 | assert not grammar.is_terminated() |
| 72 | |
| 73 | default.fill_bitmask(bm_default, 0) |
| 74 | override.fill_bitmask(bm_override, 0) |
| 75 | |
| 76 | # The extra stop token may now terminate under the override, never under |
| 77 | # the default grammar -- and the tokenizer's own eos still terminates both, |
| 78 | # so default termination is preserved. |
| 79 | assert not _token_allowed(bm_default[0], LETTER) |
| 80 | assert _token_allowed(bm_override[0], LETTER) |
| 81 | assert _token_allowed(bm_default[0], EOS) |
| 82 | assert _token_allowed(bm_override[0], EOS) |