From 5ff5aa6923b914b03b4592b841ee9d69bfde664e Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 19 Feb 2026 11:38:25 -0800 Subject: [PATCH] [spec v2]Fix torch gc of future indices (#18958) --- python/sglang/srt/managers/overlap_utils.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 97e9289c8..9246c2098 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/python/sglang/srt/managers/overlap_utils.py @@ -127,6 +127,13 @@ class FutureMap: # FIXME(lsyin): No future exists, only for prefill batch, not compatible with mixed mode return 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_index = self.topk_index_buf[indices] draft_input.verified_id = self.verified_id_buf[indices]