From 2b1524bd8ce009d24b7118d7655b46c698258b54 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 29 May 2026 07:10:53 +0800 Subject: [PATCH] Apply CP HiCache page-floor policy before backup and insert Backed CP HiCache tail splits were floored during match, but write preparation and radix insertion could still observe the raw sub-page prefix. That could skip one token during pre-forward reservation or call CpHiCacheNodeMetadata.split() with a non-page boundary. The shared helper now floors backed partial-tail prefixes before probe, match, and insert. Constraint: CP HiCache uses page-granular physical ownership while radix keys may keep valid-tail lengths. Rejected: Let probe and insert keep token-precise tail prefixes | it reintroduces half-page ownership through a different entry point. Confidence: high Scope-risk: moderate Directive: Any CP radix path that can split or reserve against backed HiCache metadata must use the same page-floor policy. Tested: Remote red tests first for probe and insert failures in g0034 container. Tested: Remote py_compile for hiradix_cache.py and test_cp_hicache_metadata.py. Tested: Remote targeted C7 tests: 4 passed, 3 warnings. Tested: Remote pytest test_cp_hicache_metadata.py test_cp_hicache_load_back_owner_lanes.py: 94 passed, 5 warnings. Not-tested: Live ETE traffic with concurrent divergent tail prompts. Co-authored-by: OmX --- ..._prefill_cp_page_aligned_cache_contract.md | 5 + python/sglang/srt/mem_cache/hiradix_cache.py | 28 ++++- .../mem_cache/test_cp_hicache_metadata.py | 102 ++++++++++++++++++ 3 files changed, 134 insertions(+), 1 deletion(-) diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index e4d1bd632..5c46a8955 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -660,6 +660,11 @@ Implemented C7 split policy: - If the floored boundary is zero, the match returns the parent/root node and sacrifices the sub-page prefix. This avoids half-page owner accounting and keeps host/device/draft metadata page-granular. +- The same floor policy is applied before forward-time write reservation and + final radix insertion. `_probe_existing_radix_prefix_len_no_split()` floors + backed partial-tail hits so backup reservation starts at the same page boundary + as `match_prefix()`, and `insert()` floors backed partial-tail splits before + calling `CpHiCacheNodeMetadata.split()`. ### C8. Owner-lane capacity must be padded-page based end to end diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py index 764228d2b..381f18c03 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -2009,6 +2009,12 @@ class HiRadixCache(RadixCache): prefix_len = self.key_match_fn(child.key, key) if prefix_len <= 0: break + if prefix_len < len(child.key): + prefix_len = self._cp_floor_backed_partial_split_len( + child, prefix_len + ) + if prefix_len <= 0: + break total_prefix_len += prefix_len if prefix_len < len(child.key): @@ -2021,6 +2027,19 @@ class HiRadixCache(RadixCache): return total_prefix_len + def _cp_floor_backed_partial_split_len( + self, child: TreeNode, prefix_len: int + ) -> int: + if ( + not self._uses_cp_hicache + or self.page_size <= 1 + or prefix_len <= 0 + or prefix_len % self.page_size == 0 + or not self._node_backuped(child) + ): + return prefix_len + return prefix_len // self.page_size * self.page_size + def prepare_write_backup_for_req(self, req) -> None: if self.disable or not self._uses_cp_hicache: return @@ -3356,7 +3375,9 @@ class HiRadixCache(RadixCache): and self._node_backuped(child) and prefix_len % self.page_size != 0 ): - prefix_len = prefix_len // self.page_size * self.page_size + prefix_len = self._cp_floor_backed_partial_split_len( + child, prefix_len + ) if prefix_len == 0: break new_node = self._split_node(child.key, child, prefix_len) @@ -3464,6 +3485,11 @@ class HiRadixCache(RadixCache): total_prefix_length += prefix_len else: # partial match, split the node + prefix_len = self._cp_floor_backed_partial_split_len( + node, prefix_len + ) + if prefix_len <= 0: + break new_node = self._split_node(node.key, node, prefix_len) # shared-prefix node should also reflect max priority new_node.priority = max(new_node.priority, priority) 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 c8e87488f..4a1b9c18a 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -2437,6 +2437,108 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): self.assertIs(root.children[(0, 1, 2, 3)], node) self.assertEqual(node.key.token_ids, list(range(6))) + def test_cp_prepare_probe_floors_backed_tail_partial_hit_to_page(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.disable = False + cache.page_size = 4 + cache.pending_host_backups = {} + cache.cache_controller = types.SimpleNamespace(has_draft_hicache=False) + 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) + root = TreeNode() + root.key = RadixKey([]) + root.children = {} + cache.root_node = root + node = TreeNode() + node.id = 143 + node.parent = root + node.key = RadixKey(list(range(6))) + node.value = None + node.host_len = 6 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=6, + padded_len=8, + owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64), + host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64), + page_owners=torch.tensor([0, 1], dtype=torch.int8), + page_size=4, + ) + root.children[(0, 1, 2, 3)] = node + + prefix_len = cache._probe_existing_radix_prefix_len_no_split( + RadixKey([0, 1, 2, 3, 4]) + ) + + self.assertEqual(prefix_len, 4) + + def test_cp_insert_floors_backed_tail_split_to_page_boundary(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.disable = False + cache.is_eagle = False + cache.page_size = 4 + cache.pending_host_backups = {} + cache.cache_controller = types.SimpleNamespace( + has_draft_hicache=False, + write_policy="write_back", + ) + 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._update_leaf_status = lambda node: None + cache._update_host_leaf_status = lambda node: None + cache._inc_hit_count = lambda *args, **kwargs: None + cache._record_store_event = lambda node: None + cache.evictable_size_ = 0 + cache.enable_storage = False + cache.enable_kv_cache_events = False + root = TreeNode() + root.key = RadixKey([]) + root.value = torch.empty((0,), dtype=torch.int64) + root.children = {} + cache.root_node = root + node = TreeNode() + node.id = 144 + node.parent = root + node.key = RadixKey(list(range(6))) + node.value = torch.arange(6, dtype=torch.int64) + node.host_len = 6 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=6, + padded_len=8, + owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64), + host_indices=torch.tensor([50, 51, 52, 53], dtype=torch.int64), + page_owners=torch.tensor([0, 1], dtype=torch.int8), + page_size=4, + ) + root.children[(0, 1, 2, 3)] = node + + result = cache.insert( + InsertParams( + key=RadixKey([0, 1, 2, 3, 4]), + value=torch.arange(5, dtype=torch.int64), + ) + ) + + self.assertEqual(result.prefix_len, 4) + parent = root.children[(0, 1, 2, 3)] + self.assertEqual(parent.key.token_ids, [0, 1, 2, 3]) + self.assertEqual(parent.cp_hicache.logical_len, 4) + self.assertEqual(parent.cp_hicache.padded_len, 4) + old_tail = parent.children[(4, 5)] + self.assertEqual(old_tail.key.token_ids, [4, 5]) + self.assertEqual(old_tail.cp_hicache.logical_len, 2) + self.assertEqual(old_tail.cp_hicache.padded_len, 4) + new_tail = parent.children[(4,)] + self.assertEqual(new_tail.key.token_ids, [4]) + self.assertEqual(new_tail.value.tolist(), [4]) + def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = False