显存复用核心技术实战 - 第2节
显存分配与回收机制：从cudaMalloc到显存池
大家好，我是你们的老朋友。今天咱们聊聊显存分配与回收这个看似基础、实则暗藏玄机的话题。
说实话，我刚开始做GPU加速时，觉得显存分配不就是调个cudaMalloc、用完再cudaFree嘛，有啥好讲的？直到有一次，我负责的一个训练任务跑了整整两天，突然在第三个epoch挂了——cudaMalloc返回了out of memory。我检查了所有张量释放逻辑，都没问题。后来一查，原来是显存碎片化惹的祸。嗯，从那以后，我再也不敢小看显存管理了。
一、CUDA显存分配API：cudaMalloc与cudaFree
1.1 cudaMalloc：显存分配的入口
cudaMalloc是CUDA编程中最基础的显存分配函数。它的签名很简单：
cudaError_t cudaMalloc(void** devPtr, size_t size);你传一个指针地址和大小，它给你分配一块显存。但这里有个坑——cudaMalloc是同步操作。什么意思？就是它调用完，GPU上的内存分配就完成了，不会异步返回。
关键点：cudaMalloc分配的是全局内存(global memory)，生命周期从分配到显式释放为止。它不会自动回收。
我在项目中遇到过一个问题：频繁调用cudaMalloc分配小块内存，性能会急剧下降。为什么？因为每次分配都要和驱动层交互，开销不小。你想想看，如果每层网络都单独分配一个张量，那光分配时间就够你喝一壶的。
1.2 cudaFree：释放的艺术
cudaError_t cudaFree(void* devPtr);cudaFree释放由cudaMalloc分配的内存。但注意，它也是同步的——释放操作会等待所有对该内存的访问完成。
警告：千万不要在GPU还在读写这块内存时就调用cudaFree。我曾经见过一个同事，在kernel还没执行完就释放了显存，结果数据全乱了。调试了整整一天才发现是释放时机的问题。
另外，cudaFree(NULL)是安全的，但cudaFree一个已经释放过的指针，会返回cudaErrorInvalidValue。所以我的习惯是：释放后立即把指针置为NULL。
1.3 分配与释放的性能开销
我做过一个测试：分配1000次1MB的显存，和分配1次1GB的显存，时间差了多少？
| 操作 | 耗时(ms) | 说明 |
|---|---|---|
| 1000次cudaMalloc(1MB) | ~45ms | 每次都有驱动交互 |
| 1次cudaMalloc(1GB) | ~0.8ms | 一次分配大块 |
| 1000次cudaFree(1MB) | ~38ms | 同样有开销 |
看到了吧？小内存频繁分配的开销，比大内存一次分配高得多。这就是为什么我们需要显存池。
二、显存池(Pool)的基本原理
2.1 为什么需要显存池？
说白了，显存池就是一个预分配机制。你一次性从GPU申请一大块显存，然后自己管理这块显存的使用。需要小内存时，从池子里切一块出去；用完了，还回池子里，而不是直接还给GPU。
这样做的好处很明显：
- 减少驱动交互：不用每次分配都调用cudaMalloc
- 降低碎片化：统一管理，可以合并空闲块
- 提升分配速度：池内分配是纯用户态操作，毫秒级变微秒级
2.2 显存池的核心设计
一个简单的显存池，通常包含这几个部分：
class MemoryPool {
// 空闲块列表，按大小排序
std::map<size_t, std::vector<void*>> free_blocks;
// 已分配块记录
std::unordered_map<void*, size_t> allocated_blocks;
// 总池大小
size_t pool_size;
// 从GPU申请的原始内存
void* base_ptr;
};分配时，从free_blocks里找大小合适的块；释放时，把块放回free_blocks，并尝试合并相邻的空闲块。
我的经验：实际项目中，我一般用buddy system或者slab allocator。buddy system适合大小不固定的分配，slab适合固定大小的分配。PyTorch的CachingAllocator用的就是类似buddy system的思路。
2.3 主流框架的显存池实现
PyTorch的CachingAllocator是我见过最成熟的显存池之一。它的核心思想是：
- 按2的幂次分桶（1KB、2KB、4KB...）
- 分配时向上取整到最近的桶
- 释放时先放回桶里，不立即还给GPU
- 当总空闲内存超过阈值时，才调用cudaFree归还
TensorFlow的BFC Allocator也是类似，但它用的是best-fit策略，找最接近需求大小的空闲块。
三、显存碎片化问题
3.1 碎片化是怎么产生的？
显存碎片化，说白了就是显存被切成了很多小块，虽然总空闲空间够，但找不到一块连续的大空间来满足分配请求。
我举个例子：你分配了A(100MB)、B(200MB)、C(100MB)，然后释放了B。这时空闲空间是200MB，但被分成了两块100MB。如果你要分配150MB，就会失败——因为没有连续150MB的空闲块。
核心问题：碎片化导致明明有足够显存，却分配不出大块内存。这是很多OOM错误的真正元凶。
3.2 碎片化的类型
碎片化分两种：
- 内部碎片：分配的内存比实际需求大。比如你要100MB，但池子只提供128MB的块，多出来的28MB就是内部碎片
- 外部碎片：空闲块分散在各处，无法合并成连续大块
内部碎片好解决，用更细的粒度分桶就行。外部碎片就麻烦了，需要做内存整理。
3.3 如何缓解碎片化？
我总结了几条实战经验：
- 预分配大块：启动时一次性分配足够大的显存池，减少动态分配
- 对齐分配：所有分配按256字节或512字节对齐，减少内部碎片
- 延迟释放：释放的内存先缓存起来，等积累到一定量再合并
- 内存整理：在空闲时做defragmentation，移动已分配块来合并空闲空间
避坑指南：我曾经在一个推理服务中，因为频繁创建和销毁不同大小的中间张量，导致显存碎片化严重。后来我改用固定大小的预分配池，把中间张量都复用起来，碎片率从30%降到了不到5%。
四、知识体系总览
下面这张图，是我自己整理的显存管理核心逻辑。你看一遍就能明白整个体系：
这张图把显存管理的四个层次串起来了。从最底层的API调用，到中间的池化策略，再到碎片化问题，最后到解决方案。你顺着箭头看，就能理解整个体系是怎么运作的。
五、实战建议
最后，我给大家几条实战建议：
- 别裸用cudaMalloc：除非你只分配一两次，否则一定要用显存池
- 监控碎片率：我习惯在代码里加个hook，定期打印空闲块数量和大小分布。碎片率超过20%就要警惕了
- 复用张量：能复用的中间结果尽量复用，别频繁创建销毁
- 调大池子：如果显存够用，把池子预分配大一点，能有效减少碎片
小技巧：用cudaMemGetInfo(&free, &total)可以查看当前显存使用情况。我经常在关键节点打印这个值，用来排查显存泄漏和碎片化问题。
好了，这一章的内容就到这里。显存管理看似简单，但真正做好并不容易。希望我的这些经验能帮你少踩一些坑。
想要更多资料，联系可可