Abstract
基于Transformer的长上下文生成模型为新兴的人工智能应用提供了动力，如小时 Level 的视频理解和项目 Level 的编码代理。与短上下文（例如4K Token ）模型版本相比，部署长上下文Transformer（例如，10万到100万个 Token ）的成本高得令人望而却步。从2024年开始，降低长上下文Transformer的成本成为迫切的研究和工程挑战。本文描述了一个并发编程框架，用于定量分析在有限的GPU高带宽内存（HBM）环境下服务多个长上下文请求的效率挑战。作者详细分析了与4K上下文相比，
所有额外计算成本的根源：单一来源，即KV缓存的大尺寸。
作者以一个34B GPT-3.5 Level 的50K上下文模型在A100 NVLink上的运行为例，描述了其大型KV缓存如何导致四种部署挑战：
(1)预填充长输入比短输入需要更长的计算时间和GPU内存；
(2)预填充后，驻留在GPU HBM上的大型KV缓存大大限制了可同时服务的并发用户数量；
(3)在解码过程中，从HBM到SM反复读取KV缓存大大增加了延迟；
(4)当KV缓存内存溢出时，从HBM交换到DDR会导致显著的上下文切换延迟。作者使用这个框架分析现有工作，并识别结合它们构建端到端系统的可能性。
总体而言，这项工作为分析长上下文Transformer部署提供了一个基础框架，并指出了降低100万个上下文推理成本的方向，使其与4K一样便宜。
1 Introduction
假设某人已成功将开源的GPT-3.5级语言模型提升到GPT-4 Level （例如LLaMA 3 [25]，DeepSeek [7]，QWen [3]和Mixtral [14]），并将其转化为长上下文变体（例如，通过持续预训练[9]和指令调优[4]），然后希望将这个模型部署在广泛的应用中，如库 Level 的编程和小时级视频理解。这些应用通常需要输入上下文为10万到1000万个标记，这与4千个短上下文的情况有本质的不同。
作者对一个雄心勃勃的目标感兴趣：
如何将部署生产 Level 的100万个上下文 Transformer 的成本降低到与4千个一样低？
为什么服务长上下文 Transformer 成本高昂？大多数成本开销都可以追溯到 _一个单一源头：KV 缓存的大小_。考虑一个拥有 30+B 参数，10 万个上下文的 GPT-3.5 质量开源模型，如 QWen 或 Yi，4K 上下文与 100K 上下文的 KV 缓存之间的差异在于：
在这里，作者使用了Yi-34B 200K [29] 的配置（60层，8个kv头和128的隐藏维度）。假设作者使用2 80G A100张量并行在bf16格式下服务这个模型，那么作者有 GB的空闲空间来存储KV缓存。从这个初步估计中，作者立即意识到在这种设置下，作者可以实现大约100+用户的4K上下文并发，但只有5个用户的100K上下文。如果作者能够像部署4K那样便宜地部署1M上下文的模型，作者实际上可以大大普及长上下文和多模态生成模型，促进新兴应用生态的发展。
这项工作从硬件到系统，从模型架构到用户偏好，全面定量分析了部署长上下文 Transformer 面临的挑战。作者首先描述了一个并发编程框架（图1），用于在有限的GPU HBM大小下服务多个长上下文用户请求。作者将会话吞吐量定义为给定时间内用户交互的轮数，作为端到端的评估，并将其分解为四个关键指标：并发水平、预填充延迟、解码延迟和上下文切换开销。作者讨论了这四个指标如何分别受到HBM大小、GPU浮点运算能力、HBM带宽和PCIE带宽的限制，以及这些挑战最终如何追溯到KV缓存的大小，导致了关于KV缓存无损压缩的核心研究问题。作者进一步讨论了硬件架构如何改变性能因素，现有工作如何从层、头、标记和隐藏状态维度压缩KV缓存，以及预填充与解码之间的相对成本如何变化。作者希望这项工作成为将1M上下文推理成本降低到与4K一样便宜的基础框架。
2 Challenges in Deploying Long-Context Transformers
图1：在有限GPU HBM大小下服务多个长上下文用户请求的并发编程框架。共有四个关键因素共同决定了用户交互会话的整体吞吐量：
（1）并发受限于HBM大小：可同时服务的并发用户数量受限于GPU高带宽内存（HBM）的大小；
（2）预填充受计算限制：生成第一个 Token 的延迟，即提示预填充时间，受限于GPU的每秒浮点运算（flops）；
（3）解码受内存限制（在关键批量大小下）：自回归解码的延迟（每秒生成 Token ）受限于HBM的带宽；
（4）上下文切换受PCIE限制：将用户1的KV缓存卸载到CPU DDR，并将用户2的KV缓存加载到HBM受限于PCIE带宽。与短上下文设置相比，长上下文设置下的并发和上下文切换挑战要严重得多。所有来自这四个关键因素的效率挑战最终都可以追溯到KV缓存的大小。
鉴于一个生产 Level 的长上下文 Transformer ，作者的目标是降低具有100万 Token 上下文的部署成本，使其与4K一样便宜。作者以50K上下文和80G A100 NVLink上的30+B模型以及分组 Query 注意力（GQA）作为运行示例，因为它们通常具有GPT-3.5的能力，具体来说，作者使用Yi-34B 200K配置，因为在本论文撰写之时，它是唯一支持长上下文的开源30+B模型。作者首先描述了一个关于服务多个长上下文用户 Query 的工作流程的并发编程框架（图1），并确定了四个关键性能指标：并发、预填充、解码和上下文切换。
作者专注于理论峰值性能分析。也就是说，作者假设作者已经有一个高效的实现，可以挤出所有硬件性能（这通常需要复杂工程努力，如cuda Kernel 编程、内存管理、模型并行等），并研究在一个足够高效的实现下，作者应该解决哪些关键挑战和限制。尽管作者没有实际的实现（其工程量相当大，例如参见vLLM [16]和InfiniteLLM [19]，因此远远超出作者当前资源的范围），但理论峰值性能（在分析LLM服务系统中广泛使用）是识别系统瓶颈并为进一步提高效率提供指导的重要工具（例如，现有的大多数服务系统都关注解码，但当上下文超过200K时，作者可能希望更多地关注预填充，正如作者在图3后面讨论的那样）。
A Concurrent Programming Framework under Limited GPU HBM Size
并发用户交互会话及偏好 如表1和图1所示，在典型的交互会话中，用户从一个长文档的提示开始，然后是一个问题，并将其输入模型。模型接收到初始提示，_预填充_它以成为KV缓存。用户等待预填充阶段直到第一个 Token 开始生成，并希望等待时间不要太长。预填充后，模型开始自动回归式地执行解码。用户在解码过程中同时阅读输出，并希望解码速度比人类阅读速度快。
模型完成解码后，用户继续阅读回应，思考，也许喝杯咖啡，然后开始输入下一个问题。后续提示通常不如第一个提示长，因为第一个提示通常包含长上下文（书籍或视频），而后续提示通常是仅问题。当第一个用户阅读模型回应并思考接下来要问什么时，模型实际上是空闲的，因此如果此时另一个用户提出了另一个问题，模型可以通过将第一个用户的KV缓存卸载到CPU DDR内存中来进行_上下文切换_，以便为存储第二个用户的KV缓存腾出HBM空间。两个用户交互式地提出后续问题，直到他们的会话结束。
基于会话的吞吐量作为端到端目标，作者考虑了多用户同时与模型交互的情况。假设平均而言，一个用户会话包括一个50K Token 的文档/视频和5轮问题。在收到上一个问题的答案后，用户大约花1分钟阅读答案并思考下一个问题。作者的目标是最大化以下定义的基于会话的吞吐量：
请注意，作者的_基于会话_的吞吐量目标与现有的_基于 Token _的吞吐量[16, 18, 23]（即在给定时间内预填充或解码的 Token 数量）不同。作者很快就会看到，基于 Token 的吞吐量只是问题的一部分。作者的基于会话的吞吐量，即在给定时间内并发用户交互的数量，是一个_端到端_目标，因为作者不仅考虑预填充和解码，还考虑内存限制和上下文切换，这些因素显著影响并发性。计算与内存限制性，算术强度和关键批处理大小 对 Transformer 推理的一个重要观察是，预填充通常受限于GPU的计算能力，即浮点运算（flops），而解码则受限于HBM内存带宽。如果一个操作符大部分完成时间是在GPU的流式多处理器（SMs，GPU在其中执行块并行计算）上计算，作者说这个操作符是计算限制的。如果一个操作符大部分完成时间是将数据从内存移到SMs（而不是实际在SMs上计算），作者说这个操作符是内存限制的。一个操作符是计算限制还是内存限制取决于其_算术强度_，定义为每个内存访问操作（IO）执行多少浮点操作（FLOP）。
算术强度
并行 Level 越高，每内存访问的浮点运算越多，操作符受计算限制的可能性越大，作者就能更好地利用硬件。在给定的GPU上，操作符从内存限制变为计算限制的关键算术强度，即并行 Level ，是其flop与内存带宽的比率。对于A100来说，它是：
对于 Transformer （transformers），其并行 Level 大约是作者输入的 Token 数，即批量大小。这意味着，在A100上，当作者的批量大小大于156个 Token 时，例如在预填充期间，用户的提示为50K个 Token ，作者受计算限制并充分利用A100的计算能力。当作者的批量大小小于156个 Token 时，例如在解码过程中，每次前向传播只解码一个 Token ，作者受内存限制并且没有充分利用A100的计算能力。
预填充 现在作者分析在A100上预填充的确切时间。由于预填充是受计算限制的，即在A100上上下文长度超过156，其理论峰值延迟为
理 论 峰 值 延 迟
预 填 充 的
的 每 秒
如 果 受 计 算 限 制
由于9.8秒是理论峰值，在图1中作者将其四舍五入到12秒，以考虑实现开销。 如果序列长度为4K，那么其相应的KV缓存仅为0.91GB，解码延迟减少到8.5秒。 然而，如果序列长度增加到200K，KV缓存变为44GB，延迟增加到14秒。 相对延迟增加与KV缓存与模型大小之间的相对大小有关，作者最终希望关闭它。
并发控制与上下文切换 另一个重要的考虑是，当KV缓存变得很大时，GPU HBM可以持有的并发用户缓存的数量是
这意味着并发受限于HBM的大小。以作者34B 50K模型为例，如果作者将其部署在一个80GB A100上，作者只能服务一个用户（图1）。但如果上下文是4K，KV缓存大约只有1GB，作者可以同时服务大约20个用户。
当第二个用户来询问关于一个长文档的问题时，为了腾出他们的KV缓存的空间，作者需要进行上下文切换：将第一个用户的KV缓存卸载到CPU内存，并加载第二个用户的KV缓存（图1）。这引起了上下文切换的开销：
这就是说，上下文切换的开销受限于PCIE带宽，即GPU HBM与CPU DDR的连接速度有多快。假设作者使用每秒20G字节的PCIE gen 4，那么两个50K上下文用户的上下文切换开销为：
在图1中，作者将1.1秒四舍五入到2秒，以考虑工程开销。如前所述，在作者的设置中，作者可以服务20个4K上下文长度的用户而无需上下文切换，因为HBM足以容纳他们的KV缓存。如果作者想将50K并发增加到20，那么整体的上下文切换开销也会随着并发性增加：
请注意，在4K上下文环境中不存在这22秒的开销。
到目前为止，作者已经讨论了在使用34B模型和50K上下文作为运行示例部署长上下文 Transformer 时的大部分细节。作者看到，由用户交互吞吐量衡量的整体性能分解为四个指标：
（1）由GPU浮点运算限制的预填充延迟；
（2）由HBM带宽限制的解码延迟；
（3）由HBM大小限制的并发程度；
（4）由GPU-CPU连接带宽，即PCIe限制的上下文切换开销。
在接下来的章节中，作者将讨论这些指标如何随着常见因素而变化，并最终确定瓶颈追溯到KV缓存的大小。
Factors that Strongly Influence the Performance Metrics
作者从两个基本因素开始：上下文长度和硬件架构。
当上下文长度从4K增加到50K时，作者展示了四个指标（预填充、解码、并发性和上下文切换）以不同的速率（线性、倒数和二次）变化。
作者进一步显示，张量并行性提高了并发性、预填充和解码，但并未改善上下文切换。将稀疏循环利用到更大的专家混合模型中减少了并发性、预填充和解码，但并未改变上下文切换。注意力的类型，即多头、多 Query 或组 Query 注意力，对性能有显著影响，因为它们直接影响KV缓存的大小。
上下文长度 在图2的第一行中，作者根据第2.1节的公式计算了对于作者的Yi 34B运行示例，四种指标相对于上下文长度的理论峰值性能。作者观察到：(1) 随着上下文长度变长，并发性倒数减少；(2) 预填充延迟随着上下文长度变长二次增加。相比之下，解码延迟和上下文切换开销仅_随着上下文长度线性增加，并且解码延迟是最不受影响的因素，因为50K上下文KV缓存仍然相对于模型参数来说较小（11GB对比68GB）。并发性和预填充是受影响最大的两个指标。
作者能否简单地通过使用更先进的硬件来提高性能？在图2的第二行中，作者展示了硬件进步时的性能改进趋势。
作者观察到：
(1) 并发性随着HBM的大小线性增加；
(2) 当设备从4090升级到A100再到H100时，预填充延迟_随着浮点运算的增加倒数减少；
(3) 解码延迟随着内存带宽的增加倒数减少；
(4) 上下文切换开销随着PCIE带宽的增加倒数减少。
请注意，这些数字是基于2024年5月最新的进展，即使作者使用最新的硬件，50K和4K之间的成本差距并未缩小。
这就是说，作者不能依赖硬件进步来降低服务长上下文模型的成，作者必须进行算法创新。
张量并行性利用多个设备以几乎为零的通信开销加速推理。将设备数量线性增加到2、4和8，引入更多的HBM空间，从而线性地提高并发性。由于作者将计算均匀分配到多个设备上，预填充和解码延迟也随着GPU数量的增加倒数减少。然而，张量并行性不能减少上下文切换开销，因为DDR到HBM之间的PCIE带宽是由所有设备共享的。
将模型升级为专家混合体假设作者将34B模型升级为8倍规模的34B专家混合体模型，各项指标会如何变化？
首先观察到的是模型大小：由于MoE模型比密集模型大得多，它们将占用更多的HBM空间，从而降低并发性。预填充和解码延迟随着激活的专家数量而变化，通常为2，因此这两种延迟将大约增加2倍。
由于MoE不改变注意力机制，KV缓存的大小不变，意味着上下文切换的开销也不变。总之，升级到MoE会改变并发性、预填充和解码延迟，但不会改变上下文切换。
注意力的类型最后一个也可能是目前最重要的因素是注意力的类型。具体来说，对于作者考虑的34B的Yi模型，它使用带有8个kv头但32个 Query 头的组 Query 注意力（GQA）。如果作者将其kv头增加到32（即原始的多头注意力，MHA），其相应的50K Token KV缓存将是：
图2：第一行：上下文长度如何改变四个关键性能指标。将上下文长度从4K增加到50K反比地降低并发性，平方级地增加预填充延迟，线性地增加上下文切换开销，以及略微（但也是线性地）增加解码延迟。第二行：不同硬件世代如何影响这些指标。
实际上，人们并不经常使用多 Query 注意力（MQA），因为其性能不佳。但与原始的MHA相比，GQA直接实现了4倍的KV缓存减少，转化为并发性和上下文切换的4倍提升。
Most Challenges Trace Back to the Size of the KV Cache
作者在第2.1节中将50K与4K进行比较时，有以下重要观察：（1）为了预填充长输入并生成KV缓存，预填充延迟从0.89秒增加到14.1秒（方程10）；（2）由于大KV缓存位于GPU内存上，并发性从大约20降低到1；（3）在解码过程中，重复加载KV缓存将延迟从8.5秒增加到9.8秒（方程13）；（4）大KV缓存导致昂贵的上下文切换开销，对于20个用户大约需要额外22秒（方程17）。这四个因素共同导致端到端会话吞吐量的显著成本。作者最终的目标是使1M上下文的部署成本与4K一样低，而4K Token 大约是1GB的KV缓存。然后，作者的观察指向了一个关键研究问题：
如何无损地将1M Token 的KV缓存压缩到1G字节？
3 Compressibility Analysis and Directions of Optimization
在本节中，作者讨论压缩KV缓存可能的维度。首先指出，在不进行任何压缩的情况下，将1M个标记存储为字节大约只需要3-10MB的磁盘存储空间（这取决于分词器词汇表的大小），因此1GB的存储空间足以保存输入标记的完整信息。问题是，如何让大型 Transformer 使用这些压缩表示。作者从层、头、标记和隐藏维度的可压缩性进行分析，然后讨论预填充和解码之间的相对成本。
Lossless Compressibility by Layer, Head, Token and Hidden
作者首先讨论“无损”概念。实践者通常在各种长上下文任务上测试模型，以检查压缩是否会有损，其中“针尖寻麦秆”测试作为一个门槛：
如果模型无法通过此测试，作者不认为它能完成更难的任务。不幸的是，看起来两个重要的模型家族，状态空间模型（例如，Mamba [11]）和线性注意力（例如，LongT5 [12]），无法通过针尖测试，因此作者不将它们纳入讨论。吴等人[27]的最新研究显示，存在一组特殊的注意力头，它们负责从上下文中检索重要信息。他们的发现表明，至少对于某些层和某些头，大多数输入标记的完整注意力应该被保留。下面作者从层、头、标记和隐藏四个维度讨论KV缓存的压缩性。表2列出了一些现有的优秀成果以及它们改进的指标。
层： 这里的基本假设是，某些任务可能不需要完整的深度计算[8]。在预填充期间跳过一些层可能对所有四个指标都有好处，因为它同时减少了预填充的浮点运算次数和KV缓存的大小。
实际上，从现有成果如吴等人[28]，孙等人[24]的结果来看，层维度可能被大幅度减少，对于长上下文任务，可能只需要保留一层的KV缓存，这是一个1/60的压缩比。
头： 这里的基本假设是，一些头专门用于检索和与长上下文相关的功能[27]，因此可能可以保留检索头并剪除其他头。
请注意，头剪枝通常发生在预填充之后，这意味着它们只改进了解码、并发性和上下文切换，而预填充仍然代价高昂。总的来说，看起来头维度具有非常高的稀疏性水平，头的数量可能被大幅度减少，例如，吴等人[28]显示最强检索头的数量少于20个，这可能转化为20/1024的压缩比。
标记： 对于标记维度，基本假设是如果一个标记的信息可以从其上下文中推理出来，作者可以通过删除它[31]或将其与邻居合并[21]来压缩这个标记。
在标记维度上的大部分压缩对预填充的改进不大，但它们通常改进了并发性、解码和上下文切换。目前，看起来标记维度可能不如层和头维度那么稀疏，因为大多数标记必须被保留以进行精确检索。作者还没有看到任何显示标记维度有超过50%压缩潜力的工作。
隐藏： 除了量化[20; 30]之外，关于减少隐藏维度的研究并不多，可能因为隐藏大小已经是128，太小了，无法进一步减少。然而，可能仍然值得尝试在KV缓存上应用维度降低，比如LoRA[13]，特别是考虑到DeepSeek V2[7]最近引入了类似LoRA的想法，有效地减少了KV头。
这里一个重要的警告是，许多现有工作可能只强调了问题的某一个方面。例如，TriForce[23]只考虑了使用推测解码的解码延迟。
它没有使KV缓存变小，甚至由于草稿模型中额外的KV缓存，还增加了GPU内存消耗的权衡。许多现有工作也是正交的，它们的优点可以从不同的方面合并力量。例如，如果作者能将KV缓存减少到只有一层或10个头，并且只保留50%的标记，作者大约会有1000倍的性能提升。这自然引出了一个研究呼吁：
作者能否将现有成果整合到一个端到端系统中，并进行全栈优化？
Relative Cost Between Prefilling and Decoding
为了端到端地优化服务系统，一个重要的考虑因素是预填充和解码之间的相对成本。图3显示了在两种模型大小下（Yi 34B和Command R+分别是开源的GPT3.5和GPT-4 Level ）的相对延迟。
作者看到，随着模型变得更大，上下文长度变得更长，成本逐渐从解码转移到预填充。
对于需要多轮用户交互的任务（例如，编码代理和AI伴侣），它们的解码成本通常高于较少轮对话的任务（例如，文档问答）。考虑到现有的推理优化，如投机解码[17; 5]和vLLM[16]，主要针对短上下文环境，它们对长上下文部署的改进可能并不足够显著。下面作者将讨论进一步的优化方向。
优化预填充 ：由于预填充受计算限制，实际上的优化空间并不太大，主要归结为减少浮点运算(flops)。一个直接的浮点运算减少方法是局部或线性注意力（例如，Mistral[14]使用的滑动窗口注意力），它去掉了方程10中的二次注意力项。然而，图3显示，对于小于50K的上下文长度，线性注意力的增益相当有限。最终，为了减少预填充延迟，可能仍然需要回归到更小、更浅的模型，如同YOCO[24]的做法。
优化解码 ：与预填充相比，解码的优化空间要大得多，这不仅因为如上所述KV缓存内部存在大量稀疏性，还因为现有技术如投机解码及其长上下文变体如TriForce[23]。因此，一般来说，优化解码可能相对容易，但重要的是要记住，对于长上下文的大型模型，预填充阶段占据了大部分时间（见图3中的Command R+ 200K案例）。
4 Conclusions
在这项工作中，作者详细分析了部署长上下文 Transformer 的挑战。作者的最终目标是降低1M上下文的服务成本，使其与4K一样低廉，从而普及像视频理解和生成代理这样的新兴AI应用。
作者描述了一个并发编程框架，以说明基于端到端用户交互会话的吞吐量，并将其分解为四个关键性能指标：并发性、预填充、解码和上下文切换。
作者讨论了常见因素如何影响这四个指标，以及现有研究如何关注不同的指标。
作者认为，将现有努力整合起来构建一个强大的端到端长上下文服务系统有着巨大的研究机遇，并相信这项工作可以为全栈长上下文推理优化提供基础。
参考
[1].Challenges in Deploying Long-Context Transformers:.