KV Cache 池 (KV Cache Pool)

KV Cache 池 (KV Cache Pool)#

为什么需要 KV Cache 池？#

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

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

因此，KV Cache 池被提议用于利用包括 HBM、DRAM 和 SSD 在内的各种存储类型来构建 KV Cache 存储池，同时使请求的前缀在所有节点间可见，从而提高所有请求的缓存命中率。

vLLM Ascend 目前支持 MooncakeStore：最受认可的 KV Cache 存储引擎之一。

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

因此，我们建议通过全新的 MooncakeStoreConnectorV1 来集成 Mooncake Store，该连接器很大程度上受到了 LMCacheConnectorV1 的启发（详见 MooncakestoreConnectorV1 是如何实现的？ 章节）。

用法#

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

有关详细的部署和配置步骤，请参阅 KV 池用户指南。

工作原理#

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

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

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

1. Combining KV Cache Pool with HBM Prefix Caching#

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

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

工作流：

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

  2. 在获取 HBM 上的命中 token 数量后，它会通过连接器查询 KV 池。如果 KV 池中有额外的命中，我们将仅从 KV 池获取额外的块，并直接从 HBM 获取其余块，以最大限度地减少数据传输延迟。

  3. 在将 KV 池中的 KV Cache 加载到 HBM 后，剩余过程与 HBM 中的前缀缓存一致。

2. Combining KV Cache Pool with Mooncake PD Disaggregation#

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

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

这样做核心优势在于：对于预填充节点，我们可以通过 HBM 和 KV 池的前缀缓存减少计算量来保持性能增益，同时利用 P2P KV 连接器在 NPU 设备间直接传输 KV Cache，不牺牲预填充和解码节点之间的数据传输效率。

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

有关详细信息，另请参阅 Mooncake 连接器存储部署指南。

MooncakestoreConnectorV1 是如何实现的？#

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

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

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

KV 连接器调度器端方法：#

get_num_new_matched_tokens: Get prefix cache hit in number of tokens through looking up into the KV pool.
update_states_after_alloc: Update KVConnector state after temporary buffer alloc.
build_connector_meta: Attach the connector metadata to the request object.
request_finished: Once a request is finished, determine whether request blocks should be freed now or will be sent asynchronously and freed later.

连接器工作端方法：#

register_kv_caches: Register KV cache buffers needed for KV cache transfer.
start_load_kv: Perform KV cache load operation that transfers KV cache from storage to device.
wait_for_layer_load: Optional; Wait for layer load in layerwise + async KV load scenario.
save_kv_layer: Optional Do layerwise KV cache put into KV Pool.
wait_for_save: Wait for KV Save to finish if async KV cache save/put.
get_finished Get request that finished KV transfer, done_sending if put finished, done_reciving if get finished.

DFX (可维护性设计)#

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

  2. 同样地，当我们尝试将块存入 KV 池失败时，我们不再存入后续块（可能会有变动）。

限制#

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

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

目录