KV 缓存池

KV 缓存池#

为什么需要 KV 缓存池？#

前缀缓存是 LLM 推理中的一个重要特性，可以显著减少预填充计算时间。

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

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

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

虽然可以通过将 MooncakeStore 设置为 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 池结合使用。

用户只需启用前缀缓存（除非设置了 --no-enable-prefix-caching 标志，否则 vLLM V1 中默认启用）并为 KV 池设置 KV 连接器（例如 MooncakeStoreConnector），即可同时启用这两个特性。

工作流程：

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

  2. 获取片上内存上的命中 token 数量后，通过连接器查询 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 connector V1 基类中定义的必需方法，可以将第三方 KV 缓存传输/存储后端集成到 vLLM 框架中。

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

需要实现的 KV Connector 方法可以分为调度器侧方法（在 V1 调度器中调用）和工作器侧方法（在 V1 工作器中调用），即：

KV 连接器调度器侧方法#

get_num_new_matched_tokens：通过查找 KV 池获取前缀缓存命中的 token 数量。 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 操作失败并继续执行；因此，该特定请求的准确性可能会受到影响。我们将通过回退请求并假设没有前缀缓存命中来重新计算所有内容来处理这种情况（或者更优的方案是仅回退一个块，并继续使用该块之前的前缀缓存）。