4、显存分配与释放：cudaMalloc/cudaFree机制、显存池化技术、显存碎片化问题
显存管理，说白了就是GPU的「内存条」怎么分、怎么还。很多同学写CUDA程序，上来就是cudaMalloc一把梭，跑完cudaFree一扔了事。嗯，小项目这么搞没问题，但一旦上了生产环境，你会发现显存分配和释放，其实是个挺讲究的活儿。
我个人习惯把显存管理分成三个层次来看：底层API怎么干活、中间层怎么优化、上层怎么避免踩坑。今天我们就从这三个角度，把显存分配与释放这件事聊透。
4.1 cudaMalloc与cudaFree：底层分配机制
先看最基础的。cudaMalloc这个函数，大家肯定都用过。但你有没有想过，它背后到底干了什么？
我简单梳理一下流程：
- 驱动层交互：cudaMalloc会调用CUDA驱动API，向GPU驱动申请一块物理连续的显存空间。
- 页表映射：驱动会在GPU的页表中建立虚拟地址到物理地址的映射。注意，这里分配的是虚拟地址连续、物理地址也连续的空间。
- 返回指针：用户拿到的是一个设备端虚拟地址，通过这个地址可以访问分配到的显存。
cudaFree的逻辑正好反过来：解除页表映射，释放物理显存，归还给驱动管理。
关键点：cudaMalloc/cudaFree是同步操作。也就是说，调用cudaFree时，GPU必须完成所有对该显存的访问，才能释放。这会导致隐式的同步点，影响性能。
我在项目中遇到过一个问题：一个训练脚本频繁地分配和释放中间结果的显存，结果每次cudaFree都卡住，GPU利用率掉到60%以下。后来一查，就是cudaFree触发了隐式同步，把流水线打断了。
避坑指南：我曾经在线上服务里看到cudaMalloc/cudaFree被频繁调用，导致显存分配耗时占了总时间的15%。后来改用显存池，这个问题就解决了。
4.2 显存池化技术：从「现用现分」到「提前囤货」
显存池化，说白了就是提前申请一大块显存，自己当「二房东」来管理。需要显存时，从池子里拿；用完了，还回池子，而不是真的还给驱动。
这样做的好处很明显：
- 减少驱动交互：cudaMalloc/cudaFree是重量级操作，池化后变成轻量级的内存块管理。
- 避免隐式同步：池内分配不涉及驱动层，不会触发同步。
- 降低碎片化：池化可以复用固定大小的块，减少小碎片产生。
常见的显存池实现策略有几种：
| 策略 | 原理 | 适用场景 |
|---|---|---|
| 固定大小块池 | 预先分配一批固定大小的块，按需分配 | 张量大小固定的推理场景 |
| 伙伴系统 | 按2的幂次划分块，分配时合并/拆分 | 大小变化较多的训练场景 |
| Slab分配器 | 为不同大小类别维护独立空闲链表 | 混合负载场景 |
你想想看，PyTorch的CUDA缓存分配器，其实就是个显存池。它默认会预分配2MB的块，用伙伴系统来管理。这也是为什么你有时候看到nvidia-smi显示显存占用很高，但实际只用了不到一半——因为池子里囤着货呢。
注意：显存池也不是万能的。池子太大，会浪费显存；池子太小，频繁扩容反而增加开销。我建议根据实际负载的峰值显存需求，再加10%-20%的余量来设置池大小。
4.3 显存碎片化问题：看不见的性能杀手
显存碎片化，是生产环境中非常头疼的问题。你明明看到还有好几GB显存空闲，但cudaMalloc就是报错out of memory。为什么会这样？
原因很简单：显存被切成了很多小块，每个小块都不连续。而cudaMalloc要求分配物理连续的空间，所以即使总空闲空间够，也凑不出一个连续的大块。
碎片化分两种：
- 内部碎片：分配了比实际需求大的块，多余部分浪费了。比如你申请100字节，但分配器给了128字节的块，那28字节就是内部碎片。
- 外部碎片：空闲块分散在显存各处，无法合并成连续大块。这是最头疼的。
我记得有一次排查一个推理服务的OOM问题，nvidia-smi显示显存用了6.8GB（总共8GB），但cudaMalloc申请一个512MB的块都失败。用cudaMemGetInfo一看，空闲显存1.2GB，但最大连续块只有300MB。这就是典型的外部碎片。
怎么解决？我总结了几条经验：
- 预分配大块：启动时一次性分配大部分显存，自己管理。PyTorch就是这么干的。
- 统一张量大小：尽量让张量的尺寸是2的幂次或对齐到某个粒度，减少内部碎片。
- 避免频繁分配释放：能用显存池就别裸调cudaMalloc。
- 定期整理：有些框架支持显存整理（defragmentation），把空闲块合并。但注意，这需要暂停所有GPU操作。
核心观点：显存碎片化不是一朝一夕形成的，而是分配释放模式长期积累的结果。好的分配策略，应该从一开始就考虑碎片化问题。
4.4 知识体系总览
下面这张图，我把显存分配与释放的核心逻辑串起来了。你可以看到从底层API到上层优化的完整链路。
这张图把三个核心模块串起来了。底层API是基础，池化技术是优化手段，碎片化问题是需要持续关注的坑。三者环环相扣，缺一不可。
个人建议：如果你刚开始接触显存管理，先搞懂cudaMalloc/cudaFree的同步语义。然后尝试用显存池替换裸分配，你会发现性能提升很明显。最后，养成监控显存碎片化的习惯——用cudaMemGetInfo看看最大连续块大小，比只看总空闲量靠谱得多。
嗯，显存分配与释放这块，说深了可以写一本书。但核心就这几条：理解底层机制、用好池化技术、警惕碎片化。做到这三点，你的GPU显存管理基本就稳了。