KV 缓存卸载
KV 缓存卸载(KV cache offloading)是指将注意力的键/值数据从 GPU 内存移动到成本更低的存储(如 CPU 内存或磁盘)的过程。它在保留无需重算即可恢复推理能力的同时释放 GPU 资源。这有助于在性能和内存使用之间取得平衡,从而高效地扩展 LLM 工作负载。
为什么 KV 缓存会成为 LLM 推理中的瓶颈?
LLM 严重依赖 KV 缓存来加速推理。缓存会为输入序列中的每个 token 存储注意力键和值,使模型在后续步骤中能够复用它们,而不是重新计算。虽然这节省了大量的计算资源并带来更快的推理,但它的内存代价很高。
随着上下文窗口的增大,KV 缓存大小随序列长度线性增长。这会很快耗尽可用的 GPU 内存,尤其是在长上下文场景中。由于 GPU 内存有限,KV 缓存常常成为运行需要长上下文的应用的瓶颈。
事实上,并非所有 KV 缓存数据都需要始终驻留在 GPU 内存中。在许多现实应用中,用户可能不会持续与 LLM 交互。例如,用户可能在打字时暂停,或离开数小时后才回来。在这种情况下,他们的 KV 缓存即使没有被主动使用,也仍然占用着 GPU 内存。同样,当多个用户/智能体在不同时间访问同一个对话、文档或会话时,同一个 KV 缓存可能会在 GPU 上闲置(而你并不想为了重新计算同样的内容而浪费 GPU 资源)。
这导致内存使用效率低下:宝贵的 GPU 内存被不活跃的会话占用,而不是用于服务新请求。随着时间推移,这会限制系统能够支持的并发用户数,并降低整体吞吐量。
为解决这些问题,KV 缓存卸载将不活跃或较少访问的缓存数据从 GPU 内存移动到成本更低、容量更大的存储,例如 CPU RAM、本地 SSD 或远程对象存储。当用户恢复交互或另一个用户访问相同内容时,缓存可以按需重新加载到 GPU 内存中。这避免了昂贵的重新计算,同时为活跃工作负载释放了 GPU 资源。
如何计算 KV 缓存大小
在卸载 KV 缓存时,了解它实际消耗多少内存会很有帮助。
在基于 Transformer 的 LLM 中,每个注意力层都需要为输入序列中的每个 token 存储两个向量(一个键和一个值)。每一层包含多个注意力头,而且所有头通常具有相同的维度。
要估算 KV 缓存消耗多少内存,请使用下面的计算器:
你通常可以在 LLM 的 Hugging Face 仓库中的 config.json 文件里找到其架构细节,包括模型架构(例如 Transformer 解码器)、层数、隐藏大小、注意力头数、词表大小,以及其他架构超参数。如果你已经知道模型的维度,就可以用 H × D 替换来简化公式(见上面的"简化计算")。
何时应该为 LLM 卸载 KV 缓存?
KV 缓存卸载在以下情况下尤其有用:
- 你部署的是长上下文窗口的 LLM,这可能导致 KV 缓存迅速超出 GPU 内存。
- 多个用户或智能体需要在会话间与相同的基础内容或上下文交互。例如,在 IDE 中使用 LLM 集成功能的开发人员,往往会反复与同一段代码片段交互。
- 你的部署受内存限制,或者需要优化基础设施成本。
- 你要在许多分布式 worker 上扩展推理,而 GPU 资源有限。
- 你的工作负载包含间歇性或空闲的用户会话,此时将 KV 缓存保留在 GPU 内存中会造成浪费。
KV 缓存卸载的好处
卸载 KV 缓存为扩展和优化 LLM 推理提供了几个重要优势:
- 更好的资源利用率。 将不活跃或共享的 KV 数据移出 GPU 内存,可以为新请求腾出空间。这使得同一块 GPU 无需触达内存上限就能服务更多并发用户或更长的输入序列。
- 更低的计算成本。 GPU 内存昂贵且有限。卸载允许工作负载利用更廉价的存储(例如 CPU RAM 或磁盘),减少仅仅为了管理缓存而过度配置高端 GPU 的需求。
- 降低延迟:卸载让模型在推理期间跳过冗余的 KV 计算,尤其是在多轮交互中存在重叠上下文时。这显著降低了 TTFT 和整体延迟。NVIDIA 报告称,对于大的输入序列,与从头重新计算 KV 缓存相比,KV 缓存卸载可以带来高达 14 倍的 TTFT 加速。
KV 缓存卸载的权衡
虽然 KV 缓存卸载可以显著改善内存效率和吞吐量,但卸载目标的速度至关重要。如果存储层(例如 CPU RAM 或磁盘)太慢,将 KV 数据传回 GPU 的开销可能会抵消收益,尤其是在对延迟敏感的应用中。
请确保数据传输的成本低于从头重新计算缓存。在长多轮对话中通常如此,因为复用之前的上下文至关重要,而重新计算的代价很高。
当系统使用选择性 KV 卸载时,还存在质量上的权衡。在解码期间,运行时可能需要决定哪些键和值应该返回 GPU。如果它遗漏了重要的上下文 token,模型就可能产生较差的答案。在上下文密集的工作负载中,这一风险很高,例如多文档问答、法律审阅和代码库推理,这些场景下提示词中的许多细节都可能很重要。
这篇论文指出了这个问题:一些 KV 卸载方法在常见的长上下文基准上表现良好,但在需要从提示词中检索大量事实的任务上会退化。实际的经验是,长上下文长度和上下文密集程度是两个不同的概念。在生产环境中启用选择性 KV 卸载之前,请在与你工作负载匹配的任务上将其与全注意力基线进行比较。在跟踪 TTFT、TPOT、吞吐量、GPU 内存使用和主机到设备传输的同时,也要跟踪答案质量。
使用 LMCache 卸载 KV 缓存
LMCache 是一个 LLM 服务引擎扩展,旨在通过降低 TTFT 和提高吞吐量来优化 LLM 推理,尤其适用于长上下文工作负载。它支持跨不同引擎实例复用重复输入内容(不仅仅是前缀)的 KV 缓存。
通过将 KV 缓存存储在 GPU、CPU DRAM 和本地磁盘等多个内存层中,LMCache 显著减少了冗余计算。这改善了响应时间并节省了 GPU 计算周期,使其非常适合多轮问答、RAG 和文档级推理等工作负载。
在基准测试中,将 LMCache 与 vLLM 结合使用,在各种用例中实现了 3 倍至 10 倍的延迟降低。
多个开源项目已经集成 LMCache 以支持高效的 KV 缓存卸载和复用:
- llm-d 使用 LMCache 将 KV 缓存数据从 GPU 内存卸载到 CPU 内存和网络磁盘等更具成本效益、更充裕的存储中。
- KServe 集成 LMCache 以降低推理成本,并确保大规模场景下延迟和吞吐量的 SLO。
- vLLM 使用 LMCache 进行 CPU 卸载、请求间缓存共享以及分离式预填充。这实现了更好的内存管理并提高了资源效率。
LMCache 目前支持将 KV 缓存数据卸载到各种存储后端,从 CPU 内存和文件系统等本地选项,到 Mooncake 和 ValKey 等分布式系统。