4、静态显存复用策略：计算图分析、张量生存期推断、静态内存规划算法
各位同学，今天我们来聊聊显存复用里最硬核的一块——静态策略。
说白了，就是让框架在运行你的模型之前，先算清楚：
哪些张量可以“住”在同一个内存地址上？
什么时候可以回收？什么时候必须保留？
我刚开始做框架优化时，总觉得动态分配就够用了。直到跑一个 70B 的大模型，显存直接爆掉。嗯，从那以后，我彻底转向了静态复用。
4.1 计算图分析：先看清全局
静态复用的第一步，是拿到一张完整的计算图。
计算图是什么？
就是你的模型在训练或推理时，所有算子（Op）和数据流（Tensor）构成的有向无环图。
我个人习惯把计算图分成三层来看：
- 算子节点：Conv、MatMul、ReLU……每个节点代表一次计算
- 张量边：节点之间的输入输出，每个张量都有大小和类型
- 控制依赖：比如同步点、梯度裁剪，这些会影响生存期
举个例子，一个简单的两层 MLP：
# 伪代码：计算图结构
x → Linear1 → relu → Linear2 → loss
↑ ↑ ↑
weight1 bias1 weight2在这个图里，x 只在第一层被用到，relu 的输出只传给第二层。这些信息，就是后续内存规划的“原材料”。
核心要点：计算图分析越细，内存复用空间越大。我见过有些框架只分析算子级别的依赖，结果漏掉了中间张量的复用机会。
4.2 张量生存期推断：谁先死，谁后死
有了计算图，下一步就是推断每个张量的生存期。
生存期怎么定义？
很简单：从张量被创建（某个算子的输出），到最后一次被使用（作为另一个算子的输入）。
我常用的方法是拓扑序扫描：
- 对计算图做拓扑排序
- 正向扫描：记录每个张量的出生点（第几个算子）
- 反向扫描：记录每个张量的死亡点（最后一次被引用的算子）
来看一个实际例子：
算子序号: 1 2 3 4 5
操作: Conv → ReLU → Pool → FC → Softmax
张量: t1 t2 t3 t4 t5
生存期推断结果：
t1: 出生=1, 死亡=2 (只被ReLU使用)
t2: 出生=2, 死亡=3 (只被Pool使用)
t3: 出生=3, 死亡=4 (只被FC使用)
t4: 出生=4, 死亡=5 (只被Softmax使用)
t5: 出生=5, 死亡=5 (输出，不再被使用)你想想看，t1 在算子2之后就死了，t2 在算子3之后也死了。那它们的内存地址，完全可以被后面的张量复用。
避坑指南：我曾经遇到过一个 case，某个张量在反向传播时被多次引用，但正向扫描只记录了第一次使用。结果内存提前释放，导致梯度计算错误。所以，一定要考虑反向传播的依赖。
4.3 静态内存规划算法：把地址“安排”得明明白白
生存期推断完了，接下来就是最核心的——内存规划算法。
目标是什么？
在满足所有张量生存期不冲突的前提下，最小化总显存占用。
这其实是一个区间着色问题：
- 每个张量是一个“区间”（出生点到死亡点）
- 每个内存地址是一种“颜色”
- 重叠的区间不能共享颜色
我常用的算法有三种，各有优劣：
| 算法 | 原理 | 复杂度 | 适用场景 |
|---|---|---|---|
| 贪心区间分配 | 按死亡点排序，先死先复用 | O(n log n) | 推理场景，图结构简单 |
| 最佳适应算法 | 找最小可用空闲块 | O(n²) | 训练场景，张量大小差异大 |
| 整数线性规划 | 全局最优解 | NP-hard | 离线编译，追求极致压缩 |
我个人最常用的是贪心区间分配，因为它快，而且效果不差。
具体做法：
- 把所有张量按死亡点从小到大排序
- 维护一个“空闲地址池”
- 遍历每个张量：如果空闲池里有地址可用，就分配；否则申请新地址
- 张量死亡后，把地址归还到空闲池
来看一个代码示例：
def greedy_interval_alloc(tensors):
# tensors: [(name, size, birth, death), ...]
sorted_tensors = sorted(tensors, key=lambda t: t[3]) # 按死亡点排序
free_pool = [] # 空闲地址列表
allocation = {} # 张量名 → 地址偏移
for name, size, birth, death in sorted_tensors:
# 尝试从空闲池中找可用地址
allocated = False
for i, (addr, addr_size) in enumerate(free_pool):
if addr_size >= size:
allocation[name] = addr
# 更新空闲块大小
if addr_size > size:
free_pool[i] = (addr + size, addr_size - size)
else:
free_pool.pop(i)
allocated = True
break
if not allocated:
# 申请新地址
new_addr = sum(t[1] for t in tensors) # 简化：末尾追加
allocation[name] = new_addr
# 实际实现中需要维护当前总大小
# 张量死亡后，归还地址
# 这里省略了死亡回调逻辑
return allocation注意：贪心算法不是最优的。我遇到过一种情况：两个大张量生存期完全错开，但中间夹了一个小张量，导致大张量无法复用。这时候最佳适应算法反而更好。
4.4 一张图看懂静态显存复用
下面我用一张 SVG 图，把整个流程串起来：
从这张图可以看得很清楚：
计算图分析 → 生存期推断 → 内存规划，三步走完，就能得到一份高效的显存复用方案。
4.5 实战中的几个坑
最后，分享几个我在项目中踩过的坑：
- 动态形状：如果你的模型输入尺寸会变（比如 NLP 里的变长序列），静态规划就不好使了。我一般会结合静态+动态混合策略，先预留最大尺寸，再动态调整。
- 梯度累积：训练时梯度张量的生存期比前向张量长得多。我曾经忘了考虑这一点，结果梯度被覆盖，模型怎么训都不收敛。
- 碎片化：频繁的分配和释放会导致显存碎片。我建议在规划时对齐到 256 字节，能有效减少碎片。
一句话总结：静态显存复用，就是用编译期的“慢”，换运行时的“省”。计算图越清晰，生存期推断越准，内存规划越优。
好了，这一章就到这里。下一章我们会聊动态显存复用，到时候再对比一下两者的优劣。