UCM增强的前缀缓存部署#

概述#

Unified Cache Management (UCM) 为 vLLM/vLLM-Ascend 中的 prefix-caching 场景提供了一个外部 KV-cache 存储层。与仅通过聚合设备内存来扩展前缀缓存容量的 KV Pooling 不同(后者仍受 HBM/DRAM 大小限制且缺乏持久性),UCM 实现了计算与存储的解耦,并采用了分层设计。每个节点使用本地 DRAM 作为快速缓存,而共享后端(如 3FS 或企业级存储)作为持久化 KV 存储。这种方法消除了设备内存带来的容量上限,实现了持久且可靠的前缀缓存,并允许缓存容量随存储系统而非计算资源扩展。

前置条件#

  • 操作系统:Linux

  • 配备 Ascend NPU 的硬件。通常为 Atlas 800 A2 系列。

  • vLLM:main 分支

  • vLLM Ascend:main 分支

UCM 安装#

请参考 针对 Ascend NPU 的官方 UCM 安装指南

为 Prefix Caching 配置 UCM#

修改 UCM 配置文件以指定使用的 UCM connector 以及 KV block 的存储位置。您可以直接编辑以下示例文件:

unified-cache-management/examples/ucm_config_example.yaml

有关更新的配置选项,请参考 针对 prefix-caching 的官方 UCM 文档

最小化配置如下所示:

ucm_connectors:
  - ucm_connector_name: "UcmNfsStore"
    ucm_connector_config:
      storage_backends: "/mnt/test"
      use_direct: false

load_only_first_rank: false

参数说明:

  • ucm_connector_name: "UcmNfsStore": 指定 UcmNfsStore 作为 UCM connector。

  • storage_backends: 指定用于存储 KV block 的目录。它可以是本地目录或 NFS 挂载路径。UCM 将在此处存储 KV block。⚠️ 请务必将 "/mnt/test" 替换为您实际的存储目录。

  • use_direct: 是否启用 Direct I/O(可选)。默认为 false

  • load_only_first_rank: 控制是否仅由 rank 0 加载 KV cache 并广播给其他 rank。此特性目前在 Ascend 上暂不支持,因此必须设置为 false(所有 rank 独立进行加载/转储)。

启动推理#

在本指南中,我们介绍使用带有 UCM connector 的 vLLM 进行在线推理,并部署为兼容 OpenAI 标准的服务。为了获得 UCM 的最佳性能,建议将 block_size 设置为 128。

使用 Qwen/Qwen2.5-14B-Instruct 模型启动 vLLM server,运行:

vllm serve Qwen/Qwen2.5-14B-Instruct \
--max-model-len 20000 \
--tensor-parallel-size 2 \
--gpu_memory_utilization 0.87 \
--block_size 128 \
--trust-remote-code \
--port 7800 \
--enforce-eager \
--no-enable-prefix-caching \
--kv-transfer-config \
'{
    "kv_connector": "UCMConnector",
    "kv_role": "kv_both",
    "kv_connector_extra_config": {"UCM_CONFIG_FILE": "/vllm-workspace/unified-cache-management/examples/ucm_config_example.yaml"}
}'

⚠️ 请务必将 "/vllm-workspace/unified-cache-management/examples/ucm_config_example.yaml" 替换为您实际的配置文件路径。

如果您看到如下日志:

INFO:     Started server process [1049932]
INFO:     Waiting for application startup.
INFO:     Application startup complete.

恭喜,您已成功启动带有 UCM connector 的 vLLM server!

评估 UCM Prefix Caching 性能#

在启动启用 UCMConnector 的 vLLM server 后,观察前缀缓存效果最简单的方法是运行内置的 vllm bench CLI。在另一个终端中执行以下命令两次,可以清晰地看到性能提升。

vllm bench serve \
--backend vllm \
--model Qwen/Qwen2.5-14B-Instruct \
--host 127.0.0.1 \
--port 7800 \
--dataset-name random \
--num-prompts 12 \
--random-input-len 16000 \
--random-output-len 2 \
--request-rate inf \
--seed 123456 \
--percentile-metrics "ttft,tpot,itl,e2el" \
--metric-percentiles "90,99" \
--ignore-eos

第一次执行后#

vllm bench 终端打印基准测试结果:

---------------Time to First Token----------------
Mean TTFT (ms):                           15323.87

检查 vLLM server 日志会发现如下条目:

INFO ucm_connector.py:228: request_id: xxx, total_blocks_num: 125, hit hbm: 0, hit external: 0

这表明对于第一个推理请求,UCM 未命中任何缓存的 KV block。因此,必须计算完整的 16K-token 预填充(prefill),导致 TTFT 相对较大。

第二次执行后#

再次运行相同的基准测试会产生:

---------------Time to First Token----------------
Mean TTFT (ms):                            1920.68

vLLM server 日志现在包含类似的条目:

INFO ucm_connector.py:228: request_id: xxx, total_blocks_num: 125, hit hbm: 0, hit external: 125

这表明在第二个请求期间,UCM 成功地从存储后端检索了所有 125 个缓存的 KV block。利用完全缓存的前缀显著降低了模型的初始延迟,与初始运行相比,TTFT 提升了约 8 倍