Skip mamba_pool_idx revert for session requests in _get_new_batch_prefill_raw (#23327)
This commit is contained in:
@@ -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)
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user