🤖 AI

cuNNQS-SCI: A Fully GPU-Accelerated Framework for High-Performance Configuration Interaction Selection withNeural Network QQantum States

The paper introduces cuNNQS-SCI, a fully GPU-accelerated framework that overcomes the scalability limitations of the existing NNQS-SCI method through distributed de-duplication, specialized CUDA kernels, and a GPU-centric runtime, achieving up to 2.32x speedup and enabling the solution of significantly larger quantum many-body problems while maintaining chemical accuracy.

Original authors: Daran Sun, Bowen Kan, Haoquan Long, Hairui Zhao, Haoxu Li, Yicheng Liu, Pengyu Zhou, Ankang Feng, Wenjing Huang, Yida Gu, Zhenyu Li, Honghui Shang, Yunquan Zhang, Dingwen Tao, Ninghui Sun, Guangming T

Published 2026-04-20

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Daran Sun, Bowen Kan, Haoquan Long, Hairui Zhao, Haoxu Li, Yicheng Liu, Pengyu Zhou, Ankang Feng, Wenjing Huang, Yida Gu, Zhenyu Li, Honghui Shang, Yunquan Zhang, Dingwen Tao, Ninghui Sun, Guangming Tan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to solve the ultimate puzzle of the universe: How do atoms stick together to form molecules?

To do this, scientists use a method called Configuration Interaction (CI). Think of the molecule as a massive library. Every possible way the electrons (the tiny particles) can arrange themselves is a different "book" in that library. To find the most stable state of the molecule (its ground state), you theoretically need to read every single book.

The problem? For even a medium-sized molecule, the library has more books than there are atoms in the universe. It's impossible to read them all.

The Old Way: The "Hybrid" Bottleneck

Recently, scientists started using AI (Neural Networks) to help. Instead of reading every book, the AI learns to predict which books are the most important. This is called NNQS-SCI.

However, the old version of this AI system had a major flaw. It was like a super-fast race car (the GPU) being driven by a slow, tired driver (the CPU).

The CPU Driver: The AI (GPU) would generate a list of "important books," but then it had to hand that list to the CPU. The CPU would then act as a librarian, checking the list to remove duplicate books (de-duplication) and organizing them.
The Traffic Jam: As the library grew, the CPU librarian got overwhelmed. The GPU race car would sit idle, waiting for the CPU to finish its paperwork. The CPU also ran out of desk space (memory) when the list got too big.

The New Solution: cuNNQS-SCI

The paper introduces cuNNQS-SCI, a system that fires the CPU librarian and puts the entire operation on the GPU. It's like turning the race car into a self-driving, autonomous vehicle that does the driving, the organizing, and the reading all at once.

Here is how they did it, using three simple analogies:

1. The "Smart Sort" (Distributed De-duplication)

The Problem: When thousands of GPUs generate lists of books, they create millions of duplicates. The old way was to dump all lists onto one central computer to sort them. This caused a massive traffic jam.
The cuNNQS-SCI Fix: Imagine a massive concert where everyone is holding a ticket. Instead of everyone running to one ticket booth to check for duplicates, the crowd is split into groups. Each group sorts their own tickets, picks a few "sample" tickets, and swaps those samples with other groups.
The Result: They can figure out exactly which section of the crowd holds which tickets without ever crowding a single booth. This keeps the traffic flowing smoothly, even with 64 GPUs working together.

2. The "Specialized Assembly Line" (GPU Kernels)

The Problem: Creating the list of "coupled configurations" (finding which books are similar) involves doing billions of tiny, repetitive math checks. The old CPU did this one by one, like a person checking off items on a list with a pen.
The cuNNQS-SCI Fix: The GPU is like a factory with millions of robots. Instead of one person checking a list, cuNNQS-SCI deploys a specialized assembly line where millions of tiny robot arms check billions of items simultaneously.
The Result: What used to take the CPU minutes to do now takes the GPU milliseconds. It's the difference between a single person digging a trench and a bulldozer doing it in a second.

3. The "Conveyor Belt" (Memory Management)

The Problem: Even with super-fast GPUs, the list of books can be so huge that it doesn't fit in the GPU's memory (its "desk"). The old system would crash if the library got too big.
The cuNNQS-SCI Fix: Imagine a conveyor belt system. The GPU doesn't try to hold the whole library at once. It pulls in a small "batch" of books, processes them, throws away the ones it doesn't need, and immediately pulls in the next batch while the previous batch is being sent back to storage.
The Result: The GPU never runs out of desk space. It can simulate molecules that are theoretically too big to fit on a single computer chip, effectively breaking the "memory wall."

The Results: Why Does This Matter?

By moving everything to the GPU and optimizing the workflow, the team achieved two massive wins:

Speed: They made the process 2.3 times faster than the best previous methods.
Scale: They can now simulate much larger, more complex molecules (like the Chromium dimer) that were previously impossible to solve accurately.

In summary: cuNNQS-SCI took a system where a slow CPU was holding back a fast AI, fired the CPU, and built a fully autonomous, self-driving AI factory. This allows scientists to solve the "puzzle of the universe" for bigger and more complex materials than ever before, paving the way for new medicines, better batteries, and advanced materials.

1. Problem Statement

The paper addresses the scalability limitations of Neural Network Quantum States (NNQS) combined with Selected Configuration Interaction (SCI), a state-of-the-art method for solving the Schrödinger equation for complex many-body systems. While NNQS-SCI offers high accuracy and better scaling than traditional methods, its practical application to large molecular systems is hindered by a hybrid CPU-GPU architecture that creates three critical bottlenecks:

Centralized CPU De-duplication (Scalability Barrier): The global de-duplication process (removing redundant electron configurations) relies on a single root CPU node to gather and process data from all GPUs. This creates a severe communication bottleneck ( $O(N)$ traffic) and memory pressure, causing parallel efficiency to drop significantly as GPU counts increase.
CPU-Bound Coupled Configuration Generation (Compute Bottleneck): The generation of "coupled configurations" (valid single and double electron excitations) is performed on the CPU. As the system scales, this CPU-bound step becomes slower than the GPU-based neural network inference, leading to severe GPU underutilization.
GPU Memory Wall (Capacity Constraint): Migrating the entire workflow to GPUs exposes the limited High-Bandwidth Memory (HBM) capacity of individual GPUs. Large chemical systems generate intermediate data (configurations and indices) that exceed the 40GB limit of a single NVIDIA A100, rendering standard in-core execution impossible.

2. Methodology: cuNNQS-SCI

The authors propose cuNNQS-SCI, a fully GPU-resident framework that redesigns the SCI workflow to eliminate host-side involvement. The methodology is built on three synergistic components:

A. Distributed Global De-duplication via Sort-Based Regular Sampling

To replace the centralized CPU reduction, the authors designed a fully distributed, GPU-resident de-duplication algorithm:

Sort-Based Strategy: Instead of hashing (which causes thread divergence and load imbalance), the algorithm uses Radix Sort on GPUs to order configurations.
Regular Sampling: Each GPU sorts its local data and samples pivots at fixed intervals. These samples are gathered to a root node to determine global partition boundaries (splitters).
AlltoAll Exchange: Data is redistributed based on these boundaries using MPI_Alltoallv, ensuring that duplicates are adjacent either locally or at partition boundaries.
Local Finalization: Each GPU performs a local merge and stream compaction to remove duplicates. This guarantees deterministic load balancing and minimizes communication overhead.

B. Fine-Grained GPU Kernels for Coupled Configuration Generation

The generation of coupled configurations is migrated entirely to the GPU using a custom CUDA kernel:

Compressed Data Layout: Instead of storing the massive Hamiltonian matrix, the system uses pre-processed Excitation Tables ( $T_{single}$ and $T_{double}$ ) containing pre-calculated integrals. This reduces memory footprint by over 15 orders of magnitude (e.g., from exabytes to <400KB).
Parallel Decomposition: The kernel employs a "one block per source configuration" and "one thread per potential excitation" mapping.
Two-Level Loop: Threads iterate over a virtualized space of excitations, performing bitwise operations to generate new configurations and indexed reads to fetch Hamiltonian elements.
Parallel Filtering: A block-wide prefix sum on a valid_mask allows for efficient, conflict-free stream compaction of valid results without atomic contention.

C. GPU Memory-Centric Execution Model

To overcome the memory wall, the framework treats GPU memory as a high-speed cache rather than a storage medium:

Three-Stage Pipeline: The workflow is split into Generation/De-duplication, Inference/Selection, and Energy Calculation/Restoration.
Mini-Batch Streaming: Large datasets are sliced into manageable batches ( $B_{size}$ ). Intermediate data is processed sequentially to cap peak memory usage.
Asynchronous Overlap: Using multiple CUDA streams, the system overlaps Host-to-Device (H2D) transfer, Kernel Computation, and Device-to-Host (D2H) transfer. This hides PCIe latency, allowing the GPU to compute batch $i$ while transferring batch $i+1$ and writing back batch $i-1$ .
Just-in-Time Indexing: Reverse indices required for energy calculation are constructed on-the-fly against the unique set streamed from host memory, avoiding the need to materialize massive indices in GPU memory.

3. Key Contributions

Novel Fully GPU-Accelerated Framework: cuNNQS-SCI is the first system to migrate the entire NNQS-SCI workflow (including generation, de-duplication, and inference) to the GPU, eliminating CPU-side scalability bottlenecks.
High-Performance Kernels & Distributed De-duplication: The paper introduces specialized CUDA kernels for exact coupled computations and a sort-based, load-balanced global de-duplication algorithm that scales efficiently across multi-GPU clusters.
Memory-Centric Execution Model: A novel runtime design that uses mini-batch streaming and asynchronous offloading to break the single-GPU memory barrier, enabling simulations of systems previously deemed intractable.
Extensive Experimental Validation: Comprehensive evaluation across diverse molecular systems (from small molecules to the 84-qubit Chromium dimer) demonstrating both numerical accuracy and system-level efficiency.

4. Experimental Results

The framework was evaluated on a cluster of 64 NVIDIA A100 GPUs (40GB each) using various chemical systems (C2, N2, Cr2, etc.):

End-to-End Speedup: cuNNQS-SCI achieved up to 2.32× speedup over the highly optimized hybrid CPU-GPU baseline (NNQS-SCI).
- For the Cr2 system (heavy generation workload), the coupled configuration generation time dropped by 432×, and de-duplication time dropped by 40×.
- For the N2 system, generation dropped by 91× and de-duplication by 22×.
Bottleneck Shift: The optimizations successfully shifted the performance bottleneck from host-side limitations back to on-device neural network inference, which is the most scalable component.
Scalability:
- Strong Scaling: On 64 GPUs, the system maintained 91.1% parallel efficiency for the N2 system.
- Weak Scaling: The system exhibited super-linear speedup (up to 163.7% efficiency on 64 GPUs). This is attributed to the increasing redundancy rate in larger configuration spaces, which the efficient de-duplication algorithm eliminates, effectively reducing the computational load per GPU as the cluster grows.
Memory Management: The framework successfully simulated systems (e.g., C2H4O, N2) that theoretically required 65–70 GB of GPU memory. By using the memory-centric model, the actual peak usage was reduced to <40 GB, staying within hardware limits.
Accuracy: The results maintained chemical accuracy (errors < 1.6 × 10⁻³ Hartree). While there were minor floating-point deviations due to non-associative GPU reductions (MAE of ~3.91 × 10⁻⁴ Hartree for Cr2), the results were consistent with the exact NNQS-SCI baseline.

5. Significance

cuNNQS-SCI represents a paradigm shift in AI-driven quantum chemistry. By resolving the system-level bottlenecks of communication, compute imbalance, and memory capacity, it enables the simulation of larger, strongly correlated quantum systems that were previously inaccessible.

Scientific Impact: It allows researchers to tackle complex molecular systems (like the 84-qubit Chromium dimer) with high precision, bridging the gap between AI-driven methods and exact quantum chemistry benchmarks.
Systems Impact: The paper provides a blueprint for "AI-for-Science" systems where non-AI components (data management, logic) are often the bottleneck. It demonstrates that fully offloading these tasks to accelerators, combined with memory-centric execution, is essential for scaling scientific simulations to exascale capabilities.