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

Support return_logprob for spec v2 (overlap safe) (#19801)

Browse Source 
Co-authored-by: Ratish1 <ratish1501@gmail.com>
Co-authored-by: Ratish1 <formula733@gmail.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
This commit is contained in:
Qiaolin Yu
2026-03-10 15:38:27 -07:00
committed by GitHub
parent 76ee4bb98c
commit 09a118fafe
6 changed files with 314 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/sglang/srt/managers/utils.py
View File

@@ -63,6 +63,21 @@ class GenerationBatchResult:
self.logits_output.input_token_logprobs = (
self.logits_output.input_token_logprobs.to("cpu", non_blocking=True)
)
if self.logits_output.next_token_top_logprobs_val is not None:
self.logits_output.next_token_top_logprobs_val = [
v.to("cpu", non_blocking=True) if torch.is_tensor(v) else v
for v in self.logits_output.next_token_top_logprobs_val
]
if self.logits_output.next_token_top_logprobs_idx is not None:
self.logits_output.next_token_top_logprobs_idx = [
x.to("cpu", non_blocking=True) if torch.is_tensor(x) else x
for x in self.logits_output.next_token_top_logprobs_idx
]
if self.logits_output.next_token_token_ids_logprobs_val is not None:
self.logits_output.next_token_token_ids_logprobs_val = [
v.to("cpu", non_blocking=True) if torch.is_tensor(v) else v
for v in self.logits_output.next_token_token_ids_logprobs_val
]
if self.logits_output.hidden_states is not None:
self.logits_output.hidden_states = self.logits_output.hidden_states.to(
"cpu", non_blocking=True