【AI大模型】显存复用:多任务共享模型显存的方法
多个任务共用同一个模型时,若每个任务各占一份显存,显存随任务数线性膨胀,很快耗尽。显存复用的核心是让同一模型的权重只加载一份,KV Cache 按任务动态划分与回收,把显存占用从任务数解耦出来。本文给出完整方法:先讲显存竞争的成因,再给显存池划分、复用策略与生命周期管理,最后用可运行代码落地。全文控制在5000字以内。
显存竞争成因
多任务共享模型时,显存竞争来自三处。一是权重重复加载,每个任务若各自加载一份模型权重,权重显存随任务数成倍增长,这是最大的浪费。二是 KV Cache 无管理,各任务预分配的缓存不释放、不复用,碎片化加剧占用。三是生命周期不清晰,请求完成后缓存未及时归还,显存长期被无效数据占据。
竞争的根源是缺乏统一的显存分配与回收机制。解法是把显存当作一个共享池管理,权重只载一份,KV Cache 按块池化,任务到达申请、完成归还。目标是让显存占用主要由并发请求数决定,而非任务数量。
显存池划分
共享显存按用途分为三区。权重区存放模型权重,只加载一份,所有任务共享,这是显存复用收益最大的一处。KV 区为并发请求服务,按任务或请求划分缓存块,动态分配与回收。余量区预留运行余量,应对批处理波动与临时分配。三区比例按负载设定,权重区优先保障,KV 区决定并发上限。
划分的关键是 KV 区不能挤占权重区,否则模型放不下;权重区也不能过大,否则无缓存可分配给请求。生产上以 gpu-memory-utilization 统一约束,权重区与 KV 区在预算内动态平衡。
复用策略
显存复用围绕三个手段展开,共同把显存占用压到最低。