Tensor superoptimizer boosts gpu program speed by up to three times
Unleashing the Power of Equality Saturation for Tensor Program Superoptimization
Distributed, Parallel, and Cluster Computing
Summary
Making programs that run on graphics cards faster is tricky because there are many ways to change the program, and trying them all takes too long. The authors created EqiForge, a tool that uses a special method called equality saturation to explore many program changes efficiently. EqiForge finds better ways to run tensor math, which is common in AI tasks, making programs up to three times faster than current methods. It can also discover new fast versions of parts used in popular AI models.
What this means in practice
- •For gpu software engineers: Accelerate tensor computations by automatically finding fast implementations that combine high-level math with low-level execution details.
- •For machine learning framework developers: Integrate advanced superoptimization techniques to improve performance of attention and transformer kernels beyond existing hand-optimized methods.
Authors
Qi Zhan, Xing Hu, Xin Xia, Shanping Li
Abstract
Efficient GPU implementations of tensor programs often require joint optimization of high-level algebraic formulations and low-level execution strategies. However, the resulting search space grows rapidly as transformations combine across operators, making joint optimization difficult to scale. We present EqiForge, a tensor program superoptimizer based on equality saturation. Its unified IR represents high-level tensor expressions and tiled computations in a single expression language. By composing equality rules, EqiForge derives fused implementations such as FlashAttention-style kernels directly from tensor expressions. Early compaction prunes redundant partial programs before completion, while subgraph composition extends the search to larger graphs. Across tensor-program benchmarks, EqiForge achieves a geometric mean speedup of 1.32x and a maximum of 2.74x over the fastest available baseline per configuration. Its attention kernels outperform FlashAttention by up to 1.87x in decode and approach its performance in prefill. EqiForge also discovers new implementations that outperform torch.compile on various Transformer layers, including QK-normalized MLA (3.16x) and mHC (5.84x).