显存浪费从 70% 降到 4%,大模型推理靠的是一套分页

跑过大模型推理的人,多半撞过显存这堵墙。

算力够用,显存不够,服务照样起不来。

卡再好,摆不下也是白搭。

PagedAttention 概念图,图中标注 LLM Memory Management 与 Paged Blocks,页脚为知乎 @deephub

真正的吃显存大户,是键值缓存。

模型是一个词一个词往下写的。

写下一个词,它得回头看前面所有的词。

注意力该落在哪些位置,全靠这些历史词算出来。

这些历史词会被投影成键和值,参与后面的计算。

为了不把这一步反复重做,投影结果被留在显存里。

留下来的这一份,就是键值缓存。

序列越长,这份缓存越大。

到了批量服务里,它就成了显存的第一大占用。

早年的服务框架,要求这些数据连续摆放。

请求一进来,先按最长的可能长度,切一整块连续显存。

假设上限是两千个词元,那就先把位置留足。

可实际输出,常常用不到那么长。

多出来的空间,就白白空着。

真正的浪费,藏在三个地方。

第一种,过量预留。

明明只写几十个词,位置却按两千个留。

第二种,内部碎片。

句子提前收尾,分到的那段还有一半没填。

第三种,外部碎片。

显存总账上还有空位,可东一块西一块,拼不出一段够大的连续空间。

传统方案里,这份缓存的可用显存最多能浪费六到八成。

八成,很吓人。

换一种摆法,局面就变了。

这种机制叫分页注意力(PagedAttention),思路借自操作系统的分页。

一个序列的缓存被切成若干块,每块固定装几个词元。

块大小可以设,经验值取十六。

这些块在显存里可以散着放,不必挨在一起。

系统另记一张块表,写清楚逻辑块对应哪块物理内存。

读写按逻辑顺序走,底下的物理位置不连续也没关系。

块,就是管理的单位。

要用多少,就申请多少。

下面看三种最典型的情形。

一个请求,被切成好几块

单序列生成动图:逻辑块经块表映射到分散的物理块,显存占用随生成过程变化

先看单序列生成。

输入一句话,模型一个字一个字往下接。

假设输入是五个词元,块大小取四。

第一个块装满四个,被指向一块随机挑的物理块。

剩下的一个,进第二个块,只占四分之一。

这个块还有三个空位。

再解出一个词,系统把它写进下一个空位,占用变成一半。

再解一个,就到了四分之三。

等这个块填满,系统再申请一块新的,挂到第三个逻辑块上。

新块只在需要时申请。

空着的槽位,于是只集中在一个没填满的块里。

前面那些块,早就装满了。

浪费被压到一个块的四分之一以内。

可用缓存的浪费,就从七成掉到四以下。

七成到四,差的不是硬件。

差的是分配思路。

这个账,靠的正是按需分配。

请求有多少,就占多少。

批处理也是同理。

一批请求各占各的块,互不牵连。

腾挪的空间,就是这么来的。

分支共享,改的时候才复制

并行采样与写时复制动图:两条输出分支共享同一批物理块

同一段输入,有时要出好几份结果。

比如编程助手,一次提问给出三份补全。

三份的开头完全一样,传统做法却要各存一份缓存。

三份就是三倍的开销。

分页的做法是共享。

两条分支先把逻辑块指向同一批物理块。

这批块的引用计数记成二。

等到两条分支开始写不同的词,系统才在块的层面复制。

把共享的那块数据拷进新申请的物理块。

原来那块的引用计数减一。

两条分支随后各写各的。

复制只落在一个块上,不是整段缓存。

并行生成这段时间,最多能共享三成的缓存显存。

改到哪,才复制到哪。

没动过的,一直共用。

路径被剪,块跟着回收

束搜索与剪枝动图:被剪枝路径的物理块按引用计数回收

束搜索会同时留几个候选句子。

每轮解码后比一比概率,高的留下,低的剪掉。

候选一路在变,传统实现要反复搬动大段内存。

分页这边,管块的方式像一棵会变的进程树。

历史相同的候选,逻辑块指向同一块物理内存。

某条路径被剪,它对应的逻辑块就销毁,相关物理块的引用计数减一。

计数掉到零,系统马上回收这块,把空间还回空闲显存池。

分叉,也只复制一块。

在束搜索这种负载下,省下来的显存能到三成七到五成五。

最后

碎片少了,服务器就能把更多请求塞进同一批处理。

同样的硬件,能同时跑的请求数最多提到四倍。

整体吞吐也就跟着往上走,翻个两到四倍。

作者还做了一个可以上手点的模拟器。

单序列生成,并行采样,束搜索,三种都覆盖。

可以一步一步走,看块表怎么更新,看复制什么时候发生。

想深挖,去翻那个开源仓库。

那里把每一步都拆开了。

说到底,这套机制解决的不是算得快不快。

它解决的是显存怎么摆,才不浪费。

浪费少了,同样的卡就能多接活。

对做 AI编程 的人来说,长任务能塞进更密的批里,跑起来更稳。

对做 AI对话 产品的人来说,同样的卡能接更多并发,单次成本更低。

对只想把活干完的人来说,这类 AI工具 最后只剩一句朴素的体验。

更快,更稳,更便宜。