ByteDance-Seed/JoltQC

ByteDance-Seed/JoltQC

Description: JIT-compiled GPU kernels for quantum chemistry

Language: Cuda

License: Apache-2.0

Stars: 35

Forks: 3

Open issues: 2

Created: 2025-08-21T01:18:44Z

Pushed: 2026-01-30T02:58:04Z

Default branch: main

Fork: no

Archived: no

README:

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
import jqc.pyscf

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 with JIT-compiled versions
mf = jqc.pyscf.apply(mf)
e_tot = mf.kernel()
print('Total energy with double precision:', e_tot)

See more examples in the examples/ directory.

Supported JIT-Compiled Kernels

| Kernel Type | Description | Precision Support | |-------------|-------------|-------------------| | J/K Potentials | Coulomb/exchange matrix construction | FP64 (default), FP32, Mixed | | DFT Density | Electron density evaluation on grids | FP64, FP32, Mixed (default) | | XC Potentials | Exchange-correlation evaluation | FP64, FP32, Mixed (default) | | VV10 Potentials | VV10 nonlocal correlation | FP64, Mixed (default) | | ECP & Derivatives | Effective Core Potentials | FP64 only |


Limitations

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

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},
}

About ByteDance Seed Team

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.