Blockbuster，第一部分：块级 AI 运算融合

Blockbuster 是一个用于 AI 运算融合的推理程序框架。Blockbuster 框架与任何具有分层内存层次结构的多处理器架构兼容，包括 GPU、多核 CPU 和一些 AI 加速器芯片。它包含一种用于 AI 工作负载的图表示方法，称为块程序 (block program)，该程序显式地建模了数据块在内存层次之间的移动方式。它还包括一种运算融合程序，由候选选择算法和融合算法组成后者融合每个单独的候选者——这种两种算法的结构使 Blockbuster 特别适用于大型 AI 程序。当前论文聚焦于融合算法，这是一种基于规则的技术。尽管文献中充满了之前的基于规则的融合算法，但我们算法的独特之处在于其直接建模内存层次之间的数据移动， resulting in uniquely powerful fusion results。作为第一次 sanity check，我们展示了我们的算法如何自动复现众所周知的 Flash Attention kernel。随后，我们展示了我们方法的强大之处，通过融合 LayerNorm 与矩阵乘法以及 RMSNorm 与 FNN-SwiGLU——后者涉及将三个矩阵乘法、一个哈达德积、一个约简和几个元素运算融合到单个巨型 kernel 中。