KV 缓存池¶
为什么需要 KV 缓存池?¶
前缀缓存是 LLM 推理中的重要特性,可以大幅降低 prefill 计算时间。
然而,前缀缓存带来的性能提升高度依赖于缓存命中率;如果仅使用片上内存来存储 KV 缓存,缓存命中率可能会受到限制。
因此,提出 KV 缓存池,以利用多种类型的存储,包括片上内存、DRAM 和 SSD,为 KV 缓存存储构建一个池,同时使请求的前缀在所有节点之间可见,从而提高所有请求的缓存命中率。
vLLM Ascend 目前支持 MooncakeStore,它是最受认可的 KV 缓存存储引擎之一。
虽然可以通过将 MooncakeStore 设置为 LMCache 的 GPU 远程后端,在 vLLM V1 引擎中使用 MooncakeStore(参见教程),但我们认为,集成一个直接支持 MooncakeStore 的连接器会更好,并且可以使用最适合华为 NPU 硬件的数据传输策略。
因此,我们提出通过全新的 MooncakeStoreConnectorV1 集成 MooncakeStore,该连接器在很大程度上借鉴了 LMCacheConnectorV1(参见MooncakeStoreConnectorV1 是如何实现的? 部分)。
使用方法¶
vLLM Ascend 目前支持使用 MooncakeStore 作为 KV 缓存池。要启用 MooncakeStore,需要配置 kv-transfer-config,并选择 MooncakeStoreConnector 作为 KV 连接器。
有关逐步部署和配置,请参阅 KV 池用户指南。
工作原理¶
KV 缓存池通过基于连接器的架构集成多个内存层级(片上内存、DRAM、SSD 等)。
每个连接器都实现了一个统一接口,用于在不同层级之间存储、检索和传输 KV 块,具体取决于访问频率和硬件带宽。
当与 vLLM 的前缀缓存机制结合使用时,该池可以实现本地(片上内存中)和全局(通过 Mooncake)的高效缓存,确保频繁使用的前缀保持热状态,而较少访问的 KV 数据可以溢出到成本更低的内存中。
1. 将 KV 缓存池与片上内存前缀缓存结合¶
vLLM V1 引擎已经支持基于片上内存的前缀缓存。 通过引入 KV 连接器 V1,用户可以无缝地将基于片上内存的前缀缓存与由 Mooncake 支持的 KV 缓存池结合起来。
用户只需启用前缀缓存,即可同时启用这两个功能。前缀缓存在 vLLM V1 中默认启用,除非设置了 --no-enable-prefix-caching 标志;同时为 KV 缓存池配置 KV 连接器(例如 MooncakeStoreConnector)。
工作流程:
-
引擎首先检查片上内存缓存中的前缀命中。
-
在获得片上内存上的命中 token 数量后,通过连接器查询 KV 缓存池。如果 KV 缓存池中还有额外命中,则仅从 KV 缓存池获取额外块,其余块直接从片上内存获取,以最小化数据传输延迟。
-
当 KV 缓存池中的 KV 缓存加载到片上内存后,后续流程与片上内存中的前缀缓存相同。
2. 将 KV 缓存池与 Mooncake PD 分离结合¶
当与 Mooncake PD(Prefill-Decode)分离一起使用时,KV 缓存池可以进一步在不同设备或节点之间解耦 prefill 和 decode 阶段。
目前,我们仅对Prefill 节点执行 KV 缓存池的 put 和 get 操作,而 Decode 节点通过 Mooncake P2P KV 连接器(即 MooncakeConnector)获取其 KV 缓存。
这样做的关键优势在于,我们可以利用片上内存和 KV 缓存池的前缀缓存减少 Prefill 节点的计算量,从而保持性能收益,同时不牺牲 Prefill 节点和 Decode 节点之间的数据传输效率,因为 P2P KV 连接器可以直接在 NPU 设备之间传输 KV 缓存。
要启用此功能,需要使用 Multi Connector 同时配置 Mooncake Connector 和 MooncakeStore Connector。Multi Connector 是 vLLM 提供的一个 KV 连接器类,可以按照特定顺序调用多个 KV 连接器。
有关详细信息,还请参见 Mooncake 连接器部署指南。
MooncakeStoreConnectorV1 是如何实现的?¶
MooncakeStoreConnectorV1 继承 vLLM V1 中的 KV 连接器 V1 类:通过实现 KV 连接器 V1 基类中定义的必需方法,可以将第三方 KV 缓存传输/存储后端集成到 vLLM 框架中。
MooncakeStoreConnectorV1 在很大程度上也借鉴了 LMCacheConnectorV1,包括用于查找 KV 缓存键的 Lookup Engine/Lookup Client 设计,以及用于将 token 处理为前缀感知哈希的 ChunkedTokenDatabase 类和其他与哈希相关的设计。在此基础上,我们还添加了自己的设计,包括 KVTransferThread,它允许通过多线程异步执行 KV 缓存的 get 和 put,以及 NPU 相关的数据传输优化,例如移除 LMCache 中的 LocalBuffer,以消除冗余数据传输。
需要实现的 KV 连接器方法可以分为在 V1 调度器中调用的调度器侧方法,以及在 V1 worker 中调用的 worker 侧方法,具体如下:
KV 连接器调度器侧方法¶
get_num_new_matched_tokens:通过查找 KV 缓存池,获取以 token 数量表示的前缀缓存命中。
update_states_after_alloc:在临时缓冲区分配后更新 KVConnector 状态。
build_connector_meta:将连接器元数据附加到请求对象。
request_finished:当请求完成后,确定请求块是立即释放,还是异步发送并在之后释放。
连接器 Worker 端方法¶
register_kv_caches: 注册 KV 缓存传输所需的 KV 缓存缓冲区。
start_load_kv: 执行 KV 缓存加载操作,将 KV 缓存从存储传输到设备。
wait_for_layer_load: 可选;在逐层 + 异步 KV 加载场景中等待层加载完成。
save_kv_layer: 可选;执行逐层 KV 缓存写入 KV Pool。
wait_for_save: 如果异步 KV 缓存保存/写入,则等待 KV 保存完成。
get_finished: 获取已完成 KV 传输的请求;如果 put 完成,则为 done_sending;如果 get 完成,则为 done_receiving。
DFX¶
- 在 KV Pool 中查找键时,如果找不到该键,则此特定块没有缓存命中;我们对该块返回未命中,并且不再为当前请求查找后续块。
- 类似地,当我们尝试将块放入 KV Pool 并失败时,我们不会放入后续块(可能会更改)。
Limitations¶
-
目前,vLLM Ascend 的 MooncakeStore 仅支持将 DRAM 作为 KV Cache Pool 的存储。
-
目前,如果我们成功查找到一个键并发现其存在,但在调用 KV Pool 的 get 函数时获取失败,我们只会输出一条日志,表明 get 操作失败并继续执行;因此,该特定请求的准确性可能会受到影响。我们将通过回退该请求并假设没有前缀缓存命中而重新计算所有内容来处理这种情况(或者更好的做法是仅回退一个块,并继续使用之前的前缀缓存)。
本页原文 Markdown:在 AtomGit 查看·内容源自开源项目 vl/vllm-ascend