chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Preserve draft KV across CP HiCache hits

Browse Source 
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.
This commit is contained in:
laoyao0822
2026-05-24 00:45:30 +08:00
committed by leavelet
parent 315eaaff56
commit fa80c6278e
2 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
test/registered/unit/managers/test_hicache_controller_cp.py
View File

@@ -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):

3
test/registered/unit/mem_cache/test_cp_hicache_metadata.py
View File

@@ -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: [