chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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

88
test/registered/unit/disaggregation/test_cp_shared_kv_transfer_mapping.py
View File

@@ -1,8 +1,11 @@
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
from sglang.srt.disaggregation.utils import (
append_cp_draft_state_buffers,
filter_kv_pages_for_cp_shared_kv,
select_pages_by_request_positions,
)
@@ -56,6 +59,91 @@ class TestCPSharedKVTransferMapping(unittest.TestCase):
self.assertEqual(logical_positions.tolist(), [2, 6])
self.assertEqual(selected_decode_pages.tolist(), [43, 47])
def test_append_cp_draft_state_buffers_registers_nsa_state_as_target_payload(self):
kv_args = SimpleNamespace(
state_type="nsa",
state_data_ptrs=[11],
state_data_lens=[12],
state_item_lens=[13],
)
appended = append_cp_draft_state_buffers(
kv_args,
draft_state_type="nsa",
draft_state_data_ptrs=[21, 22],
draft_state_data_lens=[23, 24],
draft_state_item_lens=[25, 26],
role="prefill",
cp_rank=2,
)
self.assertTrue(appended)
self.assertEqual(kv_args.draft_state_type, "nsa")
self.assertEqual(kv_args.draft_state_buffer_start, 1)
self.assertEqual(kv_args.draft_state_buffer_count, 2)
self.assertEqual(kv_args.state_data_ptrs, [11, 21, 22])
self.assertEqual(kv_args.state_data_lens, [12, 23, 24])
self.assertEqual(kv_args.state_item_lens, [13, 25, 26])
def test_append_cp_draft_state_buffers_rejects_state_mismatch_under_shared_kv(self):
kv_args = SimpleNamespace(
state_type="nsa",
state_data_ptrs=[11],
state_data_lens=[12],
state_item_lens=[13],
)
with patch(
"sglang.srt.disaggregation.utils.envs.SGLANG_CP_DRAFT_SHARED_KV.get",
return_value=True,
):
with self.assertRaisesRegex(
RuntimeError,
"prefill.*cp_rank=2.*target_state_type=nsa.*draft_state_type=mamba",
):
append_cp_draft_state_buffers(
kv_args,
draft_state_type="mamba",
draft_state_data_ptrs=[21],
draft_state_data_lens=[22],
draft_state_item_lens=[23],
role="prefill",
cp_rank=2,
)
self.assertFalse(hasattr(kv_args, "draft_state_buffer_count"))
self.assertEqual(kv_args.state_data_ptrs, [11])
def test_append_cp_draft_state_buffers_keeps_legacy_skip_when_shared_kv_disabled(
self,
):
kv_args = SimpleNamespace(
state_type="nsa",
state_data_ptrs=[11],
state_data_lens=[12],
state_item_lens=[13],
)
with patch(
"sglang.srt.disaggregation.utils.envs.SGLANG_CP_DRAFT_SHARED_KV.get",
return_value=False,
):
appended = append_cp_draft_state_buffers(
kv_args,
draft_state_type="mamba",
draft_state_data_ptrs=[21],
draft_state_data_lens=[22],
draft_state_item_lens=[23],
role="decode",
cp_rank=3,
)
self.assertFalse(appended)
self.assertEqual(kv_args.draft_state_type, "mamba")
self.assertEqual(kv_args.draft_state_buffer_start, 1)
self.assertEqual(kv_args.draft_state_buffer_count, 0)
self.assertEqual(kv_args.state_data_ptrs, [11])
if __name__ == "__main__":
unittest.main()

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

@@ -251,6 +251,24 @@ class DummyLayerDoneCounter:
return 0
class RecordingProducerEvent:
def __init__(self, order):
self.start_event = DummyEvent()
self.finish_event = DummyEvent()
self.order = order
def complete(self, layer_id):
self.order.append(("complete", layer_id))
class RecordingLayerDoneCounter:
def __init__(self, order):
self.events = [RecordingProducerEvent(order)]
def update_producer(self):
return 0
class TestHiCacheControllerCPWrite(CustomTestCase):
def setUp(self):
self.device_module_patcher = patch(
@@ -344,6 +362,26 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
self.assertEqual(host_pool.alloc_calls, [])
self.assertEqual(len(controller.ack_write_queue), 1)
def test_cp_write_zero_owned_with_draft_returns_empty_draft_metadata(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
draft_host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
controller = self.make_controller(
host_pool,
cp_rank=3,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=FakeDevicePool("draft"),
)
logical_locs = torch.arange(4, 8, dtype=torch.int64)
result = controller.write(logical_locs, node_id=18)
self.assertEqual(result.metadata.logical_len, 4)
self.assertEqual(result.metadata.host_indices.tolist(), [])
self.assertEqual(result.metadata.draft_host_indices.tolist(), [])
self.assertEqual(host_pool.alloc_calls, [])
self.assertEqual(draft_host_pool.alloc_calls, [])
self.assertEqual(len(controller.ack_write_queue), 1)
def test_cp_write_allocation_failure_reports_required_host_slots(self):
host_pool = FakeHostPool(None)
controller = self.make_controller(host_pool, cp_rank=1)
@@ -519,6 +557,50 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
self.assertEqual(len(controller.ack_load_queue), 1)
self.assertEqual(controller.ack_load_queue[0].node_ids, [14])
def test_cp_start_loading_loads_draft_before_target_layer_ready(self):
order = []
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
draft_host_pool = FakeHostPool(
torch.tensor([200, 201, 202, 203], dtype=torch.int64)
)
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))
controller = self.make_controller(
host_pool,
allocator=allocator,
cp_rank=1,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=FakeDevicePool("draft"),
)
controller.layer_done_counter = RecordingLayerDoneCounter(order)
def record_target_load(
device_pool, host_indices, device_indices, layer_id, io_backend
):
order.append(("target", layer_id))
def record_draft_load(
device_pool, host_indices, device_indices, layer_id, io_backend
):
order.append(("draft", layer_id))
host_pool.load_to_device_per_layer = record_target_load
draft_host_pool.load_to_device_per_layer = record_draft_load
node = TreeNode()
node.host_len = 16
node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
controller.load_cp([node], node_id=114)
controller.start_loading()
self.assertEqual(order, [("draft", 0), ("target", 0), ("complete", 0)])
def test_cp_load_zero_owned_returns_full_logical_locs_and_noop_ack(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 68, dtype=torch.int64))
@@ -634,6 +716,51 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
with self.assertRaisesRegex(ValueError, "host_indices.*contiguous page spans"):
controller.load_cp([node], node_id=32)
def test_cp_evict_host_frees_target_and_draft_host_indices(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
draft_host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
controller = self.make_controller(
host_pool,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=FakeDevicePool("draft"),
)
metadata = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
draft_host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
page_owners=torch.tensor([0, 1, 2, 3], dtype=torch.int8),
page_size=4,
)
freed = controller.evict_cp_host(metadata)
self.assertEqual(freed, 4)
self.assertEqual(host_pool.frees[0].tolist(), [100, 101, 102, 103])
self.assertEqual(draft_host_pool.frees[0].tolist(), [200, 201, 202, 203])
def test_cp_evict_host_rejects_missing_draft_metadata_before_target_free(self):
host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
draft_host_pool = FakeHostPool(torch.tensor([], dtype=torch.int64))
controller = self.make_controller(
host_pool,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=FakeDevicePool("draft"),
)
metadata = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([0, 1, 2, 3], dtype=torch.int64),
host_indices=torch.tensor([100, 101, 102, 103], dtype=torch.int64),
page_owners=torch.tensor([0, 1, 2, 3], dtype=torch.int8),
page_size=4,
)
with self.assertRaisesRegex(RuntimeError, "draft.*evict.*draft_host_indices"):
controller.evict_cp_host(metadata)
self.assertEqual(host_pool.frees, [])
self.assertEqual(draft_host_pool.frees, [])
if __name__ == "__main__":
main()

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