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,