[Fix] memory leak by overlap + retract (#11981)

Co-authored-by: Liangsheng Yin <lsyincs@gmail.com>
This commit is contained in:
cctry
2025-10-23 07:59:23 -07:00
committed by GitHub
parent 6c18addb6f
commit b0b4f71679
9 changed files with 132 additions and 25 deletions

View File

@@ -32,6 +32,8 @@ class ChunkCache(BasePrefixCache):
else:
self.device = torch.device("cpu")
self.protected_size_ = 0
# NOTE (csy): this is to determine if a cache has prefix matching feature.
# Chunk cache always return True to indicate no prefix matching.
# TODO (csy): Using a prefix cache trait to replace this
@@ -57,11 +59,13 @@ class ChunkCache(BasePrefixCache):
]
self.req_to_token_pool.free(req.req_pool_idx)
self.token_to_kv_pool_allocator.free(kv_indices)
self.protected_size_ -= len(req.prefix_indices)
def cache_unfinished_req(self, req: Req, chunked=False):
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, : len(req.fill_ids)
]
self.protected_size_ += len(kv_indices) - len(req.prefix_indices)
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
@@ -75,6 +79,9 @@ class ChunkCache(BasePrefixCache):
def dec_lock_ref(self, node: Any, swa_uuid_for_lock: Optional[str] = None):
return 0
def protected_size(self):
return self.protected_size_
def pretty_print(self):
return ""