混合 KV Cache 管理器

警告

本文档基于 commit 458e74 编写。该功能仍处于早期阶段，内容可能会有所变动。

什么是混合模型？

许多近期的“混合”大语言模型在单个模型中结合了多种注意力类型。例如：

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

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

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

定义

1. kv hidden size：为单层存储一个 token 的 KV Cache 所需的字节数。
2. block (块)：预留给 KV Cache 的内存被划分为多个具有相同 page size（定义如下）的块。
3. block size：一个块内包含的 token 数量。

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

   \[ \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} \]

分配

核心思路

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

核心挑战在于确保每种层类型使用相同的 page size。对于仅全注意力的模型，page size 很简单，定义为：

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

然而，在混合模型中，num_hidden_layers 会根据注意力类型而变化，这通常会导致 page size 不匹配。以下情况展示了我们如何统一它们。

情况 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。此时 page size 变为：

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

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

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

参见下方的正式定义。这些层被划分为多个 KV Cache 组，以确保：

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

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

• 组 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} \]

作为它们的 page size。

情况 3：相同的 kv_hidden_size 但没有规律模式

不幸的是，并非所有模型都有如此完美的比例，情况 2 中的方法会产生太多小分组。例如，Gemma-3-27b 有 52 个滑动窗口注意力层和 10 个全注意力层。在情况 2 的约束下，它将有 26 个滑动窗口组和 5 个全注意力组，每组包含 2 层。这种分配依然低效。为了减少 KV Cache 组的数量，我们使用所有注意力类型中最小的层数对层进行分组。例如，在 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 个填充层 (padding layers)

如果这种启发式方法在处理新模型时导致结果不佳（例如 20 个 full + 30 个 sw，组大小应为 10 而不是 20），我们将更新该算法。

这种情况出现在 Gemma-3 系列模型中，以及在情况 2 的模型中加入 eagle 推测解码（引入一个全注意力层）时。该解决方案存在一些内存浪费，并不完美。如果发现填充开销无法接受的情况，请报告，以便我们优化算法。

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

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

当前的算法是：

1. 增加注意力层的 block_size，直到 $$ \text{block_size} \times \text{kv_hidden_size}_{\text{att}} \ge \text{state_size}_{\text{mamba}} $$
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 Cache，例如 gemma-3n。在这些模型中，KVCacheManager 会忽略所有进行 KV 共享的层，仅为需要 KV Cache 的层分配缓存，并由模型运行器（model runner）进行部分修补，将分配结果应用到共享 KV 的层上。

前缀缓存

为简单起见，本节假设 block_size=1。

核心思路

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

当新请求到来时，我们检查每个组的缓存命中前缀，并将这些组的交集作为该请求的缓存前缀返回。关于检查一组缓存命中及执行交集运算的详细算法，请见下文。

情况 0：仅全注意力模型

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

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

情况 1：仅滑动窗口注意力模型

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

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

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

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

我们可以从右到左检查缓存命中，并在找到匹配项时提前退出。这与全注意力（从左到右检查，在匹配失败时提前退出）相反。一个潜在的缺点（与全注意力相比）是当没有匹配时，我们最终会遍历整个 token 列表，而这种情况往往很常见。这可能会导致不可忽略的开销，但在全注意力 + 滑动窗口注意力结合的情况下是没问题的，详见下文。

情况 2：滑动窗口注意力 + 全注意力模型

第一个问题是如何找到缓存命中前缀。我们需要通过以下方式“求交集”：

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

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

该算法适用于恰好有两种注意力类型的模型：全注意力 + X，其中 X 可以是任何高效的注意力算法，如滑动窗口、Llama 4 局部注意力和 Mamba。它不支持没有全注意力层的模型，也不支持拥有超过 2 种注意力类型的模型。这对于撰写本文时的大多数混合模型来说已经足够了。

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

情况 3：Mamba 模型

Mamba 模型的前缀缓存支持仍在开发中。一旦实现，带有 Mamba 层 + 全注意力层的模型就可以通过情况 2 中的“全注意力 + X”算法来支持。

实现

概述

KVCacheManager 组织为 3 个层级：

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

上图中的蓝色框展示了包含 10 个全注意力层和 20 个滑动窗口注意力层的情况，因此：

• 使用 HybridKVCacheCoordinator
• 对于 3 个 KVCacheGroup，使用 1 个 FullAttentionManager 和 2 个 SlidingWindowManager。

内存布局

对于一个具有 n 个 KVCacheGroup 的模型（每个组有 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 Cache 根据已分配的 block_ids 保存到缓冲区的不同片段中。

注意

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