Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Yifan Zhao, Egan Johnson, Prasanth Chatarasi, Vikram Adve, Sasa Misailovic

TL;DR

Neptune tackles the challenge of fast fusion for sequences of reduction operators, such as attention, on GPUs and introduces a repair-term paradigm with two fusion instantiations to handle complex loop-carried dependencies.It proposes a template-guided optimization pipeline that combines scheduling-level fusion with a tile-based backend by translating loop-scalar IR to tile IR through tensorization, enabling high-performance, attention-like kernels.Empirical results across four GPU architectures and 10 attention-based operators show Neptune achieving a geomean speedup of $1.35\times$ over the next-best baselines and outperforming manually optimized libraries in many configurations.The work demonstrates a practical, compiler-native path to automatic, high-performance fusion for complex reductions and suggests a broader integration of such advanced fusion into tensor compilers.

Abstract

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.

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

TL;DR

Neptune tackles the challenge of fast fusion for sequences of reduction operators, such as attention, on GPUs and introduces a repair-term paradigm with two fusion instantiations to handle complex loop-carried dependencies.It proposes a template-guided optimization pipeline that combines scheduling-level fusion with a tile-based backend by translating loop-scalar IR to tile IR through tensorization, enabling high-performance, attention-like kernels.Empirical results across four GPU architectures and 10 attention-based operators show Neptune achieving a geomean speedup of $1.35\times$ over the next-best baselines and outperforming manually optimized libraries in many configurations.The work demonstrates a practical, compiler-native path to automatic, high-performance fusion for complex reductions and suggests a broader integration of such advanced fusion into tensor compilers.

Abstract

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 over the next best alternative, demonstrating its effectiveness for deep learning workloads.

Paper Structure

This paper contains 28 sections, 4 theorems, 19 equations, 19 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Motivation
  3. System Overview
  4. Template-Guided Optimizer
  5. Rolling Update
  6. Reduction Privatization
  7. Split-K Update
  8. Translating from Loop-scalar IR to Tile IR
  9. Implementation
  10. Experimental Methodology
  11. Evaluation
  12. Neptune vs. Tensor Compilers
  13. Neptune vs. Manually Optimized Libraries
  14. Scalability Test
  15. Ablation Studies
  16. ...and 13 more sections

Key Result

theorem 1

If a function $h$ satisfies the following two conditions, then it is tag-updating:

Figures (19)

  • Figure 1: An example program requiring reduction fusion, incorrect result using naive loop fusion, and a correctly fused program, shown on the first row. The second row ("materialized example") shows the values of each matrix in the programs when applied to a concrete inp matrix. xmax and xsum are each shown four times since their values vary over j iterations.
  • Figure 2: The main components of Neptune.
  • Figure 3: The result of applying split-k update to Figure \ref{['fig:example-orig']}.
  • Figure 4: An example program in Neptune's tile IR. Neptune translates the program in loop-scalar IR (above) to the program in tile IR (below).
  • Figure 5: Performance of Neptune kernels vs. kernels by other tensor compilers. The y-axis shows relative performance of all compilers normalized to Neptune for each setup. The plots for the remaining benchmarks are in Appendix \ref{['sec:appendix-eval']}.
  • ...and 14 more figures

Theorems & Definitions (5)

  • definition 1
  • theorem 1
  • theorem 2
  • theorem 3
  • theorem 4: Correctness of the Repaired Program