一句话：提出张量编译器Neptune，通过打破依赖并构造代数修正表达式，实现归约算子序列的高级融合与高效注意力。
现有张量编译器难以融合带循环携带依赖的复杂归约，如注意力机制。Neptune提出高级算子融合：有意打破部分依赖，再构造代数修正表达式补偿，使kernel仍输出正确结果。将Neptune用于普通注意力算子，可生成等价于FlashAttention与FlashDecoding的算子。在十个注意力基准上，从普通注意力代码和高层调度模板出发，Neptune优于Triton、TVM、FlexAttention，包括基于Triton的FlashAttention实现；在四种NVIDIA与AMD GPU上平均较次优方案加速1.35×，NVIDIA最高2.65×，AMD最高3.32×。
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× over the next best alternative, with up to 2.65 × speedup on Nvidia GPUs and up to 3.32 × on AMD GPUs, demonstrating its effectiveness for deep learning workloads.
DOI 原文 · arXiv · PDF（开放获取） · DBLP
@article{abs-2510-08726,
title = {Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUs},
author = {Yifan Zhao and Egan Johnson and Prasanth Chatarasi and Vikram Adve and Sasa Misailovic},
journal = {Proceedings of the ACM on Programming Languages},
volume = {10},
number = {PLDI},
year = {2026},
doi = {10.1145/3808298}
}