[mem_cache] Release up to owned_kv_len on radix cache insert (#40075)

python/sglang/srt/managers/schedule_batch.py

@@ -1445,7 +1445,7 @@ class Req(ReqDllmMixin):
         kv, self.kv = self.kv, ReqKvInfo()
         return kv
 
-    def effective_kv_committed_len(self) -> int:
+    def owned_kv_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_serving().strip_thinking_cache and self.reasoning_tokens > 0:

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -434,8 +434,18 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
         return None
 
     @abstractmethod
-    def cache_finished_req(self, req: Req, is_insert: bool = True, **kwargs):
-        pass
+    def cache_finished_req(
+        self, req: Req, is_insert: bool = True, *, owned_kv_len: int, **kwargs
+    ):
+        """Dispose of a finished request's KV.
+
+        ``[0, req.kv.cache_protected_len)`` is cache-owned and must survive.
+        Every slot in ``[req.kv.cache_protected_len, owned_kv_len)`` is this
+        call's to account for: insert what can be keyed, release the rest.
+        Slicing the kv row by the token-id count instead strands whatever
+        lies between -- no caller releases those. ``release_kv_cache`` frees
+        everything past ``owned_kv_len``.
+        """
 
     @abstractmethod
     def cache_unfinished_req(self, req: Req, **kwargs):

python/sglang/srt/mem_cache/cache_init_params.py

@@ -30,7 +30,6 @@ class CacheInitParams:
     # Keyword arguments for the eviction policy's constructor; see the strategy
     # classes in evict_policy.py for what each policy accepts.
     eviction_policy_config: Optional[dict[str, Any]] = None
-    disable_finished_insert: bool = False
 
     enable_metrics: bool = False
     enable_kv_cache_events: bool = False

python/sglang/srt/mem_cache/chunk_cache.py

@@ -77,11 +77,11 @@ class ChunkCache(BasePrefixCache):
         return InsertResult(prefix_len=0)
 
     def cache_finished_req(
-        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+        self, req: Req, is_insert: bool = True, *, owned_kv_len: 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.
-        self.free_kv_row(req.kv, [(req.kv.cache_protected_len, kv_len_to_handle)])
+        self.free_kv_row(req.kv, [(req.kv.cache_protected_len, owned_kv_len)])
 
     def cache_unfinished_req(self, req: Req, chunked=False):
         kv_indices = self.req_to_token_pool.req_to_token[
@@ -151,11 +151,10 @@ class PureSWAChunkCache(SWAChunkCache):
     """
 
     def cache_finished_req(
-        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+        self, req: Req, is_insert: bool = True, *, owned_kv_len: int
     ):
-        kv_committed_len = kv_len_to_handle
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.kv.req_pool_idx, :kv_committed_len
+            req.kv.req_pool_idx, :owned_kv_len
         ]
         # The cache_protected_len prefix is not this req's to free.
         protected_len = req.kv.cache_protected_len
@@ -165,9 +164,9 @@ class PureSWAChunkCache(SWAChunkCache):
             parts = []
             if evict_floor > protected_len:
                 parts.append(kv_indices[protected_len:evict_floor])
-            if evicted_seqlen < kv_committed_len:
+            if evicted_seqlen < owned_kv_len:
                 parts.append(
-                    kv_indices[max(evicted_seqlen, protected_len) : kv_committed_len]
+                    kv_indices[max(evicted_seqlen, protected_len) : owned_kv_len]
                 )
             if parts:
                 self.token_to_kv_pool_allocator.free(torch.cat(parts))

python/sglang/srt/mem_cache/common.py

@@ -284,11 +284,11 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
             req.kv.mamba_pool_idx = None
         return
 
-    effective_kv_committed_len = req.effective_kv_committed_len()
+    owned_kv_len = req.owned_kv_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,
+        owned_kv_len=owned_kv_len,
     )
 
     # StreamingSession.cache_finished_req handles speculative tail trim
@@ -297,7 +297,7 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
     if not req.kv.holds_kv:
         return
 
-    start_p, end_p = effective_kv_committed_len, req.kv.kv_allocated_len
+    start_p, end_p = owned_kv_len, req.kv.kv_allocated_len
     _release_overallocated_kv_indices(req, start_p, end_p, tree_cache)
 
     # If the prefix cache doesn't manage mamba states, we must free them here.

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -544,20 +544,20 @@ class MambaRadixCache(BasePrefixCache):
         return InsertResult(prefix_len=prefix_len, mamba_exist=mamba_exist)
 
     def cache_finished_req(
-        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+        self, req: Req, is_insert: bool = True, *, owned_kv_len: int
     ) -> None:
         """Cache request when it finishes."""
         if self.disable:
             kv_indices = self.req_to_token_pool.req_to_token[
-                req.kv.req_pool_idx, :kv_len_to_handle
+                req.kv.req_pool_idx, :owned_kv_len
             ]
             self.token_to_kv_pool_allocator.free_segment(kv_indices, start_pos=0)
             self.req_to_token_pool.free_mamba_cache(req)
             return
 
-        token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
+        token_ids = (req.origin_input_ids + req.output_ids)[:owned_kv_len]
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.kv.req_pool_idx, :kv_len_to_handle
+            req.kv.req_pool_idx, :owned_kv_len
         ]
 
         if is_insert:
@@ -607,7 +607,7 @@ class MambaRadixCache(BasePrefixCache):
                 page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True)
 
             assert cache_len == page_aligned_len, (
-                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"
+                f"It is required {cache_len=}, {page_aligned_len=}, {owned_kv_len=}, {len(req.origin_input_ids)=}, {len(req.output_ids)=} ping @yizhang2077 if you see this"
             )
 
             # Radix Cache takes one ref in memory pool

python/sglang/srt/mem_cache/pure_swa_radix_cache.py

@@ -63,7 +63,7 @@ class PureSWARadixCache(RadixCache):
         return super().evict(EvictParams(num_tokens=num_tokens))
 
     def cache_finished_req(
-        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+        self, req: Req, is_insert: bool = True, *, owned_kv_len: int
     ):
         """Cache request when it finishes.
 
@@ -72,20 +72,16 @@ class PureSWARadixCache(RadixCache):
         to the allocator. The range [evict_floor, swa_evicted_seqlen) was already
         freed by _evict_swa during decode — we skip it to avoid double-free.
         """
-        if self.disable_finished_insert:
-            is_insert = False
-
-        kv_committed_len = kv_len_to_handle
         if self.disable:
             kv_indices = self.req_to_token_pool.req_to_token[
-                req.kv.req_pool_idx, :kv_committed_len
+                req.kv.req_pool_idx, :owned_kv_len
             ]
             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)[:owned_kv_len]
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.kv.req_pool_idx, :kv_committed_len
+            req.kv.req_pool_idx, :owned_kv_len
         ]
 
         radix_key = RadixKey(

python/sglang/srt/mem_cache/radix_cache.py

@@ -325,7 +325,6 @@ class RadixCache(BasePrefixCache):
         self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
         self.page_size = params.page_size
         self.is_eagle = params.is_eagle
-        self.disable_finished_insert = params.disable_finished_insert
         self.eviction_policy = params.eviction_policy.lower()
 
         self.kv_events = KVCacheEventRecorder(
@@ -477,17 +476,13 @@ class RadixCache(BasePrefixCache):
         return InsertResult(prefix_len=prefix_len, last_device_node=last_node)
 
     def cache_finished_req(
-        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+        self, req: Req, is_insert: bool = True, *, owned_kv_len: int
     ):
         """Cache request when it finishes."""
-        # In deterministic mode, disable finished request insertion to radix cache
-        if self.disable_finished_insert:
-            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.kv.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle
+                req.kv.req_pool_idx, req.kv.cache_protected_len : owned_kv_len
             ]
             self.token_to_kv_pool_allocator.free_segment(
                 kv_indices, start_pos=req.kv.cache_protected_len
@@ -498,7 +493,7 @@ class RadixCache(BasePrefixCache):
             # 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
+                req.kv.req_pool_idx, req.kv.cache_protected_len : owned_kv_len
             ]
             self.token_to_kv_pool_allocator.free_segment(
                 kv_indices, start_pos=req.kv.cache_protected_len
@@ -507,9 +502,9 @@ class RadixCache(BasePrefixCache):
                 self.dec_lock_ref(req.last_node)
             return
 
-        token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
+        token_ids = (req.origin_input_ids + req.output_ids)[:owned_kv_len]
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.kv.req_pool_idx, : len(token_ids)
+            req.kv.req_pool_idx, :owned_kv_len
         ]
 
         radix_key = RadixKey(

python/sglang/srt/mem_cache/radix_cache_cpp.py

@@ -181,20 +181,20 @@ class RadixCacheCpp(BasePrefixCache):
         return self.tree.total_size()
 
     def cache_finished_req(
-        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+        self, req: Req, is_insert: bool = True, *, owned_kv_len: int
     ):
         """Cache request when it finishes."""
         self._reject_cache_salt(req.cache_salt)
         assert req.kv.holds_kv
-        token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
+        token_ids = (req.origin_input_ids + req.output_ids)[:owned_kv_len]
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.kv.req_pool_idx, :kv_len_to_handle
+            req.kv.req_pool_idx, :owned_kv_len
         ].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_len_to_handle // self.page_size * self.page_size
+        page_aligned_overall_len = owned_kv_len // self.page_size * self.page_size
 
         if is_insert:
             new_prefix_len = self._insert(
@@ -213,7 +213,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_len_to_handle:
+        if page_aligned_overall_len < owned_kv_len:
             # 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

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

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

@@ -439,13 +439,11 @@ class LMCRadixCache(RadixCache):
             )
 
     def cache_finished_req(
-        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+        self, req: Req, is_insert: bool = True, *, owned_kv_len: int
     ) -> None:
         """On request completion, insert device KV into radix and store to LMCache."""
 
-        super().cache_finished_req(
-            req, is_insert=is_insert, kv_len_to_handle=kv_len_to_handle
-        )
+        super().cache_finished_req(req, is_insert=is_insert, owned_kv_len=owned_kv_len)
         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

@@ -460,16 +460,16 @@ class SWARadixCache(BasePrefixCache):
         return InsertResult(prefix_len=prefix_len)
 
     def cache_finished_req(
-        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int
+        self, req: Req, is_insert: bool = True, *, owned_kv_len: int
     ) -> None:
         """Cache request when it finishes."""
         if self.disable:
-            self.free_kv_row(req.kv, [(0, kv_len_to_handle)])
+            self.free_kv_row(req.kv, [(0, owned_kv_len)])
             return
 
-        token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
+        token_ids = (req.origin_input_ids + req.output_ids)[:owned_kv_len]
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.kv.req_pool_idx, :kv_len_to_handle
+            req.kv.req_pool_idx, :owned_kv_len
         ]
 
         radix_key = RadixKey(
@@ -497,7 +497,7 @@ class SWARadixCache(BasePrefixCache):
             self.free_kv_row(req.kv, [(old_prefix_len, page_aligned_len)])
 
         # free the unaligned tail
-        self.free_kv_row(req.kv, [(page_aligned_len, kv_len_to_handle)])
+        self.free_kv_row(req.kv, [(page_aligned_len, owned_kv_len)])
 
         # Remove req slot release the cache lock
         self.dec_lock_ref(

python/sglang/srt/mem_cache/unified_cache/components/README.md

@@ -228,14 +228,14 @@ receipt proves were taken. The eventual full release must pass
 
 ---
 
-### `cache_finished_req(req: Req, is_insert: bool = True, *, kv_len_to_handle: int)`
+### `cache_finished_req(req: Req, is_insert: bool = True, *, owned_kv_len: 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); `kv_len_to_handle` — committed KV length supplied by the caller |
+| **Inputs** | `req` — the finished request; `is_insert` — insert the owned range into the tree (True) or free it (False); `owned_kv_len` — end of the request-owned KV range; slots past it are freed by `release_kv_cache` |
 | **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

@@ -955,22 +955,22 @@ class UnifiedRadixCache(BasePrefixCache):
             return DecLockRefResult()
         return self.tree_core.dec_host_lock_ref(node_id, params)
 
-    @rank_consensus(same_params=["req.rid", "is_insert", "kv_len_to_handle"])
+    @rank_consensus(same_params=["req.rid", "is_insert", "owned_kv_len"])
     def cache_finished_req(
-        self, req: Req, is_insert: bool = True, *, kv_len_to_handle: int, **kwargs
+        self, req: Req, is_insert: bool = True, *, owned_kv_len: int, **kwargs
     ) -> None:
         if self.session.try_cache_finished_req(req, is_insert=is_insert, **kwargs):
             return
 
         if self.disable:
-            self.free_kv_row(req.kv, [(0, kv_len_to_handle)])
+            self.free_kv_row(req.kv, [(0, owned_kv_len)])
             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_len_to_handle]
+        token_ids = (req.origin_input_ids + req.output_ids)[:owned_kv_len]
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.kv.req_pool_idx, :kv_len_to_handle
+            req.kv.req_pool_idx, :owned_kv_len
         ]
 
         result = None
@@ -1078,7 +1078,7 @@ class UnifiedRadixCache(BasePrefixCache):
                     ranges.append((tail_free_start, len(kv_indices_full)))
             self.free_kv_row(req.kv, ranges)
         else:
-            self.free_kv_row(req.kv, [(req.kv.cache_protected_len, kv_len_to_handle)])
+            self.free_kv_row(req.kv, [(req.kv.cache_protected_len, owned_kv_len)])
 
         # Synthetic profiling requests may own KV without locking a tree node.
         if req.last_node is not None: