arXiv:2510.08726cs.PLcs.LG2025-10被引 6

提出新融合方法,让GPU高效运行复杂注意力计算。

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

  • 打破依赖关系,用代数修正保证结果正确。
  • 在十项注意力任务上平均快1.35倍,最高达3.32倍。
  • 适合需要极致性能的深度学习推理与训练场景。

算子融合是深度学习优化的关键技术,通过合并多个算子提升数据复用、减少全局内存传输。然而,现有张量编译器难以融合含循环依赖的复杂归约计算,如注意力机制。本文提出Neptune,一种面向序列归约算子的高级算子融合编译器。Neptune采用新方法:主动打破部分依赖关系,并构建代数修正表达式以确保内核输出正确。将Neptune应用于基础注意力代码和高层调度模板,生成的算子等价于FlashAttention与FlashDecoding。在十个基于注意力的基准测试中,Neptune超越Triton、TVM、FlexAttention等现有编译器,包括基于Triton的FlashAttention实现。在四款不同厂商(NVIDIA与AMD)的GPU上,Neptune生成内核平均加速1.35倍,其中NVIDIA GPU最高达2.65倍,AMD GPU最高达3.32倍,充分验证其对深度学习工作负载的有效性。

原文摘要 · Abstract (English)

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.

算子融合注意力机制GPU优化编译器

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。