点击下方卡片，关注
「集智书童」
公众号
点击加入👉
「集智书童」
交流群
导读
随着大语言模型（LLMs）中上下文窗口大小的增加，它们处理复杂、长文本任务的能力得到了提升。然而，随着对话轮次的持续，需要在GPU内存中存储大量的KV缓存，这显著影响了模型服务系统的效率甚至可用性。本文分析了来自真实用户的对话数据，并发现LLM推理表现出一个分水岭层，在此之后，轮次 Level 的注意力分布显示出显著的相似性。作者提出了轮次注意力（Round Attention），一种新颖的轮次 Level 注意力机制，它只召回和计算最相关轮次的KV缓存。实验表明，Round-Attention在不影响模型性能的情况下，节省了56%的内存使用。
1 引言
近年来，大语言模型的发展促进了语言模型服务在日常问题解决任务中的更广泛采用。然而，长时间的交互暴露出两个显著挑战。首先，由于自注意力机制的二次扩展，上下文长度的快速扩展导致巨大的计算开销。其次，尽管键值（KV）缓存可以缓解冗余计算，但它大幅增加了GPU内存需求，导致推理批次大小有限和GPU利用率低下。例如，配备40GB内存的NVIDIA A100只能容纳单个上下文长度为128K的LLaMA请求，其近50%的处理时间用于KV缓存访问。
为了提高推理效率，先前的研究探讨了LLMs的KV缓存驱逐和Sparse注意力技术，指出注意力本质上是Sparse的。这些方法要么将整个KV缓存存储在GPU内存中，在自回归过程中选择关键 Token 以减少交叉注意力计算时间，要么在CPU内存中维护KV缓存，在推理过程中逐个将 Token 传输到GPU内存中。前者并未减少GPU内存使用，而后者则产生了显著的通信开销。此外，当前方法通常需要计算每层最相关的 Token ，这需要昂贵的计算。
上述研究分析了词汇层面的上下文关系。Sun等人[2024]的研究分析表明，大多数RoPE之后的键与相邻词汇具有高余弦相似度，这为选择重要词汇提供了块级近似。LONGMEMEVAL基准测试探讨了记忆增强聊天助手的内存设计选项，并发现“轮次”是“最佳”的存储和利用交互历史的粒度。这与现实生活中个体交流的沟通习惯相一致，因为轮次比词汇更能清晰地表达更接近的上下文关系。
因此，作者通过多轮分析注意力矩阵，并识别出两种有趣的模式。首先，在目前流行的开源大型模型中，注意力分数的分布在不同层级的初始层表现出相当大的可变性；然而，从某一层开始，层与层之间的分布变得非常相似。其次，在单个对话轮次内，针对“问题”与先前对话轮次相关的注意力分数与针对相同先前轮次的“答案”计算出的注意力分数非常相似。
基于这些发现，作者提出了轮次注意力（Round Attention）方法，该方法利用了注意力矩阵的Sparse性。在推理过程中，它仅将最相关的轮次的键值（KV）缓存纳入注意力计算，并将完整的KV缓存卸载到CPU内存。轮次注意力在轮次粒度上存储和传输KV缓存，将每个轮次的KV缓存分割成两个完整的张量。由于第一个识别出的模式，作者只需要在特定层计算一次前k个
轮次，然后执行一次主机到设备（h2d）操作，将相应的KV缓存张量传输到GPU内存。这种方法与其他需要在每个层计算前k个
并按 Token 粒度传输KV缓存的方法形成对比，显著减少了与其他方法中提到的top-k计算和卸载相关的延迟开销。
作者的主要贡献如下：
- • 作者在LLM部署后以轮次粒度剖析了注意力模式，并在实际应用中揭示了注意力矩阵中的两个启发性特征。
- • 基于这些特性，作者设计了一种新颖的方法，称为“轮次注意力”，并关联了一系列用于长上下文对话的技术。这种方法以轮次粒度存储和传递KV缓存。
- • 作者对所Round-Attention进行了广泛的实验。结果表明，该方法可以将GPU内存占用减少55%，且不会损失精度。更重要的是，得益于一次性的top-k选择和主机到设备（h2d）传输，Round-Attention相比标准的非卸载Flash Attention实现了更低的延迟。
2 相关工作
2.1 注意力矩阵分析
预训练大语言模型（LLM）中注意力权重的Sparse性，尤其是在长上下文场景中，已被广泛记录。马等[2024]研究了上下文中重要 Token 的分布，发现最近 Token 比远距离 Token 更重要。他们还发现，连续层之间的注意力分数相似，这一现象也曾在较小模型中观察到（肖等[2019]，博贾纳帕利等[2021]）。
穆等[2024]报告称，在Transformer层之间，尤其是相邻层之间，注意力权重非常相似。门等[2024]发现LLM层之间存在显著的冗余，其中某些层对模型贡献微乎其微。范等[2024]表明，对于某些任务，LLM可以在某些中间层达到与最终输出相当的结果。
2.2 KV缓存
许多先前的研究致力于KV缓存压缩以加速注意力机制并减少内存使用。H2O Zhang等人[2023]为历史注意力分数总和的重要KV缓存保留了一定的预算。FastGen Ge等人[2023]进一步对 Token 进行分类，并使用更复杂的策略仅保留部分KV缓存。TOVA Oren等人[2024]通过确定当前 Query 中永久丢弃的 Token 来简化策略。StreamingLLM Xiao等人[2023]使用注意力汇聚和有限的KV缓存处理无限长的文本。SparQ Ribar等人[2023]通过通道剪枝计算近似注意力分数，并通过这些分数选择重要的 Token 。Tang等人[2024]得出结论， Token 的重要性高度依赖于 Query ，并提出了Quest方法，该方法记录KV缓存页面中的最小和最大键值，并使用 Query 向量估计页面的重要性。
然而，这些方法面临着几个挑战。首先，识别topk注意力是非常昂贵的。例如，应用一个简单的搜索算法，如IVF Douze等人[2024]，需要访问超过30%的关键状态才能获得topk结果 Liu等人[2024]，这非常耗费计算资源。其次，这些方法将KV缓存保存在GPU内存中，以避免从CPU内存中加载它们，但这并没有减少KV缓存的总内存消耗，因此限制了最大上下文窗口和推理批大小。
一些论文尝试将KV缓存卸载到CPU内存中，以减少活跃GPU内存的使用。刘等人[2024]提出在CPU内存中为KV向量构建近似最近邻搜索（ANNS）索引，并在生成过程中通过向量搜索检索最相关的结果。孙等人[2024]存储低秩键缓存，并将值缓存卸载，以减少更大批量和更长时间序列的内存占用。陈等人[2024]存储LSH哈希表，并在CPU上运行注意力计算，这显著减少了注意力计算的工作量。然而，这些工作在 Token Level 传输键值（KV）缓存，在某些方法中，top-k选择是在每层基础上计算的，这意味着KV缓存也是逐层传输的，导致h2d传输产生显著开销。
3 方法论
本节介绍了圆注意力（Round Attention），这是一种新颖的方法，通过将注意力矩阵在圆 Level 上进行分解，以
对作为基本分析单元，用于多轮对话任务。目标是减少LLMs的内存占用和推理延迟，同时不牺牲其准确性。本节还将讨论圆注意力背后的技术套件。
3.1 注意力分布
给定输入序列
，标准的Transformer网络（Vaswani等，[2023]）通过在
上应用线性变换来计算一组 Query
、键
和值
。然后，它按照以下方式计算自注意力分数：
为
了
研
究
各
轮
次
之
间
的
注
意
力
模
式
，
作
者
定
义
第
层
第
轮
中
的
以
及
前
轮
中
的
的
注
意
力
分
数
之
和
为
：
\mathsf{q A t t e n t i o n}
{k}^{l}=\sum
{\substack{i\in\mathbf{q}
{n}}}\mathrm{Attention}(Q
{i}^{l},K_{j}^{l})
表示qAttention
和 aAttention
的分布为：
作者考察了同一层内
和
的分布模式，以及
在不同层之间的分布模式。采用由真实用户与 ChatGPT 对话产生的数据集 ShareGPT52K 来分析分布模式。选取 Qwen2.5-0.5B进行演示。
####
观察一：
和
的注意力分布相似。
例如，在图2(a)中，作者选取了一段包含85轮对话进行分析，以研究第85轮的注意力概率分布与之前各层的关系。如图所示，每个层中
和
的趋势高度相似，表明与第85轮问题高度相关的轮次也与答案高度相关。因此，在对第85轮的问题进行预填充后，作者可以根据轮次注意力分布识别出与AR计算最相关的历史轮次的KV缓存，而不是利用所有轮次的KV缓存。作者想强调的是，这种模式不仅适用于这个特定例子。作者在分析了大量对话后得出了这种模式，作者使用这个例子作为分析的主题。
####
观察2：各层之间的注意力分布相似
接下来，作者分析了不同层之间
的相关性。对于给定的一层，作者计算该层与每个后续层之间的Kullback-Leibler（KL）散度，并平均这些值以获得该层与所有后续层之间的平均KL散度。然后，作者绘制所有层的这些值，结果如图2(b)所示。可以观察到，对于几乎所有目前主流的开源模型，无论其规模大小，都呈现出相似的规律。
最初的几层与后续层相比存在显著差异；然而，达到某一层后，差异突然大幅减少。作者将这个层称为“分水岭层”
，并在附录B中列出几个开源模型的这个层。从这个层开始，后续层的
值非常接近。尽管偶尔会出现散度轻微增加的情况，但这些差异远小于早期层观察到的差异。这表明作者可以在分水岭层选择与问题最相关的轮次进行后续的注意力计算，从而避免在每个层都进行这种选择计算，这将产生额外的计算时间成本。基于这两个观察结果，作者提出了作者的推理流程——轮次注意力。作者将在接下来的章节中详细描述这个流程。
###
3.2 轮式注意力推理流程
当用户变得活跃时，例如向大语言模型（LLM）提出第
个问题
，将执行以下步骤以进行本轮推理。
* • 步骤1：将
从主机内存加载到GPU内存中。
* • 步骤2：使用层1至L\_w对
进行预填充计算
* • 步骤3：通过第3.3节中提出的策略，使用qAttention选择最相关的TopK对话轮次，并加载层
的KV缓存
* • 步骤4：完成剩余层的预填充。
* • 步骤5：解码
与之前在 Token Level 工作的研究不同，因此需要在主机和设备内存之间进行多个碎片化的KV缓存传输，Round-Attention在对话轮 Level 工作，将前几轮中所有 Token 的单一代数张量一次性传输到GPU。在第n轮中完成层
的
计算后，新的KV缓存也以单一代数张量的形式保存到主机内存中。因此，Round-Attention减少了昂贵的宿主到设备（H2D）和设备到宿主（D2H）数据传输次数。此外，以大块数据移动可以更好地利用PCIe带宽。算法总结见算法1。
###
3.3 轮询策略
在发现
并计算 qAttention
后，考虑了三种策略来确定最相关的
个对话回合。
策略1：固定轮数策略通过预定义的阈值选择满意的轮数，例如 qAttentior
。分析注意力分数在轮次间的分布，作者发现注意力值集中在有限的几个轮次中，大多数轮次表现出最低的注意力分数。基于这种数据分布，作者选择了
。
策略2：TopK轮数策略选择与最高
qAttention
对应的轮次。分析轮次注意力分数的分布揭示，最高
的轮次占据了超过
的累积注意力。
策略3：自适应轮次根据qAttention分布自适应地选择轮次。条件定义为：qAttentiontu
其中mean和std分别是qAttentionk·的均值和标准差。
###
3.4 KV缓存丢弃
作者观察到ShareGPT中某些对话回合的KV缓存从未激活，即使移除它们进行注意力计算也不会影响推理质量。对于这些回合，作者在相应的 Token 中删除KV缓存，以避免将它们保存在内存中。此外，当这些KV缓存激活时，从主机内存中加载它们是不必要的，从而节省了存储和内存复制成本。
###
3.5 内存与性能分析
本节探讨了圆周注意力机制的内存和性能。
内存占用分析。给定一个具有上下文长度
、隐藏层大小
、总层数
和推理批大小
的语言模型，KV 缓存的内存占用通过公式 6 计算。
原
始
公
式
翻
译
为
简
体
中
文
为
：
原
其中前两个2代表
和
，第二个2表示float16占用2字节。
对于轮次注意力机制，假设从总对话轮次
中选择了
个最相关的轮次，每个轮次平均使用的内存量由公式 7 表示。
这是因为在第1层到第
层中使用了整个KV缓存，而后续的层
仅计算与最相关
轮的注意力。
圆周注意力机制的内存节省率可以表示为公式8。
由于在实际应用中
远小于
，例如
与数十到数百相比，方程8的上界近似为
。当
等于
时，即所有对话轮都被选中，Round Attention退化到原始推理，几乎不产生内存成本。如表B所示，对于目前主流的大模型，
的比值在0.18到0.46之间，表明内存节省的百分比在54%到82%之间，这是一个相当大的节省。
延迟分析。LLM推理的延迟由预填充和解码组成。预填充是计算密集型，总FLOPs可以近似为
，其中
是模型参数的数量，
是计算出的token数量。可以通过TTFT（第一个token的时间）来衡量，这通常由方程（9）Korthikanti等人[2022]计算得出。由于自回归解码是内存密集型，
token的生成时间接近从GPU内存中加载整个KV缓存（见方程6）所需的时间。假设GPU内存带宽为
，原始推理 Pipeline 的TPOT（每个输出token的时间）由方程10计算得出。
因此，具有
个输出 Token 的原始LLM的延迟计算为
。回顾一下，Round Attention首先从主机内存（带宽
）中加载层
中最相关的前
个轮次的 Token ，然后使用这些 Token 进行注意力操作。延迟由四个主要部分贡献：前
层的预填充、使用TopK相关轮次对剩余层进行预填充、将TopK相关KV缓存传输到设备内存以及剩余层的AR操作。由于数据传输可以是异步的，并且可以被预填充计算所隐藏，作者只考虑其他三个部分，如公式11所示。
此处，
和
分别表示 GPU 的 FP16 TFLOPs 和内存带宽。
延迟节省率为
，为简便起见，在公式12中用
表示。
以下结论可以安全得出。
* • 上下文长度越长（
越小），可以节省的延迟越多。例如，当
可以忽略不计时，Round Attention 的延迟大约等于原始推理 Pipeline 的
。
* • 水平集层
出现得越早（
越小），实现的延迟就越短。
4 实验与分析
--------
###
4.1 实验设置
本节评估了Round Attention的模型精度、内存占用和推理性能。
数据。本文使用了两个广泛使用的数据集，ShareGPTShareGPT52K [2024] 和 LONGMEMEVAL Wu et al. [2024]，以评估 Round Attention 的有效性。ShareGPT 包含了通过 ShareGPT API 抓取的大约 52K 个用户共享的对话。这些对话是多轮的，包括用户 Prompt 和 ChatGPT 的响应。LONGMEMEVAL 是一个综合性的基准，旨在评估商业聊天助手的五个核心长期记忆能力：信息提取、多会话推理、时间推理、知识更新和弃权。此基准还记录了平均 250 轮的历史用户-助手对话。这是一个难度较大的数据集，GPT-4o 在其上的准确率仅为 0.5773。
Baseline 。一套最新的开源LLM，例如Qwen2.5、LLaMA3和LLaMA3.2（Grattafiori等，2024年），在上述数据集上进行了测试，但Round-Attention可以应用于任何其他长上下文LLM。作者使用PyTorch和FlashAttention（Dao等，2022年）作为默认推理框架，称为Flash：所有测试均在配备PCIe的单个Nvidia A100 GPU（80GB内存）上执行。使用的CPU是Intel(R) Xeon(R) Gold 6346 CPU，运行频率为3.10GHz（1.16/3.60GHz），系统内存为1TB。
###
4.2 准确性评估
作者将ShareGPT根据对话轮数分为四类：微型（0-10轮）、小型（10-30轮）、中型（30-50轮）和大型（50-100轮）。作者将每个类别中的最后一个 Prompt 视为
，然后使用默认推理框架和轮注意力计算
。GPT-4o被用作评判员来评估生成结果的品质。每个
评估5次，取平均值作为该回答的最终得分。 Prompt 语来自AlignBenchLiu等[2023a]，样本可参见附录A。
如表1所示，TopK策略最为有效，平均得分超过了标准推理引擎。这表明仅使用TopK轮次的KV缓存对模型响应质量没有显著影响，因为注意力矩阵非常Sparse，尤其是在极大规模的对话轮次中。此外，与标准推理引擎相比，作者处理的大规模轮次中的 Token 数量减少了88%，这表明注意力计算大幅减少，GPU内存节省显著。
为了验证该方法的一般化能力，作者在不同规模的Qwen2.5模型以及一些Llama3和Llama3.2模型上测试了其准确性。所有实验均采用了TopK轮次策略。结果如表2所示。显然，Round-Attention在这些模型上生成的响应的整体得分与标准推理引擎产生的得分相当。然而，作者也观察到，对于Qwen-2.5系列模型，随着轮次的减少，Round Attention的准确性可能会降低。相反，当轮次超过50时，Round Attention始终优于Flash Attention。
由于ShareGPT的响应具有主观性，作者还利用客观数据集LONGMEMEVAL作为测试平台，进一步验证Round-Attention的有效性。原始LONGMEMEVAL基准评估了Llama3-8B产生的结果。为了保持一致性，作者也对该模型进行了测试。为了展示Round-Attention的一般化能力，作者还在Qwen2.5-7B上进行了相同的实验。
如表3所示，尽管LONGMEMEVAL具有挑战性，但Round Attention的表现非常出色。对于Llama 3-8B模型，Round Attention与Flash Attention相当，而对于Qwen 2.5-7B模型，Round Attention的准确率是Flash Attention的两倍。
有趣的是，在时间推理任务中，轮式注意力（Round Attention）始终优于闪式注意力（Flash Attention），这表明过多的信息可能导致推理任务中的干扰。识别关键轮次可以使得推理结果更加准确。
###
4.3 显卡内存减少与延迟减少
为了实证评估Round Attention与Flash Attention的延迟，作者从前面提到的四个类别中各选取了20个对话样本。每个样本运行了10次，计算了平均延迟并绘制在图3中。
观察发现，对于所有不同的圆形类别，Round Attention的延迟都低于Flash Attention。这种改进归因于作者的KV缓存存储和传输策略，该策略使h2d传输时间可控，以及作者的top-
选择仅计算一次而不是每一层都计算。作者在附录C中提供了延迟的详细分解。由于h2d传输和top-
选择，
预填充阶段的延迟在层
处出现轻微峰值；然而，这个峰值很小，并且只出现一次。相比之下，在
解码阶段，KV缓存的减少导致注意力计算时间减少。随着解码步骤数量的增加，这种减少累积并最终超过h2d传输和top-
选择的一次性开销，从而使得整体延迟优于Flash Attention。
###
4.4 轮次与 Token
在本节中，作者比较了两种粒度。对于这两种粒度，作者都采用了相同的top-
计算策略来检索KV缓存。具体来说，作者首先计算了每个粒度下所有 Token /轮次的平均注意力分数，并选择了那些注意力分数高于平均值的 Token /轮次。测试所使用的模型是Llama3.1-8B。结果如表4所示。很明显，在相同的top-
计算策略下，轮度粒度的召回准确率超过了 Token 度粒度。这在单会话助手任务中尤为明显，其中答案位于几个会话中。轮度粒度的召回有效地检索了最相关的会话，而 Token 度粒度的召回则分散在多个会话中，与轮度粒度相比，准确率显著降低。
5 结论与讨论
--------
在现实世界中，大语言模型（LLMs）提供服务的背景下，随着用户参与越来越长的对话交流，历史键值（KV）缓存会不断累积。作者提出，在这样扩展的对话轮次中进行推理时，采用基于轮次的策略是管理KV缓存和处理与历史信息交互的更有效方法。通过对轮次粒度下的注意力矩阵模式进行分析，作者发现当代大型模型存在一个分水岭层，在此层之后，基于轮次的注意力分布变得显著相似。这一观察使作者能够在分水岭层一次性计算后续层注意力计算的最相关轮次。
因此，作者可以显著减少GPU内存使用，同时有效限制选择所需的时间。通过基于轮次存储KV缓存，作者可以在单个主机到设备（h2d）操作中将所有必要的KV缓存数据传输到GPU内存中，从而最小化与h2d传输相关的开销。作者通过实验验证了该方法的有效性，结果表明，它能够显著降低推理延迟，同时推理精度与使用完整KV缓存时保持基本一致。
###
局限性
本节讨论了本文有意留待未来工作中进一步改进的一些局限性。
局限性1：卸载到内存会产生额外的内存开销。尽管内存的价格远低于GPU内存且容量更大，但它仍然会增加系统的额外开销。
局限性2：虽然圆周注意力机制可以减少GPU内存使用，但当对话轮数达到一定阈值时，仍然可能发生内存不足（OOM）问题。这表明仅靠圆周注意力机制无法从根本上解决与非常长对话相关的GPU内存问题。它需要与其他技术结合，以有效解决内存问题，例如，与第3.4节中提到的持续丢弃不常用键值缓存。此外，还可以结合使用各种其他键值缓存压缩和丢弃策略，以解决实际对话系统中的GPU内存问题。
参考
---
[1]. Round Attention: A Novel Round-Level Attention Mechanism to Accelerate LLM Inference
扫码加入👉
「集智书童」
交流群
（备注：
方向
+
学校/公司+
昵称
）
