Unleashing the Power of Equality Saturation for Tensor Program Superoptimization

2026-09-10

AI Summary

本文通过基于等价饱和的EqiForge方法，解决了高效GPU实现张量程序时高、低级优化难以联合扩展的问题，实现了显著的性能提升。

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).

Problem

• tensor program
• joint optimization
• equality saturation
• search space

Innovation

• Equality Saturation
• Tensor Program Superoptimization
• Unified IR
• Early Compaction
• Subgraph Composition