PaperlayerView exact revision on arXiv

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

arXiv:2603.09229v2cs.DC

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.

    Source: §1 · p. 3, §3.2 · p. 4

  • 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.

    Source: §1 · p. 3

  • 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.

    Source: §1 · p. 3, §6 · p. 10

  • Takeaways & Limitations

    Flash-kmeans makes exact k-means practical for modern GPU workloads, including large-scale and dynamically shaped deployments.

    Source: §1 · p. 3, §3.3 · p. 5, §6 · p. 10

Abstract

from arXiv · show

$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.Source: §1 · pp. 1–2

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.Source: §1 · p. 2
  • 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.Source: §1 · p. 2

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.Source: §1 · pp. 2–3
  • 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.Source: §1 · p. 3

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.Source: §2 · p. 3

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.Source: §2 · pp. 3–4
  • 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.Source: §2 · p. 4

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.Source: §2 · p. 4

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.Source: §3.1 · p. 4

Assignment bottleneck

  • Standard GPU assignment computes distances, writes the full matrix to HBM, then rereads rows for argmin selection.Source: §3.1 · p. 4, §3.2 · p. 5

System motivation

  • Together, the assignment memory wall and update-stage atomic serialization severely limit standard k-means performance on modern GPUs.Source: §3.2 · pp. 4–5

Assignment bottleneck

  • Materializing the N × K distance matrix forces Θ(NK) writes and reads, making assignment latency bandwidth-limited by 2 · Θ(NK) HBM traffic.Source: §3.2 · pp. 4–5

Update bottleneck

  • Scatter-style centroid updates issue token-granularity atomic additions, causing irregular destinations, serialization, cache-line thrashing, and O(Nd) atomic operations.Source: §3.2 · p. 5

System constraints

  • Large batches can exceed VRAM, making chunked execution communication- and synchronization-intensive and shifting pressure toward PCIe bandwidth.Source: §3.3 · p. 5

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.Source: §4.1 · p. 6
  • FlashAssign reduces ideal assignment IO complexity from O(NK) to O(Nd + Kd) by reading inputs once and writing assignments once.Source: §4.1 · p. 6

Sort-inverse update

  • Sort-inverse update argsorts assignments, groups identical cluster IDs into contiguous segments, and performs localized reductions using gathered token features.Source: §4.2 · pp. 6–7
  • Figure 2 contrasts irregular scatter updates with sorted, segment-level localized reductions and their corresponding execution timelines.Source: Figure 2 · p. 7
  • Sort-inverse update issues global atomic merges at segment boundaries rather than token granularity, reducing contention and atomic work.Source: §4.2 · p. 7

System co-design

  • Chunked stream overlap processes VRAM-exceeding inputs in chunks while overlapping host-to-device transfers with k-means computation.Source: §4.3 · p. 7
  • A cache-aware compile heuristic uses hardware cache sizes and problem shape to select near-best configurations without exhaustive tuning.Source: §4.3 · pp. 7–8

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.Source: §5.2 · p. 8

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.Source: §5.2 · p. 8, Figure 4 · p. 10
  • 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.Source: §5.2 · p. 9, Figure 4 · p. 10

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.Source: §5.3 · p. 9

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.Source: §5.3 · p. 10, Figure 5 · p. 10

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.Source: §6 · p. 10
  • 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.Source: §6 · p. 10
  • 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.Source: §6 · p. 10

Related papers

Loading 2603.09229v2…