Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUs
This work addresses a bottleneck in optimizing deep learning workloads on GPUs, offering incremental improvements for performance in domains like attention-based models.
The paper tackled the problem of fusing complex reduction computations with loop-carried dependencies in deep learning operators, such as attention mechanisms, and introduced Neptune, a tensor compiler that achieved an average speedup of 1.35x over existing compilers on ten attention-based benchmarks across multiple GPU architectures.
Operator fusion has become a key optimization for deep learning, which combines multiple deep learning operators to improve data reuse and reduce global memory transfers. However, existing tensor compilers struggle to fuse complex reduction computations involving loop-carried dependencies, such as attention mechanisms. The paper introduces Neptune, a tensor compiler for advanced operator fusion for sequences of reduction operators. Neptune presents a new approach for advanced operator fusion, which intentionally breaks some existing dependencies and compensates by constructing algebraic correction expressions that allow the kernel to produce the correct result. On ten attention-based benchmarks, Neptune, starting from simple attention code and a high-level scheduling template, outperforms existing compilers like Triton, TVM, and FlexAttention, including Triton-based implementations of FlashAttention. Across four different GPU architectures from NVIDIA and AMD, Neptune-generated kernels have average speedup of $1.35\times$ over the next best alternative, demonstrating its effectiveness for deep learning workloads.