From fa80c6278edbea8e289b6851d4553e2d82170fd5 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sun, 24 May 2026 00:45:30 +0800 Subject: [PATCH] Preserve draft KV across CP HiCache hits Cache-hit prefill can skip draft forward for the prefix while PD transfer still reads draft KV for that same prefix. CP HiCache therefore needs to persist draft/MTP KV alongside target KV instead of relying on whatever remains in the draft GPU pool. Constraint: CP HiCache is host-only here; storage backends remain unsupported for CP shared KV. Constraint: CP shared KV must keep owner-page semantics and avoid falling back to full KV on every rank. Rejected: Recompute cached-prefix draft KV during prefill | loses the HiCache benefit and reintroduces the large hidden/KV footprint. Rejected: Change PD transfer to skip draft prefix KV | decode still needs draft cache continuity for MTP acceptance. Confidence: medium Scope-risk: moderate Directive: Keep target and draft CP HiCache metadata/load/write/evict paths in lockstep; changing one without the other can silently reduce MTP accept length. Tested: Remote g0034 container /sgl-workspace/sglang-tai: python3 -m pytest -q test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_metadata.py => 58 passed, 3 warnings Not-tested: Full multi-node HiCache+MTP serving benchmark and accept-length recovery. --- test/registered/unit/managers/test_hicache_controller_cp.py | 4 +++- test/registered/unit/mem_cache/test_cp_hicache_metadata.py | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index e898462f0..49d25d9a8 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -285,7 +285,9 @@ class TestHiCacheControllerCPWrite(CustomTestCase): controller = self.make_controller(host_pool, cp_rank=1) logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64) - with self.assertRaisesRegex(ValueError, "host_indices.*whole pages"): + with self.assertRaisesRegex( + ValueError, "(host_indices|physical_device_indices).*whole pages" + ): controller.write(logical_locs, node_id=21) def test_cp_write_rejects_non_contiguous_owned_physical_page(self): 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 117c2acf6..dc8d16fa9 100644 --- a/test/registered/unit/mem_cache/test_cp_hicache_metadata.py +++ b/test/registered/unit/mem_cache/test_cp_hicache_metadata.py @@ -668,6 +668,7 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): cache._uses_cp_hicache = True cache.tp_world_size = 2 cache.tp_group = object() + cache._tp_group_rank = 0 cache.root_node = TreeNode() cache.root_node.key = RadixKey([]) cache.evictable_host_leaves = set() @@ -700,7 +701,7 @@ class TestHiRadixCacheCPSplitEvict(CustomTestCase): cache.evictable_host_leaves.add(node) all_done_states = iter([False, True]) - cache._cp_all_ranks_true = lambda done: next(all_done_states) + cache._cp_all_ranks_true = lambda done: next(all_done_states, True) cache._cp_broadcast_node_ids = lambda node_ids, max_ids: node_ids[:max_ids] cache._cp_filter_all_ranks_safe_node_ids = ( lambda node_ids, is_safe, **_kwargs: [