重新计算CPU卸载指南#

概述#

RecomputeCPUOffloadConnector保留被解码端重新计算调度器抢占的请求的KV缓存。当HBM KV块不足时,RecomputeScheduler可能会抢占正在运行的解码请求。如果没有此连接器,请求将回退到原始重新计算路径,并可能被发送回预填充节点再次执行预填充。启用重新计算CPU卸载后,在HBM块被重用之前,已计算的KV块会从HBM复制到CPU DRAM,并在请求再次被调度时复制回HBM。

此功能专为在线P/D分离部署工作负载设计,其中解码节点针对解码吞吐量进行了调优,无法在抢占后高效地重新计算长预填充。该功能为可选启用,并首先关注正确性。

在典型的解码部署中,max_num_batched_tokens通常设置为max_num_seqs * (1 + num_spec_tokens)左右,这样解码节点主要每个请求调度一个新token,加上如果启用MTP时的推测token。这对解码来说是高效的,但对于抢占后重新计算长提示来说太小了。重新计算CPU卸载避免了在KV状态可以本地保留时将此类请求发送回完整的提示重新计算路径。

关键概念#

  • 重新计算抢占:当解码端HBM KV缓存耗尽时,RecomputeScheduler可以抢占正在运行的请求,并在之后恢复它。

  • CPU DRAM保留RecomputeCPUOffloadConnector将被抢占请求的已计算KV块存储在CPU内存中。

  • H2D恢复:当被抢占的请求再次被调度时,连接器在模型前向传播之前恢复保留的KV块。

  • 回退行为:如果未配置连接器、未启用重新计算调度器或CPU卸载容量不足,vLLM将回退到原始的重新计算行为。

  • MultiConnector集成:在P/D分离部署中,使用MultiConnector将P/D连接器(例如MooncakeConnectorV1)与解码节点上的RecomputeCPUOffloadConnector组合。

配置参数#

RecomputeCPUOffloadConnector通过kv-transfer-config进行配置。

参数

描述

kv_connector

必须设置为RecomputeCPUOffloadConnector

kv_role

在解码节点上设置为kv_consumer

cpu_bytes_to_use_per_rank

可选且推荐。每个rank/卡用于重新计算卸载的CPU内存预算(字节)。如果设置,它将覆盖cpu_bytes_to_use / world_size

cpu_bytes_to_use

可选。此vLLM实例的总CPU内存预算(字节)。连接器将其除以world_size以获得每个rank的预算。这不如cpu_bytes_to_use_per_rank直接,并且在DP部署中更容易配置错误。默认总大小为8 GiB。

enable_offload_prefix_caching

可选。为具有相同前缀缓存哈希的完整哈希块启用CPU块共享。默认值为false;除非明确测试前缀共享,否则保持禁用状态。

当您希望每个rank/卡使用相同的卸载容量时,首选cpu_bytes_to_use_per_rank。例如,将cpu_bytes_to_use_per_rank设置为17179869184,即每个rank/卡16 GiB。

如果您使用cpu_bytes_to_use,请记住它会被world_size除。在典型的P/D解码部署中,这意味着该值除以活动的解码DP大小。例如,当启动DP2TP8解码服务时,将cpu_bytes_to_use设置为16 GiB,每个DP rank的卡将获得大约8 GiB的重新计算卸载空间。为避免歧义,新部署请使用cpu_bytes_to_use_per_rank

在P/D分离的解码节点上,还必须在additional-config中启用recompute_scheduler_enable

--additional-config '{"recompute_scheduler_enable":true}'

备注

recompute_scheduler_enable仅在P/D分离模式(kv_rolekv_producerkv_consumer)下有效。请勿在PD混合模式(kv_rolekv_both)下启用它。

Docker共享内存#

重新计算CPU卸载会为卸载的KV块分配固定的CPU张量。在Docker中运行时,请确保容器的共享内存足够大。如果--shm-size太小,较大的每卡卸载预算(例如每卡16 GiB)可能无法分配CPU张量,并可能导致服务在启动时OOM或挂起。

对于具有较大重新计算卸载预算的典型A3部署,在启动容器时设置--shm-size=1024g。以下代码片段遵循DeepSeek-V4-Flash教程中使用的Docker风格:

export IMAGE=quay.io/ascend/vllm-ascend:|vllm_ascend_version|-a3
docker run --rm \
    --name vllm-ascend \
    --shm-size=1024g \
    --net=host \
    --device /dev/davinci0 \
    --device /dev/davinci1 \
    --device /dev/davinci2 \
    --device /dev/davinci3 \
    --device /dev/davinci4 \
    --device /dev/davinci5 \
    --device /dev/davinci6 \
    --device /dev/davinci7 \
    --device /dev/davinci8 \
    --device /dev/davinci9 \
    --device /dev/davinci10 \
    --device /dev/davinci11 \
    --device /dev/davinci12 \
    --device /dev/davinci13 \
    --device /dev/davinci14 \
    --device /dev/davinci15 \
    --device /dev/davinci_manager \
    --device /dev/devmm_svm \
    --device /dev/hisi_hdc \
    -v /usr/local/dcmi:/usr/local/dcmi \
    -v /usr/local/Ascend/driver/tools/hccn_tool:/usr/local/Ascend/driver/tools/hccn_tool \
    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
    -v /etc/ascend_install.info:/etc/ascend_install.info \
    -v /etc/hccn.conf:/etc/hccn.conf \
    -it $IMAGE bash

与P/D分离部署一起使用#

在解码节点上,配置MultiConnector,同时包含P/D连接器和RecomputeCPUOffloadConnector。P/D连接器处理从预填充到解码的KV传输,而RecomputeCPUOffloadConnector处理解码端的抢占保留和恢复。

以下示例使用MooncakeConnectorV1进行P/D KV传输。

python3 -m vllm.entrypoints.openai.api_server \
    --model /path/to/model \
    --port 8200 \
    --trust-remote-code \
    --enforce-eager \
    --tensor-parallel-size 1 \
    --data-parallel-size 1 \
    --max-model-len 32768 \
    --block-size 128 \
    --max-num-batched-tokens 4096 \
    --additional-config '{"recompute_scheduler_enable":true}' \
    --kv-transfer-config \
    '{
      "kv_connector": "MultiConnector",
      "kv_role": "kv_consumer",
      "kv_connector_extra_config": {
        "connectors": [
          {
            "kv_connector": "MooncakeConnectorV1",
            "kv_role": "kv_consumer",
            "kv_port": "28000",
            "kv_connector_extra_config": {
              "prefill": {
                "dp_size": 1,
                "tp_size": 1
              },
              "decode": {
                "dp_size": 1,
                "tp_size": 1
              }
            }
          },
          {
            "kv_connector": "RecomputeCPUOffloadConnector",
            "kv_role": "kv_consumer",
            "kv_connector_extra_config": {
              "cpu_bytes_to_use_per_rank": 17179869184,
            }
          }
        ]
      }
    }'

对于预填充节点,保持正常的P/D连接器配置,例如将MooncakeConnectorV1kv_role设置为kv_producer

当P/D连接器也需要kv_load_failure_policy时,请将其配置在顶层MultiConnectorkv-transfer-config中,而非子连接器内部。

独立连接器示例#

RecomputeCPUOffloadConnector必须与vLLM-Ascend的RecomputeScheduler配合使用。由于RecomputeScheduler仅支持P/D分离的Decode节点,因此重计算CPU卸载当前也仅适用于P/D分离的Decode节点。

以下独立连接器配置仅作为在P/D分离的Decode节点上验证重计算卸载路径的最小配置片段。它并非PD混合部署模式。在PD混合或普通非P/D部署中,请勿启用重计算CPU卸载;引擎将使用正常的vLLM重计算行为。

from vllm.config import KVTransferConfig

kv_transfer_config = KVTransferConfig(
    kv_connector="RecomputeCPUOffloadConnector",
    kv_role="kv_consumer",
    kv_connector_extra_config={
        "cpu_bytes_to_use_per_rank": 17179869184,
        "enable_offload_prefix_caching": False,
    },
)

在线服务场景:

vllm serve /path/to/model \
    --additional-config '{"recompute_scheduler_enable":true}' \
    --kv-transfer-config '{
        "kv_connector": "RecomputeCPUOffloadConnector",
        "kv_role": "kv_consumer",
        "kv_connector_extra_config": {
            "cpu_bytes_to_use_per_rank": 17179869184,
            "enable_offload_prefix_caching": false
        }
    }'

工作原理#

  1. RecomputeScheduler无法分配足够的HBM KV块时,它会选择一个Decode端正在运行的请求进行抢占。

  2. 在请求的HBM块被重用之前,调度器会调用连接器的抢占钩子。如果CPU块足够,连接器会记录HBM块到CPU块的映射。

  3. 模型运行器在更新工作节点状态前调用handle_preemptions()。工作节点将选中的KV块从HBM复制到固定的CPU张量。

  4. 在所有工作节点报告存储完成后,被抢占的请求变为可恢复状态。

  5. 当请求再次被调度时,连接器报告可从CPU内存恢复的token数量。

  6. 调度器分配新的HBM块,连接器构建CPU块到HBM块的重载映射。

  7. 在模型前向计算前,工作节点将恢复的KV块复制回HBM。前向计算随后从恢复的KV状态继续执行。

滑动窗口与MTP支持#

滑动窗口模型可能包含注意力窗口外token的空块ID逻辑块表。连接器保留逻辑块位置而非压缩非零块ID,因此像[0, 0, 20, 21]这样的块表在CPU端保持对齐为[0, 0, cpu_20, cpu_21]。仅传输非零块。

启用MTP或推测解码时,调度器可能分配预读块。重载路径将H2D范围裁剪为实际分配给恢复请求的GPU块表。

注意事项与限制#

  • 此功能需要vLLM V1引擎和vLLM-Ascend重计算调度器。

  • 该功能仅适用于P/D分离场景中的Decode节点。PD混合或普通非P/D部署不支持。

  • 请勿在PD混合部署中启用recompute_scheduler_enable。没有P/D分离的Decode端RecomputeScheduler,重计算CPU卸载不会生效,vLLM将遵循正常重计算路径。

  • 在Docker中运行且每rank卸载预算较大时,请预留足够的共享内存。对于每卡16 GiB卸载的典型A3部署,使用--shm-size=1024g

  • enable_offload_prefix_caching为实验性功能,默认关闭。

  • 当前的D2H和H2D传输使用基础的torch复制操作。这些操作以正确性为导向,尚未针对传输吞吐量进行优化。

  • Qwen3.5的异步调度尚未完全支持。使用Qwen3.5模型进行重计算CPU卸载时,请禁用异步调度。

  • 如果CPU内存容量不足,连接器会跳过该请求的卸载,调度器回退到原始重计算行为。

常见问题#

应配置多少CPU内存?#

从能容纳预期被抢占请求块数量的预算开始。优先使用cpu_bytes_to_use_per_rank,因为它直接控制每个rank/卡的卸载预算。例如,cpu_bytes_to_use_per_rank=17179869184为每个rank/卡提供16 GiB。

cpu_bytes_to_useworld_size均分。在DP2TP8的Decode服务中,设置cpu_bytes_to_use=17179869184为每个DP rank的卡提供约8 GiB的卸载空间。如果日志显示CPU缓存空闲块不足,请增加预算。

如何确认卸载路径已激活?#

查找类似以下日志:

Recompute preemption offload enabled for request ...
Created recompute offload state for request ...
Prepared recompute offload H2D load for request ...

如果无法准备卸载,日志可能显示CPU缓存空闲块不足,请求将回退到原始重计算路径。

这与KV缓存CPU卸载相同吗?#

不同。KV Cache CPU Offload是针对非活跃KV缓存块的前缀缓存卸载路径。RecomputeCPUOffloadConnector专门用于保留被Decode端重计算调度器抢占的请求的KV块。

参考#

设计背景和实现细节请参见#10820: Recompute CPU Offload Connector