Keep CP HiCache draft KV invisible until joint readiness
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
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user