4. 显存池预分配策略

显存池怎么分配，这事儿我琢磨了好几年。

说白了，预分配策略就是决定「什么时候向GPU要内存，要多少」。策略选对了，你的训练任务能快30%以上。选错了，要么显存炸了，要么性能拉胯。

我个人习惯把主流策略分成三类：固定大小池、动态增长池、混合策略。咱们一个一个聊。

4.1 固定大小池：简单粗暴，但够稳

固定大小池的思路很直接——程序启动时，一次性向GPU申请一大块显存。比如8GB，切分成若干固定大小的块。后面所有张量分配，都从池子里拿。

我在项目中遇到过这样一个场景：一个推理服务，每次请求的输入尺寸固定。用固定大小池，分配延迟从微秒级降到了纳秒级。你想想看，这差别有多大。

// 固定大小池的简化实现
class FixedPool {
    char* base;          // 显存基地址
    size_t block_size;   // 每个块大小，比如 4MB
    int num_blocks;      // 块数量，比如 2048
    std::vector<int> free_list;  // 空闲块索引
    
public:
    void* allocate() {
        if (free_list.empty()) return nullptr;
        int idx = free_list.back();
        free_list.pop_back();
        return base + idx * block_size;
    }
    
    void deallocate(void* ptr) {
        int idx = (char*)ptr - base / block_size;
        free_list.push_back(idx);
    }
};

但缺点也很明显——灵活性差。一旦遇到超大张量，池子里的块不够大，就得走fallback路径。我曾经因为这个踩过坑，一个模型升级后参数变大了，固定池直接罢工。

4.2 动态增长池：灵活，但有代价

动态增长池的策略是「按需分配，用多少要多少」。池子初始很小，比如256MB。当分配请求超过当前容量时，自动向GPU申请更大的块。

嗯，这里要注意。动态增长不是无限增长。一般会设置一个上限，比如总显存的80%。超过上限就抛出OOM错误。

动态增长池的实现，核心是「增长因子」。我见过两种主流做法：

我个人习惯用指数增长，但增长因子设为1.5而不是2。为什么？2倍太激进，容易浪费。1.5倍是个不错的折中。

// 动态增长池的核心逻辑
class DynamicPool {
    size_t current_size = 256 * 1024 * 1024;  // 初始256MB
    size_t max_size = 8LL * 1024 * 1024 * 1024; // 上限8GB
    float growth_factor = 1.5f;
    
    void grow(size_t required_size) {
        size_t new_size = current_size;
        while (new_size < required_size) {
            new_size = static_cast<size_t>(new_size * growth_factor);
        }
        if (new_size > max_size) {
            throw std::runtime_error("OOM: 超过显存上限");
        }
        // 向GPU申请新显存
        cudaMalloc(&new_ptr, new_size);
        current_size = new_size;
    }
};

4.3 混合策略：取长补短

实际工程中，纯固定或纯动态都不够完美。混合策略才是王道。

我常用的混合策略是这样的：

1. 预分配一个固定大小的核心池，覆盖90%的常见分配请求
2. 预留一个动态增长区，处理那些超大或超小的异常请求
3. 设置水位线，当固定池使用率超过80%时，提前扩容动态区

说白了，就是用固定池保证性能，用动态池兜底。

我在一个推荐系统项目中实践过混合策略。训练时batch size偶尔会波动，固定池处理大部分情况，动态池处理峰值。最终显存利用率从65%提升到了92%。

4.4 三种策略对比

做个总结，方便你选型：

你想想看，如果你的场景是「模型固定、batch固定」，固定池就够了。如果是「多模型混跑、负载波动大」，混合策略更靠谱。

4.5 知识体系图

下面这张图，把三种策略的核心逻辑串起来了：

最后说一句，没有银弹。我见过有人把固定池玩出花来，也见过混合策略搞得太复杂反而性能下降。关键还是理解你的业务场景，选最合适的策略。