diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md index 5d9417f7d..d237e9894 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md @@ -745,3 +745,149 @@ runtime / kernel 都消费 CPSharedKVBatchPlan descriptors 4. **MLA current/partial-current 和 direct store 已有 flatten kernel 基础，但还需用 bs>1 ETE 验证没有 fallback hot path。** - 特别关注 FP8/BF16 两种 dtype、draft/EAGLE 路径、prefetch 关闭时 full/current reuse 是否仍启用。 + +--- + +## 16. 2026-06-03 W6 HiCache load/backup batch 化设计补充 + +### 当前结论 + +1. **L2->L1 load-back 的传输层已经接近 batch 化。** + - `HiRadixCache.load_back()` 仍按 request 调 `cache_controller.load_cp()`，但 `HiCacheController.start_loading()` 会对 `load_queue` / `draft_load_queue` 做 `CacheOperation.merge_ops()`。 + - 因此只要 scheduler 允许同一个 prefill batch 内多个 request 都在 `ready_to_load_host_cache()` 前完成 `init_load_back()`，目标 KV 和 draft KV 的 H2D per-layer transfer 会按 layer 合并成一次 `load_to_device_per_layer()` 调用。 + - 已有低风险验证点：补一个单测，先连续调用两次 `load_cp()`，再调用一次 `start_loading()`，断言每层只有一次 target load、一次 draft load，且 host/device indices 为两个 request 的 concat。 + +2. **D2H backup 仍没有吃到 bs>1 的收益。** + - `scheduler._prepare_hicache_write_backups_before_forward()` 当前仍然 `for req in batch.reqs: prepare_write_backup_for_req(req)`，这是正确的 per-request reservation/metadata 生命周期。 + - 真正的问题在 `HiCacheController.on_layer_end()`：它遍历 `pending_layer_writes`，对每个 reservation 单独调用 `submit_write_cp_layer()`。 + - `submit_write_cp_layer()` 又对每个 reservation / layer 分别调用 `mem_pool_host.backup_from_device_per_layer()` 和 draft host pool backup。 + - 底层 TAI direct PF/LF transfer kernel 接收 flatten indices，本身可以处理多个 request 的 concat descriptor；当前没有收益是 runtime 没有把多个 reservation 合并为 layer-level descriptor。 + +3. **scheduler 仍有 CP bs>1 admission gate。** + - `schedule_policy.py` 里仍有 `# TODO support cp with multiple requests`，当 `nsa_prefill_cp_in_seq_split` 或 `prefill_context_parallel_enabled` 打开时，`len(self.can_run_list) >= 1` 会拒绝继续加 request。 + - 所以 W6 可以先把 load/backup runtime 改成 batch-ready 并用 unit test 锁住；真正 ETE 吃到收益还依赖后续移除/替换 admission gate。 + +### 推荐实现：保持 radix/metadata per request，传输按 layer 聚合 + +目标不是把 HiCache node 合并，而是只合并底层 transfer descriptor： + +- reservation、host slot、draft host slot、`CpHiCacheNodeMetadata` 仍每个 request 独立； +- radix tree attach/rollback/ack 仍按 node_id 独立； +- layer-end hook 对同一 source/同一 layer 下所有 pending reservation 做一次 concat transfer； +- final ack 可以继续 per node append，避免改变 `writing_check()` 的队列语义； +- 不新增 collective，不改变 owner-lane 分布语义。 + +### W6A：为现有 load-back batch 行为补测试 + +文件：`test/registered/unit/managers/test_hicache_controller_cp.py` + +新增测试目标： + +1. `load_cp(req_a)` + `load_cp(req_b)` + 单次 `start_loading()`： + - `host_pool.loads` 每个 layer 只有一次； + - 传入的 host/device indices 是两个 request 的 concat； + - `ack_load_queue[0].node_ids == [node_a, node_b]`； + - draft host pool 同样是每层一次。 +2. zero-owned request 参与 batch 时仍不触发实际 host load，但 node_id 保留在 ack 中。 + +### W6B：backup layer-level grouping + +文件：`python/sglang/srt/managers/cache_controller.py` + +计划改造： + +1. 抽出 helper：`_submit_write_cp_layer_states(states, layer_id, submit_target, submit_draft)`。 +2. `submit_write_cp_layer(reservation, ...)` 保持公开 API，但内部只传单个 state 给 helper，保证 bs=1 行为不变。 +3. `on_layer_end(layer_id, source)` 不再对每个 reservation 单独 submit，而是收集当前 pending states，一次调用 helper。 +4. helper 内部： + - 计算每个 state 是否需要 target/draft； + - 记录一个 layer event，并挂到所有参与 states 上保活； + - 在 `write_stream` 上等待每个 state 的 `start_event` 和 group layer event； + - 对 target states concat `state.host_indices` / `state.physical_device_indices`，调用一次 target `backup_from_device_per_layer()`； + - 对 draft states concat `state.draft_host_indices` / 同一 physical indices，调用一次 draft backup； + - 标记每个 state 的 target/draft layer completed； + - 对完成全部层的 state 逐个 record finish_event、record_stream、append ack。 + +### W6B 单测要求 + +文件：`test/registered/unit/managers/test_hicache_controller_cp.py` + +1. 两个 target-only reservation 同一 layer-end： + - `allocator.device_pool.notify_layer_end_for_backup(0)` 后 `host_pool.layer_backups` 只新增一条； + - indices 为两个 reservation concat； + - 未完成所有 layer 前 ack 为空； + - 最后一层后两个 node_id 都释放为 ack。 +2. target + draft： + - target source layer-end 只备 target；draft source layer-end 只备 draft； + - 每个 source 每层只产生一次 backup call； + - final ack 等 target/draft 都完成。 +3. zero-owned reservation： + - 不增加实际 backup call； + - 仍按层完成并最终 ack。 +4. 直接调用 `submit_write_cp_layer()` 的旧路径仍保持单 reservation 行为。 + +### W6C：后续优化，不作为第一步 + +1. `load_cp()` 的 owner-lane allocation 仍按 request 调用；如果 admission 放开后 CPU overhead 明显，再做 batch owner allocation planner。 +2. `CacheOperation.merge_ops()` 仍会 `torch.cat` 二次 concat；如果 profile 显示 CPU cat 明显，再引入预分配 descriptor workspace。 +3. backup grouping 首版不把 ack 合并成单个 `HiCacheAck`，避免影响 `writing_check()` 与 rollback unattached prepared backup 的语义。 +4. 不引入新的 all-reduce/all-gather；W6 只改变 descriptor 聚合粒度。 + +### W6A 完成状态 + +已补 characterization tests，锁住现有 L2->L1 load-back batch 行为： + +- `test_cp_start_loading_batches_multiple_load_cp_requests_with_draft` + - 两个 `load_cp()` request 在一次 `start_loading()` 中合并； + - target 每层一次 `load_to_device_per_layer()`；draft 每层一次； + - host/device descriptors 为两个 request 的 concat； + - ack 保留两个 node_id。 +- `test_cp_start_loading_keeps_zero_owned_load_ack_in_batched_load` + - zero-owned request 不产生实际 host transfer； + - 但 zero-owned node_id 仍保留在 batched load ack 中。 + +远端验证： + +```text +PYTHONPATH=python python -m pytest -q \ + test/registered/unit/managers/test_hicache_controller_cp.py::TestHiCacheControllerCPLoad::test_cp_start_loading_batches_multiple_load_cp_requests_with_draft \ + test/registered/unit/managers/test_hicache_controller_cp.py::TestHiCacheControllerCPLoad::test_cp_start_loading_keeps_zero_owned_load_ack_in_batched_load +=> 2 passed, 3 warnings +``` + +本地限制：本地 pytest 仍因缺少 `orjson` 无法收集该测试文件；本地只做了 `py_compile`。 + +### W6B 完成状态 + +已实现 CP HiCache per-layer backup 的 layer-level grouping： + +- `HiCacheController.on_layer_end()` 不再逐 reservation 调 `submit_write_cp_layer()`；现在对当前 pending states 做一次 grouped submit。 +- reservation、radix metadata、host slot、draft host slot、ack node_id 仍保持 per request。 +- 每个 layer/source 的底层 transfer descriptor 会按 pending reservation concat： + - target source：一次 target `backup_from_device_per_layer()`； + - draft source：一次 draft `backup_from_device_per_layer()`； + - zero-owned state 不参与实际 transfer，但继续按 layer 完成并最终 ack。 +- `submit_write_cp_layer()` 公开 API 保留，内部走单 state grouped helper，bs=1 兼容路径不变。 +- 没有新增 collective；没有修改 owner-lane 语义；没有修改 tai-kernel。TAI direct transfer kernel 已支持 flatten indices，runtime concat 后即可吃到 bs>1 descriptor 收益。 + +新增测试： + +- `test_cp_layer_hook_groups_target_backups_across_pending_reservations` +- `test_cp_layer_hook_groups_target_and_draft_backups_by_source` +- `test_cp_layer_hook_keeps_zero_owned_ack_in_grouped_backup` + +远端验证： + +```text +PYTHONPATH=python python -m pytest -q \ + test/registered/unit/managers/test_hicache_controller_cp.py +=> 75 passed, 3 warnings +``` + +补充验证： + +```text +PYTHONPATH=python python -m pytest -q \ + test/registered/unit/mem_cache/test_cp_hicache_load_back_owner_lanes.py +=> 10 passed, 3 warnings +``` diff --git a/python/sglang/srt/managers/cache_controller.py b/python/sglang/srt/managers/cache_controller.py index f238531ae..bde98a514 100644 --- a/python/sglang/srt/managers/cache_controller.py +++ b/python/sglang/srt/managers/cache_controller.py @@ -452,18 +452,12 @@ class HiCacheController: return if source not in ("target", "draft"): raise ValueError(f"Unknown CP HiCache layer backup source: {source}") - reservations = [ - state.reservation for state in list(self.pending_layer_writes.values()) - ] - for reservation in reservations: - if reservation.node_id not in self.pending_layer_writes: - continue - self.submit_write_cp_layer( - reservation, - layer_id, - submit_target=(source == "target"), - submit_draft=(source == "draft"), - ) + self._submit_write_cp_layer_states( + list(self.pending_layer_writes.values()), + layer_id, + submit_target=(source == "target"), + submit_draft=(source == "draft"), + ) def _start_storage_threads(self): """Start storage prefetch/backup threads and their queues. @@ -1148,6 +1142,193 @@ class HiCacheController: state.start_event.record() return state + def _layer_write_state_done(self, state: HiCacheLayerWriteState) -> bool: + target_done = len(state.completed_target_layers) >= self.layer_num + draft_done = ( + state.draft_host_indices is None + or len(state.completed_draft_layers) + >= self.draft_mem_pool_device.layer_num + ) + return target_done and draft_done + + @staticmethod + def _record_tensor_on_stream(tensor: Optional[torch.Tensor], stream) -> None: + if tensor is not None and tensor.is_cuda: + tensor.record_stream(stream) + + @staticmethod + def _concat_layer_write_tensors(tensors: List[torch.Tensor]) -> torch.Tensor: + if len(tensors) == 1: + return tensors[0] + return torch.cat(tensors) + + def _append_layer_write_ack(self, state: HiCacheLayerWriteState) -> None: + if state.ack_appended: + return + state.ack_appended = True + reservation = state.reservation + self.pending_layer_writes.pop(reservation.node_id, None) + self.ack_write_queue.append( + HiCacheAck(state.start_event, state.finish_event, [reservation.node_id]) + ) + logger.info( + "[CacheCtrl-write] submit_write_cp_layer final ack: node_id=%d logical_len=%d layers=%d draft=%s", + reservation.node_id, + reservation.metadata.logical_len, + state.total_layers, + reservation.draft_host_indices is not None, + ) + + def _submit_write_cp_layer_states( + self, + states: List[HiCacheLayerWriteState], + layer_id: int, + *, + submit_target: bool = True, + submit_draft: bool = True, + ) -> None: + """Submit one layer for a group of reserved CP host backups. + + Radix/HiCache metadata remains per node. Only the transfer descriptor is + grouped so bs>1 batches issue one target D2H and one draft D2H per layer + instead of one D2H per request per layer. + """ + + if not states: + return + + unique_states: List[HiCacheLayerWriteState] = [] + seen_state_ids = set() + for state in states: + state_id = id(state) + if state_id in seen_state_ids: + continue + seen_state_ids.add(state_id) + unique_states.append(state) + + for state in unique_states: + if layer_id < 0 or layer_id >= state.total_layers: + raise ValueError( + f"layer_id={layer_id} is outside CP HiCache backup layer range " + f"[0, {state.total_layers}) for node_id={state.reservation.node_id}" + ) + + target_states = [ + state + for state in unique_states + if submit_target + and layer_id < self.layer_num + and layer_id not in state.completed_target_layers + ] + draft_states = [ + state + for state in unique_states + if submit_draft + and state.draft_host_indices is not None + and layer_id < self.draft_mem_pool_device.layer_num + and layer_id not in state.completed_draft_layers + ] + if not target_states and not draft_states: + return + + active_states: List[HiCacheLayerWriteState] = [] + seen_active_ids = set() + for state in target_states + draft_states: + state_id = id(state) + if state_id in seen_active_ids: + continue + seen_active_ids.add(state_id) + active_states.append(state) + + layer_event = device_module.Event() + layer_event.record() + for state in active_states: + state.layer_events.append(layer_event) + + target_transfer_states = [ + state + for state in target_states + if state.physical_device_indices is not None + and len(state.physical_device_indices) > 0 + ] + draft_transfer_states = [ + state + for state in draft_states + if state.physical_device_indices is not None + and len(state.physical_device_indices) > 0 + ] + grouped_tensors: List[torch.Tensor] = [] + final_states: List[HiCacheLayerWriteState] = [] + + with device_module.stream(self.write_stream): + for state in active_states: + state.start_event.wait(self.write_stream) + layer_event.wait(self.write_stream) + + if target_transfer_states: + target_host_indices = self._concat_layer_write_tensors( + [state.host_indices for state in target_transfer_states] + ) + target_device_indices = self._concat_layer_write_tensors( + [ + state.physical_device_indices + for state in target_transfer_states + ] + ) + self.mem_pool_host.backup_from_device_per_layer( + self.mem_pool_device, + target_host_indices, + target_device_indices, + layer_id, + self.io_backend, + ) + grouped_tensors.extend([target_host_indices, target_device_indices]) + + if draft_transfer_states: + draft_host_indices = self._concat_layer_write_tensors( + [state.draft_host_indices for state in draft_transfer_states] + ) + draft_device_indices = self._concat_layer_write_tensors( + [ + state.physical_device_indices + for state in draft_transfer_states + ] + ) + self.draft_mem_pool_host.backup_from_device_per_layer( + self.draft_mem_pool_device, + draft_host_indices, + draft_device_indices, + layer_id, + self.io_backend, + ) + grouped_tensors.extend([draft_host_indices, draft_device_indices]) + + for tensor in grouped_tensors: + self._record_tensor_on_stream(tensor, self.write_stream) + + for state in target_states: + state.completed_target_layers.add(layer_id) + for state in draft_states: + state.completed_draft_layers.add(layer_id) + + final_states = [ + state + for state in active_states + if self._layer_write_state_done(state) and not state.ack_appended + ] + for state in final_states: + state.finish_event.record() + self._record_tensor_on_stream(state.host_indices, self.write_stream) + self._record_tensor_on_stream( + state.physical_device_indices, self.write_stream + ) + self._record_tensor_on_stream( + state.draft_host_indices, self.write_stream + ) + + for state in final_states: + self._append_layer_write_ack(state) + def submit_write_cp_layer( self, reservation: HiCacheWriteReservation, @@ -1163,92 +1344,11 @@ class HiCacheController: """ state = self._get_or_create_layer_write_state(reservation) - if layer_id < 0 or layer_id >= state.total_layers: - raise ValueError( - f"layer_id={layer_id} is outside CP HiCache backup layer range " - f"[0, {state.total_layers}) for node_id={reservation.node_id}" - ) - needs_target = ( - submit_target - and layer_id < self.layer_num - and layer_id not in state.completed_target_layers - ) - needs_draft = ( - submit_draft - and state.draft_host_indices is not None - and layer_id < self.draft_mem_pool_device.layer_num - and layer_id not in state.completed_draft_layers - ) - if not needs_target and not needs_draft: - return - - layer_event = device_module.Event() - layer_event.record() - state.layer_events.append(layer_event) - - if len(reservation.physical_device_indices) > 0: - with device_module.stream(self.write_stream): - state.start_event.wait(self.write_stream) - layer_event.wait(self.write_stream) - if needs_target: - self.mem_pool_host.backup_from_device_per_layer( - self.mem_pool_device, - state.host_indices, - state.physical_device_indices, - layer_id, - self.io_backend, - ) - if needs_draft: - self.draft_mem_pool_host.backup_from_device_per_layer( - self.draft_mem_pool_device, - state.draft_host_indices, - state.physical_device_indices, - layer_id, - self.io_backend, - ) - - if needs_target: - state.completed_target_layers.add(layer_id) - if needs_draft: - state.completed_draft_layers.add(layer_id) - - target_done = len(state.completed_target_layers) >= self.layer_num - draft_done = ( - state.draft_host_indices is None - or len(state.completed_draft_layers) - >= self.draft_mem_pool_device.layer_num - ) - if not target_done or not draft_done: - return - if state.ack_appended: - return - - with device_module.stream(self.write_stream): - state.finish_event.record() - if state.host_indices is not None and state.host_indices.is_cuda: - state.host_indices.record_stream(self.write_stream) - if ( - state.physical_device_indices is not None - and state.physical_device_indices.is_cuda - ): - state.physical_device_indices.record_stream(self.write_stream) - if ( - state.draft_host_indices is not None - and state.draft_host_indices.is_cuda - ): - state.draft_host_indices.record_stream(self.write_stream) - - state.ack_appended = True - self.pending_layer_writes.pop(reservation.node_id, None) - self.ack_write_queue.append( - HiCacheAck(state.start_event, state.finish_event, [reservation.node_id]) - ) - logger.info( - "[CacheCtrl-write] submit_write_cp_layer final ack: node_id=%d logical_len=%d layers=%d draft=%s", - reservation.node_id, - reservation.metadata.logical_len, - state.total_layers, - reservation.draft_host_indices is not None, + self._submit_write_cp_layer_states( + [state], + layer_id, + submit_target=submit_target, + submit_draft=submit_draft, ) def submit_write_cp_per_layer( diff --git a/test/registered/unit/managers/test_hicache_controller_cp.py b/test/registered/unit/managers/test_hicache_controller_cp.py index 7199ebc52..447bc5c1a 100644 --- a/test/registered/unit/managers/test_hicache_controller_cp.py +++ b/test/registered/unit/managers/test_hicache_controller_cp.py @@ -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))