Operator fusion, as a key performance optimization technique in the deployment of AI models, significantly improves execution efficiency and has been widely adopted in modern AI compilers. However, for cascaded reduction operations involving multiple loops with inter-loop data dependencies, such as the safe softmax followed by GEMM within attention mechanisms, existing compilers lack effective automated fusion and kernel generation capabilities. Although some works have addressed specific instances through hand-crafted fusion strategies, their solutions are limited in generality and difficult to extend to other similar structures. Given the prevalence of such computational patterns in deep learning models, there remains significant untapped potential in achieving general and automated fusion optimization. In this paper, we present a formal theoretical methodology for analyzing cascaded reductions which can fuse them into a single loop and introduce an incremental computation form. Based on this methodology, we design Reduction Fuser (RedFuser), a framework that automatically identifies supported cascaded reduction patterns and generates optimized fused kernels. Experiments show that RedFuser successfully fuses diverse workloads, achieving up to 2$\times$ to 5$\times$ speedup over state-of-the-art AI compilers and matching the performance of highly optimized hand-written kernels. The code is available at https://github.com/alibaba/redfuser
翻译:算子融合作为AI模型部署中的关键性能优化技术,能显著提升执行效率,并已被现代AI编译器广泛采用。然而,针对涉及多循环且循环间存在数据依赖的级联规约操作(例如注意力机制中安全Softmax后接GEMM的计算模式),现有编译器缺乏有效的自动化融合与内核生成能力。尽管部分工作通过人工设计的融合策略解决了特定实例,但其解决方案泛化性有限,难以推广至其他类似结构。鉴于此类计算模式在深度学习模型中的普遍性,实现通用且自动化的融合优化仍存在巨大潜力。本文提出一种用于分析级联规约的形式化理论方法,该方法能将级联规约融合为单个循环,并引入增量计算形式。基于此方法论,我们设计了RedFuser框架,该框架可自动识别支持的级联规约模式,并生成优化后的融合内核。实验表明,RedFuser成功融合了多种负载,相较于最先进的AI编译器实现高达2倍至5倍的加速,性能可媲美高度优化的人工编写内核。代码已开源至https://github.com/alibaba/redfuser