# FlashPDE: A Drop-In Fused Triton Operator Library for Neural PDE Solvers

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 20 of its 36 references that are in the map.
- First posted on arXiv 20 Jul 2026.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Fourier Neural Operator for Parametric Partial Differential Equations](https://mlsanity.com/paper/2010.08895/) (2020)
- [DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators](https://mlsanity.com/paper/1910.03193/) (2019)
- [Learning to Simulate Complex Physics with Graph Networks](https://mlsanity.com/paper/2002.09405/) (2020)
- [Characterizing possible failure modes in physics-informed neural networks](https://mlsanity.com/paper/2109.01050/) (2021)
- [DeepXDE: A deep learning library for solving differential equations](https://mlsanity.com/paper/1907.04502/) (2019)
- [Scientific Machine Learning through Physics-Informed Neural Networks: Where we are and What's next](https://mlsanity.com/paper/2201.05624/) (2022)
- [DiffTaichi: Differentiable Programming for Physical Simulation](https://mlsanity.com/paper/1910.00935/) (2019)
- [FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving](https://mlsanity.com/paper/2501.01005/) (2025)
- [CAN-PINN: A Fast Physics-Informed Neural Network Based on Coupled-Automatic-Numerical Differentiation Method](https://mlsanity.com/paper/2110.15832/) (2021)
- [Physics-guided Convolutional Neural Network (PhyCNN) for Data-driven Seismic Response Modeling](https://mlsanity.com/paper/1909.08118/) (2019)
- [Spline-PINN: Approaching PDEs without Data using Fast, Physics-Informed Hermite-Spline CNNs](https://mlsanity.com/paper/2109.07143/) (2021)
- [Accelerated Training of Physics-Informed Neural Networks (PINNs) using Meshless Discretizations](https://mlsanity.com/paper/2205.09332/) (2022)
- [Efficient physics-informed neural networks using hash encoding](https://mlsanity.com/paper/2302.13397/) (2023)
- [Physics-informed ConvNet: Learning Physical Field from a Shallow Neural Network](https://mlsanity.com/paper/2201.10967/) (2022)
- [FastVPINNs: Tensor-Driven Acceleration of VPINNs for Complex Geometries](https://mlsanity.com/paper/2404.12063/) (2024)
- [OpenSBLI: A framework for the automated derivation and parallel execution of finite difference solvers on a range of computer architectures](https://mlsanity.com/paper/1609.01277/) (2016)
- [Accelerating High-Order Stencils on GPUs](https://mlsanity.com/paper/2009.04619/) (2020)

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

- [A Computational Comparison of Fourier Spectral Differentiation and Spatial Automatic Differentiation in Periodic Physics-Informed Neural Networks](https://mlsanity.com/paper/2609.02110/) (2026)

## Abstract (arXiv)

Physics-Informed Neural Networks (PINNs) solve PDEs by incorporating physical constraints into neural-network training, but large-scale problems are limited by automatic-differentiation memory overhead and inefficient execution of grid-based PDE operators. We present FlashPDE, a drop-in fused operator library for grid-based scientific machine learning. FlashPDE replaces fragmented PyTorch finite-difference execution with differentiable Triton kernels. Each operator integrates fused stencil evaluation, an analytic discrete-adjoint backward pass, and boundary-gradient correction within a unified PyTorch autograd Function interface. The library provides 14 differentiable PDE operators covering 17 configurations across 1D--3D elliptic, parabolic, and Navier--Stokes systems, while remaining independent of neural architectures and training strategies. Experiments on an NVIDIA A100 GPU show that FlashPDE reduces peak memory usage by up to 37.0x compared with coordinate-based automatic differentiation and reduces CUDA kernel launches by up to 3.5x compared with eager PyTorch finite-difference implementations. Across six representative PDE benchmarks, FlashPDE achieves up to 2.30x end-to-end time-to-solution speedup and up to 19.2x kernel-level acceleration while maintaining numerical agreement with PyTorch finite-difference references. FlashPDE provides a hardware-efficient execution layer that bridges differentiable PDE solvers and GPU-optimized numerical computation within the PyTorch ecosystem.