flash-kmeans and the return of exact K-Means on the GPU

Exact K-Means, with the assignment step written in Triton

flash-kmeans is described as IO-aware batched K-Means clustering implemented with Triton GPU kernels, and the repository description calls it fast and memory-efficient exact K-Means. The word exact is the important one, because a large share of the K-Means implementations in circulation are approximations.

It is also the official K-Means implementation of Sparse VideoGen2, so the design is shaped by a real workload: video generation pipelines that need to cluster large batches of latent vectors under tight memory limits. The repository carries a paper link, a blog link and a Twitter post, and the description is marked as a NeurIPS 2026 work.

The API has two layers. There is a direct functional entry point for batched Euclidean clustering, and there is a class based interface described as similar to `faiss/sklearn`, with the details in `flash_kmeans/interface.py`. The functional form is the shortest path to a result:

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)

The package is pure Python plus Triton, with `torch>=2.0`, `triton>=2.0`, `numpy` and `tqdm` as its dependencies and Python 3.8 or newer required. There is no compiled extension to build, which makes installation unremarkable, and the package data section of the build configuration explicitly includes the Triton kernels.

The point of the kernel is that the distance matrix never exists

Every K-Means implementation has the same structural problem. To assign a point to a cluster you need its distance to every centroid, so the naive formulation wants an array of shape, batch times points times clusters. With a batch of 32, 75600 points and 1000 clusters in the example above, that array is not a rounding error, it is the whole problem.

flash-kmeans avoids it by streaming. The centroids are read in chunks of `BLOCK_K` and the points are processed in tiles of `BLOCK_N`, so the kernel never materialises the full distance tensor. The README is explicit that the cross accumulator of shape `(BN, BK)` is held in registers across the D loop, and that this preserves the K-streaming property, meaning no `(B, N, K)` distance matrix is materialised.

The assign kernel ships in two flavours and dispatches automatically. The small-D path handles `D` of 512 or less, where one program loads a tile of points once and streams over the centroids. The split-D path handles `D > 512`, or any case where the small-D kernel cannot fit in shared memory at the minimum tile size, and it tiles the dimension loop with an outer K loop and an inner D loop.

The configuration is tuned per architecture. The README lists hand-tuned tables for H200, H100, A100 and GB10, choosing `BLOCK_N`, `BLOCK_K`, `num_warps` and `num_stages` from grid sweeps performed with a separate tuning project, `flash-kmeans-tune`.

Unknown GPUs get a conservative fallback, and slower

The dispatch story has a catch, and the README does not hide it. When the GPU matches none of the tuned families, the wrapper unconditionally dispatches to the split-D kernel with a conservative fallback configuration of `BLOCK_N=32`, `BLOCK_K=32`, `BLOCK_D=32`, `num_warps=4` and `num_stages=1`.

The stated reason is sensible. Rather than relying on tuning data that does not exist for that architecture, the wrapper picks a configuration that is safe, and a post-process step named `_fit_config_to_smem_split_d` guarantees the launch fits in shared memory whatever the actual budget is.

The stated cost is equally clear: performance will be suboptimal on unfamiliar architectures. The remedy is to re-tune with `flash-kmeans-tune` so the corresponding per-architecture heuristic function is populated.

The good news is that this is invisible in the API. The README states that the dispatch is automatic with no API changes, and shows the large-D plus float32 case working through the same entry point:

from flash_kmeans import FlashKMeans
km = FlashKMeans(d=2048, k=1024, dtype=torch.float32)  # large-D + fp32 ok
km.train(data)

For a reader weighing adoption, this is the single most important practical detail. On one of the four tuned architectures you get a hand-tuned launch. On anything else you get a correct result with a configuration nobody profiled, and the only way to know which case you are in is to know your GPU model.

Multi-GPU scaling without a NCCL dependency

For large-N workloads, the `kmeans_largeN` path supports automatic multi-GPU scaling. When `device=None`, all available GPUs are used, and specifying a single device such as `device="cuda:0"` falls back to single-GPU mode. No new API parameters are needed.

The design is worth reading because it avoids the usual tool. The N data points are split into contiguous block partitions, one per GPU, and each GPU runs its own double-buffered host-to-device plus compute pipeline over its partition. The stated consequence is that PCIe bandwidth scales linearly with the GPU count.

The reduction step is the interesting part. Instead of NCCL, the implementation does a manual gather-reduce-broadcast: partial centroid sums of roughly 4 MB and counts of roughly 32 KB are gathered to GPU 0 with direct device-to-device copies over NVLink, reduced, and the new centroids broadcast back. The README's justification is that the data is small enough that manual copies are faster and everything stays in a single process.

The final trick is overlap. Host-to-device transfers run on PCIe while the reduction runs on NVLink, and the first transfer block of the next iteration is prefetched during the current iteration's reduction, hiding the reduce latency behind the transfer.

The class based entry point shows how little of this you have to know:

from flash_kmeans import FlashKMeans
km = FlashKMeans(d=128, k=8192, niter=100)

The benchmark compares against four other implementations, on one GPU

The README lists the baselines, and the choice of them tells you what the authors considered competitive. There is `fast_pytorch_kmeans`, a PyTorch implementation, and `fastkmeans` in both a Triton flavour and a PyTorch fallback flavour, from a different author. Then there is flash-kmeans itself in its Triton implementation and in its PyTorch fallback, and finally a naive batched PyTorch implementation that does not consider out-of-memory situations.

The setup is stated: an NVIDIA H200 GPU, FP16 precision, 128 dimensional data, varying the number of clusters, the number of data points and the batch size. The results are given as two images, and the README reports significant performance improvements for the Triton implementation. One specific note is included: the Triton flavour of `fastkmeans` errors when k is 100 or 1000 in the first figure.

A separate large tensor benchmark covers the case where the input cannot fit in GPU memory. There, flash-kmeans and fastkmeans are compared with FP32 precision, 128 dimensions, the number of points scaling from 256K to 268M, that is N from 2^18 up to 2^28, and cluster counts following K as the square root of N, from 512 up to 16384. The input tensor is generated randomly in CPU pinned memory, and both implementations transfer data from CPU to GPU in chunks and compute.

Two caveats about reading those numbers. The comparisons are on one tuned architecture, an H200, and the tuned heuristics exist for H200, H100, A100 and GB10. And the large-N comparison is against a single competitor rather than the full field.

When a GPU K-Means is the wrong choice

There are several cases where this library is the wrong tool, and they are worth naming before anyone reaches for it.

If your data is small, a CPU implementation is the right answer and the GPU version is overhead. K-Means on a modest number of points in low dimensions finishes quickly enough in PyTorch or in a CPU library, and the Triton kernel's advantage only appears once the work is large enough to matter.

If you do not have an NVIDIA GPU with a working Triton installation, the only path is the PyTorch fallback, and the performance case for the project mostly disappears. The dependencies are `torch>=2.0` and `triton>=2.0`, and Triton is not a universal install.

If you need approximate clustering, this is the wrong tool, because it is explicitly exact. Approximate methods such as MiniBatch K-Means trade a little accuracy for a large speed gain, and there is no approximate mode here to switch on.

There is one more operational limit that matters for shared clusters. The README documents per-architecture tuning and a conservative fallback, and the tuning data exists for four GPU families. On an unfamiliar accelerator the code is correct but untuned, and fixing that means running a separate tuning project and populating the heuristics yourself.

Version, packaging and a licence that needs reading

The package version in the build configuration is 0.3.1, and the repository has no GitHub releases, so the version in `pyproject.toml` is the only marker. The last push was on 2026-08-28 and the repository is not archived.

The build backend is setuptools with a `requires = ["setuptools>=64", "wheel"]` requirement, the package is `flash-kmeans`, and the description in the manifest is fast batched K-Means clustering with Triton GPU kernels. The project also ships a `MANIFEST.in`, and the `examples/` directory is unusually rich for a library of this size. It contains benchmark scripts for the flash implementation, GPU scaling, large inputs and overlap, plus a patch for the fastkmeans baseline and a shell script that runs a benchmark grid. There are also focused test scripts for assign with non-power-of-two dimensions, for the split-D path, for a centroid out-of-memory reproduction, for compilation and for the API surface.

The licence needs a note, because the repository metadata and the package disagree. The repository description records the licence as Apache-2.0, while the classifier in `pyproject.toml` reads `License :: OSI Approved :: MIT License`, and the manifest points at a `LICENSE.txt` file for the terms. If you are redistributing the package or shipping it inside a product, read that file rather than trusting either label.

Editorial conclusion

flash-kmeans is the right tool when the dataset does not fit the naive approach, when you need exact K-Means rather than an approximation such as MiniBatch, and when you are already running NVIDIA GPUs with PyTorch and Triton in the environment. It is the wrong tool for small data that a CPU implementation clusters in milliseconds, and for a team without a Triton-capable NVIDIA GPU, where the PyTorch fallback is the only path. Verify first which GPU you are on, because the tuned heuristics cover H200, H100, A100 and GB10 and anything else falls back to a conservative split-D configuration that the README says will be suboptimal, and check the licence file yourself, since the repository metadata and the package classifiers disagree on whether the licence is Apache-2.0 or MIT.

Frequently asked questions

Official sources

1. Issues
2. License: Apache-2.0
3. README
4. svg-project/flash-kmeans on GitHub