4、显存池设计哲学：预分配策略、复用策略、回收策略、扩容策略
聊到显存池，很多人第一反应就是「哦，不就是搞个内存池嘛，把显存提前申请好放着」。嗯，这话对了一半。显存池的核心确实是个池子，但GPU显存这玩意儿，比CPU内存娇贵多了。我刚开始做分布式训练框架时，就吃过一次大亏——显存碎片化导致一个8卡任务死活跑不起来，最后发现是池子策略没设计好。
今天咱们就把显存池的四个核心策略掰开揉碎讲清楚。说白了，就是四个问题：什么时候拿？拿完怎么用？用完怎么还？不够了怎么办？
核心观点：显存池不是简单的「先申请后释放」，而是一套围绕GPU显存特性的资源调度系统。预分配、复用、回收、扩容这四个策略，缺一个都会出问题。
4.1 预分配策略：不打无准备之仗
预分配，顾名思义，就是在真正需要显存之前，先把坑占好。为什么要这么做？因为cudaMalloc这个操作，开销比你想象的大得多。我在项目中测过，一次cudaMalloc大概要几十微秒到几百微秒，如果训练过程中频繁申请释放，光这个开销就能吃掉5%~10%的吞吐。
预分配策略通常有两种做法：
- 启动时全量预分配：程序启动时，一次性把最大可能用到的显存全申请好。比如模型训练需要24GB，那就直接cudaMalloc 24GB。优点是后续零延迟，缺点是显存利用率低——万一别的进程也需要呢？
- 按需预分配：先申请一个基础池（比如8GB），后续根据实际使用量逐步扩大。我比较推荐这种方式，尤其是在多租户场景下。
我的经验：预分配的大小建议设置为「模型参数+优化器状态+最大batch梯度」的1.2倍。留出20%的余量，避免频繁触发扩容。我曾经因为余量只留了5%，结果一个特殊batch就把池子撑爆了，排查了半天才发现是输入数据尺寸波动导致的。
代码层面，预分配通常长这样：
// 伪代码：显存池预分配
class MemoryPool {
void* base_ptr;
size_t pool_size;
size_t used_size;
void init(size_t initial_size) {
// 预分配基础池
cudaMalloc(&base_ptr, initial_size);
pool_size = initial_size;
used_size = 0;
}
void* allocate(size_t size) {
// 从池中分配，不触发cudaMalloc
void* ptr = base_ptr + used_size;
used_size += align_up(size);
return ptr;
}
};4.2 复用策略：让每一字节显存都物尽其用
复用是显存池的灵魂。你想想看，训练过程中张量的生命周期有长有短：权重是长生命周期的，梯度是中等生命周期的，中间激活值可能只存活几个算子。如果每个张量都独立申请释放，碎片化会非常严重。
复用策略的核心思路就一个：大小匹配，类型匹配。具体来说：
- 大小匹配：释放的显存块，优先分配给大小最接近的请求。比如释放了一个4MB的块，下一个请求是3.8MB，那就直接复用。如果请求是1MB，那就找个更合适的。
- 类型匹配：float32的块尽量给float32用，int8的块给int8用。虽然显存本身不关心数据类型，但对齐要求不同，混用可能导致性能下降。
| 复用策略 | 适用场景 | 优点 | 缺点 |
|---|---|---|---|
| 最佳匹配（Best Fit） | 大小变化频繁 | 碎片少 | 查找慢 |
| 首次匹配（First Fit） | 大小相对固定 | 速度快 | 碎片多 |
| 伙伴系统（Buddy System） | 2的幂次大小 | 合并方便 | 内部碎片 |
避坑指南：我曾经在复用策略里踩过一个坑——没有考虑显存对齐。GPU显存要求64字节对齐，如果复用时不检查对齐，分配出去的地址可能不对齐，导致cudaMemcpy报错。所以复用前一定要做对齐检查。
4.3 回收策略：优雅地归还资源
回收策略听起来简单——不就是把显存还回池子里吗？但这里有个关键问题：什么时候还？
我见过两种极端做法：
- 立即回收：张量一销毁，立刻把显存标记为空闲。优点是池子利用率高，缺点是频繁操作链表，性能开销大。
- 延迟回收：张量销毁后，显存块先进入一个「待回收队列」，等队列满了或者池子压力大了再统一回收。我比较推荐这种做法，尤其是在训练循环中，张量创建销毁非常频繁。
延迟回收还有个额外好处——可以避免「ABA问题」。什么意思呢？就是一块显存刚被释放，又被另一个请求分配走，然后原来的持有者还在用这块显存（比如异步操作没完成），就会出数据竞争。延迟回收可以保证显存块在被真正复用前，有一段「冷却期」。
我的做法：延迟回收的冷却期设置为一个kernel launch的周期。也就是说，一块显存被释放后，至少要等当前stream上的所有操作完成，才能被重新分配。这样既保证了安全，又不会太保守。
4.4 扩容策略：当池子不够用时
扩容是显存池的最后一道防线。预分配再精确，也难免遇到突发的大显存需求。这时候怎么办？
扩容策略通常有三种：
- 原地扩容：如果池子后面还有空闲显存，直接扩展池子大小。这是最理想的情况，但现实中很少见。
- 搬迁移容：申请一块更大的显存，把池子里所有数据拷贝过去，然后释放旧池子。代价是拷贝开销大，而且需要暂停所有正在进行的GPU操作。
- 溢出分配：池子不够的部分，直接调用cudaMalloc额外申请。这是最简单的做法，但会导致显存碎片化加剧。
我个人习惯的做法是：先尝试溢出分配，如果溢出分配导致碎片率超过30%，就触发一次搬迁移容。这样既避免了频繁迁移的开销，又能在必要时整理碎片。
小技巧：扩容阈值不要设得太死。比如池子用了80%就开始准备扩容，而不是等到100%才动手。因为cudaMalloc本身可能失败（显存不足），提前扩容可以留出缓冲时间。
4.5 四种策略的协同工作
说了这么多，四种策略不是孤立的。它们需要协同工作，才能形成一个健壮的显存池。下面这张图展示了它们之间的关系：
从图中可以看到，四种策略形成了一个闭环：预分配提供初始资源 → 复用策略高效分配 → 回收策略归还资源 → 扩容策略应对突发需求 → 扩容结果反馈给预分配策略调整下次的初始大小。
嗯，这里要注意一点：扩容策略不应该频繁触发。如果发现扩容很频繁，说明预分配策略的初始大小设得太小了，需要调大。反过来，如果池子一直用不满，说明预分配太大了，可以适当减小。
总结一下：显存池设计没有银弹。预分配要留余量但别浪费，复用要快且准，回收要安全且高效，扩容要兜底但不频繁。这四个策略调优好了，你的显存池就能在性能和利用率之间找到最佳平衡点。
我在实际项目中，通常会用一套监控系统来观察这四个策略的运行情况。比如记录预分配的使用率、复用的命中率、回收的延迟时间、扩容的触发频率。这些指标能帮你快速定位瓶颈——是预分配太小了，还是复用策略不够智能？有了数据，调优就不再是拍脑袋了。