计算图执行优化：缓冲区复用与关键路径调度

先说结论计算图这玩意儿把节点连起来只是第一步真正跑起来之后内存占用和执行效率完全是另一回事。这个Python程序是我在调一个推理管线的运行时问题时动手写的——图里有两百多个算子按默认的拓扑顺序直接执行峰值内存逼近2GB关键路径上一有重计算节点卡住后面一堆无关节点全在空等。折腾了大半个月最后发现瓶颈不在算子实现而在缓冲区怎么分配、节点什么时候调度。于是就有了这个专门做计算图缓冲区管理和调度优化的程序把原本盲目无序的执行过程变成一套有统一内存管家、有动态优先级的执行框架。1. 为什么要单独做一套缓冲区和调度层很多用计算图的场景执行方式比我一开始还粗暴解析完整张图按拓扑排序的节点顺序一个个跑下去每个节点的输出缓冲区现场分配用完等整个图结束再统一清理。能跑确实能跑。只要图规模小、数据量小这方案一点问题都没有。但一旦节点数上百、单节点输出buf从几MB涨到几十MB三个问题就会接踵而来。第一是内存无序爆炸。每个节点的输出缓冲从生成那一刻起要到全图结束时才随整体清理释放。而实际上大部分buffer的生命周期是短暂且错开的一批节点的buffer早就可以回收了只是因为没人做回收动作就一直占着内存。最后的结果是明明活着的buffer只有几百MB实测峰值却吃掉了几个GB。第二是分配开销被严重低估。Python里如果用numpy array现场分配大块内存一次分配本身不慢但次数一多就出问题。两百个节点的图跑一遍产生上千次大块buffer分配和释放malloc开销、页面触页开销叠加在一起执行时间被白白拖长10%到20%。这个在机器监控里很难看因为CPU看着不高但墙钟时间就是下不去。第三是调度没考虑依赖之外的客观耗时。拓扑排序保证的只是“每个节点执行时依赖都已就绪”它不保证“整体等待时间最短”。实际跑起来以后某个重计算节点还在执行后方一大批相互不依赖的轻量节点已经完全就绪却在排队等着。默认的静态顺序不会管这些只会按序往下走把本可以并行的节点一个一个处理完。这个Python程序处理的就是这三件事做一个BufferPool统一管buffer按引用计数延迟释放、按大小分桶复用做一个Scheduler动态决定当前时刻跑哪个节点关键路径上的节点能插队、短任务能快速清场并用调度窗口限制同时存活的buffer数量。核心目标只有三条降峰值内存、升buffer复用率、缩整体执行时间。这套东西对写手写计算图执行器、或者想优化推理管线内存延迟的人最有用如果用的是TensorFlow这类现成框架内部已经有类似机制但理解这套思路对读懂框架底层也很有帮助。1.1 先明确这个程序到底处理什么问题要讲清楚先给计算图下个定义节点代表一次运算边代表数据依赖。应用场景比很多人想的多——不光是深度学习模型还有ETL数据管道、流式计算、编译器的中间表示层。我的程序针对的是“一次性执行完一张静态图”的场景也就是图结构固定、每个节点只跑一遍。动态图、循环图不适用那几个场景需要完全不同的策略后面我会讲到为什么。插这一层之后程序要回答的问题就非常明确了当一个节点执行完毕哪些buffer可以立刻回收哪些节点已经具备启动条件在并行执行资源的限制下先跑哪一个能获得最佳的整体收益调度器每做一次决策既要问内存又要问时间合在一起就是这套运行时策略引擎的核心。这也带出一个原则调度器不能只看依赖是否满足还要看buffer是否允许复用。比如同一块存储空间只要前一个使用者已经不再需要新的节点就可以把它当作自己的输出缓冲。这是缓冲区管理和调度编排能协同工作的根本前提。把这两件事分开做效果的提升会非常有限。1.2 为什么选Python而不是C这是项目一开始被问得最多的一个问题。如果目标是极限性能那一定是C或者Rust更合适。但这个程序的定位是“运行时策略引擎”不是被反复调用百万次的算子内核。真正的算子计算可以是numpy向量化、C扩展、外部子进程Python层只负责决策和调度。在这个前提下Python的开发效率就是最大优势。实际使用中Python做这种工作的短板也确实存在主要是GIL和单线程吞吐。我的解决方案是把调度器里所有纯计算部分——比如关键路径深度、就绪队列优先级——全部收敛到几个自包含函数里用numba编译或直接C扩展加速执行部分用线程池重算子本身不在Python循环里空转实测8个线程下加速比能到4到5倍这个后面实测章节会讲。我的建议是如果你的算子也是numpy或C扩展这种能释放GIL的可以放心用Python做调度层如果算子本身就是慢吞吞的Python循环那并行优化基本没救不如先优化算子实现。这一点在设计架构的时候就要想清楚不然吃力不讨好。2. 架构拆解三个模块怎么协同程序整体分成三个模块Graph负责图的解析和静态信息计算BufferPool负责所有缓冲区的分配、复用、回收Scheduler负责生成动态执行顺序。三者各干各的活边界非常清晰调试问题的时候能快速定位故障层。2.1 从节点定义到依赖关系节点用dataclass来表示这是最常见的做法。每个节点需要记录四类信息名字、算子类型、输入依赖列表、输出buffer大小。算子的预估耗时一开始没有我会让它带上一个可选的est_duration属性从上一轮运行的滑动平均里取值。dataclass class Node: name: str op_type: str inputs: list[str] output_size: int est_duration: float 0.0图本身就是一个nodes列表加一个邻接表。依赖关系的解析只需要一遍拓扑排序用Kahn算法逐层剥离入度为0的节点。但这里有个关键点拓扑排序的结果只用于确认图结构合法和计算关键路径真正的执行顺序不归它管。如果谁把拓扑序直接当成调度策略那这篇文章前面说的调度优化就白搭了。2.2 缓冲区池的模块边界BufferPool的定位是整个图生命周期内唯一的内存管家。它对外暴露的接口就三个acquire申请一块输出缓冲、acquire_for_consumer给下游节点登记一次读引用、release将一个buffer还给空闲池。模块边界上有一条硬规矩池不负责buffer内容的初始化。只有上层明确要求清零时才做fill_zero池本身只判断大小和引用状态不碰数据。这样省掉了大量无意义的memset也让池的逻辑足够单纯。池最外层统一销毁内部不逐块去GC避免了反复分配释放的开销。另外池内部要维护的东西比表面多不少。一是空闲分桶二是忙碌表三是每个buffer的引用计数。池的峰值上限也要在这里控制否则就会出现“为了省内存反而撑爆内存”的怪圈。这个我会在第三节详细展开。2.3 调度器的模块边界Scheduler要回答的是当前时刻哪些节点可以开始执行哪个节点优先级最高。它不关心算子的具体实现也不负责执行算子只负责把“就绪且最高优先级”的节点交给Executor。每个调度决策点出现在两类时机一是有节点执行完毕二是有新的buffer被释放回池。第一种时机很直观一个节点跑完以后它的下游节点入度可能会清零。第二种时机很容易被忽略某个大buffer被释放后一个等待了半天的节点可能因为“目标输出buffer终于有空间”而变得可执行。这类节点本质上是受了内存限制不是受依赖限制调度器必须把这种约束也纳入判断。这也就是为什么我有意把Scheduler和Executor解耦。执行器只负责跑节点调度器只负责选节点。否则项目一复杂两边逻辑缠在一起出了问题根本不知道应该查调度决策还是查执行实现。3. 缓冲区池实现分配、复用、回收的完整闭环缓冲区管理是这套程序里收益最直接的部分。很多图的节点输出buffer大小是有规律分布的只要把生命周期错开的buffer复用起来峰值内存能省一半以上。3.1 从“分配即新建”到“复用即分配”最朴素的写法是这种buf np.zeros((1024, 1024), dtypenp.float32) # 用完就丢等GC问题在于对一个有几百个节点的图来说这句话会被执行几百次每次分配的都是新的内存用完之后引用消失又被垃圾回收。下次再有节点需要同样大小的buffer就再来一轮分配、回收。分配的开销在Python里比想象中大因为numpy还要做数组对象初始化和类型检查。buffer pool的版本是反过来设计class BufferPool: def __init__(self): self._free_buckets defaultdict(list) self._busy {} self._refcount defaultdict(int) self._bucket_limits (4096, 65536) def _bucket_for(self, size: int) - int: for limit in self._bucket_limits: if size limit: return limit return size def acquire(self, owner: str, size: int) - np.ndarray: bucket self._bucket_for(size) if self._free_buckets[bucket]: buf self._free_buckets[bucket].pop() else: buf np.empty(size, dtypenp.float32) self._busy[owner] buf self._refcount[owner] 0 return buf def release(self, owner: str): if owner not in self._busy: return buf self._busy.pop(owner) bucket self._bucket_for(buf.nbytes) self._free_buckets[bucket].append(buf) self._refcount.pop(owner, None)acquire的时候优先从空闲桶里拿现成buffer没有才新建。release的时候不销毁buffer直接扔回空闲桶。这样同一块内存可以被图中不同节点在生命周期错开的情况下反复使用。这里有个细节要特别说明为什么用np.empty而不是np.zeros。因为很多算子根本不在乎输出缓冲的旧内容直接覆盖写用了np.zeros等于白白多一次memset。只有那些确实需要初值零的算子上层才自己调用buf.fill(0)。这种按需清零的策略在一张图动辄几百个buffer时省下来的时间非常可观。3.2 引用计数与提前回收如果图是一根直线节点执行完立刻释放buffer没问题。但实际图里到处都是分支结构一个输出buffer可能要喂给多个下游节点。最安全的回收时机是所有依赖它的节点都已经读完。我给每个buffer单独维护一个consumer_count初始值等于它的下游节点数。每个下游节点在开始执行前通过acquire_for_consumer登记一次读引用读完一次就减一。只有引用计数归零buffer才真正进入空闲池。def attach_consumer(self, producer: str): self._refcount[producer] 1 def mark_consumed(self, producer: str): self._refcount[producer] - 1 if self._refcount[producer] 0: self.release(producer)这个“延迟释放”我是踩过坑之后才重视的。早期版本图省事节点跑完就直接release结果就是同一个buffer在某个下游节点还没读完的时候被重新分配给了别的节点产生脏数据。当时排查了很久最后靠打印每个buffer的分配时间戳才发现问题。从那时候起我坚持所有回收操作都必须走引用计数宁可多写几行代码也不为图省事埋炸弹。3.3 大小分桶与峰值内存控制如果只是简单地把空闲buffer丢进一个list运行一段时间就会发现一个问题空闲list里积压了大量size各异的buffer而真正有用的却拿不到。比如一个1MB的buffer释放后接下来所有节点只要不是恰好1MB就不会复用它。最后池子里养着一堆“看起来能用、实际用不上”的空间峰值内存还是高。所以我学内存分配器的思路搞了大小分桶。把buffer按size分成几档小于4KB的进小桶4KB到64KB的中桶更大的进大桶。申请时向上取整到自己所在的桶释放时按实际大小入桶。这样小块内存不会被大buffer长期占用大buffer也不会被零碎的小请求反复折腾。峰值内存控制则靠一条硬规则整个池子累计分配出去的buffer字节数不能超过max_pool_bytes。超过之后release的buffer不入空闲池直接丢弃交还给系统。这样内存峰值被硬生生顶在一个上限不会出现“省内存省到反而撑爆内存”的荒诞局面。max_pool_bytes这里怎么取我的经验是先按所有节点buffer的一起生命周期重叠部分估算通常取总和的一半到一个三分之一再跑两轮实测微调。4. 调度器实现从拓扑排序到关键路径优先调度是整套程序里最考验设计能力的部分。它不只是把合法顺序排出来而是在“合法”之外再优化“高效”。4.1 拓扑排序只是起点不是终点Kahn算法做拓扑排序本质是“哪个节点入度先清零就先输出哪个”。这个顺序合法但是很傻它不知道节点之间的耗时差距也不知道哪条链对整体延迟影响最大。我做一个最简单的假设场景来说明问题一条关键链上有一个耗时80ms的重计算节点A旁边有一百个0.1ms的轻节点B都处于就绪状态。拓扑排序如果先把A排到后面那所有依赖A的节点都要多等80ms如果先把A排前面轻节点在后面跑整体等待时间就会短很多。这还只是单核情况。如果做成并行调度错乱导致的等待浪费更明显。调度器要做的是对每个节点计算两个值一是关键深度也就是从它的下游链路中最长的那条链到终点的耗时估算二是最晚开始时间基于全图的总关键路径长度反推。关键路径上的节点也就是那些“一旦延迟就会让整图延迟”的节点必须获得最高优先级。4.2 动态就绪队列关键路径优先短作业前置整个调度循环用事件驱动方式实现核心简化版本是这样ready_queue PriorityQueue() def on_ready(node): # 第一优先级是否关键路径节点第二优先级预估耗时短的先跑 ready_queue.put((0 if node.on_critical else 1, node.est_duration, node.name)) while not all_done: if executor.has_idle_slot() and not ready_queue.empty(): _, _, name ready_queue.get() node nodes[name] # 先挂上游buffer的引用再执行 for inp in node.inputs: pool.attach_consumer(inp) executor.submit(node, pool) else: # 等待执行完成事件或者buffer释放事件 event wait_event() if event.type NODE_COMPLETED: for dep in out_edges[event.node_name]: dec_in_degree(dep) if in_degree[dep] 0: on_ready(dep) elif event.type BUFFER_RELEASED: for waiter in waiting_for_memory: if pool.has_space(waiter): on_ready(waiter)为什么优先级是“关键路径节点优先 短作业优先”关键路径优先很好理解它直接压缩了整图的理论最短时间。短作业优先则是为了让就绪队列快速腾出位置让更多节点进入可并行状态避免一个0.1ms的小任务卡住后面一堆节点。这两者组合在一起能明显提升调度窗口内的任务饱满度。4.3 调度窗口卡住内存和线程池的边界调度窗口这个参数的用处很多人一开始想不到。它不是时间窗口而是“同时执行的最大节点数”我通常设为CPU核数的1到1.5倍。窗口有两个作用第一限制同时存在的buffer数量配合BufferPool的max_pool_bytes一起卡峰值内存第二避免一次性往线程池提交上百个节点让任务全在排队等调度造成无意义的时间分片。窗口大小怎么设我会先跑一遍profiling把每个节点的预估耗时和并行阶段并发数算出来。如果某个阶段平均并发只有3窗口设8就是浪费如果平均并发已经到7窗口设4就会人为串行化得不偿失。我推荐的流程是先设一个偏小的窗口比如4跑一遍然后逐步往上加观察墙钟时间和峰值内存的变化。一般来说在这个数据到达拐点之前停止增加就够了。我自己的压测图里8是最合适的值再往上内存飙升时间却几乎不变。还有一个坑要提醒调度窗口如果做得太激进比如窗口内都是重节点可能出现“窗口被占满但在等一个还没就绪的节点”的假死状态。我的处理方式是窗口的空位不固定分配给某个节点而是每次做决策时重新评估如果当前就绪队列里最好的节点预估耗时太长而其他可等待任务能在窗口内穿插完成就允许插入少量轻节点。本质上是贪心策略但加了关键路径这个强制优先级来防止本末倒置。5. 实测数据与排查实录这些坑必须提前知道工具做得好不好最终要看实测。我给自己造了一个模拟推理图尺寸不小240个节点3条并行主链每条链40个节点中间穿插一些共享节点和分支结构节点buffer从几KB到几十MB不等。所有buffer需求总和有5.2GB但因为生命周期重叠没那么高同时存活的理论峰值约1.9GB。测试机是8核CPU16GB内存。5.1 三个策略的量化对比我把三种情况放在同一台机器上跑执行方式峰值内存总耗时拓扑序直接执行每节点即时分配1.87GB328ms只加缓冲池不改调度1.02GB301ms缓冲池 关键路径优先调度0.84GB214ms这个结果有几个点值得细看。缓冲池的收益非常扎实因为生命周期不重叠的buffer大量复用之后峰值直接从1.87GB砍到1.02GB这个效果是“把大块内存反复用”带来的和算子本身快慢无关。调度优化的收益更隐蔽它把总耗时压低了接近35%但这个收益不是凭空来的它是把关键路径上的重计算节点从“被轻节点插队”的困境中解救出来让整图的理论最短时间真正兑现了。我也试过不设关键路径优先级、只开缓冲池加并行窗口的组合总耗时只到272ms说明优先级策略对整个结果贡献很大。只做内存优化不做调度优化的结果则说明内存搞好了如果执行顺序还是一团糟时间上的浪费一样很可观。5.2 常见问题排查表实操中一定会有问题我把踩过的和帮别人看过的坑整理成了一张速查表现象根因排查/解法某个节点的输出被后续节点改写数据变得莫名其妙缓冲区提前释放引用计数漏加检查每个input是否在上游启动前attach_consumer打印buffer分配时间戳定位内存不降反升池子越养越大空闲分桶队列里积压了长期不用的buffer给空闲buffer加时间戳超过阈值定期销毁多线程调度时偶发卡死就绪队列空转条件变量等待丢失用事件驱动循环重写wait加上超时重查性能比拓扑序还慢关键路径计算有误或窗口设太大导致资源被低价值任务抢占打印关键路径节点列表核对窗口减半重测图有环时程序死循环没做环检测构造阶段用Tarjan算法判环遇到环直接给出节点列表报错大buffer一直没有被复用分桶粒度不合理比如把1MB和100MB放同一档细化桶的档位或者按2的幂次分桶5.3 几个值得注意的工程细节多线程执行节点时Python的GIL影响比想象中大。如果节点算子是纯Python循环那开八个线程也拿不到实际并行度反而线程切换拖慢单节点执行。我用的是numpy向量化和C扩展算子在执行期间能释放GIL所以线程池才有效果。这是设计这套程序时就埋下的前提如果你的算子做不到这一点整个并行策略要推翻重来。关于节点预估耗时的获取不要在真实环境里临时profiling开销太大。我为每个算子类型维护了一个滑动平均时长oatype - ema_duration每执行完一个节点就用真实耗时更新一次。调度器的优先级依赖这个数据所以估得越准调度越聪明。第一轮跑的时候数据还是空的我会用默认值第二轮开始基本就收敛了。还有一个很容易被忽略的点输出缓冲区的大小必须按实际值上报。有的算子输出大小依赖输入的形状比如padding、concat、slice预估值很可能不对。我在节点执行完的回执里显式带上实际buffer大小BufferPool按真实尺寸登记这样后面释放和复用才不至于错位。最后说一个我个人体会比较深的地方缓冲区管理和调度优化这类工作很容易被当成性能优化的“边角料”但实际它在工程里的收益往往比打磨单个算子更大。这个程序后续想扩展的话方向也挺明确——把分桶策略改成TLSF算法、把调度器接入异步IO事件循环、或者把池内buffer换成分布式共享内存都是可以深挖的点。但不管怎么改核心思路都是那句内存要复用调度要动态关键路径要优先生跑。希望这篇记录能帮正在做计算图执行器的人少走几个来回。