KV 缓存池

KV 缓存池#

为什么需要 KV 缓存池？#

前缀缓存是大语言模型推理中的一项重要特性，可以显著减少预填充计算时间。

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

因此，我们提出了 KV 缓存池，旨在利用包括片上内存、DRAM 和 SSD 在内的多种存储类型，构建一个 KV 缓存存储池，同时使请求的前缀在所有节点间可见，从而提高所有请求的缓存命中率。

vLLM Ascend 目前支持 MooncakeStore，这是最受认可的 KV 缓存存储引擎之一。

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

因此，我们提议将 MooncakeStore 与全新的 MooncakeStoreConnectorV1 集成，该连接器的设计在很大程度上受到了 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 Connector V1，用户可以无缝地将基于片上内存的前缀缓存与 Mooncake 支持的 KV 池结合起来。

The user can enable both features simply by enabling Prefix Caching, which is enabled by default in vLLM V1 unless the --no-enable-prefix-caching flag is set, and setting up the KV Connector for KV Pool (e.g., the MooncakeStoreConnector).

工作流程：

  1. 引擎首先检查片上内存缓存中的前缀命中情况。

  2. 获取片上内存上的命中令牌数量后，引擎通过连接器查询 KV 池。如果在 KV 池中有额外的命中，我们仅从 KV 池获取额外的块，其余块则直接从片上内存获取，以最小化数据传输延迟。

  3. 将 KV 池中的 KV 缓存加载到片上内存后，剩余过程与片上内存中的前缀缓存相同。

2.将 KV 缓存池与 Mooncake PD 解耦结合#

当与 Mooncake PD（预填充-解码）解耦功能结合使用时，KV 缓存池可以进一步在设备或节点间解耦预填充和解码阶段。

目前，我们仅对预填充节点执行 KV 池的 put 和 get 操作，解码节点则通过 Mooncake P2P KV 连接器（即 MooncakeConnector）获取其 KV 缓存。

这样做的主要好处是，我们可以通过为预填充节点使用来自片上内存和 KV 池的前缀缓存来减少计算量，从而保持性能增益，同时又不牺牲预填充节点与解码节点之间的数据传输效率，因为 P2P KV 连接器直接在 NPU 设备间传输 KV 缓存。

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

详情请参阅 Mooncake Connector Store 部署指南。

MooncakeStoreConnectorV1 是如何实现的？#

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

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

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

KV 连接器调度器端方法#

get_num_new_matched_tokens：通过查询 KV 池，获取以令牌数表示的前缀缓存命中数。 update_states_after_alloc：临时缓冲区分配后更新 KVConnector 状态。 build_connector_meta：将连接器元数据附加到请求对象。 request_finished：请求完成后，确定请求块是应立即释放，还是将异步发送并稍后释放。

连接器工作器端方法#

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

DFX（可诊断性、可维护性、可服务性）#

  1. 在 KV 池中查找键时，如果找不到该键，则此特定块没有缓存命中；我们返回此块未命中，并且不再为当前请求查找后续块。

  2. 类似地，当我们尝试将一个块放入 KV 池但失败时，我们不会放入后续块（可能更改）。

限制#

  1. 目前，vLLM-Ascend 的 MooncakeStore 仅支持 DRAM 作为 KV 缓存池的存储介质。

  2. 目前，如果我们成功查找到一个键并确认其存在，但在调用 KV 池的 get 函数时获取失败，我们仅输出一条日志表明 get 操作失败并继续执行；因此，该特定请求的准确性可能会受到影响。我们将通过回退该请求并假设没有前缀缓存命中来重新计算所有内容（或者更优的方案是，仅回退一个块并继续使用该块之前的前缀缓存）来处理这种情况。