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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user