Keep CP HiCache valid tails page-owned during extension

CP shared KV HiCache now treats non-page-aligned valid-tail nodes as page-owned when a later request extends beyond them. The prefix probe, match path, insert path, and prepared backup start now agree on flooring the reusable prefix to the previous physical page boundary, so prepared backup metadata cannot start mid-page or fail to attach at insertion.

Duplicate frees under CP HiCache now go through a page-safe free helper. Insert and unfinished duplicate ranges free only fully unprotected pages; no-insert completion still releases the right tail owned only by the finishing request.

Constraint: CP HiCache allocator frees whole physical pages even when called with token-granular locs.

Rejected: Partial-page sharing/refcounting | too complex for the current page-as-minimum-unit contract.

Rejected: Fix only prepare_write_backup_for_req | match_prefix and insert would still expose exact valid-tail hits and desynchronize prepared backup length.

Confidence: medium

Scope-risk: moderate

Directive: Do not expose non-page-aligned CP valid-tail hits to extending requests unless partial-page ownership is explicitly implemented end-to-end.

Tested: remote g0034 py_compile for touched files

Tested: remote g0034 test_cp_hicache_metadata.py 97 passed

Tested: remote g0034 test_cp_shared_kv_runtime.py 73 passed

Not-tested: test_cp_shared_kv_layout.py aborts during installed sgl_kernel architecture-specific op loading before assertions

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -48,6 +48,7 @@ from sglang.srt.mem_cache.radix_cache import (
     TreeNode,
     ceil_to_page_len,
     compute_node_hash_values,
+    floor_to_page_len,
     split_node_hash_value,
 )
 from sglang.srt.mem_cache.utils import convert_to_bigram_key
@@ -2008,6 +2009,9 @@ class HiRadixCache(RadixCache):
         while len(key) > 0 and child_key in node.children:
             child = node.children[child_key]
             prefix_len = self.key_match_fn(child.key, key)
+            prefix_len = self._cp_floor_exact_valid_tail_extension_len(
+                child, prefix_len, len(key)
+            )
             if prefix_len <= 0:
                 break
             if prefix_len < len(child.key):
@@ -2040,6 +2044,20 @@ class HiRadixCache(RadixCache):
             return prefix_len
         return prefix_len // self.page_size * self.page_size
 
+    def _cp_floor_exact_valid_tail_extension_len(
+        self, child: TreeNode, prefix_len: int, request_len: int
+    ) -> int:
+        if (
+            not self._uses_cp_hicache
+            or self.page_size <= 1
+            or prefix_len <= 0
+            or prefix_len != len(child.key)
+            or prefix_len >= request_len
+            or prefix_len % self.page_size == 0
+        ):
+            return prefix_len
+        return floor_to_page_len(prefix_len, self.page_size)
+
     def prepare_write_backup_for_req(self, req) -> None:
         if self.disable or not self._uses_cp_hicache:
             return
@@ -2077,9 +2095,10 @@ class HiRadixCache(RadixCache):
                 is_bigram=self.is_eagle,
             )
         )
-        start = min(max(req.cache_protected_len, existing_prefix_len), len(keys))
-        if len(keys) <= start:
+        raw_start = min(max(req.cache_protected_len, existing_prefix_len), len(keys))
+        if len(keys) <= raw_start:
             return
+        start = floor_to_page_len(raw_start, self.page_size)
 
         kv_indices = self.req_to_token_pool.req_to_token[
             req.req_pool_idx, start : len(keys)
@@ -3428,6 +3447,11 @@ class HiRadixCache(RadixCache):
             if self._is_pinned(child):
                 child.pin_expiry = time.monotonic() + child.pin_ttl
             prefix_len = self.key_match_fn(child.key, key)
+            prefix_len = self._cp_floor_exact_valid_tail_extension_len(
+                child, prefix_len, len(key)
+            )
+            if prefix_len <= 0:
+                break
             if prefix_len < len(child.key):
                 if (
                     self._uses_cp_hicache
@@ -3528,10 +3552,19 @@ class HiRadixCache(RadixCache):
         total_prefix_length = 0
 
         while len(key) > 0 and child_key in node.children.keys():
+            parent_node = node
             node = node.children[child_key]
             node.last_access_time = time.monotonic()
             node.priority = max(node.priority, priority)
-            prefix_len = self.key_match_fn(node.key, key)
+            raw_prefix_len = self.key_match_fn(node.key, key)
+            prefix_len = self._cp_floor_exact_valid_tail_extension_len(
+                node, raw_prefix_len, len(key)
+            )
+            stop_after_page_floor = prefix_len != raw_prefix_len
+            if prefix_len <= 0:
+                if stop_after_page_floor:
+                    node = parent_node
+                break
 
             if prefix_len == len(node.key):
                 if node.evicted:
@@ -3575,6 +3608,8 @@ class HiRadixCache(RadixCache):
 
             if len(key):
                 child_key = self.get_child_key_fn(key)
+            if stop_after_page_floor:
+                break
 
         if len(key):
             use_prepared_cp_backup = (

python/sglang/srt/mem_cache/radix_cache.py

@@ -123,6 +123,12 @@ def ceil_to_page_len(length: int, page_size: int) -> int:
     return ((length + page_size - 1) // page_size) * page_size
 
 
+def floor_to_page_len(length: int, page_size: int) -> int:
+    if page_size <= 1:
+        return length
+    return length // page_size * page_size
+
+
 class TreeNode:
 
     counter = 0
@@ -472,6 +478,22 @@ class RadixCache(BasePrefixCache):
         prefix_len = self._insert_helper(self.root_node, key, value, priority, chunked)
         return InsertResult(prefix_len=prefix_len)
 
+    def _free_kv_indices_range(
+        self,
+        kv_indices: torch.Tensor,
+        start: int,
+        end: int,
+        *,
+        include_partial_tail: bool = False,
+    ):
+        if getattr(self, "_uses_cp_hicache", False):
+            start = ceil_to_page_len(start, self.page_size)
+            if not include_partial_tail:
+                end = floor_to_page_len(end, self.page_size)
+        if end < start:
+            end = start
+        self.token_to_kv_pool_allocator.free(kv_indices[start:end])
+
     def cache_finished_req(self, req: Req, is_insert: bool = True):
         """Cache request when it finishes."""
         # In deterministic mode, disable finished request insertion to radix cache
@@ -518,16 +540,19 @@ class RadixCache(BasePrefixCache):
                 self._rollback_prepared_cp_backup(prepared_cp_backup, "insert_miss")
             new_prefix_len = result.prefix_len
             # Free the duplicates that were already in the tree
-            self.token_to_kv_pool_allocator.free(
-                kv_indices[req.cache_protected_len : new_prefix_len]
+            self._free_kv_indices_range(
+                kv_indices, req.cache_protected_len, new_prefix_len
             )
         else:
             if prepared_cp_backup is not None and hasattr(
                 self, "_rollback_prepared_cp_backup"
             ):
                 self._rollback_prepared_cp_backup(prepared_cp_backup, "no_insert")
-            self.token_to_kv_pool_allocator.free(
-                kv_indices[req.cache_protected_len : len(keys)]
+            self._free_kv_indices_range(
+                kv_indices,
+                req.cache_protected_len,
+                len(keys),
+                include_partial_tail=True,
             )
         if prepared_cp_backup is not None:
             req.cp_hicache_prepared_backup = None
@@ -587,9 +612,7 @@ class RadixCache(BasePrefixCache):
             req.cp_hicache_prepared_backup = None
         new_prefix_len = result.prefix_len
 
-        self.token_to_kv_pool_allocator.free(
-            kv_indices[req.cache_protected_len : new_prefix_len]
-        )
+        self._free_kv_indices_range(kv_indices, req.cache_protected_len, new_prefix_len)
 
         # The prefix indices could be updated, reuse it
         match_result = self.match_prefix(MatchPrefixParams(key=radix_key))