Efficient Memory Management for Large Language Model Serving with PagedAttention

* 精简概括

*.1 What

提出了：

算法 PagedAttention，是一种关注算法
vLLM：基于 PagedAttention 算法，是一个 LLM 服务系统，实现了以下两个目标：
（1）在 KV 缓存内存方面几乎没有浪费
（2）在请求之间和内部灵活共享 KV 缓存以进一步减少内存使用

vLLM 在相同的延迟水平下提高了流行的 LLMs 的吞吐量

*.2 Why

现有系统存在问题，因为每个请求的键值缓存（KV 缓存）内存庞大且动态增长和收缩。当管理不当时，这种内存可能会因碎片化和冗余复制而被浪费，从而限制了批处理大小。

*.3 How

提出了 PagedAttention，这是一种受操作系统解决内存碎片化和共享问题的启发的关注算法

PagedAttention 将请求的 KV 缓存划分为块，每个块可以包含固定数量的标记的注意力键和值。在 PagedAttention中，KV 缓存的块不一定存储在连续的空间中。

这种设计通过使用相对较小的块并根据需要分配它们来减轻内部碎片化。此外，它消除了外部碎片化，因为所有块具有相同的大小。最后，它允许在块的粒度上进行内存共享，跨与同一请求相关的不同序列，甚至跨不同请求。

构建了 vLLM，这是一个基于 PagedAttention 的高吞吐量分布式 LLM 服务引擎，实现了 KV 缓存内存的近零浪费。

vLLM 使用块级内存管理和抢占式请求调度。

vLLM 采用了一个集中式调度器来协调分布式 GPU 工作节点的执行

KV 缓存管理器有效地以分页方式管理 KV 缓存：KV 缓存管理器通过由集中式调度器发送的指令来管理 GPU 工作节点上的物理 KV 缓存内存。

KV 块管理器还维护块表，记录每个请求的逻辑和物理 KV 块之间的映射关系。

0 Abstract

将大型语言模型（LLMs）进行高吞吐量的服务，需要一次性批量处理足够多的请求。

提出了：

算法 PagedAttention，是一种关注算法
vLLM：基于 PagedAttention 算法，是一个 LLM 服务系统，实现了以下两个目标：
（1）在 KV 缓存内存方面几乎没有浪费
（2）在请求之间和内部灵活共享 KV 缓存以进一步减少内存使用

vLLM 在相同的延迟水平下提高了流行的 LLMs 的吞吐量

1 Introduction

大型语言模型（LLMs）的出现已经启用了新的应用程序，运行这些应用程序非常昂贵，需要大量的硬件加速器，如 GPU。鉴于这些高昂的成本，提高 LLM 服务系统的吞吐量，从而降低每个请求的成本，变得越来越重要.

LLMs 的核心是一个自回归 Transformer 模型。这个模型基于输入（prompt）和它到目前为止生成的输出标记序列，逐个生成单词（tokens）。对于每个请求，这个昂贵的过程会重复，直到模型输出终止标记。

这个顺序生成过程使工作负载受限于内存，未充分利用 GPU 的计算能力，并限制了服务的吞吐量

通过将多个请求进行批处理可以提高吞吐量。在部署的 LLM 的内存分布的三部分中，模型权重是恒定的，激活值只占用 GPU 内存的一小部分，因此 KV 缓存的管理方式对于确定最大批量大小至关重要。

观察到现有的 LLM 服务系统在有效管理 KV 缓存内存方面存在不足。这主要是因为它们将请求的 KV 缓存存储在连续的内存空间中。KV 缓存具有独特的特点：随着模型生成新标记，它会动态增长和收缩，其生命周期和长度事先是未知的。这些特点使得现有系统的方法在两个方面显着低效：

现有系统受到内部和外部内存碎片化的影响
现有系统无法利用内存共享的机会

提出了 PagedAttention，这是一种受操作系统解决内存碎片化和共享问题的启发的关注算法

PagedAttention 将请求的 KV 缓存划分为块，每个块可以包含固定数量的标记的注意力键和值。在 PagedAttention中，KV 缓存的块不一定存储在连续的空间中。

构建了 vLLM，这是一个基于 PagedAttention 的高吞吐量分布式 LLM 服务引擎，实现了 KV 缓存内存的近零浪费。

vLLM 使用块级内存管理和抢占式请求调度。

2 Background

2.1 基于 Transformer 的 LLM

基于 Transformer 的语言模型最重要的组件是其自注意力层（self-attention）

一些计算过程，对于输入的隐藏状态序列，做线性变换，获取查询 q，键 k 和值 v 向量。
之后计算注意力分数，再计算输出

2.2 LLM 服务与自回归生成 LLM Service & Autoregressive Generation

LLM 只能逐个采样和生成新的 token，每个新标记的生成过程依赖于序列中所有先前的标记，具体来说是它们的键和值向量。

给定一个请求 prompt，LLM 服务中的生成计算可以分解为两个阶段：

The prompt phase：以整个用户 prompt 作为输入，计算第一个新 token 的概率 𝑃(𝑥𝑛+1

𝑥1, . . . , 𝑥𝑛)。在此过程中，还生成键向量 𝑘1，…，𝑘𝑛 和值向量 𝑣1，…，𝑣𝑛。由于提示标记 𝑥1，…，𝑥𝑛 都是已知的，因此提示阶段的计算可以使用矩阵乘法操作进行并行化。因此，这个阶段可以有效地利用 GPU 中固有的并行性。

The autoregressive generation phase：生成其余的新 token. 在第 𝑡 次迭代中，模型将一个标记 𝑥𝑛+𝑡 作为输入，并使用键向量 𝑘1，…，𝑘𝑛+𝑡 和值向量 𝑣1，…，𝑣𝑛+𝑡 计算概率 𝑃(𝑥𝑛+𝑡+1

𝑥1, . . . , 𝑥𝑛+𝑡)。由于数据依赖性，不同迭代的计算无法并行化，通常使用矩阵-向量乘法，效率较低。因此，这个阶段严重未充分利用了 GPU 计算，并且成为了内存限制，导致大部分单个请求的延迟。

2.3 LLM 的批处理 Batching Techniques for LLMs

通过对多个请求进行批处理可以提高 LLMs 的计算利用率。因为这些请求共享相同的模型权重，所以将移动模型权重的开销分摊到批处理中的请求中.

将请求批处理到 LLM 服务中有两个非常重要的问题：

这些请求可能在不同的时间到达，导致排队延迟。
这些请求的输入和输出长度可能差异很大，一个简单的批处理技术会填充请求的输入和输出以使它们的长度相等，浪费 GPU 计算和内存资源。

提出了细粒度的批处理机制，如细胞批处理和迭代级调度。在每个迭代之后，已完成的请求会从批处理中移除，然后添加新的请求。此外，通过使用特殊的 GPU 内核，这些技术消除了填充输入和输出的需求。

3 Memory Challenges in LLM Serving

可以一起批处理的请求数量仍然受到 GPU 内存容量的限制，特别是用于存储 KV 缓存的空间.

需要解决内存管理中的以下挑战：

大的 KV 缓存：随着请求数量的增加，KV 缓存的大小迅速增长。
复杂的解码算法：LLM 服务为用户提供了一系列解码算法可供选择，每种算法对内存管理复杂性都有不同的影响。KV 缓存共享的程度取决于所采用的具体解码算法。
未知输入和输出长度的调度：随着请求的输出长度在解码过程中增长，其 KV 缓存所需的内存也会扩展，可能会耗尽用于新请求或正在生成现有提示的内存。系统需要做出调度决策，例如从 GPU 内存中删除或交换出一些请求的 KV 缓存。

先前的 LLM 服务系统会根据请求的最大可能序列长度静态分配一个内存块，而不考虑请求的实际输入或最终输出长度。

有三个主要的内存浪费来源：为未来标记保留的插槽、由于过度为潜在的最大序列长度提供而导致的内部碎片以及来自内存分配器（如伙伴分配器）的外部碎片

4 Method

vLLM 的架构：

vLLM 采用了一个集中式调度器来协调分布式 GPU 工作节点的执行

KV 缓存管理器有效地以分页方式管理 KV 缓存：KV 缓存管理器通过由集中式调度器发送的指令来管理 GPU 工作节点上的物理 KV 缓存内存

4.1 PagedAttention

PagedAttention 允许在非连续的内存空间中存储连续的键和值

PagedAttention 将每个序列的 KV 缓存分成 KV blocks。每个块包含了一定数量的标记的键和值向量，记 block size 为 B.

在注意力计算过程中，PagedAttention 内核单独识别并提取不同的 KV 块

4.2 KV Cache Manager

vLLM 的内存管理器的核心思想类似于操作系统中的虚拟内存.

一个请求的 KV 缓存被表示为一系列逻辑 KV 块，随着新标记及其 KV 缓存的生成，从左到右填充。

KV 块管理器还维护块表，记录每个请求的逻辑和物理 KV 块之间的映射关系。每个块表条目记录了逻辑块的相应物理块以及填充位置的数量。

4.3 Decoding with PagedAttention and vLLM

vLLM 在生成更多标记及其 KV 缓存时会动态分配新的物理块给逻辑块.

由于所有块都是从左到右填充的，并且仅在所有先前的块都已满时才分配新的物理块，vLLM 将请求的所有内存浪费限制在一个块内.

4.4 Application to Other Decoding Scenarios

在许多成功的 LLM 应用中，LLM 服务必须提供更复杂的解码场景，这些场景具有复杂的访问模式和更多的内存共享机会.

4.4.1 并行抽样 Parallel sampling

在基于 LLM 的程序助手中，LLM 为单个输入提示生成多个抽样输出；用户可以从不同的候选中选择一个喜欢的输出.

在并行抽样中，一个请求包含多个样本，它们共享相同的输入提示，允许共享提示的 KV 缓存。通过其 PagedAttention 和分页内存管理，vLLM 可以轻松实现这种共享并节省内存.

vLLM 允许在多个输出样本之间共享用于存储提示的大部分空间，除了最后一个逻辑块，它由写时复制机制管理。

通过在多个样本之间共享物理块，内存使用可以大大减少，尤其是对于长输入提示而言。

4.4.2 束搜索 Beam search

束搜索被广泛用于从 LLM 中解码最有可能的输出序列.

以前的 LLM 服务系统需要频繁复制束候选者之间的 KV 缓存，vLLM 的物理块共享显著减少了这种频繁的内存复制开销。

在 vLLM 中，不同束候选者的大部分块可以共享。仅当新生成的标记位于旧的共享块内时，如在并行解码中，才会应用写时复制机制。这只涉及复制一个数据块。

4.4.3 共享前缀 Shared prefix

许多用户提示共享一个前缀，因此 LLM 服务提供者可以提前存储前缀的 KV 缓存，以减少在前缀上的冗余计算.

在 vLLM 中，LLM 服务提供者可以通过为一组预定义的共享前缀保留一组物理块来方便地实现这一点，具有共享前缀的用户输入提示可以简单地将其逻辑块映射到缓存的物理块（最后一个块标记为写时复制）。

4.4.4 混合解码方法 Mixed decoding mathods

vLLM 支持同时处理具有不同解码偏好的请求，因为 vLLM 通过一个通用的映射层将逻辑块翻译成物理块，从而隐藏了不同序列之间复杂的内存共享。

4.5 Scheduling and Preemption

vLLM 采用了先到先服务（FCFS）

当 vLLM 需要抢占请求时，它确保 earliest 请求先得到服务，latest 请求先被抢占.

随着请求数量和它们的输出增加，vLLM 可能会耗尽 GPU 的物理块来存储新生成的 KV 缓存，在这种情况下，vLLM 需要回答两个经典问题：（1）应该驱逐哪些块？（2）如果需要再次使用，如何恢复已驱逐的块？

4.5.1 应该驱逐哪些块？

因为序列的所有块都一起被访问，因此实施了一个全有或全无的淘汰策略，即要么淘汰一个序列的所有块，要么不淘汰。

在一个请求中的多个序列被作为一个序列组进行批处理。同一序列组中的序列总是一起抢占或重新调度，因为这些序列之间可能存在内存共享。

4.5.2 如果需要再次使用，如何恢复已驱逐的块？

两种方法：

交换 Swapping：

将被淘汰的块复制到 CPU 内存中。

一旦一个请求完成，它的块就会从内存中释放出来，而一个被抢占的序列的块则被重新引入，以继续处理该序列。

重新计算 Recomputation：

被抢占的序列被重新调度后，简单地重新计算 KV 缓存。

4.6 Distributed Execution

许多 LLM 的参数大小超过了单个 GPU 的容量。因此，有必要将它们分布在多个 GPU 上，并以模型并行的方式执行它们。

即使采用模型并行执行，每个模型分片仍然处理相同的输入标记集，因此需要相同位置的 KV 缓存。因此，vLLM 在集中式调度器内具有一个单一的 KV 缓存管理器。不同的 GPU 工作节点共享管理器，以及从逻辑块到物理块的映射。

GPU 工作节点无需在内存管理上进行同步，因为它们只需要在每个解码迭代的开始时接收所有的内存管理信息以及步骤输入。

# 概念学习

#.1 KV Cache

KV cache 是键值缓存的缩写，是一种用于存储和检索键值对数据的缓存系统。

在 KV cache 中，每个数据项都与一个唯一的键相关联，通过这个键可以快速访问和检索数据。这种缓存系统通常用于加速对数据的读取操作，特别是在需要频繁访问相同数据的情况下，可以大大提高性能。

在本论文中，KV cache 是指用于存储和管理大型语言模型（LLMs）服务所需的键值对数据的缓存，其中每个请求都需要访问这个缓存以提高服务的吞吐量。文中提到了问题是由于 KV cache 内存的动态增长和收缩以及管理不当而导致内存浪费和碎片化，因此需要一种新的算法（PagedAttention）和系统（vLLM）来解决这个问题。

#.2 Beam search

束搜索的主要目标是在给定一个输入序列后，找到生成的输出序列中最有可能的若干候选序列。

预测：在每个生成步骤，模型考虑所有可能的下一个标记（词汇表中的词汇），并为每个标记计算其生成的概率。
剪枝：束搜索会保留前 𝑘 个概率最高的候选序列，称为”束”。这个 𝑘 是束宽参数，决定了每一步保留的候选数量。
扩展：对于每个候选序列，将其扩展成多个可能的候选序列，每个都考虑了一个可能的下一个标记。然后，计算每个扩展后的候选序列的概率。
选择：从扩展后的候选序列中选择前 𝑘 个概率最高的序列作为下一步的候选。

束搜索的优点在于它能够显著降低搜索空间的复杂度，因为它只保留了前 𝑘 个最有可能的序列

Paper Reading vLLM