From d021990bf5aff7f0cbcc7a024aa17dbfaee824d0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?H=E1=BB=93=20S=E1=BB=B9=20Th=E1=BA=BF?= <95731191+hsthe29@users.noreply.github.com> Date: Sat, 25 Jul 2026 18:36:46 +0700 Subject: [PATCH] [DFLASH] Support grammar-constrained decoding in speculative verify (#30096) Co-authored-by: hnyls2002 --- python/sglang/srt/managers/scheduler.py | 4 +- python/sglang/srt/speculative/dflash_info.py | 4 ++ python/sglang/srt/speculative/dflash_utils.py | 11 +++-- .../srt/speculative/dflash_worker_v2.py | 33 ++++++++++++++- python/sglang/srt/speculative/spec_info.py | 2 +- python/sglang/srt/speculative/spec_utils.py | 14 +++++++ test/registered/spec/dflash/test_dflash.py | 18 ++++++-- .../spec/test_spec_utils_traverse_tree.py | 42 ++++++++++++++++++- 8 files changed, 118 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 3fc35ee2e..489678af1 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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) diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index 8ec5420bc..6758878b5 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/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: diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 032c29e90..1527cf0d0 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/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." ) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index a1095828f..58a3a11d5 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/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 ( diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 488608577..bb20dff4f 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/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 diff --git a/python/sglang/srt/speculative/spec_utils.py b/python/sglang/srt/speculative/spec_utils.py index 5ebac1c60..c05d2dfc5 100644 --- a/python/sglang/srt/speculative/spec_utils.py +++ b/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() diff --git a/test/registered/spec/dflash/test_dflash.py b/test/registered/spec/dflash/test_dflash.py index c4f764cad..44da2500f 100644 --- a/test/registered/spec/dflash/test_dflash.py +++ b/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 diff --git a/test/registered/unit/spec/test_spec_utils_traverse_tree.py b/test/registered/unit/spec/test_spec_utils_traverse_tree.py index 69c5c66c2..ac0ec087f 100644 --- a/test/registered/unit/spec/test_spec_utils_traverse_tree.py +++ b/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()