Support grammar + spec + reasoning (#14163)

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:

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()

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,

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),

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()