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

Keep CP HiCache draft KV invisible until joint readiness

Browse Source 
CP HiCache now treats draft KV as a strict target-owned payload through pending write visibility, host eviction, and state-buffer registration. Host metadata created before async D2H ack is no longer request-visible, so match_prefix cannot select an in-flight host node. Draft host eviction now fails before target cleanup when draft metadata is missing, and prefill/decode share one helper for draft NSA state buffers so shared-KV mode cannot silently skip mismatched draft state.

Constraint: CP shared KV + HiCache + EAGLE/MTP must not expose target-only host hits or skipped draft state as valid cache hits

Rejected: Rely on event-loop ordering and lock_ref to hide in-flight writes | match_prefix does not consult lock_ref and can observe host_len/cp_hicache directly

Rejected: Keep draft state mismatch as debug-only skip | it can poison speculative acceptance while looking like a successful cache hit

Confidence: high

Scope-risk: moderate

Directive: Do not reintroduce silent draft/target fallback in CP shared-KV HiCache paths; malformed strong-sync metadata should fail fast

Tested: python -m py_compile targeted modified files

Tested: remote g0034 container pytest test/registered/unit/mem_cache/test_cp_hicache_metadata.py test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py -q (91 passed)

Not-tested: Full CP shared KV + HiCache + EAGLE/MTP ETE server run after this commit
This commit is contained in:
laoyao0822
2026-05-27 05:46:34 +08:00
parent 71c4f66968
commit 367dff06f3
9 changed files with 448 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -620,6 +620,29 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
self.assertTrue(cache._node_backuped(node))
def test_node_backuped_excludes_inflight_cp_write_until_ack(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.cache_controller = types.SimpleNamespace(has_draft_hicache=True)
cache.ongoing_write_through = {}
node = TreeNode()
node.id = 125
node.host_len = 64
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=64,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
draft_host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=64,
)
cache.ongoing_write_through[node.id] = node
self.assertFalse(cache._node_backuped(node))
cache.ongoing_write_through.clear()
self.assertTrue(cache._node_backuped(node))
def test_single_node_write_lock_updates_device_evictable_leaf_set(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache.disable = False
@@ -672,6 +695,34 @@ class TestHiRadixCacheCPBackup(CustomTestCase):
self.assertEqual(node.hit_count, 1)
def test_inc_hit_count_does_not_duplicate_inflight_cp_write(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.write_through_threshold = 1
cache.ongoing_write_through = {}
cache.cache_controller = type(
"Controller", (), {"write_policy": "write_through"}
)()
cache.write_backup = lambda node: (_ for _ in ()).throw(
AssertionError("must not launch a second write")
)
node = TreeNode()
node.id = 127
node.host_len = 4
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=4,
owned_positions=torch.tensor([], dtype=torch.int64),
host_indices=torch.tensor([], dtype=torch.int64),
draft_host_indices=torch.tensor([], dtype=torch.int64),
page_owners=torch.zeros(max(4, 0), dtype=torch.int8),
page_size=1,
)
cache.ongoing_write_through[node.id] = node
cache._inc_hit_count(node)
self.assertEqual(node.hit_count, 1)
def test_write_backup_retries_by_required_physical_slots(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
@@ -1217,6 +1268,49 @@ class TestHiRadixCacheCPLoadBack(CustomTestCase):
self.assertIs(result.last_host_node, node)
self.assertIs(result.last_device_node, root)
def test_cp_match_prefix_does_not_admit_inflight_write_as_host_hit(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True
cache.device = "cpu"
cache.disable = False
cache.page_size = 1
cache.ongoing_write_through = {}
cache.cache_controller = types.SimpleNamespace(has_draft_hicache=True)
cache.get_child_key_fn = lambda key: key.token_ids[0]
cache.key_match_fn = lambda child_key, key: sum(
1 for lhs, rhs in zip(child_key.token_ids, key.token_ids) if lhs == rhs
)
cache.maybe_bigram_convert = lambda key: (key, None)
root = TreeNode()
root.key = RadixKey([])
root.value = torch.empty((0,), dtype=torch.int64)
root.host_len = 0
cache.root_node = root
node = TreeNode()
node.id = 126
node.parent = root
node.key = RadixKey(list(range(8)))
node.value = torch.arange(8, dtype=torch.int64)
node.host_value = None
node.host_len = 8
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=8,
owned_positions=torch.tensor([0, 1], dtype=torch.int64),
host_indices=torch.tensor([50, 51], dtype=torch.int64),
draft_host_indices=torch.tensor([150, 151], dtype=torch.int64),
page_owners=torch.zeros(max(8, 0), dtype=torch.int8),
page_size=1,
)
root.children[0] = node
cache.ongoing_write_through[node.id] = node
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(list(range(8)))))
self.assertEqual(result.device_indices.tolist(), list(range(8)))
self.assertEqual(result.host_hit_length, 0)
self.assertIs(result.last_device_node, node)
self.assertIs(result.last_host_node, root)
def test_cp_match_prefix_shorter_than_page_returns_empty_root_match(self):
cache = HiRadixCache.__new__(HiRadixCache)
cache._uses_cp_hicache = True