PagedAttention原理详解:GPU显存优化与模型路由技术

引言:AI推理背后的显存工程
随着大语言模型(LLM)的规模持续膨胀,如何高效利用GPU显存已经成为AI基础设施领域最核心的挑战之一。近日,一位技术研究者(LinkedIn: gustavkeller)在Reddit上分享了一份关于GPU显存机制、PagedAttention以及模型路由的完整技术白皮书,引发了ML与AI Infra社区的广泛讨论。

这份白皮书虽然作者自谦"语法可能不够完美",但其技术价值不容忽视。本文将围绕这三大核心主题,梳理现代大模型推理系统在显存管理层面的关键工程实践与优化思路。
GPU显存机制:大模型推理的隐形瓶颈
显存消耗的三大来源
在大模型推理场景中,GPU显存的消耗主要来自三个部分:模型权重、激活值以及KV Cache(键值缓存)。其中,模型权重是固定开销,而KV Cache则会随着序列长度和并发请求数量的增加而急剧膨胀。
以一个典型的Transformer模型为例,每处理一个token,模型都需要为每一层缓存对应的Key和Value向量。当上下文窗口达到数万token、同时服务成百上千个并发请求时,KV Cache所占用的显存往往超过模型权重本身。
KV Cache是Transformer架构在自回归生成过程中的核心优化机制。在标准的自注意力计算中,每生成一个新token都需要重新计算所有历史token的Key和Value向量,计算复杂度为O(n²)。KV Cache通过缓存已计算过的Key和Value矩阵,使得每一步生成只需计算当前token的Query与缓存的Key做点积注意力,将增量计算降为O(n)。然而,这一优化是以空间换时间的典型策略——对于一个具有L层、隐藏维度为d、注意力头数为h的模型,单个序列长度为n时,KV Cache的显存占用为2×L×n×h×d_head×sizeof(dtype)字节。以LLaMA-70B为例,80层、64个注意力头、每头128维度,FP16精度下单个序列4096长度的KV Cache约需5GB显存,这解释了为何KV Cache管理成为推理工程的核心挑战。
值得深入理解的是Transformer自注意力机制与自回归生成的本质。Transformer架构的核心是自注意力(Self-Attention)机制,其通过Query、Key、Value三组向量的交互来捕捉序列中任意两个位置之间的依赖关系。在自回归生成模式下,模型逐token生成输出,每生成一个新token都需要"看到"之前所有已生成的token。这意味着注意力计算的规模随序列长度线性增长(在缓存Key/Value的情况下),而如果不做缓存则为二次增长。这种自回归特性使得推理过程本质上是序列化的,每个token的生成依赖于前一个token的输出,形成了推理延迟的根本瓶颈。
GPU显存层次结构与带宽瓶颈
理解显存消耗问题还需要认识GPU的存储层次结构。GPU的存储层次从快到慢依次为寄存器、共享内存/L1缓存(数十TB/s带宽)、L2缓存和全局显存(HBM,约2-3TB/s带宽)。现代推理工作负载的核心瓶颈往往不是计算能力(FLOPS),而是显存带宽——即从HBM读取模型权重和KV Cache的速度。这被称为"内存带宽受限"(memory-bandwidth bound)问题。以NVIDIA A100为例,其提供80GB HBM2e显存和2TB/s的带宽,但在自回归解码阶段,每生成一个token几乎需要读取整个模型权重一次,导致GPU计算单元大部分时间在等待数据加载。这也解释了为什么批处理(batching)对推理吞吐如此重要——多个请求共享同一次权重读取可以大幅提升计算效率。
传统显存分配方式的两大缺陷
传统的推理框架通常为每个请求预先分配一块连续的显存区域来存储KV Cache,且分配大小往往按照最大可能序列长度来预留。这种做法带来了两个严重问题:
- 内部碎片化:实际序列长度远小于预留空间,导致大量显存被浪费。
- 外部碎片化:连续内存块的分配与释放会在显存中留下无法有效利用的空隙。
研究表明,这种静态分配方式下,显存的实际利用率可能低至20%-40%,极大地限制了系统能够承载的并发量。这正是PagedAttention要解决的核心问题。
PagedAttention原理:借鉴操作系统的分页思想
核心设计理念与工作机制
PagedAttention的灵感来自操作系统中的虚拟内存与分页机制。它将KV Cache划分为固定大小的"块"(block),每个块可以存储固定数量token的键值对。这些块在物理显存中无需连续存储,而是通过一张"块表"(block table)来维护逻辑序列与物理块之间的映射关系。
要理解PagedAttention的设计精妙之处,有必要回顾其所借鉴的操作系统分页机制——这是计算机科学中最经典的抽象之一。在操作系统中,物理内存被划分为固定大小的页帧(通常4KB),进程的虚拟地址空间通过页表(Page Table)映射到不连续的物理页帧上。这一设计解决了早期连续内存分配导致的外部碎片问题,使得进程无需关心物理内存的实际布局。此外,操作系统还引入了写时复制(Copy-on-Write, COW)技术——当多个进程共享同一物理页时,只有当某个进程尝试写入时才会真正复制该页,从而节省内存并降低fork等操作的开销。PagedAttention将这些已被验证数十年的操作系统设计模式迁移到GPU显存管理领域,体现了系统设计中跨领域知识迁移的工程智慧。
这一设计带来了根本性的改变:显存不再需要为每个请求预留连续的大块空间,而是按需分配、动态增长。当一个序列生成新token时,系统只需分配新的物理块并更新映射即可。
PagedAttention的三大核心优势
采用分页机制后,推理系统能够获得以下收益:
- 近乎消除显存碎片化:由于以固定块为单位分配,内部碎片仅限于最后一个未填满的块,外部碎片则被完全消除。
- 支持显存共享:在并行采样、Beam Search等场景下,多个序列可以共享相同的前缀块,通过写时复制(Copy-on-Write)机制避免重复存储。
- 显著提升推理吞吐量:更高的显存利用率意味着系统能够容纳更多并发请求,从而大幅提升整体吞吐量。
关于第二点,Beam Search是一种广泛应用于机器翻译、文本生成等任务的搜索算法,它在每一步保留概率最高的k个候选序列(beam),最终从中选取最优结果。传统实现中,k个beam各自维护独立的KV Cache副本,但实际上它们在分叉之前共享完全相同的前缀历史。类似地,在温度采样等场景下,同一个prompt可能需要生成多个不同的候选回复,这些回复的KV Cache在prompt阶段完全相同。PagedAttention通过引用计数和写时复制机制,允许多个序列共享相同的物理块,只有当某个序列产生分歧时才分配新块。这种设计在beam_width=4的典型配置下可节省约55%的KV Cache显存开销。
这一技术正是vLLM等高性能推理引擎能够实现数倍吞吐提升的关键基础。vLLM是由UC Berkeley的Sky Computing Lab于2023年开源的高性能大模型推理引擎,其核心创新正是PagedAttention算法。在发布时的基准测试中,vLLM相比当时主流的HuggingFace Transformers推理方案实现了2-4倍的吞吐量提升。vLLM通过连续批处理(Continuous Batching)和PagedAttention的结合,能够动态地将新请求插入正在执行的批次中,避免了传统静态批处理中短序列等待长序列完成的资源浪费。同类竞品包括NVIDIA的TensorRT-LLM(侧重于编译优化和量化加速)、DeepSpeed-FastGen以及近期兴起的SGLang等,它们各有侧重但都在显存管理方面借鉴了类似的分页思想。
连续批处理与PagedAttention的协同
连续批处理(Continuous Batching)是与PagedAttention配合使用的另一项关键技术。传统的静态批处理(Static Batching)将一组请求打包后统一处理,必须等待批次中最长的序列完成后才能释放整个批次的资源。这导致短序列生成完毕后仍需等待,造成严重的GPU空闲浪费。连续批处理(也称为迭代级批处理,Iteration-level Batching)则在每个解码步骤后检查哪些序列已经完成,立即将其移出并插入等待队列中的新请求。这种细粒度的调度将GPU利用率从传统方案的30-50%提升至80%以上。Orca系统(2022年发表于OSDI)是最早提出这一概念的学术工作,随后vLLM、TensorRT-LLM等系统均采纳了这一设计。连续批处理与PagedAttention是天然互补的——前者解决了调度层面的效率问题,后者解决了显存层面的碎片问题。
量化技术:与PagedAttention正交的显存优化路径
除了PagedAttention等显存管理层面的优化,量化(Quantization)是另一条减少显存占用的重要路径。量化通过降低模型权重和/或KV Cache的数值精度(例如从FP16降至INT8或INT4)来直接压缩存储空间。GPTQ、AWQ、SqueezeLLM等方法可以将模型权重压缩至4-bit而仅损失极小的精度。针对KV Cache的量化(如KV Cache INT8量化)则可以在不修改模型权重的情况下将KV Cache显存减半。这些技术与PagedAttention并不冲突,而是正交互补的——量化减小每个block中存储的数据量,PagedAttention优化block的分配与管理方式。在实际生产部署中,两者通常结合使用以实现最大化的成本效益。
它把AI推理从粗放的资源占用,推向了精细化的显存工程。
模型路由技术:多模型服务的智能调度策略
从单一模型到多模型编排
随着AI应用日趋复杂,单一模型往往难以在成本、延迟与质量之间取得最佳平衡。模型路由(Model Routing)应运而生,它的核心思想是根据请求的特征,将其动态分发到最合适的模型上处理。
例如,一个简单的问答请求可以交给参数量较小、响应更快的模型;而涉及复杂推理的任务则路由到更强大但成本更高的模型。这种策略在保证服务质量的同时,能够显著降低平均推理成本。
模型路由的概念在工业界已有多种成熟实践。OpenAI的GPT-4系统据推测内部采用了混合专家(Mixture of Experts, MoE)架构,本质上是token级别的路由。而在服务层面,Martian、Unify等创业公司提供了请求级别的模型路由服务,根据输入的语义复杂度、领域特征等自动选择最合适的LLM API。学术界也有RouteLLM(由LMSys团队提出)等研究,通过训练专门的路由模型在强弱模型之间做最优分配,在保持95%以上质量的情况下将成本降低40-60%。
混合专家架构(MoE)与服务级路由的内在联系
理解模型路由的全貌,有必要认识混合专家(Mixture of Experts)架构与服务级路由之间的思想一脉相承。MoE是一种稀疏激活的神经网络设计,其核心思想是将模型的前馈网络(FFN)层拆分为多个"专家"子网络,每个token通过一个门控网络(Router/Gate)仅激活其中的Top-K个专家(通常K=1或K=2)。这使得模型可以在保持极大参数规模(提升模型容量)的同时,每次前向传播只使用一小部分参数(控制计算成本)。例如,Mixtral 8x7B拥有约46.7B总参数,但每个token仅激活约12.9B参数,推理成本接近一个13B的稠密模型。MoE架构中的token级路由与服务层面的请求级路由在哲学上一致:都是通过智能分发来实现效率与质量的平衡。但MoE的路由发生在模型内部的每一层,决策粒度为单个token;而服务级路由发生在系统入口,决策粒度为整个请求。
路由策略设计的关键维度
有效的模型路由通常需要综合考虑以下因素:
- 任务复杂度评估:通过轻量级分类器或启发式规则判断请求的难度。典型的路由分类器可以使用BERT等轻量模型在数毫秒内完成难度评估,其计算开销相比路由决策带来的成本节约几乎可以忽略不计。
- 成本与延迟约束:在满足SLA的前提下选择性价比最优的模型。
- 负载均衡:避免请求过度集中于某个模型实例,导致显存与算力瓶颈。
模型路由与前述的显存管理技术形成互补:PagedAttention等显存优化技术提升了单个模型的服务能力,而模型路由则在系统层面实现了资源的最优配置。二者共同构成了现代AI推理基础设施的技术骨架。
总结:AI推理基础设施的工程价值
这份白皮书所涵盖的三大主题——GPU显存机制、PagedAttention与模型路由,实际上勾勒出了大模型推理系统优化的完整脉络:从底层的显存物理管理,到中层的KV Cache分页算法,再到上层的模型调度策略。
对于身处ML工程或AI基础设施领域的从业者而言,理解这些机制不仅有助于选择合适的推理框架(如vLLM、TensorRT-LLM等),更能在成本敏感的生产环境中做出关键的架构决策。在大模型走向规模化落地的今天,AI Infra的工程能力正成为决定产品竞争力的核心要素之一。
注:本文基于Reddit社区分享的技术白皮书整理。原作者非专业写作者,如需深入交流可通过其LinkedIn(gustavkeller)联系。
相关推荐

Claude自主设计蛋白质成功率35%,远超人类专家水平
Anthropic的Claude模型在自主设计靶向疾病蛋白质任务中取得35%实验成功率,远超人类专家10%-15%的平均水平。本文深入解析这一湿实验验证成果对生物医药行业的潜在影响。

Perplexity Discover多语言支持突然消失,国际用户为何不满?
Perplexity Discover新闻资讯功能突然取消多语言支持,仅保留英文内容,引发国际用户强烈不满。本文分析功能回退的可能原因,探讨AI产品国际化面临的资源权衡与用户信任挑战。
