LLM 推理优化 Prefix Caching 及其实现

Prefix Caching

上文中提到，Prompt 计算（Prefill 阶段）和生成阶段的计算特性很不相同。为避免重复计算，所有框架的 prefill 阶段主要作用就是给迭代的生成阶段准备 KV Cache。但这些 KV Cache 仅仅是为单次生成式请求服务的，那很自然的一种想法就是，KV Cache 能不能跨请求复用？

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

限制

但 KV Cache 跟其它服务缓存不一样的地方是，它太大了，以至于（目前）很难通过 Redis/Memcache 这种分布式缓存服务存取。比如对 13B LLM 模型来说，在 FP16 精度下单 token 的 KV Cache 大约是 1MB，假设要缓存的前缀有 500 个 token（大约800多个汉字），那就是 500MB。一般来说，我们不会每次请求去从分布式系统里读取/传输 500MB 的缓存，甚至都不会每次请求从内存往显存中拷贝 500MB 的缓存，所以大部分情况下，prefix cache 都会放在显存里。

这也就意味着，如果你想命中 prefix cache，必须把相同 prefix 的请求发到同一张 GPU卡上才行。

实现

由于不是普遍需求，加上前面说的限制，prefix caching 作为一个加速特性，不是很受关注，一般也不是默认开启的。各框架的实现和配置略有差异，这里简单做下记录，便于回顾。

刚开始 vLLM 的实现是给 generate 接口增加一个 prefix_pos 参数，通过 prefix_pos 输入参数为每个请求指定 prefix 长度，为 prefix 建一个带淘汰的哈希缓存。后来觉得这样做使用上不够便利，升级成了自动前缀缓存，即将 prompt 的 kv cache 分成 block，然后为 block 建设 LRU 缓存机制，这样就不必在接口上使用 prefix_pos 指定哪部分是 prefix 了。自动前缀缓存功能默认是不开启的，开启的配置项为 --enable-prefix-caching 。

TensorRT-LLM 与 vLLM 后来的实现类似，也是实现了 block kv cache，配置项是 enableBlockReuse，默认也是不开启的。代码未开源，无法看到实现。

Lmdeploy 的 Python 和 TurboMind C++ 版本的 prefix caching 功能都已经有了 PR，但现在（20240425）看还没有合入主干。有意思的是它没有使用 hash block 对应 token_id 子串的所有 token_id 前缀然后组成哈希表的方式，而是用 hash 当前 block 对应的 token_id 子串然后组成 trie 树的缓存管理结构。默认的参数名与 vLLM 相同，也叫做 --enable-prefix-caching。

HuggingFace TGI 现在看起来还没实现 Prefix Caching 功能。

Prompt Caching

除了 Prefix Caching 这种比较直观的工程优化，现在也有一些研究在看 Prompt 的其它缓存机制。比如设计一种机制让 prompt 模块化，不仅可以复用 prefix，还能复用中间的部分；或者通过 query 的相似性复用其它 query 的 prompt。

但目前看实现上都过于复杂，比如第一种要求模型使用不连续的 poition_id，这样就可以插入 token，但这种方式对 attention 的计算机制有一定的影响，难以说明它对效果的影响。