diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index e7bee63c0..05c31b6c3 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -494,6 +494,19 @@ class RadixCache(BasePrefixCache): ) return + if not is_insert: + # Frees committed slots that no token id names, which the insert + # path below cannot reach; the protected prefix stays with the cache. + kv_indices = self.req_to_token_pool.req_to_token[ + req.kv.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle + ] + self.token_to_kv_pool_allocator.free_segment( + kv_indices, start_pos=req.kv.cache_protected_len + ) + if req.last_node is not None: + self.dec_lock_ref(req.last_node) + return + token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle] kv_indices = self.req_to_token_pool.req_to_token[ req.kv.req_pool_idx, : len(token_ids) @@ -509,39 +522,36 @@ class RadixCache(BasePrefixCache): values = kv_indices[:key_len].to(dtype=torch.int64, copy=True) # Radix Cache takes one ref in memory pool - if is_insert: - priority = getattr(req, "priority", 0) or 0 - result = self.insert( - InsertParams(key=radix_key, value=values, priority=priority) - ) - # A request that was never cached while unfinished can add its - # whole prompt and generated output as one leaf. Split that leaf at - # the prompt boundary so LRU eviction can discard output KV without - # also losing the reusable prompt KV. Reinserting a prefix only - # changes radix topology; it reuses the indices inserted above. - prompt_key = RadixKey( - token_ids[: len(req.origin_input_ids)], - req.extra_key, - is_bigram=self.is_eagle, - cache_salt=req.cache_salt, - ).page_aligned(self.page_size) - if 0 < len(prompt_key) < key_len: - self.insert( - InsertParams( - key=prompt_key, - value=values[: len(prompt_key)], - priority=priority + 1, - # Topology-only re-insert: this request created these - # nodes moments ago, so counting it as a hit is the - # same self-referencing inflation `chunked` exists to - # suppress. hit_count drives eviction order, so an - # extra bump would silently promote every prompt node. - chunked=True, - ) + priority = getattr(req, "priority", 0) or 0 + result = self.insert( + InsertParams(key=radix_key, value=values, priority=priority) + ) + # A request that was never cached while unfinished can add its + # whole prompt and generated output as one leaf. Split that leaf at + # the prompt boundary so LRU eviction can discard output KV without + # also losing the reusable prompt KV. Reinserting a prefix only + # changes radix topology; it reuses the indices inserted above. + prompt_key = RadixKey( + token_ids[: len(req.origin_input_ids)], + req.extra_key, + is_bigram=self.is_eagle, + cache_salt=req.cache_salt, + ).page_aligned(self.page_size) + if 0 < len(prompt_key) < key_len: + self.insert( + InsertParams( + key=prompt_key, + value=values[: len(prompt_key)], + priority=priority + 1, + # Topology-only re-insert: this request created these + # nodes moments ago, so counting it as a hit is the + # same self-referencing inflation `chunked` exists to + # suppress. hit_count drives eviction order, so an + # extra bump would silently promote every prompt node. + chunked=True, ) - freed_end = result.prefix_len - else: - freed_end = key_len + ) + freed_end = result.prefix_len # duplicates / uninserted range, then the unaligned tail self.token_to_kv_pool_allocator.free_segments( diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index 601175449..172d44187 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -47,6 +47,7 @@ from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.srt.utils.common import Range +from sglang.test.test_utils import CustomTestCase def _make_cache_with_pools(page_size=1): @@ -104,7 +105,7 @@ def _make_req(fill_ids, req_pool_idx=0, cache_protected_len=0, last_node=None): return MockReq(fill_ids, req_pool_idx, cache_protected_len, last_node) -class TestDecodeLockRefScenarios(unittest.TestCase): +class TestDecodeLockRefScenarios(CustomTestCase): def setUp(self): # The decode queue reads its config from the bags. reset_context() @@ -290,7 +291,10 @@ class TestDecodeLockRefScenarios(unittest.TestCase): self.assertEqual(cache.evictable_size(), len(full_ids)) def test_incremental_transfer_failure(self): - """Scenario 3: prefix match > 0, transfer fails. + """Scenario 3: prefix match > 0, transfer fails after KV commit. + + The final committed KV slot has no corresponding output token. Cleanup + must preserve the matched prefix and release the full request-owned suffix. Flow: inc_lock_ref(pop_preallocated) -> dec_lock_ref(cache_finished_req via release_kv_cache is_insert=False) @@ -322,13 +326,20 @@ class TestDecodeLockRefScenarios(unittest.TestCase): cache_protected_len=prefix_len, last_node=matched_node, ) + req.output_ids = array("q") # Transfer fails -> cache_finished_req with is_insert=False - # This frees delta tokens and dec_lock_ref on last_node + cache.token_to_kv_pool_allocator.reset_mock() cache.cache_finished_req( req, is_insert=False, kv_len_to_handle=req.kv.kv_committed_len ) + free_call = cache.token_to_kv_pool_allocator.free_segment.call_args + torch.testing.assert_close( + free_call.args[0], torch.tensor(full_vals[prefix_len:]) + ) + self.assertEqual(free_call.kwargs["start_pos"], prefix_len) + # The prefix node should be unlocked (back to evictable) self.assertEqual(cache.root_node.lock_ref, 1) self.assertEqual(cache.protected_size(), 0)