vLLM 自动前缀缓存

自动前缀缓存（Automatic Prefix Caching，APC） 通过缓存已有查询的 KV Cache，使新查询在与已有查询共享相同前缀时可直接复用对应 KV Cache，从而跳过共享部分的重复计算。

1. 实现细节

PagedAttention 将每个请求的 KV Cache 划分为多个固定大小的 KV 块。每个块包含固定数量 Token 的注意力键值对，并存储在非连续的物理内存中以消除碎片。

APC 的关键观察是：每个 KV 块可通过块内 Token 与该块之前的 prefix token 唯一标识。

                    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 ---->|

• Block 1 可由块内 Token 序列 "A gentle breeze stirred" 唯一标识。
• Block 3 需同时包含块内 Token "laughed in the distance" 与前缀 Token "A gentle breeze stirred the leaves as children"。

由此建立映射：

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

2. vLLM 的 KV Cache 管理优化

通过引入该映射，vLLM 在逻辑 KV 块与物理块之间增加了一层间接性：

• 逻辑 KV 块映射到其哈希值。
• 全局哈希表存储所有物理块。
• 相同哈希值的 KV 块（如不同请求共享的前缀块）可映射到同一物理块，实现内存共享。

该设计无需维护树形结构即可实现自动前缀缓存。所有 KV 块彼此独立，可单独分配与释放，管理方式类似操作系统普通缓存。

3. 通用缓存策略

将 KV 块存储在哈希表中，使 vLLM 可缓存早期请求的 KV 块。例如，若新请求与之前请求共享相同的 System Prompt，则可直接复用该 Prompt 对应的 KV Cache，无需重新计算。

当缓存空间满时，vLLM 采用以下淘汰策略：

1. 优先淘汰引用计数为 0 的 KV 块（即当前无请求使用）。
2. 若存在多个引用计数为 0 的块，优先淘汰最近最少使用（LRU）的块。
3. 若最近访问时间相同，优先淘汰位于最长前缀末尾的块。

该策略在完整注意力模型中等效于 RadixAttention 的“优先淘汰引用计数为 0 且最近未使用的前缀树叶节点”策略。

4. 扩展能力

基于哈希的 KV Cache 管理为更复杂的在线服务场景提供了灵活性：

4.1 多 LoRA 模型服务

同时服务多个 LoRA 适配器时，可在计算 KV 块哈希时加入 LoRA ID，使不同适配器的 KV 块正确区分并缓存，统一管理与提高全局缓存命中率。

4.2 多模态模型

对于图像、音频等非离散输入，可采用不同哈希方式。例如，图像输入可使用感知哈希（Perceptual Hashing）缓存相似图像。

5. 启用 APC

在 vLLM 引擎中设置 enable_prefix_caching=True 即可启用 APC。

import time
from vllm import LLM, SamplingParams

LONG_PROMPT = """You are a helpful assistant..."""  # 长文本或表格

def get_generation_time(llm, sampling_params, prompts):
    start = time.time()
    output = llm.generate(prompts, sampling_params=sampling_params)
    elapsed = time.time() - start
    print(f"Output: {output[0].outputs[0].text}")
    print(f"Generation time: {elapsed:.2f} s")

llm = LLM(
    model="lmsys/longchat-13b-16k",
    enable_prefix_caching=True,
)
sampling_params = SamplingParams(temperature=0, max_tokens=100)

# 首次查询，需计算 LONG_PROMPT 的 KV Cache
get_generation_time(
    llm, sampling_params,
    LONG_PROMPT + "Question: what is the age of John Doe? Your answer: The age of John Doe is ",
)

# 第二次查询共享前缀，可直接复用缓存
get_generation_time(
    llm, sampling_params,
    LONG_PROMPT + "Question: what is the age of Zack Blue? Your answer: The age of Zack Blue is ",
)

6. 典型应用场景

1. 长文档查询：用户对同一份长文档多次查询时，APC 只需处理一次文档，后续请求复用其 KV Cache。
2. 多轮对话：同一会话中多次交互时，APC 可跨后续轮次复用历史对话的 KV Cache。

7. 限制说明

APC 通常不会降低 vLLM 性能，但需注意：

• APC 仅优化 Prefill 阶段耗时，不会减少 Decode 阶段耗时。
• 当大部分时间用于生成长答案，或新查询与现有查询无共享前缀时，APC 不会带来性能提升。

参考

• vLLM Prefix Caching Documentation