[DSPARK] Grammar-constrained decoding, incl. tool_choice=auto (#31753)

Co-authored-by: shanemort1982 <shanemort1982@users.noreply.github.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>

python/sglang/srt/managers/scheduler.py

@@ -2253,9 +2253,7 @@ class Scheduler(
         self._maybe_namespace_elastic_radix_cache(req)
 
         if self.spec_algorithm.is_dflash_family():
-            error_msg = validate_dflash_request(
+            error_msg = validate_dflash_request(req, self.enable_overlap)
-                req, self.enable_overlap, self.spec_algorithm
-            )
             if error_msg is not None:
                 req.set_finish_with_abort(error_msg)
                 self.init_req_max_new_tokens(req)

python/sglang/srt/speculative/dflash_info_v2.py

@@ -6,6 +6,7 @@ from typing import List, Optional
 
 import torch
 
+from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
 from sglang.srt.environ import envs
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.mem_cache.allocation import alloc_for_spec_decode
@@ -56,6 +57,9 @@ class DFlashDraftInputV2(SpecInput):
 
     verify_token_budget: Optional[int] = None
 
+    # Stamped by generate_token_bitmask during verify, read back to apply the mask.
+    grammar: Optional[BaseGrammarObject] = None
+
     def __post_init__(self):
         super().__init__(spec_input_type=SpecInputType.DFLASH_DRAFT)
         # Spec v2 draft state itself does not change token accounting.

python/sglang/srt/speculative/dflash_utils.py

@@ -12,7 +12,6 @@ import torch.nn.functional as F
 from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
 from sglang.srt.layers.sampler import apply_custom_logit_processor
 from sglang.srt.managers.schedule_batch import Req
-from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
 from sglang.srt.utils import is_cuda, is_musa
 
 DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>"
@@ -794,26 +793,11 @@ def build_dflash_verify_target_probs(
     return target_probs.view(bs, draft_token_num, -1).contiguous()
 
 
-def validate_dflash_request(
+def validate_dflash_request(req: Req, enable_overlap: bool) -> Optional[str]:
-    req: Req, enable_overlap: bool, spec_algorithm: SpeculativeAlgorithm
-) -> Optional[str]:
     if req.return_logprob:
         return "DFLASH speculative decoding does not support return_logprob yet."
 
     if enable_overlap and req.return_hidden_states:
         return "DFLASH speculative decoding does not support return_hidden_states yet."
 
-    # Grammar support in this family is the verify-time bitmask plus the grammar
-    # barrier, so the capability that gates the barrier also gates admission.
-    if not spec_algorithm.supports_grammar_overlap() and (
-        req.sampling_params.json_schema is not None
-        or req.sampling_params.regex is not None
-        or req.sampling_params.ebnf is not None
-        or req.sampling_params.structural_tag is not None
-    ):
-        return (
-            f"{spec_algorithm.name} speculative decoding does not support "
-            "grammar-constrained decoding yet."
-        )
-
     return None

python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py

@@ -53,7 +53,11 @@ from sglang.srt.speculative.dspark_components.dspark_verify import (
     TargetVerifyExecutor,
     verify_logits_adjustments_are_noop,
 )
-from sglang.srt.speculative.spec_utils import draft_tp_context
+from sglang.srt.speculative.spec_utils import (
+    GrammarTree,
+    build_grammar_vocab_mask,
+    draft_tp_context,
+)
 from sglang.srt.utils import get_available_gpu_memory, is_cuda
 
 logger = logging.getLogger(__name__)
@@ -358,6 +362,7 @@ class DSparkWorkerV2(BaseSpecWorker):
         self,
         batch: ScheduleBatch,
         on_publish=None,
+        grammar_barrier=None,
     ) -> GenerationBatchResult:
         if getattr(batch, "return_logprob", False):
             raise ValueError(
@@ -369,7 +374,7 @@ class DSparkWorkerV2(BaseSpecWorker):
             self._observers.note_prefill_step()
             return self._forward_prefill(batch, on_publish)
 
-        return self._forward_decode(batch, on_publish)
+        return self._forward_decode(batch, on_publish, grammar_barrier)
 
     def _forward_prefill(
         self, batch: ScheduleBatch, on_publish
@@ -475,7 +480,7 @@ class DSparkWorkerV2(BaseSpecWorker):
         )
 
     def _forward_decode(
-        self, batch: ScheduleBatch, on_publish
+        self, batch: ScheduleBatch, on_publish, grammar_barrier=None
     ) -> GenerationBatchResult:
         if batch.spec_info is None:
             batch.spec_info = DFlashDraftInputV2.create_idle_input(device=self.device)
@@ -568,11 +573,19 @@ class DSparkWorkerV2(BaseSpecWorker):
             [draft_block_ids[:, :1], draft_tokens], dim=1
         ).contiguous()
 
+        # Must stay ahead of the target verify launch below.
+        grammar_tree = (
+            GrammarTree.from_linear_chain(verify_ids_2d) if batch.has_grammar else None
+        )
+
+        # A live grammar forces the eager path: the folded epilogue accepts inside
+        # the cuda graph off its own buffers, where the mask below never lands.
         fold_eligible = (
             self._verify_executor.verify_epilogue is not None
             and proposal.folded
             and verify_logits_adjustments_are_noop(sampling_info)
             and self._simulate_acc_len <= 0
+            and not batch.has_grammar
         )
         with self._observers.segment(InfoSegment.TARGET_VERIFY):
             if run_compact:
@@ -598,6 +611,25 @@ class DSparkWorkerV2(BaseSpecWorker):
         logits_output = target_verify.logits_output
         can_run_cuda_graph = target_verify.can_run_cuda_graph
 
+        if batch.has_grammar:
+            # Both the FSM advance over the previous batch's committed tokens and
+            # the traversal below are host work, so they overlap the launch above.
+            if grammar_barrier is not None:
+                grammar_barrier()
+            # run_compact scatters its rows back to (bs * chain_len), so the mask
+            # lines up with the logits on both verify paths.
+            vocab_mask = build_grammar_vocab_mask(
+                reqs=batch.reqs,
+                verify_input=draft_input,
+                tree=grammar_tree,
+                sampling_info=sampling_info,
+                device=logits_output.next_token_logits.device,
+            )
+            if vocab_mask is not None:
+                draft_input.grammar.apply_vocab_mask(
+                    logits=logits_output.next_token_logits, vocab_mask=vocab_mask
+                )
+
         epilogue = self._verify_executor.verify_epilogue
         folded_accept = fold_eligible and run_compact and can_run_cuda_graph
         accept = self._verify_executor.accept_and_finalize(

python/sglang/srt/speculative/spec_info.py

@@ -139,7 +139,7 @@ class SpeculativeAlgorithm(Enum):
         # Needs a GPU draft phase to hide the grammar CPU work under: NGRAM drafts
         # from a host corpus lookup, so it stays synchronous by design.
         # STANDALONE inherits the EAGLE V2 worker's verify path, barrier included.
-        return self.is_eagle() or self.is_standalone() or self.is_dflash()
+        return self.is_eagle() or self.is_standalone() or self.is_dflash_family()
 
     def has_draft_kv(self) -> bool:
         """Whether the draft phase writes KV chains. NGRAM does not (its tree

test/registered/core/test_basic_sanity_dspark.py

@@ -7,6 +7,8 @@ from sglang.test.kits.basic_decode_correctness_kit import BasicDecodeCorrectness
 from sglang.test.kits.basic_scheduler_stress_kit import BasicSchedulerStressMixin
 from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
 from sglang.test.kits.fwd_occupancy_kit import FwdOccupancyMixin
+from sglang.test.kits.json_constrained_kit import JSONConstrainedMixin
+from sglang.test.kits.spec_server_kits import SpecGrammarKit
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
@@ -14,7 +16,7 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=120, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=180, stage="base-b", runner_config="1-gpu-large")
 
 TARGET_MODEL = "Qwen/Qwen3-14B"
 DRAFT_MODEL = "deepseek-ai/dspark_qwen3_14b_block7"
@@ -34,6 +36,8 @@ class TestBasicSanityDSpark(
     BasicSchedulerStressMixin,
     FwdOccupancyMixin,
     GSM8KMixin,
+    JSONConstrainedMixin,
+    SpecGrammarKit,
     CustomTestCase,
 ):
     served_model_name = TARGET_MODEL
@@ -84,6 +88,10 @@ class TestBasicSanityDSpark(
             },
         )
 
+    @unittest.skip("DSPARK rejects return_logprob at admission")
+    def test_grammar_logprob_count_matches_completion_tokens(self):
+        pass
+
     @classmethod
     def tearDownClass(cls):
         if cls.process is not None: