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

发布时间:2026/7/6 4:59:10
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 或动态 masktensor 生命周期不再是简单一维区间。分支 A 中 tensor_x 在 step5 释放分支 B 延续到 step12——life range 应取所有执行路径并集first_write 到最大 last_read这是保守且安全的。更精细的方案是为每条分支独立规划在合并点插入 copy 或 remap代价是增加运行时复杂度。在 Rust 中可用 BitSet 表示每个 tensor 在各 step 的活跃状态bit i1 表示在 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 更难排查。