diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 1f619132c..bfcb1e9a1 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -790,6 +790,12 @@ class HybridReqToTokenPool(ReqToTokenPool): ] ) self.mamba_pool.free(mamba_ping_pong_track_buffer_to_free) + # Match the req.mamba_pool_idx=None clear above so the next + # alloc() doesn't see a stale ping-pong reference on the req + # and skip allocation (which would silently reuse a freed + # tensor on the req side while the new pool slot leaks). + req.mamba_ping_pong_track_buffer = None + req.mamba_next_track_idx = None def clear(self): logger.info("Reset HybridReqToTokenPool") diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index 17602c3b3..475e9cf83 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -85,8 +85,20 @@ class SessionSlot: self.mamba_last_track_seqlen = req.mamba_last_track_seqlen self.mamba_branching_seqlen = req.mamba_branching_seqlen + # Ownership has transferred to the slot. Null *all* of the req's + # references so any later alloc()/free path that inspects the req + # (e.g. the alloc-skip check on `req.mamba_ping_pong_track_buffer + # is None`, or the retract cleanup) sees no dangling pointers + # into slot-owned tensors. Without this the alloc path can decide + # the req still has a ping-pong buffer and skip alloc, causing + # the slot's tensor to be reused by a new req and leaked when + # the slot is later freed. req.req_pool_idx = None req.mamba_pool_idx = None + req.mamba_ping_pong_track_buffer = None + req.mamba_next_track_idx = None + req.mamba_last_track_seqlen = None + req.mamba_branching_seqlen = None def restore_to_req(self, req: Req): """Restore KV state from this slot into an incoming request.""" @@ -288,14 +300,23 @@ class StreamingSession(BasePrefixCache): # slot from req state so release_session handles cleanup. # Include last_node/cache_protected_len from the req so # release_session calls dec_lock_ref on the tree lock. + # Also carry the mamba refs over so _free_slot_mamba can + # return the (possibly extra_buffer ping-pong) slots to + # the mamba pool; otherwise the abort orphans them. slot = SessionSlot( req_pool_idx=req.req_pool_idx, kv_allocated_len=req.kv_allocated_len, last_node=req.last_node, cache_protected_len=req.cache_protected_len, swa_uuid_for_lock=req.swa_uuid_for_lock, + mamba_pool_idx=req.mamba_pool_idx, + mamba_ping_pong_track_buffer=req.mamba_ping_pong_track_buffer, ) self.slots[session_id] = slot + # Slot now owns the mamba state — drop the req's refs so + # the abort fall-through doesn't double-free. + req.mamba_pool_idx = None + req.mamba_ping_pong_track_buffer = None slot.kv_allocated_len = max(slot.kv_allocated_len, req.kv_allocated_len) self.release_session(session_id) req.req_pool_idx = None