一、显存池概述：从“内存碎片”到“资源复用”

大家好，我是你们的老朋友。今天咱们聊聊显存池。

说实话，我刚开始做深度学习框架适配那会儿，最头疼的就是显存管理。你想想看，一张显卡就那么多显存，十几个模型轮着跑，动不动就 OOM（显存溢出）。后来我慢慢摸索，发现问题的根源在于——每次张量创建都去申请显存，用完就释放，这跟频繁开关水龙头一样，浪费又低效。

显存池，说白了就是提前申请一大块显存，自己当“管家”。谁要用，从池子里拿；用完还回来，不还给操作系统。这样既避免了频繁调用 cudaMalloc 的开销，又能把碎片化的显存整合起来。

1.1 什么是显存池

显存池是一个预分配的显存缓冲区。它维护着一组空闲块和已分配块。当框架需要显存时，池子先查有没有合适的空闲块；有就直接给，没有才去系统申请新的。

我习惯把它比作“图书馆的座位”。你不需要每次学习都去建一个新座位，而是从已有的座位里找一个空的。用完离开，座位留给下一个人。

1.2 为什么需要显存池

这个问题，我当年在适配 PyTorch 和 TensorFlow 时深有体会。没有显存池，你会遇到三个大坑：

• 分配开销大：每次 cudaMalloc 都要走内核态，延迟在微秒级。模型训练时成千上万次分配，累积起来就是几十毫秒的浪费。
• 显存碎片化：频繁分配释放，显存会变得像“蜂窝煤”一样。明明总容量够，但连续大块显存就是找不到。我遇到过最夸张的一次，8GB 显存只剩 200MB 碎片，但一个 1GB 的张量就是分配不出来。
• 框架兼容性差：不同框架对显存的管理策略完全不同。PyTorch 用缓存分配器，TensorFlow 用 BFC 算法，MindSpore 又另搞一套。没有统一的显存池，适配工作就是“打地鼠”——修完一个框架，另一个又崩了。

1.3 显存池的核心设计目标

设计一个显存池，不是简单地把显存堆在一起就完事了。我总结了三个核心目标：

目标	说明	我的经验
低延迟分配	分配和释放操作要在微秒级完成	我习惯用空闲链表 + 哈希表，O(1) 查找
高内存利用率	尽量减少碎片，提高显存使用率	合并相邻空闲块，定期整理碎片
框架无关性	同一套池子能适配 PyTorch、TensorFlow 等	用统一的接口封装，内部实现可插拔

嗯，这里要注意：低延迟和高利用率有时候是矛盾的。你追求极致的分配速度，可能就得牺牲一些内存整理的时间；你追求零碎片，分配时就得花时间做合并。我个人习惯是“先保证延迟，再优化利用率”。因为推理场景下，延迟是硬指标，内存多点少点还能忍。

1.4 显存池的工作流程

为了让你更直观地理解，我画了一张流程图。这张图展示了显存池从初始化到分配、释放、回收的完整链路。

你看，流程其实不复杂。核心就两步：先查池子，池子有就直接用；池子没有，才去系统申请。释放的时候也一样，把块还回池子，而不是还给系统。

1.5 显存池的常见实现策略

市面上主流的显存池实现，我归纳为三种：

• 伙伴系统（Buddy System）：把显存按 2 的幂次分成块。分配时找最小满足的块，释放时合并相邻块。优点是碎片少，缺点是内部碎片（比如申请 3KB，实际给 4KB 块，浪费 1KB）。
• SLAB 分配器：预分配固定大小的对象池。适合频繁创建相同大小张量的场景（比如卷积核）。我曾在某个推理框架里用 SLAB，分配延迟从 2μs 降到了 0.3μs。
• BFC（Best-Fit with Coalescing）：TensorFlow 用的算法。维护一个按地址排序的空闲块列表，分配时找最合适的块，释放时合并相邻块。优点是内存利用率高，缺点是分配时 O(n) 查找。

我个人习惯是：训练场景用 BFC，因为内存利用率更重要；推理场景用 SLAB，因为延迟更重要。当然，你也可以混合使用——大块用 BFC，小块用 SLAB。

1.6 避坑指南：显存池的常见误区

最后，分享几个我踩过的坑：

• 误区一：池子越大越好。 错！池子太大会占用显存，导致其他进程无法使用。我建议初始大小设为模型峰值显存的 1.2 倍，不够再动态扩展。
• 误区二：释放后立即整理碎片。 整理碎片需要移动数据，开销很大。我建议只在池子空闲率低于 50% 时才触发整理。
• 误区三：所有框架用同一套参数。 不同框架的张量大小分布不同。PyTorch 偏爱小张量，TensorFlow 偏爱大张量。我建议为每个框架单独配置池子参数。