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

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

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_amd_ci(est_time=302, stage="stage-b", runner_config="1-gpu-small-amd")
+register_cuda_ci(est_time=420, stage="base-b", runner_config="1-gpu-small")
+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()