From c01ad10d1670f21ef4b544d723596f38bca2104c Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Wed, 24 Jun 2026 12:27:13 -0700 Subject: [PATCH] [perf] tiny optimize select_index op for draft extend (#29078) --- python/sglang/srt/speculative/eagle_worker_v2.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) 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 )