[Bugfix][Structured Output] Mask request stop tokens in xgrammar until grammar terminates (#50595)
Signed-off-by: Yifan Zong <yzong@redhat.com>
yzong-rh committed
Aug 13, 2026 at 14:07 UTC
e6b2a8ad5eb1757b211475cdd57fdf1f37e94331
2 files changed
+41
-14
tests/parser/test_harmony.py
+38
-9
@@ -12,10 +12,9 @@ from openai_harmony import (
12
RenderConversationConfig,
13
Role,
14
)
15
-from transformers import AutoTokenizer
16
-from xgrammar import Grammar
17
-from xgrammar.testing import _is_grammar_accept_string
15
+from transformers import AutoTokenizer, GenerationConfig
16
17
+from vllm.config import StructuredOutputsConfig, VllmConfig
18
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
19
from vllm.entrypoints.openai.engine.protocol import FunctionCall
20
from vllm.entrypoints.openai.parser.harmony_utils import (
@@ -25,6 +24,8 @@ from vllm.entrypoints.openai.responses.protocol import ResponsesRequest
24
from vllm.parser.harmony import HarmonyParser
25
from vllm.parser.parser_manager import ParserManager
26
from vllm.sampling_params import StructuredOutputsParams
27
+from vllm.v1.structured_output.backend_types import StructuredOutputOptions
28
+from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend
29
30
REASONING_MODEL_NAME = "openai/gpt-oss-20b"
31
@@ -34,6 +35,26 @@ def gpt_oss_tokenizer():
35
return AutoTokenizer.from_pretrained(REASONING_MODEL_NAME)
36
37
38
+@pytest.fixture(scope="module")
39
+def gpt_oss_stop_token_ids() -> set[int]:
40
+ eos_token_id = GenerationConfig.from_pretrained(REASONING_MODEL_NAME).eos_token_id
41
+ if isinstance(eos_token_id, int):
42
+ return {eos_token_id}
43
+ return set(eos_token_id)
44
+
45
+
46
+@pytest.fixture(scope="module")
47
+def xgrammar_backend(gpt_oss_tokenizer) -> XgrammarBackend:
48
+ vllm_config = VllmConfig(
49
+ structured_outputs_config=StructuredOutputsConfig(backend="xgrammar")
50
+ )
51
+ return XgrammarBackend(
52
+ vllm_config,
53
+ tokenizer=gpt_oss_tokenizer,
54
+ vocab_size=len(gpt_oss_tokenizer),
55
+ )
56
+
57
+
58
@pytest.fixture
59
def harmony_parser(gpt_oss_tokenizer):
60
parser_cls = ParserManager.get_parser(
@@ -1036,21 +1057,25 @@ class TestAdjustRequest:
1057
cls,
1058
adjusted_request: ChatCompletionRequest | ResponsesRequest,
1059
expected_admission: Sequence[str],
1060
+ xgrammar_backend: XgrammarBackend,
1061
+ stop_token_ids: set[int],
1062
) -> None:
1063
structured_outputs = adjusted_request.structured_outputs
1064
assert structured_outputs is not None
1065
assert structured_outputs.structural_tag is not None
1066
assert structured_outputs.all_non_structural_tag_constraints_none()
1067
1045
- grammar = Grammar.from_structural_tag(structured_outputs.structural_tag)
1068
+ grammar = xgrammar_backend.compile_grammar(
1069
+ StructuredOutputOptions.STRUCTURAL_TAG,
1070
+ structured_outputs.structural_tag,
1071
+ stop_token_ids=stop_token_ids,
1072
+ )
1073
expected_admission_set = set(expected_admission)
1074
1075
for sample_name in cls.ADMISSION_SAMPLES:
1049
- admitted = _is_grammar_accept_string(
1050
- grammar,
1051
- getattr(cls, sample_name),
1052
- require_termination=False,
1053
- )
1076
+ tokens = encode_output(getattr(cls, sample_name))
1077
+ accepted = grammar.validate_tokens(tokens)
1078
+ admitted = accepted == tokens
1079
should_admit = sample_name in expected_admission_set
1080
assert admitted is should_admit, (
1081
f"Expected structured_outputs admission for {sample_name} "
@@ -1183,6 +1208,8 @@ class TestAdjustRequest:
1208
def test_adjust_request(
1209
self,
1210
harmony_parser,
1211
+ xgrammar_backend,
1212
+ gpt_oss_stop_token_ids,
1213
request_kind,
1214
request_kwargs,
1215
expected_admission,
@@ -1193,4 +1220,6 @@ class TestAdjustRequest:
1220
self._assert_structured_outputs_admission(
1221
adjusted_request,
1222
expected_admission,
1223
+ xgrammar_backend,
1224
+ gpt_oss_stop_token_ids,
1225
)
vllm/parser/harmony.py
+3
-5
@@ -388,11 +388,9 @@ class HarmonyParser(DelegatingParser):
388
389
390
# Harmomy's stop tokens are <|return|>, <|call|>, <|endoftext|>
391
-# <|return|> is represented as "" since it's the default stop token, which xgrammar
392
-# disallows under constraints, leading to bad or infinite generation.
393
-# StreamableParser doesn't consider <|endoftext|> as a message end, so it's excluded
394
-# TODO: Remove <|call|> once #50595 lands.
395
-_END_TAG = ["<|end|>", "<|call|>", ""]
391
+# They are represented as "" since xgrammar disallows stop tokens while
392
+# under constraints, leading to bad or infinite generations.
393
+_END_TAG = ["<|end|>", ""]
394
_FINAL_BEGIN = "<|channel|>final{constrain}<|message|>"
395
_TOOL_CALL_CHANNELS = [
396
"<|channel|>commentary",