跳至内容

自动前缀缓存

前缀缓存kv-cache块是LLM推理中一种流行的优化技术,用于避免冗余的提示计算。其核心思想很简单——我们缓存已处理请求的kv-cache块,当新请求与先前请求具有相同前缀时复用这些块。由于前缀缓存几乎是零成本且不会改变模型输出,它已被许多公共端点(如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是:

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

注 1

我们只缓存完整的块。

注2

上述哈希键结构并非100%无冲突。理论上,不同的前缀标记仍有可能具有相同的哈希值。为避免多租户环境中的哈希冲突,我们建议使用SHA256作为哈希函数而非默认的内置哈希。vLLM自v0.8.3版本起支持SHA256,需通过命令行参数启用。该方案会导致每个标记约100-200纳秒的性能损耗(5万个标记的上下文约损耗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支持通过可选的请求级盐值隔离前缀缓存复用。在请求中包含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"
}

通过这种设置,缓存共享仅限于明确约定使用相同salt的用户或请求,从而在信任组内实现缓存复用,同时隔离其他组。

注意

引擎V0版本不支持缓存隔离。

数据结构

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: Optional["KVCacheBlock"] = None
    next_free_block: Optional["KVCacheBlock"] = None

有两个设计要点需要强调:

  1. 我们在初始化KV缓存管理器时分配所有KVCacheBlock作为块池。这样可以避免Python对象创建开销,并能够轻松持续追踪所有块。
  2. We introduce doubly linked list pointers directly in the KVCacheBlock, so that we could construct a free queue directly. This gives us two benefits:
    1. 我们可以以O(1)的复杂度将中间元素移动到尾部。
    2. 我们可以避免引入另一个Python队列(例如deque),它会对元素进行包装。

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

Component Overview

  • 块池:一个KVCacheBlock的列表。
  • 空闲块队列:仅存储用于操作的头尾块指针。
  • 缓存块:从哈希键到块ID的映射。
  • 请求块:从请求ID到已分配块ID的映射关系。

操作

块分配

新请求: 调度器为新请求分配KV缓存块的工作流程:

  1. 调度器调用kv_cache_manager.get_computed_blocks()来获取已计算完成的块序列。这是通过对请求中的提示令牌进行哈希处理并查找缓存块来实现的。
  2. The scheduler calls kv_cache_manager.allocate_slots(). It does the following steps:
    1. 计算所需新块的数量,如果没有足够的块可供分配则返回。
    2. "触碰"已计算的内存块。它会将该计算块的引用计数加一,如果该块未被其他请求使用,则将其从空闲队列中移除。这是为了避免这些已计算块被驱逐。具体示例请参阅下一章节说明。
    3. 通过弹出空闲队列头部来分配新块。如果头部块是缓存块,这也会"驱逐"该块,使得其他请求从此无法再重用该块。
    4. 如果分配的块已完全填满令牌,我们会立即将其添加到缓存块中,以便同一批次中的其他请求可以重复使用该块。

运行请求: 调度器为带有KV缓存块分配的运行请求安排工作流程:

  1. The scheduler calls kv_cache_manager.allocate_slots(). It does the following steps:
    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中的块表是仅追加的,这意味着不允许将块表从[0, 3]更改为[0, 1]。因此,我们将为哈希键E-H保留重复的块。当请求被释放时,这种重复将被消除。

免费

当一个请求完成时,如果没有任何其他请求正在使用相关块(引用计数=0),我们会释放该请求占用的所有块。在本示例中,我们释放了请求1及其关联的块2、3、4、8。可以看到被释放的块会以逆序方式添加到空闲队列尾部。这是因为请求的最后一个块通常包含更多哈希标记,被其他请求重复利用的可能性较低,因此应该优先被置换出去。

Free queue after a request us freed

淘汰策略 (LRU)

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

  1. 从空闲队列头部弹出块。这是将被淘汰的最近最少使用(LRU)块。
  2. 从缓存块中移除块ID。
  3. 移除区块哈希。

示例

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

时间点1:缓存为空时收到新请求。 我们分配了4个块。其中3个块已填满并被缓存。第4个块部分填充了4个token中的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

优云智算