Source-linked AI summary
Flash-KMeans: Fast and Memory-Efficient Exact K-Means
Shuo Yang, Haocheng Xi, Yilong Zhao, Muyang Li, Xiaoze Fan, Jintao Zhang, Han Cai, Yujun Lin, Xiuyu Li, Kurt Keutzer, Song Han, Chenfeng Xu, Ion Stoica
TL;DR
Existing GPU k-means implementations are limited by hardware-level memory and synchronization bottlenecks, including distance-matrix materialization and atomic contention. Flash-kmeans restructures the execution path without changing Lloyd k-means or adding approximations, achieving up to 17.9× end-to-end speedup over baselines and up to 33× and over 200× over cuML and FAISS.
Problem
GPU k-means implementations face hardware-level bottlenecks from explicit N × K distance-matrix materialization and irregular atomic centroid-update operations.
Method
Flash-kmeans preserves exact Lloyd k-means while combining FlashAssign, sort-inverse update, chunked stream overlap, and cache-aware compilation to restructure GPU execution.
Results
Up to 17.9× end-to-end speedup over the strongest baseline, with up to 33× and over 200× speedups over NVIDIA cuML and FAISS, respectively.
Takeaways & Limitations
Flash-kmeans makes exact k-means practical for modern GPU workloads, including large-scale and dynamically shaped deployments.
Abstract
from arXiv · showhide
$k$-means has historically been positioned primarily as an offline processing primitive, typically used for dataset organization or embedding preprocessing rather than as a first-class component in online systems. In this work, we revisit this classical algorithm under the lens of modern AI system design and enable $k$-means as an online primitive. We point out that existing GPU implementations of $k$-means remain fundamentally bottlenecked by low-level system constraints rather than theoretical algorithmic complexity. Specifically, the assignment stage suffers from a severe IO bottleneck due to the massive explicit materialization of the $N \times K$ distance matrix in High Bandwidth Memory (HBM). Simultaneously, the centroid update stage is heavily penalized by hardware-level atomic write contention caused by irregular, scatter-style token aggregations. To bridge this performance gap, we propose flash-kmeans, an IO-aware and contention-free $k$-means implementation for modern GPU workloads. Flash-kmeans introduces two core kernel-level innovations: (1) FlashAssign, which fuses distance computation with an online argmin to completely bypass intermediate memory materialization; (2) sort-inverse update, which explicitly constructs an inverse mapping to transform high-contention atomic scatters into high-bandwidth, segment-level localized reductions. Furthermore, we integrate algorithm-system co-designs, including chunked-stream overlap and cache-aware compile heuristics, to ensure practical deployability. Extensive evaluations on NVIDIA H200 GPUs demonstrate that flash-kmeans achieves up to 17.9$\times$ end-to-end speedup over best baselines, while outperforming industry-standard libraries like cuML and FAISS by 33$\times$ and over 200$\times$, respectively. Our code is open-sourced at https://github.com/svg-project/flash-kmeans.
1 Introduction
Modern GPU k-means is limited by memory traffic, atomic contention, and deployment constraints rather than only theoretical computation. Flash-kmeans preserves exact Lloyd k-means while restructuring its kernels and execution path for these bottlenecks.
- Motivation: Modern GPU k-means shifts toward online AI workloads, but traditional algorithmic optimizations often fail to deliver end-to-end speedups.The workload is moving from offline processing toward frequent use in training and inference pipelines.
- Kernel-level bottlenecks: O(NKd) assignment explicitly materializes an N × K distance matrix, creating substantial HBM traffic.Standard implementations write the matrix to HBM and immediately read it back before applying arg min.
- Kernel-level bottlenecks: Atomic scatter updates serialize centroid aggregation, especially when many threads target unbalanced or hot clusters.The standard update reaches only 50 GB/s effective bandwidth on an NVIDIA H200 GPU.
- Contribution: Flash-kmeans preserves the standard Lloyd formulation and avoids approximations while restructuring assignment, centroid updates, and system execution.Its design targets the three hardware bottlenecks identified in the standard pipeline.
- Contribution: FlashAssign fuses streaming distance computation with online argmin, while sort-inverse update converts atomic scatters into regular segment-level merges.These techniques avoid distance-matrix construction and organize assignments by cluster id before aggregation.
2 Related Work
Related work reduces k-means arithmetic or optimizes hardware data paths, but the paper positions implementation-level IO and synchronization optimization as necessary for modern GPU performance. This reflects k-means’ transition from offline analysis toward latency-sensitive online invocation.
- Workload evolution: Modern AI pipelines increasingly invoke k-means online and repeatedly, shifting evaluation emphasis from offline throughput to invocation latency.The paper cites applications including semantic deduplication, embedding quantization, token permutation, and KV-cache quantization.
- Algorithmic optimizations: Prior k-means methods reduce distance calculations, dataset size, or convergence cost through algorithmic optimizations.Examples include triangle-inequality pruning, summarization or sampling, and dual-distance metrics.
- Algorithmic optimizations: These reductions in arithmetic do not by themselves overcome the IO-bound behavior of standard GPU implementations.Implementation-level optimization is therefore presented as necessary for translating FLOP reductions into wall-clock speedups.
- Hardware-aware primitives: Hardware-aware systems optimize memory dataflow and synchronization without changing mathematical formulas.FlashAttention exemplifies avoiding explicit matrix materialization, while sorting followed by segmented operations addresses irregular scatter writes.
3 Preliminary and Motivation
Standard GPU k-means is limited by implementation-level memory traffic and atomic serialization rather than distance-computation complexity alone. These bottlenecks become more consequential when k-means is used online with large, dynamically shaped workloads.
- Lloyd’s algorithm: Lloyd’s algorithm alternates between assigning each point to its nearest centroid and aggregating assigned points to update centroids.
- Assignment bottleneck: Standard GPU assignment computes distances, writes the full matrix to HBM, then rereads rows for argmin selection.
- System motivation: Together, the assignment memory wall and update-stage atomic serialization severely limit standard k-means performance on modern GPUs.
- Assignment bottleneck: Materializing the N × K distance matrix forces Θ(NK) writes and reads, making assignment latency bandwidth-limited by 2 · Θ(NK) HBM traffic.
- Update bottleneck: Scatter-style centroid updates issue token-granularity atomic additions, causing irregular destinations, serialization, cache-line thrashing, and O(Nd) atomic operations.
- System constraints: Large batches can exceed VRAM, making chunked execution communication- and synchronization-intensive and shifting pressure toward PCIe bandwidth.
4 Methodology
Flash-kmeans restructures assignment, centroid updates, and execution around GPU bottlenecks without changing Lloyd k-means mathematically or introducing approximations. Its kernels remove distance-matrix materialization and regularize aggregation, while system co-design addresses large-scale execution and compilation overhead.
- FlashAssign: FlashAssign fuses streaming distance computation with online argmin, so the full N × K distance matrix is never explicitly constructed.
- FlashAssign: FlashAssign reduces ideal assignment IO complexity from O(NK) to O(Nd + Kd) by reading inputs once and writing assignments once.
- Sort-inverse update: Sort-inverse update argsorts assignments, groups identical cluster IDs into contiguous segments, and performs localized reductions using gathered token features.
- Sort-inverse update: Figure 2 contrasts irregular scatter updates with sorted, segment-level localized reductions and their corresponding execution timelines.
- Sort-inverse update: Sort-inverse update issues global atomic merges at segment boundaries rather than token granularity, reducing contention and atomic work.
- System co-design: Chunked stream overlap processes VRAM-exceeding inputs in chunks while overlapping host-to-device transfers with k-means computation.
- System co-design: A cache-aware compile heuristic uses hardware cache sizes and problem shape to select near-best configurations without exhaustive tuning.
5 Experiments
Experiments on NVIDIA H200 GPUs show that flash-kmeans accelerates end-to-end and kernel-level execution across diverse workload regimes, while supporting memory-constrained and dynamic-shape deployments.
- End-to-end speedup benchmark: 17.9× speedup: flash-kmeans reduces end-to-end latency by 94.4% over fast_pytorch_kmeans in the compute-intensive large-N, small-K regime.In the memory-intensive large-N, large-K regime, it outperforms fastkmeans by over 5.4× and avoids the out-of-memory failures of standard PyTorch.
- Kernel-level efficiency breakdown: 21.2× speedup: FlashAssign reduces assignment latency from 122.5ms to 5.8ms on N = 1M, K = 8192.The kernel removes HBM distance materialization by fusing the assignment computation with the nearest-cluster selection.
- Kernel-level efficiency breakdown: 6.3× speedup: sort-inverse update accelerates centroid reduction on massive-scale workloads such as B = 1, N = 33M, K = 4096.It replaces per-token scatter atomic adds with sorted, localized merges.
- Large-scale out-of-core data processing: One billion points: flash-kmeans is evaluated in out-of-core workloads where standard PyTorch immediately fails with out-of-memory errors.The experiment compares flash-kmeans with fastkmeans, the most robust available out-of-core baseline.
- Fast time-to-first-run for dynamic shapes: Less than 0.3% performance difference: heuristic-driven kernel iteration latency matches the exhaustively tuned oracle across tested shapes.This preserves runtime performance while avoiding expensive configuration search.
6 Conclusion
The paper concludes that flash-kmeans makes mathematically exact k-means highly deployable on modern GPU workloads by restructuring execution around memory, synchronization, and compilation bottlenecks. It reports substantial end-to-end acceleration, extreme-scale support, and sharply reduced compilation overhead with near-zero performance degradation.
- 6 Conclusion: FlashAssign eliminates massive distance-matrix materialization, while Sort-Inverse Update resolves write-side atomic contention without altering the underlying mathematics.The implementation is coupled with an asynchronous out-of-core pipeline and shape-aware compile heuristics.
- 6 Conclusion: 17.9× end-to-end speedup: flash-kmeans outperforms the best baselines, while exceeding cuML by 33× and FAISS by over 200×.These are the reported comparisons for the proposed implementation.
- 6 Conclusion: 175× lower compilation overhead: flash-kmeans scales to workloads of one billion points with near-zero performance degradation.The conclusion presents these results as evidence of a robust, mathematically exact, deployable clustering primitive.