arXiv cs.AI· David Pissarra, Jinkun Lin, Haitian Jiang, Aurojit Panda, Jinyang Li·· 10 天前AI 评分27
Cleave：通过解耦代数搜索与算子调度扩展张量程序优化
Cleave: Scaling Tensor Program Optimization via Decoupled Algebraic Search and Operator Scheduling
AI 导读
机器学习编译器 Cleave 通过符号化解耦，在符号形状计算图上做超优化发现变换，再对具体形状调度，生成的 kernel 比最优基线快最高 2.8 倍（平均 1.6 倍），编译时间比 Mirage 平均减少 5.9 倍。在生产服务动态负载上，Cleave 对每个算子仅编译一次，相对 FlashInfer 手写 FA2 和 FA3 后端分别取得 1.4 倍和 1.7 倍的几何平均加速。代码已开源。
正文 · 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 的代码可在以下地址获取：this https URL
| 主题： | 编程语言（cs.PL）；人工智能（cs.AI） |
| 引用为： | arXiv:2610.07742 [cs.PL] |
| （或 arXiv:2610.07742v1 [cs.PL] 用于此版本） | |
| https://doi.org/10.48550/arXiv.2610.07742 由 DataCite 颁发的 arXiv DOI（待注册） |
提交历史
来自：David Pissarra [查看电子邮件]
[v1]
2026年10月6日星期二 04:37:22 UTC（722 KB）
来源：arXiv cs.AI · arxiv.org