跳转至

KV缓存池

为什么需要KV缓存池？

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

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

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

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

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

因此，我们提出集成MooncakeStore与全新的**MooncakeStoreConnectorV1**，这实际上很大程度上受到了**LMCacheConnectorV1**的启发（参见How is MooncakeStoreConnectorV1 Implemented?部分）。

使用方法

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缓存池结合起来。

用户只需启用前缀缓存（除非设置了--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缓存。

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

有关详细信息，请同时参考Mooncake连接器部署指南。

MooncakeStoreConnectorV1是如何实现的？

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

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

需要实现的 KV Connector 方法可分为调度器端方法和工作器端方法，前者在 V1 调度器中调用，后者在 V1 工作器中调用，具体如下：

KV Connector 调度器端方法

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 传输的请求，done_sending 完成时返回 put，done_receiving 完成时返回 get。

DFX

  1. 在 KV 池中查找键时，如果找不到该键，则此特定块无缓存命中；我们返回此块未命中，并且不再为当前请求查找后续块。
  2. 类似地，当尝试将块放入 KV 池但失败时，我们不会继续放入后续块（可能变更）。

局限性

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

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