Title: Forge: Compiling a Unified Abstraction into Scalable Kernels for Linear Attention
Primary area: infrastructure, software libraries, hardware, systems, etc.
Keywords: linear attention, efficiency, compiler, kernels

Abstract: The quadratic complexity of softmax attention poses a major bottleneck for long-context modeling, motivating a surge of linear attention variants with linear complexity. Unlike softmax attention, which benefits from optimized kernels, linear attention lacks general-purpose, hardware-efficient support and scalable distributed implementations. We introduce Forge, a domain-specific compiler that automates the generation of high-performance, scalable kernels for a wide range of linear attention models directly from high-level PyTorch code. At its core, Forge employs an intuitive programming abstraction that decomposes any linear attention algorithm into three canonical phases: intra-chunk computation, inter-chunk state propagation, and output merging. This unified abstraction enables Forge to perform domain-specific optimizations, automatically generating kernels that fuse computation and communication at a fine-grained tile level and eliminating. Our evaluation demonstrates that Forge combines programmability with performance: a wide range of linear attention variants can be implemented in just a few dozen lines of code, while the generated kernels deliver 1.01x-4.9x the performance of sate-of-the-art expert-optimized library and achieve near-linear scalability in a 128-GPUs distributed system, surpassing state-of-the-art distributed baselines by up to 7.2x.

TL;DR: Forge is a domain-specific compiler that generates scalable kernels for diverse linear attention variants, from a simple unified abstraction, providing comparable or even better performance than expert-tuned libraries.