Cache completed chunked prefill pages in CP HiCache

Chunked prefill was skipping prepared CP HiCache backup entirely while still inserting unfinished valid-tail radix state. That combination caused middle chunks to miss reusable host cache and repeatedly hit stale-tail split fallback paths.\n\nThis keeps CP HiCache page-granular: middle chunks now reserve/write only completed pages, and unfinished radix inserts expose only page-complete prefixes while preserving the request-owned sub-page tail in prefix_indices for later chunks.\n\nConstraint: CP HiCache radix and host residency are managed at page granularity, but chunked prefill can end on a sub-page tail.\nRejected: Continue skipping chunked backup | loses middle-chunk reuse and triggers fallback storms.\nRejected: Insert sub-page chunk tails into CP HiCache | creates stale-tail prune/split conflicts against request-owned KV.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not reintroduce chunked_req backup skip; keep CP HiCache chunked inserts page-complete unless the radix/host cache contract becomes sub-page aware.\nTested: python -m py_compile python/sglang/srt/mem_cache/radix_cache.py python/sglang/srt/mem_cache/hiradix_cache.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py\nTested: g0034 targeted chunked CP HiCache regression tests, 2 passed\nTested: g0034 PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py, 107 passed\nNot-tested: Fresh chunked-prefill ETE traffic run after process restart

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -2325,14 +2325,6 @@ class HiRadixCache(RadixCache):
                 rid=getattr(req, "rid", "<unknown>"),
             )
             return
-        if getattr(req, "is_chunked", 0) > 0:
-            self._warn_cp_hicache_fallback(
-                "prepare_write_backup_skipped",
-                "chunked_req",
-                rid=getattr(req, "rid", "<unknown>"),
-                is_chunked=getattr(req, "is_chunked", 0),
-            )
-            return
         if getattr(req, "cp_hicache_prepared_backup", None) is not None:
             prepared = getattr(req, "cp_hicache_prepared_backup", None)
             self._warn_cp_hicache_fallback(
@@ -2345,20 +2337,40 @@ class HiRadixCache(RadixCache):
 
         token_ids = req.fill_ids
         keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids
+        backup_end = len(keys)
+        if getattr(req, "is_chunked", 0) > 0 and self.page_size > 1:
+            # Chunked prefill should still populate CP HiCache, but only for
+            # pages whose KV has become fully available.  Leaving the current
+            # sub-page tail out of the prepared write avoids creating a
+            # scheduler-visible stale tail that the next chunk has to prune
+            # while the same request still owns/locks it.
+            backup_end = floor_to_page_len(backup_end, self.page_size)
+            if backup_end <= 0:
+                logger.debug(
+                    "[HiCache-write] skip chunked CP backup with no completed page: rid=%s key_len=%d",
+                    getattr(req, "rid", "<unknown>"),
+                    len(keys),
+                )
+                return
+            keys_for_probe = keys[:backup_end]
+        else:
+            keys_for_probe = keys
         existing_prefix_len = self._probe_existing_radix_prefix_len_no_split(
             RadixKey(
-                keys,
+                keys_for_probe,
                 getattr(req, "extra_key", None),
                 is_bigram=self.is_eagle,
             )
         )
-        raw_start = min(max(req.cache_protected_len, existing_prefix_len), len(keys))
-        if len(keys) <= raw_start:
+        raw_start = min(max(req.cache_protected_len, existing_prefix_len), backup_end)
+        if backup_end <= raw_start:
             return
         start = floor_to_page_len(raw_start, self.page_size)
+        if backup_end <= start:
+            return
 
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, start : len(keys)
+            req.req_pool_idx, start:backup_end
         ].to(dtype=torch.int64, copy=True)
         if len(kv_indices) == 0:
             self._warn_cp_hicache_fallback(
@@ -2366,7 +2378,7 @@ class HiRadixCache(RadixCache):
                 "empty_kv_indices",
                 rid=getattr(req, "rid", "<unknown>"),
                 start=start,
-                key_len=len(keys),
+                key_len=backup_end,
             )
             return
 

python/sglang/srt/mem_cache/radix_cache.py

@@ -635,6 +635,18 @@ class RadixCache(BasePrefixCache):
         keys = convert_to_bigram_key(token_ids) if self.is_eagle else token_ids
         if not getattr(self, "_uses_cp_hicache", False):
             keys = page_align_keys(keys, self.page_size)
+        elif chunked and self.page_size > 1:
+            # CP HiCache owns KV at page granularity.  For an unfinished
+            # chunked request, cache only completed pages and keep the current
+            # sub-page tail in the request-owned KV slots until a later chunk
+            # or the final insert can make it page-complete.  This preserves
+            # intermediate chunk caching without repeatedly creating/pruning
+            # stale valid-tail radix nodes.
+            completed_len = floor_to_page_len(len(keys), self.page_size)
+            if completed_len <= req.cache_protected_len:
+                req.prefix_indices = kv_indices
+                return
+            keys = keys[:completed_len]
         values = kv_indices[: len(keys)].to(dtype=torch.int64, copy=True)
         radix_key = RadixKey(keys, req.extra_key, is_bigram=self.is_eagle)
         prepared_cp_backup = getattr(req, "cp_hicache_prepared_backup", None)