大模型推理显存优化：KV Cache与OOM排查实战

上周我调试一个本地部署的生成式模型任务折腾到晚上十一点问题卡在同一个地方显存不够。模型权重加载进去还剩几个G结果一跑推理直接CUDA OOM查看占用曲线发现显存峰值根本不是加载模型时出现的而是在生成长序列的过程中一路爬升最后把整张卡撑爆了。这其实就是典型的显存瓶颈问题。很多人一说到大模型推理第一反应是“参数量太大权重放不下”但实际上真正限制并发、限制上下文长度的往往不是权重而是推理过程中动态膨胀的KV Cache以及prefill阶段瞬间抬高的激活值。我这一篇就想把这件事彻底讲清楚显存到底被谁吃掉了推理过程中内存是怎么变化的有哪些调整手段可以压住峰值以及我从一次次OOM里总结出来的排查顺序。不讲复杂的概念堆砌就用一次真实任务的内存轨迹来拆适合正在本地部署、调优推理性能、或者做服务端并发预估的工程师参考。1. 显存到底被谁吃掉了四大内存块的完整拆解很多人在排查显存问题的时候习惯性只看“模型占用”和“输入长度”但这个思路是错的。一次大模型推理的显存占用至少由四块构成模型权重、算子运行时的预分配空间、KV Cache、以及激活值和临时缓冲区。你不把这四块分开看就永远找不到真正的瓶颈在哪。1.1 固定开销权重、CUDA上下文与运行时预分配先说最直观的模型权重。权重显存的计算公式很简单参数量乘以每个参数占用的字节数。一个70亿参数的模型用BF16精度加载每个参数占2字节那么权重本身需要14GB显存。如果是FP32直接翻倍到28GB如果是INT4量化可以压到接近3.5GB左右。这个计算本身不难难的是很多人忽略了一个“隐藏的固定开销”——CUDA上下文和运行时预分配。每次进程启动、加载CUDA驱动和推理框架的时候GPU会预留一部分显存作为上下文、算子调度器和内存池。这块空间依框架和GPU驱动版本不同通常在500MB到2GB不等。多卡场景下更明显每张卡都要一份。如果你的模型权重刚好占13.5GB而卡是16GB的可能一张卡根本跑不起来因为运行时预分配已经把剩余空间吞了。这类固定开销在显存计算器里通常不会给你算进去所以我的建议是给固定开销预留至少1GB到2GB的余量尤其在推理服务常驻的场景下内存碎片会在运行数小时后进一步推高实际占用量。1.2 KV Cache推理过程中最容易被忽视的显存黑洞接下来是这篇的重头戏——KV Cache。我见过很多人第一次听到这个词的时候一脸茫然但KV Cache恰恰是“推理过程中的显存黑洞”也是标题里“内存管理机制”最核心的部分。自回归模型生成文本时是一个token接一个token地往外蹦的。每生成一个新token模型都要重新计算注意力分数而注意力计算需要用到当前序列里所有历史token的Key和Value向量。如果不做缓存那么每生成一个新token就得把前面所有token的K、V全部重新算一遍计算量随序列长度线性爆炸根本没有可用的服务体验。所以主流推理框架的做法是把历史token的Key和Value向量缓存在显存里每生成一个token就追加一段新的K、V进去。这个不断累积的缓存就是KV Cache。KV Cache的显存占用可以精确算出来显存占用 2Key和Value两份 × 层数 × KV头数 × 头维度 × 序列长度 × 批量大小 × 单元素字节数举个例子一个7B模型假设32层、每层KV头数32、头维度128用BF16精度那么单个token对应的KV Cache就是2 × 32 × 32 × 128 × 2字节 512KB。注意这是一个token的开销。当你把上下文推到1024个token时单条序列的KV Cache是512MB如果你同时并行处理32条请求KV Cache总占用直接到16GB。这时候哪怕权重只占14GB整张80GB的卡也已经快被吃满了。这也是为什么本地部署时只要并发一上来OOM几乎是必然的。还有一个关键认知KV Cache是推理过程中动态增长的。模型权重是静态的加载一次就不变了但KV Cache会随着生成token数量的增加持续膨胀。很多人只看“加载模型还剩多少显存”却忽略了接下来生成阶段的内存爬坡这就是OOM频发的核心原因。1.3 激活值与临时缓冲区峰值显存的隐形推手模型权重和KV Cache已经解释了大量显存但还差一块不能忽略——激活值Activations。激活值是前向传播过程中每一层的中间计算结果。在推理阶段大部分激活值用完即弃不需要像训练那样保存下来做反向传播但由于transformer层数深、张量尺寸大在prefill阶段也就是一次性处理整段输入提示词的阶段激活值的瞬时峰值可能高得吓人。举个直观的例子输入序列是2048个token模型隐藏层维度是4096那么单层的激活张量size就是 2048×4096×2字节 ≈ 16MB。这只是一个中间张量transformer层里还有注意力矩阵、全连接层的中间结果、LayerNorm等操作一层跑下来可能就有几百MB临时分配几十层叠加瞬时峰值冲到好几GB并不罕见。激活值的峰值有几个影响因素输入长度、批量大小、隐藏层维度、以及算子融合程度。输入越长、batch越大prefill阶段的激活峰值越高。所以在部署推理服务时如果允许用户输入很长的提示词prefill阶段的显存突然跳变就是脏显存这会导致后续decode阶段没有足够的空间给KV Cache扩展。临时缓冲区这块还包括推理框架为了高效执行而预分配的工作区workspace。比如某些矩阵乘法的高性能算子会在初始化时预留几十MB到几百MB的临时空间。这部分在推理框架文档里标注得比较少但在显存告急时它就是压死骆驼的那根稻草。2. 一次完整推理中的显存动态变化prefill 与 decode 阶段的差异剖析搞清楚显存构成之后再看一次完整推理过程中的显存变化曲线很多问题就清晰了。大模型推理通常划分为两个阶段prefill预填充和decode逐token解码。这两个阶段的显存表现截然不同。2.1 prefill阶段计算密集与激活值峰值prefill阶段模型一次性接收整段用户输入并行计算所有输入token的注意力分数输出第一个生成token的logits。这个阶段是典型的高计算量、高激活值占用、KV Cache刚开始积累。注意prefill阶段并不需要为全部输入token预先分配KV Cache——虽然逻辑上它会产生一批K、V。但实际框架在预填充过程中会为这整段输入一次性计算出所有层的K、V然后把它们写入KV Cache。所以prefill结束的瞬间输入长度对应的KV Cache已经占据了显存。同时因为一次性处理整个输入序列中间激活值的张量尺寸也很大两者叠加导致显存会有一个明显的“瞬时抬升”。很多显存监控工具显示的内存尖峰就出现在这个时刻。如果你配置的显存上限比较紧张OOM往往不是在decode阶段慢慢涨上来的而是在prefill一上来就爆掉。优化prefill阶段显存的一个常见手段是chunked prefill也就是把长输入切成若干块分多次计算注意力而不是一次性处理全部。这样激活值的峰值显著降低代价是KV Cache写入和位置编码处理更复杂需要框架层面的支持。2.2 decode阶段KV Cache逐步膨胀decode阶段就是咱们看到的“一个字一个字往外蹦”的过程。这个过程每一步只处理一个新token计算量比prefill小很多激活值占用也低但是KV Cache在每步都会追加一个新token的K、V值。如果上下文长度上限是4096那么decode阶段的KV Cache就会从输入长度一直扩展到4096越到后边越紧。一个直观的比喻就是prefill阶段像是往杯子里倒了一杯水水位瞬间抬高decode阶段像是水龙头一直在滴水看着不猛但滴久了照样溢出。所以decode阶段的显存管理重点在于KV Cache的上限设定和容纳多少条并发序列。你设的上下文长度越长、并发batch越大KV Cache的压力就越大。这也是很多推理框架默认并发并不高的原因——不是算力不够是KV Cache给显存拖垮了。2.3 推理解码生命周期里的显存分配策略差异不同推理框架对KV Cache的分配策略差异很大这也是“内存管理机制”里最有门道的地方。早期做法是“静态分配”在初始化阶段就为最大序列长度、最大batch数一次性预留KV Cache空间。优点是不会碎片化缺点特别明显——无论实际跑多短的序列显存都被锁得死死的利用率极低。如果你设置最大上下文为4096、最大并发32那么KV Cache在启动那一刻就可能占据十几GB哪怕你实际上只在跑一个短对话。后来主流做法转向“动态分配”甚至“分页分配”。动态分配的逻辑是只有实际需要缓存K、V时才申请对应大小的显存块。这样做显存利用率高了但容易产生碎片——频繁申请和释放不同大小的显存块导致剩余显存虽然总量够却没有连续区域可用。更近一步的方案是“分页显存管理”类似操作系统的虚拟内存分页。KV Cache不是一大块连续空间而是一个个等大小的页block。每个页可以存固定数量的token的K、V。序列长度增长时按需申请新的页不需要时整页释放。页的大小统一碎片问题大幅减轻也让显存分配变得非常灵活。这套设计也让“beam search”束搜索这类需要多条候选序列并行推理的场景受益明显因为多条序列可以共享某些前缀的KV Cache页内存复用率更高。3. 实测有效的显存优化手段与配置参考聊完了原理少不了实操手段。我在本地部署和调优推理服务时验证过一批确实可行的方案按收益从大到小排序。3.1 量化用精度换显存见效最快量化是压缩模型权重最直接的手段。把BF16换成INT8权重直接砍半再激进一点用INT4权重降到原来的四分之一左右。但量化不是免费的午餐。INT4量化带来的精度损失在实际任务里大部分时候感知不强但在代码生成、数学推理这些对精确性要求很高的任务上退化挺明显的。我的做法是分场景取舍对话聊天、摘要生成这类任务用INT4没问题代码生成、逻辑推理至少用INT8或者BF16。量化影响的不只是权重显存某些实现里KV Cache也能一起量化。KV Cache量化通常用8bit甚至4bit存储K、V向量推理时反量化回来参与计算。这在长上下文的场景下收益极其明显因为KV Cache的占用和上下文长度是线性关系量化后直接压缩。不过需要注意显存优化的本质是“空间换时间换精度”量化后部分算子需要额外执行反量化操作推理速度可能略微下降。实测中INT8量化对速度的影响很小INT4在某些GPU上有加速效果但兼容性要提前验证。3.2 KV Cache层面的优化不止量化一条路KV Cache优化的路线挺多的我把实测有用的几条列出来。第一是上限约束。推理框架一般允许你设置最大上下文长度和最大batch数。显存不够时优先压这两个参数。上下文从4096压到2048KV Cache直接减半。代价是用户能聊的轮数变短但这往往比OOM崩溃好得多。第二是分页显存管理。这个前面提到过等大小分块动态分配KV Cache。它的核心收益是显存利用率高不会因为预分配过多而白白浪费也能容纳更多并发请求。使用这类推理框架后我遇到过不少“同样的卡之前并发8都难现在并发20多都稳”的情况。当然具体能提多少要看平均请求长度和峰值请求长度的差异。第三是前缀缓存复用。如果服务场景有大量重复的系统提示词或固定指令前缀这部分前缀对应的KV Cache可以直接复用。也就是说多条请求只要前缀相同前边一段的K、V就不用重复计算和存储直接共享。这在多轮对话、Agent场景里效果很明显能显著降低重复计算也降低了KV Cache写入压力。第四是动态KV Cache淘汰策略。像缓存淘汰一样在内存不够时优先删掉最不活跃的序列的缓存或者把某些序列的KV Cache从显存换到CPU内存。代价是换回时需要重新计算延迟会高一些。这个方案适合内存带宽充足、显存极紧张的场景。3.3 调度层面的优化投机采样与chunked prefill显存管理不只是静态压缩调度策略也能帮上大忙。投机采样Speculative Decoding的思路是用一个小模型先生成多个候选token再由大模型一次性验证。这样做的好处是在保证最终输出分布不变的前提下把大模型的decode步数降下来。显存上KV Cache占用的增长速度不会变但单位时间内生成token更快了等价于同样的显存可以服务更多的并发请求周转率提升。chunked prefill前面提过核心是避免prefill阶段激活值峰值过高。做法是把一个长输入拆成小段分别计算注意力中间状态合理释放。框架层面实现较好的话prefill的峰值显存可能下降30%以上代价是prefill的总计算时间略微增加因为部分中间结果需要重新计算。持续批处理continuous batching也值得提。传统的静态批处理要等一个batch所有请求完成后才开始下一批空闲等待多、显存利用率低。持续批处理在单条序列生成完成后立刻插入新请求让GPU一直处于满负荷状态。对显存的影响是间接的同样的显存能跑更高的吞吐因为空闲期被压缩了。3.4 一个可参考的显存估算流程我把实操中用来估算显存的一套流程写下来你可以拿来自查。先列模型参数参数量P、精度字节数B。权重显存 W P × B。再测运行时预分配在你自己的环境里启动一个空模型进程观察空闲显存被吃掉多少这个值记为F。不同框架、驱动版本差异很大必须以实测为准。然后估算KV Cache需要知道层数L、KV头数H、头维度D、单元素字节数B_kv、最大序列长度S、最大并发数C。KV Cache 2 × L × H × D × S × C × B_kv。再估算激活值峰值这个最不精确但可以用一个经验比例。实测中prefill阶段激活值峰值大约是KV Cache的20%到50%具体取决于输入长度和batch。保守起见KV Cache按最大值估算后再额外加20%给激活值和临时缓冲区。所以需要的总显存 ≈ W F KV_Cache 0.2 × KV_Cache。如果算出来超过你的卡容量先压缩并发数C和上下文长度S这两个参数是调节显存最有效的旋钮。4. 常见问题与排查技巧我从OOM中踩过的坑最后这部分把我调试过程中的血泪教训直接整理成问题速查。这里的每一条都是我真的在显存告急时被卡过的地方。4.1 一启动就OOM先确认是不是固定开销在作怪现象明明算出来模型权重加KV Cache不超过显存容量但一启动就OOM。排查方向先看运行时预分配和CPU内存映射。GPU驱动、推理框架版本不同预分配相差很大。另外有些框架支持“按需分配”显存有些则会一次性预占大量显存。如果你用的是预占型框架试试开启显存按需增长模式或者直接换一个更省显存的推理后端。还有一个隐蔽问题CPU内存不够也会触发错误但日志被误判为显存问题。权重映射到CPU内存、再按需搬到GPU的场景很常见CPU RAM不足同样会报错。4.2 运行一段时间后才OOM重点盯KV Cache和显存碎片现象启动正常刚开始几条请求也稳但跑了半小时后开始频繁OOM。这个情况九成是KV Cache在高并发下持续膨胀加上显存碎片累积。先用显存监控工具观察KV Cache占用曲线如果它在持续爬升直到接近上限那就是上下文长度或并发数设置过高。把最大上下文缩短、batch调小一般能解决。如果占用曲线没有触顶但显存分配失败那就是碎片问题。频繁申请释放大小不一的显存块导致碎片化。解决办法是优先用分页式KV Cache的推理框架如果框架不支持可以定时重启服务或者把batch size调成固定值减少不同尺寸显存块的申请频率。4.3 长上下文下必炸激活值可能是元凶现象短序列没问题一旦用户输入很长立刻OOM。这个典型是prefill阶段激活值峰值过高。chunked prefill是首选解法其次可以限制单次请求的最大输入长度超长输入直接返回提示信息。注意长上下文的KV Cache也需要看。如果输入长度直接拉满KV Cache的初始占用就很高留给decode阶段扩展的空间就不够了。所以服务上线前我会专门想清楚预设上下文长度是给“最终输出长度”用的不是给“输入长度”用的。两者要留足缓冲。4.4 批量大小总是挤不上去从请求长度分布下手现象想在推理卡上把batch调大但无论怎么调都OOM。很多人忽略了一个事实batch占用的KV Cache取决于这batch里每个请求的实际长度而不是平均长度。如果一批请求里恰好有几个特别长的那么KV Cache占用会被那几个长请求大幅推高。所以调batch时不仅要看平均请求长度更要看P99长度或最大长度。另外长度不均衡的请求混在一个batch里短请求先结束的话传统批处理框架还要等长请求跑完才释放显存浪费更严重。开启持续批处理短请求结束就立刻释放显存并插入新的请求能有效提升吞吐也会让显存回收更及时。4.5 多卡并行但显存仍然不够注意通信缓冲和负载不均衡现象从单卡升级到多卡以为显存翻倍就够了结果还是OOM。多卡并行时除了权重切分和KV Cache分布每张卡还要预留通信缓冲区、AllReduce的临时张量空间。某些张量并行的实现在运行时额外占用更多的激活值显存。所以多卡场景的可用显存不等于把所有卡显存简单相加再减去模型占用。建议先把单卡跑通、记录各项峰值再考虑多卡扩展避免在集群里反复试错浪费时间。5. 一条实操建议从最小可跑配置开始调最后我再分享一个调整显存配置时很常用的经验习惯从保守配置起步。很多人在部署时倾向于把参数调到“看起来最优”——上下文长度设置到模型上限并发设置到参考值上限。结果一跑就OOM然后开始怀疑是代码问题浪费大量时间。我的做法是反过来先给一个非常保守的配置比如上下文512、并发1、量化后权重尽量小确保能完整跑通一次推理记录显存占用基线。确认基线没问题后再逐步往上加先加并发观察KV Cache峰值再加上下文长度观察是否触发prefill峰值每次只改一个变量记录数据。这样两三轮下来你就能得到一张“显存与并发、上下文的映射表”后面调参会轻松很多。在最终给服务定型时不要踩在极限值上运行。显存用到90%以上很多时候只是看起来没事一旦请求长度波动或者内存碎片累积就会掉到OOM的坑里。留出至少10%-20%的余量像跑长途不把油箱完全清空一样这是我踩过多次坑后最想强调的一条经验。