From 5534380d462147a2cbd24e180f2ddb6055c7eab2 Mon Sep 17 00:00:00 2001 From: Ethan <59432975+QAQEthan@users.noreply.github.com> Date: Mon, 17 Aug 2026 06:05:13 +0800 Subject: [PATCH] [Spec] Support logprobs with DSpark speculative decoding (#34696) Co-authored-by: QAQEthan --- python/sglang/srt/managers/scheduler.py | 6 +--- .../dspark_components/dspark_worker_v2.py | 28 +++++++++++++++---- .../core/test_basic_sanity_dspark.py | 7 ++--- 3 files changed, 26 insertions(+), 15 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 12ae8603c..58f9e4eb6 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2510,11 +2510,7 @@ class Scheduler( self._maybe_namespace_elastic_radix_cache(req) if self.spec_algorithm.is_dflash_family(): - error_msg = ( - "DSpark speculative decoding does not support return_logprob yet." - if self.spec_algorithm.is_dspark() and req.return_logprob - else validate_dflash_request(req, self.enable_overlap) - ) + 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/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 5b59f9f62..d2aff43ac 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -11,6 +11,7 @@ from sglang.kernels.ops.attention.dsv4.unified_kv_kernels.env_gate import ( from sglang.srt.configs.hybrid_arch import mambaish_config from sglang.srt.distributed.parallel_state_wrapper import ParallelState from sglang.srt.environ import envs +from sglang.srt.layers.logprob_processor import compute_spec_v2_logprobs from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker @@ -405,17 +406,25 @@ class DSparkWorkerV2(BaseSpecWorker): def note_request_finished(self, *, rid: str, natural_stop: bool) -> None: self._observers.note_request_finished(rid=rid, natural_stop=natural_stop) + def _linear_accept_indices(self, bs: int) -> torch.Tensor: + num_indices = bs * self.verify_num_draft_tokens + if ( + self._linear_accept_index_cache is None + or self._linear_accept_index_cache.numel() < num_indices + ): + self._linear_accept_index_cache = torch.arange( + num_indices, dtype=torch.int64, device=self.device + ) + return self._linear_accept_index_cache[:num_indices].view( + bs, self.verify_num_draft_tokens + ) + def forward_batch_generation( self, batch: ScheduleBatch, on_publish=None, grammar_barrier=None, ) -> GenerationBatchResult: - if getattr(batch, "return_logprob", False): - raise ValueError( - "DSpark speculative decoding does not support return_logprob yet." - ) - if batch.forward_mode.is_extend() or batch.is_extend_in_batch: self._verify_planner.note_non_decode_step() self._observers.note_prefill_step() @@ -708,6 +717,15 @@ class DSparkWorkerV2(BaseSpecWorker): prefix_lens=prefix_lens, draft_tokens=draft_tokens, ) + if batch.return_logprob: + compute_spec_v2_logprobs( + batch, + logits_output, + accept.out_tokens.reshape(-1), + self._linear_accept_indices(bs), + self.verify_num_draft_tokens - 1, + ) + if on_publish is not None: if confidence is not None: on_publish(accept.new_seq_lens, confidence=confidence) diff --git a/test/registered/core/test_basic_sanity_dspark.py b/test/registered/core/test_basic_sanity_dspark.py index 12baba85a..0b07364b7 100644 --- a/test/registered/core/test_basic_sanity_dspark.py +++ b/test/registered/core/test_basic_sanity_dspark.py @@ -8,7 +8,7 @@ from sglang.test.kits.basic_scheduler_stress_kit import BasicSchedulerStressMixi 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.kits.spec_server_kits import SpecGrammarKit, SpecLogprobKit from sglang.test.test_utils import ( DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, DEFAULT_URL_FOR_TEST, @@ -38,6 +38,7 @@ class TestBasicSanityDSpark( GSM8KMixin, JSONConstrainedMixin, SpecGrammarKit, + SpecLogprobKit, CustomTestCase, ): served_model_name = TARGET_MODEL @@ -88,10 +89,6 @@ 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: