[Spec] Support logprobs with DSpark speculative decoding (#34696)

Co-authored-by: QAQEthan <QAQEthan@users.noreply.github.com>

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)

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)

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: