真实面经题目 · 原创解析

LLM 推理引擎中 GPU 内存管理机制应如何设计,如何管理 KV Cache、显存碎片、并发 batch 和 OOM 降级?

这题考的是推理引擎的显存资源治理能力:不能只会调用 CUDA malloc,而要能把 KV Cache、临时 workspace、权重、并发请求、碎片控制和 OOM 降级统一成可预测、可观测、可调度的内存系统。

出现于:阿里巴巴 · C/C++

60 秒回答模板

我会把 LLM 推理引擎的 GPU 内存管理设计成分层的资源管理系统,而不是在算子里随用随申请。首先要做显存预算,把常驻权重、CUDA context、通信 buffer、算子 workspace、KV Cache、输入输出 buffer 和预留安全水位分开统计。推理阶段最核心的是 KV Cache,因为它随 batch size、序列长度、层数、hidden size、KV head 数、dtype 线性增长,并且会在生成过程中持续扩张。工程上通常会做 paged KV cache 或 block-based cache,把 KV 按固定大小 block 管理,通过 free list、引用计数和 request-to-block 映射支持动态 batch、prefix 复用和请求释放,避免每个请求申请一段连续大内存导致碎片。其次要控制碎片:大块内存预分配、按 size class 建内存池、区分长生命周期 KV 和短生命周期临时张量,尽量不要混用同一 allocator;对临时 workspace 可以用 arena 或 CUDA graph 友好的静态 buffer。并发 batch 上要由 scheduler 基于 token 预算而不是只基于请求数做准入控制,例如限制总 KV block、prefill token、decode token、最大上下文和最大并发序列,并根据剩余显存做 continuous batching。OOM 降级要分层:准入前拒绝或排队超预算请求,运行中可以缩小 batch、降低 max new tokens、截断低优先级上下文、迁移或释放可重算缓存,必要时把部分请求退回 CPU 或返回可解释的繁忙错误。最后要有监控和保护:记录 reserved、allocated、free blocks、碎片率、KV 命中、OOM 原因、峰值 workspace 和每请求显存账单,用压测覆盖长输入、高并发、混合长度和取消请求场景。核心目标是让显存使用可预测、释放可及时、失败可降级,而不是靠 OOM 后重启进程。

考点 显存预算
难度 真实面经题
回答目标 让候选人能从 C++ 推理引擎视角设计 GPU 显存管理:明确显存账本,重点管理 KV Cache,用块化和池化降低碎片,用 token-aware 调度控制并发,并提供可观测的 OOM 降级机制。

深入解析

01

先建立显存账本

推理引擎需要先把显存拆成几类:模型权重和量化权重是常驻内存,CUDA context、NCCL 通信 buffer 和运行时库开销是基础开销,KV Cache 是随请求动态增长的主要占用,算子 workspace 和激活临时张量是峰值来源,输入输出 buffer、采样 buffer 和日志统计也会占一部分。没有这张账本,就无法判断 OOM 来自并发太高、上下文太长、workspace 峰值太大,还是 allocator 碎片。

02

KV Cache 是内存管理核心

自回归推理中,每生成一个 token 都会追加每层的 key 和 value。KV Cache 占用大致和层数、batch 内活跃序列数、历史 token 数、KV head 数、head dimension 和 dtype 字节数成正比。多轮对话、长上下文和连续 batch 会让 KV 成为显存瓶颈。设计时要把 KV 作为一等资源纳入调度,而不是把它当成普通临时张量。

03

用块化管理减少连续内存依赖

更稳的做法是把 KV Cache 划成固定大小 block 或 page,每个请求维护逻辑 token 到物理 block 的映射。分配时从 free list 取块,释放时按请求归还;如果支持 prefix cache,还需要引用计数避免共享前缀被提前释放。这样连续的逻辑上下文不要求物理显存连续,长短请求混合时也不容易被外部碎片拖垮。

04

碎片控制要区分生命周期

显存碎片通常来自不同大小、不同生命周期的 buffer 混在一起申请释放。KV Cache 生命周期跟请求绑定,可能跨很多 decode step;workspace 和激活张量生命周期短,常在一次 forward 内结束;权重几乎不释放。工程上应使用预分配大池、size class、arena、独立 KV pool 和临时 workspace pool,避免长生命周期小块夹在短生命周期大块中间。

05

并发 batch 要按 token 和 block 调度

LLM 推理不能只按请求数限制并发,因为一个长输入请求可能比多个短请求更吃显存。scheduler 应同时考虑 prefill token 数、decode 活跃序列数、剩余 KV block、最大上下文长度、最大输出 token、优先级和延迟目标。continuous batching 的关键是让新请求只在显存预算允许时进入,并在请求完成、取消或超时后立即释放其 KV block。

06

prefill 和 decode 的内存形态不同

prefill 阶段一次处理大量输入 token,attention 和算子 workspace 峰值更高;decode 阶段每步只处理新增 token,但活跃序列多、KV Cache 压力持续存在。设计调度时可以分离 prefill 队列和 decode 队列,对 prefill 做 token 上限和微批切分,对 decode 做活跃序列合批,避免 prefill 峰值把 decode 请求挤到 OOM。

07

OOM 降级要发生在准入和运行两端

准入阶段可以根据显存预算拒绝、排队或要求缩短上下文;运行阶段可以缩小 micro batch、降低 max new tokens、暂停低优先级请求、释放已取消请求、禁用部分 cache 复用,或把不关键的缓存转移到 CPU。降级要明确优先级和用户可见语义,不能让所有请求一起失败。

08

监控和压测决定可靠性

生产系统需要记录 reserved 显存、实际 allocated 显存、KV block 使用率、碎片率、峰值 workspace、OOM 类型、请求长度分布、取消释放延迟和降级次数。压测要覆盖长短混合、持续生成、突发 prefill、请求取消、prefix 复用和极限 max context,验证显存能否稳定回收。

易错点

  • 只说使用 cudaMalloc 或框架 allocator,没有设计显存预算、池化和调度策略。
  • 忽略 KV Cache 的线性增长,把 OOM 归因成模型权重太大。
  • 用请求数限制并发,却不考虑输入长度、输出长度、prefill token 和 KV block 数。
  • 把 KV Cache 和临时 workspace 放在同一个无约束内存池里,导致生命周期混杂和碎片加剧。
  • 只在 OOM 后返回失败,没有准入控制、排队、缩 batch 或上下文降级策略。
  • 没有处理请求取消、超时和异常路径的 KV 释放,导致显存缓慢泄漏。
  • 没有记录显存峰值、碎片率和 OOM 原因,线上只能靠重启恢复。
  • 认为总 free memory 足够就一定能分配成功,忽略连续大块需求、对齐和 allocator 保留内存。

面试官追问

KV Cache 大小怎么估算?

可以按层数乘以活跃 token 数乘以 KV head 数乘以 head dimension,再乘以 key 和 value 两份以及 dtype 字节数估算。实际还要考虑 padding、block 内未用 token、对齐开销、beam search 或多副本缓存,以及 tensor parallel 下每张卡保存的分片比例。

为什么 paged KV Cache 比连续分配更适合在线服务?

在线请求长度差异大,完成时间也不同,连续分配容易出现总空闲显存足够但没有足够连续大块的问题。paged KV 把逻辑连续上下文映射到多个物理块,释放和复用更细粒度,能更好支持 continuous batching 和 prefix 共享。

显存池和 CUDA 自带 allocator 的关系怎么处理?

可以让框架 allocator 负责普通张量,但推理引擎最好对 KV Cache 和大 workspace 建专用池,因为它们的生命周期、大小分布和调度语义都很明确。关键是减少运行中 cudaMalloc 和 cudaFree,避免同步开销和碎片不可控。

如何处理请求取消后的内存释放?

请求取消要进入调度器状态机,停止后续 decode,把该请求持有的 KV block 引用计数减一,释放独占 block,并清理采样状态和输出 buffer。释放最好在安全的 stream/event 边界之后执行,避免 GPU kernel 仍在使用时被复用。

OOM 时为什么不能只捕获异常后重试?

GPU OOM 往往说明当前调度已经超过显存水位,简单重试可能继续失败,还会破坏延迟和稳定性。更好的做法是准入前预测,运行中有可控降级,并把失败原因记录到请求和系统指标中。