跳到内容

NIXL KV 缓存租约更新

在预填充/解码分离（disaggregated prefill/decode）部署中，预填充实例 (P) 在完成预填充后，必须在 GPU 内存中保留 KV 缓存块，等待解码实例 (D) 通过 RDMA 读取它们。当 D 无法检索这些块时，需要一种机制来确定何时可以安全地释放它们。该机制已在 PR #41383 中引入。

动机

单一超时问题

最初的设计使用单一的长超时时间（VLLM_NIXL_ABORT_REQUEST_TIMEOUT，默认 480 秒）来控制 P 保留 KV 块的时长。当 D 崩溃或断开连接时，P 会持有长达 8 分钟可能达数 GB 的“死”块后才回收。在此窗口期间，后续到达 P 的请求会发现缓存容量减少，从而导致性能下降。

过载问题

单纯降低超时时间会引入另一种故障模式。在流量激增时，请求可能会在 D 的等待队列中停留很长时间才被调度。如果 P 上的固定超时时间太短，KV 块会在 D 还没机会读取之前就被释放，从而导致不必要的重复计算和预填充工作的浪费。

解决方案：通过心跳更新租约

租约更新机制同时解决了这两个问题。当预填充完成时，P 授予一个简短的初始租约（默认 30 秒）。当请求在 D 上排队或执行中时，D 会定期向 P 发送心跳以延长租约。如果 D 崩溃并停止发送心跳，P 会在最后一次心跳后的几秒内回收块，而不是等待数分钟。如果 D 仅仅是过载，心跳会只要在需要时就保持块的存活。

工作原理

租约生命周期

当 P 完成预填充时，它会为 KV 块设定一个初始租约时长（kv_lease_duration，默认 30 秒）。从那时起，块将被保留，直到：

  1. D 完成 KV 传输 —— P 收到读取完成通知并立即释放块。
  2. D 持续发送心跳 —— 每次心跳将租约延长 lease_duration * 2/3（约 20 秒），只要 D 健康，就能无限期保持块的存活。
  3. 没有心跳到达 —— 租约过期，P 回收这些块。

利用 NIXL 通知机制

心跳复用了 NIXL 现有的通知系统（send_notif / get_new_notifs），而不是引入新的传输通道。通知介质与后端相关，NIXL 已经处理了从 IB/RoCE 到 TCP 的自动回退。从 D 发送到特定 P 的每个单一心跳消息都会代表该 D 更新 P 中保留的所有请求——换句话说，每次迭代发送一条批处理消息即可更新多个请求的租约。

调度器端跟踪 (D)

一个关键的见解是，心跳必须在请求进入 D 的调度器时立即开始，而不是在它被调度执行时。在重负载下，请求在等待队列中的驻留时间可能远超初始租约时长，且到达与调度之间的间隔是无界的。

为了实现这一点，D 的连接器 (NixlConnectorScheduler) 通过 on_new_request() 挂载到调度器中。当 do_remote_prefill=True 的请求到达时，连接器立即开始对其进行心跳跟踪。请求按 remote_engine_id 分组以实现高效批处理。在每个调度步骤中，心跳元数据被封装到 NixlConnectorMetadata 中并发送给工作进程，并由心跳间隔 lease_duration // 6（约 5 秒）进行节流。

当 KV 传输完成（通过 update_connector_output）或请求完成/中止（通过 request_finished）时，跟踪停止。

时机与简洁性

心跳的发送和处理发生在前向传播循环中，而不是在后台线程中。这意味着时机并非毫秒级精确——长时间的模型前向传递会延迟心跳。然而，租约时长的配置留有足够的余量：在默认设置下，心跳间隔（约 5 秒）和租约延长（约 20 秒）至少比典型的前向传递大一个数量级。这避免了线程间的锁复杂性，同时保持设计简单且可扩展。

正常流程

sequenceDiagram
    participant R as Routing Proxy
    participant P as Prefill Instance
    participant D as Decode Instance

    R->>P: Request (do_remote_decode=True)
    P->>P: Run prefill
    P->>P: Grant lease (30s)
    P->>R: Response (with kv_transfer_params)

    R->>D: Request (do_remote_prefill=True)
    note over D: Request enters waiting queue
    D->>D: on_new_request() starts tracking

    loop Every ~5s (heartbeat interval)
        D->>P: Heartbeat (extend lease)
        P->>P: Lease extended by ~20s
    end

    note over D: Request scheduled for execution
    D->>P: KV transfer (RDMA read)
    P-->D: Transfer complete
    D->>D: Stop heartbeating
    P->>P: Free KV blocks

解码实例崩溃

sequenceDiagram
    participant R as Routing Proxy
    participant P as Prefill Instance
    participant D as Decode Instance

    R->>P: Request (do_remote_decode=True)
    P->>P: Run prefill (holds onto KVs with lease)
    P->>R: Response

    R->>D: Request (do_remote_prefill=True)
    D->>P: Heartbeat (extend lease)
    D->>P: Heartbeat (extend lease)
    note over D: D crashes
    note over P: No heartbeat received
    P->>P: Lease expires (~20s, not 480s)
    P->>P: Free KV blocks

工作端发送与接收

在 D 端（发送）： 在 start_load_kv() 期间（每次前向传递都会调用），工作进程读取 metadata.heartbeat_by_engine 并向每个远程 P 引擎发送批处理的心跳通知。如果 D 尚未与给定的 P 引擎建立握手（常见于仍在等待队列中的请求），它会在后台线程触发主动握手。心跳会在握手完成后推迟到下一步发送——提前握手还能加速最终的 KV 传输。

在 P 端（接收）： 在 _get_new_notifs() 中，P 的工作进程检查传入的 NIXL 通知。以 "HB:" 开头的消息被路由到 _handle_heartbeat()，该函数使用 max(old_expiry, now + lease_extension) 延长每个引用请求的租约过期时间。这确保了租约永远不会被意外缩短。

双向 KV 传输

对于多轮对话，双向 KV 传输 允许 D 缓存 KV 块，以便 P 在后续回合中提取。由于下一轮对话的时机取决于客户端（不受系统控制），基于心跳的租约机制在此不适用。相反，一个单独的 decoder_kv_blocks_ttl（默认 480 秒）为 D 上缓存的块提供了简单的固定超时。如果客户端等待太久才继续对话，块将过期。D 会传回过期时间，以便 P 知道块何时过期并进行重新计算。由于截止时间是 D 上生成的 perf_counter 值，且两个引擎运行在不同的进程中（时钟互不相关），P 通过握手往返估计与 D 的时钟偏移，并在将其与自己的 perf_counter 比较之前进行修正。未来的工作可能会将对称的心跳机制扩展到此情况。

核心设计决策

  • 基于请求的租约，而非基于实例。 P 不知道其 KV 块属于哪个 D——块的所有权仅在预填充完成且路由器选择了 D 之后才确定。在请求级别进行租约避免了负载均衡器中 P/D 选择的耦合。在实践中，D 通过将具有相同 remote_engine_id 的请求分组，来向同一个 P 批量发送租约延长。

  • 使用 NIXL 通知作为传输。 心跳复用了现有的 send_notif/get_new_notifs 系统，而不是增加 ZMQ 连接或更改 API。通知介质与后端相关，且已处理 IB/RoCE 到 TCP 的回退，使心跳能在任何 NIXL 支持的传输方式上工作。

  • 无后台线程。 心跳的发送和处理发生在前向循环中（start_load_kv / get_finished）。这避免了线程间的锁复杂性。租约时长相对于前向传递延迟（秒级对比毫秒级）提供了足够的余量。

  • 主动握手。 当 D 需要向尚未连接的 P 引擎发送心跳时（常见于仍在等待队列中的请求），它会在后台线程触发早期握手。这也能加速最终的 KV 传输。

  • 异构 TP 支持。 当 P TP > D TP（例如 P TP=4，D TP=2）时，单个 D 工作进程会从多个 P 工作进程提取。心跳必须发送给给定引擎的所有 P 工作进程。反之，当 D TP > P TP 时，单个 P 会收到来自多个 D 的通知，这只是多次刷新 TTL，没有副作用。

配置

租约机制通过 --kv-transfer-config 中的 kv_connector_extra_config 进行控制

参数 默认值 描述
kv_lease_duration 30s P 上的初始租约时长。心跳间隔和延长量是自动派生的（interval = duration // 6，extension = duration * 2 // 3）。
decoder_kv_blocks_ttl 480s 在双向传输模式下缓存于 D 的 KV 块的生存时间 (TTL)。简单的固定超时，不通过心跳更新。 
vllm serve <MODEL> \
  --kv-transfer-config '{
    "kv_connector": "NixlConnector",
    "kv_role": "kv_producer",
    "kv_connector_extra_config": {"kv_lease_duration": 60}
  }'

有关 NixlConnector 配置的完整详情，请参阅 NixlConnector 使用指南。