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

RedFuser is an automatic framework that employs a formal theoretical methodology to identify and fuse cascaded reduction operations into optimized single-loop kernels, achieving significant speedups over state-of-the-art AI compilers while matching the performance of hand-written solutions.

The Gist

Imagine you are running a massive, high-speed kitchen (an AI accelerator) trying to prepare a complex dish for millions of customers (AI models). The recipe involves a series of steps where you have to taste, adjust, and combine ingredients repeatedly.

In the world of AI, these steps are called Cascaded Reductions. They are like a chain of tasks where you can't start the next one until the previous one is finished. For example, to make a "Safe Softmax" (a common math operation in AI), you first have to find the biggest number in a list, then subtract that number from everything else, then add up the results, and finally divide.

The Problem: The "Stop-and-Go" Traffic Jam

Currently, most AI compilers (the chefs' managers) handle this chain of tasks inefficiently. They treat each step as a separate trip to the pantry.

1. Trip 1: Go get the ingredients, find the biggest number, write it down, and go back to the pantry.
2. Trip 2: Go get the ingredients again, subtract the number, write down the sum, and go back.
3. Trip 3: Go get the ingredients again to do the final division.

This is like a delivery driver who drops off a package, drives all the way back to the warehouse to pick up the next package, and drives back out again, even though the next package was sitting right next to the first one. It wastes time (latency) and clogs up the roads (memory bandwidth).

Furthermore, because each step depends on the result of the last, the kitchen has to store all the intermediate results on the counter. If the list of ingredients is huge, the counter gets cluttered, and the chef runs out of space to work.

The Solution: RedFuser (The "Super-Organizer")

The paper introduces RedFuser, a new framework that acts like a genius sous-chef who realizes, "Wait, we don't need to make three separate trips! We can do all these steps in one smooth motion."

RedFuser uses a clever mathematical trick to fuse (merge) these separate steps into a single, continuous loop.

The Creative Analogy: The Assembly Line vs. The Relay Race

• Old Way (Relay Race): Imagine a relay race where Runner A runs the whole track, stops, hands a baton to Runner B, who then runs the whole track again, stops, and hands it to Runner C. They are all running the same distance, but they keep stopping and starting.
• RedFuser Way (Assembly Line): Now imagine an assembly line where the product moves down a conveyor belt. As it passes Station A, a worker adds a screw. As it immediately passes Station B, another worker adds a nut. As it passes Station C, a worker tightens it. The product never stops moving, and the workers don't have to run back and forth.

How RedFuser Works (The Magic Tricks)

1. The "One-Time Load" Trick
Instead of loading the data from the main pantry (memory) three times, RedFuser loads it once. It keeps the data in a small, super-fast "pocket" (on-chip memory) and performs all the math operations on it while it's there. This eliminates the traffic jams.

2. The "Incremental Update" Trick (The Rolling Calculator)
This is the paper's most brilliant insight. Usually, to do these steps, you need to know the entire result of the previous step before you can start the next one. If you have a million numbers, you'd need a huge counter to hold the sum of the first million before you can do the next step.

RedFuser uses an Incremental Computation method. Think of it like a rolling calculator:

• Instead of waiting for the whole list to finish adding up, you add the first number, update your total, add the second number, update the total, and so on.
• You don't need a giant counter; you just need a tiny pocket to hold the current total.
• This allows the system to handle massive lists of data without running out of "counter space" (memory), even on small devices.

The Results: Speeding Up the Kitchen

The authors tested RedFuser on real-world AI tasks like:

• Attention Mechanisms: The core of how AI models (like Chatbots) understand context.
• MoE Routing: How models decide which "expert" to ask for help.
• FP8 Quantization: A way to shrink models to make them faster.

The Outcome:

• RedFuser made these tasks run 2 to 5 times faster than the best existing AI compilers (like TVM or PyTorch).
• It performed just as well as hand-written code by expert engineers (who usually spend weeks manually optimizing these specific tasks).
• Crucially, it did this automatically. You don't need a PhD in math to use it; the framework figures out the fusion for you.

In a Nutshell

RedFuser is a tool that stops AI computers from making unnecessary trips to the memory pantry. It combines multiple math steps into one smooth, continuous flow, allowing AI models to run significantly faster and more efficiently, bringing us closer to having super-smart, lightning-fast AI assistants.

Technical Summary

Here is a detailed technical summary of the paper "RedFuser: An Automatic Operator Fusion Framework for Cascaded Reductions on AI Accelerators."

1. Problem Statement

Modern AI models, particularly Large Language Models (LLMs), rely heavily on cascaded reductions—sequences of reduction operations with strict data dependencies (e.g., the safe softmax followed by GEMM in attention mechanisms, or MoE routing pipelines).

• The Bottleneck: In current AI compilers (e.g., TVM, PyTorch Dynamo), these operations are executed as separate kernels. This leads to:
  • Redundant Memory Access: Input data and intermediate results must be reloaded from global memory for each reduction stage.
  • Limited Parallelism: The sequential dependency between reduction trees prevents parallel execution and hides memory latency.
  • Cache Pressure: Intermediate results must be cached entirely before the next stage begins, limiting the length of reductions that can be processed on-chip.
• Limitations of Existing Solutions: While hand-optimized kernels (e.g., FlashAttention, FlashDecoding) achieve high performance, they are specific to certain patterns and require expert-level manual tuning. Existing compilers lack the formal methodology to automatically generalize and fuse these complex, dependent reduction chains.

2. Methodology

The authors propose a formal theoretical framework to analyze and fuse cascaded reductions into a single loop, introducing an incremental computation form to overcome cache constraints.

A. Formal Definition & Fusion Conditions

The paper defines cascaded reductions as a "Chain of Reduction Trees." To fuse them, the $i$-th reduction function $F_i(\mathbf{X}, \mathbf{D}_i)$ must satisfy three conditions:

1. Decomposability: $F_i$ can be split into $G_i(\mathbf{X}) \otimes_i H_i(\mathbf{D}_i)$, where $\otimes_i$ is a binary operator.
2. Algebraic Structure: The set and operator $(\mathcal{S}, \otimes_i)$ must form a commutative monoid (associative, commutative, identity element exists).
3. Distributivity: The reduction operator $\oplus_i$ must distribute over $\otimes_i$.

If these hold, the authors derive a Fused Expression where the output of a segment at level $k$ depends only on the partial results of the same level from preceding reductions, rather than waiting for the final output of the previous reduction tree.

B. Incremental Computation Form

A major challenge in fusion is that standard fused expressions require caching all intermediate results from the previous level, which exceeds on-chip memory for long sequences.

• Solution: The authors derive an incremental computation form. Instead of waiting for all inputs, the system updates the current partial result immediately as new data arrives.
• Benefit: This reduces storage complexity from $O(L)$ to $O(1)$, allowing reductions of arbitrary length to be processed within a single thread block (CTA) without spilling to global memory.

C. The RedFuser Framework

RedFuser is a two-stage automatic fusion framework built on the TVM compiler stack:

1. Symbolic Deduction Engine:
   • Parses the computational graph (Relax/TIR) to extract mathematical expressions.
   • Uses a Fixed-Point Analysis Method to automatically verify decomposability and derive the functions $G_i$ and $H_i$.
   • Generates the fused mathematical expressions and their incremental forms.
2. Hardware-Aware Code Generation:
   • Strategies: Supports Single-Segment (entire reduction in one CTA) and Multi-Segment (parallel processing across CTAs) strategies.
   • Tensorization: Converts scalar loops into tile-level IR using TileLang.
   • Optimization: Applies thread mapping, software pipelining, and hardware-specific instructions (e.g., TMA, WGMMA on Hopper GPUs).
   • Auto-Tuning: Optimizes parameters like block size and pipeline depth.

3. Key Contributions

1. Theoretical Methodology: A formal proof and derivation for fusing cascaded reductions into a single loop, including a novel incremental computation scheme that eliminates on-chip memory bottlenecks.
2. Automatic Framework (RedFuser): The first compiler framework capable of automatically identifying, analyzing, and fusing diverse cascaded reduction patterns without manual intervention.
3. Generalization: Unlike previous hand-tuned solutions (FlashAttention), RedFuser generalizes to various patterns including Attention, MoE routing, and FP8 Quantization + GEMM.
4. Performance: Demonstrates that automated fusion can match or exceed the performance of expert-written kernels.

4. Experimental Results

The framework was evaluated on NVIDIA A10 and H800 GPUs across four workloads: Multi-Head Attention (MHA), Multi-Latent Attention (MLA), MoE routing, and FP8 Quant + GEMM.

• Speedup: RedFuser achieves 2× to 5× speedup over state-of-the-art compilers (TVM, PyTorch Dynamo).
• Comparison to Hand-Optimized Kernels:
  • MHA: Matches or slightly exceeds FlashAttention2 (avg 1.09× performance).
  • MLA: Achieves 102% of FlashMLA performance.
  • MoE & Quant: Outperforms general compilers by 1.7× to 12.1×, as no specialized hand-tuned libraries existed for these specific cascades.
• Scalability: The incremental computation mode allows RedFuser to handle long sequences (e.g., LLM decoding) where non-incremental fused approaches fail due to memory constraints.
• Portability: Validated across A10, A100, H800, and AMD MI308X, showing consistent gains.

5. Significance

• Bridging the Gap: RedFuser bridges the gap between high-performance, ad-hoc manual optimizations and general-purpose compiler automation. It proves that complex, dependent reductions can be optimized automatically.
• Efficiency: By eliminating redundant memory traffic and enabling fine-grained parallelism, it significantly reduces the computational cost and latency of training and inference for modern AI models.
• Future-Proofing: As AI models grow more complex with diverse reduction patterns (e.g., in MoE and quantization), frameworks like RedFuser are essential for maintaining efficiency without requiring a new hand-written kernel for every architectural change.

In conclusion, RedFuser provides a rigorous mathematical foundation and a practical implementation for automating one of the most critical performance bottlenecks in modern AI accelerators.