diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 489678af1..3fc35ee2e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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( - req, self.enable_overlap, self.spec_algorithm - ) + error_msg = validate_dflash_request(req, self.enable_overlap) 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_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index 72162ba51..29b593fd6 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/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. diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 1527cf0d0..246b7a1ca 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/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( - req: Req, enable_overlap: bool, spec_algorithm: SpeculativeAlgorithm -) -> Optional[str]: +def validate_dflash_request(req: Req, enable_overlap: bool) -> 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 diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index dfbff5730..2ac5d2b2c 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/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( diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index bb20dff4f..898c6384e 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() 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 diff --git a/test/registered/core/test_basic_sanity_dspark.py b/test/registered/core/test_basic_sanity_dspark.py index c27d9056d..12baba85a 100644 --- a/test/registered/core/test_basic_sanity_dspark.py +++ b/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: