Batch CP HiCache backup submits across requests

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

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
+```

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(

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