diff --git 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 new file mode 100644 index 000000000..64507a072 --- /dev/null +++ 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 层。 +