From 0df69849bacf0d63ebf6c94a1d205254c6359cb0 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Fri, 28 Aug 2026 18:10:41 -0700 Subject: [PATCH] [Fix] Drop the duplicated DSpark draft sample_block call (#36934) --- .../dspark_components/dspark_draft_sampler.py | 11 +---------- 1 file changed, 1 insertion(+), 10 deletions(-) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py b/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py index d5853e5fb..a5b2a9f7f 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py @@ -157,7 +157,7 @@ class DsparkDraftSampler: draft_tokens, corrected_logits = self.markov_head.sample_block( base_logits, first_prev_tokens=anchor, - hidden_states=hidden_states.view(bs, self.gamma, -1), + hidden_states=sample_hidden, sampler=sampler, collect_corrected=self.folded_sampling, ) @@ -166,15 +166,6 @@ class DsparkDraftSampler: corrected_logits.reshape(bs * self.gamma, -1) ) - else: - sampler = greedy_step_sampler - - draft_tokens, corrected_logits = self.markov_head.sample_block( - base_logits, - first_prev_tokens=anchor, - hidden_states=sample_hidden, - sampler=sampler, - ) self.out[: draft_tokens.numel()].copy_(draft_tokens.reshape(-1)) if self.confidence_out is not None: confidence = self.confidence_fn(