[Fix] Respect cache_protected_len in ChunkCache and disabled-radix release paths (#31662)

This commit is contained in:
Liangsheng Yin
2026-07-18 12:04:25 -07:00
committed by GitHub
parent b3a0185cab
commit 10908a6793
3 changed files with 32 additions and 10 deletions
+14 -8
View File
@@ -80,8 +80,9 @@ class ChunkCache(BasePrefixCache):
self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
):
# For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids
# The protected prefix is not this req's to free.
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_len_to_handle
req.req_pool_idx, req.cache_protected_len : kv_len_to_handle
]
self.token_to_kv_pool_allocator.free(kv_indices)
@@ -146,9 +147,10 @@ class PureSWAChunkCache(SWAChunkCache):
explicitly skip the range already freed by ``free_swa_out_of_window_slots``
(a.k.a. _evict_swa) during decode.
``req.swa_evict_floor`` only protects the prompt/image KV while the request
is active. ChunkCache does not retain finished prefixes, so the protected
prefix is released here when the request finishes.
``req.swa_evict_floor`` shields the prompt/image KV from window eviction
only while the request is active, so that range IS released here on
finish. Distinct from the ``cache_protected_len`` prefix, which is owned
elsewhere and never freed by this path.
"""
def cache_finished_req(
@@ -158,15 +160,19 @@ class PureSWAChunkCache(SWAChunkCache):
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_committed_len
]
# The cache_protected_len prefix is not this req's to free.
protected_len = req.cache_protected_len
evict_floor = req.swa_evict_floor
evicted_seqlen = req.kv.swa_evicted_seqlen
if evicted_seqlen > evict_floor:
parts = []
if evict_floor > 0:
parts.append(kv_indices[:evict_floor])
if evict_floor > protected_len:
parts.append(kv_indices[protected_len:evict_floor])
if evicted_seqlen < kv_committed_len:
parts.append(kv_indices[evicted_seqlen:kv_committed_len])
parts.append(
kv_indices[max(evicted_seqlen, protected_len) : kv_committed_len]
)
if parts:
self.token_to_kv_pool_allocator.free(torch.cat(parts))
else:
self.token_to_kv_pool_allocator.free(kv_indices)
self.token_to_kv_pool_allocator.free(kv_indices[protected_len:])
+2 -1
View File
@@ -443,8 +443,9 @@ class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache):
is_insert = False
if self.disable:
# The protected prefix is not this req's to free.
kv_indices = self.req_to_token_pool.req_to_token[
req.req_pool_idx, :kv_len_to_handle
req.req_pool_idx, req.cache_protected_len : kv_len_to_handle
]
self.token_to_kv_pool_allocator.free(kv_indices)
return