K-means clustering has long been treated as an offline preprocessing step: run it once on a static dataset, record the cluster assignments, and move on to the downstream task. That assumption no longer holds. Modern AI pipelines increasingly call k-means inside training loops and inference passes, where latency per call matters far more than theoretical FLOP counts. A team of researchers from UC Berkeley and UT Austin has released Flash-KMeans, an open-source GPU library that rethinks how exact k-means moves data through memory, delivering end-to-end speedups of up to 17.9× over the best baseline on an NVIDIA H200, 33× over NVIDIA cuML, and over 200× against FAISS.
What Is Flash-KMeans
Flash-KMeans is a batched k-means library implemented entirely in Triton GPU kernels and distributed under the Apache 2.0 license. It installs with a single pip command and produces output that is mathematically identical to standard Lloyd’s k-means. The speedup does not come from approximations, triangle-inequality pruning, coreset sampling, or any algorithmic shortcut. It comes entirely from restructuring how data flows through GPU memory during the two core stages of each Lloyd iteration: the assignment step and the centroid update step.
A standard Lloyd iteration computes the distance from every point to every centroid (assignment), then averages the points assigned to each cluster to produce new centroids (update). Both stages are arithmetic-light and memory-heavy on GPUs. Flash-KMeans attacks the memory bottleneck directly, leaving the math unchanged.
The Two Bottlenecks Flash-KMeans Attacks
Assignment: FlashAssign Eliminates the N×K Distance Matrix
The assignment stage in standard k-means constructs a full distance matrix of shape N×K in high-bandwidth memory (HBM). For a typical workload with N=65,536 points, K=1,024 centroids, d=128 dimensions, and a batch size of 32, the actual distance arithmetic takes about 2.6 milliseconds. Writing that N×K matrix to HBM and reading it back to compute the argmin consumes roughly 23 milliseconds. The matrix itself is the dominant cost, not the arithmetic.
FlashAssign borrows the tiling strategy from FlashAttention. It streams tiles of points and centroids from HBM into on-chip SRAM, fuses the distance computation with an online argmin, and never materializes the full N×K matrix. This reduces the dominant IO complexity from O(NK) to O(Nd + Kd). At the kernel level, FlashAssign achieves up to 21.2× speedup over standard implementations. In one measured case, assignment time dropped from 122.5 milliseconds to 5.8 milliseconds.
Update: Sort-Inverse Update Replaces Contended Scatter Atomics
The centroid update stage in standard k-means uses scatter-style atomic adds. Each thread adds its point into a shared sum buffer indexed by cluster ID. When many threads hit the same “hot” centroid simultaneously, atomic contention forces hardware serialization. The research team measured effective bandwidth of only 50 GB/s on an H200 for this operation.
Flash-KMeans replaces this with a Sort-Inverse Update. It sorts the 1D assignment vector by cluster ID using argsort, so identical cluster IDs form contiguous segments. Each thread block reduces one segment on-chip and then issues a single atomic add per segment. The heavy point matrix is never physically permuted. Atomic operations drop from O((K + N/B_N)d) to O((K + N/B_N)d) — the key improvement is that per-point atomics become per-segment atomics, eliminating contention. The update kernel reaches up to 6.3× speedup.
Benchmark Results
The research team benchmarked Flash-KMeans on an NVIDIA H200 with CUDA 12.8, FP16 data, and d=128 dimensions, sweeping N, K, and batch size B. They compared against four optimized baselines: fast_pytorch_kmeans, fastkmeans, cuML, and FAISS.
- End-to-end vs best baseline: up to 17.9× (N=8M, K=1,024, large N, small K)
- vs NVIDIA cuML: 33×
- vs FAISS: over 200×
- FlashAssign kernel (assignment): up to 21.2× (N=1M, K=8,192)
- Sort-Inverse Update kernel: up to 6.3× (N=33M, K=4,096)
- Out-of-core, large scale: up to 10.5× (N=400M, K=16,384 vs fastkmeans)
One important failure mode provides context: standard PyTorch implementations run out of memory in large-K regimes because they cannot materialize the N×K distance matrix. FAISS, the industry-standard library underlying many production vector-search systems, is outperformed by over 200×.
Flash-KMeans also runs out-of-core. On one billion points with K=32,768 and d=128, it completes an iteration in 41.4 seconds against 261.8 seconds for the baseline. It uses chunked stream overlap to hide PCIe transfer behind compute. A cache-aware compile heuristic reduces tuning overhead by up to 175× while staying within 0.3% of fully tuned performance.
Use Cases: Online K-Means in AI Pipelines
Faster exact k-means shifts the algorithm from an offline batch step to an online building block that can run inside training and inference loops.
- Vector search indexing: FAISS builds its search indices with k-means. Faster k-means enables re-indexing as data shifts instead of rebuilding overnight.
- Sparse attention routing: Routing Transformers and Tactic cluster tokens to route attention. Millisecond-scale k-means makes this viable inside the inference loop.
- KV-cache compression: ClusterKV clusters tokens in semantic space to compress the cache. Cheaper clustering makes per-layer, per-step compression practical.
- Low-bit KV quantization: Recent methods cluster KV entries into codebooks repeatedly. Faster clustering shrinks that preprocessing cost.
- Diffusion Transformers: Sparse VideoGen2 calls batched k-means during forward passes to permute tokens by semantic similarity and exploit sparsity.
Using Flash-KMeans
The API mirrors FAISS and scikit-learn. The following call clusters a batched (B, N, d) tensor on GPU:
import torch
from flash_kmeans import batch_kmeans_Euclid
x = torch.randn(32, 75600, 128, device="cuda", dtype=torch.float16)
cluster_ids, centers, _ = batch_kmeans_Euclid(
x, n_clusters=1000, tol=1e-4, verbose=True
)
preprepre
A scikit-learn-style interface is also available:
from flash_kmeans import FlashKMeans km = FlashKMeans(d=128, k=8192, niter=100) labels = km.fit_predict(large_cpu_tensor) # device=None uses all visible GPUs
preprepre
The kernel auto-dispatches by shape and dtype. A small-D path handles d ≤ 512. A split-D path handles larger d without materializing the distance matrix. Multi-GPU runs trigger automatically for large-N data held in CPU memory.
What This Means for Developers
Flash-KMeans demonstrates that for memory-bound operations like k-means, dataflow architecture on the GPU can deliver order-of-magnitude speedups without altering the underlying algorithm. The library is exact, open-source, and ready to install today. For teams running k-means inside training loops, inference pipelines, or large-scale indexing jobs, the practical implication is straightforward: the bottleneck you previously accepted as a hardware limitation was actually a software architecture limitation. Flash-KMeans removes it, and you can test it on your own workloads right now by running pip install flash-kmeanscodecodecode.
