MakoXC: Rearchitecting DFT Exchange-Correlation with Matrix-Aligned and Knowledge-Organized Sparsity

Haozhi Han; Fusong Ju; Jing Bai; Ruge Zhang; Xiang Zhao; Liang Yuan; Yunquan Zhang; Ting Cao; Liu Yunxin; Yifeng Chen; Kun Li

Source-linked AI summary

MakoXC: Rearchitecting DFT Exchange-Correlation with Matrix-Aligned and Knowledge-Organized Sparsity

Haozhi Han, Fusong Ju, Jing Bai, Ruge Zhang, Xiang Zhao, Liang Yuan, Yunquan Zhang, Ting Cao, Liu Yunxin, Yifeng Chen, Kun Li

arXiv:2609.01025v2Sep 1, 2026cs.DCphysics.chem-ph

TL;DR

XC evaluation is a cubic-scaling bottleneck in DFT, and existing linear-scaling methods leave implicit sparsity difficult to use while producing irregular GPU workloads. MakoXC restructures that sparsity into matrix-aligned, structured execution through three co-designed techniques, achieving large speedups and scaling an integrated DFT calculation to a 1,231-atom ubiquitin system on 64 GPUs.

Problem
XC evaluation is a major DFT bottleneck, while existing linear-scaling methods leave implicit AO-product sparsity difficult to exploit and generate irregular workloads that map inefficiently to GPUs.
Method
MakoXC combines Matrix-Aligned Cells, Sparsity-Guided Activation, and Kernel-Fused Pipeline to organize nearsightedness-induced sparsity into structured accelerator execution.
Results
67.8× average speedup over DenseXC and 4.7× over GauXC were achieved, while ubiquitin with 1,231 atoms completed end-to-end DFT in under five minutes on 64 GPUs.
Takeaways & Limitations
MakoXC demonstrates practical linear-scaling XC evaluation with high parallel efficiency on modern accelerators and supports larger accurate DFT calculations.

Abstract

from arXiv · show

Density Functional Theory (DFT) is indispensable for materials science and drug discovery, yet the exchange--correlation (XC) evaluation remains a major bottleneck due to its cubic scaling. Although linear-scaling methods exploit electronic nearsightedness to reduce asymptotic complexity, they produce irregular sparse workloads that hide implicit sparsity and prevent efficient use of modern AI accelerators. We present MakoXC, a modular matrix-aligned XC evaluation engine that rearchitects nearsightedness-induced sparsity into regular, accelerator-friendly computations. MakoXC co-designs three key techniques: (1) Matrix-Aligned Cells reorganize nearsightedness-induced interactions into dense, accelerator-aligned data clusters; (2) Sparsity-Guided Activation translates deeper implicit sparsity into numerically correct structured execution for practical linear scaling; and (3) Kernel-Fused Pipeline consolidates fragmented workloads into a unified, compute-intensive execution path that fully unleashes accelerator throughput. Extensive evaluations show that MakoXC achieves average speedups of 67.8$\times$ speedup over standard XC evaluation and 4.7$\times$ over state-of-the-art linear-scaling methods. When integrated into a production-grade commercial DFT package, MakoXC scales XC evaluation to ubiquitin (1,231 atoms, def2-SVP) on 64 GPUs, enabling the end-to-end DFT calculation to complete in under five minutes. By restructuring XC evaluation into a unified, structured computation, MakoXC demonstrates how scientific workloads can achieve genuine low complexity while maximizing parallel efficiency on AI accelerators.

I. INTRODUCTION

MakoXC addresses the bottleneck and hardware inefficiency of XC evaluation by restructuring nearsightedness-induced sparsity into matrix-aligned, accelerator-friendly computation. Its three techniques expose implicit sparsity, fuse fragmented workloads, and deliver strong performance across molecular systems and basis sets.

Motivation: XC evaluation dominates large-scale DFT because conventional computation has cubic-like complexity and repeatedly integrates over large grids and basis functions.The bottleneck lies on the critical path of self-consistent-field iterations.
Challenges: Linear-scaling methods reduce asymptotic complexity toward O(N), but their irregular sparsity hides AO-product sparsity and misaligns execution with accelerator tiles.These limitations create fragmented execution and inefficient utilization of modern AI accelerators.
Approach: MakoXC reorganizes electronic-nearsightedness-induced sparsity into matrix-aligned structures that connect algorithmic efficiency with accelerator utilization.The design exploits global irregularity together with locally dense structure.
Approach: Matrix-Aligned Cells cluster local interactions into dense, fixed-shape accelerator-aligned units while adapting to basis-set and grid resolution.This preserves locality and exposes regular computation for parallel execution.
Approach: Sparsity-Guided Activation exposes implicit AO-product sparsity through value filtering and AO product–density correspondence, while Kernel-Fused Pipeline unifies GEMM and reduction stages.Together, these techniques target structured execution and reduced fragmentation.
Results: 67.8× average speedup over DenseXC, 5.15× over GPU4PySCF, and 4.7× over GauXC were achieved across three molecular series and two basis-set configurations.On 64 GPUs, the integrated engine completed an end-to-end DFT calculation for ubiquitin with 1,231 atoms in under five minutes.
Impact: MakoXC is implemented as a modular XC engine compatible with mainstream DFT software and deployed in a production-grade commercial package.The work also presents matrix-aligned formulation as a design principle for translating irregular scientific sparsity into accelerator-native execution.

II. BACKGROUND AND CHALLENGES

DFT expresses ground-state energy through density-dependent terms, with XC evaluation requiring numerical grid integration. The dominant density-construction and potential-assembly steps scale cubically-like, while electronic nearsightedness offers a route toward near-linear computation.

A. Density Functional Theory and Exchange–Correlation Evaluation: DFT describes interacting-electron ground-state properties through electron density, and the energy functional includes kinetic, external, Coulomb, and exchange–correlation terms.The XC contribution represents many-body exchange and correlation effects and requires approximation by numerical density-function evaluation.
A. Density Functional Theory and Exchange–Correlation Evaluation: XC contributions are evaluated on atom-centered quadrature grids by constructing density, evaluating the functional pointwise, and assembling energy and Fock contributions.The pointwise functional evaluation is computationally light relative to the integration steps.
A. Density Functional Theory and Exchange–Correlation Evaluation: Steps 1 and 3 dominate runtime because dense basis–grid integrations have complexity O(NgridN^2), making large systems increasingly expensive.Here N is basis size and Ngrid is the number of grid points.
B. Linear-Scaling XC Evaluation: Electronic nearsightedness makes each grid point interact mainly with nearby basis functions because atomic orbitals decay rapidly.Shell-dependent cutoff radii discard distant AO contributions and reduce density construction and potential assembly toward near-linear scaling.

B. Memory Hierarchy and Tensor Cores on GPU

NVIDIA GPUs combine CUDA Cores, Tensor Cores, and hierarchical memory across multithreaded streaming multiprocessors. Tensor Cores provide high-throughput fixed-tile matrix operations, but rigid tile geometry limits efficiency for irregular scientific sparsity.

Memory Hierarchy: Each GPU streaming multiprocessor hosts CUDA Cores, Tensor Cores, registers, shared on-chip SRAM, and access to larger off-chip HBM.Warps contain 32 threads, while thread blocks share programmer-managed SRAM.
Tensor Cores: Tensor Cores execute fixed-shape matrix multiply–accumulate operations at several-fold higher throughput than CUDA Cores.Their FP64 example uses tile dimensions m = 8, n = 8, and k = 4.
Tensor Cores: Rigid Tensor Core tile geometry prevents effective utilization when scientific workloads contain irregular and fragmented sparsity patterns.This creates a hardware-mapping challenge for linear-scaling XC evaluation.

C. Opportunity: Bridge Low Complexity to High Parallelism with Sparsity

MakoXC addresses the tension between near-linear XC complexity and accelerator efficiency by reorganizing physically induced irregular sparsity into regular, tile-aligned computation while exploiting implicit sparsity.

Opportunity: MakoXC treats physically induced sparsity as locally dense structure that can be reorganized into regular computational units.This connects nearsightedness-driven algorithmic efficiency with hardware utilization.
Opportunity: Matrix-Aligned Cells map nearsightedness-induced locality into accelerator-aligned units for highly parallel execution.The design targets the first challenge: converting irregular workloads into hardware-friendly structures.
Opportunity: Sparsity-Guided Activation identifies physically effective cells, while Kernel-Fused Pipeline consolidates the activated operations into a unified high-throughput execution path.The two stages operate before and within the SCF loop, respectively.

A. Matrix-Aligned Cells

Matrix-Aligned Cells convert globally sparse but locally dense Grid–AO interactions into regular accelerator-oriented structures, with granularity and storage adapted to the workload and hardware.

A. Matrix-Aligned Cells: Matrix-Aligned Cells map electronic-nearsightedness locality into dense micro-matrices tailored for efficient accelerator utilization.The design exploits local clustering and fixed-size MMA tile alignment.
A. Matrix-Aligned Cells: MAC size is selected from the quadrature grid and basis set while remaining within available on-chip memory.This adapts computation granularity to both scientific inputs and hardware capacity.
A. Matrix-Aligned Cells: The MAC-based sparse format organizes local cells globally through CSR-like row and column indices for parallel XC execution.RowPtr enumerates MACs by row, while ColumnIdx records their column locations.
A. Matrix-Aligned Cells: Sparsity-Guided Activation builds on MACs to filter insignificant cells and exploit both explicit and implicit sparsity.Its one-time activation is performed before the SCF loop because nearsightedness-induced sparsity is static for a fixed basis set and grid.
A. Matrix-Aligned Cells: Value-based filtering discards MACs whose entries are all below ε and stores MACs with more than 90% subthreshold entries in CSR format.The default threshold is ε = 10^-12; value-significant MACs remain fully activated.

2) Correspondence Activation:

Correspondence Activation extends value-based filtering by using mathematically justified shell-pair significance to activate density-matrix MACs corresponding to effective AO products, preserving numerical correctness while reducing work.

2) Correspondence Activation: Value-based filtering exposes Grid–AO explicit sparsity but does not reveal the deeper implicit sparsity of AO products.This limitation motivates correspondence-based activation.
2) Correspondence Activation: AO-product magnitude is exponentially suppressed by shell-center distance, and its global significance is bounded by orbital self-integrals.These properties provide the mathematical basis for shell-pair screening.
2) Correspondence Activation: The shellpair significance criterion uses large inter-shell distances or small Schwarz bounds to identify negligible AO products.Its proof connects Gaussian-product decay and Schwarz bounds to uniform insignificance.
2) Correspondence Activation: Screened shellpairs conservatively approximate AO-product significance: negligible products may remain, but genuinely significant contributions are not discarded.This no-false-negative guarantee preserves numerical correctness while reducing explicit AO-product evaluation.
2) Correspondence Activation: Correspondence Activation maps basis-function positions from screened shellpairs to unique, symmetric density-matrix MAC indices and activates those cells.Algorithm 1 constructs the index set S and applies activation to the density matrix.
2) Correspondence Activation: SGA combines value-based filtering with correspondence activation to translate implicit sparsity into structured execution.This two-stage design targets effective density-matrix MACs rather than only individual AO basis functions.

C. Kernel-Fused Pipeline

KFP fuses XC evaluation into a unified on-chip pipeline by exploiting MAC regularity, SRAM reuse, and redundancy elimination. This design reduces fragmented execution overhead while mapping computation efficiently onto Tensor Core and CUDA Core hardware.

KFP fuses GEMM, DOT, and related operators within SRAM to maximize on-chip data reuse and minimize memory traffic.
The COMPUTERHO kernel computes grid density as ρ = (Φ1×D)·Φ2 using Grid–AO and density matrices.
KFP assigns work across valid MACs, executing dense aligned cases on Tensor Cores and sparse cases on CUDA Cores before SRAM-resident DOT reduction.
The unified pipeline reduces memory traffic and kernel-launch overhead while exposing inefficiencies that isolated operators obscure.
Shape-aligned MACs make operator fusion possible without costly data reordering, while symmetry folding and norm-cache gating eliminate redundant work.

A. Evaluation Setup

The evaluation uses single- and multi-GPU systems, scalable molecular and materials benchmarks, common basis sets and functional settings, and established comparison packages. Performance and accuracy are assessed under explicitly defined timing and threshold protocols.

Single-GPU experiments use an NVIDIA A100 paired with an AMD EPYC 7V13 CPU, while multi-GPU tests use 8-A100 nodes connected by 200 GB/s InfiniBand.
Standalone performance measures average XC evaluation time within SCF cycles, while end-to-end tests measure average SCF iteration time over ten post-initialization iterations.
Comparisons include GauXC, GPU4PySCF, and DenseXC, with GauXC integrated through the same DFT interface for fairness.
Benchmarks cover polyglycine chains, boron-nitride sheets, and water clusters spanning one-, two-, and three-dimensional scaling behaviors.
Experiments primarily use def2-SVP and def2-TZVP basis sets, the B3LYP functional, and a common PySCF-compatible grid strategy.

B. Overall Performance

Across molecular systems and basis sets, MakoXC substantially outperforms dense and linear-scaling baselines while exhibiting practical linear scaling and agreement with established software. It also enables large end-to-end DFT calculations on multi-GPU hardware.

87.1× over DenseXC, 3.6× over GauXC, and 6.7× over GPU4PySCF are achieved with def2-SVP on one A100 GPU.
48.4× over DenseXC, 2.9× over GauXC, and 3.6× over GPU4PySCF are sustained with def2-TZVP on one A100 GPU.
MakoXC exhibits essentially perfect agreement with the theoretical linear-scaling trend as molecular size increases.
MakoXC shows excellent agreement with established software under per-atom energy-error validation.

C. Ablation Study

Ablation results show complementary gains from value filtering, full SGA, and KFP, while SGA adds little runtime overhead. KFP’s fusion and redundancy-elimination mechanisms become increasingly beneficial as systems grow.

9.0× speedup over the dense baseline comes from value-based filtering that exploits Grid–AO sparsity alone.
1.71× further improvement comes from full SGA by additionally exploiting Density Matrix sparsity.
4.1× additional acceleration comes from combining SGA with the KFP implementation.
Less than 3% of total runtime is spent in SGA across globular systems and boron-nitride sheets, with its relative cost diminishing as size increases.
Fusion yields the largest KFP gain, while symmetry folding and norm-cache gating provide increasing benefits with molecular size.

D. End-to-End Evaluation and Scalability

MakoXC delivers substantial end-to-end gains in an integrated DFT package and supports multi-GPU scaling for large molecular systems. Its evaluation also distinguishes XC-specific challenges from prior matrix-aligned ERI techniques.

End-to-End Evaluation: 3.94× average speedup over GPU4PySCF is achieved by the integrated in-house DFT package with MakoXC under def2-SVP.The integration combines MakoXC with a fully featured software stack and optimized ERI kernels.
Scalability Evaluation: Over 74% parallel efficiency is sustained when scaling ubiquitin from 1 to 8 nodes, or 1–64 GPUs, with def2-SVP.XC integration assigns disjoint grid-point sets to GPUs and communicates only for the final reduction.
Evaluation Context: Prior matrix-centric techniques mainly target ERI evaluation, whereas XC calculations require new techniques for geometry-dependent, dynamically evolving sparsity.This distinction motivates treating XC as a separate accelerator-mapping problem.
System Design: MakoXC converts fragmented and irregular XC workloads into structured, matrix-aligned computation for practical linear scaling on AI accelerators.The design uncovers implicit sparsity and consolidates execution into a high-throughput path.