FlashLib: GPU-Accelerated Classical Machine Learning Operators via Triton and CuteDSL

What FlashLib Does and Who It Is For

Classical machine learning operations such as k-means clustering, nearest-neighbor search, and principal component analysis were designed for CPU execution. Running them on datasets with millions of rows and hundreds of dimensions is slow on CPU, and GPU libraries purpose-built for these primitives are relatively sparse compared to deep learning frameworks. FlashLib fills that gap: it provides GPU implementations of 18 high-level classical ML primitives, built from scratch on NVIDIA Triton and CuteDSL rather than wrapping existing libraries.

The primary audience is data scientists and ML engineers who work with large-scale embeddings, feature matrices, and clustering tasks and who need GPU throughput without writing CUDA C or understanding the inner workings of NVIDIA's kernel language. The library exposes a familiar API: each primitive is available as a top-level flash_* function and as a scikit-learn-style class with .fit() and .transform() methods, so the learning curve for teams already using scikit-learn is low.

FlashLib is licensed under Apache 2.0, requires Python 3.9 or later, and currently targets NVIDIA GPUs. The pyproject.toml lists Linux and Windows support and NVIDIA CUDA as the environment. A version 0.3.0 is the current release.

Installing FlashLib and Running a First Clustering Job

FlashLib installs from PyPI:

pip install flashlib

Installing from source is straightforward with pip's editable mode:

git clone https://github.com/FlashML-org/flashlib.git
cd flashlib
pip install -e .

The README gives this example for k-means on a one-million-row, 128-dimension float32 tensor:

import torch
from flashlib import flash_kmeans

x = torch.randn(1_000_000, 128, device="cuda", dtype=torch.float32)
labels, centroids, n_iter = flash_kmeans(x, n_clusters=1024, max_iters=20)

All inputs must be CUDA tensors; the library does not move data to the device automatically. The function returns a tuple of labels, cluster centroids, and the number of iterations taken, which is more explicit than scikit-learn's API and makes it easy to log convergence behavior. The sklearn-style class (KMeans) is also available when you prefer the .fit() interface used in existing pipelines.

Approximate Nearest-Neighbor Indexes: IVF-Flat, IVF-PQ, and CAGRA

FlashLib ships three ANN index types with different trade-offs. The choice between them depends on dataset size, available GPU memory, and the recall target.

IVFFlat partitions the dataset into nlist inverted lists and probes nprobe of them at query time. nprobe is the recall control: higher values inspect more candidates and improve recall at the cost of throughput. The README states that IVFFlat wins batched search above roughly 0.99 recall, because its GEMM fine-scan shares list reads across the batch on tensor cores, yielding approximately 1.5x throughput over cuVS at recall 0.999 on SIFT-1M:

import torch
from flashlib import IVFFlat

db = torch.randn(1_000_000, 128, device="cuda")
queries = torch.randn(10_000, 128, device="cuda")

index = IVFFlat(nlist=1024, nprobe=16).fit(db)
distances, indices = index.kneighbors(queries, n_neighbors=10)  # squared L2

IVFPQ adds product-quantization compression. Each vector is stored as m one-byte codes rather than as a full float32 vector, reducing memory by a factor of 8 to 32. The README gives this example for 32x compression on 128-dimensional vectors:

from flashlib import IVFPQ

# 128-dim fp32 (512 B/vec) -> m=16 PQ codes (16 B/vec) = 32x compression
index = IVFPQ(nlist=1024, m=16, nprobe=16).fit(db)
distances, indices = index.kneighbors(queries, n_neighbors=10)  # ADC squared L2
print(index.compression_ratio)  # 32.0

IVFPQ trades recall ceiling for a dramatic reduction in GPU memory, which makes billion-scale indexes feasible on cards that could not hold the full float32 vectors. The recall ceiling is lower because ADC (Asymmetric Distance Computation) reconstructs distances from compressed codes rather than from exact vectors.

CAGRA builds a proximity graph using exact k-NN construction followed by detour pruning and reverse edges, then answers queries with a fused greedy traversal that keeps the priority buffer in registers:

from flashlib import CAGRA

index = CAGRA(graph_degree=32, itopk_size=64).fit(db)
distances, indices = index.kneighbors(queries, n_neighbors=10)  # squared L2

The README reports that CAGRA's fused traversal outperforms cuVS CAGRA on H100 in the 0.9 to 0.99 recall band. For online serving with small batches, CAGRA wins at every recall level; for large batched search above 0.99 recall, IVFFlat's tensor-core GEMM advantage takes over. The itopk_size parameter controls recall: raising it (alongside graph_degree) improves recall at the cost of query speed.

The Informative API for CPU-Side Cost Estimation

FlashLib includes a flashlib.info submodule that predicts runtime, FLOPs, and HBM bytes for any primitive in approximately five microseconds on CPU. The README notes that it does not import torch, triton, or cutlass, which means it can run in GPU-free environments such as a developer laptop or a CI runner:

import flashlib.info as info

est = info.estimate("kmeans",
                    shape=(100_000, 64),
                    params={"K": 256, "max_iters": 20},
                    device="H200")
print(est.summary_line())

This is useful for pipeline design: before provisioning GPU resources, you can estimate whether a given operation will fit within a memory or compute budget. The README says it is small enough for an LLM agent to call in a GPU-less environment, which points at its intended use in automated pipeline construction or cost-aware scheduling.

The API returns structured estimates rather than just a scalar, giving callers access to FLOPs and bandwidth numbers that can feed into higher-level optimization decisions.

Full Primitive Coverage: Clustering, Decomposition, Manifold, and Regression

Beyond ANN indexes, FlashLib covers four other primitive families. The clustering family includes flash_kmeans, flash_dbscan, flash_hdbscan, and flash_spectral_clustering. DBSCAN and HDBSCAN are density-based methods that do not require specifying the number of clusters ahead of time, which makes them useful for exploratory analysis on high-dimensional embeddings where the cluster structure is not known.

The decomposition family includes flash_pca and flash_truncated_svd. For dimensionality reduction before clustering or visualization, these primitives reduce the feature space on the GPU before passing the compressed representation to k-NN or k-means, keeping the entire pipeline on device without round-trips through CPU memory.

The manifold family provides flash_umap and flash_tsne for nonlinear dimensionality reduction. Both are commonly used for visualizing high-dimensional embeddings, and GPU implementations are essential for large datasets where CPU-based UMAP or t-SNE would take hours.

The regression family includes flash_linear_regression, flash_ridge, and flash_logistic_regression, plus a flash_standard_scaler for preprocessing. These are the primitives that connect classical supervised learning into a GPU-native workflow alongside the unsupervised methods.

Low-level linear algebra primitives (cov_gemm, gram_gemm, eigh, cholqr2) and a Pareto-frontier set of GEMM variants covering fp16, bf16, TF32, and int8 are also available for teams that need custom numerical kernels without full CUDA programming.

Limitations: CUDA-Only, Alpha Status, and the cuML Comparison

FlashLib is CUDA-only. The pyproject.toml explicitly lists NVIDIA CUDA as the environment and does not list ROCm or any CPU backend. Teams running on AMD GPUs or relying on CPU fallback for development and CI have no supported path with the current release.

The library carries Development Status Alpha (3 - Alpha) in its classifiers, and the version is 0.3.0. The benchmark comparisons in the README are specific to H100 hardware; results on older NVIDIA architectures (V100, A100 series) or consumer GPUs (RTX cards) may differ. The benchmark methodology for CAGRA comparisons is in benchmarks/vs_cuml/cagra.py, which the README references for recall and QPS frontier methodology; teams should run those benchmarks on their target hardware rather than applying the H100 numbers directly.

RAPIDS cuML is the most direct alternative: it provides GPU-accelerated scikit-learn-compatible algorithms including k-means, DBSCAN, UMAP, PCA, and ANN indexes, and it is backed by NVIDIA with production support. FlashLib's stated advantage is kernel-level optimizations (Triton and CuteDSL implementations) that it claims outperform cuML's cuVS in specific recall bands, but cuML's RAPIDS ecosystem, documentation depth, and multi-GPU support make it the safer default for production workloads. FlashLib is a better fit for teams already using Triton for custom kernels and wanting consistent tooling across classical ML and deep learning primitives.

Maintenance, Citation, and License

The last push to the FlashLib repository was on 2026-08-31. The project has no GitHub releases; distribution is through PyPI as of v0.3.0. The Apache 2.0 license permits commercial use and modification with attribution. The project is affiliated with researchers at UC Berkeley and lists a citation entry in the README pointing to the FlashLib technical paper on the flashml-org.github.io website.

The repository includes a Slack workspace and a Discord server for community questions. Given the alpha status and the absence of a formal release process (no GitHub Releases page), teams that adopt FlashLib for production use should pin the PyPI version and monitor the repository for breaking changes between minor versions.

Editorial conclusion

FlashLib is a practical choice for data scientists and ML engineers who need GPU-scale classical ML without writing CUDA kernels. Its scikit-learn-compatible API reduces migration friction, and the CAGRA, IVF-Flat, and IVF-PQ implementations address different recall and memory trade-offs for approximate nearest-neighbor search at scale. The library requires a CUDA-capable NVIDIA GPU and is at version 0.3.0 Alpha; verify that the primitives you need are included in the 18 shipped with the current release, since some planned operators listed in documentation may not yet be implemented. Check benchmarks/vs_cuml/ for the recall and throughput methodology before adopting the ANN indexes in production.

Frequently asked questions

Official sources

1. FlashML-org/flashlib on GitHub
2. Issues
3. License: Apache-2.0
4. README