LLM 推理优化 Prefix Caching 及其实现

LLM 推理优化 Prefix Caching 及其实现

一、Prefix Caching 概述

Prefix Caching 是一种针对大型语言模型（LLM）推理过程中的优化技术。在 LLM 的推理过程中，Prompt 计算（Prefill 阶段）和生成阶段的计算特性存在显著差异。为避免重复计算，Prefill 阶段的主要作用是为迭代的生成阶段准备 KV Cache（键值缓存）。然而，传统的 KV Cache 仅为单次生成式请求服务，无法跨请求复用。

在某些场景下，多次请求的 Prompt 可能会共享同一个前缀（Prefix），如拟人 Agent 的人物设定、文档阅读理解时的文档内容等。这些情况下，许多请求的前缀的 KV Cache 计算结果是相同的。因此，可以像互联网服务的请求缓存一样，将这些相同的 KV Cache 缓存起来，供下一个请求复用。

二、Prefix Caching 的特点与限制

Prefix Caching 与其他服务缓存的主要区别在于其庞大的体积。对于大型 LLM 模型（如 13B 参数模型），在 FP16 精度下，单 token 的 KV Cache 大小约为 1MB。若需缓存的前缀包含 500 个 token（约 800 多个汉字），则 KV Cache 的大小将达到 500MB。

由于这种庞大的体积，Prefix Cache 很难通过 Redis/Memcache 等分布式缓存服务进行存取。同时，每次请求从分布式系统读取/传输如此大的缓存，或从内存往显存中拷贝如此大的缓存，都是不现实的。因此，Prefix Cache 通常会被放置在显存中。这也意味着，若想命中 Prefix Cache，必须将具有相同前缀的请求发送到同一张 GPU 卡上。

三、Prefix Caching 的实现

由于 Prefix Caching 并非普遍需求，且受到上述限制，因此它通常作为加速特性存在，并非默认开启。各框架的实现和配置略有差异，以下是一些主要框架的实现方式：

1. vLLM：

   初始实现：给 generate 接口增加一个 prefix_pos 参数，通过该参数为每个请求指定前缀长度，并为前缀建立一个带淘汰机制的哈希缓存。

   升级实现：后来升级为自动前缀缓存，即将 prompt 的 KV Cache 分成 block，并为 block 建设 LRU（最近最少使用）缓存机制。这样就不必在接口上使用 prefix_pos 指定哪部分是前缀了。自动前缀缓存功能默认不开启，可通过 --enable-prefix-caching 配置项开启。

2. TensorRT-LLM：

   实现方式：与 vLLM 后来的实现类似，实现了 block KV Cache。配置项为 enableBlockReuse，默认不开启。由于代码未开源，无法查看具体实现。

3. Lmdeploy：

   Python 和 TurboMind C++ 版本：已经有了关于 Prefix Caching 功能的 PR（Pull Request），但截至 2024 年 4 月 25 日，尚未合入主干。

   实现特点：没有使用哈希 block 对应 token_id 子串的所有 token_id 前缀然后组成哈希表的方式，而是使用哈希当前 block 对应的 token_id 子串然后组成 trie 树的缓存管理结构。默认的参数名与 vLLM 相同，也叫做 --enable-prefix-caching。

4. HuggingFace TGI：

   目前尚未实现 Prefix Caching 功能。

四、Prompt Caching 的其他研究方向

除了 Prefix Caching 这种直观的工程优化外，目前还有一些研究在探索 Prompt 的其他缓存机制。例如：

• 设计一种机制让 prompt 模块化，不仅可以复用前缀，还能复用中间的部分。
• 通过 query 的相似性复用其他 query 的 prompt。

然而，这些实现方式目前都过于复杂。例如，第一种方式要求模型使用不连续的 position_id，以便插入 token。但这种方式对 attention 的计算机制有一定影响，且难以说明其对效果的影响。

综上所述，Prefix Caching 是一种针对 LLM 推理过程中的优化技术，它通过缓存相同前缀的 KV Cache 来避免重复计算。然而，由于其庞大的体积和特定的存储要求，Prefix Caching 的实现和配置在不同框架中存在差异。同时，也有一些研究在探索 Prompt 的其他缓存机制，但目前尚未有成熟的实现方案。