From 6abc540faef14cb3278ab3d5c6d0ae894b6f14fd Mon Sep 17 00:00:00 2001 From: sfeng33 <4florafeng@gmail.com> Date: Tue, 18 Aug 2026 13:51:23 -0400 Subject: [PATCH 1/3] Stop XGrammar token batches at termination Signed-off-by: sfeng33 <4florafeng@gmail.com> --- .../spec_decode/test_mtp_structured_output.py | 39 +++++++++++++++++++ vllm/v1/structured_output/backend_xgrammar.py | 10 ++++- 2 files changed, 47 insertions(+), 2 deletions(-) diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 619f3ad6fded..d5b2e42b75bb 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -262,6 +262,45 @@ def test_validate_tokens_then_bitmask_round_trip(backend): assert not grammar.is_terminated() +def test_xgrammar_accept_tokens_stops_at_termination(capfd): + """Tokens after a terminating EOS do not reach the matcher.""" + tokenizer, _, request, prompt = _make_manager_and_request("xgrammar") + grammar = request.structured_output_request.grammar + + assert grammar.accept_tokens(request.request_id, prompt) + + eos = tokenizer.eos_token_id + trailing = tokenizer.encode("\n")[0] + processed_before = grammar.num_processed_tokens + + assert grammar.accept_tokens(request.request_id, [eos, trailing]) + assert grammar.is_terminated() + assert grammar.num_processed_tokens == processed_before + 1 + assert "trying to accept new token" not in capfd.readouterr().err + + grammar.reset() + assert not grammar.is_terminated() + assert grammar.num_processed_tokens == 0 + + +def test_xgrammar_validate_tokens_stops_at_termination(capfd): + """Validation rolls back after reaching a terminating EOS.""" + tokenizer, _, request, prompt = _make_manager_and_request("xgrammar") + grammar = request.structured_output_request.grammar + + assert grammar.accept_tokens(request.request_id, prompt) + + eos = tokenizer.eos_token_id + trailing = tokenizer.encode("\n")[0] + + assert grammar.validate_tokens([eos, trailing]) == [eos] + assert "trying to accept new token" not in capfd.readouterr().err + assert not grammar.is_terminated() + + assert grammar.accept_tokens(request.request_id, [eos]) + assert grammar.is_terminated() + + class _MarkerReasoner: """Stub reasoner whose reasoning-end marker is a single fixed token.""" diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 258b1dff32f1..447487108e95 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -164,6 +164,7 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: return False for token in tokens: if not self.matcher.accept_token(token): + self._is_terminated = self.matcher.is_terminated() logger.error( "Failed to advance FSM for request %s " "for tokens %s. Please file an issue.", @@ -172,7 +173,9 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: ) return False self.num_processed_tokens += 1 - self._is_terminated = self.matcher.is_terminated() + self._is_terminated = self.matcher.is_terminated() + if self._is_terminated: + break return True def validate_tokens(self, tokens: list[int]) -> list[int]: @@ -185,6 +188,8 @@ def validate_tokens(self, tokens: list[int]) -> list[int]: for token in tokens: if self.matcher.accept_token(token): accepted_tokens.append(token) + if self.matcher.is_terminated(): + break else: break if len(accepted_tokens) > 0: @@ -204,8 +209,9 @@ def is_terminated(self) -> bool: return self._is_terminated def reset(self): - self.num_processed_tokens = 0 self.matcher.reset() + self.num_processed_tokens = 0 + self._is_terminated = False # cf https://github.com/mlc-ai/xgrammar/blob/a32ac892676d2eedc0327416105b9b06edfb94b2/cpp/json_schema_converter.cc From a5492df6c2d1a210929b1a392126f8eca87944bd Mon Sep 17 00:00:00 2001 From: sfeng33 <4florafeng@gmail.com> Date: Tue, 18 Aug 2026 14:20:01 -0400 Subject: [PATCH 2/3] fix Signed-off-by: sfeng33 <4florafeng@gmail.com> --- tests/v1/spec_decode/test_mtp_structured_output.py | 8 ++++++++ vllm/v1/structured_output/backend_xgrammar.py | 6 ++++-- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index d5b2e42b75bb..984703a5a288 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -278,6 +278,11 @@ def test_xgrammar_accept_tokens_stops_at_termination(capfd): assert grammar.num_processed_tokens == processed_before + 1 assert "trying to accept new token" not in capfd.readouterr().err + processed_after_eos = grammar.num_processed_tokens + assert grammar.accept_tokens(request.request_id, [trailing]) + assert grammar.num_processed_tokens == processed_after_eos + assert "trying to accept new token" not in capfd.readouterr().err + grammar.reset() assert not grammar.is_terminated() assert grammar.num_processed_tokens == 0 @@ -300,6 +305,9 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd): assert grammar.accept_tokens(request.request_id, [eos]) assert grammar.is_terminated() + assert grammar.validate_tokens([trailing]) == [] + assert "trying to accept new token" not in capfd.readouterr().err + class _MarkerReasoner: """Stub reasoner whose reasoning-end marker is a single fixed token.""" diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 447487108e95..15ea285b7b69 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -161,10 +161,9 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: Returns False if the FSM failed to advance. """ if self._is_terminated: - return False + return True for token in tokens: if not self.matcher.accept_token(token): - self._is_terminated = self.matcher.is_terminated() logger.error( "Failed to advance FSM for request %s " "for tokens %s. Please file an issue.", @@ -184,6 +183,9 @@ def validate_tokens(self, tokens: list[int]) -> list[int]: Returns the prefix list of tokens that are accepted by the FSM. """ + if self._is_terminated: + return [] + accepted_tokens = [] for token in tokens: if self.matcher.accept_token(token): From dd25c941229b1dc74b802aad28e8927cb83237ae Mon Sep 17 00:00:00 2001 From: sfeng33 <4florafeng@gmail.com> Date: Tue, 18 Aug 2026 15:37:12 -0400 Subject: [PATCH 3/3] address review comments Signed-off-by: sfeng33 <4florafeng@gmail.com> --- tests/v1/spec_decode/test_mtp_structured_output.py | 3 ++- vllm/v1/structured_output/backend_xgrammar.py | 5 +++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 984703a5a288..8bd599733a24 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -300,7 +300,8 @@ def test_xgrammar_validate_tokens_stops_at_termination(capfd): assert grammar.validate_tokens([eos, trailing]) == [eos] assert "trying to accept new token" not in capfd.readouterr().err - assert not grammar.is_terminated() + # Check matcher state directly to verify validation rolled it back. + assert not grammar.matcher.is_terminated() assert grammar.accept_tokens(request.request_id, [eos]) assert grammar.is_terminated() diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 15ea285b7b69..5b24a19780aa 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -157,8 +157,9 @@ class XgrammarGrammar(StructuredOutputGrammar): def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: """Accepts a list of tokens and advances the FSM. - Returns True if the FSM was advanced successfully. - Returns False if the FSM failed to advance. + Returns True if all grammar-constrained tokens were accepted. + Tokens after termination are ignored. Returns False if the FSM + failed to advance. """ if self._is_terminated: return True