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

Batch CP HiCache backup submits across requests

Browse Source 
CP HiCache write reservations must stay per radix node, but the transfer descriptor does not need to be per request. This changes the layer-end hook to group pending write states for the same source and layer, so bs>1 prefill emits one target D2H descriptor and one draft D2H descriptor per layer while preserving per-node metadata, rollback, and ack semantics.\n\nConstraint: CP shared-KV HiCache metadata, host slots, and radix acknowledgements remain per request/node.\nConstraint: TAI direct transfer kernels already accept flattened page descriptors, so no tai-kernel change is required.\nRejected: Merge HiCache reservations or radix nodes | would complicate rollback and split handling.\nRejected: Add collective synchronization for grouped backup | grouping is local descriptor construction and must not add rank-level sync.\nConfidence: high\nScope-risk: moderate\nDirective: Keep target and draft source notifications separate; final ack must wait for both when draft HiCache is attached.\nTested: local py_compile for cache_controller.py and test_hicache_controller_cp.py\nTested: local git diff --check\nTested: remote pytest test/registered/unit/managers/test_hicache_controller_cp.py test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py => 85 passed, 3 warnings\nNot-tested: full ETE bs>1 CP HiCache replay with admission gate removed\nNot-tested: Nsight/throughput validation of reduced D2H submit count
This commit is contained in:
laoyao0822
2026-06-03 10:23:30 +08:00
parent 19dcd6c4dc
commit b3913046b6
3 changed files with 655 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/managers/test_hicache_controller_cp.py
+311
View File
@@ -1157,6 +1157,177 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
self.assertEqual(controller.ack_write_queue[0].node_ids, [84])
self.assertEqual(host_pool.backups, [])
def test_cp_layer_hook_groups_target_backups_across_pending_reservations(self):
class SequentialHostPool(FakeHostPool):
def __init__(self, alloc_results):
super().__init__(torch.empty((0,), dtype=torch.int64))
self.alloc_results = [result.clone() for result in alloc_results]
def alloc(self, need_size):
self.alloc_calls.append(need_size)
if not self.alloc_results:
return None
return self.alloc_results.pop(0).clone()
host_pool = SequentialHostPool(
[
torch.tensor([100, 101, 102, 103], dtype=torch.int64),
torch.tensor([200, 201, 202, 203], dtype=torch.int64),
]
)
allocator = FakeAllocator()
allocator.device_pool = FakeDevicePool("target", layer_num=2)
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
reservation_a = controller.reserve_write_cp(
torch.arange(4, 20, dtype=torch.int64), node_id=501
)
reservation_b = controller.reserve_write_cp(
torch.arange(20, 36, dtype=torch.int64), node_id=502
)
controller.submit_write_cp_per_layer(reservation_a, catch_up_all_layers=False)
controller.submit_write_cp_per_layer(reservation_b, catch_up_all_layers=False)
allocator.device_pool.notify_layer_end_for_backup(0)
self.assertEqual(len(host_pool.layer_backups), 1)
host_indices, device_indices, layer_id, device_pool = host_pool.layer_backups[0]
self.assertEqual(
host_indices.tolist(), [100, 101, 102, 103, 200, 201, 202, 203]
)
self.assertEqual(device_indices.tolist(), [4, 5, 6, 7, 8, 9, 10, 11])
self.assertEqual(layer_id, 0)
self.assertIs(device_pool, allocator.device_pool)
self.assertEqual(controller.ack_write_queue, [])
allocator.device_pool.notify_layer_end_for_backup(1)
self.assertEqual(len(host_pool.layer_backups), 2)
self.assertEqual([backup[2] for backup in host_pool.layer_backups], [0, 1])
self.assertEqual(len(controller.ack_write_queue), 2)
self.assertEqual(
[ack.node_ids for ack in controller.ack_write_queue], [[501], [502]]
)
def test_cp_layer_hook_groups_target_and_draft_backups_by_source(self):
class SequentialHostPool(FakeHostPool):
def __init__(self, alloc_results):
super().__init__(torch.empty((0,), dtype=torch.int64))
self.alloc_results = [result.clone() for result in alloc_results]
def alloc(self, need_size):
self.alloc_calls.append(need_size)
if not self.alloc_results:
return None
return self.alloc_results.pop(0).clone()
host_pool = SequentialHostPool(
[
torch.tensor([100, 101, 102, 103], dtype=torch.int64),
torch.tensor([200, 201, 202, 203], dtype=torch.int64),
]
)
draft_host_pool = SequentialHostPool(
[
torch.tensor([300, 301, 302, 303], dtype=torch.int64),
torch.tensor([400, 401, 402, 403], dtype=torch.int64),
]
)
allocator = FakeAllocator()
allocator.device_pool = FakeDevicePool("target", layer_num=2)
draft_device_pool = FakeDevicePool("draft", layer_num=2)
controller = self.make_controller(
host_pool,
allocator=allocator,
cp_rank=1,
draft_host_pool=draft_host_pool,
draft_mem_pool_device=draft_device_pool,
)
reservation_a = controller.reserve_write_cp(
torch.arange(4, 20, dtype=torch.int64), node_id=601
)
reservation_b = controller.reserve_write_cp(
torch.arange(20, 36, dtype=torch.int64), node_id=602
)
controller.submit_write_cp_per_layer(reservation_a, catch_up_all_layers=False)
controller.submit_write_cp_per_layer(reservation_b, catch_up_all_layers=False)
allocator.device_pool.notify_layer_end_for_backup(0)
self.assertEqual(len(host_pool.layer_backups), 1)
self.assertEqual(len(draft_host_pool.layer_backups), 0)
self.assertEqual(
host_pool.layer_backups[0][0].tolist(),
[100, 101, 102, 103, 200, 201, 202, 203],
)
self.assertEqual(
host_pool.layer_backups[0][1].tolist(),
[4, 5, 6, 7, 8, 9, 10, 11],
)
draft_device_pool.notify_layer_end_for_backup(0)
self.assertEqual(len(draft_host_pool.layer_backups), 1)
self.assertEqual(
draft_host_pool.layer_backups[0][0].tolist(),
[300, 301, 302, 303, 400, 401, 402, 403],
)
self.assertEqual(
draft_host_pool.layer_backups[0][1].tolist(),
[4, 5, 6, 7, 8, 9, 10, 11],
)
self.assertEqual(controller.ack_write_queue, [])
allocator.device_pool.notify_layer_end_for_backup(1)
self.assertEqual(len(host_pool.layer_backups), 2)
self.assertEqual(controller.ack_write_queue, [])
draft_device_pool.notify_layer_end_for_backup(1)
self.assertEqual(len(draft_host_pool.layer_backups), 2)
self.assertEqual(len(controller.ack_write_queue), 2)
self.assertEqual(
[ack.node_ids for ack in controller.ack_write_queue], [[601], [602]]
)
def test_cp_layer_hook_keeps_zero_owned_ack_in_grouped_backup(self):
class SequentialHostPool(FakeHostPool):
def __init__(self, alloc_results):
super().__init__(torch.empty((0,), dtype=torch.int64))
self.alloc_results = [result.clone() for result in alloc_results]
def alloc(self, need_size):
self.alloc_calls.append(need_size)
if not self.alloc_results:
return None
return self.alloc_results.pop(0).clone()
host_pool = SequentialHostPool(
[torch.tensor([100, 101, 102, 103], dtype=torch.int64)]
)
allocator = FakeAllocator()
allocator.device_pool = FakeDevicePool("target", layer_num=2)
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
zero_owned = controller.reserve_write_cp(
torch.arange(4, 8, dtype=torch.int64), node_id=701
)
owned = controller.reserve_write_cp(
torch.arange(8, 24, dtype=torch.int64), node_id=702
)
controller.submit_write_cp_per_layer(zero_owned, catch_up_all_layers=False)
controller.submit_write_cp_per_layer(owned, catch_up_all_layers=False)
allocator.device_pool.notify_layer_end_for_backup(0)
allocator.device_pool.notify_layer_end_for_backup(1)
self.assertEqual(len(host_pool.layer_backups), 2)
for host_indices, device_indices, _, _ in host_pool.layer_backups:
self.assertEqual(host_indices.tolist(), [100, 101, 102, 103])
self.assertEqual(device_indices.tolist(), [4, 5, 6, 7])
self.assertEqual(len(controller.ack_write_queue), 2)
self.assertEqual(
[ack.node_ids for ack in controller.ack_write_queue], [[701], [702]]
)
def test_cp_layer_hook_waits_for_draft_source_before_final_ack(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
draft_host_pool = FakeHostPool(
@@ -1227,6 +1398,146 @@ class TestHiCacheControllerCPWrite(CustomTestCase):
class TestHiCacheControllerCPLoad(TestHiCacheControllerCPWrite):
def test_cp_start_loading_batches_multiple_load_cp_requests_with_draft(self):
class SequentialOwnerAllocator(FakeAllocator):
def __init__(self, alloc_results):
super().__init__()
self.alloc_results = [result.clone() for result in alloc_results]
self.device_pool = FakeDevicePool("target", layer_num=2)
def alloc_pages_with_owners(self, page_owners):
owners = list(page_owners)
self.owner_alloc_calls.append(owners)
if not self.alloc_results:
return None
return self.alloc_results.pop(0).clone()
host_pool = FakeHostPool(torch.empty((0,), dtype=torch.int64))
draft_host_pool = FakeHostPool(torch.empty((0,), dtype=torch.int64))
allocator = SequentialOwnerAllocator(
[
torch.arange(64, 80, dtype=torch.int64),
torch.arange(80, 96, dtype=torch.int64),
]
)
draft_device_pool = FakeDevicePool("draft", layer_num=2)
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_a = TreeNode()
node_a.host_len = 16
node_a.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([300, 301, 302, 303], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
node_b = TreeNode()
node_b.host_len = 16
node_b.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
host_indices=torch.tensor([200, 201, 202, 203], dtype=torch.int64),
draft_host_indices=torch.tensor([400, 401, 402, 403], dtype=torch.int64),
page_owners=torch.tensor([3, 0, 1, 2], dtype=torch.int8),
page_size=4,
)
device_indices_a = controller.load_cp([node_a], node_id=201)
device_indices_b = controller.load_cp([node_b], node_id=202)
controller.start_loading()
self.assertEqual(device_indices_a.tolist(), list(range(64, 80)))
self.assertEqual(device_indices_b.tolist(), list(range(80, 96)))
self.assertEqual(
allocator.owner_alloc_calls,
[[3, 0, 1, 2], [3, 0, 1, 2]],
)
self.assertEqual(len(host_pool.loads), 2)
self.assertEqual([load[2] for load in host_pool.loads], [0, 1])
for host_indices, device_indices, _, device_pool in host_pool.loads:
self.assertEqual(
host_indices.tolist(),
[100, 101, 102, 103, 200, 201, 202, 203],
)
self.assertEqual(device_indices.tolist(), [20, 21, 22, 23, 24, 25, 26, 27])
self.assertIs(device_pool, allocator.device_pool)
self.assertEqual(len(draft_host_pool.loads), 2)
self.assertEqual([load[2] for load in draft_host_pool.loads], [0, 1])
for host_indices, device_indices, _, device_pool in draft_host_pool.loads:
self.assertEqual(
host_indices.tolist(),
[300, 301, 302, 303, 400, 401, 402, 403],
)
self.assertEqual(device_indices.tolist(), [20, 21, 22, 23, 24, 25, 26, 27])
self.assertIs(device_pool, draft_device_pool)
self.assertEqual(len(controller.ack_load_queue), 1)
self.assertEqual(controller.ack_load_queue[0].node_ids, [201, 202])
def test_cp_start_loading_keeps_zero_owned_load_ack_in_batched_load(self):
class SequentialOwnerAllocator(FakeAllocator):
def __init__(self, alloc_results):
super().__init__()
self.alloc_results = [result.clone() for result in alloc_results]
self.device_pool = FakeDevicePool("target", layer_num=2)
def alloc_pages_with_owners(self, page_owners):
owners = list(page_owners)
self.owner_alloc_calls.append(owners)
if not self.alloc_results:
return None
return self.alloc_results.pop(0).clone()
host_pool = FakeHostPool(torch.empty((0,), dtype=torch.int64))
allocator = SequentialOwnerAllocator(
[
torch.arange(64, 68, dtype=torch.int64),
torch.arange(80, 96, dtype=torch.int64),
]
)
controller = self.make_controller(host_pool, allocator=allocator, cp_rank=1)
zero_owned_node = TreeNode()
zero_owned_node.host_len = 4
zero_owned_node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=4,
owned_positions=torch.empty((0,), dtype=torch.int64),
host_indices=torch.empty((0,), dtype=torch.int64),
page_owners=torch.tensor([0], dtype=torch.int8),
page_size=4,
)
owned_node = TreeNode()
owned_node.host_len = 16
owned_node.cp_hicache = CpHiCacheNodeMetadata(
logical_len=16,
owned_positions=torch.tensor([4, 5, 6, 7], dtype=torch.int64),
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,
)
zero_visible_indices = controller.load_cp([zero_owned_node], node_id=301)
owned_visible_indices = controller.load_cp([owned_node], node_id=302)
controller.start_loading()
self.assertEqual(zero_visible_indices.tolist(), [64, 65, 66, 67])
self.assertEqual(owned_visible_indices.tolist(), list(range(80, 96)))
self.assertEqual(allocator.owner_alloc_calls, [[0], [3, 0, 1, 2]])
self.assertEqual(len(host_pool.loads), 2)
self.assertEqual([load[2] for load in host_pool.loads], [0, 1])
for host_indices, device_indices, _, _ in host_pool.loads:
self.assertEqual(host_indices.tolist(), [200, 201, 202, 203])
self.assertEqual(device_indices.tolist(), [24, 25, 26, 27])
self.assertEqual(len(controller.ack_load_queue), 1)
self.assertEqual(controller.ack_load_queue[0].node_ids, [301, 302])
def test_cp_load_allocates_full_logical_locs_and_transfers_owned_physical_locs(self):
host_pool = FakeHostPool(torch.tensor([100, 101, 102, 103], dtype=torch.int64))
allocator = FakeAllocator(alloc_result=torch.arange(64, 80, dtype=torch.int64))