跳到内容

混合 KV 缓存管理器

警告

本文档基于提交 458e74 编写。此功能仍处于早期阶段，可能会发生变化。

什么是混合模型？

许多最近的“混合”LLM 在一个模型中结合了多种注意力类型。例如：

  1. 滑动窗口注意力 (sw) + 全注意力 (full)：gpt-oss, Gemma 2/3, Ministral, cohere 等。
  2. Mamba + full：Bamba, Jamba, Minimax 等。
  3. 局部块注意力 + full：Llama4

为了高效地服务这些模型，我们的 KVCacheManager 必须

  1. 为不同层类型分配不同的槽位，例如：
    • 全注意力层：为**所有** token 保留槽位。
    • 滑动窗口层：仅为最近的 **sliding_window_size** 个 token 保留槽位。
  2. 支持特定于层的 the prefix-cache 规则，例如：
    • 全注意力：缓存命中的前缀要求**所有** token 保留在 KV 缓存中。
    • 滑动窗口：缓存命中的前缀仅要求最后 **sliding_window_size** 个 token 保留在 KV 缓存中。

定义

  1. kv 隐藏大小：存储单个 token 的 KV 缓存对单个层所需的字节数。
  2. 块：为 kv 缓存保留的内存被划分为多个具有相同*页大小*（定义如下）的*块*。
  3. 块大小：块内的 token 数量。

  4. 页大小：块的物理内存大小，定义为

    \[ \text{num_layers} \times \text{block_size} \times \text{kv_hidden_size} \]

    num_layers 并不意味着模型中的总层数。确切数字取决于本文档的上下文。

    注意

    这与代码中的 KVCacheSpec.page_size_bytes 不同，后者定义为

    \[ \text{block_size} \times \text{kv_hidden_size} \]

分配

高层思路

我们为所有层类型使用单个内存池。内存池被划分为多个具有相同页大小的块。 KVCacheManager 根据其注意力类型为不同层分配不同数量的块。

核心挑战是确保每种层类型使用相同的**页大小**。对于仅全注意力模型，页大小的定义很直接：

\[ \text{page_size} = \text{block_size} \times \text{num_hidden_layers} \times \text{kv_hidden_size} \]

然而，在混合模型中，num_hidden_layers 因注意力类型而异，这通常会产生不匹配的页大小。下面的案例展示了我们如何统一它们。

案例 1：玩具模型

让我们从一个玩具示例开始：一个模型有 1 个全注意力层和 3 个滑动窗口注意力层。所有层具有相同的 kv_hidden_size。

我们让每个块持有 block_size 个 token，用于一层，因此：

\[ \text{page_size} = \text{kv_hidden_size} \times \text{block_size} \]

KVCacheManager 为每个层分配不同数量的块。

此案例仅为玩具示例。对于真实模型，请参阅以下案例。

案例 2：相同的 kv_hidden_size 和规则模式

当模型有更多层时，例如，20 个滑动窗口注意力层和 10 个全注意力层，且具有相同的 kv_hidden_size。为每层调用一次分配器（30 次调用）是可以的，但效率低下。作为一种解决方案，我们将需要相同数量块的层分组，以减少调用次数。

分组是可行的，因为不同类型层之间的数量通常存在漂亮的比例。例如：

  • Gemma-2：1 sw : 1 full
  • Llama 4：3 local : 1 full

我们的示例可以看作是 2 sw : 1 full。我们可以像模型中有 2 个 sw 和 1 个 full 一样分配块，然后将结果重复 10 次来生成 30 个层的 block_ids。页大小变为：

\[ 10 \times \text{kv_hidden_size} \times \text{block_size} \]

假设 block_size 为 16，滑动窗口大小为 32，请求长度为 112，那么对于上述示例模型，我们需要分配 11 个块（0-6 用于 full，7-8 用于 sw 组 1，9-10 用于 sw 组 2）。

Allocation Result

此处，“/”表示不需要块（滑动窗口层不需要早期 token 的槽位）。

见下文正式定义。层被划分为多个*KV 缓存组*，以便：

  1. 每个组内具有相同的注意力类型：每个组仅包含具有相同注意力类型的层，因此对于给定的请求，它们需要相同数量的块。这使得同一组中的层可以共享相同的块 id，而不会造成内存浪费。
  2. 组之间的页大小相同：因为我们的内存池只有一个页大小。

我们的示例模型被划分为 3 个 KV 缓存组：

  • 组 0：10 个全注意力层（full.0 - full.9）
  • 组 1：10 个滑动窗口注意力层（sw.0 - sw.9）
  • 组 2：10 个滑动窗口注意力层（sw.10 - sw.19）

显然，它满足规则 1。对于规则 2，所有 3 个组都有

\[ 10 \times \text{kv_hidden_size} \times \text{block_size} \]

作为它们的页大小。

案例 3：相同的 kv_hidden_size 和无规则模式

不幸的是，并非所有模型都具有如此优美的比例，并且案例 2 中的方法会产生太多的小组。例如，Gemma-3-27b 有 52 个滑动窗口注意力层和 10 个全注意力层。在案例 2 的约束下，它将是 26 个滑动窗口组和 5 个全注意力组，每个组包含 2 层。分配仍然效率低下。为了减少 KV 缓存组的数量，我们使用所有注意力类型中最小的层数来分组层。例如，Gemma-3-27b 中每组为 min(52, 10) = 10 个层。然后分组结果为：

  • 组 0：10 个全注意力层（full.0 - full.9）
  • 组 1：10 个滑动窗口注意力层（sw.0 - sw.9）
  • 组 2：10 个滑动窗口注意力层（sw.10 - sw.19）
  • ...
  • 组 6：10 个滑动窗口注意力层（sw.40 - sw.49）
  • 组 7：2 个滑动窗口注意力层（sw.50 - sw.51）和 8 个填充层

如果这种启发式方法在出现新模型时导致了糟糕的结果（例如，20 个 full + 30 个 sw，组大小应为 10 而不是 20），我们将更新此算法。

这种情况发生在 Gemma-3 系列模型以及案例 2 中的模型，但带有引入一个全注意力层的 eagle 投机解码。该解决方案存在一些内存浪费，并不完美。请报告任何填充开销无法接受的情况，以便我们改进算法。

案例 4：不同的 kv_hidden_size（主要是混合 Mamba 模型）

一些架构（例如，Bamba、Jamba、Minimax）将标准注意力层与 Mamba 层交织在一起，其中每个 Mamba 层的每个 token 的状态大小可能远大于注意力层的 kv_hidden_size。因为我们只支持跨所有组的单个页大小，所以我们必须协调这些不同的隐藏大小。

当前的算法是：

  1. 增加注意力层的 block_size，直到 $$ \text{block_size} \times \text{kv_hidden_size}{\text{att}} \ge \text{state_size} $$}
  2. 填充每个层的 mamba 状态为 $$ \text{block_size} \times \text{kv_hidden_size}_{\text{att}} $$
  3. 应用案例 3 中的分组策略。

注意

这可能导致注意力层的 block_size 超过 400，这太大了。另一种填充策略是增加 block_size，直到

\[ \text{block_size} \times \text{kv_hidden_size}_{\text{att}} \times \text{num_attn_layers} \ge \text{state_size}_{\text{mamba}} \]

此填充策略仍在进行中。

案例 5：KV 共享

KV 共享是指一个层使用另一个层的 KV 缓存，例如 gemma-3n。在这些模型中，KVCacheManager 忽略所有具有 KV 共享的层，仅为需要 KV 缓存的层分配 KV 缓存，并在模型运行器中进行一些修补以将分配结果应用于 KV 共享层。

前缀缓存

为简单起见，我们在本节中假设 block_size=1。

高层思路

块池使用类似于 tuple(block_hash, group_id) -> block 的字典来捕获完整的块。这意味着不同组的相同 token 会独立缓存和逐出。

当有新请求进入时，我们检查每个组的缓存命中前缀，并将这些组的交集作为请求的缓存前缀返回。有关检查单个组的缓存命中和执行交集的详细算法，请参阅下文。

案例 0：仅全注意力模型

对于全注意力层，为请求中的所有 token 分配块。有关底层设计的详细信息，请参阅 前缀缓存。

要找到请求的最长缓存命中前缀，我们从左（第一个块）到右（最后一个块）枚举，检查块是否被缓存，并在缓存未命中时退出。例如，在下面的示例中，我们将返回前 7 个 token（0-6）作为缓存命中前缀（蓝色块是已缓存的）。

Prefix Caching of Full Attention

案例 1：仅滑动窗口注意力模型

对于滑动窗口注意力层，内存分配的朴素实现是分配 sliding_window_size 个块，并以循环方式填充块。但这种朴素实现与前缀缓存不兼容，因此我们没有选择这种设计。在 vLLM 中，我们为不同的 token 分配不同的块，并释放超出滑动窗口的块。

对于新请求，缓存命中前缀仅需要缓存最后 sliding_window_size - 1 个 token。假设 sliding_window_size = 4 且 block_size = 1，请求是一个 15 个 token 的提示（蓝色块是已缓存的）。

Prefix Caching of Sliding Window Attention

有 3 种可能的缓存命中前缀：

  • 缓存命中长度 5，使用 [2, 3, 4] 进行预填充计算 → [5, 6, …, 14]
  • 缓存命中长度 6，使用 [3, 4, 5] 进行预填充计算 → [6, 7, …, 14]
  • 缓存命中长度 14，使用 [11, 12, 13] 进行预填充计算 → [14]（最高效）

我们可以从右到左检查缓存命中，并在找到匹配时提前退出。这与全注意力模型相反，后者从左到右检查并在匹配失败时提前退出。一个潜在的缺点（与全注意力相比）是，当没有匹配时，我们最终会遍历整个 token 列表，而这通常是常见情况。这可能会导致不可忽略的开销，但与 full + swa 结合使用效果很好，如下所述。

案例 2：滑动窗口注意力 + 全注意力模型

第一个问题是如何找到缓存命中前缀。我们需要通过以下方式“相交”全局和滑动窗口注意力层的缓存命中：

  1. 获取全注意力模型的最长缓存命中（从左到右扫描）。
  2. 获取滑动窗口注意力模型的最长缓存命中，该命中长度小于上述长度。通过从右到左开始检查缓存命中来实现，从全注意力模型的缓存命中长度开始。

可以确保滑动窗口注意力层的最终缓存命中也是全注意力层的缓存命中。这比查找每个组的所有可能前缀并执行交集更有效，因为我们的方法可以在没有缓存命中时提前退出。

该算法适用于具有完全两种注意力类型的模型：全注意力 + X，其中 X 可以是任意高效的注意力算法，如滑动窗口、Llama 4 局部注意力、Mamba。它不支持没有全注意力层的模型，也不支持超过 2 种注意力类型的模型。这足以满足本文档编写时的大多数混合模型。

第二个问题是缓存逐出策略。目前，我们为所有 KV 缓存组使用一个 LRU 队列。当块被释放时（因为请求完成或块超出滑动窗口），它们会被添加到 LRU 队列。

案例 3：Mamba 模型

Mamba 模型的前缀缓存支持仍在开发中。一旦实现，具有 Mamba 层 + 全注意力层的模型可以通过案例 2 中的 full attention + X 算法来支持。

实现

概述

Overview of Hybrid KV Cache Manager

KVCacheManager 被组织成 3 层：

  • KVCacheManager：调度程序和 KV 缓存管理系统之间的接口。
  • KVCacheCoordinator：协调每个组的 SingleTypeKVCacheManagers 来生成请求的分配结果。根据模型的配置，会选择以下协调器之一：
    • KVCacheCoordinatorNoPrefixCache：在禁用前缀缓存时使用。
    • UnitaryKVCacheCoordinator：如果只有一个 KV 缓存组。前缀缓存逻辑被简化，无需交集。
    • HybridKVCacheCoordinator：处理恰好两个 KV 缓存组（必须包含一个全注意力组加上另一个高效注意力组）。其他情况尚未实现。您可以禁用前缀缓存以使用 KVCacheCoordinatorNoPrefixCache。
  • SingleTypeKVCacheManager：每个实例管理一个 KV 缓存组的分配和前缀缓存，实现注意力类型特定的逻辑（例如，全注意力、滑动窗口、Mamba）。

上图中的蓝色框显示了 10 个全注意力层和 20 个滑动窗口注意力层的案例，因此：

  • 使用 HybridKVCacheCoordinator。
  • 为 3 个 KVCacheGroups 使用 1 个 FullAttentionManager 和 2 个 SlidingWindowManager。

内存布局

对于一个具有 n 个 KVCacheGroups 的模型，每个组有 m 层，我们分配 m 个缓冲区。每个缓冲区由 n 层共享，每层来自一个组。

下图显示了一个具有 10 个全注意力层（full.0 - full.9）和 20 个滑动窗口注意力层（sw.0-sw.19）的模型。它遵循“分配”部分中的“案例 2”，并被划分为 3 个组。

  • 组 0：10 个全注意力层（full.0 - full.9）
  • 组 1：10 个滑动窗口注意力层（sw.0 - sw.9）
  • 组 2：10 个滑动窗口注意力层（sw.10 - sw.19）

对于一个请求，我们分配 11 个块，其中 block_id 0-6 分配给组 0，7-8 分配给组 1，9-10 分配给组 2。

在这样一个示例中，物理内存被划分为 10 个缓冲区（KVCacheTensor 0 - KVCacheTensor 9）。每个缓冲区由 3 层共享（例如，KVCacheTensor 0 由组 0 的 full.0、组 1 的 sw.0 和组 2 的 sw.10 共享），并被划分为大小为 block_size * kv_hidden_size 的块。这些 3 个注意力层的 KV 缓存根据分配的 block_ids 被保存在缓冲区的不同块中。

Example Memory Layout

注意

一个逻辑“块”映射到物理内存的 10 个缓冲区中的 10 个块。