论文标题: 借助缓存注意力机制实现经济高效的多轮对话大规模语言模型服务
作者/机构: Bin Gao (新加坡国立大学), Zhuomin He (上海交通大学), Puru Sharma (新加坡国立大学), Qingxuan Kang (新加坡国立大学), Djordje Jevdjic (新加坡国立大学), Junbo Deng (华为云), Xingkun Yang (华为云), Zhou Yu (华为云), and Pengfei Zuo (华为云)

A1 主要贡献

本文旨在解决现有大型语言模型（LLM）服务引擎在执行多轮对话时效率低下的问题。核心问题在于，这些引擎需要重复计算历史token的键值（KV）缓存，从而导致高昂的服务成本。为了解决这一问题，本文提出了CachedAttention，这是一种新的注意力机制，它能够在多轮对话中复用KV缓存，从而显著减少重复计算的开销。

本文的主要贡献如下：
* 调查了LLM在多轮对话中KV缓存的重计算开销，并识别了在多轮对话中保留KV缓存所面临的挑战。
* 提出了CachedAttention，一种新的注意力机制，允许在同一会话的任何后续对话轮次中重用KV缓存，从而显著减少了LLM中KV缓存的重计算开销。
* 设计了多种优化方案以提高CachedAttention的效率，包括重叠KV缓存访问、分层KV缓存放置以及位置编码解耦的KV缓存截断方案。
* 使用真实数据集对CachedAttention进行了全面评估，以证明其有效性和效率。

A3 背景知识与动机

2.1 生成式LLM推理基础

Transformer架构与KV缓存。Transformer已成为生成式LLM推理的公认标准，像GPT和LLaMA这类广泛使用的LLM都基于自回归Transformer架构【索引16，Transformer in transformer，2021，NeuIPS；索引47，Attention is all you need，2017，NeuIPS】。在推理过程中，模型处理用户提示并生成响应。模型由一系列Transformer层组成，每层包含自注意力和前馈网络（FFN）两个阶段。对于输入token列表 $X = [x_1, x_2, ...x_s]$，每层使用权重 $W_Q, W_K, W_V$ 对每个token进行投影，生成查询（Q）、键（K）和值（V）。

图片描述
图片描述

随后，通过Q、K、V计算注意力分数：

图片描述
图片描述

其中 $d_K$ 是键向量K的维度。最后，投影操作对注意力分数进行线性变换，其结果被传递给FFN层，并作为输入进入下一个Transformer层。这个过程会为每个token生成中间的K和V张量。在生成后续token时，所有先前token的KV张量对于计算自注意力是必需的，这些张量通常被缓存在GPU中，称为KV缓存。KV缓存通常占用大量空间，例如GPT-3【索引10，Gpt3. int8 (): 8-bit matrix multiplication for transformers at scale，2022，NeuIPS；索引29，https://openai.com/blog/chatgpt，2024】为每个token生成4.5MB的KV缓存，其大小随提示token数量线性增加。


图1: 预填充（Prefilling）和解码（decoding）阶段。延迟是在4个A100 GPU上对LLaMA-70B模型（批量大小为8）测量的。

2.2 自回归生成

两个不同的生成阶段。如图1a所示，基于Transformer的生成过程可逻辑上分为两个阶段【索引1，Sarathi: Efficient llm inference by piggybacking decodes with chunked prefills，2023，arXiv】。预填充（prefilling）阶段接收提示token列表 $X = [x_1, x_2, ...x_s]$ 作为输入，计算下一个token $x_{s+1}$，并生成从1到s的KV缓存，供解码阶段使用。解码（decoding）阶段则迭代地生成输出token，它利用前一阶段的token $s+1$ 和KV缓存[1:s]来计算KV缓存 $s+1$ 和token $s+2$，直至生成<eos>标记或达到最大生成长度。这两个阶段表现出不同的执行时间特性：预填充阶段并行计算KV缓存，其持续时间与输入提示token的数量密切相关；而解码阶段每次迭代只计算一个token，因此每次迭代的计算时间相对恒定，如图1b所示。

2.3 多轮对话推理

多轮对话的普遍性与低效性。与人类进行多轮对话是现代LLM的一项基本功能。一个多轮对话会话由一系列连续对话 $D = [d_1, d_2, ...d_N]$ 组成。在每次对话 $d_j$ 中，用户输入一个新问题 $q_j$ 并等待LLM的回答 $a_j$。为了保持上下文连贯，LLM会基于所有历史对话轮次 $d_{[1:N]}$ 和当前轮次的输入 $q_{N+1}$ 来生成回答 $a_{N+1}$。根据对真实数据集ShareGPT【索引33，https://huggingface.co /datasets/philschmid/sharegpt-raw/tree/main /sharegpt_90k_raw_dataset；索引38，https://sharegpt.com/】的分析，73%的对话是多轮的（图2a），并且30%的对话token数超过4K（图2b）。然而，当前LLM服务引擎在处理多轮对话时效率低下，因为它们需要重复计算KV缓存。如图3a所示，在第一轮对话后，引擎会丢弃KV缓存以回收HBM空间；在第二轮对话时，会重新生成第一轮的KV缓存；第三轮时，又会重新生成前两轮的KV缓存。随着会话的扩展，历史token不断累积，重复计算量显著增加。如图4a所示，随着对话轮次增加，历史token的比例会超过99%。这种重复计算时间占据了新对话中预填充时间（即首个token生成时间）的99%，如图4b所示。


图2: (a) ShareGPT中对话轮次的分布。(b) ShareGPT的会话长度分布。为更好地展示，统计数据排除了超过40轮的对话或长度超过32K的会话。

图3: 重计算与CachedAttention的比较。

图4: 重计算的低效性。(a) ShareGPT中不同轮次的平均历史token和新token数量。(b) 使用Mistral-7B在1个A100 GPU上对ShareGPT进行所有token和仅新输入token的预填充GPU时间。

2.4 机会与挑战

机会与挑战并存。分析表明，如果在多轮对话中复用KV缓存，可以减少高达99%的预填充成本。具体方法是将历史对话的KV缓存保存在GPU之外的缓存系统中，当会话再次激活时，GPU从该系统加载并复用KV缓存。然而，构建一个高效的KV缓存系统面临以下重大挑战：
1. 高昂的KV缓存访问开销：GPU的计算可能会因为等待从缓存系统加载KV缓存而被阻塞。例如，在使用4个NVIDIA A100 GPU的LLaMA-65B模型上，预填充2K tokens耗时约360毫秒，而从主机内存向GPU加载这2K tokens的KV缓存（5GB）就需要约192毫秒。
2. 巨大的KV缓存存储容量需求：存储每个请求的KV缓存需要大量空间。例如，在运行LLaMA-65B的4个A100 GPU（每个80GB HBM）上，KV缓存的生成速度约为13.9 GB/s。剩余的190 GB空闲HBM空间将在14秒内被占满。若将KV缓存溢出到主机内存（如512GB），也会在不到1分钟内被填满。
3. 在不同存储层级中合理放置KV缓存：磁盘容量远大于主机内存（数十TB vs. 数百GB），因此大部分KV缓存会保存在磁盘上。但磁盘访问带宽较低（<5 GB/s），随机到达的对话请求其对应的KV缓存很可能位于磁盘上，导致推理性能不佳。必须确保即将被访问的KV缓存总是位于主机内存中。
4. 已保存KV缓存的意外失效：当对话轮次增加，历史token数可能超过上下文窗口限制。LLM服务引擎通常会进行token截断【索引15，Lm-infinite: Simple on-the-fly length generalization for large language models，2023，arXiv；索引30，https://platform.openai.com/docs/a ssistants/how-it-works/managing-threads-a nd-messages】，这会改变每个token的位置信息，从而使保存在AttentionStore中的KV缓存失效，因为旧的嵌入式位置编码无法匹配。上下文窗口溢出的概率很高，例如，对于上下文窗口为4K的LLaMA-2，30%的对话会话发生溢出；对于窗口为2K的OPT系列，47%的会话会发生溢出。

A2 方法细节

3.1 概述

CachedAttention核心思想。本文提出了一种名为CachedAttention的新型注意力机制，它能够在多轮对话中重用KV缓存。与传统注意力机制在预填充时使用所有提示token不同，CachedAttention仅使用新对话轮次中输入的新token和历史token的KV缓存进行预填充，如图3所示。具体来说，当相关对话会话变为非活动状态时，CachedAttention会将KV缓存保存在一个名为AttentionStore的KV缓存系统中，而不是像传统机制那样丢弃它们。如果同一对话会话将来被激活，其KV缓存将从AttentionStore中取出并用于推理。通过这种方式，CachedAttention仅对部分提示token（即新一轮对话中输入的新token）执行预填充，而不是对所有提示token进行预填充。如图3b所示，在执行第三轮推理时，q1、a1、q2和a2的KV缓存被重用，只需对q3进行预填充。CachedAttention有效消除了历史token的重复计算开销，从而降低了预填充成本。

系统架构。图5展示了CachedAttention的架构概览。它维护一个分层的KV缓存系统，即AttentionStore，并采用高效的KV缓存访问、放置和截断技术来应对第2.4节中提到的挑战。
* 针对挑战1（高访问开销）：为了减少从AttentionStore加载KV缓存到HBM的开销，CachedAttention利用分层预加载方案来重叠KV缓存加载与推理计算。为了减少从HBM保存KV缓存到主机内存的开销，CachedAttention利用异步保存方案来重叠保存与推理计算（§3.2）。
* 针对挑战2和3（高存储需求和放置问题）：为了扩大可用的KV缓存存储空间，CachedAttention在AttentionStore中采用了多层级、高成本效益的存储介质，即主机内存和磁盘。为了减少访问慢速磁盘对推理性能的影响，我们提出了一种调度器感知的获取方案，利用作业调度器的提示，将待访问的KV缓存从磁盘预取到主机内存。同时，为了有效利用有限的主机内存空间，我们提出了一种调度器感知的驱逐方案，识别价值最低的KV缓存并将其驱逐到磁盘或系统外（§3.3）。
* 针对挑战4（KV缓存失效）：为了处理因上下文窗口溢出导致的已保存KV缓存失效问题，我们利用位置编码解耦的截断方案来保存不含嵌入位置编码的KV缓存，从而支持直接对KV缓存进行截断。加载KV缓存时，CachedAttention会重新嵌入新的位置编码（§3.4）。


图5: CachedAttention的系统架构。

3.2 重叠的KV缓存访问

减少KV缓存访问开销。使用较慢的内存/存储层级会带来显著的访问开销，因为KV缓存需要在HBM和较慢介质之间传输，这会阻塞推理并浪费计算资源。为了减少从主机内存到HBM的KV缓存加载开销，CachedAttention使用分层预加载方案，将KV缓存的加载与推理计算逐层重叠（§3.2.1）。为了减少KV缓存的保存开销，CachedAttention开发了一种异步保存方案，将保存KV缓存与推理计算重叠（§3.2.2）。

3.2.1 从内存到HBM的分层预加载

分层预加载机制。CachedAttention从主机内存加载KV缓存到HBM，这带来了很高的数据访问开销，并且该过程位于推理执行的关键路径上，如图6a所示。为消除此开销，CachedAttention采用了分层预加载方案。其主要思想是将KV缓存的加载与对话中新输入token的预填充计算重叠。具体而言，LLM模型由多个Transformer层链接而成，每层都有自己的KV缓存。当GPU执行某一层时，后续层所需的KV缓存可以从主机内存并发加载。这样，当GPU开始计算某一层的自注意力时，该层对应的KV缓存已经位于HBM执行缓冲区中。

带缓冲区的优化预加载。图6b以一个3层模型为例，展示了分层预加载方案如何将KV缓存获取时间与计算时间重叠。在开始计算第1层之前，该层的KV缓存必须先准备好。读取流首先发出一个KV缓存加载操作，将第1层的KV缓存读入HBM执行缓冲区。然后执行流开始计算第1层。当执行流计算某一层时，读取流会并发加载下一层的KV缓存，从而实现加载与计算的重叠。然而，上一个作业与当前作业的第一层之间仍然存在一个间隙，因为加载操作只有在HBM执行缓冲区可用（即上一个作业完成）后才能开始。为了进一步减小这个间隙，CachedAttention预留了一个HBM读取缓冲区来消除它。具体来说，如图6c所示，有了读取缓冲区，读取流不必等待上一个作业释放执行缓冲区，它可以在上一个作业仍在运行时就开始预加载。

处理不完美重叠。然而，如果KV缓存的加载时间长于预填充的计算时间，预加载可能无法完全与计算重叠。如图7a所示，层与层之间的计算存在多个间隙，因为每层KV缓存的获取时间超过了每层的计算时间，导致不完美的重叠。这个开销可以通过使用一个定制的更大预加载缓冲区来进一步最小化。有了更大的缓冲区，预加载可以更早开始。例如，如图7b所示，更大的缓冲区允许预加载更多层的KV缓存，从而可以重叠层间的间隙。设 $T_{load}$、 $T_{pref}$、 $L_{hist}$ 和 $L_{new}$ 分别表示一个token的KV缓存访问时间、一个token的预填充时间、会话中历史token的长度和对话中新输入token的长度。不完美的重叠发生在 $T_{load}L_{hist} > T_{pref}L_{new}$ 时，这表示传输时间大于部分预填充时间。缓冲区用于填补时间差 $T_{load}L_{hist} - T_{pref}L_{new}$。结合PCIe带宽B，缓冲区大小可由以下公式设定：$S_{buf} = B(T_{load}L_{hist} - T_{pref}L_{new})$。


(a) 基线：无并发操作的KV缓存加载。

图7: 分层KV缓存预加载。
(b) 使用自定义更大缓冲区的完美预加载。

(c) 带缓冲区的分层预加载。

图6: 分层KV缓存预加载。蓝色块表示每个transformer层的执行。红色块表示每个transformer层的KV缓存加载。

3.2.2 从HBM到内存的异步保存

异步保存机制。CachedAttention需要将KV缓存保存到主机内存，以实现跨对话的重用。一种基线方法是在一轮对话结束后，将所有生成的KV缓存一起写入。然而，这种方法可能会延迟下一个调度作业的执行，因为KV保存时间位于推理的关键路径上，如图8a所示。为了减少这个开销，CachedAttention采用了一种异步KV缓存保存方案，将KV缓存的回写与计算重叠，并考虑了预填充和解码阶段的不同特性来执行不同的重叠机制。

针对不同阶段的异步保存策略。具体来说，预填充和解码阶段的KV缓存生成速度不同。预填充阶段并发处理token，因此在有限的时间内生成大量的KV缓存。而解码阶段一次只生成一个token的KV缓存。如图8b所示，对于预填充阶段，由于每个自注意力操作都能产生大量KV缓存，写入流会逐层地保留KV缓存。预填充阶段产生的KV缓存可以与解码阶段重叠。对于解码阶段，由于KV缓存是迭代产生的，写入流会在解码的同时逐层回写KV缓存。为了避免在解码已完成而KV缓存尚未完全回写时发生阻塞，我们还预留了一个HBM写入缓冲区，类似于KV缓存预取中使用的读取缓冲区。未完成的KV缓存会暂时移到写入缓冲区，以避免阻塞下一个作业的执行。


图8: 异步KV缓存保存。

3.3 分层KV缓存放置

利用分层存储。CachedAttention利用主机内存和磁盘来扩展可用的KV缓存存储空间。主机内存（DRAM）的访问速度远高于磁盘（SSD）（数十GB/s vs. 几GB/s）。如果待访问的KV缓存总是在主机内存中而不是磁盘中找到，那么KV缓存的访问性能将是最佳的。为实现这一目标，CachedAttention应用了调度器感知的获取方案，将KV缓存从磁盘预取到主机内存，确保以最佳速度访问KV缓存（§3.3.1），以及调度器感知的驱逐方案，将合适的KV缓存从主机内存驱逐到磁盘（§3.3.2）。


图9: 调度器感知的KV缓存获取和驱逐。

3.3.1 从磁盘到内存的调度器感知获取

调度器感知的预取。由于磁盘的容量远大于主机内存（数十TB vs. 数百GB），大多数KV缓存都保存在磁盘中。由于对话请求是随机到达的，它们对应的KV缓存更有可能位于磁盘上，导致访问性能不佳。为了解决这个问题，我们利用一种调度器感知的KV缓存获取方案，通过利用来自推理作业调度器的提示，将待访问的KV缓存从磁盘预取到主机内存。具体来说，作业调度器维护一个作业队列，因此完全了解等待中的作业。CachedAttention应用一个前瞻预取窗口来观察待执行的等待作业。如果等待作业的KV缓存在磁盘中命中，CachedAttention将在这些等待作业被执行之前，将其KV缓存从磁盘预取到主机内存。前瞻预取窗口的长度由主机内存中的可用容量决定。给定用于预取的可用内存容量 $C_{mem}$ 和一个会话的平均KV大小 $S_{kv}$，预取窗口长度为 $L_{pw} = C_{mem}/S_{kv}$。

预取示例。图9展示了一个调度器感知的获取示例。当作业1正在执行时，KV缓存管理器应用一个大小为2的前瞻窗口（主机内存中有2个KV缓存槽位用于KV缓存获取）来检查等待中的作业2-3的KV缓存命中状态。作业2的KV缓存在主机内存中命中，但作业3的KV缓存不在主机内存中。然后，KV缓存获取线程开始将作业3的KV缓存从磁盘获取到主机内存。值得注意的是，CachedAttention包含一个主机内存缓冲区，允许从磁盘到内存的无缝KV缓存获取，防止在主机内存满时发生任何延迟。当可用内存容量达到定义的阈值时，CachedAttention会触发从主机内存到磁盘的KV驱逐，以确保持续可用主机内存缓冲区。

3.3.2 从内存到磁盘的调度器感知驱逐

调度器感知的驱逐策略。当主机内存中的空闲空间耗尽时，我们需要将一些KV缓存从主机内存驱逐到磁盘。同时，如果磁盘也满了，我们还需要将一些存储在磁盘中的KV缓存驱逐出系统。因此，仔细选择合适的KV缓存候选项进行驱逐对于实现高缓存命中率至关重要。与现有的缓存驱逐策略（如最近最少使用（LRU）【索引46，Fast and exact analysis for lru caches，2019，POPL】、先进先出（FIFO）【索引8，An approximate analysis of the lru and fifo buffer replacement schemes，1990，ACM SIGMETRICS】及其变体）仅依赖KV缓存的历史访问信息不同，CachedAttention提出了一种调度器感知的驱逐方案，可以利用KV缓存的未来访问信息来实现更高的缓存命中率。作业调度器中的作业队列为我们提供了实现这一目标的机会。具体来说，CachedAttention在作业队列中维护一个前瞻驱逐窗口。前瞻驱逐窗口的最大长度由AttentionStore的总存储容量决定。假设磁盘中的总可用容量为 $C_{disk}$，前瞻驱逐窗口长度为 $(C_{mem} + C_{disk})/S_{kv}$。当CachedAttention试图从AttentionStore中驱逐一个项目时，如果在前瞻驱逐窗口中找到该项目，则该项目被豁免。当CachedAttention从主机内存向磁盘驱逐一个项目时，位于前瞻驱逐窗口尾部的项目具有更高的被驱逐优先级。注意，一个项目对应于与一个对话会话相关的所有KV缓存，这是CachedAttention中最小的驱逐和获取粒度。

驱逐示例。图9展示了一个调度器感知的驱逐示例。当作业3的KV缓存被选择迁移到主机内存时，将使用缓冲区。为了在主机内存中维持一个缓冲区，CachedAttention需要将KV缓存从主机内存驱逐到磁盘。CachedAttention使用一个大小为6的前瞻驱逐窗口来监控作业的KV缓存状态。首先，它发现主机内存中的所有KV缓存都在作业队列中有一个关联的作业。然后，它从尾到头扫描前瞻驱逐窗口，优先驱逐靠近尾部的作业。因此，作业4的KV缓存被选择从主机内存驱逐到磁盘。由于磁盘也已满，扫描过程识别出作业队列中最后到达的作业9是最合适的被驱逐候选项。最终，作业4的KV缓存被移动到先前由作业9占用的位置。

3.4 解耦的KV缓存截断

上下文窗口溢出问题。当历史token超过上下文窗口的限制时，LLM服务引擎通常会执行token截断【索引30，https://platform.openai.com/docs/a ssistants/how-it-works/managing-threads-a nd-messages】。如图10a所示，上下文窗口大小为4K。一旦上下文窗口溢出，LLM服务引擎会切掉提示的前2K个token。截断对以前的LLM服务引擎没有影响，因为它们总是根据输入的提示重新计算KV缓存，无论是否截断。然而，截断会使存储在CachedAttention中的KV缓存失效，显著降低了CachedAttention的效率。这是由于KV缓存中嵌入了位置编码。在对提示进行token截断时，每个token的位置都发生了变化。KV缓存中嵌入的位置编码无法修改以匹配提示中token的位置，从而使KV缓存失效。

位置编码解耦方案。为了解决这个问题，CachedAttention通过解耦位置编码，使得截断后的KV缓存仍然有效。CachedAttention需要与相对位置编码（RPE）【索引42，Roformer: Enhanced transformer with rotary position embedding，2024，Neurocomputing；索引44，Llama: Open and efficient foundation language models，2023，arXiv；索引54，Efficient streaming language models with attention sinks，2023，arXiv】协同工作。与绝对位置编码（APE）将位置编码添加到输入中不同，RPE直接将位置编码嵌入到查询（Q）和键（K）向量中，如图11b所示。大量研究表明，RPE比APE能让LLM从更长的数据序列中学习【索引11，Bert: Pre-training of deep bidirectional transformers for language understanding，2018，arXiv；索引47，Attention is all you need，22017，NeuIPS；索引60，Ernie: Enhanced language representation with informative entities，2019，arXiv】，因此RPE被广泛用于现代LLM，如LLaMA、T5、Falcon、Mistral、Mixtral和Transformer-XL【索引7，Transformerxl: Attentive language models beyond a fixed-length context，2019，arXiv】。通过简单地将缓存KV的时间点移到RPE中嵌入位置编码之前，如图11c所示，CachedAttention可以在AttentionStore中存储没有嵌入位置编码的KV。当在CachedAttention中重用这些KV时，它们会被嵌入新的位置编码，并用于后续的推理。

KV缓存截断示例。图12提供了一个CachedAttention如何支持KV缓存截断的示例。CachedAttention存储没有位置编码的KV缓存。在需要进行KV缓存截断的情况下，LLM引擎检索被截断的KV缓存（即KV[0:1536]）并将其加载到HBM。随后，新的位置编码被应用于该KV缓存。值得注意的是，CachedAttention还允许选择性地保留某些KV缓存以进行压缩，例如具有重要分数的初始token【索引54，Efficient streaming language models with attention sinks，2023，arXiv】或重要token【索引12，Model tells you what to discard: Adaptive kv cache compression for llms，2023，arXiv；索引24，Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time，2023，arXiv；索引61，H2o: Heavyhitter oracle for efficient generative inference of large language models，2023，arXiv】，以进一步提高LLM的生成质量。具体来说，一种给定的KV缓存压缩技术实质上提供了一种在提示中创建token丢弃列表（TDL）的方法。CachedAttention直接遵守TDL，在AttentionStore中丢弃与TDL相关的KV缓存。


图10: 管理上下文窗口溢出的图示。上下文窗口大小：4K，截断比例：2K。(a) 基线。Token截断。(b) KV缓存截断。

图11: (a) 绝对位置编码。(b) 相对位置编码。(c) 解耦位置编码的KV缓存。

图12: 使用CachedAttention进行KV缓存截断的图示。

A4 实验环境

  • 硬件配置:
    • GPU: 4块NVIDIA A100 GPU，每块配备80GB HBM。
    • 主机内存: 128GB DRAM。
    • 存储: 10TB SSD。
    • 连接: GPU通过PCIe Gen 4连接到主机。
  • 软件配置:
    • 实现: 在Pytorch和Python中实现。
    • 依赖库: 基于Transformers库【索引51，Huggingface’s transformers: State-of-the-art natural language processing，2019，arXiv】集成了LLaMA【索引44，Llama: Open and efficient foundation language models，2023，arXiv】和Falcon【索引32，The RefinedWeb dataset for Falcon LLM: outperforming curated corpora with web data, and web data only，2023，arXiv】等流行LLM的实现。
    • 特性: 使用专用的CUDA流在GPU和主机内存之间移动数据，实现计算与主动交换的重叠。独立的IO线程在主机内存和磁盘之间迁移数据，实现执行与KV缓存迁移的重叠。实验中启用了连续批处理【索引57，Orca: A distributed serving system for transformer-based generative models，2022，OSDI】。
  • 模型架构:
    • 评估了开源的LLaMA-1 65B【索引44，Llama: Open and efficient foundation language models，2023，arXiv】、LLaMA-2 13B和70B【索引45，Llama 2: Open foundation and fine-tuned chat models，2023，arXiv】以及Falcon 40B【索引32，The RefinedWeb dataset for Falcon LLM: outperforming curated corpora with web data, and web data only，2023，arXiv】。
    • 中间激活使用FP16。
    • 还实现了具有32K上下文窗口的Mistral-7B【索引19，Mistral 7b，2023，arXiv】。
    • LLaMA-13B在2个GPU上运行，批处理大小为24；LLaMA-65B、LLaMA-70B和Falcon-40B在4个GPU上运行，批处理大小均为24。
  • 数据集:
    • 名称: ShareGPT数据集【索引33，https://huggingface.co /datasets/philschmid/sharegpt-raw/tree/main /sharegpt_90k_raw_dataset；索引62，Judging llm-as-ajudge with mt-bench and chatbot arena，2023，arXiv】。
    • 规模与用途: 使用了9000个对话会话。请求到达时间根据泊松分布（λ=1.0）生成。

A4 实验结果

端到端性能

实验使用ShareGPT的9000个对话，平均轮次为5.75，总计约52000个对话轮次。使用前10000个对话轮次预热AttentionStore，然后在接下来的42000个轮次上评估性能。

  • 缓存命中率：如图13所示，CachedAttention (CA) 在LLaMA-13B、LLaMA-65B、LLaMA-70B和Falcon-40B上分别达到了约86%、71%、89%和90%的高命中率。LLaMA-65B的命中率相对较低，因为其KV缓存占用空间更大，导致在相同存储空间下能容纳的会话更少。
  • 首个token生成时间（TTFT）：如图14所示，与重计算（RE）相比，CA显著降低了TTFT，降幅分别为85%、61%、87%和86%。这是因为CA消除了历史token的大量重复计算。
  • 预填充吞吐量：如图15所示，CA的预填充吞吐量比RE高出6.8倍、2.6倍、7.8倍和7.2倍。这得益于预填充时间的减少和与KV缓存加载重叠的优化。
  • GPU时间：如图16所示，CA完成所有推理任务的总GPU时间比RE快4.0倍、1.9倍、3.3倍和3.4倍。性能提升来自于减少了历史token的重计算以及上下文溢出后的重计算。同时，通过减少预填充时间，CA也缩短了连续批处理中解码作业的等待时间，从而间接减少了解码时间。
  • 推理成本：基于AWS EC2实例价格计算，如图17所示，CA相比RE节省了70%、43%、66%和68%的成本。成本节约主要源于GPU时间的减少。存储成本（DRAM和SSD）在CA的总成本中占比为9.0%至16.4%。


图13: 缓存命中率。

图14: 首个token生成时间。

图15: 预填充吞吐量。

图16: GPU时间。

图17: 推理成本。

消融研究

  • 重计算 vs. CachedAttention：如图18所示，在不同历史/新token比例下，CA始终优于RE。随着历史token比例增加，CA的优势更明显，且分层预加载方案能有效隐藏KV缓存加载时间。
  • 重叠KV缓存访问：如图19和20所示，分层预加载方案将预填充时间减少了高达61%（与无预加载相比）。异步保存方案将整体执行时间减少了13%至15%。
  • 调度器感知的获取和驱逐：如图21所示，CA的驱逐算法在整体缓存命中率上比LRU和FIFO高出27%和31%，带来了高达2.7倍的GPU时间加速。由于调度器感知的预取，CA命中在DRAM的比例超过99.6%，而LRU和FIFO的DRAM命中率极低（约0.5%）。
  • 解耦KV缓存截断的性能：如图22所示，与在上下文溢出时需要重计算的基线方法（OF）相比，CA通过直接截断KV缓存避免了缓存失效，从而将命中率提高了17.6%至41.5%，并减少了GPU时间。
  • 解耦位置编码的准确性：如表1和表2所示，CA在困惑度（PPL）和各种基准测试的准确性上与标准的token截断（TT）相当，而朴素的KV缓存截断（NKVT）由于位置信息错乱导致性能大幅下降。
  • 缓存容量需求：如图23所示，当缓存容量达到理论最大需求（CCpUT）的25%时，CA即可实现98%的缓存命中率和峰值吞吐量。
  • 缓存存储介质的影响：如图24所示，仅使用HBM缓存的命中率接近0%。增加DRAM能有所改善，但只有结合大容量SSD才能实现CA的高命中率（71%-90%）和显著的性能提升。
  • 会话到达率的影响：如图25所示，当会话到达率从0.5/s增加到2.0/s时，缓存命中率从82%轻微下降到77%，对TTFT、吞吐量和GPU时间有一定影响，但总体影响较小，CA仍然保持高效。


图18: 重计算 vs. CachedAttention。

图19: CA无预加载 vs. CA预加载（不同缓冲区大小）。
图20: 使用写重叠的性能影响。

图21: 不同存储设置下驱逐算法的比较。

图22: 上下文溢出影响。
表1: 不同方法的PPL比较。

表2: 不同方法的准确性。


图23: 存储容量和不同会话数量的影响。

图24: 不同缓存配置下的性能。

图25: 会话到达率的影响。

A5 结论

本文提出了CachedAttention，一种新的注意力机制，通过在同一对话的后续轮次中重用KV缓存，显著减少了LLM中的重计算开销。为了提高CachedAttention的效率，我们设计了重叠的KV缓存访问、分层的KV缓存放置以及位置编码解耦的KV缓存截断方案。大量的实验结果表明，CachedAttention在多轮对话中，可将首个token生成时间（TTFT）最多降低87%，将提示预填充吞吐量提高7.8倍，并将端到端推理成本最多降低70%。

方法细节中的引用汇总

  • 引用 [7] Zihang Dai, et al. "Transformer-xl: Attentive language models beyond a fixed-length context." (2019, arXiv).
    • 引用位置: §3.4 Decoupled KV Cache Truncation.
    • 原文描述: 用于说明相对位置编码（RPE）被广泛用于现代LLM，Transformer-XL是使用RPE的例子之一。
  • 引用 [8] Asit Dan and Don Towsley. "An approximate analysis of the lru and fifo buffer replacement schemes." (1990, ACM SIGMETRICS).
    • 引用位置: §3.3.2 Scheduler-aware Eviction from Memory to Disks.
    • 原文描述: 作为现有缓存驱逐策略（如FIFO）的例子，与本文提出的调度器感知驱逐方案进行对比。
  • 引用 [11] Jacob Devlin, et al. "Bert: Pre-training of deep bidirectional transformers for language understanding." (2018, arXiv).
    • 引用位置: §3.4 Decoupled KV Cache Truncation.
    • 原文描述: 用于支持“大量研究表明RPE比APE能让LLM从更长的数据序列中学习”这一论点。
  • 引用 [12] Suyu Ge, et al. "Model tells you what to discard: Adaptive kv cache compression for llms." (2023, arXiv).
    • 引用位置: §3.4 Decoupled KV Cache Truncation.
    • 原文描述: 作为一种可以保留重要token以进行压缩的技术示例，说明CachedAttention可以与之兼容。
  • 引用 [24] Zichang Liu, et al. "Scissorhands: Exploiting the persistence of importance hypothesis for llm kv cache compression at test time." (2023, arXiv).
    • 引用位置: §3.4 Decoupled KV Cache Truncation.
    • 原文描述: 同上，作为保留重要token的压缩技术示例。
  • 引用 [30] OpenAI. "Managing Threads and Messages." (2024, OpenAI Platform Documentation).
    • 引用位置: §3.4 Decoupled KV Cache Truncation.
    • 原文描述: 用于说明当历史token超过上下文窗口限制时，LLM服务引擎通常会执行token截断。
  • 引用 [42] Jianlin Su, et al. "Roformer: Enhanced transformer with rotary position embedding." (2024, Neurocomputing).
    • 引用位置: §3.4 Decoupled KV Cache Truncation.
    • 原文描述: 作为相对位置编码（RPE）的一个例子，说明CachedAttention需要与这类编码协同工作。
  • 引用 [44] Hugo Touvron, et al. "Llama: Open and efficient foundation language models." (2023, arXiv).
    • 引用位置: §3.4 Decoupled KV Cache Truncation.
    • 原文描述: 作为使用RPE的现代LLM的例子之一。
  • 引用 [46] Valentin Touzeau, et al. "Fast and exact analysis for lru caches." (2019, POPL).
    • 引用位置: §3.3.2 Scheduler-aware Eviction from Memory to Disks.
    • 原文描述: 作为现有缓存驱逐策略（如LRU）的例子，与本文提出的调度器感知驱逐方案进行对比。
  • 引用 [47] Ashish Vaswani, et al. "Attention is all you need." (2017, NeuIPS).
    • 引用位置: §3.4 Decoupled KV Cache Truncation.
    • 原文描述: 用于支持“大量研究表明RPE比APE能让LLM从更长的数据序列中学习”这一论点。
  • 引用 [54] Guangxuan Xiao, et al. "Efficient streaming language models with attention sinks." (2023, arXiv).
    • 引用位置: §3.4 Decoupled KV Cache Truncation.
    • 原文描述: 引用了两次。第一次作为RPE的例子。第二次作为保留初始重要token以进行压缩的技术示例。
  • 引用 [60] Zhengyan Zhang, et al. "Ernie: Enhanced language representation with informative entities." (2019, arXiv).
    • 引用位置: §3.4 Decoupled KV Cache Truncation.
    • 原文描述: 用于支持“大量研究表明RPE比APE能让LLM从更长的数据序列中学习”这一论点。
  • 引用 [61] Zhenyu Zhang, et al. "H2o: Heavyhitter oracle for efficient generative inference of large language models." (2023, arXiv).
    • 引用位置: §3.4 Decoupled KV Cache Truncation.
    • 原文描述: 作为保留重要token的压缩技术示例。