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

Revert "[perf][spec decoding] Skip full-vocab softmax in EAGLE draft when topk == 1 (#26235)" (#26358)

Browse Source
This commit is contained in:
Michael
2026-05-26 02:47:52 -07:00
committed by GitHub
parent d9c82934c8
commit 9409969fd5
2 changed files with 6 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+2 -8
View File
@@ -401,14 +401,8 @@ class EAGLEDraftExtendCudaGraphRunner:
forward_batch.positions,
forward_batch,
)
if self.topk == 1:
ret.topk_index = torch.argmax(
ret.next_token_logits, dim=-1, keepdim=True
)
ret.topk_p = torch.ones_like(ret.topk_index, dtype=torch.float32)
else:
probs = torch.softmax(ret.next_token_logits, dim=-1)
ret.topk_p, ret.topk_index = fast_topk(probs, self.topk, dim=-1)
probs = torch.softmax(ret.next_token_logits, dim=-1)
ret.topk_p, ret.topk_index = fast_topk(probs, self.topk, dim=-1)
forward_batch.out_cache_loc = output_cache_loc_backup
forward_batch.spec_info.hidden_states = hidden_states_backup
python/sglang/srt/speculative/eagle_worker_v2.py
+4 -18
View File
@@ -483,16 +483,8 @@ class EagleDraftWorker(BaseDraftWorker):
forward_batch, skip_attn_backend_init=True
).logits_output
maybe_detect_nan(logits_output.next_token_logits, f"draft_forward step {i}")
if self.topk == 1:
# topk=1 → degenerate single-path tree; `topk_p` is unused
# downstream, so skip softmax and just argmax over logits.
topk_index = torch.argmax(
logits_output.next_token_logits, dim=-1, keepdim=True
)
topk_p = torch.ones_like(topk_index, dtype=torch.float32)
else:
probs = torch.softmax(logits_output.next_token_logits, dim=-1)
topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
probs = torch.softmax(logits_output.next_token_logits, dim=-1)
topk_p, topk_index = fast_topk(probs, self.topk, dim=-1)
maybe_detect_oob(
topk_index,
0,
@@ -659,14 +651,8 @@ class EagleDraftWorker(BaseDraftWorker):
draft_logits_output.hidden_states = draft_logits_output.hidden_states[
select_index
]
if self.topk == 1:
ret_topk_index = torch.argmax(
draft_logits_output.next_token_logits, dim=-1, keepdim=True
)
ret_topk_p = torch.ones_like(ret_topk_index, dtype=torch.float32)
else:
probs = torch.softmax(draft_logits_output.next_token_logits, dim=-1)
ret_topk_p, ret_topk_index = fast_topk(probs, self.topk, dim=-1)
probs = torch.softmax(draft_logits_output.next_token_logits, dim=-1)
ret_topk_p, ret_topk_index = fast_topk(probs, self.topk, dim=-1)
ret_hidden_states = draft_logits_output.hidden_states
# Construct the return values