From d2212125656a84bb6e77f2028571163549404ae3 Mon Sep 17 00:00:00 2001 From: leavelet Date: Wed, 10 Jun 2026 05:07:51 +0000 Subject: [PATCH] Batch EAGLE spec output D2H copies in disagg prefill results MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- python/sglang/srt/disaggregation/prefill.py | 40 +++++++++++++++++---- 1 file changed, 34 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py index 9298130f8..7e7b63cec 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py @@ -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: