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× to 5× 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
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.
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.
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.
Prod 1 1 1 ∏Fi can be transformed into summation via ∏Fi=sgn(⋅)2∑log2∣Fi∣ .
+
∗
Table 1. Common reduction operations in machine learning and their corresponding binary operators
Figure 5. The normalized performance of fusing four selected modules on GPUs. The normalized performance of fusing four selected modules on GPUs.
Figure 6. Performance comparison results. (a) fusion at different levels. (b) incremental and non-incremental computation. Performance comparison results at different levels.
Figure 7. Illustration of the impact of fusion at different levels on memory access patterns. (a) No fusion: dK needs to be loaded L0 times to compute Fi(⋅) . (b) Fusion at level k : reduces accesses to the final result dK to Lk 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
Figure 8. The normalized performance of non-ML workloads across multiple platforms. The normalized performance of non-ML workloads across multiple platforms.
Figure 9. The normalized performance of ML workloads across multiple platforms. The normalized performance of ML workloads across multiple platforms.
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× 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.
Varun Kumar Dasoju, Tian Zhao
Department of Computer Science University of Wisconsin–Milwaukee
We present Forge-UGC (FX Optimization and Register-Graph Engine for Universal Graph Compilation), a four-phase compiler for transformer deployment on heterogeneous accelerator hardware, validated on Intel AI Boost NPU. Existing frameworks such as OpenVINO and ONNX Runtime often use opaque compilation pipelines, limited pass-level visibility, and weak buffer management, which can lead to higher compilation cost and runtime overhead. Forge-UGC addresses this with a hardware-agnostic design that separates graph capture, optimization, intermediate representation lowering, and backend scheduling. Phase 1 captures graphs with torch.export at the ATen operator level, supporting modern transformer components such as rotary position embeddings, grouped-query attention, and SwiGLU without manual decomposition. Phase 2 applies six optimization passes: dead code elimination, common subexpression elimination, constant folding, attention fusion, operator fusion, and layout optimization, reducing graph node count by 14.2 to 21.9%. Phase 3 lowers the optimized graph into a typed intermediate representation with explicit virtual register assignments. Phase 4 performs liveness analysis, linear-scan buffer allocation, reducing peak buffer count by 30 to 48%, and device-affinity scheduling, reducing NPU-CPU transitions by 42 to 65%. Across six model families ranging from 125M to 8B parameters, evaluated on WikiText-103 and GLUE, Forge-UGC delivers 6.9 to 9.2x faster compilation than OpenVINO and ONNX Runtime, 18.2 to 35.7% lower inference latency, and 30.2 to 40.9% lower energy per inference. Fidelity is preserved, with max absolute logit differences below 2.1e-5 and KL divergence below 8.4e-9. We also introduce Fusion Gain Ratio, Compilation Efficiency Index, and per-pass execution profiling for systematic evaluation of NPU compilation pipelines.
Optimized kernels such as FlashAttention and FlashDecoding are crucial for accelerating today's large models. Most of them are handwritten by experts because existing ML compilers cannot match their efficiency. Producing such kernels requires fusing computations with multiple reductions, which requires both algebraic transformation of the computation graph and operator scheduling of the transformed graph. Unfortunately, searching the two jointly yields a space too large to navigate. We propose Cleave, an ML compiler built on symbolic decoupling: Cleave discovers transformations by performing superoptimization on a graph with symbolic shapes, and then schedules each resulting graph on concrete shapes. Representing shapes as symbols makes equivalence checking cheap and lets a new Split operator, with a symbolic split count, parallelize along a reduction dimension. Cleave's scheduler fuses graphs with multiple reductions through iterative tiling and horizontal fusion. Evaluation on common LLM subgraphs shows that Cleave generates kernels up to 2.8x faster than the best baseline (1.6x on average) and reduces compilation time by 5.9x on average compared to Mirage. For dynamic workloads captured from production serving traces, Cleave compiles each operator once and achieves geometric mean speedups of 1.4x and 1.7x over FlashInfer's handwritten FA2 and FA3 backends. Cleave's code is available at: https://github.com/nyu-systems/cleave