ShadowKV: KV Cache in Shadows for High-Throughput Long-Context LLM Inference

📄 arXiv: 2410.21465v3 📥 PDF

作者: Hanshi Sun, Li-Wen Chang, Wenlei Bao, Size Zheng, Ningxin Zheng, Xin Liu, Harry Dong, Yuejie Chi, Beidi Chen

分类: cs.LG, cs.CL

发布日期: 2024-10-28 (更新: 2025-04-25)

🔗 代码/项目: GITHUB

💡 一句话要点

ShadowKV：面向高吞吐长文本LLM推理的低秩KV缓存与动态重构

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长文本LLM推理 KV缓存 低秩分解 稀疏注意力 高吞吐量 内存优化 动态重构

📋 核心要点

  1. 长文本LLM推理面临KV缓存随序列长度增长带来的内存瓶颈和访问延迟问题，现有稀疏注意力方法未能有效降低内存或引入显著延迟。
  2. ShadowKV通过存储低秩键缓存和卸载值缓存来减少内存占用，并采用精确的KV选择策略动态重构稀疏KV对，降低解码延迟。
  3. 实验表明，ShadowKV在多种长文本基准测试和模型上，显著提升了吞吐量并支持更大的批量大小，同时保持了模型精度。

📝 摘要（中文）

随着长文本大型语言模型（LLM）的广泛部署，对高效支持高吞吐量推理的需求日益增长。然而，随着键值（KV）缓存随着序列长度的扩展，不断增长的内存占用以及每次token生成时都需要访问它，导致在服务长文本LLM时吞吐量较低。虽然已经提出了各种动态稀疏注意力方法来加速推理，同时保持生成质量，但它们要么未能充分减少GPU内存消耗，要么通过将KV缓存卸载到CPU而引入了显著的解码延迟。我们提出了ShadowKV，一种高吞吐量长文本LLM推理系统，它存储低秩键缓存并卸载值缓存，以减少更大批量大小和更长序列的内存占用。为了最大限度地减少解码延迟，ShadowKV采用了一种精确的KV选择策略，可以动态地重构最小的稀疏KV对。通过在广泛的基准测试（包括RULER、LongBench和Needle In A Haystack）以及Llama-3.1-8B、Llama-3-8B-1M、GLM-4-9B-1M、Yi-9B-200K、Phi-3-Mini-128K和Qwen2-7B-128K等模型上评估ShadowKV，我们证明了它可以支持高达6倍的更大批量大小，并在A100 GPU上将吞吐量提高高达3.04倍，而不会牺牲准确性，甚至超过了在无限GPU内存假设下可实现的无限批量大小的性能。

🔬 方法详解

问题定义：长文本LLM推理过程中，KV缓存的大小随着上下文长度线性增长，导致GPU内存占用过高，限制了批量大小和吞吐量。现有的动态稀疏注意力方法虽然试图缓解这个问题，但要么无法充分降低GPU内存消耗，要么需要将KV缓存卸载到CPU，从而引入显著的解码延迟，影响推理效率。

核心思路：ShadowKV的核心思路是将KV缓存分解为低秩的Key缓存和Value缓存，只在GPU上存储低秩的Key缓存，并将Value缓存卸载到CPU。在需要时，通过精确的KV选择策略，动态地从CPU重构出最小的稀疏KV对，从而在降低内存占用的同时，最大限度地减少解码延迟。

技术框架：ShadowKV的整体框架包括以下几个主要模块：1) 低秩Key缓存存储模块：负责在GPU上存储低秩的Key缓存。2) Value缓存卸载模块：负责将Value缓存卸载到CPU。3) KV选择模块：根据当前token的上下文信息，从CPU的Value缓存中选择最相关的Value，并与GPU上的Key缓存进行组合，重构出稀疏的KV对。4) 推理模块：使用重构的稀疏KV对进行LLM推理。

关键创新：ShadowKV的关键创新在于其精确的KV选择策略。该策略能够根据当前token的上下文信息，准确地选择出最相关的Value，从而保证重构的稀疏KV对能够尽可能地保留原始KV缓存的信息，避免精度损失。此外，通过只在需要时才进行KV重构，ShadowKV能够最大限度地减少解码延迟。

关键设计：ShadowKV的关键设计包括：1) 低秩Key缓存的秩的选择：需要根据具体的模型和数据集进行调整，以在内存占用和精度之间取得平衡。2) KV选择策略的设计：需要考虑如何根据上下文信息准确地选择出最相关的Value。3) KV重构的优化：需要尽可能地减少KV重构的计算量，以降低解码延迟。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

ShadowKV在A100 GPU上，能够支持高达6倍的更大批量大小，并将吞吐量提高高达3.04倍，同时保持模型精度。在RULER、LongBench和Needle In A Haystack等长文本基准测试中，ShadowKV的表现优于现有方法，甚至超过了在无限GPU内存假设下可实现的无限批量大小的性能。例如，在Llama-3-8B-1M模型上，ShadowKV实现了显著的吞吐量提升。

🎯 应用场景

ShadowKV可应用于各种需要处理长文本的LLM推理场景，例如长篇文档摘要、长对话生成、代码生成等。通过降低内存占用和提高吞吐量，ShadowKV可以帮助企业更高效地部署和使用长文本LLM，从而提升用户体验并降低运营成本。未来，ShadowKV可以进一步扩展到支持更多类型的LLM和硬件平台。

📄 摘要（原文）

With the widespread deployment of long-context large language models (LLMs), there has been a growing demand for efficient support of high-throughput inference. However, as the key-value (KV) cache expands with the sequence length, the increasing memory footprint and the need to access it for each token generation both result in low throughput when serving long-context LLMs. While various dynamic sparse attention methods have been proposed to speed up inference while maintaining generation quality, they either fail to sufficiently reduce GPU memory consumption or introduce significant decoding latency by offloading the KV cache to the CPU. We present ShadowKV, a high-throughput long-context LLM inference system that stores the low-rank key cache and offloads the value cache to reduce the memory footprint for larger batch sizes and longer sequences. To minimize decoding latency, ShadowKV employs an accurate KV selection strategy that reconstructs minimal sparse KV pairs on-the-fly. By evaluating ShadowKV on a broad range of benchmarks, including RULER, LongBench, and Needle In A Haystack, and models like Llama-3.1-8B, Llama-3-8B-1M, GLM-4-9B-1M, Yi-9B-200K, Phi-3-Mini-128K, and Qwen2-7B-128K, we demonstrate that it can support up to 6$\times$ larger batch sizes and boost throughput by up to 3.04$\times$ on an A100 GPU without sacrificing accuracy, even surpassing the performance achievable with infinite batch size under the assumption of infinite GPU memory. The code is available at https://github.com/bytedance/ShadowKV.