作者: Yihua Cheng, Kuntai Du, Jiayi Yao, Junchen Jiang
机构: The University of Chicago

A1 主要贡献

随着大型语言模型（LLM）的快速发展，为各种LLM查询补充知识的需求也随之增长。因此，实现模块化、高效地在LLM推理中注入新知识至关重要。本文从系统架构的角度探讨了不同的知识注入方法，并提出了一个新的系统组件。

  • 核心问题：现有的知识注入方法，如微调（fine-tuning）和上下文学习（in-context learning），在模块化（添加新知识的便捷性和指定注入知识的灵活性）和效率（单次查询成本和响应延迟）之间存在固有的权衡。微调效率高但模块化差，而上下文学习模块化好但效率低。

  • 研究目标与创新点：

    1. 提出KV缓存学习作为更优的折衷方案：本文主张，使用KV缓存作为知识媒介的“KV缓存学习”方法，可以在低成本和快速响应的情况下，同时提高知识注入的模块化和LLM服务的效率。
    2. 构想知识交付网络（KDN）：为了使KV缓存学习变得实用，本文构想了一个名为知识交付网络（KDN）的新系统组件。KDN旨在动态优化KV缓存在LLM引擎及其他计算和存储资源之间的存储、传输和组合。
    3. KDN的核心理念：类似于内容分发网络（CDN）通过高效的数据交付促进了互联网生态系统的成功，KDN旨在通过高效的知识交付，成为LLM应用成功的关键。其核心设计是将KV缓存管理与LLM服务引擎解耦，从而实现更高的模块化和效率。
    4. 原型验证：通过一个开源的KDN原型LMCache，初步验证了该方法的可行性，展示了其在模块化和效率方面相较于传统方法的显著优势。

图1. 知识注入方法在模块化和效率之间进行权衡。借助知识交付网络（KDN），与上下文学习和微调相比，KV缓存学习可以在两个维度上都得到改善。
图1. 知识注入方法在模块化和效率之间进行权衡。借助知识交付网络（KDN），与上下文学习和微调相比，KV缓存学习可以在两个维度上都得到改善。

A3 从系统角度看LLM知识注入

在人工智能文献中，知识增强生成，特别是微调和上下文学习，已得到充分研究。微调将文本语料库嵌入LLM的权重中，使得微调后的模型能以较低的响应延迟直接回应用户查询。然而，由于整个文本语料库必须一同嵌入模型中，微调缺乏添加新知识和指定应使用或不应使用哪些知识的灵活性。上下文学习则与微调相反，它允许操作者通过将文本放入LLM的输入中，轻松指定应使用哪些外部知识。然而，预填充（prefilling）的计算开销会随着输入长度超线性增长，导致当更多外部数据添加到输入中时，响应延迟会很长。

2.1 效率-模块化权衡

两种方法的质量相似，但系统特性不同。如果使用得当，两种方法都可以达到相似的文本生成质量【索引11, Simone Alghisi等人, Should We Fine-Tune or RAG? Evaluating Different Techniques to Adapt LLMs for Dialogue, arXiv preprint arXiv:2406.06399, 2024】【索引21, Aman Gupta等人, RAG vs Fine-tuning: Pipelines, Tradeoffs, and a Case Study on Agriculture, arXiv preprint arXiv:2401.08406, 2024】【索引44, Tianjun Zhang等人, Raft: Adapting language model to domain specific rag, arXiv preprint arXiv:2403.10131, 2024】。本文不从准确性的机器学习视角出发，而是从以下两个以系统为中心的指标来比较知识注入方法：模块化和效率。

  • 模块化（Modularity）：在知识增强LLM的背景下，方法的模块化包括两个方面。首先，模块化的方法应允许服务提供商轻松指定使用哪些知识并进行组合。其次，将新知识注入LLM的开销（如时间、成本）应最小化。上下文学习将新知识置于模型输入而非模型本身，这种知识与模型的分离是模块化的关键，服务提供商可以指定使用哪些知识并轻松组合不同知识片段，从而帮助LLM避免知识冲突并提高生成质量。相比之下，微调的模块化很差，用户无法指定微调模型中的哪些知识将用于生成答案。此外，每次有新知识和新模型时都需要进行微调，这可能需要数小时到数天才能完成。

  • 效率（Efficiency）：另一方面，知识增强LLM系统的效率通过LLM推理期间的单次查询成本和响应延迟来衡量。成本是LLM处理一个请求所用的计算量，响应延迟定义为LLM接收到请求到生成第一个token之间的时间。从这个角度看，上下文学习并不理想，因为它必须花费很长时间预填充包含知识的输入文本，然后才能生成第一个token。相比之下，微调在效率方面更好，因为知识已嵌入模型权重中，微调后的模型可以跳过漫长的预填充过程。

总结。简而言之，上下文学习更具模块化，但牺牲了效率；而微调虽然实现了更好的效率，但由于模块化差，在知识的整合和控制方面开销较大。

2.2 KV缓存学习

一种替代方案。KV缓存学习将知识存储在LLM生成的KV缓存中，并通过将KV缓存提供给LLM来注入知识，而无需修改模型。KV缓存以模型处理文本后生成的注意力状态形式存储知识，因此一旦生成，如果后续请求使用相同的上下文，LLM可以重用该KV缓存以跳过预填充。当许多查询重用相同的上下文时，重用其KV缓存可以减少单次查询的延迟和计算使用，同时仍保持与上下文学习相当的模块化。KV缓存学习的思想在LLM服务中已获得越来越多的关注（例如，【索引7, Perplexity, https://lmcache.github.io/2024-09-17-release/】【索引36, Ruoyu Qin等人, Mooncake: Kimi’s kvcache-centric architecture for llm serving, arXiv preprint arXiv:2407.00079, 2024】【索引46, Lianmin Zheng等人, Efficiently programming large language models using sglang, arXiv preprint arXiv:2312.07104, 2023】）。

为何在KV缓存中存储知识是划算的？。表面上看，使用KV缓存似乎只是一个空间换时间的权衡（用KV缓存存储空间换取更短的预填充时间），但这种权衡是有利的，原因有二：
- KV缓存被大量重用。许多上下文，特别是长上下文，经常被不同查询和不同用户重用。这可以看作是LLM版本的帕累托法则：LLM像人一样，80%的时间使用20%的知识，这意味着知识被频繁重用。例如，如果用户让LLM阅读一本书，他们不太可能只问一个与书相关的问题。
- KV缓存的大小增长慢于预填充延迟。随着上下文长度的增加，KV缓存大小呈线性增长，而预填充延迟则呈超线性增长。并且，随着LLM变得更大，更多的计算发生在不影响KV缓存大小的前馈层（feedforward layers）。

A2 方法细节

3. 用于高效知识注入的知识交付网络

现有KV缓存系统的局限性。尽管前景广阔，但现有的KV缓存学习系统仍存在一些技术局限性。
- KV缓存的存储空间有限：目前，许多服务引擎仅使用单个LLM实例本地可访问的GPU/CPU内存来存储KV缓存。这种仅限本地内存的存储方式极大地限制了可存储和重用的KV缓存数量。例如，一个64GB的CPU内存仅能为一个小型LLM（Llama3-4B）存储160K tokens（相当于两份pdf报告）的KV缓存。然而，将KV缓存的存储扩展到磁盘或远程服务器会严重限制将KV缓存加载到GPU内存的带宽。
- 仅限前缀的重用：为了重用KV缓存，大多数系统要求KV缓存的文本必须是LLM输入的前缀。尽管LLM天然支持重用前缀的KV缓存，但这种“仅共享前缀”的假设严重限制了其应用场景。例如，检索增强生成（RAG）会在LLM输入中拼接多个检索到的文本块，因此大多数重用的文本不会是前缀。
- 长上下文导致质量下降：最后，随着更多文本作为LLM的上下文被添加到输入中（即长上下文），LLM捕捉重要信息的能力可能会下降，从而降低推理质量。因此，当KV缓存被更多查询重复使用时，质量下降也会影响到更多查询。

3.1 知识交付架构

核心洞察与KDN架构。乍一看，先前基于KV缓存的系统面临的这些挑战可能看起来各不相同。然而，我们的洞察是，它们共享一个共同的需求——一个独立的KV缓存管理系统，该系统能够动态地压缩、组合和修改KV缓存，以优化KV缓存的存储和交付，以及基于KV缓存的LLM推理。我们将这样的系统称为知识交付网络（KDN）。如图2所示，我们构想的KDN架构包含三个主要模块：
- 存储模块（storage module）：该模块存储以各种文本为键的KV缓存，并离线修改KV缓存内容，以便在将KV缓存提供给LLM时提高推理质量。
- 交付模块（delivery module）：该模块将压缩后的KV缓存从存储设备传输到运行LLM的服务器，并在GPU中解压以供LLM服务引擎使用。
- 混合模块（blending module）：当多个文本被一起放入LLM查询的上下文中时，该模块动态地组合与这些不同文本相对应的多个KV缓存。

图2. 知识交付网络（KDN）的架构。
图2. 知识交付网络（KDN）的架构。

KDN的核心理念是解耦。现有的LLM服务系统【索引23, HuggingFace, text-generation-inference, 2024】【索引26, Woosuk Kwon等人, Efficient memory management for large language model serving with pagedattention, SOSP 2023】【索引46, Lianmin Zheng等人, Efficiently programming large language models using sglang, arXiv preprint arXiv:2312.07104, 2023】将KV缓存与LLM引擎深度耦合。相比之下，KDN的概念是将KV缓存的管理与LLM服务引擎分离。这种分离将促进KV缓存的存储、共享和使用方面的创新，而无需与快速发展的LLM服务引擎深度耦合。通过独立实现这些优化，KDN成为LLM服务生态系统中一个关键的新系统模块。它实现了KV缓存管理与LLM服务引擎的解耦，使LLM服务引擎能够专注于快速的查询驱动推理，而KDN则可以独立地专注于与KV缓存相关的优化。

3.2 技术路线图

KDN的可行性。KDN的架构本身并不能直接解决与KV缓存相关的挑战；它仅仅是将一个潜在的解决方案分解为三个模块。幸运的是，我们观察到新兴的研究成果可以为每个模块提供充分的设计，从而使KDN变得实用。

表1. 在RAG设置下不同知识注入方法的比较。借助KDN，KV缓存学习在整合新知识方面比微调快40倍，并且在推理过程中比上下文学习快3.7倍、成本低2.5倍。
表1. 在RAG设置下不同知识注入方法的比较。借助KDN，KV缓存学习在整合新知识方面比微调快40倍，并且在推理过程中比上下文学习快3.7倍、成本低2.5倍。

KV缓存交付。最近的KV缓存压缩技术使得在GPU和CPU内存之外廉价存储和快速加载KV缓存成为可能。例如，CacheGen【索引33, Yuhan Liu等人, Cachegen: Fast context loading for language model applications, arXiv preprint arXiv:2310.07240, 2023】通过量化然后将KV缓存编码为二进制字符串来对其进行压缩。LLMLingua【索引24, Huiqiang Jiang等人, Llmlingua: Compressing prompts for accelerated inference of large language models, arXiv preprint arXiv:2310.05736, 2023】引入一个较小的语言模型来识别并移除知识文本中的非必要tokens，从而减小相应KV缓存的大小。H2O【索引45, Zhenyu Zhang等人, H2O: Heavy-Hitter Oracle for Efficient Generative Inference of Large Language Models, Workshop on Efficient Systems for Foundation Models @ ICML2023, 2023】根据推理过程中计算的重要性直接移除KV缓存中的元素。通过结合上述技术，KV缓存的内存占用可以减少超过10倍，从而极大地提高了加载速度和KV缓存的存储成本。

KV缓存混合。最近的一些工作也提高了KV缓存的可组合性。例如，CacheBlend【索引40, Jiayi Yao等人, Cacheblend: Fast large language model serving with cached knowledge fusion, arXiv preprint arXiv:2405.16444, 2024】通过重新计算KV缓存之间的交叉注意力（cross-attention）来实现任意组合不同的KV缓存，而这种重新计算仅需预填充全文所需计算量的10%。PromptCache【索引19, In Gim等人, Prompt cache: Modular attention reuse for low-latency inference, 2023】允许用户定义一个包含不同片段的提示模板，从而使每个片段的KV缓存可以在不同位置被重用，而不仅仅是作为前缀。

离线KV缓存编辑。通过将KV缓存管理与LLM服务引擎分离，KDN为提高推理质量开辟了新的可能性。最近的研究表明，如果适当地修改LLM输入的注意力矩阵，推理质量可以显著提高【索引13, Anonymous, Model tells itself where to attend: Faithfulness meets automatic attention steering, Submitted to ACL Rolling Review - June 2024, 2024】【索引42, Qingru Zhang等人, Tell your model where to attend: Post-hoc attention steering for llms, arXiv preprint arXiv:2311.02262, 2023】。换句话说，当一个KV缓存被“存入”KDN时，KDN不仅存储它，还可以通过离线编辑KV缓存并在下次检索时返回编辑后的版本来主动影响LLM的推理质量，所有这些都无需对模型本身或输入提示进行任何更改。

与LLM服务引擎的接口。目前，大多数LLM服务引擎，如vLLM【索引26, Woosuk Kwon等人, Efficient memory management for large language model serving with pagedattention, SOSP 2023】、HuggingFace TGI【索引23, HuggingFace, text-generation-inference, 2024】和SGLang【索引46, Lianmin Zheng等人, Efficiently programming large language models using sglang, arXiv preprint arXiv:2312.07104, 2023】，并不直接支持将外部提供的KV缓存作为LLM输入的一部分注入。相反，它们内部的KV缓存管理实现（例如vLLM中的分页内存）与模型推理实现深度耦合。部署KDN的一种方法是为每个LLM引擎增加一个KDN作为其子模块。然而，如3.2节所述，开发一个高性能的KDN是一项庞大的工程，如果每个LLM引擎都维护和发展自己的KDN，将导致大量的重复工程工作。为避免重复造轮子，LLM服务引擎可以通过两个API与一个独立的、共享的KDN提供商接口：（i）LLM将KV缓存及其关联文本存储到KDN，以及（ii）LLM使用某个文本从KDN检索KV缓存。鉴于大多数流行的LLM引擎都是开源的，暴露这些API是可行的。尽管如此，仍有几个设计问题待解决。如何利用异构链接（如NVLink、RDMA或PCIe）来传输KV缓存？这些API是否可以与其他LLM功能（如解耦的预填充）共享？

初步验证。基于这些技术，我们实现了一个KDN原型 LMCache (https://github.com/LMCache/LMCache)。我们在一个RAG用例中，比较了使用KDN的KV缓存学习与微调和上下文学习在模块化和效率方面的表现，设置如下：
- 知识总量（文本）: 200万 tokens。
- 每个请求: 8K tokens 的知识加上 2K tokens 的用户聊天记录。
- 语言模型: Llama 3.1 70B 【索引18, Abhimanyu Dubey等人, The llama 3 herd of models, arXiv preprint arXiv:2407.21783, 2024】。
- 硬件: 2块Nvidia A40 GPU。

模块化通过将知识库注入LLM的时间来衡量，效率通过每次请求的推理成本和响应延迟来衡量。表1显示，借助KDN，KV缓存学习在注入新知识时比微调快40倍，并且在推理时比上下文学习的成本低3.7倍、速度快2.5倍。

A4 实验环境与结果

实验环境

  • 知识库与数据集：实验基于一个检索增强生成（RAG）用例。知识库总大小为200万tokens。每个请求包含8000 tokens的知识和2000 tokens的用户聊天历史。
  • 模型架构：使用了Llama 3.1 70B模型【索引18, Abhimanyu Dubey等人, The llama 3 herd of models, arXiv preprint arXiv:2407.21783, 2024】。
  • 硬件配置：实验在配备2块Nvidia A40 GPU的服务器上进行。
  • 软件配置：实现了一个名为LMCache的KDN原型，该原型已开源：https://github.com/LMCache/LMCache。

实验结果

实验旨在比较三种知识注入方法：KV缓存学习（使用KDN）、微调和上下文学习。评估指标包括模块化（注入知识库所需的时间）和效率（单次请求的推理成本和响应延迟）。

  • 实验内容：

    1. 模块化评估：测量将200万tokens的知识库整合到LLM所需的时间。
    2. 效率评估：测量处理一个包含8K知识和2K历史的请求的成本和端到端响应延迟。

  • 实验结论（见表1）：

    • 模块化优势：与微调相比，使用KDN的KV缓存学习在注入新知识方面的速度快40倍。
    • 效率优势：与上下文学习相比，使用KDN的KV缓存学习在推理期间的成本便宜3.7倍，响应速度快2.5倍。
    • 这一结果证明，KDN可以有效地结合模块化和效率，克服传统方法的局限性。

A5 结论

简而言之，本文主张（1）将以KV缓存形式存在的知识管理与LLM服务引擎分离，以及（2）提出知识交付网络（KDN）作为一个新的LLM系统组件，该组件利用最新的研究进展来优化基于KV缓存的知识注入的效率（在速度和成本方面）。我们希望本文能够激励更多研究来解决上述问题。

方法细节中的引用汇总

参考文献编号 内容 引用段落及描述
[13] Anonymous. Model tells itself where to attend... (ACL Rolling Review 2024) 3.2 技术路线图 - 离线KV缓存编辑：引用此文献证明，通过修改注意力矩阵可以提高推理质量，这为KDN通过离线编辑KV缓存来提升性能提供了理论依据。
[19] In Gim et al. Prompt cache... (2023) 3.2 技术路线图 - KV缓存混合：引用此文献作为KV缓存组合性的一个示例，其允许对提示的不同片段进行缓存重用，打破了仅限前缀重用的限制。
[23] HuggingFace. text-generation-inference (2024) 3.1 知识交付架构 & 3.2 技术路线图 - 与LLM服务引擎的接口：引用HuggingFace TGI作为现有LLM服务引擎的一个例子，这些引擎将KV缓存管理与推理引擎深度耦合，缺乏外部KV缓存注入的接口。
[24] Huiqiang Jiang et al. Llmlingua... (arXiv 2023) 3.2 技术路线图 - KV缓存交付：引用此文献介绍了一种KV缓存压缩技术，该技术使用一个小型语言模型去除不重要的tokens，从而减小KV缓存大小，使得在GPU外部存储和快速传输KV缓存成为可能。
[26] Woosuk Kwon et al. Efficient memory management... (SOSP 2023) 3.1 知识交付架构 & 3.2 技术路线图 - 与LLM服务引擎的接口：引用vLLM作为现有LLM服务引擎的一个例子，其内部的KV缓存管理（如paged attention）与推理实现紧密耦合，不支持外部缓存注入。
[33] Yuhan Liu et al. Cachegen... (arXiv 2023) 3.2 技术路线图 - KV缓存交付：引用此文献介绍了一种KV缓存压缩技术（量化和二元编码），该技术可以大幅减小KV缓存的内存占用，从而支持高效的存储和交付。
[40] Jiayi Yao et al. Cacheblend... (arXiv 2024) 3.2 技术路线图 - KV缓存混合：引用此文献介绍了一种先进的KV缓存组合技术，它通过以较低的计算成本（仅为预填充的10%）重新计算交叉注意力，实现了对不同KV缓存的任意组合。
[42] Qingru Zhang et al. Tell your model where to attend... (arXiv 2023) 3.2 技术路线图 - 离线KV缓存编辑：与[13]一同被引用，以支持KDN通过离线编辑注意力矩阵来提升推理质量的潜力。
[45] Zhenyu Zhang et al. H2O... (ICML 2023 Workshop) 3.2 技术路线图 - KV缓存交付：引用此文献介绍了一种基于重要性直接移除KV缓存元素的压缩技术，是实现高效KV缓存交付的关键技术之一。
[46] Lianmin Zheng et al. Efficiently programming... (arXiv 2023) 2.2 KV缓存学习 & 3.1 知识交付架构 & 3.2 技术路线图 - 与LLM服务引擎的接口：引用SGLang作为现有LLM服务系统的一个例子，说明KV缓存学习已获得关注，同时也指出其KV缓存管理与引擎深度耦合的局限性。