跳到主要内容
返回时间线
arXiv来源发表:

Cleave 用符号化解耦把代数变换与算子调度分开,在 LLM 子图上生成比最佳基线快 2.8 倍的融合内核

核心概要

Cleave 是一个基于符号化解耦的 ML 编译器:它在符号形状图上做超优化发现代数等价变换,再在具体形状上对每个结果图做调度,通过迭代分块与水平融合来融合含多个归约的图;在常见 LLM 子图上生成的内核比最佳基线快至多 2.8 倍(平均 1.6 倍),编译时间比 Mirage 平均减少 5.9 倍。

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

深度剖析

Cleave 提出符号化解耦:代数阶段在符号形状图上搜索等价变换,调度阶段再把符号绑定到具体形状并调度每个结果图,使代数搜索不再枚举循环结构与分块大小,调度器也不再推理代数等价性。 此前的调度型编译器只优化循环、分块与内存布局,代数优化器只搜索等价图,Mirage 则联合搜索两者导致搜索空间过大;Cleave 沿“变换有效性只取决于图本身”这条界线把两者分开。 论文以 FlashAttention(把 softmax 的除法移到 V-Matmul 之后)与 Split-K 两个例子论证该分离的合理性,并在 9 个子图上做消融:禁用代数变换性能下降 44%,不生成迭代分块调度下降 40%,两者同时禁用下降 47%。

Cleave 引入带符号切分数的 Split 算子,让超优化自行决定在哪里切分归约维以及如何合并部分结果,具体切分值由调度器稍后绑定;符号图还让等价性检查可以用满足形状约束的小规模实例做概率测试,从而降低测试成本。 Split-K 原本以具体切分数表述,看起来必须先定调度才能做变换;把切分数变成又一个符号后,代数阶段只需决定切分位置与合并方式。 论文用基于可归约性的剪枝限制哪些维度可被切分或归约,并在超优化结束时仍在用户图的原始形状上做概率测试以保持正确性保证;解码设定下启用 Split-K 比不启用最多快数倍(原文该处数值在正文中缺失)。

Cleave 的调度器用迭代分块把归约维切成小块、逐块累加,并用水平融合把共享上游生产者的多个归约放进同一循环,从而在不物化完整中间张量的前提下融合多归约图。 现有调度型编译器只对输出维分块,即使拿到变换后的图也会把完整中间张量物化到 SMEM,因而对注意力只能支持较短序列。 论文以注意力为例说明 Rowsum 与第二个 Matmul 共享 Exp 因而被分到同一组;Welder 在序列长度达到 512 后无法融合连续归约子图,而 Cleave 仍能融合。

在动态负载上,Cleave 对每个算子只编译一次,代码生成把动态轴的范围保留为符号与运行时循环边界,一个内核即可服务该算子所有捕获形状。 调度需要具体维度来选分块形状、线程数与缓冲区大小,Cleave 用代表性具体形状调度、再在代码生成时把动态轴符号化。 九个分页与不规则注意力算子在 FlashInfer-Bench 捕获的 303 个形状上各编译一次、无需重编译,相对 FlashInfer 手写 FA2 后端几何平均快 1.4 倍、相对 FA3 快 1.7 倍;仅在两个 MLA 算子上慢于 FA3,因为 FA3 在单个协作内核内合并部分结果而 Cleave 需要第二次启动。

启示与展望

该结果面向以归约型算子为主的 LLM 子图与 transformer 层,在 A100-SXM4-40GB 与 H200-SXM5-141GB 上、FP16 输入与 FP32 累加、FP16 输出的设定下测得;动态算子部分针对分页与不规则注意力,形状取自 FlashInfer-Bench 捕获的 303 个生产轨迹形状。可直接复用的对象是编译器与内核工程实践:把代数搜索与调度解耦、用符号形状降低等价性检查成本、用迭代分块与水平融合处理多归约图。论文还显示,把 Cleave 的代数变换结果交给 Welder 能提升其性能,说明该变换阶段可独立为其他调度型编译器提供输入。

正文中若干数值在加载文本里缺失(例如解码设定下 Split-K 带来的加速倍数、动态算子相对 FA2/FA3 的逐算子加速区间、以及结论中编译时间与加速比的倍数),因此这些具体幅度只能按论文给出的几何平均与总体区间理解。可归约性标注目前需要用户提供,论文指出它可由数据依赖分析自动化,但自动化后的效果未在本文评估。代数变换对基于规则的编译器(torch.compile)反而降低了性能,论文将其归因于变换后编译器无法识别原有融合模式,这提示该变换阶段与不同编译器的配合方式仍是一个开放问题。此外,动态算子中两个 MLA 算子慢于 FA3,原因被归为 FA3 在单个协作内核内合并部分结果而 Cleave 需要第二次启动,这一结构性差异在何种负载下会改变结论值得继续观察。

来源