AI 技术
#PagedAttention#vLLM#KV Cache#推理服务#显存管理

PagedAttention 的块表机制:按块管理 KV Cache 如何减少显存碎片并提升在线服务吞吐

面向 vLLM 在线服务场景,解释 PagedAttention 为何把 KV Cache 切成固定大小的逻辑块、用块表映射到非连续物理块,并按需分配。文章对比连续缓存与块式缓存的显存占用和调度灵活性,分析块大小对碎片率与吞吐的影响,并给出可观测信号、失败模式与适用边界。

一个请求只生成几十个 token,为什么显存却先满了

设想一台 A100 40GB 的在线推理机,跑一个 13B 模型。模型权重常驻显存,剩下的一部分留给 KV Cache 和中间激活。现在有 200 个用户同时提问,每个问题的输入长度从几十 token 到几千 token 不等,输出长度也无法预估。运维看到的不是算力打满,而是显存先耗尽,调度器被迫拒绝新请求,GPU 利用率反而掉下来。

直觉上的解释是“并发太多”。但把并发降到一半,吞吐并没有线性回升,因为真正被浪费的是 KV Cache 的预留空间。传统推理系统在请求开始时,就按模型支持的最大长度(例如 2048 token)为每个请求分配一整块连续显存。请求实际只用了几百个 token,剩下的预留空间既不能给别的请求,也不能被回收。资料 1 的论文摘要指出,KV Cache 随请求动态增长和收缩,管理不当时会被碎片和冗余复制大量浪费,从而限制批大小。资料 2 和资料 3 引用的论文实验显示,传统系统中真正用于存放 KV Cache 的有效内存占比最低约 20.4%。

这篇文章以 vLLM 在线服务为贯穿场景,拆解 PagedAttention 的块表机制:它把 KV Cache 切成固定大小的逻辑块,通过块表映射到非连续物理块,按需分配。我们要回答的是它为什么能减少碎片、系统如何实现、以及块大小、抢占和 Continuous Batching 在什么条件下会让收益退化。

连续缓存的三类浪费与基线为什么失效

要理解块式管理的价值,先要看连续缓存到底浪费在哪里。资料 2 把传统预分配策略的浪费归为三类。

第一类是已预留但尚未使用的空间。请求开始时按最大长度预留,但生成是逐 token 进行的,大部分预留空间在请求生命周期内一直空着。第二类是内部碎片:请求实际 token 数小于预留长度,块内剩余空间无法被其他请求使用。第三类是外部碎片:不同请求长度不一,显存被切成许多不连续的小块,新来的长请求即使总剩余显存足够,也可能找不到一块足够大的连续空间。

这三类浪费叠加,导致基线系统的有效 KV Cache 占比很低。资料 4 给出的量级是内存浪费可高达 60%~80%,与资料 2、资料 3 引用的约 20.4% 有效占比在方向上一致。

连续缓存的另一个问题是难以共享。在 Parallel Sampling 或 Beam Search 中,一个请求会分出多个候选分支,它们共享同一段 prompt。理论上这段 prompt 的 KV Cache 可以只存一份,但连续内存要求每个分支各占一块独立区域,物理上无法共享,只能冗余复制。资料 2 指出,这显著增加了内存开销。

还有一个容易被忽略的约束:连续缓存让抢占变得昂贵。显存不足时,调度器需要暂停某个请求。如果它的 KV Cache 是一整块连续区域,移动或换出这块区域的操作粒度很大,恢复时也要重新找到同样大的连续空间。粒度粗,调度器就不敢频繁抢占,只能保守地限制并发。

块表如何把逻辑连续映射到物理分散

PagedAttention 的核心思路借鉴操作系统的虚拟内存分页。资料 2 和资料 4 都强调这一点:程序看到的是连续的虚拟地址空间,实际物理页可以分散在内存任意位置,由页表负责映射。PagedAttention 把这套机制搬到 KV Cache 上。

它做了三件事。第一,把每个序列的 KV Cache 在逻辑上切成固定大小的块,每个块存固定数量 token 的 Key 和 Value 向量。资料 2 和资料 3 提到 vLLM 默认块大小为 16 个 token。第二,这些块不要求物理连续,可以放在显存中任意空闲位置。第三,为每个请求维护一张块表,记录它的第几个逻辑块对应哪个物理块。

这里的关键术语是块表。它解决的是“逻辑顺序”和“物理位置”分离的问题。模型在计算注意力时,需要按 token 顺序访问历史 KV,但物理上这些 KV 分散在不同块里。块表让计算核先根据 token 位置算出它属于哪个逻辑块,再查表找到物理块地址,最后从该地址取出 Key 和 Value。资料 4 把这一步描述为定制 CUDA kernel 的间接寻址:接收块表作为额外输入,按逻辑块查物理地址,再 gather 数据。

分配策略是按需的。资料 2 和资料 3 用同一个例子说明:prompt 有 7 个 token,块大小为 4,prefill 阶段分配两个逻辑块,分别映射到物理块 7 和 1。第一个块存 4 个 token,第二个块存 3 个 token 加第一个生成 token,填充计数为 4。生成下一个 token 时,如果当前块未满,直接写入;块满后,才为新的逻辑块分配物理块并更新映射表。

这种按需分配把浪费压到最小。唯一的内部碎片只出现在每个序列最后一个未填满的块里,量级是“块大小减一”个 token 的空间,而不是“最大长度减实际长度”的空间。

flowchart TD
    A[请求到达: prompt 7 token] --> B[prefill 计算 KV]
    B --> C[逻辑块 0 映射物理块 7]
    B --> D[逻辑块 1 映射物理块 1]
    C --> E[块 0 存 4 token 已满]
    D --> F[块 1 存 3 token 填充计数 3]
    F --> G[decode 生成 token: 写入块 1]
    G --> H{块 1 是否已满}
    H -- 未满 --> I[填充计数 3 变 4]
    H -- 已满 --> J[申请新物理块 3]
    J --> K[逻辑块 2 映射物理块 3]
    K --> L[更新块表并继续 decode]

图中展示的是单个请求的块分配路径。真正让吞吐提升的,是这张块表可以被多个请求共享,以及物理块可以被调度器灵活回收。

共享、写时复制与抢占如何依赖块粒度

块表机制不只是省空间,它改变了调度器能做什么。

先看共享。在 Parallel Sampling 中,同一 prompt 生成多个候选。资料 2 和资料 4 说明,这些序列的块表在初始阶段指向同一组存储 prompt KV 的物理块,系统用引用计数追踪有多少序列在共享。当某个分支生成与其它分支不同的 token 时,系统只新分配一个块存这个独有的 KV,并只更新该分支自己的块表。资料 4 把这种成本变化描述为从“乘法”关系变成“加法”关系。

写时复制(copy-on-write)是共享的配套机制。多个序列共享同一物理块时,块内容只读。一旦某个序列要往共享块里写新数据,系统先把该块复制到新物理位置,再让这个序列的块表指向新块,从而保证数据隔离。资料 2 和资料 3 都指出,只有当序列出现不同分支时才触发复制。

再看抢占。当显存不足、请求量超过处理能力时,vLLM 需要回收部分请求的 KV Cache。资料 3 提到 vLLM 采用 FCFS 策略,优先服务最早到达的请求,优先抢占最近到达的请求。回收粒度是 All-or-Nothing:一个序列的所有 KV Cache 块要么全部回收,要么不回收;Beam Search 这类含多个子序列的请求,必须作为一个 sequence group 整体抢占或恢复,以免破坏共享关系。

恢复有两种方式。资料 3 列出 Swapping 和 Recomputation:Swapping 把被抢占请求的 KV Cache 块从 GPU 移到 CPU 内存,需要时再换回;Recomputation 则丢弃 KV,等请求重新调度时用 prompt 重新计算。资料 4 提到 vLLM 在启动时会同时管理 GPU 块和 CPU 块,为 Swapping 提供基础。

块粒度让这些操作变轻。换出一整块连续 KV 区域,需要处理大块内存拷贝和重新寻找连续空间;换出若干固定大小的物理块,则可以逐块搬运,块表只需更新映射。调度器因此敢更频繁地抢占和恢复,从而在同样的显存下维持更多活跃请求。

块大小、碎片率与吞吐之间的权衡

块大小是这套机制里最直接的调节旋钮。资料 2 和资料 3 提到默认值是 16 个 token,但这不是唯一选择。

块越小,最后一个未填满块的内部碎片越小。极端情况下块大小为 1,内部碎片几乎为零,但每个块都要在块表里占一项,块表变长,查表和间接寻址的开销上升,物理块数量也变多,分配器管理成本增加。块越大,块表项越少,寻址开销低,但每个序列最后一个块的浪费变大,而且一个请求要凑满一个大块才能高效利用,短请求的碎片率反而上升。

这里存在一个和 Continuous Batching 的耦合。Continuous Batching 允许新请求在任意 decode 步加入批次,也允许完成的请求随时退出。批次里每个请求的序列长度不同,块分配和回收在每一步都可能发生。块大小如果太大,短请求占不满块,批次里就会出现大量半空块;块大小如果太小,每步的块表更新和物理块分配次数增加,调度开销上升。

下面的表格对比连续缓存与块式缓存在几个工程维度上的差异。表格中的定性判断来自资料 1 至资料 4 的机制描述,具体数值只在资料明确给出处标注。

维度连续缓存(预分配)块式缓存(PagedAttention)
分配粒度按最大长度整块预留按固定大小块按需分配
内部碎片实际长度与预留长度之差,可达很大仅最后一个未满块,上限约为块大小减一
外部碎片不同长度请求切割显存,易产生物理块等大,可填充任意空闲块
有效内存占比资料 2、3 引用约 20.4% 下限资料 4 称可提升至约 96%
前缀共享物理连续,难以共享块表可指向同一物理块,配引用计数与写时复制
抢占粒度整块连续区域,操作重逐物理块换出或重算,块表更新轻
寻址开销直接按连续偏移访问需查块表间接寻址,定制 kernel
与 Continuous Batching 配合批大小受连续空间限制可按块动态增减,批大小更灵活

表格里最值得注意的一行是寻址开销。块式缓存不是免费的,它把内存管理的复杂度转移到了计算核和块表维护上。资料 4 指出,标准 Attention 库是为连续内存设计的,PagedAttention 必须实现定制 CUDA kernel 来处理间接寻址。这意味着块式缓存的收益要以 kernel 实现复杂度和一定的查表开销为代价。

在线服务中该观察哪些信号

把 PagedAttention 部署到在线服务后,光看 GPU 利用率不够,需要盯住几类和块管理直接相关的信号。

第一类是块利用率。可以观察已分配的物理块中有多少 token 槽位真正存了 KV,以及有多少块处于半满状态。如果半满块比例长期偏高,说明块大小相对请求长度分布偏大,或者批次里短请求占比高。

第二类是抢占频率和恢复方式。资料 3 描述的 Swapping 和 Recomputation 有不同代价:Swapping 消耗 CPU-GPU 带宽,Recomputation 消耗算力重算 prefill。如果抢占频繁且大量走 Recomputation,说明显存长期紧张,或者调度策略过于激进。

第三类是块表规模和查表开销。块越小,块表项越多。如果 decode 阶段每步的调度时间随活跃序列数明显上升,需要检查块表维护和物理块分配是否成为瓶颈。

第四类是共享命中情况。在 Parallel Sampling、Beam Search 或共享 system prompt 的场景中,如果引用计数显示共享块很少,说明请求之间没有形成可复用的前缀,块表共享机制没有发挥作用。

这些信号要结合请求长度分布一起看。长输入短输出的请求和短输入长输出的请求,对块大小和抢占策略的压力完全不同。

什么时候块式缓存也会退化

PagedAttention 不是无条件提升吞吐。资料 1 的论文摘要指出,改进在更长序列、更大模型和更复杂解码算法下更明显。反过来说,以下条件会让收益缩小甚至消失。

请求普遍很短且长度接近时,连续缓存的内部碎片本来就不大,块式缓存省下的空间有限,反而多出块表维护和间接寻址的开销。

块大小设置与请求长度分布严重错配时,碎片率会回升。如果绝大多数请求只生成十几个 token,而块大小远大于此,每个请求都会留下一个几乎全空的尾块,块利用率下降。

共享前缀少、请求之间没有公共 prompt 时,引用计数和写时复制机制没有用武之地,收益主要来自按需分配,而不是共享。

定制 kernel 在小批量或低并发下可能体现不出优势。块式缓存的吞吐收益来自“同样显存装下更多请求”,如果并发本来就低、显存不紧张,间接寻址的额外开销没有足够多的请求来摊薄。

抢占策略过于频繁时,Swapping 的 CPU-GPU 拷贝或 Recomputation 的重算会吃掉吞吐。块粒度让抢占变轻,但不等于抢占免费。

还有一个边界是模型架构。资料 4 提到 vLLM 用不同的 KVCacheSpec 子类支持 FullAttention、SlidingWindow 和 Mamba 等结构,滑动窗口注意力只需要为窗口内 token 分配缓存,其内存计算方式与全局注意力不同。块表机制要适配这些差异,不能假设所有层都按同一种方式增长。

块表机制真正的工程价值

回到开头的场景:200 个并发请求把显存打满,不是因为算力不够,而是因为连续缓存把大量显存锁在预留空间里。PagedAttention 用固定大小的块和块表,把“逻辑连续”和“物理连续”解耦,让显存可以按 token 实际增长逐块分配,把内部碎片压到最后一个未满块,把外部碎片消解为等大物理块的填充问题。

它同时改变了调度的可能性空间。块表让前缀共享、写时复制、逐块抢占和恢复成为常规操作,这些操作又反过来支撑 Continuous Batching 在动态批次下维持较高并发。资料 1 的论文摘要报告 vLLM 在相同延迟水平下相比 FasterTransformer 和 Orca 等系统提升吞吐 2~4 倍,并指出改进在长序列、大模型和复杂解码下更明显。

工程上的判断可以落在一句话上:块式缓存的收益取决于请求长度分布、共享前缀比例和并发压力。三者都强时,块表机制把显存从“预留”变成“按需”,吞吐收益最大;请求短、无共享、并发低时,块表维护和间接寻址的成本可能盖过碎片节省。部署时先量出请求长度分布和共享前缀比例,再决定块大小和抢占策略,比直接套用默认值更可靠。

资料来源

  1. vLLM: Easy, Fast, and Cheap LLM Serving with PagedAttention
  2. 详解vLLM PagedAttention优化大模型推理KV Cache内存的原理-开发者社区-阿里云
  3. vLLM 核心技术 PagedAttention 原理详解-腾讯云开发者社区-腾讯云
  4. 从KV-Cache到PagedAttention,揭秘LLM推理性能的全部细节 - SIo_2 - 博客园