From 512fe92a83d66e95fcf5e9635e5c5564935b3854 Mon Sep 17 00:00:00 2001 From: leavelet Date: Thu, 18 Jun 2026 11:54:24 +0000 Subject: [PATCH] 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 --- python/sglang/srt/managers/scheduler.py | 17 ++ python/sglang/srt/mem_cache/hiradix_cache.py | 22 +- .../test_cp_hicache_chunked_prepare_repro.py | 263 ++++++++++++++++++ .../mem_cache/test_cp_hicache_metadata.py | 55 ++++ 4 files changed, 351 insertions(+), 6 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_cp_hicache_chunked_prepare_repro.py diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 0fdcb8d60..9aae15d6d 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 ) diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index dfd9a90df..76e27d51c 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/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( diff --git a/test/registered/unit/mem_cache/test_cp_hicache_chunked_prepare_repro.py b/test/registered/unit/mem_cache/test_cp_hicache_chunked_prepare_repro.py new file mode 100644 index 000000000..34d580c45 --- /dev/null +++ b/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) diff --git a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py index 6e7e62450..4884bd311 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/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