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

[Perf] Make latest_output_ids H2D non-blocking in prepare_for_decode (#28491)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-06-16 23:46:35 -07:00
committed by GitHub
parent 27291118b9
commit 3bc618485a
1 changed files with 7 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/schedule_batch.py
+6 -10
View File
@@ -2595,17 +2595,13 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# real token is relayed via future_map and resolved at forward # real token is relayed via future_map and resolved at forward
# entry. So take the last output token from Req directly # entry. So take the last output token from Req directly
# (origin_input_ids[-1] on the first decode, before any output). # (origin_input_ids[-1] on the first decode, before any output).
latest_output_ids = torch.tensor( last_tokens = [
[ req.output_ids[-1] if len(req.output_ids) else req.origin_input_ids[-1]
(
req.output_ids[-1]
if len(req.output_ids)
else req.origin_input_ids[-1]
)
for req in self.reqs for req in self.reqs
], ]
dtype=torch.int64, # Non-blocking H2D so this per-step copy doesn't sync behind the forward.
device=self.device, latest_output_ids = torch.tensor(last_tokens, dtype=torch.int64).to(
self.device, non_blocking=True
) )
self.sampling_info.penalizer_orchestrator.cumulate_output_tokens( self.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
latest_output_ids latest_output_ids