3. 显存池核心设计：预分配策略、块管理算法、回收复用机制

好，咱们直接进入正题。显存池这东西，说白了就是给GPU显存搞了个「内存管理器」。你想想看，每次训练模型都手动调用cudaMalloc，那效率得多低？我早年做分布式训练时就吃过这个亏——频繁申请释放显存，结果大部分时间都花在等待上了。

今天我就把显存池最核心的三个设计点掰开揉碎讲清楚：预分配策略、块管理算法、回收复用机制。这三个东西搞明白了，你基本就能自己手写一个轻量级显存池了。

3.1 预分配策略：别等用的时候再申请

预分配，顾名思义，就是提前把显存「占上」。为什么要这么做？因为cudaMalloc是个同步操作，它要等GPU返回确认。你频繁调用，程序就卡在那干等。

我个人的习惯是：启动时一次性申请一个大块显存，比如总显存的80%。剩下的20%留给系统和其他进程。这样后面所有张量分配都在池子里完成，速度能快一个数量级。

预分配的核心参数：
  • 初始大小：一般设为总显存的70%-90%，具体看你的模型峰值需求
  • 增长策略：如果不够用，是按固定大小增长（比如每次加1GB），还是按比例增长（比如每次翻倍）？我建议用指数增长+上限截断，避免频繁扩容
  • 对齐要求：显存分配通常要256字节或512字节对齐，否则某些算子会报错

举个例子，PyTorch的CachingAllocator就是这么干的。它启动时申请2MB的小块和20MB的大块，后面按需分配。不够了再申请更大的块，但不会超过max_split_size的限制。

// 伪代码：预分配策略
class MemoryPool {
    size_t total_size = 0.8 * get_device_memory();  // 占80%
    size_t used_size = 0;
    void* base_ptr = cudaMalloc(total_size);  // 一次性申请
    
    void* allocate(size_t size) {
        // 对齐到256字节
        size = align_up(size, 256);
        // 从池子里找空闲块...
    }
};
注意：预分配不是越多越好。我见过有人直接申请95%的显存，结果系统显存不足，CUDA直接报OOM。留点余量，给临时变量和中间结果用。

3.2 块管理算法：怎么找、怎么分、怎么记

显存池拿到一大块连续显存后，怎么管理里面的小块？这就涉及到块管理算法了。常见的就三种：伙伴算法、SLAB算法、最佳适应算法。

我直接说结论：深度学习场景下，最佳适应算法最实用。为什么？因为模型训练时，张量大小相对固定（比如权重、梯度、激活值），不会频繁出现大小波动。最佳适应能找到最匹配的空闲块，碎片最少。

3.2.1 伙伴算法

把显存分成2的幂次大小的块。比如1MB、2MB、4MB...分配时找最小的满足条件的块，如果不够就分裂。释放时合并相邻的伙伴块。

  • 优点：合并快，外部碎片少
  • 缺点：内部碎片严重（比如申请3MB，给你4MB的块，浪费1MB）

3.2.2 SLAB算法

预先创建固定大小的对象池。比如专门给4KB、16KB、64KB的张量准备独立的池子。分配时直接从对应池子里取。

  • 优点：分配速度极快，O(1)复杂度
  • 缺点：灵活性差，大小不匹配时浪费严重

3.2.3 最佳适应算法

维护一个按大小排序的空闲块列表。分配时遍历列表，找大小最接近的块。释放时插入到合适位置，并合并相邻空闲块。

  • 优点：碎片最少，利用率高
  • 缺点：分配速度稍慢（O(n)），但n通常不大
我的经验：实际项目中，我通常用最佳适应+哈希索引。用红黑树或跳表维护空闲块，查找复杂度降到O(log n)。对于大小固定的张量（比如Transformer的权重），再叠加SLAB思想，单独开一个池子。

3.3 回收复用机制：别让显存闲着

回收复用，就是张量用完了，显存别急着还给系统，先放池子里等着复用。这里面有两个关键点：什么时候回收和怎么复用。

3.3.1 延迟回收

张量析构时，不调用cudaFree，而是把块标记为「空闲」。这样下次分配时直接复用，省去了与GPU的交互开销。

我遇到过一个问题：如果延迟回收太久，显存会被「死块」占满。比如训练循环中，每次迭代都创建临时张量，但回收不及时，池子满了，新分配只能去申请系统显存。结果就是显存越占越多，最后OOM。

解决方案是设置回收阈值：当空闲块总大小超过某个值（比如池子大小的20%），就触发一次真正的cudaFree，把多余显存还给系统。

3.3.2 碎片整理

频繁分配释放会产生碎片。比如你申请了3个1GB的块，释放了中间那个，结果池子里有两个1GB的空闲块，但中间隔着一个已分配的块。这时候你想分配一个2GB的块，就失败了——尽管总空闲空间够，但不连续。

碎片整理的思路：移动已分配块，合并空闲块。但GPU显存不像CPU内存那样支持「内存移动」，因为张量指针是固定的，移动后所有引用都得更新。所以实际中很少做真正的碎片整理。

更实用的做法是避免碎片产生：

  • 分配时尽量从大块的一端开始，释放时从另一端开始（类似栈式分配）
  • 对大小相近的张量，尽量分配在连续区域
  • 使用地址排序：分配时优先从低地址开始，释放时合并相邻块
回收复用的核心数据结构：
结构 作用 复杂度
空闲块列表 按大小排序，快速查找最佳匹配 O(log n)
已分配块哈希表 根据指针快速找到块元信息 O(1)
延迟回收队列 暂存待释放的块，批量处理 O(1)

3.4 一张图看懂显存池核心流程

说了这么多，不如直接看张图。下面是我手绘的显存池核心流程，从分配到回收，一目了然。

显存池核心流程 预分配大块显存 收到分配请求 (size, align) 查找最佳匹配空闲块 未找到 找到 扩容 / 触发碎片整理 分割块，标记已分配 返回指针 + 记录元信息 释放请求 标记为空闲 合并相邻空闲块 检查回收阈值 图：显存池分配与回收核心流程

3.5 避坑指南：我踩过的那些坑

最后分享几个实战中容易翻车的地方：

  • 对齐问题：我曾经因为没做256字节对齐，结果在调用cuBLAS时直接crash。排查了半天才发现是地址不对齐。记住：cudaMalloc返回的地址天然对齐，但你自己管理的池子不一定。
  • 多流同步：多个CUDA流同时分配释放，不加锁的话，空闲列表会被破坏。我建议用细粒度锁，每个大小类一个锁，别一把大锁锁死全部。
  • 显存泄漏检测：池子里的块如果忘记释放，不会像cudaMalloc那样直接报错。你需要在池子析构时检查是否有未释放的块，打印警告。
  • 峰值显存控制：预分配太多，模型加载时可能OOM。我一般先跑一次小批量，估算峰值显存，再动态调整池子大小。
一句话总结：显存池的核心就是「提前占、快速找、延迟还」。预分配减少系统调用，块管理降低碎片，回收复用提升吞吐。把这三点做好，你的训练速度至少能快20%。

好了，这一章就到这里。显存池的设计其实还有很多细节，比如跨设备管理、异步回收、内存压缩等，但核心骨架就是上面这些。你先把预分配、块管理、回收复用这三个机制吃透，后面再深入就轻松了。

← 上一章 📖 返回目录 下一章 →