大家好，今天想来介绍下当红推理框架vLLM的核心技术PagedAttention。 PagedAttention的设计灵感来自操作系统的虚拟内存分页管理技术。 vLLM的论文是在假设读者对这项分页管理技术非常熟悉的情况下，对PagedAttention进行介绍的，这对一些非计算机专业出身，或者对操作系统相关知识有所遗忘的读者来说并不友好。

因此，本文进行介绍时，会对照着操作系统的相关知识，和大家一起来看vLLM是如何“一步步”从传统方法进化到PagedAttention的，同时本文会尽量将抽象的显存优化知识通过图解的方式向大家说明 。

全文目录如下：

一、LLM推理的两阶段

二、为KV cache分配存储空间的传统方法

三、PagedAttention原理
3.1 操作系统的虚拟内存
3.2 PagedAttention

四、示例：PagedAttention在不同decoding场景下的运作流程和优势
4.1 Parallel Sampling
4.2 Beam Search

五、vLLM的调度与抢占
5.1 总原则
5.2 终止和恢复被抢占的请求

六、分布式管理

一、LLM推理的两阶段

一个常规的LLM推理过程通常分为两个阶段： prefill和decode 。 通常会使用KV cache技术加速推理 。

1.1 Prefill

预填充阶段。在这个阶段中，我们把整段prompt喂给模型做forward计算。如果采用KV cache技术，在这个阶段中我们会把prompt过 后得到的 保存在cache_k和cache_v中。 这样在对后面的token计算attention时，我们就不需要对前面的token重复计算 了，可以帮助我们节省推理时间。

在上面的图例中，我们假设prompt中含有3个token，prefill阶段结束后，这三个token相关的KV值都被装进了cache。

1.2 Decode

生成response的阶段 。在这个阶段中， 我们根据prompt的prefill结果，一个token一个token地生成response。

同样，如果采用了KV cache，则每走完一个decode过程，我们就把对应response token的KV值存入cache中，以便能加速计算。例如对于图中的t4，它与cache中t0~t3的KV值计算完attention后，就把自己的KV值也装进cache中。对t6也是同理。

由于Decode阶段的是逐一生成token的，因此它不能像prefill阶段那样能做大段prompt的并行计算，所以在LLM推理过程中，Decode阶段的耗时一般是更大的。

从上述过程中，我们可以发现使用KV cache做推理时的一些特点：

• 随着prompt数量变多和序列变长，KV cache也变大，对gpu显存造成压力
• 由于输出的序列长度无法预先知道，所以我们很难提前为KV cache量身定制存储空间

下图展示了一个13B的模型在A100 40GB的gpu上做推理时的显存占用分配（others表示forward过程中产生的activation的大小，这些activation你可以认为是转瞬即逝的，即用完则废，因此它们占据的显存不大），从这张图中我们可以直观感受到推理中KV cache对显存的占用。 因此，如何优化KV cache，节省显存，提高推理吞吐量，就成了LLM推理框架需要解决的重点问题。

二、为KV cache分配存储空间的常规方式

对于训练好的模型，一种常用的部署方式是将其打包成一个推理服务（server），它接收客户端发送来的请求（request），读取请求中的数据（prompt）来做推理。一个请求中可以只有1个prompt，也可以包含多个prompt。

在常规的推理框架中，当我们的服务接收到一条请求时，它会为这条请求中的prompts分配gpu显存空间，其中就包括对KV cache的分配。由于推理所生成的序列长度大小是无法事先预知的，所以大部分框架会按照(batch_size, max_seq_len)这样的固定尺寸，在gpu显存上预先为一条请求开辟一块连续的矩形存储空间。然而，这样的分配方法很容易引起“gpu显存利用不足”的问题，进而影响模型推理时的吞吐量。你可能觉得这个描述有点抽象，别着急，我们来具体看一个例子。

下图展示了一个常规的推理框架是如何为请求中的prompt在gpu显存上分配KV cache的。在本例中，我们假设一个请求只发送1条prompt（本例中共有3条请求）：

我们假设 max_seq_len = 8 ，所以当第1条请求(prompt1)过来时，我们的推理框架为它安排了(1, 8)大小的连续存储空间。

当第2条请求（prompt2）过来时，同样也需要1块(1, 8)大小的存储空间。但此时prompt1所在的位置上，只剩3个空格子了，所以它只能另起一行做存储。对prompt3也是同理。

仔细观察这3条prompt的KV cache排布，你是不是隐约觉得这种排布似乎没有充分利用起gpu的显存？ ：

• 浅色块 ：观察图中的浅色块，它是prefill阶段prompt的KV cache，是无论如何都会被使用的空间，它不存在浪费。

• 中色块 ：观察图中的中色块，它是decode阶段的KV cache，其中 表示序列生成的截止符。虽然这些中色块最终都会被我们用上，但是在decode阶段一个个token生成时，我们并不能预知哪些块会被最终用上。例如对于prompt2，当你生成when的时候，你无法知道下一个会生成 ，还是会生成别的词。 所以这些中色块都是一种“潜在的浪费”，我们称中色块的部分为预留碎片（reservation fragment）。

• 深色块 ：观察图中的深色块，它也是decode阶段的KV cache，但直到序列生成完毕，它都没有被用上。 由于这些深色块是预留的KV cache的一部分，所以我们称其为内部碎片（internal fragment）。

• 灰色块 ：观察图中的灰色块，它不是我们预留的KV cache的一部分，且最终也没有被用上， 我们称这些灰色块为外部碎片（external fragment） 。想象一下，此时新来了一条prompt4，它也要求显存中的8个格子作为KV cache。 此时你的显存上明明有9个空格子，但因为它们是不连续的碎片，所以无法被prompt4所使用。 这时prompt4的这条请求只好在队列中等待，直到gpu上有足够显存资源时再进行推理，这不就对模型推理的吞吐量造成显著影响了吗？

观察整个KV cache排布，你会发现它们的毛病在于太过“静态化”。 当你无法预知序列大小时，你为什么一定要死板地为每个序列预留KV cache空间呢？ 为什么不能做得更动态化一些，即“用多少占多少”呢？ 这样我们就能减少上述这些存储碎片，使得每一时刻推理服务能处理的请求更多，提高吞吐量，这就是vLLM在做的核心事情，我们先通过一张实验图来感受下vLLM在显存利用上的改进效果（VS 其它推理框架）：

不难发现，相比于别的推理框架，vLLM几乎能做到将显存完全打满。

读到这里，你可能会有以下疑问：

• 问题1：vLLM是通过什么技术，动态地为请求分配KV cache显存，提升显存利用率的？

• 问题2: 当采用动态分配显存的办法时，虽然明面上同一时刻能处理更多的prompt了，但因为没有为每个prompt预留充足的显存空间，如果在某一时刻整个显存被打满了，而此时所有的prompt都没做完推理，那该怎么办？

在后文的第三～四章，我们将回答问题1。第五章回答问题2。

三、PagedAttention原理

在本节中， 我们先来回答问题1：vLLM通过一种名为PagedAttention的技术，动态地为请求分配KV cache显存，提升显存利用率。

整体上来说，PagedAttention的设计灵感来自操作系统中虚拟内存的分页管理技术 。所以本节会先通过通俗易懂的方式，和大家一起快速回顾操作系统的虚拟内存技术，在这个过程中和大家一起具象化感受PagedAttention的设计思想。然后再来详细介绍PagedAttention的运作流程。

3.1 操作系统的虚拟内存

（1）不使用虚拟内存

我们知道程序运行时，会将代码、数据等内容存放在物理内存上。 在最原始的做法中（没有操作系统，例如单片机），程序直接对物理内存进行操作，决定使用它的哪些存储地址。

如果你只跑一个进程，那还好说。但如果需要运行多个进程时，麻烦就来了 ：由于我直接操作了物理内存地址，所以我在为自己的进程分配物理内存时，还要考虑别的进程是如何分配物理内存的（别人已经占用的我不能用）。这样不同进程间的耦合性太高了，给开发带来难度。

有没有一种办法，让各个进程间的开发能够相互独立呢？一种直觉的做法是：

• 给每个进程分配一个虚拟内存 。每个进程在开发和运行时，可以假设这个虚拟内存上只有自己在跑，这样它就能大胆操作。
• 虚拟内存负责统一规划代码、数据等如何在物理内存上最终落盘 。这个过程对所有进程来说都是透明的，进程无需操心

虚拟内存的核心思想可简化成下图：

（2）虚拟内存的分段管理

在分段式内存管理中，虚拟内存会尽量为每个进程在物理内存上找到一块连续的存储空间，让进程加载自己的全部代码、数据等内容 。我们来看一个具体的例子：

在这个例子中，3个进程的虚拟内存各自为它们在物理内存上映射了一块连续的存储空间。在某一时刻，我释放了进程2，同时想运行进程4。这时我尴尬地发现， 虽然物理内存上有640M的空间剩余，但因为是碎片化的，我的进程4无法加载进去 ，因此它只能等待（回想一下本文第二部分对传统KV cache显存分配的分析）。

在这个情况下，如果我硬要运行进程4，也是有办法的：我可以先把进程3从物理内存上 交换（swap） 到磁盘上，然后把进程4装进来，然后再把进程3从磁盘上加载回来。通过这种方法我 重新整合了碎片 ，让进程4能够运行。

但这种办法的显著缺点是： 如果进程3过大，同时内存到磁盘的带宽又不够，整个交换的过程就会非常卡顿。这就是分段式内存管理的缺陷。

这时，我自然而然会想到： 我为什么要给所有进程都预分配一个固定的存储块（段）呢？ 假设这个进程是一个浏览器，我难道会一下就用到这个进程里所有的功能吗？就不能进程运行到哪里，或者我想用哪个具体功能时，再加载这部分相关的内容去内存，以此让整个内存分配更加动态？

（3）虚拟内存的分页管理

我们可以将进程1、进程2想成是两本书。代码分布在书的不同page上。我们希望想读哪一页，就加载哪一页，而不是一下把两本书都加载进来。 同时，当我们不想读某些页的时候，我们也能根据页码将其清空。

现在，我们希望读进程1和进程2的page1，我们就将其加载到物理内存上。虚拟内存会帮我们做好映射，把来自不同进程的这两页分别加载到物理内存对应位置。

虚拟内存的分业管理技术总结起来就是：

• 将物理内存划分为固定大小的块，我们称每一块为页（page） 。从物理内存中模拟出来的虚拟内存也按相同的方式做划分

• 对于1个进程，我们不需要静态加载它的全部代码、数据等内容。我们想用哪部分，或者它当前跑到哪部分，我们就动态加载这部分到虚拟内存上，然后由虚拟内存帮我们做物理内存的映射。

• 对于1个进程，虽然它在物理内存上的存储不连续（可能分布在不同的page中），但它在自己的虚拟内存上是连续的。 通过模拟连续内存的方式，既解决了物理内存上的碎片问题，也方便了进程的开发和运行。

3.2 PagedAttention

（1）处理单个请求

现在，你已经知道虚拟内存分页管理的基本原理和优势，趁热打铁，我们马上来看以其为灵感的PagedAttention技术是如何操作的。我们还是从具体的例子讲起。

假设现在你向模型server发送一条请求，prompt为 Four score and seven years ago our ，你希望模型能做续写。PagedAttention的运作流程如下图：

在图中：

• 请求（request）可理解为操作系统中的一个进程

• 逻辑内存（logical KV blocks）可理解为操作系统中的虚拟内存，每个block类比于虚拟内存中的一个page。每个block的大小是固定的，在vLLM中默认大小为16，即可装16个token的K/V值

• 块表（block table）可理解为操作系统中的虚拟内存到物理内存的映射表

• 物理内存（physical KV blocks）可理解为操作系统中的物理内存，物理块在gpu显存上，每个block类比于虚拟内存中的一个page

图中带圈的序号表示操作步骤，我们就按这个顺序来看看。

(i) Prefill阶段

• 划分逻辑块 ：vLLM拿到这条prompt，先按照设定好的block大小B（本例中B=4），为prompt划分逻辑块（Logical KV blocks）。由于prompt中有7个token，所以vLLM用2个逻辑块（block 0， block 1）来装它们的KV值。其中，在逻辑块1中目前只装了"years", "ago", "hour"这3个token的KV值，有1个位置是空余的。这个位置就被称为保留位（reservation）

• 划分物理块 ：划分好逻辑块后，我们就可以将其映射到物理块中去了。物理块是实际存放KV值的地方。我们通过一张block table来记录逻辑块和物理块的映射关系，block table的主要内容包括：

• 逻辑块和物理块的映射关系（physical block number） ：例如逻辑块0对应物理块7
• 每个物理块上被填满的槽位（# filled） ：例如在prefill阶段，对物理块7，其4个槽位都被填满；对物理块1，其3个槽位被填满。

• 正常计算prompt的KV值，并通过划分好的关系填入物理块中。

（ii）Decode阶段-生成第1个词

• 使用KV cache计算attention，生成第1个词fathers 。不难发现，当我们计算时，我们使用的是逻辑块，即形式上这些token都是连续的。与此同时，vLLM后台会通过block table这个映射关系，帮我们从物理块上获取数据做实际计算。 通过这种方式，每个request都会认为自己在一个连续且充足的存储空间上操作，尽管物理上这些数据的存储并不是连续的。

• 基于新生成的词，更新逻辑块、物理块和block table 。对于block table，vLLM将它filled字段由3更新至4。

• 分配新的逻辑块和物理块 。当fathers更新进去后，逻辑块已装满。所以vLLM将开辟新的逻辑块2，并同时更新对应的block table和物理块。

（iii）Deocde阶段-生成第2个词

类比步骤（2）来进行。

（2）处理多个请求

有了（1）的解释，大家看懂这张图应该不难。通过多个请求（prompt）同时做推理的例子，大家可以更好感受到PagedAttention是如何通过动态存储KV cache的方式，来更充分利用gpu显存的。

四、PagedAttention在不同解码场景下的例子

通过前文的解释，我们已经基本掌握了PagedAttention的设计思想、运作流程。你可能隐隐能感受到它在显存管理上的“灵活性”，和减少碎片化显存的能力。 但可能你觉得还不够具象，所以在本节中，我们通过更具体的场景，再假设一下对PagedAttention优势的理解。

我们知道，根据实际需求，大模型的解码方式也比较复杂 ，例如：

• Parallel Sampling ：我给模型发送一个请求，希望它对prompt做续写，并给出三种不同的回答。我们管这个场景叫parallel sampling。在这个场景中，我们可以将prompt复制3次后拼接成1个batch喂给模型，让它做推理。但我们也需注意到，这种方式会产生prompt部分KV cache的重复存储。

• Beam Search ：束搜索，这是LLM常用的deocde策略之一，即在每个decode阶段，我不是只产生1个token，而是产生top k个token（这里k也被称为束宽）。top k个token必然对应着此刻的top k个序列。我把这top k个序列喂给模型，假设词表的大小为|V|，那么在下一时刻，我就要在k*|V|个候选者中再选出top k，以此类推。不难想象每一时刻我把top k序列喂给模型时，它们的前置token中有大量的KV cache是重复的。

• Shared prefix ：在某些大模型中，所有请求可能都会共享一个前置信息（比如system message: “假设你是一个有帮助的AI助手...."），这些前置信息没有必要重复存储KV cache

• 其余一般场景 ：在一些更通用的场景中，虽然两个prompt可能完全没有关系，但它们中某些KV cache却是可以共用的。例如两个prompt的相同位置（position）恰好出现了完全一样的序列，比如它们的结尾都是好想下班。假设这个相同序列已经存在于KV cache中，那也没有必要重复计算和存储了。

在下文里，我们会详细解释PagedAttention在Parallel Sampling和Beam Search场景上的优势。剩余两个场景读者可以自行做类比分析。

4.1 Parallel Sampling

下面说明在parallel sampling的场景下，vLLM（PagedAttention）是怎么做到节省显存的。

传统KV cache怎么做： 假设模型的max_seq_len = 2048。传统KV cache可能在显存中分配两块长度是2048的空间。由于prompt一致，这两块2048的空间中存在大量重复的KV cache。

vLLM怎么做：

假定我们发给模型1个request，这个request中包含2个prompt/sample，记为Sample A1和Sample A2，这两个prompt完全一致，都为 Four score and seven years ago our ，我们希望模型对这两个prompt分别做续写任务。

（1）首先，Prefill阶段，vLLM拿到Sample A1和Sample A2，根据其中的文字内容，为其分配逻辑块和物理块。

• 分配逻辑块 ：对于A1，vLLM为其分配逻辑块block0和block1；对于A2，vLLM为其分配逻辑块block0和block1。 需要注意的是，A1的逻辑块和A2的逻辑块是独立的（尽管它们都叫block0和block1） ，你可以将A1和A2视作操作系统中两个独立运行的进程。