Let cache backend do not couple with owned committed kv details and avoid kv_committed_freed/kv_overallocated_freed fields (#29428)

python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py

@@ -255,7 +255,7 @@ class DecodeKVCacheOffloadManager:
         if req.req_pool_idx is None or req.req_pool_idx == -1:
             return
 
-        kv_committed_len = req.pop_committed_kv_cache()
+        kv_committed_len = req.effective_kv_committed_len()
 
         # Free the prefill-aligned slots. Previously this was done
         # eagerly in offload_kv_cache (mid-decode), which raced with
@@ -276,7 +276,7 @@ class DecodeKVCacheOffloadManager:
 
         # Free over-allocated KV cache slots (e.g. from speculative decoding v2).
         # Without spec v2, start_p == end_p so this is a no-op.
-        start_p, end_p = req.pop_overallocated_kv_cache()
+        start_p, end_p = kv_committed_len, req.kv.kv_allocated_len
         if self.page_size > 1:
             start_p = ceil_align(start_p, self.page_size)
         if start_p < end_p:

python/sglang/srt/managers/schedule_batch.py

@@ -749,8 +749,6 @@ class Req(ReqDllmMixin):
         # For req-level memory management
         self.kv_committed_len = 0
         self.kv: ReqKvInfo = ReqKvInfo(kv_allocated_len=0, swa_evicted_seqlen=0)
-        self.kv_committed_freed = False
-        self.kv_overallocated_freed = False
 
         # for cross-encoder model
         self.token_type_ids = token_type_ids
@@ -1076,33 +1074,13 @@ class Req(ReqDllmMixin):
             or self.mamba_host_hit_length > 0
         )
 
-    def _cache_commit_len(self) -> int:
+    def effective_kv_committed_len(self) -> int:
         # Report only the prompt prefix so thinking + answer fall into the
         # overallocated range and are reclaimed by release_kv_cache. #22373.
         if get_server_args().strip_thinking_cache and self.reasoning_tokens > 0:
             return min(self.kv_committed_len, len(self.origin_input_ids))
         return self.kv_committed_len
 
-    def pop_committed_kv_cache(self) -> int:
-        """Return the length of committed KV cache and mark them as freed."""
-        assert (
-            not self.kv_committed_freed
-        ), f"Committed KV cache already freed ({self.kv_committed_len=})"
-        self.kv_committed_freed = True
-        return self._cache_commit_len()
-
-    def pop_overallocated_kv_cache(self) -> Tuple[int, int]:
-        """Return the range of over-allocated KV cache and mark them as freed."""
-
-        # NOTE: This function is called when there is over-allocation of KV cache.
-        # Over-allocation: we allocate more KV cache than the committed length.
-        # e.g., speculative decoding may allocate more KV cache than actually used.
-        assert (
-            not self.kv_overallocated_freed
-        ), f"Overallocated KV cache already freed, {self.kv_committed_len=}, {self.kv.kv_allocated_len=}"
-        self.kv_overallocated_freed = True
-        return self._cache_commit_len(), self.kv.kv_allocated_len
-
     def update_spec_correct_drafts_histogram(self, num_correct_drafts: int):
         """Update the speculative decoding acceptance histogram.
 
@@ -1525,8 +1503,6 @@ class Req(ReqDllmMixin):
         self.already_computed = 0
         self.kv.kv_allocated_len = 0
         self.kv_committed_len = 0
-        self.kv_committed_freed = False
-        self.kv_overallocated_freed = False
         self.kv.swa_evicted_seqlen = 0
         self.extend_batch_idx = 0
         self.decode_batch_idx = 0

python/sglang/srt/managers/scheduler.py

@@ -2564,9 +2564,7 @@ class Scheduler(
             req.pending_bootstrap = False
         if self.enable_hicache_storage:
             self.tree_cache.release_aborted_request(req.rid)
-        if (
+        if req.req_pool_idx is not None or self.tree_cache.supports_mamba():
-            req.req_pool_idx is not None or self.tree_cache.supports_mamba()
-        ) and not req.kv_committed_freed:
             release_kv_cache(req, self.tree_cache, is_insert=False)
 
         self.chunked_req = None

python/sglang/srt/managers/scheduler_components/invariant_checker.py

@@ -246,8 +246,7 @@ class SchedulerInvariantChecker:
         swa_uncached = 0
         for batch in batches:
             for req in batch.reqs:
-                assert req.kv_committed_freed == req.kv_overallocated_freed
+                if req.req_pool_idx is None:
-                if req.kv_committed_freed or req.req_pool_idx is None:
                     continue
 
                 allocated_len = req.kv.kv_allocated_len

python/sglang/srt/mem_cache/chunk_cache.py

@@ -76,11 +76,12 @@ class ChunkCache(BasePrefixCache):
         # ChunkCache does not support prefix caching, so insert is a no-op
         return InsertResult(prefix_len=0)
 
-    def cache_finished_req(self, req: Req, is_insert: bool = True):
+    def cache_finished_req(
-        kv_committed_len = req.pop_committed_kv_cache()
+        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
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, :kv_committed_len
+            req.req_pool_idx, :kv_len_to_handle
         ]
         self.token_to_kv_pool_allocator.free(kv_indices)
 
@@ -150,8 +151,10 @@ class PureSWAChunkCache(SWAChunkCache):
     prefix is released here when the request finishes.
     """
 
-    def cache_finished_req(self, req: Req, is_insert: bool = True):
+    def cache_finished_req(
-        kv_committed_len = req.pop_committed_kv_cache()
+        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+    ):
+        kv_committed_len = kv_len_to_handle
         kv_indices = self.req_to_token_pool.req_to_token[
             req.req_pool_idx, :kv_committed_len
         ]

python/sglang/srt/mem_cache/common.py

@@ -640,17 +640,19 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
             req.mamba_pool_idx = None
         return
 
+    effective_kv_committed_len = req.effective_kv_committed_len()
     tree_cache.cache_finished_req(
         req,
         is_insert=is_insert and not getattr(req, "skip_radix_cache_insert", False),
+        kv_len_to_handle=effective_kv_committed_len,
     )
 
     # StreamingSession.cache_finished_req handles speculative tail trim
-    # and bookkeeping flag sync internally, then sets req_pool_idx = None.
+    # internally, then sets req_pool_idx = None.
     if req.req_pool_idx is None:
         return
 
-    start_p, end_p = req.pop_overallocated_kv_cache()
+    start_p, end_p = effective_kv_committed_len, req.kv.kv_allocated_len
 
     global_server_args = get_server_args()
     page_size = global_server_args.page_size

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -523,20 +523,21 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
         )
         return InsertResult(prefix_len=prefix_len, mamba_exist=mamba_exist)
 
-    def cache_finished_req(self, req: Req, is_insert: bool = True) -> None:
+    def cache_finished_req(
+        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+    ) -> None:
         """Cache request when it finishes."""
-        kv_committed_len = req.pop_committed_kv_cache()
         if self.disable:
             kv_indices = self.req_to_token_pool.req_to_token[
-                req.req_pool_idx, :kv_committed_len
+                req.req_pool_idx, :kv_len_to_handle
             ]
             self.token_to_kv_pool_allocator.free(kv_indices)
             self.req_to_token_pool.free_mamba_cache(req)
             return
 
-        token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
+        token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, :kv_committed_len
+            req.req_pool_idx, :kv_len_to_handle
         ]
 
         if is_insert:
@@ -572,7 +573,7 @@ class MambaRadixCache(KVCacheEventMixin, BasePrefixCache):
 
             assert (
                 cache_len == page_aligned_len
-            ), f"It is required {cache_len=}, {page_aligned_len=}, {kv_committed_len=}, {len(req.origin_input_ids)=}, {len(req.output_ids)=} ping @yizhang2077 if you see this"
+            ), f"It is required {cache_len=}, {page_aligned_len=}, {kv_len_to_handle=}, {len(req.origin_input_ids)=}, {len(req.output_ids)=} ping @yizhang2077 if you see this"
 
             # Radix Cache takes one ref in memory pool
             # insert the token_ids and kv_indices into the radix tree

python/sglang/srt/mem_cache/pure_swa_radix_cache.py

@@ -62,7 +62,9 @@ class PureSWARadixCache(RadixCache):
         num_tokens = max(params.num_tokens, params.swa_num_tokens)
         return super().evict(EvictParams(num_tokens=num_tokens))
 
-    def cache_finished_req(self, req: Req, is_insert: bool = True):
+    def cache_finished_req(
+        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+    ):
         """Cache request when it finishes.
 
         Only inserts the prefill portion [0, evict_floor) into the radix tree.
@@ -73,7 +75,7 @@ class PureSWARadixCache(RadixCache):
         if self.disable_finished_insert:
             is_insert = False
 
-        kv_committed_len = req.pop_committed_kv_cache()
+        kv_committed_len = kv_len_to_handle
         if self.disable:
             kv_indices = self.req_to_token_pool.req_to_token[
                 req.req_pool_idx, :kv_committed_len

python/sglang/srt/mem_cache/radix_cache.py

@@ -434,21 +434,22 @@ class RadixCache(SessionRadixCacheMixin, KVCacheEventMixin, BasePrefixCache):
         )
         return InsertResult(prefix_len=prefix_len, last_device_node=last_node)
 
-    def cache_finished_req(self, req: Req, is_insert: bool = True):
+    def cache_finished_req(
+        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+    ):
         """Cache request when it finishes."""
         # In deterministic mode, disable finished request insertion to radix cache
         if self.disable_finished_insert:
             is_insert = False
 
-        kv_committed_len = req.pop_committed_kv_cache()
         if self.disable:
             kv_indices = self.req_to_token_pool.req_to_token[
-                req.req_pool_idx, :kv_committed_len
+                req.req_pool_idx, :kv_len_to_handle
             ]
             self.token_to_kv_pool_allocator.free(kv_indices)
             return
 
-        token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
+        token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
         kv_indices = self.req_to_token_pool.req_to_token[
             req.req_pool_idx, : len(token_ids)
         ]

python/sglang/srt/mem_cache/radix_cache_cpp.py

@@ -169,19 +169,20 @@ class RadixCacheCpp(BasePrefixCache):
     def total_size(self):
         return self.tree.total_size()
 
-    def cache_finished_req(self, req: Req, is_insert: bool = True):
+    def cache_finished_req(
+        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+    ):
         """Cache request when it finishes."""
         assert req.req_pool_idx is not None
-        kv_committed_len = req.pop_committed_kv_cache()
+        token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
-        token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, :kv_committed_len
+            req.req_pool_idx, :kv_len_to_handle
         ].to(dtype=torch.int64, copy=True)
 
         # NOTE: our C++ implementation don't need `token_ids` and `kv_indices` to be page-aligned
         # it will automatically align them, but length of them should be equal
         old_prefix_len = len(req.prefix_indices) // self.page_size * self.page_size
-        page_aligned_overall_len = kv_committed_len // self.page_size * self.page_size
+        page_aligned_overall_len = kv_len_to_handle // self.page_size * self.page_size
 
         if is_insert:
             new_prefix_len = self._insert(
@@ -200,7 +201,7 @@ class RadixCacheCpp(BasePrefixCache):
             )
 
         # need to free the unaligned part, since it cannot be inserted into the radix tree
-        if page_aligned_overall_len < kv_committed_len:
+        if page_aligned_overall_len < kv_len_to_handle:
             # NOTE: sglang PagedAllocator support unaligned free (which will automatically align it)
             self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_overall_len:])
 

python/sglang/srt/mem_cache/storage/flexkv/flexkv_radix_cache.py

@@ -379,10 +379,12 @@ class FlexKVRadixCache(RadixCache):
     # ------------------------------------------------------------------
 
     def cache_finished_req(  # type: ignore[override]
-        self, req: Req, is_insert: bool = True
+        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
     ) -> None:
         """Base cache_finished_req then fire an async FlexKV store."""
-        super().cache_finished_req(req, is_insert=is_insert)
+        super().cache_finished_req(
+            req, is_insert=is_insert, kv_len_to_handle=kv_len_to_handle
+        )
         if not is_insert:
             self._load_markers.pop(req.rid, None)
             return

python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py

@@ -428,10 +428,14 @@ class LMCRadixCache(RadixCache):
                 )
             )
 
-    def cache_finished_req(self, req: Req, is_insert: bool = True) -> None:
+    def cache_finished_req(
+        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+    ) -> None:
         """On request completion, insert device KV into radix and store to LMCache."""
 
-        super().cache_finished_req(req, is_insert=is_insert)
+        super().cache_finished_req(
+            req, is_insert=is_insert, kv_len_to_handle=kv_len_to_handle
+        )
         if not is_insert:
             if self._mode is LMCacheMode.MP:
                 self._mp_load_back_markers.pop(req.rid, None)

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -456,19 +456,20 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
         )
         return InsertResult(prefix_len=prefix_len)
 
-    def cache_finished_req(self, req: Req, is_insert: bool = True) -> None:
+    def cache_finished_req(
+        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+    ) -> None:
         """Cache request when it finishes."""
-        kv_committed_len = req.pop_committed_kv_cache()
         if self.disable:
             kv_indices = self.req_to_token_pool.req_to_token[
-                req.req_pool_idx, :kv_committed_len
+                req.req_pool_idx, :kv_len_to_handle
             ]
             self.token_to_kv_pool_allocator.free(kv_indices)
             return
 
-        token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
+        token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, :kv_committed_len
+            req.req_pool_idx, :kv_len_to_handle
         ]
 
         radix_key = RadixKey(

python/sglang/srt/mem_cache/unified_cache_components/README.md

@@ -199,14 +199,14 @@ Unlock a previously locked node path.
 
 ---
 
-### `cache_finished_req(req: Req, is_insert: bool = True)`
+### `cache_finished_req(req: Req, is_insert: bool = True, *, kv_len_to_handle: int)`
 
 Cache a completed request's KV data into the tree.
 
 | Aspect | Detail |
 |--------|--------|
 | **Purpose** | After a request finishes, insert its token/KV data into the tree for future reuse |
-| **Inputs** | `req` — the finished request; `is_insert` — whether to insert (True) or just release locks (False) |
+| **Inputs** | `req` — the finished request; `is_insert` — whether to insert (True) or just release locks (False); `kv_len_to_handle` — committed KV length supplied by the caller |
 | **Output** | `None` |
 | **Mutation** | Calls component hooks → `insert` → `dec_lock_ref` → component cleanup. Frees unaligned tail KV indices; frees non-inserted KV indices when `is_insert=False`. |
 | **Complexity** | **O(K + D·C)** — insert O(K + D·C) + lock release O(D). Simplifies to **O(K)**. |

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -712,24 +712,24 @@ class UnifiedRadixCache(KVCacheEventMixin, BasePrefixCache):
         self._update_evictable_leaf_sets(node)
         return DecLockRefResult()
 
-    def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs) -> None:
+    def cache_finished_req(
+        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int, **kwargs
+    ) -> None:
         if self.session.try_cache_finished_req(req, is_insert=is_insert, **kwargs):
             return
 
-        kv_committed_len = req.pop_committed_kv_cache()
-
         if self.disable:
             kv_indices = self.req_to_token_pool.req_to_token[
-                req.req_pool_idx, :kv_committed_len
+                req.req_pool_idx, :kv_len_to_handle
             ]
             self.token_to_kv_pool_allocator.free(kv_indices)
             for comp in self._components_tuple:
                 comp.cleanup_after_caching_req(req, is_finished=True)
             return
 
-        token_ids = (req.origin_input_ids + req.output_ids)[:kv_committed_len]
+        token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, :kv_committed_len
+            req.req_pool_idx, :kv_len_to_handle
         ]
 
         result = None

python/sglang/srt/session/streaming_session.py

@@ -325,7 +325,6 @@ class StreamingSession(BasePrefixCache):
             self.release_session(session_id)
             req.req_pool_idx = None
             req.session.abort_req()
-            self._mark_kv_freed(req)
             return True
 
         if is_first:
@@ -348,7 +347,6 @@ class StreamingSession(BasePrefixCache):
         # Update req_nodes to this successfully finished request.
         req.session.finish_req(req)
 
-        self._mark_kv_freed(req)
         return True
 
     def try_cache_unfinished_req(
@@ -570,14 +568,6 @@ class StreamingSession(BasePrefixCache):
             tail = self.req_to_token_pool.req_to_token[pool_idx, start:end]
             self.token_to_kv_pool_allocator.free(tail)
 
-    @staticmethod
-    def _mark_kv_freed(req: Req) -> None:
-        """Set bookkeeping flags so busy check skips this finished req."""
-        if not req.kv_committed_freed:
-            req.pop_committed_kv_cache()
-        if not req.kv_overallocated_freed:
-            req.pop_overallocated_kv_cache()
-
     # -- Pass-through methods --
 
     def evictable_size(self):

test/registered/disaggregation/test_specv2_kvcache_offloading.py

@@ -36,22 +36,8 @@ def _make_mock_req(
     req.req_pool_idx = req_pool_idx
     req.kv_committed_len = kv_committed_len
     req.kv = SimpleNamespace(kv_allocated_len=kv_allocated_len)
-    req.kv_committed_freed = False
-    req.kv_overallocated_freed = False
     req.prefix_indices = list(range(prefix_indices_len))
-
+    req.effective_kv_committed_len = lambda: req.kv_committed_len
-    def pop_committed():
-        assert not req.kv_committed_freed
-        req.kv_committed_freed = True
-        return req.kv_committed_len
-
-    def pop_overallocated():
-        assert not req.kv_overallocated_freed
-        req.kv_overallocated_freed = True
-        return req.kv_committed_len, req.kv.kv_allocated_len
-
-    req.pop_committed_kv_cache = pop_committed
-    req.pop_overallocated_kv_cache = pop_overallocated
     return req
 
 

test/registered/unit/mem_cache/test_decode_radix_lock_ref.py

@@ -83,18 +83,10 @@ class MockReq:
         self.priority = 0
         self.kv_committed_len = len(fill_ids)
         self.kv = SimpleNamespace(kv_allocated_len=len(fill_ids))
-        self.kv_committed_freed = False
 
     def get_fill_ids(self):
         return self.full_untruncated_fill_ids[: self.extend_range.end]
 
-    def pop_committed_kv_cache(self):
-        self.kv_committed_freed = True
-        return self.kv_committed_len
-
-    def pop_overallocated_kv_cache(self):
-        return (self.kv_committed_len, self.kv.kv_allocated_len)
-
 
 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)
@@ -152,7 +144,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
         cache.cache_unfinished_req(req)
 
         # Step 3: cache_finished_req with is_insert=True (dec lock)
-        cache.cache_finished_req(req)
+        cache.cache_finished_req(req, kv_len_to_handle=req.kv_committed_len)
 
         # Verify: all non-root nodes should have lock_ref == 0
         # (root always has lock_ref == 1)
@@ -201,7 +193,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
         cache.cache_unfinished_req(req)
 
         # Step 3: cache_finished_req (dec leaf)
-        cache.cache_finished_req(req)
+        cache.cache_finished_req(req, kv_len_to_handle=req.kv_committed_len)
 
         # Root lock unchanged, all nodes unlocked
         self.assertEqual(cache.root_node.lock_ref, root_lock_before)
@@ -244,7 +236,9 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
 
         # Transfer fails -> cache_finished_req with is_insert=False
         # This frees delta tokens and dec_lock_ref on last_node
-        cache.cache_finished_req(req, is_insert=False)
+        cache.cache_finished_req(
+            req, is_insert=False, kv_len_to_handle=req.kv_committed_len
+        )
 
         # The prefix node should be unlocked (back to evictable)
         self.assertEqual(cache.root_node.lock_ref, 1)
@@ -289,7 +283,9 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
 
         # Transfer fails -> cache_finished_req with is_insert=False
         # dec_lock_ref(root) is a no-op
-        cache.cache_finished_req(req, is_insert=False)
+        cache.cache_finished_req(
+            req, is_insert=False, kv_len_to_handle=req.kv_committed_len
+        )
 
         # Root lock unchanged, nothing protected or evictable
         self.assertEqual(cache.root_node.lock_ref, root_lock_before)
@@ -400,7 +396,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
             )
 
             cache.cache_unfinished_req(req)
-            cache.cache_finished_req(req)
+            cache.cache_finished_req(req, kv_len_to_handle=req.kv_committed_len)
 
         # After all iterations, root lock should be 1, no protected nodes
         self.assertEqual(cache.root_node.lock_ref, 1)

test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py

@@ -37,7 +37,7 @@ class TestPureSWAChunkCache(CustomTestCase):
         )
         cache.token_to_kv_pool_allocator = _FakeAllocator()
 
-        cache.cache_finished_req(_FakeReq())
+        cache.cache_finished_req(_FakeReq(), kv_len_to_handle=8)
 
         self.assertEqual(len(cache.token_to_kv_pool_allocator.freed), 1)
         freed = cache.token_to_kv_pool_allocator.freed[0]

test/registered/unit/mem_cache/test_streaming_session_unit.py

@@ -61,8 +61,6 @@ class _FakeReq:
             kv_allocated_len=allocated,
             swa_evicted_seqlen=0,
         )
-        self.kv_committed_freed = False
-        self.kv_overallocated_freed = False
         self.origin_input_ids = list(range(committed))
         self.output_ids = []
         self.extra_key = None
@@ -74,22 +72,10 @@ class _FakeReq:
         self.mamba_next_track_idx = None
         self.mamba_last_track_seqlen = None
         self.mamba_branching_seqlen = None
-        self.pop_overallocated_calls = 0
         self.to_finish = None
         self.finished_reason = None
         self.finished_len = None
 
-    def pop_committed_kv_cache(self):
-        assert not self.kv_committed_freed
-        self.kv_committed_freed = True
-        return self.kv_committed_len
-
-    def pop_overallocated_kv_cache(self):
-        assert not self.kv_overallocated_freed
-        self.pop_overallocated_calls += 1
-        self.kv_overallocated_freed = True
-        return self.kv_committed_len, self.kv.kv_allocated_len
-
 
 def test_preabort_detaches_session_and_preserves_slot():
     """Pre-aborted req (to_finish set before match_prefix) is detached from
@@ -161,9 +147,6 @@ def test_first_mid_abort_nukes_ephemeral_slot():
     assert req_to_token_pool.free_slots == [0]
     assert len(allocator.freed) == 1
     assert allocator.freed[0].tolist() == list(range(20))
-    # Bookkeeping flags set.
-    assert req.kv_committed_freed is True
-    assert req.kv_overallocated_freed is True
 
 
 def test_nth_mid_abort_nukes_session_slot():
@@ -200,9 +183,6 @@ def test_nth_mid_abort_nukes_session_slot():
     # Pool slot returned.
     assert req_to_token_pool.free_slots == [0]
     assert req.req_pool_idx is None
-    # Bookkeeping flags set.
-    assert req.kv_committed_freed is True
-    assert req.kv_overallocated_freed is True
 
 
 # Shrink tests removed: streaming sessions are append-only after the

test/registered/unit/mem_cache/test_swa_eviction_boundary.py

@@ -116,7 +116,6 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree):
         prefix_indices=torch.tensor([], dtype=torch.int64, device=tree.device),
         _kv_committed_len=len(token_ids),
     )
-    req.pop_committed_kv_cache = lambda: req._kv_committed_len
     return req
 
 
@@ -191,7 +190,9 @@ class TestSWAEvictionBoundary(unittest.TestCase):
             insert_len = seq_len // page_size * page_size
             self.assertLess(req.kv.swa_evicted_seqlen, insert_len)
 
-            tree.cache_finished_req(req, is_insert=True)
+            tree.cache_finished_req(
+                req, is_insert=True, kv_len_to_handle=req._kv_committed_len
+            )
             tree.sanity_check()
 
     # -- Eviction formula: page_size == 1 --
@@ -216,7 +217,9 @@ class TestSWAEvictionBoundary(unittest.TestCase):
                 req.kv.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size))
             )
 
-            tree.cache_finished_req(req, is_insert=True)
+            tree.cache_finished_req(
+                req, is_insert=True, kv_len_to_handle=req._kv_committed_len
+            )
             tree.sanity_check()
 
     # -- Eviction formula: no-op when seq too short --
@@ -253,7 +256,9 @@ class TestSWAEvictionBoundary(unittest.TestCase):
         kv1 = _swa_alloc(allocator, first_len)
         pool.write((0, slice(0, first_len)), kv1)
         req1 = _make_req(0, list(range(first_len)), 0, tree)
-        tree.cache_finished_req(req1, is_insert=True)
+        tree.cache_finished_req(
+            req1, is_insert=True, kv_len_to_handle=req1._kv_committed_len
+        )
         tree.sanity_check()
 
         # Second request: 24 tokens, first 16 overlap with tree
@@ -269,7 +274,9 @@ class TestSWAEvictionBoundary(unittest.TestCase):
         self.assertLessEqual(req2.kv.swa_evicted_seqlen, first_len)
 
         swa_evictable_before = tree.swa_evictable_size_
-        tree.cache_finished_req(req2, is_insert=True)
+        tree.cache_finished_req(
+            req2, is_insert=True, kv_len_to_handle=req2._kv_committed_len
+        )
 
         # New tokens [16, 24) should all be non-tombstone
         new_tokens = second_len // page_size * page_size - first_len
@@ -300,7 +307,9 @@ class TestSWAEvictionBoundary(unittest.TestCase):
         self.assertGreater(req.kv.swa_evicted_seqlen, 0, "Should have some eviction")
         self.assertLess(req.kv.swa_evicted_seqlen, insert_len, "Should be partial")
 
-        tree.cache_finished_req(req, is_insert=True)
+        tree.cache_finished_req(
+            req, is_insert=True, kv_len_to_handle=req._kv_committed_len
+        )
 
         non_tombstone = insert_len - req.kv.swa_evicted_seqlen
         self.assertEqual(tree.swa_evictable_size_, swa_evictable_before + non_tombstone)
@@ -338,7 +347,9 @@ class TestSWAEvictionBoundary(unittest.TestCase):
         req.kv.swa_evicted_seqlen = old_evicted
         swa_evictable_before = tree.swa_evictable_size_
 
-        tree.cache_finished_req(req, is_insert=True)
+        tree.cache_finished_req(
+            req, is_insert=True, kv_len_to_handle=req._kv_committed_len
+        )
 
         self.assertEqual(tree.swa_evictable_size_, swa_evictable_before)
 
@@ -367,7 +378,9 @@ class TestSWAEvictionBoundary(unittest.TestCase):
             insert_len = seq_len // page_size * page_size
             self.assertLess(req.kv.swa_evicted_seqlen, insert_len, f"turn {turn}")
 
-            tree.cache_finished_req(req, is_insert=True)
+            tree.cache_finished_req(
+                req, is_insert=True, kv_len_to_handle=req._kv_committed_len
+            )
             tree.sanity_check()
 
     # -- Integration: page_size=1 full flow --
@@ -391,7 +404,9 @@ class TestSWAEvictionBoundary(unittest.TestCase):
                 req.kv.swa_evicted_seqlen, max(0, seq_len - 1 - max(window, page_size))
             )
 
-            tree.cache_finished_req(req, is_insert=True)
+            tree.cache_finished_req(
+                req, is_insert=True, kv_len_to_handle=req._kv_committed_len
+            )
             tree.sanity_check()
 
 

test/registered/unit/mem_cache/test_swa_unittest.py

@@ -38,9 +38,6 @@ class _DummyReq:
         self.swa_prefix_lock_released = False
         self.kv = SimpleNamespace(swa_evicted_seqlen=0)
 
-    def pop_committed_kv_cache(self):
-        return self._kv_committed_len
-
 
 def _build_swa_tree(
     is_eagle: bool,
@@ -592,7 +589,9 @@ class TestSWA(unittest.TestCase):
             return original_insert(params)
 
         tree.insert = wrapped_insert
-        tree.cache_finished_req(req, is_insert=True)
+        tree.cache_finished_req(
+            req, is_insert=True, kv_len_to_handle=req._kv_committed_len
+        )
 
         self.assertEqual(captured["prev_prefix_len"], req.cache_protected_len)
         self.assertTrue(captured["is_bigram"])
@@ -624,7 +623,9 @@ class TestSWA(unittest.TestCase):
             return original_free(indices)
 
         allocator.free = wrapped_free
-        tree.cache_finished_req(req2, is_insert=False)
+        tree.cache_finished_req(
+            req2, is_insert=False, kv_len_to_handle=req2._kv_committed_len
+        )
 
         # EAGLE + page_size=1 => page_aligned_len = committed_len - 1 = 5
         # Expected frees:

test/registered/unit/mem_cache/test_unified_radix_cache_bench.py

@@ -643,7 +643,6 @@ def bench_cache_finished(
         req.last_node = node
         req.cache_protected_len = matched_len
         req.kv_committed_len = len(seq)
-        req.kv_committed_freed = False
         if hasattr(lr, "swa_uuid_for_lock"):
             req.swa_uuid_for_lock = lr.swa_uuid_for_lock
         env.rtp.req_to_token[req.req_pool_idx, : len(kv_indices)] = kv_indices
@@ -656,7 +655,9 @@ def bench_cache_finished(
     return bench_api(
         "cache_finished",
         lambda: req_items,
-        lambda req: env.tree.cache_finished_req(req, is_insert=True),
+        lambda req: env.tree.cache_finished_req(
+            req, is_insert=True, kv_len_to_handle=req.kv_committed_len
+        ),
         len(req_items) - warmup,
         env.avg_tokens,
         warmup,

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -854,7 +854,9 @@ class UnifiedRadixCacheSuite:
         if self.cfg.has_mamba:
             req.mamba_last_track_seqlen = kv_len
 
-        cache.cache_finished_req(req, is_insert=True)
+        cache.cache_finished_req(
+            req, is_insert=True, kv_len_to_handle=req.effective_kv_committed_len()
+        )
 
         all_ids = input_ids + output_ids
         aligned_len = (len(all_ids) // ps) * ps
@@ -893,8 +895,10 @@ class UnifiedRadixCacheSuite:
         get_server_args().strip_thinking_cache = True
         try:
             avail_before = allocator.available_size()
-            cache.cache_finished_req(req, is_insert=True)
+            cache.cache_finished_req(
-            start_p, end_p = req.pop_overallocated_kv_cache()
+                req, is_insert=True, kv_len_to_handle=req.effective_kv_committed_len()
+            )
+            start_p, end_p = req.effective_kv_committed_len(), req.kv.kv_allocated_len
         finally:
             get_server_args().strip_thinking_cache = False
         if ps > 1:
@@ -936,7 +940,9 @@ class UnifiedRadixCacheSuite:
         )
 
         avail_before = allocator.available_size()
-        cache.cache_finished_req(req, is_insert=False)
+        cache.cache_finished_req(
+            req, is_insert=False, kv_len_to_handle=req.effective_kv_committed_len()
+        )
 
         self.assertEqual(allocator.available_size(), avail_before + kv_len)
         m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
@@ -1108,7 +1114,9 @@ class UnifiedRadixCacheSuite:
             req.mamba_last_track_seqlen = kv_len
 
         avail_before = allocator.available_size()
-        cache.cache_finished_req(req, is_insert=True)
+        cache.cache_finished_req(
+            req, is_insert=True, kv_len_to_handle=req.effective_kv_committed_len()
+        )
 
         self.assertEqual(allocator.available_size(), avail_before + tail_extra)
         aligned = input_ids[: (len(input_ids) // ps) * ps]