Fix CP HiCache catch_up_all_layers fallback on chunked-prefill final chunk

Under overlap scheduling a chunked request's final-chunk write-backup prepare
read a stale `is_chunked` (>0): that per-tick counter is decremented in
process_batch_result, which the overlap loop runs AFTER the run_batch prepare
hook. So prepare floored `backup_end` to a page boundary (the intermediate-chunk
rule) and dropped the now-complete final tail page, while the final non-chunked
insert builds the radix node at full unaligned length. The exact-equality attach
predicate (prepared.logical_len == len(value)) then failed by
(num_tokens-1) mod page_size, dropped the per-layer overlap backup, and forced
the serial all-layer catch-up (~89% of large chunked requests in prod).

Decouple the floor decision from the stale counter: the scheduler marks
`req.cp_backup_is_intermediate_chunk = (req is self.chunked_req)` in
`_prepare_hicache_write_backups_before_forward` (the live chunked_req identity
is the authoritative "will be chunked further" signal, set this tick before
run_batch), and the prepare candidate builder floors on that flag instead of
`is_chunked`. Intermediate chunks still floor; only the misclassified final
chunk now backs up its full tail and attaches the overlap backup.

Tests: update the chunked prepare test to the new flag; add a regression test
that a final chunk with stale is_chunked reserves the full tail; add an offline
repro that drives the real prepare/insert/probe and sweeps unaligned tails.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

python/sglang/srt/managers/scheduler.py

@@ -2990,6 +2990,23 @@ class Scheduler(
         ):
             return
 
+        # Tell the CP HiCache backup prepare which req (if any) is a genuine
+        # INTERMEDIATE chunk, i.e. will be prefilled further next round, so it
+        # floors off its still-incomplete tail page; every other req in this
+        # batch completes here and must back up its full (now-complete) tail.
+        #
+        # We mark it from the live `chunked_req` identity rather than letting
+        # prepare read `req.is_chunked`, because `is_chunked` is a per-tick
+        # counter whose decrement lags one iteration under overlap scheduling
+        # (it is decremented in process_batch_result, which runs AFTER this
+        # run_batch). On a request's final chunk it therefore reads stale (>0),
+        # making prepare floor off the now-complete tail page -> prepared
+        # logical_len != final radix node len -> the overlap backup is dropped
+        # and the node falls back to the serial all-layer catch-up.
+        carried_chunked_req = self.chunked_req
+        for req in batch.reqs:
+            req.cp_backup_is_intermediate_chunk = req is carried_chunked_req
+
         prepare_batch_fn = getattr(
             self.tree_cache, "prepare_write_backups_for_reqs", None
         )

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -2581,12 +2581,22 @@ 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.
+        if getattr(req, "cp_backup_is_intermediate_chunk", False) and self.page_size > 1:
+            # Only a genuine INTERMEDIATE chunk floors off its tail page: a later
+            # chunk of the same request will extend it, so backing up that still
+            # incomplete page now would create a scheduler-visible stale tail the
+            # next chunk has to prune while this request still owns/locks it. The
+            # FINAL chunk's tail page is complete and must be backed up in full,
+            # to match the radix node that the final (non-chunked) insertion
+            # builds -- otherwise prepared logical_len != len(value) and the
+            # exact-equality attach predicate drops the overlap backup.
+            #
+            # NOTE: we must NOT key this off `req.is_chunked`. It is a per-tick
+            # counter whose decrement lags one iteration under overlap scheduling
+            # and so reads stale (>0) on the final chunk. The scheduler sets
+            # `cp_backup_is_intermediate_chunk` from the live `chunked_req`
+            # identity in `_prepare_hicache_write_backups_before_forward`, which
+            # runs in the same run_batch immediately before this.
             backup_end = floor_to_page_len(backup_end, self.page_size)
             if backup_end <= 0:
                 logger.debug(

test/registered/unit/mem_cache/test_cp_hicache_chunked_prepare_repro.py

@@ -0,0 +1,263 @@
+"""Offline repro for the CP HiCache `catch_up_all_layers` fallback.
+
+Isolates the *length arithmetic* that drives the fallback: the prepared
+write-backup's ``logical_len`` (computed at prepare, before the forward) vs the
+final radix node's ``len(value)`` (computed at insert, after the forward). The
+production attach predicate is exact equality
+``prepared.logical_len == len(value)`` (hiradix_cache.py); a mismatch drops the
+prepared overlap backup and falls back to the serial all-layer catch-up.
+
+We drive the REAL fix-output methods (``insert``,
+``_build_prepare_write_backup_candidate_for_req``,
+``_probe_existing_radix_prefix_len_no_split``) with a minimal ``__new__``
+fixture, scaled to page_size=4, and replay a 2-chunk request under two
+orderings:
+
+  * Order A  (insert chunk1  -> prepare chunk2 -> insert chunk2): the prepare
+    probe sees chunk1, so the prefix matched at prepare == prefix matched at
+    insert. Predicate should hold -> overlap backup attaches.
+  * Order B  (prepare chunk2 -> insert chunk1 -> insert chunk2): the prepare
+    probe runs against a tree that does NOT yet contain chunk1, so the prefix
+    grows between prepare and insert. Predicate should FAIL -> fallback.
+
+Run in the dev-cu13 container (needs torch + sgl_kernel; the model-stack CUDA
+coupling means this will not import on a CPU-only box).
+"""
+
+import sys
+import types
+import unittest
+from functools import partial
+from types import SimpleNamespace
+
+import torch
+
+# Prefer the real sgl_kernel; fall back to a stub on CPU-only hosts.
+try:
+    import sgl_kernel  # noqa: F401
+except (ImportError, RuntimeError):
+    if "sgl_kernel" not in sys.modules:
+        stub = types.ModuleType("sgl_kernel")
+        stub.__path__ = []
+
+        def _getattr(name):
+            if name.startswith("__"):
+                raise AttributeError(name)
+            fn = lambda *a, **k: None
+            setattr(stub, name, fn)
+            return fn
+
+        stub.__getattr__ = _getattr
+        sys.modules["sgl_kernel"] = stub
+
+from sglang.srt.mem_cache.base_prefix_cache import InsertParams  # noqa: E402
+from sglang.srt.mem_cache.hiradix_cache import HiRadixCache  # noqa: E402
+from sglang.srt.mem_cache.radix_cache import (  # noqa: E402
+    RadixKey,
+    TreeNode,
+    _key_match_page_size1,
+    _key_match_paged,
+    convert_to_bigram_key,
+    floor_to_page_len,
+    get_child_key,
+)
+
+PAGE = 4
+MAXLEN = 8192
+
+
+def make_cache(page_size=PAGE, is_eagle=False):
+    c = HiRadixCache.__new__(HiRadixCache)
+    c.page_size = page_size
+    c.is_eagle = is_eagle
+    c._uses_cp_hicache = True
+    c.disable = False
+    c.enable_storage = False
+    c.enable_kv_cache_events = False
+    c.evictable_size_ = 0
+    root = TreeNode()
+    root.children = {}
+    c.root_node = root
+    if page_size == 1:
+        c.key_match_fn = _key_match_page_size1
+        c.get_child_key_fn = get_child_key
+    else:
+        c.key_match_fn = partial(_key_match_paged, page_size=page_size)
+        c.get_child_key_fn = partial(get_child_key, page_size=page_size)
+    c.ongoing_write_through = {}
+    c.pending_host_backups = {}
+    c.cache_controller = SimpleNamespace(write_policy="write_through")
+    # Stubs for the non-arithmetic bits insert() touches.
+    c._node_device_resident_len = lambda node: (
+        len(node.value) if node.value is not None else 0
+    )
+    c._update_leaf_status = lambda node: None
+    c._update_host_leaf_status = lambda node: None
+    c._record_store_event = lambda node: None
+    c._inc_hit_count = lambda node, chunked=False: None
+    # req_to_token pool: row 0 maps token position -> a distinct KV slot index.
+    req_to_token = torch.arange(MAXLEN, dtype=torch.int64).unsqueeze(0)
+    c.req_to_token_pool = SimpleNamespace(req_to_token=req_to_token)
+    return c
+
+
+def _bigram_len(n_tokens, is_eagle):
+    if not is_eagle:
+        return n_tokens
+    return max(0, n_tokens - 1)
+
+
+def insert_chunk(cache, tokens, *, chunked, floor_for_chunked):
+    """Insert a chunk the way cache_unfinished_req does for the CP path.
+
+    For chunked intermediate inserts the CP path floors the key length to a
+    completed page boundary (``completed_len = floor_to_page_len(len(keys))``)
+    and only inserts the completed pages; the final (non-chunked) insert keeps
+    the full unaligned key length.
+    """
+    key = RadixKey(list(tokens), None, is_bigram=False)
+    key, _ = cache.maybe_bigram_convert(key)  # -> bigram if eagle
+    keys = list(key.token_ids)
+    if chunked and floor_for_chunked and cache.page_size > 1:
+        completed_len = floor_to_page_len(len(keys), cache.page_size)
+        keys = keys[:completed_len]
+    value = torch.arange(len(keys), dtype=torch.int64)
+    radix_key = RadixKey(keys, None, is_bigram=cache.is_eagle)
+    res = cache.insert(
+        InsertParams(key=radix_key, value=value, chunked=chunked)
+    )
+    return res, len(keys)
+
+
+def make_req(tokens, *, intermediate_chunk):
+    return SimpleNamespace(
+        fill_ids=list(tokens),
+        req_pool_idx=0,
+        cache_protected_len=0,
+        extra_key=None,
+        # is_chunked is deliberately set stale (>0) to mimic the overlap lag on
+        # a request's final chunk; the FIXED prepare must ignore it and key off
+        # cp_backup_is_intermediate_chunk instead.
+        is_chunked=1,
+        cp_backup_is_intermediate_chunk=intermediate_chunk,
+        rid="repro",
+    )
+
+
+def prepare_logical_len(cache, req):
+    """Return the prepared-backup logical_len (= len(candidate.kv_indices)),
+    or None if no candidate is produced."""
+    cand = cache._build_prepare_write_backup_candidate_for_req(req)
+    if cand is None:
+        return None
+    return int(len(cand.kv_indices))
+
+
+def final_insert_value_len(cache, tokens):
+    """Run the final (non-chunked) insert and return the new node's value len."""
+    key = RadixKey(list(tokens), None, is_bigram=False)
+    bigram_key, _ = cache.maybe_bigram_convert(
+        RadixKey(list(tokens), None, is_bigram=False)
+    )
+    full_keys_len = len(bigram_key.token_ids)
+    value = torch.arange(full_keys_len, dtype=torch.int64)
+    radix_key = RadixKey(list(key.token_ids), None, is_bigram=False)
+    res = cache.insert(
+        InsertParams(key=radix_key, value=value, chunked=False)
+    )
+    # new node value len = total key len - matched prefix len
+    return full_keys_len - int(res.prefix_len), int(res.prefix_len), full_keys_len
+
+
+def run_scenario(is_eagle, chunk1_tokens, total_tokens, intermediate_chunk):
+    """Production ordering (overlap): the prior chunk is already inserted when
+    the final chunk's prepare runs (process_prefill_chunk precedes run_batch),
+    so the prepare probe and the final insert match the SAME prefix. The only
+    question is whether prepare floors the tail page -- controlled by
+    intermediate_chunk (cp_backup_is_intermediate_chunk).
+
+    Returns dict with prepared logical_len, final node value len, predicate."""
+    cache = make_cache(is_eagle=is_eagle)
+    TreeNode.counter = 0
+    chunk1 = list(range(chunk1_tokens))
+    full = list(range(total_tokens))
+    req = make_req(full, intermediate_chunk=intermediate_chunk)
+
+    # insert prior chunk (chunked) -> prepare final chunk -> final insert
+    insert_chunk(cache, chunk1, chunked=True, floor_for_chunked=True)
+    logical_len = prepare_logical_len(cache, req)
+    val_len, matched, full_keys = final_insert_value_len(cache, full)
+
+    predicate = logical_len is not None and logical_len == val_len
+    return {
+        "is_eagle": is_eagle,
+        "chunk1_tokens": chunk1_tokens,
+        "total_tokens": total_tokens,
+        "intermediate_chunk_flag": intermediate_chunk,
+        "prepared_logical_len": logical_len,
+        "final_node_value_len": val_len,
+        "matched_prefix_at_insert": matched,
+        "full_bigram_keys": full_keys,
+        "predicate_holds(attach)": predicate,
+    }
+
+
+class TestChunkedPrepareRepro(unittest.TestCase):
+    def _report(self, label, r):
+        print(
+            f"\n[{label}] eagle={r['is_eagle']} "
+            f"chunk1={r['chunk1_tokens']} total={r['total_tokens']} "
+            f"intermediate_flag={r['intermediate_chunk_flag']}\n"
+            f"    prepared_logical_len = {r['prepared_logical_len']}\n"
+            f"    final_node_value_len = {r['final_node_value_len']} "
+            f"(full_bigram_keys={r['full_bigram_keys']} - matched={r['matched_prefix_at_insert']})\n"
+            f"    PREDICATE attach = {r['predicate_holds(attach)']}  "
+            f"-> {'ATTACH (overlap ok)' if r['predicate_holds(attach)'] else 'FALLBACK (catch_up_all_layers)'}"
+        )
+
+    def test_fix_final_chunk_attaches(self):
+        # THE FIX: a request's final chunk is flagged NOT intermediate
+        # (cp_backup_is_intermediate_chunk=False) even though its is_chunked is
+        # stale (>0) under overlap. Prepare must reserve the FULL tail -> matches
+        # the final radix node -> ATTACH (no catch_up_all_layers).
+        for is_eagle in (False, True):
+            r = run_scenario(
+                is_eagle, chunk1_tokens=16, total_tokens=22, intermediate_chunk=False
+            )
+            self._report("FIX final-chunk", r)
+            self.assertTrue(
+                r["predicate_holds(attach)"],
+                f"final chunk must ATTACH after fix, got: {r}",
+            )
+
+    def test_bug_misflagged_intermediate_falls_back(self):
+        # Documents the pre-fix behaviour: if the final chunk is (wrongly) treated
+        # as intermediate -- exactly what the stale is_chunked>0 did -- prepare
+        # floors the tail page and the predicate fails -> FALLBACK.
+        for is_eagle in (False, True):
+            r = run_scenario(
+                is_eagle, chunk1_tokens=16, total_tokens=22, intermediate_chunk=True
+            )
+            self._report("BUG misflagged-intermediate", r)
+            self.assertFalse(
+                r["predicate_holds(attach)"],
+                f"misflagged-intermediate should FALLBACK (documents the bug): {r}",
+            )
+
+    def test_fix_holds_for_various_tail_remainders(self):
+        # The mismatch the bug produced was (num_tokens-1) mod page_size; sweep a
+        # range of unaligned totals to ensure the fix attaches for all of them.
+        for is_eagle in (False, True):
+            for total in range(17, 32):
+                r = run_scenario(
+                    is_eagle, chunk1_tokens=16, total_tokens=total,
+                    intermediate_chunk=False,
+                )
+                self.assertTrue(
+                    r["predicate_holds(attach)"],
+                    f"final chunk must ATTACH for total={total} eagle={is_eagle}: {r}",
+                )
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=2)

test/registered/unit/mem_cache/test_cp_hicache_metadata.py

@@ -1539,6 +1539,9 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
         req = types.SimpleNamespace(
             rid="rid-chunked",
             is_chunked=1,
+            # The scheduler marks a genuine intermediate chunk via this flag
+            # (from the live chunked_req identity), not via is_chunked.
+            cp_backup_is_intermediate_chunk=True,
             fill_ids=list(range(10)),
             req_pool_idx=0,
             extra_key=None,
@@ -1555,6 +1558,58 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
         )
         self.assertEqual(req.cp_hicache_prepared_backup.logical_len, 8)
 
+    def test_prepare_write_backup_for_req_final_chunk_reserves_full_tail(self):
+        # Regression: under overlap scheduling a request's FINAL chunk can still
+        # carry a stale is_chunked>0 (its decrement lags one iteration), but it
+        # is NOT an intermediate chunk -- the scheduler leaves
+        # cp_backup_is_intermediate_chunk False. The prepared backup must then
+        # reserve the FULL (page-unaligned) tail so it matches the radix node
+        # the final non-chunked insertion builds; otherwise the exact-equality
+        # attach predicate drops the overlap backup -> serial catch_up_all_layers.
+        cache = HiRadixCache.__new__(HiRadixCache)
+        cache.disable = False
+        cache._uses_cp_hicache = True
+        cache.cache_controller = FakeReserveWriteController([])
+        cache.is_eagle = False
+        cache.page_size = 4
+        cache.root_node = TreeNode()
+        cache.root_node.key = RadixKey([])
+        cache.root_node.children = {}
+        cache.get_child_key_fn = lambda key: tuple(key.token_ids[:4])
+        cache.key_match_fn = lambda child_key, key: _key_match_paged(
+            child_key, key, page_size=4
+        )
+        cache.maybe_bigram_convert = lambda key, value=None: (key, value)
+        cache.req_to_token_pool = types.SimpleNamespace(
+            req_to_token=torch.arange(16, dtype=torch.int64).view(1, 16)
+        )
+        cache._reserve_write_cp_indices_no_collective = (
+            lambda indices, node_id: make_write_reservation(
+                indices, node_id=node_id, host_start=300
+            )
+        )
+
+        req = types.SimpleNamespace(
+            rid="rid-final",
+            is_chunked=1,  # stale under overlap -- must be ignored here
+            cp_backup_is_intermediate_chunk=False,
+            fill_ids=list(range(10)),
+            req_pool_idx=0,
+            extra_key=None,
+            cache_protected_len=0,
+            cp_hicache_prepared_backup=None,
+        )
+
+        cache.prepare_write_backup_for_req(req)
+
+        self.assertEqual(len(cache.cache_controller.submitted), 1)
+        # Full 10 tokens reserved (not floored to 8): matches the final node.
+        self.assertEqual(
+            cache.cache_controller.submitted[0].physical_device_indices.tolist(),
+            list(range(10)),
+        )
+        self.assertEqual(req.cp_hicache_prepared_backup.logical_len, 10)
+
     def test_cp_cache_unfinished_chunked_inserts_only_completed_pages(self):
         cache = HiRadixCache.__new__(HiRadixCache)
         cache.disable = False