3. 显存峰值问题
做AI推理优化这些年，我见过太多团队在吞吐压测上栽跟头。大家往往盯着算力利用率、batch size调优这些显性指标，却忽略了那个藏在暗处的家伙——显存峰值。
说白了，显存峰值就是你的模型在推理过程中，某一瞬间占用的最大显存量。它不像平均显存那样温和，而是像心电图上的尖峰，稍不留神就捅破显存天花板。
显存峰值产生的原因
为什么会突然冒出个峰值来？我总结下来，主要有三个推手：
- 中间激活的堆积：Transformer模型里，每一层的前向计算都会产生中间激活值。这些值在反向传播时才会释放，但推理场景下，它们会一直攒着直到整个序列处理完。你想想看，一个7B模型，序列长度2048，光中间激活就能吃掉好几个GB。
- KV Cache的膨胀：自回归生成时，每生成一个，KV Cache就膨胀一圈。我在项目中遇到过，一个70B模型跑长文本生成，KV Cache从几百MB一路飙到十几GB，最后直接OOM。
- 框架的内存管理策略：PyTorch的缓存分配器、TensorRT的内存池，这些框架为了减少频繁申请释放的开销，往往会预分配一大块显存。嗯，这里要注意，预分配策略虽然提升了速度，但也推高了显存峰值。
核心观点：显存峰值不是某个单一因素造成的，而是模型结构、推理策略、框架机制三者共同作用的结果。
显存峰值对吞吐的影响
吞吐和显存峰值之间，存在一个微妙的平衡关系。我习惯用一个公式来理解：
有效吞吐 = batch_size × 序列长度 / 单次推理延迟但这里有个陷阱——显存峰值限制了batch_size的上限。你想想看，为了追求高吞吐，你拼命加大batch size，结果显存峰值先爆了，系统直接OOM，吞吐归零。
我在实际调优中遇到过这样的场景：
| batch size | 平均显存 | 显存峰值 | 吞吐(tokens/s) | 状态 |
|---|---|---|---|---|
| 1 | 4.2 GB | 5.1 GB | 120 | 正常 |
| 4 | 12.8 GB | 18.3 GB | 380 | 正常 |
| 8 | 22.5 GB | 35.6 GB | 620 | OOM |
看到没？batch size从4调到8，平均显存只涨了不到10GB，但显存峰值直接翻倍。这就是典型的峰值陷阱——你以为显存够用，结果峰值一来直接崩盘。
个人经验：我建议在压测前，先用profiler跑一次小batch，摸清显存峰值的增长曲线。这样你就能预判出安全batch size的上限，而不是盲目往上加。
显存OOM的典型场景
显存OOM不是随机发生的，它往往出现在几个特定场景里。我整理了一下，最常见的有三种：
- 长序列推理：输入序列长度超过模型训练时的最大长度。比如模型训练时最大支持2048，你硬塞了个4096的序列进去。KV Cache直接翻倍，中间激活也成倍增长，OOM几乎是必然的。
- 高并发请求：多个请求同时到达，每个请求都分配独立的KV Cache。我曾经在线上遇到过，平时单请求显存峰值只有6GB，结果10个请求同时涌进来，显存峰值瞬间飙到60GB+，直接OOM。
- 动态batch的显存抖动：有些框架支持动态batch，但不同长度的序列混在一起时，显存分配会出现剧烈波动。短的序列还没释放，长的序列又申请了更多显存，峰值就上去了。
避坑指南：我曾经在压测时犯过一个低级错误——只关注了平均显存，没看峰值。结果线上部署后，用户请求一多就频繁OOM。后来我养成了习惯：压测报告里必须包含显存峰值曲线，而且峰值要留出至少20%的余量。
显存峰值的可视化分析
为了让你更直观地理解显存峰值的形成过程，我画了一张图：
从这张图你能看到，显存峰值往往出现在KV Cache快速膨胀的阶段，也就是t1到t2这个区间。平均显存虽然平稳上升，但峰值已经悄悄逼近OOM临界点了。
如何定位显存峰值问题
定位显存峰值，我有一套自己的三板斧：
- 使用NVIDIA Nsight Systems：它能精确记录每个CUDA kernel的显存分配和释放。我习惯在压测时跑一次profiling，然后盯着显存分配曲线看，哪个时间点分配量最大，那个点就是峰值来源。
- PyTorch的memory_stats：在代码里插入torch.cuda.memory_stats()，可以实时监控显存分配情况。我一般会在每个推理步骤前后都打一次日志，这样就能画出完整的显存变化曲线。
- 手动计算理论峰值：根据模型参数量、序列长度、batch size，手动估算出理论上的显存峰值。比如一个7B的FP16模型，参数占14GB，KV Cache占2GB，中间激活占3GB，那理论峰值就是19GB左右。如果实际峰值远超这个数，那肯定有内存泄漏或者分配策略的问题。
小技巧：我习惯在压测脚本里加一个显存告警阈值，比如设置成显存总量的80%。一旦超过就自动打印堆栈信息，这样能快速定位到是哪个操作在吃显存。
显存峰值这个问题，说白了就是压测中的隐形杀手。它不像算力利用率那样直观，也不像延迟那样容易测量，但一旦爆发，整个系统就崩了。我见过太多团队在压测时只盯着吞吐数字，结果上线第一天就OOM，然后通宵排查。
记住一句话：压测不测峰值，等于白测。