显存碎片化：NVIDIA CUDA内存池的运作机制与调优

做YOLO部署的朋友，十有八九都遇到过这个场景：明明显存还有好几个G，但一分配新张量就报out of memory。我刚开始搞部署时也被这个问题折磨过，后来才明白——这锅多半得甩给显存碎片化。

说白了，显存碎片化就像你搬家时往屋里塞各种形状的箱子。大箱子先进去，小箱子再塞缝隙。等大箱子搬走了，留下的空洞往往不够下一个大箱子用。CUDA的内存池，就是管理这些"箱子"的管家。

CUDA内存池长什么样？

CUDA从11.0开始引入了cudaMemPool机制。默认情况下，每个CUDA设备都有一个全局内存池。你调用cudaMalloc时，不是直接找驱动要显存，而是先问内存池有没有合适的空闲块。

我画了张图，帮你理解这个结构：

你看，内存池里既有已分配的块，也有空闲块。碎片化就发生在这些空闲块不连续的时候。

碎片化是怎么产生的？

YOLO部署中，碎片化的元凶主要有三个：

• 频繁分配释放不同大小的张量——预处理、后处理阶段尤其明显
• 流(stream)之间的异步分配——不同流上的分配顺序不可控
• 临时张量的生命周期交错——比如特征图、中间计算结果

我记得有一次排查一个YOLOv8部署的性能抖动问题。模型推理时间忽高忽低，最高能差30%。后来用nvidia-smi一看，显存分配次数在抖动时暴增。原因就是碎片化导致内存池频繁触发整理(compaction)，甚至回退到直接向驱动申请新显存。

怎么诊断碎片化？

别靠猜。CUDA Toolkit提供了几个实用工具：

工具/API	用途	我常用的场景
cudaMemPoolGetAttribute	查询池属性，如已用大小、空闲大小	快速判断池内是否有足够空闲
cuMemGetAllocGranularity	查看分配粒度（通常2MB/64KB）	确认小分配是否浪费了空间
nsys profile --cuda-memory-usage true	可视化显存分配时间线	定位碎片化高峰时段
nvidia-smi -q -d MEMORY	查看全局显存使用情况	初步排查是否整体显存不足

调优手段：从池参数到代码策略

调优不是一锤子买卖。我一般按优先级从高到低来：

1. 调整内存池的释放阈值

默认情况下，内存池不会立刻把空闲块还给驱动。你可以通过cudaMemPoolSetAttribute控制：

cudaMemPool_t pool;
cudaDeviceGetMemPool(&pool, 0);

// 设置释放阈值：当空闲内存超过200MB时，开始归还给驱动
uint64_t threshold = 200 * 1024 * 1024;
cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReleaseThreshold, &threshold);

// 设置保留比例：保留最近使用的80%块，其余可释放
int keepRatio = 80;
cudaMemPoolSetAttribute(pool, cudaMemPoolAttrReservedMemHigh, &keepRatio);

阈值设得太小，频繁与驱动交互，性能下降。设得太大，碎片化加剧。我个人习惯从模型峰值显存占用的1.5倍开始试。

2. 使用专用池隔离不同大小的分配

YOLO部署中，特征图张量通常很大（几十MB），而后处理中的边界框、置信度数组很小（几KB）。混在一个池里，大块释放后留下的空洞很难被小块填满。

CUDA 11.3+支持创建多个内存池：

cudaMemPool_t pool_large, pool_small;
cudaMemPoolCreate(&pool_large, 0);
cudaMemPoolCreate(&pool_small, 0);

// 大池：分配粒度2MB，用于特征图
uint64_t granularity = 2 * 1024 * 1024;
cudaMemPoolSetAttribute(pool_large, cudaMemPoolAttrMinAllocSize, &granularity);

// 小池：分配粒度64KB，用于后处理
granularity = 64 * 1024;
cudaMemPoolSetAttribute(pool_small, cudaMemPoolAttrMinAllocSize, &granularity);

// 分配时指定池
cudaMallocFromPoolAsync(&d_tensor, 32 * 1024 * 1024, pool_large, stream);
cudaMallocFromPoolAsync(&d_boxes, 4096, pool_small, stream);

3. 预分配 + 张量复用

这是最直接、最有效的办法。YOLO推理的输入输出尺寸通常是固定的。提前把需要的显存一次性分配好，后面只复用，不释放。

// 预分配输入、输出、中间特征图
void* d_input, *d_output, *d_feat;
size_t input_size = 3 * 640 * 640 * sizeof(float);
size_t output_size = 8400 * 85 * sizeof(float);
size_t feat_size = 256 * 80 * 80 * sizeof(float);

cudaMalloc(&d_input, input_size);
cudaMalloc(&d_output, output_size);
cudaMalloc(&d_feat, feat_size);

// 推理循环中只传指针，不分配
for (int i = 0; i < num_frames; i++) {
    cudaMemcpy(d_input, frames[i], input_size, cudaMemcpyHostToDevice);
    infer(d_input, d_output, d_feat, stream);
    // 注意：不要在这里调用cudaFree！
}

我曾经把一个YOLOv5s的部署从每帧分配释放改成预分配复用后，显存分配次数从每秒2000+降到了个位数，帧率直接翻倍。

避坑指南

几个实战建议：

• 别在热点路径上分配显存——推理循环内不要出现cudaMalloc
• 统一分配粒度——所有张量尺寸对齐到2MB或64KB的整数倍
• 用cudaMallocAsync替代cudaMalloc——前者能更好地与流配合，减少碎片
• 定期重置内存池——如果模型需要长时间运行，每隔N帧调用cudaMemPoolTrimTo清理一下

碎片化调优没有银弹。你得结合自己的YOLO模型结构、输入尺寸、后处理逻辑来试。但记住一个原则：减少分配次数，控制分配粒度，隔离大小块。做到这三点，90%的碎片化问题都能解决。