Blockbuster, Part 1: Block-level AI Operator Fusion

TL;DR

Blockbuster introduces a flexible framework for AI operator fusion across various hardware architectures, utilizing a graph-based workload representation and a novel rule-based fusion algorithm that models data movement between memory tiers.

Contribution

It presents a new rule-based fusion algorithm that explicitly models memory tier data movement, enabling powerful fusion of complex AI operations.

Findings

Successfully rediscovered Flash Attention kernel

Fused LayerNorm with matrix multiplication

Fused multiple operations into a single mega-kernel

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.