[PD / HiCache]fix deocde kvcache offload manager memory leak (#12774)

This commit is contained in:
huangtingwei
2025-11-14 03:31:53 +08:00
committed by GitHub
parent 0779c3d148
commit 4ac65e3c7d
2 changed files with 15 additions and 9 deletions
@@ -184,14 +184,15 @@ class DecodeKVCacheOffloadManager:
finish_count -= 1 finish_count -= 1
def _release_finished_req(self, req: Req, prefill_offloaded_len: int): def _release_finished_req(self, req: Req, prefill_offloaded_len: int):
# FIXME: not sure which length to use here: kv_allocated_len or kv_committed_len kv_committed_len = req.pop_committed_kv_cache()
kv_indices = self.req_to_token_pool.req_to_token[ kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, prefill_offloaded_len : req.kv_allocated_len req.req_pool_idx, prefill_offloaded_len:kv_committed_len
] ]
# Free the incremental part of the request # Free the incremental part of the request
self.token_to_kv_pool_allocator.free(kv_indices) self.token_to_kv_pool_allocator.free(kv_indices)
self.req_to_token_pool.free(req.req_pool_idx) self.req_to_token_pool.free(req.req_pool_idx)
self.tree_cache.protected_size_ -= len(req.prefix_indices)
def _check_backup_progress(self, finish_count): def _check_backup_progress(self, finish_count):
"""Check the progress of backup from host to storage.""" """Check the progress of backup from host to storage."""
+12 -7
View File
@@ -1762,13 +1762,18 @@ class ServerArgs:
"and cannot be used at the same time. Please use only one of them." "and cannot be used at the same time. Please use only one of them."
) )
if ( if self.disaggregation_decode_enable_offload_kvcache:
self.disaggregation_decode_enable_offload_kvcache if self.disaggregation_mode != "decode":
and self.disaggregation_mode != "decode" raise ValueError(
): "The argument disaggregation-decode-enable-offload-kvcache is only supported for decode side."
raise ValueError( )
"The argument disaggregation-decode-enable-offload-kvcache is only supported for decode side." if (
) self.disaggregation_mode == "decode"
and envs.SGLANG_ENABLE_SPEC_V2.get()
):
raise ValueError(
"Spec v2 and decode offload kv cache are incompatible and cannot be enabled together."
)
def _handle_metrics_labels(self): def _handle_metrics_labels(self):
if ( if (