[Scheduler] Align RadixCache no-insert cleanup with kv_len_to_handle (#35204)

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,7 +522,6 @@ 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)
@@ -540,8 +552,6 @@ class RadixCache(BasePrefixCache):
                 )
             )
         freed_end = result.prefix_len
-        else:
-            freed_end = key_len
 
         # duplicates / uninserted range, then the unaligned tail
         self.token_to_kv_pool_allocator.free_segments(

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)