chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Batch EAGLE spec output D2H copies in disagg prefill results

Browse Source 
process_batch_result_disagg_prefill did one synchronous
hidden_states[i].cpu().clone() per finished request, and stored
topk_p/topk_index as GPU row views whose D2H then happened one by one
inside MetadataBuffers.set_buf — 3*bs device syncs per finished batch
on the scheduler thread. Gather the finishing rows once and do one
D2H per tensor (3 total); requests hold CPU row views, which set_buf
copies CPU-to-CPU into the metadata buffers.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
孙远航
2026-06-10 05:07:51 +00:00
parent 87a22b17ce
commit d221212565
1 changed files with 34 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files

40
python/sglang/srt/disaggregation/prefill.py
View File

@@ -933,6 +933,35 @@ class SchedulerDisaggregationPrefillMixin:
logits_output.input_token_logprobs.tolist()
)
# Batch the D2H copies of the EAGLE spec outputs: one gathered copy per
# tensor for all finishing requests instead of one synchronous device
# sync per request here (hidden_states) and per request inside
# MetadataBuffers.set_buf (topk_p/topk_index, which were stored as GPU
# row views and copied into the CPU metadata buffers one by one).
spec_cpu_row = None
if self.spec_algorithm.is_eagle() and batch.spec_info is not None:
finished_idx = [i for i, r in enumerate(batch.reqs) if r.is_chunked <= 0]
if finished_idx:
spec_cpu_row = {i: k for k, i in enumerate(finished_idx)}
if len(finished_idx) == len(batch.reqs):
gather = lambda t: t
else:
gather_idx = torch.tensor(
finished_idx,
dtype=torch.int64,
device=batch.spec_info.hidden_states.device,
)
gather = lambda t: t.index_select(0, gather_idx)
# .to(copy=True) detaches from spec_info even when the batch is
# already on CPU (set_buf reads these synchronously below).
spec_topk_p_cpu = gather(batch.spec_info.topk_p).to("cpu", copy=True)
spec_topk_index_cpu = gather(batch.spec_info.topk_index).to(
"cpu", copy=True
)
spec_hidden_cpu = gather(batch.spec_info.hidden_states).to(
"cpu", copy=True
)
for i, (req, next_token_id) in enumerate(
zip(batch.reqs, next_token_ids, strict=True)
):
@@ -943,12 +972,11 @@ class SchedulerDisaggregationPrefillMixin:
req.output_ids.append(next_token_id)
self.tree_cache.cache_unfinished_req(req) # update the tree and lock
self.disagg_prefill_inflight_queue.append(req)
if self.spec_algorithm.is_eagle() and batch.spec_info is not None:
req.output_topk_p = batch.spec_info.topk_p[i]
req.output_topk_index = batch.spec_info.topk_index[i]
req.hidden_states_tensor = (
batch.spec_info.hidden_states[i].cpu().clone()
)
if spec_cpu_row is not None:
k = spec_cpu_row[i]
req.output_topk_p = spec_topk_p_cpu[k]
req.output_topk_index = spec_topk_index_cpu[k]
req.hidden_states_tensor = spec_hidden_cpu[k]
else:
req.hidden_states_tensor = None
if req.return_logprob: