Cleave：通过解耦代数搜索与算子调度实现张量程序优化的扩展

中文摘要 AI 辅助

诸如FlashAttention和FlashDecoding等优化内核对于加速当今的大模型至关重要。其中大多数是由专家手工编写的，因为现有的机器学习编译器无法达到其效率。生成此类内核需要融合具有多个归约的计算，这既需要对计算图进行代数变换，也需要对变换后的图进行算子调度。不幸的是，联合搜索这两者会产生一个过于庞大的搜索空间，难以导航。我们提出了Cleave，一个基于符号解耦的机器学习编译器：Cleave通过在具有符号形状的图上执行超级优化来发现变换，然后在具体形状上调度每个生成的图。将形状表示为符号使得等价性检查变得廉价，并允许一个新的Split算子（具有符号分割计数）沿归约维度进行并行化。Cleave的调度器通过迭代分块和水平融合来融合具有多个归约的图。对常见LLM子图的评估表明，Cleave生成的内核比最佳基线快高达2.8倍（平均快1.6倍），并且与Mirage相比，编译时间平均减少5.9倍。对于从生产服务轨迹捕获的动态工作负载，Cleave仅编译每个算子一次，并在FlashInfer的手写FA2和FA3后端上实现了1.4倍和1.7倍的几何平均加速。Cleave的代码可在以下网址获取：此https URL

英文摘要

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

发表机构

• New York University（纽约大学）
• Cornell University（康奈尔大学）

机构由 AI 辅助整理，请以论文原文为准。