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

• URL: http://arxiv.org/abs/2510.08726v1
• Date: Thu, 09 Oct 2025 18:33:52 GMT
• Title: Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUs
• Authors: Yifan Zhao, Egan Johnson, Prasanth Chatarasi, Vikram Adve, Sasa Misailovic,
• Abstract summary: The paper introduces Neptune, a tensor compiler for advanced operator fusion for sequences of reduction operators.<n>N Neptune presents a new approach for advanced operator fusion, which intentionally breaks some existing dependencies.<n>On ten attention-based benchmarks, Neptune outperforms existing compilers like Triton, TVM, and FlexAttention.
• Score: 7.912341462809207
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• CUDA Agent: Large-Scale Agentic RL for High-Performance CUDA Kernel Generation [51.72529978689561]
  Agent is a large-scale agentic reinforcement learning system that develops kernel expertise through three components.<n>Agent delivers 100%, 100%, and 92% faster rate over torchcompile on KernelBench.
  arXiv  Detail & Related papers  (2026-02-27T18:58:05Z)
• Eliminating Multi-GPU Performance Taxes: A Systems Approach to Efficient Distributed LLMs [61.953548065938385]
  We introduce the ''Three Taxes'' (Bulk Synchronous, Inter- Kernel Data Locality, and Kernel Launch Overhead) as an analytical framework.<n>We propose moving beyond the rigid BSP model to address key inefficiencies in distributed GPU execution.<n>We observe a 10-20% speedup in end-to-end latency over BSP-based approaches.
  arXiv  Detail & Related papers  (2025-11-04T01:15:44Z)
• AutoTriton: Automatic Triton Programming with Reinforcement Learning in LLMs [87.8306870967343]
  We introduce AutoTriton, the first model dedicated to Triton programming powered by reinforcement learning (RL)<n>AutoTriton performs supervised fine-tuning (SFT) to be equipped with essential Triton programming expertise using a high-quality data gathering pipeline.<n> Experiments across five evaluation channels of TritonBench and KernelBench illustrate that our 8B model AutoTriton achieves performance comparable to mainstream large models.
  arXiv  Detail & Related papers  (2025-07-08T05:38:24Z)
• Hexcute: A Tile-based Programming Language with Automatic Layout and Task-Mapping Synthesis [8.742879659920643]
  Hexcute is a tile-based programming language that exposes shared memory and register abstractions to enable fine-grained optimization for mixed-type operators.<n>It automates layout and task mapping synthesis with a novel type-inference-based algorithm.<n>Our evaluation shows that Hexcute generalizes to a wide range of DL operators, achieves 1.7-11.28$times$ speedup over existing DL compilers for mixed-type operators, and brings up to 2.91$times$ speedup in the end-to-end evaluation.
  arXiv  Detail & Related papers  (2025-04-22T19:01:28Z)
• TritonBench: Benchmarking Large Language Model Capabilities for Generating Triton Operators [59.625889531331815]
  Triton is a high-level Python-like language designed for building efficient GPU kernels.<n>Despite advances in large language models (LLMs) for conventional code generation, these models struggle to generate accurate, performance-optimized Triton code.<n>In this work, we introduce TritonBench, the first comprehensive benchmark for Triton operator generation.
  arXiv  Detail & Related papers  (2025-02-20T17:21:27Z)
• An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networks [0.5737287537823071]
  Rotation equivariant graph neural networks yield state of the art performance on spatial deep learning tasks.<n>Key to these models is the Clebsch-Gordon (CG) tensor product, a kernel that contracts two dense feature vectors with a highly-structured sparse tensor to produce a dense output vector.<n>We introduce a GPU sparse kernel generator for the CG tensor product that provides significant speedups over the best existing open and closed-source implementations.
  arXiv  Detail & Related papers  (2025-01-23T08:20:47Z)
• TensorGRaD: Tensor Gradient Robust Decomposition for Memory-Efficient Neural Operator Training [91.8932638236073]
  We introduce textbfTensorGRaD, a novel method that directly addresses the memory challenges associated with large-structured weights.<n>We show that sparseGRaD reduces total memory usage by over $50%$ while maintaining and sometimes even improving accuracy.
  arXiv  Detail & Related papers  (2025-01-04T20:51:51Z)
• FusionLLM: A Decentralized LLM Training System on Geo-distributed GPUs with Adaptive Compression [55.992528247880685]
  Decentralized training faces significant challenges regarding system design and efficiency. We present FusionLLM, a decentralized training system designed and implemented for training large deep neural networks (DNNs) We show that our system and method can achieve 1.45 - 9.39x speedup compared to baseline methods while ensuring convergence.
  arXiv  Detail & Related papers  (2024-10-16T16:13:19Z)
• Flash-LLM: Enabling Cost-Effective and Highly-Efficient Large Generative Model Inference with Unstructured Sparsity [12.663030430488922]
  We propose Flash-LLM for enabling low-cost and highly-efficient large generative model inference on high-performance Cores. At SpMM kernel level, Flash-LLM significantly outperforms the state-of-the-art library, i.e., Sputnik and SparTA by an average of 2.9x and 1.5x, respectively.
  arXiv  Detail & Related papers  (2023-09-19T03:20:02Z)
• INR-Arch: A Dataflow Architecture and Compiler for Arbitrary-Order Gradient Computations in Implicit Neural Representation Processing [66.00729477511219]
  Given a function represented as a computation graph, traditional architectures face challenges in efficiently computing its nth-order gradient. We introduce INR-Arch, a framework that transforms the computation graph of an nth-order gradient into a hardware-optimized dataflow architecture. We present results that demonstrate 1.8-4.8x and 1.5-3.6x speedup compared to CPU and GPU baselines respectively.
  arXiv  Detail & Related papers  (2023-08-11T04:24:39Z)
• Harnessing Deep Learning and HPC Kernels via High-Level Loop and Tensor Abstractions on CPU Architectures [67.47328776279204]
  This work introduces a framework to develop efficient, portable Deep Learning and High Performance Computing kernels. We decompose the kernel development in two steps: 1) Expressing the computational core using Processing Primitives (TPPs) and 2) Expressing the logical loops around TPPs in a high-level, declarative fashion. We demonstrate the efficacy of our approach using standalone kernels and end-to-end workloads that outperform state-of-the-art implementations on diverse CPU platforms.
  arXiv  Detail & Related papers  (2023-04-25T05:04:44Z)
• Tensor Slicing and Optimization for Multicore NPUs [2.670309629218727]
  This paper proposes a compiler optimization pass for Multicore NPUs, called Slicing Optimization (TSO) TSO identifies the best tensor slicing that minimizes execution time for a set of CNN models. Results show that TSO is capable of identifying the best tensor slicing that minimizes execution time for a set of CNN models.
  arXiv  Detail & Related papers  (2023-04-06T12:03:03Z)
• Hidet: Task Mapping Programming Paradigm for Deep Learning Tensor Programs [11.338285393619042]
  We propose to embed the scheduling process into tensor programs and use dedicated mappings, called task mappings, to define the computation assignment and ordering. With the proposed paradigm, we implement a deep learning compiler - Hidet.
  arXiv  Detail & Related papers  (2022-10-18T05:32:13Z)
• VersaGNN: a Versatile accelerator for Graph neural networks [81.1667080640009]
  We propose textitVersaGNN, an ultra-efficient, systolic-array-based versatile hardware accelerator. textitVersaGNN achieves on average 3712$times$ speedup with 1301.25$times$ energy reduction on CPU, and 35.4$times$ speedup with 17.66$times$ energy reduction on GPU.
  arXiv  Detail & Related papers  (2021-05-04T04:10:48Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.