Skip to content

ByteDance-Seed/JoltQC

Repository files navigation

👋 Hi, everyone!
We are ByteDance Seed team.

You can get to know us better through the following channels👇

seed logo

JoltQC: a JIT-compiled GPU backend for quantum chemistry calculations

formerly xQC

⚠️ Experimental Project
This project is under active development and may change significantly. Use at your own risk.

Overview

JoltQC is not intended to be a complete quantum chemistry package, but rather a collection of standalone, optimized CUDA kernels designed specifically for quantum chemistry workloads. Please see GPU4PySCF installation instructions for full capabilities of quantum chemistry calculations.

Key Features

  • High-performance GPU kernels for quantum chemistry
  • Completely JIT compilation, without any pre-compiled kernels
  • Support FP64, FP32, and mixed-precision schemes
  • Python interface to GPU4PySCF

Recommendations

  • Use PySCF/GPU4PySCF for general-purpose workflows; JoltQC is a JIT backend
  • NVIDIA Ampere or newer GPUs are recommended
  • Newer CUDA versions (e.g., CUDA 12.4+) improve JIT compilation speed

Installation

pip3 install -e .

Python Interface

This example shows how to use JoltQC as a JIT backend with GPU4PySCF:

import numpy as np
import pyscf
from gpu4pyscf import scf
from jqc.pyscf import jk

atom = '''
O       0.0000000000    -0.0000000000     0.1174000000
H      -0.7570000000    -0.0000000000    -0.4696000000
H       0.7570000000     0.0000000000    -0.4696000000
'''

mol = pyscf.M(atom=atom, basis='def2-tzvpp')
mf = scf.RHF(mol)
mf.verbose = 1
mf.conv_tol = 1e-10
mf.max_cycle = 50

# In-place overwrite PySCF kernels
mf_jit = jqc.pyscf.apply(mf)
e_tot = mf_jit.kernel()
print('Total energy with double precision:', e_tot)

See more examples in the examples/ directory.

Limitations

  • No support for density-fitting (DF); DF does not benefit significantly from JIT at this stage.
  • First runs may be slow due to JIT compilation (especially with large basis sets).
  • Only RHF and RKS are currently supported.
  • The performance of small systems is bounded by Python overhead and kernel launch overhead.
  • Support up to 65535 atomic basis

Disclaimer

  • The integral evaluation kernels are based on the GPU4PySCF v1.4 project.

Citation

If you use this project in your research, please cite:

@misc{wu2025designingquantumchemistryalgorithms,
      title={Designing quantum chemistry algorithms with just-in-time compilation}, 
      author={Xiaojie Wu, Qiming Sun and Yuanheng Wang},
      year={2025},
      eprint={2507.09772},
      archivePrefix={arXiv},
      primaryClass={physics.comp-ph},
      url={https://arxiv.org/abs/2507.09772}, 
}

Founded in 2023, ByteDance Seed Team is dedicated to crafting the industry's most advanced AI foundation models. The team aspires to become a world-class research team and make significant contributions to the advancement of science and society.

About

JIT-compiled GPU kernels for quantum chemistry

Topics

Resources

License

Contributing

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published