跳到内容

自动前缀缓存

前缀缓存 KV 缓存块是 LLM 推理中一种流行的优化方法,旨在避免冗余的提示词(prompt)计算。其核心思想很简单——我们缓存已处理请求的 KV 缓存块,当收到与之前请求具有相同前缀的新请求时,复用这些块。由于前缀缓存几乎是“免费午餐”且不会改变模型输出,它已被许多公共端点(如 OpenAI、Anthropic 等)和大多数开源 LLM 推理框架(如 SGLang)广泛使用。

虽然实现前缀缓存的方法有很多,但 vLLM 选择了一种基于哈希的方法。具体来说,我们根据块中的 token 以及该块之前的前缀 token 对每个 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 缓存可以通过 token “A gentle breeze stirred” 进行唯一标识。第三个块可以通过该块中的 token “laughed in the distance” 以及前缀 token “A gentle breeze stirred the leaves as children” 进行唯一标识。因此,我们可以构建块哈希 hash(tuple[components]),其中 components 为

  • 父哈希值:父哈希块的哈希值。
  • 块 token:该块中 token 的元组。包含精确 token 的原因是为了减少潜在的哈希碰撞。
  • 额外哈希:使该块唯一所需的其他值,例如 LoRA ID、多模态输入哈希(参见下文示例)以及用于隔离多租户环境缓存的缓存盐值(cache salts)。

注 1

我们仅缓存完整的块。

注 2

在之前的版本中,哈希键不能保证无碰撞。自 v0.11 起,默认哈希算法为 sha256,解决了碰撞风险。

对于 vllm serve,您可以通过 --prefix-caching-hash-algo 控制哈希算法: - sha256 (默认):使用 Python 的 pickle 进行序列化。哈希值在不同 Python 或 vLLM 版本之间可能无法复现。 - sha256_cbor:使用 cbor2 进行序列化,提供可复现、跨语言兼容的哈希值。推荐用于跨环境的确定性缓存。 - xxhash:使用带有 xxHash (128-bit) 的 Pickle 序列化,实现更快的非加密哈希。需要安装可选的 xxhash 包。重要提示:使用不被视为加密安全的哈希算法在理论上增加了哈希碰撞的风险,这可能导致未定义的行为,甚至在多租户环境中泄露私有信息。即使碰撞概率极低,在开启此功能前,请务必权衡安全风险承受能力与性能收益。 - xxhash_cbor:结合规范的 CBOR 序列化与 xxHash 进行可复现哈希。需要安装可选的 xxhash 包。

多模态输入下的哈希示例
在此示例中,我们说明前缀缓存如何处理多模态输入(例如图像)。假设我们有一个包含以下消息的请求

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]

正如我们所见,在分词(tokenization)之后,[IMG] 将被替换为一系列占位符 token,这些占位符在预填充(prefill)期间会被替换为图像嵌入(image embeddings)。前缀缓存支持此场景的挑战在于我们需要区分图像与占位符。为了解决这个问题,我们对前端图像处理器生成的图像哈希进行编码。例如,上述提示词中块的哈希值将为(假设块大小为 16,且有 41 个占位符 token)

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 支持通过可选的每请求盐值(per-request salting)来隔离前缀缓存复用。通过在请求中包含 cache_salt,该值被注入到第一个块的哈希中,确保只有具有相同盐值的请求才能复用缓存的 KV 块。这防止了攻击者通过观察延迟差异来推断缓存内容的计时攻击。这在不损害性能的前提下提供了保护。

{
  "messages": [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": "Here is a document with details about the world series: ..."},
    {"role": "user", "content": "Who won the world series in 2020?"}
  ],
  "cache_salt": "your-cache-salt"
}

通过此设置,缓存共享仅限于明确同意使用共同盐值的用户或请求,从而在信任组内实现缓存复用,同时隔离其他用户。

数据结构

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: BlockHash
    # 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: "KVCacheBlock | None" = None
    next_free_block: "KVCacheBlock | None" = 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() 来获取已经计算出的块序列。这是通过对请求中的提示词 token 进行哈希并查找缓存块来完成的。
  2. 调度器调用 kv_cache_manager.allocate_slots()。它执行以下步骤
    1. 计算所需新块的数量,如果剩余块不足以分配则返回。
    2. “触摸”(Touch)计算出的块。它将该计算块的引用计数增加 1,如果该块未被其他请求使用,则将其从空闲队列中移除。这是为了防止这些已计算块被逐出。详情见下一节的示例。
    3. 通过弹出空闲队列的头部来分配新块。如果头块是缓存块,这也意味着“逐出”该块,从而使其他请求从现在起无法再复用它。
    4. 如果一个已分配的块已存满 token,我们立即将其添加到缓存块映射中,以便该块可以被同一批次中的其他请求复用。

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

  1. 调度器调用 kv_cache_manager.allocate_slots()。它执行以下步骤
    1. 计算所需新块的数量,如果剩余块不足以分配则返回。
    2. 通过弹出空闲队列的头部来分配新块。如果头块是缓存块,这也意味着“逐出”该块,从而使其他请求从现在起无法再复用它。
    3. 将 token ID 追加到现有块以及新块的插槽中。如果一个块已满,我们将其添加到缓存块中以进行缓存。

重复块
假设块大小为 4,您发送了一个请求(请求 1),其提示词为 ABCDEF,解码长度为 3

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 中的块表是仅追加的(append-only),这意味着不允许将块表从 [0, 3] 更改为 [0, 1]。因此,对于哈希键 E-H,我们将拥有重复的块。当请求被释放时,这种重复将被消除。

释放

当一个请求完成时,如果没有任何其他请求在使用其块(引用计数 = 0),我们就释放所有关联的块。在此示例中,我们释放请求 1 以及与之关联的块 2、3、4、8。我们可以看到被释放的块以相反的顺序被添加到空闲队列的尾部。这是因为请求的最后一个块包含更多的 token 哈希,且被其他请求复用的可能性较低。因此,它应该被优先逐出。

Free queue after a request us freed

逐出 (LRU)

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

  1. 从空闲队列头部弹出该块。这是要被逐出的 LRU 块。
  2. 从缓存块映射中移除该块 ID。
  3. 移除该块的哈希值。

示例

在此示例中,我们假设块大小为 4(每个块可缓存 4 个 token),并且 KV 缓存管理器中总共有 10 个块。

时间 1:缓存为空,收到新请求。 我们分配 4 个块。其中 3 个已存满并缓存。第四个块存有 4 个 token 中的 3 个,处于部分满状态。

Example Time 1

时间 2:请求 0 将块 3 填满并请求新块以继续解码。 我们缓存块 3 并分配块 4。

Example Time 2

时间 3:请求 1 带着 14 个提示词 token 到来,前 10 个 token 与请求 0 相同。 可以看到只有前 2 个块(8 个 token)命中缓存,因为第 3 个块只匹配了 4 个 token 中的 2 个。

Example Time 3

时间 4:请求 0 完成并释放。 块 2、3 和 4 以相反顺序加入空闲队列(但块 2 和 3 仍被缓存)。块 0 和 1 因为正在被请求 1 使用,所以没有被加入空闲队列。

Example Time 4

时间 5:请求 1 完成并释放。

Example Time 5

时间 6:请求 2 带着 29 个提示词 token 到来,前 12 个 token 与请求 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 6