Source-linked AI summary
Unleashing the Power of Equality Saturation for Tensor Program Superoptimization
Qi Zhan, Xing Hu, Xin Xia, Shanping Li
TL;DR
Jointly optimizing tensor algebra and execution strategy is difficult to scale as transformations combine across operators. EqiForge uses equality saturation over a unified IR and compacts redundant partial programs. It reports speedups across tensor-program benchmarks, including a 1.32× geometric mean over the fastest available baseline.
Problem
Jointly exploring algebraic formulations and execution strategies is difficult to scale because their combinations create a large search space and complete candidate construction becomes expensive.
Source: §1 · p. 1, §5.1 · p. 7Method
EqiForge uses equality saturation over a unified IR of tensor expressions and tiled computations, composing rewrite rules and pruning equivalent partial programs early.
Source: §1 · p. 2Results
EqiForge achieves a 1.32× geometric mean speedup and up to 2.74× over the fastest available baseline across tensor-program benchmarks.
Source: §1 · p. 2, §7.2 · p. 9Takeaways & Limitations
Composed equality rules derive fused implementations such as online attention, while subgraph composition extends optimization to larger programs.
Source: §9 · p. 12The algebraic rewrites assume real arithmetic, and floating-point implementations may differ because of reassociation and internal precision choices.
Source: §4 · p. 5
Abstract
from arXiv · showhide
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).
1 Introduction
EqiForge addresses the difficulty of scaling joint optimization across algebraic formulations and execution strategies. It combines these transformations in an equality-saturation search and reports speedups across tensor-program benchmarks.
1 Introduction
- FlashAttention illustrates why both mathematical formulation and execution strategy may need to change: online softmax is combined with tiled attention matrix multiplications.The resulting fused kernel computes incrementally without storing full score and probability matrices in global memory.Source: §1 · p. 1
- Prior approaches explored algebraic rewriting, scheduling, fusion, and equality-saturation optimization, but the paper frames scalable joint exploration as the remaining challenge.The cited work includes tensor-graph rewrites and loop and memory transformations.Source: §1 · pp. 1–2
- The central challenge is preserving optimization opportunities across operators while keeping candidate construction and evaluation tractable.Restricting search scope can separate operations that could benefit from joint optimization.Source: §1 · p. 2
- EqiForge jointly searches high-level tensor expressions and tiled implementations, composing equality rules to derive fused kernels such as FlashAttention-style implementations.Its typed IR combines tensor expressions and tiled computations, while compaction prunes equivalent partial programs before complete candidates are constructed.Source: §1 · p. 2
- A 1.32× geometric mean speedup and up to 2.74× speedup over the fastest available baseline summarize EqiForge’s tensor-program results.The paper also reports up to 1.87× over FlashAttention in decode, similar performance in prefill, and 3.16× and 5.84× over torch.compile for QK-normalized MLA and mHC, respectively.Source: §1 · p. 2
2 Equality Saturation
An e-graph compactly stores equivalent expressions so that successive rewrites can expose further optimization opportunities. Extraction then selects a finite program, whose execution cost can differ even when its value is unchanged.
2 Equality Saturation
- An e-graph groups equivalent expressions into e-classes, with e-nodes recording operators and references to child e-classes.Factoring AB + AC adds A(B + C) to the same e-class without duplicating the enclosing Exp node.Source: §2 · p. 2
- Repeated valid rewrites add equivalent expressions while preserving congruence, allowing newly exposed expressions to trigger further rules.Rewriting continues until saturation or a resource limit is reached.Source: §2 · p. 2
- Extraction resolves choices from the result e-class to form a finite program, even when the e-graph contains cycles.Equivalent forms can require different numbers of matrix multiplications and therefore have different execution costs.Source: §2 · pp. 2–3
- EqiForge applies equality saturation to one IR containing high-level tensor expressions and tiled computations, then extracts complete programs for performance evaluation.Algebraic rewrites and parallel refinements add equivalent expressions to that same e-graph.Source: §2 · p. 3
3 A Unified IR for Tensor Programs
EqiForge uses one pure expression language for tensor computations and their parallel implementations. Its types track tensor shape and parallel state so that transformations can compose across tensor operations, tiles, and reductions.
3 A Unified IR for Tensor Programs
- The unified IR represents tensor expressions and parallel implementations in one pure language, with operators for tensor values, tuples, and field selection.The IR includes parallel-view primitives, reductions, and intermediate storage; x names a function argument, c a scalar literal, and ρ a reducer.Source: §3 · p. 3
3.1 Parallelism
- Parallel state distinguishes partitioned slices, replicated operands, and partial reduction contributions, even when their logical shapes and dtypes match.Partition divides a tensor across an axis, Combination rejoins matching slices, Replicate shares an operand, and Reduce combines partial contributions.Source: §3.1 · p. 3
3.2 Computation on Tiles
- Tile computations reuse pointwise operators, while refinement relates high-level MatMul to TileMatMul and the operations that combine its output tiles.A row-tiled product can apply Exp locally and then use Combination to join the output tiles.Source: §3.2 · pp. 3–4
3.3 Reduction State and Intermediate Storage
- Reduce can operate within a tile along a tensor axis or across parallel coordinates, with tensor or tuple values as reduction state.In the row-sum example, per-tile partial sums are combined across p to produce the complete result.Source: §3.3 · p. 4
- Collection converts a parallel axis into a leading tensor dimension, allowing partial values to be stored before a later reduction.This offers a multi-kernel alternative to keeping intermediates within a kernel, trading reduced intermediate memory traffic against potentially limited parallelism.Source: §3.3 · pp. 4–5
4 Equality Rules for Tensor Programs
EqiForge uses equality rules to connect tensor-level formulations with tiled implementations, including fused reductions. These rules construct parallel implementations, propagate tile views, and optionally store reduction inputs to increase parallelism.
4 Equality Rules for Tensor Programs
- Algebraic rewriting, refinement, propagation, and repair fuse dependent reductions into a tiled rolling-attention implementation, assuming real arithmetic despite possible floating-point differences.The repair rescales the accumulated normalizer and weighted sum when the running maximum changes.Source: §4 · p. 5, §4.2 · p. 6
4.1 Constructing Parallel Implementations
- MatMul and reductions can partition output dimensions into complete tiles or reduction dimensions into partial results that require combination; batch refinement also handles broadcast operands.Partitioning a reduction dimension requires an associative reducer, with local and final reductions initialized to the identity.Source: §4.1 · pp. 5–6
4.2 Composing Reduction State
- Independent reductions fuse by pairing their states, while dependent reductions require repair updates before one tuple-state Reduce can represent both computations.The same Reduce representation applies across tensor dimensions and parallel axes.Source: §4.2 · p. 6
4.3 Propagating Parallelism
- Propagation rules move Partition and Replicate through pointwise operations and align compatible parallel views, expressing attention score tiles from query and key tiles.The resulting tiles supply the rolling reduction, and final division computes normalized output tiles before Combination.Source: §4.3 · p. 6
4.4 Storing Reduction Inputs
- Storing reduction inputs can expose more thread blocks when fused kernels have few output tiles, at the cost of intermediate memory traffic and a separate reduction kernel.For attention, stored tuples include scores, local maxima, and value tiles, which a later kernel consumes with the same rolling reducer.Source: §4.4 · p. 6
5 Scalable Program Extraction
EqiForge extracts representative complete programs from an equivalent-expression space, using early compaction to avoid redundant extraction and subgraph composition to search larger graphs.
5 Scalable Program Extraction
- Extraction fixes computation and kernel boundaries while leaving scheduling parameters symbolic, then selects tiled MatMul implementations and forms finite acyclic programs.The section develops compaction and subgraph combination to make extraction scalable.Source: §5 · p. 6
5.1 Compacting the Search Space
- Enumerating complete programs before grouping still constructs exponentially many candidates; early compaction instead prunes equivalent partial branches before their holes are resolved.The equivalence relation groups commutative pointwise operand orders and equivalent parallel views, but keeps distinct MatMul structures, reduction organizations, and intermediate storage separate.Source: §5.1 · p. 7
5.2 Searching Larger Graphs
- Larger subgraphs permit more joint transformations, while smaller ones reduce search cost but may require more intermediate storage; EqiForge ranks partitions by estimated work and boundary data transfer.It composes candidates lazily through typed interfaces rather than enumerating the full Cartesian product.Source: §5.2 · p. 7
6 Code Generation and Tuning
EqiForge maps implementation families to executable GPU kernels and tunes schedules. It ranks candidates with optimistic latency estimates before progressively allocating measurement effort.
6.1 Execution Families and Code Generation
- Each implementation family combines an extracted program with its legal schedule domain; the backend generates kernels under legal schedules and measures end-to-end latency.Execution structures are canonicalized, and a family is discarded when another with the same structure covers its entire schedule domain.Source: §6 · p. 7
6.2 Static Family Ranking
- Candidate ranking estimates optimistic latency by minimizing over legal schedules the maximum modeled work or data movement divided by hardware throughput or bandwidth.A linear-programming relaxation avoids enumerating every schedule, and its solution is rounded into an initial tuning schedule.Source: §6.2 · p. 8
6.3 Progressive Schedule Tuning
- Successive halving allocates larger measurement budgets to faster candidates while retaining some with low estimated latency; schedule search combines local changes, evolutionary mutations, and random sampling.For multi-kernel programs, buffer-shape parameters are fixed and kernels are tuned one at a time while validating and timing the complete program.Source: §6.3 · p. 8
7 Evaluation
EqiForge combines tensor-level algebraic rewriting with tiled implementations and evaluates the resulting search across tensor programs, Transformer layers, and search costs. Results include faster selected implementations and reduced extraction overhead, with performance varying by workload.
7.1 Experimental Setup
- The implementation captures PyTorch programs, expands composite operators for optimization, and generates Triton kernels using egg for saturation and extraction.The IR supports 18 pointwise operators plus reshape, transpose, and tensor construction, with 120 rewrite rules; search and tuning use a four-hour budget per configuration.Source: §7.1 · pp. 8–9
7.2 Tensor Program Performance
- 2.74× over torch.compile on RMSNorm–MLP at B=8, where EqiForge moves normalization factors after projections and fuses both projections with the final elementwise product.This avoids intermediate writes to global memory.Source: §7.2 · p. 9
- 3–6% is the prefill latency gap to FlashAttention; EqiForge also outperforms all compared baselines in the three decode cases.The generated prefill kernels combine tiled matrix multiplications with online softmax, without storing full score and probability matrices.Source: §7.2 · p. 9, Figure 8 · p. 9
7.3 Transformer Layer Optimization
- With one query token, EqiForge achieves 3.16× over torch.compile for QK-normalized MLA (1, 4096, 16K) and 5.84× for mHC (16, 512, 1K).For mHC, EqiForge combines the first 9 residual update with the next connection’s normalization in one kernel, reusing the updated residual within that kernel.Source: §7.3 · pp. 9–10
- Over torch.compile, sliding-window GQA reaches 1.06×, attention sinks 1.48×, and differential attention 1.51×.The attention-sink implementation processes four initial sink tokens and the most recent 4096 tokens in one softmax reduction, skipping the masked middle.Source: §7.3 · p. 10
7.4 Search Cost
- Attention e-graphs exceed 22,000 e-nodes in decode and 65,000 in prefill, as interacting tiling, fusion, and storage choices increase saturation and extraction costs.Most non-attention workloads produce fewer than 1,200 e-nodes and complete saturation and extraction within tens of seconds.Source: §7.4 · p. 10, Table 2 · p. 11
- Early compaction reduces expanded states by factors of 128–2,465 while returning the same number of distinct candidates across four workloads where both strategies finish.For nGPT, enumeration falls from 106 s to 0.47 s; for RMSNorm–SwiGLU, early compaction returns 330 candidates in 16.97 s, while compaction after completion times out after two hours with 72.Source: §7.4 · p. 11, Figure 10 · p. 11
8 Related Work
Related work spans tensor-compiler scheduling and code generation, tensor-program superoptimization, fusion, and equality saturation. EqiForge’s approach combines algebraic forms and parallel refinements in one equality search, with extraction that prunes redundant partial programs.
Tensor compilers and optimization
- Tensor compilers optimize schedules, code generation, tiles, or latency models, while EqiForge complements schedule optimization by jointly searching algebraic forms and parallel decompositions [2] [6] [39].The passage also describes TVM, TensorIR, Ansor, ROLLER, and Felix as examples of these approaches.Source: §8 · p. 11
Tensor program superoptimization
- Tensor-program superoptimizers search graph rewrites or equivalent implementations; EqiForge uses equality saturation and early compaction during extraction [32].Mirage [28] and Prism [29] are described as enumeration-based superoptimizers.Source: §8 · p. 11
Operator fusion and equality saturation
- Unlike Trinity’s tile-level IR and two-pass extraction, EqiForge uses a pure expression IR where tensor rewrites compose with parallel refinement; its equality rules also express Neptune’s repair relation [36].The related-work discussion places these approaches alongside operator fusion systems and equality-saturation methods.Source: §8 · p. 11
9 Conclusion
EqiForge brings algebraic reformulation and parallel implementation generation into one equality search, using a unified typed IR and compaction to prune equivalent branches before completion. Its evaluation reports a 1.32× geometric mean speedup and up to 2,465 fewer expanded extraction states.
9 Conclusion
- 1.32× geometric mean speedup over the fastest baseline for each tensor-program configuration; early compaction cuts extraction states by up to 2,465.The unified typed IR composes algebraic and parallel transformations, while subgraph composition extends search to larger programs and across Transformer-layer components.Source: §9 · p. 12