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

This commit is contained in:
Liangsheng Yin
2026-08-28 18:10:41 -07:00
committed by GitHub
parent 7276a30c45
commit 0df69849ba
@@ -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(