Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUs

Understand

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.
• This 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.