AI

Rapid transfer_kv：只为更快的 PD 分离

Posted by w@hidva.com on March 27, 2026

关于目前线上 PD 分离实现, 我之前其实已经零散写过不少:

但这些文章里, 对 transfer_kv 这个控制面动作本身讲得都不算特别细. 最近正好把这条链路又往前推了一大步, 于是准备单独写一篇, 讲下我们是如何一步步让 transfer_kv 变得更及时的.

先补一点背景知识

如果之前完全没接触过 HybridConnector / blade-kvt, 可以先建立一个足够粗糙, 但够用的心智模型. HybridConnector 运行在 vLLM 这一层, 它不是传输内核, 更像一个 “异步 KV 搬迁运行时”. 它负责:

  • 请求生命周期管理,
  • block refcnt,
  • abort / bypass,
  • 把 backend 的异步状态翻译成 scheduler / worker 可以消费的 metadata.

blade-kvt 则是更靠近数据面的那一层. 对于 P 节点来说, 正常情况下上层会:

  1. 先用 submit_req_send2(...) / submit_delta_send(...) 提交本轮待发送的请求和 token 范围,
  2. 再调用 start_send_step() 开启当前 step 的发送窗口,
  3. 在每层 attention 计算完成之后调用 record_event(layer_idx),
  4. 底层后台线程等 layer ready 之后按层发送,
  5. 最后通过 SEND_DONE_REQ 告诉上层 “这个请求我真的不再访问它的 kvcache 了”.

也就是说, blade-kvt 的正常增量发送路径是强依赖 step 和 layer event 的. 而 transfer_kv 本身其实是一个控制面动作. 它的意思非常简单:

目标端已经把 block 和 worker 信息准备好了, 源端你现在可以开始往我这里搬 KV 了.

它最终当然会变成一次真实的数据发送, 但它自己并不等于发送本身.

transfer_kv 的使用场景

在早期单请求 remote prefill 模式里, D 端会直接把整个请求交给 P, 然后通过 PREFILL_REQ 完成 “prefill + kvcache 传输 + first token 返回”. 那时还不太需要单独的 transfer_kv.

后来为了更灵活的调度, 我们引入了双请求模式. transfer_kv 最初其实是为了双请求调度引入的. 在双请求模式中请求 R 会同时发给 P, D 两侧:

  1. P 侧收到 R 之后立刻开始 prefill.
  2. D 侧收到 R 之后先分配目标端 kvcache block.
  3. D 侧构造好 DInfo（包含了 D block ids, instance id 等信息），再通过 TRANSFER_KV_REQ 把这些目标端信息发给 P.
  4. P 侧在合适的时机开始把已经算出来的 kvcache 传给 D.

后来做请求迁移时我又发现, 这个接口本质上只是: 告诉源端, 目标端 ready 了, 你现在可以开始搬 KV. 所以迁移链路也完全可以复用 transfer_kv.

也正因为如此, transfer_kv 的及时性非常重要. 在双请求模式里, 它影响 D 端什么时候能尽快接上 decode; 在迁移场景里, 它影响请求什么时候能在新节点真正恢复执行. 如果这个动作总要晚一拍, 整条链路的收益就会被吃掉不少.

最初实现: 严格对齐到 step

最初的实现很自然: 让 transfer_kv 严格对齐到正常 step.

P 侧 TRANSFER_KV_REQ 到达之后, 并不会立刻去直接调用 kvt 发送接口. 它只是先把 DInfo 放进 PBackend 的状态里. 后续 PBackend.build_backend_meta() 在下次正常 step 被调用时, 才会根据当前 SchedulerOutput、收到的 DInfo 等信息生成 PReqMeta; worker 再把这些 meta 翻译成 submit_req_send2(...) / submit_delta_send(...) 调用; 最后 start_send_step() 打开这一轮 step 的发送窗口. 这条链路的好处很明显:

  • 逻辑顺,
  • 正确性边界清楚,
  • 完全贴合 blade-kvt 现有的 step / layer-ready 抽象.

但它也有个同样明显的问题: transfer_kv 只能在 step 边界上生效. 也就是说, 如果 P 正在执行某个 step, 这时 D 端或者迁移控制面发来一个 TRANSFER_KV_REQ, P 虽然已经知道目标端 ready 了, 但这次 transfer_kv 没法立刻转化成一次真正的发送任务. 因为:

  • 当前 step 的 metadata 已经下发给 worker 了,
  • 当前 step 的 Step 已经建立了,
  • 当前 step 的发送仍然严格依赖 worker 主线程里的 record_event().

结果就是, 这次 transfer_kv 往往只能等当前 step 结束, 在下一个 step 的 build_backend_meta() 里才真正被看到. 这就是最早版本的语义: transfer_kv 对齐到 step.

第一次优化: bypass step

后来我们引入了 bypass 小循环. 其基本想法是: 在 engine core 等待 gpu worker execute_model/sample_tokens 返回时, 不要只是傻等, 而是趁这个空档继续让 connector 往前推进. 具体来说, core 侧会在等待期间反复执行一个小循环:

  • kvconn.step()
  • build_connector_meta(SchedulerOutput.make_empty())
  • send_bypass_task(meta)

然后把 bypass metadata 直接广播给 worker. worker 侧的 bypass 线程收到 metadata 之后, 会走 _do_bypass_meta(), 也就是:

  • bind_backend_metadata(meta.reqs)
  • clear_backend_metadata()

对于 load 型 backend, 还可以进一步启动 async_load_kv(). 对 kvt 来说, bypass 最安全的路径是 full-send / freeze-send. 因为 start_req_send() 本身就是 thread-safe 的, 而且底层会直接创建一个 “所有 layer 已 ready” 的 Step, 完全不依赖当前 forward 的 layer event.

这一版优化首先解决的是 finished req. 对于已经 finished 的请求, 或者双请求/迁移这种 “待发送的 KV cache 已经计算完毕” 的场景, 如果在 step 执行期间收到了 transfer_kv, 那么这次传输不必再等当前 step 结束, 也不必等下一次真实调度, 而是可以在 bypass 小循环里立刻触发 full-send.

这个收益已经很可观了, 但问题还没有完全解决. 因为 bypass step 只安全地覆盖了 full-send. 对 unfinished req 来说, 我们想发送的并不是 “补发一整份 KV”, 而是:

当前这个还在 forward 的请求, 已经算完的那部分 KV, 能不能尽快送出去?

这种发送本质上还是普通 step 内增量发送. 它依赖:

  • 当前 step 的 layer-ready event,
  • worker 主线程中已经建立好的那个 Step,
  • 以及 record_event() 的原有时序.

所以 unfinished req 仍然会卡在老问题上: step 中途收到 transfer_kv, 还是得等 step 结束, 再在下一个 step 才开始真正传输.

第二次优化: substep

真正把 unfinished req 也拉快的, 是 substep. substep 的核心思路其实不复杂:

既然普通增量发送强依赖当前 step, 那就不要在 bypass 线程里重新造一个新 step, 而是把新的发送任务追加到当前正在执行的那个 step 上.

也就是说, bypass 线程并不拥有自己独立的一条数据面时间线. 它只是获得了一种能力: 可以把新来的发送任务, 重新塞回当前主时间线.

start_send_substep()

首先映入眼帘的是 blade-kvt 这一侧的改造. 以前正常发送路径大概是这样:

  1. submit_req_send2(...) / submit_delta_send(...)
  2. start_send_step()
  3. record_event(layer_idx)
  4. flush_send_step()

一个 step 对应一个 Step, 这事很清楚. 现在为了支持 substep, blade-kvt 新增了 start_send_substep(stepid, substepid, metas) 接口. 它的语义不是 “创建一个新的 step”, 而是:

向当前 step 附加新的发送任务.

这背后其实有几个挺重要的约束.

  1. 真实 step 仍然由 worker 主线程创建，start_send(stepid, sched_tokens) 仍然是创建真实 Step / StepGuard 的入口. 这条路径没变. 换句话说, 负责建立当前 step 时间线的人, 仍然是 worker 主线程.

  2. bypass 线程只能往当前 step 上追加任务，如果 bypass 小循环在 step 执行期间收到了 unfinished req 的 transfer_kv, worker 侧 bypass 线程不会去调用 start_req_send() 另起炉灶, 而是调用 start_send_substep(stepid, substepid, metas)，把这批 nonfreeze_metas 追加到当前 step 上.

  3. 早到和晚到都要处理，这个接口在 blade-kvt C++ 里做了三种分支处理:

    • 如果 stepid < 当前 coord_step_id, 说明这批 substep 来晚了, 当前主线程已经去到更后面的 step, 直接丢弃.
    • 如果 stepid > 当前 coord_step_id, 说明 bypass 比主线程还快, 当前 step 还没真正建立起来, 那就先把 metas 暂存在 pending_step_metas_ 里, 等之后主线程真正 start_send(stepid, ...) 时再统一附加.
    • 如果 stepid == 当前 coord_step_id 且 step 还没 freeze, 那就直接把这批 metas 转成 StepTasks, 挂到当前 step 上.

    所以它不是一个简单的 “提交任务” 接口, 而是一个带时序协调语义的接口.

  4. substep 当前只处理一类更容易守正确性的 meta，这里还有一个很务实的小约束: 当前 start_send_substep() 处理的 nonfreeze_metas, 都要求 has_last_token=True.

    原因不复杂. 在 blade-kvt 里, 如果 has_last_token=False, 往往意味着你要把请求状态继续留在 reqs_ 里, 让后续的 delta send 继续接着发. 但 start_send_substep() 是在 bypass 线程里调的, 而 reqs_ 这类状态原本更适合由 worker 主线程管理. 所以这里选择了一个更收敛的设计: substep 只处理那些已经能直接收口的发送任务.

为什么不直接再造一个 step?

一开始我也想过一个更直观的方案: bypass 期间每来一个新的 transfer_kv, 就临时创建一个新的 Step, 像 start_req_send() 那样各发各的. 但很快就发现这条路并不顺.

对于 unfinished req 来说, 你真正想复用的是当前 forward 已经在发生的 record_event(layer_idx) 和 layer-ready 时序. 如果临时再造一个 step, 那么这些 event 到底应该通知给谁? 当前真实 step? 还是这些临时 step? 还是两边都通知? 这相当于把原本很清楚的一条 step 时间线硬生生拆成了几条, 正确性会立刻变得很难守.

所以最后的 substep 方案反而更克制: 不是在 bypass 线程里“复制一条 step”, 而是想办法把新任务重新并回当前 step.

ack, rollback

如果只是 “把任务挂上去”, substep 还不算特别难. 真正麻烦的是:

scheduler 侧怎么知道哪些 substep 真的被 worker 当前 step 接纳了?

因为在 scheduler 侧, 你一旦根据 bypass substep 生成了 KVTPMeta.nonfreeze_metas, 通常也就顺手更新了 PBackend 内部状态, 比如 _infly_kvt、untouched 等. 但这些副作用的生成时点, 往往是早于 worker 真正接纳该 substep 的. 于是会出现这样一种时序:

step 1:     build_backend_meta -> kvmeta1
substep 11: build_backend_meta -> kvmeta11
substep 12: build_backend_meta -> kvmeta12
step 1 结束
step 2 开始

这里 kvmeta11 也许真的已经被 worker 附加到 step 1 上了, 但 kvmeta12 也许只是 scheduler 侧生成过, 实际上根本没赶上 step 1. 如果不把 kvmeta12 对应的副作用回滚掉, 后面状态就脏了. 所以我们又补了一层 ack 机制.

worker 在当前 step 收口前会调用 freeze_send_step(). 这个接口除了把当前 step 冻结住, 还有一个很重要的返回值: ack_substep_id，它表示:

当前 step 至少确认接纳到了哪个 substep.

更准确一点说, 所有 <= ack_substep_id 的 substep, 都已经进入了当前 step. worker 侧把这个信息包装成 SubstepAckEvent(stepid, substepid). 然后通过 get_kv_connector_kv_cache_events() -> get_event() 这条链路, 和 normal model output 一起回传给 scheduler 侧.

scheduler update_connector_output -> PBackend.update_events() 会先校验:

  • 收到的 event 数是否等于 TP world size,
  • 所有 worker 返回的 stepid/substepid 是否一致.

确认无误之后, 再调用 _rollback_substep_side_effects(stepid, ack_substepid)，把所有 > ack_substepid 的 substep 副作用回滚掉. 例如某个请求在 substep 中把 untouched 改成了 False, 但这个 substep 最终没被当前 step 接纳, 那么这里就会把它还原回去.

我个人觉得这一层其实才是整个 substep 设计里最关键的部分. 没有它, substep 只是 “看起来更及时”; 有了它, 才是真正把及时性和正确性一起兜住了.

后记

这类优化我一直很喜欢. 它不是什么宏大重构, 但会很明显地改善整条链路的体感. 你把控制面和数据面之间那一点点别扭的时序抹平之后, 系统就会顺很多.

最后再记一笔题外话: 这次整个 substep 开发我几乎没碰 IDE, 就是写好 design guide 之后让 AI 帮我读代码、补 patch、改接口、对调用时序. 某种意义上, 这应该算是一次相当彻底的 vibe coding 了.

CATALOG