Papers for

machine learning framework developers

Papers whose findings have a practical use for this group, as judged from the abstract. Open a paper to read what it means in practice.

Tensor superoptimizer boosts gpu program speed by up to three times

Unleashing the Power of Equality Saturation for Tensor Program Superoptimization

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

Fri 11 SeptDistributed, Parallel, and Cluster Computing
The gist
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.
Open 2609.12330v1