LLM2D
blockbuster,第1部分:块级AI操作融合
Blockbuster, Part 1: Block-level AI Operator Fusion
作者: Ofer Dekel
发布日期: 5/14/2025
arXiv ID: oai:arXiv.org:2505.07829v1

摘要

arXiv:2505.07829v1 Announce Type: 跨越式 摘要:Blockbuster是一种用于推理程序的AI操作融合框架。Blockbuster框架适用于具有分层内存体系结构的任何多处理器架构,包括GPU、多核CPU和某些AI加速器芯片。它包含一种称为块程序的基于图的表现形式,该表现形式明确地表示数据块在内存层次之间的移动方式。它还包含一种操作融合程序,该程序由候选选择算法和一个将每个单独候选操作融合的融合算法组成——这种双算法结构使Blockbuster特别适合大型AI程序。当前的论文集中在融合算法上,该算法是一种基于规则的技术。尽管文献中充满了基于规则的融合算法,但我们的算法的独特之处在于它直接对内存层次之间的数据移动进行了建模,从而产生了独特而强大的融合结果。作为第一次合理性检查,我们展示了我们的算法如何自动重新发现广为人知的Flash Attention内核。然后,我们通过将LayerNorm与矩阵乘法融合,以及将RMSNorm与FNN-SwiGLU融合,展示了我们方法的实际威力——后者涉及将三个矩阵乘法、哈达玛积、归约和几种元素操作融合为一个巨核。