Skip to main content
Back to timeline
arXivSource publication:

Cleave decouples algebraic search from operator scheduling to generate fused kernels up to 2.8x faster than the best baseline on LLM subgraphs

Synopsis

Cleave is an ML compiler built on symbolic decoupling: it discovers algebraic transformations by superoptimizing over graphs with symbolic shapes, then schedules each resulting graph on concrete shapes, fusing graphs with multiple reductions through iterative tiling and horizontal fusion; on common LLM subgraphs it generates kernels up to 2.8x faster than the best baseline (1.6x on average) and cuts compilation time by 5.9x on average versus Mirage.

Source-provided article image: Cleave: Scaling Tensor Program Optimization via Decoupled Algebraic Search and Operator Scheduling
Figure 1 ·

Figure 1 . Symbolic decoupling in Cleave , illustrated by the fusion of RMSNorm followed by Matmul. The first phase performs algebraic transformation that finds an equivalent graph which orders the row-wise scaling after Matmul. The second phase performs operator scheduling which makes each slice of X X feed both reductions in a single loop over K K .

arXiv

Interpretation

Cleave introduces symbolic decoupling: the algebraic phase searches for equivalent transformations on a graph with symbolic shapes, and the scheduling phase binds those symbols to concrete shapes and schedules each resulting graph, so algebraic search never enumerates loop structures or tile sizes and the scheduler never reasons about algebraic equivalence. Prior scheduling-based compilers optimize loops, tile sizes and memory layouts but do not discover the needed algebraic transformations; algebraic optimizers search equivalent graphs but do not decide how to tile and execute them; Mirage searches both jointly and suffers a search space too large to navigate. Cleave separates them along the line that a transformation's validity is a property of the graph alone. The paper argues the separation through FlashAttention (moving softmax's division after the V-Matmul) and Split-K, and ablates on 9 subgraphs: disabling algebraic transformation degrades performance by 44%, not generating iterative tiling schedules by 40%, and disabling both by 47%.

Cleave adds a Split operator with a symbolic split count, letting superoptimization decide where to split a reduction dimension and how to combine partial results while the concrete split value is bound later by the scheduler; symbolic graphs also let equivalence checking use any small instantiation satisfying the graph's shape constraints, lowering testing cost. Split-K is normally stated with a concrete partition size, which makes it appear that the schedule must be fixed before the transformation can be applied; making the split count one more symbol lets the algebraic phase focus on where to split and how to combine. The paper prunes the search with reducibility-based constraints on which dimensions may be split or reduced, and still performs probabilistic testing on the original user-graph shapes at the end of superoptimization to preserve the correctness guarantee; in the decode setting, enabling Split-K yields kernels up to several times faster than without it (the specific figure is absent from the loaded text).

Cleave's scheduler uses iterative tiling to break reduction dimensions into smaller tiles accumulated step by step, and horizontal fusion to compute reductions sharing upstream producers in the same loop, fusing multi-reduction graphs without materializing full intermediate tensors. Existing scheduling-based compilers tile only output dimensions, so even given the transformed graph they materialize full intermediate tensors in SMEM, limiting attention fusion to short sequence lengths. The paper's attention example groups Rowsum and the second Matmul because they share the Exp node; Welder fails to fuse consecutive reduction-based subgraphs once sequence length reaches 512, while Cleave still fuses them.

For dynamic workloads, Cleave compiles each operator once and keeps dynamic axis extents symbolic with runtime loop bounds during code generation, so a single kernel serves all captured shapes of that operator. Scheduling needs concrete dimensions to pick tile shapes, thread counts and buffer sizes, so Cleave schedules against one representative concrete shape and then symbolizes the dynamic axes at code generation. Nine paged and ragged attention operators are each compiled once and run all 303 shapes captured in FlashInfer-Bench with no recompilation, reaching geometric mean speedups of 1.4x over FlashInfer's handwritten FA2 backend and 1.7x over FA3; it is slower only on the two MLA operators, where FA3 merges partial results inside a single cooperative kernel while Cleave pays a second launch.

Perspective

The results target LLM subgraphs and transformer layers dominated by reduction-based operators, measured on an NVIDIA A100-SXM4-40GB and an H200-SXM5-141GB with FP16 inputs, FP32 accumulation and FP16 outputs; the dynamic-operator portion covers paged and ragged attention with shapes taken from 303 production-trace shapes captured in FlashInfer-Bench. The directly reusable object is compiler and kernel engineering practice: decoupling algebraic search from scheduling, using symbolic shapes to cheapen equivalence checking, and using iterative tiling plus horizontal fusion for multi-reduction graphs. The paper also shows that feeding Cleave's algebraic transformation output to Welder improves its performance, indicating the transformation phase can independently supply input to other scheduling-based compilers.

Several figures are absent from the loaded text (for example the decode-setting Split-K speedup, the per-operator speedup ranges over FA2 and FA3, and the multipliers for compilation time and speedup in the conclusion), so those magnitudes can only be read through the geometric means and overall ranges the paper reports. Reducibility annotation currently requires user input; the paper notes it can be automated with data-dependency analysis, but the automated version is not evaluated here. Algebraic transformation generally worsened performance for the rule-based torch.compile, which the paper attributes to the compiler no longer recognizing fusion patterns it found in the original graph, suggesting that how this transformation phase pairs with different compilers remains an open question. In addition, two MLA operators are slower than FA3, attributed to FA3 merging partial results inside a single cooperative kernel while Cleave pays a second launch; under which workloads that structural difference would change the conclusion is worth watching.

Sources