最近在做AI推理系统优化的时候我一直在琢磨一个问题当一台机器上同时挂着GPU、NPU、FPGA甚至还有一堆专用加速卡时深度学习模型这堆算子到底应该怎么摆才能既跑得快又省电还不至于把显存和内存占爆单纯追求单个算子最快、或者用手工经验把几个算子塞到某个卡上在真实异构环境里往往撑不过一轮线上流量波动。这也是我做HeteroOpt这套框架的初衷把深度学习计算图视为一个有向图问题以全局视角做多目标调度一次搜索给出多组可落地的折中方案而不是只给一个“看起来最快”的答案。HeteroOpt本质上是一个面向异构新型硬件的深度学习图调度框架核心解决的是算子/子图级别的设备分配、算子融合、kernel实现选择、循环变换参数、以及内存规划等调度变量的联合优化问题。它面对的目标不只有延迟还包括峰值内存、功耗、吞吐等多种指标最终输出一组Pareto最优解供实际场景选择。适合AI Infra工程师、模型部署工程师、编译器优化方向研究者以及想往深度学习系统方向进阶的算法同学参考。这篇东西我会把设计思路、关键拆解、落地实现和踩坑实录都摊开讲不会只停留在概念层面。1. 为什么需要“全局”又“多目标”的调度框架1.1 异构硬件环境下的调度困境先说个很现实的场景。我手上的推理服务里一台服务器同时有NVIDIA GPU和一个国产NPU加速卡偶尔某些自定义算子还得回退到CPU。模型是那种典型的Transformer结构几十个算子算子之间是有依赖关系的DAG。如果只看单个算子矩阵乘在GPU上很快LN和Softmax这种内存密集型算子在NPU上反而更省能自定义的Gather算子在CPU上也有优势——每个算子单独挑最优硬件的话听起来很美好。但真正落地时问题就来了。你让矩阵乘跑GPULayerNorm跑NPUGather跑CPU每个算子都选“当地最快”结果跨设备数据搬运开销直接把收益吃掉有些融合机会也因为算子被拆到三个设备上而彻底消失。更头疼的是当所有算子都挑最省时间的设备时整张图跑下来延迟可能不如把所有算子都放到同一个设备上。异构调度的难点恰恰在于此单个算子的最优不等于全局最优局部贪心策略在设备间通信、显存占用、融合约束存在时几乎必然失效。这也是我在早期尝试里踩得最深的坑——我用一个还挺复杂的模型做了个贪心调度的版本实测延迟反而比“全都放GPU”慢了将近一倍问题就出在通信和融合交互上。1.2 多目标调度意味着什么调度如果只看延迟其实是相对简单的优化问题把图上每个算子都绑定到延迟最低的设备、选择最快的kernel参数就行。但生产环境从来不是单目标问题。线上推理服务要控制P99延迟、GPU和NPU的功耗上限、内存峰值不能爆离线训练batch任务则更在意吞吐和显存占用。不同用户、不同场景对目标的偏好是完全不同的。多目标调度就是在这样一个背景下产生的。它不是简单地把延迟、内存、功耗加权成一个标量因为一旦换成另一批流量特征权重就得重新调而这种调权过程往往没有理论指导。更合理的做法是用Pareto支配关系在调度空间里找一组非支配解集合Pareto front让用户在时延、功耗、内存之间按实际需求做权衡。比如你的服务是面向C端用户的在线推理那就选延迟最低的那个解如果是内部离线批量任务选吞吐优先的那个解如果机房租期和电费是主要成本功耗优先的解就更合适。这种思路和传统编译优化里“-O2固定优化级别”的模式很不一样它把“什么样的结果算好”这个问题从框架里“解耦”了出去。框架负责充分探索候选调度最终拍板权交还给工程师或上层调度器而不是用一个固定的贪心策略替所有人做决定。1.3 为什么“全局”比“局部贪心”重要前面提到过贪心的失效这里把底层逻辑说透。深度学习计算图是一个有向无环图每一个调度决策其实都会影响相邻算子的执行条件。你决定把Conv融合进一个kernel里那BatchNorm就不存在了你决定把Attention的QKV分别放到不同设备那个场景下就无法再使用FlashAttention这类高效的融合实现你决定把某个算子放在NPU上数据就必须经过PCIe或片间互联接口搬运这个搬运时间在几十微秒到几毫秒不等会直接影响流水线排布和带宽占用。这些决策之间的耦合关系决定了调度只能作为一个整体图优化问题来做。HeteroOpt把设备分配、融合决策、kernel选择、循环变换参数、内存规划统一抽象成一张“调度决策表”然后在这个巨大的组合空间里搜索让优化器同时看到所有决策的交互影响而不是像贪心那样一个个算子孤立地下判断。这个“全局性”看起来只是换了个搜索方法实际上是把问题从“局部贪心选择”重构成“组合优化”。每次调整一个算子的设备绑定都会连带影响后续所有算子的通信代价和可选融合方案只有全局目标函数才能捕获这种连锁反应。这也是HeteroOpt能给出优于手工调优结果的根本原因。2. HeteroOpt的整体设计与关键技术拆解2.1 计算图抽象与调度变量设计做调度框架第一步是把计算图变成优化器可以理解和操作的结构。HeteroOpt里我沿用了主流的DAG表示法节点是算子或子图边是张量依赖。和纯推理框架相比这里的图节点需要附带更多信息——算子类型、输入输出shape、数据布局、是否支持融合、是否可切分到子设备以及该算子在不同硬件上的候选实现列表。调度变量是整个搜索空间的“原子”单位。我在设计时把它们归纳为四类设备分配每个算子绑定的目标硬件GPU、NPU、CPU或专有加速器。融合决策哪些相邻算子合并为一个编译单元例如ConvBNReLU合并成一个kernel或者Transformer里QKV投影和Attention结构合并处理。kernel实现选择同一个算子在不同硬件上有多种底层实现。GPU上matmul可能选择cuBLAS、CUTLASS或FlashAttention专用kernelNPU上可能有自研的高性能算子库。不同实现对应不同性能和功耗特征。循环变换与内存参数分块大小、向量化宽度、并行线程数、双缓冲/提前分配策略等。把这四类调度变量组合起来就形成了每个算子的候选集合整张图的全部候选组合构成全局调度空间。这个空间通常是指数级膨胀的但它是所有优化搜索的前提。没有候选空间优化器就无从谈起“选择”。2.2 多目标目标函数的数学化表达有了调度空间下一步得定义目标函数。HeteroOpt的多目标优化可以形式化这样描述。给定计算图G(V,E)、设备集合D、算子调度候选集合S目标是寻找一个调度解s使得下面这几个目标函数向量在Pareto意义下最小化F(s) (F_latency(s), F_memory(s), F_power(s), F_throughput(s))满足约束条件C(s) C_max。其中F_latency(s)表示整图端到端执行延迟由每个算子在所选设备上的执行时间、设备间通信时间、kernel切换开销、以及融合后产生的加速效应共同估算。F_memory(s)表示执行过程中的峰值内存/显存占用需要考虑张量生命周期、设备内存大小限制、以及跨设备拷贝时额外内存开销。F_power(s)是设备功耗之和通常需要从硬件计数器或外部功耗仪获取没有功耗计的话可以基于利用率做近似建模。F_throughput(s)在流水线执行方式下通常由最慢阶段决定的吞吐瓶颈。约束条件一般包括单设备显存上限、最大可接受延迟、某些算子必须落在特定设备如只能在NPU上运行的量化算子、可用的通信带宽上限等。这些约束在搜索中起到剪枝作用可以显著缩小搜索范围。在实际实现中我们不要求每个目标都有绝对精确的模型因为预算有限。关键是目标函数需要能够区分不同候选调度的优劣趋势。这也引出了后面要讲的“两级评估”机制——快速代价模型和真实Profile验证。2.3 搜索算法的选型从启发式到进化计算全局搜索空间是离散组合空间每个算子的候选数量少则几个、多则几十个整图组合量动辄10的20次方以上。这种空间里梯度方法基本没法用最常见的有效方案就是进化算法或启发式局部搜索。HeteroOpt里我把NSGA-II非支配排序遗传算法作为主搜索器并且做了几处工程化扩展。NSGA-II的思想是维护一个种群每代通过选择、交叉、变异产生子代然后基于Pareto支配关系做非支配排序再按拥挤度保留多样性较好的解。它天然支持多目标不需要归一化权重非常适合调度这种离散、非线性、多峰的空间。但纯NSGA-II也有一些问题。当候选空间特别大时初始种群容易陷在局部区域而调度解的分布有较强的结构性——例如某些融合选择直接决定后续设备分配是否合理。为了处理这一点我在初始化种群时加入了“专家启发式解”纯GPU解、纯NPU解、按算子类型优先分配的设备绑定方案、默认融合方案等保证初始种群覆盖不同倾向的调度策略加快收敛速度。对于候选规模较小、或者在线调度延迟敏感的场景我会切换到基于模拟退火多目标切比雪夫聚合的变体每次只做局部扰动比如交换一个算子的设备绑定、调整一次分块大小然后根据加权切比雪夫距离决定是否接受。这种方式收敛快适合几十毫秒内必须给出调度结果的控制场景但牺牲了全局搜索能力需要根据实际需求取舍。关于贝叶斯优化我也试过。如果每次真实评估调度方案的成本很高例如必须执行kernel才能得到耗时数据贝叶斯优化确实能以较少次评估找到较好解。但难点在于高斯过程或随机森林代理模型对高维离散空间的拟合精度有限容易把搜索引导到不理想区域。所以我的使用策略是先用小规模随机采样代价模型快速筛选出TopK候选再用贝叶斯优化做精细搜索而不是一上来就全场BO。2.4 约束处理与Pareto前沿判定约束处理是调度搜索里极易被忽视但影响巨大的一环。如果只是简单把不满足约束的解作废可能会把大量“接近合法”但很有潜力的解白白扔掉。比如某个解显存超了8MB修改一下内存分配策略就合法了直接淘汰很可惜。HeteroOpt的做法是约束分级 惩罚函数。硬约束如某些算子的强制设备绑定、完全不可突破的显存上限直接用惩罚项使该解在任何目标下都不会支配合法解同时仍在种群中保留利用它的结构信息指导交叉变异。软约束如目标延迟预算、功耗上限有一定容忍度则转化为一个额外的“约束违规度”子目标参与非支配排序。这样设计之后解集里即包含完全合法的方案也包含少量违规度较小的“接近方案”后续再通过修复算子例如调整分块大小、改变内存重叠策略尝试把接近方案变为合法方案。Pareto前沿判定方面标准做法是非支配排序中比较不同目标向量的支配关系。A支配B的判定条件是A在所有目标上都不差于B且至少有一个目标严格优于B。每轮进化我都保留当前代的非支配解集合并通过自适应网格归档器维持解的多样性避免Pareto前沿拥挤在某个区域。最终搜索结束框架输出完整的前沿解列表前面可能有几十个解每个解都对应一组具体调度参数和预期性能指标。3. 实战落地面向真实硬件的HeteroOpt实现要点3.1 硬件能力建模先给目标机器打标这个框架不是纯软件模拟是要跑到真实异构硬件上出结果的所以第一步必须把目标机器的“底细”摸清楚。我会为每一种目标设备建立一张能力矩阵主要包括峰值算力、内存带宽、片上SRAM/缓存容量、支持的算子库版本、可用算子集合、以及典型延迟特征。设备模型可以直接从厂商文档拿理论值但真实跑起来差异很大。我的习惯是先用一段标准的micro-benchmark脚本跑一遍全算子覆盖对每个算子类型取不同shape和不同kernel实现在设备上执行并记录耗时、功耗、内存峰值。这块数据会成为后续代价模型和搜索评估的校准基准。实测中发现一个细节GPU和NPU在内存带宽特性上差异巨大。GPU对连续大块数据搬运优化很好而某些NPU在特定layout下容易出现碎片化访问导致执行时间成倍增加。这类特征很难从理论参数里预判必须靠基准数据说话。我会把这类异常行为记录成一张“设备行为画像”在搜索时一旦遇到匹配条件就自动调整算子候选的优先级。3.2 调度候选集合的生成规则调度候选集合不能无脑生成。如果每个算子都列二十个候选整图的组合空间会爆炸到无法搜索。HeteroOpt里我采用了两级候选生成策略。第一级由规则引擎生成“结构候选”根据算子类型和硬件能力自动推断可用的设备绑定和融合模式。例如ONNX里常见的ConvBNReLU三连算子在GPU和NPU上通常都支持融合成单算子如果模型已经是融合版本就只保留融合后的选项。第二级由算子库探测生成“参数候选”对特定算子在目标设备上枚举kernel选择、分块大小、向量化宽度等参数组合。一个关键的经验是候选生成必须结合显式约束做“预剪枝”特别是要考虑设备内存上限。你在搜索之后才发现某组组合导致显存爆掉代价太高。我会在候选生成阶段就估算每个算子组合的内存区间若某个设备上总内存需求超过上限直接不生成该绑定候选。这项预剪枝能减少至少一半的搜索空间。另一件值得注意的事是融合决策的粒度。不是所有融合都值得尝试过大的融合粒度会降低调度灵活性过小又无法充分发挥kernel融合的收益。我一般从模型原生结构出发把每个子图单元控制在“包含2-5层融合深度”的范围内。对于Attention结构QKV投影作为一个单元、Attention计算作为一个单元、MLP作为一个单元这样既能识别出FlashAttention这种跨算子融合机会也保留了把不同单元分配到不同设备上的灵活性。3.3 目标评估器的构建快速代理模型与精确Profile相结合目标评估器是整个搜索循环的引擎。我设计了“两级评估”机制来平衡效率和精度。第一级是快速代理模型用于在搜索早期海量过滤候选。这个模型我最初试过用图神经网络结构直接预测整图延迟效果一般原因是输入特征里包含了太多算子级细节GNN很难捕捉设备分配和融合导致的非线性变化。后来换成更务实的做法离线训练多个“算子-设备-形状”查询表在线按算子粒度查表累加同时对融合收益做修正。Fusion带来的收益通常不是线性可加的所以我会对常见融合模式ConvBNReLU、Attention整体等直接做一次性整kernel Profile把融合后的耗时缓存下来这样在线评估时直接使用融合单元的实际数据而不是算子累加。第二级是精确Profile只对搜索阶段选出的TopK候选调度执行真实运行把真正的时间和功耗数据反馈回搜索过程。这个反馈机制是整个框架收敛质量的核心只依靠代理模型搜索出来的最优解可能并不真实必须有一个真实执行的“裁判”。实际使用时这两级评估的比例大约是1000:1。快速代理模型负责把包含几十万候选的搜索空间粗筛到几百个精确Profile负责给最终解集做验证级评估。两级之间会有个校正环节我会定期用最近一次的精确Profile数据去微调代理模型的偏差系数避免不准确的预测带偏搜索方向。3.4 与TVM/MLIR等编译器栈的对接方式HeteroOpt不是从零写一个编译器它是建立在现有深度学习编译器栈之上的一个调度决策层。我的实现思路是把HeteroOpt作为一个Pass插入到TVM Relay 或 MLIR图优化流程里图输入经过设备分配和融合决策后转换为目标设备上的算子配置再由底层编译器完成kernel生成和代码编译。对接的时候有几个关键点。第一是插入时机必须在图级别IR完成Shape推导之后、降低到具体算子实现之前插入。太早插入设备相关特征不完整太晚插入就失去了全局图视角。第二是策略映射HeteroOpt产生的高层调度决策需要翻译成TVM的Schedule配置或MLIR的TargetAttribute。例如一个分块大小参数在TVM里对应Tile Size在MLIR里可能对应Linalg Op的TileTransform参数。第三是跨设备数据通路当两个相邻算子被分配给不同设备时HeteroOpt需要自动插入设备间拷贝节点并选择合适的通信方式比如纯PCIe搬运、通过主机smem做中转、或是使用设备直通P2P通道。对接完成后HeteroOpt的输出是一份可执行的最终Schedule IR可以是JSON格式、TVM的OperatorConfig列表也可以是MLIR的#transform策略序列。保留中间表示很重要的原因是可以做后续回归验证和审计调度决策不是黑盒输出出了问题能回溯到是哪个决策导致的性能劣化这个能力在生产环境里特别有价值。3.5 调度结果的验证与回归保护调度框架上线后最怕的是“在评测环境优化得挺好上了生产就性能回退”。我在这块下了比较重的功夫。验证环节分为三层单元验证、模型验证、服务验证。单元验证是针对每个调度方案里的核心融合结构和设备分配组合用真实Profile确认是否达到预期性能模型验证是把整个调度方案编译一遍跑完整的模型推理对比基准版本的延迟和内存消耗服务验证则是放到一个模拟线上流量的压测环境里验证多路并发、不同batch size下的表现。回归保护我做成了一套自动化测试集。每次框架代码更新、搜索算法参数调整、或者目标机器驱动/算子库升级都要跑一遍测试集检查结果是否存在明显劣化。性能容忍度需要设定合理范围异构硬件上的Profile结果天然有波动我一般把同配置重复运行5次取中位数作为基准回归阈值设为5%。如果超出阈值自动触发报告定位是搜索逻辑改变还是硬件波动导致的。4. 常见问题与排查技巧实录4.1 搜索空间过大搜索迟迟不收敛这个问题我接手第一个版本时几乎每天都会遇到。一次搜索跑几小时都出不了可用的解原因无非是候选集太大、评估次数太多、种群退化。解决思路有几个方向可以叠加使用。一是收紧预剪枝规则特别是内存上限剪枝往往能砍掉60%以上的无效候选。二是对模型做“结构等价值检测”如果图中有多个结构完全相同的子图块Transformer里很常见只需要为其中一个子图做完整的调度搜索其他子图复制配置即可。三是给进化算法增加早停条件如果连续20代Pareto前沿的平均目标值变化小于1%就停止迭代用当前topk解作为最终结果。注意早停阈值不要设得太激进。异构执行时间有噪声有时前进代会让前沿提升0.5%下一轮波动后又跌回来。我建议观察3代或5代平滑后的变化率而不是只看单代之间的差值。4.2 多目标冲突时如何取舍多目标问题的核心难点不是求Pareto前沿而是用户最终只要一个解。每次我把几十个Pareto解摆在业务方面前对方都会愣住到底该选哪个我自己总结了一套比较实用的选择流程。第一步是硬性约束过滤例如线上延迟要求小于50ms、显存低于16GB先把不满足约束的解过滤掉。第二步按场景选目标权重做TOPSIS排序线上推理场景延迟占0.6、功耗占0.3、内存峰值占0.1离线训练场景吞吐占0.5、显存占0.3、功耗占0.2。权重不需要特别精细TOPSIS对权重敏感度有限排序结果大体能给出可接受的推荐。第三步是抽样验证把待选的Top3解编译运行两三遍用真实数据做最终决策。有一点经验值得分享输出的解未必每一个都真的优于手工方案。有时Pareto前沿里某些解在目标上“不差”但实际工程实现复杂度很高比如换了一种特殊kernel依赖维护成本巨大。所以我在HeteroOpt的输出里增加了“工程复杂度”标注字段把算子库依赖、驱动版本要求、编译时间等因素计入右侧提供给下游做综合判断。工程师选型时不仅要看目标值也要看落地代价。4.3 换了硬件版本或驱动调度结果突然失准这是做调度框架最无奈的坑。同一张卡驱动从530升到535原本很快的某个kernel可能变慢了20%导致之前搜索出来的调度解不再是Pareto最优。异构环境下这类情况更频繁NPU的算子库迭代快每个版本性能变化都很大。我的应对是两条线并行。第一是给框架增加“硬件软件栈指纹”概念把GPU/CPU/加速卡型号、驱动版本、算子库版本、CUDA版本、固件版本组合成一个指纹搜索出来的调度方案连同指纹一起归档。环境变化后先对比指纹如果指纹变了自动触发增量重搜用旧解作为初始种群的一部分而不是从零开始。第二是搜索时加入“鲁棒性目标”即不只考虑当前环境下的延迟还要考虑参数在一定范围内的扰动下延迟的稳定性。比如对分块大小做微调采样如果某个调度的延迟方差特别大就适当降低它在Pareto排序里的优先级。4.4 过拟合Profile数据导致部署漂移这是个很容易被忽视但实际影响很大的问题。因为搜索阶段的目标评估大多基于离线Profile数据如果你做Profile用的输入shape是固定的一组而线上真实流量是动态shape那搜索出来的“最优调度”很可能在真实流量下并不是最优。我在一次线上优化中吃过亏一个视觉模型评估时固定用1024x1024输入shape搜出来的调度上线后发现线上大量请求是512x768这种不规则分辨率延迟直接恶化30%。后来我把目标评估器改成多shape采样方式——每次评估一个调度方案时不仅跑固定shape还按线上shape分布跑一个采样集取加权平均作为目标值。搜索阶段和验证阶段用同一套shape分布策略确保优化目标和真实流量对齐。注意加权平均的权重需要按线上真实流量统计来设定而不是拍脑袋。否则优化方向仍然和真实需求有偏差。4.5 跨设备通信开销被严重低估异构调度里设备间通信往往是最容易拖后腿的部分。模型切分到多设备后如果两个算子分属GPU和NPU数据就得跨过PCIe或其他互联链路搬运这个代价在早期评估器里我只给了带宽估算值结果导致搜索出来的很多解在真实环境下因为通信瓶颈变得毫无优势。解决方式是单独建一张“通信延迟查询表”按照数据量大小、设备组合、传输方向、是否支持P2P来查询实际通信开销。这张表不需要很精细按每档数据量做bucket统计即可。搜索阶段的目标函数必须显式加入通信项而且要把它当作“可优化”变量——可以通过调度决定通信和计算的重叠策略把某个算子的数据搬运提前或延后隐藏通信延迟。实测下来加入了通信模型之后搜索出来的异构调度解质量明显提升尤其是设备间划分比较均衡的场景延迟能再降低15%-25%。写在最后调度框架的价值不是“一步到位”从我个人的实践体验来说做HeteroOpt最有价值的收获不是最后的性能数字而是把“深度学习模型部署到异构硬件”这件事从手工试错变成了系统化优化。过去调一个异构模型的性能主要靠经验、trial-and-error、以及对着profiler的结果反复改配置效率和可复现性都很难保证。有了图级全局多目标调度框架之后每次部署新模型到异构环境我可以先把整套搜索流程跑一遍在Pareto前沿里挑一个适合业务场景的解心里会踏实得多。最后再分享一个实用技巧别一开始就追求覆盖全硬件、全算子的完整框架。先把最简闭环跑通——一个小模型、两个异构设备、延迟和内存两个目标从图解析、候选生成、搜索、编译到最终部署验证串起来再逐步增加功耗目标、更多设备和更复杂的融合策略。这个“小而全”的迭代路径比一开始就设计一个庞大系统要靠谱得多。调度搜索本身就是一个需要持续打磨的系统工程先把整个流程走顺后面再往里面填细节就能少走很多弯路。