Flash-KMeans: Fast and Memory-Efficient Exact K-Means

TL;DR

Flash-KMeans is a hardware-aware, mathematically exact GPU implementation of Lloyd’s k-means that targets k-means as an online primitive inside training/inference loops (semantic deduplication, embedding/KV quantization, sparse routing, DiT token permutation), rather than as an offline batch tool. It replaces standard implementations with two kernel rewrites — FlashAssign (online argmin fused with distance computation, no distance-matrix materialization) and sort-inverse update (argsort the assignment vector, then do segment-level on-chip reductions instead of scatter atomics) — plus stream-overlapped out-of-core execution and a cache-aware compile heuristic. Reported speedups on H200: up to 17.9× end-to-end over the best PyTorch baseline, 33× vs cuML, >200× vs FAISS, and ~10.5× at one-billion-point out-of-core scale, with 175× less tuning overhead. Code: https://github.com/svg-project/flash-kmeans.

Key claims

k-means is shifting from an offline analytics tool to an online primitive — semantic dedup, VQ for late-interaction search (PLAID/ColBERT), sparse-attention token routing, KV-cache quantization, and semantic-aware token permutation inside DiTs — so the relevant metric becomes per-invocation latency, not aggregate throughput [§1, §2 Related].
Standard GPU k-means is bottlenecked by two implementation-level constraints rather than algorithmic FLOPs: (1) explicit HBM materialization of the N×K distance matrix dominates the assignment stage, and (2) scatter-style atomic adds keyed by cluster id cause severe write contention in the centroid-update stage, limiting effective bandwidth to ~50 GB/s on H200 [§3.2].
FlashAssign fuses streaming distance computation with an online argmin that maintains per-point (min_dist, min_idx) registers and scans centroid tiles with double-buffered async prefetch, eliminating the distance matrix entirely; ideal IO complexity drops from O(NK) to O(NF + KF + N) [§4.1].
Sort-Inverse Update argsorts only the 1-D assignment vector (the heavy N×F point matrix is never permuted), then has each CTA do on-chip segment-level partial sums and issue only one atomic_add per contiguous cluster-id segment; worst-case atomic count drops from O(N) to O(K + chunks) [§4.2].
Algorithm-system co-designs: chunked stream overlap (double-buffer host-to-device transfers) for out-of-core execution, plus a cache-aware (L1/L2 size) compile heuristic that picks near-optimal kernel configs analytically instead of via exhaustive autotune [§4.3].
Kernel-level speedups on H200: up to 21.2× over standard assignment and up to 6.3× over standard centroid update [§5.2, Fig. 4].
End-to-end speedups across (N, K, F, B) sweep: ≥5.4× over fastkmeans in the large-N/large-K regime where PyTorch baselines OOM; 17.9× over fast_pytorch_kmeans in large-N/small-K; 15.3× in small-N/small-K highly-batched regime [§5.2, Fig. 3].
Vs industry standards: up to 33× faster than NVIDIA cuML and >200× faster than FAISS at end-to-end iteration latency [§5.2, Abstract].
Scales to one billion points out-of-core (data exceeds VRAM): 41.4 s vs 261.8 s for fastkmeans at N=1B (6.3× speedup); 10.5× at the (N=2B, F=64) config — standard PyTorch baselines OOM and cannot run [§5.3].
Cache-aware compile heuristic finishes config selection in <2.5 s vs >325 s for exhaustive tuning (up to 175× faster time-to-first-run) with <0.3% kernel-latency regression vs the autotuned oracle [§5.3, Fig. 5].
The output is mathematically exact Lloyd’s k-means, not an approximation — the contribution is dataflow restructuring, not algorithm change [§1, §4].

Method

A standard Lloyd iteration has two stages: (1) assignment computes distances D = ||X − C||² and takes the row-wise argmin to assign each point to its nearest centroid; (2) update averages the points assigned to each cluster to form new centroids. Naive GPU implementations write the full N×K distance matrix to HBM and read it back (assignment), and use atomic scatter writes keyed by cluster id (update).

Flash-KMeans replaces both stages. FlashAssign uses 2D tiling over points and centroids, keeps a running (min_dist, min_idx) per point in registers, and computes distances tile-by-tile on chip, folding the argmin into the same kernel via an online update; an async prefetch double-buffer hides centroid-tile HBM loads behind the current tile’s compute. The distance matrix is never materialized. Sort-Inverse Update argsorts the 1-D assignment vector to construct an inverse mapping sorted_idx; each CTA then processes a contiguous chunk of the sorted sequence, gathers the corresponding features from the original (un-permuted) X via sorted_idx, accumulates segment-wise partial (sum, count) buffers in shared memory/registers, and emits one atomic_add per segment boundary to HBM — converting scatter contention into regular segmented reduction. Worst-case atomics drop from O(N) to O(K + W) where W is the number of CTA chunks.

For deployments where N exceeds VRAM, a CUDA-stream-overlapped pipeline transfers the next data chunk while the current chunk is processed. To avoid autotune blow-up under dynamic shapes, a cache-aware heuristic selects tile dimensions directly from L1/L2 cache sizes and (N, K, F, B) — no exhaustive sweep.

Results

On a single NVIDIA H200, evaluated against fast_pytorch_kmeans, fastkmeans (AnswerDotAI), cuML, and FAISS:

Kernel: FlashAssign up to 21.2× over standard assignment; Sort-Inverse Update up to 6.3× over standard scatter update [§5.2, Fig. 4].
End-to-end: 17.9× over the best PyTorch baseline in compute-intensive regimes; ≥5.4× over fastkmeans in memory-intensive regimes where PyTorch OOMs; 15.3× in small-N/small-K highly-batched regimes [§5.2, Fig. 3].
vs industry: 33× vs cuML, >200× vs FAISS [§5.2, Abstract].
Out-of-core (N=1B, F=64, K=4096): Flash-KMeans completes one iteration in 41.4 s; fastkmeans takes 261.8 s — 6.3× speedup. At (N=2B, F=64) the speedup is 10.5× (8.4 s vs 88.4 s). PyTorch baselines OOM and cannot run [§5.3].
Time-to-first-run: cache-aware heuristic picks configs in <2.5 s vs >325 s for exhaustive tuning (175× faster) with <0.3% latency regression [§5.3, Fig. 5].

Why it’s interesting

Terrance’s Slack note asks the right question — whether k-means is still load-bearing for anything Luma does. The paper’s pitch is that k-means is no longer just offline batch — it’s a recurring online primitive inside modern pipelines. Two filed papers on the wiki already lean on it directly. Nemotron-CLIMB: CLustering-based Iterative Data Mixture Bootstrapping for Language Model Pre-training (Nemotron-CLIMB) embeds Nemotron-CC + smollm-corpus and runs FAISS k-means at K=1000 to construct ~21 super-clusters as the axes of an automated data-mix search — Flash-KMeans’s claimed 200× speedup over FAISS would collapse that preprocessing step into something cheap enough to re-cluster every search iteration rather than once per project. The Smol Training Playbook (GPU MODE talk on SmolLM3) talk describes semantic deduplication of trillion-token corpora as a load-bearing pretraining step where k-means clustering is in the critical path.

The deeper interest is the recipe, not the algorithm. Flash-KMeans is a clean second application of the same IO-aware kernel-design playbook — now anchored at IO-Aware Kernel Design — that FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling applies to attention, and that FlashSampling: Fast and Memory-Efficient Exact Sampling applies to LM-head sampling: identify HBM materialization of an intermediate matrix as the bottleneck, fuse the produce/consume pair into a streaming kernel with on-chip state, and rewrite write-side contention into sorted segment reductions. The fact that this playbook generalizes off attention onto a completely classical algorithm — without changing the math — is the most useful signal for Luma: any time a kernel materializes an N×K intermediate and then reduces along K, the same restructuring probably applies.