Blockbuster, Part 1: Block-level AI Operator Fusion

TL;DR

Blockbuster introduces a block-program representation and a rule-based fusion engine to optimize inference AI workloads on memory-tiered hardware. By explicitly modeling block-level data movement and applying a structured set of substitution rules, the framework discovers aggressive fusions, including rediscovery of Flash Attention and mega-kernel shapes like LayerNorm+Matmul and RMSNorm+FFN-SwiGLU. A separate fusion-candidate selection component guides which fused versions to deploy for the full program, enabling scalability to large Transformer blocks. The approach emphasizes memory-locality-aware fusion rather than exhaustive search, offering practical improvements and a path toward further optimization and numerical-safety considerations.

Abstract

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.