大模型推理引擎 vLLM(2)：PagedAttention论文学习以及原理解析

文章目录

[1 PagedAttention 论文解读](#1 PagedAttention 论文解读)
- [1.0 Abstract](#1.0 Abstract)
- [1.1 Introduction](#1.1 Introduction)
- [1.2 Background](#1.2 Background)
- [1.3 Memory Challenges in LLM Serving](#1.3 Memory Challenges in LLM Serving)
- - [1.3.1 KV Cache 占用迅速增长，极易耗尽显存](#1.3.1 KV Cache 占用迅速增长，极易耗尽显存)
  - [1.3.2 预分配导致内存碎片严重](#1.3.2 预分配导致内存碎片严重)
  - [1.3.3 KV Cache 难以共享，内存复用受限](#1.3.3 KV Cache 难以共享，内存复用受限)
- [1.4 Method](#1.4 Method)
- - [1.4.1 PagedAttention](#1.4.1 PagedAttention)
  - [1.4.2 KV Cache Manager](#1.4.2 KV Cache Manager)
  - [1.4.x 个人阶段总结](#1.4.x 个人阶段总结)
  - [1.4.3 DecodingwithPagedAttention and vLLM](#1.4.3 DecodingwithPagedAttention and vLLM)
  - [1.4.4 Application to Other Decoding Scenarios](#1.4.4 Application to Other Decoding Scenarios)
  - - [1.4.4.1 并行采样](#1.4.4.1 并行采样)
    - [1.4.4.2 束搜索（Beam Search）](#1.4.4.2 束搜索（Beam Search）)
    - [1.4.4.3 共享前缀（Shared Prefix）](#1.4.4.3 共享前缀（Shared Prefix）)
    - [1.4.4.4 混合解码（Mixed Decoding）](#1.4.4.4 混合解码（Mixed Decoding）)
- [1.5 调度和抢占](#1.5 调度和抢占)
- - [1.5.1 总原则](#1.5.1 总原则)
  - [1.5.2 终止和恢复被抢占的请求](#1.5.2 终止和恢复被抢占的请求)
- [1.6 分布式管理](#1.6 分布式管理)
[2 总结与整理](#2 总结与整理)
参考文献：

abstract
高占用、内存碎片和内存不能复用
按需分配不预留
划分固定大小block，block table映射，连续的KV保存在不连续的内存空间中，
为了解决 KV Cache 内存管理中的 高占用、内存碎片和内存不能复用 等问题。

借鉴操作系统中虚拟内存的分页技术，将KV缓存划分为多个block块，每个块保存固定数量token的KV数据，并通过block table将逻辑block和物理block进行映射。连续的逻辑页可以对应到非连续的物理内存页，并且物理的内存空间并不需要提前去预留，而是按需分配，

PagedAttention就是允许将连续的KV保存在不连续的内存空间中

块视为内存页：block-page，KV Cache被按块组织，每块包含多个token的KV对，就像内存分页一样。
token视为字节：token-byte，每个token对应一个模型的输入单位，相当于内存中的最小数据单位。
请求视为进程：request-process，每个推理请求就像一个进程，需要他自己的地址空间来存储token的KV Cache.

1 PagedAttention 论文解读

论文解读这一章节主要是看了B站上的一个学习视频，视频地址是： AI INFRA 学习 02 - vLLM PagedAttention 论文精读

1.0 Abstract

所谓的PagedAttention就是从操作系统的虚拟内存和分页管理得到灵感而发明的，我看其实就是分页管理。

虚拟内存倍划分为固定大小的块，称为页(Pages).
物理内存被划分为与页大小相同的块，称为页框。
页表(Page Table)是一个数据结构，用于记录虚拟页与物理页框之间的映射关系。
当程序访问某个虚拟地址时，系统会通过页表将其转换为对应的物理地址，如果所需的页不在物理内存中，就会发生缺页中断(Page Fault)，操作系统会将所需页从磁盘加载到内存中。

摘要里面还提到，这个PagedAttention能够实现几乎是零损耗的一个KV Cache的一个浪费，另外，他还能使用一个sharing机制进一步减少KV Cache的使用。

1.1 Introduction

里面提到LLMs是一个自回归的Transformer模型，也就是都是基于之前已经生成的Token去预测下一个Token，

这里是说对于一个13B的模型，他的显存使用占用情况，可以看到KV Cche占用了大约30%的消耗。

In this paper, we observe that existing LLM serving systems
$31, 60\] fall short of managing the KV cache memory efficiently. This is mainly because they store the KV cache of a request in contiguous memory space, as most deep learning frameworks \[33, 39\] require tensors to be stored in contiguous memory. However, unlike the tensors in the traditional deep learning workloads, the KV cache has unique characteristics: it dynamically grows and shrinks over time as the model generates new tokens, and its lifetime and length are not known a priori. These characteristics make the existing systems' approach significantly inefficient in two ways:$

在本文中，我们观察到现有的大语言模型（LLM）服务系统[31, 60]在高效管理KV缓存内存方面存在不足。这主要是因为它们将请求的KV缓存存储在连续的内存空间中，而大多数深度学习框架[33, 39]都要求张量存储在连续内存中。然而，与传统深度学习工作负载中的张量不同，KV缓存具有独特的特性：它会随着模型生成新token而动态增长和收缩，且其生命周期和长度在事先并不可知。这些特性使得现有系统的处理方式在两个方面显著低效：

First, the existing systems [31, 60] suffer from internal and

external memory fragmentation. To store the KV cache of

a request in contiguous space, they pre-allocate a contiguous

chunk of memory with the request's maximum length

(e.g., 2048 tokens). This can result in severe internal fragmentation,

since the request's actual length can be much

shorter than its maximum length (e.g., Fig. 11). Moreover,

even if the actual length is known a priori, the pre-allocation

is still inefficient: As the entire chunk is reserved during the

request's lifetime, other shorter requests cannot utilize any

part of the chunk that is currently unused. Besides, external

memory fragmentation can also be significant, since the preallocated

size can be different for each request. Indeed, our

profiling results in Fig. 2 show that only 20.4% - 38.2% of the

KV cache memory is used to store the actual token states in

the existing systems.

首先，现有系统[31, 60]面临内部和外部内存碎片化问题。为了将请求的KV缓存存储在连续内存空间中，它们会预先分配一块与请求最大长度（例如2048个token）相等的连续内存块。这可能导致严重的内部碎片化，因为请求的实际长度可能远小于其最大长度（例如图11所示）。此外，即使实际长度事先已知，预分配仍然效率低下：由于整个内存块在整个请求生命周期内都被保留，其他较短的请求也无法利用当前未使用的部分。同时，外部内存碎片化也可能很严重，因为每个请求的预分配大小可能不同。事实上，如图2所示的我们的性能分析结果表明，在现有系统中，仅有20.4%至38.2%的KV缓存内存被用于存储实际的token状态。

个人理解：这段就是说现在的那些方法会有内存碎片问题，预分配2048个token，会浪费，还有预分配一大块，那么空闲的部分其他的也没法用，

Second, the existing systems cannot exploit the opportunities

for memory sharing. LLM services often use advanced decoding algorithms, such as parallel sampling and beam

search, that generate multiple outputs per request. In these

scenarios, the request consists of multiple sequences that can

partially share their KV cache. However, memory sharing is

not possible in the existing systems because the KV cache of

the sequences is stored in separate contiguous spaces

其次，现有系统无法利用内存共享的机会。LLM服务通常使用高级解码算法，例如并行采样（parallel sampling）和束搜索（beam search），这些算法会为每个请求生成多个输出。在这些场景中，单个请求包含多个序列，这些序列的KV缓存可以在某种程度上共享。然而，在现有系统中无法实现内存共享，因为这些序列的KV缓存被存储在各自独立的连续内存空间中。
个人理解：这就是说没法内存共享。

To address the above limitations, we propose PagedAttention,

an attention algorithm inspired by the operating

system's (OS) solution to memory fragmentation and sharing:

virtual memory with paging. PagedAttention divides the

request's KV cache into blocks, each of which can contain

the attention keys and values of a fixed number of tokens. In

PagedAttention, the blocks for the KV cache are not necessarily

stored in contiguous space. Therefore, we can manage

the KV cache in a more flexible way as in OS's virtual memory:

one can think of blocks as pages, tokens as bytes, and

requests as processes. This design alleviates internal fragmentation

by using relatively small blocks and allocating

them on demand. Moreover, it eliminates external fragmentation

as all blocks have the same size. Finally, it enables

memory sharing at the granularity of a block, across the

different sequences associated with the same request or even

across the different requests.

为了解决上述限制，我们提出了PagedAttention（分页注意力机制），这是一种受操作系统（OS）解决内存碎片化和共享方案启发的注意力算法：使用分页技术的虚拟内存。PagedAttention将请求的KV缓存划分为多个块，每个块可以包含固定数量token的注意力键和值。在PagedAttention中，KV缓存的块不必存储在连续的内存空间中。因此，我们可以像操作系统管理虚拟内存一样更灵活地管理KV缓存：可以将块视为内存页，将token视为字节，将请求视为进程。这种设计通过使用相对较小的块并按需分配来减轻内部碎片问题。此外，由于所有块大小相同，它完全消除了外部碎片。最后，它能够在块的粒度级别实现内存共享，不仅可以在同一请求的不同序列之间共享，甚至可以在不同请求之间共享内存。

个人理解：就是说受到操作系统中分页技术管理虚拟内存方法的启发，然后这里也是将KV 缓存划分成多个快，然后这里还有个比喻

块视为内存页：block-page，KV Cache被按块组织，每块包含多个token的KV对，就像内存分页一样。
token视为字节：token-byte，每个token对应一个模型的输入单位，相当于内存中的最小数据单位。
请求视为进程：request-process，每个推理请求就像一个进程，需要他自己的地址空间来存储token的KV Cache.

然后这种设计，可以减少内部内存碎片和外部内存碎片。

1.2 Background

论文的Background这一章节主要就是先介绍了注意力公式，

然后介绍了KV Cache，介绍了自回归，这些知识之前都看过了就不用看了。

推理过程可以划分为两个阶段：

prefill 阶段：模型接收完整的用户输入 prompt，并一次性并行计算所有 token 的 Query、Key 和 Value 向量。这一阶段是高度并行的，能够充分利用 GPU 的算力资源，因此属于 compute-bound（计算受限），瓶颈在于算力而非内存。
decode 阶段：此阶段模型开始逐个 token 地生成输出。每一步仅处理一个新的 token，模型首先计算其对应的 Query、Key 和 Value 向量。当前 token 的 Query 会与历史及当前的 Key 向量进行点积计算，通过 softmax 得到 attention 权重后，再与历史及当前的 Value 向量进行加权求和，得到当前 token 的 attention 输出。随后，当前步骤产生的 Key 和 Value 向量会被追加写入 KV Cache，以供后续使用。由于该过程是串行的、难以并行加速，且频繁读写缓存数据，因此整体表现为 memory-bound（内存受限），瓶颈主要在 KV Cache 的存储与访问效率。

可以看出，prefill 阶段更侧重于并行计算，而 decode 阶段的性能很大程度上取决于 KV Cache 的内存管理能力。随着生成的 token 数量增加，KV Cache 会线性增长，并逐步占据大量 GPU 显存。

1.3 Memory Challenges in LLM Serving

当前主流的 LLM 推理系统在 KV Cache 的内存管理上普遍存在 3 类结构性问题：显存占用增长快、内存碎片严重，以及缓存难以复用。

1.3.1 KV Cache 占用迅速增长，极易耗尽显存

以 OPT-13B 为例，论文指出其每个 token 的 KV Cache 占用约为 800KB，计算方式为：

2（Key 和 Value）× 5120（hidden size）× 40（Transformer 层数）× 2 字节（FP16） = 819,200 字节 ≈ 800KB

如果生成完整的 2048 个 token，单个请求的 KV Cache 占用就可达约 1.6GB。在一块 40GB 显存的 A100 GPU 上，这种增长速度意味着仅同时处理少量请求，就可能达到内存瓶颈，直接限制了批处理规模和系统吞吐量。

1.3.2 预分配导致内存碎片严重

传统推理系统（如 FasterTransformer 和 Orca）通常采用预分配策略：请求开始时，就按最大可能生成的长度（如 2048 token）为每个请求分配一整块连续内存空间。这种方式带来 3 类显著的内存浪费：

已预留但尚未使用的空间（Reserved）：虽然未来会使用，但当前暂未使用到。
内部碎片（Internal Fragmentation）：实际 token 数小于预留长度，剩余空间浪费。
外部碎片（External Fragmentation）：不同请求所需内存大小不一，造成内存块间不连续，产生碎片。

论文中的实验结果显示，再传统系统中真正用于存放 KV Cache 的有效内存占比最低仅约 20.4%，其余全为浪费。

1.3.3 KV Cache 难以共享，内存复用受限

在 Parallel Sampling 或 Beam Search 等复杂推理模式中，一个请求可能包含多个生成分支（如多个候选答案），这些分支共享相同的 prompt，因此理论上可以共享其 KV Cache。但传统系统将每个分支的 KV Cache 存放在独立的连续内存中，物理上无法实现共享，只能进行冗余复制，进而显著增加了内存开销。

1.4 Method

In this work, we develop a new attention algorithm, PagedAttention,

and build an LLM serving engine, vLLM, to tackle

the challenges outlined in §3. The architecture of vLLM is

shown in Fig. 4. vLLM adopts a centralized scheduler to

coordinate the execution of distributed GPU workers. The

KV cache manager effectively manages the KV cache in a

paged fashion, enabled by PagedAttention. Specifically, the

KV cache manager manages the physical KV cache memory

on the GPU workers through the instructions sent by the

centralized scheduler.

在本工作中，我们开发了一种新的注意力算法PagedAttention，并构建了一个LLM服务引擎vLLM，以应对§3中概述的挑战。vLLM的架构如图4所示。vLLM采用集中式调度器来协调分布式GPU worker的执行。KV缓存管理器通过PagedAttention以分页方式有效管理KV缓存。具体来说，KV缓存管理器通过集中式调度器发送的指令来管理GPU worker上的物理KV缓存内存。

1.4.1 PagedAttention

PagedAttention allows storing continuous

keys and values in non-contiguous memory space

PagedAttention就是允许将连续的KV保存在不连续的内存空间中，这句话就是PagedAttention的核心。

每个内存块都存储了固定数量的token，这个大小称为block size(B),论文中的block size是16.

这里就是说计算注意力机制的时候直接按block计算了，而不是一个token一个token的去计算。

1.4.2 KV Cache Manager

这里就是说参考了操作系统中的虚拟内存的方式，然后连续的逻辑页可以对应到非连续的物理内存页，并且物理的内存空间并不需要提前去预留，而是按需分配，然后我们就可以用固定大小的KV blocks管理KV Cache ,就像虚拟地址的页那样。

一个request的KV Cache是可以代表一系列的逻辑KV blocks,GPU的blcok engine会一次性的申请一大块的显存，然后把这一大块显存分成多块物理的KV blocks，并且会映射成逻辑的地址，然后必要的时候还会swap到cpu的内存。

然后这里还有个KV blockmanager,他维护了一个baock talbe，也就是逻辑和物理的KVblock的映射关系。

1.4.x 个人阶段总结

为了解决 KV Cache 内存管理中的高占用、严重碎片和复用困难等问题，vLLM 提出了一种全新的注意力机制 ------ PagedAttention。它的核心思想，借鉴自操作系统中广泛应用的虚拟内存分页机制（Virtual Memory & Paging）。

在操作系统中，虚拟内存将程序的地址空间划分为固定大小的"页"（Page），这些页可以映射到物理内存中任意位置，实现非连续分配，从而有效解决了内存碎片和共享问题。PagedAttention 将这一思路引入 KV Cache 的管理中，并带来了 3 项关键改进：

KV Cache 被切分为固定大小的 block：PagedAttention 将每个序列的 KV Cache 切分为固定大小的 block（默认是 16 个 token），每个 block 存储若干个 token 的 Key 和 Value 向量。这种设计统一了内存分配粒度，使系统能够以更标准化的方式管理 KV Cache 的分配与回收，从而提升内存复用效率，并有效减少内存碎片。
block 可以存放在非连续的物理内存中：与传统的 Attention 不同，PagedAttention 不再要求这些 KV 向量在内存中连续排列，而是通过逻辑 block 与物理 block 的映射，实现非连续存储。映射关系由 block table 维护，它类似于操作系统中的页表，用于记录每个逻辑 block 对应的物理内存位置，确保模型在推理过程中可以正确读取所需的 KV Cache。
支持灵活的分配与释放，以及共享机制：PagedAttention 支持按需分配和回收 block，并允许多个序列共享 block。PagedAttention 使用了 copy-on-write（CoW）机制，允许多个生成样本共享大部分输入 prompt 的 KV Cache，只有当某个分支需要写入新 token 的 KV 数据时，系统才会将相关 block 复制到新的物理位置，从而在保证数据隔离的同时极大地节省显存资源，提升推理效率与吞吐量。

1.4.3 DecodingwithPagedAttention and vLLM

prefill 阶段：vLLM 为前 7 个 token 分配两个逻辑块 block 0 和 block 1，分别映射到物理块 7 和 1。block 0 存储前 4 个 token，block 1 存储后 3 个 token 及第一个生成 token "fathers"，填充数为 4。
decode 阶段 - 生成第 1 个词：生成 token "brought"，由于 block 1 尚未填满（最多容纳 4 个 token），因此直接将新 KV Cache 写入该块，填充计数从 3 更新为 4。
decode 阶段 - 生成第 2 个词：生成下一个 token，此时 block 1 已满，系统为逻辑块 block 2 分配新的物理块 block 3，并写入 KV Cache，同时更新映射表。

1.4.4 Application to Other Decoding Scenarios

PagedAttention 中采用的基于分页的 KV Cache 管理机制，不仅在常规单序列生成中表现出色，也天然适合多种复杂的解码策略。

1.4.4.1 并行采样

在 Parallel Sampling 中，同一个 prompt 会生成多个候选输出，便于用户从多个备选中选择最佳响应，常用于内容生成或模型对比测试。

在 vLLM 中，这些采样序列共享相同的 prompt，其对应的 KV Cache 也可以共用同一组物理块。PagedAttention 通过引用计数和 block-level 的 copy-on-write 机制实现共享与隔离的平衡：只有当序列出现不同分支时，才会触发复制操作。

1.4.4.2 束搜索（Beam Search）

Beam Search 是机器翻译等任务中常见的解码策略。它会维护多个"beam"路径，每轮扩展最优候选并保留 top-k 序列。

在 vLLM 中，多个 beam 可以共享公共前缀部分的 KV Cache，不仅包括输入 prompt，还包括生成过程中的公共前缀 token。只要这些 beam 的生成路径尚未分叉，它们就会复用相同的物理块。当路径分叉发生后，vLLM 才通过 copy-on-write 机制对共享块进行拆分，从而保证每个 beam 的独立性。

1.4.4.3 共享前缀（Shared Prefix）

在许多提示工程实践中，多个请求会以相同的系统提示或 few-shot 示例开头（例如翻译任务中的多个例句）。这些共享前缀同样可以被 vLLM 缓存并复用。

1.4.4.4 混合解码（Mixed Decoding）

前面提到的几种解码方式（如 Parallel Sampling、Beam Search、Shared Prefix 等）在内存的共享和访问方式上存在差异，传统系统很难同时高效地处理这些不同策略的请求。而 vLLM 则通过一个通用的映射层（block table）屏蔽了这些差异，让系统能够在统一框架下处理多样化的请求。

1.5 调度和抢占

到目前为止，我们已经回答了"vLLM是如何优化KV cache显存分配"的问题，现在我们来回答另一个重要的问题：

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

1.5.1 总原则

当有一堆请求来到vLLM服务器上时，vLLM需要一个调度原则来安排如何执行这些请求，这个调度原则概括如下：

先来的请求先被服务（First-Come-First-Serve, FCFS）

如有抢占的需要，后来的请求先被抢占（preemption）

（1）先来的请求先被服务这个很好理解，当有一堆请求到达vLLM服务器时，vLLM肯定优先处理来得早的请求

（2）后来的请求先被抢占想象一下，当一堆请求来到vLLM服务器做推理，导致gpu显存不足时，vLLM会怎么做呢？

最直接的办法，就是暂停这堆请求中最后到达的那些请求的推理，同时将它们相关的KV cache从gpu上释放掉，以便为更早到达的请求留出足够的gpu空间，让它们完成推理任务。如果不这样做的话，各个请求间相互争夺gpu资源，最终将导致没有任何一个请求能完成推理任务。等到先来的请求做完了推理，vLLM调度器认为gpu上有足够的空间了，就能恢复那些被中断的请求的执行了。

在资源不足的情况下，暂时中断一些任务的执行，这样的举动就被称为"抢占（preemption）"。

1.5.2 终止和恢复被抢占的请求

对于这些因gpu资源不足而被抢占的任务，vLLM要完成两件事：

暂停它们的执行，同时将与之相关的KV cache从gpu上释放掉

等gpu资源充足时，重新恢复它们的执行

针对这两件事，vLLM分别设计了Swapping（交换策略）和Recomputation（重计算策略）来解决。我们来细看这两个策略。

（1）Swapping

对于被抢占的请求，vLLM要将其KV cache从gpu上释放掉，那么：

问题1：该释放哪些KV cache？

问题2：要把这些KV cache释放到哪里去？

先看问题1。由前文PagedAttention原理可知，一个请求可能对应多个block。我们既可以选择释放掉部分block，也可以选择释放掉全部block，或者更科学地，我们可以预测一下哪些block被使用的频率最低，然后释放掉这些低频block（但这种方式实现起来难度较大，性价比不是很高）。在vLLM中，采取的是all-or-nothing策略，即释放被抢占请求的所有block。

再来看问题2。对于这些被选中要释放的KV block，如果将它们直接丢掉，那未免过于浪费。vLLM采用的做法是将其从gpu上交换（Swap）到cpu上。这样等到gpu显存充份时，再把这些block从cpu上重载回来。

（2）Recomputation

知道了Swapping机制，重计算的过程也很好理解了：当vLLM调度器任务gpu资源充足时，对于那些被抢占的请求，它会将其卸载到cpu上的KV block重新加载进gpu中，继续完成推理任务。

（3）总结

好，到这里，我们总结一下vLLM对请求的调度处理流程：

当一堆请求来到vLLM服务器上时，按照First-Come-First-Serve（FCFS）原则，优先处理那些最早到来的请求。

当gpu资源不足时，为了让先来的请求能尽快做完推理，vLLM会对那些后到来的请求执行"抢占"，即暂时终止它们的执行。

一旦vLLM决定执行抢占操作，它会暂停处理新到来的请求。在此期间，它会将被抢占的请求相关的KV block全部交换（swap）至cpu上。等交换完成后，vLLM才会继续处理新到来的请求。

当vLLM认为gpu有足够资源时，它会将cpu上的KV block重新加载回gpu，恢复被抢占请求的执行（recomputation）

1.6 分布式管理

在LLM推理实操中，某些场景下单卡是完成不了推理的，需要多卡。那么对于多gpu这种更普适性的情况，vLLM是怎么处理的呢？

上图显示了在分布式场景下，vLLM的整体运作流程：

首先，vLLM有一个中央调度器（Scheduler），它负责计算和管理每张卡上KV cache从逻辑块到物理块的映射表(block tables)
在做分布式计算时，Schedular会将映射表广播到各张卡上，每张卡上的Cache engine接收到相关信息后，负责管理各卡上的KV block

上图中给出的例子，是用张量模型并行（megatron-lm）做分布式推理时的情况，所以图中每个worker上写的是model shard。在张量并行中，各卡上的输入数据相同，只是各卡负责计算不同head的KV cache。所以这种情况下，各卡上的逻辑块-物理块的映射关系其实是相同的（用的同一张block table），只是各卡上物理块中实际存储的数据不同而已。

2 总结与整理

为了解决 KV Cache 内存管理中的高占用、内存碎片和内存复用等问题。