Plug mamba_extra_buffer ping-pong slot leaks (#26941)
This commit is contained in:
@@ -790,6 +790,12 @@ class HybridReqToTokenPool(ReqToTokenPool):
|
|||||||
]
|
]
|
||||||
)
|
)
|
||||||
self.mamba_pool.free(mamba_ping_pong_track_buffer_to_free)
|
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):
|
def clear(self):
|
||||||
logger.info("Reset HybridReqToTokenPool")
|
logger.info("Reset HybridReqToTokenPool")
|
||||||
|
|||||||
@@ -85,8 +85,20 @@ class SessionSlot:
|
|||||||
self.mamba_last_track_seqlen = req.mamba_last_track_seqlen
|
self.mamba_last_track_seqlen = req.mamba_last_track_seqlen
|
||||||
self.mamba_branching_seqlen = req.mamba_branching_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.req_pool_idx = None
|
||||||
req.mamba_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):
|
def restore_to_req(self, req: Req):
|
||||||
"""Restore KV state from this slot into an incoming request."""
|
"""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.
|
# slot from req state so release_session handles cleanup.
|
||||||
# Include last_node/cache_protected_len from the req so
|
# Include last_node/cache_protected_len from the req so
|
||||||
# release_session calls dec_lock_ref on the tree lock.
|
# 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(
|
slot = SessionSlot(
|
||||||
req_pool_idx=req.req_pool_idx,
|
req_pool_idx=req.req_pool_idx,
|
||||||
kv_allocated_len=req.kv_allocated_len,
|
kv_allocated_len=req.kv_allocated_len,
|
||||||
last_node=req.last_node,
|
last_node=req.last_node,
|
||||||
cache_protected_len=req.cache_protected_len,
|
cache_protected_len=req.cache_protected_len,
|
||||||
swa_uuid_for_lock=req.swa_uuid_for_lock,
|
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
|
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)
|
slot.kv_allocated_len = max(slot.kv_allocated_len, req.kv_allocated_len)
|
||||||
self.release_session(session_id)
|
self.release_session(session_id)
|
||||||
req.req_pool_idx = None
|
req.req_pool_idx = None
|
||||||
|
|||||||
Reference in New Issue
Block a user