大模型推理时出现 OOM，如何从 batch、输入长度、KV Cache、临时张量峰值和算子实现排查？｜月之暗面算法面经解析

60 秒回答模板

我会先把推理 OOM 当成可复现的显存账本问题来排查，而不是直接调小一个参数。第一步固定模型、dtype、并发策略和输入样本，记录 OOM 发生在加载模型、prefill、decode 还是采样阶段，同时采集 GPU allocated、reserved、峰值显存、请求长度、batch 内 token 数和 CUDA kernel 日志。第二步拆 batch：分别测试 batch size、并发序列数、prefill token 数和 decode 活跃序列数，因为 LLM 的内存压力不是简单等于请求数，长输入会显著放大 attention 临时张量和 KV Cache。第三步看输入长度和输出长度：输入越长，prefill 峰值和初始 KV 越大；生成越长，KV Cache 持续增长。可以画显存随 token 增长曲线，看是否符合 KV 线性增长。第四步核对 KV Cache 公式和实现：层数、KV head 数、head dim、dtype、beam 或多候选、padding、block 未利用率、prefix cache 引用计数、释放时机都可能导致超预期。第五步排查临时张量峰值和算子实现：普通 attention、FlashAttention、GQA/MQA、fused kernel、workspace 配置、contiguous 拷贝、mask 构造、logits buffer、top-k/top-p 采样都可能制造额外峰值；用 profiler 看内存峰值对应哪个 op。最后看 allocator 和碎片：reserved 远大于 allocated、总 free 足够但申请失败、长短请求混跑后失败，都提示池化或碎片问题。修复上按根因选择：准入限 token、分离 prefill/decode、paged KV、缩 micro batch、优化 attention kernel、复用 workspace、释放中间张量、降低 max context 或做 OOM 降级，并用回归压测证明同样流量不再 OOM。

考点 阶段定位

难度 真实面经题

回答目标让候选人能给出可执行的推理 OOM 排查路径：先定位阶段，再拆解 batch 与长度，核对 KV Cache，使用 profiler 找临时峰值和算子退化，最后用针对性修复和压测闭环验证。

深入解析

先定位 OOM 阶段

OOM 可能发生在模型加载、首轮 prefill、后续 decode、采样、通信或请求释放之后。不同阶段根因不同：加载阶段多与权重、dtype、tensor parallel 切分有关；prefill 多与输入长度和 attention 临时张量有关；decode 多与 KV Cache 和并发序列有关；释放后仍不降则要怀疑引用、缓存或 allocator 保留。

batch 要拆成多个维度

面试里说降低 batch size 太粗。推理 batch 至少包括请求数、batch 内总 token、prefill micro batch、decode 活跃序列、beam 或候选数。两个请求数相同的 batch，如果一个是 128 token，另一个是 16K token，显存完全不同。排查时要分别固定输入长度改变并发，固定并发改变输入长度，建立显存曲线。

输入长度影响 prefill 峰值

prefill 阶段要一次处理提示词 token，长输入会增加 attention 计算、中间激活、mask、position 相关 buffer 和初始 KV 写入。某些实现如果没有使用内存高效 attention，可能生成很大的临时 attention score 或 workspace。输入长度导致的 OOM 常表现为首 token 前失败，缩短 prompt 或拆 prefill 后显著改善。

生成长度影响 KV 线性增长

decode 阶段每生成一个 token 都追加 KV Cache，因此 OOM 可能不是一开始发生，而是在生成到某个长度后发生。排查时要记录每步已生成 token、活跃序列数和 KV block 使用量。如果显存随生成 token 近似线性增长，首先核对 max new tokens、early stop、请求取消释放和 KV block 回收。

KV Cache 要核对公式和实现细节

KV Cache 理论占用和层数、token 数、KV head 数、head dimension、key/value 两份和 dtype 字节数相关。但实际实现还会受 block 大小、padding、对齐、beam search、prefix cache 共享、GQA/MQA、tensor parallel 分片和重复缓存影响。理论值和实测差距大时，要查是否重复保存、释放延迟或 block 利用率低。

临时张量峰值要靠 profiler 对齐

很多 OOM 不是常驻内存太大，而是某个算子临时峰值把水位顶破。常见来源包括 attention score、softmax buffer、workspace、contiguous 转换、hidden state 拷贝、logits 全量保留、top-k buffer 和通信聚合 buffer。用 profiler 或内存 hook 对齐 op 时间线，才能知道峰值来自哪个算子。

算子实现会改变内存复杂度

同样是 attention，不同实现的内存行为差别很大。FlashAttention 类实现通过分块减少 attention score 的物化，GQA/MQA 减少 KV head 数，fused kernel 减少中间张量落地，in-place 或 workspace 复用降低峰值。排查时要确认当前路径是否真的启用了预期 kernel，是否因 shape、dtype 或 mask 退回到低效实现。

碎片和泄漏要用长压测验证

如果单次请求没问题，长时间混合长度请求后 OOM，可能是碎片、缓存未释放、引用计数错误或 allocator reserved 过高。要做循环压测，观察每轮结束后的显存基线是否回落，reserved 与 allocated 的差距是否扩大，以及取消、超时、异常路径是否释放 KV 和 workspace。

易错点

看到 OOM 就只建议降低 batch size，没有区分请求数、总 token、prefill 和 decode。
只检查模型权重大小，忽略 KV Cache 随上下文和生成长度持续增长。
没有记录 OOM 发生阶段，导致把 prefill 临时峰值误判成 decode 并发问题。
用单个短样本验证修复，却没有覆盖长输入、长输出和混合长度流量。
忽略 attention 或采样算子的临时 buffer，把所有显存都算成常驻内存。
没有确认高效 kernel 是否实际生效，shape 不匹配导致 fallback 也没发现。
请求取消或异常时不检查 KV 释放路径，长期运行后显存基线越来越高。
只看 nvidia-smi 的 free memory，不看 allocator reserved、allocated、峰值和碎片。

面试官追问

怎么区分 KV Cache OOM 和临时张量峰值 OOM？

KV Cache OOM 通常随生成 token 持续增长，可能在 decode 到某个长度后失败；临时张量峰值 OOM 往往发生在 prefill 或某个特定算子执行时，峰值过后本应释放。可以用逐步显存曲线和 profiler 的 op 时间线区分。

为什么输入长度比 batch size 更容易误导排查？

因为一个长输入请求可能比很多短请求消耗更多 prefill 内存和初始 KV。只看请求数会把不同 token 规模的 batch 视为相同，导致准入控制和压测结论都不准。

如果 profiler 显示 reserved 很高但 allocated 不高，说明什么？

这通常说明 allocator 为了复用保留了大量显存，或者出现了碎片和缓存池未归还。需要继续看最大连续可分配块、内存池策略、长短生命周期是否混用，以及循环压测后基线是否持续抬高。

算子 fallback 为什么会导致 OOM？

高效 kernel 可能只支持特定 dtype、head dimension、mask 类型或 shape。条件不满足时回退到普通实现，可能物化更大的 attention score 或中间张量，使同样模型在某些输入长度下突然 OOM。

线上不能直接复现时怎么办？

先从日志记录请求长度、max new tokens、并发序列、峰值显存、OOM 阶段和模型配置，构造离线最小复现。再用流量回放或合成压测模拟长短混合和持续并发，逐步缩小到触发 OOM 的长度、batch 和算子组合。