arXiv:2607.18020cs.LGcs.MS2026-07被引 1

用融合内核加速神经PDE求解,大幅降低内存与计算开销。

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

论文配图:FlashPDE: A Drop-In Fused Triton Operator Library for Neural PDE Solvers
图 1 · 摘自论文原文
  • 将有限差分运算融合为可微的Triton内核,统一自动求导接口
  • 在A100上内存占用减少37倍,内核调用减少3.5倍,速度提升2.3倍
  • 适配各类神经网络结构,适合大规模科学计算场景

物理信息神经网络(PINNs)通过将物理约束融入神经网络训练来求解偏微分方程(PDEs),但大规模问题受限于自动微分的内存开销和基于网格的PDE算子执行效率低下。本文提出FlashPDE,一个面向网格型科学机器学习的即插即用融合算子库。FlashPDE用可微的Triton内核替代分散的PyTorch有限差分执行,每个算子集成融合模板计算、解析离散伴随反向传播和边界梯度修正,封装于统一的PyTorch autograd Function接口中。该库提供14个可微分的PDE算子,覆盖1D至3D的椭圆型、抛物型及纳维-斯托克斯系统共17种配置,且独立于神经架构与训练策略。在NVIDIA A100 GPU上的实验表明,FlashPDE相较坐标基自动微分将峰值内存使用降低37.0倍,相较急切式PyTorch有限差分实现减少3.5倍CUDA内核调用。在六个代表性PDE基准测试中,端到端求解时间最高提速2.30倍,内核级别加速达19.2倍,同时保持与PyTorch有限差分参考结果的数值一致性。FlashPDE构建了一个硬件高效的执行层,连接可微分PDE求解器与GPU优化数值计算,嵌入PyTorch生态。

原文摘要 · Abstract (English)

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.

PDE求解Triton神经网络加速计算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。