diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index a6f46ea19..23589830f 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -821,8 +821,12 @@ class EagleDraftWorker(EagleDraftWorkerBase): num_tokens_for_logprob_per_req=self.speculative_num_draft_tokens, ) select_index = ( - torch.arange(len(batch.seq_lens), device=self.device) - * self.speculative_num_draft_tokens + torch.arange( + 0, + len(batch.seq_lens) * self.speculative_num_draft_tokens, + self.speculative_num_draft_tokens, + device=self.device, + ) + batch_result.accept_lens - 1 )