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

arXiv: https://arxiv.org/abs/2510.08726 · HTML: https://mlsanity.com/paper/2510.08726/

## Key facts

- Cited by 4 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 58 references that are in the map.
- First posted on arXiv 9 Oct 2025.

## Builds on (most cited first, 16 of 16)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025)
- [Ansor: Generating High-Performance Tensor Programs for Deep Learning](https://mlsanity.com/paper/2006.06762/) (2020)
- [Flex Attention: A Programming Model for Generating Optimized Attention Kernels](https://mlsanity.com/paper/2412.05496/) (2024)
- [DNNFusion: Accelerating Deep Neural Networks Execution with Advanced Operator Fusion](https://mlsanity.com/paper/2108.13342/) (2021)
- [TileLang: A Composable Tiled Programming Model for AI Systems](https://mlsanity.com/paper/2504.17577/) (2025)
- [Mirage: A Multi-Level Superoptimizer for Tensor Programs](https://mlsanity.com/paper/2405.05751/) (2024)
- [TensorIR: An Abstraction for Automatic Tensorized Program Optimization](https://mlsanity.com/paper/2207.04296/) (2022)
- [ThunderKittens: Simple, Fast, and Adorable AI Kernels](https://mlsanity.com/paper/2410.20399/) (2024)
- [Tensor Program Optimization with Probabilistic Programs](https://mlsanity.com/paper/2205.13603/) (2022)
- [Equality Saturation for Tensor Graph Superoptimization](https://mlsanity.com/paper/2101.01332/) (2021)
- [UNIT: Unifying Tensorized Instruction Compilation](https://mlsanity.com/paper/2101.08458/) (2021)
- [Automatic Kernel Generation for Volta Tensor Cores](https://mlsanity.com/paper/2006.12645/) (2020)

## Built on by (most cited first, 4 of 4)

- [Hexagon-MLIR: An AI Compilation Stack For Qualcomm's Neural Processing Units (NPUs)](https://mlsanity.com/paper/2602.19762/) (2026)
- [Nautilus: An Auto-Scheduling Tensor Compiler for Efficient Tiled GPU Kernels](https://mlsanity.com/paper/2604.14825/) (2026)
- [Ada-MK: Adaptive MegaKernel Optimization via Automated DAG-based Search for LLM Inference](https://mlsanity.com/paper/2605.11581/) (2026)
- [Cleave: Scaling Tensor Program Optimization via Decoupled Algebraic Search and Operator Scheduling](https://mlsanity.com/paper/2610.07742/) (2026)

## Abstract (arXiv)

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.