Tensor 生命周期分析:复用内存之前,先证明不会重叠

📅 2026/7/6 5:07:49 👁️ 阅读次数 📝 编程学习
Tensor 生命周期分析:复用内存之前,先证明不会重叠

Tensor 生命周期分析:复用内存之前,先证明不会重叠

一、内存复用不是把 buffer 反复借出去

AI 推理引擎为了降低峰值内存,会复用中间 tensor 的 buffer。理论上,只要两个 tensor 生命周期不重叠,就可以共享内存。问题是图优化、分支、动态形状和异步执行都会让生命周期分析变复杂。复用错一次,输出就会被悄悄写坏。

Memory Planner 的核心不是聪明地省内存,而是证明安全。每个 tensor 从最后一次写入到最后一次读取之间,都不能被其他写操作覆盖。这个边界必须来自图依赖分析,而不是靠经验。

二、先构建读写区间,再做复用决策

可以把每个 tensor 的生命周期抽象成[first_write, last_read]。区间不重叠才允许复用。

flowchart TD A[计算图拓扑排序] --> B[记录 Tensor 写入点] B --> C[记录最后读取点] C --> D[生成生命周期区间] D --> E{区间是否重叠} E -->|否| F[允许复用 buffer] E -->|是| G[分配独立 buffer]

动态执行时,拓扑顺序可能不是唯一。Planner 需要和执行器的调度策略保持一致。否则分析基于一种顺序,运行时按另一种顺序执行,就会出错。

三、用 Rust 类型表达规划结果

规划结果应是不可变的执行计划。运行时只按计划索引 buffer,不再临时猜测。

#[derive(Clone, Copy, Debug)] pub struct LifeRange { pub first_write: usize, pub last_read: usize, pub bytes: usize, } pub fn can_share(a: LifeRange, b: LifeRange) -> bool { a.last_read < b.first_write || b.last_read < a.first_write }

真实 planner 还要考虑对齐、设备内存类型和 inplace 算子。这个函数只是最小边界:生命周期重叠就绝不能共享。

区间分析的一个工程难点是跨分支处理。当计算图有条件分支(如 if/else 或动态 mask),tensor 生命周期不再是简单一维区间。分支 A 中 tensor_x 在 step5 释放,分支 B 延续到 step12——life range 应取所有执行路径并集(first_write 到最大 last_read),这是保守且安全的。更精细的方案是为每条分支独立规划,在合并点插入 copy 或 remap,代价是增加运行时复杂度。在 Rust 中可用 BitSet 表示每个 tensor 在各 step 的活跃状态:bit i=1 表示在 step i 活跃,两个 tensor 可共享 buffer 当且仅当其 BitSet 按位 AND 为零。这种方案复杂度 O(n²×s),适合离线编译场景。折中做法:编译期用 BitSet 精确分析,运行期用区间做快速回退。

四、异步执行会让生命周期更难

如果算子提交到 GPU/NPU 后异步返回,CPU 侧认为某个 tensor 已读完,设备侧可能还在用。此时必须插入同步点或事件依赖。否则 buffer 被复用后,硬件还在读旧数据。

inplace 算子也要特别标记。某些算子允许输入输出同 buffer,某些不允许。不能只看 shape 相同就复用。算子 schema 应明确 alias 规则。

最后,Memory Planner 要有 debug 模式。可以给 buffer 填充哨兵值,或在复用前后校验 checksum。性能模式可以关闭,但开发阶段必须能抓出错误复用。

动态 shape 需要分桶规划。每个请求都重新规划会增加延迟,所有 shape 共用一份计划又不安全。常见做法是按 batch、sequence length 和 dtype 建立 plan cache。cache key 与编译缓存类似,也要包含影响内存布局的字段。

还要把峰值内存写入执行计划。运行前先判断设备剩余内存是否足够,不够就拒绝或降级。不要等分配失败时才报错。推理服务的内存错误如果发生在执行中,恢复成本会更高。

五、总结

Tensor 内存复用的前提是生命周期分析成立。Planner 要从图依赖中计算读写区间,只有不重叠的 tensor 才能共享 buffer。异步执行、inplace 算子和动态形状都会增加风险。省内存很重要,但错误复用会让结果无声损坏,这比 OOM 更难排查。