显存分配与释放机制：CUDA显存分配API、显存碎片化问题、显存池的预分配策略

聊到显存优化，很多人第一反应是「模型太大放不下怎么办」。但我在实际项目中踩过最多的坑，反而不是容量不够，而是分配和释放的时机不对。你想想看，显存这东西，分配得快，释放得慢，碎片一多，明明还有空间，就是申请不到一块连续的——这感觉，跟租房时看到满屏「已租」却找不到一套整租的差不多。

CUDA显存分配API：你用的到底是哪个？

先说说最基础的。CUDA 提供了几个显存分配的接口，我习惯把它们分成两类：显式分配和隐式分配。

  • cudaMalloc：最常用的，分配一块连续的显存。返回的是设备指针。
  • cudaMallocPitch：处理二维数组时用，会自动对齐行宽，提升访存效率。
  • cudaMalloc3D：三维数组专用，适合体素、3D卷积这类场景。
  • cudaHostAlloc：分配锁页内存，配合流和异步拷贝用。

嗯，这里要注意一个细节：cudaMalloc 是同步的。什么意思？就是调用返回时，显存已经真的分配好了。但 cudaMallocAsync（CUDA 11.2+）是异步的，它把分配操作丢到流里，不阻塞主线程。我个人习惯在推理场景用 cudaMallocAsync，训练场景反而更倾向 cudaMalloc——因为训练时分配次数少，异步带来的收益不明显。

// 传统方式
float* d_ptr;
cudaMalloc(&d_ptr, 1024 * sizeof(float));

// 异步方式（需要支持CUDA 11.2+）
cudaStream_t stream;
cudaStreamCreate(&stream);
cudaMallocAsync(&d_ptr, 1024 * sizeof(float), stream);
我的小经验： 如果你在写一个长期运行的服务，比如推荐系统或在线推理，尽量用 cudaMallocAsync。它配合 cudaMemPool 使用，能自动复用之前释放的显存块，减少向驱动申请的开销。

显存碎片化问题：看不见的「内存泄漏」

碎片化这个问题，说实话，我入行第三年才真正重视起来。当时做一个多模型并行的推理服务，显存总量 32GB，每个模型只占 4GB，按理说跑 7 个没问题。结果跑到第 5 个就报 out of memory。查了半天，发现是碎片把显存切得七零八落。

碎片化分两种：

  • 外部碎片：多次分配释放后，空闲块不连续。比如你分配了 1GB、释放了 500MB、又分配了 800MB，那 500MB 的空隙可能就用不上了。
  • 内部碎片：分配时对齐导致的浪费。比如你只要 100 字节，但 cudaMallocPitch 给你对齐到 512 字节，多出来的 412 字节就浪费了。

为什么会这样？说白了，CUDA 驱动层的分配器是通用的，它不知道你的业务模式。你今天分配一个 1.2GB 的 tensor，明天释放一个 0.8GB 的，后天又分配一个 1.5GB 的——驱动只能按「先来后到」的方式管理，久而久之就碎了一地。

我曾经踩过的坑： 在 PyTorch 里频繁调用 torch.cuda.empty_cache()。这个操作会清空缓存池，把显存还给驱动。但下次分配时又得重新申请，反而加剧了碎片化。我的建议是：不要主动调 empty_cache，让缓存池自己管理。

显存池的预分配策略：把碎片扼杀在摇篮里

解决碎片化最有效的手段，就是显存池。说白了，就是提前跟驱动申请一大块显存，然后自己在用户态管理分配和释放。这样驱动层面看到的只有一块连续的大显存，碎片问题就转移到了我们自己手里。

我常用的预分配策略有三种：

策略 适用场景 优点 缺点
固定大小池 模型尺寸固定的推理服务 零碎片，分配 O(1) 灵活性差，无法应对动态尺寸
伙伴系统 训练场景，tensor 尺寸变化大 碎片少，分配快 内部碎片最多浪费 50%
slab 分配器 多模型混合部署 按常用尺寸分类，复用率高 实现复杂，需要预知尺寸分布

我个人最常用的是伙伴系统。它的核心思想是：把显存按 2 的幂次分成块。比如你申请 300MB，它给你分配 512MB 的块。释放时，如果相邻的块也是空闲的，就合并成 1024MB。这样碎片率能控制在 10% 以内。

// 一个简化的伙伴系统分配器思路
class BuddyAllocator {
    // 假设总显存 4GB，按 2^12=4KB 为最小块
    static const int MIN_ORDER = 12;   // 4KB
    static const int MAX_ORDER = 22;   // 4GB
    
    // 每个 order 维护一个空闲链表
    std::vector<std::list<void*>> free_lists;
    
    void* allocate(size_t size) {
        int order = ceil(log2(size));
        // 如果当前 order 有空闲块，直接返回
        if (!free_lists[order].empty()) {
            void* ptr = free_lists[order].front();
            free_lists[order].pop_front();
            return ptr;
        }
        // 否则从更高 order 分裂
        return split_and_allocate(order);
    }
    
    void deallocate(void* ptr, size_t size) {
        int order = ceil(log2(size));
        // 尝试合并相邻块
        merge_with_buddy(ptr, order);
    }
};
核心要点： 显存池的预分配不是越大越好。我见过有人一上来就预分配 90% 的显存，结果模型加载时反而因为碎片化报错。建议预分配 60%-70%，留一些余量给临时变量和中间结果。

实战中的避坑指南

最后分享几个我实际项目中总结出来的经验：

  • 不要频繁分配释放：每次 cudaMalloc 都有几百微秒的开销。如果 tensor 尺寸固定，复用是最好的策略。
  • 注意流同步：异步分配时，一定要确保在释放前所有使用该显存的流都已完成。否则会出现「释放了还在用的显存」这种诡异 bug。
  • 监控碎片率：可以用 cudaMemGetInfo 获取空闲总量和最大连续块大小。如果空闲总量很大但最大连续块很小，说明碎片严重。
// 检查碎片率的实用代码
size_t free_bytes, total_bytes;
cudaMemGetInfo(&free_bytes, &total_bytes);

// 分配一个最大可能的连续块
void* test_ptr;
cudaError_t err = cudaMalloc(&test_ptr, free_bytes);
if (err == cudaSuccess) {
    // 没有碎片，可以分配全部空闲显存
    cudaFree(test_ptr);
} else {
    // 有碎片，尝试二分查找最大连续块
    size_t low = 0, high = free_bytes;
    while (low < high) {
        size_t mid = (low + high + 1) / 2;
        err = cudaMalloc(&test_ptr, mid);
        if (err == cudaSuccess) {
            cudaFree(test_ptr);
            low = mid;
        } else {
            high = mid - 1;
        }
    }
    float fragmentation = 1.0f - (float)low / free_bytes;
    printf("碎片率: %.2f%%\n", fragmentation * 100);
}

嗯，显存分配这块，说复杂也复杂，说简单也简单。核心就一句话：减少向驱动申请的次数，尽量在用户态管理。你想想看，驱动层的分配器是为通用场景设计的，而你的模型有固定的模式——那为什么不自己管呢？

一个小技巧： 如果你用 PyTorch，可以设置 torch.cuda.set_per_process_memory_fraction(0.8) 来限制每个进程的显存使用上限。这样即使有碎片，也不会影响到其他进程。

显存池的预分配策略，说白了就是「用空间换时间，用规划换碎片」。我刚开始做优化时，总觉得预分配浪费显存。后来发现，浪费 10% 的显存换来 50% 的吞吐提升，这笔账怎么算都划算。

想要更多资料，联系可可

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