mlxmolkit — GPU-accelerated molecular clustering on Apple Silicon

Recently a guest post from NVIDIA described GPU-Accelerated Clustering with nvMolKit that uses CUDA.

A recent post no describes a port of the nvMolKit (CUDA) molecular clustering pipeline to Apple Metal via MLX. All code is on GitHub https://github.com/guillaume-osmo/mlxmolkit?tab=readme-ov-file

Implements the same 3-step workflow as the RDKit blog post:

1. Morgan Fingerprinting — RDKit GetMorganGenerator (CPU, multi-threaded)
2. Pairwise Tanimoto Similarity — Custom Metal kernel (GPU)
3. Butina Clustering — Greedy algorithm on CSR neighbour list (CPU)

Key results

Tested on Enamine REAL 10.4M subset (same dataset as the blog), Apple M3 Max: