vLLM推理框架｜用"内存分页术"榨干GPU，让KV缓存不再"爆仓"！

高效的批处理策略，使得更多的请求可以组成batch并行处理，但是batch组的请求数仍受到GPU内存的限制，如何的突破内存瓶颈，最大化batch中请求数量，是提高推理吞吐量的关键，本文主要围绕分页注意力高效管理KV cache缓存机制，介绍以下几个问题：

1. 1. 分页管理KV cache的基本流程
2. 2. 区别于传统注意力计算，分页注意力核心公式推导过程
3. 3. 在不同解码方式下，分页注意力是如何实现页共享，进行高效内存管理。

1，操作系统分页内存

1.1 早期内存管理的局限

• • 地址冲突

早期系统中进程直接使用物理内存，单个进程可以高枕无忧，但系统往往是多个进程并行，此时就存在不同的进程，可能读写相同的物理地址，造成地址冲突，进而引发数据覆盖或进程崩溃。

• • 内存碎片

为此引入了分段式虚拟内存管理，尽量为每个进程在物理内存上找到一块连续的存储空间，让进程加载自己的全部代码、数据等内容。但是在分段式内存管理中，虽然空闲内存大于新进程需求内存，但是空闲内存不连续，依然不能加载新进程。

即：不连续的内存空间，不能满足新进程的内存要求。

物理内存总大小：20 KB （假设地址范围 0x0000~0x4FFF)。

当前内存分配状态：进程A-E共5个请求，每个请求占4KB内存空间。

内存释放：进程B和进程D退出，释放其占用的内存，此时空闲内存为两段，总空闲8 KB，但不连续。

新进程F请求加载失败：进程F需要5 KB 的连续内存，虽然总空闲内存足够 (8 KB>5KB)，但两段空闲块均小于5KB （每段仅4 KB)，因此操作系统无法分配连续5 KB的空间，导致进程F加载失败。

1.2 分页内存管理

分页是把整个虚拟和物理内存空间，切成一段段固定大小的内存，也就是页(Page)。这样做的好处有两点：

• • 在虚拟地址中的连续地址可以映射到离散的物理地址中，防止了内存碎片的产生。
• • 如果内存空间不够，操作系统会把正在运行的进程中“最近没有被使用”的内存页面给换出(Swap Out) 至磁盘上，并释放掉该页面。一旦需要的时候，再从磁盘换入(Swap In) 该页面。

2，KV cache显存浪费问题

2.1 KV cache 显存分配

大模型推理系统倾向于将KV cache 存放到一块连续的内存中，且由于每个请求的输出不定，所以推理系统预先分配了一块连续的内存给每个请求，而不管请求的实际输入和输出token的长度。这就造成了三种主要内存的浪费。

• • 保留槽位，用于实际生成的新token。虽然最终被使用，但是可以提前被其他请求使用，提高吞吐量。
• • 内部碎片，超过实际产生token而提前过量分配的槽位。
• • 外部碎片，不同请求的长度不同，而内存分配而产生的内存气泡。

在上图中，请求A预留了2048个KV cache的槽位，请求B预留了512个KV cache 槽位，系统内存分配时产生的外部内存碎片。

再看请求A内部，实际请求A当前时刻，7个槽位已经用于prompt，1个用于槽位用于已经解码的token（“fathers”），1个槽位用于当前进行的token（“brought”），假设请求还会生成2个的token（“forth”和 “eos”），此时共2038个槽位浪费。

2.2 预分配KV cache 定量计算

KV cache 的大小随着请求量的增加而激增。

举例：对于 13B 的OPT模型，隐藏层维度h=5120，层数 layer=40，假设其生成token的最大数量为2048。

单个token的KV cache 显存占用为：2×5120×40×2=800K

总KV cache 显存占用为：采用预分配的方式，预留2048个KV缓存的槽位，共800×2048=1.6G，即使GPU所有显存都用于存储KV cache，也仅仅能够处理几十条的请求。

故，高效管理KV cache，减少内存占用成为推理加速的关键之一！

3, 分页注意力paged Attention

核心： 通过分页机制，实现KV cache可以被存放在不连续的GPU显存中。

page attention 核心组件：

• • 每条请求相当于操作系统中的一个进程，需要内存和计算资源。
• • 逻辑内存(Logical KV Blocks) 相当于操作系统中的虚拟内存，每个page包含一定数量的 插槽 Block。
• • 块表(Block Table)相当于操作系统中的虚拟内存到物理内存的映射表。
• • 物理内存 (Physical KV Blocks) 相当于操作系统中的物理内存，这里为GPU的 DRAM。

显存不足的时候的换入换出机制，即VLLM可以将部分请求的KV Cache 先swap out到其他内存上，以腾出空间完成其余请求。当有请求完成并释放显存后，再将那部分请求swap in回来继续推理。

3.1 分页注意力推理流程

逻辑内存(Logical KV Blocks) 中包含多个Block，一个Block相当于一页，每个Block中又包含一定数量的槽位，每个槽位用于存放一个token的所有KV cache缓存值。如上图中，此时逻辑内存中包含4个Block (0-3) 每个Block包含4个槽位，可以容纳4个token的KV缓存值。

假如：当前请求A，prompt有7个token，将要产生的token依次为"fathers”，“brough” 具体分页注意力计算过程如下：

prefill阶段：

使用传统的self-attention操作，计算提示词阶段7个token的KV cache，并存储到逻辑内存的Block 0和Block 1中，此时Block 0已满 #filled 的状态为4。而Block1还空闲一个槽位，将要用于存放解码阶段"fathers”所产生的KV cache，当前其#filled状态为3。

对于逻辑内存中的 Block 0 和Bolck 1，通过映射表，将对应的KV cache值存储到不连续的物理内存 Block 7和Block 1中。进而实现了在虚拟的逻辑内存中每个token 的内存是连续的，而在物理内存中却是不连续的。

decode 阶段：

此时使用page attention 利用"fathers"的KVcache和已经缓存prompt的KV值计算新的token (brought)，并将新计算“fathers”的KV值缓存到 Block 1中的最后一个槽位，将#filled 状态更新 (3->4)。

再继续使用page attention 生成新的token时，并将 (brough) 的KV 值放到新Block 2的第一个槽位中，更新其#filled状态为1。通过映射表其对应的物理内存为Block 3。依次类推直到最后一个token生成完毕。

3.2 分页注意力公式推导

分页注意力有别与传统的注意力，需要将当前解码token 和不同Block页内的token计算注意力，传统的注意力机制如下：

使用分页注意力，公式1需要改写为:

公式中Block中包含插槽的数量记为B，即每个页的大小，下标 j 代表Block的计数，i 代表当前token的计数。

对应的Key值为： ，对应的value值为： 。

注意K和V的下标是连续的，比如Block 2 中的第一个插槽的k值表示为： ，Block 2 中所有槽位的k值可以表示为：

举个例子：

当前第10个token为 “forth"，其query向量 ，分别和三个Block 0、1和2 进行注意力计算，得到 ，最终将其进行拼接输入到之后的前向计算中得到新的token。

4，多种解码策略下的page attention

4.1 并行采样

定义：并行采样是指LLM针对同一输入提示(prompt)，同时输出多个不同的结果，用户可以从这些候选结果中选择最符合需求的输出。

主要特点

• • 共享提示状态

多个输出序列共享同一份prompt的 KV cache，避免重复存储。如上图请求A1和请求A2有共同的提示词，所以A1和A2的逻辑内存中Block 0 和Block 1 共享物理内存Block 7 和 Block 1。因为物理内存中的同一个Block可以对应逻辑内存中多个不同的Block，所以采用引用计数（reference count），即此时物理内存中共享的 Block 7和Block 1 的引用计数都为2。

• • 按需复制 (Copy-on-Write)

请求A1和A2开始产生不同的token时，根据引用计数按需复制。即请求A产生新token（”fathers"）并放置在逻辑内存Block 1 的最后一个插槽中，当映射到物理内存Block 1时，发现其引用计数值为2，此时将Block 1复制到Block 3，并将“fathers”产生的KV cache放到Block 3的最后一个插槽中。当请求A2产生新token（“mothers”）时，映射到物理内存Block1中，此时其引用计数值为1（2->1）正常放入。

4.2 束搜索

定义：从初始prompt生成首字开始，保留前k个最可能的候选序列，对当前每个候选序列，生成所有可能的下一词（共 个候选，V为词表大小) 用大语言模型计算每个新序列的联合概率（即当前序列概率乘以下一词概率），仅保留全局概率最高的k个新序列，其余丢弃。

以图中虚线处为当前时刻，前后分为两部分，自开始到虚线处，共经历了4次迭代。束搜索大小为k=4，即在符合置信度阈值情况下，最多每次保留top4的候选，每个Beam candidate 在每次迭代时都会有一个Block，其中Block 0, Block 1和Block 3 被共享。比如Beam candidate 2 当前时刻的页路径为：Block 0 ->Block 1 ->Block 3 ->Block 7。

在生成第4个Block时候，候选 Beam candidate 0 和 5，不再符合阈值要求被剪枝，此时其占用的Block（Block2，4，8和5）被释放。在虚线之后又有新的token产生，Block 9 至Block 12 为生成新的候选。

4.3 共享前缀树

定义：用户通常会在请求中附加一个长系统提示（如任务指令、示例输入输出)，与具体任务输入拼接后形成完整的提示。

列如：[ 系统提示：翻译为英文 ] + [ 任务输入：今天天气很好 ] -- 生成结果

分页存储

• • 将系统提示词存储在只读页 (Read-Only Pages) ，所有请求共享
• • 用户输入和生成序列存储在可写页 (Writable Pages)， 每个请求独立。

5，总结

并行采样通过分页隔离实现高吞吐，适合独立请求批量处理。束搜索通过动态共享候选序列状态，平衡生成质量与内存开销。前缀树通过只读共享页最大化复用，适合多请求共享系统提示的场景。

vllm中分页注意力的核心价值：

• • 内存效率：分页机制避免KV Cache的预分配浪费。
• • 计算复用：共享公共部分（前缀候选）减少冗余计算。
• • 灵活性：适配不同解码策略 （并行、束搜索、前缀共享）

这些优化使得vLLM在长序列，高并发，高质量生成的场景下显著优于传统方法。

参考：

[1] Woosuk et al. Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180

历史文章：

大模型推理-极致化的批处理策略介绍

大模型推理- PD分离部署，势在必行！

大模型推理-高效推理必备KV cache

大模型训练-混合专家系统MoE

大模型训练-Nvidia GPU 互联技术全景图

大模型训练-流水线并行PP

大模型训练-张量并行TP