PaperlayerView exact revision on arXiv

Source-linked AI summary

Cleave: Scaling Tensor Program Optimization via Decoupled Algebraic Search and Operator Scheduling

David Pissarra, Jinkun Lin, Haitian Jiang, Aurojit Panda, Jinyang Li

arXiv:2610.07742v1cs.PLcs.AI

TL;DR

Cleave addresses the difficulty of generating expert-quality fused kernels for reduction-heavy computations, where algebraic graph changes and operator scheduling must work together. It decouples symbolic graph search from concrete-shape scheduling, reporting faster kernels and shorter compilation than baselines.

  • Problem

    Automatically generating fused kernels that match expert implementations remains difficult, in part because reduction dependencies can strain intermediate storage and limit GPU parallelism.

    Source: §1 · p. 1

  • Method

    Cleave searches for equivalent computation graphs with symbolic shapes, then schedules each graph on concrete shapes using iterative tiling and horizontal fusion.

    Source: §1 · p. 2, §2.3 · p. 4

  • Results

    Kernels are up to 2.8× faster than the best baseline on each subgraph, averaging 1.6×, while compilation time is reduced by 5.9× on average versus Mirage.

    Source: §1 · p. 2

  • Takeaways & Limitations

    Cleave combines symbolic-shape transformation with iterative tiling and horizontal fusion to generate high-performance kernels within reasonable time.

    Source: §9 · p. 13

    Users must annotate which input-tensor dimensions are eventually reduced so Cleave can prune Split insertions.

    Source: §4 · p. 5

Abstract

from arXiv · show

Optimized kernels such as FlashAttention and FlashDecoding are crucial for accelerating today's large models. Most of them are handwritten by experts because existing ML compilers cannot match their efficiency. Producing such kernels requires fusing computations with multiple reductions, which requires both algebraic transformation of the computation graph and operator scheduling of the transformed graph. Unfortunately, searching the two jointly yields a space too large to navigate. We propose Cleave, an ML compiler built on symbolic decoupling: Cleave discovers transformations by performing superoptimization on a graph with symbolic shapes, and then schedules each resulting graph on concrete shapes. Representing shapes as symbols makes equivalence checking cheap and lets a new Split operator, with a symbolic split count, parallelize along a reduction dimension. Cleave's scheduler fuses graphs with multiple reductions through iterative tiling and horizontal fusion. Evaluation on common LLM subgraphs shows that Cleave generates kernels up to 2.8x faster than the best baseline (1.6x on average) and reduces compilation time by 5.9x on average compared to Mirage. For dynamic workloads captured from production serving traces, Cleave compiles each operator once and achieves geometric mean speedups of 1.4x and 1.7x over FlashInfer's handwritten FA2 and FA3 backends. Cleave's code is available at: https://github.com/nyu-systems/cleave

1 Introduction

Cleave separates algebraic graph transformation from concrete-shape operator scheduling, addressing the difficulty of fusing reduction-heavy computations. Its evaluation reports faster kernels and shorter compilation times than baselines, including dynamic-workload and transformer-layer results.

1 Introduction

  • Modern fused kernels are difficult to generate because dependencies among reductions can exceed available intermediate storage, while straightforward output partitioning may expose too little GPU parallelism.Existing scheduling-based compilers do not discover needed algebraic transformations, algebraic optimizers do not determine execution schedules, and joint search can become too large.Source: §1 · pp. 1–2
  • Algebraic transformation can reorder RMSNorm scaling after Matmul so both reductions consume the same input tile, and can expose partial computations through Split-K for parallelism.Scheduling then determines the tile sizes and maps partial computations across threadblocks.Source: §1 · p. 2
  • Symbolic decoupling searches algebraically equivalent graphs with symbolic dimensions, then schedules each graph for concrete shapes, separating equivalence reasoning from execution choices.Algebraic validity depends on the graph rather than its schedule or tensor shapes; profitability depends on scheduling and concrete shapes.Source: §1 · p. 2
  • Cleave uses superoptimization to enumerate equivalent graphs, then a dedicated scheduler performs fused tile-by-tile execution and selects promising candidates with an analytical cost model.The symbolic Split operator partitions a reduction dimension into s parts, with the scheduler binding s later.Source: §1 · p. 2
  • Iterative tiling breaks reduction dimensions into small tiles, while horizontal fusion computes reductions sharing upstream producers in the same loop and reuses their tiles.This avoids materializing full intermediate tensors across reductions.Source: §1 · p. 2
  • Across commonly used subgraphs including GQA, MLA, and SwiGLU, Cleave kernels are up to 2.8× faster than the best baseline per subgraph, averaging 1.6×, and reduce compilation time versus Mirage by 5.9× on average.Source: §1 · p. 2
  • For dynamic workloads, one compilation per operator serves all 303 captured shapes across nine operators, with geometric mean speedups of 1.4× over FlashInfer’s FA2 and 1.7× over its FA3 backend.The workloads vary in batch size and sequence length.Source: §1 · p. 2

2 Motivation and Our Approach

Cleave separates algebraic graph rewrites from execution scheduling, making it possible to discover equivalent graphs independently of concrete schedules. Its motivating examples show why both stages are needed to fuse reduction-heavy computations.

2.1 Background

  • Tiling reuses data in on-chip memory, and fusion keeps intermediate results there rather than sending them through global memory.These techniques reduce data movement, a central requirement for efficient GPU kernels.Source: §2.1 · p. 3

2.2 Motivating Examples

  • FlashAttention moves softmax division after the V-Matmul, then processes the sequence reduction tile by tile so only one tile of each intermediate remains live.The original graph requires materializing the large softmax input and output, which makes direct fusion costly.Source: §2.2 · p. 3
  • Split-K partitions a reduction into partial dot products that can run on separate threadblocks, then sums their results in a combine kernel.A concrete partition size appears to couple the rewrite to scheduling, but a symbolic split count separates those decisions.Source: §2.2 · p. 4

2.3 Our Approach

  • Symbolic decoupling searches for equivalent graphs with symbolic dimensions, while iterative tiling and horizontal fusion schedule each transformed graph without materializing full intermediates.The symbolic split count lets algebraic search choose where to split without choosing a concrete partition size.Source: §2.3 · p. 4

3 System Design Overview

Cleave optimizes a computation graph in separate algebraic-transformation and operator-scheduling phases, followed by code generation and profiling.

3 System Design Overview

  • Cleave first symbolizes tensor dimensions, searches for mathematically equivalent graphs, validates them with probabilistic testing, and restores the original shapes.The algebraic phase returns validated equivalent graphs with concrete shapes.Source: §3 · p. 4
  • Its scheduler applies iterative tiling and horizontal fusion, then uses an analytical cost model to select top-k tile sizes.The schedule transforms the validated graph’s loop structure.Source: §3 · p. 4
  • For top-k candidates, Cleave chooses tensor placement and layout, generates kernels through TileLang, and profiles them to pick the best candidate.A final pass handles memory placement, code generation, and profiling.Source: §3 · p. 4

4 Algebraic Transformation Phase

Cleave performs superoptimization on symbolic computation graphs, including a symbolic Split operator, then recovers concrete shapes and prunes transformations that violate reducibility constraints.

4 Algebraic Transformation Phase

  • Symbolic dimensions let Cleave test equivalence on small tensors, reducing the cost of random testing while preserving broadcasting structure and operator shape constraints.Dimensions of size 1 are retained; other dimensions become symbolic variables.Source: §4 · p. 5
  • The Split operator divides a dimension into s parts, allowing algebraic search to find where to split without searching for a concrete split size.For example, splitting dimension 1 of A with shape [m,k] yields shape [s,m,k/s].Source: §4 · p. 5
  • In attention, splitting the key-value reduction yields partial row sums and value products, which a combine kernel reduces across s before final division.The Split operator creates a two-kernel split-and-combine decomposition at reductions over the split dimension.Source: §4 · p. 5
  • Reducibility-based pruning rejects Split or reduction insertions that change which input dimensions are eventually reduced in the output.Cleave currently requires users to annotate input dimensions that are eventually reduced.Source: §4 · p. 5
  • Cleave enumerates candidate graphs up to a size limit and tests equivalence by numerically comparing outputs on small concrete shapes that meet operator constraints.The tests use probabilistic testing; concrete shapes for each equivalent graph are recovered by propagating dimensions and validating possible split factors.Source: §4 · p. 5

5 Operator Scheduling Phase

Cleave schedules transformed graphs by partitioning outputs across threadblocks, then iteratively tiling reductions and horizontally fusing shared work within each threadblock. It represents these nested execution loops as itGraphs and selects among candidate tile programs.

5.1 Motivating Example and Overview

  • Iterative tiling reduces intermediate storage by staging reduction slices, while horizontal fusion lets Rowsum and the second Matmul reuse Exp’s output in the same loop.Tiling N with bR reduces the Exp and first Matmul buffers from [bO, N] to [bO, bR].Source: §5.1 · p. 6
  • Cleave partitions output tiles across threadblocks, then repeatedly tiles reductions and horizontally fuses shared work within each threadblock to build candidate schedules.It ranks candidate itGraphs with a cost model and compiles and profiles the most promising candidates.Source: §5.1 · pp. 6–7

5.2 itGraph Semantics

  • An itGraph represents loop nests as nested computation graphs, with inner loops accumulating into the same output tile and outer parallel loops writing distinct output tiles.Each graph uses a tensor map to describe the slices read and written at each iteration.Source: §5.2 · p. 7

5.3 Inter- and Intra-threadblock Scheduling

  • Inter-threadblock scheduling builds an itGraph from the transformed graph, partitions the output into tiles, and propagates slice mappings backward through data dependencies.A tile-size candidate is rejected when multiple paths give a node mismatching tensor maps.Source: §5.3 · pp. 7–8
  • Intra-threadblock scheduling groups last-wave reductions that share upstream nodes, wraps each group and its dependencies in an inner itGraph, and repeats tiling until all reductions are tiled.This addresses large intermediate buffers caused by untiled reduction dimensions; fused reductions must traverse shared inputs with matching tensor maps.Source: §5.3 · p. 8

6 Memory Planning and Code Generation

Cleave plans memory placement and layouts for scheduled itGraphs, then lowers them to TileLang. For dynamic operators, it specializes scheduling to a representative shape and emits symbolic dynamic extents for runtime use.

6.1 Memory Layout Planning

  • Cleave jointly plans memory placement and layouts for the scheduled itGraph, choosing RMEM or SMEM based on operator needs and reuse opportunities.It can hoist loop-invariant inputs to avoid redundant global-memory loads while maintaining tensor-core layout constraints.Source: §6.1 · p. 8, §6 · p. 8

6.2 TileLang Code Generation

  • Cleave lowers the annotated multi-level itGraph into TileLang by generating wrapper and iteration functions that separate tensor slicing from execution.The wrapper uses tensor maps to supply the appropriate slices to each iteration function.Source: §6.2 · pp. 8–9
  • The current code-generation path targets TileLang; support for other tile-level DSL backends is future work.Source: §6.2 · p. 8

6.3 Compiling Dynamic Operators

  • For dynamic operators, Cleave schedules one representative concrete shape, then replaces dynamic-axis extents with symbols whose sizes are read at runtime.Predicated loads and stores handle tiles larger than the actual input, and one compilation serves shapes varying along the declared dynamic axes.Source: §6.3 · p. 9

7 Evaluation

Cleave is evaluated on fixed-shape and dynamic LLM workloads, end-to-end transformer layers, compilation scalability, and component ablations. The results show gains across these settings, with different boundaries and trade-offs for each.

7.1 Experimental Setup

  • The evaluation covers nine common transformer subgraphs, three model architectures, and nine dynamic attention operators, using FP16 inputs and outputs with FP32 accumulation.Dynamic operator shapes come from 303 workloads captured in FlashInfer-Bench; the evaluated baselines include compiler-generated kernels and FlashInfer's cuDNN, FA2, and FA3 backends.Source: §7.1 · p. 9

7.2 End-to-end Benchmark

  • Across evaluated architectures, end-to-end comparisons measure kernels on a single transformer layer; subgraph evaluation covers prefill and decode, using large sequence and context lengths for prefill and small sequence but large context for decode.For prefill, both input sequence and context lengths are ≥256, and speedup is reported as a geometric mean across shapes versus torch.compile.Source: §7.2 · p. 9, §7.3 · p. 10
  • Cleave is fastest on GPT-2, Llama-3 8B, and Qwen3-32B, with 2.45×, 1.47×, and 1.43× speedup over torch.compile, respectively.Gains are smaller on the larger models because linear layers account for much of their computation time.Source: §7.2 · pp. 9–10

7.3 Subgraph-level Benchmark

  • Cleave beats the best baseline by up to 2.8× on each fixed-shape subgraph; on decode shapes, Split-K yields kernels up to 2.1× faster than without it.Welder stops fusing consecutive reduction-based subgraphs at sequence length 512, leaving intermediates to be materialized in GMEM.Source: §7.3 · pp. 10–11

7.4 Dynamic Operator Benchmark

  • Across 303 captured workloads on nine dynamic operators, one compilation per operator serves every shape; geometric-mean speedups are 1.4× over FA2 and 1.7× over FA3.Cleave is faster than FA2 on all nine operators and faster than FA3 on seven; it is slower on two MLA operators, where FA3 combines partial results in one cooperative kernel while Cleave incurs a second launch. Split-K adds 1.15×–2.99× on four decode operators but does not benefit the five prefill operators.Source: §7.4 · p. 11

7.5 Search Space

  • With Split-K enabled, Cleave compiles subgraphs in 8–14 minutes on average, versus 121–224 minutes for Mirage on the four subgraphs where it generates valid code.Mirage runs out of memory on two other subgraphs and cannot generate code for GQA, MQA, or MLA; with Split-K disabled, Cleave's 1–5 minutes is comparable to Welder's 1–4 minutes.Source: §7.5 · p. 11, Table 1 · p. 11

7.6 Ablation Study

  • Disabling algebraic transformation reduces performance by 44%, while disabling iterative tiling reduces it by 40%, across nine subgraphs and varied shapes.Disabling both scheduling and algebraic transformation reduces performance by 47%; not explicitly maintaining tensor-core layout constraints reduces it by 6%.Source: Figure 12 · p. 12, §7.6 · p. 12

8 Related Work

Related work addresses kernel fusion through algebraic transformation, scheduling, joint optimization, and agentic code search. The paper distinguishes Cleave by its automated combination of general algebraic transformations and scheduling at scalable search cost.

Graph and algebraic transformation

  • Algebraic graph optimizers search equivalent computation graphs [10] [36], but their rules generally ignore scheduling and miss optimizations found by Cleave.TASO searches for substitution rules between equivalent graphs, a large search space that limits its rules to small graphs; Cleave searches for equivalents of a single input graph.Source: §8 · p. 12

Scheduling-based tensor compilers

  • Scheduling-based compilers automate tiling and fusion [41] [42] [44], but none combines general algebraic transformations with scalable scheduling search.Some methods lack iterative tiling, while attention-oriented online-update methods rely on expert-provided algebraic rules.Source: §8 · p. 12

Joint optimization

  • Full-stack compilers such as TVM [4], ONNXRuntime, and TorchInductor perform both transformation and scheduling, but rely mostly on rewrite rules and expert-written templates.Mirage [31] jointly searches algebraic structure and schedules, making its search space too large for common kernels; Prism [32] makes tile sizes symbolic but still jointly searches tiling choices and algebraic structure.Source: §8 · p. 12
  • Trinity can find fine-grained joint optimizations such as online softmax, but requires handwritten rules and a user-supplied tiled program, and does not search for loop tiling.The paper says this requires extra user effort to recover Cleave's optimizations.Source: §8 · p. 12

Agentic kernel optimization

  • LLM-based kernel agents propose, revise, compile, test, and time code, but repeated optimization runs can be expensive; CAKE reports 80 million tokens for one run [38].The paper notes that this cost recurs for each new optimization problem.Source: §8 · pp. 12–13

9 Conclusion

Cleave separates algebraic transformation from operator scheduling to scale automated kernel optimization. Its symbolic-shape graph search and iterative-tiling scheduler produce faster kernels with less compilation time than the joint-search baseline described.

9 Conclusion

  • Cleave reduces compilation time by 5.9× on average versus Mirage and generates kernels up to 2.8× faster than the best baseline.Its symbolic-shape algebraic search and iterative tiling with horizontal fusion target graph transformation and execution across reductions.Source: §9 · p. 13

A Implementation Details

The implementation converts concrete graph dimensions into symbolic groups while preserving broadcasting, then searches equivalent graphs and schedules candidates. Its scheduler proposes iterative tile configurations and ranks them with an analytical cost model.

Symbolic graph conversion

  • Symbolization groups dimensions constrained to match, assigns symbols to non-broadcast groups, and keeps extent-one axes in a separate broadcast group.For example, a concrete batched Matmul becomes Matmul([b,m,k] × [1,k,n] → [b,m,n]).Source: A Implementation Details · p. 15, Table 2 · p. 15

Supporting algebraic graph superoptimization

  • The superoptimizer validates symbolic graph operators across symbol instantiations, tests concrete small instantiations probabilistically, and restricts search to one-level graphs.It adapts Mirage's search to symbolic shapes and uses the Split operator to resolve the scheduling search needed for Split-K.Source: A Implementation Details · p. 15

Tile Size Proposal

  • The scheduler analyzes hardware-imposed reduction constraints, generates candidate output tiles, and groups compatible reductions to explore shared reduction tile sizes.Candidate configurations form iterative graphs for later scheduling.Source: A Implementation Details · p. 15

itGraph cost model

  • Cleave's cost model estimates cumulative GMEM traffic and peak SMEM use while tracking buffer liveness and reuse across operators.For simplicity, the model currently assumes all buffers are in SMEM.Source: A Implementation Details · p. 16

B Maintaining Tensor Core Constraints

Cleave propagates tensor-core layout constraints through the itGraph from MatMul nodes, across element-wise operators, and stops at non-element-wise boundaries. The identified layouts guide memory transfers and avoid bank conflicts.

B Maintaining Tensor Core Constraints

C Additional Evaluation

Additional evaluations find that algebraic transformations improve Welder’s performance but can disrupt TorchInductor fusion, while Cleave’s scheduling addresses workloads Welder cannot fuse at longer sequence lengths. The section also reports evaluation configurations and shape ranges.

C Additional Evaluation

  • Evaluation uses TileLang v0.1.5 and compares Cleave with Mirage, Welder, TorchInductor, TensorRT, FlashAttention, FlashMLA, and cuBLAS.The passage specifies versions or commit hashes for the compared systems.Source: C Additional Evaluation · p. 16
  • Prefill benchmarks sweep sequence length s from 256 to 4096 with attention context length L=4096; decode benchmarks sweep batch size b from 16 to 128, with MLA context lengths from 1024 to 4096.Attention dimensions also vary by subgraph: prefill uses d in {64,128} for MHA/GQA/MQA and d=128, d_r=192 for MLA; decode uses d in {64,128} and L=4096 for MHA/GQA/MQA, and d_c=512, d_r=64 for MLA.Source: Table 3 · p. 16, Table 4 · p. 16, C Additional Evaluation · p. 16
  • The benchmark notation uses X for activations, W for weights, Q/K/V for query/key/value, s for prefill length, b for decode batch size, and L for attention context length.For MLA, d_c and d_r denote content and positional/rotary dimensions, while Q_pe and K_pe are positional-encoding components.Source: C Additional Evaluation · p. 16
  • Welder performs better with Cleave-transformed graphs, though its kernels remain slower than Cleave-optimized kernels.Figure 14 reports normalized performance relative to Cleave for four subgraphs and different shapes.Source: C Additional Evaluation · p. 17, Figure 14 · p. 17
  • Algebraic transformations generally worsen TorchInductor performance because they can obscure fusion patterns and prevent it from fusing all kernels.The authors identify this as a challenge for applying algebraic transformation with a rule-based compiler.Source: C Additional Evaluation · p. 17
  • Welder cannot fuse attention or similar consecutive-reduction workloads at sequence lengths of 512 or larger, causing significant performance degradation.Figure 15 compares runtime as sequence length increases for self-attention, MatMul + MatMul, and MatMul + RowSum.Source: C Additional Evaluation · p. 17, Figure 15 · p. 17
Loading 2610.07742v1…