feat: Naive support Spec V2 + Constrained Decoding (#13425)

Signed-off-by: Ubospica <ubospica@gmail.com>
Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>

python/sglang/srt/managers/schedule_batch.py

@@ -1924,6 +1924,8 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             dimensions=self.dimensions,
             dllm_block_offsets=[req.dllm_block_offset for req in self.reqs],
             dllm_config=self.dllm_config,
+            reqs=self.reqs,
+            has_grammar=self.has_grammar,
         )
 
     def copy(self):
@@ -2041,3 +2043,8 @@ class ModelWorkerBatch:
     # Diffusion LLM
     dllm_block_offsets: Optional[List[int]] = None
     dllm_config: Optional[DllmConfig] = None
+
+    # For constrained decoding
+    # FIXME(lsyin): remove this after fully overlap grammar
+    reqs: Optional[List[Req]] = None
+    has_grammar: bool = False

python/sglang/srt/managers/scheduler.py

@@ -1016,7 +1016,16 @@ class Scheduler(
                 and self.last_batch.forward_mode.is_extend()
             )
 
-            if disable_overlap_for_batch:
+            # FIXME(lsyin): remove this grammar sync
+            need_grammar_sync = (
+                batch is not None
+                and batch.forward_mode.is_decode()
+                and batch.has_grammar
+                and batch.is_v2_eagle
+                and len(self.result_queue) > 0
+            )
+
+            if disable_overlap_for_batch or need_grammar_sync:
                 pop_and_process()
 
             batch_result = None
@@ -1025,7 +1034,7 @@ class Scheduler(
                 self.result_queue.append((batch.copy(), batch_result))
 
             if self.last_batch:
-                if not disable_overlap_for_batch:
+                if not disable_overlap_for_batch and not need_grammar_sync:
                     pop_and_process()
             elif batch is None:
                 # When the server is idle, do self-check and re-init some states

python/sglang/srt/managers/scheduler_output_processor_mixin.py

@@ -395,10 +395,16 @@ class SchedulerOutputProcessorMixin:
                     logits_output.hidden_states[i].cpu().clone().tolist()
                 )
 
-            if req.grammar is not None and batch.spec_algorithm.is_none():
+            if req.grammar is not None:
                 # FIXME: this try-except block is for handling unexpected xgrammar issue.
                 try:
-                    req.grammar.accept_token(next_token_id)
+                    if batch.spec_algorithm.is_none():
+                        # Normal decode: single token
+                        req.grammar.accept_token(next_token_id)
+                    elif batch.is_v2_eagle:
+                        # Speculative decode: next_token_id is a list of accepted tokens
+                        for token_id in next_token_id:
+                            req.grammar.accept_token(token_id)
                 except ValueError as e:
                     # Grammar accept_token can raise ValueError if the token is not in the grammar.
                     # This can happen if the grammar is not set correctly or the token is invalid.

python/sglang/srt/speculative/eagle_info_v2.py

@@ -256,6 +256,7 @@ class EagleVerifyInputV2Mixin:
         self: EagleVerifyInput,
         batch: ModelWorkerBatch,
         logits_output: LogitsProcessorOutput,
+        vocab_mask: torch.Tensor = None,
     ):
         """
         Verify and find accepted tokens based on logits output and batch
@@ -276,6 +277,13 @@ class EagleVerifyInputV2Mixin:
         next_token_logits = logits_output.next_token_logits
         device = batch.input_ids.device
 
+        # Apply grammar mask if provided
+        if vocab_mask is not None:
+            assert self.grammar is not None
+            self.grammar.apply_vocab_mask(
+                logits=next_token_logits, vocab_mask=vocab_mask
+            )
+
         candidates = self.draft_token.reshape(bs, self.draft_token_num)
         predict_shape = list(next_token_logits.shape)[:-1]
         predict = torch.zeros(predict_shape, dtype=torch.int32, device=device).flatten()

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -36,6 +36,7 @@ from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.speculative.spec_utils import (
     detect_nan,
     draft_tp_context,
+    generate_token_bitmask,
     load_token_map,
 )
 from sglang.srt.utils.common import (
@@ -667,7 +668,15 @@ class EAGLEWorkerV2(BaseSpecWorker):
                 ),
             )
 
-        # Run target verify batch in the main compute stream
+        # Prepare grammar data on CPU if needed
+        if batch.has_grammar:
+            retrieve_next_token_cpu = verify_input.retrive_next_token.cpu()
+            retrieve_next_sibling_cpu = verify_input.retrive_next_sibling.cpu()
+            draft_tokens_cpu = verify_input.draft_token.view(
+                verify_input.retrive_next_token.shape
+            ).cpu()
+
+        # Run target verify batch in the main compute stream (GPU compute)
         forward_batch_output = self.target_worker.forward_batch_generation(
             model_worker_batch=None,
             forward_batch=verify_forward_batch,
@@ -676,6 +685,26 @@ class EAGLEWorkerV2(BaseSpecWorker):
         )
         logits_output = forward_batch_output.logits_output
 
+        # Generate vocab mask for constrained decoding
+        vocab_mask = None
+        if batch.has_grammar:
+            # Generate the logit mask for structured output.
+            vocab_mask = generate_token_bitmask(
+                batch.reqs,
+                verify_input,
+                retrieve_next_token_cpu,
+                retrieve_next_sibling_cpu,
+                draft_tokens_cpu,
+                batch.sampling_info.vocab_size,
+            )
+
+            if vocab_mask is not None:
+                assert verify_input.grammar is not None
+                vocab_mask = vocab_mask.to(verify_input.retrive_next_token.device)
+                # NOTE: otherwise, this vocab mask will be the one from the previous extend stage
+                # and will be applied to produce wrong results
+                batch.sampling_info.vocab_mask = None
+
         # Sample
         if self.enable_nan_detection:
             detect_nan(logits_output)
@@ -683,7 +712,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
             predict,
             accept_length,
             accept_index,
-        ) = verify_input.sample(batch, logits_output)
+        ) = verify_input.sample(batch, logits_output, vocab_mask)
         new_seq_lens = batch.seq_lens + accept_length
         verify_done = torch.get_device_module(self.device).Event()
         verify_done.record()

python/sglang/test/kits/json_constrained_kit.py

@@ -24,7 +24,10 @@ class TestJSONConstrainedMixin:
         response = requests.post(
             self.base_url + "/generate",
             json={
-                "text": "The capital of France is",
+                "text": (
+                    "Introduce the capital of France. Return in a JSON format. The JSON Schema is: "
+                    + json.dumps(json_schema)
+                ),
                 "sampling_params": {
                     "temperature": 0 if n == 1 else 0.5,
                     "max_new_tokens": 128,
@@ -69,7 +72,8 @@ class TestJSONConstrainedMixin:
                 {"role": "system", "content": "You are a helpful AI assistant"},
                 {
                     "role": "user",
-                    "content": "Introduce the capital of France. Return in a JSON format.",
+                    "content": "Introduce the capital of France. Return in a JSON format. "
+                    "The JSON Schema is: " + json.dumps(self.json_schema),
                 },
             ],
             temperature=0,

test/srt/run_suite.py

@@ -7,6 +7,7 @@ from sglang.test.ci.ci_utils import TestFile, run_unittest_files
 # NOTE: please sort the test cases alphabetically by the test file name
 suites = {
     "per-commit-1-gpu": [
+        TestFile("test_eagle_constrained_decoding.py", 100),
         TestFile("debug_utils/test_tensor_dump_forward_hook.py", 15),
         TestFile("hicache/test_hicache_storage.py", 127),
         TestFile("hicache/test_hicache_variants.py", 393),

test/srt/test_eagle_constrained_decoding.py

@@ -0,0 +1,77 @@
+import unittest
+
+from sglang.srt.environ import envs
+from sglang.srt.utils import kill_process_tree
+from sglang.test.kits.json_constrained_kit import TestJSONConstrainedMixin
+from sglang.test.kits.regex_constrained_kit import TestRegexConstrainedMixin
+from sglang.test.test_utils import (
+    DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST,
+    DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST,
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+
+class TestEagleConstrainedDecoding(
+    CustomTestCase, TestRegexConstrainedMixin, TestJSONConstrainedMixin
+):
+    max_running_requests = 64
+    attention_backend = "triton"
+    spec_steps = 5
+    spec_topk = 1
+    spec_draft_tokens = 6
+    page_size = 1
+    other_launch_args = []
+    model = DEFAULT_EAGLE_TARGET_MODEL_FOR_TEST
+    draft_model = DEFAULT_EAGLE_DRAFT_MODEL_FOR_TEST
+    grammar_backend = "xgrammar"
+    eagle_v2 = False
+
+    @classmethod
+    def setUpClass(cls):
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        launch_args = [
+            "--trust-remote-code",
+            "--attention-backend",
+            cls.attention_backend,
+            "--speculative-algorithm",
+            "EAGLE",
+            "--speculative-draft-model",
+            cls.draft_model,
+            "--speculative-num-steps",
+            cls.spec_steps,
+            "--speculative-eagle-topk",
+            cls.spec_topk,
+            "--speculative-num-draft-tokens",
+            cls.spec_draft_tokens,
+            "--page-size",
+            str(cls.page_size),
+            "--mem-fraction-static",
+            "0.75",
+            "--max-running-requests",
+            str(cls.max_running_requests),
+            "--grammar-backend",
+            cls.grammar_backend,
+        ]
+        launch_args.extend(cls.other_launch_args)
+        with envs.SGLANG_ENABLE_SPEC_V2.override(cls.eagle_v2):
+            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)
+
+
+class TestEagleConstrainedDecodingV2(TestEagleConstrainedDecoding):
+    eagle_v2 = True
+
+
+if __name__ == "__main__":
+    unittest.main()