显存分配与释放机制：从cudaMalloc到显存池

做深度学习框架开发，绕不开一个核心问题——显存管理。

我记得刚入行那会儿，写一个简单的训练脚本，用cudaMalloc分配显存，用完再cudaFree释放。看起来天经地义对吧？但跑大模型时，光分配显存就能卡住几百毫秒。更别说频繁分配释放带来的碎片问题，简直让人抓狂。

这一讲，我们就来聊聊显存分配与释放的那些坑，以及显存池怎么解决这些问题。

cudaMalloc/cudaFree的痛点

先说说cudaMalloc。这个API的作用是向GPU申请一块连续显存。听起来简单，但实际开销不小。

痛点一：分配速度慢

cudaMalloc内部需要调用驱动层的函数，涉及内核态切换。每次分配，GPU驱动都要维护一个空闲块列表，查找合适大小的连续空间。我实测过，一次cudaMalloc大概耗时几十到几百微秒。单次看不多，但训练过程中动辄成千上万次分配，累积起来就是秒级的延迟。

痛点二：释放开销不可忽略

cudaFree也不是免费的。它需要通知驱动回收显存，更新空闲列表。更麻烦的是，cudaFree是同步操作——它会等待GPU上所有之前提交的任务完成，确保没有正在使用这块显存。这就导致了一个问题：你明明只是想释放显存，结果却阻塞了计算流水线。

痛点三：显存碎片化

这是最隐蔽的坑。频繁分配释放不同大小的显存块，会导致显存空间变得支离破碎。比如你分配了4个1GB的块，释放了第2个和第4个，剩下两个1GB的块中间夹着2GB的空闲空间。这时候你想分配一个2GB的连续块，明明总空闲显存够，但cudaMalloc就是返回失败。

显存池的预分配策略

既然cudaMalloc/cudaFree有这么多问题，那怎么办？答案就是——显存池。

显存池的核心思想很简单：提前向GPU申请一大块显存，然后自己管理这块显存的分配和释放。说白了，就是把cudaMalloc/cudaFree的调用次数降到最低。

预分配策略一：固定大小池

最朴素的做法。预先分配固定大小的显存块，比如每个块256MB。当框架需要显存时，从池中取出一个块；用完再还回去。这种策略实现简单，但缺点也很明显——如果某个张量需要300MB，你就得分配两个块，造成内部碎片。

预分配策略二：按需增长池

我比较推荐的做法。初始分配一个较小的池，比如2GB。当请求的显存超过当前池大小时，再调用cudaMalloc追加分配。这样既避免了浪费，又不会频繁调用cudaMalloc。

预分配策略三：伙伴系统

这是Linux内核中常用的内存管理算法。把显存块按2的幂次大小划分。比如最小块64KB，然后128KB、256KB...直到最大块。分配时，找到大小最合适的块；释放时，如果相邻块也是空闲的，就合并成更大的块。这种策略能有效减少外部碎片。

回收与复用机制

显存池的另一个关键能力是回收与复用。说白了，就是张量用完显存后，不真正释放给GPU驱动，而是留在池里给下一个张量用。

回收策略一：引用计数

每个显存块记录一个引用计数。当张量引用这块显存时，计数加1；张量销毁时，计数减1。当计数降到0时，显存块回到空闲列表。这是最基础的回收方式，PyTorch的CachingAllocator就是这么干的。

回收策略二：延迟回收

有时候张量销毁了，但马上又有新的张量需要分配。如果每次计数归零就回收，反而增加了管理开销。延迟回收的做法是：计数归零后，先把块标记为「待回收」，等空闲列表不足时再真正回收。这样能减少不必要的回收操作。

复用机制：大小匹配

当新张量请求显存时，显存池会从空闲列表中找一个大小最接近的块。如果找不到完全匹配的，就找一个更大的块，然后切分。如果所有块都太小，就触发新的分配。

这里有个细节：大小匹配的阈值。我习惯设置一个容忍度，比如允许分配比请求大10%的块。如果超过这个比例，就重新分配，避免浪费。

核心逻辑流程图

下面这张图展示了显存池的核心工作流程。从张量请求显存开始，到最终回收复用，每一步都清晰可见。

不同框架的显存池实现对比

了解了基本原理，我们来看看主流框架是怎么做的。下面这个表格是我整理的核心差异。

嗯，这一讲的内容就到这里。显存池的设计看似简单，但实际落地时有很多细节。比如多线程安全、显存泄漏检测、与CUDA Stream的交互等等。这些我们后面会陆续讲到。

想要更多资料，联系可可