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 <omx@oh-my-codex.dev>
This commit is contained in:
laoyao0822
2026-05-29 05:43:49 +08:00
parent 0043037f78
commit 7e06eaebbf
2 changed files with 78 additions and 0 deletions

View File

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

View File

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