通过显式建模数据移动，实现大模型推理中算子融合的高效优化
Blockbuster, Part 1: Block-level AI Operator Fusion
- 基于块级数据流动建模，设计可适配多架构的算子融合框架
- 自动融合层归一化与矩阵乘，实现三重矩阵乘+哈达玛积的超大核
- 适用于大规模AI程序，尤其适合有复杂内存层级的GPU/加速器
Blockbuster 是一个面向推理程序的 AI 算子融合框架，兼容具有分层内存结构的多处理器架构（如 GPU、多核 CPU 及部分 AI 加速芯片）。其核心是基于图的「块程序」表示，显式建模数据块在不同内存层级间的移动。框架包含候选选择算法与融合算法两部分，其中融合算法为规则驱动。不同于以往方法，该算法直接建模跨内存层级的数据流动，实现更优融合效果。初步验证中，算法自动重构出著名的 Flash Attention 内核；进一步展示其强大能力：将 LayerNorm 与矩阵乘融合，并成功将 RMSNorm 与 FNN-SwiGLU 融合——后者涉及三个矩阵乘、一个哈达玛积、一次归约及若干逐元素操作，合并为单一超大内核。
原文摘要 · Abstract (English)
Blockbuster is a framework for AI operator fusion in inference programs. The Blockbuster framework is compatible with any multiprocessor architecture that has a tiered memory hierarchy, including GPUs, multi-core CPUs, and some AI accelerator chips. It includes a graph-based representation for AI workloads, called a block program, which explicitly models how blocks of data move between the memory tiers. It also includes an operator fusion procedure, which is made up of a candidate selection algorithm and a fusion algorithm that fuses each individual candidate - this two-algorithm structure makes Blockbuster especially suitable for large AI programs. The current paper focuses on the fusion algorithm, which is a rule-based technique. While the literature is full of previous rule-based fusion algorithms, what sets our algorithm apart is its direct modeling of data movement between memory tiers, resulting in uniquely powerful fusion results. As a first sanity check, we demonstrate how our algorithm automatically rediscovers the well-known Flash Attention kernel. Then, we demonstrate the real power of our approach by fusing LayerNorm with matrix multiplication and RMSNorm with FNN-SwiGLU - the latter involves fusing three matrix multiplications, a Hadamard product, a reduction, and a few elementwise operations into a single mega-kernel.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。