用CPU缓存大模型注意力键值对，高效支持长文本生成
SpeCache: Speculative Key-Value Caching for Efficient Generation of LLMs
- 将完整KV缓存移至CPU，仅在显存保留重要部分
- 10倍压缩比下仍保持准确率，无需重新训练
- 预取机制减少通信延迟，适合长序列生成场景
基于Transformer的大型语言模型在长文本任务中表现优异，但随着序列长度增加，键值（KV）缓存需求线性增长，受限于有限的GPU显存（VRAM），已成为应用瓶颈。现有方法通过淘汰、合并或量化KV缓存来减小体积，但会导致不可逆的信息丢失，影响后续解码精度。本文提出SpeCache，充分利用大容量且易扩展的CPU内存，将完整的KV缓存移出显存，并在每一步解码中根据显存内低比特副本的重要性动态加载所需键值对。为避免CPU-GPU通信带来的推理延迟，SpeCache采用推测策略，预判下一标记可能关注的键值对，在解码前提前预取，实现预取与计算并行。在LongBench和Needle-in-a-Haystack基准测试中验证，SpeCache在不重训练的前提下，有效降低显存占用，避免信息遗忘，即使达到10倍的高压缩比依然表现良好。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。