From 7e06eaebbf818e18dd68ad623d629a813903935a Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 29 May 2026 05:43:49 +0800 Subject: [PATCH] Prove draft HiCache mirrors target valid-tail pages Draft HiCache already piggybacks target CP reservations, but the controller coverage only exercised page-aligned nodes. The new tests pin the intended contract for valid-tail nodes: target and draft reserve/load the same padded physical page while scheduler-visible load results keep only valid logical locs. Constraint: Draft KV is a target mirror; it must not choose an independent valid-only host path. Rejected: Re-enable draft partial-current reuse in this slice | current safe contract keeps EAGLE/NextN cache-hit draft on full materialization until same-layer padded visibility is proven. Confidence: medium Scope-risk: narrow Directive: Keep draft and target host metadata coupled; do not add draft-only capacity or prefetch decisions. Tested: local py_compile for test_hicache_controller_cp.py. Tested: remote g0034 test_hicache_controller_cp.py: 59 passed, 3 warnings. Not-tested: CUDA E2E runtime for this commit. Co-authored-by: OmX --- ..._prefill_cp_page_aligned_cache_contract.md | 12 ++++ .../managers/test_hicache_controller_cp.py | 66 +++++++++++++++++++ 2 files changed, 78 insertions(+) diff --git a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md index 56135aa07..5b3c6dd18 100644 --- a/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md +++ b/docs/advanced_features/nsa_prefill_cp_page_aligned_cache_contract.md @@ -697,6 +697,9 @@ Current state: `draft_host_indices` are present. - Draft partial-current reuse/prefetch has been disabled or pushed to older materialization paths in some cases to avoid EAGLE hangs. +- Existing draft HiCache controller tests cover page-aligned write/load/evict, + but they do not yet prove that a valid-tail node mirrors the target padded + physical span while exposing only the valid length. Correction: @@ -712,6 +715,15 @@ Tests: `owned_positions` and padded physical length. - EAGLE one-layer path does not prefetch or consume an unmasked padded tail. +Implemented C9 slice: + +- Added target/draft write coverage for a valid-tail suffix. Target and draft + reserve the same padded local positions and transfer the same physical device + span while metadata keeps the valid `logical_len`. +- Added target/draft load coverage for a valid-tail host hit. Both pools load + the padded physical page, but `load_cp()` returns only valid logical locs to + the scheduler. + ### C10. Disaggregated transfer has not been audited for valid/padded split Current state: diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index e0816b672..6e7fd205a 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -805,6 +805,35 @@ class TestHiCacheControllerCPWrite(CustomTestCase): self.assertEqual(len(controller.ack_write_queue), 1) self.assertEqual(controller.ack_write_queue[0].node_ids, [77]) + def test_cp_write_valid_tail_with_draft_mirrors_target_padded_page(self): + 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) + ) + draft_device_pool = FakeDevicePool("draft") + controller = self.make_controller( + host_pool, + cp_rank=1, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=draft_device_pool, + ) + logical_locs = torch.tensor([8, 9, 10], dtype=torch.int64) + + result = controller.write(logical_locs, node_id=177) + + self.assertEqual(result.metadata.logical_len, 3) + self.assertEqual(result.metadata.padded_len, 4) + self.assertEqual(result.metadata.owned_positions.tolist(), [0, 1, 2, 3]) + self.assertEqual(result.metadata.host_indices.tolist(), [100, 101, 102, 103]) + self.assertEqual( + result.metadata.draft_host_indices.tolist(), [200, 201, 202, 203] + ) + self.assertEqual(host_pool.alloc_calls, [4]) + self.assertEqual(draft_host_pool.alloc_calls, [4]) + self.assertEqual(host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7]) + self.assertEqual(draft_host_pool.layer_backups[0][1].tolist(), [4, 5, 6, 7]) + self.assertIs(draft_host_pool.layer_backups[0][3], draft_device_pool) + def test_cp_write_draft_allocation_failure_rolls_back_target_host(self): host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64)) draft_host_pool = FakeHostPool(None) @@ -1138,6 +1167,43 @@ class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite): self.assertEqual(len(controller.ack_load_queue), 1) self.assertEqual(controller.ack_load_queue[0].node_ids, [14]) + def test_cp_load_valid_tail_with_draft_returns_valid_locs_only(self): + 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) + ) + draft_device_pool = FakeDevicePool("draft") + allocator = FakeAllocator(alloc_result=torch.arange(8, 12, dtype=torch.int64)) + controller = self.make_controller( + host_pool, + allocator=allocator, + cp_rank=1, + draft_host_pool=draft_host_pool, + draft_mem_pool_device=draft_device_pool, + ) + node = TreeNode() + node.host_len = 3 + node.cp_hicache = CpHiCacheNodeMetadata( + logical_len=3, + padded_len=4, + 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([1], dtype=torch.int8), + page_size=4, + ) + + device_indices = controller.load_cp([node], node_id=178) + controller.start_loading() + + self.assertEqual(device_indices.tolist(), [8, 9, 10]) + self.assertEqual(allocator.owner_alloc_calls, [[1]]) + self.assertEqual(host_pool.loads[0][1].tolist(), [4, 5, 6, 7]) + self.assertEqual(draft_host_pool.loads[0][1].tolist(), [4, 5, 6, 7]) + self.assertIs(draft_host_pool.loads[0][3], draft_device_pool) + self.assertEqual(len(controller.ack_load_queue), 1) + self.assertEqual(controller.ack_load_queue[0].node_ids, [178]) + 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))