资讯详情

PagedAttention原理浅析:KV Cache显存优化的核心机制

📅 2026/10/19 10:34:53 | 华诺云谱 👁 阅读
PagedAttention原理浅析:KV Cache显存优化的核心机制
1. 从显存焦虑说起大模型推理为什么卡在内存上做过大模型推理部署的工程师大概率都经历过类似的场景模型加载只占了几个GB但一跑起来并发一高显存直接拉满然后就是OOM崩溃。很多人第一反应是模型太大但其实更隐蔽的凶手是KV Cache——也就是推理过程中缓存下来的Key和Value张量。用GPT类模型做生成时每个新token都需要对之前所有的token做注意力计算。为了不重复计算业界通行做法是把之前每层的Key和Value都存下来下次直接用。这个缓存的增长速度非常惊人。以7B参数的模型为例4096长度下单条请求的KV Cache大约需要几十MB显存如果是13B、70B参数配合更长的序列几百MB甚至上GB都是常有的事。显存本身就是稀缺资源KV Cache和激活值、模型权重争抢空间一旦并发上来碎片化、浪费、溢出问题一个接一个。大家通常的缓解方式很朴素要么batch开小一点要么限制最大序列长度。这些办法确实有效但代价是吞吐上不去本来能同时处理10个请求的卡可能要降到4个。真正解决这个问题的关键思路其实是借鉴操作系统的一个经典设计——虚拟内存的分页管理。把KV Cache切成固定大小的块逻辑上连续、物理上分隔用时分配、不用释放、需要共享时还能通过引用来复用。这套机制在推理框架里的落地形态就是PagedAttention。我第一次看到这个思路的时候第一反应是这其实不是在搞什么全新数学而是把几十年前操作系统的内存管理策略搬到GPU上。但恰恰是这种“老方案新场景”的迁移解决了大模型推理落地中最棘手的一段隐性内存困境。这篇文章不打算给你复述论文结论我会拆开块映射、动态分配、写时复制这三个核心机制聊清楚它们到底解决了什么问题、内部是怎么运转的以及在真实部署中你会碰到的坑和调优经验。2. 块映射用一张表解决“逻辑连续物理离散”的问题2.1 为什么连续显存是奢望碎片化才是常态推理过程里KV Cache的分配是按序列增长的。每个请求的序列都在变长每个请求的长度也各不相同如果按传统方式给每个请求预先分配整块连续显存会出现两个问题。第一是内存浪费。为了不频繁扩容通常的做法是预留最大可能长度。比如预设最大长度2048那每个batch槽位都会预留2048长度的KV Cache哪怕这个请求实际只生成了200个token就结束了剩下1848个token的空间全部闲置。支撑数百并发时显存浪费比例往往会超过一半。第二是碎片化。处理完的请求释放显存后空闲区域是被切割成大小不等的不规则区块的。新请求要分配大块连续内存的时候即便总体剩余显存足够也会因为没有足够大的连续区域而失败。PagedAttention没有选择“尽量找连续大块”的路线而是反其道而行之放弃物理连续性只保证逻辑连续性。它的做法是把KV Cache切成固定大小的块block每个块能存固定数量token的Key和Value。比如块大小设为16个token那么每16个token的KV就放进一个块里。序列的KV Cache不再是一整块内存而是一串块列表。这些块在物理内存上的位置不需要相邻。某序列的第二个块可能在第3号物理块上第三个块在第97号物理块上完全错落。为了让逻辑顺序正确每生成一批token系统就把新分配的物理块编号追加到该序列的块索引表里。读取时按逻辑顺序遍历索引表逐个访问对应的物理块。提示块大小是个关键超参。块越小内部碎片越少但索引表和分配开销越大块越大管理开销小但每条序列尾部未用完部分浪费的内存越多。常见默认在16到32个token之间实际取舍需要结合平均生成长度来评估。2.2 块表结构一张表如何支撑任意长度的序列从工程实现上看块映射的本质是一张“逻辑块号 → 物理块号”的翻译表。这张表维护了每个序列的块列表、每块的引用计数、物理块的空闲状态等元数据。访问某个块时通过索引查表一次就能定位到物理地址。类似CPU里面TLB和进程页表的语义。GPU上访问KV Cache需要频繁执行这个查表动作所以设计上必须保证这个翻译过程尽可能快。在实现层面一个好的技巧是把块表放在共享内存或者专用的高速缓存区域避免每次访问都走一遍全局内存查表。这里还有个小细节物理块是全局共享的。一个物理块在某一个时刻可能只属于一个序列但也可能被多个序列同时引用——这为后面讲到的写时复制Copy-on-Write打下了基础。因为块是共享的所以每个块都要维护一个引用计数表示当前有多少序列在使用这个块。引用计数归零时这个块才会被真正回收。来看一个典型的KV Cache块表结构序列ID逻辑块索引物理块编号引用计数token数最后一块Req_ABlock 0Phy 12116Req_ABlock 1Phy 45116Req_ABlock 2Phy 2129Req_BBlock 0Phy 2129这三个序列共享了同一个物理块Phy 21——它们在某一段历史时刻生成的token是完全一样的所以显存里只需要存一份。这也是PagedAttention最精妙的地方它把KV Cache的存储粒度从“整条序列”降到了“块”共享和复用的可能性被大幅打开了。2.3 从虚拟内存借用过来的设计哲学虚拟内存为什么能对应用“撒谎”告诉进程你有无限连续的内存因为中间加了一层映射关系。物理内存随便怎么碎片化只要映射表能把不连续的地方串成连续的逻辑地址进程就感知不到底层的碎片。PagedAttention对KV Cache做的事情本质上一样。一个token序列在逻辑上依然是连续的——从第1个token到第N个token顺序推下去不会有断点。但在物理层面这些token的KV可以分布在任意物理位置。映射关系完全由块表决定。这就把“显存碎片化”这个难以根治的问题转化为“在块表里查一个编号”这种极低成本的操作。这种“加中间层解决不连续问题”的思路在计算机系统里并不少见。文件系统用块号映射解决文件碎片问题数据库用缓冲池解决磁盘块和内存页的映射问题。PagedAttention只是把这个成熟思想代入到GPU内存管理场景里。它最大的意义在于当你不再需要连续物理内存时内存的调度灵活性就剧增了。这为后面要说的动态分配和写时复制扫清了机制上的障碍。3. 动态分配按需取用拒绝静态预留3.1 静态分配的显存黑洞效应传统推理框架对KV Cache最省事的管理方式是什么预分配。比如最大序列长度定义为2048batch大小定义为8那KV Cache总量就是8×2048×层数×头数×每个值占用字节数一次性用torch.empty或者CUDA API分配好。之后不管请求实际是否用到这些显存都被永久“锁住”了。这个设计在并发低、长度固定的场景下很简单但真实业务场景下请求长度差异极大。有的请求只生成50个token就结束有的却要写一篇2000字的文章。按最长的来预留短请求的显存就白白浪费按平均长度留长请求又有被截断的风险。静态分配本质上是让用户为最坏情况买单。PagedAttention把分配的粒度从“序列”级别降到了“块”级别。每个序列开始时只分配第一个块。当生成的新token超过当前块的容量时才从全局块池里申请下一个块。块用完了就还给全局池其他序列可以马上复用。这套机制带来的最直观收益就是显存只跟着“实际生成的token数”走而不是跟着“最大可能的token数”走。在真实业务中请求的实际生成长度往往远小于上限显存利用率通常能提升2到4倍。这意味着同样的GPU可以同时跑更多的并发请求整体吞吐自然就上去了。3.2 全流程从请求进入到块回收动态分配虽然听起来只是“按需分配”但工程上必须处理一整套状态流转。这里简单拆几条主线路径。分配路径新请求进入时从全局空闲块池中取一个块初始化引用计数为1挂到请求的块表尾部。当生成循环发现最后一个块已满比如块大小16已经存了16个token就再向空闲块池申请新块。如果空闲块池耗尽系统进入阻塞或调度状态等待其他请求结束、释放块或者实施一些强制回收策略。释放路径请求结束后遍历其块表把每个块的引用计数减1。减到0时物理块真正归还给空闲池。如果引用计数仍然大于0那说明还有其他序列在共享这个块不能释放只需要断开当前序列的引用即可。扩容路径如果整个全局块池不够用了CUDA malloc新的显存并扩容空闲块池是一种常见的做法。但要小心这类动态显存扩张可能引入额外的地址不连续设计上要做两层映射块表只保存块编号具体物理地址由块编号二次解析。这样即使底层显存地址变动上层也只是改一个地址表项的事。整个生命周期涉及块状态的变化空闲→已分配→共享引用计数1→可回收→空闲。状态机的鲁棒性很重要最容易出bug的是共享计数和块释放的顺序问题。如果先减引用计数再清理地址映射刚好有其他请求在读这个块数据就出现悬空指针问题。正确的顺序应该是先摘除映射表里的对应项再去递减引用计数。注意我在写分布式推理服务的时候踩过类似坑。对引用计数做原子操作还不够必须配合内存屏障确保可见性。多线程并发分配/释放场景下一个块的引用计数状态对别的线程不可见会导致野指针和显存泄漏风险这类问题排查起来极度痛苦。3.3 GPU显存里还有多少利用空间批处理如何从块收益动态分配还有一个隐藏收益就是它直接重构了批处理batching的方式而批处理正是GPU利用率的核心来源。传统连续KV Cache约束下一个batch内的所有请求往往被强行限制成长度差不多否则短请求生成完后它所占据的那段预留给它的KV Cache空间就只能空着利用率极低。在PagedAttention的块级分配下每个请求独立管理自己的块列表生成结束的请求可以立即释放空间新的请求可以随时进入batch这个在工程上叫作连续批处理Continuous Batching。GPU上每一轮迭代选一批当前需要计算的请求它们各自的KV Cache分散在不同物理块里但计算逻辑不变对每个请求按它的块表顺序逐个块读取KV完成注意力计算。正因为没有“空闲槽位必须固定在某个位置”的约束调度器可以按需组合任意请求进入batch理论上每个迭代都能保持GPU算力满负荷运行。在这个层面上PagedAttention对于GPU利用率的意义远远大于单纯省显存。4. 核心妙手写时复制如何把显存需求指数级压缩4.1 一次采样生成多个候选序列的显存噩梦大模型生成过程里有一种非常常见但显存开销极大的场景并行采样生成多条候选序列。用户在做beam search集束搜索、best-of-n采样、或者self-consistency自洽性采样时要从同一个prompt出发在某个分支点同时生成多个不同的后续序列。从同一个prompt前缀开始这些序列前面的KV Cache其实是一模一样的。如果给每条候选序列都单独存一份完整的KV Cache显存开销会按候选数量线性放大。假设beam size是4完整的KV Cache原本就要占1GB那4条序列就要4GB实际上一大半都是重复存储。在没有块级映射机制的时候这个问题很难解。要么就接受成倍的显存浪费要么就限制beam size效果都有代价。但通过PagedAttention的块级共享写时复制这个问题做到了几乎“免费”的水平。4.2 Copy-on-Write的块级实现方案要理解这里面的机制关键在于“何时复制”和“复制什么”。假设起点是一个已经生成的prompt前缀它对应的一系列物理块已经被分配好了。现在要从这个前缀分出4条候选序列。正常的做法是给这4条序列各自分配一份完整的KV Cache空间然后把prompt前缀的KV逐块复制过来。这个过程又慢又浪费。PagedAttention的做法截然不同这4条序列的前缀部分直接共享同一组物理块不做任何数据复制。在每条序列的块表里引用相同物理块编号每个块的引用计数变成4。因为大家引用的数据完全一样所以没有正确性问题。接下来每条候选序列继续独立生成各自的token。新token的KV Cache需要申请新的物理块来写。在这个点上就遇到了“写时复制”的核心判断当一条序列需要写入某个块但这个块被多个序列共享时它必须先检查这个块是否只属于我如果不是我就需要复制一份物理块数据然后只更新自己的块表指向新的物理块再在新块的副本上写入新数据。这整个过程可以拆解成几个关键判断点读取前缀块直接读不需要复制因为所有序列读到的都是同样的KV。写入新块这个块单独属于当前序列直接写。写入共享块需要先触发复制然后把原来共享块的引用计数减1再把当前序列的块表指向新副本。值得强调的是这种复制是懒惰的lazy——复制只发生在实际要写入那个块的时候而不是从分支一开始就把所有前缀块复制一遍。如果一条候选序列生成了两步就终止了那它可能从头到尾只复制了一次或两次块。相比“创建分支时全量复制”这种按需复制策略在候选序列寿命不长时节省的显存和时间都非常可观。4.3 引用计数与安全回收写时复制最容易出错的一环写时复制机制里最需要小心的地方除了何时复制还有何时释放。前面提到引用计数决定了物理块何时可以被回收。写时复制场景下引用计数会频繁地增减一个块的引用计数可能在1、2、4、3之间反复跳动任何一步出错都可能导致严重后果。一类典型错误是某条序列在被调度器淘汰时遍历自己的块表把每个块的引用计数都减1。如果某块引用计数减到了0但这个块实际上仍被其他序列引用着那么其他序列在后续迭代中会读取被回收的内存数据随机、计算崩溃。这类问题在多序列并发场景里极其隐蔽。工程上的建议是对所有引用计数的修改统一集中到一个模块管理提供inc_ref(block_id)和dec_ref(block_id)两个原语接口所有读写路径都通过这些接口操作。不要在业务逻辑代码里零散地修改引用计数否则迟早会漏掉某个路径。经验之谈对于写时复制强烈建议在调试模式下定期做显存一致性检查——把每个物理块的引用计数遍历一遍与实际块表中的引用关系做交叉验证。我在实际项目里就是靠这种一致性校验抓到了一个非常隐蔽的双重释放问题。5. 实操视角PagedAttention的典型实现路径与调优5.1 从零实现时的核心模块拆解如果要在自己的推理引擎里实现类似PagedAttention的机制大概会涉及三个核心模块第一个是显存分配器。它维护一个物理块池对外提供alloc_block()和free_block()接口。分配器的目标是最小化空闲碎片同时保证分配效率。实际实现时可以用自由链表维护空闲块也可以用红黑树支持按地址查找复杂度可控。第二个是块表管理器。每个正在处理的请求拥有一张独立的块表维护逻辑块到物理块的映射。块表需要在请求生命周期内动态扩展所以底层通常是一个指针数组或者链表结构。读取时通过逻辑块编号查物理块地址写入时可能需要触发写时复制。第三个是核心注意力kernel。标准的注意力kernel假设KV Cache是连续内存可以直接通过token位置索引。PagedAttention下的kernel必须改为“先查块表再访问块数据”的模式。这一步其实是最影响性能的查表的开销必须足够小否则多级间接寻址带来的延迟会抵消省下显存带来的吞吐收益。业界通常会把块表预加载到共享内存或者直接让kernel在读取下一个块时做预取prefetch隐藏掉查表延迟。实现过程中最容易被忽视但也最关键的是块表的缓存友好性。GPU上的内存访问模式对性能影响极大。如果按照逻辑token顺序读取但物理块地址跳来跳去会打乱GPU的合并内存访问coalesced memory access模式导致实际带宽利用率暴跌。针对这个问题的一个常用优化是对同一个序列的块尽量在物理内存中保持连续分配让大部分访问路径仍然是线性的。5.2 关键参数块大小应该怎么定块大小block size是PagedAttention最重要的超参数直接影响内存利用率和kernel效率。从论文和工程实践来看需要在这两者之间做一个清晰的权衡。块大小越大管理开销越小kernel访问KV Cache时的局部性越好但尾部浪费会增大——每个序列最后一个块往往装不满造成平均约半个块的浪费。块大小太小分配频率变高块表查找开销更大而且写时复制频繁。一般来说如果生产环境里序列长度的分布整体较长块大小可以适度调大如果很多短请求块大小又应该调小。我见过一些团队直接把块大小设成16或者32再用实际流量回放做A/B对比这比反复读论文里的推荐值靠谱得多。根据我的经验16到64之间通常能找到当前业务的最佳点不要相信一个设置能通吃所有场景。5.3 显存扩容与空闲块水位线设计一个真实的推理服务块池的总量并不会一成不变。系统需要根据显存压力动态扩容或者收缩。这里有一个设计问题空闲块池的水位线应该怎么设置如果空闲块太少高并发场景下新请求可能等待块分配而排队产生不必要的延迟。如果空闲块太多说明大量显存被“预留”但却没有实际使用资源利用不充分。一个常见做法是设置两个阈值低水位和高水位。当空闲块数低于低水位时尝试扩容块池高于高水位时可以考虑收缩。还可以根据预测算法预估未来请求量提前扩容到目标水位。GPU显存不像系统内存可以随时换页所以水位线策略要和调度器配合。一种常用的技巧是如果空闲块不足不立即报错或阻塞而是把调度器改为“等待模式”让当前批次的请求继续生成直到释放出可用块再吸收新请求。这样能极大缓解显存突发压力代价是某些请求的排队延迟变大。5.4 实测调优案例一个并发翻倍的优化过程我在做某跨平台系统部署时最初按传统静态分配方式配置batch size最大开到16显存就逼近上限。每次请求都要预留2048序列长度的KV Cache但实际平均生成长度只有300多token显存浪费比例接近80%。切换到PagedAttention方案后同样的GPUbatch size提升到48还能维持稳定运行。显存利用率的提升主要来自两个环节一是块级分配让短请求不再拖累整体空间二是写时复制让beam search场景的显存成本几乎降到了原来的四分之一。整个调优过程其实没有太多高深技巧就是几个动作把块大小设为32空闲块低水位设为总量的15%写时复制开关打开然后反复压测。最有意思的一次问题是明明总占用显存下降了吞吐却反而下跌了。排查半天发现是kernel的块表查表逻辑写得太蠢每一轮都从全局内存里重新读块表导致显存访问量在查表上暴增。后来把块表缓存到共享内存里吞吐立刻恢复了。这提醒我一个道理省显存只是第一步把省下来的资源真正转化为计算效率才是PagedAttention落地的关键。6. 常见问题速查与实战避坑指南6.1 显存泄漏排查引用计数为什么会越滚越大显存泄漏是PagedAttention实现里最让人头疼的问题之一。块池的空闲块数量一点点减少最后服务OOM。我遇到过的典型泄漏场景包括请求被异常中断时没有走正常的块表释放逻辑导致已分配的块没有归还到空闲池连续批处理调度中请求在某一轮被提前移出batch但块表的引用计数没有同步更新多线程环境下引用计数被并发修改丢失了一次decrement操作。排查思路很简单在分配器模块里加一个统计计数记录当前已分配物理块总数和引用计数总和。如果引用计数总合大于块表里实际引用数的总和就说明有引用泄漏。还可以为每个块记录分配时的调用栈一旦发现异常情况直接看是哪条路径只增不减。6.2 吞吐上不去查表开销是隐形瓶颈不少人以为只要上了PagedAttention吞吐就自动高了。实际上如果kernel实现得粗糙可能比连续内存版本还慢。块表每层KV Cache都要做一次逻辑块号到物理块号的翻译翻译过程涉及一次访存。如果翻译路径设计不好等于每个token多了一次额外访存延迟累积起来非常可观。解决方法把当前批次所有请求的块表预取到共享内存采用两级块表结构让常用块尽量命中L1/L2对多次重复读取的块数据做缓存从kernel层面预测下一个物理块并提前发起加载。如果这些优化都做了吞吐还是上不去那大概率是块大小或调度策略的问题要用profiler对照实际瓶颈再调。6.3 显存分配不出来的死循环低显存场景下空闲块池耗尽但块池上方还压着部分并不需要立即释放的空间这时候如果调度器逻辑写得不严谨可能进入死循环——一直在等待某个块释放但所有请求都在等待分配块继续生成谁也等不到谁。正确的设计是当空闲块不足时优先对当前batch内部分请求执行强制终止或者临时降级策略释放显存而不是死等。还有一种方案是把当前batch中已完成全部生成步骤的请求立即踢出释放KV Cache而不是等到下一轮调度时再处理。提示设置一个硬性上限当某个请求的空闲块等待时间超过阈值就要做出取舍宁可牺牲一部分请求的完成度也要保住系统整体不陷入僵局。6.4 一个容易忽略的点前缀共享的边界写时复制能省显存但要小心前缀的KV Cache并不总是完全可共享的。如果两条序列在某个位置的分支点已经不同比如一个接着写“今天天气很好”另一个接着写“今天的股价波动”那么从这个分支点往后的物理块引用关系就完全不同。在实际实现时块表里的引用关系必须精确到逻辑块而不是整个前缀都一刀切共享。一个简单稳妥的策略是前缀完全相同的那段KV Cache块共享进入分支后新逻辑块一律独立分配不再共享。块级粒度让这种“精确到块”的共享成为可能——整个前缀里只有还没产生分歧的那一段才会被引用计算覆盖到。7. 顺带聊几句这套机制能扩展到哪些更远的场景PagedAttention的设计里最核心的思想也就是“块级映射共享引用”和“懒惰复制”不只适用于大模型的KV Cache。类似的显存优化技术在多模态模型的视觉token缓存、长文本的向量索引缓存、以及各类LLM Agent式的长上下文推理服务里都有直接或间接的移植价值。比如在长文本对话场景里系统经常需要维护同一用户的多轮历史上下文。不同轮次的对话之间历史前缀重叠度很高。用块级映射管理多轮对话的KV Cache不同会话可以安全共享前缀块按需增量保存后续token不仅省显存还能显式控制历史上下文的保留与淘汰这比以前的“按轮次整段缓存”要优雅得多。另一个方向是跨请求的prompt复用。如果多个用户同时请求同一个热度很高的prompt前缀比如公共的system prompt或工具定义这个前缀块的引用计数可以支持超大规模共享。只要前缀块未被写入并行请求之间就能够完全复用同一份显存也因此省去了重复计算这部分前缀的时间开销。这些思路本质上已经超越了PagedAttention论文本身的范围变成了一套“显存资源池化 引用管理”的通用方法论。我个人的实际体会是这套机制越深入到工程细节越能理解它最开始为什么要从“块”出发做设计。正是靠块这个粒度分配、映射、共享、复制、回收这些操作才能同时成立。用连续内存那一套思路KV Cache共享根本无从谈起。你现在如果正在做长上下文推理、高并发推理服务、或者多候选序列生成我建议认认真真把块映射、动态分配和写时复制这套机制从原理到代码级别吃透它对系统整体收益的影响短时间内很难被其他优化方式替代。
📝

华诺云谱内容团队

资深建站顾问 · 行业研究员

10年+企业数字化服务经验,专注智能建站、SEO优化与品牌营销,持续输出建站技巧、行业洞察与营销干货,已帮助5000+企业实现数字化增长。

你可能需要的服务

订阅华诺云谱资讯周报

每周一封,精选建站技巧、SEO与营销干货,直达邮箱。已有 8,000+ 企业主订阅,助你少走弯路。

↑