Robin's AI Lab

主题切换

Automatic Prefix Caching ​

简介 ​

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

Automatic Prefix Caching 实现细节 ​

PagedAttention的核心思想是将每个请求的KV缓存划分为多个KV块。每个块包含固定数量token的注意力键值对。PagedAttention 算法允许这些块存储在非连续的物理内存中，从而通过按需分配内存来消除内存碎片。

我们基于以下关键观察实现KV缓存的自动化管理：每个KV块可以通过块内token和该块之前的prefix token唯一标识。

shell
                    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"唯一标识
  • 第三个KV块则需要同时包含：
    • 块内token序列"laughed in the distance"
    • 前缀token序列"A gentle breeze stirred the leaves as children"

由此可建立严格的映射关系：

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

vLLM 的键值缓存（KV Cache）管理优化 ​

通过引入这种映射机制，我们在 vLLM 的键值（KV）缓存管理中增加了一层间接性。以往，vLLM 中的每个序列都维护一个从逻辑 KV 块到物理块的映射。为了实现 KV 块的自动缓存，我们将逻辑 KV 块映射到其哈希值，并维护一个全局的哈希表，用于存储所有的物理块。通过这种方式，所有哈希值相同的 KV 块（例如，不同请求中共享的前缀块）可以映射到相同的物理块，从而共享内存空间。

该设计无需在 KV 块之间维护树形结构即可实现自动前缀缓存。具体而言，所有 KV 块都是彼此独立的，可以单独分配和释放，这使我们可以将 KV 缓存管理方式类比为操作系统中的普通缓存管理。

通用缓存策略（Generalized Caching Policy） ​

将所有 KV 块存储在哈希表中，使得 vLLM 可以缓存来自早期请求的 KV 块，从而节省内存并加速后续请求的计算。例如，如果一个新请求与之前的请求共享相同的系统提示词（system prompt），则可以直接复用共享提示词对应的 KV 缓存，无需重新计算。

然而，由于 KV 缓存空间有限，我们必须在缓存空间满时决定保留哪些 KV 块、淘汰哪些块。

使用哈希表管理 KV 缓存使我们可以实现灵活的缓存策略。例如，在当前 vLLM 中，我们实现了如下的淘汰策略：

  1. 当没有可用空闲块时，我们优先淘汰引用计数为 0 的 KV 块（即当前无任何请求使用该块）。
  2. 如果存在多个引用计数为 0 的块，则优先淘汰最近最少使用（LRU）的块。
  3. 如果存在多个最近访问时间相同的块，则优先淘汰位于最长前缀末尾的块（即该块前面拥有最多块的情况）。

值得注意的是，该淘汰策略在应用于完整注意力（full attention）模型时，实际上等效于 RadixAttention 中的策略：优先淘汰引用计数为 0 且最近未被使用的前缀树叶节点。

哈希表机制的扩展能力 ​

基于哈希的 KV 缓存管理机制为我们处理更复杂的在线服务场景提供了更大的灵活性，也支持实现比上述更复杂的淘汰策略，例如：

多 LoRA 模型服务（Multi-LoRA Serving） ​

当同时服务多个 LoRA 适配器请求时，我们可以在计算 KV 块哈希值时加入 LoRA ID，使每个请求对应的 KV 块能够正确区分并缓存。通过这种方式，我们可以统一管理所有适配器的 KV 块，简化系统实现，并提高全局缓存命中率与使用效率。

多模态模型（Multi-modal Models） ​

当用户输入包含非离散的模态（例如图像、音频等）时，我们可以采用不同的哈希方式来缓存不同模态的输入。例如，对于图像输入，可以使用感知哈希（perceptual hashing）方法，以便缓存相似的输入图像。

在vLLM中启用APC ​

在vLLM引擎中设置enable_prefix_caching=True即可启用APC。以下为示例代码：

python
import time
from vllm import LLM, SamplingParams


# A prompt containing a large markdown table. The table is randomly generated by GPT-4.
LONG_PROMPT = "You are a helpful assistant in recognizes the content of tables in markdown format. Here is a table as follows.\n# Table\n" + """
| ID  | Name          | Age | Occupation | Country     | Email                  | Phone Number | Address                         |
| --- | ------------- | --- | ---------- | ----------- | ---------------------- | ------------ | ------------------------------- |
| 1   | John Doe      | 29  | Engineer   | USA         | john.doe@example.com   | 555-1234     | 123 Elm St, Springfield, IL     |
| 2   | Jane Smith    | 34  | Doctor     | Canada      | jane.smith@example.com | 555-5678     | 456 Oak St, Toronto, ON         |
| 3   | Alice Johnson | 27  | Teacher    | UK          | alice.j@example.com    | 555-8765     | 789 Pine St, London, UK         |
| 4   | Bob Brown     | 45  | Artist     | Australia   | bob.b@example.com      | 555-4321     | 321 Maple St, Sydney, NSW       |
| 5   | Carol White   | 31  | Scientist  | New Zealand | carol.w@example.com    | 555-6789     | 654 Birch St, Wellington, NZ    |
| 6   | Dave Green    | 28  | Lawyer     | Ireland     | dave.g@example.com     | 555-3456     | 987 Cedar St, Dublin, IE        |
| 7   | Emma Black    | 40  | Musician   | USA         | emma.b@example.com     | 555-1111     | 246 Ash St, New York, NY        |
| 8   | Frank Blue    | 37  | Chef       | Canada      | frank.b@example.com    | 555-2222     | 135 Spruce St, Vancouver, BC    |
| 9   | Grace Yellow  | 50  | Engineer   | UK          | grace.y@example.com    | 555-3333     | 864 Fir St, Manchester, UK      |
| 10  | Henry Violet  | 32  | Artist     | Australia   | henry.v@example.com    | 555-4444     | 753 Willow St, Melbourne, VIC   |
| 11  | Irene Orange  | 26  | Scientist  | New Zealand | irene.o@example.com    | 555-5555     | 912 Poplar St, Auckland, NZ     |
| 12  | Jack Indigo   | 38  | Teacher    | Ireland     | jack.i@example.com     | 555-6666     | 159 Elm St, Cork, IE            |
| 13  | Karen Red     | 41  | Lawyer     | USA         | karen.r@example.com    | 555-7777     | 357 Cedar St, Boston, MA        |
| 14  | Leo Brown     | 30  | Chef       | Canada      | leo.b@example.com      | 555-8888     | 246 Oak St, Calgary, AB         |
| 15  | Mia Green     | 33  | Musician   | UK          | mia.g@example.com      | 555-9999     | 975 Pine St, Edinburgh, UK      |
| 16  | Noah Yellow   | 29  | Doctor     | Australia   | noah.y@example.com     | 555-0000     | 864 Birch St, Brisbane, QLD     |
| 17  | Olivia Blue   | 35  | Engineer   | New Zealand | olivia.b@example.com   | 555-1212     | 753 Maple St, Hamilton, NZ      |
| 18  | Peter Black   | 42  | Artist     | Ireland     | peter.b@example.com    | 555-3434     | 912 Fir St, Limerick, IE        |
| 19  | Quinn White   | 28  | Scientist  | USA         | quinn.w@example.com    | 555-5656     | 159 Willow St, Seattle, WA      |
| 20  | Rachel Red    | 31  | Teacher    | Canada      | rachel.r@example.com   | 555-7878     | 357 Poplar St, Ottawa, ON       |
| 21  | Steve Green   | 44  | Lawyer     | UK          | steve.g@example.com    | 555-9090     | 753 Elm St, Birmingham, UK      |
| 22  | Tina Blue     | 36  | Musician   | Australia   | tina.b@example.com     | 555-1213     | 864 Cedar St, Perth, WA         |
| 23  | Umar Black    | 39  | Chef       | New Zealand | umar.b@example.com     | 555-3435     | 975 Spruce St, Christchurch, NZ |
| 24  | Victor Yellow | 43  | Engineer   | Ireland     | victor.y@example.com   | 555-5657     | 246 Willow St, Galway, IE       |
| 25  | Wendy Orange  | 27  | Artist     | USA         | wendy.o@example.com    | 555-7879     | 135 Elm St, Denver, CO          |
| 26  | Xavier Green  | 34  | Scientist  | Canada      | xavier.g@example.com   | 555-9091     | 357 Oak St, Montreal, QC        |
| 27  | Yara Red      | 41  | Teacher    | UK          | yara.r@example.com     | 555-1214     | 975 Pine St, Leeds, UK          |
| 28  | Zack Blue     | 30  | Lawyer     | Australia   | zack.b@example.com     | 555-3436     | 135 Birch St, Adelaide, SA      |
| 29  | Amy White     | 33  | Musician   | New Zealand | amy.w@example.com      | 555-5658     | 159 Maple St, Wellington, NZ    |
| 30  | Ben Black     | 38  | Chef       | Ireland     | ben.b@example.com      | 555-7870     | 246 Fir St, Waterford, IE       |
"""


def get_generation_time(llm, sampling_params, prompts):
    # time the generation
    start_time = time.time()
    output = llm.generate(prompts, sampling_params=sampling_params)
    end_time = time.time()
    # print the output and generation time
    print(f"Output: {output[0].outputs[0].text}")
    print(f"Generation time: {end_time - start_time} seconds.")


# set enable_prefix_caching=True to enable APC
llm = LLM(
    model='lmsys/longchat-13b-16k',
    enable_prefix_caching=True
)

sampling_params = SamplingParams(temperature=0, max_tokens=100)

# Querying the age of John Doe
get_generation_time(
    llm,
    sampling_params,
    LONG_PROMPT + "Question: what is the age of John Doe? Your answer: The age of John Doe is ",
)

# Querying the age of Zack Blue
# This query will be faster since vllm avoids computing the KV cache of LONG_PROMPT again.
get_generation_time(
    llm,
    sampling_params,
    LONG_PROMPT + "Question: what is the age of Zack Blue? Your answer: The age of Zack Blue is ",
)

典型应用场景 ​

我们描述两种APC能显著提升性能的场景：

  1. 长文档查询

    用户对同一份长文档（如软件手册或年度报告）进行多次不同查询时，APC允许vLLM仅需处理一次长文档，后续所有请求均可通过复用其KV缓存来避免重复处理。这使得vLLM能以更高吞吐量和更低延迟服务后续请求。

  2. 多轮对话

    用户在同一会话中与应用进行多次交互时，APC允许vLLM跨所有后续对话轮次复用历史对话的处理结果，从而显著提升后续请求的吞吐量和降低延迟。

限制说明 ​

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

  • APC仅优化查询处理阶段（预填充阶段）耗时，不会减少新token生成阶段（解码阶段）耗时
  • 当vLLM大部分时间用于生成长答案时，或新查询与现有查询无共享前缀时，APC不会带来性能提升

Maintained by Robin

© 2017 — 2025. Powered by VitePress