通过复用文档的键值缓存,显著加速大模型推理。
KVLink: Accelerating Large Language Models via Efficient KV Cache Reuse

- 预计算文档的键值缓存,推理时直接拼接复用,避免重复编码。
- 在7个数据集上平均提升问答准确率4%,首字生成延迟降低96%。
- 适合需要频繁复用上下文的场景,如检索增强生成系统。
我们提出KVLink,一种高效复用大语言模型(LLM)键值(KV)缓存的方法。在许多LLM应用中,不同输入可能共享重叠上下文(如多个查询使用同一文档)。但传统方法仍需对每个查询重新编码完整上下文,造成冗余计算。本文提出新策略:预先独立计算每个文档的KV缓存,在推理时将其拼接使用,实现缓存复用。为缓解独立编码导致的性能下降,KVLink引入两项关键技术:调整拼接后全局位置的相对位置编码,以及使用可学习特殊标记恢复跨文档自注意力。在7个数据集上的实验表明,相比现有最优方法,KVLink平均提升问答准确率4%。同时,借助预计算缓存,首字生成时间最多减少96%,显著提升效率。此外,该方法可与KV缓存压缩结合,进一步降低加载与存储开销,且优于基线方案。
原文摘要 · Abstract (English)
We describe KVLink, an approach for efficient key-value (KV) cache reuse in large language models (LLMs). In many LLM applications, different inputs can share overlapping context, such as the same retrieved document appearing in multiple queries. However, the LLMs still need to encode the entire context for each query, leading to redundant computation. In this paper, we investigate a new strategy to eliminate such inefficiency, where the KV cache of each document is precomputed independently. During inference, the KV caches of retrieved documents are concatenated, allowing the model to reuse cached representations instead of recomputing them. To mitigate the performance degradation when using KV caches computed independently for each document, KVLink introduces two key techniques: adjusting positional embeddings of the KV cache at inference to match the global position after concatenation, and using trainable special tokens to restore self-attention across independently encoded documents. Experiments across 7 datasets demonstrate that KVLink improves question answering accuracy by an average of 4% over state-of-the-art methods. Furthermore, by leveraging precomputed KV caches, our approach reduces time-to-first-token by up to 96% compared to standard LLM inference, making it a scalable and efficient solution for context reuse. Additionally, KVLink can be combined with KV cache compression to further save cache loading and storage overhead while outperforming the baselines.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。