具有成本效益的大型语言模型，通过 CachedAttention 服务于多轮对话

摘要

通过多轮对话与人类交互是大语言模型的基本特征。 然而，现有的执行多轮对话的大语言模型服务引擎由于需要重复计算历史 Token 的键值（KV）缓存而效率低下，导致服务成本较高。 为了解决这个问题，本文提出了 CachedAttention，这是一种新的注意力机制，可以在多轮对话中重用 KV 缓存，从而显着减少重复计算开销。 CachedAttention 维护一个分层的 KV 缓存系统，利用经济高效的内存/存储介质来保存所有请求的 KV 缓存。 为了减少慢速介质的 KV 缓存访问开销，CachedAttention 采用分层预加载和异步保存方案，将 KV 缓存访问与 GPU 计算重叠。 为了确保要访问的 KV 缓存放置在最快的层次结构中，CachedAttention 采用调度程序感知的获取和驱逐方案，根据推理作业调度程序的提示有意识地将 KV 缓存放置在不同的层中。 为了避免上下文窗口溢出导致已保存的 KV 缓存失效，CachedAttention 通过解耦位置编码并有效截断 KV 缓存，使已保存的 KV 缓存保持有效。 大量的实验结果表明，CachedAttention 将第一个词符 (TTFT) 的时间显着缩短了 87%，将多轮对话的提示预填充吞吐量提高了 7.8$\times$，并减少了结束时间- 端到端推理成本降低高达 70%。

1简介

大型语言模型（大语言模型）在各种任务上表现出色，开创了生成式应用的新时代[29,44,43]。 然而，由于大语言模型推理使用大量 GPU，因此使用大语言模型为这些生成应用程序提供服务非常昂贵。 鉴于对生成应用程序的高需求，降低推理成本变得至关重要。

与人类进行多轮对话是大语言模型[53, 49]的必备能力。 这些多轮对话有助于大语言模型理解上下文、用户意图和情感细微差别，从而增强他们做出适当反应的能力。 根据从 ChatGPT 收集的广泛使用的真实数据集 ShareGPT 数据[33]，73% 的对话涉及多次对话，如第 2.3 节中分析的那样。

然而，在现有的大语言模型服务引擎中执行多轮对话效率很低，需要大量的重复计算，服务成本很高。 在单轮会话期间，大语言模型引擎将中间数据、键值 (KV) 对 [3, 34, 21] 存储在 GPU 上有限的高带宽内存 (HBM) 中。 当该会话结束且会话会话变为非活动状态时，大语言模型引擎通常会丢弃与该会话关联的 KV 缓存，以便为其他活动会话释放 HBM 中的空间。 当会话再次激活时，即用户在会话中发送下一条消息时，大语言模型引擎会再次计算整个 KV 缓存。 这会导致同一个KV缓存的重复计算，浪费宝贵的GPU计算资源。 随着会话次数的增加，重复计算开销线性增加。 我们基于ShareGPT的分析表明，高达99%的预填充成本来自于KV缓存的重复计算，如2.3节所示。

为了降低服务成本并提高推理性能，本文提出了CachedAttention，这是一种新的注意力机制，可以在多轮对话中重用 KV 缓存，而不是丢弃它们。 当对话会话变为非活动状态时，CachedAttention 将相应的 KV 缓存保存在 KV 缓存系统中，名为 AttentionStore。 当同一会话恢复时，CachedAttention 会加载并重用 AttentionStore 中保存的 KV 缓存，从而消除重复计算的开销。 然而，为多轮对话构建如此高效的 KV 缓存系统面临着巨大的挑战。

首先，AttentionStore作为GPU的外部存储，通过低速链路连接到GPU。 由于需要在 HBM 和 AttentionStore 之间传输 KV 缓存，因此使用 AttentionStore 会带来很大的访问开销。 KV缓存的访问开销位于推理执行的关键路径上。 这是因为 GPU 只有在成功将相应的 KV 缓存加载到 HBM 中后才能执行推理作业的计算。 同样，如果HBM空间不够，后续的推理作业也需要等待，直到之前作业的KV缓存被移出HBM。 为了减少 KV 缓存加载开销，CachedAttention 使用分层预加载方案，将加载 KV 缓存的时间与推理计算的时间逐层重叠。 为了减少 KV 缓存保存开销，CachedAttention 开发了一种异步保存方案，将保存 KV 缓存的时间与推理计算重叠。

其次，KV缓存占用大量存储空间，并且在会话过程中不断扩展。 先前的工作试图通过在 HBM [18, 63] 中保留跨多轮对话的 KV 缓存来减少重复 KV 计算的低效率。 然而，这很快就会耗尽有限的 HBM 容量。 我们在第 2.3 节中介绍了 LLaMA-65B 的示例，该示例显示 KV 缓存在 14 秒内完全占用了 HBM 内的可用空间。 为了应对这一挑战，CachedAttention探索并利用比HBM更慢但容量更大的存储层次结构，包括主机内存和磁盘，为缓存KV缓存提供足够的存储空间。

第三，由于磁盘的容量比主机内存大得多（数十 TB vs. 几百GB），大多数KV缓存都保留在磁盘中用于CachedAttention。 由于会话请求随机到达，其对应的KV缓存更有可能位于磁盘中，导致访问性能较差。 为了解决这个问题，CachedAttention 使用了调度程序感知的 KV 缓存获取方案。 该方案利用从推理作业调度程序接收到的提示，将可能从磁盘访问的 KV 缓存预取到主机内存。 当主机内存的可用空间不足时，CachedAttention还采用调度程序感知的驱逐方案，高效地识别内存中最合适的KV缓存，并将其驱逐到磁盘或系统之外。

最后，当会话会话超过大语言模型上下文窗口的限制时，例如LLaMA-2 [45]中的4K，大语言模型通常会截断最旧的标记并将上下文限制为最旧的标记。最新的 Token [30]。 此截断使得 CachedAttention 中该对话的所有保存的 KV 缓存无效，因为嵌入在 KV 缓存中的所有 Token 的位置信息都发生了更改。 为了解决这个问题，CachedAttention 在保存 KV 缓存时将位置编码与 KV 缓存解耦。 加载时，它会将位置编码重新嵌入到 KV 缓存中。 解耦后，可以直接对KV缓存进行截断，从而保证保存的KV缓存的可重用性。

我们实现 CachedAttention 并使用真实的 ShareGPT 数据集 [33] 对其进行评估。 大量的实验结果表明，对于多轮对话，CachedAttention 将到达第一个词符 (TTFT) 的时间显着缩短了 87%，并将提示预填充吞吐量提高了 7.8$\times$。 它还将端到端推理成本降低高达 70%。 总而言之，本文做出了以下贡献：

• 我们研究了大语言模型中跨会话轮次的 KV 缓存的重新计算开销，并确定了与跨多轮会话保留 KV 缓存相关的挑战。

• 我们提出了 CachedAttention，这是一种新的注意力机制，允许在同一会话的任何后续会话中重用 KV 缓存，从而显着减少大语言模型中 KV 缓存的重新计算开销。

• 为了提高CachedAttention的效率，我们设计了重叠KV缓存访问、分层KV缓存放置和位置编码解耦KV缓存截断方案。

• 我们使用真实数据集彻底评估 CachedAttention，以证明其功效和效率。

2 背景和动机

本节首先介绍生成式大语言模型推理的基础知识，然后探讨大语言模型在多轮对话中的低效率问题。 本节最后讨论了处理这些低效率问题的设计机会以及设计此类系统所面临的挑战。

2.1 生成式大语言模型推理基础

Transformer 架构。 Transformer 已成为生成大语言模型推理中广泛接受的标准。 广泛使用的大语言模型，如 GPT [29] 和 LLaMA [44, 45] 是建立在自回归 Transformer 架构[47, 16] 之上>。 在推理过程中，这些模型处理用户的提示并生成响应。 提示被作为输入标记序列进行处理，并且响应是由模型使用所有先前标记的上下文预测后续标记的概率来生成的。 Transformer 模型由一系列 $l$ Transformer 层组成，每个层由两个阶段组成：自注意力和前馈网络 (FFN)。

对于输入词符列表$X=[x_1,x_2,\mathrm{\dots }{x}_s]$，每个层使用权重$W_Q,W_K,W_V$对$X$中的每个词符应用一系列投影。 这会生成查询、键和值集中的元素，分别称为 $Q$、$K$ 和 $V$：

$$Q=W_{Q}X,K=W_{K}X,V=W_{V}X$$

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

$$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d_K}})V$$

其中 $\sqrt{d_k}$ 是键向量 $k$ 的维度。最后，投影操作对注意力分数应用线性变换。 该投影结果被传递给 FFN 层。 FFN 的结果作为输入传递到下一个 Transformer 层。 最后，在输入经过所有 $l$ Transformer 层处理后，输出是一个概率向量，标记出最可能的输出标记。

KV缓存： 上述过程为每个词符生成中间 K 和 V 张量。 当生成后续 token 时，前面 token 的所有 KV 张量都是计算 self-attention 所必需的。 这些K和V张量通常缓存在GPU中，称为KV缓存。 KV 缓存通常占用较大的空间。 例如，GPT-3 [29, 10] 为每个词符生成 4.5MB KV 缓存。 KV缓存的大小随着提示 Token 的数量线性增加。 包含数千个 Token 的会话会话将产生数 GB 的 KV 缓存。

2.2 自回归生成

(a) Two-phase illustration.

(b) Execution latency.

图1： 预填充和解码阶段。 在 4 个 A100 GPU 上测量批量大小为 8 的 LLaMA-70B 的延迟。

如图1(a)所示，基于 Transformer 的发电在逻辑上可以被识别为两个不同的阶段[1]。

预填充阶段。 给定请求提示，生成器将提示词符列表$X=[x_1,x_2,\mathrm{\dots }{x}_s]$作为输入，然后继续计算词符$x_{s+1}$。 这个过程会生成一系列的KV，具体形成从1到$s$的KV缓存，用于解码阶段。

解码阶段。 解码阶段迭代地生成输出 Token 。 解码阶段将词符$\text{s+}1$和预填充阶段的KV缓存$[1$:$\text{s}]$作为输入来计算KV缓存$\text{s+}1$和词符$\text{s+}2$。 生成过程不断迭代，直到生成的词符为<eos>或者迭代次数达到允许的最大生成次数。 由于对先前迭代的严重数据依赖性，解码阶段按顺序发生。

这两个阶段呈现出不同的执行时间特征。 预填充阶段并行计算 KV 缓存，其持续时间与输入提示标记的数量密切相关。 如图1(b)所示，预填充阶段的执行时间随着输入 Token 的增加而增加。 相反，解码阶段每次迭代计算单个词符，从而导致每次迭代的计算时间相对恒定。

图2： (a) ShareGPT [33] 中对话轮次的分布。 (b) ShareGPT 的会话长度分布。 为了更好的显示效果，统计时不包括超过40次的会话或超过32K的会话。

2.3 多轮对话推理

图3： 重新计算和CachedAttention的比较。

让人类参与多轮对话是现代大语言模型的基本特征。 多轮对话会话由一系列连续的对话组成，表示为$D=[d_1,d_2,\mathrm{\dots }{d}_N]$。 在每个对话$d_j$中，用户输入新的问题或命令$q_j$，然后等待大语言模型的响应$a_j$。 为了保持上下文的连贯性和对对话会话的理解，大语言模型根据所有先前对话回合的历史标记生成 $a_{N+1}$ $d[1$:$N]$以及当前回合的输入标记，表示为$q_1{a}_1{q}_2{a}_2\mathrm{\dots }{q}_N{a}_N{q}_{N+1}$。

图4： 重新计算效率低下。 (a) ShareGPT 不同轮次历史 Token 和新 Token 的平均数量[33]。 (b) 在 1 个 A100 GPU 上使用 Mistral-7B [19] 预填充 ShareGPT 中的所有 Token 和仅新输入 Token 的 GPU 时间。

通过对从 ChatGPT 收集的包含超过 90K 对话的真实数据集 ShareGPT [33, 38] 的分析，我们发现 73% 的对话是多轮的，如图 标签：图：dist_turnnum。 此外，30% 的对话拥有超过 4K 的 Token ，如图 LABEL:fig:dist_tokennum 所示。

然而，由于跨多个会话轮次的KV缓存的重复计算，在当前的大语言模型服务引擎中执行多轮会话效率低下。 如图LABEL:fig:conversation所示，在会话轮1中，大语言模型服务引擎生成$q_1$和$a_1$的KV缓存。 完成第1轮后，大语言模型服务引擎丢弃KV缓存以回收HBM空间。 步骤2，大语言模型服务引擎重新生成$q_1$和$a_1$的KV缓存。 第3步，重新生成$q_1$、$a_1$、$q_2$、$a_2$的KV缓存。 随着会话的扩展，历史 Token 不断积累，重复计算量显着增加。 如图LABEL:fig:recompute-tokens所示，随着会话次数的增加，新会话中历史token的比例将超过99%。 重复计算时间占据了新会话中预填充时间（即到第一个词符的时间）的 99%，如图LABEL:fig:recompute-costs所示。

2.4机遇与挑战

根据上面的分析，我们观察到，如果 KV 缓存可以在多轮会话中重用，则可以减少高达 99% 的预填充成本。 具体地，可以将历史会话的KV缓存保存在GPU外的KV缓存系统中。 重新激活会话会话后，GPU 会从 KV 缓存系统加载关联的 KV 缓存，并将其重新用于新一轮会话。 然而，要构建一个高效的KV缓存系统，还存在许多重大挑战。

1）KV缓存访问开销高。 在推理过程中，GPU 的计算可能会由于等待从 KV 缓存系统加载 KV 缓存而被阻塞。 与KV缓存的重复计算时间相比，出块时间是不可忽略的，使得KV缓存解决方案失去效力。 例如，我们使用 4 个 NVIDIA A100 GPU 评估 LLaMA-65B 模型的推理时间，并观察到预填充提示的 2K Token 大约消耗 360 毫秒。 相比之下，将 2K Token （5 GB）的 KV 缓存从主机内存加载到 GPU 大约需要 192 ms（具有 16 通道 PCIe Gen4 的 GPU 系统的有效数据传输带宽约为 26 GB/s）。

2）KV缓存的存储容量要求高。 为每个请求存储KV缓存会消耗大量的存储空间。 例如，当使用 4 个 A100 GPU（每个 GPU 具有 80GB HBM）来运行 LLaMA-65B 时，预填充 2K Token 大约需要 360 毫秒。 该过程生成了5GB的KV缓存，表明KV缓存的生成速度约为13.9GB/s。 由于分配了 130 GB 的 HBM 空间来存储模型，因此剩余的 190 GB 空闲 HBM 空间将在 14 秒内被 KV 缓存完全占用。 如果将KV缓存溢出到主机内存（例如512GB空间），主机内存将在不到1分钟的时间内被填满。 使用磁盘保存KV缓存可以扩展存储空间，但访问性能会变差，如下所示。

3）不同层次结构中KV缓存的适当放置。 磁盘提供比主机内存大得多的容量（数十 TB 与 10 TB）。 数百 GB）。 因此，大多数 KV 缓存都保留在磁盘中。 然而，磁盘的访问带宽小于5GB/s。 由于会话请求随机到达，访问时其对应的KV缓存更有可能位于磁盘中，导致推理性能较差。 必须确保近期要访问的KV缓存始终放置在主机内存中而不是磁盘中。

4) 保存的KV缓存意外失效。 随着会话次数的增加，历史 Token 可能会超出上下文窗口限制。 大语言模型服务引擎一般会进行词符截断[30, 15]来减少输入提示。 截断对以前的大语言模型服务引擎没有影响，因为它们总是根据截断的输入提示重新计算 KV 缓存。 然而，截断使得AttentionStore中保存的KV缓存无效，因为截断后每个词符的位置都发生了变化。 因此它无法匹配保存的 KV 缓存中旧的嵌入位置编码。 这种上下文窗口溢出很有可能发生。 如图LABEL:fig:dist_tokennum所示，分别有 47% 和 30% 的对话会话的上下文长度超过 2K 和 4K。 这意味着，当使用具有 4K 上下文窗口 [45] 的 LLaMA-2 系列时，30% 的对话会话会发生上下文窗口溢出。 当使用具有 2K 上下文窗口 [59] 的 OPT 系列时，47% 的对话会话会发生上下文窗口溢出。

3 CachedAttention 设计

3.1概述

在本文中，我们提出了一种新的注意力机制，称为CachedAttention，它可以在多轮对话中重用KV缓存。 与传统注意力机制使用所有提示 token 进行预填充不同，CachedAttention 使用新会话轮次输入的新 token 和历史 token 的 KV 缓存进行预填充，如图 3 所示。 具体来说，当关联的会话会话处于非活动状态时，CachedAttention 将 KV 缓存保存在名为 AttentionStore 的 KV 缓存系统中，而不是像传统的注意力机制那样丢弃它们。 如果将来激活相同的会话会话，则会从 AttentionStore 中获取其 KV 缓存并重新用于推理。 这样，CachedAttention 只执行部分提示标记的预填充，即新一轮对话中输入的新标记，而不是预填充所有提示标记。 如图 LABEL:fig:cachedattention 所示，在执行第 3 转的推理时，$q_1$、$a_1$、$q_2$ 和 $a_2$ 的 KV 缓存被重复使用，只有 $q_3$ 需要预填充。 CachedAttention有效消除了历史token的重复计算开销，从而降低了预填充成本。

图5展示了CachedAttention的架构概览。 它维护了一个分层的KV缓存系统，即AttentionStore，具有高效的KV缓存访问、放置和截断技术，以解决第2.4节中提到的挑战。

图5： CachedAttention的系统架构。

对于挑战 1，为了减少 KV 缓存从 AttentionStore 加载到 HBM 的开销，CachedAttention 利用分层预加载方案将 KV 缓存加载与推理计算重叠。 为了减少从 HBM 到主机内存的 KV 缓存保存开销，CachedAttention 利用异步保存方案将保存与推理计算重叠。 （§3.2)。

针对挑战2和3，为了扩大缓存KV缓存的可用存储空间，CachedAttention在AttentionStore中采用了多层经济高效的存储介质，即主机内存和磁盘。 为了减少访问慢速磁盘对推理性能的影响，我们提出了一种调度程序感知的获取方案，该方案利用作业调度程序的提示来预取要从磁盘访问的 KV 缓存到主机内存。 同时，为了有效地利用有限的主机内存空间，我们提出了一种调度程序感知的驱逐方案，该方案可以识别最不有价值的 KV 缓存并将它们驱逐到磁盘或从缓存系统中驱逐（§3.3)。

对于挑战4，为了解决由于上下文窗口溢出而导致CachedAttention中保存的KV缓存失效的问题，我们利用位置编码解耦截断方案来保存不嵌入位置编码的KV缓存，从而支持直接在KV缓存上进行截断。 加载 KV 缓存时，CachedAttention 会将新的位置编码重新嵌入到 KV 缓存中（§3.4)。

3.2 重叠KV缓存访问

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

3.2.1 从内存分层预加载到 HBM

(a) Baseline: KV cache loading without concurrent operations.

(b) Layer-wise pre-loading without buffer.

(c) Layer-wise pre-loading with buffer.

图6： 逐层KV缓存预加载。 蓝色块表示每个 Transformer 层的执行。 红色块表示每个Transformer层的KV缓存加载。

(a) Layer-wise pre-loading with imperfect overlapping.

(b) Perfect pre-loading with a customized larger buffer.

图7： 逐层KV缓存预加载。

CachedAttention 将 KV 缓存从主机内存加载到 HBM，导致较高的数据访问开销。 如图6(a)所示，访问过程处于推理执行的关键路径，因为GPU必须依赖KV缓存来执行推理计算。 随着 KV 缓存大小的增加，这种开销变得更加显着，如 2.4 节中所述。 为了消除这种开销，CachedAttention 采用分层预加载方案来减轻影响。 主要思想是将 KV 缓存的加载与会话新输入 Token 的预填充计算重叠。 特别是，大语言模型由多个 Transformer 层链接起来，每个 Transformer 层都有自己的 KV 缓存。 当GPU执行一个层时，后续层所需的KV缓存可以同时从主机内存加载。 这样，当 GPU 开始计算某个层的 self-attention 时，该层对应的 KV 缓存就已经在 HBM 执行缓冲区中了。

图6(b)展示了分层预加载方案如何将KV缓存获取时间与计算时间重叠，为简单起见，使用3层模型。 在开始第1层的计算之前，必须首先在HBM中准备好该层的KV缓存。 读取流首先发出 KV 缓存加载操作，将第 1 层的 KV 缓存读入 HBM 执行缓冲区。 然后执行流开始计算第 1 层。 当执行流正在计算一层时，读取流会同时加载下一层的 KV 缓存，从而使加载与计算重叠。 然而，最后一个作业和当前作业的第一层之间仍然存在间隙，因为只有当 HBM 执行缓冲区可用时（即最后一个作业完成）才能开始加载。 为了进一步缩小最后一个作业与当前作业第一层之间的差距，CachedAttention 保留了 HBM 读取缓冲区来消除差距。 具体来说，如图6(c)所示，有了读缓冲区，读流就不必等待上一个作业释放执行缓冲区。 读流可以在最后一个作业运行时启动预加载。

然而，如果KV缓存加载时间长于预填充计算时间，则预加载可能无法与计算完全重叠。 如图7(a)所示，各层的计算之间存在多个间隙，因为每层的KV缓存获取时间超过了每层的计算时间，导致不完美的重叠。 通过采用定制的更大的预加载缓冲区可以进一步最小化开销。 通过更大的缓冲区，可以更早地启动预加载。 例如，如图7(b)所示，使用更大的缓冲区，允许预加载为更多层预加载KV缓存，从而可以重叠层与层之间的间隙。 设$T_{load}$、$T_{pref}$、$L_{hist}$和$L_{new}$表示词符的KV缓存的访问时间，词符的预填充时间词符，分别是会话中历史 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}).$

3.2.2 从 HBM 异步保存到内存

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

(a) Baseline: KV cache saving without concurrent operations.

(b) Asynchronous KV cache saving with overlapping.

图8： 异步KV缓存保存。

具体来说，KV缓存在预填充和解码阶段的生成速度是不同的。 预填充阶段同时处理 Token ，从而在有限的时间内生成大量的 KV 缓存。 相比之下，解码阶段一次生成一个词符的KV缓存。 如图8(b)所示，对于预填充阶段，由于每个self-attention操作都会产生大量的KV缓存，因此写入流会逐层保留KV缓存。 预填充阶段产生的KV缓存可以与解码阶段重叠。 对于解码阶段，由于KV缓存是迭代产生的，写流在解码的同时逐层写回KV缓存。 为了避免解码完成后缓存的 KV 未完全写回而卡住，我们还保留了一个 HBM 写入缓冲区来覆盖这种情况，类似于 KV 缓存预取中使用的读取缓冲区。 未完成的KV缓存会暂时移至写缓冲区，以避免阻塞下一个作业的执行。

3.3 分层KV缓存放置

CachedAttention 利用主机内存和磁盘来扩展 KV 缓存存储的可用空间。 主机内存（即 DRAM）的访问速度远高于磁盘（即 SSD）（数十 GB/s vs 10 GB/s）。 数 GB/秒）。 如果要访问的KV缓存总是在主机内存而不是磁盘中找到，那么KV缓存的访问性能将是最优的。 为了实现这一点，CachedAttention 应用了调度程序感知的获取方案，将 KV 缓存从磁盘预取到主机内存，确保以最佳速度访问 KV 缓存（§3.3.1)，并使用调度程序-感知驱逐方案，将合适的 KV 缓存从主机内存驱逐到磁盘（§3.3.2)。

3.3.1 调度程序感知从磁盘获取到内存

由于磁盘的容量比主机内存大得多（数十 TB 与 10 TB）。 几百GB），大多数KV缓存都保留在磁盘中用于CachedAttention。 由于会话请求随机到达，其对应的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从主机内存逐出到磁盘，以确保主机内存缓冲区的持续可用性。 下一小节将介绍从主机内存到磁盘的逐出过程。

图9： 调度程序感知的 KV 缓存获取和驱逐。

3.3.2 调度程序感知从内存到磁盘的逐出

当主机内存中的可用空间耗尽时，我们需要将一些KV缓存从主机内存中驱逐到磁盘上。 同时，如果磁盘已满，我们还需要将一些存储在磁盘中的KV缓存移出系统。 因此，仔细选择要驱逐的合适的 KV 缓存候选者以实现高缓存命中率非常重要。

与现有的缓存驱逐策略不同，如最近最少使用（LRU）[46]、先进先出（FIFO）[8]及其变体仅依赖于 KV 缓存的历史访问信息，CachedAttention 提出了一种调度程序感知的驱逐方案，可以利用 KV 缓存的未来访问信息来实现更高的缓存命中率。 作业调度程序中的作业队列使我们有机会实现这一目标。 具体来说，CachedAttention 在作业队列中维护一个前瞻驱逐窗口。 Look-ahead eviction window的最大长度由AttentionStore的总存储容量决定。 假设磁盘中的总可用容量为$C_{disk}$。 前瞻逐出窗口长度为$(C_{mem}+C_{disk})/S_{kv}$。 当 CachedAttention 尝试从 AttentionStore 中逐出一项时，如果在前瞻逐出窗口中找到要逐出的项，则该项将被豁免。 当 CachedAttention 将一项从主机内存驱逐到磁盘时，位于前瞻驱逐窗口尾部的该项具有更高的优先级被驱逐。 请注意，一项对应于与会话会话关联的所有 KV 缓存，这是 CachedAttention 中的最小驱逐和获取粒度。 这是因为同一个会话会话中的KV缓存要么全部被使用，要么完全没有被使用。

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

3.4 解耦的KV缓存截断

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

当历史token超过上下文窗口的限制时，大语言模型服务引擎一般会进行词符截断[30]。 如图10所示，上下文窗口大小为4K。 一旦上下文窗口溢出，大语言模型服务引擎就会切断提示符的前 2K 个标记。 此次截断对之前的大语言模型服务引擎没有影响，因为无论截断如何，它们都会根据输入提示重新计算KV缓存。 但截断使得CachedAttention中存储的KV缓存失效，显着降低了CachedAttention的效率。 这是由于 KV 缓存中嵌入了位置编码。 在提示符上执行词符截断时，每个词符的位置都会改变。 无法修改 KV 缓存中嵌入的位置编码以匹配提示中标记的位置，从而导致 KV 缓存无效。

为了解决这个问题，CachedAttention通过解耦位置编码，使得截断后的KV缓存仍然有效。 CachedAttention 需要与相对位置编码 (RPE) [42, 54, 44] 配合使用。 与绝对位置编码（APE）将位置编码添加到输入不同，RPE 直接将位置编码嵌入查询（Q）和密钥（K）向量中，如图11所示。 大量研究表明，RPE 允许大语言模型从比 APE [47, 11, 60] 更长的数据序列中学习。 因此，RPE在现代大语言模型中得到广泛应用，例如LLaMA [44]、T5 [55]、Falcon [32]、 Mistral [19]、Mixtral [20] 和 Transformer-XL [7]。 通过简单地移动在 RPE 中嵌入位置编码之前缓存 KV 的时间（如图11所示），CachedAttention 可以将不嵌入位置编码的 KV 存储在 AttentionStore 中。 当在 CachedAttention 中重用 KV 时，KV 会嵌入新的位置编码，并进一步用于以下推理。

图12提供了CachedAttention如何支持KV缓存截断的示例。 CachedAttention 存储不带位置编码的 KV 缓存。 在需要截断 KV 缓存的情况下，大语言模型引擎会检索截断的 KV 缓存（即 KV [0:1536]）并将其加载到 HBM。 新的位置编码随后被应用到 KV 缓存。

请注意，CachedAttention 还允许选择性保留某些 KV 缓存以进行压缩，例如，具有重要分数 [54] 的初始标记或重要标记 [12, 61, 24] ，进一步提高大语言模型的生成质量。 具体来说，给定的KV缓存压缩技术本质上提供了一种在提示中创建词符丢弃列表（TDL）的方法，例如丢弃最不重要的 Token 。 CachedAttention 直接遵守 TDL，丢弃 AttentionStore 中与 TDL 关联的 KV 缓存。 AttentionStore丢弃选定的KV缓存后，在推理执行时将剪枝后的KV缓存传递给推理引擎。

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

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

Figure 13: Cache hit rate.

Figure 14: Time to first token.

Figure 15: Prefill throughput.

Figure 16: GPU time.

4绩效评估

4.1 实验设置

测试平台。 我们所有的实验均在 4 个 NVIDIA A100 GPU 上进行，每个 GPU 具有 80GB HBM。 该系统配备128GB DRAM和10TB SSD。 GPU 通过 PCIe Gen 4 连接到主机。

我们在 Pytorch 和 Python 中实现了 CachedAttention。 主机内存和磁盘以块的形式进行管理，以提高存储利用率，类似于[21]。 我们的内部存储分配器根据需要分配和取消分配存储块。 模型执行器方面，CachedAttention基于Transformers [51]集成了LLaMA [44]、Falcon [32]等流行大语言模型的实现t2>。 专用 CUDA 流用于在 GPU 和主机内存之间移动数据，将计算与主动交换重叠。 单独的 IO 线程在主机内存和磁盘之间迁移数据，与 KV 缓存迁移重叠执行。 连续批处理[57]是通过实验启用的。

型号。 实验评估了使用 65B [44] 的开源 LLaMA-1 模型，使用 13B、70B 和 Falcon 40B [ 的 LLaMA-2 模型 [45]。 32]。 中间激活使用 FP16，与之前的系统[57, 48]保持一致。 我们还实现了具有 32K 上下文窗口的 Mistral-7B [19]。 除非另有说明，LLaMA-13B 在两个 GPU 上运行，处理 24 个批次，而 LLaMA-65B、LLaMA-70B 和 Falcon-40B 在四个 GPU 上运行，每个 GPU 处理 24 个批次。

工作负载。 工作负载是从 ShareGPT 数据集 [33, 62] 集成的。 由于数据集中没有可用的公共请求到达时间戳，因此我们遵循先前的工作[21, 52]，根据不同到达率的泊松分布生成请求到达时间。 我们根据泊松分布($\lambda$ = 1.0）设置每秒到达的不同会话数。 实验中使用了 9K 会话会话。

基线。 我们将 CachedAttention (CA) 与重新计算 (RE) 进行比较。 RE 仅保留对话会话的历史标记。 它在提供会话后丢弃 KV 缓存，并且在会话会话不活动时不保留 KV 缓存。 当与特定会话关联的对话再次激活时，RE 利用该会话中的历史 Token 来重新计算其 KV 缓存。 当历史token超过上下文窗口限制时，RE将应用词符截断，与一般大语言模型服务[30]相同。 为了简单起见，词符截断率设置为0.5，意味着当发生溢出时，系统将丢弃最早的一半 Token 。

4.2 端到端性能

在端到端实验中，我们使用来自 ShareGPT [33] 的 9K 对话，这些对话的平均轮数为 5.75。 因此总的对话轮数约为52K。 我们使用前 10K 轮对话来预热 AttentionStore，然后评估接下来 42K 轮的性能。

缓存命中率。 我们首先展示 CA 中 AttentionStore 的缓存命中率，因为其他性能指标与其密切相关。 图16显示了各种大语言模型的总KV缓存命中率，包括DRAM和磁盘命中率。 CA 对 LLaMa-13B、LLaMA-65B、LLaMA-70B 和 Falcon-40B 的命中率分别约为 86%、71%、89% 和 90%。 相比之下，我们观察到 LLaMA-65B 的命中率相对较低。 出现这种差异的原因是 LLaMA-65B 需要更大的存储空间来保存 KV 缓存。 在相同的可用存储空间的情况下，CA 为 LLaMA-65B 容纳的会话较少，从而限制了命中率。 具体来说，LLaMA-65B 的 KV 缓存中的每个词符需要 2.5MB 的空间。 LLaMA-13B 需要 0.78MB。 由于使用 GQA 因子为 8 和 16 的组查询注意力，LLaMA-70B 和 Falcon-40B 每个词符仅需要 0.31MB 和 0.12MB 的空间。

到第一个词符（TTFT）的时间了。 TTFT 是服务 [31, 2] 的大语言模型中服务质量的重要指标。 它指示用户输入提示后多快开始看到大语言模型的输出。 如图16所示，CA使LLaMA-13B、LLaMA-65B、LLaMA-70B和Falcon-40B的TTFT分别显着降低了85%、61%、87%和86%，与 RE 相比。 这是因为CA消除了预填充阶段生成历史token的KV缓存的大量重复计算。 当缓存命中时，CA的TTFT仅依赖于新会话轮次中新输入的 Token 的数量。

预填充吞吐量。 预填充吞吐量是评估处理提示速度的指标。 图16显示了测量的预填充吞吐量。 我们观察到 CA 为 LLaMA-13B、LLaMA- 提供了 6.8$\times$、2.6$\times$、7.8$\times$ 和 7.2$\times$ 的显着加速与 RE 相比，分别为 65B、LLaMA-70B 和 Falcon-40B。 CA对预填充吞吐量的提高来自于预填充时间的减少。 CA仅预填充新对话的新输入。 此外，CA可以通过分层预加载优化从AttentionStore加载和重用历史KV缓存。 历史 KV 缓存加载与新输入 Token 的预填充同时发生。

Figure 17: Inference cost.

Figure 18: Recomputation v.s. CachedAttention.

GPU 时间。 图 16 显示了完成工作负载中所有推理作业的端到端 GPU 时间。 我们观察到 CA 对于 LLaMA-13B、LLaMA- 实现了 4.0$\times$、1.9$\times$、3.3$\times$ 和 3.4$\times$ 的加速分别与 RE 相比，分别为 65B、LLaMA-70B 和 Falcon-40B。 CA的性能提升来自两个方面，一是缓解历史token的KV缓存重新计算，二是缓解上下文溢出后重新计算KV缓存。 对于重新预填充的缓解，CA有效地将KV缓存保存在AttentionStore中，并在需要历史 Token 时加载它。 另一方面，一旦作业完成，RE 就会丢弃 KV 缓存，需要对每个作业重新进行预填充以重新生成 KV 缓存。 为了减轻上下文流后 KV 缓存的重新计算，RE 应用词符截断，由于 KV 缓存中嵌入的位置编码，每次截断都会使 KV 缓存失效。 这会提示 RE 根据截断的历史 Token 重新计算 KV 缓存。 相比之下，CA将位置信息与KV缓存解耦，允许直接截断KV缓存。 这种方法避免了 RE 需要的 KV 缓存的重新计算。 请注意，CachedAttention 不仅有助于最大限度地减少预填充时间，还有助于减少解码时间[57]。 具体来说，在连续批处理下，每个新到达的作业必须完成预填充才能加入其他解码作业。 此过程会阻止解码作业的执行，从而导致解码时间延长。 然而，CachedAttention 通过最小化新到达作业的预填充时间来缓解这个问题。 预填充时间的减少反映了 RE 中观察到的解码时间的改进，如图 16 所示。

推理成本。 我们根据AWS EC2实例的按需价格[4, 5]来评估资源成本，即每个A100 GPU为5美元/小时，DRAM为0.0088美元/小时/GB，DRAM为0.000082美元/小时/ GB 用于 SSD。 图18显示了完成工作负载的不同方法的总成本。 与 RE 相比，CA 为 LLaMA-13B、LLaMA-65B、LLaMA-70B 和 Falcon-40B 实现了显着的成本节省，分别达到 70%、43%、66% 和 68%。 这些成本节省主要源于 GPU 时间的减少，因为 CA 有效减少了历史 Token 的冗余预填充以及上下文溢出期间的重新计算成本，如图 16 所示。 CA 采用经济高效的存储介质（包括主机内存和磁盘）来在非活动会话会话期间缓存 KV 缓存。 对于 LLaMA-13B、LLaMA-65B、LLaMA-70B 和 Falcon-40B，主机内存和磁盘的存储成本分别占 CA 总成本的 16.4%、9.0%、9.0% 和 9.0%。

4.3消融研究

4.3.1 重新计算v.s. 缓存注意力

我们研究了不同方法在不同历史和新词符比率下的预填充性能。 不同的方法在 LLaMA-13B 的 A100 GPU 上以 16 的批量大小预填充相同的 1K Token 。 RE 计算所有 token 的 KV 缓存，而 CA 从 AttentionStore 加载历史 token 的 KV 缓存，并部分预填充新的输入 token。 例如，设置600/400表示CA加载600个 Token 的KV缓存并计算400个 Token 的KV缓存。 如图18所示，CA 在所有测试设置中始终优于 RE。 随着新输入 Token 的百分比减少（从 500 到 100），这一优势变得更加明显，如每个条组的中间条所示。 尽管CA的KV缓存加载时间随着历史 Token 的百分比（从500到900）逐渐增加，但分层预加载方案有效地消除了该加载时间，如每个条组的第三个条所示。 请注意，当 KV 缓存加载时间超过预填充时间（例如，设置 900/100 的第二条）时，CA 可以通过启用读缓冲区来隐藏 KV 缓存加载时间。 读缓冲区的影响将在下一小节中评估。

4.3.2 重叠KV缓存访问

Figure 19: CA with no pre-loading v.s. CA pre-loading with various buffer sizes.

Figure 20: Performance impact of using write overlap.

本小节评估了所提出的用于加载和保存 KV 缓存的重叠访问技术的有效性。 使用的模型是具有单GPU的LLaMA-13B，批量大小设置为16。

逐层KV缓存预加载。 我们将历史 Token 的长度设置为1K，将新输入的 Token 的长度设置为100，以研究lay-wise预加载方案的有效性。 图20中的第一条，即NO-PL，显示了没有预加载方案的预填充时间，包括两部分：KV缓存加载时间和新输入 Token 的计算时间。 图20中的以下条显示了分层预加载方案具有不同大小的读缓冲区时的预填充时间。 为了清楚起见，我们使用层数来表示缓冲区大小，例如，PL-B0表示无读缓冲区，PF-B5表示5层KV缓存的读缓冲区大小。 我们观察到，尽管没有读缓冲区，即 PL-B0，但与 NO-PL 相比，预加载方案将预填充时间减少了 35%。 PF-B15完美重叠了KV缓存加载时间，与NO-PL相比，预填充时间减少了61%。

异步KV缓存保存。 我们设置不同的提示长度从1K到1.6K，解码步骤数为20，以研究异步保存方案的有效性。 如图20所示，保存时间随着提示长度的增加而增加，并且随着要保存的KV缓存的大小的增加而增加。 为了减少保存开销，CachedAttention采用异步保存方案，允许KV缓存保存与推理执行重叠，从而将整体执行时间减少13％至15％。

4.3.3 调度程序感知的获取和驱逐

(a) Impact on hit rate.

(b) Impact on GPU time.

图21： 各种存储设置下驱逐算法的比较。

我们评估了 CachedAttention 中调度程序感知的获取和驱逐在提高缓存命中率方面的有效性。 我们比较了各种存储配置中 CA 的总体缓存命中率、DRAM 命中率和磁盘命中率以及现有驱逐策略（包括 LRU 和 FIFO）。

如图21(a)所示，对于128GB DRAM和2TB SSD的128G/2T配置，CA在整体缓存命中率上分别优于LRU和FIFO 27%和31% 。 随着SSD容量的增加（128G/10T），CA的命中率达到了86%，超越了LRU（58%）和FIFO（48%）。 CA 实现了较高的总体命中率，因为 CA 知道未来的 KV 缓存访问信息，以避免驱逐将来将使用的 KV 缓存。 更高的命中率转化为更少的 GPU 时间，如图 21(b) 所示，其中 CA 实现了高达 2.7$\times$ 的加速。 从命中率细分来看，对于128G/2T的配置，LRU和FIFO仅实现了0.5%和0.5%（太小无法显示）的DRAM命中率，剩下的12.4%和9.0%是磁盘命中率。 即使对于 128G/10T 的较大容量，整体命中率分别增加到 58% 和 48%，LRU 和 FIFO 仍然表现出有限的 DRAM 命中率，分别约为 0.6% 和 0.5%。 这是因为LRU和FIFO缺乏对未来KV缓存信息的感知，无法将KV缓存从磁盘预取到主机内存，从而限制了它们提高DRAM命中率的能力。 相比之下，CA 的缓存命中率高达 86%，由于其调度程序感知策略，超过 99.6% 的命中发生在 DRAM 中。

4.3.4 解耦KV缓存截断性能

(a) Impact on hit rate.

(b) Impact on GPU time.

图22： 上下文溢出影响。

当上下文窗口超过其限制时，CachedAttention会直接截断KV缓存，从而避免了重新计算的需要并减少了开销。 我们评估了 CA 中管理上下文溢出的方法的有效性。 具体来说，我们比较了一种基线方法溢出（OF），该方法在 KV 缓存中嵌入位置编码，导致 AttentionStore 中的 KV 缓存失效。 OF 将上下文溢出归咎于重新计算来管理。 本实验使用128GB DRAM和10TB SSD。 从图22(a)可以看出，OF与CA相比，LLaMA-13B、LLaMA-65B、LLaMA-70B的命中率分别下降了17.6%、41.5%、18.1%、18.4%。和Falcon-40B，分别。 这种下降的原因是，如果应用 OF，上下文溢出的每个实例都需要上下文截断，从而使 AttentionStore 中的 KV 缓存失效。 命中率的降低随后会转化为更长的 GPU 时间，如图 22(b) 所示。

CachedAttention 保证了上下文溢出时系统中保存的 KV 缓存的有效性，并保证更高的命中率和减少 GPU 时间。 由于 2K 上下文窗口有限，LLaMA-65B 的 OF 命中率较低。 在服务第一轮会话后，会话很容易达到上下文窗口限制，从而导致关联的 KV 缓存失效。 随后，同一会话中的后续会话将面临 KV 缓存未命中的情况。

4.3.5 解耦位置编码的准确性

为了保持存储的 KV 缓存的有效性，CachedAttention 将位置编码与 KV 缓存解耦，并在重用存储的 KV 缓存时嵌入新的位置编码，如 3.4 节中所示。 我们利用广泛使用的基准评估了不同方案（包括 CA、词符截断（TT）和朴素 KV 缓存截断（NKVT））对困惑度（PPL）和大语言模型准确性的影响。 在历史 Token 数量超过上下文窗口限制的情况下，TT 会删除历史 Token 并重新计算剩余 Token 的 KV 缓存，NKVT 直接丢弃与位置编码相关的 KV 缓存并使用截断的 KV 缓存。

PPL。 PPL 是用于评估模型生成标记[54, 13]质量的指标。 较低的 PPL 值表明模型能够更好地预测文本并表现出对语言的更好理解。 表 1 显示使用数据集 WikiText-2 [14]、C4 [ 在 CA、TT 和 NKVT 设置中 LLaMA-7B 和 LLaMA-13B 的 PPL 比较35] 和 PTB [26]。 TT 通过在上下文窗口溢出后重新计算 KV 缓存来始终实现较低的 PPL。 CA 还通过在截断后将新的位置编码合并到 KV 缓存中来保持与 TT 相当的低 PPL（差异为 )。 相比之下，NKVT 由于其 KV 缓存内位置编码的耦合而表现出较高的 PPL ($>{10}^3$)。 直接截断 KV 缓存会扰乱耦合的位置信息，导致模型无法维持较低的 PPL。

准确度。 为了分析模型在截断后回答问题的准确性，我们使用 MMLU [17]、LongEval [22, 54] 和 PIQA [ 6] 基准。 具体来说，我们首先输入一个长文本来模拟历史输入的溢出以触发截断操作，然后将基准测试中的问题附加为新的输入。 如表2所示，CA 和 TT 都提供了较高的可比精度。 TT通过支付上下文窗口溢出的重新计算成本来实现高精度，而我们的CA避免了这种成本并且仍然保持高精度。 相比之下，NKVT 的精度比 CA 和 TT 低得多，因为 KV 缓存截断后耦合的位置编码是错误编码的，这会导致对新输入的更多干扰。

表1： 不同方法的 PPL 比较。

Dataset	Model	CA	TT	NKVT
WikiText-2	LLaMA-7B	5.47	5.48	2198.7
	LLaMA-13B	4.91	4.90	1647.7
PTB	LLaMA-7B	8.48	8.49	2543.5
	LLaMA-13B	7.61	7.60	1865.8
C4	LLaMA-7B	6.96	6.98	2343.5
	LLaMA-13B	6.44	6.45	1745.6

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

Benchmark	Model	CA	TT	NKVT
MMLU	LLaMA-7B	43.7%	43.4%	21.8%
	LLaMA-13B	52.3%	53.2%	29.6%
LongEval	LLaMA-7B	66.0%	65.9%	12.0%
	LLaMA-13B	68.0%	68.0%	14.0%
PIQA	LLaMA-7B	77.1%	77.2%	48.9%
	LLaMA-13B	80.5%	80.4%	50.2%

4.3.6 缓存容量要求

(a) Impact on hit rate.

(b) Impact on throughput.

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

在本小节中，我们将研究 CachedAttention 需要多少缓存容量才能实现显着的缓存命中率。 所需的缓存容量与大语言模型服务系统单位时间内服务的不同会话会话的最大数量（记为DSpUT)有关。 DSpUT值越大，系统单位时间内处理的不同会话越多，需要的KV缓存存储空间也越大。 而且，由于最大上下文窗口的限制，一个保存会话所需的最大KV缓存容量是固定的，即等于最大上下文窗口的长度乘以每个词符的KV大小，记为CCpS。 因此单位时间所需的最大缓存容量为$\text{CCpUT}=\text{DSpUT}\ast \text{CCpS}$。 在CachedAttention中，每个会话的KV缓存都有一个TTL（time to live），表示它自上次访问以来的最大保存时间。 TTL设置为上述的单位时间。 通过将CachedAttention的缓存容量配置为CCpUT，在不考虑新到达会话的情况下，我们可以实现100%的缓存命中率。 不过，在实际场景中，为了达到较高的缓存命中率，我们不需要配置这么大的容量，因为缓存项的热度不同。

为了弄清楚所需的缓存容量($RCC$)和$CCpUT$之间的关系，我们评估了不同$RCC$比率下的缓存命中率和解码吞吐量。 $CCpUT$。 在实验中，我们将TTL设置为一小时。 如图23(a)所示，当比率$RCC/CCpUT$为0.1时，我们实现了51%的缓存命中率。 当比率$RCC/CCpUT$为0.25时，我们实现了98%的缓存命中率。 随着命中率达到峰值，吞吐量也达到峰值，如图23(b)所示。

4.3.7 缓存存储介质的影响

(a) Impact of caching storage mediums on hit rates.

(b) Impact of caching storage mediums on GPU time.

图24： 各种缓存配置下的性能。

一些现有的工作[63,18,28]仅使用HBM空间来缓存多轮会话的KV缓存。 我们在这里比较了 HBM 上的机制缓存 KV 与 DRAM 和 SSD 上的 CachedAttention 缓存 KV 的性能。 在实验中，我们将HBM缓存的大小配置为10GB，DRAM的大小为128GB，SSD的大小为10TB。 图24(a)展示了不同机制的缓存命中率和推理性能。 由于HBM容量有限，所有模型的仅HBM缓存方法的命中率接近0%。 将 HBM​​ 与 DRAM 结合使用，型号 LLaMA-13B、LLaMA-65B、LLaMA-70B 和 Falcon-40B 的缓存命中率分别提高到 3.4%、1.7%、7.7% 和 19.1%。 相比之下，通过SSD进一步扩展缓存容量，CachedAttention将LLaMA-13B、LLaMA-65B、LLaMA-70B和Falcon-40B模型的缓存命中率分别提高到86%、71%、89%和90%，分别。 与仅 HBM/HBM+DRAM 策略相比，CachedAttention 具有更高的命中率，显着提高了推理性能，如图24(b)所示。

(a) Hit rates.

(b) TTFT.

(c) Prefill TP.

(d) GPU time.

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

4.3.8 会话到达率的影响

我们从命中率、TTFT、预填充吞吐量和 GPU 时间方面评估会话到达率如何影响 CachedAttention 的性能。 在实验中，我们将会话到达率配置为 0.5/s 到 2.0/s。 我们像以前一样在预热后使用 128GB/10TB 存储测试 LLaMA-13B。 如图25(a)所示，随着会话到达率从0.5增加到2，缓存命中率略有下降，从82%下降到77%。 如果会话到达率较低，则特定时间窗口将包含较少的会话，这表明 4.3.6 节中讨论的 DSpUT 较小。 低会话到达率时的DSpUT小于高会话到达率时的DSpUT。 在命中率相同的情况下，较低的 DSpUT 需要较少的缓存存储空间。 换句话说，在给定相同的缓存存储的情况下，较低的DSpUT会导致较高的命中率，如图25(a)中观察到的结果。 因此，由于KV缓存命中较少，更多的会话需要重新计算KV缓存，导致TTFT从0.122s增加到0.154s，如图25(b)所示。 由于 KV 缓存命中率略有下降，图 25(c) 中显示的预填充吞吐量也出现下降，从 858K/s 降至 681K/s。 图25(d)所示的GPU时间从6.25H升级到7.01H。 总体而言，不同的会话到达率对有效性的影响很小。 请注意，即使会话到达率有所增加，KV 缓存命中率仍然很高，这意味着 CachedAttention 仍然可以有效地利用 KV 缓存，始终比重新计算提供显着的性能改进。

5相关工作

KV 缓存管理。 在单轮会话中，KV缓存被广泛用于提高解码阶段[48,57,63,9,41,50]的性能。 为了减少 HBM 上 KV 缓存的存储开销，现有工作对 KV 缓存采用量化和压缩技术[54,12,61,24]。 为了减少碎片带来的内存浪费，vLLM [21]从虚拟内存中汲取灵感，允许KV缓存使用细粒度的非连续内存。 这些技术与 CachedAttention 正交，后者专注于多轮对话。

LMDeploy [18]是一个大语言模型推理框架，将多轮对话的KV缓存缓存在HBM上。 RadixAttention [63]、ChunkAttention [56] 和 Pensieve [58] 是与 CachedAttention 同时开发的推理技术。 RadixAttention 和 ChunkAttention 优化了共享提示前缀的推理任务。 具有相同提示前缀的任务共享相同的KV缓存，以减少KV计算。 Pensieve 利用 GPU 和 CPU 内存来存储多轮对话的 KV 缓存。 与所有这些工作不同的是，CachedAttention 利用更慢但更大的存储层次结构来保存 KV 缓存，以实现高缓存命中率（如 4.3.7 节所示），并专注于设计系统技术来解决以下挑战：卸载到较慢的介质。

推理参数卸载。 FlexGen [39] 将模型权重和 KV 缓存卸载到 DRAM 和磁盘，以支持大语言模型的离线推理。 DeepSpeed Inference [3, 37, 36] 将模型权重卸载到 DRAM 和磁盘并按需获取。 Lina [23]将大语言模型中不常用的专家权重卸载到主机内存中，以提高内存效率。 PowerInfer [40] 和大语言模型 in a flash[2] 利用 FFN 计算中的稀疏性 [25, 27] 来卸载大部分主机内存或磁盘的非活动权重可减少内存使用和计算。 FastServe [52] 将 KV 缓存调度到主机内存以优化作业完成时间。 相比之下，CachedAttention 利用 KV 缓存卸载来减少多轮对话中的重新计算开销。

6结论

本文提出了 CachedAttention，这是一种新的注意力机制，它可以在同一会话的任何后续回合中重用 KV 缓存，从而显着减少大语言模型中 KV 缓存的重新计算开销。 为了提高CachedAttention的效率，我们设计了重叠KV缓存访问、分层KV缓存放置和位置编码解耦KV缓存截断方案。 大量实验结果表明，对于多轮对话，CachedAttention 可将 TTFT 显着降低高达 87%，并将提示预填充吞吐量提高 7.8$\times$。 它将端到端推理成本降低高达 70%。

致谢

这项工作得到了华为云、新加坡国立大学高级研究与技术创新中心 (ARTIC) 的部分支持，资助金额为 FCT-RP1 A-0008129-00-00，以及新加坡教育部资助金额 A-0008143- 00-00 和 A-0008024-00-00。