温馨提示：
本篇文章已同步至"AI专题精讲 " CacheGen：用于快速大语言模型推理服务的 KV 缓存压缩与流式传输
摘要
随着大语言模型（LLMs）处理的任务日益复杂，其输入通常需要补充更长的上下文以整合领域知识。然而，使用长上下文面临挑战：在整个上下文被模型处理完之前，模型无法开始生成任何内容。虽然可以通过在不同输入间复用上下文的 KV 缓存来减少处理延迟，但由于 KV 缓存包含大型张量，跨网络获取这些缓存会引入较高的额外网络延迟。
CacheGen 是一个用于 LLM 系统的快速上下文加载模块。首先，CacheGen 利用 KV 缓存的分布特性，通过自定义张量编码器将 KV 缓存编码为更紧凑的比特流表示，在几乎无解码开销的前提下大幅节省带宽。其次，CacheGen 会根据可用带宽的变化，自适应调整 KV 缓存不同部分的压缩等级，从而保持较低的上下文加载延迟和较高的生成质量。我们在多个主流 LLM 模型和数据集上对 CacheGen 进行了测试。与现有 KV 缓存复用系统相比，CacheGen 将 KV 缓存大小压缩了 3.5--4.3 倍 ，将上下文的获取与处理总延迟降低了 3.2--3.7 倍 ，并且对生成质量影响可以忽略。我们的代码已开源，地址为：https://github.com/UChi-JCL/CacheGen
CCS 分类概念
- 计算方法论 → 自然语言生成
- 网络 → 应用层协议
- 信息系统 → 信息系统应用
关键词
大语言模型、KV 缓存、压缩
1 引言
得益于其出色的生成质量，大语言模型（LLMs）已被广泛应用于个人助理、AI 医疗和市场营销等场景中 [22, 38, 46, 128]。LLM API（如 OpenAI GPT-4 [108]）的普及、行业级开源模型（如 Llama [129]）的发展，以及应用框架（如 HuggingFace [10] 和 Langchain [83]）的广泛使用，进一步推动了 LLM 的普及。
为了完成复杂任务，用户或应用程序通常会在 LLM 输入前添加包含数千个 token 的长上下文。例如，某些上下文会在用户提示前补充领域知识文本，从而使 LLM 能够使用其未显式掌握的特定知识生成响应。另一个例子是，在用户与 LLM 交互过程中产生的对话历史，会被追加作为上下文补充用户提示。尽管短输入仍然有用 [94, 124]，但更长的输入通常能提升回答质量与连贯性 [31, 32, 35, 45, 67, 116, 130, 141]，这也催生了对能处理更长上下文的 LLM 的竞赛，从 ChatGPT 的 2K token 提升到 Claude 的 100K token [24]。
然而，使用长上下文会带来响应延迟的问题：在整个上下文被加载并处理完之前，模型无法生成任何响应。处理长上下文所需的计算量会随着上下文长度呈超线性增长 [31, 47, 116, 131, 150]。尽管一些最新研究致力于提升长上下文处理的吞吐量 [17]，但处理延迟仍可能高达数秒（例如处理一个 3K token 的上下文约需 2 秒）[17, 58]。为应对这一问题，许多系统尝试通过存储并复用上下文的 KV 缓存，跳过重复计算以减少上下文处理延迟（例如 [23, 58, 82, 156]）。
然而，当下一个输入到来时，所需上下文的 KV 缓存未必仍保存在本地 GPU 显存中，而是可能需要从其他机器获取，这就引入了额外的网络延迟（如图 1a 所示）。例如，背景文档数据库可能保存在独立的存储服务中，当接收到相关查询时，系统才会选取并加载这些文档（即上下文）以辅助 LLM 推理 [27, 31, 36, 84, 110]。
目前，对于从网络中加载 KV 缓存所引发的延迟问题，尚未受到足够关注。现有系统往往假设上下文的 KV 缓存在多个请求之间始终保存在同一 GPU 显存中 [58]，或者假设 KV 缓存体积足够小，可以通过高速互联设备快速传输 [111, 157]。然而，正如第 §3 节将详细讨论的那样，KV 缓存本质上是高维浮点张量，其大小随上下文长度与模型规模增长，容易达到几十 GB，因而网络传输延迟可达数百毫秒甚至超过 10 秒，严重影响交互体验 [1, 2, 87]。
简而言之，在从其他机器加载 KV 缓存的场景下，若仅优化计算延迟，可能反而会加剧整体响应延迟，因为加载 KV 缓存会引入额外的网络延迟。
近年来出现了一些旨在减少 GPU 显存中 KV 缓存运行时尺寸的研究工作，目的是为了适应显存限制或 LLM 的输入限制。一些方法通过从 KV 缓存或上下文文本中移除不重要的 token来实现 [71, 72, 95, 153]，另一些方法则对 KV 缓存张量应用智能量化技术 [62, 78, 97]。与此不同，我们的目标是减少 KV 缓存的传输时间大小，从而降低网络延迟。因此，我们不需要保留 KV 缓存的张量格式，而是可以将其编码为更紧凑的比特流。
我们提出了 CacheGen，这是一个用于加速 LLM 系统上下文加载的模块，专注于减少在获取和处理长上下文时的网络延迟（如图 1b 所示）。它包含两个关键技术：
KV 缓存的编码与解码：
CacheGen 将预先计算的 KV 缓存编码为更紧凑的比特流表示，而不是保留原始张量的形状。这样在传输 KV 缓存时能显著节省带宽并减少延迟。我们设计的 KV 缓存编码器采用了自定义的量化与算术编码策略，利用了 KV 缓存的分布特性，例如：
- 邻近 token 之间的 KV 张量具有局部性；
- KV 缓存的不同层对量化误差的敏感度不同。
此外，KV 缓存的解码（解压缩）通过 GPU 加速实现，并与传输过程流水线并行执行，进一步减少了解码对整体推理延迟的影响。
KV 缓存的流式传输：
CacheGen 以适应网络条件变化的方式流式传输 KV 缓存的编码比特流。在用户查询到达之前，CacheGen 会将长上下文划分为多个 chunk（块），并对每个块在不同压缩级别下分别编码（类似于视频流媒体技术）。当需要发送某段上下文的 KV 缓存时，CacheGen 会逐块获取并动态调整每块的压缩率，以在保持高生成质量的同时，控制网络延迟在服务等级目标（SLO）范围内。
当网络带宽太低时，CacheGen 也可以回退为以文本格式发送该块，并交由 LLM 自行重新计算该块的 KV 缓存。
简而言之，不同于以往只关注 GPU 显存中 KV 缓存优化的系统，CacheGen 专注于优化 KV 缓存的网络传输延迟。我们将 CacheGen 与多个基线方法进行了对比，包括 KV 量化 [120]、加载文本形式的上下文、以及最新的上下文压缩技术 [72, 153]。实验涵盖了三种不同规模的流行 LLM（从 7B 到 70B）以及四个包含长上下文的数据集（共 662 个上下文，长度在 1.4K 到 16K tokens 之间）。表 1 展示了部分实验结果。我们的主要发现如下：
-
在上下文传输与处理延迟方面（即首 token 响应时间），CacheGen 相比量化基线快3.2--3.7 倍，且生成质量（F1 分数和困惑度）相当；相比加载文本上下文，CacheGen 快 3.1--4.7 倍，且精度下降不超过 2%。值得注意的是，与 8-bit 量化（几乎无损压缩）相比，CacheGen 仍能将上下文加载延迟降低 1.67--1.81 倍。
-
在带宽使用方面，CacheGen 在保持相同生成质量的情况下，所用带宽比量化基线少 3.5--4.3 倍。
-
结合最新上下文压缩技术 [72, 153] 使用时，CacheGen 还能进一步将其 KV 缓存传输带宽需求降低 3.3--4.2 倍。
本研究未涉及任何伦理问题。
2 背景与动机
2.1 大语言模型基础知识
Transformer 模型 [37, 44, 131] 是当前几乎所有大语言模型（LLM）服务的事实标准。从高层次来看，Transformer 接收一系列输入 token，并通过两个阶段生成输出 token 序列。
在 prefill 阶段，一个注意力神经网络处理输入 token。注意力模块中的每一层（共 𝑙 层）会生成两个二维张量：一个 key（K）张量和一个 value（V）张量。这些 K 和 V 张量中包含了 LLM 后续使用上下文所需的重要信息。所有层中的 KV 张量统称为KV 缓存（KV cache）。
在 生成阶段（也称解码阶段），LLM 使用 KV 缓存来计算 token 之间的注意力得分，从而生成注意力矩阵，并以自回归的方式逐步生成输出 token。为了性能考虑，KV 缓存通常保存在 GPU 显存中（因为其占用内存较大 [82]），并在生成结束后释放。一些新兴优化策略支持 跨 LLM 请求保存和复用 KV 缓存，我们将在后文中详细解释。
在所有主流模型中，prefill 阶段的计算开销会随着输入长度呈超线性增长。由于必须在 prefill 完成后才能生成首个输出 token，因此其耗时被称为 Time-to-First-Token（TTFT）。本文的目标是在不改变解码过程的前提下，降低 prefill 阶段的 TTFT。
2.2 LLM 输入中的上下文
当响应内容所需的知识未嵌入在模型中时，LLM 可能会生成质量较低或出现幻觉的回答。因此，许多 LLM 应用和用户会在输入中补充额外文本，称为上下文（context） [53, 89]。LLM 可通过先阅读上下文，利用**上下文学习（in-context learning）**能力生成更高质量的回答。
LLM 输入中的上下文具有多种用途：
- (i) 用户问题可以通过包含特定领域知识的文档补充，以生成更准确的回答 [3, 7, 117]，如使用最新新闻进行事实核查 [8, 9]，或使用判例法、法规文件进行法律辅助 [118, 125]；
- (ii) 代码分析应用可从代码仓库中检索上下文，以回答问题或生成摘要 [30, 69, 73]；金融公司也采用类似方式，基于详细的财务文档生成摘要或回答问题 [105]；
- (iii) 游戏应用中，使用特定角色的描述作为上下文，使 LLM 能生成符合角色性格的对话或动作 [110, 121, 140]；
- (iv) 在小样本学习（few-shot learning）中，一组问答对被用作上下文，指导 LLM 学会回答某类问题 [18, 99, 123]；
- (v) 在聊天应用中，用户的对话历史通常会作为上下文附加在新输入前，以保证回复的一致性与连贯性 [26, 76]。
我们观察到，在实际应用中，上下文往往很长，且经常被复用来补充不同的用户输入。
长上下文在实践中变得越来越常见。例如前述的各种文档（如案例法、财务报告、新闻文章、代码文件或聊天记录）都可以轻易地包含上千个 token。直观来看，更长的上下文更有可能包含相关信息，因此有助于提升响应质量。
确实，FiD [67] 显示当上下文长度从 1K token 增加到 10K 时，模型准确率从 40% 提升到 48%；Retro [35] 同样表明上下文从 6K 增至 24K 时，生成质量（困惑度）显著提升。
本文关注的上下文包括：如聊天会话中积累的对话历史，或用户提供的用于完成任务所需的单个文档。
此外，这些长上下文通常会被不同输入反复使用。例如，在财务分析中，针对两个不同查询："请基于公司上季度财报写一个简短总结"和"公司上季度的主要收入来源是什么？"，补充的上下文很可能是同一份财报。类似地，在法律助手或事实核查应用中，相同的执法文件或新闻报道可以用于多个不同问题。又如在一次聊天过程中，早期的对话内容会反复作为后续输入的上下文。
简而言之，更长的上下文会导致更高的 prefill 延迟，从而拉长 Time-to-First-Token（TTFT）。但由于同一段上下文经常被复用，因此通过缓存中间结果（即 KV 缓存）来避免重复计算 prefill 阶段，从而降低 TTFT，是一个很有前景的优化方向。这个思路在近期的一些工作中 [23, 58, 82] 已经得到了探索，并展示出其潜力。但该方案仍存在一个关键问题，我们将在下一节中详细讨论。
3 隐藏的网络瓶颈
尽管复用长上下文的 KV 缓存可以显著降低 TTFT（首次生成时间），但这个优势存在一个前提条件------复用的 KV 缓存必须首先保存在本地 GPU 显存中 [23, 58, 82, 156]。
为什么 KV 缓存需要加载？
实际上，复用的 KV 缓存往往需要从其他机器中取回。这是因为 GPU 显存不足以长期保存多个复用上下文的 KV 缓存。例如，在金融助手应用中，LLM 需要对长篇的财务报告进行分析 [107]，这些报告常包含成千上万个 token，导致 KV 缓存非常庞大。具体来说，处理 Amazon 2023 年度报告（大约 80,000 个 token）时，使用 Llama-34B 模型会生成 19 GB 大小的 KV 缓存，这一规模已经和 LLM 模型本身相当。
由于不同查询对 KV 缓存的复用可能间隔数小时，为了给新的聊天会话腾出空间，旧的 KV 缓存可能被提前卸载。此外，随着新一代 LLM 能处理更长的上下文 [51, 56, 63, 91, 138]，将 KV 缓存存储在专用的存储服务器（而非 CPU 或 GPU 内存）将更加实用和经济。另外，不同的请求可能不会命中同一个 GPU，也意味着 KV 缓存必须跨机器迁移。从其他机器加载 KV 缓存会带来显著的延迟，但这一"网络延迟"问题尚未得到充分重视。
这是一个新问题吗？
虽然近期一些工作确实提出了通过 GPU 之间共享 KV 缓存来实现多 GPU 推理，但它们假设使用高速互联（如 NVLink）[111, 157]，这些链路的带宽高达数百 Gbps，因此网络延迟几乎可以忽略不计。
但现实中，KV 缓存也常常需要通过普通云服务器之间的低速链路进行传输，而这些链路的带宽通常只有几 Gbps [70]。如图 2b 所示，在这种环境下，将 KV 缓存加载到 GPU 显存的延迟可能和不用 KV 缓存时执行 prefill 的耗时一样长，甚至更长。
我们的方法：
本文关注于降低从其他机器加载 KV 缓存的网络延迟。为此，我们提出将 KV 缓存编码为更紧凑的比特流表示形式（如图 2c 所示），以此压缩其大小，降低传输延迟。
这一定程度上看似类似于最近的一些研究，它们采用的方法包括：
- 从文本上下文中丢弃部分 token；
- 对 KV 缓存张量进行量化 [62, 78, 95, 97, 153]。
但关键区别在于：
- 上述技术主要是为了减少运行时在 GPU 中的 KV 缓存占用，保留其张量形状；
- 而我们的方法是为了减少传输时的大小 ，通过编码为比特流以减少网络传输延迟。
此外，这两类方法还可以天然协同：那些经过缩减的 KV 缓存，仍然可以通过我们的编码方式进一步压缩，从而进一步降低网络延迟。
4 CacheGen：KV Cache 编码与流式传输
减少 KV cache 传输延迟的需求促使我们在 LLM 系统中引入一个新的模块，我们称之为 KV cache 流处理器（KV cache streamer）。该模块承担三个角色：
- (1) 将给定的 KV cache 编码为更紧凑的比特流表示形式 ------ KV bitstreams。该过程可以在离线进行。
- (2) 通过吞吐量变化的网络连接流式传输编码后的 KV bitstream。
- (3) 将接收到的 KV bitstream 解码为 KV cache。
乍一看，我们的 KV cache 流处理器可能看起来与近期的一些技术（例如 [72, 95, 153]）类似，这些技术通过丢弃不太重要的 token 来压缩长上下文。然而，它们之间存在关键区别：
这些最新技术旨在减少运行时 KV cache 的大小，以适应 GPU 的内存限制或 LLM 的输入窗口限制；而我们则旨在减少 KV cache 的传输时大小，以降低网络延迟。因此，前述技术必须保留 KV cache 的大型浮点张量的张量形状，以便在运行时直接被 LLM 使用；同时，它们可以利用生成阶段的信息，判断当前 query 中哪些上下文 token 更为重要。而相反，我们不需要保留原始的张量形状，可以将其编码为更紧凑的比特流，并根据网络带宽调整表示形式。同时，我们必须在特定 query 被处理之前就决定使用哪种压缩方案，因此无法使用生成阶段的信息。
本文提出了 CacheGen，这是 KV cache 流处理器的一个具体实现设计。首先，CacheGen 使用定制的 KV cache 编解码器（encoder 和 decoder），通过利用 KV cache 张量的若干分布属性（§5.1），最大程度地压缩 KV bitstream 的大小。这大大减少了传输 KV cache 所需的带宽，从而直接降低了 TTFT。其次，在动态带宽环境下进行 KV bitstream 流式传输时，CacheGen 会在不同编码等级之间动态切换，或根据需要重新计算 KV cache，以在保持高响应质量的同时，将 TTFT 控制在给定的截止时间内。KV 的编码与解码计算开销可以忽略不计，并与网络传输流水线并行执行，以最小化其对端到端延迟的影响。
5 CacheGen 设计
我们现在介绍 CacheGen 的设计，首先从 KV cache 的经验性洞见开始（§5.1），它启发了 KV cache 编码器的设计（§5.2），随后介绍 CacheGen 如何适应带宽变化（§5.3）。
5.1 KV cache 的经验性洞见
我们强调了关于 KV cache 数值特性的三个观察。尽管从理论上难以证明这些观察对任意 LLM 和任意上下文都成立，但我们使用具有代表性的工作负载，通过实证方式展示这些观察的普遍性。该工作负载包括两个不同容量的 LLM（Llama-7B 和 Llama-13B）以及 LongChat 数据集 [90]（包含从总共 200 个上下文中随机抽取的 100 个长上下文，长度在 9.2K 到 9.6K tokens 之间），这是最大规模的长上下文数据集之一。该工作负载的详细信息见 §7.1。
5.1.1 Token 级的局部性
第一个观察涉及上下文中 token 间 K 和 V 张量数值的变化。具体而言，我们观察到：
洞见 1：在相同的 layer 和 channel 中，相邻 token 的 K/V 张量值要比距离较远的 token 更相似。
对于每个模型，我们对比了 K（或 V）张量的原始值分布和 delta 值的分布 ------ 即上下文中每对相邻 token 在同一层和同一通道上的 K（或 V）张量值之差。图 3 展示了所有上下文中某一层的原始张量绝对值分布与其 delta 值分布³。在两个模型的所有上下文中，我们都可以看到 delta 的值更加集中于零附近，而不像原始值那样分散。因此，delta 的方差比原始值的方差低 2.4-2.9 倍。K 和 V 张量的 token 级局部性启发了 CacheGen：使用 delta 编码而不是原始值。
这种 token 级的局部性可以通过 transformer 的自注意力机制直观地解释，该机制负责计算 KV 张量。从数学角度来看，该机制等价于基于前一个 token 的 KV 张量来计算当前 token 的 KV 张量。这意味着一个 token 的 KV 张量在本质上与前一个 token 的 KV 张量相关联。
5.1.2 按层的丢失敏感性
第二个观察涉及 K 和 V 张量中不同数值对数据丢失的敏感程度。我们的观察如下：
洞见 2：LLM 的输出质量对浅层 KV cache 数值的丢失更敏感，而对深层的丢失不那么敏感。
不同层之间对丢失的敏感性不一致，这表明我们的 KV cache 编码器应对不同的层采用不同的压缩方式。图 4 展示了在 K 和 V 张量中某一特定层组施加数据丢失时对准确率的影响程度。这里，我们使用四舍五入作为数据丢失的手段，并计算该数据集中 100 个上下文所对应的平均响应准确率（定义见 §7.1）。可以看出，当数据丢失施加在模型的早期层时，平均响应准确率显著下降；而同样的丢失作用在较深层时，对准确率的影响则要小得多。这个结果在我们测试的不同模型中都表现得非常一致。
从直觉上看，KV cache 的深层提取的是更高层次的结构和知识，而浅层则包含更基础的信息 [119, 132]。因此，如果在浅层 cache 上移除精度造成信息丢失，这种影响可能会向后传播并影响后续层的 cache，从而削弱模型获取用于生成高质量响应所需的高层次结构信息的能力。
5.1.3 沿层、通道和 token 分布的情况
最后，关于 KV cache 在三个维度------层（layers）、通道（channels）和 token 位置（tokens）上的数值分布，我们有如下观察：
洞见 3 ：KV cache 中的每一个数值都由其对应的通道、层和 token 位置索引。将这些数值按通道和层进行分组所带来的信息增益，远高于按 token 位置分组所带来的信息增益。
直观上，这可以粗略地解释为：同一通道（channel）或层（layer）中的不同 KV 值之间相互之间比同一 token 位置下的不同 KV 值更相似。一种可能的解释是，不同的通道或层会捕捉输入中的不同特征 [49, 92]。例如，一些通道可能专注于主谓宾关系，而另一些则关注形容词等修饰成分。至于不同的层，以往的研究指出，较深的层比浅层捕捉到更抽象的语义信息 [49, 92]。另一方面，在给定层和通道的情况下，不同 token 的 KV 值之间也更相似，这很可能是由于自注意力机制的影响------每个 token 的 KV 都是由其前面所有 token 推导而来的。对于这些现象的更深入分析，我们将留待未来的工作中探讨。
为了从实证上验证这一洞见，我们将由两个模型和 100 个上下文所产生的 KV cache 中的数值，分别按照其层、通道或 token 位置进行分组，然后计算每组的熵。图 5 展示了在不同分组策略下的平均熵（以 bit/元素 为单位），包括：不分组、按 token 位置分组、按通道分组、按层分组。结果显示，按 token 位置分组所带来的熵下降远不如按通道或按层分组带来的熵下降效果明显。
温馨提示：
阅读全文请访问"AI深语解构 " CacheGen：用于快速大语言模型推理服务的 KV 缓存压缩与流式传输