☰
PagedAttention原理,KV缓存也能分页
2026/10/7 19:31:57 网站建设 项目流程

摘要

大模型推理的瓶颈常常不是算力,而是装不下更多请求的 KV 缓存。vLLM 的 PagedAttention 借用操作系统的虚拟内存思路,把 KV 缓存切成固定大小的块,按需分配、用块表映射,几乎消除了显存碎片,还能在多个请求之间共享块。原论文(SOSP 2023)的实验里,相同延迟下吞吐比 FasterTransformer 和 Orca 高 2 到 4 倍。我还写了一个小模拟器,把这套机制跑了一遍。

背景与问题

自回归生成时,每个已经处理过的 token 都要保存它在每一层的 Key 和 Value,这就是 KV 缓存。它很大:论文里 13B 的 OPT 模型,单个 token 就要 800KB(2 个向量 × 5120 隐层维度 × 40 层 × 2 字节),一个 2048 token 的请求最多占 1.6GB。在 40GB 的 A100 上,约 65% 的显存给了权重,接近 30% 用来存 KV 缓存(其余是激活等),能同时服务多少请求,几乎由这 30% 决定。

早期系统把一个请求的 KV 缓存放在一段连续显存里,并且按最大长度(比如 2048)提前预留。论文把浪费分成三类:

  • 预留浪费:为将来要生成的 token 先占着位置,但这些位置在整个请求期间别人用不了
  • 内部碎片:预留的最大长度远大于实际长度,用不上的部分白白浪费
  • 外部碎片:不同请求预留的大小不同,分配器留下一堆零散缝隙

论文在实验里测到,现有系统里只有 20.4% 到 38.2% 的 KV 缓存显存真正存着 token 状态。另外,连续存放还让共享变得很难:同一个请求的多个采样结果、多个请求共用的系统提示词,它们的 KV 本来可以共用,却各存一份。

核心思路与优势

像操作系统一样分页

PagedAttention 的类比很直接:块是页,token 是字节,请求是进程。

  • 把 KV 缓存切成固定大小的块,vLLM 论文默认每块 16 个 token
  • 每个请求有一张块表,把它的逻辑块(按顺序排的)映射到物理块(显存里的实际位置)
  • 物理块不必连续,用到哪里分配到哪里,新 token 写满一个块才再要一个新块

这样三类浪费就都被压住了:不再提前预留最大长度;浪费只可能出现在每个请求最后一个没写满的块里;所有块一样大,也就没有外部碎片。论文的说法是 KV 缓存显存接近零浪费。

代价是注意力内核要先查块表再取数据。论文测得注意力内核延迟比高度优化的 FasterTransformer 高 20% 到 26%,但这只影响注意力算子,不影响线性层,端到端仍然大幅领先:论文 2023 年的实验里,相同延迟下吞吐比 FasterTransformer 和 Orca 高 2 到 4 倍,原因是显存省下来之后,一次能塞进更多请求。

块级共享和写时复制

块表还带来第二个好处:共享。物理块上有引用计数,多个序列的块表可以指向同一个物理块。

  • 并行采样:同一个提示词生成多个候选,提示词部分的块全部共用
  • 束搜索:不同候选的前缀大部分相同,共享比例更高
  • 共享的块如果某个序列要往里写新内容,就用写时复制:先复制一份给它,再改,其他序列不受影响

论文给出的显存节省:并行采样 6.1% 到 9.8%,束搜索 37.6% 到 55.2%(Alpaca 数据);换成对话更长的 ShareGPT 数据,分别是 16.2% 到 30.5% 和 44.3% 到 66.3%。

显存不够时的调度

块是按需增长的,所以显存可能中途耗尽。vLLM 的做法是:先来先服务,显存不够就抢占,而且整个序列的块要么全部驱逐,要么都不驱逐,同一个请求里的多个序列(比如束搜索的候选)一起抢占、一起恢复。被驱逐的块有两种恢复方式:

  • 交换:搬到 CPU 内存,需要时再搬回来
  • 重算:丢掉,之后重新计算这些 token 的 KV

论文发现:块太小时交换要做大量零碎的 CPU-GPU 小传输,开销很大;重算不读 KV 块,开销与块大小无关。所以小块时重算更划算,大块时交换更划算,但即便在交换占优的情况下,重算也最多比交换慢 20% 左右;块大小在 16 到 64 之间时,两者的端到端性能相当。

块大小怎么选

块太小,GPU 读 KV 缓存的效率下降;块太大,内部碎片变多,共享的机会也变少。论文在 ShareGPT 数据上测到 16 到 128 都不错,在较短序列的 Alpaca 数据上 16 和 32 好、更大的块性能明显变差,最终默认 16。

后来的演进:自动前缀缓存

分块之后,还能做一件事:把已经算好的块缓存下来,新请求只要前缀相同就直接复用。vLLM 当前的实现是基于哈希的自动前缀缓存:每个块的哈希由前一个块的哈希、本块的 token,以及 LoRA 编号、多模态输入哈希、缓存盐值等附加信息共同决定,所以哈希能唯一标识"这个块 + 它前面的全部内容"。几个要点:

  • 只缓存写满的块
  • 从 v0.11 起默认哈希算法是 sha256,降低碰撞风险;也可以通过--prefix-caching-hash-algo换成可跨环境复现的 sha256_cbor,或更快但非加密、碰撞风险理论上更高的 xxhash / xxhash_cbor
  • 多租户场景可以给请求加cache_salt,只有盐值相同的请求才能复用缓存,避免通过延迟差异推测别人缓存了什么内容

另外要注意:vLLM 官方仓库里paged_attention.md那篇讲内核实现的设计文档,开头就标明是基于原论文的历史文档,已不再描述今天 vLLM 的代码。想读当前实现,看注意力后端和 KV 缓存管理器的文档更靠谱。

面向人群

  • 想弄清 vLLM 为什么比朴素推理吞吐高的工程师
  • 需要给线上服务估算显存、调max-model-len之类参数的人
  • 学过操作系统,想看虚拟内存思想怎么用到大模型上的人
  • 面试或做技术分享,需要讲清 PagedAttention 的人

实践步骤

第一步:先算清楚 KV 缓存多大

每个 token 的 KV 大小 = 2(K 和 V)× 层数 × KV 头数 × 头维度 × 每个数的字节数。

论文里 OPT-13B 用的是传统多头注意力,KV 头数等于注意力头数,所以大。现在的新模型多用分组查询注意力(GQA),KV 头数少得多。我用本机的 Qwen2.5-1.5B-Instruct 配置算了一下:28 层、2 个 KV 头、头维度 128(1536 隐层 ÷ 12 个注意力头)、BF16 每个数 2 字节,每个 token 只要 28KiB,2048 个 token 才 56MiB,比 OPT-13B 单个 token 800KB 小了约 28 倍。所以你的模型到底有多吃 KV 缓存,要看它的配置,不能套用论文里的数字。

第二步:用一个小模拟器看清分页的效果

下面的数字是我写的一个玩具模拟器算出来的,不是 vLLM 的实测,目的是把机制跑一遍。设定:沿用论文的 OPT-13B 数字(每 token 800KB),给 KV 缓存 12GB(约 15000 个 token 位置),最大长度 2048,块大小 16,5000 个合成请求(提示词和输出长度服从对数正态分布,平均总长约 480 个 token),每个请求取生命周期中随机一刻的快照。

方案真正存着 token 的显存占比同时能容纳的请求数
连续预留 2048 个位置17.1%7
分页(块大小 16)97.9%40

分页方案里平均每个请求浪费 7.4 个位置,就是最后一块没写满的部分。要说明的是,这里"能容纳 40 个"是按当前长度算的快照,没有模拟后续增长、抢占和调度,真实系统不可能一直满载。但量级的差异能说明问题:连续预留把绝大部分显存都耗在了用不上的位置上。

核心的块管理只需要引用计数和写时复制:

BLOCK=16classBlockManager:def__init__(self):self.ref={}# 物理块编号 -> 引用计数self.next_id=0defnew_block(self):b=self.next_id self.next_id+=1self.ref[b]=1returnbdeffork(self,table):# 复制块表,只增加引用计数forbintable:self.ref[b]+=1returnlist(table)defappend_token(self,table,n_tokens):# 给序列再放一个 tokenifn_tokens%BLOCK==0:# 最后一块写满了 -> 新开一块table.append(self.new_block())elifself.ref[table[-1]]>1:# 要写的块被共享 -> 写时复制self.ref[table[-1]]-=1table[-1]=self.new_block()

用它模拟"提示词 300 个 token,每个采样再生成 100 个 token"的并行采样:

采样数不共享(块)共享(块)节省
2503236.0%
41004654.0%
61506060.0%

采样数越多省得越多(提示词在总长里占比越高,同理也越省),趋势和论文一致;具体数字和论文不同,因为这只是固定长度的玩具设定,工作负载完全不同。

第三步:在 vLLM 里和分页相关的参数

下面的参数和默认值取自 2026 年 10 月初 vLLM 主分支的缓存配置,你装的版本可能不同,以vllm serve --help为准:

  • --gpu-memory-utilization:这个 vLLM 实例能用的显存比例,默认 0.92。启动时在这个额度内扣掉权重、剖析得到的峰值激活等开销,剩下的大体划给 KV 缓存的块池。和别的进程共卡时要调低
  • --max-model-len:单个请求(提示词加输出)的最大长度,不指定就取模型配置里的上下文长度。块是按需分配的,它并不改变块池大小;它限制的是单个请求最多占多少块,启动时 vLLM 还会检查块池至少装得下一个满长度的请求,装不下就报错并给出估算的可用长度(设成 -1 或 auto 会自动选一个装得下的长度)
  • --block-size:块大小,不指定时默认 16,个别平台或注意力后端会自行调整,一般不需要改
  • 自动前缀缓存:当前版本默认开启,系统提示词很长、多轮对话多的场景收益最大
  • --prefix-caching-hash-algo:前缀缓存的哈希算法,默认 sha256

启动时 vLLM 会先剖析模型的显存占用,再算出能分出多少个 KV 块,日志里会打印类似GPU KV cache size: N tokens, Maximum concurrency for L tokens per request: X x的一行。后半句是按每个请求都用满最大长度算的,比较保守;用 N 除以你的平均请求长度,可以得到一个粗略的并发上限。

我的看法

PagedAttention 的价值不在于某个精巧的算法,而在于换了一个视角:KV 缓存的问题本质上是内存管理问题,而内存管理问题操作系统几十年前就解决过。把页、页表、引用计数、写时复制、换入换出这一整套搬过来,论文在 2023 年的实验里就换到了 2 到 4 倍的吞吐。

几点提醒:论文的 2 到 4 倍是在 2023 年的模型、基线和负载上测的,今天的推理引擎也在吸收分页的思想,别把它当成对任何场景都成立的加速比;我的模拟器是玩具,只证明机制,不证明 vLLM 的实际收益;新的混合注意力模型(滑动窗口、Mamba 等)对 KV 缓存管理提出了新要求,vLLM 官方已有专门的混合 KV 缓存管理器,但文档标注这个功能还处于早期阶段。

需要专业的网站建设服务?

联系我们获取免费的网站建设咨询和方案报价,让我们帮助您实现业务目标

立即咨询