KV 缓存池

KV Cache Pool#

为什么需要 KV Cache Pool？#

前缀缓存是 LLM 推理中的一项重要功能，可以大幅减少预填充计算时间。

然而，前缀缓存带来的性能提升高度依赖于缓存命中率，而如果仅使用 HBM 进行 KV 缓存存储，缓存命中率可能会受到限制。

因此，我们提出了 KV Cache Pool，它利用 HBM、DRAM 和 SSD 等各种类型的存储器，构建一个 KV Cache 存储池，同时使请求前缀在所有节点之间可见，从而提高所有请求的缓存命中率。

vLLM Ascend 目前支持 MooncakeStore：这是最知名的 KV Cache 存储引擎之一；

虽然可以通过将 Mooncake Store 设置为 LMCache 的远程后端（GPU）来在 vLLM V1 引擎中使用它（参见 教程），但我们发现集成一个直接支持 Mooncake Store 并能利用最适合华为 NPU 硬件的数据传输策略的连接器会更好。

因此，我们提议集成 Mooncake Store，并引入一个全新的 **MooncakeStoreConnectorV1**，它很大程度上借鉴了 **LMCacheConnectorV1**（参见“MooncakestoreConnectorV1 是如何实现的？”部分）。

用法#

vLLM Ascend 目前支持 Mooncake Store 用于 KV Cache Pool。要启用 Mooncake Store，需要配置 `kv-transfer-config` 并选择 `MooncakeStoreConnector` 作为 KV Connector。

有关分步部署和配置，请参阅 KV Pool 用户指南。

它是如何工作的？#

KV Cache Pool 通过连接器架构集成了多个内存层（HBM、DRAM、SSD 等）。

每个连接器都实现了统一的接口，用于根据访问频率和硬件带宽在不同层之间存储、检索和传输 KV 块。

当与 vLLM 的前缀缓存机制结合使用时，该池可以实现本地（HBM 中）和全局（通过 Mooncake）的高效缓存，确保常用前缀保持热度，而不太常用的 KV 数据可以溢出到成本较低的内存中。

1. KV Cache Pool 与 HBM 前缀缓存的结合#

vLLM V1 引擎已支持 HBM 前缀缓存。通过引入 KV Connector V1，用户可以无缝地将基于 HBM 的前缀缓存与基于 Mooncake 的 KV Pool 结合起来。

用户可以通过启用前缀缓存来同时启用这两项功能，前缀缓存默认在 vLLM V1 中启用，除非设置了 –no_enable_prefix_caching 标志，并为 KV Pool 设置 KV Connector（例如 MooncakeStoreConnector）。

工作流程:

  1. 引擎首先检查 HBM 缓存中的前缀命中。

  2. 在获得 HBM 中的命中 token 数量后，通过连接器查询 KV Pool。如果 KV Pool 中有额外的命中，我们**仅从 KV Pool 获取额外的块**，其余块直接从 HBM 获取，以最小化数据传输延迟。

  3. 在将 KV Cache 加载到 HBM 后，其余过程与 HBM 中的前缀缓存相同。

2. KV Cache Pool 与 Mooncake PD 分离的结合#

与 Mooncake PD（Prefill-Decode）分离结合使用时，KV Cache Pool 可以进一步将预填充和解码阶段跨设备或节点解耦。

目前，我们仅对**预填充节点**执行 KV Pool 的 put 和 get 操作，解码节点通过 Mooncake P2P KV Connector（即 MooncakeConnector）获取其 KV Cache。

这样做的一个关键好处是，我们可以通过预填充节点上 HBM 和 KV Pool 的较少计算来保持性能收益，同时又不牺牲预填充和解码节点之间通过 P2P KV Connector 进行数据传输的效率，该连接器直接在 NPU 设备之间传输 KV Cache。

要启用此功能，我们需要使用 Multi Connector 设置 Mooncake Connector 和 Mooncake Store Connector，Multi Connector 是 vLLM 提供的一个 KV Connector 类，可以按特定顺序调用多个 KV Connector；

有关详细信息，请参阅 Mooncake Connector Store 部署指南。

MooncakestoreConnectorV1 是如何实现的？#

**MooncakestoreConnectorV1** 继承自 vLLM V1 中的 KV Connector V1 类：通过实现 KV connector V1 基类中定义的必需方法，可以将第三方 KV 缓存传输/存储后端集成到 vLLM 框架中。

MooncakeStoreConnectorV1 在 `Lookup Engine`/`Lookup Client` 的设计（用于查找 KV 缓存键）以及 `ChunkedTokenDatabase` 类（用于将 token 处理成感知前缀的哈希以及其他与哈希相关的设计）方面，很大程度上也借鉴了 LMCacheConnectorV1。在此基础上，我们还增加了自己的设计，包括允许异步 KV 缓存 `get` 和 `put` 的 `KVTransferThread`（使用多线程），以及 NPU 相关的数据传输优化，例如移除 LMCache 中的 `LocalBuffer` 以消除冗余数据传输。

需要实现而被调用的 KV Connector 方法可以分为在 V1 调度器中调用的调度器端方法和在 V1 工作端中调用的工作端方法，即

KV Connector 调度器端方法：#

`get_num_new_matched_tokens`: 通过查找 KV Pool 来获取新的匹配 token 数量（前缀缓存命中）。
`update_states_after_alloc`: 在临时缓冲区分配后更新 KVConnector 状态。
`build_connector_meta`: 将连接器元数据附加到请求对象。
`request_finished`: 请求完成后，确定请求块是现在释放还是异步发送后释放。

Connector 工作端方法：#

`register_kv_caches`: 注册 KV 缓存传输所需的 KV 缓存缓冲区。
`start_load_kv`: 执行 KV 缓存加载操作，将 KV 缓存从存储传输到设备。
`wait_for_layer_load`: 可选；在层式 + 异步 KV 加载场景中等待层加载。
`save_kv_layer`: 可选；执行层式 KV 缓存 put 操作到 KV Pool。
`wait_for_save`: 如果 KV 缓存保存/put 是异步的，则等待 KV 保存完成。
`get_finished`: 获取 KV 传输完成的请求，如果 `put` 完成则为 `done_sending`，如果 `get` 完成则为 `done_reciving`。

DFX#

  1. 在 KV Pool 中查找键时，如果找不到键，则表示此特定块没有缓存命中；我们返回此块没有命中，并且不对当前请求继续查找其他块。

  2. 同样，当我们尝试将一个块 put 到 KV Pool 并失败时，我们也不会 put 其他块（可能更改）。

限制#

  1. 目前，vLLM-Ascend 的 Mooncake Store 仅支持 DRAM 作为 KV Cache Pool 的存储。

  2. 目前，如果我们成功查找了某个键并发现它存在，但在调用 KV Pool 的 get 函数时未能获取到它，我们只输出一条日志表明 get 操作失败并继续；因此，该特定请求的准确性可能会受到影响。我们将通过回退请求来处理这种情况，并假设没有前缀缓存命中（甚至更好，只回退一个块，并保留之前的 Prefix Caches）。

目录