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. Applying Neptune's advanced operator fusion to a plain attention operator generates operators equivalent to FlashAttention and FlashDecoding. On ten attention-based benchmarks, Neptune, starting from a plain 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 an average speedup of $1.35\times$ over the next best alternative, with up to $2.65\times$ speedup on Nvidia GPUs and up to $3.32\times$ on AMD GPUs, demonstrating its effectiveness for deep learning workloads.

翻译：算子融合已成为深度学习的关键优化技术，它通过合并多个深度学习算子来提高数据复用率并减少全局内存传输。然而，现有张量编译器难以处理涉及循环携带依赖的复杂归约计算（如注意力机制）。本文提出Neptune，一种针对归约算子序列实现高级算子融合的张量编译器。Neptune提出了一种新型高级算子融合方法，该方法有意打破部分现有依赖关系，并通过构建代数校正表达式进行补偿，从而确保内核能生成正确结果。将Neptune的高级算子融合应用于朴素注意力算子，可生成等价于FlashAttention与FlashDecoding的算子。在十个基于注意力的基准测试中，Neptune从朴素注意力代码和高层次调度模板出发，其性能优于Triton、TVM及FlexAttention等现有编译器，包括基于Triton实现的FlashAttention。在来自NVIDIA与AMD的四种不同GPU架构上，Neptune生成的内核相较于次优方案平均加速达$1.35\times$，其中NVIDIA GPU最高加速$2.65\times$，AMD GPU最高加速$3.32\times$，验证了其对深度学习工作负载的有效性。