自动融合复杂循环依赖的计算操作,提升AI加速器性能
RedFuser: An Automatic Operator Fusion Framework for Cascaded Reductions on AI Accelerators
- 提出理论分析方法,将多层归约操作合并为单循环
- 实测速度提升2到5倍,媲美手工优化内核
- 适合需要高效部署注意力机制的AI工程师
算子融合是AI模型部署中的关键性能优化技术,显著提升执行效率,已被现代AI编译器广泛采用。然而,对于涉及多层循环且存在循环间数据依赖的级联归约操作(如注意力机制中的安全softmax后接GEMM),现有编译器缺乏有效的自动化融合与内核生成能力。尽管部分工作通过手工设计策略解决特定场景,但其通用性差且难以扩展至其他类似结构。鉴于此类计算模式在深度学习模型中普遍存在,仍存在巨大的自动化融合优化潜力。本文提出一种形式化理论方法,用于分析级联归约,可将其融合为单个循环并引入增量计算形式。基于该方法,我们设计了RedFuser框架,能自动识别支持的级联归约模式并生成优化后的融合内核。实验表明,RedFuser成功融合多种工作负载,在速度上相较当前最优编译器提升2×至5×,并达到高度手工优化内核的性能水平。代码已公开于https://github.com/alibaba/redfuser。
原文摘要 · Abstract (English)
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
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。