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
This commit is contained in:
laoyao0822
@@ -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
|
||||
```
|
||||
|
||||
Reference in New Issue
Block a user