minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Fix] Drop the duplicated DSpark draft sample_block call (#36934)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-08-28 18:10:41 -07:00
committed by GitHub
parent 7276a30c45
commit 0df69849ba
1 changed files with 1 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py
+1 -10
View File
@@ -157,7 +157,7 @@ class DsparkDraftSampler:
draft_tokens, corrected_logits = self.markov_head.sample_block( draft_tokens, corrected_logits = self.markov_head.sample_block(
base_logits, base_logits,
first_prev_tokens=anchor, first_prev_tokens=anchor,
hidden_states=hidden_states.view(bs, self.gamma, -1), hidden_states=sample_hidden,
sampler=sampler, sampler=sampler,
collect_corrected=self.folded_sampling, collect_corrected=self.folded_sampling,
) )
@@ -166,15 +166,6 @@ class DsparkDraftSampler:
corrected_logits.reshape(bs * self.gamma, -1) 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)) self.out[: draft_tokens.numel()].copy_(draft_tokens.reshape(-1))
if self.confidence_out is not None: if self.confidence_out is not None:
confidence = self.confidence_fn( confidence = self.confidence_fn(