EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

[Eagle] reduce one draft forward (#3468)

Browse Source
This commit is contained in:
Ying Sheng
2025-02-10 04:21:49 -08:00
committed by GitHub
parent 2d61132374
commit d23cb9a01e
2 changed files with 5 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
python/sglang/srt/layers/attention/flashinfer_backend.py
View File

@@ -947,7 +947,7 @@ class FlashInferMultiStepDraftBackend:
triton.next_power_of_2(bs), triton.next_power_of_2(bs),
) )
for i in range(self.speculative_num_steps): for i in range(self.speculative_num_steps - 1):
forward_batch.spec_info.kv_indptr = self.kv_indptr[i, : bs + 1] forward_batch.spec_info.kv_indptr = self.kv_indptr[i, : bs + 1]
forward_batch.spec_info.kv_indices = kv_indices_buffer[i][ forward_batch.spec_info.kv_indices = kv_indices_buffer[i][
: seq_lens_sum * self.topk + bs * (i + 1) : seq_lens_sum * self.topk + bs * (i + 1)

4
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -234,6 +234,10 @@ class EAGLEWorker(TpModelWorker):
token_list.append(tree_info[1]) token_list.append(tree_info[1])
parents_list.append(tree_info[2]) parents_list.append(tree_info[2])
# we don't need to run the last forward. we get 1 token from draft prefill and (#spec steps - 1) tokens here
if i == self.speculative_num_steps - 1:
break
# Set inputs # Set inputs
forward_batch.input_ids = input_ids forward_batch.input_ids = input_ids
forward_batch.out_cache_loc = out_cache_loc[ forward_batch.out_cache_loc = out_cache_loc[