4、显存池核心架构：预分配机制、回收复用机制、碎片整理策略

显存池这东西，说白了就是给GPU提前准备好一块“内存仓库”。你想想看，每次训练模型都要频繁申请和释放显存，如果每次都走cudaMalloc，那性能损耗可不是一星半点。我在早期做分布式训练时，就吃过这个亏——模型没跑起来，光显存分配就占了小一半时间。

今天咱们就拆开显存池的三大核心机制，看看它们到底是怎么工作的。

4.1 预分配机制：提前把坑占好

预分配，顾名思义，就是在真正需要显存之前，先向GPU申请一大块连续空间。我习惯把这个过程叫做“圈地运动”。

具体怎么做呢？一般有两种策略：

• 固定大小预分配：启动时就申请一个固定大小的显存块，比如4GB、8GB。适合显存需求稳定的场景。
• 动态增长预分配：先申请一个较小的块，比如1GB，当发现不够用时再逐步扩展。我建议用这种方式，因为很多模型的实际显存需求是跑起来才知道的。

代码层面，预分配通常长这样：

// 伪代码示例：显存池预分配
class MemoryPool {
    void* base_ptr;
    size_t pool_size;
    
    void init(size_t size_mb) {
        pool_size = size_mb * 1024 * 1024;
        cudaMalloc(&base_ptr, pool_size);
        // 初始化空闲链表
        free_list.push_back({base_ptr, pool_size});
    }
};

4.2 回收复用机制：用完别扔，洗洗还能用

回收复用，是显存池的灵魂。你想想，如果每次用完显存都直接还给GPU，那预分配的意义就没了。

我遇到过最典型的场景：训练一个BERT模型，前向传播时申请了2GB显存，反向传播时释放了。如果没有回收复用，这2GB就白白浪费了。但有了显存池，这块显存会被标记为“空闲”，下一个算子可以直接拿来用。

回收复用的核心数据结构，我推荐用空闲链表：

• 每个空闲块记录起始地址和大小
• 申请时，从链表中找到大小最合适的块（最佳适配）
• 释放时，把块重新插入链表，并尝试合并相邻的空闲块

回收时还有一个坑：内存泄漏。我曾经调试过一个bug，模型跑着跑着显存就涨上去了，查了两天才发现是某个算子释放时没有正确归还给池子。所以，我建议在回收接口里加一个校验：

void deallocate(void* ptr) {
    // 校验指针是否属于本池
    if (!is_in_pool(ptr)) {
        // 直接走cudaFree，避免泄漏
        cudaFree(ptr);
        return;
    }
    // 正常回收
    free_list.insert(ptr);
}

4.3 碎片整理策略：给显存做个“大扫除”

碎片，是显存池最头疼的问题。你想想看，频繁申请和释放不同大小的块，空闲空间会变得支离破碎。明明总空闲空间还有2GB，但就是申请不到一个连续的1GB块。

碎片整理，说白了就是把分散的小块合并成大块。我常用的策略有三种：

策略名称	原理	适用场景
立即合并	释放时检查相邻块，能合并就合并	碎片不严重的场景
延迟整理	当碎片率达到阈值时，触发一次全量整理	高并发、碎片严重的场景
内存拷贝整理	把已分配块拷贝到一起，腾出连续空间	需要大块连续显存的场景

碎片率的计算，我一般用这个公式：

碎片率 = 1 - (最大连续空闲块大小 / 总空闲大小)

当碎片率超过0.3时，就该考虑整理了。

4.4 核心架构总览

说了这么多，咱们用一张图来总结显存池的核心架构：

嗯，这张图基本把显存池的三大模块串起来了。预分配是基础，回收复用是核心，碎片整理是保障。三者缺一不可。

4.5 实战中的避坑指南

最后，分享几个我在实战中踩过的坑：

• 不要盲目预分配：我曾经在一个8卡服务器上，每张卡预分配了6GB，结果模型加载时直接OOM。后来改成动态增长，问题就解决了。
• 回收时一定要校验指针：有一次，一个第三方库偷偷调了cudaMalloc，结果释放时传给了我们的池子，导致内存泄漏。加个校验就稳了。
• 碎片整理要谨慎触发：我建议只在碎片率超过0.4时触发，而且最好在模型加载或推理空闲时做。否则，整理期间的性能抖动会让你怀疑人生。

好了，显存池的核心架构就聊到这儿。下一节咱们会深入压测工具链，看看怎么用实际数据来验证这些机制的效果。