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

2026-07-20

AI Summary

This work addresses the high memory overhead and inefficient grid operations in large-scale physics-informed neural networks (PINNs) for solving partial differential equations, which stem from automatic differentiation. The authors propose the first differentiable PDE operator library that unifies forward stencil computation, analytically derived discrete adjoint-based backward propagation, and boundary gradient correction within a single torch.autograd.Function interface. Implemented with Triton for high-performance GPU kernels, the library supports 14 operators across 17 one- to three-dimensional elliptic, parabolic, and Navier–Stokes problems and is fully decoupled from network architecture and training strategies. Experiments on an A100 GPU demonstrate up to 37.0× lower peak memory usage, 3.5× fewer CUDA kernel launches, end-to-end speedups of up to 2.30×, and kernel-level acceleration of up to 19.2×, all while preserving numerical accuracy identical to a PyTorch reference implementation.

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.

Problem

• Physics-Informed Neural Networks
• automatic differentiation
• memory overhead
• PDE operators
• grid-based scientific computing

Innovation

• FlashPDE
• differentiable operators
• Triton kernels
• physics-informed neural networks
• finite-difference fusion