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

Avoid relaying per-step outputs through ScheduleBatch fields in disagg prefill and PP (#30677)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-15 14:33:37 +08:00
committed by GitHub
parent 01343d2759
commit 201ddeaba1
2 changed files with 9 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/prefill.py
+6 -4
View File
@@ -620,6 +620,8 @@ class SchedulerDisaggregationPrefillMixin:
result.indexer_topk_output = None result.indexer_topk_output = None
logprob_pt = 0 logprob_pt = 0
assert batch.spec_info is result.next_draft_input
draft_input = result.next_draft_input
# Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue # Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue
next_token_ids = result.next_token_ids.tolist() next_token_ids = result.next_token_ids.tolist()
self.batch_result_processor.move_logprobs_to_cpu( self.batch_result_processor.move_logprobs_to_cpu(
@@ -651,11 +653,11 @@ class SchedulerDisaggregationPrefillMixin:
req.output_ids.append(next_token_id) req.output_ids.append(next_token_id)
maybe_cache_unfinished_req(req, self.tree_cache) maybe_cache_unfinished_req(req, self.tree_cache)
self.disagg_prefill_inflight_queue.append(req) self.disagg_prefill_inflight_queue.append(req)
if self.spec_algorithm.is_eagle() and batch.spec_info is not None: if self.spec_algorithm.is_eagle() and draft_input is not None:
req.output_topk_p = batch.spec_info.topk_p[i] req.output_topk_p = draft_input.topk_p[i]
req.output_topk_index = batch.spec_info.topk_index[i] req.output_topk_index = draft_input.topk_index[i]
req.hidden_states_tensor = ( req.hidden_states_tensor = (
batch.spec_info.hidden_states[i].cpu().clone() draft_input.hidden_states[i].cpu().clone()
) )
dsa_topk_indices = batch.spec_info.dsa_topk_indices dsa_topk_indices = batch.spec_info.dsa_topk_indices
if dsa_topk_indices is not None: if dsa_topk_indices is not None:
python/sglang/srt/managers/scheduler_pp_mixin.py
+3 -3
View File
@@ -1101,7 +1101,6 @@ class SchedulerPPMixin:
# next_pp_outputs = None so non-last ranks skip forwarding # next_pp_outputs = None so non-last ranks skip forwarding
# (pp_outputs is None gate). Placeholder carried in # (pp_outputs is None gate). Placeholder carried in
# batch_result.next_token_ids for process_batch_result_prefill. # batch_result.next_token_ids for process_batch_result_prefill.
batch.output_ids = placeholder
batch_result = GenerationBatchResult( batch_result = GenerationBatchResult(
logits_output=None, logits_output=None,
pp_hidden_states_proxy_tensors=None, pp_hidden_states_proxy_tensors=None,
@@ -1133,13 +1132,14 @@ class SchedulerPPMixin:
extend_input_len_per_req, extend_input_len_per_req,
extend_logprob_start_len_per_req, extend_logprob_start_len_per_req,
) = get_logprob_from_pp_outputs(pp_outputs) ) = get_logprob_from_pp_outputs(pp_outputs)
batch.input_ids = pp_outputs["next_token_ids"].to(torch.int64) next_token_ids = pp_outputs["next_token_ids"].to(torch.int64)
# PP rank 0 also relays into output_tokens_buf so the next iter's # PP rank 0 also relays into output_tokens_buf so the next iter's
# resolve_forward_inputs finds these tokens for the decode portion # resolve_forward_inputs finds these tokens for the decode portion
# of mixed-chunk batches (which gather via mix_running_indices). # of mixed-chunk batches (which gather via mix_running_indices).
self.future_map.stash( self.future_map.stash(
batch.req_pool_indices, RelayPayload(bonus_tokens=batch.input_ids) batch.req_pool_indices, RelayPayload(bonus_tokens=next_token_ids)
) )
batch.input_ids = None
output_result = GenerationBatchResult( output_result = GenerationBatchResult(
logits_output=logits_output, logits_output=logits_output,
pp_hidden_states_proxy_tensors=None, pp_hidden_states_proxy_tensors=None,