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

Abstract: 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.

@inproceedings{scmatrixquantumchemistry,
  title={Matrix Is All You Need: Rearchitecting Quantum Chemistry to Scale on AI Accelerators},
  author={Haozhi Han and Kun Li and Fusong Ju and Yifeng Chen and Yunquan Zhang and Ting Cao and Mao Yang},
  booktitle={International Conference for High Performance Computing, Networking, Storage, and Analysis (SC)},

  year={2025}
}