标题：面向 PagedAttention 服务的大型语言模型的高效内存管理
1.摘要
大型语言模型（LLM）的高吞吐量服务需要一次处理足够多的请求。
然而，现有的系统很难做到这一点，因为每个请求的键值缓存（KV 缓存）内存都很大，并且动态地增长和收缩。当管理效率低下时，碎片和冗余复制会严重浪费此内存，从而限制批处理大小。
为了解决这个问题，我们提出了 PagedAttention，这个注意力算法的灵感来自经典的虚拟内存和操作系统中的分页技术。在此基础上，我们构建了 vLLM，这是一个 LLM 服务系统，它实现了
（1）KV 缓存内存的近零浪费，以及
（2）在请求内部和请求之间灵活共享 KV 缓存，以进一步减少内存使用。
我们的评估表明，与 FasterTransformer 和 Orca 等最先进的系统相比，vLLM 在相同延迟水平下将流行 LLM 的吞吐量提高了 2-4 倍。对于更长的序列、更大的模型和更复杂的解码算法，改进更明显。
2.引言
很多云公司争相提供大模型的托管服务，但是托管大模型十分昂贵，需要很多的硬件加速器。所以为大模型服务减少每个请求的代价很重要。
自回归模型生成过程中，一次一个 token，存在顺序性。这种顺序生成过程使工作负载受到内存限制，无法充分利用GPU的计算能力，并限制了服务吞吐量
通过将多个请求放在一起，可以提高吞吐量，但是内存受不了，所以需要高效管理内存。
在图一(左)中，13 B 参数的 LLM 在 A100 中，
65% 的内存用于存储参数，服务中静态不变。
接近 30% 的内存用于存储动态的请求状态。
对于 Transformers 来说，这些状态由 key 和 value 张量组成，通常称为 KV cache，他表示来自之前 token 的上下文，去生成新的输出 token。
还有一小部分用于其他数据，包括激活值。
由于模型的参数不变，而激活值只占小部分内存，那么 KV cache 的内存管理在决定最大批次大小的时候变成更重要了。
作者观察到现有的大模型服务系统在 KV cache 内存管理上比较短缺，这主要是因为它们将请求的 KV cache 存储在连续的内存空间中，因为大多数深度学习框架要求张量存储在连续的内存中。然而，KV cache 具有独特的特性(与传统张量不同)：随着模型生成新 token，它会随着时间的推移动态增长和收缩，并且其生命周期和长度事先未知。
这些特征使得现有系统的方法在两个方面显著低效：
- 现有的系统有内部和外部碎片问题(图二)。为了存储 KV cache，需要预先分配一个请求的最大长度(比如 2048 个 token)的内存块
- 现有的系统不能利用存储器共享的能力。请求由多个序列组成，这些序列可以部分共享它们的 KV cache。然而，存储器共享在现有系统中是不可能的，因为序列的 KV cache 存储在单独的连续空间中
为了避免这些限制，作者们提出了 PagedAttention，灵感来自于操作系统的内存碎片和共享的解决方案：分页虚拟内存
PagedAttention 将请求的 KV cache 划分为多个块，每个块可以包含固定数量的 token 的 key 和 value。
在 PagedAttention 中，KV cache 的块不一定存储在连续空间中。因此，我们可以像在操作系统的虚拟内存中那样以更灵活的方式管理KV缓存：
可以将 块 视为页面，token 视为字节，请求 视为 进程。
这种设计通过使用相对较小的块并按需分配它们来消除内部碎片。
此外，它消除了外部碎片，因为所有块都具有相同的大小。
最后，它支持以 块 为粒度的内存共享，跨与同一请求相关联的不同序列，甚至跨不同请求。
3.背景
1.基于 Transformer 的 LLM
2.LLM 服务 和 自回归生成
一旦训练完成，LLM 通常部署为一个条件生成服务。
LLM 一般一个接一个地采样和生成新 token，并且每个新 token 的生成过程取决于该序列中的所有先前 token ，特别是它们的 key 和 value 向量。
在这个顺序生成过程中，现有 token 的 key 和 value 向量通常被缓存以生成未来的令牌，称为 KV cache。请注意，一个 token 的 KV cache 取决于其之前的所有 token。这意味着出现在序列中不同位置的相同 token 的 KV cache 将不同。
给定请求提示，LLM服务中的生成计算可以分解为两个阶段：
- 提示阶段。该阶段充分利用了 GPU 的并行性
- 自回归生成阶段。该阶段由于数据依赖性(当前 token 依赖之前的 token)，不同迭代的计算不能并行化，并且通常使用效率较低的矩阵向量乘法。因此，此阶段严重未充分利用GPU计算，并且成为内存受限阶段，这是单个请求的大部分延迟的原因。
3.LLM 的批次技术
可以通过批处理多个请求来提高服务 LLM 中的计算利用率。
由于请求共享相同的模型权重，因此移动权重的开销在批处理中的请求之间分摊，并且当批处理大小足够大时，可能会被计算开销所淹没。
但是，由于两个原因，对 LLM 服务的请求进行批处理并不容易。
首先，请求可能在不同的时间到达。简单的批处理策略会使较早的请求等待较晚的请求，或者将传入的请求延迟到较早的请求完成，从而导致显著的排队延迟。
第二，请求的输入和输出长度可能有很大的不同（图11）。直接的批处理技术将填充请求的输入和输出以均衡其长度，从而浪费GPU计算和内存。
4.内存挑战
服务系统的吞吐量是受内存限制的。克服这种内存限制需要解决内存管理中的以下挑战
-
巨大的 KV cache
大量的请求会导致 KV cache 增长很快。
-
复杂的解码算法
-
未知输入和输出长度的调度
1.现有系统的内存管理
图3 示出了两个请求：具有 2048 个最大可能序列长度的请求 A 和具有 512 个最大可能序列长度的请求 B。
现有系统中的块预分配方案有三个主要的内存浪费来源：
为未来 token 预留的插槽，
由于潜在最大序列长度的过度供应而导致的内部碎片，
以及来自内存分配器（如伙伴分配器）的外部碎片。
外部碎片永远不会用于生成的 token，这在服务请求之前是已知的。内部碎片也保持未使用状态，但这仅在请求完成采样后才实现。
虽然最终会使用预留的内存，但在整个请求持续时间内预留该空间，特别是当预留空间很大时，会占用原本可用于处理其他请求的空间。
我们在图 2 中可视化了我们实验中内存浪费的平均百分比，揭示了以前系统中的实际有效内存可以低至 20.4%。
5.方法
vLLM 的架构如图 4 所示。
vLLM 采用集中式调度器来协调分布式 GPU 的执行。
KV cache 管理器以分页方式有效地管理 KV cache，由 PagedAttention 启用。
具体地，
KV cache 管理器通过集中式调度器发送的指令来管理 GPU 工作器上的物理 KV cache。
接下来，我们在 §1 中描述 PagedAttention 算法。
有了这些，我们在 §2 中展示了 KV cache 管理器的设计，在 §3 中展示了它如何促进 PagedAttention。
然后，我们展示了这种设计如何促进各种解码方法的有效内存管理（§4），并处理可变长度的输入和输出序列（§5）。
最后，我们展示了vLLM的系统设计如何在分布式环境中工作
1.PagedAttention
图5.PagedAttention算法的图示，其中注意键和值向量作为非连续块存储在内存中。
为了解决内存挑战，作者们引入了 PagedAttention。
PagedAttention 允许去存储连续的 key 和 value 在非连续的内存中。具体的，PagedAttention 划分每个序列的 KV cache 为 KV 块。每个块包含固定 token 数量的键和值的向量，表示为 KV 块大小(B)。表示 key block K j K_j Kj 和 value block V j V_j Vj。
注意力计算被转化成下面的按块计算(左为原来，右为新)
A i j A_{ij} Aij 是第 j 个 KV 块的注意力分数的行向量。
在注意力计算期间，PagedAttention 内核分别识别和获取不同的 KV 块。我们在图 5 中展示了一个 PagedAttention 的例子：
键和值向量分布在三个块上，这三个块在物理内存上是不连续的。
在每一次，内核将查询 token（“forth”）的查询向量 q i q_i qi 乘以块中的 key 向量 K j K_j Kj （例如，块 0 的“Four score and seven”的 key 向量）来计算注意力分数 A i j A_{ij} Aij，并且稍后将注意力分数 A i j A_{ij} Aij 与块中的 value 向量 V j V_j Vj 相乘以导出最终的注意力输出 o i o_i oi
总之，PagedAttention 算法允许 KV 块存储在非连续物理存储器中，这使得 vLLM 中的分页存储器管理更加灵活。
2.KV Cache 管理器
vLLM 的内存管理器背后的关键思想类似于操作系统中的虚拟内存。
操作系统将内存划分为固定大小的页面，并将用户程序的逻辑页面映射到物理页面。连续的逻辑页可以对应于非连续的物理内存页，允许用户程序访问内存，就好像它是连续的一样。此外，物理存储器空间不需要预先完全保留，从而使 OS 能够根据需要动态地分配物理页面。
vLLM 使用虚拟内存背后的思想来管理 LLM 服务中的 KV cache。通过启用 PagedAttention ，我们将 KV cache 组织为固定大小的 KV 块，就像虚拟内存中的页面一样。
请求的 KV cache 表示为一系列逻辑 KV block，从左到右填充为新 token 及其 KV cache。最后一个 KV 块的未填充位置是为将来的生成保留的。
在 GPU worker 中，一个 block engine 分配一个 GPU DRAM 的连续块 并且 将其划分为物理 KV 块。
KV 块管理器 还维护 block tables——每个请求的逻辑和物理 KV 块之间的映射。
每个块表条目记录逻辑块的相应物理块和填充位置的数量。
分离逻辑和物理 KV 块允许 vLLM 动态地增长 KV 高速缓存存储器，而无需提前为所有位置保留它，这消除了现有系统中的大部分存储器浪费，如图2所示。
3.使用 PagedAttention 和 vLLM 解码
接下来，我们通过一个示例，如图 6 所示，演示 vLLM 如何在单个输入序列的解码过程中执行 PagedAttention 并管理内存：
-
与 OS 的虚拟内存一样，vLLM 最初不需要为最大可能生成的序列长度保留内存。相反，它只保留必要的 KV 块，以容纳在提示计算期间生成的 KV 缓存。
在这种情况下，提示符有 7 个 token，因此 vLLM 将前 2 个逻辑 KV 块（0和1）映射到 2 个物理 KV 块（分别为7和1）。
在预填充步骤中，vLLM 利用常规自注意算法。
vLLM 然后将前 4 个 token 的 KV cache 存储在逻辑块 0 中，并将随后的 3 个 token 存储在逻辑块 1 中。
剩余的卡槽被保留用于随后的自回归生成阶段。
-
在第一自回归解码步骤中，vLLM 在物理块 7 和 1 上利用 PagedAttention 算法生成新令牌。由于在最后一个逻辑块中仍有一个槽可用，因此新生成的 KV cache 存储在那里，并且 块表 的 #filled 记录被更新。
-
在第二解码步骤，由于最后一个逻辑块已满，vLLM 将新生成的 KV cache 存储在新的逻辑块中;
vLLM 为其分配新的物理块（物理块3），并将该映射存储在块表中。
4.应用其他的解码情景
-
并行采样
LLM为单个输入提示生成多个采样输出;用户可以从各种候选中选择最喜欢的输出。
在并行采样中，一个请求涉及多个共享同一个输入提示的样本，允许提示的 KV cache 也被共享
途径 PagedAttention 和 分页内存管理器，vLLM 能够轻易实现这个共享和节约内存。
如图 8 展示的一个两个输出的并行解码的例子。
由于两个输出共享相同的提示，所以在 prompt 阶段，我们只为 prompt 状态的一个副本保留空间;两个序列的提示逻辑快被映射到了相同的物理块上：块 0 和 块 1。由于一个单个的物理块能被映射到多个逻辑块，所以作者们引入了一个 reference counts(引用计数) 给每个物理块。
在这种情况下，物理块7和1的参考计数都是2
在生成阶段，两个输出采样不同的输出 token 并且需要单独存储 KV cache。vLLM 实现一个 copy-on-write(写时复制)的机制 在需要由多个序列修改的物理块的块粒度上，类似于 OS 虚拟存储器中的写时复制技术。
如图 8，当样本 A1 需要写他最后一个逻辑块(逻辑块 1)的时候，vLLM 识别出对应的物理块（物理块1）的引用计数大于1; 他分配一个新的物理快(物理块 3),命令块引擎去从物理块 1 复制信息，并且将引用计数器减到 1.
接下来，当样本 A2 写入物理块 1 的时候，引用计数器已经是 1 了；因此 A2 直接写入他新生成的 KV cache 到 物理块 1
-
束搜索
-
共享前缀
一般来说，LLM 用户提供一个长度任务描述(也称之为 system prompt)，通常与实际任务拼接到一起，形成任务提示。LLM 的输出基于全部的提示。如图 10。此外共享前缀能够进一步调优，途径一个提示工程，去改善下游任务的准确率。
对于这类应用，许多用户提示共享一个前缀，因此 LLM 服务提供商可以提前存储前缀的 KV 缓存，以减少前缀上的冗余计算。在 vLLM 中，这可以通过由 LLM 服务提供商为一组预定义的共享前缀保留一组物理块来方便地实现，如 OS 如何跨进程处理共享库。
具有共享前缀的用户输入提示符可以简单地将其逻辑块映射到缓存的物理块（最后一个块标记为写时复制）。提示阶段计算只需要在用户的任务输入上执行。
-
混合解码方法
vLLM 有助于同时处理具有不同解码偏好的请求，而现有系统无法有效地做到这一点。这是因为 vLLM 通过将逻辑块转换为物理块的公共映射层隐藏了不同序列之间的复杂内存共享。LLM 及其执行内核只看到每个序列的物理块 ID 列表，不需要处理跨序列的共享模式。与现有的系统相比，这种方法扩大了具有不同采样要求的请求的处理机会，最终提高了系统的整体吞吐量。
5.调度和抢占
在 vLLM 中，我们对所有请求采用先来先服务（FCFS）调度策略，确保公平性并防止饥饿。当 vLLM 需要抢占请求时，它确保最早到达的请求首先得到服务，最新的请求首先被抢占。
LLM服务面临着一个独特的挑战：LLM的输入提示在长度上可能会有很大的变化，并且由此产生的输出长度是先验未知的，取决于输入提示和模型。随着请求及其输出数量的增长，vLLM 可能会耗尽 GPU 的物理块 来存储新生成的 KV cache.
这引发了两个问题：
(1)它应该驱逐哪些块？
(2)如果再次需要，如何恢复被驱逐的块？
通常情况下，驱逐策略使用优先级来预测哪个块将在未来被访问得最远，并驱逐该块。
由于在我们的情况下，我们知道一个序列的所有块都是一起访问的，我们实现了一个全有或全无的驱逐策略，即，或者驱逐序列的所有块，或者不驱逐序列的任何块。
为了回答第二个问题，即如何恢复被驱逐的块，我们考虑两种技术：
-
交换
这是大多数虚拟内存实现所使用的经典技术，它们将收回的页复制到磁盘上的交换空间。在本例中，我们将逐出的块复制到 CPU 内存中。如图 4 所示，除了 GPU 块分配器之外，vLLM 还包括 CPU 块分配器，用于管理交换到 CPU RAM 的物理块。当 vLLM 耗尽新令牌的可用物理块时，它会选择一组序列来收回并将其 KV 缓存传输到 CPU。一旦它抢占了一个序列并收回了它的块，vLLM 就会停止接受新的请求，直到所有被抢占的序列都完成为止。一旦请求完成，它的块就从内存中释放出来，被抢占序列的块就被带回，以继续处理该序列。
-
重新计算
在这种情况下，当被抢占的序列被重新调度时，我们只需重新计算 KV cache。请注意，重新计算等待时间可以显著低于原始等待时间，因为在解码时生成的令牌可以与原始用户提示连接作为新的提示–它们在所有位置的 KV cache 可以在一个提示阶段迭代中生成。
6.分布式执行
vLLM在分布式设置中是有效的，它支持广泛使用的Megatron-LM风格的张量模型并行策略。
6.实现
vLLM是一个端到端服务系统，具有 FastAPI 前端和基于 GPU 的推理引擎。
前端扩展了OpenAI API 接口，允许用户为每个请求自定义采样参数，例如 maximum sequence length 和 beam width k。vLLM 引擎是用 8.5K 行 Python 和 2K 行 C++/CUDA 代码编写的。我们使用 NCCL 在分布式 GPU 工作器之间进行张量通信。
7.评估
在本节中，我们将评估vLLM在各种工作负载下的性能。
1.实验设置
-
模型和服务配置
我们使用参数为 13B、66B 和 175B 的 OPT 模型和参数为 13B 的 LLaMA 模型进行评估。13B 和 66B 是 LLM 的流行尺寸，而 175B 是著名的 GPT-3 型号的尺寸。在我们的所有实验中，我们使用了 A2 实例和 Google Cloud Platform 上的NVIDIA A100 GPU。详细的型号大小和服务器配置如表 1 所示。
-
工作负载
我们基于 ShareGPT 和 Alpaca 数据集合成工作负载，这些数据集包含真实的LLM服务的输入和输出文本。
ShareGPT 数据集是用户与 ChatGPT 共享的对话的集合。
Alpaca数据集是由 GPT3.5 生成的具有自我指令的指令数据集。
我们对数据集进行标记化，并使用它们的输入和输出长度来合成客户端请求。如图 11 所示，与 Alpaca 数据集相比，ShareGPT 数据集的输入提示平均长 8.4 倍，输出平均长 5.8 倍，方差更高。
-
OPT 在 ShareGPT 和 Alpaca 数据集上的单个序列生成
-
OPT-13B 在 ShareGPT(2 reqs/s) 和 Alpaca(30 reqs/s) 追踪上的平均批次请求数
-
OPT-13B 在 Alpaca 数据集上的并行生成和束搜索
-
OPT 14B 在 Alpaca 追踪上的平均内存节省(共享 KV 块)
