Record CPU-source limits for CP HiCache transfer

CP HiCache write-through keeps a host copy, but current disaggregation registers and addresses prefill GPU KV pages. The document records why CPU-source transfer needs a distinct source contract, host backup readiness, host-slot pinning, and CP logical-page mapping before it can safely release prefill GPU L1 pages early. Constraint: Current Mooncake KVArgs register GPU token_to_kv_pool buffers as the prefill transfer source Rejected: Reuse device page_indices as host page ids | host/L2 allocator uses separate physical slots and would corrupt decode destinations Confidence: medium Scope-risk: narrow Directive: Do not switch prefill-to-decode transfer to host source without explicit backup-ack and host-slot lifetime contracts Tested: Documentation-only change; inspected current prefill/decode/Mooncake/HiCache paths Not-tested: CPU-source transfer implementation (cherry picked from commit 01d08e3526ad95c0ab924e6ba859139615fc65c9)
2026-06-12 08:11:14 +08:00
parent e6ab699f8f
commit da8bd7f813
1 changed files with 496 additions and 0 deletions
--- a/docs/advanced_features/nsa_prefill_cp_hicache_cpu_based_transfer_exploration.md
+++ b/docs/advanced_features/nsa_prefill_cp_hicache_cpu_based_transfer_exploration.md
@@ -0,0 +1,496 @@
+# NSA Prefill CP HiCache CPU-based KV transfer 探索记录
+
+> 日期：2026-06-11  
+> 分支：`cjy-cp-refactor`  
+> 范围：CP shared-KV + HiCache + write-through 场景下，探索 prefill -> decode KV transfer 是否可以从 CPU/host HiCache 发起，而不是从 prefill GPU L1 KV pool 发起，以便更早释放 prefill GPU KV 显存。
+
+## 0. 结论
+
+CPU-based transfer **可行，但当前代码不能直接做到**。
+
+当前 disaggregation/Mooncake transfer 的源端注册和地址计算都绑定在 prefill GPU KV pool：
+
+```text
+prefill GPU L1 KV slot -> Mooncake transfer -> decode GPU KV slot
+```
+
+HiCache write-through 确实会把 KV 备份到 CPU pinned host pool，但这是另一个异步 D2H 流水线：
+
+```text
+prefill GPU L1 KV slot -> HiCache D2H backup -> host/L2 slot
+```
+
+要改成 CPU-based transfer，需要引入新的 source contract：
+
+```text
+host/L2 slot -> Mooncake transfer -> decode GPU KV slot
+```
+
+并且必须保证：
+
+1. host 数据已经完成 backup ack；
+2. host slot 在 transfer 完成前不能被 evict/free；
+3. source index 语义不能把 CP logical page、device physical page、host physical page 混用；
+4. target KV 和 draft/EAGLE KV 必须继续强一致；
+5. NSA state / index extra transfer 不能 silent fallback 或漏传。
+
+因此它更像是一个新的 transfer source mode，而不是把现有 `send_kvcache()` 换几个指针。
+
+## 1. 当前代码事实
+
+### F1. Prefill KV manager 注册的是 GPU device KV buffer
+
+`python/sglang/srt/disaggregation/prefill.py:319-365`
+
+```python
+kv_data_ptrs, kv_data_lens, kv_item_lens = (
+    self.token_to_kv_pool.get_contiguous_buf_infos()
+)
+...
+if self.draft_token_to_kv_pool is not None:
+    draft_kv_data_ptrs, draft_kv_data_lens, draft_kv_item_lens = (
+        self.draft_token_to_kv_pool.get_contiguous_buf_infos()
+    )
+    kv_data_ptrs += draft_kv_data_ptrs
+...
+kv_args.kv_data_ptrs = kv_data_ptrs
+```
+
+`token_to_kv_pool` 是 prefill GPU L1 KV pool。draft pool 也按同样方式追加 GPU KV buffer。
+
+结论：当前 Mooncake prefill source buffer 是 GPU KV，不是 HiCache host pool。
+
+### F2. Decode KV manager 注册的是 decode GPU KV buffer
+
+`python/sglang/srt/disaggregation/decode.py:378-426`
+
+decode 端同样通过：
+
+```python
+self.token_to_kv_pool.get_contiguous_buf_infos()
+```
+
+注册 decode GPU KV pool 作为目的端。
+
+结论：CPU-based transfer 只改变 prefill source；decode destination 仍然是 GPU KV pool。
+
+### F3. Mooncake engine 当前只注册 `kv_args.kv_data_ptrs`
+
+`python/sglang/srt/disaggregation/mooncake/conn.py:286-321`
+
+```python
+if self.kv_args.kv_data_ptrs and self.kv_args.kv_data_lens:
+    self.engine.batch_register(
+        self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens
+    )
+```
+
+结论：如果要让 Mooncake 从 host pool 读数据，需要把 host pool ptr/lens 也注册进 engine，或者引入新的 host-source KVArgs。
+
+### F4. Prefill `send_kv_chunk()` 传的是 device KV page indices
+
+`python/sglang/srt/disaggregation/prefill.py:1223-1373`
+
+```python
+page_indices = _kv_locs_to_page_indices_cpu(
+    self.req_to_token_pool.req_to_token[req.req_pool_idx, start_idx:end_idx],
+    page_size,
+)
+...
+req.disagg_kv_sender.send(page_indices, state_indices)
+```
+
+这里的 `req_to_token` 记录的是当前 request 在 GPU KV pool 中的 token slots。`page_indices` 因此是 GPU L1 page id。
+
+结论：CPU-based transfer 不能直接复用这个 `page_indices` 作为 host page id；必须显式记录 source type 和对应 index 映射。
+
+### F5. CP shared-KV transfer 当前按 logical request page 做过滤
+
+`python/sglang/srt/disaggregation/mooncake/conn.py:1673-1782`
+
+CP shared-KV 下：
+
+```python
+kv_indices, logical_page_positions = filter_kv_pages_for_cp_shared_kv(
+    layout=layout,
+    logical_pages=kv_indices,
+    chunk_page_start=index_slice.start,
+)
+```
+
+后续 decode worker 用 `logical_page_positions` 从 decode 分配好的 `dst_kv_indices` 中选目的页：
+
+`python/sglang/srt/disaggregation/mooncake/conn.py:1225-1233`
+
+```python
+chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.logical_page_positions]
+```
+
+结论：`logical_page_positions` 是 request 内 logical page 位置，不是源物理 page。CPU source 也必须保留这一层语义，否则 decode destination 会对错页。
+
+### F6. 当前 `send_kvcache()` 地址计算假设 source index 是 registered source buffer 的 page offset
+
+`python/sglang/srt/disaggregation/mooncake/conn.py:466-609`
+
+```python
+src_addr = src_ptr + int(prefill_index[0]) * item_len
+dst_addr = dst_ptr + int(decode_index[0]) * item_len
+length = item_len * len(prefill_index)
+```
+
+`send_kvcache()` 固定传入：
+
+`python/sglang/srt/disaggregation/mooncake/conn.py:661-679`
+
+```python
+src_data_ptrs=self.kv_args.kv_data_ptrs
+item_lens=self.kv_args.kv_item_lens
+prefill_data_indices=prefill_kv_indices
+```
+
+结论：CPU source 可以复用 `_send_kvcache_generic()` 的前提是 host layout 能表示成：
+
+```text
+per-layer base ptr + host_page_id * page_bytes
+```
+
+否则需要新的 address builder。
+
+### F7. HiCache write-through 已有 CP host reservation 与 per-layer D2H backup
+
+`python/sglang/srt/managers/cache_controller.py:998-1089`
+
+`reserve_write_cp()` 会生成：
+
+```text
+CpHiCacheNodeMetadata(
+  logical_len,
+  padded_len,
+  owned_positions,
+  host_indices,
+  page_owners,
+  draft_host_indices,
+)
+```
+
+并把 logical/device indices 映射为 CP rank-local physical device indices：
+
+```python
+owned_logical_indices = padded_device_indices[owned_mask]
+physical_device_indices = layout.logical_locs_to_physical(owned_logical_indices)
+host_indices = host_alloc(len(physical_device_indices))
+```
+
+`python/sglang/srt/managers/cache_controller.py:1268-1405`
+
+`_submit_write_cp_layer_states()` 会 batch 多个 reservation，在 `write_stream` 上执行：
+
+```python
+self.mem_pool_host.backup_from_device_per_layer(
+    self.mem_pool_device,
+    target_host_indices,
+    target_device_indices,
+    layer_id,
+    self.io_backend,
+)
+```
+
+draft pool 也走同样的 per-layer D2H backup。
+
+结论：write-through 已经提供 CPU source 所需的数据来源，但只有在最终 write ack 之后才是安全可读的完整 host copy。
+
+### F8. HiCache L2->L1 load 已有 host->device per-layer 路径
+
+`python/sglang/srt/managers/cache_controller.py:1524-1744`
+
+`load_cp()` 会根据 node metadata 的 `page_owners` 重新分配匹配 CP owner pattern 的 GPU slots，并把 `host_indices` / `physical_device_indices` 送入 load queue。
+
+`python/sglang/srt/managers/cache_controller.py:1770-1870`
+
+`start_loading()` 在 `load_stream` 上逐层调用：
+
+```python
+self.mem_pool_host.load_to_device_per_layer(
+    self.mem_pool_device,
+    host_indices,
+    device_indices,
+    i,
+    self.io_backend,
+)
+```
+
+结论：host pool 的 page 生命周期和 transfer descriptor 体系已经存在，可以复用一部分 metadata，但它当前服务的是 L2->L1 load，不服务 prefill->decode transfer。
+
+### F9. `page_first_direct` 不适合直接作为 `_send_kvcache_generic()` 的 CPU source
+
+`python/sglang/srt/mem_cache/memory_pool_host.py:1628-1658`
+
+MLA host layout：
+
+```text
+page_first_direct: [page, layer, page_size, 1, kv_cache_dim]
+layer_page_first: [layer, page, page_size, 1, kv_cache_dim]
+```
+
+当前 `_send_kvcache_generic()` 的地址模型是：
+
+```text
+src_layer_base + page_id * item_len
+```
+
+这天然匹配 `layer_page_first`，但不匹配 `page_first_direct` 的 per-layer view；`page_first_direct` 下固定 layer 的相邻 page 被其他 layer 隔开。
+
+结论：如果目标是 CPU source + per-layer Mooncake transfer，`layer_page_first` 是更自然的 host layout。`page_first_direct` 需要 page-whole transfer 或新的 strided/page-first address builder。
+
+## 2. 为什么 write-through 不等于“可以立即释放 GPU”
+
+write-through 的语义是“最终会写到 host”，不是“forward 结束立即有完整、安全、可被 decode 读取的 host 副本”。
+
+当前 per-layer backup 是异步的：
+
+1. forward 每层结束后触发 D2H；
+2. D2H 在 `write_stream` 上排队；
+3. 所有 target/draft layer 完成后 append write ack；
+4. radix/HiCache 的 `writing_check()` drain ack 后，节点才从 pending/in-flight 状态变成 backed 状态。
+
+所以 CPU-based transfer 的安全释放点至少是：
+
+```text
+layer data produced
+-> D2H backup completed
+-> host slot pinned for PD transfer
+-> decode transfer no longer needs prefill GPU source
+-> GPU token slot may be released/reused
+```
+
+如果只等 forward 完成而不等 host backup ack，CPU transfer 可能读到未写完的 host page。
+
+## 3. 关键设计分支
+
+### 方案 A：只对已 backed/cache-hit prefix 使用 host source
+
+适用场景：
+
+```text
+request 命中 L2/host prefix
+prefix node 已有 cp_hicache metadata
+host slot 当前 resident 且未被 evict
+```
+
+行为：
+
+1. prefix pages 从 host pool 直接 transfer 到 decode GPU；
+2. current/new extend pages 仍走当前 GPU source transfer；
+3. 不等待新 write-through backup。
+
+优点：
+
+- 风险最低；
+- 可以验证 Mooncake CPU pinned memory source 是否可行；
+- 对 cache-hit prefix 大的请求能减少 prefill GPU source 依赖。
+
+缺点：
+
+- 对新生成的 current pages 不能立即释放 GPU；
+- 需要一次 transfer chunk 中支持 mixed source，或者拆成 host-prefix chunk + gpu-current chunk；
+- draft/state 仍需同步补齐。
+
+### 方案 B：完整 write-through host-source transfer
+
+适用场景：
+
+```text
+所有需要发给 decode 的 KV page 都先写入 host/L2
+PD transfer 只从 host/L2 source 发起
+```
+
+行为：
+
+1. forward / per-layer backup 正常运行；
+2. transfer worker 等待对应 node 的 host backup ack；
+3. host slots 加 transfer ref/pin，防止 evict/free；
+4. Mooncake 从 registered host ptr + host page indices 发起 transfer；
+5. transfer 完成后释放 host transfer ref；
+6. GPU L1 page 可以在 host backup 完成后提前 release，不必等 RDMA 从 GPU 读完。
+
+优点：
+
+- 符合“write-through 后让 CPU 承担 prefill->decode source”的目标；
+- 能真正缩短 prefill GPU KV slot 占用周期；
+- 可以把 decode transfer 与 GPU compute/allocator 压力解耦。
+
+缺点：
+
+- 控制面复杂度明显增加；
+- 新增 host slot transfer pin/ref 生命周期；
+- 需要处理 backup ack、split/evict、draft、state、chunked prefill 的一致性；
+- raw latency 可能比 GPUDirect GPU->GPU 慢，需要靠 overlap 和显存释放收益抵消。
+
+### 方案 C：per-layer streaming CPU transfer
+
+行为：
+
+```text
+layer k forward done
+-> layer k D2H backup done
+-> layer k host->decode transfer
+```
+
+优点：
+
+- overlap 空间最大；
+- transfer 可以尽早开始。
+
+缺点：
+
+- 当前 GPU KV allocator 是 token/page slot 粒度，一个 token slot 覆盖所有 layer；
+- 即使 layer k 已经传走，也无法只释放该 token slot 的第 k 层显存；
+- 真正做到 per-layer free 需要改 device KV pool / allocator 语义，风险大。
+
+结论：第一阶段不建议做 C。可以做 per-layer host->decode overlap，但不要承诺 per-layer 释放 GPU token slot。
+
+## 4. 推荐路线
+
+### P0：只做能力验证，不改线上默认
+
+目标：确认 Mooncake engine 能否稳定从 CPU pinned host memory 作为 source 传到 decode GPU。
+
+验证内容：
+
+- register host `kv_buffer` ptr/lens；
+- 单 rank / 单 layer / 少量 page host->decode GPU；
+- 校验 decode GPU bytes；
+- 对比 GPU source 与 host source latency/bandwidth；
+- 验证 host pinned memory 的 NUMA/NIC 亲和影响。
+
+### P1：实现 host-source transfer descriptor，不接调度
+
+新增概念：
+
+```text
+TransferSourceType:
+  GPU_DEVICE_KV
+  HOST_HICACHE_KV
+```
+
+新增 metadata：
+
+```text
+source_type
+source_page_indices
+logical_page_positions
+source_pool_kind: target/draft/state
+```
+
+要求：
+
+- source page id 必须显式标注是 GPU page 还是 host page；
+- CP logical positions 继续用于 decode destination 选择；
+- 不允许 silent fallback。
+
+### P2：prefix-only host source
+
+先只支持已 backed prefix：
+
+```text
+host-backed prefix -> host source transfer
+current/new pages -> current GPU source transfer
+```
+
+这是最适合暴露 source index / draft / state 错配的阶段，因为新写入路径不参与。
+
+### P3：full write-through host source
+
+在 P2 稳定后支持：
+
+```text
+new/current pages wait backup ack -> host source transfer
+```
+
+需要新增：
+
+- host slot transfer ref/pin；
+- pending split 与 in-flight host transfer 的关系；
+- transfer completion 后 release pin；
+- backup ack 等待的限频/非 collective 设计；
+- host evict fail-fast 合同。
+
+### P4：GPU slot early release
+
+只有 P3 正确后才考虑释放策略：
+
+```text
+backup ack complete + host transfer source pinned
+=> prefill GPU source 不再被 PD transfer 依赖
+=> 可以 release L1 pages
+```
+
+注意：这只能按 token/page slot 粒度释放，不能按 layer 粒度释放，除非未来重构 device KV allocator。
+
+## 5. 必须避免的错误
+
+### E1. 混淆 source page id 与 logical page position
+
+当前 CP transfer 有两套 id：
+
+```text
+source page id:       用于 source address = base + source_page * item_len
+logical page position: 用于 dst_pages = req.dst_kv_indices[logical_page_positions]
+```
+
+CPU source 会引入第三套：
+
+```text
+host physical page id
+```
+
+任何隐式复用 `kv_indices` 都容易造成串页。
+
+### E2. host slot 未 pin 就发起 transfer
+
+如果 host evict/free 在 transfer 完成前复用同一 host page，decode 会收到错误 KV。
+
+需要 transfer ref/pin，而不是只依赖 radix node 存在。
+
+### E3. backup pending 时读取 host
+
+write-through backup 是异步 per-layer D2H。读取未 ack 的 host page 等价于读未完成 DMA。
+
+### E4. 忽略 draft/EAGLE mirror
+
+当前 prefill KV manager 会把 draft GPU KV buffer 追加进 `kv_args`，说明 PD transfer 已经包含 draft KV。CPU source 也必须包含 draft host pool，且 target/draft 的 page id 生命周期必须一致。
+
+### E5. 忽略 NSA state/index extra
+
+`send_kv_chunk()` 在 NSA pool 下还会准备 `state_indices`，transfer worker 在 final chunk 里通过 `maybe_send_extra()` 发送 state/extra pool。CPU source 如果只改主 KV，不处理 state/extra，可能出现主 KV 正确但 NSA state 错误的隐性精度问题。
+
+### E6. 在 `page_first_direct` 上强行套 per-layer base+page address
+
+`page_first_direct` 的 per-layer page 不连续。要么：
+
+1. 切 `layer_page_first`；
+2. 新增 page-first host address builder；
+3. 按 page-whole transfer，但这会传所有 layer，和 per-layer transfer 目标冲突。
+
+## 6. 初步判断
+
+CPU-based transfer 对“立即空出 GPU 显存”的帮助取决于采用到哪一层：
+
+| 阶段 | 能否更早释放 prefill GPU KV | 风险 | 说明 |
+|---|---:|---:|---|
+| P2 prefix-only host source | 部分 | 中 | 对已 backed prefix 有帮助；current pages 仍依赖 GPU source |
+| P3 full host source after backup ack | 是 | 高 | transfer 不再读 GPU，但要等 host backup 完成 |
+| P4 early release | 是 | 高 | 需要严格证明 decode transfer 不再依赖 GPU page |
+| C per-layer free | 暂不现实 | 很高 | 当前 allocator 不是 layer 粒度 |
+
+推荐先做 P0/P1/P2。等确认 host source transfer 的吞吐、正确性、host pin 生命周期后，再进入 P3/P4。
+
+## 7. 需要补充验证的问题
+
+1. Mooncake 当前 `batch_register()` 对 CPU pinned host memory 作为 RDMA source 的支持和性能。
+2. host source 在 8 CP ranks 并发时是否会被 CPU NUMA / PCIe / NIC 拓扑限制。
+3. `layer_page_first` 下 host ptr/lens 如何注册才能和 `_send_kvcache_generic()` 的 per-layer address builder 对齐。
+4. draft host pool 是否已经具备完整 target mirror metadata，可直接作为 source。
+5. NSA state/index extra 是否需要独立 host source，还是第一阶段继续 GPU source 并显式记录限制。
+6. host transfer pin/ref 应该挂在 radix node、CpHiCacheNodeMetadata，还是 host allocator 层。
+