Skip to main content

优化大模型推理:PagedAttention 如何实现 GPU 显存高效管理

·933 words·5 mins
一幅抽象的、高效自动化工厂车间景象,象征着优化后的大型语言模型(LLM)推理系统。高速运转的 GPU 单元如同生产线上的机器,由动态、模块化的存储单元(键值缓存 KV cache)无缝供料,形象地展示了 PagedAttention 和连续批处理等技术如何最大化吞吐量并减少资源浪费。

最近看了硅谷 101 的一期节目——《“榨”出硅的极限:怎么让 GPU 不“闲着?》,其中对 SGLang(Structured Generation Language,结构化生成语言)、RadixArk 与人工智能基础设施(Artificial Intelligence Infrastructure,简称 AI Infra)的讨论,让我重新思考了一个看似简单的问题:当图形处理器(Graphics Processing Unit,GPU)已经足够昂贵、模型也已经足够大时,为什么推理系统仍然会让大量算力处于等待状态?

直觉上,我们很容易把大模型推理理解成一个纯粹的算力问题:模型越大,就需要越多 GPU;GPU 越快,生成速度就越快。但在真实的在线服务中,GPU 更像是一条昂贵的生产线。生产线是否高效,不只取决于机器的运转速度,还取决于原料能否及时送达、半成品如何存放,以及不同订单如何排队。

在大模型推理里,“半成品”就是键值缓存(Key–Value Cache,KV cache)。如果显存分配不合理,即使计算单元仍有余力,系统也可能因为无法容纳更多请求而无法扩大批次(batch)。此时,GPU 不是无法有效计算,而是被低效的内存管理“饿住了”。

vLLM 的重要性正在这里。它没有修改 Transformer 的数学定义,也没有创造新的模型结构,而是借用了操作系统中已经存在了几十年的虚拟内存与分页思想,重新设计了 KV cache 的管理方式。这个看似底层的变化,后来又与连续批处理(continuous batching)、前缀缓存和集群调度结合,逐渐形成了现代大语言模型(Large Language Model,LLM)推理服务的基础框架。

先给结论:现代 LLM Serving 在解决什么 #

如果把一套在线推理系统看成一座工厂,那么几项代表性技术分别解决了不同层面的问题:

下表中的 P/D 指预填充与解码(Prefill/Decode)两个阶段;TTFT 指首个词元(token)延迟(Time to First Token),TPOT 指每个输出 token 的平均生成时间(Time per Output Token)。

技术主要问题可以类比为直接影响
continuous batching哪些请求在这一轮进入 GPU动态调整生产队列并发、排队时间
PagedAttention不断增长的 KV cache 放在哪里用标准货箱按需占用仓位显存利用率、batch size
RadixAttention已计算的公共前缀能否复用复用相同订单的半成品Prefill 计算量、缓存命中率
P/D 分离两类不同工序是否应共享设备将备料与装配分到不同车间TTFT、TPOT、尾延迟

它们不是同一个技术的不同名称,也不是简单的前后替代关系。准确地说,它们分别优化调度、存储、复用和集群资源配置,只有组合起来,节省出的显存和计算才能真正转化为更高的有效吞吐。

在理解 PagedAttention 之前,先理解 KV cache #

大模型生成文本通常分为两个阶段。

第一个阶段是 Prefill。模型一次处理整个输入提示(prompt),计算所有输入 token 的中间状态,并产生用于采样首个输出 token 的输出分数(logits)。因为许多 token 可以并行处理,这个阶段包含大量矩阵运算,通常更偏向计算密集型。

第二个阶段是 Decode。在标准自回归解码中,每轮通常为每个活跃序列生成一个 token。每产生一个新 token,都需要参考此前的上下文。如果每一步都重新计算所有历史 token,代价会非常高,因此系统会保存各层注意力机制(attention)已经计算出的 Key 与 Value,供后续步骤直接读取。这部分状态就是 KV cache

可以把它想象成厨师做一道不断追加配料的菜。KV cache 就像厨师放在手边的操作记录:前面用了哪些原料、处理到了哪一步,都已经记好。每加入一种新配料,只需读取记录并继续,而不必从头再做一次整道菜。

但这份“记录”并不小。单个序列的 KV cache 大小可以粗略表示为:

$$ \text{KV cache} \approx 2 \times L \times H_{kv} \times D \times T \times B_{dtype} $$

其中,$2$ 代表 Key 与 Value,$L$ 是模型层数,$H_{kv}$ 是 KV 头数量,$D$ 是每个头的维度,$T$ 是上下文 token 数,$B_{dtype}$ 是每个数值占用的字节数。并发请求越多、上下文越长,需要保存的 KV cache 就越多。

vLLM 论文给出了一个很直观的例子:OPT-13B(Open Pretrained Transformer,约 130 亿参数)在 FP16(16-bit floating-point format,16 位浮点格式)下,每个 token 的 KV cache 约占 800 KiB(kibibyte,二进制千字节);如果序列长度达到 2048,单个请求的 KV cache 就可能占用约 1.6 GiB(gibibyte,二进制吉字节,约合 1.7 GB)的显存 [1],这不包含对齐、元数据及其他运行时开销。当几十个用户同时生成长文本,KV cache 很快就会成为一种比模型权重更动态、也更难管理的显存负担。

这里真正棘手的地方,不只是“大”,而是长度未知且不断增长。一个请求可能只回答 20 个 token,也可能生成 2000 个 token;系统在请求开始时通常无法准确知道生成会在哪里结束。

连续预分配为什么会浪费显存 #

在早期推理框架中,一个常见做法是按照请求可能达到的最大长度,提前申请一段连续显存。这种方式容易实现,因为每个请求的 KV cache 像一个固定大小的数组,访问路径也相对简单。

问题在于,它相当于要求酒店在客人入住时,就按照“最多可能住 30 天”锁定整整 30 天的房间。即使客人第二天退房,其他客人在此之前也无法使用那些被预留的位置。

可以看一个高度简化的例子。假设三个请求的最大生成长度都设置为 2048 token,但它们最后分别生成了 32、256 和 1024 token:

  • 系统预留的容量:$3 \times 2048 = 6144$ 个 token 位置;
  • 实际使用的容量:$32 + 256 + 1024 = 1312$ 个 token 位置;
  • 在这些输出上,实际使用率只有约 21%。

这种连续预分配的结构性浪费也是 PagedAttention 设计目标之一 [1]。

这个例子忽略了输入 prompt、对齐方式和运行时回收,其目的仅为说明结构性问题:只要输出长度难以预测,按最大长度预留就很容易让大量空间在请求生命周期中暂时无法被其他请求使用。

具体来说,传统连续分配会产生三类浪费:

  • 预留空间: 请求运行期间为未来 token 保留、暂时不能供其他请求使用的容量;
  • 内部碎片: 分配粒度或最大长度预估造成的、最终没有存放有效 token 的已分配空间;
  • 外部碎片: 空闲容量被分散成不连续区域,无法满足新的连续内存申请。

外部碎片很像停车场里零散空着许多车位,但一辆大巴需要四个连续车位,最终仍然停不进去。

传统连续预分配产生预留空间、内部碎片和外部碎片,而分页式 KV cache 可以按需复用空闲 block

图 1:传统连续预分配与分页式 KV cache 管理的对比。固定大小的 block 使空闲显存可以更快地交给其他请求使用。

vLLM 论文的实验显示,在被比较的传统方案下,真正存放有效 token 状态的 KV cache 显存比例只有 20.4%–38.2% [1]。这并不意味着 GPU 其余显存全部永久“消失”,但在动态在线负载下,大量容量会因预留和碎片,无法及时容纳新请求。

结果是一个反直觉的现象:GPU 可能还有计算余力,却因为无法放入更多请求的 KV cache 而不能扩大 batch。生产线仍能运转,仓库管理却阻止了新订单进场。

PagedAttention:把连续上下文与连续显存解绑 #

PagedAttention 的核心思想可以用一句话概括:逻辑上的连续,不必等于物理上的连续。

操作系统早已在使用类似方法。程序看到的是连续的虚拟地址,但这些数据在物理内存中可以分散存放;操作系统通过页表(page table),把虚拟页(page)映射到不同的物理页。

vLLM 把 KV cache 切分为固定大小的块(block),再通过块表(block table)建立逻辑 block 与物理 block 之间的映射:

逻辑 KV block 0  ──→  物理 block 7
逻辑 KV block 1  ──→  物理 block 1
逻辑 KV block 2  ──→  物理 block 12

PagedAttention 通过 block table 将逻辑连续的 KV cache 映射到不连续的 GPU 物理显存

图 2:PagedAttention 解除了逻辑上下文连续性与物理显存连续性之间的绑定,并在新 token 到来时按需申请 block。

对于模型而言,这些 token 仍然属于同一段连续上下文;对于 GPU 显存而言,它们不必挨在一起。PagedAttention 计算内核(kernel)在执行 attention 时,根据 block table 找到相应的物理块并完成计算。

如果继续沿用仓库类比,传统方案要求每个订单占用一整片连续仓位;PagedAttention 则把货物放进统一规格的标准箱,哪里有空位就放到哪里,再由清单记录每个箱子的位置。订单在业务上仍然完整,物理存储却获得了更大的灵活性。

按需分配:用多少,申请多少 #

vLLM 不再为未知的最大输出长度一次性预留整段显存,而是在生成过程中逐步申请物理 block。

假设 block size 是 16,一个请求当前共有 35 个 token 的 KV cache,它只需要 3 个 block:前两个已经装满,第三个装了 3 个 token。此时,最多只有最后一个 block 中的 13 个位置暂时未使用,而非为可能达到的 2048 token 提前保留剩余空间。

原始论文将 16 个 token 作为当时实践中的默认 block size。block 太小,会增加映射、调度和 kernel 访问开销;block 太大,又会增加最后一块中的内部碎片,降低共享机会。因此,block size 本身就是硬件并行效率与内存利用率之间的折中 [1]。

写时复制(Copy-on-Write):相同部分不必重复保存 #

分页还让 KV cache 共享变得自然。

假设同一个 1000 token 的 prompt 需要生成四个候选答案。朴素方案可能为四条生成路径分别保留一份 prompt 的 KV cache。实际上,在第一个输出 token 产生之前,四条路径的历史完全相同。vLLM 可以让它们引用同一组物理 block;只有当答案开始分叉时,才为各自的新 token 分配不同 block。

这种写时复制机制会先共享同一份只读内容,真正发生修改时再复制。对于并行采样(parallel sampling)和束搜索(beam search),这能显著减少重复存储 [1]。

分页不是没有代价 #

PagedAttention 并不是一种“只赚不亏”的抽象。block table 引入间接寻址,kernel 还需要处理非连续内存和不同序列长度。vLLM 原始论文测得,在其评测的模型和工作负载下,完整 vLLM 系统相比 FasterTransformer 和 Orca,在维持相近延迟水平时实现了 2–4 倍吞吐提升;在长序列、大模型和复杂解码场景下,提升效果更明显 [1]。论文报告的是完整系统相对特定基线的端到端结果,不能完全归因于单个内核的优化。

这是一个很典型的系统工程结论:局部更快,不等于整体更快;局部增加一点开销,也可能换来更大的全局收益。 如果只比较单个 kernel 的延迟,就会错过 PagedAttention 真正要解决的问题。

只有“放得下”还不够:请求还要排得好 #

更高的显存利用率只是创造了容纳更多请求的可能性。要让 GPU 真正忙起来,还需要调度器及时把请求送进计算队列。

在讨论系统设计时,也应区分 GPU 利用率与有效吞吐(goodput)这两个层次:前者侧重资源被利用的程度,后者侧重在满足服务等级目标(Service-Level Objective,SLO)的前提下,系统单位时间完成的请求数 [4]。

传统静态批处理(static batching)像一辆必须等所有乘客到终点才允许换人的班车。只要同一 batch 中还有一个长请求没有结束,即使其他位置已经空出,新请求也只能继续等待。

Orca 在 2022 年提出迭代级调度(iteration-level scheduling):调度器在每一轮生成后重新组织 batch,让已完成的请求退出,让等待中的新请求进入 [2]。今天常说的 continuous batching,正是沿着这种细粒度调度思想发展而来。

例如,一个 batch 中有四个请求,其中三个在本轮生成后结束,第四个还要继续生成 500 个 token。静态批处理可能让三个位置一直空着;continuous batching 则可以在下一轮立即补入三个新请求。

static batching 会让已经完成的请求留下空槽,而 continuous batching 能在下一轮补入新请求

图 3:continuous batching 在每轮生成后重新组织 batch,使已完成请求退出,并让等待中的请求及时进入。

因此,更严谨的说法不是“PagedAttention 发明了 continuous batching”,而是两者解决了互补问题:

  • continuous batching 决定“哪些请求应该参加这一轮计算”;
  • PagedAttention 决定“这些不断进入、退出和增长的请求如何使用显存”。

如果没有灵活调度,节省出来的显存难以及时转化为吞吐;如果没有高效内存管理,调度器即使想加入更多请求,也可能找不到足够空间。vLLM 的价值,正在于把 attention 计算内核、块管理器与请求调度器共同设计成一个完整的推理服务引擎。

从“放好缓存”到“复用缓存”:SGLang 的 RadixAttention #

PagedAttention 首先解决“KV cache 如何高效存放”。但在聊天、检索增强生成(Retrieval-Augmented Generation,RAG)和智能体(Agent)工作流中,另一个问题越来越重要:已经计算过的 KV cache,能不能被后续请求再次使用?

考虑一个零售客服 Agent。每个请求前面可能都有三段相同内容:

  1. 一段固定的系统提示词(system prompt);
  2. 相同的退换货政策与商品目录;
  3. 当前用户自己的问题。

前两部分可能有几千个 token,并且被大量用户重复使用。如果每个请求都从头执行 Prefill,就像每接待一位顾客,都重新打印并阅读一次完整员工手册。模型得到的结果并无变化,计算却被一遍遍重复。

SGLang 提出的 RadixAttention,把 token 序列与 KV cache 组织成基数树(radix tree),并配合最近最少使用(Least Recently Used,LRU)淘汰和缓存感知调度管理缓存机制 [3]。它不仅复用预定义的固定 prompt,还能自动发现多层、动态的公共前缀结构。

例如,下列三个请求具有逐层分叉的前缀:

system prompt → 退货政策 → 商品 A → 问题 1
                         └→ 问题 2
              └→ 送货政策 → 问题 3

RadixAttention 使用 radix tree 组织并复用不同请求之间完全相同的 token 前缀

图 4:RadixAttention 对公共 token 前缀进行分层复用,并结合命中路径、引用状态和淘汰策略管理 KV cache。

Radix tree 可以让相同的 system prompt 只保存一次,让“退货政策”和“商品 A”等内容可以继续被相应分支共享。缓存不足时,系统再根据 LRU 等策略淘汰较长时间未使用、且没有被运行中请求引用的分支。

可以把 PagedAttention 与 RadixAttention 的侧重点理解为:

  • PagedAttention 侧重解决 KV cache 的分页布局、按需分配以及生成分支之间的块共享;
  • RadixAttention 在分页存储之上,用 radix tree 保留和索引跨调用、跨请求的公共 token 前缀,并结合淘汰与缓存感知调度提高复用率。

对于单轮短问答,前缀缓存的收益可能有限;但对于多轮对话、少样本提示(few-shot prompting)、RAG、自洽性采样(self-consistency)和 Agent,重复前缀越长、分支越多,缓存复用的价值就越明显 [3]。此时,调度问题也从“谁先到就先运行”升级为“怎样排序请求,才能在延迟公平性与缓存命中率之间取得平衡”。

这意味着 AI Infra 正在从无状态的模型调用,走向对推理状态的系统化管理。

从单机引擎到集群:为什么要分离 Prefill 与 Decode #

当服务规模继续扩大,单张 GPU 内的显存优化还不够。Prefill 和 Decode 虽然属于同一次推理,却具有明显不同的资源特征 [4]:

  • Prefill 一次性处理整个 prompt,矩阵运算规模较大,通常更偏向计算密集型;
  • Decode 在标准自回归解码中,每轮通常为每个活跃序列生成一个 token,却需要反复读取模型权重和持续增长的 KV cache,因此更容易受到内存带宽和调度抖动的影响。

可以想象两个用户同时使用服务。用户 A 上传一份很长的合同并要求总结,触发大型 Prefill;用户 B 正在进行实时聊天,希望每个 token 都稳定输出。如果两者共用同一组 GPU,大型 Prefill 可能插入用户 B 的 Decode 过程,使其突然“停顿一下再继续说”。平均吞吐也许仍然不错,但用户感受到的尾延迟已经恶化。

DistServe 等工作由此提出预填充–解码分离(Prefill–Decode disaggregation):把两个阶段放到不同 GPU 或资源池中,分别优化 TTFT 和 TPOT,再通过高速互联传输 KV cache [4]。

在标准的自回归 Decode 中,TPOT 更能反映请求级别的平均解码速度,而 token 间延迟(Inter-Token Latency,ITL)则用于衡量相邻输出 token 的实际间隔及其抖动。

Prefill 和 Decode 分别运行在不同 GPU 资源池,并通过网络将 KV cache 从 Prefill 节点传递给 Decode 节点

图 5:P/D 分离允许 Prefill 与 Decode 独立配置和扩缩容,但延迟隔离收益必须覆盖 KV cache 的传输成本。

它类似于把餐厅的备料区与出餐区分开:备料可以追求批量效率,出餐则更重视稳定节奏。分开后,两类工作可以采用不同的并行策略并独立扩缩容,也能减少大型 Prefill 对 Decode 的直接干扰。

但 P/D 分离不是免费的午餐。Prefill 完成后,KV cache 必须传给 Decode 节点,这会消耗带宽并增加同步、路由及故障处理的复杂度。服务规模不大、prompt 较短、网络较慢或负载不稳定时,传输成本可能抵消其隔离收益 [4]。

vLLM 的文档也明确指出,其实验性分离式预填充(Disaggregated Prefilling)实现并不提高整体吞吐,主要用于分别调节 TTFT 与 ITL,并控制尾部 ITL [5]。但这并不否定 DistServe 等系统通过资源配置优化在 SLO 约束下提升 goodput 的结果 [4]。

从 vLLM 看 AI Infra 的演进逻辑 #

从 Orca 的 iteration-level scheduling,到 vLLM 的 PagedAttention,再到 SGLang 的 RadixAttention 与集群级 P/D 分离,可以看到一条相当清晰的演进路径:

  1. 让 batch 动起来: 请求可以在 token 迭代之间动态进入和退出;
  2. 让显存变灵活: KV cache 从连续预分配变为分页、按需分配;
  3. 让状态可复用: 公共前缀不再反复计算,缓存开始影响调度;
  4. 让阶段可拆分: Prefill 与 Decode 可以独立配置和扩缩容;
  5. 让系统围绕 SLO 优化: 目标从追求峰值每秒浮点运算次数(Floating-Point Operations per Second,FLOPS)转向 TTFT、TPOT、吞吐、goodput 与单位 token 成本的综合权衡。

这里有一个容易被忽略的变化:优化对象正在不断上移。

最初,人们关注的是一个算子能否更快;随后关注单张 GPU 能否容纳更大 batch;再往后,问题变成跨请求缓存怎样复用、不同 GPU 怎样分工、网络如何传输 KV cache,以及整个集群能否在延迟约束内完成更多有效请求。

这也解释了为什么相同的模型、相同数量的 GPU,在不同推理服务系统上可能呈现出完全不同的并发能力、尾延迟与成本结构。拥有 GPU 不等于能够高效使用 GPU;模型能够运行,也不等于它能稳定、经济地服务真实用户。

工程实践中,不要只看 GPU 利用率 #

“怎么让 GPU 不闲着”很容易被简化成追求更高的 GPU 利用率,但这个概念本身包含多种衡量口径:GPU 忙碌时间(busy time)表示设备执行任务的时间比例;流式多处理器活跃度(Streaming Multiprocessor Activity,常记作 SM Active)表示采样区间内至少有一个线程束(warp)处于活跃状态的时间比例,但即使等待内存的线程束也可能被计为活跃 [6];模型浮点运算利用率(Model FLOPs Utilization,MFU)则衡量实际模型计算相对硬件理论峰值的比例 [7]。MFU 通常更常用于训练或受控基准测试;而在在线推理场景下,还需要明确模型 FLOPs 的估算方法和具体的负载口径。任何单一指标都可能产生误导。

例如,把 batch 做得很大,通常能提高吞吐量和计算利用率,却也可能让新请求等待更久;尽量保留所有前缀缓存,可以增加命中率,却可能挤占活跃请求所需的 KV cache 空间;强行进行 P/D 分离,可能改善 Decode 尾延迟,也可能让网络传输成为新的瓶颈。

因此,一个生产级推理系统至少需要同时观察:

  • TTFT(Time to First Token): 用户何时能看到第一个 token;
  • TPOT(Time per Output Token): 一个请求在解码阶段生成后续 token 的平均耗时;
  • ITL(Inter-Token Latency): 相邻输出 token 之间的实际时间间隔,其尾部延迟分布可以反映卡顿和抖动情况;
  • throughput: 单位时间内总共处理的 token 或请求数量;
  • goodput: 在满足 TTFT、TPOT 或尾部 ITL 等服务等级目标(SLO)的前提下,系统单位时间能够完成的请求数量 [4];
  • prefix cache hit rate(前缀缓存命中率): 查询前缀缓存时,命中的 token 数量占被查询 token 总数的比例;命中率越高,通常意味着可以避免的预填充计算越多 [8];
  • 单位 token 成本: 吞吐量的提升是否真正转化为资源成本的下降。

真正成熟的 AI Infra,不是堆叠最多技术名词,也不是让某个监控数字长期保持 100%,而是根据请求长度、并发模式、缓存命中率、网络拓扑和业务 SLO 选择合适的组合。

结语:前沿系统创新,常常来自经典思想的重新组合 #

在我看来,PagedAttention 最有启发性的地方,不只是它把显存利用率提高了多少,而是它展示了一种在计算机科学中反复出现的创新方式:真正改变前沿系统的思想,未必来自全新的数学公式,也可能来自对经典抽象的重新理解。

虚拟内存、分页、写时复制、缓存、调度和资源隔离都不是新概念。但当它们被重新放进 LLM 推理的上下文后,就构成了现代 AI Infra 的基础。

所谓“榨出硅的极限”,并不是让每一个 kernel 永远满载,也不是把 GPU 利用率这个单一数字推到 100%。准确地说,是让计算、显存、带宽与请求调度彼此匹配,在真实延迟约束下,用更少资源完成更多有效工作。

下一阶段的竞争,不只是谁能训练出更大的模型,也是谁能把模型的每一次计算,变成更可靠、更低成本、更可扩展的服务。

参考资料 #

[1] Kwon, W. et al. Efficient Memory Management for Large Language Model Serving with PagedAttention, Symposium on Operating Systems Principles (SOSP 2023).

[2] Yu, G.-I. et al. Orca: A Distributed Serving System for Transformer-Based Generative Models, USENIX Symposium on Operating Systems Design and Implementation (OSDI 2022).

[3] Zheng, L. et al. SGLang: Efficient Execution of Structured Language Model Programs, Conference on Neural Information Processing Systems (NeurIPS 2024).

[4] Zhong, Y. et al. DistServe: Disaggregating Prefill and Decoding for Goodput-optimized Large Language Model Serving, USENIX Symposium on Operating Systems Design and Implementation (OSDI 2024).

[5] vLLM Documentation. Disaggregated Prefilling (experimental), v0.26.0.

[6] NVIDIA Data Center GPU Manager (DCGM) Documentation. Profiling Metrics.

[7] NVIDIA NeMo Documentation. Performance Summary.

[8] vLLM Documentation. Metrics: Prefix Cache Metrics, v0.11.0.

延伸观看 #