KV缓存池

为什么需要KV缓存池？

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

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

因此，我们提出了KV缓存池的概念，它利用包括HBM、DRAM和SSD在内的多种存储介质，构建一个统一的KV缓存存储池。该池使请求的前缀能够在所有节点间可见，从而提升所有请求的缓存命中率。

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

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

因此，我们提议将Mooncake Store与一个全新的 MooncakeStoreConnectorV1 集成，该连接器的设计在很大程度上受到了 LMCacheConnectorV1 的启发（详见MooncakestoreConnectorV1是如何实现的？章节）。

使用方法

vLLM Ascend 当前支持使用Mooncake Store作为KV缓存池。要启用Mooncake Store，需要配置 kv-transfer-config 并选择 MooncakeStoreConnector 作为KV连接器。

关于分步部署与配置，请参阅 KV缓存池用户指南。

工作原理

KV缓存池通过基于连接器的架构，整合了HBM、DRAM、SSD等多级存储介质。

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

当与vLLM的前缀缓存机制结合时，该缓存池支持在本地（HBM）和全局（通过Mooncake）进行高效缓存。这确保了高频使用的前缀保持热状态，而较少访问的KV数据则可以溢出到成本更低的存储中。

1.将KV缓存池与HBM前缀缓存结合

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

用户只需启用前缀缓存（vLLM V1中默认启用，除非设置了 --no_enable_prefix_caching 标志）并为KV缓存池（例如MooncakeStoreConnector）设置好KV连接器，即可同时启用这两个功能。

工作流程:

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

2. 在获取HBM中的命中词元数量后，引擎通过连接器查询KV缓存池。如果池中有额外的命中，则仅从KV缓存池获取这些额外的块，其余块直接从HBM获取，以最大限度地减少数据传输延迟。

3. 将KV缓存池中的KV缓存加载到HBM后，其余处理过程与使用HBM前缀缓存时相同。

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

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

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

这样做的主要好处在于，我们可以在预填充节点上利用HBM和KV缓存池的前缀缓存减少计算量，从而保持性能优势；同时，通过使用P2P KV连接器在NPU设备间直接传输KV缓存，不会牺牲预填充节点与解码节点之间的数据传输效率。

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

详情请参阅Mooncake连接器存储部署指南。

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: 临时缓冲区分配后，更新KV连接器状态。 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缓存保存/存储操作，则等待保存完成。 get_finished: 获取已完成KV传输的请求，done_sending 表示 put 完成，done_reciving 表示 get 完成。

可维护性（DFX）

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

2. 类似地，当我们尝试将块存入KV缓存池失败时，我们将不再存入后续块（此行为可能更改）。

限制

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

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