NPU DeepSeek-V4 AutoFuse算子自动融合优化【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer随着AI模型结构日益复杂化特别是MoE(混合专家)及多模态架构中动态、细粒度的小算子组合已成为主流设计模式。DeepSeek本次发布的最新模型中诸如hcPre、hcPost均是灵活小算子结构未来在多模态场景类似结构也会成为常态。这类算子结构灵活多变组合方式多样若依赖传统手写融合算子方式进行优化其开发效率已难以跟上模型快速迭代的需求。因此实现对算子的自动化融合能力具有显著的技术必要性与紧迫性。PyTorch框架在图模式下的Inductor组件为此提供了基础的算子融合与编译能力其设计兼容现有PyTorch生态具备良好的软件继承性。然而由于NPU与GPU在内存架构、计算单元及并行机制等方面存在本质差异直接沿用面向GPU的优化策略难以充分发挥NPU的硬件潜力。为在继承Inductor整体流程与接口的基础上实现针对NPU的高效融合需具备面向NPU的即时编译JIT算子生成能力。为此CANN框架继承的AutoFuse组件针对NPU架构特征进行了深度优化实现了从计算图到高效NPU算子代码的自动化映射与生成。该组件不仅保持与Inductor技术的协同与兼容更进一步通过架构感知的融合策略、内存访问优化与指令调度等手段显著提升了算子融合在NPU上的执行效率从而支撑复杂模型在NPU平台上的高性能部署。在本篇技术报告中我们将对NPU的vector算子自动融合技术以及与inductor的对接工作进行阐述同时简单介绍模型收益与使能方式。Highlights使用简单PyTorch前端增加一行代码即可使能性能提升依赖NPU亲和的算子生成技术提升模型开箱性能亲和NPU模板规约的schedule技术基于硬件建模动态求解的算子tiling技术基于AscendLoopIR表达的AscendC算子kernel代码生成技术泛化与完备度当前已支持152个LoopIR中的46个未来将逐步补齐与LoopIR完整对等OutlineNPU DeepSeek-V4 AutoFuse算子自动融合优化HighlightsOutlineinductor与AutoFuse对接AutoFuse JIT的算子生成能力亲和NPU模板规约的schedule技术基于硬件建模动态求解的算子tiling技术基于AscendLoopIR表达的AscendC算子kernel代码生成技术AscendLoopIR简介inductor与AutoFuse对接我们扩展了inductor在NPU上的Codegen后端通过将InductorLoopIR转换为AscendLoopIR调用AutoFuse模块进行Schedule模板选择、UB复用与Tiling调优生成高效的AscendC Kernel实现。Inductor对接Autofuse的参考实现inductor_npu_ext脚本使能方式在脚本开头导入 inductor_npu_ext 以注册基于 Autofuse 的 inductor 编译后端扩展。import inductor_npu_ext通过在函数上添加 torch.compile 装饰以对函数进行编译建议基于以下规则选择 inductor 编译范围编译范围内图结构稳定不包含执行时分支选择、.item()内存同步等行为避免过多的断图Graph Break。避免编译范围过大过大的编译范围由于 dynamo guard 叠加更容易触发重新编译re-compile。选择包含较多 Pintwise/Reduce 计算的范围进行编译inductor 对该类计算的融合效果较好能以较小的编译代价获取性能收益。典型片段收益在 DeepSeek-V4 模型中hc_post 片段包含多个小算子组合适合通过 AutoFuse 进行融合优化。通过 torch.compile 装饰对 hc_post 函数进行编译可以显著提升其执行性能。torch.compile def hc_post(x: torch.Tensor, residual: torch.Tensor, post: torch.Tensor, comb:torch.Tensor) - torch.Tensor: y post.unsqueeze(-1) * x.unsqueeze(-2) torch.sum(comb.unsqueeze(-1) * residual.unsqueeze(-2), dim2) return y.type_as(x)性能数据在昇腾 A3 基于 hc_post 执行时的典型输入 shape 进行测试。模型部分未编译耗时编译后耗时(dynamicFalse)性能提升编译后耗时(dynamicTrue)性能提升hc_post30 us7.5 us4x10 us3xAutoFuse JIT的算子生成能力算子的核心交付件主要包含kernel源码与tiling源码AutoFuse模块将算子代码的生成过程抽象为几个阶段AscendLoopIR是抽象算子计算过程的表达IR描述计算过程、切分轴、UB内存等信息schedule组件负责对Hint graph进行处理。根据实际硬件的能力提供能够应对所有shape场景的全量的轴切分策略。每一个具体的切分策略由一张称为impl graph的AscendLoopIR承载组合成一组impl graphs作为后续ATT与codegen组件的输入ATT组件负责根据每个impl graph生成tiling代码tiling代码在执行时会根据具体shape评估每个impl graph对应kernel模板的性能选择最优模板计算切分参数通过tiling_data传递给kernelcodegen组件负责根据每个impl graph生成kernel代码每个impl graph对应一份模板实现在执行时根据tiling_data参数决定实际要执行的模板以及循环搬入的具体参数AutoFuse模块的源码已开源可在此获取亲和NPU模板规约的schedule技术在NPU上硬件单元一次运行可以完成批数据的向量运算为了充分发挥硬件特点需要准确选择进行向量运算的轴。同时面对硬件的多核能力以及内存有限的状态还需要进一步选择进行分核切分的轴、进行UB切分的轴。这三种轴的正确选择对功能与执行性能都有巨大影响选择错误就容易出现UB超载的功能问题或性能严重劣化的问题。不同的输入shape规格通常亲和不同的切分策略特定切分策略的算子实现我们通常称为模板。提供通用的模板可以简化问题复杂度但无法提供更好的性能效果提供精细、繁多的模板库可以面对不同场景提供极致性能但会让模板选择、代码维护变得困难。AutoFuse作为JIT的算子融合技术在模板设计上选择了折中的方案提供多模板但对适用场景做抽象规约控制模板数量。通过有限的模板数在绝大多数shape规格下都能提供优秀的性能。Schedule模块负责根据HintGraph产生一组ImplGraph每个ImplGraph对应一个模板实现。以reduce多模板举例Reduce融合是深度学习中最常见的场景之一涉及Sum、Max、Min、Mean、Prod、Any、All等归约操作。其核心挑战在于归约轴与非归约轴的差异化处理多核负载均衡归约后的继续计算优化为应对不同shape下的切分策略归纳如下三类主要模板应对不同的适用场景。模板类型触发条件调度特点适用场景全载模板Reduce后节点≤4且为AR/ARA模式不需要切分Reduce轴小数据量、简单图结构通用模板默认选择UB切分Block切分中等规模、通用场景R轴分核模板数据规模大且可切分Reduce轴作为多核切分轴Reduce轴较大在大R轴场景若只设计通用模板会导致UB超载的功能问题在小数据场景选择通用模板则又会因为核启动开销大而影响性能AutoSchedule实现介绍AutoSchedule的核心思想是将融合图中的各种计算节点抽象为几大类ComputeType根据每种类型在NPU上的特定切分方式通过AxisGroup对轴进行语义分组并利用Merge规则生成适用于所有节点的公共AxisGroup最终生成多种切分候选方案TilingCase。ComputeType抽象与NPU切分方式将融合图中的各种计算节点抽象为几大类ComputeType每种类型在NPU上有其特定的最优切分方式ComputeType典型算子NPU切分特点轴分组特征ElementwiseAdd, Mul, Relu等可任意维度切分支持向量化全部归入Y组Broadcast广播算子需处理不同shape的对齐计算轴归入Y组ReduceSum, Max, Min等归约轴需特殊处理归约轴R组非归约轴Y组Concat拼接拼接维度不能切分非拼接轴归入Y组AxisGroup轴分组与Merge规则针对不同ComputeType的切分需求定义了AxisGroup结构对轴进行语义分组X组需要完整数据的轴如Concat拼接轴切分后需保持数据完整性Y组元算子计算轴如Elementwise可灵活进行UB切分R组Reduce归约轴采用专用的多核归约策略由于融合图中可能包含多种ComputeType通过Merge规则生成一个适用于所有节点的公共AxisGroup规则类型合并条件结果说明YY两节点均为Y组保持Y组计算特性相同可直接合并YRY组遇到R组升级为R组Reduce优先级更高YYRY组遇到YR组扩展为YR组兼容归约和非归约计算TilingCase候选方案生成基于公共AxisGroup生成多种切分候选方案TilingCase。每个TilingCase定义了一组具体的切分参数UB切分指定哪个轴进行UB缓冲区切分Block切分指定哪个轴用于多核并行Reduce特殊处理是否启用R轴切多核策略Schedule调度变换Schedule对每个TilingCase执行具体的调度变换将抽象的切分方案转换为具体的调度操作变换类型说明效果TileSplitUB级切分将数据分块加载到UB缓冲区BlockSplitBlock级切分实现多核并行ReduceBlockTilingReduce专用切分优化归约操作的多核执行向量化识别可向量化轴提升计算吞吐基于硬件建模动态求解的算子tiling技术schedule技术在编译时给出了应对不同shape的规约模板集合在执行时就需要有配套技术能够正确选择模板输出核切分UB切分结果。计算结果错误会直接导致UB越界错误或性能急剧劣化。在AutoFuse组件此工作由ATT(Auto tiling)技术来完成。通过预置的性能建模以及多样的求解策略来保证泛化场景的功能正确性与普适性能同时寻求求解准度与求解速度之间的平衡。自动Tiling的挑战性能评估的困难schedule输出的不同模板在GM非连续性、UB非连续性、是否向量化宽度对齐、循环次数、计算瓶颈和搬运瓶颈等方面存在较大差异因此在不同的融合场景和Shape下会表现出不同的性能优劣。这就要求自动Tiling能够准确地评估不同模板的执行性能。分核与分块的权衡增加分核数量会带来额外的核启动开销包括硬件开销和软件开销而减少分核则会导致单核负荷过大。因此如何合理权衡分核数和分Tile块大小成为我们面临的关键难题。自动Tiling核心功能自动Tiling的目标是找到能够实现Kernel最佳执行性能的Tiling策略其核心功能包括寻找最优切分结果为每个ImplGraph找到其最优切分结果确定Kernel执行时需要使用的核数量、每次从GM搬运到UB的数据量、UB内的循环次数、每次UB内计算的数据量以及每次从UB搬运到GM的次数选出最优模板基于每个ImplGraph的最优切分方案尽可能准确地评估其Kernel执行性能从而选出最优模板。自动Tiling核心流程Kernel的执行逻辑会在IR中表达。自动Tiling会根据IR图的表达提取关键信息包括约束的符号化表达:LocalBuffer占用约束自动Tiling求解时需确保每一级LocalBuffer的占用都在硬件允许的范围内。例如Kernel申请的TQue/TBuf及临时Buf的大小之和不能超过硬件的UB大小限制。IR会表达出每个Tensor的location是在GM还是UB上以及Tensor间的复用关系自动Tiling根据这些信息将各级LocalBuffer的约束进行符号化表达。核数占用约束自动Tiling会根据轴的切分方式生成实际核数的符号化表达在Tiling运行时根据实际的硬件约束和确定的符号值判断是否满足约束。切分约束例如保证子轴大小小于等于父轴的约束、切分轴需要32B对齐等。各流水耗时符号化表达自动Tiling会对各个API进行性能建模通过符号化的形式表达这些API在各个流水线上的性能。根据IR表达的循环轴来确定API的调用次数从而推导出该IR的所有API在各个流水线上的总耗时。此处假设瓶颈流水线的执行可以较好地掩盖非瓶颈流水线的执行瓶颈流水线简化理论因此任务执行的总耗时主要体现在瓶颈流水线的时间上。如下图所示图中包含三个pipe流水其中瓶颈流水为AIV_MTE2AIV_MTE2在执行时会掩盖AIV_MTE3和AIV_VEC的执行任务执行的总耗时体现在AIV_MTE2的执行耗时上。pipe流水示意图自动Tiling根据提取的关键信息生成Tiling求解器代码默认生成轴排序Tiling求解器。该求解器会以算子实现过程中的存储占用不超过NPU各级物理存储大小为约束条件根据不同的IR表达选择最合适的Tiling算法保证基础解的性能再以最小化瓶颈执行单元的耗时为优化目标通过性能建模进行求解提升解的性能上限。此外系统还提供可选配置的启发式迭代求解器。该求解器采用启发式求解方法从初始解开始在满足约束条件的可行域内沿着性能公式建模的梯度下降方向搜索满足内存要求的解空间直到公式建模的值达到最小值时退出寻优流程并返回最优Tiling。关键技术性能公式建模定位性能公式的仿真精度决定了对Schedule生成的模板选择的准确性轴排序核数调整及核内循环选择的准确性启发式迭代求解的准确性。因此需要对各类API进行尽可能精确的性能建模。搬运类API建模实现识别影响搬运性能的主要因素。例如数据量、GM非连续性、UB非连续性、是否向量化宽度对齐、分核数、Pipe头开销、API头开销、API调用次数等。基于影响搬运性能的主要因素及其对性能的影响设计性能公式。例如假设MTE2的开销为Cost(MTE2)其中API头开销为h数据量为DataSizeAPI调用次数为CountPipe头开销为H带宽为TCost(MTE2) ((DataSize / T h) * Count) H其中带宽是与是否向量化宽度对齐、分核数、GM非连续性、UB非连续性相关的因变量。基于假设的性能模型通过实验验证不同主因素变化下的实际搬运性能。建立搬运性能与主因素之间的关系并通过符号化技术表达出来。计算类API建模实现针对AscendC的稳定基础API自动Tiling模块会根据不同的输入采集性能数据得到性能与输入的关系并建立性能模型通过符号化技术表达出来。针对调用AscendC的易变API自动Tiling模块会基于API调用的逻辑计算其调用入参和调用次数生成该API的完整性能模型从而得到相对准确的性能建模。基于搬运类API建模和计算类API建模以及瓶颈流水线简化理论可以将自动Tiling的两个核心问题简化为如何调整Tiling使得瓶颈流水公式值最小选择瓶颈流水公式值最小的模板作为最优模板。轴排序求解器定位 Tiling求解器的目标是基于输入Shape确定合适的分核及分块大小以获得尽可能好的Kernel性能。轴排序求解器支持的启发式优先级排序规则优先是指优先不切的轴及基础算法如UB占用核数占用权衡算法、UB优先贪心算法、对称切分可以很好地保证基础解的性能另一方面结合性能建模的表达可以提升Kernel性能的上限。实现确定切分轴的优先级首先需要基于API来确定切分轴的优先级顺序如下父轴优先级高于子轴功能性硬约束规约化类轴高于非规约化类轴启发式规则广播轴高于非广播轴启发式规则非最内轴高于最内轴启发式规则搬运API的尾轴具有同等优先级启发式规则。分核与核内切分其次分成两个部分进行切分包括核内Tiling和分核Tiling。以UB优先贪心算法为例适合API头开销较大的场景核内Tiling按照轴排序的逆序依次遍历优先将变量调整至最大值判断是否符合硬件约束条件。若不满足则通过二分法调整该变量直到符合硬件约束条件为止随后调整下一个核内Tiling变量直至所有变量均满足硬件约束条件。例如s1tt2、s1tt、s1t、s2t是Tiling相关轴其切分流程如下分核Tiling识别与多核相关的变量按从大到小的顺序遍历这些变量找到核数占用更大的记录若超出物理核数则返回。如下图所示bngs1T是多核切分轴其切分流程如下。选择策略为当核数占用不同时优先选择占用核数更大的记录。根据上述策略最终选定的占用核数为47。具体流程为优先遍历s2t调至最大值1024符合硬件约束条件然后遍历s1t调至最大值256符合硬件约束条件然后依次调整下个变量s1tts1tt2直到所有变量均满足硬件约束条件。其他算法UB占用核数占用权衡算法适合需要精细控制核数占用及UB占用的场景有预设值可根据性能公式动态调整核数。对称切分适合尾轴转置场景需要保证作为搬运类API的多个尾轴可以按照同等优先级切分。基于AscendLoopIR表达的AscendC算子kernel代码生成能力Codegen的功能是解析Schedule生成的impl_graph生成融合算子的host_code和kernel_code其中host_code包括tiling_func、infer_shape以及get_kernel代码kernel_code包括kernel和tiling_data代码最后完成代码的编译。Codegen主要流程Codegen的核心功能是生成在device侧执行的kernel代码Codegen的输入是schedule生成的impl_graph该图上承载了shape信息、节点信息、轴信息等生成代码的基本要素。impl_graph示例如下图上的具体信息如下所示Sizes: z0z1t_size: VAR z0z1Tb_size: VAR Axis: z0(0) : 200, ORIGINAL, align: -1, allow_oversize_axis: 0, allow_unaligned_tail: 1 //轴信息 z1(1) : 200, ORIGINAL, align: -1, allow_oversize_axis: 0, allow_unaligned_tail: 1 z0z1(2) : 40000, ORIGINAL, align: -1, allow_oversize_axis: 0, allow_unaligned_tail: 1 //z0z1合轴后的信息 z0z1T(3) : Ceiling((40000 / (z0z1t_size))), TILE_OUT, from: {z0z1, }, align: 1, allow_oversize_axis: 0, allow_unaligned_tail: 1 //z0z1T对应Tile切分的外轴, 对应所有核UB内for循环的总数量 z0z1t(4) : z0z1t_size, TILE_IN, from: {z0z1, }, align: 1, allow_oversize_axis: 0, allow_unaligned_tail: 1 //z0z1T对应Tile切分的内轴也就是每次AscendC api处理的数据量 z0z1TB(5) : Ceiling((Ceiling((40000 / (z0z1t_size))) / (z0z1Tb_size))), BLOCK_OUT, from: {z0z1T, }, align: 1, allow_oversize_axis: 0, allow_unaligned_tail: 1 //z0z1TB对应Block切分的外轴对应需要的核数 z0z1Tb(6) : z0z1Tb_size, BLOCK_IN, from: {z0z1T, }, align: 1, allow_oversize_axis: 0, allow_unaligned_tail: 1 //z0z1T对应Tile切分的内轴对应单个核UB内for循环的数量 Nodes: ...... abs_test/load_0: Load (1) ...... abs_test/abs_0: Abs (2) .axis {z0z1TB, z0z1Tb, z0z1t, } //Abs节点所有轴信息 .loop_axis z0z1Tb //Abs节点循环轴的分界该轴之前的轴都是循环轴之后的轴是向量化轴对应api一次处理的数据量 .api: .compute_type elewise //计算类型 .type Compute .unit Vector //计算单元 .x abs_test/gather_0.y .y.dtype float32 //数据类型 .y.axis {z0z1TB, z0z1Tb, z0z1t, } //输出tensor的轴信息 .y.repeats {(40000 / (z0z1Tb_size * z0z1t_size)), z0z1Tb_size, z0z1t_size, } //轴大小 .y.strides {(z0z1Tb_size * z0z1t_size), z0z1t_size, 1, } //轴stride .y.vectorized_axis {z0z1t, } //向量化轴 .y.vectorized_strides {1, } //向量化轴的stride .y.mem: .tensor_id 3 //输出tensor的id .alloc_type Queue //分配类型Queue表示在UB的TQueue中分配内存 .hardware UB //硬件类型UB表示在UB中分配内存 .position TPosition::VECOUT //内存位置VECOUT表示输出tensor的内存位置 .y.que: .id 1 //表示tensor对应的queue_id .depth 2 //queue的深度 .buf_num 2 //queue中buffer的数量 .reuse_id 1 //表示queue是否可以复用1表示可以复用 abs_test/store: Store (3) ......根据impl_graph解析如下所示是Codegen生成DataCopy api的示例代码可以看出核心流程是解析图上的切分策略组装api的入参生成在device侧执行的代码。ss DataCopyPadExtend( ub , gm [ gm_offset tpipe.tiler.Size(api_attr.offset) ], dma_param.block_count , dma_param.block_len , dma_param.src_stride , dma_param.dst_stride ); std::endl;与传统算子代码类似针对不同的切分策略schedule会生成多个模板对应不同的impl_graphCodegen需要依次解析生成模板函数在核函数入口的位置根据tiling_key调用不同的模板函数。示例如下extern C __global__ __aicore__ void abs_test(GM_ADDR abs_test_Data_0, GM_ADDR abs_test_Output_0, GM_ADDR workspace, AutofuseTilingData param) { const AutofuseTilingData t; if (TILING_KEY_IS(0)) { abs_test_0_general_0_nil_2_nil(abs_test_Data_0, abs_test_Output_0, workspace, t); } else if (TILING_KEY_IS(1)) { abs_test_0_general_0_nil_2_nil_unaligned(abs_test_Data_0, abs_test_Output_0, workspace, t); } }Codegen生成的kernel代码是以AscendC api为基础的。不同的模板函数生成的for循环以及每次处理的数据量有所不同示例如下tiling_key0对z1轴进行切分每次在UB中完成计算的向量化轴为z1t。for (int z0z1Tb 0; z0z1Tb z0z1Tb_loop_size; z0z1Tb) { Abs(y_local[0], xlocal[0], z1t_actual_size); }tiling_key1对z0轴进行切分每次在UB中完成计算的向量化轴为z0t以及z1。for (int z0Tb 0; z0Tb z0Tb_loop_size; z0Tb) { Abs(y_local[0], xlocal[0], z0t_actual_size * z1_axis_size); }总体而言Codegen是依据impl_graph生成kernel代码的但是要持续提升kernel代码性能仍需要在Codegen框架优化和api内部优化两方面不断挖掘。AscendLoopIR简介AutoFuse组件中的AscendLoopIR是一套基于Loop定义的IR其类型(type)映射到的明确的vector计算如AddMul等。多个AscendLoopIR组成的DAG图称做AscGraph表达算子多层循环中的多步计算。每一步计算对应途中的一个节点节点则由AscendLoopIR实例化而来节点之间的有向边表示数据的传递关系每个节点还包含属性用于指定该计算所处的循环层级。在AscGraph中节点通常嵌套在多层循环中被循环多次执行。例如假设Foo节点的调度信息与0号输出信息如下Foo.sched.axis [z0, z1, z2, z3]; Foo.output[0].axis [z0, z1, z2, z3]; Foo.output[0].repeats [s0, s1, s2, s3]; Foo.output[0].stride [s1 * s2 * s3, s2 * s3, s3, 1];这意味着Foo节点有四个相关轴z0, z1, z2, z3每个轴对应的大小为s0, s1, s2, s3根据stride可以判断是连续输出。Foo节点将在s0*s1*s2*s3次迭代中被执行等价如如下C代码for (uint64_t i0 0; i0 s0; i0) { for (uint64_t i1 0; i1 s1; i1) { for (uint64_t i2 0; i2 s2; i2) { for (uint64_t i3 0; i3 s3; i3) { // 执行Foo计算 } } } }上述代码每个轴都展开成循环了因此Foo在最内层做的就是一个scalar数的计算如果把最内轴s3向量化Foo上提一个循环那么Foo每次进行的就是一次向量计算了。根据sched与output轴的灵活组合就能够表达出广播 、规约、非连续搬运等各种复杂的计算形式【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考