Matrix Is All You Need: Rearchitecting Quantum Chemistry to Scale on AI Accelerators

Published in International Conference for High Performance Computing, Networking, Storage, and Analysis (SC), 2025

Scientific computing remains fundamentally misaligned with the execution paradigm of modern AI accelerators, which rely on structured, low-precision matrix operations for performance and scalability. Quantum chemistry exemplifies this gap through three core scalability limits: irregular computational patterns, fragmented hardware utilization, and limited scientific reach. In this work, we present Mako, a quantum chemistry system that rearchitects first-principles electronic structure computations as high-performance matrix-aligned kernels to scale on modern AI accelerators. Mako integrates three co-designed components: KernelMako reformulates ERI evaluation into structured matrix operations and leverages CUTLASS to enable transparent, composable MatMul pipelines; QuantMako introduces physics-informed, stage-aware quantization to exploit low-precision compute potential while preserving scientific fidelity; CompilerMako captures static execution patterns across angular momentum classes and automates kernel fusion and architecture-tuned specialization. Mako achieves up to 20× end-to-end speedup on high-angular-momentum basis sets. It sustains over 90% parallel efficiency on a single node and 70% across 64 GPUs, completing the accurate energy calculation of ubiquitin (1,231 atoms, def2-TZVP) from days to just 58 minutes. By restructuring quantum chemistry to align with the AI software-hardware stack, Mako demonstrates how scientific workloads can inherit deep learning–style scalability—scaling beyond the long-standing limits of irregularity, fragmentation, and complexity.

Recommended citation: Haozhi Han, Kun Li, Fusong Ju, Yifeng Chen, Yunquan Zhang, Ting Cao, Mao Yang. (2025). "Matrix Is All You Need: Rearchitecting Quantum Chemistry to Scale on AI Accelerators." SC.
Download Paper