资讯动态

大模型推理显存优化:KV Cache才是并发OOM的真凶

发布时间:2026/10/20 8:59:10 来源:尧图企业网站定制
先说一个真实场景某公司把一个 7B 模型接到线上做长文档问答prompt 动不动就几千字。上线前测试单条请求一切正常一旦并发跑到二三十路GPU 直接 OOM把容器都带崩了。当时所有人第一反应都是“模型太大了权重装不下了”把 batch 调成 1、甚至怀疑量化不到位。结果折腾了一圈发现真正的问题根本不在这——权重占用的显存是固定的真正随着请求并发和上下文长度疯狂膨胀的是每一路请求各自持有的 KV Cache。这篇文章就想把大模型推理时的显存账算清楚显存到底被谁吃了、KV Cache 为什么是瓶颈、主流推理框架的内存管理策略各有什么取舍以及在不换卡的前提下哪些压低显存的手段真正有效。适合正在做推理部署、被显存和并发折磨的工程师参考。1. 同一条提示词为什么线上系统总在几十并发就OOM——显存消耗的各环节拆解推理跟训练的内存画像完全不同。训练时显存大头是权重、梯度和优化器状态推理没有反向传播梯度不占地方理论上内存压力应该小很多。但实际跑起来人很容易懵一个 7B 模型的 FP16 权重撑死 14GB单卡 24GB 明明放得下为什么并发一高就挂因为推理时显存实际上被四个东西同时占着模型权重、中间激活值、KV Cache以及推理引擎自己预留的调度和通信缓冲。模型权重这部分很好理解加载到显存后基本不动。中间激活值指 forward 过程中每一层的输入输出、归一化统计量、注意力分数这些临时张量随着 batch size 和序列长度变化且在计算完成后会释放。KV Cache 是我们今天的重点它是 Transformer 在自回归解码阶段为了不重复计算历史 tokens 的 K、V 矩阵而建立的缓存业务并发和输入长度越大它占得越夸张。最后还有一块容易被忽略的“隐形成本”CUDA context、cuDNN 算法空间、推理引擎的显存池预分配任何一次推理请求都得先给引擎留出几 GB 的地盘这部分在小显存卡上非常致命。把账算得更细一点关键差别在于各部分的“生命周期”。权重是一次性成本与并发无关激活值是瞬态成本算完就还回去但 prefill 阶段因为要一次性处理整段 prompt峰值非常陡峭只有 KV Cache 是持续成本从第一个 token 开始分配直到整个序列生成完毕才释放。所以并发从 1 增加到 32权重不涨、激活值跟着 batch 线性涨、KV Cache 也线性涨但很多部署场景里 KV Cache 本身的单请求占用就已经很大线性放大后直接就把显存捅穿了。1.1 权重只是“一次性成本”真正波动的是激活和缓存我经常拿高速公路来类比。权重好比收费站本身建好后放多少车过来都不再加钱但激活值是在修临时匝道每来一辆车都得现搭一块KV Cache 则是停在服务区里的所有车每辆车占了车位就得停到它走为止。并发一高服务区车位先不够用这就是 OOM 的第一现场。具体到数值上以常见 7B 模型为例FP16 权重约 14GB4K 上下文的单请求激活值在 prefill 阶段可能瞬间冲到几 GB但 attention 计算完成后大部分释放decode 阶段激活值很小而 KV Cache 在 4K 上下文下约占 2GB32K 上下文下约 16GB——大家注意这个数字它已经超过了权重本身。很多人在线下测试时只用短 prompt、低并发KV Cache 没显示存在感一到线上长上下文场景立刻原形毕露。1.2 监控显存时最容易漏掉的那几 GB排查 OOM 时nvidia-smi看到的“已用显存”跟 PyTorch 内部分配的显存往往对不上。原因在于 CUDA 的缓存分配器PyTorch 从 CUDA 申请显存后并不会立刻还给驱动而是留在自己的缓存池里复用下次申请更小张量时直接从池子里拿。这本来是为了避免频繁cudaMalloc的性能损耗但也意味着你看nvidia-smi会看到一个滞留的高水位即使当前计算已经不再需要那么多显存。实际操作中我习惯三步走第一步用nvidia-smi --query-gpumemory.used --formatcsv看整体水位第二步在代码里打点观察torch.cuda.memory_allocated()和torch.cuda.memory_reserved()的差值差值越大说明缓存池越空转第三步单独统计 KV Cache 的分配量——很多框架比如 HuggingFace 的past_key_values都会返回缓存张量的 shape把各层的 shape 乘起来就能精确算出它在总显存里的占比。只要看到 KV Cache 占比超过权重就该立刻反思上下文长度和并发策略而不是急着加显卡。2. KV Cache是吃掉显存的最大变量——从“每个Token都要存K/V”到“该留多大的缓存”为什么 Transformer 解码需要 KV Cache这份缓存又为什么大到离谱根子在 Attention 的计算方式上。生成第 N 个 token 时它要去“回看”前面 N-1 个 token 的 K、V 矩阵算一个加权求和。如果没有缓存每一步都重新从前面的 hidden state 出发把 QKV 全部重新算一遍理论正确但计算量翻着跟头涨。KV Cache 的思路是前面的 K、V 反正算过一次了存起来生成第 N 个 token 时只要算当前 token 的 Q、K、V再用当前 Q 去跟已缓存的全部 K 做相似度、跟全部 V 做加权问题就解决了。这个“存起来”就是显存黑洞的来源。每生成一个 token每个注意力层的 K 和 V 各需要往缓存里追加一份向量这个向量长度等于注意力头的维度乘上头数或者分组查询下的 KV 头数。累积到长序列时缓存大小 token 数 × 层数 × 2K和V × 向量维度 × 精度字节数。看着是几个小数字相乘但 Transformer 动辄几十层、向量维度几千乘起来就是天文数字。2.1 一个能背下来的估算公式每 token 每层到底吃多少显存我不建议每次都用 PyTorch 打点算工程上更需要快速估一个数量级。以 Llama 2 风格 7B 模型为例hidden size 4096共 32 层FP16 精度下每个 token 的 KV Cache 占用是 2 × 4096 × 2 字节 16KBK 和 V 各一份4096 是隐藏层维度。单层 16KB 看着不起眼乘上 32 层就是 512KB/token。这个数字可以直接背下来用4K 上下文约 2GB8K 上下文约 4GB32K 上下文约 16GB。我在线下用这个公式快速估模型能不能扛住某个长度的输入基本不会翻车。这里还有一个隐藏变量GQA分组查询注意力和 MQA多查询注意力。这些结构让多个 Q 头共享同一组 K、V 头KV Cache 直接从 hidden size 缩水到“KV 头数 × 头维度”比例可观。如果模型报告了num_key_value_heads参数一定要代入计算很多新版模型的 KV Cache 就是这么被砍下来的。2.2 prefill 和 decode 两个阶段显存性格完全不一样KV Cache 的分配策略还要分阶段看。prefill 阶段也就是用户输入整段 prompt、模型并行处理的那一段需要一次性把整段输入的所有 K、V 算出来并写入缓存对显存是峰值冲击但好处是计算高度并行GPU 能吃饱。decode 阶段则是逐个 token 蹦每步只往缓存尾巴追加一点点但每一步都依赖全部历史缓存所以缓存不能释放得一直占着直到 EOS。这个差异直接决定了配置策略。如果用户输入的是 10 分钟长的会议记录prefill 瞬间就要为 5000 个 token 分配 KV 空间跟后续生成 200 个 token 的增长完全不成比例。很多推理引擎就是用max_seq_len一次性把最长序列的 KV Cache 空间都预留出来跟实际增长解耦。所以你会看到某些框架配置里显存分配量跟“最大允许的序列长度”强相关而不是跟“实际平均长度”强相关。预留少了长文本请求直接拒绝预留多了并发又上不去。这是后文所有工程策略都要围绕的一对矛盾。3. PagedAttention只是把“大块不够用”转成“小块分页管理”——推理框架的内存策略与trade-off既然 KV Cache 这么大省显存的第一反应自然是动态分配用多少、分配多少。但这里有个底层麻烦显存碎片化。假设每路请求长短不一生成的 token 数也各不相同如果每个请求都在 GPU 显存里随机大小地申请一块连续空间跑一段时间后哪怕总空闲量足够也可能找不到一整块能塞下新请求的空间就像磁盘碎片导致文件写不进去。于是 vLLM 那套 paged attention 的思路登场了它把操作系统里虚拟分页的做法搬到了 GPU 显存上KV Cache 按固定大小的 block 切分逻辑上连续的 K/V 序列可以放在物理上不连续的块里通过块表映射访问。这样碎片不再是问题新的需求只需要凑够若干空闲块即可显存利用率上了一个台阶。但用“分页”换来的也不是免费的午餐。块表映射本身增加了解索引开销block 大小直接影响匹配效率块太大会浪费尾块空间一个请求最后一小块不足一整块也要整块占住块太小则块表膨胀、按块查找的开销上升。工程里常见的 block 大小在 16~32 个 token 之间需要在碎片率和寻址开销之间反复调。另一层开销是调度逻辑的复杂度框架要维护“哪些块空闲、哪些块被哪路请求占着、哪些块可以复用”这些元数据虽然在 CPU 侧但频繁分配和释放也会抬高延迟极低延迟场景下反而可能成为瓶颈。3.1 连续缓存 vs 分页缓存的取舍框架为什么各做各的有意思的是不同推理引擎在这个问题上做出了不同的工程取舍。一类是 vLLM、SGLang 这类彻底分页化换显存利用率的极致另一类是 TensorRT-LLM 这类采用“连续 KV Cache 剩余增量用池化”的混合思路理论上先为请求预留一段连续空间如果不够再按块追加。这种折衷保留了连续空间顺序读取的性能优势又避免了每个请求一开始就预留全长度的浪费。我个人的经验是短请求、高并发的场景下连续缓存命中率高延迟稳长上下文、高波动场景下纯分页更省显存但求稳的时候不要把 block 调太大。没有绝对优劣只有匹配不匹配。3.2 实践观察长短上下文的并发峰值差一个数量级这类机制带来的实际影响是不要拿一组 benchmark 数据直接推算线上容量。我在某次压测中见过一个配置用 2048 长度 prompt 测出来可容纳 128 并发但线上真实请求平均 8000 长度实际并发降到 16 就濒临 OOM。这个数量级差异背后就是 KV Cache 的线性爆炸。排障时建议按“最大允许上下文 × 目标并发 × 单 token 缓存大小”先算上界再决定要不要调低max_seq_len或者上模型并行。很多团队上线后才发现问题往往就是漏算了这层“长尾请求的显存上界”。4. 压显存还可以从算子、量化、并行三个方向下手——比直接换卡更现实的路径说完机制聊聊真正能落地的优化手段。我的建议是不要一上来就换 A100/H100显存优化有非常成熟的三个方向组合起来往往能把容量翻一倍以上算子层面的激活优化、权重/ KV Cache 量化、模型并行。这三个方向各有侧重也各有代价理解了代价才能选得准。4.1 FlashAttention 究竟省了什么它绕开了“注意力矩阵”但不是 KV Cache 的救星很多人误以为 FlashAttention 能压 KV Cache这其实是概念混淆。FlashAttention 做的是把标准注意力中的完整 N×N 分数矩阵分段计算不再显式物化保存整个注意力分数省的是中间激活值让 batch × 序列长度的乘积可以做得更大而不爆显存。但 KV Cache 是另一回事它存在于 decode 阶段是跨 token 复用的持久状态越使用反而越多FlashAttention 根本不碰这部分。这就解释了为什么上了 FlashAttention 之后 batch 能拉大、prefill 峰值下降了但长序列的 KV Cache 压力依然在。不过算子级优化并非全然无用。有些实现会把注意力计算和设备内存带宽的节奏配合得更好GPU 利用率上去了同等显存下同样时间能跑完更多请求间接缓解了“显存不足”的拥堵感。选框架版本时留意 Attention 后端的实现版本往往小版本更新就能带来几个百分点的吞吐提升这是性价比极高的第一刀。4.2 量化权重压缩是常规做法KV Cache 量化才是长上下文的关键量化是压缩权重大小的立竿见影手段。FP16 的 7B 模型约 14GBINT8 降到 7GBINT4 进一步到 4GB 左右。可问题在于很多团队量化完权重发现显存确实宽裕了但跑长上下文时还是爆——因为前面算过32K 上下文的 KV Cache 约 16GB权重就算压缩到 4GB总占用依然很高。这时更该关心的是 KV Cache 本身能不能量化。KV Cache 量化通常是把缓存里的 K、V 从 FP16 压到 FP8甚至更低。好处是长上下文场景的显存直接减半坏处是精度敏感尤其在 attention 分数对数值差异敏感的任务上量化不当会带来可感知的生成质量退化。操作上有两个经验一是分通道量化比全局量化稳得多因为不同注意力头对 K、V 的数值范围差异很大二是有些框架默认不对 KV Cache 做量化需要显式打开配置项用之前先跑一组同种子对比测试看困惑度变化。我在实际使用中发现对大多数文档摘要类任务FP8 缓存几乎无损性价比很高。4.3 模型并行把“塞不下”拆成“多张卡一起扛”但要小心通信放大模型并行是终极手段。Tensor Parallel 把权重和 KV Cache 按头切到多张卡上单卡显存压力直线下降非常适合大模型Pipeline Parallel 则是按层切分但每层依然需要保留完整的 KV Cache多卡并行时前半段层的缓存会拖住整体。所以单纯为了扩上下文长度上 TP 是对的上 PP 要慎重。不过并行不是免费的。TP 在每一层 attention 后都要做 all-reduce 同步通信量不小跨机用 PCIe 互联时单次通信延迟可能比计算本身还高。很多场景里双卡 TP 甚至比单卡跑得慢尤其小 batch 下。我的经验是如果一张卡装得下权重但装不下长上下文的 KV Cache优先选量化 分页缓存而非直接上多卡只有模型本身大到单卡装不下时才考虑 TP并尽量保证卡间用 NVLink 等高带宽互联而不是走普通 PCIe。5. 给不同显存的卡一份可落地的配置参考——我实测下来的参数颗粒度最后直接给出可抄作业的配置思路。这里基于常见的 24GB / 48GB / 80GB 三档单卡模型按 7B、13B、70B 三个规模给典型配置覆盖大多数线上推理场景。需要说明的是这只是起点具体还要结合业务请求长度分布和并发要求微调。单卡显存模型规模建议精度最大上下文参考并发关键策略24GB7BFP16 或 INT84K8~16优先限制max_seq_len不要贪长上下文24GB7BINT48K~16K4~8权重省下的显存留给 KV Cache适合长文档场景48GB7BFP1632K8~12KV Cache 吃重建议开启量化和分页缓存48GB13BINT48K8~12权重约 7GB主要显存预算仍压在 KV 上80GB70BINT48K4~8权重约 40GB长上下文仍紧张必要时上 TP80GB×270BINT416K~32K8~16TP 切分 KV Cache通信走 NVLink 才划算这里有个比较容易忽略的细节很多推理引擎允许设置max_seq_len但 KV Cache 的预分配是按这个上限来的。你把上限调成 32K即使业务里 99% 请求只有 2K显存也按 32K 占满预留空间并发直接腰斩。引擎的“按需分配”能力差异很大一定要确认框架是在请求到来时才分配缓存还是启动时就整体预留一块按最大长度算的池子。如果属于后者宁可把上限调小一点宁可拒绝超长请求也要保住并发吞吐。实际操作中我的配置顺序是先根据业务统计得到 P95 的最大请求长度留 20% 余量设为max_seq_len再按单 token 缓存大小算出单请求 KV Cache 峰值接着用“空闲显存 总显存 - 权重 - CUDA 开销 - 激活峰值”除以单请求峰值估算出并发上限最后压测到 80% 水位作为线上配置。这套流程下来基本能把每张卡的容量压榨到比较干净的水平。最后再分享一个小经验如果线上同时有“短 prompt 高并发”和“长 prompt 低并发”两类业务不要混在同一个引擎实例里跑。长请求会把短请求的 KV Cache 空间挤占掉导致短请求也跟着被拒。把长上下文业务单独拆一个实例把max_seq_len调大、并发压低短上下文实例调高并发、限死长度整体吞吐往往比混跑高出一截。踩过几次坑之后我对这件事特别敏感显存瓶颈很多时候不是硬件不够而是策略上把不同性格的请求塞进了同一个笼子。

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

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

免费获取报价 →
↑