diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 78970d058..51b365b3d 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/python/sglang/srt/mem_cache/chunk_cache.py @@ -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:]) diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index b6d416168..b766be298 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -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 diff --git a/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py b/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py index aff31561e..dfe8cedb2 100644 --- a/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py +++ b/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py @@ -23,6 +23,7 @@ class _FakeAllocator: class _FakeReq: req_pool_idx = 0 swa_evict_floor = 3 + cache_protected_len = 0 kv = SimpleNamespace(swa_evicted_seqlen=6) def pop_committed_kv_cache(self): @@ -30,12 +31,16 @@ class _FakeReq: class TestPureSWAChunkCache(CustomTestCase): - def test_finished_req_skips_already_evicted_swa_range(self): + def _make_cache(self): cache = PureSWAChunkCache.__new__(PureSWAChunkCache) cache.req_to_token_pool = SimpleNamespace( req_to_token=torch.arange(10, dtype=torch.int64).unsqueeze(0) ) cache.token_to_kv_pool_allocator = _FakeAllocator() + return cache + + def test_finished_req_skips_already_evicted_swa_range(self): + cache = self._make_cache() cache.cache_finished_req(_FakeReq(), kv_len_to_handle=8) @@ -43,6 +48,16 @@ class TestPureSWAChunkCache(CustomTestCase): freed = cache.token_to_kv_pool_allocator.freed[0] self.assertTrue(torch.equal(freed, torch.tensor([0, 1, 2, 6, 7]))) + def test_finished_req_skips_protected_prefix(self): + cache = self._make_cache() + req = _FakeReq() + req.cache_protected_len = 2 + + cache.cache_finished_req(req, kv_len_to_handle=8) + + freed = cache.token_to_kv_pool_allocator.freed[0] + self.assertTrue(torch.equal(freed, torch.tensor([2, 6, 7]))) + if __name__ == "__main__": unittest.main()