真实面经题目 · 原创解析
LLM 推理引擎中 GPU 内存管理机制应如何设计,如何管理 KV Cache、显存碎片、并发 batch 和 OOM 降级?
这题考的是推理引擎的显存资源治理能力:不能只会调用 CUDA malloc,而要能把 KV Cache、临时 workspace、权重、并发请求、碎片控制和 OOM 降级统一成可预测、可观测、可调度的内存系统。
真实面经题目 · 原创解析
这题考的是推理引擎的显存资源治理能力:不能只会调用 CUDA malloc,而要能把 KV Cache、临时 workspace、权重、并发请求、碎片控制和 OOM 降级统一成可预测、可观测、可调度的内存系统。
我会把 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 后重启进程。
推理引擎需要先把显存拆成几类:模型权重和量化权重是常驻内存,CUDA context、NCCL 通信 buffer 和运行时库开销是基础开销,KV Cache 是随请求动态增长的主要占用,算子 workspace 和激活临时张量是峰值来源,输入输出 buffer、采样 buffer 和日志统计也会占一部分。没有这张账本,就无法判断 OOM 来自并发太高、上下文太长、workspace 峰值太大,还是 allocator 碎片。
自回归推理中,每生成一个 token 都会追加每层的 key 和 value。KV Cache 占用大致和层数、batch 内活跃序列数、历史 token 数、KV head 数、head dimension 和 dtype 字节数成正比。多轮对话、长上下文和连续 batch 会让 KV 成为显存瓶颈。设计时要把 KV 作为一等资源纳入调度,而不是把它当成普通临时张量。
更稳的做法是把 KV Cache 划成固定大小 block 或 page,每个请求维护逻辑 token 到物理 block 的映射。分配时从 free list 取块,释放时按请求归还;如果支持 prefix cache,还需要引用计数避免共享前缀被提前释放。这样连续的逻辑上下文不要求物理显存连续,长短请求混合时也不容易被外部碎片拖垮。
显存碎片通常来自不同大小、不同生命周期的 buffer 混在一起申请释放。KV Cache 生命周期跟请求绑定,可能跨很多 decode step;workspace 和激活张量生命周期短,常在一次 forward 内结束;权重几乎不释放。工程上应使用预分配大池、size class、arena、独立 KV pool 和临时 workspace pool,避免长生命周期小块夹在短生命周期大块中间。
LLM 推理不能只按请求数限制并发,因为一个长输入请求可能比多个短请求更吃显存。scheduler 应同时考虑 prefill token 数、decode 活跃序列数、剩余 KV block、最大上下文长度、最大输出 token、优先级和延迟目标。continuous batching 的关键是让新请求只在显存预算允许时进入,并在请求完成、取消或超时后立即释放其 KV block。
prefill 阶段一次处理大量输入 token,attention 和算子 workspace 峰值更高;decode 阶段每步只处理新增 token,但活跃序列多、KV Cache 压力持续存在。设计调度时可以分离 prefill 队列和 decode 队列,对 prefill 做 token 上限和微批切分,对 decode 做活跃序列合批,避免 prefill 峰值把 decode 请求挤到 OOM。
准入阶段可以根据显存预算拒绝、排队或要求缩短上下文;运行阶段可以缩小 micro batch、降低 max new tokens、暂停低优先级请求、释放已取消请求、禁用部分 cache 复用,或把不关键的缓存转移到 CPU。降级要明确优先级和用户可见语义,不能让所有请求一起失败。
生产系统需要记录 reserved 显存、实际 allocated 显存、KV block 使用率、碎片率、峰值 workspace、OOM 类型、请求长度分布、取消释放延迟和降级次数。压测要覆盖长短混合、持续生成、突发 prefill、请求取消、prefix 复用和极限 max context,验证显存能否稳定回收。
可以按层数乘以活跃 token 数乘以 KV head 数乘以 head dimension,再乘以 key 和 value 两份以及 dtype 字节数估算。实际还要考虑 padding、block 内未用 token、对齐开销、beam search 或多副本缓存,以及 tensor parallel 下每张卡保存的分片比例。
在线请求长度差异大,完成时间也不同,连续分配容易出现总空闲显存足够但没有足够连续大块的问题。paged KV 把逻辑连续上下文映射到多个物理块,释放和复用更细粒度,能更好支持 continuous batching 和 prefix 共享。
可以让框架 allocator 负责普通张量,但推理引擎最好对 KV Cache 和大 workspace 建专用池,因为它们的生命周期、大小分布和调度语义都很明确。关键是减少运行中 cudaMalloc 和 cudaFree,避免同步开销和碎片不可控。
请求取消要进入调度器状态机,停止后续 decode,把该请求持有的 KV block 引用计数减一,释放独占 block,并清理采样状态和输出 buffer。释放最好在安全的 stream/event 边界之后执行,避免 GPU kernel 仍在使用时被复用。
GPU OOM 往往说明当前调度已经超过显存水位,简单重试可能继续失败,还会破坏延迟和稳定性。更好的做法是准入前预测,运行中有可控降级,并把失败原因记录到请求和系统指标中。