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

Remove sync when enabling return_logprob (#20972)

Browse Source
This commit is contained in:
Qiaolin Yu
2026-03-27 16:36:28 -07:00
committed by GitHub
parent e2b8463c80
commit a27651d5e0
2 changed files with 19 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/sampler.py
+5 -3
View File
@@ -327,13 +327,15 @@ class Sampler(nn.Module):
(
logits_output.next_token_top_logprobs_val,
logits_output.next_token_top_logprobs_idx,
) = get_top_logprobs(logprobs, top_logprobs_nums)
) = get_top_logprobs(logprobs, top_logprobs_nums, no_copy_to_cpu=True)
if any(x is not None for x in token_ids_logprobs):
(
logits_output.next_token_token_ids_logprobs_val,
logits_output.next_token_token_ids_logprobs_idx,
) = get_token_ids_logprobs(logprobs, token_ids_logprobs)
) = get_token_ids_logprobs(
logprobs, token_ids_logprobs, no_copy_to_cpu=True
)
logits_output.next_token_logprobs = logprobs[
torch.arange(len(batch_next_token_ids), device=sampling_info.device),
@@ -397,7 +399,7 @@ class Sampler(nn.Module):
(
logits_output.next_token_top_logprobs_val,
logits_output.next_token_top_logprobs_idx,
) = get_top_logprobs(logprobs, top_logprobs_nums)
) = get_top_logprobs(logprobs, top_logprobs_nums, no_copy_to_cpu=True)
# Handle token_ids logprobs if requested
if needs_token_ids_logprobs:
python/sglang/srt/managers/scheduler_output_processor_mixin.py
+14 -2
View File
@@ -152,6 +152,18 @@ class SchedulerOutputProcessorMixin:
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
]
hidden_state_offset = 0
@@ -377,7 +389,7 @@ class SchedulerOutputProcessorMixin:
if batch.return_logprob:
next_token_logprobs = logits_output.next_token_logprobs.tolist()
if batch.is_spec_v2 and logits_output.next_token_top_logprobs_val:
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
]
@@ -385,7 +397,7 @@ class SchedulerOutputProcessorMixin:
x.tolist() for x in logits_output.next_token_top_logprobs_idx
]
if batch.is_spec_v2 and logits_output.next_token_token_ids_logprobs_val:
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