AI

转折中的 PD 分离

Posted by w@hidva.com on November 28, 2025

就在昨天晚上, 我们线上 PD 分离部署规模又上了一个新的台阶. 这次没有想往常几次爬坡出现新的问题, 一时间心有所感, 就是想写点什么记录下我们整个 PD 分离的开发历程.

kvt 1.0

看过我历史几篇关于 PD 分离的文章, 比如 PD 分离中的 GDR 等的同学都有了解, 我们 PD 分离在早期设计上就期望做到旁路, 不需要对 step 主流程做大的改动. ZeroOverhead, 不会因为 PD 分离在 Step 执行链路引入太重的负载! 基于这些要求, 我们设计了 KVT 模块, 其负责在两个节点之间完成 kvcache transfer 任务. 大体上 kvt 分为如下几个模块:

  • 接入层, 与 python 对接. 会在实例启动时, 将各个 layer 对应信息, 比如 layer 地址, layer cache shape 等注册到 kvt 上. 之后对于 P 节点, 会根据每次 step 对应 ScheduleOutput 生成 KVTMeta, 记录着待发送的请求相关信息, 比如源端目的端信息, 当前 step 处理 token 范围等. 会在每次 attention layer 计算完毕之后通知 kvt 本层数据可以发送了, 在 python 侧这里 “通知” 就是 record 一个固定 cuda event, 所以我们 PD 分离 P 节点也是 full-cuda-graph 兼容的.

  • ParseBlock; 根据 layer 信息, 以及请求相关信息, 计算 layer 待发送的 IpcBlocks: Vec<IpcBlock>. 一个 IpcBlock 就记录了 srcoff, dstoff, len 语义是将源端某个 layer srcoff, len 范围内的数据写入到对端对应 layer dstoff 处.

  • 控制层, 负责对端链接维护, 监听 layer 计算完成信号, 在每个 layer 计算完成之后调用传输层接口完成数据传输; 以及传输出错之后的容错处理等.

  • 传输层, 负责传输 Vec<IpcBlock>. 见 PD 分离中的 kvcache 传输优化.

kvt 提供的传输抽象是: 每个 layer 对应着一块显存. 每块显存划分为若干 block, block 具有相同的字节大小. 每个 block 内划分为若干 token, token 具有相同的字节大小; 在实例启动初始化 kvt 时会传递 block_size_bytes, token_size_bytes, num_blocks 这些描述 layer 物理布局的信息. 这是因为早期我们的 cache shape 是 (num_blocks, block_size, 2, num_heads, head_dim), 所以就不由自主地在设计上往这方面靠拢了. 后来在接入 vllm 之后发现 vllm fa cache shape (2, num_blocks, block_size, num_kv_heads, head_size), 也即对于一个 token, 其 K/V 并不是一个 block 内了; 当时还以为怕是要对 kvt 一番大改, 后来灵机一动发现只需要新增一个 parse block 实现就行: 在初始化 kvt 时, block_size_bytes, token_size_bytes 仍假装 token kv 在一起; 在新增的 parse block 那里在解释 block_size_bytes, token_size_bytes 字段时再分开就可以了. 从这以后 kvt 在迭代时才开始有意识地把 ParseBlock 抽离出来. 这样的抽离使得后续我们在支持 Qwen3-Next GDN 时也只是新增了个 gdn 对应的 parse block.

hybrid connector

在 kvt 实现之后, 很显然接下来就是要在 vllm 中把 kvt 集成进来, 但正如 文章; 社区 vllm v1 connector 接口并不是很符合我们一开始的设想: 零侵入式旁路与零额外开销. 社区 vllm v1 kv connector 设计中:

  1. 把 kvcache 未完成 load/save 的请求, 也放入 waiting 队列中, 依赖于一些 “空” 的 step 来更新 kvcache 的状态, 混淆了 kvcache load/save 和计算任务, 即带来复杂度又影响性能.

  2. 同步接口会阻塞 step, 如 get_num_new_matched_tokens 这种可能需要调用外部服务 api, 比如目前 lmcache get_num_new_matched_tokens 实现. 并且当前 step 链路也引入了太多 PD 分离相关的逻辑处理, 可能会影响模型运行效能.

  3. 容错支持不完善. 目前 vllm 代码库中 PD 分离 connector 对请求 abort 这些异常路径要么是不支持, 要么是支持的不完善.

因此我们期望分离 “计算” 和 “kvcache load/save”; kvcache load/save 未 ready 的请求, 对于 scheduler 来说是 zero-overhead 的. 进一步地我们将这套组件分为 connector, backend 两个模块, 这里 backend 只需要负责 kvcache 的传输, 加载, 存储; 编写者只需要了解 kvcache layout 以及对应后端存储相关知识即可, 对 vllm scheduler 相关细节完全不需要感知. connector 则负责为 backend 提供运行环境; 以及动态扩缩, 容错, 请求生命周期管理等.

传输与请求生命周期解偶; 对于待传输 kvcache 的请求 R, hybrid connector 会先增加请求 R kvcache block ref cnt. 之后在传输完成之后调用 kvcache manager free block 递减 R kvcache block ref cnt 并在需要时将 block 放入 free list; 通过管理 ref cnt 使得传输本身与请求生命周期解偶了, 这样 Scheduler 完全不需要为 PD 分离引入额外的逻辑. 具体来说 hybrid connector 会在 EngineCore 进程单独起一个线程运行着 RPC Server, 该 RPC Server 运行着不同的 Backend 注册各种 method.

# PD 分离中, 运行在 P 节点的 PBackend 注册的 rpc method.
rpcsrv.register_method(TRANSFER_KV_REQ, self._on_transfer_kv)
rpcsrv.register_method(PREFILL_REQ, self._on_prefill)
rpcsrv.register_method(SEND_DONE_REQ, self._on_send_done)
rpcsrv.register_method(ABORT_REQS_REQ, self._on_abort_reqs)

# 请求迁移链路中, 迁移 Backend 注册的 RPC method.
rpcsrv.register_method(NEW_REQ_REQ, self._on_new_req)
rpcsrv.register_method(MIGRATE_TO_REQ, self._on_migrate_to)
rpcsrv.register_method(SUSPEND_REQ, self._on_suspend)

位于每个 Worker 的 kvt 会在完成传输任务发起 SEND_DONE_REQ rpc, hybrid connector 会在收到所有 rank 发来的 send done rpc 之后, 汇总各个 worker 发送状态; 并往 EngineCore 递交任务, 该任务中执行 free block 动作.

单边的 put; 当前我们 PD 分离在 D 节点上不需要运行任何逻辑, 这样我们 D 节点也是 full-cuda-graph 兼容的. 早期我们只支持单请求模式, 请求 R 只会发给 D 节点, D 节点运行着的 DBackend 会先为 R 分配 kvcache, 之后劫持请求 R, 在 hybrid connector loop 中为 R 选择 P 节点, 发送 PREFILL_REQ rpc, P 节点收到 R 之后, 会开始 R 的 prefill 过程并在每层计算完毕之后发送对应层的 kvcache. PREFILL_REQ rpc 会在 R prefill 以及 kvcache 传输完毕之后返回结果. DBackend 会在 PREFILL_REQ rpc 结束之后才将 R 放入 Scheduler, 此时会调整 R.num_computed_tokens 字段为传输 kvcache token 数. 即对于 Scheduler 来说, R 等同于命中了 prefix cache 的请求. 还是那句话在我们的设想中, Scheduler 完全不需要为 PD 分离引入额外的逻辑, 但现在 Scheduler 充斥着大量的 PD 分离逻辑, 看得让人头秃.

后来为了支持更大的调度灵活性, 我们引入了双请求模式, 此时请求 R 会同时发送给 P, D 两个节点; P 节点收到请求 R 之后会立刻开始 prefill, 由于此时尚未收到 R 对应的 DInfo, 不会触发传输动作. D 节点在收到请求 R 之后, DBackend 会调用 TRANSFER_KV_REQ rpc, P 节点在 TRANSFER_KV_REQ rpc 中接收到 R DInfo, 会在下一个 step 开始为 R 传输 kvcache. 本来一开始 TRANSFER_KV_REQ rpc 只是为了双请求模式引入的, 后来在支持请求迁移开发过程中, 发现其也可以直接调度 TRANSFER_KV_REQ rpc 来完成 kvcache 传输.

Abort 的处理; 既然我们劫持了新增请求的处理链路, 相应地 abort 链路也需要处理. hybrid connector 会通过 abort 请求信息告知 backend, 不同 backend 采取不同的处理策略. 比如 PD 分离中 PBackend 在收到 abort 请求之后会往 kvt 下发一个终止传输的信号, 之后 kvt 便认为这个请求传输完成了便发送 SEND_DONE_REQ rpc 带上 cached_token, 即实际传输的 token; connector 会发现 cached_token 小于待传输的 token, num_prompt_tokens, connector 便认为请求传输失败, 对应的 PREFILL_REQ/TRANSFER_KV_REQ rpc 会返回错误码. 在 DBackend, 其收到 abort 请求之后, 会立刻让请求结束, 并往 P 节点发送 ABORT_REQS_REQ rpc, 请求对应的 kvcache block 会通过上面提到的 ref cnt 技术并不会真正释放, 会直至收到对应 PREFILL_REQ/TRANSFER_KV_REQ rpc 结束时才会真正释放 kvcache block, 因为这个期间 P 可能仍在通过 GDR 写入 D kvcache block.

kvcache store backend; 如 这篇文章 所示, 我们当前也基于 connector/backend 这套架构实现了 kvsbackend, 负责将 kvcache 保存到 global kvcache store 以及从 global kvcache store 中加载 kvcache. 实际上当前我们会运行着 PBackend + DBackend + MigrationBackend + KVSBackend, 分别负责 PD 分离, 请求迁移, kvs 功能.

kvt 2.0

随着模型结构在飞快演进，比如 Qwen3-Next 的 GDN、DeepSeek 的 DSA，以及 vLLM Hybrid KVCache Manager 融合了更多样的 KVCacheSpec，对 kvt 也提出了更高的要求：它要在抽象层面更灵活、更通用。

另一方面，随着 MoE 架构的增多，我们需要特别避免 kvcache transfer 对 EP all2all 的影响。我们要考虑通过 TCP 来绕开现有的 GPU/GDR 通路，从而把不同类型的流量彻底隔离开。

这些需求也意味着：kvt 还得继续演进=

CATALOG