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

[PD] Deduplicate PD logprob normalization (#27085)

Browse Source
This commit is contained in:
Ye (Charlotte) Qi
2026-06-03 19:08:15 +08:00
committed by GitHub
parent 44d4a25a07
commit 03c77dc33d
3 changed files with 8 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/prefill.py
+4 -20
View File
@@ -529,26 +529,10 @@ class SchedulerDisaggregationPrefillMixin:
logprob_pt = 0
# Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue
next_token_ids = result.next_token_ids.tolist()
if batch.return_logprob:
if logits_output.next_token_logprobs is not None:
logits_output.next_token_logprobs = (
logits_output.next_token_logprobs.tolist()
)
if logits_output.input_token_logprobs is not None:
logits_output.input_token_logprobs = tuple(
logits_output.input_token_logprobs.tolist()
)
if logits_output.next_token_top_logprobs_val:
logits_output.next_token_top_logprobs_val = [
v.tolist() for v in logits_output.next_token_top_logprobs_val
]
logits_output.next_token_top_logprobs_idx = [
x.tolist() for x in logits_output.next_token_top_logprobs_idx
]
if logits_output.next_token_token_ids_logprobs_val:
logits_output.next_token_token_ids_logprobs_val = [
v.tolist() for v in logits_output.next_token_token_ids_logprobs_val
]
self.batch_result_processor.move_logprobs_to_cpu(
batch=batch,
logits_output=logits_output,
)
def advance_logprob_pt(i: int, req: Req) -> None:
nonlocal logprob_pt
python/sglang/srt/managers/scheduler_components/batch_result_processor.py
+2 -2
View File
@@ -206,7 +206,7 @@ class SchedulerBatchResultProcessor:
# Move next_token_ids and logprobs to cpu
next_token_ids = next_token_ids.tolist()
self._move_logprobs_to_cpu(batch=batch, logits_output=logits_output)
self.move_logprobs_to_cpu(batch=batch, logits_output=logits_output)
self._validate_pp_skip_output_comm(batch, result)
@@ -356,7 +356,7 @@ class SchedulerBatchResultProcessor:
embeddings = [tensor.tolist() for tensor in embeddings]
return embeddings
def _move_logprobs_to_cpu(
def move_logprobs_to_cpu(
self,
*,
batch: ScheduleBatch,