fix mamba cache leak when adder fails to add a matched req. (#21404)
This commit is contained in:
@@ -2416,6 +2416,13 @@ 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
|
||||||
|
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:
|
||||||
|
self.tree_cache.req_to_token_pool.mamba_pool.free(
|
||||||
|
req.mamba_pool_idx.unsqueeze(-1)
|
||||||
|
)
|
||||||
|
req.mamba_pool_idx = None
|
||||||
break
|
break
|
||||||
|
|
||||||
# Update waiting queue
|
# Update waiting queue
|
||||||
|
|||||||
Reference in New Issue
Block a user