minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 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: