CuPy-KNN ⚡

GPU-Accelerated K-Nearest Neighbors using NVIDIA CUDA and CuPy

🚀 Overview

CuPy-KNN is a GPU accelerated implementation of the K-Nearest Neighbors algorithm powered by CuPy.

It provides a scikit-learn compatible API while exploiting NVIDIA GPUs for large-scale nearest neighbor search.

Designed for:

• Machine Learning
• Deep Learning embeddings
• Computer Vision features
• Bioinformatics embeddings
• Large-scale similarity search

✨ Features

✅ GPU acceleration with CUDA
✅ Batch-based distance computation
✅ Multiple distance metrics:

• Euclidean
• Cosine
• Manhattan

✅ sklearn-like API

• fit()
• predict()
• predict_proba()
• kneighbors()

✅ Supports large datasets
✅ Fully implemented with CuPy tensors

📦 Installation

Install from PyPI:

pip install cuknn

For CUDA-enabled systems install CuPy first:

pip install cupy-cuda12x

or select the version matching your CUDA installation.

⚡ Quick Start

from cuknn import KNN

from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split


X, y = make_classification(
    n_samples=100000,
    n_features=512,
    n_classes=5
)


X_train, X_test, y_train, y_test = train_test_split(
    X,
    y
)


model = KNN(
    k=5,
    distance="euclidean",
    batch_size=1024
)


model.fit(
    X_train,
    y_train
)


prediction = model.predict(
    X_test
)

📚 Citation

If you use cuKNN in your research:

@software{cupy_knn,
 author = {Calderaro, Salvatore},
 title = {cuKNN: GPU Accelerated K-Nearest Neighbors},
 year = {2026}
}