[Bugfix][Structured Output] Stop XGrammar token batches at termination (#52805)
Signed-off-by: sfeng33 <4florafeng@gmail.com>
Flora Feng committed
Aug 18, 2026 at 17:21 UTC
12f64b39d29282437e35be9aa5db432fb2a1a6e6
2 files changed
+62
-5
tests/v1/spec_decode/test_mtp_structured_output.py
+48
@@ -262,6 +262,54 @@ def test_validate_tokens_then_bitmask_round_trip(backend):
262
assert not grammar.is_terminated()
263
264
265
+def test_xgrammar_accept_tokens_stops_at_termination(capfd):
266
+ """Tokens after a terminating EOS do not reach the matcher."""
267
+ tokenizer, _, request, prompt = _make_manager_and_request("xgrammar")
268
+ grammar = request.structured_output_request.grammar
269
+
270
+ assert grammar.accept_tokens(request.request_id, prompt)
271
+
272
+ eos = tokenizer.eos_token_id
273
+ trailing = tokenizer.encode("\n")[0]
274
+ processed_before = grammar.num_processed_tokens
275
+
276
+ assert grammar.accept_tokens(request.request_id, [eos, trailing])
277
+ assert grammar.is_terminated()
278
+ assert grammar.num_processed_tokens == processed_before + 1
279
+ assert "trying to accept new token" not in capfd.readouterr().err
280
+
281
+ processed_after_eos = grammar.num_processed_tokens
282
+ assert grammar.accept_tokens(request.request_id, [trailing])
283
+ assert grammar.num_processed_tokens == processed_after_eos
284
+ assert "trying to accept new token" not in capfd.readouterr().err
285
+
286
+ grammar.reset()
287
+ assert not grammar.is_terminated()
288
+ assert grammar.num_processed_tokens == 0
289
+
290
+
291
+def test_xgrammar_validate_tokens_stops_at_termination(capfd):
292
+ """Validation rolls back after reaching a terminating EOS."""
293
+ tokenizer, _, request, prompt = _make_manager_and_request("xgrammar")
294
+ grammar = request.structured_output_request.grammar
295
+
296
+ assert grammar.accept_tokens(request.request_id, prompt)
297
+
298
+ eos = tokenizer.eos_token_id
299
+ trailing = tokenizer.encode("\n")[0]
300
+
301
+ assert grammar.validate_tokens([eos, trailing]) == [eos]
302
+ assert "trying to accept new token" not in capfd.readouterr().err
303
+ # Check matcher state directly to verify validation rolled it back.
304
+ assert not grammar.matcher.is_terminated()
305
+
306
+ assert grammar.accept_tokens(request.request_id, [eos])
307
+ assert grammar.is_terminated()
308
+
309
+ assert grammar.validate_tokens([trailing]) == []
310
+ assert "trying to accept new token" not in capfd.readouterr().err
311
+
312
+
313
class _MarkerReasoner:
314
"""Stub reasoner whose reasoning-end marker is a single fixed token."""
315
vllm/v1/structured_output/backend_xgrammar.py
+14
-5
@@ -157,11 +157,12 @@ class XgrammarGrammar(StructuredOutputGrammar):
157
def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
158
"""Accepts a list of tokens and advances the FSM.
159
160
- Returns True if the FSM was advanced successfully.
161
- Returns False if the FSM failed to advance.
160
+ Returns True if all grammar-constrained tokens were accepted.
161
+ Tokens after termination are ignored. Returns False if the FSM
162
+ failed to advance.
163
"""
164
if self._is_terminated:
164
- return False
165
+ return True
166
for token in tokens:
167
if not self.matcher.accept_token(token):
168
logger.error(
@@ -172,7 +173,9 @@ class XgrammarGrammar(StructuredOutputGrammar):
173
)
174
return False
175
self.num_processed_tokens += 1
175
- self._is_terminated = self.matcher.is_terminated()
176
+ self._is_terminated = self.matcher.is_terminated()
177
+ if self._is_terminated:
178
+ break
179
return True
180
181
def validate_tokens(self, tokens: list[int]) -> list[int]:
@@ -181,10 +184,15 @@ class XgrammarGrammar(StructuredOutputGrammar):
184
185
Returns the prefix list of tokens that are accepted by the FSM.
186
"""
187
+ if self._is_terminated:
188
+ return []
189
+
190
accepted_tokens = []
191
for token in tokens:
192
if self.matcher.accept_token(token):
193
accepted_tokens.append(token)
194
+ if self.matcher.is_terminated():
195
+ break
196
else:
197
break
198
if len(accepted_tokens) > 0:
@@ -204,8 +212,9 @@ class XgrammarGrammar(StructuredOutputGrammar):
212
return self._is_terminated
213
214
def reset(self):
207
- self.num_processed_tokens = 0
215
self.matcher.reset()
216
+ self.num_processed_tokens = 0
217
+ self._is_terminated = False
218
219
220
# cf https://github.com/mlc-ai/xgrammar/blob/a32ac892676d2eedc0327416105b9b06edfb94b2/cpp/json_schema_converter.cc