Implementation of the multigrid Gaussian-Plane-Wave algorithm with GPU acceleration in PySCF

TL;DR

This paper presents a GPU-accelerated multigrid Gaussian-Plane-Wave density fitting method integrated into PySCF, significantly speeding up quantum chemistry calculations on large molecules and solids with high efficiency.

Contribution

It introduces a novel GPU-accelerated multigrid Gaussian-Plane-Wave density fitting approach in PySCF, achieving high performance and scalability for large-scale quantum chemistry computations.

Findings

Up to 25x speedup on NVIDIA H100 GPUs compared to CPU implementations.

Achieves up to 80% of FP64 peak performance with no efficiency loss for high angular momentum functions.

Ground-state energy and nuclear gradients for large molecules computed in ~30 seconds on a single GPU.

Abstract

We introduce a GPU-accelerated multigrid Gaussian-Plane-Wave density fitting (FFTDF) approach for efficient Fock builds and nuclear gradient evaluations within Kohn-Sham density functional theory, as implemented in the GPU4PySCF module of PySCF. Our CUDA kernels employ a grid-based parallelization strategy for contracting Gaussian basis function pairs and achieve up to 80% of the FP64 peak performance on NVIDIA GPUs, with no loss of efficiency for high angular momentum (up to f-shell) functions. Benchmark calculations on molecules and solids with up to 1536 atoms and 20480 basis functions show up to 25x speedup on an H100 GPU relative to the CPU implementation on a 28-core shared memory node. For a 256-water cluster, the ground-state energy and nuclear gradients can be computed in ~30 seconds on a single H100 GPU. This implementation serves as an open-source foundation for many applications, such as ab initio molecular dynamics and high-throughput calculations.