3. 静态内存复用：计算图分析、内存池化技术
好，咱们进入第三章。静态内存复用，这名字听着挺唬人，对吧？
说白了，就是让显存里的数据「住得更挤一点」。你想想看，深度学习模型跑起来，中间会产生多少临时张量？有些张量用完就扔，有些要留到后面用。如果我们能提前规划好，让不同时间点产生的张量共用同一块显存，那就能省下大把空间。
3.1 计算图分析：先画地图再出发
静态复用的第一步，是搞清楚张量的生命周期。我个人习惯把计算图想象成一张地铁线路图——每个节点是一个算子，每条边是一个张量。张量什么时候产生？什么时候被最后一次使用？这就是它的「生存区间」。
举个例子，一个简单的计算图：
# 伪代码示意
a = Input() # 生存区间 [0, 3]
b = Conv(a) # 生存区间 [1, 4]
c = ReLU(b) # 生存区间 [2, 5]
d = Add(a, c) # a 最后一次使用在 3，c 最后一次使用在 5
你看，a 在步骤 3 之后就没用了，b 在步骤 4 之后也没用了。那 a 和 b 能不能共用同一块显存？嗯，这里要注意——它们的生存区间有重叠（步骤 1-3 两者都活着），所以不能复用。但 a 和 c 呢？a 在 3 结束，c 在 2 才开始，完全不重叠，可以复用。
核心思想：张量生存区间不重叠 → 可以复用同一块显存。
我在项目中遇到过一个大模型，显存占用死活降不下来。后来一分析，发现是某个中间张量的生存区间被拉长了——因为后面有个分支还在引用它。其实那个分支根本不需要原始张量，只需要它的 shape 信息。改了一行代码，显存直接降了 30%。
3.2 内存池化技术：别频繁找系统要内存
搞清楚了哪些张量可以复用，接下来就是怎么实现。你想想看，如果每次分配显存都去调用 cudaMalloc，那性能会惨不忍睹。系统调用是有开销的，而且显存分配本身也可能导致碎片。
内存池化，就是提前申请一大块显存，然后自己管理。我习惯把它比作「包工头」——先跟系统要一整块地，然后自己画格子，谁要用就分一小块，用完收回。
常见的池化策略有几种：
| 策略 | 原理 | 适用场景 |
|---|---|---|
| Buddy System | 将内存按 2 的幂次分割，合并时找 buddy | 大小规整的张量，碎片少 |
| Slab Allocator | 预分配固定大小的 slab，每个 slab 存同类张量 | 大量同 shape 的张量（如 Transformer 中的 QKV） |
| Best Fit / First Fit | 从空闲链表中找最合适/第一个可用的块 | 通用场景，但可能有碎片 |
我曾经在一个推理框架里用 Buddy System，结果发现小张量太多，碎片率高达 40%。后来换成 Slab Allocator，专门为 attention 层的 QKV 张量预分配固定大小的 slab，碎片率直接降到 5% 以下。
我的建议：不要迷信某一种策略。实际项目中，往往是多种策略组合使用。大张量用 Buddy System，小张量用 Slab Allocator，中间层用 Best Fit。
3.3 静态复用 vs 动态复用
这里有个容易混淆的点。静态复用是在编译期（或图构建期）就确定好的，而动态复用是在运行时根据实际情况调整的。
静态复用的好处是确定性高——你提前算好了每块显存归谁用，运行时直接按计划执行，没有额外开销。坏处是灵活性差——如果计算图变了，或者输入 shape 变了，之前的复用计划可能就废了。
动态复用则相反，它可以在运行时根据当前内存压力动态调整。比如 PyTorch 的 caching allocator，就是动态的——它会缓存释放的显存块，下次分配时优先复用。
避坑指南：我曾经在一个动态 shape 的模型上用了纯静态复用，结果输入尺寸一变，显存直接爆了。后来改成「静态规划 + 动态兜底」——先按最大 shape 做静态复用，如果实际 shape 更小，多余的空间留给动态分配。
3.4 实战：手写一个简单的内存池
光说不练假把式。咱们写个极简的内存池，感受一下核心逻辑。
class SimpleMemoryPool:
def __init__(self, total_size):
self.total_size = total_size
self.free_blocks = [(0, total_size)] # (offset, size)
self.allocated = {} # name -> (offset, size)
def allocate(self, name, size):
for i, (offset, block_size) in enumerate(self.free_blocks):
if block_size >= size:
# 分配
self.allocated[name] = (offset, size)
# 更新空闲块
if block_size == size:
self.free_blocks.pop(i)
else:
self.free_blocks[i] = (offset + size, block_size - size)
return offset
raise MemoryError("Out of memory")
def free(self, name):
offset, size = self.allocated.pop(name)
# 合并相邻空闲块
self.free_blocks.append((offset, size))
self.free_blocks.sort()
# 合并逻辑省略...
你看，核心就两个操作：分配时找空闲块，释放时合并空闲块。当然，实际生产级的池子要复杂得多——要考虑对齐、碎片整理、多线程安全等等。
3.5 计算图分析 + 内存池 = 黄金组合
把这两者结合起来，就是静态内存复用的完整方案：
- 图分析阶段：遍历计算图，计算每个张量的生存区间。
- 区间着色：把生存区间不重叠的张量分配到同一个「颜色」（即同一块显存）。
- 池化分配：根据着色结果，从内存池中分配显存块。
- 运行时执行：按计划使用显存，用完立即归还到池中。
我参与过的一个项目，用这套流程把 BERT 的显存占用从 12GB 降到了 7GB。嗯，省下来的显存刚好可以塞进一个更大的 batch size，训练速度直接快了 40%。
一句话总结：静态内存复用 = 算清楚谁什么时候死 + 让活着的人挤一挤。
3.6 本章知识体系
下面这张图展示了本章的核心逻辑：
嗯，这张图把咱们今天讲的内容串起来了。从计算图分析开始，到区间着色，再到内存池化，最后落到策略选择。每一步都有坑，但也都有对应的解法。
想要更多资料，联系可可