From 8b78e0888cce593aa178685bcd0f52fd5a416d20 Mon Sep 17 00:00:00 2001 From: Shenxiu Liu Date: Wed, 22 Apr 2026 07:28:06 -0700 Subject: [PATCH] Skip mamba_pool_idx revert for session requests in _get_new_batch_prefill_raw (#23327) --- python/sglang/srt/managers/scheduler.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index a9719e5a4..8c7c63959 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2575,9 +2575,16 @@ class Scheduler( ) > 0 or (not self.running_batch.is_empty()) else: 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] - 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( req.mamba_pool_idx.unsqueeze(-1) )