自动前缀缓存

自动前缀缓存#

前缀缓存 kv-cache 块是 LLM 推理中一种流行的优化方法，用于避免冗余的提示计算。核心思想很简单——我们缓存已处理请求的 kv-cache 块，并在新请求与之前的请求具有相同前缀时重用这些块。由于前缀缓存几乎是免费的午餐，并且不会更改模型输出，因此已被许多公共端点（例如，OpenAI、Anthropic 等）和大多数开源 LLM 推理框架（例如，SGLang）广泛使用。

虽然有很多方法可以实现前缀缓存，但 vLLM 选择了基于哈希的方法。具体来说，我们通过块中的令牌和块之前的前缀中的令牌来哈希每个 kv-cache 块

                    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(tuple[components])，其中组件是

  • 父哈希值：父哈希块的哈希值。

  • 块令牌：此块中令牌的元组。包含确切令牌的原因是为了减少潜在的哈希值冲突。

  • 额外哈希：使此块唯一的其他值，例如 LoRA ID 和多模态输入哈希（请参见下面的示例）。

注意 1： 我们仅缓存完整块。

注意 2： 上述哈希键结构并非 100% 无冲突。理论上，不同的前缀令牌仍然可能具有相同的哈希值。为了避免多租户设置中出现任何哈希冲突，我们建议使用 SHA256 作为哈希函数，而不是默认的内置哈希。vLLM v0.8.3 及更高版本支持 SHA256，必须使用命令行参数启用。它会带来大约 100-200ns/令牌的性能影响（对于 50k 上下文令牌约为 6 毫秒）。

多模态输入的哈希示例
在此示例中，我们说明了前缀缓存如何与多模态输入（例如，图像）一起使用。假设我们有一个包含以下消息的请求

messages = [
    {"role": "user",
     "content": [
         {"type": "text",
          "text": "What's in this image?"
         },
         {"type": "image_url",
          "image_url": {"url": image_url},
         },
    ]},
]

它将变成以下提示

Prompt:
    <s>[INST]What's in this image?\n[IMG][/INST]

Tokenized prompt:
    [1, 3, 7493, 1681, 1294, 1593, 3937, 9551, 10, 4]

Prompt with placeholders (<P>):
    [1, 3, 7493, 1681, 1294, 1593, 3937, 9551, <P>, <P>, ..., <P>, 4]

正如我们所看到的，在令牌化之后，[IMG] 将被一系列占位符令牌替换，并且这些占位符将在预填充期间被图像嵌入替换。前缀缓存支持这种情况的挑战是我们需要区分图像和占位符。为了解决这个问题，我们对前端图像处理器生成的图像哈希进行编码。例如，上述提示中块的哈希将是（假设块大小为 16，并且我们有 41 个占位符令牌）

Block 0
    Parent hash: None
    Token IDs: 1, 3, 7493, 1681, 1294, 1593, 3937, 9551, <p>, ..., <p>
    Extra hash: <image hash>
Block 1
    Parent hash: Block 0 hash
    Token IDs: <p>, ..., <p>
    Extra hash: <image hash>
Block 2
    Parent hash: Block 1 hash
    Token IDs: <p>, ..., <p>
    Extra hash: <image hash>
Block 3
    Parent hash: Block 2 hash
    Token IDs: <p>, ..., <p>, 4
    Extra hash: <image hash>

在本文档的其余部分，我们首先介绍 vLLM v1 中用于前缀缓存的数据结构，然后介绍主要 KV 缓存运算符（例如，分配、追加、释放、驱逐）的前缀缓存工作流程。最后，我们使用一个示例来说明端到端的前缀缓存工作流程。

数据结构#

vLLM v1 中的前缀缓存在 KV 缓存管理器中实现。基本构建块是“Block”数据类（简化版）

class KVCacheBlock:
    # The block ID (immutable)
    block_id: int
    # The block hash (will be assigned when the block is full,
    # and will be reset when the block is evicted).
    block_hash: BlockHashType
    # The number of requests using this block now.
    ref_cnt: int

    # The pointers to form a doubly linked list for the free queue.
    prev_free_block: Optional["KVCacheBlock"] = None
    next_free_block: Optional["KVCacheBlock"] = None

以下是需要强调的两个设计要点

  1. 我们在初始化 KV 缓存管理器时分配所有 KVCacheBlock 以作为块池。这避免了 Python 对象创建的开销，并且可以轻松地始终跟踪所有块。

  2. 我们在 KVCacheBlock 中直接引入了双向链表指针，以便我们可以直接构建空闲队列。这为我们带来了两个好处

    1. 我们可以将中间的元素以 O(1) 的复杂度移动到尾部。

    2. 我们可以避免引入另一个 Python 队列（例如，deque），该队列具有元素的包装器。

因此，当 KV 缓存管理器初始化时，我们将具有以下组件

Component Overview

  • 块池：KVCacheBlock 的列表。

  • 空闲块队列：仅存储头部和尾部块的指针以进行操作。

  • 缓存块：从哈希键到块 ID 的映射。

  • 请求块：从请求 ID 到已分配块 ID 的映射。

操作#

块分配#

新请求： 调度器调度带有 KV 缓存块分配的新请求的工作流程

  1. 调度器调用 kv_cache_manager.get_computed_blocks() 以获取已计算的块序列。这是通过哈希请求中的提示令牌并查找缓存块来完成的。

  2. 调度器调用 kv_cache_manager.allocate_slots()。它执行以下步骤

    1. 计算新的所需块的数量，如果没有足够的块可分配，则返回。

    2. “触摸”已计算的块。它将已计算块的引用计数增加 1，如果该块未被其他请求使用，则从空闲队列中删除该块。这是为了避免这些已计算的块被驱逐。有关说明，请参见下一节中的示例。

    3. 通过弹出空闲队列的头部来分配新块。如果头部块是缓存块，则这也“驱逐”该块，以便从现在开始其他请求不再可以重用它。

    4. 如果分配的块已经满了令牌，我们会立即将其添加到缓存块，以便同一批次中的其他请求可以重用该块。

正在运行的请求： 调度器调度带有 KV 缓存块分配的正在运行的请求的工作流程

  1. 调度器调用 kv_cache_manager.append_slots()。它执行以下步骤

    1. 计算新的所需块的数量，如果没有足够的块可分配，则返回。

    2. 通过弹出空闲队列的头部来分配新块。如果头部块是缓存块，则这也“驱逐”该块，以便从现在开始其他请求不再可以重用它。

    3. 将令牌 ID 追加到现有块以及新块中的槽位。如果某个块已满，我们会将其添加到缓存块以缓存它。

重复块
假设块大小为 4，并且您发送一个带有提示 ABCDEF 和解码长度 3 的请求（请求 1）

Prompt: [A, B, C, D, E, F]
Output: [G, H, I]

Time 0:
  Tokens: [A, B, C, D, E, F, G]
  Block Table: [0 (ABCD), 1 (EFG)]
  Cache Blocks: 0
Time 1:
  Tokens: [A, B, C, D, E, F, G, H]
  Block Table: [0 (ABCD), 1 (EFGH)]
  Cache Blocks: 0, 1
Time 2:
  Tokens: [A, B, C, D, E, F, G, H, I]
  Block Table: [0 (ABCD), 1 (EFGH), 2 (I)]
  Cache Blocks: 0, 1

现在块 0 和块 1 已缓存，我们再次发送相同的请求（请求 2）并使用贪婪采样，以便它将生成与请求 1 完全相同的输出

Prompt: [A, B, C, D, E, F]
Output: [G, H, I]

Time 0:
  Tokens: [A, B, C, D, E, F, G]
  Block Table: [0 (ABCD), 3 (EFG)]
  Cache Blocks: 0, 1
Time 1:
  Tokens: [A, B, C, D, E, F, G, H]
  Block Table: [0 (ABCD), 3 (EFGH)]
  Cache Blocks: 0, 1, 3

可以看出，块 3 是一个新的完整块并且已缓存。但是，它是冗余的，因为块 1，这意味着我们缓存了相同的块两次。在 v0 中，当检测到块 3 是重复的时，我们释放块 3 并让请求 2 改为使用块 1，因此其块表在时间 1 中变为 [0, 1]。但是，vLLM v1 中的块表是仅追加的，这意味着不允许将块表从 [0, 3] 更改为 [0, 1]。因此，我们将为哈希键 E-H 获得重复的块。当请求被释放时，将消除此重复。

释放#

当请求完成时，如果没有其他请求正在使用其块（引用计数 = 0），我们将释放其所有块。在此示例中，我们释放请求 1 以及与其关联的块 2、3、4、8。我们可以看到，释放的块以相反的顺序添加到空闲队列的尾部。这是因为请求的最后一个块必须哈希更多的令牌，并且不太可能被其他请求重用。因此，应首先驱逐它。

Free Queue after Free a Request

驱逐 (LRU)#

当空闲队列的头部块（最近最少使用的块）被缓存时，我们必须驱逐该块，以防止其他请求使用它。具体来说，驱逐涉及以下步骤

  1. 从空闲队列的头部弹出块。这是要驱逐的 LRU 块。

  2. 从缓存块中删除块 ID。

  3. 删除块哈希。

示例#

在此示例中，我们假设块大小为 4（每个块可以缓存 4 个令牌），并且 KV 缓存管理器中总共有 10 个块。

时间 1：缓存为空，并且新请求进入。 我们分配 4 个块。其中 3 个已满并已缓存。第四个块部分已满，包含 4 个令牌中的 3 个。

Example Time 1

时间 3：请求 0 使块 3 满，并要求新块以继续解码。 我们缓存块 3 并分配块 4。

Example Time 3

时间 4：请求 1 进入，带有 14 个提示令牌，其中前 10 个令牌与请求 0 相同。 我们可以看到只有前 2 个块（8 个令牌）命中缓存，因为第 3 个块仅匹配 4 个令牌中的 2 个。

Example Time 4

时间 5：请求 0 已完成并释放。 块 2、3 和 4 以相反的顺序添加到空闲队列（但块 2 和 3 仍被缓存）。块 0 和 1 未添加到空闲队列，因为它们正在被请求 1 使用。

Example Time 5

时间 6：请求 1 已完成并释放。

Example Time 6

时间 7：请求 2 进入，带有 29 个提示令牌，其中前 12 个令牌与请求 0 相同。 请注意，即使空闲队列中的块顺序为 7 - 8 - 9 - 4 - 3 - 2 - 6 - 5 - 1 - 0，缓存命中块（即 0、1、2）在分配之前被触摸并从队列中删除，因此空闲队列变为 7 - 8 - 9 - 4 - 3 - 6 - 5。因此，分配的块是 0（已缓存）、1（已缓存）、2（已缓存）、7、8、9、4、3（已驱逐）。

Example Time 7