资讯动态

KV Cache显存优化:PagedAttention的块映射与写时复制

发布时间:2026/10/19 10:34:34 来源:尧图企业网站定制
我先说说为什么会专门挑这个题目来写。如果跑过大模型推理服务尤其是自己搭过基于KV Cache生成的长对话、长文档总结这类场景你大概率遇到过这样的情况显存明明还剩下不少但只要并发稍微拉高一点就会出现out of memory或者同一个模型换了一批输入长度分布完全不同的业务数据显存占用曲线突然变得极不规律。我在排查这类问题的过程中一步步把PagedAttention这套机制吃透了。这套机制的牛X之处在于它不是靠“压缩”或者“量化”省显存而是靠“更聪明地管理KV Cache的摆放位置”来提升显存利用率顺带解决传统连续缓存方案里的碎片化、预分配浪费、并行采样时的共享复制成本这几个棘手问题。这篇文章我就把块映射、动态分配、写时复制这三个核心机制拆开讲顺便聊聊我在实际工程中观察到的收益和一些容易被忽略的细节。1. KV Cache为什么是“显存黑洞”先把这笔账算清楚想理解PagedAttention为什么有效先得明白KV Cache在传统实现里是怎么存的以及它到底吃掉多少显存。1.1 KV Cache的规模不是线性增长的是指数级的心头大患Transformer结构里每个token在做attention计算时会产生各自的Key和Value向量。自回归生成下一个token时新token要与已有所有token做attention已有的Key和Value必然要被反复使用。与其每次都重新计算一遍不如把之前算好的K、V缓存起来这就是KV Cache的由来。但问题在于KV Cache不是只有一层Transformer有多少层每一层都要缓存一份当前对话有多长就缓存多长的历史。公式可以写成KV Cache显存占用 ≈ 2K和V两份 × num_layers层数 × max_seq_len序列长度 × hidden_size每层维度 × precision_bytes精度fp16下为2字节拿一个7B量级的模型举例大致配置是32层、hidden size是4096、fp16精度、上下文长度支持4096。如果此时并发拉满到16条请求另外每个请求都接近最大长度这个峰值占用算出来大概是2 × 32 × 4096 × 4096 × 2 × 16 ≈ 32GB也就是说在推理这种规模的模型时光KV Cache一项就可以吃掉接近模型权重本身约14GB fp16两倍以上的显存。你以为模型是显存大头实际上KV Cache才是真正的吞金兽。1.2 传统“连续大数组”缓存的致命伤预分配和碎片化更麻烦的是传统实现给KV Cache分配显存的方式也很粗暴按最大序列长度给每条请求预留一块连续空间。也就是说只要模型配置里写了context length是4096就按4096个token的KV大小提前占好不管你实际只生成了20个token。这里有几个非常现实的问题预分配浪费线上业务的实际平均生成长度可能只有最大长度的20%到30%但显存已经被“透支”了这导致能同时服务的并发数被严重压低。长度波动造成的碎片就算你取消预分配动态给不同序列分不同大小的连续空间因为每条请求的长度完全不一样用完释放以后显存间会出现很多大小不一的空洞。新的请求可能刚好找不到足够大的连续空间明明总剩余空间足够却依然分配失败。无法共享一旦某条请求的KV Cache被放在一块连续区域里你想在并行采样时让多条子请求共享相同的前缀历史几乎不可能除非抄一遍数据。这种“按最大长度、要连续地址、各管各”的策略和操作系统里内存管理早期的“连续分配”方式几乎一模一样而操作系统早就用虚拟内存和分页解决了这些问题。PagedAttention正是把这一套类似页表的做法搬进了GPU显存管理。2. 块映射KV Cache的“虚拟地址表”是怎么搭的PagedAttention的核心思想是放弃“一条序列必须占用一段连续显存”的约束把KV Cache切分成等长的块可以散落地存放在任意物理位置。2.1 把大数组切成定长块每块只装固定数量的token假设一个块固定装16个token对应的KV数据。一条长度是50的序列原本需要分配一个容量为50的连续缓存现在只需要分配4个容量为16的块3个满块加1个装了2个token的尾部块。每个块是一个固定大小的显存区域。比如在某种实现里每个物理块的形状大致是block_size例如16个token × num_kv_heads × head_dim不同层之间的KV一般是独立的块集合但块的大小和分配逻辑是一致的。这样设计的好处是显存不再按“序列”这种不可预期的大小来切分而是按“块”这种固定颗粒度来切分分配和释放的管理粒度变得非常简单和平均。2.2 Block Table逻辑块号到物理块号的“页表”块切好了下一个问题是怎么找到一个序列的第0个块、第1个块放在物理显存的哪个位置答案是一张映射表也就是PagedAttention里的block table。每条序列都会维护自己的一张逻辑块号表。逻辑块号就是这条序列自己的“第0块”“第1块”“第2块”的概念编号而表里记录的值是对应物理块在全局显存中的实际块号。看个直观例子逻辑块号序列维度物理块号全局显存维度块内token数Block 0物理块 816/16Block 1物理块 2516/16Block 2物理块 316/16Block 3物理块 4712/16尾部块这条逻辑上连续、看起来是“从第0个token到第50个token”的KV序列实际在显存里被拆到物理块8、25、3、47这四个完全不相邻的位置。计算attention时只需要按block table找到对应物理块地址取出KV数据即可。对上层模型代码来说它感知不到这种离散存储只会觉得KV数据是按序可取的。2.2对应的内容其实可以继续展开两句物理地址到真实GPU显存地址的映射一般由块分配器或缓存分配器维护。每个物理块在分配时会记录对应的显存偏移量。当block table里记了“逻辑块2指向物理块3”实际读取时就会从物理块3的起始地址开始连续读块内部的数据仍然是连续的。这正是PagedAttention和“纯页式存储”的相似之处块内连续块间离散。2.3 块映射带来的直接收益碎片不再致命有了块映射之前说的“找不到连续大空间”的问题基本消失。因为每个块都是一样大的释放任意一块它就可以被任意其他序列使用。整个显存池被划分成等量的小格子只要剩余的空闲块数量足够任何序列都能获得存储空间不必关心这些空闲块是连在一起的还是散落在各个角落。这就像你有一堆统一规格的储物柜而不是一个大仓库。任何物品都能拆开放在任意多个柜子里柜子之间不需要彼此相连。碎片化的概念被大大稀释了——只要柜子够用大件物品也能塞进去。3. 动态分配显存不再“死占”而是用到哪分到哪块映射解决了“KV Cache怎么散落摆放”的问题动态分配解决的是“KV Cache什么时候拿、什么时候还”的问题。3.1 从“按最大长度预分配”到“按需逐块增长”传统做法在请求开始时就按max_seq_len一次性分配整块显存。PagedAttention的做法是请求开始时只分配容纳现有token数量的块。随着自回归生成每个新token产生后需要写入新的KV数据此时如果当前最后一个块已满就再申请一个新的物理块挂到block table末尾。这种“写满了再申请”的策略非常像操作系统里的请求调页demand paging按实际使用情况懒惰分配。效果很直接一个只生成50个token的请求就只占用4个块4×1664个token容量而不会占用最大长度对应的256个块。线上业务里大量短文本场景显存利用率可以瞬间翻倍。3.2 分配器与引用计数谁在用哪个块一目了然动态分配需要一套块管理器来记录全局所有物理块的状态。通常每个物理块有以下字段状态空闲 / 已分配。引用计数当前有多个少序列的逻辑块映射到这个物理块。绝大多数情况是1在做并行采样共享KV时可能大于1。物理地址对应的显存偏移或者缓存块张量中的行id。分配流程大致是模型发来“我需要一个新的物理块来存下一段KV”的请求块管理器从空闲列表中弹出一个块把引用计数设为1把物理块号返回给调用方调用方把物理块号追加到序列的block table末尾。释放流程则反过来序列终止时遍历它的block table每块引用计数减1减到0的块归还空闲列表。这里有一个容易踩的坑显存里的物理块池本身是怎么创造的常见的做法是在推理服务初始化阶段先预留一块连续显存然后把它均匀地切成物理块。但预留多少是个问题——预留太多本身就是浪费预留太少高并发时还是会OOM。成熟实现里一般会根据总显存、模型权重占用、最大并发数综合估算也有一些实现支持一定程度的动态扩容但复杂度和稳定性风险都会上升。3.3 动态分配对性能的副作用以及怎么消除动态分配不是没有成本。每次申请新块都是一次管理器层面的操作如果极端情况下每个token都触发一次新的块分配这个开销会被放大。为了减少分配频率块设计得越大分配次数越少但尾部浪费越多。我见过不少实际调优场景block size从8调到32分配次数能下降好几倍但显存碎片浪费率也会上升两者本身就是一个需要根据业务折中的参数。另一个副作用是物理块池预留过大时GPU显存里会有一块“冻结”的缓存区。如果你同时跑推理和训练这个区域不能随便让出去容易造成整体显存不灵活。所以如果要在自己的推理引擎里实现类似的机制建议把物理块池做成可配置的而不是写死。4. 写时复制并行采样里“共享前缀”的优雅解法块映射和动态分配解决的是单条请求的存储问题写时复制解决的是多条请求之间的公共KV数据要不要各存一份的问题。4.1 并行采样场景里共享前缀是天然的优化点做beam search、并行采样或者把一个prompt跑多个随机结果时通常会先生成一个共享的前缀序列然后在某个位置分裂出多条分支。比如用同一个prompt生成3条不同的回答前100个token的prompt本身是一样的这些token对应的KV Cache也完全一样。传统实现对每个生成样本都独立维护一份完整的KV Cache哪怕前面100个token一模一样也要抄三份。如果生成长度长、采样条数多这部分的显存浪费和拷贝开销非常可观。4.2 COW的触发条件与实现流程写时复制的出发点很朴素既然数据和前缀部分完全相同那就先让它们共用同一个物理块谁也别拷贝。只有在某个样本要写入这个共享块时才真正复制一份各改各的。具体流程可以简化成三步并行采样开始时多个序列的block table前几项都指向同一批物理块物理块的引用计数变成3。某个样本继续生成需要往一个共享的、还没写满的块里写入新token。这时发现它的block table里这个物理块引用计数大于1不能直接改原块触发复制。分配一个新物理块把原块内容拷贝过去引用计数减1修改触发样本的block table里对应项指向新块。新块的引用计数是1后续可以随意写入。这里有个容易被忽略的细节块是否写满决定了COW的触发时机。如果多个样本共享的块已经写满了后续每个样本各自要写新的token时是申请全新的空块来写这个新块本来就是各自独占的完全不需要复制。只有当共享块还有剩余槽位、而又不能再共用时才需要复制一整块出去。理解这一点很关键不然在断点调试时会觉得COW的行为忽隐忽现。4.3 COW的权衡省显存但不省计算写时复制省的是显存和初始化阶段的KV计算量。前缀共享时每个物理块只算了一次多个样本复用节省了大量重复attention计算。但要注意触发复制的那一次确实要多做一轮块拷贝这个带宽开销在block size较大时会比较明显。所以COW也不是万能的它最受益的场景是分裂点靠后、原始前缀足够长、同一块被大量样本共享。如果序列非常短或者分裂点非常靠前共享节省的显存可能还抵不上维护引用计数和块表映射的复杂度。还有个工程细节引用计数必须谨慎处理尤其是并发环境。多条样本同时触发对同一块的COW时如果引用计数的增减操作不是原子的轻则计数错乱导致块被提前回收重则显存被重复分配出现数据串写。我自己在调试中遇到过这种问题表现非常奇怪——有的请求生成内容突然变乱码排查半天才意识到是共享块被另一条线程改写了。5. 实际动手观察块大小、物理块总量与收益验证把机制讲清楚了还得落到工程选型上。这一节聊聊我在实际优化推理服务时怎么考虑这些参数以及如何验证这套机制到底带来了多少收益。5.1 块大小怎么选一个典型的权衡块大小block_size直接决定块表查多少项、块内部有几条token的KV、尾部浪费的期望值。如果块大小为16一条长度为150的序列分10个块9个满块 1个6 token的尾块尾部浪费10/160 ≈ 6.3%。如果块大小为32尾部浪费约3.1%但块表项更少分配次数也更少。如果块大小为8尾部浪费更低但每块只装8个token分配频率大幅上升block table也会更长且块级传输效率变低。从我的实践经验看小模型几B级别用32比较舒服大模型和超高并发场景用8或者16更合适。没有绝对最优只看你的平均生成长度和并发分布。建议在服务里把block size做成动态可配置项跑不同业务数据时分别压测一轮。5.2 物理块总量一个可以提前算出来的数物理块池能做多大其实受限于你愿意让KV Cache占用多少显存。计算公式大约是物理块总量 你为KV Cache预留的显存字节数 / 单块字节数 单块字节数 block_size × 层数 × KV头数 × head_dim × 2假设你决定KV Cache最多占20GBblock size是16单块字节数如果算出来是64KB物理块总量就是20GB / 64KB ≈ 320,000块。这个数字看起来很大但几千个并发序列共享这32万个块时实际每个序列平均能占到的块就变得很清楚了。时刻要记住物理块池是所有序列公用的不是每条一份。5.3 验证收益在线上业务中看几个关键指标怎么知道这套机制对你有效不用看理论看下面几个指标显存分配率过去按最大长度预分配时分配率可能长期只有30%到40%。引入按需分配后这个数值应该能提升到70%以上。并发上限同样一块卡如果之前只能同时跑4路长对话改动后能跑到8路甚至更多说明动态分配省的显存确实转化成了吞吐。OOM频率短请求和长请求混跑时之前的碎片化很容易导致OOM块映射之后OOM频率应该大幅下降。我在实际场景里测过一组很典型的业务数据prompt平均长度450生成平均长度120max_seq_len配置为2048。旧方案按2048预分配时显存利用率只有27%左右切换到按需分块之后利用率拉到60%以上单卡并发量翻了近一倍。这就是实打实的收益不需要复杂的基准测试也能感受到。5.4 如果要在自己的代码里模仿这套机制建议从哪里入手如果你不打算直接引入大型推理框架而是想在项目里复刻这套思想我的建议是动手做一个简单的模拟器用一个字典模拟物理块池键是物理块号值是KV数据索引。用另一个字典模拟block table键是序列ID值是物理块号列表。用一个空闲列表记录可用物理块。写三个函数allocate_block、free_sequence_blocks、copy_on_write_block。只需要几十行Python就能完全复现分配、释放、COW的流程。跑通这个模拟器之后你会对block table为什么需要引用计数、为什么COW只在未满块上触发、为什么并行采样能省这么多显存有更直观的理解。写在最后的实践建议我在实际调优中最大的体会是这套机制的核心不是“省显存”这个结果而是“吞掉显存中的不确定性”。连续缓存的最大问题不是空间不够而是明明有空间却用不上分块动态管理则让空间利用率变得可控、可预测。如果你准备把类似思路引入自己的推理服务有两个小建议。第一个是优先做块映射和按需分配COW可以后置。前两者几乎不改变模型计算逻辑改动风险最低收益却立竿见影。第二个是把block size和物理块池大小做成启动参数别写死。没有一种配置能适应所有业务留出调节空间比追求单点最优更重要。我每次调优都是靠线上数据反推参数而不是靠公式硬算。这机制的细节其实比表面上看起来要多也正因如此它才值得花时间彻底吃透。如果你正在被KV Cache的显存问题困扰希望这篇拆解能让你少走一些弯路。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑