FlashPDE: A Drop-in Fused Triton Operator Library for Neural PDE Solvers
2026-07-20 • Machine Learning
Machine LearningMathematical Software
AI summaryⓘ
The authors created FlashPDE, a tool that helps neural networks solve physics problems involving equations on grids more efficiently. Their tool combines several calculations into one step, saving memory and speeding up processing on GPUs compared to older methods. FlashPDE works with many types of equations and fits smoothly into the PyTorch machine-learning framework. Tests showed it uses much less memory and runs faster while giving results similar to standard approaches.
Physics-Informed Neural NetworksPartial Differential EquationsFinite Difference MethodTriton KernelsAutomatic DifferentiationPyTorchGPU AccelerationElliptic PDEParabolic PDENavier-Stokes Equations
Authors
Peiyu Zang, Bosen Xie, Ruoxiang Xu, Yongqiang Cai
Abstract
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 torch.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.