[Bugfix][Structured Output] Mask request stop tokens in xgrammar until grammar terminates (#49227)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
Flora Feng committed
Aug 10, 2026 at 14:29 UTC
05f0a80016a4a12c5f4e241e92e233acd1c2c0b0
7 files changed
+114
-6
tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py
new
+82
@@ -0,0 +1,82 @@
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)
vllm/v1/structured_output/__init__.py
+8
-1
@@ -184,7 +184,14 @@ class StructuredOutputManager:
184
try:
185
request_type, grammar_spec = struct_request.structured_output_key
186
assert self.backend is not None
187
- return self.backend.compile_grammar(request_type, grammar_spec)
187
+ stop_token_ids = (
188
+ request.sampling_params.all_stop_token_ids
189
+ if request.sampling_params is not None
190
+ else None
191
+ )
192
+ return self.backend.compile_grammar(
193
+ request_type, grammar_spec, stop_token_ids=stop_token_ids
194
+ )
195
except Exception:
196
logger.exception(
197
"Failed to compile grammar for request %s", request.request_id
vllm/v1/structured_output/backend_guidance.py
+4
-1
@@ -106,7 +106,10 @@ class GuidanceBackend(StructuredOutputBackend):
106
)
107
108
def compile_grammar(
109
- self, request_type: StructuredOutputOptions, grammar_spec: str
109
+ self,
110
+ request_type: StructuredOutputOptions,
111
+ grammar_spec: str,
112
+ stop_token_ids: set[int] | None = None,
113
) -> StructuredOutputGrammar:
114
self.serialized_grammar = serialize_guidance_grammar(
115
request_type,
vllm/v1/structured_output/backend_lm_format_enforcer.py
+4
-1
@@ -99,7 +99,10 @@ class LMFormatEnforcerBackend(StructuredOutputBackend):
99
)
100
101
def compile_grammar(
102
- self, request_type: StructuredOutputOptions, grammar_spec: str
102
+ self,
103
+ request_type: StructuredOutputOptions,
104
+ grammar_spec: str,
105
+ stop_token_ids: set[int] | None = None,
106
) -> StructuredOutputGrammar:
107
character_level_parser: lmformatenforcer.CharacterLevelParser
108
if request_type == StructuredOutputOptions.JSON:
vllm/v1/structured_output/backend_outlines.py
+4
-1
@@ -71,7 +71,10 @@ class OutlinesBackend(StructuredOutputBackend):
71
return index
72
73
def compile_grammar(
74
- self, request_type: StructuredOutputOptions, grammar_spec: str
74
+ self,
75
+ request_type: StructuredOutputOptions,
76
+ grammar_spec: str,
77
+ stop_token_ids: set[int] | None = None,
78
) -> StructuredOutputGrammar:
79
if request_type == StructuredOutputOptions.JSON:
80
regex = json_schema.build_regex_from_schema(grammar_spec)
vllm/v1/structured_output/backend_types.py
+7
-1
@@ -105,7 +105,10 @@ class StructuredOutputBackend(ABC):
105
106
@abstractmethod
107
def compile_grammar(
108
- self, request_type: StructuredOutputOptions, grammar_spec: str
108
+ self,
109
+ request_type: StructuredOutputOptions,
110
+ grammar_spec: str,
111
+ stop_token_ids: set[int] | None = None,
112
) -> StructuredOutputGrammar:
113
"""
114
Compiles a grammar specification into a structured output grammar.
@@ -114,6 +117,9 @@ class StructuredOutputBackend(ABC):
117
request_type (StructuredOutputOptions): The type of structured
118
output request.
119
grammar_spec (str): The grammar specification to compile.
120
+ stop_token_ids (set[int] | None): The request's EOS and user
121
+ stop-token ids (``SamplingParams.all_stop_token_ids``), masked
122
+ until the grammar terminates.
123
124
Returns:
125
StructuredOutputGrammar: The compiled structured output grammar.
vllm/v1/structured_output/backend_xgrammar.py
+5
-1
@@ -76,7 +76,10 @@ class XgrammarBackend(StructuredOutputBackend):
76
)
77
78
def compile_grammar(
79
- self, request_type: StructuredOutputOptions, grammar_spec: str
79
+ self,
80
+ request_type: StructuredOutputOptions,
81
+ grammar_spec: str,
82
+ stop_token_ids: set[int] | None = None,
83
) -> StructuredOutputGrammar:
84
if request_type == StructuredOutputOptions.JSON:
85
ctx = self.compiler.compile_json_schema(
@@ -119,6 +122,7 @@ class XgrammarBackend(StructuredOutputBackend):
122
return XgrammarGrammar(
123
matcher=xgr.GrammarMatcher(
124
ctx,
125
+ override_stop_tokens=list(stop_token_ids) if stop_token_ids else None,
126
max_rollback_tokens=self.num_speculative_tokens,
127
),
128
vocab_size=self.vocab_size,