KernelForge - Optimized Kernels for ML

I really only care about writing optimized kernel code, so this project will be completed as I find additional time... XD

I'm reviving this project to finish an old project using random Fourier features for kernel ML.

Installation

Quick Start (Recommended)

For most users, install from PyPI:

pip install kernelforge

This installs pre-compiled wheels with optimized BLAS libraries:

• Linux: OpenBLAS
• macOS: Apple Accelerate framework

Requirements: Python 3.10+

Development Installation

Linux

# Create virtual environment with uv
uv venv
source .venv/bin/activate

# Install in editable mode with test dependencies
make install-linux

# Or manually:
CMAKE_ARGS="-DKF_USE_NATIVE=ON" uv pip install -e .[test] --verbose

macOS

macOS requires Homebrew LLVM for OpenMP support:

# Install dependencies
brew install llvm libomp

# Create virtual environment
uv venv
source .venv/bin/activate

# Install in editable mode
make install-macos

# Or manually:
CMAKE_ARGS="-DCMAKE_C_COMPILER=/opt/homebrew/opt/llvm/bin/clang -DCMAKE_CXX_COMPILER=/opt/homebrew/opt/llvm/bin/clang++ -DKF_USE_NATIVE=ON" uv pip install -e .[test] --verbose

Note: The -DKF_USE_NATIVE=ON flag enables -march=native/-mcpu=native optimizations for maximum performance on your specific CPU.

Linux with CUDA (development)

Planned PyPI layout: a CPU kernelforge wheel plus a Linux companion kernelforge-cuda that drops cuda_*.so into site-packages/kernelforge/. Until the companion is published, use a local monorepo build or build both wheels yourself:

# Day-to-day monorepo (one editable tree with CPU+CUDA):
source /opt/intel/oneapi/setvars.sh
make install-linux-mkl-ilp64-cuda   # passes -DKF_WITH_CUDA=ON

# Or build the split wheels and smoke-test imports in a fresh venv:
make demo-cuda-wheels

Once the companion is on PyPI (release workflow builds it on GHA with manylinux_cuda; no GPU needed to compile), Linux users can install with:

uv pip install kernelforge                 # CPU
uv pip install 'kernelforge[cuda]'         # + torch + kernelforge-cuda

Note: pip/uv install torch from the default index may be CPU-only; install a CUDA build of PyTorch (e.g. from the PyTorch CUDA wheel index) if you need GPU execution. Also requires an NVIDIA driver compatible with that build.

Requires a CUDA toolkit + PyTorch for local builds. Release CUDA wheels are cross-compiled on GitHub Actions (build-wheels-cuda, CPython 3.14 for now); claim the kernelforge-cuda project on PyPI and add a Trusted Publisher matching the kernelforge release workflow before the first publish.

Advanced: Custom BLAS/LAPACK Libraries

Intel MKL (Linux)

# Install Intel oneAPI Base Toolkit
sudo apt install intel-oneapi-base-toolkit

# Set up environment
source /opt/intel/oneapi/setvars.sh

# Install (MKL will be auto-detected by CMake)
uv pip install -e .[test] --verbose

# Optional: Use Intel compilers
CC=icx CXX=icpx uv pip install -e .[test] --verbose

Note: In practice, GCC/G++ with OpenBLAS performs similarly to (or better than) Intel compilers with MKL. On macOS, LLVM with Accelerate framework is highly optimized for Apple Silicon.

ASE molecular dynamics

Install the optional ASE extra:

pip install 'kernelforge[ase]'

Wrap a fitted force-capable model (for example CudaLocalKRRModel) with KernelForgeCalculator, or use run_md() / the kernelmd CLI:

kernelmd --model model.npz --structure start.xyz --steps 5000 --dt 0.5 \
         --temperature 300 --output traj.extxyz

See examples/run_md_rmd17_*.py and examples/run_opt_rmd17_*.py.

Timings

I've rewritten a few of the kernels from the original QML code completely in C++. There are performance gains in most cases. These are primarily due to better use of BLAS routines for calculating, for example, Gramian sub-matrices with chunked DGEMM/DSYRK calls, etc. In the gradient and Hessian matrices there are also some algorithmic improvements and pre-computed terms. Memory usage might be a bit higher, but this could be optimized with more fine-grained chunking if needed. More is coming as I find the time ...

Some speedups vs the original QML code are shown below:

TODO list

The goal is to remove pain-points of existing QML libraries

• Removal of Fortran dependencies
  • No Fortran-ordered arrays
  • No Fortran compilers needed
• Simplified build system
  • No cooked F2PY/Meson build system, just CMake and Pybind11
• Improved use of BLAS routines, with built-in chunking to avoid memory explosions
• Better use of pre-computed terms for single-point inference/MD kernels
• Low overhead with Pybind11 shims and better aligned memory?
• Simplified entrypoints that are compatible with RDKit, ASE, Scikit-learn, etc.
  • A few high-level functions that do the most common tasks efficiently and correctly
• Efficient FCHL19 out-of-the-box
  • Fast training with random Fourier features
  • With derivatives

Priority list for the next months:

• Finish the inverse-distance kernel and its Jacobian

• Make Pybind11 interface

  • Finalize the C++ interface

• Finish the Gaussian kernel

• Notebook with rMD17 example

• Finish the Jacobian and Hessian kernels

• Notebook with rMD17 forces example

• FCHL19 support:

  • Add FCHL19 descriptors
  • Add FCHL19 kernels (local/elemental)
  • Add FCHL19 descriptor with derivatives
  • Add FCHL19 kernel Jacobian
  • Add FCHL19 kernel Hessian (GDML-style)
  • Improve FCHL19 kernel Jacobian performance (it's poor)

• Finish the random Fourier features kernel and its Jacobian

  • Parallel random basis sampler
  • RFF kernel for global descriptors
  • SVD and QR solvers for rectangular matrices
  • RFF kernel for local descriptors (FCHL19)
  • RFF kernels with Cholesky solver and chunked DSYRK kernel updates
  • RFF kernels with RFP format with chunked DSFRK kernel updates
  • RFF kernel Jacobian for global descriptors
  • RFF kernel Jacobian for local descriptors (FCHL19)

• Notebook with rMD17 random Fourier features examples

• High-level model API + ASE:

  • Local/global KRR and RFF models (CPU)
  • CUDA local/global KRR and RFF models
  • CUDA FCHL18 kernels + CudaFCHL18KRRModel
  • ASE interface (KernelForgeCalculator, run_md(), kernelmd CLI)
  • Optional kernelforge[cuda] PyPI packaging (companion CUDA wheel)

• Science:

  • Benchmark full kernel vs RFF on rMD17 and QM7b and QM9
  • Both FCHL19 and inverse-distance matrix

Todos:

• Housekeeping:
  • Pybind11 bindings and CMake build system
  • Setup CI with GitHub Actions
  • Rewrite existing kernels to C++ (no Fortran)
  • Setup GHA to build PyPI wheels
  • Test Linux build matrices
  • Test MacOS build matrices
  • Test Windows build matrices No.
  • Add build for all Python version >=3.10
  • Model save/load via .npz (BaseModel.save / load)
  • Add mkdocs documentation site
• Ensure correct linking with optimized BLAS/LAPACK libraries:
  • OpenBLAS (Linux) <- also used in wheels
  • MKL (Linux)
  • Accelerate (MacOS)
• Add global kernels:
  • Gaussian kernel
  • Jacobian/gradient kernel
  • Optimized kernel for single inference (for MD)
  • Hessian kernel
  • GDML-like kernel
  • Full GPR kernel
  • All kernels in RFP format
• Add local kernels:
  • Gaussian kernel
  • Jacobian/gradient kernel
  • Optimized Jacobian kernel for single inference
  • Hessian kernel (GDML-style)
  • Full GPR kernel
  • Optimized GPR kernel with pre-computed terms for single inference/MD
• Add random Fourier features kernel code:
  • Fourier-basis sampler in C++ with OpenMP parallelization
  • RFF kernel
  • RFF gradient kernel
  • RFF chunked DSYRK kernel
  • Optimized RFF gradient kernel for single inference/MD
  • The same as above, just for Hadamard features when I find the time?
• GDML and sGDML kernels:
  • Inverse-distance matrix descriptor
  • Packed Jacobian for inverse-distance matrix
  • GDML kernel (brute-force implemented)
  • sGDML kernel (brute-force implemented)
  • Full GPR kernel
  • Optimized GPR kernel with pre-computed terms for single inference/MD
• FCHL18 support:
  • Complete rewrite of FCHL18 analytical scalar kernel in C++
  • Stretch goal 1: Add new analytical FCHL18 kernel Jacobian
  • Stretch goal 2: Add new analytical FCHL18 kernel Hessian (+GPR/GDML-style)
  • CUDA FCHL18 kernels + CudaFCHL18KRRModel (scalar / Jacobian / Hessian / full + RFP)
  • Stretch goal 3: Attempt to optimize hyperparameters and cut-off functions
• Add standard solvers:
  • Cholesky in-place solver
    • L2-reg kwarg
    • Toggle destructive vs non-destructive
  • RFP format in-place Cholesky solver
  • QR and/or SVD for non-square matrices
• Add molecular descriptors with derivatives:
  • Coulomb matrix + misc variants without derivatives
  • FCHL19 + derivatives
  • GDML-like inverse-distance matrix + derivatives

Stretch goals:

• Plan RDKit interface
• Plan Scikit-learn interface
• ASE interface (KernelForgeCalculator, run_md(), kernelmd CLI)