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: