KVLINK: Accelerating Large Language Models via Efficient KV Cache Reuse

文章标题：KVLINK: 通过高效的KV缓存重用加速大型语言模型
作者/机构：
- Jingbo Yang, Bairu Hou, Shiyu Chang (加州大学圣巴巴拉分校 计算机科学系)
- Wei Wei, Yujia Bao (埃森哲 先进人工智能中心)

A1 主要贡献

核心问题：大型语言模型（LLM）在处理许多应用（如检索增强生成 RAG）时，需要处理由多个片段组成的上下文。传统架构要求LLM将整个拼接后的上下文作为一个单元进行编码，这导致了高昂的预填充成本，并妨碍了模型独立编码和重用每个片段的预计算表示（即键值状态）。例如，在RAG中，当不同的查询共享相同的检索文档时，模型会冗余地重新编码这些内容未变的文档，造成计算资源的浪费。

研究目标：本文旨在探索一种替代策略，通过独立地为每个文档或文本片段预计算键值（KV）状态，并在推理时重用这些预计算的状态，从而消除冗余计算，显著提高效率。然而，简单地独立编码并拼接KV缓存会导致性能下降，因为这与LLM在完整连续序列上的训练方式存在差异，主要体现在位置嵌入的不匹配和文档间交叉注意力的缺失。

创新点：为了解决上述挑战，本文提出了KVLINK方法，旨在弥合独立编码片段与完整上下文编码之间的差距。其核心创新点包括：
1. KV缓存位置重编码 (KV cache positional re-encoding)：在推理时调整存储的KV缓存的位置嵌入，确保它们在拼接后能与给定查询的全局位置正确对齐。
2. 可训练的跨段特殊令牌 (Trainable cross-segment special tokens)：为了有效地“链接”独立编码的片段，KVLINK在每个片段的预计算KV状态之间附加一小组可训练的“链接令牌”（link tokens）。这些令牌的KV表示在推理时计算，它们可以关注所有在它们之前的令牌，从而在只引入极小计算开销的情况下，帮助恢复跨片段的自注意力机制。

通过在7个数据集上的实验证明，KVLINK相比现有SOTA方法平均提升了4%的问答准确率，并将首个令牌生成时间（time-to-first-token）最多减少了96%。此外，KVLINK还可以与KV缓存压缩技术结合，进一步节省缓存加载和存储开销。


图 1: 标准方法（顶部）对每个文档的KV缓存进行编码时，会依赖于前面的令牌，导致共享文档（如 Docb）的KV缓存编码产生冗余且不可重用。相比之下，我们的设置（底部）分别对文档进行编码，从而允许在不同查询之间重用KV缓存。

A2 方法细节

2.1 问题描述

重用KV缓存的目标。我们旨在解决在LLM中重用键值（KV）缓存的问题，而无需重复编码相同的上下文片段。在许多应用中，相同的文档或上下文片段会出现在不同的推理查询中，但当前的LLM流程每次都会冗余地重新编码这些片段。

检索增强生成中的问题示例。以检索增强生成问题为例，输入问题表示为q，我们为每个问题检索N个文档。这些文档与查询拼接在一起，形成完整的LLM输入。在标准的LLM推理流程中，整个输入作为单个连续序列传递给LLM。如图1顶部所示，每个文档的KV缓存与其前面的文档纠缠在一起。因此，即使Docb在处理q1时已经被编码到KV缓存中，其产生的KV缓存也无法在处理q2时直接重用，因为它是在q1中以其前面的文档为条件生成的。

上下文无关的KV缓存预计算方案。为了解决这个问题，我们考虑一种情景，即知识库中每个文档的KV缓存都以一种上下文无关的方式预先计算，如图1底部所示。具体来说，对于知识库中的每个文档，我们只将该文档的令牌输入LLM，并记录产生的KV缓存。在推理时，当我们为给定查询检索到一组文档后，我们将它们预先计算的KV缓存拼接起来。这种设计使我们能够为不同查询中的重叠文档重用相同的缓存，从而消除冗余计算。

独立编码带来的挑战与解决方案。然而，上述方法通常会导致明显的性能下降，因为LLM通常在完全拼接的序列上进行训练，每个令牌都会关注前面的上下文。当我们以上下文无关的方式编码每个文档时，模型会丢失跨文档的依赖关系。先前的工作也通过经验证明，当每个检索到的文档被单独编码到KV缓存中时，问答任务的准确率相对下降高达35%【7, Block-attention for efficient rag. arXiv, 2024】。我们的方法旨在通过引入两个关键组件，即❶ KV缓存位置重编码和❷ 使用链接令牌进行跨文档重连接，使LLM在每个文档的KV缓存被独立预计算时仍能产生高质量的输出。

2.2 KV缓存位置重编码

位置不匹配问题。独立编码文档的第一个问题是推理过程中的位置不匹配。现代LLM通常使用旋转位置编码（Rotary Positional Encoding, RoPE）【14, Roformer: Enhanced transformer with rotary position embedding. Neurocomputing, 2024】，其中每个令牌根据其在完整序列中的位置被分配一个独特的位置嵌入。然而，当文档被独立编码时，令牌是在其各自的文档内部进行索引的，忽略了它们在完整拼接输入中的实际位置。以图1为例，由于我们为每个文档分别预计算KV缓存，Docb中的第二个令牌在计算其KV缓存时被分配了位置索引2。但是，当我们在推理时将Docb与Doca拼接（例如，处理查询q1时），该令牌的实际位置应该是|Doca| + 2，其中|Doca|是第一个文档的长度。由于Docb的KV缓存在预计算时没有意识到这种全局位置偏移，LLM会应用错误的位置嵌入，导致错误的注意力计算。

将键值状态与位置嵌入解耦。为了解决这个问题，我们在存储键值状态时将其与位置嵌入解耦。

保留局部自注意力。在编码每个片段以进行局部自注意力计算时，我们仍然对令牌应用旋转位置编码。

存储无位置信息的KV缓存并在推理时应用全局位置编码。然而，在保存片段的KV缓存之前，我们会排除这些位置变换。更具体地说，我们将某个transformer层中位置i处令牌的隐藏状态表示为$x_i \in R^d$，其中d是隐藏维度。键向量计算为$R_i W_k x_i$，值向量计算为$W_v x_i$。这里，$W_{\{k,v\}} \in R^{d \times d}$是表示将隐藏状态投影到键和值空间的权重矩阵，$R_i \in [-1, 1]^{d \times d}$是RoPE中使用的位置相关的旋转矩阵。由于RoPE通过这些旋转直接编码位置信息，KV缓存可以在没有位置嵌入的情况下存储，即存储$W_{\{k,v\}}x_i$。在推理时，所有文档的键值状态被拼接起来，然后我们为每个令牌的KV状态应用与其在完整序列中位置相适应的全局旋转嵌入。这个操作只引入了微不足道的时间开销，确保了我们方法的效率。这也与先前的方法【6, Cacheblend: Fast large language model serving with cached knowledge fusion. arXiv, 2024】【8, Attention entropy is a key factor: An analysis of parallel context encoding with full-attention-based pre-trained language models. arXiv, 2024】一致，这些方法在推理时调整了独立编码的KV缓存的位置编码。

2.3 使用链接令牌进行跨文档重连接

恢复跨段自注意力连接的机制。当文档被独立编码和缓存时，后面文档中的令牌无法再关注前面文档中的令牌。这在我们的推理方案和标准的LLM训练之间造成了差距，并可能限制模型性能。为了缓解这个问题，我们设计了一种新颖的机制，以在保留KV缓存重用计算效率的同时，恢复跨片段的自注意力连接。具体来说，KVLINK引入了一组可训练的“链接令牌”。对于每个长度为L的文档c，我们附加K个（例如，K=5）链接令牌：

$$\boldsymbol{c}=\left(c_{1}, \ldots, c_{L}, c_{\text {link } 1}, \ldots, c_{\text {link } 5}\right),$$

其中$c_1, . . . , c_L$是原始文档中的令牌。一个定制的注意力图确保每个文档的链接令牌可以关注 (i) 前面文档中的所有令牌（包括链接令牌）以及 (ii) 当前文档内的令牌。这些令牌充当片段之间的接口，允许模型重建因独立编码而丢失的依赖关系。

链接令牌的工作原理。图2阐释了这一机制，其中每个文档附加一个链接令牌（K=1），分别是令牌3、6和8。文档内的令牌只能保持局部的因果注意力，即它们只能关注同一文档内较早的令牌。每个链接令牌可以关注所有前面的令牌。因此，在推理过程中，不同的文档通过这些链接令牌被隐式地连接起来。例如，令牌6（文档2的链接令牌）关注前两个文档，而令牌8（文档3的链接令牌）关注令牌6，从而混合了所有检索文档的信息。在检索到的文档之后，问题和用户令牌遵循标准的因果自注意力机制。在训练期间，我们使用这种注意力机制对LLM进行微调，共同优化模型和新引入的链接令牌。

图 2: 左图：所有令牌的注意力图。link1令牌仅关注Doca中的令牌；link2关注Doca、Docb和link1中的令牌；link3关注所有重用上下文和其他链接令牌中的令牌。右图：推理期间三个链接令牌和第一个用户输入令牌的注意力图。这两个注意力图是相同的。

推理策略。图2阐释了我们的推理策略，该策略避免了为检索到的文档重新计算KV缓存。在实践中，我们首先加载并拼接所有检索文档的预计算KV缓存。由于这些缓存是预计算的，它们不需要被重新编码。然后，我们只将新引入的链接令牌（对应于三个文档的令牌3、6和8）附加到重新排序序列的末尾。我们使用定制的注意力图对这些令牌执行一次前向传播，确保它们的注意力行为与训练阶段相匹配。

2.4 压缩KV缓存链接

存储开销问题。虽然重用预计算的KV缓存可以极大地加速推理，但它可能会在数据库中存储这些缓存以及将它们加载到GPU上时引入相当大的开销。这种开销是将会缓存重用集成到实际RAG部署中的一个主要挑战。例如，一个以UTF-8文本存储的1000个令牌的文档仅占用约5KB，而相应的Llama3-8B KV缓存则需要大约131MB。为了解决这个问题，我们进一步探索将我们的方法与KV缓存压缩技术相结合。

压缩技术选项。具体来说，我们考虑以下选择：❶ LLMLINGUA【12, Llmlingua-2: Data distillation for efficient and faithful task-agnostic prompt compression. arXiv, 2024】，它通过令牌丢弃来压缩文档；以及❷ ANLLMS【13, Anchorbased large language models. Findings of the Association for Computational Linguistics ACL 2024, 2024】，它通过微调LLM将文档中的每个句子压缩成一个带有特殊注意力掩码的锚点令牌。

对ANLLMS的改进。在使用ANLLMS时，我们根据经验观察到LLM在问答基准测试上的性能显著下降。为了解决这个问题，我们引入了以下旨在提高其有效性的修改。首先，我们不将每个句子压缩成单个锚点令牌并存储其KV缓存，而是将每个文档分成多个固定长度的连续块，并将每个块压缩成多个锚点令牌。其次，我们修改了注意力掩码，使得每个锚点令牌只关注其自身块内的令牌以及前面的锚点令牌。

改进后的训练方法。更具体地说，假设x是一个来自预训练数据集的训练样本，由一个令牌序列组成。我们将前N个令牌分成N/s个固定大小为s的块：

$$\boldsymbol{x}^{(n)}=x_{(n-1) s+1}, \ldots, x_{n s}, n=1, \ldots, N,$$

其中s表示块大小。在我们的实验中，我们默认设置s=100。然后，我们向每个块附加几个锚点令牌。LLM被训练来对位置N之后的剩余令牌进行下一令牌预测，这些令牌被限制为只能关注锚点令牌。这种设置鼓励模型将输入压缩成少量锚点令牌，使我们能够只存储这些令牌的KV缓存，从而显著减少存储开销。此外，通过增加锚点令牌的数量，我们可以在使用压缩KV缓存的情况下更好地保持模型性能。

应用流程。为了应用这种方法，我们首先在预训练数据上使用这个目标对LLM进行预训练。然后，我们使用KVLINK在QA数据集上对模型进行微调。更多的实现细节在附录A.6中提供。

A4 实验环境

  • 对比基线:

    1. PROMPTCACHE 【9, Prompt cache: Modular attention reuse for low-latency inference. Proceedings of Machine Learning and Systems, 2024】: 直接重用预计算的KV缓存。我们通过应用本文的位置重编码方法对其进行了增强。
    2. CACHEBLEND 【6, Cacheblend: Fast large language model serving with cached knowledge fusion. arXiv, 2024】: 拼接预计算的KV缓存后，重新计算一小部分（18%）选中令牌的KV缓存。
    3. BLOCKATTENTION 【7, Block-attention for efficient rag. arXiv, 2024】: 通过在独立编码的QA任务上微调，训练模型处理分离的KV缓存。

  • 模型架构:

    • 主干模型为 Llama-3.2-1B-Instruct, Llama-3.2-3B-Instruct, 和 Llama-3.1-8B-Instruct 【15, Llama 3 and vision: Bringing next-gen ai to edge and mobile devices, 2024】。
    • 本文方法训练了三个版本：KVLINK0, KVLINK1, KVLINK5，分别对应在每个文档后附加0, 1, 或 5个链接令牌。

  • 硬件配置:

    • GPU: 8块 H100 GPU。

  • 软件配置与训练设置:

    • 使用全局批大小64进行6000步微调。
    • 训练数据集: 混合了多个数据集，包括2WikiMQA【16, Constructing a multi-hop qa dataset for comprehensive evaluation of reasoning steps. In Proceedings of the 28th International Conference on Computational Linguistics, 2020】、TriviaQA【17, Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension. In Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), 2017】的训练集，以及来自FineWeb【18, The fineweb datasets: Decanting the web for the finest text data at scale. In The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks Track, 2024】的预训练数据和TÜLU 3【19, T\" ulu 3: Pushing frontiers in open language model post-training. arXiv, 2024】。

  • 评估数据集:

    • 问答 (QA): NaturalQuestions【10, Natural questions: a benchmark for question answering research. Transactions of the Association for Computational Linguistics, 2019】, 2WikiMQA, TriviaQA, HotpotQA【11, Hotpotqa: A dataset for diverse, explainable multi-hop question answering. arXiv, 2018】, MuSiQue【20, Musique: Multihop questions via single-hop question composition. Transactions of the Association for Computational Linguistics, 2022】。
    • 文本摘要: MultiNews【23, Multi-news: A large-scale multi-document summarization dataset and abstractive hierarchical model. arXiv, 2019】, Samsum【24, Samsum corpus: A human-annotated dialogue dataset for abstractive summarization. arXiv, 2019】。
    • 通用能力: IFEval【26, Instruction-following evaluation for large language models. arXiv, 2023】, GSM8K【27, Training verifiers to solve math word problems. arXiv, 2021】, MMLU【28, Measuring massive multitask language understanding. arXiv, 2020】, ARC-Challenge & ARC-Easy【29, Think you have solved question answering? try arc, the ai2 reasoning challenge. arXiv, 2018】, PiQA【30, Piqa: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on artificial intelligence, 2020】, SciQ【31, Crowdsourcing multiple choice science questions. arXiv, 2017】, Winogrande【32, Winogrande: An adversarial winograd schema challenge at scale. Communications of the ACM, 2021】, HellaSwag【33, Hellaswag: Can a machine really finish your sentence? arXiv, 2019】。

A4 实验结果

主要性能结果

  • 实验内容：评估KVLINK在使用独立编码的KV缓存时的性能，并与多个基线以及使用完整上下文编码的Llama模型（作为性能上限）进行比较。

  • 实验结果（Table 1）:

    • KVLINK全面领先：在所有QA和摘要任务中，KVLINK始终优于所有基线方法。例如，在HotpotQA上，KVLINK5 (Llama-3.2-3B) 达到了69.5%的准确率，而表现最好的基线BLOCKATTENTION仅为64.3%。
    • 优越的泛化能力：在训练数据未覆盖的数据集上，KVLINK的优势更加明显，表明其具有出色的域外泛化性能。
    • 接近性能上限：KVLINK的性能非常接近使用完整上下文编码进行微调的Llama模型，表明它在大幅提升效率的同时，性能损失极小。
    • 链接令牌的有效性：随着链接令牌数量的增加（从KVLINK0到KVLINK5），性能持续提升，验证了链接令牌在恢复文档间连接方面的关键作用。

    Table 1: KVLINK与其他方法的性能比较。最佳结果以粗体突出显示。KVLINK0、KVLINK1、KVLINK5分别指使用0、1和5个链接令牌。

推理效率评估

  • 实验内容：测量在重用10个不同长度（100到500个令牌）文档的预计算KV缓存时，生成首个令牌的时间（TTFT），并与标准的完整上下文编码进行对比。实验考虑了将KV缓存从CPU加载到GPU的开销。

  • 实验结果（Figure 3）:

    • 显著降低延迟：与标准解码相比，KVLINK将TTFT降低了85%至96%。这得益于加载KV缓存、应用新位置编码和编码链接令牌这三个操作的极低延迟。
    • 良好的可扩展性：随着上下文长度从1000增加到5000个令牌，KVLINK与标准解码之间的TTFT差距持续扩大。当上下文长度达到5000时，KVLINK将推理延迟降低了96%，证明了其处理长上下文的可扩展性。


    图 3: 与十个不同长度的重用上下文进行推理速度比较。随着上下文大小的增长，KVLINK1和KVLINK5的首个令牌生成时间（TTFT）都显著低于标准解码。

通用能力保持评估

  • 实验内容：在一系列推理和指令遵循基准测试上评估KVLINK，以验证其在提升效率的同时是否保持了模型的通用能力。

  • 实验结果（Table 2）:

    • 性能保持稳定：KVLINK在所有通用能力测试中均表现出与微调后Llama模型相当的性能，证明该方法没有显著损害模型的通用推理和指令遵循能力。
    • 差异在可接受范围内：尽管在某些基准（如ARC-C和Winogrande）上相比原始模型有轻微下降，但差异通常小于3%，在可接受范围内。

    Table 2: KVLINK、原始Llama和在标准解码下微调的Llama的性能比较。WG指WinoGrande，HS指HellaSwag。

结合KV缓存压缩的评估

  • 实验内容：评估KVLINK与两种KV缓存压缩技术（LLMLINGUA和改进的ANLLMS）结合使用时的性能，压缩率分别为50%和75%。

  • 实验结果（Table 3）:

    • 压缩方法对比：改进的ANLLMS压缩方法比LLMLINGUA保留了更多信息，在所有QA基准上表现更优。
    • KVLINK缓解性能下降：在使用压缩KV缓存时，KVLINK的链接令牌机制能够有效缓解压缩带来的精度下降，其性能始终优于同样使用压缩缓存的BLOCKATTENTION基线。例如，在使用本文方法（改进的ANLLMS）压缩75%时，KVLINK在NQ上的准确率（40.9%）高于BLOCKATTENTION（37.5%）。

    Table 3: 在不同压缩方案下对KVLINK的评估。我们报告了使用压缩缓存时的QA准确率。斜线（/）左侧的值对应75%的压缩率，右侧的值对应50%的压缩率。

消融研究

  • 实验内容：探究不同训练数据混合对KVLINK性能的影响。
  • 实验结果（Appendix A.4）:
    • 鲁棒性：KVLINK在不同的数据混合下仍能取得有竞争力的性能，表现出良好的鲁棒性。
    • 数据选择的重要性：特定的任务和领域覆盖组合会轻微影响下游性能，表明平衡的数据选择对于优化性能至关重要。

A7 补充细节 (相关工作)

LLM的高效推理。随着模型规模的增长，高效地服务LLM变得越来越具挑战性。先前的工作通过模型剪枝【35-38】、量化【39-44】以及优化的解码算法（如非自回归解码【45】、推测解码【46-48】或早退【49, 50】）来解决这个问题。与这些方法类似，我们的方法也旨在通过减少推理时的计算来提高LLM的推理效率。

KV缓存重用。近期的工作探索了在多个查询中重用预计算的KV缓存，通常侧重于仅前缀的重用【51-53】。少数方法将重用扩展到任意位置，以追求更高的效率。PROMPTCACHE【9, Prompt cache: Modular attention reuse for low-latency inference. Proceedings of Machine Learning and Systems, 2024】允许使用不连续位置ID进行KV缓存重用，但忽略了跨块的注意力，导致性能下降。它还依赖于提示标记语言，限制了其在不同任务间的灵活性。CACHEBLEND【6, Cacheblend: Fast large language model serving with cached knowledge fusion. arXiv, 2024】通过选择性地为值状态方差最大的令牌重新计算缓存来减少位置不匹配；然而，它存在性能下降和重新计算10%-20%令牌（以及第一层的所有令牌）的不可忽视的成本。BLOCKATTENTION【7, Block-attention for efficient rag. arXiv, 2024】通过重新编码KV缓存来移除位置嵌入，并直接在块注意力数据上进行微调，但缺乏更好地连接重用上下文或优化数据处理的策略，导致性能未达理想水平。TurboRAG【54, Turborag: Accelerating retrieval-augmented generation with precomputed kv caches for chunked text. arXiv, 2024】引入了两个特殊令牌来标记重用缓存的边界。然而，它在有效重连重用缓存方面仍然不足，与BLOCKATTENTION类似。详细讨论见附录A.8。

KV压缩。虽然我们的方法侧重于通过重用KV缓存来加速推理，但一个互补的研究方向旨在通过缓存驱逐和量化来减少内存开销。例如，【55-58】引入了一种驱逐策略，显著降低了生成过程中的内存占用，而【59, 42, 39, 41, 60-62】研究了量化KV缓存以最小化存储需求，且性能损失很小。值得注意的是，我们的方法可以与这些最先进的驱逐和量化策略无缝集成，从而同时解决速度和内存效率问题。

检索增强生成。检索增强生成（RAG）被广泛用于通过检索支持文档来增强语言模型的生成能力。通常使用密集检索器或BM25【22, 63, 64】等检索器来为给定的用户查询或上下文找到最相关的文档。除了简单地将检索到的文本与输入拼接【65】之外，还有其他将检索知识集成到模型中的方法。例如，LongMem【66】将文档编码到缓存中，并使用一个训练好的旁路网络将这些知识缓存并入上下文，以用于长上下文任务。随着LLM在各种NLP任务中的进步，某些研究专注于用LLM改进RAG【67, 68】。一些研究还探讨了如何增强LLM在处理检索知识时的鲁棒性。例如，RAFT【69】提供了一种微调策略，可以增强LLM在领域特定RAG中对抗嘈杂上下文的能力。

A5 结论

本文介绍了一种通过重用预计算的KV缓存来提升LLM效率的方法。通过为检索到的文档或上下文片段预计算KV缓存，我们的方法可以避免在不同查询中对重叠上下文进行冗余计算。未来，我们将优化我们的数据混合，并研究最佳的微调策略以进一步提升性能。此外，我们计划探索真实世界的部署场景，以最大限度地提高LLM的推理效率。局限性和社会影响部分包含在附录中。

A6 附录

A.1 数据混合

数据构成。我们使用一个混合了不同数据集的数据集来微调我们的模型。

  • 检索增强QA：这是最直接相关的任务。我们从TriviaQA【17】和2WikiMQA【16】中采样问题，并使用Contriever【22】为每个问题检索十个相关的维基百科段落，每个段落都单独编码。为了获得高质量的监督信号，我们使用GPT-4生成参考答案作为训练的基准。
  • 多轮对话：许多现有的指令遵循数据集【70】包含用户与助手之间的多轮对话，这为在不相连的上下文上训练模型提供了一种自然的方式。我们随机将早期的对话轮次转换为独立编码的KV缓存，然后训练LLM在后续轮次中生成适当的响应。这确保了模型能学会整合分段的上下文，同时保持其指令遵循能力。
  • 摘要：摘要是另一个有用的场景，模型必须从独立编码的文档块中聚合信息以生成连贯的摘要。我们采用了XSum【71】数据集。每个文档被随机分割成多个连续的片段，每个片段被独立编码到KV缓存中。模型被训练来生成基准摘要，确保它学会跨独立片段整合信息。

  • 保留原始能力的数据：为了更好地保留模型的原始能力，我们还包括了标准版本的TriviaQA和2WikiMQA，其中检索到的文档是作为一个整体而不是分开编码的。在这种设置下，模型被训练以在完全拼接的上下文条件下生成基准答案。最后，我们加入了T"ulu 3数据集【19】以保留指令遵循能力，以及一小部分来自Fineweb【18】的预训练数据以保留语言建模能力。

    Table 4: 任务、数据集以及有无上下文重用的样本数量概览。

数据处理细节。表4概述了我们训练中使用的数据。对于多轮对话数据，每个先前的用户-助手对话都被独立编码为重用上下文，我们仅在助手的最终响应上计算训练损失。在摘要任务中，我们将源文档分割成100个令牌的片段，每个片段作为一个重用上下文。所有训练样本都被截断至最大长度4096个令牌。

A.2 实现细节

训练。我们建议使用从不同来源抽取的混合数据集对LLM进行微调，使模型能够在整合不相连的上下文片段的同时，保留其原始能力。图5展示了我们在训练期间为每个任务类别使用的系统提示。这些提示由GPT-4【72】生成，并在训练期间随机选择。此外，对于QA训练数据，我们还在上下文A.5中对检索到的重用文档进行洗牌。

数据预处理。图4展示了在KVLINK中重用上下文时如何处理每个数据样本。特别地，我们还包括了两个特殊令牌，KV-START和KV-END，以指定重用上下文的边界。对于每个重用的文档或上下文片段，链接令牌取决于该文档在所有重用文档中的索引，这意味着，在不同的提示中，对于第n个重用文档，其链接令牌始终是linkn-1, linkn-2, ..., linkn-K（如果使用K个链接令牌）。


图 4: 上下文重用的数据预处理。


图 5: 用于训练的系统提示。我们为三种主要任务类型——SFT、QA和摘要——采用量身定制的系统提示，以反映不同的目标并相应地指导模型的响应。

评估。对于使用独立编码KV缓存的QA评估，我们采用准确率作为评估指标，遵循【7, 21, 73, 74】的做法，即如果预测的任何子字符串与任何标准答案完全匹配，则认为预测正确。对于通用能力保持的评估，我们使用LM Evaluation Harness【75】。我们对GSM8K和MMLU使用少样本示例（GSM8K为8-shot，MMLU为5-shot）。

A.3 基线实现的更多细节

BLOCKATTENTION的数据处理。对于BLOCKATTENTION，我们通过分别编码每个重用的文档或上下文片段，并直接拼接缓存而不插入任何特殊令牌来处理带有上下文重用的训练数据（如图4所示）。对于没有上下文重用的数据，数据处理方式与KVLINK相同。

PROMPTCACHE的优化。对于PROMPTCACHE，在其原始实现中，当在新提示中重用时，位置编码并未进行调整。它在重用的缓存之间使用不连续的位置信息，这是不准确的。我们通过在评估期间为所有重用的上下文提供黄金标准的位置信息，来最大化PROMPTCACHE的性能。

A.4 数据混合的消融研究

Table 5: 不同训练数据混合下的性能。

实验设置。我们在KVLINK5设置下试验了多种数据混合，每种混合都从我们原始的训练组合中省略了一个不同的子集（见附录A.1）。具体来说，我们探索了三种配置：无摘要、无多轮对话和仅QA（见表5）。对于前两种配置，我们保持其余任务的相对比例不变，并训练6000步，而仅QA配置则训练两个epoch以防过拟合。

结论。我们的结果表明，需要跨块推理的任务对于稳健的缓存重用至关重要。此外，整合多种类型的缓存重用数据可以提高泛化能力。然而，为使LLM完全具备通用缓存重用能力而设计的最佳训练数据配方仍然是一个悬而未决的问题，我们将其留给未来的工作。

A.5 答案文档位置的影响

问题观察。根据经验，我们观察到直接使用Contriever【22】检索到的相关文档来构建QA训练数据并不理想，因为检索器通常会根据相关性得分将包含答案的文档放在前面。结果，微调后的模型“学会了”在上下文的开头寻找答案，而不是在整个重用的上下文中进行推理。

实验验证。为了验证我们的说法，我们在与BLOCKATTENTION相同的未打乱训练数据上微调了Llama-3.2-3B基础模型，并在【21】的设置下进行评估。在【21】中，NaturalQuestions的测试集也是用Contriever检索的上下文构建的，但将包含正确答案的文档放在了不同的位置。如图6所示，当基准真相文档位于上下文中较靠后的位置时，模型的性能显著下降。


图 6: 答案文档位置对准确率的影响。当正确文档距离开头越远时，准确率显著下降，这表明在QA训练数据中打乱检索到的文档是必要的。

A.6 改进版ANLLM压缩的训练

训练流程。我们对改进版ANLLM压缩的训练包含两个阶段：持续预训练和QA微调。


图 7: 训练压缩令牌时的注意力机制。

持续预训练。该阶段涉及训练压缩令牌以捕捉来自较长前序序列的信息。如图7所示，给定一个文本样本，我们首先将其分成两半。我们压缩前半部分，并在后半部分计算交叉熵损失。在压缩过程中，我们根据100个令牌的固定长度将前半部分分割成多个文本段。在每个段之后，我们附加几个压缩令牌。这些压缩令牌关注所有前面的压缩令牌以及它们所在段内的原始令牌，这是我们方法与ANLLM的区别。每个段内的原始令牌保持局部注意力，而后半部分的令牌不能关注这些原始令牌，只能关注压缩令牌。

QA微调。QA微调的目标是使压缩令牌适应于压缩文档，并进一步训练模型基于压缩信息执行QA任务。QA微调期间的注意力机制与持续预训练中的类似。唯一的区别是，每个文档的压缩令牌不能相互关注。

A.7 KV缓存重用的存储需求

成本权衡。存储KV缓存的成本高于存储原始文本。下面，我们讨论存储与效率之间的权衡以及如何管理它。首先，将KVLink与现有的KV缓存压缩方法相结合有助于减少存储使用。在2.4节中，我们介绍了两种此类方法，表明精心设计的压缩策略只会导致最小的性能损失。压缩技术的持续改进可以进一步减少存储需求。

成本分析。其次，GPU使用成本通常高于存储成本。如我们的实验所示，在A100 80GB GPU上运行Llama3.1-8B模型时，对于5000个令牌的输入，KVLink可以将延迟降低96%。在固定的GPU小时数下，这使得KVLink能够处理比标准解码多25倍的请求。对于每百万次这样的请求，KVLink需要9个GPU小时（16美元），而标准解码需要246个GPU小时（440美元），基于当前费率（A100 80GB为1.79美元/小时）。另一方面，存储仍然便宜。例如，亚马逊S3的标准计划每月每GB仅收费0.023美元。

系统级解决方案。最后，我们可以通过系统级解决方案降低存储成本。可以使用两种设计策略：
❶ 仅缓存高命中率文档。像LRU或LFU这样的策略可以帮助识别哪些文档应该保留在缓存中。其他的可以像往常一样以纯文本形式存储。
❷ 使用分层KV缓存存储。我们可以使用分层KV缓存存储来进一步节省存储成本。例如，虽然一个1000个令牌的文档以UTF-8文本存储约占5KB，但其Llama3-8B KV缓存需要大约131MB。如果它很少被使用，我们可以将此缓存放在更便宜的存储中。总的来说，低访问率的缓存应该保存在像SSD这样的更便宜的存储上，而重要的缓存则保留在像CPU RAM这样的更快的内存中。

A.8 与TurboRAG的比较

方法对比。KV缓存重用的另一个基线是TurboRAG【54】，我们在此处讨论它，因为它与我们实验中的另一个基线BlockAttention【7】相似，都是直接拼接重用的缓存。更具体地说：❶ TurboRAG引入了两个额外的令牌：在文档前加上<doc_start>，在文档后加上<doc_end>。与BlockAttention类似，TurboRAG计算并存储文档和这两个令牌的KV缓存，并采用局部自注意力。这意味着这两个令牌是作为文档的一部分离线预计算的，并且只在文档内部保持局部注意力以标记文档边界。❷ 因此，这两个令牌主要用于指示文档的边界。相比之下，我们的方法引入了链接令牌，并在测试时重新计算它们的KV缓存，以重新连接独立编码的文档。❸ 在训练阶段，我们还引入了链接令牌，并以连接独立计算的上下文为目标进行训练，这带来了比TurboRAG和BlockAttention更好的性能。

实验对比。我们还进行了实验，以经验性地比较KVLink和TurboRAG。我们复现了TurboRAG，并使用与我们方法相同的训练数据和训练设置对其进行训练。结果如表6所示。

Table 6: KVLink5与TurboRAG在QA数据集上的比较。

A.9 局限性

尽管我们的方法在提高推理效率的同时取得了最先进的性能，但它仍然存在一些局限性。首先，尽管KVLINK可以在使用压缩缓存时恢复性能，但在压缩后仍然存在一些性能下降。一个可能的解决方案是进行更大规模的训练以获得更好的压缩和链接效果。其次，KV缓存的大小差异很大，因为它部分取决于文档的长度。因此，将这些缓存存储在向量数据库中会很困难。有效地组织和存储预计算的KV缓存仍然是大规模部署面临的挑战。

A.10 社会影响

在本文中，我们的主要目标是通过重用KV缓存来提高LLM的效率。我们的方法旨在提高推理效率，并同时保持模型在下游任务上的性能。本文使用的训练数据是众所周知的，并被广泛用于其他项目中，我们验证了其质量和安全性，以确保在训练或评估过程中不包含任何私人或敏感信息。此外，本文提出的KV缓存重用方法不会给LLM带来额外的风险或偏见。虽然我们承认任何LLM都可能存在偏见或潜在的滥用，但我们提出的技术产生风险或被滥用的可能性已大大降低。

A.11 许可证

训练和评估中包含了各种数据集，其许可证列表如下。IFEval, NaturalQuestions, 2WikiMQA, TriviaQA, HotpotQA 数据集在 Apache License 2.0下发布。HumanEval, MMLU, GSM8K, HellaSwag, WinoGrande, XSum 和 LM-EvaluationHarness 在 MIT 许可证下。ARC 数据集在 Creative Commons Attribution Share Alike 4.0 许可证下提供，SciQ 在 Creative Commons Attribution NonCommercial 3.0 许可证下。DaringAnteater 数据集在 Creative Commons Attribution 4.0 许可证下。MuSiQue 数据集在 Creative Commons Attribution 4.0 International 许可证下。Samsum 数据集在 Creative Commons Attribution Non Commercial No Derivatives 4.0 许可证下。MultiNews 数据集受与 LILY LAB 的法律协议约束。PiQA 在 Academic Free License v. 3.0 下获得许可。Tulu3-sft-mixture 和 fineweb 数据集在 Open Data Commons License Attribution family 许可证下。本研究中的所有数据集和软件包均严格用于其预期目的——即训练和评估 LLM。我们确认我们的材料中不包含任何个人身份信息或冒犯性内容。

A.12 LLM使用情况

我们使用GPT-4来协助全文的校对和提高清晰度。尽管如此，研究概念、分析和原始写作仍然完全由我们作者完成。