自动前缀缓存

自动前缀缓存#

PagedAttention 的核心思想是将每个请求的 KV 缓存划分为 KV 块。每个块包含固定数量令牌的注意力键和值。 PagedAttention 算法允许将这些块存储在非连续的物理内存中，以便我们可以通过按需分配内存来消除内存碎片。

为了自动缓存 KV 缓存，我们利用了以下关键观察：每个 KV 块都可以通过块内的令牌和块之前的前缀中的令牌来唯一标识。

                    Block 1                  Block 2                  Block 3
         [A gentle breeze stirred] [the leaves as children] [laughed in the distance]
Block 1: |<--- block tokens ---->|
Block 2: |<------- prefix ------>| |<--- block tokens --->|
Block 3: |<------------------ prefix -------------------->| |<--- block tokens ---->|

在上面的示例中，第一个块中的 KV 缓存可以通过令牌 “A gentle breeze stirred” 唯一标识。第三个块可以通过块中的令牌 “laughed in the distance” 以及前缀令牌 “A gentle breeze stirred the leaves as children” 唯一标识。因此，我们可以构建以下一对一映射

hash(prefix tokens + block tokens) <--> KV Block

通过此映射，我们可以在 vLLM 的 KV 缓存管理中添加另一个间接层。 以前，vLLM 中的每个序列都维护从其逻辑 KV 块到物理块的映射。 为了实现 KV 块的自动缓存，我们将逻辑 KV 块映射到其哈希值，并维护所有物理块的全局哈希表。 这样，所有共享相同哈希值的 KV 块（例如，跨两个请求的共享前缀块）都可以映射到同一个物理块并共享内存空间。

这种设计实现了自动前缀缓存，而无需维护 KV 块之间的树结构。 更具体地说，所有块都是彼此独立的，可以自行分配和释放，这使我们能够像操作系统中的普通缓存一样管理 KV 缓存。

通用缓存策略#

将所有 KV 块保存在哈希表中使 vLLM 能够缓存来自较早请求的 KV 块，以节省内存并加速未来请求的计算。 例如，如果新请求与之前的请求共享系统提示，则共享提示的 KV 缓存可以直接用于新请求，而无需重新计算。 但是，KV 缓存的总空间是有限的，我们必须决定在缓存已满时保留或逐出哪些 KV 块。

使用哈希表管理 KV 缓存使我们能够实现灵活的缓存策略。 例如，在当前的 vLLM 中，我们实现了以下逐出策略

  • 当没有剩余的空闲块时，我们将逐出一个引用计数（即，使用该块的当前请求数）等于 0 的 KV 块。

  • 如果存在多个引用计数等于 0 的块，我们将优先逐出最近最少使用的块 (LRU)。

  • 如果存在多个上次访问时间相同的块，我们将优先逐出位于最长前缀末尾的块（即，在其之前具有最大块数的块）。

请注意，当应用于具有完全注意力的模型时，此逐出策略有效地实现了与 RadixAttention 中完全相同的策略，该策略优先逐出引用计数为零和前缀树中最近最少使用的叶节点。

但是，基于哈希的 KV 缓存管理使我们能够灵活地处理更复杂的服务场景，并实现超出上述策略的更复杂的逐出策略

  • Multi-LoRA 服务。 在为多个 LoRA 适配器服务请求时，我们可以简单地让每个 KV 块的哈希值也包含请求查询的 LoRA ID，以便为所有适配器启用缓存。 这样，我们可以联合管理不同适配器的 KV 块，从而简化系统实现并提高全局缓存命中率和效率。

  • 多模态模型。 当用户输入不仅仅包含离散令牌时，我们可以使用不同的哈希方法来处理不同模态输入的缓存。 例如，图像的感知哈希用于缓存相似的输入图像。