显存池架构设计：层级结构、内存块管理策略、分配与回收机制

好，咱们直接进入正题。

显存池这东西，说白了就是给GPU搞一个「内存管家」。你想想看，每次训练模型都要频繁申请释放显存，那抖动得多严重。我早期做分布式训练时就吃过这个亏——模型没跑起来，先被显存碎片搞崩溃了。

一、显存池的层级结构

我个人习惯把显存池分成三层。不是拍脑袋定的，是踩坑踩出来的。

层级	名称	作用	典型大小
L1	线程级缓存	每个CUDA线程独享，零延迟	几十KB
L2	块级池	CTA（线程块）内共享，减少跨块通信	几MB
L3	全局池	整个进程共享，兜底分配	几百MB到几GB

为什么这么分？因为GPU的访存模式跟CPU不一样。L1层我用的是__shared__内存，配合寄存器，基本能做到零开销。L2层我一般用cudaMalloc预分配一大块，然后自己切。L3层嘛，就是兜底的——当L1和L2都拿不到内存时，才去全局池里捞。

核心原则：能局部就别全局，能预分配就别动态申请。

二、内存块管理策略

这块我踩过不少坑。最早我用的就是简单的「链表+首次适应」，结果跑大模型时碎片化严重。后来改成了「伙伴系统+空闲链表」的组合拳。

2.1 伙伴系统

伙伴系统说白了就是把内存按2的幂次切块。比如你要申请3KB，我就给你4KB的块。浪费一点空间，但换来的是极快的合并速度。

// 伪代码示例
struct Block {
    size_t size;      // 块大小（2的幂次）
    bool   free;      // 是否空闲
    Block* next;      // 同大小链表的下一个
    Block* buddy;     // 伙伴指针
};

// 分配：从对应大小的空闲链表中取
Block* alloc(size_t req_size) {
    size_t actual = next_power_of_two(req_size);
    int idx = log2(actual);
    if (free_list[idx] != nullptr) {
        return pop_block(free_list[idx]);
    }
    // 从更大块分裂
    return split_block(idx + 1);
}

// 回收：检查伙伴是否空闲，是则合并
void free(Block* block) {
    Block* buddy = get_buddy(block);
    if (buddy->free && buddy->size == block->size) {
        merge(block, buddy);
    } else {
        insert_to_free_list(block);
    }
}

我在项目中遇到过一个问题：频繁分配释放小内存块时，伙伴系统的内部碎片率会飙到30%以上。后来我加了个「小对象缓存池」，专门处理64字节以下的请求，这才压下来。

2.2 空闲链表优化

纯伙伴系统有个毛病——大块分配时分裂次数太多。我一般会再加一层「大块缓存」：超过1MB的请求直接走cudaMalloc，不走池子。为什么？因为大块分配本身开销占比小，不值得用池子去管理。

我的经验：池子管理的块大小范围建议在64B~1MB之间。太小了管理开销占比大，太大了不如直接走驱动。

三、分配与回收机制

分配和回收，说白了就是「借」和「还」。但GPU的借还跟CPU不一样——你得考虑异步执行的问题。

3.1 分配策略

我常用的分配策略有三种，按场景选：

最快适配：从空闲链表中找第一个大小够的块。适合小请求频繁的场景。
最佳适配：找大小最接近的块。碎片率最低，但查找慢。适合大块分配。
缓存优先：优先从L1/L2层分配，不够再往上层要。这是默认策略。

你可能会问：「为什么不直接用最佳适配？」嗯，因为GPU的访存延迟敏感，查找时间长了反而得不偿失。我一般只在L3层用最佳适配，L1和L2层用最快适配。

3.2 回收机制

回收比分配更讲究。我曾经因为回收时机不对，导致显存泄漏查了两天。

核心要点就两个：

延迟回收：不立即把内存还给驱动，而是先放回空闲链表。等池子里的空闲块超过阈值（比如80%），再真正释放给驱动。
异步安全：GPU的kernel是异步执行的。你回收内存时，得确保没有正在使用的流还在访问这块内存。我一般用cudaEvent做同步点。

// 延迟回收示例
void deferred_free(Block* block, cudaStream_t stream) {
    // 记录回收时的流状态
    cudaEvent_t event;
    cudaEventCreate(&event);
    cudaEventRecord(event, stream);
    
    // 把块放入延迟队列
    deferred_queue.push({block, event});
}

// 真正回收的时机
void process_deferred_queue() {
    for (auto& item : deferred_queue) {
        if (cudaEventQuery(item.event) == cudaSuccess) {
            // 流执行完毕，可以安全回收
            insert_to_free_list(item.block);
            cudaEventDestroy(item.event);
        }
    }
}

注意：千万别在kernel还没跑完时就回收显存。我见过有人用cudaDeviceSynchronize硬等，结果性能直接腰斩。用事件查询才是正解。

四、架构总览图

下面这张图是我自己项目里用的显存池架构，画出来给大家参考：

这张图我画得比较简洁。实际项目中，每层之间还有「水位线」控制——比如L1的空闲率低于20%时，提前从L2预取一批块过来。这个后面讲调优时会细说。

五、避坑指南

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

别用全局锁：GPU的并发度很高，全局锁会直接干掉性能。我改用cuda::atomic做无锁队列。
注意对齐：显存分配要256字节对齐，否则某些GPU会报错。我一般直接按512字节对齐，省心。
别忽略碎片整理：跑长任务时，显存碎片会慢慢累积。我每1000次分配做一次碎片整理，把空闲块合并成大块。

一个小技巧：在显存池初始化时，按2的幂次预分配好所有块。虽然启动慢一点，但运行时零分配开销。适合推理场景。

好了，显存池的架构设计就聊到这儿。记住一句话：好的架构是「预则立，不预则废」。提前规划好层级、策略、回收机制，后面调优时能省一半力气。

想要更多资料，联系可可

📚 课程目录

显存池架构设计：层级结构、内存块管理策略、分配与回收机制

一、显存池的层级结构

二、内存块管理策略

2.1 伙伴系统

2.2 空闲链表优化

三、分配与回收机制

3.1 分配策略

3.2 回收机制

四、架构总览图

五、避坑指南