From 0a9d64530dba85724f164b010dfc0c73ca38ec98 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 30 Nov 2025 21:19:57 +0800 Subject: [PATCH] Support grammar + spec + reasoning (#14163) --- .../constrained/reasoner_grammar_backend.py | 33 ++++-- .../srt/constrained/xgrammar_backend.py | 36 +++++++ python/sglang/srt/speculative/spec_utils.py | 2 - test/srt/run_suite.py | 1 + ...est_constrained_decoding_spec_reasoning.py | 100 ++++++++++++++++++ 5 files changed, 164 insertions(+), 8 deletions(-) create mode 100644 test/srt/test_constrained_decoding_spec_reasoning.py diff --git a/python/sglang/srt/constrained/reasoner_grammar_backend.py b/python/sglang/srt/constrained/reasoner_grammar_backend.py index 57fd55a3b..5fe7b0c8c 100644 --- a/python/sglang/srt/constrained/reasoner_grammar_backend.py +++ b/python/sglang/srt/constrained/reasoner_grammar_backend.py @@ -29,14 +29,35 @@ class ReasonerGrammarObject(BaseGrammarObject): super().__init__() self.grammar = grammar self.think_end_id = think_end_id - self.is_in_reasoning = True + # -1 means thinking has not ended yet + # 0 means just ended thinking in the last token + # + means number of tokens after thinking ended + self.tokens_after_think_end = -1 + + def transfer_state(self, token: int) -> int: + if self.tokens_after_think_end == -1 and token == self.think_end_id: + self.tokens_after_think_end = 0 + elif self.tokens_after_think_end >= 0: + self.tokens_after_think_end += 1 + + def rollback_state(self): + if self.tokens_after_think_end == 0: + self.tokens_after_think_end = -1 + elif self.tokens_after_think_end > 0: + self.tokens_after_think_end -= 1 def accept_token(self, token: int): - if token == self.think_end_id: - self.is_in_reasoning = False - - if not self.is_in_reasoning and token != self.think_end_id: + if self.tokens_after_think_end >= 0: self.grammar.accept_token(token) + self.transfer_state(token) + + def rollback(self, k): + steps_after_think = min(k, self.tokens_after_think_end) + if steps_after_think > 0: + self.grammar.rollback(steps_after_think) + + for _ in range(k): + self.rollback_state() def allocate_vocab_mask( self, vocab_size: int, batch_size: int, device @@ -44,7 +65,7 @@ class ReasonerGrammarObject(BaseGrammarObject): return self.grammar.allocate_vocab_mask(vocab_size, batch_size, device) def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: - if not self.is_in_reasoning: + if self.tokens_after_think_end >= 0: self.grammar.fill_vocab_mask(vocab_mask, idx) def move_vocab_mask(self, vocab_mask: torch.Tensor, device) -> torch.Tensor: diff --git a/python/sglang/srt/constrained/xgrammar_backend.py b/python/sglang/srt/constrained/xgrammar_backend.py index 3cb9d3841..5f2806e3b 100644 --- a/python/sglang/srt/constrained/xgrammar_backend.py +++ b/python/sglang/srt/constrained/xgrammar_backend.py @@ -305,3 +305,39 @@ class XGrammarGrammarBackend(BaseGrammarBackend): def reset(self): self.grammar_compiler.clear_cache() + + +def demo_test(): + from transformers import AutoConfig, AutoTokenizer + + from sglang.test.test_utils import DEFAULT_MODEL_NAME_FOR_TEST + + tokenizer = AutoTokenizer.from_pretrained(DEFAULT_MODEL_NAME_FOR_TEST) + hf_config = AutoConfig.from_pretrained(DEFAULT_MODEL_NAME_FOR_TEST) + + # Should use vocab size from model config + vocab_size = hf_config.vocab_size + eos_token_id = tokenizer.eos_token_id + + backend = XGrammarGrammarBackend( + tokenizer, vocab_size=vocab_size, model_eos_token_ids=[eos_token_id] + ) + regex = r"hello (world|there)" + grammar = backend.dispatch_regex(regex) + tokens = [ + tokenizer.encode(t, add_special_tokens=False)[0] for t in ["hello", " world"] + ] + + # Test termination + grammar.accept_token(tokens[0]) # accept "hello" + grammar.accept_token(tokens[1]) # accept " world" + grammar.accept_token(eos_token_id) # accept EOS + assert grammar.is_terminated() + + # Test rollback the terminated state + grammar.rollback(1) + assert not grammar.is_terminated() + + +if __name__ == "__main__": + demo_test() diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 5e215156b..108b4f0f8 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/python/sglang/srt/speculative/spec_utils.py @@ -581,8 +581,6 @@ def traverse_tree( retrieve_next_token.shape == retrieve_next_sibling.shape == draft_tokens.shape ) - allocate_token_bitmask.fill_(0) - def dfs( curr: int, retrieve_next_token: torch.Tensor, diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index e55a1a4dc..f427562f3 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -144,6 +144,7 @@ suites = { TestFile("models/test_glm4_moe_models.py", 100), TestFile("models/test_kimi_linear_models.py", 90), TestFile("rl/test_update_weights_from_distributed.py", 103), + TestFile("test_constrained_decoding_spec_reasoning.py", 60), TestFile("test_data_parallelism.py", 73), TestFile("test_disaggregation_basic.py", 400), TestFile("test_dp_attention.py", 350), diff --git a/test/srt/test_constrained_decoding_spec_reasoning.py b/test/srt/test_constrained_decoding_spec_reasoning.py new file mode 100644 index 000000000..1370cd0d0 --- /dev/null +++ b/test/srt/test_constrained_decoding_spec_reasoning.py @@ -0,0 +1,100 @@ +import json +import unittest + +import openai + +from sglang.srt.utils import kill_process_tree +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + popen_launch_server, +) + + +class ServerWithGrammar(CustomTestCase): + json_schema = json.dumps( + { + "type": "object", + "properties": { + "name": {"type": "string", "pattern": "^[\\w]+$"}, + "population": {"type": "integer"}, + "languages": { + "type": "array", + "items": {"type": "string"}, + "minItems": 1, + }, + "has_held_olympics": {"type": "boolean"}, + }, + "required": ["name", "population", "languages", "has_held_olympics"], + "additionalProperties": False, + } + ) + + @classmethod + def setUpClass(cls): + cls.model = "openai/gpt-oss-120b" + cls.base_url = DEFAULT_URL_FOR_TEST + launch_args = [ + "--trust-remote-code", + "--tp=2", + "--reasoning-parser=gpt-oss", + "--speculative-algorithm=EAGLE3", + "--speculative-draft-model-path=lmsys/EAGLE3-gpt-oss-120b-bf16", + "--speculative-num-steps=5", + "--speculative-eagle-topk=4", + "--speculative-num-draft-tokens=8", + ] + + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=launch_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_json_openai(self): + client = openai.Client(api_key="EMPTY", base_url=f"{self.base_url}/v1") + + response = client.chat.completions.create( + model=self.model, + messages=[ + {"role": "system", "content": "You are a helpful AI assistant"}, + { + "role": "user", + "content": "Introduce the capital of France. Return in a JSON format. " + "The JSON Schema is: " + json.dumps(self.json_schema), + }, + ], + temperature=0, + max_tokens=1024, + response_format={ + "type": "json_schema", + "json_schema": {"name": "foo", "schema": json.loads(self.json_schema)}, + }, + ) + text = response.choices[0].message.content + + print("\n=== Reasoning Content ===") + reasoning_content = response.choices[0].message.reasoning_content + assert reasoning_content is not None and len(reasoning_content) > 0 + print(reasoning_content) + + try: + js_obj = json.loads(text) + print("\n=== Parsed JSON Content ===") + print(json.dumps(js_obj)) + except (TypeError, json.decoder.JSONDecodeError): + print("JSONDecodeError", text) + raise + + self.assertIsInstance(js_obj["name"], str) + self.assertIsInstance(js_obj["population"], int) + + +if __name__ == "__main__": + unittest.main()