Scalable High-Fidelity Macromolecular Docking for GPU-Accelerated Supercomputers

Overview and Motivation

The paper presents SparkleDock, a scalable Glowworm Swarm Optimization (GSO)-based flexible macromolecular docking framework designed for GPU-accelerated supercomputers. The authors address a fundamental tension in computational biology: flexible macromolecular docking provides high-fidelity predictions of biomolecular interactions but is computationally prohibitive at scale. The paper states: "Flexible macromolecular docking is essential for accurately modeling biomolecular interactions, including proteins, antibody–antigen, membrane receptor–soluble protein interactions, and drug design. Compared to rigid-body methods, it improves the docking success rate by about 20%–30% as it accounts for conformational flexibility that traditional rigid-body docking cannot capture."

The existing state-of-the-art flexible docking method, LightDock, adopts GSO for high accuracy, but the paper identifies three key bottlenecks preventing efficient GPU execution: (1) Limited parallelism—GSO exposes parallelism primarily at the swarm level (typically hundreds to thousands of swarms), far below GPU concurrency; (2) Irregular compute pattern—the energy score calculation accounts for 89% of runtime and over 95% of FLOPs but involves irregular pairwise Euclidean distance computations that cannot be directly mapped onto Tensor Core Units (TCUs); (3) Poor scalability—GSO exhibits inherent load imbalance due to heterogeneous exploration, and large tasks (e.g., PDB entry 4GAM) have memory footprints of hundreds of gigabytes, exceeding single-GPU memory.

Key Contributions

The paper claims four main contributions: (1) a novel fine-grained, glowworm agent-level parallelization scheme matching GPU compute hierarchy; (2) restructuring the dominant energy scoring computation into a TCU-compatible formulation—"We restructure the dominant energy scoring computation into a TCU–compatible formulation (i.e. 89.1% of the total execution time), transforming it into structured matrix computations to fully exploit TCUs while preserving the original scoring semantics"; (3) performance-model-guided scaling methods for load balancing and out-of-core chunk division; and (4) demonstration of significant speedups on A100/H100 GPUs and multi-GPU systems.

Algorithmic Redesign

The original GSO algorithm organizes optimization as interacting glowworm agents within each swarm, where each glowworm agent updates its state based on its local neighborhood. SparkleDock performs loop reordering and merges the swarm and agent loops, mapping them onto CUDA threads: We first perform loop reordering between the swarm and simulation loops. We then merge and unroll the swarm and agent loops, mapping them onto CUDA threads. This agent-level parallelization increases parallelism by more than Ng times compared to swarm-level parallelization. The simulation is implemented as five CUDA kernels: prepare rec pose, prepare lig pose, calc dfire, calc neighbors, and movement.

TCU Reformulation of Energy Scoring

The core technical innovation is reformulating the DFIRE energy score calculation for Tensor Cores. The pairwise distance computation in Eqn. 1 is expanded via the identity:

D x,y = sqrt sum k=1 cubed(A x,k(i,j)) squared + sum k=1 cubed(B y,k(i,j)) squared - 2 times sum k=1 cubed A x,k(i,j) times B y,k(i,j)

where the underlined term is a tall-skinny matrix multiplication adapted for TCU, and the remaining terms are squared norms computable on CUDA cores. The decomposition enables simultaneously execution of (1) on TCUs and (2) on CUDA cores. Since the k-dimension (3) is lower than the minimum FP64 mma fragment dimension (k=4), the implementation allocates 16×4 shared memory tiles and pads zeros along k.

Register Remapping and Pipeline Optimization

The paper identifies that "in FP64 mma operation, fragA follows a row-major layout and fragB follows a column-major layout, where each thread holds one element... The accumulator fragC is organized in a row-major layout, where each thread holds two elements. These heterogeneous layout designs introduce mismatches across fragments." To avoid costly shared memory traffic (23 cycles per shared memory load vs. 2 cycles per warp shuffle), SparkleDock devises shuffle-based register remapping: for fragA, shufl sync broadcasts values within a row to compute row-wise sums of squares; for fragB, a column-packing and warp-level transpose strategy converts the column-major layout to match fragC. The paper reports: Only 12 warp shuffle instructions are required in register remapping with 24 cycles... This optimization theoretically accelerates by 12×.

Pipeline optimization uses cp.async on Ampere+ GPUs to overlap global-to-shared memory transfers with TCU computation via a four-stage pipeline (T0–T3) organized into three asynchronous-copy groups.

Performance Model

The authors construct a memory budget model Mem total and performance budget model T total with parameters including swarm/agent counts, receptor/ligand atom counts, kernel throughputs, and bandwidth. The model decomposes runtime into DFIRE calculation, pose preparation, neighbor/movement, and data transfer components. The total runtime for a multi-GPU environment is the maximum over MPI ranks of summed per-swarm times. The model serves three purposes: performance prediction, MPI load balancing (using a greedy algorithm assigning each swarm to the currently-least-loaded MPI rank), and automatic out-of-core chunk division for large workloads exceeding GPU memory.

Experimental Results

The evaluation uses the BM5.2 benchmark with nine representative complexes (2VXT, 3VLB, 2A1A, 2GTP, 2X9A, 1RKE, 4GAM, 4JCV, 4LW4) spanning all categories. Key results:

• Single GPU speedups: SparkleDock achieves an average speedup of 9.7× and 18.9×, with corresponding throughputs of 1017.7 and 2005.9 agents/s on A100 and H100 GPUs, respectively, compared to LightDock-Rust. Table IV reports Max/Avg speedups of 10.6×/9.7× on A100 and 20.0×/18.9× on H100.

• TCU vs. CUDA cores: The TCU-accelerated calc dfire kernel achieves around 1.99× speedup on single A100 GPU and 2.43× acceleration on single H100 GPU compared to a CUDA-core implementation. Step-wise breakdown shows register remapping gives up to 2.7× (A100) and 3.2× (H100) additional speedup, and pipeline overlapping provides further 1.3× and 1.2× speedups, reaching 2.27 TFLOP/s on A100 and 4.60 TFLOP/s on H100.

• Strong scaling: On up to 512 A100 GPUs, SparkleDock achieves 183.1×, 94.4×, and 67.0× speedups for 4GAM, 4JCV, and 4LW4 for 512 GPU scaling, respectively, and a 32.1× speedup for 2VXT for 256 GPU scaling, relative to a single A100 GPU. Load balance evaluation shows standard deviations of 0.0035, 0.0078, 0.0041, and 0.0072 across MPI ranks for the four datasets.

• Accuracy preservation: Docking success rates at Top-1, Top-5, Top-10, Top-20, Top-50, and Top-100 are 55.6%, 88.9%, 88.9%, 88.9%, 88.9%, and 88.9%, respectively, identical to those obtained by the LightDock implementation. L-RMSD values for 2GTP and 2VXT are below 2.0 Å.

• Model validation: Predicted runtime closely matches measured time with mean absolute percentage error (MAPE) of 12.55%.

• Roofline analysis: prepare lig pose and calc dfire are compute-bound kernels achieving over 25% of measured peak performance on A100; prepare rec pose is memory-bound achieving 92% memory occupancy.

Conclusion

The paper claims SparkleDock is the first flexible docking framework capable of multi-GPU scaling while utilizing TCUs, enabling large-scale virtual screening. On 512 GPUs, it reduces docking time from hours to seconds, enabling large-scale, high-fidelity virtual screening previously impractical with flexible docking. Future work includes extending to ARM-based platforms like Fugaku and applying TCU reformulation to other domains.

1. Tensor-core-compatible reformulation of AI kernels with pairwise distance/energy terms

• For AI systems that compute pairwise interactions (attention scores, graph message passing, equivariant neural networks, molecular force fields), replace scattered Euclidean-distance computations with the SparkleDock-style expansion:

D = sqrt(A2 + B2 − 2·A·BT)

Then run the A·BT term as a tall-skinny matrix multiplication on Tensor Cores, and compute the squared norms on CUDA cores concurrently.

• The improved AI system can train and run graph transformers, point-cloud models, and AI-driven molecular-docking predictors with 2–4× faster pair-interaction layers on A100/H100 GPUs, while preserving exact semantics — no approximation or loss of accuracy.

2. Fine-grained agent-level parallelization for population-based AI training

• Apply the glowworm-agent-level mapping: merge swarm and agent loops and map them to CUDA threads/kernels instead of treating each algorithm instance as one GPU thread block.

• Improved AI systems using evolutionary strategies, genetic algorithms, population-based reinforcement learning, or MCTS self-play can scale to thousands of simultaneous agents on a single GPU, increasing training throughput by 10–20× compared to instance-level parallel implementations.

3. Register-remapping and warp-shuffle data layout transformation for tensor-core AI operations

• Use the fragment-layout mismatch solution: for matrix-multiplication operands whose output layout differs from input layout (e.g., FP64 mma fragA row-major vs fragB column-major vs fragC row-major), replace shared-memory round-trips with 12 warp-shuffle instructions.

• Improved AI systems running FP64/FP32 mixed-precision linear layers, attention projections, or convolution via implicit GEMM can reduce kernel latency and improve FLOP/s by up to 3× per tensor-core operation, especially for small hidden dimensions common in graph neural networks and protein language models.

4. Four-stage cp.async pipeline overlapping memory transfer and tensor-core compute

• Apply SparkleDock’s asynchronous-copy, multi-stage pipeline to AI training/inference data-loading and tensor-core computation.

• Improved AI systems (e.g., large language model training, drug screening pipelines) can hide global-memory-to-shared-memory transfer time behind GEMM/attention computation, reducing kernel stall time by 1.3× and improving sustained throughput on A100/H100, particularly for memory-bound embedding lookups and long-sequence attention.

5. Performance-model-guided load balancing for distributed AI training

• Use SparkleDock’s Mem total and T total performance budget model with a greedy least-loaded assignment for sharding data, experts, or population members across GPUs.

• Improved AI systems with heterogeneous inputs (variable sequence lengths, varying molecular sizes, unbalanced RL environments) can achieve near-zero load imbalance across MPI ranks — the paper reports standard deviations of 0.0035–0.0078 across ranks — enabling up to 183× scaling on 512 GPUs instead of being bottlenecked by stragglers.

6. Automatic out-of-core chunk division for AI workloads exceeding GPU memory

• Use the model-driven chunking approach to split large tensors, graphs, or docking complexes into chunks that fit in GPU memory while preserving the exact computation order and overlap communication with execution.

• Improved AI systems can process protein complexes, high-resolution point clouds, or long-document transformers that currently require hundreds of gigabytes of GPU memory on a single GPU, with automatic spillover and recomputation scheduling that keeps utilization high and prevents OOM failures.

7. Swarm-optimization-based AI inference for non-differentiable objectives

• Use the GSO algorithm as an alternative AI “inference” layer when gradients are unavailable (e.g., black-box optimization, molecule design, policy search with discrete actions). SparkleDock’s scaling framework makes GSO practical at supercomputer scale.

• Improved AI system can perform large-scale flexible docking and drug screening faster than LightDock: on 512 A100 GPUs, it reduces docking time from hours to seconds, enabling AI-driven virtual screening of millions of candidates with flexible receptor/ligand conformations — something previously impractical.

8. Specialized AI kernels for molecular/RL simulation environments

• Port the five-kernel decomposition (prepare rec pose, prepare lig pose, calc dfire, calc neighbors, movement) as reusable, TCU-accelerated primitives for AI simulation environments in physics-based reinforcement learning and generative molecular design.

• Improved AI system can generate physically accurate protein-ligand conformations at 2005.9 agents/s on a single H100 and maintain identical docking success rates (Top-1 55.6%, Top-5 88.9%, Top-10+ 88.9%) compared to baseline LightDock, providing both speed and scientific fidelity.