PagedAttention是vLLM团队提出的一种内存管理机制，它借鉴了现代操作系统中虚拟内存(Virtual Memory)和分页(Paging)的设计思想。

1 为什么需要PagedAttention?

在传统的KV Cache实现中，我们通常需要预先为每个序列分配一块连续的显存空间，其大小通常基于模型能够处理的最大上下文长度（比如32k/128k/256k），这会导致以下问题：

  • 显存浪费：即使推理一个序列只需要1K的显存，在推理的过程中也会为它预留最大上下文长度的显存空间，除了真正使用的1K空间，剩下的空间完全是闲置的
  • 内存碎片：当多个请求并存，显存中散落着各种被预留但未填满的显存空隙，无法为新来的长序列推理分配连续的显存空间

2 PagedAttention的原理

针对于上述KV Cache的缺点，PagedAttention将KV Cache切分成一个个固定大小的“块”，类似于操作系统中将物理内存切分为固定大小的“页”（Pages）。

PagedAttention的关键部分：

  • Block Table：块表，每一个序列在模型推理时都有一个自己的Block Table，Block Table记录了该推理请求对应的逻辑块(Logical Blocks)与物理块(Physical Blocks)的映射关系
  • Physical Block Pool：物理块池，显存中存在一个公共的物理块池，当模型生成一个新的Token时，系统只需要动态的从这个池子里分配一个空的物理块，并将其填入该推理请求的Block Table

举个例子，假设模型的推理流程是从图书馆借阅书籍的过程，传统的KV Cache的内存分配要求是图书馆必须留出一整排书架给一个人借书和放书，而PagedAttention实际上是给每个人一张图书借阅卡，这张卡就是PagedAttention的Block Table，而书可以放在图书馆的任意的空书架上(物理块上)。

PagedAttention的核心优点：

  • 非连续存储：KV Cache不需要再占用物理上连续的显存空间，逻辑上相邻的Token在物理上可以存储在完全不同的显存块中；
  • 几乎零显存闲置：只有在真正需要存储Token时才会分配物理块，这使得显存的利用率接近100%
  • 灵活的共享机制：因为是通过Block Table进行索引的，PagedAttention还可以轻松实现Prompt共享，例如，多个用户发送相同的提示词，在推理时，可以指向同一个物理块，进一步节省大量显存

3 PagedAttention的工作流程

在处理一个推理请求时，PagedAttention的工作流程如下：

  1. 物理块池化初始化
    在模型启动时，系统会将所有可用的GPU显存预先划分为固定大小的物理块，假设每个物理块可以存储16个Token，这些物理块组成了一个空闲列表。

  2. 推理请求处理流程
    当一个生成任务（prompt）进入推理流程，需要经过以下两个阶段：

（1）提示词阶段

  • 初始化：根据输入的Prompt的长度，系统计算出需要多少个物理块
  • 分配与映射：从空闲列表中拿走对应数量的物理块，并将它们的物理地址写入到该请求的Block Table中
  • 计算：模型计算所有Token的KV值，并将其填入到这些物理块中

（2）生成阶段
依次生成Token，每生成一个Token，都按照以下流程：

  • 检查剩余空间：检查最后一个物理块是否还有空位
  • 动态分配：
    • 如果有空位，直接写入
    • 如果没有空位，则说明当前块已满，系统立刻从空闲列表中申请一个新的物理块，更新该请求的Block Table，并将新Token写入到物理块中
      KV访问：当计算Attention时，模型通过读取Block Table，知道要去哪几个不连续的物理池地址把KV读取出来，并送入GPU计算新的Token

4 PagedAttention与CPU的内存分页的区别？

在现代操作系统中，CPU的内存分页本质上是为了让进程访问比实际物理内存更大的空间。
而PagedAttention本质上是为了提高显存的分配效率和吞吐量，使得模型可以在显存总量不变的情况下支持更多的并发请求或者更长的上下文。