实现

实现#

PagedAttention 的核心思想是将每个请求的 KV 缓存划分为 KV 块。每个块包含固定数量的 token 的注意力键和值。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缓存管理为我们提供了灵活性，以处理更复杂的提供场景，并实现上述策略之外的更复杂的驱逐策略：

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

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