minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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
+7 -11
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]
( for req in self.reqs
req.output_ids[-1] ]
if len(req.output_ids) # Non-blocking H2D so this per-step copy doesn't sync behind the forward.
else req.origin_input_ids[-1] latest_output_ids = torch.tensor(last_tokens, dtype=torch.int64).to(
) self.device, non_blocking=True
for req in self.reqs
],
dtype=torch.int64,
device=self.device,
) )
self.sampling_info.penalizer_orchestrator.cumulate_output_tokens( self.sampling_info.penalizer_orchestrator.cumulate_output_tokens(
latest_output_ids latest_output_ids