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 <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-29 07:10:53 +08:00
co-authored by OmX
parent 8001c4ae8e
commit 2b1524bd8c
3 changed files with 134 additions and 1 deletions
@@ -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