SpeCache: Speculative Key-Value Caching for Efficient Generation of LLMs

📄 arXiv: 2503.16163v1 📥 PDF

作者: Shibo Jie, Yehui Tang, Kai Han, Zhi-Hong Deng, Jing Han

分类: cs.CL

发布日期: 2025-03-20

💡 一句话要点

SpeCache：通过推测性键值缓存加速长文本LLM生成，解决VRAM瓶颈。

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

关键词: 键值缓存 长文本处理 大型语言模型 GPU显存优化 推测执行

📋 核心要点

  1. 长文本任务中，LLM的KV缓存需求随序列长度线性增长，导致GPU显存成为瓶颈，限制了LLM的应用。
  2. SpeCache将完整KV缓存卸载到CPU内存，并根据重要性动态地从CPU取回KV对，避免信息丢失。
  3. SpeCache通过推测性预取KV对，隐藏CPU-GPU通信延迟，实验表明其有效降低VRAM占用，且无需重新训练。

📝 摘要（中文）

基于Transformer的大型语言模型(LLM)在长文本任务上取得了显著成果，但随着序列长度的增加，键值(KV)缓存的需求呈线性增长，导致有限的GPU显存(VRAM)资源难以满足需求，成为LLM应用于长序列的瓶颈。现有的KV缓存压缩方法包括驱逐、合并或量化KV缓存以减小其大小。然而，压缩会导致不可逆的信息丢失，可能影响后续解码的准确性。本文提出了SpeCache，它充分利用了大型且易于扩展的CPU内存来卸载完整的KV缓存，并基于VRAM中低比特KV缓存副本所衡量的每个KV对的重要性，在每个解码步骤中动态地将KV对取回。为了避免CPU-GPU通信引起的推理延迟，SpeCache推测性地预测下一个token可能关注的KV对，从而允许我们在下一个解码步骤之前预取它们，从而实现预取和计算的并行化。在LongBench和Needle-in-a-Haystack基准测试上的实验验证了SpeCache有效地减少了VRAM的使用，同时避免了长序列的信息丢失，无需重新训练，即使KV缓存压缩比高达10倍。

🔬 方法详解

问题定义：论文旨在解决大型语言模型在处理长文本时，由于KV缓存线性增长导致的GPU显存(VRAM)瓶颈问题。现有方法如KV缓存压缩（驱逐、合并、量化）会导致信息不可逆丢失，影响解码精度。

核心思路：SpeCache的核心思路是将完整的KV缓存卸载到容量更大的CPU内存中，并在每个解码步骤中，根据KV对的重要性动态地从CPU取回。通过这种方式，既避免了信息丢失，又降低了VRAM占用。同时，采用推测性预取机制，隐藏CPU-GPU通信带来的延迟。

技术框架：SpeCache的整体框架包含以下几个主要模块：1) KV缓存卸载：将完整的KV缓存存储在CPU内存中。2) 重要性评估：在VRAM中维护一个低比特的KV缓存副本，用于评估每个KV对的重要性。3) 推测性预取：预测下一个token可能关注的KV对，并提前从CPU内存加载到VRAM中。4) 解码：使用VRAM中的KV缓存进行解码，同时并行进行预取操作。

关键创新：SpeCache的关键创新在于：1) 利用CPU内存扩展KV缓存容量，避免信息丢失。2) 提出基于低比特KV缓存副本的重要性评估方法，指导KV对的动态加载。3) 引入推测性预取机制，隐藏CPU-GPU通信延迟，提升推理速度。与现有方法相比，SpeCache无需压缩KV缓存，避免了信息损失，同时通过预取机制降低了延迟。

关键设计：SpeCache的关键设计包括：1) 低比特KV缓存副本的比特数选择，需要在VRAM占用和重要性评估精度之间进行权衡。2) 推测性预取的预测算法，需要保证较高的预测准确率，避免无效的预取操作。3) CPU-GPU之间的数据传输优化，例如采用异步传输和数据压缩等技术。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，SpeCache在LongBench和Needle-in-a-Haystack基准测试上，能够在不损失精度的前提下，有效降低VRAM的使用。即使在10倍的KV缓存压缩比下，SpeCache仍然能够保持与原始模型相当的性能。此外，SpeCache无需重新训练模型，即可直接应用于现有的LLM。

🎯 应用场景

SpeCache可应用于需要处理长文本的各种LLM应用场景，例如长篇文档摘要、机器翻译、代码生成、对话系统等。通过降低VRAM需求，SpeCache使得LLM能够在资源受限的设备上运行，并能够处理更长的输入序列，从而提升LLM的实用性和适用范围。未来，SpeCache可以进一步扩展到多GPU环境，以支持更大规模的LLM训练和推理。

📄 摘要（原文）

Transformer-based large language models (LLMs) have already achieved remarkable results on long-text tasks, but the limited GPU memory (VRAM) resources struggle to accommodate the linearly growing demand for key-value (KV) cache as the sequence length increases, which has become a bottleneck for the application of LLMs on long sequences. Existing KV cache compression methods include eviction, merging, or quantization of the KV cache to reduce its size. However, compression results in irreversible information forgetting, potentially affecting the accuracy of subsequent decoding. In this paper, we propose SpeCache, which takes full advantage of the large and easily expandable CPU memory to offload the complete KV cache, and dynamically fetches KV pairs back in each decoding step based on their importance measured by low-bit KV cache copy in VRAM. To avoid inference latency caused by CPU-GPU communication, SpeCache speculatively predicts the KV pairs that the next token might attend to, allowing us to prefetch them before the next decoding step which enables parallelization of prefetching and computation. Experiments on LongBench and Needle-in-a-Haystack benchmarks verify that SpeCache effectively reduces VRAM usage while avoiding information forgetting for long sequences without re-training, even with a 10x high KV cache compression ratio.