Prevent CP HiCache pending-backup splits from killing chunked prefill
Chunked prefill can revisit a sub-page CP HiCache tail while a per-layer backup is still in flight. The old insert path split the radix node first and only then tried to prune the stale tail, so an unprunable pending backup raised after tree mutation and propagated to the scheduler. This makes split/prune atomic from the radix-tree perspective: drain completed write acks only on the conflict path, preflight pending/unprunable state before split, and return a deferred insert result when the backup is still in flight. Unfinished requests keep their KV ownership for transfer/release instead of rematching or freeing pages under a stale tree state. Constraint: CP HiCache backup metadata is node/page owned and cannot be repartitioned while per-layer D2H is pending Rejected: Split the pending backup node | would require repartitioning in-flight backup metadata and host reservations Rejected: Delete the stale tail unconditionally | risks freeing device/host pages still owned by pending backup Confidence: high Scope-risk: moderate Directive: Do not move stale-tail prune after split without a preflight; pending backup split conflicts must remain non-mutating Tested: remote g0034 container py_compile for mem_cache files Tested: remote g0034 PYTHONPATH=python pytest -q test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py => 115 passed Tested: local git diff --check Not-tested: full chunked prefill ETE after service restart Co-authored-by: OmX <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
@@ -1612,6 +1612,96 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
|
||||
self.assertEqual(allocator.freed, [])
|
||||
self.assertEqual(req.cache_protected_len, 6)
|
||||
|
||||
def test_cache_unfinished_req_deferred_insert_preserves_request_kv(self):
|
||||
cache = HiRadixCache.__new__(HiRadixCache)
|
||||
cache.disable = False
|
||||
cache._uses_cp_hicache = True
|
||||
cache.is_eagle = False
|
||||
cache.page_size = 4
|
||||
deferred_node = TreeNode()
|
||||
|
||||
class Pool:
|
||||
req_to_token = torch.arange(8, dtype=torch.int64).view(1, 8)
|
||||
|
||||
def write(self, index, values):
|
||||
raise AssertionError("deferred insert must not rewrite req pool")
|
||||
|
||||
cache.req_to_token_pool = Pool()
|
||||
cache.token_to_kv_pool_allocator = RecordingTokenAllocator()
|
||||
cache.insert = lambda params: types.SimpleNamespace(
|
||||
prefix_len=1,
|
||||
pending_backup_deferred_node=deferred_node,
|
||||
)
|
||||
cache.match_prefix = lambda params: (_ for _ in ()).throw(
|
||||
AssertionError("deferred insert must not rematch")
|
||||
)
|
||||
cache._free_kv_indices_range = lambda *args, **kwargs: (_ for _ in ()).throw(
|
||||
AssertionError("deferred unfinished insert must keep request KV")
|
||||
)
|
||||
old_last_node = TreeNode()
|
||||
req = types.SimpleNamespace(
|
||||
fill_ids=list(range(6)),
|
||||
req_pool_idx=0,
|
||||
extra_key=None,
|
||||
cache_protected_len=1,
|
||||
last_node=old_last_node,
|
||||
prefix_indices=torch.empty((0,), dtype=torch.int64),
|
||||
cp_hicache_prepared_backup=None,
|
||||
)
|
||||
|
||||
cache.cache_unfinished_req(req, chunked=True)
|
||||
|
||||
self.assertEqual(req.cache_protected_len, 1)
|
||||
self.assertIs(req.last_node, old_last_node)
|
||||
self.assertEqual(req.prefix_indices.tolist(), list(range(6)))
|
||||
self.assertEqual(cache.token_to_kv_pool_allocator.freed, [])
|
||||
|
||||
def test_cache_unfinished_req_deferred_insert_rolls_back_unattached_backup(self):
|
||||
cache = HiRadixCache.__new__(HiRadixCache)
|
||||
cache.disable = False
|
||||
cache._uses_cp_hicache = True
|
||||
cache.is_eagle = False
|
||||
cache.page_size = 4
|
||||
deferred_node = TreeNode()
|
||||
rollbacks = []
|
||||
|
||||
class Pool:
|
||||
req_to_token = torch.arange(8, dtype=torch.int64).view(1, 8)
|
||||
|
||||
cache.req_to_token_pool = Pool()
|
||||
cache.token_to_kv_pool_allocator = RecordingTokenAllocator()
|
||||
cache.insert = lambda params: types.SimpleNamespace(
|
||||
prefix_len=1,
|
||||
pending_backup_deferred_node=deferred_node,
|
||||
)
|
||||
cache.match_prefix = lambda params: (_ for _ in ()).throw(
|
||||
AssertionError("deferred insert must not rematch")
|
||||
)
|
||||
cache._free_kv_indices_range = lambda *args, **kwargs: (_ for _ in ()).throw(
|
||||
AssertionError("deferred unfinished insert must keep request KV")
|
||||
)
|
||||
cache._rollback_prepared_cp_backup = lambda prepared, reason: rollbacks.append(
|
||||
(prepared, reason)
|
||||
)
|
||||
prepared = types.SimpleNamespace(attached=False)
|
||||
req = types.SimpleNamespace(
|
||||
fill_ids=list(range(6)),
|
||||
req_pool_idx=0,
|
||||
extra_key=None,
|
||||
cache_protected_len=1,
|
||||
last_node=TreeNode(),
|
||||
prefix_indices=torch.empty((0,), dtype=torch.int64),
|
||||
cp_hicache_prepared_backup=prepared,
|
||||
)
|
||||
|
||||
cache.cache_unfinished_req(req, chunked=False)
|
||||
|
||||
self.assertEqual(
|
||||
rollbacks, [(prepared, "unfinished_pending_backup_split_deferred")]
|
||||
)
|
||||
self.assertIsNone(req.cp_hicache_prepared_backup)
|
||||
self.assertEqual(req.prefix_indices.tolist(), list(range(6)))
|
||||
|
||||
def test_prepare_write_backup_for_req_skips_existing_insert_prefix(self):
|
||||
cache = HiRadixCache.__new__(HiRadixCache)
|
||||
cache.disable = False
|
||||
@@ -3011,6 +3101,60 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
|
||||
self.assertEqual(new_tail.key.token_ids, [4])
|
||||
self.assertEqual(new_tail.value.tolist(), [4])
|
||||
|
||||
def test_cp_insert_defers_stale_tail_prune_while_backup_pending(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.ongoing_write_through = {}
|
||||
cache.cache_controller = types.SimpleNamespace(
|
||||
has_draft_hicache=False,
|
||||
write_policy="write_through",
|
||||
ack_write_queue=[],
|
||||
)
|
||||
cache.tp_world_size = 1
|
||||
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 = 152
|
||||
node.parent = root
|
||||
node.key = RadixKey(list(range(6)))
|
||||
node.value = torch.arange(6, dtype=torch.int64)
|
||||
node.host_len = 0
|
||||
node.cp_hicache = None
|
||||
root.children[(0, 1, 2, 3)] = node
|
||||
cache.ongoing_write_through[node.id] = node
|
||||
|
||||
result = cache.insert(
|
||||
InsertParams(
|
||||
key=RadixKey(list(range(10))),
|
||||
value=torch.arange(10, dtype=torch.int64),
|
||||
)
|
||||
)
|
||||
|
||||
self.assertEqual(result.prefix_len, 0)
|
||||
self.assertIs(result.pending_backup_deferred_node, node)
|
||||
self.assertIs(root.children[(0, 1, 2, 3)], node)
|
||||
self.assertEqual(node.key.token_ids, list(range(6)))
|
||||
self.assertEqual(node.value.tolist(), list(range(6)))
|
||||
|
||||
def test_cp_insert_extends_from_page_boundary_after_exact_valid_tail(self):
|
||||
cache = HiRadixCache.__new__(HiRadixCache)
|
||||
cache._uses_cp_hicache = True
|
||||
|
||||
Reference in New Issue
Block a user