抢占卸载指南¶
概述¶
PreemptOffloadConnector 会保留被 Decode 侧重算调度器抢占的请求的 KV 缓存。当 HBM KV 块不足时,RecomputeScheduler 可能会抢占一个正在运行的 Decode 请求。如果没有此连接器,该请求会回退到原始重算路径,并可能被发回 Prefill 节点重新执行 prefill。启用抢占卸载后,在 HBM 块被复用之前,已计算的 KV 块会从 HBM 复制到 CPU DRAM,并在请求再次被调度时复制回 HBM。
此功能面向在线 P/D 分离工作负载,其中 Decode 节点针对 decode 吞吐进行了调优,在抢占后无法高效地重算长 prefill。该功能为可选启用,并优先保证正确性。
在典型的 Decode 部署中,max_num_batched_tokens 通常设置为大约 max_num_seqs * (1 + num_spec_tokens),以便 Decode 节点主要为每个请求调度一个新 token,如果启用 MTP,则还包括投机 token。这对 decode 很高效,但对于抢占后重算长 prompt 来说太小。当请求的 KV 状态可以在本地保留时,重算 CPU 卸载可以避免将这些请求重新通过完整的 prompt 重算路径发送。
关键概念¶
- 重算抢占:当 Decode 侧 HBM KV 缓存耗尽时,
RecomputeScheduler可以抢占一个正在运行的请求,并在之后恢复它。 - CPU DRAM 保留:
PreemptOffloadConnector将抢占请求已计算的 KV 块存储在 CPU 内存中。 - H2D 恢复:当被抢占的请求再次被调度时,连接器会在模型前向之前恢复保留的 KV 块。
- 回退行为:如果未配置连接器、未启用重算调度器,或 CPU 卸载容量不足,vLLM 会回退到原始重算行为。
MultiConnector集成:在 P/D 分离中,使用MultiConnector在 Decode 节点上将 P/D 连接器(例如MooncakeConnectorV1)与PreemptOffloadConnector组合使用。
配置参数¶
PreemptOffloadConnector 通过 kv-transfer-config 进行配置。
| 参数 | 描述 |
|---|---|
kv_connector |
必须设置为 PreemptOffloadConnector。 |
kv_role |
在 Decode 节点上设置为 kv_consumer。 |
cpu_bytes_to_use_per_rank |
可选。每个 rank/卡使用的 CPU 内存预算(以字节为单位)。此选项优先级最高。 |
cpu_bytes_to_use |
可选。此 vLLM 实例的总 CPU 内存预算(以字节为单位)。连接器会将其除以 world_size。当未设置 cpu_bytes_to_use_per_rank 时使用。 |
offload_host_memory_ratio |
可选。CPU 与 GPU KV 块数量之间的比例。CPU 缓存包含 int(ratio * gpu_num_blocks) 个块。仅当两个基于字节的选项都未设置时使用,默认值为 1。 |
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 Decode 部署中,这意味着该值会除以活动的 Decode DP 大小。例如,当启动 DP2TP8 Decode 服务时,将 cpu_bytes_to_use 设置为 16 GiB,会使每个 DP rank 的卡大约有 8 GiB 的重算卸载空间。为避免歧义,新部署请使用 cpu_bytes_to_use_per_rank。
如果两个基于字节的选项都未设置,默认的 offload_host_memory_ratio=1 会为每个 GPU KV 块分配一个 CPU KV 块。可以设置不同的比例来扩展此容量,而无需计算物理字节大小。
在 P/D 分离的 Decode 节点上,还必须通过 additional-config 启用 scheduler_config.recompute_scheduler_enable:
`scheduler_config.recompute_scheduler_enable` is only valid in P/D-disaggregated mode
(`kv_role` is `kv_producer` or `kv_consumer`). Do not enable it in PD-mixed
mode (`kv_role` is `kv_both`).
Docker 共享内存¶
重算 CPU 卸载会为卸载的 KV 块分配 pinned CPU 张量。在 Docker 中运行时,请确保容器的共享内存足够大。如果 --shm-size 太小,较大的每卡卸载预算(例如每卡 16 GiB)可能无法分配 CPU 张量,并可能导致服务在启动期间 OOM 或挂起。
对于具有较大重算卸载预算的典型 A3 部署,在启动容器时设置 --shm-size=1024g。以下代码片段遵循 DeepSeek-V4-Flash 教程中使用的 Docker 风格:
export IMAGE=quay.io/ascend/vllm-ascend:v0.23.0-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 分离¶
在 Decode 节点上,使用 P/D 连接器和 PreemptOffloadConnector 配置 MultiConnector。P/D 连接器负责从 Prefill 到 Decode 的 KV 传输,而 PreemptOffloadConnector 负责 Decode 侧的抢占保留与恢复。
以下示例使用 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 '{"scheduler_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": "PreemptOffloadConnector",
"kv_role": "kv_consumer",
"kv_connector_extra_config": {
"cpu_bytes_to_use_per_rank": 17179869184,
}
}
]
}
}'
对于 Prefill 节点,保留常规的 P/D 连接器配置,例如将 kv_role 设置为 kv_producer 的 MooncakeConnectorV1。
如果 P/D 连接器还需要 kv_load_failure_policy,请将其配置在顶层 MultiConnector 的 kv-transfer-config 中,而不是子连接器内部。
独立连接器示例¶
PreemptOffloadConnector 必须与 vLLM-Ascend 的 RecomputeScheduler 一起使用。由于 RecomputeScheduler 仅支持 P/D 分离的 Decode 节点,因此抢占卸载目前也仅可用于 P/D 分离的 Decode 节点。
以下独立连接器配置仅作为最小配置片段,用于在 P/D 分离的 Decode 节点上验证重计算卸载路径。它不是 PD 混合部署模式。在 PD 混合或常规非 P/D 部署中,请勿启用抢占卸载;引擎将改用常规的 vLLM 重计算行为。
from vllm.config import KVTransferConfig
kv_transfer_config = KVTransferConfig(
kv_connector="PreemptOffloadConnector",
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 '{"scheduler_config":{"recompute_scheduler_enable":true}}' \
--kv-transfer-config '{
"kv_connector": "PreemptOffloadConnector",
"kv_role": "kv_consumer",
"kv_connector_extra_config": {
"cpu_bytes_to_use_per_rank": 17179869184,
"enable_offload_prefix_caching": false
}
}'
工作原理¶
- 当
RecomputeScheduler无法分配足够的 HBM KV 块时,它会选择一个 Decode 侧正在运行的请求进行抢占。 - 在请求的 HBM 块被重用之前,调度器会调用连接器的抢占钩子。如果有足够的 CPU 块可用,连接器会记录 HBM 块到 CPU 块的映射。
- 模型运行器在更新 worker 状态之前调用
handle_preemptions()。worker 会将选中的 KV 块从 HBM 复制到固定的 CPU 张量。 - 在所有 worker 报告存储完成后,被抢占的请求变为可恢复状态。
- 当请求再次被调度时,连接器会报告可以从 CPU 内存恢复的 token 数量。
- 调度器分配新的 HBM 块,连接器构建 CPU 块到 HBM 块的重载映射。
- 在模型前向之前,worker 将恢复的 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 混合部署中启用
scheduler_config.recompute_scheduler_enable。如果没有 P/D 分离 Decode 侧的RecomputeScheduler,抢占卸载不会生效,vLLM 将遵循常规重计算路径。 - 在 Docker 中运行且每个 rank 的卸载预算较大时,请预留足够的共享内存。对于典型的每卡 16 GiB 卸载的 A3 部署,请使用
--shm-size=1024g。 enable_offload_prefix_caching是实验性功能,默认禁用。- 当前的 D2H 和 H2D 传输使用基本的 torch 复制操作。它们以正确性为导向,尚未针对传输吞吐量进行优化。
- Qwen3.5 配合异步调度尚未完全支持。在使用 Qwen3.5 模型进行抢占卸载时,请禁用异步调度。
- 如果 CPU 内存容量不足,连接器将跳过该请求的卸载,调度器会回退到原始重计算行为。
常见问题¶
我应该配置多少 CPU 内存?¶
从一个能够容纳预期数量的被抢占请求块的预算开始。优先使用 cpu_bytes_to_use_per_rank,因为它直接控制每个 rank/卡的卸载预算。例如,cpu_bytes_to_use_per_rank=17179869184 会为每个 rank/卡提供 16 GiB。
cpu_bytes_to_use 会除以 world_size。在 DP2TP8 Decode 服务中,设置 cpu_bytes_to_use=17179869184 会为每个 DP rank 的卡提供约 8 GiB 的卸载空间。如果日志显示 CPU 缓存空闲块不足,请增加预算。
或者,省略这两个基于字节的选项,并设置 offload_host_memory_ratio。其默认值 1 会创建相同数量的 CPU 和 GPU KV 块。
如何知道卸载路径已激活?¶
查找类似以下内容的日志:
Recompute preemption offload enabled for request ...
Created preempt offload state for request ...
Prepared preempt offload H2D load for request ...
如果无法准备卸载,日志可能会显示 CPU 缓存空闲块不足,并且请求将回退到原始重算路径。
这与 KV Cache CPU Offload 相同吗?¶
不。KV Cache CPU Offload 是用于非活跃 KV 缓存块的前缀缓存卸载路径。PreemptOffloadConnector 专门用于保留被 Decode 端重算调度器抢占的请求的 KV 块。
参考¶
有关设计背景和实现细节,请参阅 #10820: Preempt Offload Connector。
本页原文 Markdown:在 AtomGit 查看·内容源自开源项目 vl/vllm-ascend