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 a87ab76c0..e4d1bd632 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 @@ -648,13 +648,18 @@ Implemented C7 slice: - `_key_match_paged()` now returns the true valid prefix length for a partial final page. -Remaining C7 limitation: +Implemented C7 split policy: -- Splitting an already-backed CP HiCache node inside a padded tail page still - needs a deliberate design. The current safe slice supports exact valid-tail - hits and writes; divergent requests that force a split inside a physical tail - page are still a follow-up because splitting one physical host page across two - radix nodes would otherwise double-count or lose ownership metadata. +- Cache management treats a page as the minimum ownership unit. If matching an + already-backed CP HiCache node would split inside a padded physical tail page, + the match is floored to the previous page boundary instead of splitting that + physical page. +- Exact valid-tail hits still report the valid length. The floor policy only + applies when a shorter/divergent request would require a new radix boundary + inside the tail page. +- 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. ### 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 706726f96..764228d2b 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py @@ -3350,6 +3350,15 @@ class HiRadixCache(RadixCache): and child.id in getattr(self, "pending_host_backups", {}) ): raise HiCachePendingBackupSplit(child) + if ( + self._uses_cp_hicache + and self.page_size > 1 + and self._node_backuped(child) + and prefix_len % self.page_size != 0 + ): + prefix_len = prefix_len // self.page_size * self.page_size + if prefix_len == 0: + break new_node = self._split_node(child.key, child, prefix_len) if not new_node.evicted: value.append(new_node.value) 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 b262e45de..c8e87488f 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -2339,6 +2339,104 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase): self.assertIs(result.last_device_node, root) self.assertIs(result.last_host_node, node) + def test_cp_backed_tail_split_floors_to_page_boundary(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.device = "cpu" + cache.disable = False + cache.page_size = 4 + cache.ongoing_write_through = {} + 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: (key, None) + root = TreeNode() + root.key = RadixKey([]) + root.value = torch.empty((0,), dtype=torch.int64) + root.host_len = 0 + root.children = {} + cache.root_node = root + node = TreeNode() + node.id = 141 + node.parent = root + node.key = RadixKey(list(range(6))) + node.value = None + node.host_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 + + result = cache.match_prefix(MatchPrefixParams(key=RadixKey([0, 1, 2, 3, 4]))) + + self.assertEqual(result.device_indices.tolist(), []) + self.assertEqual(result.host_hit_length, 4) + self.assertIs(result.last_device_node, root) + self.assertEqual(result.last_host_node.key.token_ids, [0, 1, 2, 3]) + self.assertEqual(result.last_host_node.host_len, 4) + self.assertEqual(result.last_host_node.cp_hicache.logical_len, 4) + self.assertEqual(result.last_host_node.cp_hicache.padded_len, 4) + child = result.last_host_node.children[(4, 5)] + self.assertEqual(child.key.token_ids, [4, 5]) + self.assertEqual(child.host_len, 2) + self.assertEqual(child.cp_hicache.logical_len, 2) + self.assertEqual(child.cp_hicache.padded_len, 4) + + def test_cp_backed_tail_split_before_first_page_returns_root_match(self): + cache = HiRadixCache.__new__(HiRadixCache) + cache._uses_cp_hicache = True + cache.device = "cpu" + cache.disable = False + cache.page_size = 4 + cache.ongoing_write_through = {} + 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: (key, None) + root = TreeNode() + root.key = RadixKey([]) + root.value = torch.empty((0,), dtype=torch.int64) + root.host_len = 0 + root.children = {} + cache.root_node = root + node = TreeNode() + node.id = 142 + node.parent = root + node.key = RadixKey(list(range(6))) + node.value = None + node.host_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 + + result = cache.match_prefix(MatchPrefixParams(key=RadixKey([0, 1, 2]))) + + self.assertEqual(result.device_indices.tolist(), []) + self.assertEqual(result.host_hit_length, 0) + self.assertIs(result.last_device_node, root) + self.assertIs(result.last_host_node, root) + self.assertIs(root.children[(0, 1, 2, 3)], node) + self.assertEqual(node.key.token_ids, list(range(6))) + def test_non_cp_match_prefix_uses_root_when_no_host_backup_exists(self): cache = HiRadixCache.__new__(HiRadixCache) cache._uses_cp_hicache = False