[DFLASH] Support grammar-constrained decoding in speculative verify (#30096)

Co-authored-by: hnyls2002 <lsyincs@gmail.com>

python/sglang/srt/managers/scheduler.py

@@ -2253,7 +2253,9 @@ class Scheduler(
         self._maybe_namespace_elastic_radix_cache(req)
 
         if self.spec_algorithm.is_dflash_family():
-            error_msg = validate_dflash_request(req, self.enable_overlap)
+            error_msg = validate_dflash_request(
+                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.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, Optional
 import torch
 
 from sglang.kernels.ops.attention.utils import create_flashinfer_kv_indices_triton
+from sglang.srt.constrained.base_grammar_backend import BaseGrammarObject
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.model_executor.forward_batch_info import (
     CaptureHiddenMode,
@@ -44,6 +45,9 @@ class DFlashVerifyInput(SpecInput):
 
     ragged_verify_layout: Optional[RaggedVerifyLayout] = 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_VERIFY)
         if self.num_tokens_per_req == -1:

python/sglang/srt/speculative/dflash_utils.py

@@ -12,6 +12,7 @@ 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|>"
@@ -793,21 +794,25 @@ def build_dflash_verify_target_probs(
     return target_probs.view(bs, draft_token_num, -1).contiguous()
 
 
-def validate_dflash_request(req: Req, enable_overlap: bool) -> Optional[str]:
+def validate_dflash_request(
+    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."
 
-    if (
+    # 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 (
-            "DFLASH speculative decoding does not support "
+            f"{spec_algorithm.name} speculative decoding does not support "
             "grammar-constrained decoding yet."
         )
 

python/sglang/srt/speculative/dflash_worker_v2.py

@@ -47,7 +47,11 @@ from sglang.srt.speculative.draft_worker_common import (
     make_draft_sampler_capture_hook,
 )
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
-from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func
+from sglang.srt.speculative.spec_utils import (
+    GrammarTree,
+    assign_req_to_token_pool_func,
+    build_grammar_vocab_mask,
+)
 from sglang.srt.utils import get_available_gpu_memory, is_cuda, is_hip, is_npu
 
 _is_npu = is_npu()
@@ -1367,6 +1371,7 @@ class DFlashWorkerV2(BaseSpecWorker):
         self,
         batch: ScheduleBatch,
         on_publish=None,
+        grammar_barrier=None,
     ) -> GenerationBatchResult:
         if getattr(batch, "return_logprob", False):
             raise ValueError(
@@ -1633,6 +1638,11 @@ class DFlashWorkerV2(BaseSpecWorker):
         draft_tokens[:, 0].copy_(block_ids[:, 0])
         draft_tokens[:, 1:].copy_(draft_next)
 
+        # Must stay ahead of the target verify launch below.
+        grammar_tree = (
+            GrammarTree.from_linear_chain(draft_tokens) if batch.has_grammar else None
+        )
+
         # --- 2) Target verify.
         # TARGET_VERIFY uses standard causal masking; custom masks are unnecessary here.
         custom_mask = None
@@ -1678,6 +1688,21 @@ class DFlashWorkerV2(BaseSpecWorker):
         logits_output = target_out.logits_output
         can_run_cuda_graph = target_out.can_run_cuda_graph
 
+        vocab_mask = None
+        if batch.has_grammar:
+            # Grammar barrier: advance the previous batch's FSM over its committed
+            # tokens before building this batch's bitmask. Runs after the target
+            # launch, so the advance and the traversal both overlap the forward.
+            if grammar_barrier is not None:
+                grammar_barrier()
+            vocab_mask = build_grammar_vocab_mask(
+                reqs=batch.reqs,
+                verify_input=verify_input,
+                tree=grammar_tree,
+                sampling_info=batch.sampling_info,
+                device=logits_output.next_token_logits.device,
+            )
+
         if sampling_info is not None:
             apply_dflash_verify_logits_adjustments(
                 next_token_logits=logits_output.next_token_logits,
@@ -1685,6 +1710,12 @@ class DFlashWorkerV2(BaseSpecWorker):
                 draft_token_num=int(self.block_size),
             )
 
+        # Constrain every chain position before accept picks from it.
+        if vocab_mask is not None:
+            verify_input.grammar.apply_vocab_mask(
+                logits=logits_output.next_token_logits, vocab_mask=vocab_mask
+            )
+
         candidates = draft_tokens
         new_seq_lens = None
         if (

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()
+        return self.is_eagle() or self.is_standalone() or self.is_dflash()
 
     def has_draft_kv(self) -> bool:
         """Whether the draft phase writes KV chains. NGRAM does not (its tree

python/sglang/srt/speculative/spec_utils.py

@@ -590,6 +590,20 @@ class GrammarTree:
     ) -> GrammarTree:
         return cls((retrieve_next_token, retrieve_next_sibling, draft_token), None)
 
+    @classmethod
+    def from_linear_chain(cls, verify_ids_2d: torch.Tensor) -> GrammarTree:
+        """Degenerate tree for chain-verify algorithms: node i's only child is i + 1.
+
+        ``verify_ids_2d`` is (bs, chain_len) with column 0 the already-committed
+        token, so mask rows line up with the target's logits rows one-for-one.
+        Only the ids need a copy; the links are fixed by the shape.
+        """
+        bs, chain_len = verify_ids_2d.shape
+        next_token = torch.full((bs, chain_len), -1, dtype=torch.int64)
+        next_token[:, :-1] = torch.arange(1, chain_len, dtype=torch.int64)
+        next_sibling = torch.full((bs, chain_len), -1, dtype=torch.int64)
+        return cls.from_device(next_token, next_sibling, verify_ids_2d)
+
     def resolve(self) -> Tuple[torch.Tensor, ...]:
         if self._done is not None:
             self._done.synchronize()

test/registered/spec/dflash/test_dflash.py

@@ -6,11 +6,13 @@ from sglang.srt.environ import envs
 from sglang.srt.utils import is_hip, kill_process_tree
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
+from sglang.test.kits.json_constrained_kit import JSONConstrainedMixin
 from sglang.test.kits.matched_stop_kit import MatchedStopMixin
 from sglang.test.kits.radix_cache_server_kit import (
     gen_radix_tree,
     run_radix_attention_test,
 )
+from sglang.test.kits.spec_server_kits import SpecGrammarKit
 from sglang.test.test_utils import (
     DEFAULT_DRAFT_MODEL_DFLASH,
     DEFAULT_TARGET_MODEL_DFLASH,
@@ -20,11 +22,17 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=302, stage="base-b", runner_config="1-gpu-small")
+register_cuda_ci(est_time=420, stage="base-b", runner_config="1-gpu-small")
-register_amd_ci(est_time=302, stage="stage-b", runner_config="1-gpu-small-amd")
+register_amd_ci(est_time=420, stage="stage-b", runner_config="1-gpu-small-amd")
 
 
-class TestDFlashServerBase(CustomTestCase, MatchedStopMixin, GSM8KMixin):
+class TestDFlashServerBase(
+    CustomTestCase,
+    MatchedStopMixin,
+    GSM8KMixin,
+    JSONConstrainedMixin,
+    SpecGrammarKit,
+):
     max_running_requests = 64
     attention_backend = "triton" if is_hip() else "flashinfer"
     page_size = 1
@@ -124,6 +132,10 @@ class TestDFlashServerBase(CustomTestCase, MatchedStopMixin, GSM8KMixin):
         self.assertEqual(outputs[0], outputs[1])
         assert self.process.poll() is None
 
+    @unittest.skip("DFLASH rejects return_logprob at admission")
+    def test_grammar_logprob_count_matches_completion_tokens(self):
+        pass
+
 
 class TestDFlashServerPage256(TestDFlashServerBase):
     page_size = 256

test/registered/unit/spec/test_spec_utils_traverse_tree.py

@@ -11,7 +11,7 @@ from unittest.mock import MagicMock
 
 import torch
 
-from sglang.srt.speculative.spec_utils import traverse_tree
+from sglang.srt.speculative.spec_utils import GrammarTree, traverse_tree
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=4, suite="base-a-test-cpu")
@@ -40,6 +40,11 @@ class TestTraverseTreePassesIntsToGrammar(unittest.TestCase):
         grammar.rollback.return_value = None
         return grammar, accept_calls, fill_calls
 
+    def _chain(self, verify_ids_2d):
+        """Row 0 of the links chain-verify algorithms actually feed traverse_tree."""
+        links = GrammarTree.from_linear_chain(verify_ids_2d).resolve()
+        return tuple(t[0] for t in links)
+
     def test_branching_tree_passes_ints(self):
         # Binary tree exercises both child recursion and sibling recursion:
         #   0 ─┬─ 1
@@ -66,6 +71,41 @@ class TestTraverseTreePassesIntsToGrammar(unittest.TestCase):
         for idx in fill_calls:
             self.assertIsInstance(idx, int)
 
+    def test_linear_chain_visits_all_positions_in_order(self):
+        # Chain-verify algorithms (DFLASH/DSPARK) have no branching, so their tree
+        # degenerates to 0 -- 1 -- 2 -- 3 with column 0 the already-committed token.
+        rnt, rns, draft_tokens = self._chain(torch.tensor([[100, 11, 22, 33]]))
+        self.assertEqual(rnt.tolist(), [1, 2, 3, -1])
+        self.assertEqual(rns.tolist(), [-1, -1, -1, -1])
+        bitmask = torch.full((4, 4), -1, dtype=torch.int32)  # all allowed
+
+        grammar, accept_calls, fill_calls = self._record_grammar()
+        traverse_tree(rnt, rns, draft_tokens, grammar, bitmask)
+
+        # Root (col 0) is never accepted; every draft token is, in chain order.
+        self.assertEqual(accept_calls, [11, 22, 33])
+        self.assertEqual(fill_calls, [0, 1, 2, 3])
+        for token in accept_calls:
+            self.assertIsInstance(token, int)
+        for idx in fill_calls:
+            self.assertIsInstance(idx, int)
+
+    def test_linear_chain_stops_at_grammar_reject(self):
+        # A draft token the grammar disallows must stop the descent: no accept/fill
+        # for that node or anything after it, so the mask rows past it stay unfilled
+        # and only the already-filled prefix can be committed.
+        rnt, rns, draft_tokens = self._chain(torch.tensor([[100, 5, 7, 9]]))
+        bitmask = torch.full((4, 4), -1, dtype=torch.int32)  # all allowed
+        # Disallow token id 7 (draft_tokens[2]) in node 1's mask (its parent).
+        bitmask[1, 7 // 32] &= ~(1 << (7 % 32))
+
+        grammar, accept_calls, fill_calls = self._record_grammar()
+        traverse_tree(rnt, rns, draft_tokens, grammar, bitmask)
+
+        # Node 1 accepted+filled; node 2 rejected -> node 2 and node 3 skipped.
+        self.assertEqual(accept_calls, [5])
+        self.assertEqual(fill_calls, [0, 1])
+
 
 if __name__ == "__main__":
     unittest.main()