Neptune: Advanced ML Operator Fusion for Locality and Parallelism on GPUs
summary
Neptune is a tensor compiler that implements advanced operator fusion by deliberately breaking loop-carried dependencies and synthesizing algebraic repair terms (generated with a SymPy-based solver) so fused kernels produce correct results. Two fusion variants map to known algorithms: rolling update yields FlashAttention-style kernels and split K update yields FlashDecoding-like kernels. The pipeline separates loop IR → tile IR, reuses Triton-style tile backends, and evaluates 10 attention variants plus other ops across four GPUs, reporting 1.35× geometric average speedup, up to 2.65× (NVIDIA) and 3.32× (AMD), and up to 6× on L2 norm. Numerical accuracy matches baselines.
AI-written summary · about
chapters (16)
- 0:00Neptune and operator fusion
- 1:26Attention escapes compilers
- 2:31Advanced fusion for attention
- 3:14Breaking programs for fusion
- 4:15Neptune's two-stage pipeline
- 5:21Flash attention fusion challenges
- 7:04Repair terms and correctness
- 8:29Fusion with scalar expressions
- 9:53Performance gains overview
- 10:35Split K update variant
- 11:41Loop and tile IR
- 12:45Math to hardware performance
- 13:26Evaluation setup
- 14:49Results across GPUs
- 15:31Conclusion and future
- 15:52Q&A
related talks
24mPortable High‑Performance LLM Serving: A Triton Backend for...
42mAnima Anadkumar, Principal Scientist, Amazon Web Services, Endowed Professor, CalTech
22mMonoNN: Enabling a New Monolithic Optimization Space for Neural Network Inference...
22mOptimizing Large MoE Inference on NVIDIA Blackwell: NVFP4, ADP, and DualPipe Strat...
14mKairox: Adaptive GPU-CPU Hybrid LLM Inference via Online Neuron Balancing
06mGlobalFusion: A Global Attentional Deep Learning Framework for Multisensor Information Fusion
29mBuild Customize and Deploy LLMs At-Scale on Azure with NVIDIA NeMo
27mAdvancements in optimizing ML Inference at CERN