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

[spec v2]Fix torch gc of future indices (#18958)

Browse Source
This commit is contained in:
Liangsheng Yin
2026-02-19 11:38:25 -08:00
committed by GitHub
parent 44ab752b7a
commit 5ff5aa6923
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/overlap_utils.py
+7
View File
@@ -127,6 +127,13 @@ class FutureMap:
# FIXME(lsyin): No future exists, only for prefill batch, not compatible with mixed mode # FIXME(lsyin): No future exists, only for prefill batch, not compatible with mixed mode
return return
indices = draft_input.future_indices.indices indices = draft_input.future_indices.indices
# The indices tensor was allocated on the default stream but is
# used here on the forward stream. Meanwhile, the old spec_info
# holding this tensor will lose all Python references (replaced at
# model_worker_batch.spec_info and batch.spec_info), so the
# caching allocator (torch GC) could reclaim the memory before
# the GPU finishes reading it.
indices.record_stream(torch.get_device_module(self.device).current_stream())
draft_input.topk_p = self.topk_p_buf[indices] draft_input.topk_p = self.topk_p_buf[indices]
draft_input.topk_index = self.topk_index_buf[indices] draft_input.topk_index = self.topk_index_buf[indices]
draft_input.verified_id = self.verified_id_buf[indices] draft_input.verified_id = self.verified_id_buf[indices]