要解决的问题
每个请求的输出长度事先未知。如果按最大长度给每个请求预留连续显存,会产生大量碎片和浪费,能同时服务的请求数就上不去。KV 缓存管理是当今推理服务系统的核心。
核心方法
- 分页:借鉴操作系统的虚拟内存分页,把 KV 缓存切成固定大小的块,按需分配,逻辑上连续、物理上可以不连续。
- 共享与写时复制(copy-on-write):并行采样、束搜索等场景下,多个序列共享相同前缀的块,只有修改时才复制。
- 前缀复用:不同请求之间如果有相同的前缀(系统提示、少样本示例、多轮对话历史),直接复用已计算的 KV 缓存。
- 淘汰策略:显存不足时决定淘汰哪些块,可以结合优先级。
代表工作
- PagedAttention(2023):提出分页式 KV 缓存管理,并据此实现了 vLLM,显著提高吞吐。
- RadixAttention(2023):由 SGLang 提出,用基数树(radix tree)组织 KV 缓存,自动在请求之间复用公共前缀。
- TensorRT-LLM:支持跨请求的 KV 缓存复用、基于优先级的淘汰,以及量化的 KV 缓存。
难点
块的分配与回收、写时复制、前缀复用、淘汰策略,以及量化或卸载之后重新恢复的成本。
关键指标
共享前缀命中率、可复用 token 数、显存占用、首 token 延迟与吞吐都需要测量。命中率高有助于减少重复预填充,但总成本还取决于输出长度、批处理和硬件利用率。
参考资料
相关专题
- KV缓存压缩与注意力变体:从架构和数值层面缩小 KV 缓存。
- 推理并行与服务调度:KV 缓存管理与连续批处理、调度紧密配合。
- 内存卸载:显存放不下时把 KV 缓存卸载到 CPU 内存或磁盘。
- 预填充与解码分离:分离式服务需要在实例之间传输 KV 缓存。