显存复用策略：内存池设计、对象池模式、显存碎片整理

大家好，我是你们的GPU系统讲师。今天咱们聊一个实战中绕不开的话题——显存复用。

说实话，我见过太多同学在训练大模型时，明明显存够用，却因为分配策略不当导致OOM。或者模型跑着跑着，显存碎片越来越多，最后卡死在分配上。嗯，这些问题，其实都可以通过合理的复用策略来解决。

为什么需要显存复用？

先问大家一个问题：每次需要显存时都去cudaMalloc，用完就cudaFree，这样不好吗？

我刚开始做CUDA开发时也是这么干的。直到有一次，我在一个实时推理系统里频繁分配释放显存，结果发现系统响应越来越慢。一查，原来每次cudaMalloc都要和驱动交互，开销巨大。更糟的是，频繁分配释放会导致显存碎片化——就像硬盘碎片一样，明明总空间够，但就是分配不出一块连续的大内存。

所以，显存复用的核心目标就两个：

• 减少分配次数——避免频繁调用cudaMalloc/cudaFree
• 降低碎片化——让显存利用率更高

内存池（Memory Pool）设计

内存池说白了就是：提前申请一大块显存，自己管理。需要时从池子里拿，用完还回去。这样就不用每次都找驱动要了。

我个人的习惯是，把内存池设计成两级结构：

这里我贴一段简化版的内存池实现，大家感受一下：

class MemoryPool {
private:
    // 空闲块列表，按大小排序
    std::map<size_t, std::vector<void*>> free_blocks;
    // 已分配块记录
    std::unordered_map<void*, size_t> allocated;
    
public:
    void* allocate(size_t size) {
        // 查找合适大小的空闲块
        auto it = free_blocks.lower_bound(size);
        if (it != free_blocks.end()) {
            void* ptr = it->second.back();
            it->second.pop_back();
            if (it->second.empty()) free_blocks.erase(it);
            allocated[ptr] = size;
            return ptr;
        }
        // 没有合适的，向驱动申请
        void* ptr;
        cudaMalloc(&ptr, size);
        allocated[ptr] = size;
        return ptr;
    }
    
    void deallocate(void* ptr) {
        size_t size = allocated[ptr];
        allocated.erase(ptr);
        // 归还到空闲列表
        free_blocks[size].push_back(ptr);
    }
};

对象池模式

对象池和内存池思路类似，但更高级一点。它不只管理内存，还管理对象的生命周期。

举个例子，在Transformer推理时，每个token都要创建attention mask、position encoding等临时对象。如果每次都new一个，用完再delete，那效率太低了。

对象池的做法是：

1. 初始化时创建一批对象
2. 需要时从池中借出一个
3. 用完归还，对象重置但不销毁

我曾经在一个BERT推理服务里用了对象池，QPS直接提升了30%。为什么？因为省去了反复构造和析构的开销。

template<typename T>
class ObjectPool {
    std::vector<T*> pool;
    std::vector<bool> in_use;
    
public:
    T* acquire() {
        for (size_t i = 0; i < pool.size(); i++) {
            if (!in_use[i]) {
                in_use[i] = true;
                return pool[i];
            }
        }
        // 池不够用了，扩容
        T* obj = new T();
        pool.push_back(obj);
        in_use.push_back(true);
        return obj;
    }
    
    void release(T* obj) {
        for (size_t i = 0; i < pool.size(); i++) {
            if (pool[i] == obj) {
                in_use[i] = false;
                // 重置对象状态
                obj->reset();
                return;
            }
        }
    }
};

显存碎片整理

碎片整理是个老话题了。在CPU内存里，我们有GC、有内存压缩。但在GPU上，事情没那么简单——你不能随便移动显存里的数据，因为GPU kernel可能正在访问它。

我总结了几种常见的碎片整理策略：

• 被动整理：在分配失败时触发，尝试合并相邻空闲块
• 主动整理：定期检查碎片率，超过阈值就整理
• 预防性分配：提前分配好常用大小的块，避免小碎片产生

这里我画了一张图，展示碎片整理的核心逻辑：

实际实现时，我建议采用延迟合并策略。什么意思呢？就是释放块时不立即合并，而是等碎片率达到阈值再统一处理。这样能避免频繁合并的开销。

最后说一句，这些策略不是孤立的。我通常会在一个系统里同时使用内存池和对象池，再配合碎片整理。比如，小张量走对象池，大张量走内存池，定期检查碎片情况。这样组合下来，显存利用率能提升20%以上。

好了，这一章就到这里。记住，显存管理是个实践活，光看理论没用。建议大家在自己的项目里试试这些策略，遇到问题欢迎交流。