RedFuser: An Automatic Operator Fusion Framework for Cascaded Reductions on AI Accelerators

Xinsheng Tang; Yangcheng Li; Nan Wang; Zhiyi Shu; Xingyu Ling; Junna Xing; Peng Zhou; Qiang Liu

cs.ARFeb 24, 2026

RedFuser: An Automatic Operator Fusion Framework for Cascaded Reductions on AI Accelerators

Authors: Xinsheng Tang, Yangcheng Li, Nan Wang, Zhiyi Shu, Xingyu Ling, Junna Xing, Peng Zhou, Qiang Liu

Organizations: Alibaba Cloud Computing, Shanghai, China · Alibaba Cloud Computing, Shenzhen, China

Abstract

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

Figures & tables

Panel of Figure 1 — Figure 1. The reduction tree structure in cascaded reductions The reduction tree structure in cascaded reductions

Figure 2. The formal definition of cascaded reductions The formal definition of cascaded reductions.

Panel of Figure 3 — Figure 3. Comparison of computation and memory access patterns in the cascaded reductions before and after fusion. (a) Unfused cascaded reductions: each reduction re-loads the input data and accesses the results from prior reductions. (b) Fused cascaded reductions: the input is loaded only once, and memory accesses to results of preceding reductions are eliminated. Cascaded reductions before and after fusion.

Panel of Figure 4 — Figure 4. Comparison between non-incremental and incremental computation. (a) Non-incremental computation: the $k$ -th level reduction must wait until all inputs are available before execution and on-chip memory consumption grows with the input length. (b) Incremental computation: the result at level $k$ can be updated immediately upon arrival of new input, maintaining constant on-chip memory footprint. Illustration of memory access pattern in non-incremental and incremental modes.

Reduction Operation $\sideset{}{{}_{i}}{\mathop{\scalebox{1.5}{\raisebox{0.0pt}{$ \mathcal{R} $}}}\limits}$	$\oplus_{i}$	$\otimes_{i}$
Max, ArgMax, TopK, etc.	$\max$	$+$
Min, ArgMin, etc.	$\min$	$+$
Sum, Inner Product, Matrix Multiply, etc.	$+$	$\ast$
Prod 1 1 1 $\prod\mathrm{F}_{i}$ can be transformed into summation via $\prod\mathrm{F}_{i}=\mathrm{sgn}(\cdot)2^{\sum\log_{2}\|\mathrm{F}_{i}\|}$ .	$+$	$\ast$

Table 1. Common reduction operations in machine learning and their corresponding binary operators

Panel of Figure 5 — Figure 5. The normalized performance of fusing four selected modules on GPUs. The normalized performance of fusing four selected modules on GPUs.

Panel of Figure 6 — Figure 6. Performance comparison results. (a) fusion at different levels. (b) incremental and non-incremental computation. Performance comparison results at different levels.

Panel of Figure 7 — Figure 7. Illustration of the impact of fusion at different levels on memory access patterns. (a) No fusion: $d^{K}$ needs to be loaded $L_{0}$ times to compute $\mathrm{F}_{i}(\cdot)$ . (b) Fusion at level $k$ : reduces accesses to the final result $d^{K}$ to $L_{k}$ times. Illustration of fusion at different levels.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Table 2. The configurations of four workloads

Table 3. The configurations of two non-ML workloads

Panel of Figure 8 — Figure 8. The normalized performance of non-ML workloads across multiple platforms. The normalized performance of non-ML workloads across multiple platforms.

Panel of Figure 9 — Figure 9. The normalized performance of ML workloads across multiple platforms. The normalized performance of ML workloads across multiple platforms.

Table 13

Figure 12. FlashAttention IR FlashAttention IR

Figure 13. FlashDecoding IR FlashDecoding IR

GPU Kernel Optimization

GPU kernel optimization improves the speed and memory efficiency of computations by tailoring code to GPU hardware and workloads.

See all 209 papers

AttnFuse: A Composable DSL for Compiling Attentions to Fused GPU Kernels
Sep 11, 2026Varun Kumar Dasoju, Tian ZhaoEfficient Attention GPU Kernel Optimization
Department of Computer Science University of Wisconsin–Milwaukee
Forge-UGC: FX optimization and register-graph engine for universal graph compiler
Apr 14, 2026Satyam Kumar, Saurabh JhaCompiler Optimization AI Accelerator Inference
Cleave: Scaling Tensor Program Optimization via Decoupled Algebraic Search and Operator Scheduling
Oct 6, 2026David Pissarra, Jinkun Lin, Haitian Jiang +2Compiler Optimization GPU Kernel Optimization
New York University New York, USA

Sep 11, 2026J/K move · Enter open · S save

AttnFuse: A Composable DSL for Compiling Attentions to Fused GPU Kernels

Varun Kumar Dasoju, Tian Zhao

Department of Computer Science University of Wisconsin–Milwaukee

Efficient Attention GPU Kernel Optimization

Open paper AskPDF arXiv

Modern AI systems are built on the Transformer architecture, whose core operation, attention, accounts for the majority of computation and memory cost. Researchers continually propose new attention variants to improve quality, efficiency, or context length, but each variant currently requires expert-written GPU code to run at usable speeds. PyTorch's recent flex_attention lets researchers describe custom attention patterns in Python and compile them to fused kernels, but its design is limited to modifications applied after the central matrix multiplication, excluding Rotary Position Embedding (RoPE), the positional encoding used by every major LLM. We introduce AttnFuse, a small DSL for attention that makes pre-multiplication transformations like RoPE first-class operations. Researchers compose ten high-level building blocks to describe a variant, and AttnFuse's compiler emits a single fused GPU kernel for the entire computation. On an RTX 3090, AttnFuse achieves a 2.10

\times

speedup over flex_attention on the RoPE+causal pattern. On an H100, it runs a full Llama-3-8B training step within 5% of PyTorch's hand-tuned backend. Our investigation reveals the Rotation Calculus: whether to fuse RoPE or apply it separately depends on the GPU's compute-to-bandwidth ratio, with a derived crossover that matches measurement. AttnFuse demonstrates that a small, attention-specific compiler can close the gap between flexible research code and production kernels.

Abstract

GPU Kernel Optimization

Explore similar work

AttnFuse: A Composable DSL for Compiling Attentions to Fused GPU Kernels

Forge-UGC: FX optimization and register-graph engine for universal graph compiler

Cleave: Scaling Tensor Program Optimization via Decoupled Algebraic Search and Operator Scheduling