大模型推理KV Cache显存优化：从PagedAttention到AI内存系统演进

大模型推理这件事真正做过线上部署的人都有一个共同体会模型权重加载完之后真正的麻烦才刚刚开始。显存里那块被KV Cache吃掉的空间往往比模型本身还大而且它随着并发数和上下文长度动态膨胀稍不留神就把服务打爆。我从最早用单卡跑7B模型做demo到后来在集群上折腾长上下文推理服务踩过的坑基本都围绕着一个核心问题——显存里的KV Cache到底该怎么管。这篇文章不打算复述Transformer的注意力公式而是想从工程视角把KV Cache到AI内存系统这条演进线索讲清楚顺带把每一步背后的取舍逻辑拆开。如果你正在做推理服务、显存优化或者单纯好奇为什么现在的推理框架都在讲PagedAttention、前缀共享、分层缓存这些概念那这篇内容应该能帮你把零散的知识点串成一条线。1. 先把KV Cache这件事说透它为什么存在又为什么让人头疼1.1 自回归解码的重复计算问题Transformer的解码是逐token生成的。每生成一个新token模型都要对当前序列做一次注意力计算。如果每次都对完整序列重新算一遍Key和Value那第n个token的生成就要做n次矩阵乘法整体复杂度是O(n²)级别的重复劳动。这在早期demo阶段无所谓因为序列短、请求少。但一旦上生产batch里几十个请求、每个请求上下文几千token重复计算直接把GPU算力浪费在已经算过的东西上。KV Cache的思路非常直接把每个已生成token对应的Key和Value张量缓存下来下一个token只需要计算自己的Query然后和缓存的K、V做注意力即可。这样每步的计算量从O(n²)降到O(n)代价是显存里多存了一份K和V。具体占多少以常见的配置估算层数L、注意力头数H、头维度d、序列长度S、batch大小B、精度字节数pFP16为2KV Cache大小约为2 × L × H × d × S × B × p。拿一个13B模型举例L40、H40、d128单条4096长度序列的KV Cache就是2×40×40×128×4096×2 ≈ 3.4GB。注意这只是单条序列batch开到16就是54GB比模型权重还夸张。这就是为什么长上下文和高并发场景下KV Cache成了显存的第一大消耗源。1.2 显存碎片与动态增长的矛盾KV Cache最麻烦的地方在于它不是静态分配的。请求来了要分配请求走了要释放序列还在不断变长。早期做法是给每个请求预分配一个最大长度的连续显存块比如按4096长度预留。问题是大部分请求根本用不到那么长预分配造成巨大浪费而如果按需增长又会产生大量不连续的碎片GPU显存分配器在碎片化严重时会出现明明总量够却分配不出来的尴尬。我最早遇到这个问题是在一个对话服务上单卡24G模型占14G剩下10G理论上能撑不少并发。结果压测时并发刚过8就开始OOM排查发现就是预分配策略太粗放每个请求都按2048预留实际平均只用300多token浪费了将近七成。这个经历让我意识到KV Cache的管理不能停留在分配一块内存的层面它需要一套真正的内存管理系统。1.3 从缓存到系统的认知转变很多人一开始把KV Cache当成一个简单的数据结构觉得无非是存K和V的数组。但当你真正要优化它时会发现它涉及的问题和操作系统的虚拟内存管理惊人地相似分配、回收、碎片整理、换入换出、共享。这就是AI内存系统这个说法出现的背景——它不再是一个被动缓存而是一个需要主动调度、分层管理、支持共享的完整系统。这个认知转变很关键。一旦你把它当系统看很多设计就有了理论依据分页机制解决碎片、引用计数解决共享、分层存储解决容量、调度策略解决吞吐。后面几节我会沿着这条线把每个机制讲清楚它解决了什么、怎么实现的、实际用起来什么感觉。2. PagedAttention把操作系统的分页思想搬进显存2.1 连续分配的浪费到底有多严重在PagedAttention出现之前主流做法是连续分配。每个序列的KV Cache在显存里占一段连续空间按最大可能长度预留。这种做法的浪费来自三个方向预留但未使用的内部碎片、序列间无法复用的外部碎片、以及为了防碎片而过度预留的安全余量。有研究统计过在真实负载下连续分配的显存有效利用率往往只有20%到40%。也就是说你买了100G的显存真正存了有用KV的只有二三十G其余都在空转。这个数字在长上下文场景下更难看因为预留长度要按最坏情况来而最坏情况可能比平均值大一个数量级。2.2 分页机制的核心设计PagedAttention的思路是把KV Cache切成固定大小的块block每个块存若干个token的K和V。序列不再要求连续而是通过一张块表block table把逻辑上连续的token映射到物理上分散的块。这和操作系统用页表把虚拟地址映射到物理页帧是同一个套路。这样做的好处立竿见影。首先块按需分配序列长到哪就分配到哪内部碎片最多浪费一个块的大小而不是整个预留长度。其次块可以来自显存中任意空闲位置外部碎片问题基本消失。第三块表让共享变得自然——两个序列如果前缀相同可以让它们的块表指向同一批物理块实现前缀共享。块大小是个需要权衡的参数。块太小块表本身占空间、管理开销大块太大内部碎片又上来了。实践中常见的块大小是16个token这个值在多数负载下能取得不错的平衡。我在自己的测试里对比过8、16、32三档16在吞吐和显存利用率上综合表现最好32在长序列上略优但短序列浪费明显。2.3 实际部署中的观察切换到PagedAttention之后最直观的变化是同样显存能撑的并发数上去了。之前那个24G卡跑8并发就OOM的服务换成分页管理后稳定跑到20以上显存利用率从三成提到七成多。但也不是没有代价块表的查找和间接寻址带来了一点额外开销在极短序列、超高并发的场景下这个开销占比会上升。提示分页机制在长序列、中等并发下收益最大。如果你的场景全是几十token的短请求分页带来的管理开销可能抵消掉部分收益这时候要实测对比再决定。还有一个容易忽略的点块表的生命周期管理。请求结束时必须正确回收所有块否则会泄漏。我在早期自己实现时漏掉了一个异常分支的回收逻辑结果服务跑几天后显存缓慢增长最后OOM。这类bug很隐蔽因为正常路径没问题只有异常路径才触发。用成熟框架能避免这类问题但理解它的回收机制仍然必要。3. 前缀共享与引用计数让相同的内容只存一份3.1 系统提示词的重复存储问题现在很多应用都会给每个请求加一段固定的系统提示词比如角色设定、格式要求、few-shot示例。这段前缀可能几百到几千token而且每个请求都一样。如果每个请求都独立存一份前缀的KV Cache那就是纯粹的浪费。前缀共享要解决的就是这个问题把相同前缀的KV块在多个序列间共享物理上只存一份。实现上当新请求的前缀和已有序列匹配时直接让新序列的块表指向已有的物理块而不是重新计算和存储。这里有个前提前缀的KV必须已经算好。如果前缀是固定的可以预先计算并缓存新请求直接复用连prefill的计算都省了。这就是所谓的前缀缓存prefix caching在系统提示词很长的场景下能把首token延迟降一大截。3.2 引用计数怎么保证安全回收共享带来一个新问题什么时候能释放一个块如果块被多个序列共享某个序列结束了不能直接释放因为别的序列还在用。解决办法是引用计数——每个物理块维护一个计数器记录有多少序列在引用它。序列结束时计数减一减到零才真正回收。这个机制听起来简单但实现时要注意原子性。在高并发下多个请求可能同时增减同一个块的计数如果操作不是原子的就会出现计数错误要么提前释放导致数据损坏要么永远不释放导致泄漏。成熟框架会用锁或者原子操作来保证自己实现时这是必须小心的地方。另一个细节是写时复制Copy-on-Write。如果两个序列共享了前缀块但其中一个序列要继续生成新token它需要往块里写数据。这时候不能直接写共享块否则会影响另一个序列。做法是复制一份出来单独写。这个机制保证了共享的安全性代价是首次写入时多一次拷贝。3.3 共享粒度与命中率的权衡共享的粒度直接影响命中率。按块共享意味着只有完整块相同才能共享如果两个前缀只差最后一个不完整的块那前面的块仍然能共享只是最后一块要各自算。这个设计在实践中命中率相当高因为系统提示词通常很长差异往往出现在末尾。我实测过一个场景系统提示词固定2000token用户输入平均200token。开启前缀共享后prefill阶段的计算量下降了约85%首token延迟从400多毫秒降到80毫秒左右。这个收益在交互式应用里体感非常明显。但要注意如果每个请求的前缀都不一样共享就无从谈起这时候收益为零反而多了查找匹配的开销。4. 分层缓存当显存装不下时往哪里放4.1 显存、内存、存储的速度容量阶梯显存再大也有上限而长上下文场景下KV Cache的需求可以无限增长。这时候就需要分层热数据放显存温数据放主机内存冷数据放本地存储。这三层的带宽和容量差异巨大——显存带宽是TB/s级主机内存是百GB/s级本地存储是GB/s级容量则反过来。分层缓存的核心问题是哪些KV该放在哪一层什么时候换入换出。这又回到操作系统的页面置换问题。LRU最近最少使用是最常用的策略但在推理场景下不一定最优因为序列的访问模式有特点正在生成的序列访问频繁已完成的序列可能再也不访问而等待中的序列处于中间状态。4.2 换入换出的时机与代价把KV从显存换出到内存需要一次数据传输这个传输本身要占用带宽和时间。如果换出太频繁传输开销会吃掉收益如果换出太保守显存又不够用。所以换出策略要预测哪些块短期内不会被访问。一个实用做法是结合序列状态正在活跃生成的序列的块留在显存被抢占preempted的序列的块换出到内存已完成的序列如果可能被复用比如多轮对话则保留一段时间。这种基于状态的策略比纯LRU更贴合推理负载。换入的代价更敏感因为它直接影响延迟。当一个被换出的序列要恢复生成时需要把它的KV块从内存搬回显存这个延迟会体现在用户感知的响应时间上。所以换入要尽量提前或者用异步传输掩盖延迟。4.3 实际系统中的分层实践在实际部署里分层缓存通常和调度器配合。调度器决定哪些序列这一轮参与计算缓存管理器负责把这些序列需要的块准备好。如果显存不够调度器会先抢占一些序列把它们的块换出腾出空间给高优先级序列。这里有个反直觉的点有时候主动换出比被动OOM更好。当显存接近满时与其等到分配失败再处理不如提前把低优先级序列换出保证高优先级序列的稳定。这种预防式换出需要监控显存水位设置合理的阈值。注意分层缓存引入了数据一致性风险。换出到内存的块如果被修改比如写时复制必须同步更新否则换回来的是旧数据。这个同步逻辑要仔细设计。我在测试分层缓存时发现主机内存的带宽往往是瓶颈。如果换入换出太频繁内存带宽被打满反而拖慢整体。所以分层不是免费的它适合显存确实不够、且序列有明显的冷热分化的场景。如果所有序列都活跃分层帮不上忙。5. 调度与批处理内存系统之上的决策层5.1 连续批处理与内存的互动连续批处理continuous batching是提升吞吐的关键技术不等一个batch里所有序列都结束而是每当有序列完成就立刻填入新序列。这让GPU尽量保持满载。但它对内存系统提出了更高要求——序列的进出变得频繁KV Cache的分配和回收压力大增。连续批处理和PagedAttention是绝配。因为分页让分配回收变得轻量连续批处理才能高效运转。如果还是连续分配频繁的进出会导致严重碎片连续批处理根本跑不起来。这两个技术是相互成就的关系。5.2 抢占策略谁该让出显存当显存不够时必须抢占一些序列。抢占的粒度可以是整个序列也可以是部分块。整序列抢占简单但可能抢占过多部分块抢占精细但实现复杂。多数框架采用整序列抢占把被抢占序列的块换出或直接丢弃如果重新计算比换出更划算。丢弃还是换出是个值得算的账。如果序列很短重新计算的代价可能小于换入换出的传输代价那就直接丢弃等它恢复时重算。如果序列很长重算代价高就换出保留。这个决策要基于序列长度和当前负载动态判断。5.3 公平性与吞吐的取舍调度还要考虑公平性。如果总是优先服务短序列长序列可能饿死如果严格按到达顺序短序列又被长序列堵住。实践中常用的是折中策略比如给每个序列分配一定的计算配额或者按序列已等待时间动态调整优先级。这个取舍没有标准答案取决于业务。交互式对话要求低延迟倾向于优先短序列和活跃序列批量处理要求高吞吐可以接受长序列等待。理解自己的业务特征才能选对调度策略。6. 从工程视角看这套系统的演进逻辑6.1 每一步都在解决前一步暴露的问题回头看这条演进线逻辑其实很清晰。最开始KV Cache只是朴素缓存暴露了显存浪费问题于是有了分页管理分页之后发现相同前缀重复存储于是有了共享和引用计数共享之后显存还是不够于是有了分层分层和连续批处理结合又对调度提出了新要求。每一步都是被前一步的问题逼出来的不是凭空设计。理解这个因果链比记住每个技术的名字更重要。因为你的场景可能只走到其中某一步就够了不需要全套上。比如短上下文、低并发的场景朴素KV Cache完全够用上分页反而是过度设计。6.2 不同规模下的选型建议场景特征推荐方案理由单请求、短上下文朴素KV Cache无并发压力管理开销不值得中等并发、变长序列PagedAttention解决碎片提升显存利用率固定长前缀、多请求分页前缀共享大幅降低prefill计算超长上下文、显存紧张分页共享分层容量不够时靠分层兜底高吞吐批处理全套连续批处理吞吐优先接受调度复杂度这张表不是绝对的实际选型要压测验证。我见过有人在小规模场景硬上全套结果管理开销比收益还大得不偿失。6.3 自己实现还是用框架如果你的目标是快速上线直接用成熟推理框架它们把这些机制都实现好了而且经过大量生产验证。自己实现的价值在于理解原理和特殊定制。我建议至少自己实现一个简化版的PagedAttention哪怕只支持单序列这个过程能让你真正理解块表、引用计数、写时复制这些概念用框架时遇到问题也能快速定位。自己实现时最容易出错的地方是并发安全和异常路径的回收。正常流程跑通不难难的是高并发下的计数正确性和各种异常分支的资源释放。这两块要重点测试。7. 几个实际踩过的坑和对应经验7.1 块大小选错导致吞吐上不去前面提过块大小要权衡但具体怎么选很多人凭感觉。我的经验是先用16跑基准然后根据负载特征调整。如果序列普遍很长几千token以上可以试32减少块表开销如果序列普遍很短几百token试8或16减少内部碎片。调整后一定要压测对比吞吐和显存利用率不要只看单一指标。有个细节块大小一旦确定已分配的块不能改大小所以调整需要重启服务。这意味着块大小是个部署时决策不能动态改。选之前想清楚。7.2 前缀共享的匹配开销被低估前缀共享要查找是否有可复用的前缀这个查找本身有开销。如果前缀库很大查找可能成为瓶颈。实践中要用哈希或者前缀树来加速匹配不能线性扫描。我早期用线性扫描前缀库上千条时匹配耗时占了prefill的相当比例后来换成前缀树才降下来。另外共享的前缀块要防止被意外修改。写时复制机制必须正确实现否则一个序列的写入会污染其他序列的KV导致输出错乱。这种bug很难复现因为依赖特定的并发时序。7.3 分层缓存的传输成为新瓶颈分层缓存听起来美好但传输带宽是硬约束。如果换入换出太频繁PCIe带宽被打满整体性能反而下降。我的经验是给换入换出设置速率限制并且优先保证换入的带宽因为换入影响延迟换出可以慢慢来。还有一个坑是内存碎片。主机内存也会碎片化频繁分配释放大块内存会导致分配失败。用内存池预分配能缓解这个问题。7.4 监控指标要选对优化这套系统监控是关键。要盯的指标包括显存利用率、块利用率、共享命中率、换入换出频率、抢占次数、首token延迟、每token延迟。这些指标能帮你判断瓶颈在哪。比如显存利用率高但块利用率低说明碎片或预留过多共享命中率低说明前缀差异大共享没收益。我习惯在压测时把这些指标都打出来对照着看。单看一个指标容易误判组合起来才能定位问题。8. 写在最后的一点个人体会这套东西我从头到尾折腾下来最大的感受是推理优化本质上是内存管理问题。算力固然重要但在大多数实际场景里瓶颈在显存、在带宽、在调度而不是在FLOPS。谁能把内存管好谁就能在同样的硬件上跑出更好的效果。另一个体会是不要盲目追新。每个机制都有适用场景用之前先问自己我的瓶颈到底在哪是显存不够还是延迟太高还是吞吐上不去定位清楚再选方案比堆技术栈有效得多。我见过太多人上来就全套PagedAttention加分层结果负载根本不匹配白折腾。最后这套系统的演进还在继续新的思路不断出现。但底层的逻辑——分页、共享、分层、调度——是相对稳定的。把这些搞懂新东西出来你也能快速判断它解决了什么问题、适不适合你的场景。这比记住某个具体框架的API有用得多。