本期论文：AttentionStore: Cost-effective Attention Reuse across Multi-turn Conversations in
Large Language Model Serving
在现代人工智能技术中，能够与人类进行多轮对话是大型语言模型（LLMs）的一项基本功能。然而，目前现有的LLM服务引擎在执行多轮对话时效率低下，原因在于它们需要反复计算历史对话中的键值（KV）缓存，这导致了高昂的服务成本。
在一次对话中，LLM引擎会将中间数据，即键值（KV）对，存储在GPU上的高带宽内存（HBM）中。然而，当对话结束且会话变得不活跃时，为了给其他活跃会话腾出空间，LLM引擎通常会丢弃与该会话相关的KV缓存。当同一个会话再次变得活跃（即用户发送下一个消息）时，LLM引擎需要重新计算整个KV缓存。这就导致了同样的KV缓存被重复计算，浪费了宝贵的GPU计算资源。随着对话轮次的增加，重复计算的开销也会线性增加。
根据对ShareGPT【
一个从ChatGPT收集的广泛使用的真实数据集
】的分析显示，高达98%的预填充成本来源于KV缓存的重复计算。这意味着在多轮对话中，绝大部分计算资源都被浪费在了重复生成已经计算过的数据上。
今天，我们要介绍一种突破性的新技术——AttentionStore。这种创新的注意力机制通过复用KV缓存（即注意力复用），大幅降低了重复计算的开销。AttentionStore引入了一个分层的KV缓存系统，利用成本效益高的存储介质来保存所有请求的KV缓存，从而显著提高了多轮对话的效率。
-
为了减少从慢速存储介质访问KV缓存的开销，
AttentionStore采用了分层预加载和异步保存方案
，将KV缓存的访问与GPU计算重叠进行。
-
同时，为确保需要访问的KV缓存始终处于最快速的层级，
AttentionStore使用了调度感知的获取和驱逐方案
，根据推理任务调度器提供的提示，智能地将KV缓存放置在不同的层级中。
-
此外，
AttentionStore通过解耦位置编码和有效截断KV缓存
，避免了因上下文窗口溢出而导致的KV缓存失效。这一创新使得保存的KV缓存始终保持有效。
大量的实验结果表明，AttentionStore显著降低了首个标记生成的时间（TTFT）多达88%，在多轮对话中提升了8.2倍的提示预填充吞吐量，并减少了高达56%的端到端推理成本。在长序列推理中，AttentionStore将TTFT减少了多达95%，并提升了22倍的提示预填充吞吐量。
在理解生成型大型语言模型（LLM）的推理机制时，特别是基于Transformer的生成模型，我们可以将其逻辑上分为两个截然不同的阶段：预填充阶段和解码阶段。
预填充阶段：预填充阶段的主要任务是处理输入提示（prompt）。当收到一个请求提示时，生成过程会将提示符的标记列表X = [x1, x2, ... xs]作为输入，然后计算下一个标记xs+1。在此过程中，生成了一系列的键值（KV），具体形成从1到s的KV缓存，这些缓存将用于解码阶段。
解码阶段：解码阶段负责生成输出标记，这个过程是自回归的，意味着每一步生成的标记都会影响后续的生成。在解码阶段，模型将第s+1个标记和预填充阶段生成的KV缓存[1 : s]作为输入，计算新的KV缓存s+1以及标记s+2。这个生成过程会不断迭代，直到生成的标记是
（结束标记）或达到允许的最大生成次数。由于每次迭代都严重依赖于前一次迭代的数据，解码阶段只能顺序执行。
这两个阶段在执行时间上表现出显著不同的特性。预填充阶段可以并行计算KV缓存，因此其持续时间与输入提示的标记数量密切相关。输入标记越多，预填充阶段的执行时间就越长，如图1b所示。
问题分析
多轮对话是现代大型语言模型（LLM）的基本功能。在这种对话中，一个多轮对话会话由一系列连续的对话组成，记作D = [d1, d2, ... dN]。在每个对话dj中，用户输入一个新的问题或命令qj，然后等待LLM的响应aj。为了保持对话的一致性和连贯性，LLM需要基于所有先前对话轮次的历史标记d[1 : N]以及当前轮次的输入标记qN+1生成响应aN+1。
根据对ShareGPT的分析（一个包含超过10万次对话的真实数据集），如图2a所示，我们发现78%的对话是多轮的，大多数对话在5轮以内就结束了。
此外如图2b所示，72%的对话包含超过4000个标记，大部分会话的token数在8,000以内。
这表明多轮对话在实际应用中非常普遍。
然而，当前LLM服务引擎在执行多轮对话时存在显著的效率问题，主要原因是多轮对话中KV缓存的重复计算。如图3a所示：在对话的第一个轮次，LLM服务引擎会生成q1和a1的KV缓存。然而，在完成第一个轮次后，为了回收HBM（高带宽内存）空间，LLM服务引擎会丢弃这些KV缓存。在第二轮次中，LLM服务引擎需要重新生成q1和a1的KV缓存。在第三轮次中，它需要重新生成q1、a1、q2和a2的KV缓存。随着会话的扩展，历史标记不断积累，重复计算的开销显著增加。
如图4a所示，随着对话轮次的增加，在新一轮对话中，历史标记的比例将超过98%。如图4b所示，在新的对话中，重复计算时间占据了预填充时间（即生成第一个标记的时间）的98%。
基于上述分析，我们观察到，如果能够在多轮对话中重用KV缓存，那么可以减少多达98%的预填充成本。具体来说，历史对话的KV缓存可以保存在GPU外的KV缓存系统中。当对话会话重新激活时，GPU从KV缓存系统中加载相关的KV缓存，并在新的对话轮次中重用它们。然而，要构建一个高效的KV缓存系统，还存在许多重大挑战。
1.高KV缓存访问开销
：在推理过程中，GPU的计算可能因等待从KV缓存系统加载KV缓存而被阻塞。与KV缓存的重复计算时间相比，这种阻塞时间不可忽略。例如，我们评估了使用4个NVIDIA A100 GPU运行LLaMA-65B模型的推理时间，发现预填充2K个令牌的提示大约需要360毫秒，而从主机内存加载2K令牌（5GB）的KV缓存到GPU大约需要192毫秒（16条PCIe Gen4通道的GPU系统具有大约26GB/s的有效数据传输带宽）。
2.KV缓存的高存储容量需求
：每个请求的KV缓存存储需要大量空间。例如，使用4个80GB HBM的A100 GPU运行LLaMA-65B时，预填充2K个令牌需要大约360毫秒，生成5GB的KV缓存，生成速度约为13.9GB/s。130GB的HBM空间分配给模型存储后，剩余的190GB HBM空间将在14秒内被KV缓存占满。如果将KV缓存溢出到主机内存（例如512GB空间），主机内存将在不到1分钟内被填满。使用磁盘可以扩展存储空间，但会导致访问性能下降。
3.KV缓存的合适分层放置
：磁盘提供比主机内存更大的容量（数十TB对比几百GB），因此大多数KV缓存保留在磁盘中。然而，磁盘的访问带宽不到5GB/s（PCIe 4.0 x4接口的NVMe SSD）。由于对话请求随机到达，它们的KV缓存更可能位于磁盘中，导致推理性能差。必须确保即将访问的KV缓存始终放在主机内存而不是磁盘中。
4.KV缓存的非预期失效
：随着对话轮数增加，历史令牌可能超过上下文窗口限制。LLM服务引擎通常执行令牌截断来减少输入提示，这对以前的LLM服务引擎没有影响，因为它们总是根据截断后的输入提示重新计算KV缓存。然而，截断会使保存在KV缓存系统中的KV缓存无效，因为截断后每个令牌的位置发生了变化，无法匹配旧的嵌入位置编码。上下文窗口溢出发生的概率很高。根据图2b，84% 的对话会话的上下文长度超过了2000个令牌（tokens），69% 的对话会话的上下文长度超过了4000个令牌，当使用4K上下文窗口的LLaMA-2系列时，69%的对话会话会发生上下文窗口溢出；当使用2K上下文窗口的OPT系列时，这一比例为84%。
AttentionStore设计与实现
为了应对多轮对话中KV缓存重复计算的低效问题，我们提出了一种新的注意力机制——AttentionStore。该机制通过保存和复用KV缓存，有效减少了多轮对话中的重复计算开销。在传统的注意力机制中，每次对话结束后，KV缓存都会被丢弃。然而，AttentionStore打破了这一惯例，它将KV缓存保存在一个KV缓存系统中。当对话会话不活跃时，相关的KV缓存会被保存下来。如果同一会话在未来重新激活，其KV缓存将从KV缓存系统中获取并重新用于推理。这样，AttentionStore只需对新一轮对话中的新输入标记进行部分预填充，而无需重新预填充所有历史标记。例如，在执行第三轮对话推理时，q1、a1、q2和a2的KV缓存被复用，只需预填充q3，如图3b所示。
系统架构
图5展示了AttentionStore的系统架构概览。它维持了一个分层的KV缓存系统，通过高效的KV缓存访问、放置和截断技术来应对多种挑战。
挑战1：KV缓存加载和保存的开销解决方案：为了减少从KV缓存系统加载KV缓存到HBM的开销，AttentionStore采用了分层预加载方案，将KV缓存加载与推理计算重叠进行。同样，为了减少从HBM保存KV缓存到主机内存的开销，AttentionStore采用了异步保存方案，将保存过程与推理计算重叠进行。
挑战2和挑战3：扩大缓存KV缓存的可用存储空间解决方案：为了扩大可用的存储空间，AttentionStore利用了多层次、成本效益高的存储介质，即主机内存和磁盘。为了减少访问慢速磁盘对推理性能的影响，我们提出了一个调度感知的预取方案，利用任务调度器的提示，将需要访问的KV缓存从磁盘预取到主机内存。同时，为了高效利用有限的主机内存空间，我们提出了一个调度感知的驱逐方案，识别最不重要的KV缓存，并将它们驱逐到磁盘或缓存系统之外。
挑战4：处理因上下文窗口溢出导致的KV缓存失效问题解决方案：为了应对上下文窗口溢出导致的KV缓存失效问题，我们使用了位置编码解耦截断方案，在保存KV缓存时不嵌入位置编码，从而支持直接在KV缓存上进行截断。当加载KV缓存时，AttentionStore重新嵌入新的位置编码。
通过这些设计，AttentionStore有效解决了多轮对话中的KV缓存管理问题，显著提高了多轮对话的效率和性能。
分层预加载机制
在AttentionStore中，KV缓存从主机内存加载到高带宽内存（HBM）过程中，数据访问开销较高。由于GPU需要依赖KV缓存来执行推理计算，这个访问过程是推理执行的关键路径，如图6a所示。随着KV缓存大小的增加，这种开销变得更加显著。为了消除这一开销，AttentionStore采用了分层预加载机制来缓解影响。
核心理念：分层预加载机制的核心理念是将KV缓存的加载与新输入标记的预填充计算重叠进行。LLM模型由多个Transformer层链接而成，每层都有自己的KV缓存。当GPU执行某一层时，可以同时从主机内存加载下一层所需的KV缓存。这样，当GPU开始计算某层的自注意力时，该层对应的KV缓存已经在HBM执行缓冲区中准备就绪。
实现细节：图6b展示了分层预加载机制如何将KV缓存的获取时间与计算时间重叠。以一个简单的三层模型为例：
1.在开始计算第一层之前，必须先在HBM中准备好该层的KV缓存。
2.读取流（read stream）首先发出KV缓存加载操作，将第一层的KV缓存读取到HBM执行缓冲区中。
3.执行流（execution stream）然后开始计算第一层。当执行流计算一层时，读取流同时加载下一层的KV缓存。
这种方式下，加载过程与计算过程是重叠进行的。然而，我们注意到在最后一个任务和当前任务的第一层之间仍然存在一个间隙，因为加载只能在HBM执行缓冲区可用（即上一个任务完成）后开始。为了进一步减少这个间隙，AttentionStore保留了一个HBM读取缓冲区（read buffer），如图6c所示。有了读取缓冲区，读取流不必等待上一个任务释放执行缓冲区，可以在上一个任务运行时开始预加载。
如果KV缓存加载时间比预填充计算时间更长，预加载可能无法完全与计算重叠，如图7a所示。多个层计算之间存在间隙，因为每层的KV缓存获取时间超过了层计算时间，导致重叠不完全。为进一步减少开销，可以采用自定义更大的预加载缓冲区。更大的缓冲区允许更早开始预加载，如图7b所示，可以预加载更多层的KV缓存，从而使层之间的间隙得到重叠。
异步保存机制
为了在会话之间复用KV缓存，AttentionStore需要将KV缓存保存到主机内存中。一种基线方法是，在对话结束后将所有生成的KV缓存一起写入内存。然而，这种方法可能会延迟下一个计划任务的执行，因为KV缓存的保存时间是推理的关键路径，如图8a所示。为了减少这一开销，AttentionStore采用了异步KV缓存保存机制，将KV缓存写回与计算重叠进行，同时考虑预填充和解码阶段的不同特性，实施不同的重叠机制。
预填充和解码阶段生成KV缓存的速度不同。预填充阶段并行处理标记，因此在有限的时间内生成大量的KV缓存。相比之下，解码阶段一次只生成一个标记的KV缓存。正如图8b所示，对于预填充阶段，由于每个自注意力操作可以生成大量的KV缓存，写入流（write stream）逐层保存KV缓存。预填充阶段生成的KV缓存可以与解码阶段重叠进行。
对于解码阶段，由于KV缓存是逐步生成的，写入流在解码时逐层写回KV缓存。为了避免在解码完成时KV缓存未完全写回而导致卡顿，我们保留了一个HBM写入缓冲区，类似于在KV缓存预取中使用的读取缓冲区。未完成的KV缓存暂时移动到写入缓冲区，以避免阻塞下一个任务的执行。
调度器感知预取机制
由于磁盘的容量远大于主机内存（数十TB对比数百GB），AttentionStore中大多数KV缓存保留在磁盘上。由于对话请求随机到达，其对应的KV缓存更有可能位于磁盘中，导致访问性能较差。
为了解决这个问题，我们利用了一种调度感知的KV缓存预取方案，将要访问的KV缓存从磁盘预取到主机内存。这是通过利用推理作业调度器提供的提示来完成的。具体来说，作业调度器维护着一个作业队列，因此可以完全了解等待执行的作业。AttentionStore应用了一个前瞻预取窗口来监视即将执行的等待作业。如果等待作业的KV缓存位于磁盘中，AttentionStore将在这些等待作业执行之前，将其KV缓存从磁盘预取到主机内存中。前瞻预取窗口的长度由主机内存中可用的容量决定。给定预取的可用内存容量Cmem和会话的平均KV大小Skv，则预取窗口的长度为Lpw = Cmem/Skv。
图9展示了一个调度感知的预取示例。当作业1正在执行时，KV缓存管理器应用一个前瞻窗口大小为2（主机内存有2个KV缓存槽用于KV缓存获取）来检查等待的作业2-3的KV缓存命中状态。作业2的KV缓存命中了主机内存，但作业3的KV缓存不在主机内存中。然后，KV缓存预取线程开始将作业3的KV缓存从磁盘预取到主机内存中。
值得注意的是，AttentionStore包括一个主机内存缓冲区，允许从磁盘无缝地预取KV缓存到内存中，防止当主机内存已满时出现任何延迟。当空闲内存的容量达到定义的阈值时，AttentionStore会触发将主机内存中的KV缓存驱逐到磁盘，以确保主机内存缓冲区的持续可用性。从主机内存到磁盘的驱逐过程将在下一小节中介绍。
调度器感知驱逐机制
