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

Skip mamba_pool_idx revert for session requests in _get_new_batch_prefill_raw (#23327)

Browse Source
This commit is contained in:
Shenxiu Liu
2026-04-22 22:28:06 +08:00
committed by GitHub
parent 4323fce82a
commit 8b78e0888c
1 changed files with 9 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+9 -2
View File
@@ -2575,9 +2575,16 @@ class Scheduler(
) > 0 or (not self.running_batch.is_empty()) ) > 0 or (not self.running_batch.is_empty())
else: else:
self.running_batch.batch_is_full = True self.running_batch.batch_is_full = True
# revert matched mamba idx to avoid memory leak, if req is not added # revert matched mamba idx to avoid memory leak, if req is not added.
# Only free if the slot was freshly allocated in this batch (not
# pre-existing from a session). Session-held slots have their own
# lifecycle and freeing them here causes double-free.
added = len(adder.can_run_list) > 0 and req is adder.can_run_list[-1] added = len(adder.can_run_list) > 0 and req is adder.can_run_list[-1]
if not added and req.mamba_pool_idx is not None: if (
not added
and req.mamba_pool_idx is not None
and not getattr(req, "session", None)
):
self.tree_cache.req_to_token_pool.mamba_pool.free( self.tree_cache.req_to_token_pool.mamba_pool.free(
req.mamba_pool_idx.unsqueeze(-1) req.mamba_pool_idx.unsqueeze(-1)
) )