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

• URL: http://arxiv.org/abs/2607.18020v3
• Date: Fri, 24 Jul 2026 02:33:15 GMT
• Title: FlashPDE: A Drop-In Fused Triton Operator Library for Neural PDE Solvers
• Abstract summary: FlashPDE is a drop-in fused operator library for grid-based machine learning.<n>It replaces fragmented PyTorch finite-difference execution with differentiable Triton kernels.<n>It reduces peak memory usage by up to 37.0x compared with coordinate-based automatic differentiation.
• Score: 6.9997407868865364
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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 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.

Related papers

• Enabling Automatic Differentiation with Mollified Graph Neural Operators [73.52999622724101]
  We propose the mollified graph neural operator ($m$GNO), the first method to leverage automatic differentiation and compute exact gradients on arbitrary geometries.<n>For a PDE example on regular grids, $m$GNO paired with autograd reduced the L2 relative data error by 20x compared to finite differences.<n>It can also solve PDEs on unstructured point clouds seamlessly, using physics losses only, at resolutions vastly lower than those needed for finite differences to be accurate enough.
  arXiv  Detail & Related papers  (2025-04-11T06:16:30Z)
• Mechanistic PDE Networks for Discovery of Governing Equations [52.492158106791365]
  We present Mechanistic PDE Networks, a model for discovery of partial differential equations from data.<n>The represented PDEs are then solved and decoded for specific tasks.<n>We develop a native, GPU-capable, parallel, sparse, and differentiable multigrid solver specialized for linear partial differential equations.
  arXiv  Detail & Related papers  (2025-02-25T17:21:44Z)
• Transolver++: An Accurate Neural Solver for PDEs on Million-Scale Geometries [67.63077028746191]
  Transolver++ is a highly parallel and efficient neural solver that can solve PDEs on million-scale geometries.<n>Transolver++ increases the single- GPU input capacity to million-scale points for the first time.<n>It achieves over 20% performance gain in million-scale high-fidelity industrial simulations.
  arXiv  Detail & Related papers  (2025-02-04T15:33:50Z)
• Unisolver: PDE-Conditional Transformers Towards Universal Neural PDE Solvers [53.79279286773326]
  We present Unisolver, a novel Transformer model trained on diverse data and conditioned on diverse PDEs.<n>Unisolver achieves consistent state-of-the-art on three challenging large-scale benchmarks, showing impressive performance and generalizability.
  arXiv  Detail & Related papers  (2024-05-27T15:34:35Z)
• Pretraining Codomain Attention Neural Operators for Solving Multiphysics PDEs [85.40198664108624]
  We propose Codomain Attention Neural Operator (CoDA-NO) to solve multiphysics problems with PDEs. CoDA-NO tokenizes functions along the codomain or channel space, enabling self-supervised learning or pretraining of multiple PDE systems. We find CoDA-NO to outperform existing methods by over 36% on complex downstream tasks with limited data.
  arXiv  Detail & Related papers  (2024-03-19T08:56:20Z)
• Deep Equilibrium Based Neural Operators for Steady-State PDEs [100.88355782126098]
  We study the benefits of weight-tied neural network architectures for steady-state PDEs. We propose FNO-DEQ, a deep equilibrium variant of the FNO architecture that directly solves for the solution of a steady-state PDE.
  arXiv  Detail & Related papers  (2023-11-30T22:34:57Z)
• Elucidating the solution space of extended reverse-time SDE for diffusion models [23.637881476921596]
  We formulate the sampling process as an Extended Reverse-Time SDE (ER SDE)<n>We offer exact solutions and approximate solutions for SDE and VE SDE, respectively.<n>We devise efficient high-quality samplers, namely ER-SDE-rs.
  arXiv  Detail & Related papers  (2023-09-12T12:27:17Z)
• Deep Learning of Delay-Compensated Backstepping for Reaction-Diffusion PDEs [2.2869182375774613]
  Multiple operators arise in the control of PDE systems from distinct PDE classes. DeepONet-approximated nonlinear operator is a cascade/composition of the operators defined by one hyperbolic PDE of the Goursat form and one parabolic PDE on a rectangle. For the delay-compensated PDE backstepping controller, we guarantee exponential stability in the $L2$ norm of the plant state and the $H1$ norm of the input delay state.
  arXiv  Detail & Related papers  (2023-08-21T06:42:33Z)
• Separable Physics-Informed Neural Networks [6.439575695132489]
  We propose a network architecture and training algorithm for PINNs. SPINN operates on a per-axis basis to significantly reduce the number of network propagations in multi-dimensional PDEs. We show that SPINN can solve a chaotic (2+1)-d Navier-Stokes equation significantly faster than the best-performing prior method.
  arXiv  Detail & Related papers  (2023-06-28T07:11:39Z)
• PDE-constrained Models with Neural Network Terms: Optimization and Global Convergence [0.0]
  Recent research has used deep learning to develop partial differential equation (PDE) models in science and engineering. We rigorously study the optimization of a class of linear elliptic PDEs with neural network terms. We train a neural network model for an application in fluid mechanics, in which the neural network functions as a closure model for the Reynolds-averaged Navier-Stokes equations.
  arXiv  Detail & Related papers  (2021-05-18T16:04:33Z)

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.