先讲一个让我印象深刻的崩溃现场。当时我在给一个自研推理引擎做性能压测模型是一个BERT类的中等规模网络输入batch调到32之后突然来了一行熟悉的报错CUDA out of memory。我看了一眼显存占用曲线模型权重只占了一半剩下的一半居然全是被中间张量撑爆的。那一刻我意识到计算图跑起来之后真正决定显存生死的往往不是参数而是那些用完即弃的中间缓冲区内存分配问题。说白了就是算子执行过程中产生的临时缓存怎么分配、怎么复用、怎么调度的问题。后来我花了几个礼拜在引擎里完整做了一套缓存分配和调度优化系统目标只有一个在保证正确性的前提下把计算图执行时的峰值显存压到接近理论下界同时不把吞吐拖垮。这篇文章就是这次实战的记录。里面没有花哨的框架没有调包式的解决方案全是关于张量生命周期、内存复用、执行顺序和碎片控制这些硬骨头。如果你也在搞推理引擎、AI编译器、图计算系统或者只是被大模型显存卡得头疼这篇东西应该能给你一些走野路子的灵感。1. 先算一笔账计算图里中间张量为什么占掉半壁江山1.1 算子输出不断叠加的朴素分配有多蠢很多人第一次写执行器的时候采用的是最直白的办法每个算子执行完给输出张量分配一块新内存然后不管不问直到整个图跑完再统一释放。这个方案在正确性上挑不出毛病但内存利用率惨不忍睹。举个简单的例子一个四节点链式图Conv - ReLU - Add - Softmax。朴素分配下输入 A、B、权重、卷积输出、ReLU 输出、Add 输出、Softmax 输出全都同时存活在显存里。虽然卷积输出在 ReLU 执行完之后就已经没人引用了但它占的那几百 MB 依然被死死攥着不释放。一个本来只需要最大两三个张量共存就能跑完的图最后峰值却等于所有张量大小之和。在深度学习里有个更阴险的地方很多算子输出规模非常大。比如卷积的输出 feature map 动辄就是几十 MB 到几百 MBTransformer 里 attention 的 K、Q、V 以及中间 scores 也都是大块头。十几个算子串下来的输出叠加起来比模型权重还大是经常的事。所以朴素分配不是浪费一点而是直接决定你一个 batch 都跑不起来。1.2 生命周期 liveness找到那些可以互相复用的张量要解决这个问题得把分配重新理解成规划。每个中间张量从被产生到被最后一次消费中间这段时间叫它的生命周期live interval。一个有向无环计算图里同时活着的张量集合决定了显存需求的下限。如果两个张量的生命周期完全不重叠那它们就能安全地用同一块内存。这个概念和酒店客房分配一模一样。一间房不可能同时住两个像但只要退房和入住时间错开一张床就能被几十拨客人轮流睡。计算图里的缓冲区也是这样一张显存块可以依次装下卷积输出、ReLU 输出和 Add 输出只要它们不会在同一时刻活着。实践里的麻烦在于大多数执行器根本没有记录生命周期这个习惯。一个张量什么时候生算子执行完的那刻容易知道但什么时候死最后一次被谁读取却要在整张图上做一次反向传播式的 use-site 扫描。我第一次做静态分析时就是遍历每个节点的输入张量把它在每个算子的最后一次使用标记出来。这一步是整个优化系统真正的地基做不准确后面全是空谈。1.3 理论极值压力测试下能算出的最理想峰值有了每个张量的生命周期理论上就能算这张图最优能做到多省把时间轴虚拟地划成无数个瞬间对每个瞬间统计活着张量的大小之和最后取一个最大值。这个值就是这张图在给定执行顺序下内存需求的绝对下界。为什么只是给定执行顺序下因为张量的生死区间不是固定的它和调度顺序强相关。比如说有两个独立的支路先跑哪个后跑哪个都会改变中间张量共存的时间。如果你把内存需求看作一个区间图染色问题在固定区间下求最少数量的颜色实际上有一个贪心最优解但只要允许你自由调整算子执行顺序这个问题就滑向了一般图染色那一类 NP-hard 的范畴。换句话说想全局最优很难但想做到在某个合理调度下接近最优是完全可行的。我把这个过程想明白之后对图计算有了点新的敬意内存规划本质上就是一个图优化问题你在图上做节点分析、区间划分、顺序调整比纯粹调 kernel 参数有意思多了。这一套思路后来我拿到图计算引擎里处理活跃顶点集的缓冲区复用也同样成立。2. 分配策略的取舍贪心、扫描线规划和运行时池子2.1 方案一size-first 贪心复用为何经常失灵最直觉的方案是贪心把所有张量按大小排序大的优先申请然后维护一个空闲块列表哪个块能装进去就装进去。我一开始就是这么干的代码写起来也快大概两三百行就跑通了。问题出在最优性上。贪心算法只看当下不看未来。举个例子有一块 300 MB 的空闲内存现在有个 280 MB 的张量要用按 size-first 肯定直接给它。但如果稍后一个 290 MB 的张量马上也要用而整个过程没有其他能容纳它的空隙那 290 MB 就得另外申请新内存峰值被生生抬高了一截。贪心还有更隐蔽的坑它容易制造碎片化而且是语义层面的碎片化。大块内存被不大不小的张量一点点切碎后面任何大请求都找不到连续地址可用。这时候你可能会想那改用 best-fit 或者 first-fit 是不是就行了实际测下来自适应效果有改善但都只是局部技巧。真正的问题在于贪心根本没有利用生命周期信息它不知道哪个块会在什么时刻空出来只知道现在谁有空闲。2.2 方案二扫描线式区间规划在固定图上逼近最优我后面换成了另一个思路既然生命周期已经静态分析出来了不如把分配过程当作区间事件处理用扫描线sweep line一张张地把内存拼起来。做法是这样的把张量生命周期看作数轴上的区间区间左端点是生产者执行完成的时刻右端点是最后一个消费者读取它的时刻。然后沿着时间轴从左往右扫维护一个 active set。每遇到一个张量要出生就从当前空闲块集合里挑一块与所有 active 区间不冲突的地址每遇到一个张量死亡就立刻回收它的地址块并和相邻块做合并。这段伪代码能表达核心逻辑events [] for t in tensors: events.append((alloc, t.birth, t.size, t)) events.append((free, t.death, t.size, t)) sort events by time for ev in events: if ev.type alloc: block free_list.find_best_fit(ev.size) if block is None: block arena.alloc_new_region(align_up(ev.size)) assign(tensor, block) else: free_list.release(tensor.block) free_list.merge_adjacent()这个方案的好处是它天然会把不重叠的张量配对到同一个物理块。由于 DAG 的拓扑执行顺序是确定的区间也都是合法的给定顺序下它能做到几乎最优的复用。而且实现也不难关键数据结构就是一个按地址排序的空闲块列表加一个记录活动张量和块映射的表。我后来在固定 batch 的推理场景里用这个方案把 BERT 类模型的峰值从所有中间输出总和压到了最大活跃集总和附近大概省了 40% 到 60% 的内存。效果非常直观。2.3 方案三运行时缓存池接住动态不确定性扫描线规划适合静态 shape 的场景但现实业务往往是动态的输入分辨率会变、batch 会变、甚至图结构会随控制流变化。这时候做纯离线规划就不够用了需要运行时缓存池兜底。运行时缓存池的核心思想其实不新鲜向设备一次性申请一大块显存内部用 free list 管理分配请求来了先在池里找找不到才向设备申请。PyTorch 的 caching allocator 也是类似思路。关键区别在于我的池子会用第一步算出的生命周期信息做预约如果规划器知道某个时刻会有大块空间释放它会在分配时更愿意把当前请求塞进一个很快会空出来的区域而不是贸然扩大池子。说白了离线静态规划和在线缓存池不是二选一而是两层配合。静态规划负责给出一个理想内存布局运行时缓存池负责接住动态请求并且尽量向理想布局靠拢。这个双轨制是目前工业界很务实的解法。2.4 我的选型结论三种方案我最后都保留了但角色不同方案核心思路优势劣势适用场景size-first 贪心按大小优先复用实现简单无需生命周期分析容易碎片化峰值不稳定原型验证、临时工具扫描线区间规划按生命周期区间复用固定图下近乎最优峰值可控依赖静态 shape动态场景失效静态 batch 推理、CUDA Graph 捕获运行时缓存池大块 arena free list适应动态 shape实现灵活碎片需长期治理在线服务、动态请求实际生产里我的默认组合是扫描线规划给出布局 缓存池执行兜底。如果只是快速验证一个想法那随便写个贪心也能跑但别指望它能扛住压测。3. 工程落地Arena内存池、对齐与动态Shape的三重考验3.1 Arena大块缓存让设备端少吞几万次内存申请规划器再好最终还是要落到内存分配这个物理操作上。这里有个容易被忽视的性能陷阱直接给设备驱动发内存申请是非常昂贵的。一次 malloc/free 调用可能只要几微秒但在一个几百节点的计算图里每个算子都调几次累积起来就是毫秒级的开销对追求低延迟的服务来说完全不能忍。所以我在实现时坚持了一个原则向设备申请内存的次数必须少剩下的内部自己消化。方案就是 arena也叫 region-based allocation。启动阶段一次性申请一块很大的显存比如 4 GB然后内部通过偏移量切块。分配一个张量只是把 arena 的游标往前推一下或者从 free list 里取一个现成的块全程零系统调用。arena 的核心结构很简单就是块描述符链表struct MemoryBlock { size_t offset; // 相对 arena 起始地址的偏移 size_t size; bool in_use; MemoryBlock *prev, *next; };释放张量时把块标记为空闲然后立刻做一次合并邻居操作如果当前块的前一块或后一块也是空闲的就合并成一个更大的块。这能显著减缓外部碎片的产生。别小看这一步没有合并的 free list 用不了多久就成了蜂窝煤。3.2 对齐和碎片64字节能省出来的坑与填法设备内存几乎都有对齐要求比如 CUDA 通常要求 64 字节或 256 字节对齐。你分配一个刚好 100 字节的张量底层实际上会占 128 字节。所以分配器必须做对齐计算否则某些算子底下跑 kernel 时直接 misaligned address 崩给你看。对齐计算本身没什么难度对齐后的偏移等于原偏移向上取整到对齐值的倍数。容易踩坑的是对齐会制造内部碎片也就是张量实际没用到、但被白白占住的内存。单个张量浪费 24 字节不算什么但假如一张图里有几万次分配累积起来就相当可观了。另一个坑是 free list 的选择策略。地址有序的 free list 适合做 first-fit内存寻址行为好合并邻居也方便按大小有序的 free list 适合做 best-fit碎片率更低但每次分配要额外维护大小索引。我的建议是如果规划器已经给出了稳定的布局用地址有序 first-fit如果请求完全动态用 best-fit 更稳。碎片实在太严重时兜底方案是整块 arena 整体重规划一次把活跃张量搬到新布局里旧 arena 直接作废。这个操作成本高但比让服务慢性崩溃好得多。3.3 动态Shape下别把规划做死了动态 shape 是内存规划的天敌。同一个算子上次输出 256 MB这次输入分辨率一改可能直接输出 512 MB。如果规划器缓存的是固定布局遇到新 shape 就只能重新规划而重规划一次可能得花十几毫秒直接卡住服务。我的做法是引入shape bucket把请求按输入尺寸、batch 大小分桶同一个桶内的中间张量大小认为是稳定的规划结果缓存下来复用。比如输入长宽在 512 到 520 之间的统一按 520 规划多出来的内存最多浪费几个百分点但省掉了反复规划的抖动。还需要用心处理的另一个点是算子内部的临时内存。很多底层算子库矩阵乘法、卷积在执行时会额外申请 workspace这些空间不体现在你的图张量生命周期里但确实会占用显存。规划器如果不把 workspace 需求登记进来就很容易出现规划时看起来很充足、真跑起来又 OOM的诡异场景。我在实现里给每个算子加了一个workspace_size(input_shape)接口所有 workspace 统一走同一个规划通道不再绕过分配器偷偷申请。4. 调度与分配的协同顺序、并发和流水线怎么影响显存4.1 拓扑排序不是终点执行顺序直接影响峰值很多执行器拿到计算图之后按拓扑序排一下就开跑。这没有错但只能保证不会出现依赖错误。从内存角度看拓扑序往往不是最优的。想象一个简单的场景A 分支会产出一个超大的中间张量B 分支拥有一个很小的中间张量两者最终汇合到节点 C。假如先执行 A那么 A 的大输出从出生到 C 之前一直活着此时 B 再执行又叠加了 B 的输出峰值就是 A 输出加 B 输出。如果反过来先执行 B让 B 的小输出先占住内存然后执行 A 时峰值只是 A 输出加 B 输出但如果 A 的输出比 B 大很多正式峰值仍是 AB。真正更优的做法是看看哪个张量死得早把死得早的大张量先产生、先消费、先释放再让别的节点进来。结论是执行顺序必须纳入内存优化的目标。这和你在给集装箱装货时先放大件还是先填缝是一个道理。4.2 峰值敏感的调度启发式给早死早超生的算子插队在静态图中我常用一个启发式来辅助调度计算每个节点的紧张度也就是如果跳过它整张图还会有多少待处理的大张量等在后面。紧张度越高优先执行。简单说就是让那些产生大输出、但下游很快就能消费掉的算子先跑抓紧把显存让出来。这种调度策略在 CUDA Graph 捕获场景特别香。因为 CUDA Graph 本身就是静态图捕获一次之后反复重放内存规划和调度可以完全离线做。峰值敏感调度配合扫描线分配能把整张图的显存需求压得很低而且重放时没有任何分配开销。一个实用的实现思路是先按拓扑序做一个基础调度然后反复尝试交换相邻的、无依赖关系的节点每次交换后重新评估峰值如果峰值下降就保留交换。几千个节点的图跑几百轮迭代也就几十毫秒换来的是几十甚至上百 MB 的显存节省。4.3 多Stream与流水线并发和复用只能二选一吗到这里必须承认一件事纯顺序执行对内存最友好因为张量只有一个主线生命周期。一旦引入多 stream 并行不同 stream 的中间张量同时活着active set 变大内存复用空间变小峰值理所当然会涨。但推理服务要的是低延迟高吞吐完全放弃并发不现实。我的折中方案是流水线化把计算图切成几个 stagestage 之间用固定数量的缓冲区传递数据stage 内部顺序执行。这样并行度上去了但跨 stage 的中间数据量是可控的因为缓冲区是固定的不再随 batch 或算子数量无限膨胀。一个典型的流水线配置是双缓冲前一个 stage 在算第 t 帧时后一个 stage 在算第 t-1 帧。中间结果只需要两份 buffer内存开销固定但吞吐能基本翻倍。这个模式说到底就是用固定的小块内存换并行度正巧也是调度优化和缓存分配的交叉地带。5. 真实战绩与三个翻车教训5.1 实测数据一个BERT类推理模型压力测试我在自己引擎上跑了一个 BERT-base 类模型hidden size 768batch 32输入长度 128做了一套三档对比配置峰值显存单次推理耗时备注朴素分配不释放4.7 GB11.2 ms基线立即释放 贪心复用2.9 GB11.1 ms生命周期分析后简单释放扫描线规划 峰值敏感调度1.8 GB10.9 ms加上调度优化可以看到光是做生命周期分析和即时释放峰值就降了接近 40%再叠加扫描线区间复用和调度优化总共降了 60% 以上。耗时没有恶化反而略有下降——原因是内存紧凑之后缓存局部性变好设备端预取压力小了一些。这组数据验证了一个判断显存瓶颈很多时候不是真缺内存而是内存没被管理起来。5.2 坑一in-place和view偷偷改写了生命周期我踩的第一个大坑是 in-place 算子和视图操作。比如 ReLU 的 in-place 版本会直接改写输入张量的存储而不是产出一个新张量。从生命周期分析角度看这个张量的死期提前了但如果执行器不知道这是 in-place规划器仍然以为输出缓存要另分配一块不仅浪费内存极端情况下还会把一块还在被其他 view 引用的内存复用掉造成数据错乱。解决办法是给每个算子做严格的读写集分析。能提供read_set和write_set规划器才能正确判断哪些张量可以提前释放、哪些块绝对不能碰。这个分析最好在图编译阶段就做完别拖到运行时。5.3 坑二算子workspace没注册规划白做有阵子我发现自己规划的峰值已经很漂亮了但实际跑起来还是偶尔 OOM。翻了好久才定位到问题某些矩阵乘法和卷积算子底层实现需要一块很大的 workspace而这部分空间根本不在计算图的中间张量列表里。规划器满以为显存空着算子一执行却突然狮子大开口。这之后我定了一条铁规矩每个算子必须显式声明本轮执行所需的最大 workspace 大小所有 workspace 必须走分配器的登记通道规划器在计算 active set 时必须把 workspace 也当成一个虚拟张量加进去。加了这一步之后规划数字和真实峰值才终于对得上。5.4 坑三动态Shape长期服务的碎片老化与兜底最后一个坑是时间维度上的。一个在线推理服务跑上几天之后即使每天请求分布一模一样分配器内部也会因为各种瞬时峰值产生零碎的空闲块。这些碎片单个只有几 KB 到几 MB但积少成多某天一个大请求进来空闲总内存明明够却找不到连续地址直接分配失败。我给这个情况准备的兜底策略有三个层次。第一层是常规的相邻块合并每次释放都做成本很低。第二层是 shape bucketing把同类请求的规划结果缓存避免反复小幅度调整布局。第三层是碎片监控每当空闲块数量超过活跃块数量两倍或者合并后最大空闲块容量低于当前请求大小就触发一次整体重规划把 arena 压缩重建。这套机制上线之后服务已经连续跑了两周多没有复发。最后分享一个小技巧不管你的图是推理引擎里的还是图计算框架里的都可以在动手写分配器之前先只做一件事——把所有中间对象按生命周期画出来。不需要什么专业工具一个最简单的排序加区间扫描就够了。算清楚理论下界是多少再决定投入多少精力去做缓存分配和调度优化。很多时候你会发现问题根本不在算法多精巧而在于你从没认真统计过那些缓冲区到底什么时候生、什么时候死。这一步补上了显存优化就已经完成了大半。